GLM-5 DSA indexer fusion on ROCm (diff vs upstream main)

[Jacob0226]

Self-review only, not for upstream. Base is the branch behind sgl-project#34394, so this shows exactly what has changed since that PR was last pushed.

Commits

1. Pass fp32 LayerNorm params to the aiter fused indexer kernel — aiter post-[Bug] Error while serving deepseek-ai/DeepSeek-V2-Lite in NVIDIA A40-48Q sgl-project/sglang#3451 requires norm_weight/norm_bias as fp32 and aborts at the host check otherwise. _k_norm_params() loses its dtype argument and always widens to fp32, so the contract lives in the code instead of in the caller remembering to pass the right thing. Without this the server SIGABRTs on the first decode against current aiter.

2. Ask the model whether the indexer merged weights_proj — replaces the module-level _use_dsa_indexer_fusion mirror with indexer_merges_weights_proj(model), which queries the built Indexer instances. The mirror read the env through get_bool_env_var() while the instance read it through envs.…get() (these disagree on values like yes), and it dropped is_neox_style, so neox models were wrongly rejected by the LoRA validation.

3. Collapse the indexer fusion flags into one — use_dsa_indexer_fusion, merge_wk_weights_proj and use_aiter_indexer_fusion all answered "is the indexer q/k path fused?" from three angles. They also disagreed: the merged GEMM was on for every ROCm device while only gfx95 had a fused kernel to feed it. Now one flag, ROCm opts in only where aiter can actually run; the platform split moves to the two call sites that need it. is_aiter_gfx950 becomes aiter_indexer_qk_fusion_available.

Reviewing tips

k_norm moved ahead of the projections (the fusion decision reads its type), which shifts a block of lines — use "Hide whitespace" and read dsa_indexer.py by commit rather than as a combined diff.

Known scope

Non-gfx95 ROCm loses the merged GEMM. That configuration was never validated end to end, but it does mean parameter layout is now tied to kernel availability, so the merged GEMM can no longer be A/B'd on its own without reinstating an override env.

Prefill still excluded

Measured on MI355X TP4, GLM-5.2-MXFP4, extend chunk 7679 tokens: the fused kernel is 228.8 us/layer against 99.5 us for the split kernels it replaces, while in decode it is 5.6 vs 30.4. The cause is aiter's launch geometry (grid(num_tokens, n_heads), block(head_dim)): one 128-thread workgroup per (token, head) with one bf16 per thread, so block_reduce goes through LDS and every wave moves 128 B per memory instruction. It sustains ~430 GB/s, about 5% of HBM peak, so this is a kernel-tuning problem on the aiter side rather than a reason prefill cannot be fused.

Made with Cursor

---

CI States

Latest PR Test (Base): ❌ Run #32316103572
Latest PR Test (Extra): ❌ Run #32316103421