feat: support GLM-5.2

[CUHKSZzxy]

Motivation

Add GLM-5.2 support to the PyTorch backend, including DSA, MTP, routed-expert replay, and serving integration. The shared DSA changes also improve DeepSeek-V3.2 correctness and decode efficiency.

Changes

• Add GLM-5.2 config/model loading, interleaved RoPE, MoE-only FP8, and tool calls emitted directly from reasoning.
• Implement DSA indexing for GLM-5.2 and DeepSeek-V3.2: radix top-k, fused indexer preparation, shared top-k reuse, and BF16/FP8 sparse FlashMLA cache paths. LMDEPLOY_DISABLE_DSA_INDEXER_FUSION=1 keeps the original path.
• Add GLM-5.2 MTP with aligned draft distribution, shared embeddings, post-norm hidden-state flow, index sharing, and multi-token query handling.
• Add routed-expert capture/replay through prefix caching and OpenAI/Anthropic serving APIs.
• Centralize Transformers config fallback and add coverage for model config, CUDA kernels, cache/attention, MTP/routed-expert replay, prefix caching, and response parsing.

Acc Benchmarks

Academic

Dataset	Metric	LMDeploy	vLLM	LMDeploy − vLLM
MMLU-Pro	Accuracy	86.86%	87.01%	-0.15 pp
AIME 2025	Accuracy, 32-run average	93.33%	93.54%	-0.21 pp
GPQA Diamond	Accuracy, 4-run average	90.03%	90.03%	0.00 pp

with MTP 5

Dataset	Metric	LMDeploy
AIME 2025	Accuracy, 32-run average	93.54%
GPQA Diamond	Accuracy, 4-run average	92.42%

Both engines were evaluated with TP8 under identical settings on AutoToolChain.

Generation used temperature=1.0, top_p=0.95, and a maximum generation length of 163,840 tokens, following the GLM-5.2 model card. Scores may differ from the published results because other evaluation details noted in the model card, such as the reasoning-task system prompt, may not have been fully aligned.

GSM8k

Engine / topology	MTP	Correct / total	Accuracy	95% Wilson CI	Errors	Empty	Cap hits
LMDeploy TP8	off	1266 / 1319	95.98%	94.78–96.91%	0	0	2
LMDeploy TP8	5 tokens	1267 / 1319	96.06%	94.87–96.98%	0	0	1
LMDeploy DP8+EP8	off	1272 / 1319	96.44%	95.29–97.31%	0	0	1
LMDeploy DP8+EP8	5 tokens	1271 / 1319	96.36%	95.21–97.24%	0	0	1
vLLM TP8	off	1271 / 1319	96.36%	95.21–97.24%	0	0	0
vLLM TP8	5 tokens	1265 / 1319	95.91%	94.70–96.85%	0	0	3

Acceptance of MTP

Aggregate results

Short aggregate per-position acceptance rates:

Engine	Position 0	Position 1	Position 2	Position 3	Position 4
LMDeploy	77.95%	49.62%	32.32%	19.58%	10.46%
vLLM	75.05%	48.42%	30.35%	18.81%	11.92%

Long per-position acceptance rates:

Workload	Engine	Position 0	Position 1	Position 2	Position 3	Position 4
>4K	LMDeploy	73.50%	41.55%	22.89%	12.91%	4.59%
>4K	vLLM	72.16%	46.57%	25.29%	12.73%	5.64%
>8K	LMDeploy	77.10%	48.93%	27.53%	14.50%	7.39%
>8K	vLLM	74.12%	45.18%	23.91%	11.81%	5.08%
>16K	LMDeploy	76.70%	48.01%	26.74%	14.92%	8.13%
>16K	vLLM	74.02%	45.07%	25.14%	13.39%	7.15%

Efficiency Benchmark

Throughput

Output length	Requests	TP8 output tok/s	DP8+EP8 output tok/s	Speedup	TP8 duration	DP8+EP8 duration
Natural	10,000	3,228.23	4,831.38	1.50x	610.01 s	407.60 s
2,048	8,000	2,582.32	6,036.53	2.34x	6,344.68 s	2,714.14 s
4,096	8,000	2,145.58	4,049.21	1.89x	15,272.31 s	8,092.45 s
8,192	4,000	1,594.37	3,035.30	1.90x	20,552.32 s	10,795.63 s
16,384	1,000	1,012.60	1,605.38	1.59x	16,180.10 s	10,205.66 s
32,768	500	629.18	991.59	1.58x	26,040.29 s	16,522.91 s

Latency

Output length	TP8 mean TTFT	DP8+EP8 mean TTFT	TP8 mean TPOT	DP8+EP8 mean TPOT
Natural	247.93 s	137.65 s	306.85 ms	399.75 ms
2,048	2,832.64 s	1,022.85 s	253.28 ms	280.64 ms
4,096	7,038.16 s	3,117.33 s	194.88 ms	315.44 ms
8,192	7,136.28 s	2,109.02 s	590.41 ms	542.26 ms
16,384	12.34 s	7.43 s	706.87 ms	452.57 ms
32,768	6.54 s	5.17 s	660.96 ms	413.83 ms

DP8+EP8 config: max batch size 256, cache 0.6

Performance Optimization

Five-layer ablations

Matched five-layer GLM-5.2 FP8 ablations on the same H200 pair: TP2, BF16 MLA KV cache, CUDA graphs enabled, and MTP disabled. Throughput is the mean of three post-warmup batch-256 runs with 128 input and 2,048 forced output tokens (CV ≤ 0.73%). Timeline measurements average both ranks from a one-second batch-128 steady-decode capture.

Optimization	Median cycle	Cycle change	CUDA nodes/cycle	Node change	Output throughput	Gain
Baseline	16.634 ms	—	268	—	14,240 tok/s	—
Zero-copy BF16 sparse MLA	4.512 ms	-72.87%	258	-3.73%	42,122 tok/s	+195.79%
Compact blocked-FP8 MoE scheduling	16.457 ms	-1.06%	280	+4.48%	14,328 tok/s	+0.62%

Both optimizations are measured independently against the same baseline.

The previous BF16 sparse path materialized the non-contiguous paged main-K cache before attention. In a representative five-layer draft forward, the whole-cache copy took 1.876 ms of the 3.027 ms forward; the matched FP8 path had no corresponding copy and took 1.394 ms. The BF16 sparse-attention kernel itself was faster—the avoidable cache materialization was the bottleneck addressed by the zero-copy strided view.

DSA indexer fusion

Full-model GLM-5.2 FP8 TP8 timeline comparison:

Metric	Fusion off	Fusion on	Change
Forward duration	27.345 ms	26.585 ms	-2.78%
CUDA-graph child kernels	4,455	4,182	-273 / -6.13%

The full model has 21 indexer layers, so fused Q/K preprocessing removes exactly 13 CUDA-graph child kernels per layer (21 × 13 = 273).

Validation

• Pre-commit hooks and targeted replay, prefix-cache, engine, and serving tests pass locally (218 tests).
• Real 5-layer FP8 smoke tests passed for GLM-5.2 and DeepSeek-V3.2.
• Indexer fusion improved GLM-5.2 5-layer MTP batch-256 median output throughput by 1.12% across three profiler-free runs; short batch-8 timelines reduced median cycle time by 3.41% for GLM-5.2 and 4.66% for DeepSeek-V3.2.

Assistance

Assisted with Codex + GPT-5.6-Sol xHigh, reviewed manually

[Copilot]

Pull request overview

Adds GLM-5.2 support to the PyTorch backend and extends the serving/replay surface to capture and return sparse-attention indexer top‑k (plus related DSA/MTP improvements shared with DeepSeek‑V3.2). The PR also introduces CUDA kernels (TileLang/Triton) and broad test coverage across parsing, cache/attention, speculative decoding, and serving APIs.

Changes:

• Add GLM‑MoE‑DSA MTP model path, interleaved RoPE support, DSA indexer fusion + shared top‑k reuse, and sparse MLA cache/index improvements.
• Add “indexer_topk” capture/replay end-to-end (sequence history, prefix cache/block trie, engine outputs, OpenAI+Anthropic serving, CLI flags).
• Add/extend unit tests for response parsing (reasoning↔tool boundary), CUDA kernels, paging replay semantics, and engine transfer behavior.

Reviewed changes

Copilot reviewed 71 out of 71 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/test_lmdeploy/serve/test_generation_config.py	Adds request flag coverage for returning indexer top‑k.
tests/test_lmdeploy/serve/parsers/test_glm47_parser.py	Adds GLM‑5.2 reasoning/tool streaming boundary tests.
tests/pytorch/spec_decode/test_strategies.py	Adds spec-decode tests ensuring indexer history trimming/shape behavior.
tests/pytorch/spec_decode/test_spec_agent.py	Generalizes dist-context reuse tests for matching draft topology.
tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py	New tests for GLM MTP + shared top‑k buffer behavior.
tests/pytorch/spec_decode/test_cudagraph_strategy.py	Ensures CUDA graph key separates skip_topk variants.
tests/pytorch/paging/test_block_trie.py	Adds cached indexer_topk replay tests and adjusts routed-expert replay semantics.
tests/pytorch/nn/test_nsa.py	New test for NSA index meta query-width handling during decoding.
tests/pytorch/kernel/test_sparse_index_topk.py	New CUDA tests for TileLang sparse_index_topk selector.
tests/pytorch/kernel/test_mla_attention.py	New tests for MLA sparse decode index updates and metadata behavior.
tests/pytorch/kernel/test_dsa_indexer.py	New tests for fused DSA indexer Q/K preparation correctness.
tests/pytorch/kernel/test_apply_rotary.py	Adds interleaved RoPE tests for apply_rotary_pos_emb.
tests/pytorch/engine/test_engine_sleep.py	Adds engine-instance transfer test for trimmed indexer_topk.
tests/pytorch/engine/test_cache_engine.py	Adds sparse MLA cache layout/policy tests.
lmdeploy/serve/parsers/response_parser.py	Allows tool-open tag to implicitly terminate reasoning in streaming/complete parsing.
lmdeploy/serve/openai/serving_generate.py	Validates return_indexer_topk against engine configuration.
lmdeploy/serve/openai/serving_chat_completion.py	Validates return_indexer_topk for chat-completions.
lmdeploy/serve/openai/protocol.py	Adds return_indexer_topk request flag and response fields.
lmdeploy/serve/openai/api_server.py	Streams/finalizes indexer_topk output and adds serving flag constraints.
lmdeploy/serve/core/async_engine.py	Plumbs indexer_topk through async generation outputs and stop-trimming.
lmdeploy/serve/anthropic/streaming.py	Adds indexer_topk to Anthropic streaming extension fields.
lmdeploy/serve/anthropic/protocol.py	Adds return_indexer_topk + indexer_topk response typing.
lmdeploy/serve/anthropic/endpoints/messages.py	Validates and returns indexer_topk for Messages endpoint.
lmdeploy/serve/anthropic/adapter.py	Maps Anthropic return_indexer_topk to generation config.
lmdeploy/pytorch/transformers/configuration_deepseek_v32.py	Switches DeepSeekV3.2 config import source (hf_configs).
lmdeploy/pytorch/transformers/init.py	Delegates config fallback/registration to lmdeploy.hf_configs.
lmdeploy/pytorch/strategies/ar/sequence.py	Adds indexer_topk plumbing into AR sequence updates and step truncation.
lmdeploy/pytorch/strategies/ar_spec/sequence.py	Adds indexer_topk plumbing for speculative sequences.
lmdeploy/pytorch/spec_decode/proposers/deepseek_mtp.py	Adds shared-topk buffer handling + target logit path for GLM-style MTP.
lmdeploy/pytorch/spec_decode/base.py	Reuses DistContext when draft dist_config matches the target’s.
lmdeploy/pytorch/paging/block_trie.py	Stores/replays indexer_topk through prefix caching & match verification.
lmdeploy/pytorch/nn/rotary_embedding.py	Allows ApplyRotaryEmb to be built with interleaved mode.
lmdeploy/pytorch/nn/nsa.py	Fixes NSA meta max_q_seqlen derivation for flattened multi-token decode queries + adds fused entrypoint.
lmdeploy/pytorch/models/utils/cudagraph.py	Avoids FlashMLA metadata buffers for BF16 sparse decode; captures all_indexer_topk outputs.
lmdeploy/pytorch/models/module_map.py	Registers GLM MTP architecture mapping.
lmdeploy/pytorch/models/glm_moe_dsa_mtp.py	New GLM MoE DSA multi-token predictor model implementation.
lmdeploy/pytorch/models/deepseek_v32.py	Adds shared DSA top‑k buffer, full-vs-shared indexer layer logic, capture of indexer top‑k, and fused indexer weight loading.
lmdeploy/pytorch/models/deepseek_v2.py	Adds routed-expert capture, DP-aware weight layout handling, and FP8 “moe_only” scope guard.
lmdeploy/pytorch/model_inputs.py	Extends BuildModelContext with enable_return_indexer_topk.
lmdeploy/pytorch/messages.py	Adds SamplingParam flag + HistoryIndexerTopK storage and sequence accessors.
lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py	New TileLang sparse index top‑k selector kernel.
lmdeploy/pytorch/kernels/cuda/dsa_indexer.py	New Triton fused DSA indexer preparation kernels.
lmdeploy/pytorch/kernels/cuda/apply_rotary_pos_emb.py	Adds interleaved pairing mode to rotary kernel.
lmdeploy/pytorch/envs.py	Adds LMDEPLOY_DISABLE_DSA_INDEXER_FUSION env flag.
lmdeploy/pytorch/engine/model_agent/agent.py	Plumbs return_indexer_topk and carries all_indexer_topk in batched outputs.
lmdeploy/pytorch/engine/inputs_maker.py	Adds detection of sequences requesting indexer_topk.
lmdeploy/pytorch/engine/engine.py	Extends InferOutput to carry indexer_topk.
lmdeploy/pytorch/engine/engine_loop.py	Appends/caches indexer_topk and includes it in responses.
lmdeploy/pytorch/engine/engine_instance.py	Transfers/validates indexer_topk like routed_experts (incl. stop trimming).
lmdeploy/pytorch/engine/config_builder.py	Adjusts deepseek_mtp draft dist config selection for GLM.
lmdeploy/pytorch/engine/cache_engine.py	Adds sparse MLA cache dtype policy selection and layout sizing updates.
lmdeploy/pytorch/configurations/glm_moe_dsa.py	New GLM-MoE-DSA config normalization + moe-only FP8 scope support.
lmdeploy/pytorch/configurations/deepseek_v32.py	Updates DeepSeekV3.2 env checks and MLA cache dtype behavior.
lmdeploy/pytorch/config.py	Replaces use_mla_fp8_cache flag with mla_kv_cache_dtype + adds enable_return_indexer_topk.
lmdeploy/pytorch/backends/nsa.py	Adds forward_fused abstract API for fused indexer preparation.
lmdeploy/pytorch/backends/dlinfer/apply_rotary_emb.py	Falls back to default impl for interleaved rotary.
lmdeploy/pytorch/backends/default/apply_rotary_emb.py	Adds interleaved rotary implementation.
lmdeploy/pytorch/backends/cuda/op_backend.py	Makes FlashMLA metadata conditional for BF16 sparse decode and enforces int32 offsets.
lmdeploy/pytorch/backends/cuda/nsa.py	Adds sparse top‑k selector integration + fused Q/K-cache prep path.
lmdeploy/pytorch/backends/cuda/graph_runner.py	Includes skip_topk in CUDA graph key.
lmdeploy/pytorch/backends/cuda/attention/mla.py	Adds BF16 sparse decode path + multi-token decode index update support.
lmdeploy/pytorch/backends/cuda/apply_rotary_emb.py	Wires interleaved flag into CUDA rotary kernel.
lmdeploy/pytorch/backends/apply_rotary_emb.py	Updates builder interface to accept interleaved flag.
lmdeploy/messages.py	Adds indexer_topk + engine config flag support in shared message types.
lmdeploy/hf_configs/configuration_deepseek_v32.py	New HF config shim for DeepSeekV3.2.
lmdeploy/hf_configs/init.py	New centralized config fallback/registration module.
lmdeploy/cli/utils.py	Adds CLI flag for enabling indexer_topk output and clarifies FP8 policy for DSA models.
lmdeploy/cli/serve.py	Plumbs enable_return_indexer_topk into api_server config.
lmdeploy/archs.py	Uses centralized config_from_pretrained for model arch detection.
docker/prepare_3rdparty_wheel.sh	Adds building fast_hadamard_transform wheel.
benchmark/benchmark_chat_completion.py	Adds benchmark plumbing for return_indexer_topk and shared-store fetching.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 41 out of 41 changed files in this pull request and generated 2 comments.

[grimoire]

I think we can create a function call in models/utils/cudagraph.py and override it in modeling to add these "model sepcialized" key.

[grimoire]

why don't we get the function in init?

[grimoire]

This value won't change in single forward, right? should we place it in meta so we only need to compute it once?

[grimoire]

I have not run the real model yet. Would the stride change for different token numbers? tl.contexpr for frequent change value might leads to repeated compilation.

[grimoire]

This is true, but this won't bring much benefits, right?