[ROCm][Model][Bugfix] Enable GLM-5.2-MXFP4 on the deepseek_v32 path and fix sparse attention correctness

[jhu960213]

Purpose

Enables GLM-5.2 (GlmMoeDsaForCausalLM) end-to-end on vllm/models/deepseek_v32/amd/
for gfx942/gfx950. Routing is opt-in via --model-class-overrides; the registry entry
is unchanged, so the default path for GLM-5.2 and DeepSeek-V3.2 is untouched.

The following issues were also fixed as a result of this overall enablement since it surfaced dormant bugs from a prior deepseek_v32 porting work:

• launch_pdl forwarded to a Triton runtime that doesn't accept it — the shared
  common/kernels.py passed this CUDA-only kwarg unconditionally; ROCm's Triton rejects
  it by name even when False, so the engine died at startup with KeyError. Now passed
  only when the platform actually supports PDL.

• fp4 bmm invoked with mismatched positional args — the output tensor landed in the
  kernel's dtype slot (Expected 'Optional[int]' for 'dtype' but found Tensor). Routed
  through the existing rocm_aiter_ops.batched_gemm_a16wfp4 wrapper rather than calling
  aiter directly, so the argument contract lives in one place.

• fp4 MLA bmm silently wrong for non-power-of-two K — batched_gemm_a16wfp4 only
  supports power-of-two K, and returns wrong values without erroring otherwise. GLM-5.2's
  qk_nope_head_dim=192 lands in the unsupported set; DeepSeek-V3.2's 128 does not.
  Guarded on both qk_nope_head_dim and kv_lora_rank to fall back to the fp8 bmm, with
  a warning. W_K/W_K_scale are also made contiguous after the transpose, which
  otherwise produced NaN.

• Indexer-K cache written in the plain layout while its reader expects the shuffled
  one — affects DeepSeek-V3.2 identically, not just GLM-5.2. The AMD path writes this
  cache from the shared fused_norm_rope, which only implemented
  [n_blocks, blk, head_dim], while the reader selects SHUFFLE above block size 1: the
  write kept storing row-major, but the reader expected shuffled, so the indexer scored
  garbage K and output degenerated once a sequence exceeded index_topk=2048, with no
  error raised. Refactored so the shared kernel stays platform-agnostic — it writes
  whichever layout the cache reports via a new uses_shuffled_layout property — and the
  AMD-specific shuffled write ([n_blocks, blk/16, head_dim/16, 16, 16], used above block
  size 1) lives in the AMD backend. The base class returns False, so CUDA is unchanged.

• Sparse-MLA block sizes narrowed instead of extended — the ported backends replaced
  the inherited get_supported_kernel_block_sizes() with [16, 32], which made
  select_common_block_size silently downgrade a requested 64 to 32 via its
  largest-divisor fallback. Now concatenated with the base list.

• Kernel tests gated on is_cuda() and running against dims that matched neither
  model — the tests were skipped entirely on CDNA. The gate is now supports_fp8(), so
  the suite runs on gfx950. The hardcoded attention dims were written as if one set served
  both GLM-5.2 and DeepSeek-V3.2, which it did not; they are now taken from each model's
  published config and parametrized over both.

Test Plan

Server:

VLLM_ROCM_USE_AITER=1
HF_HUB_OFFLINE=1 SAFETENSORS_FAST_GPU=1
vllm serve /path-to-your/GLM-5.2-MXFP4
--model-class-overrides '{"GlmMoeDsaForCausalLM": "vllm.models.deepseek_v32:DeepseekV32ForCausalLM"}'
--tensor-parallel-size 8
--kv-cache-dtype auto
--max-model-len 65536
--no-enable-prefix-caching
--trust-remote-code
--port 8000

Variants: --kv-cache-dtype fp8 for the fp8 KV row; --block-size {16,64} for the
block-size matrix. Drop --model-class-overrides to run the legacy deepseek_v2 path
for comparison.

Accuracy:

lm_eval --model local-completions
--model_args "model=/models/GLM-5.2-MXFP4,base_url=http://localhost:8000/v1/completions,num_concurrent=32,tokenized_requests=False,max_retries=3,timeout=3600"
--tasks gsm8k --num_fewshot 5 --gen_kwargs max_gen_toks=2048 --limit 200

Throughput:

vllm bench serve --backend vllm --model /shareddata/models/GLM-5.2-MXFP4
--dataset-name random --random-input-len 4096 --random-output-len 1024
--num-prompts 200 --max-concurrency 32 --ignore-eos --port 8000
--save-result --result-filename glm52_v32amd_bench.json

Kernel unit tests:

pytest tests/kernels/test_fused_deepseek_v32_norm_rope.py -v

Test Result

gsm8k — GLM-5.2 200 samples, 5-shot, exact_match, deepseek_v32 AMD path

KV cache	Before (flexible / strict)	After (flexible / strict)
auto (bf16)	0.0000 / 0.0000	0.9400 / 0.9400 ± 0.0168
fp8	0.0000 / 0.0000 *	0.9350 / 0.9350 ± 0.0175

* Not separately measured. fp8 KV goes through the same absorbed-weight bmm where the
bug was, so it fails identically; only the bf16 run was evaluated pre-fix.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jhu960213.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Rohan138]

there's another GLM bug rn btw: #56604

Blocks merging or not really?

eh not really, plus #57252 is already set up for auto-merge to fix it.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jhu960213.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @jhu960213, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[mergify]

Hi @jhu960213, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[dllehr-amd]

/ci run

[github-actions]

✅ @jhu960213, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[github-actions]

✅ Triggered Buildkite CI #90294 for commit 933d45cba075.

[dllehr-amd]

Thanks for this one @jhu960213 ! Sorry for the delay!