w4a16 Triton prefill tuning - #923
Merged
Merged
Conversation
…=4096
Replace four shape-specific tile configs in the M<=128 branch with a
single universal config (BLOCK_M=64, BLOCK_N=32, BLOCK_K=128,
num_warps=4) that wins on every projection shape across Qwen3-8B and
Llama-3.1-8B. Previous dispatch had three sub-cases (K>=2N, N>K, N=K)
each picking different BLOCK sizes, all underperforming for K=4096
shapes. Sweeping all M<=128 K=4096 shapes converged on the same winning
config; the earlier shape-specific configs were dead-ends.
Microbench (Strix Halo, gfx1151), Qwen3-8B and Llama-3.1-8B all
projections converge to:
gate_up (24576-28672, 4096): ~14.4 -> 19.1-19.4 TFLOPS (+24-35%)
down (4096, 12288-14336): ~14.9 -> 18.4-18.5 TFLOPS (+23%)
qkv (6144, 4096): ~10.2 -> 19.5-19.7 TFLOPS (+90%)
o_proj (4096, 4096): ~12.2 -> 18.0-18.1 TFLOPS (+47%)
Smaller-N or smaller-K shapes (Qwen3-VL-4B / Qwen3-4B with N<4096 or
K<4096) keep the legacy shape-specific configs - at small N the BN=32
grid drops below the 40-CU saturation threshold.
End-to-end Qwen3-8B-AWQ TTFT: 185 -> 125 ms (-32.4%).
Why this config wins (cross-checked with rdna35-expert wmma.md):
- BK=128 (=group_size) doubles inner-loop work per K-tile, escaping
WMMA's 1-wave/SIMD latency-bound regime (~48% of peak) and
approaching the >=2-waves/SIMD saturation regime (~95% of peak).
- BN=32 keeps the workgroup grid large enough to saturate all 40
CUs even at N up to 28k; BN=64 would halve grid size and drop
below saturation.
- NW=4 hits the ">=2 waves/SIMD" sweet spot per the rocWMMA
blocks/CU sweep.
Changes:
- Replaces three nested if/elif sub-cases with a single guarded
universal config plus three legacy fallbacks for K<4096 / N<4096.
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>
ISA inspection of the unified M<=128 dispatch (BN=32 BK=128 NW=4) showed
the kernel uses 256 VGPRs/wave with 3 spills, capping occupancy at 6
waves/SIMD on gfx1151's 1536-VGPR file. Per the AMD-Triton compiler
comment ("Specifying N, N forces LLVM to focus on a single register
count, simplifies some heuristics and may improve scheduling"), pinning
the target to the natural occupancy yields better LLVM scheduling.
Sweep (waves_per_eu=N) on all four Qwen3-8B M=128 shapes:
N=0 (default) -> baseline
N=4 -> ~neutral
N=6 -> +4-8% per shape (winner)
N=8 -> -57% (LLVM forced to spill heavily to fit 192 VGPRs)
N>=10 -> -75% to -80%
Microbench, all four projection shapes (Strix Halo, gfx1151):
gate_up: 19.1 -> 20.4 TFLOPS (+6.8%)
down: 18.4 -> 19.0 TFLOPS (+3.3%)
qkv: 19.7 -> 20.3 TFLOPS (+3.0%)
o_proj: 18.0 -> 18.8 TFLOPS (+4.5%)
End-to-end Qwen3-8B-AWQ TTFT: 125 -> 117 ms (-6.6%); total -36.9% vs
the 185 ms pre-campaign baseline.
Hint scoped to the universal config branch only - other branches keep
default behaviour (waves_per_eu=0 = no constraint).
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>
mgehre-amd
changed the title
mgehre-amd
marked this pull request as ready for review
May 6, 2026 15:45
eble-amd
approved these changes
May 7, 2026
eble-amd
left a comment
eble-amd
left a comment
There was a problem hiding this comment.
Code LGTM and I can't argue with the results.
mgehre-amd
added a commit
that referenced
this pull request
Sep 7, 2026
…ction The Triton prefill path of RDNAHybridW4A16LinearKernel is VALU-issue-bound on RDNA3, and its tile ladder was tuned before the dequant changed shape. Packed fp16 dequant. OR-ing a 4-bit code n into the low mantissa of fp16 1024.0 (0x6400) bitcasts to exactly 1024+n. Applying that to a full 32-bit lane with one v_and_or_b32 dequants a nibble PAIR per instruction, against the scalar v_and_b16 + v_or_b16 pair Triton emits from the elementwise form. The ExLlama shuffle already stores val[2p] at bits [4p:4p+4] and val[2p+1] at [16+4p:20+4p], so a single pre-shift by 4p yields both halves of a half2 in K order and the downstream affine packs into v_pk_fma_f16. Dtype-aware gfx1151 tile selection. The packed fp16 path and the scalar bf16 path want different tiles -- most visibly BLOCK_N at deep M, 256 against 64 -- so a single M-ladder cannot serve both. Measured on gfx1151, 3 interleaved reps per arm, 2048-token prefill: Qwen3-8B-quantized.w4a16 fp16 asymmetric +7.8% prefill throughput / -7.3% TTFT; Qwen3-4B-quantized.w4a16 fp16 symmetric +13.0% / -11.5%. Both fp16 rep ranges are disjoint. The bf16 control lands at +0.02% with overlapping rep ranges, i.e. the null it should be. Changes: - The packed dequant is gated from the launcher (a plain tl.constexpr) rather than from tl.target_info. This module is imported on every platform, and vLLM's Triton placeholder shim exposes no target_info attribute, so a module-scope @tl.target_info.constexpr_function would raise AttributeError on builds without Triton. - fp16 is a hard requirement of the magic constant, not a tuning preference, so it bounds the packed_dequant override as well. Tests pass packed_dequant=False to force the scalar unpack and assert the two agree bit-for-bit. - num_stages=1 is applied to the fp16 arm ONLY. Applying it to bf16 as well measured -15.8% end-to-end prefill on an asymmetric bf16 model: the packed fp16 path issues one per-group load and does not miss the pipelining, while the scalar asymmetric path issues two (scale and zero point) and needs the software pipeline to hide the second gather. - The bf16 arm is therefore byte-for-byte the pre-existing scalar-tuned ladder, its per-shape override table, and its pipeline depth. Verified identical to the previous selection across all 576 shape x group-size combinations swept. - Scope is gfx1151 only. Widening the tuned table to the whole gfx11 family was considered and rejected: it was swept on gfx1151, and gfx1100 is a 96-CU discrete part that cannot be measured here. - #923's K>=4096 and N>=4096 bf16 tile branch was evaluated and deliberately NOT carried over. At group_size 128 it is shadowed by the later, more specific per-shape override table for every shape that table covers; at group_size 32/64 its distinguishing BLOCK_K=128 collapses to the group size anyway. It would only ever fire on shapes nobody measured, in a region where the later sweep disagreed with it. - Profiler scope labels gain a g=<group_size> sym/asym suffix. Per-group metadata adds ~25% on top of the weight bytes at group_size 32 asymmetric, so a label carrying only MxNxK cannot be turned into a bandwidth number. - The kernel benchmark gains --dtype: it was fp16-hardcoded and so could not reach the scalar bf16 path at all. Numerics. The dequant arithmetic is unchanged: 1024+n is integer-exact in fp16, whose 11-bit mantissa covers every integer below 2048, so folding the 1024 into the subtrahend -- (b_raw - (1024 + zp)) == (nibble - zp) -- is exact and the multiply that follows rounds once, as before. End-to-end output is nonetheless not GUARANTEED identical on fp16, because the new tiles move BLOCK_K (64 -> 32 on most shapes) and so change the fp32 accumulation order. Measured: greedy decode over 6 fixed prompts is byte-identical on all three configurations, and GSM8K 5-shot over 500 questions is unchanged on the asymmetric model (0.884 -> 0.884) and moves by one question on the symmetric one (0.848 -> 0.850), against a ~1.5pp single-arm sampling error. Testing: pytest tests/kernels/quantization/test_rdna_hybrid_w4a16.py tests/kernels/quantization/test_w4a16_kernel_selection.py -- 109 passed on gfx1151 (Radeon 8060S, torch 2.11.0+rocm7.15, Triton 3.8.0).
mgehre-amd
added a commit
that referenced
this pull request
Sep 7, 2026
…ction The Triton prefill path of RDNAHybridW4A16LinearKernel is VALU-issue-bound on RDNA3, and its tile ladder was tuned before the dequant changed shape. Packed fp16 dequant. OR-ing a 4-bit code n into the low mantissa of fp16 1024.0 (0x6400) bitcasts to exactly 1024+n. Applying that to a full 32-bit lane with one v_and_or_b32 dequants a nibble PAIR per instruction, against the scalar v_and_b16 + v_or_b16 pair Triton emits from the elementwise form. The ExLlama shuffle already stores val[2p] at bits [4p:4p+4] and val[2p+1] at [16+4p:20+4p], so a single pre-shift by 4p yields both halves of a half2 in K order and the downstream affine packs into v_pk_fma_f16. Dtype-aware gfx1151 tile selection. The packed fp16 path and the scalar bf16 path want different tiles -- most visibly BLOCK_N at deep M, 256 against 64 -- so a single M-ladder cannot serve both. Measured on gfx1151, 3 interleaved reps per arm, 2048-token prefill: Qwen3-8B-quantized.w4a16 fp16 asymmetric +7.8% prefill throughput / -7.3% TTFT; Qwen3-4B-quantized.w4a16 fp16 symmetric +13.0% / -11.5%. Both fp16 rep ranges are disjoint. The bf16 control lands at +0.02% with overlapping rep ranges, i.e. the null it should be. Changes: - The kernel resolves the unpack itself, from its own compile target and the activation dtype -- there is no flag to pass and no way for the host's view of the GPU to disagree with what is being compiled. That needs a @triton.constexpr_function helper, which in turn needs constexpr_function added to vLLM's Triton placeholder shim: this module is imported on every platform, and the shim exposed jit/autotune/heuristics/Config but not constexpr_function, so a module-scope use raised AttributeError on builds without Triton. A dummy decorator is sufficient there, since such bodies only run while a kernel is being compiled. - The packed dequant is enabled for the whole gfx11 family on fp16: fp16 is a hard requirement of the magic constant, and given that, the packed form is a pure instruction-count reduction producing bit-identical values, so there is nothing to tune per part. Only the tile table, which IS tuned, stays gfx1151-gated. Verified in the generated ISA: the fp16 kernel contains 65 v_and_or_b32, the bf16 one none from the dequant. - num_stages=1 is applied to the fp16 arm ONLY. Applying it to bf16 as well measured -15.8% end-to-end prefill on an asymmetric bf16 model: the packed fp16 path issues one per-group load and does not miss the pipelining, while the scalar asymmetric path issues two (scale and zero point) and needs the software pipeline to hide the second gather. - The bf16 arm is therefore byte-for-byte the pre-existing scalar-tuned ladder, its per-shape override table, and its pipeline depth. Verified identical to the previous selection across all 576 shape x group-size combinations swept. - Scope is gfx1151 only. Widening the tuned table to the whole gfx11 family was considered and rejected: it was swept on gfx1151, and gfx1100 is a 96-CU discrete part that cannot be measured here. - #923's K>=4096 and N>=4096 bf16 tile branch was evaluated and deliberately NOT carried over. At group_size 128 it is shadowed by the later, more specific per-shape override table for every shape that table covers; at group_size 32/64 its distinguishing BLOCK_K=128 collapses to the group size anyway. It would only ever fire on shapes nobody measured, in a region where the later sweep disagreed with it. - Profiler scope labels gain a g=<group_size> sym/asym suffix. Per-group metadata adds ~25% on top of the weight bytes at group_size 32 asymmetric, so a label carrying only MxNxK cannot be turned into a bandwidth number. - The kernel benchmark gains --dtype: it was fp16-hardcoded and so could not reach the scalar bf16 path at all. Numerics. The dequant arithmetic is unchanged: 1024+n is integer-exact in fp16, whose 11-bit mantissa covers every integer below 2048, so folding the 1024 into the subtrahend -- (b_raw - (1024 + zp)) == (nibble - zp) -- is exact and the multiply that follows rounds once, as before. End-to-end output is nonetheless not GUARANTEED identical on fp16, because the new tiles move BLOCK_K (64 -> 32 on most shapes) and so change the fp32 accumulation order. Measured: greedy decode over 6 fixed prompts is byte-identical on all three configurations, and GSM8K 5-shot over 500 questions is unchanged on the asymmetric model (0.884 -> 0.884) and moves by one question on the symmetric one (0.848 -> 0.850), against a ~1.5pp single-arm sampling error. Testing: pytest tests/kernels/quantization/test_rdna_hybrid_w4a16.py tests/kernels/quantization/test_w4a16_kernel_selection.py -- 103 passed on gfx1151 (Radeon 8060S, torch 2.11.0+rocm7.15, Triton 3.8.0).
mgehre-amd
added a commit
that referenced
this pull request
Sep 7, 2026
…ction The Triton prefill path of RDNAHybridW4A16LinearKernel is VALU-issue-bound on RDNA3, and its tile ladder was tuned before the dequant changed shape. Packed fp16 dequant. OR-ing a 4-bit code n into the low mantissa of fp16 1024.0 (0x6400) bitcasts to exactly 1024+n. Applying that to a full 32-bit lane with one v_and_or_b32 dequants a nibble PAIR per instruction, against the scalar v_and_b16 + v_or_b16 pair Triton emits from the elementwise form. The ExLlama shuffle already stores val[2p] at bits [4p:4p+4] and val[2p+1] at [16+4p:20+4p], so a single pre-shift by 4p yields both halves of a half2 in K order and the downstream affine packs into v_pk_fma_f16. Dtype-aware gfx1151 tile selection. The packed fp16 path and the scalar bf16 path want different tiles -- most visibly BLOCK_N at deep M, 256 against 64 -- so a single M-ladder cannot serve both. Measured on gfx1151, 3 interleaved reps per arm, 2048-token prefill: Qwen3-8B-quantized.w4a16 fp16 asymmetric +7.8% prefill throughput / -7.3% TTFT; Qwen3-4B-quantized.w4a16 fp16 symmetric +13.0% / -11.5%. Both fp16 rep ranges are disjoint. The bf16 control lands at +0.02% with overlapping rep ranges, i.e. the null it should be. Changes: - The kernel resolves the unpack itself, from its own compile target and the activation dtype -- there is no flag to pass and no way for the host's view of the GPU to disagree with what is being compiled. That needs a @triton.constexpr_function helper, which in turn needs constexpr_function added to vLLM's Triton placeholder shim: this module is imported on every platform, and the shim exposed jit/autotune/heuristics/Config but not constexpr_function, so a module-scope use raised AttributeError on builds without Triton. A dummy decorator is sufficient there, since such bodies only run while a kernel is being compiled. - The packed dequant is enabled for the whole gfx11 family on fp16: fp16 is a hard requirement of the magic constant, and given that, the packed form is a pure instruction-count reduction producing bit-identical values, so there is nothing to tune per part. Only the tile table, which IS tuned, stays gfx1151-gated. Verified in the generated ISA: the fp16 kernel contains 65 v_and_or_b32, the bf16 one none from the dequant. - num_stages=1 is applied to the fp16 arm ONLY. Applying it to bf16 as well measured -15.8% end-to-end prefill on an asymmetric bf16 model: the packed fp16 path issues one per-group load and does not miss the pipelining, while the scalar asymmetric path issues two (scale and zero point) and needs the software pipeline to hide the second gather. - The bf16 arm is therefore byte-for-byte the pre-existing scalar-tuned ladder, its per-shape override table, and its pipeline depth. Verified identical to the previous selection across all 576 shape x group-size combinations swept. - Scope is gfx1151 only. Widening the tuned table to the whole gfx11 family was considered and rejected: it was swept on gfx1151, and gfx1100 is a 96-CU discrete part that cannot be measured here. - #923's K>=4096 and N>=4096 bf16 tile branch was evaluated and deliberately NOT carried over. At group_size 128 it is shadowed by the later, more specific per-shape override table for every shape that table covers; at group_size 32/64 its distinguishing BLOCK_K=128 collapses to the group size anyway. It would only ever fire on shapes nobody measured, in a region where the later sweep disagreed with it. - The kernel benchmark gains --dtype: it was fp16-hardcoded and so could not reach the scalar bf16 path at all. Numerics. The dequant arithmetic is unchanged: 1024+n is integer-exact in fp16, whose 11-bit mantissa covers every integer below 2048, so folding the 1024 into the subtrahend -- (b_raw - (1024 + zp)) == (nibble - zp) -- is exact and the multiply that follows rounds once, as before. End-to-end output is nonetheless not GUARANTEED identical on fp16, because the new tiles move BLOCK_K (64 -> 32 on most shapes) and so change the fp32 accumulation order. Measured: greedy decode over 6 fixed prompts is byte-identical on all three configurations, and GSM8K 5-shot over 500 questions is unchanged on the asymmetric model (0.884 -> 0.884) and moves by one question on the symmetric one (0.848 -> 0.850), against a ~1.5pp single-arm sampling error. Testing: pytest tests/kernels/quantization/test_rdna_hybrid_w4a16.py tests/kernels/quantization/test_w4a16_kernel_selection.py -- 103 passed on gfx1151 (Radeon 8060S, torch 2.11.0+rocm7.15, Triton 3.8.0).
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
2 participants
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.
Uh oh!
There was an error while loading. Please reload this page.