Repository navigation
[DO NOT MERGE] Integration: #1371 (gfx1150 ROCM_RDNA_ATTN) + #1401 (skinny GEMM spills) - #1402
Draft
roberteg16 wants to merge 6 commits into
Draft
roberteg16 wants to merge 6 commits into
roberteg16 wants to merge 6 commits into
Conversation
The prefill kernel asks for kWavesPerSimd waves a SIMD (amdgpu_waves_per_eu), sized for archs with 1536 VGPRs a SIMD lane: 6 waves leave a wave 256 VGPRs. On archs with 1024 VGPRs the same hint caps a wave at 160 and LLVM spills: 394 of the 430 prefill entries spill, 352 by 32 VGPRs or more (up to 133), and prefill runs 12-16x slower than on a 1536-VGPR arch instead of ~2.5x. On 1024-VGPR archs the hint becomes 2/3 of the waves, in whole workgroups: the same VGPRs a wave as with 1536. launch_prefill() sizes the resident workgroups (firstRound) with the same target, telling the archs apart by hipDeviceAttributeMaxRegistersPerMultiprocessor. The hardware numbers involved are named (kSimdsPerWgp, kWgpLdsBytes, kBigLaneVgprs, kSmallLaneVgprs). Code for 1536-VGPR archs is unchanged: identical resource usage for all 430 entries. The kernels are still built and served only where they were; this makes a 1024-VGPR build of the same entries usable. On a 1024-VGPR arch (Radeon 890M), entries built through the tools' JIT: - 430 entries: 165 spill, none by 32 or more (max 8). - Every entry of 6 shapes (D 64..512, fp16/bf16, sliding window) matches Triton (D=512 bf16 to one bf16 ulp, as close to an fp32 reference as Triton); so do 456 more knob sets of 10/10/128 and 16/2/256. - 7 shapes x 33 cells (M 16..16k after 0/4k/16k cached tokens), rows as tuned on a 1536-VGPR arch: faster than TRITON_ATTN, ROCM_SEGMENTED_ATTN and ROCM_ATTN in every cell (geomean 6.6x, 3.1x, 9.0x; worst 1.08x, 1.04x, 1.38x); 2.1-6.5x faster than without this change, except D=64 (1.0x), whose entries spilled little. AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
The runner rotates the inputs over num_layers (sized so the KV working set leaves the cache) but every layer wrote the same output buffer, so its lines stayed in the write-back L2 and the output writes mostly never reached DRAM. Small prefills timed faster than their DRAM traffic allows: above the roofline, which counts the output. The output now rotates like the inputs. On a 1024-VGPR GPU, prefills of 32..64 tokens without a prefix: 4-44% slower with the output rotated, none above the roofline any more. AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
…x1153 The configs are tuned on gfx1151, with 1536 VGPRs a SIMD lane. gfx1150, gfx1152 and gfx1153 have 1024 and a smaller GPU, and their best launch knobs differ; this gives them tables of their own, tuned on gfx1150, and selects the backend by default on gfx1150 too (gfx1152 and gfx1153, not tested yet, can select it). The batched tables start as gfx1151's (prefill batch cells pointing at gfx1150's prefill entries). - Per-family tables. The configs headers hold one complete table per Arch (rdna_attn_types.h): kGfx1151*Configs (the existing entries) and kGfx1150*Configs, the same shapes, which gfx1152 and gfx1153 also use, for the decode, batched decode, prefill and prefill batch tables. They are reached through decode_configs<A>() etc. and the k*Configs[kNumArches] spans. CMake builds one object library per family for its archs in the arch list (rdna_attn_targets), and _rocm_C lists and launches the tables of the current GPU's family. - An entry in both tables is a distinct decode kernel per family (decode_attn<C, kBatched, OutT, Arch>); prefill's already have internal linkage. Same instructions on gfx1151. - Decode launch: the LDS residency correction (128 KiB a WGP) covers gfx1150..gfx1153, with their VGPR file (1024 in blocks of 16). - Tools: jit.py builds for the current GPU (RDNA_JIT_ARCH overrides it); tune.py records the family it measured on and --apply writes that family's table; utils.py has the roofline peaks of each GPU. - tune.py: a failed run of the current row is no baseline for accepting a decode candidate; reused points are built before the A/B; the decisions keep every A/B sample; prefill --reuse matches its points; a prefill step's winner loses at most --prefill-loss-limit (2%) to the current entry at any cell, and --reselect chooses stored decisions again without timing; prefill steps go up to M=8192 (vLLM's default chunk), --steps tunes some of them and --apply keeps the others. gfx1150's entries were tuned from gfx1151's on the gfx1150 boards whose BIOS keeps kernel launches fast (BIOS 1.08 doubles the launch latency and favours other knobs), each candidate checked against a float reference and timed through the backend, fp16 (bf16 shares the entries). Against gfx1151's entries in the same A/B: - prefill: 33 configurations x 9 steps of M (16..4096) x prefixes 0/1024/16384, 891 cells: 1.124x, none below 0.98x. The M=8192 step, tuned on its own, gains 3% or less on 17 configurations, which keep their M>=4096 entry; 16/8/256, 32/16/256 w1024 and 32/32/96 get one (1.44x, 1.06x, 1.03x). 334 entries (328 builds) against gfx1151's 434. - decode: 193 of 264 (configuration, M) accepted (1.01x-1.46x weighted, within the per-context loss limits); 1848 cells: 1.065x. Re-timed against gfx1151's entries in the kernel matrix, each twice on each of five gfx1150 boards: prefill 1.095x over 660 cells (M 16..8192, prefixes 0/4096), decode 1.061x over 1848 cells (M 1..8, contexts 128..32k). Tests: the backend's 398 tests pass on gfx1150 and gfx1151; gfx1150's tables cover gfx1151's shapes. gfx1151's code is the same as before (decode, batched decode and prefill units compared). AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
gfx1150's batched decode and prefill batch tables were gfx1151's. Tuned on the low-latency gfx1150 boards (fp16 measured, bf16 shares the knobs) and written with tools/tune.py --apply: - batched decode: 171 of 264 (configuration, M) accepted against the current entry, 1.087x geomean over the accepted. - prefill batch: per (mean M, keys) cell, the fastest prefill entry and head group on uniform batches. D=512 keeps the previous cells: in mixed batches (many decodes and one long prefill) its tuned cells were up to 2.8x slower, as a cell's batches are uniform. Against the previous tables in the kernel matrix (matrix.py --batch, 33 configurations, each table set timed twice on three gfx1150 boards, median spread 0.1%): - batched decode (10 batches x M 1..8): 2640 cells, 1.047x, 916 faster and 16 slower by more than 2% (down to 0.95x); up to 18x where gfx1151's entries were far off (16/1/512 and 16/2/512, 64 sequences, M 6..8). - prefill batches, D < 512: 140 cells, 1.112x (0.90x-3.77x). - mixed batches, D < 512: 168 cells, 1.013x (0.83x-1.82x). AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
wvSplitK_hf_sml_, wvSplitK_hf_ and wvSplitK_hf_big_ held every activation chunk of an unrolled step in registers next to the weights', and bf16 widened the weights to f32 before multiplying. With 1024-thread workgroups a wave gets 128 VGPRs on gfx1150 (192 on gfx1151), and the large tiles spilled: 102 of the 388 kernels of skinny_gemms.cu on gfx1150, 78 on gfx1151, among them the 32/1/32/16/8/1 bf16 tile most decode GEMMs of a 2B model use (139 VGPRs spilled). A spilling GEMM is slow, and its speed depends on the scratch the process allocated before: a kernel run before with 256 B/lane of scratch made it 1.3x slower, 4 KiB 2.1x. On gfx1x: - bf16 multiplies with v_dot2_f32_bf16, as fp16 with v_dot2_f32_f16, keeping the weights packed. - A tile holds the activation chunks while (N + YTILE) * UNRL * A_CHUNK / 2 operand VGPRs plus 64 fit the wave's VGPRs (1024 or 1536 a lane, split by the workgroup's waves per SIMD); else it reads each chunk from LDS where it is used, and only the weights' loads stay in flight. gfx9 code is unchanged. All 388 kernels of skinny_gemms.cu on gfx1150, gfx1151 and gfx1201 compile without VGPR spills; the ones left with scratch are the assert(false) stubs of the gfx950/MI3xx-only wvSplitKrc_ and wvSplitKQ_ (fp8) kernels. bench_rocm_skinny_gemm.py (76 shape x N cells per dtype, N 1..4, best of two runs per build), against gfx11: - gfx1150 bf16 1.050x (0.985x-2.63x, none slower by more than 2%), fp16 1.008x (0.984x-1.51x); N=1 and N=2 bf16 1.09x and 1.11x. - gfx1151 bf16 1.004x, fp16 1.000x. The 4096x2048 GEMM keeps 186-210 us whatever scratch a kernel run before used (242-651 us before). Cosmos-Reason2-2B, document-qa-long-context, decode on gfx1150 with TRITON_ATTN in both builds: 19.5 -> 22.6 tokens/s. test_rocm_skinny_gemms.py: 172 passed on gfx1150 and gfx1151 (the others are gfx950- or fp8-only). AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
…only reachable tiles built The int4 dispatch macros expand every tile for every N = 1..5, so tiles the dispatch only picks for N = 1 were also built for N = 2..5, where they are never launched: 480 of the 1684 int4 kernels, and most of the ones that spill (310 of 334 on gfx1150, 285 of 293 on gfx1151). The N = 1 arms are now `if constexpr (__N == 1)`, so those are not built; the dispatch is unchanged. Of the kernels left, the only one that spills on gfx1150 is the MoE (YTILE=4, W=32, AC=32, UNRL=4) tile for N = 1 at K = 2048 (Qwen3.5-A3B gate_up): a 1024-thread workgroup leaves a wave 128 VGPRs on a 1024-VGPR arch. gfx1102, gfx1103 and gfx1150..gfx1153 take the (4, 4, 16, 16) tile there instead; gfx1151 keeps it (no spills with 192 VGPRs). Left spilling: that tile on gfx1150 (built for gfx1151, not launched on gfx1150), and on gfx1151 eight tiles at their VGPR limit as before: the dense (YTILE=2, W=16, AC=32, UNRL=8) N = 1 tile for K % 8192 == 0 (mostly with zero points) and the MoE (2, 16, 16, 4) N = 5 bf16 tile. The kernels' code is unchanged; the compiler emits different code for 15 (gfx1150) and 17 (gfx1151) of the 1204 kernels left in the smaller translation unit. Timed against gfx11 (int4 W4A16, group size 128, M = 1..4, best of two runs per build; dense: Qwen3-4B, Llama-8B and Qwen3.5-A3B shapes, fp16 and bf16; MoE: Qwen3.5-A3B, 256 experts, top 8, K = 2048): - gfx1150: MoE 1.19x (1.16x-1.24x); dense 0.99x-1.00x (the M = 3..4 cells vary by up to 20% between processes with either build). - gfx1151: dense 1.00x, MoE 1.00x (all cells within 1.5%). test_rdna_hybrid_w4a16.py (102 passed), test_hybrid_w4a16_moe.py (70 passed) and test_rocm_skinny_gemms.py (172 passed) on gfx1150 and gfx1151; the test_rdna_hybrid_w4a16.py tests needing dist_init were not run. AI assistance was used for this patch. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
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.
Integration branch, not for merging: builds a wheel with both changes for the nightly regression on Strix Point (gfx1150).
ROCM_RDNA_ATTNtables for gfx1150 (used by gfx1152/gfx1153) and the backend selected by default on gfx1150 and gfx1151. Includes the latest local revision of [ROCm][RDNA] RDNA single-seq: occupancy and tables of their own for 1024-VGPR archs #1371 (batched tables), not yet pushed there.wvSplitK*, int4wvSplitK_int4*).Base
gfx110402a20; the commits are those of the two PRs, unchanged.AI assistance (Claude Code) was used; the submitter reviews every line.