Repository navigation
[ROCm] wvSplitK: launch rows of their own for gfx1151 and gfx1150 - #1406
Closed
roberteg16 wants to merge 3 commits into
Closed
roberteg16 wants to merge 3 commits into
roberteg16 wants to merge 3 commits into
Conversation
The bigType union in wvSplitK_hf_sml_, wvSplitK_hf_ and wvSplitK_hf_big_
defines half4 as (A_CHUNK / 2) x bf16, so h4[A_CHUNK / 4] grows as
A_CHUNK^2 / 4 bytes while every other member is A_CHUNK halves:
A_CHUNK | h[A_CHUNK] | h4[A_CHUNK/4] | sizeof(bigType)
--------|------------|---------------|----------------
8 | 16 B | 16 B | 16 B
16 | 32 B | 64 B | 64 B
With A_CHUNK = 8 the two agree; with 16 every bigType copy moves 64 B
where the kernel assumes 32 B: it reads past the end of an activation
row (hipErrorIllegalAddress when the row ends at a page boundary) and
adjacent threads overwrite each other's LDS. No launch uses A_CHUNK = 16
today; the next commit does.
half4 is now 4 x bf16, the operand of its only consumer
(__builtin_amdgcn_mfma_f32_4x4x4bf16_1k), which keeps sizeof(bigType)
at A_CHUNK halves for every A_CHUNK.
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>
Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
On gfx11 and gfx12 the fp16 DOT2C already issues v_dot2_f32_f16, while the bf16 one widens both operands to f32, multiplies and adds. Issue v_dot2_f32_bf16 instead (VOP3-P, available on gfx11 and gfx12): the weights stay packed, as with fp16, and the bf16 path no longer needs twice the registers for them. AI assistance was used for this patch. Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
The wave32 wvSplitK launches took either a gfx1151 chain or, on every other gfx11/gfx12 arch, the generic YTILE-by-M rows, all with A_CHUNK = 8. A sweep of every non-spilling launch configuration on gfx1151 and gfx1150 (YTILE 1-4 x UNRL 1/2/4/8 x A_CHUNK 8/16/32 x WvPrGrp 16/32, N = 1..5, bf16 and fp16) finds one tile per arch within ~1% of the best for nearly every shape: - gfx1151: (YTILE 1, UNRL 4, A_CHUNK 16); (2, 4, 8) for K < 1024. - gfx1150, gfx1152, gfx1153: (2, 1, 16); (1, 2, 16) when M <= 32 * CUs. - A_CHUNK 8 when K % 16 != 0. WVSPLITK_CFG_AC takes A_CHUNK; WVSPLITK_CFG keeps 8. Other archs keep the generic rows. AI assistance was used for this patch. Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
Author
|
Already in upstream vllm-project#63835 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Port to upstream of the wvSplitK changes in #1401. The base branch is upstream/main (9e3db48), so the diff is what an upstream PR would contain.
Purpose
Faster wave32
wvSplitKon gfx1151 and gfx1150/1152/1153 with one launch tile per arch, from a sweep of every non-spilling launch configuration.Changes
bigTypesized by A_CHUNK: with A_CHUNK = 16 it was twice as large, reading past an activation row and overlapping LDS. No launch used A_CHUNK = 16 before commit 3.DOT2Con gfx11/gfx12 issuesv_dot2_f32_bf16, as fp16 issuesv_dot2_f32_f16.The dispatch code is identical to #1401's.
Results
Against upstream/main. Interleaved A/B, cold caches, padded weights, median of 5 rounds; 43 shapes x N = 1..5, bf16/fp16 (13 shapes not used for tuning, odd M, K 768/1000/2880), plus bias.
Cells slower by more than 2%: 0-1 per set (at 0.98). 99.3-99.8% of the per-shape best tile, against 96-98% for upstream/main.
Tests
torch.nn.functional.linear: correct on gfx1150 and gfx1151.test_rocm_skinny_gemms.pygains K % 16 != 0 and N = 5 cases; with the same dispatch code in [ROCm][RDNA] Skinny GEMMs: no VGPR spills on gfx11/gfx12, wvSplitK launch rows for gfx1150/gfx1151 #1401: 190 passed on gfx1150 and gfx1151.Related
AI assistance was used; the submitter reviews every line.