[AMD/MI300A] DeepSeek-V4 DSA indexer + attention: validate/provide gfx942 forward path

[xzyaoi]

Context

DeepSeek-V4 (deepseek_v4) on AMD MI300A (gfx942). V4's attention is a hybrid Compressed Sparse Attention (CSA) + Heavily Compressed Attention (HCA) with a DSA indexer that selects top-512 (Flash) / 1024 (Pro) KV per query (index_n_heads: 64, index_head_dim: 128), plus per-layer compress_ratios and sliding_window: 128.

Status (single-node probe, MI300A, 2026-06-11)

Good: the V4 model imports and constructs cleanly on AMD — the attention + indexer modules build without error (the probe reached MoE weight allocation past them). No import wall.

Unknown: the forward path is unvalidated — the probe OOM'd at MoE weight alloc before running a forward (blocked on the mxfp4-EP gap; see companion issue). The indexer forward uses NVIDIA-only kernels:

• deep_gemm.fp8_fp4_mqa_logits(...) (thirdparty/deep_gemm, CUDA) for the indexer logits, and
• indexer_mxfp4_paged_gather(...) (tokenspeed_kernel/ops/attention/cuda/deepseek_v4.py, CUDA-only — the Triton variant ops/attention/triton/deepseek_v4.py does not have the mxfp4 paged gather).

There is a fallback gate _deepseek_v4_deepgemm_fp4_indexer_available() in models/deepseek_v4.py, but its AMD branch is unverified.

Ask

Once the mxfp4-EP blocker is resolved and a V4 model fits on MI300A, run a forward and:

1. Validate the indexer + CSA/HCA attention produce correct output on gfx942 (the fallback path may already work via Triton).
2. If the deep_gemm fp8/fp4 MQA-logits and the cuda mxfp4 paged-gather are required, provide gfx942 equivalents (Triton/Gluon or AITER) for: the indexer logits (fp8/fp4 MQA) and the mxfp4 paged KV gather.

Repro / HW

jobs/serve-v4pro-1node-probe2.sbatch on beverin; V4-Pro staged at infra01/hf_models/.../DeepSeek-V4-Pro. MI300A gfx942 / ROCm 7.2 / torch 2.11. Companion to the mxfp4-EP issue + the V4 tracking issue.

[xzyaoi]

Validated the DSA-indexer forward path on MI300A / gfx942 and shipped a portable gfx942 logits op in #31 (draft).

What shipped: xkernels.dsa_indexer_logits (REFERENCE oracle + Triton gfx942 backend) + dsa_indexer_topk. It computes the indexer score that selects which KV survive — logits[t,j] = sum_h weights[t,h] * relu(q[t,h,:]·k[j,:]) (q [T,H,D], shared MQA k [K,D], per-head weights, optional causal window) — numerically identical to tokenspeed's torch oracle _indexer_topk_reference. This is the portable replacement for the NVIDIA-only deep_gemm.fp8_fp4_mqa_logits.

On-device (beverin, MI300A gfx942, job 381969):

• 7/7 GPU pytest (real Triton bf16 compile, incl. per-query causal masking).
• V4 shape H=64 D=128 K=4096 bf16: max|err| = 6.10e-05 (rel 1.76e-07) vs fp32 oracle.
• Flash top-512 selection: mean Jaccard = 1.0000 (perfect parity).

Local CPU gate (TRITON_INTERPRET=1): 7/7 + full suite 101 passed; ruff clean.

Caveats / deferred (called out in the PR + docs/issue-27-dsa-indexer.md):

• The fp8/fp4 packing in the upstream CUDA kernel is an encoding detail that does not change the selection, so the gfx942 op uses the dequantized-math equivalent (bf16/fp16 in, fp32 logits out). A native fp8/fp4 logits kernel + the mxfp4 paged-gather are a bandwidth optimization on top — their own measured follow-up, not needed for a correct forward path. With dequantized math the gather is an ordinary k[block_table] index.
• Wiring into models/deepseek_v4.py lives in the tokenspeed runtime, not this kernels package; this PR provides + validates the gfx942 op the runtime can call.
• Gotcha caught only on-device: a float("-inf") module global referenced inside @triton.jit compiles fine under TRITON_INTERPRET=1 but is rejected by the real Triton compiler — inlined the literal.

PR: #31

[xzyaoi]

DSA indexer gfx942 forward path delivered + parity-validated by #31 (dsa_indexer_logits / dsa_indexer_topk — Triton MQA weighted-ReLU score, numerically identical to tokenspeed's _indexer_topk_reference) and #29 (Triton mxfp4 paged KV gather). The remaining V4 attention compute (CSA/HCA/SWA sparse-MLA flash kernel that consumes these top-k indices) is tracked separately in #32. Closing the indexer/selection half.