Repository navigation
[Qwen3.5] asm FP8 FMHA for prefix-hit prefill; fuse RMSNorm+FP8 quant and tiled sigmoid-gate+FP8 quant - #2534
Conversation
_can_use_fp8_prefill_attention refused any batch with a cached prefix (`not has_cached`, and `q.shape[0] == k.shape[0]` evaluated after the prefix gather), so every chunk after the first, and every prefix-cache hit, fell back to aiter.flash_attn_varlen_func in BF16. With GQA 8 per rank (Qwen3.5 at TP4) the BF16 hd256 asm kernel does not apply, so that meant the CK FmhaFwdKernel. The asm fmha_fwd_hd256_fp8_causal_group kernel already takes the gathered cached+new K/V with cu_seqlens_k longer than cu_seqlens_q and masks bottom-right causal, which is the same call SGLang's aiter backend makes for context-chunk prefill. Microbench on gfx950 (8 q heads, 1 kv head, hd256): 16k new / 32k ctx 5878 -> 1408 us, 16k / 128k 29653 -> 6154 us, including the three FP8 casts.
triton_fused_norm_rope_cache picked each rotary frequency's position by contiguous mrope_section ranges (f < s0 -> t, f < s0+s1 -> h, else w). Qwen3.5 sets mrope_interleaved, where frequency f takes h when f % 3 == 1 and f < 3*section_h, w when f % 3 == 2 and f < 3*section_w, else t -- the rule triton_mrope.py and aiter's MRotaryEmbedding already use. Text-only positions (t == h == w) gave the same result either way; image and video tokens did not. The kernel now takes MROPE_INTERLEAVED / MROPE_SH / MROPE_SW from rotary_emb.mrope_interleaved and mrope_section. Against GemmaRMSNorm + aiter MRotaryEmbedding with t != h != w positions: max |dq| 0.031, |dk| 0.031 (bf16 rounding), versus 8.97 / 8.00 with the old rule.
Qwen3_5DecoderLayer built input_layernorm without a quant config, so the norm wrote BF16 and qkv_proj / in_proj_qkvz quantized it again in a separate dynamic_per_group_scaled_quant launch. Qwen3NextDecoderLayer already passes the consuming projection's layer quant config; the inherited forward routes the FP8 output and scales to qkv_proj and to the GDN layer's x_fp8 / x_scale. Do the same: full-attention layers take qkv_proj's config, linear-attention layers take in_proj_qkvz's with write_bf16=True for the BF16 in_proj_ba. GemmaRMSNorm only fuses for per_1x128, so BF16 and MXFP4 checkpoints are unchanged.
For a pure prefill (no spec tokens) the non-spec rows are rows [0, T) of the custom op's output, but chunk_gated_delta_rule wrote into a fresh `o` and the merge step copied it over: a 64 MiB hipMemcpy per GDN layer at T=16384, 28.7 us in the c16 trace (labelled gdn_state_copy there, but it is the output). Pass that slice as `o=`, which chunk_gated_delta_rule already supports and checks for shape, dtype and contiguity, and skip the copy when it was used.
Each GDN layer's prefill path cast the int32 read and write slot indices to int64 for advanced indexing, ran `~has_initial_state`, and zeroed fresh states through boolean-mask indexing: three extra launches per layer (about 12 us at c16) that only depend on the step. The metadata builder now also carries no_initial_state and int64 copies of the read / write slots; the layer gathers with index_select, zeroes with masked_fill_, and writes back with index_copy_.
… count With ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 both one-kernel fusions are slower than their unfused twins at prefill sizes, though they save launches at decode sizes. Qwen3.5 TP4 on MI355X, fused vs unfused (us): tokens 16 2048 4096 8192 16384 gate + quant 21.5 21.6 41.6 78.5 147.4 vs 27.9 26.6 28.5 47.3 87.2 norm+rope+cache 38.9 38.7 39.2 60.2 113.1 vs 64.7 66.4 66.4 67.1 82.5 Above ATOM_QK_NORM_ROPE_CACHE_FUSION_MAX_TOKENS (8192) the attention impl now runs fused_qk_norm, the MRoPE Triton kernel and the per-token FP8 cache write separately, writing the same SHUFFLE layout and scales. Above ATOM_SIGMOID_MUL_QUANT_FUSION_MAX_TOKENS (2048) the gate runs as o * sigmoid(gate) followed by per_group_quant_hip, the same scales o_proj would produce itself. The gate choice sits in a custom op (sigmoid_mul_fp8_group_quant) so it is made per call; inside the compiled graph a branch on M would be fixed by the traced size. prefill[bs=1 tok=16368 ctx=32752] per full-attention layer: 6532 -> 6417 us; decode[bs=6 tok=24] keeps the fused kernels (316 -> 317 us). GSM8K 3-shot, first 256: 0.8867 flexible vs 0.8711 on main (McNemar p=0.48).
The per-group kernel runs one program per (row, 128-column group), which scales badly with rows: 6.3 us of GPU time at 384 rows and 147 us at 16384 (Qwen3.5 TP4, N=2048, MI355X). The previous commit therefore sent prefill-sized batches to sigmoid*o + per_group_quant_hip, which still took 95 us. Add _fused_sigmoid_mul_fp8_group_quant_tiled_kernel: one program per BLOCK_M x BLOCK_N tile, with the same arithmetic in the same order, so its FP8 bytes and scales are bit-identical to the per-group kernel's at every size tested. sigmoid_mul_fp8_group_quant now always uses it, picking the tile from the row count: 8 x 256 on 4 warps below 4096 rows, 8 x 1024 on 8 warps from there (BLOCK_N halves until it divides N). GPU time, per-group vs tiled: rows 24 384 1024 16384 per-group 3.06 6.34 12.44 147 tiled 3.17 3.66 4.92 30.6 With no fallback left to choose, ATOM_SIGMOID_MUL_QUANT_FUSION_MAX_TOKENS is removed. In the prefill[bs=1 tok=16368 ctx=32752] trace the gate costs 34 us per full-attention layer, down from 95 us; decode is 4.0 us, down from 4.2 us.
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
|
@junhaha666 Could you please help review the PR? |
|
Reviewed commit
Validation scope: I inspected all nine changed files and the relevant callers, and syntax checks passed for all eight changed Python files. The GitHub checks that ran passed, but those results do not rule out the two scenarios above. These findings are based on code-path and interface inspection; I did not run GPU or end-to-end reproductions because the local Python environment lacks PyTorch. |
Conflict in atom/models/qwen3_5.py: #2534 added the same per_1x128 input-norm quant config and write_bf16 to Qwen3_5DecoderLayer that this branch already has; kept this branch's version, which also carries fused_allreduce. attention_gdn.py and qwen3_next.py merged cleanly: #2534's prefill direct write (wrote_core_attn_out) and sigmoid_mul_fp8_group_quant rename sit next to this branch's spec-decode o= path and reduce_results plumbing. Signed-off-by: zhuyuhua-v <yuhzhu@amd.com>
Summary
This PR speeds up Qwen3.5 on gfx950, mainly in prefill.
Measured on MI355X, Qwen3.5-397B-A17B-FP8, TP4, against main (details under Results):
New attention kernel for prefix-hit prefill.
fmha_fwd_hd256_fp8_causal_group_gfx950. Before, it ran on CKFmhaFwdKernel<bf16>._can_use_fp8_prefill_attentionblocked it:not attn_metadata.has_cachedturned it off whenever a prefix was cached.q.shape[0] == k.shape[0]was checked after the prefix gather, whenkholds cached + new rows, so it never held.flash_attn_varlen_func. Qwen3.5 at TP4 has GQA 8 per rank (8 q heads, 1 kv head), and the BF16 hd256 asm kernel needsnhead_q == nhead_k. So the fallback ran the CK kernel.cu_seqlens_klonger thancu_seqlens_qand applies a bottom-right causal mask. SGLang's aiter backend uses the same kernel and the same shape of call for context-chunk prefill.Fusions.
Qwen3NextDecoderLayer, and it drops onedynamic_per_group_scaled_quantper layer.sigmoid(gate) * o+ FP8 per-group quant foro_projis a new tiled kernel,_fused_sigmoid_mul_fp8_group_quant_tiled_kernel.sigmoid_mul_fp8_group_quant, so the tile is picked from each call's row count. A branch inside the compiled graph would keep the size it was traced with.triton_fused_norm_rope_cache) stays fused up toATOM_QK_NORM_ROPE_CACHE_FUSION_MAX_TOKENS(default 8192). Above that it runs as separate kernels, because the fused kernel is slower at prefill sizes. The fused kernel also now applies Qwen3.5's interleaved M-RoPE rule (f % 3). Text-only positions were already correct; image and video positions were not.ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1, which still defaults to off. The FMHA change, the input-norm fusion and the GDN changes below are on by default.GDN prefill.
o=). This removes a 64 MiBhipMemcpyper GDN layer at T=16k.~has_initial_stateare now built once per step in the metadata builder instead of once per layer.Results
Setup.
recipes/Qwen3.5.mddefaults.AITER_QUICK_REDUCE_QUANTIZATIONis not set, so the prefill all-reduce runs on RCCL.Same-shape kernel breakdown
How it was measured.
config.jsonlayer_types: 15 full-attention layers and 45 GDN layers.--state-checkpoint-interval-tokens -1on both sides; without it, the 8192-token checkpoint rungs split the second chunk. Its second chunk is the annotationprefill[bs=1 tok=16368 ctx=32752].decode[bs=6 tok=24 d=6 spec=3]steps.prefill[bs=1 tok=16368 ctx=32752]decode[bs=6 tok=24 d=6 spec=3]Prefill, by change. Per-layer values are kernel times taken from the two traces.
.to(fp8)casts 32.4sigmoid*o36.9 +o_projquant 19.8 -> 34.1Decode, by change. Per full-attention layer, derived from the same traces.
o_projquant 4.1_fused_qkv_norm_rope_cache_kernel4.5, tiled gate 4.0qkvquant 4.1Trace breakdown CSVs and the xlsx are on the author's machine and are available on request.
Kernel microbenchmarks
Per-rank shapes.
FMHA, prefix-hit prefill.
Gate + quant.
[M, 2048]bf16; the gate is a strided view of the qkv output, as in the model.Norm + RoPE + cache: fused vs unfused.
End-to-end, AgentX c48
Run setup.
inferencex-agentx-mvpscenario onsemianalysis_cc_traces_weka_062126_256k.--apply-chat-template.ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1.Per request, paired on
(source_trace_id, turn_index).Aggregate, InferenceX metrics.
Caveat on the size of the gain.
Accuracy
All accuracy runs use GSM8K 3-shot via
local-completions, as inrecipes/Qwen3.5.md, with forced acceptance off and paired McNemar tests on the same questions.