Conversation
|
/tag-and-rerun-ci |
…ernel ROCm ran the DSA indexer q/k preparation as a chain of separate kernels: k LayerNorm, RoPE, fp8 quant into the paged index-k cache, then q Hadamard and quant, then the head-gate scale. aiter's indexer_qk_rope_quant_and_cache covers all of it in one launch, and fusing lets the Hadamard rotation go entirely -- the indexer logits are invariant to it, so dropping it is safe as long as the write and the read back agree on the basis. ROCm now shares upstream's use_dsa_indexer_fusion flag rather than carrying a separate opt-in. On gfx95 the flag additionally requires a LayerNorm k_norm, head_dim 128, rope_head_dim 64 and block_size == head_dim, which is what the kernel accepts. The flag also folds wk and weights_proj into one bf16 wk_weights_proj parameter, so the weight loader fills it from either shard and indexer LoRA, which wraps modules that no longer exist, now raises instead of being silently dropped. Two places needed the platform test that the flag used to imply: the k-only prefill path, which has no aiter counterpart, and the graph head-gate scale, whose custom op is CUDA-only. Measured on MI355X, GLM-5.2-MXFP4 TP4, 8k prefill, per indexer layer per forward: 189.3 -> 55.3 us, a 3.4x cut worth ~2.8 ms per forward across the 21 indexer layers. Needs an aiter carrying the vectorized form of the kernel; the earlier scalar one runs at 228 us and is slower than the unfused path. Co-authored-by: Cursor <cursoragent@cursor.com>
249b762 to
82bd89f
Compare
|
/rerun-failed-ci |
CI Status for PR #34394Merge verdict: ❌ Do not merge on green — PR CI is incomplete and does not exercise this PR's changed code. Every executed failure is unrelated to this PR (0 related failures out of 6 executed), but the AMD pipeline fast-failed at stage-a on an unrelated Caution PR CI is incomplete AND this PR's core code path is unexercised. The AMD run (PR Test ROCm 7.2 (AMD)) fast-failed at Changed files: Executed CI failure attribution: AMD: 1 executed failure (0 related) · Others: 5 executed failures (0 related). NVIDIA base run overall cancelled but 52/65 jobs passed — see notes. AMD Executed Failures
Other Executed Failures
Not real failures / not tested (excluded from the count):
Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
1 similar comment
|
/rerun-failed-ci |
|
/rerun-failed-ci |
4 participants
Summary
On ROCm the DSA indexer spends 11 kernels per decode layer on q/k prep. This PR collapses most of that into one aiter kernel.
indexer_qk_rope_quant_and_cachedoes k LayerNorm + RoPE + fp8 quant + paged index-k store, q RoPE + fp8 quant, and the head-gate scale in a single launch — six kernels become one. Like the CUDA fused path it drops the Hadamard rotation, which is logit-preserving and only shifts fp8 quantisation error. Prefill drops it too, since prefill writes the index-K cache decode reads back and the two must agree on basis.wkandweights_projare folded into onewk_weights_projGEMM. This needs only the concatenated weight, no fused kernel, but was gated on CUDA; it now runs on CUDA or HIP and is numerically identical to the split projections.SGLANG_USE_AITER, and an aiter build exporting the symbol — a capability probe, not a preference: anything outside the envelope falls back to the existing path. CUDA and NPU are unchanged.Decode indexer on GLM-5.2 / MI355X: 11 kernels / 70.8 us per layer → 5 kernels / 41.4 us (−41.5%), which closes the MI355X-vs-B200 indexer gap from 2.55x to 1.49x.
Drive-by fix: define the module-level
_use_dsa_indexer_fusionthatlora_manager.pyalready imports but nothing defined, so an indexer-targeted LoRA adapter raisedImportErrorinstead of the intendedValueError.Test plan
MI355X TP4, GLM-5.2-MXFP4, fp8 KV,
chunked_prefill=16384, DSA backend triton, i8192 / o1024, one run per point. Baseline is the same AMD GLM-5.2 stack without this PR (#30519, #30575, #30715, #30808, #31323, #31324, #32888, plus ROCm/aiter#4453).ITL improves −4.2% to −5.5% at every concurrency, matching the 29.4 us/layer the kernel table predicts over the 21 layers that run the indexer (0.62 ms per decode forward). TTFT is flat, which is the check that matters given the kernel is decode-only.
GSM8K (1200 questions): 0.928 baseline, 0.934 this PR. The index-K cache is bit-exact against the unfused path across both
preshuffleandscale_fmt; the q fp8 payload differs by 1 ULP on 0.10% of elements, because aiter multiplies by the reciprocal rather than dividing.Per-layer decode indexer kernels, and why the aiter kernel is decode-only
hgemm_16x64x128x3(wk)hgemm_16x64x128x3(merged wk_weights_proj)hgemm_16x64x64x6(weights_proj)Layernorm2dFwdindexer_qk_rope_quant_and_cachekn_entry_2c_sbhd_cached_indirect(RoPE)fast_hadamard_transform_kernel_hadamard_quant_kernelindexer_k_quant_and_cache_kerneltriton_poi_fused_mul_unsqueeze_0wv_splitk_small(q proj)wv_splitk_smallfp8_paged_mqa_logits_preshufflefp8_paged_mqa_logits_preshuffletopk_transform_decode_kerneltopk_transform_decode_kernelOf the 29.4 us/layer saved, 5.4 is the merged projection and 21.3 is six kernels becoming one; downstream topk and logits move by 0.1–0.3 us, which is noise.
With the aiter kernel enabled for prefill too, the prefill indexer went from 1614.5 up to 1780.5 us/layer: the fused kernel alone costs 300.4 us/layer against the 162.8 of the six it replaces, and drags downstream topk and logits ~8% each. Gating it to decode gives prefill 1519.6 us/layer (−5.9%), since prefill still benefits from the merged projection and from dropping the Hadamard.
CI States
Latest PR Test (Base): ✅ Run #34937420770
Latest PR Test (Extra): ❌ Run #34937420593
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.