[AMD] [GLM5] Fuse the DSA indexer Q/K rope + quant + K-cache write in…

[namle-213]

Motivation

The DSA indexer Q/K fusion (sgl-project#27705) is CUDA-only. On ROCm the indexer still runs
the unfused path every layer, every step: rope on q and k, LayerNorm on k, an fp8
quant of each, the head-gate scale, and the index-K cache write. On GLM-5.2 decode
that work is launch-bound, not arithmetic-bound.

aiter ships indexer_qk_rope_quant_and_cache, which does all of it in one launch.
This wires it up by extending the existing fusion switch to ROCm rather than adding
a second one -- the fused wk_weights_proj GEMM, the no-Hadamard invariant, the
weight loader and the graph split-op plumbing are already platform-agnostic. No
kernel changes; the kernel is already in the aiter that ships in the ROCm image.

Modifications

• Extend use_dsa_indexer_fusion to ROCm, gated on probing aiter for the kernel at
  import (aiter is pinned per image; an older one imports fine and would only fail
  at the first forward, so it falls back with a warning).
• Add _aiter_fused_qk_prepare_and_store, the single-kernel Q/K path. It returns
  the same (q_fp8, weights) contract as the CUDA path, including the num_tokens
  slicing the graph split-op relies on, so the call site is unchanged.
• Build k_norm in fp32 whenever fusion is on. The kernel requires fp32 norm params,
  exactly as the CUDA kernels already do. bf16 stays for the unfused ROCm path, where
  matching x.dtype is what selects aiter's CK layernorm over the native upcast.
• Read cos/sin off aiter's rope module, which keeps them apart as
  [max_position, 1, 1, rope_dim/2] rather than in one cos_sin_cache.
• Give _fused_k_prepare_and_store a ROCm branch. Its CUDA fallback is a JIT kernel,
  and the k-only decode fast path ([AMD] [GLM5] Skip DSA decode indexer when kv_len <= index_topk (dense k-only fast path) sgl-project/sglang#31324) reaches it whenever fusion is on.
• Gate the JIT fused-store branch on _is_cuda so non-CUDA stops paying for a compile
  attempt that can only fail.

Two adjacent fixes, easy to split out if preferred:

• lora_manager imports _use_dsa_indexer_fusion, which [Fix] Fix DSA indexer fusion for NeoX RoPE sgl-project/sglang#30111 deleted when it moved
  the flag into Indexer.__init__. The indexer-LoRA guard therefore raises
  ImportError instead of its intended error, on CUDA today. Restored as
  dsa_indexer_fusion_supported().
• The ROCm branch of _store_index_k_cache read forward_batch.out_cache_loc,
  ignoring the sliced out_cache_loc its caller passes under the graph split-op
  contract.

Accuracy

GLM-5.2-FP8, TP8, 8x MI355X (gfx950).

python3 -m sglang.test.run_eval --port 30000 --eval-name gsm8k \
  --num-examples 500 --num-threads 128

Two independent runs per side, measured on two different bases (this PR's base and
upstream main, whose dsa_indexer.py is byte-identical):

run	baseline	fused
1	0.948	0.962
2	0.958	0.952
mean	0.953	0.957

sigma is ~0.0093 for n=500, and the per-run difference changes sign between the two
pairs, so this is no measurable change in either direction -- not a regression and not
an improvement.

GPQA-Diamond

python3 -m sglang.test.run_eval --port 30000 --eval-name gpqa \
  --num-examples 198 --max-tokens 128000 --repeat 8 --num-threads 16 \
  --top-p 0.95 --temperature 1.0 --thinking-mode glm-45

Fused build only; there is no on-box control run.

	value
mean of 8 repeats	0.911
per-repeat	0.919, 0.919, 0.909, 0.899, 0.904, 0.904, 0.919, 0.914

All 1584 requests returned 200, zero dropped or retried, and zero
retract_decode: aborted last request (ordinary retraction preserves output_ids,
so it costs throughput, not correctness).

The 8 repeats reuse the identical 198 questions with identical answer orderings
(random.Random(0) in simple_eval_gpqa.py), so they are correlated and estimate
temperature-1.0 sampling variance only -- not 1584 independent trials. Observed stdev
across repeats is 0.0080 against a 0.0202 single-run binomial SE, which is consistent
with that. Absent a control run this is an absolute quality check, not evidence about
the fusion; the paired evidence for the kernel is the gsm8k table above.

RULER / NIAH long context

lm_eval RULER niah_single_2, thinking on, 500 samples per context length. The
server is served at --context-length 262144 because the 131072 tier needs room for
the prompt plus an 8192-token generation budget.

python3 -m lm_eval run --model local-chat-completions \
  --model_args model=<model>,base_url=http://127.0.0.1:30000/v1/chat/completions,\
tokenizer=<model>,num_concurrent=8,max_length=262144,timeout=3600,tokenized_requests=False \
  --tasks niah_single_2 --metadata '{"max_seq_lengths":[65536,131072]}' \
  --gen_kwargs max_gen_toks=8192 --apply_chat_template \
  --seed 1234 --log_samples --output_path ./niah

context	niah_single_2
65536	0.998
131072	1.000

1000 records, 1 empty generation, 64 min wall clock. --model local-chat-completions
rather than lm_eval's in-process vllm backend, so the eval runs against this PR's
SGLang server instead of loading a second engine.

Cross-checked against NVIDIA/RULER driven through its native sglang client at n=100
with thinking off and RULER's own 128-token budget: 100.0 at both lengths. Two
harnesses, two prompting regimes and a 5x sample difference agree, so the result is a
property of the model and kernel rather than of either harness.

The two K-cache writers stay live at once, since the k-only decode path keeps writing
unfused, and they are not byte-identical: the unfused path rounds to bf16 before
quantizing while the fused kernel goes fp32 -> fp8 in one step, so a value near an fp8
midpoint can land on either neighbour. Measured over 41x128 elements: ue8m0 scales
identical, 15/5248 elements differ, each by exactly one fp8 code, and the fused result
is the closer of the two to an fp32 reference. test_dsa_indexer_qk_fuse_rocm.py pins
that bound.

Benchmark

python3 -m sglang.benchmark.serving --backend sglang-oai --port 30000 \
  --dataset-name random --random-input-len 16384 --random-output-len 512 \
  --random-range-ratio 1 --max-concurrency 1 --num-prompts 8

BS=1 at 16K context, where kv_len > index_topk so the full indexer runs and decode
is launch-bound. Two runs per side, both stable to two decimals:

	baseline	fused
output tok/s	70.86 / 70.90	75.24 / 75.25	+6.2%
mean TPOT	12.47 ms	11.65 ms	-6.6%
mean TTFT	847 / 851 ms	848 / 849 ms	flat

The same comparison against upstream main as the base gives 70.74/70.76 -> 75.32/75.33
(+6.5%), so the win does not depend on which of the two bases it is measured from.

TTFT is flat because prefill runs the same indexer but is compute-bound; the whole
gain lands in decode, which is the signature a launch-count fix should have.

At --random-input-len 8192 --random-output-len 1024 --max-concurrency 64,
1278.73 -> 1301.31 output tok/s (+1.8%). Launch cost amortizes across a large batch,
so that point reads as "no throughput regression" rather than as the fusion's
contribution.

A short-context BS=1 point is not a measurement of this change: with
index_topk=2048, a 32-token prompt never exceeds kv_len 2048, so every decode step
takes the k-only dense path and never calls the fused kernel.

Profile

Torch profiler, GPU activities, decode at 16K context, measured on this PR's base and
this PR's head. Both captures cover exactly 2.0 decode steps, normalising by
_fused_qk_rope_cat_and_cache_mla_kernel, which fires once per layer per step across
all 78 layers. GLM-5.2 runs a full indexer on 21 of those 78 (indexer_types: 21
full, 57 shared).

Launches per full-indexer layer per decode step (raw counts in parentheses):

kernel	baseline	fused
kn_entry_2c_sbhd_cached_indirect_inplace (rope q+k)	1.00 (42)	0.00 (0)
Layernorm2dFwd (aiter CK)	1.00 (42)	0.00 (0)
act_quant_kernel__kernel	1.00 (42)	0.00 (0)
indexer_k_quant_and_cache_kernel	1.00 (42)	0.00 (0)
indexer_qk_rope_quant_and_cache_kernel	0.00 (0)	1.00 (42)
total	4.00	1.00

84 launches per decode step across the 21 full-indexer layers become 21.

Checklist

• test/registered/kernels/ops/attention/test_dsa_indexer_qk_fuse_rocm.py, 6 cases,
  all passing on MI355X: K-cache agreement with the unfused writer in both paged
  layouts, the one-fp8-step numeric bound, q/head-gate agreement against
  rope + act_quant + _scale_head_gates, strided-vs-contiguous byte equality, and
  the cos/sin 2-D view tracking a replaced rope cache.
• Registered under the jit-kernel-unit-test-amd suite.