[JaredforReal]

Purpose

Glm5NextLinearAttention.get_state_dtype() and Glm5NextForCausalLM.get_mamba_state_dtype_from_config() call MambaStateDtypeCalculator.kda_state_dtype() without the SSM cache dtype. As a result, the KDA recurrent state is always float32, and --mamba-ssm-cache-dtype bfloat16 is silently ignored for GLM-5.3-Flash. Kimi-K3's KDA already passes it through.

This PR passes cache_config.mamba_ssm_cache_dtype in both places. Glm5NextForConditionalGeneration delegates to the causal-LM class, so it is fixed too.

Why it matters: with bfloat16 the recurrent state page halves. Mamba align-mode prefix caching sizes the attention block from that page, so at TP2 the block shrinks from 2304 to 1280 tokens. With speculative decoding, each running request holds 2 + num_speculative_tokens state blocks per group, so on decode instances this state is what bounds concurrency.

Test Plan

• New test tests/models/glm5next/test_kda_state_dtype.py:
  • auto / float32 / bfloat16 must reach both the model-level and the layer-level recurrent-state dtype;
  • the bf16 case fails on main.
• Existing tests: tests/models/glm5next.
• End to end: GLM-5.3-Flash, TP2, B300, with --mamba-ssm-cache-dtype bfloat16:
  • the KV-cache spec now reports the bf16 state page;
  • teacher-forced NLL on long-context prompts is unchanged within run-to-run noise vs the float32 state.

Test Result

• tests/models/glm5next: 14 passed.
• With bf16 state, a spec-decoding decode instance (K=4, P/D serving) fits ~1.25× more concurrent requests before running out of state blocks (≈64 → 80+).

AI assistance (Claude) was used; every line was reviewed and the tests above were run by the submitter.

[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.

[lesj0610]

LGTM. The model and layer now agree on the requested recurrent-state dtype, and the default remains FP32. The auto, FP32, and BF16 dtype tests pass.

[gcanlin]

I think this is a good feature for performance. We also enable --mamba-ssm-cache-dtype bfloat16 for Kimi K3. cc @ZJY0516

[gcanlin]

I tested the same change (the same two call sites, as a local patch) on nightly 8cbd5d03, and it helps a long-context agentic workload. Sharing the numbers in case they are useful for review.

Setup: GLM-5.3-Flash, 4× B200, TP1 × DP4 + EP, FP8 KV cache, MTP K=5, --prefix-match-unit 16, SimpleCPUOffloadConnector (300 GB per rank). MoE ran through the MegaMoE path from #60895, so the absolute numbers are not stock-main numbers. One run per arm.

Startup, per rank:

	fp32 state	--mamba-ssm-cache-dtype bfloat16
Attention block size (tokens)	8,960	4,864
GPU KV cache size (tokens)	8,821,237	9,676,485 (+9.7%)

InferenceX AgentX trace replay (aiperf), 96 concurrent sessions, 3,600 s, mean prompt about 120K tokens:

	fp32 state	bf16 state
Output tok/s	3,332	3,454 (+3.7%)
Total tok/s	402,628	417,912 (+3.8%)
TTFT mean / p50 / p99 (s)	0.81 / 0.49 / 5.99	0.75 / 0.44 / 6.02
ITL mean (ms)	19.5	17.9
Prompt tokens served from the GPU prefix cache	71.2%	86.8%
Prompt tokens reloaded from host offload	22.8%	7.5%
Preemptions	2	0

vllm bench serve, random 1024 in / 1024 out at concurrency 128: 11,691 → 12,289 output tok/s (+5.1%).

Accuracy, with real MTP acceptance:

• GSM8K (lm-eval, 5-shot, 1,319 questions): bf16 0.9735 ± 0.0044, fp32 0.9704 ± 0.0047.
• Key-to-value retrieval from a ~120K-token prompt (24 prompts × 4 positions): 96/96 for both.

The throughput arms used synthetic acceptance (length 3.61). The measurements were run with AI assistance (Claude Code).

[ZJY0516]

I'm okay to change it, but according to my experience, it will have accuracy issue when using bf16

[github-actions]

✅ @JaredforReal, 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.

[ZJY0516]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94083 for commit 817d5d4a075f.

[Copilot]

🟢 Approval recommended

The focused changes correctly use the existing KDA dtype API and include appropriate regression coverage.

0 open findings

What changed in this PR

Fixes GLM-5.3-Flash KDA recurrent-state dtype selection to honor the configured SSM cache dtype.

Changes:

• Passes mamba_ssm_cache_dtype through model- and layer-level dtype calculation.
• Adds coverage for automatic, float32, and bfloat16 configurations.

File	Description
vllm/​models/​glm5next/​common/​model.py	Propagates the configured SSM cache dtype at model level.
vllm/​models/​glm5next/​common/​kda.py	Propagates the dtype when constructing KDA layers.
tests/​models/​glm5next/​test_kda_state_dtype.py	Verifies consistent dtype selection across model wrappers and layers.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.