feat: support GLM-5.2#4737
Open
CUHKSZzxy wants to merge 33 commits into
Open
Conversation
# Conflicts: # docker/Dockerfile.jetson # lmdeploy/pytorch/engine/inputs_maker.py # lmdeploy/pytorch/paging/block_trie.py
# Conflicts: # lmdeploy/pytorch/engine/engine_instance.py
Contributor
There was a problem hiding this comment.
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.
CUHKSZzxy
force-pushed
the
feat/glm-5-2-support
branch
from
July 23, 2026 12:50
9c55e4e to
7196130
Compare
# Conflicts: # lmdeploy/hf_configs/__init__.py # lmdeploy/pytorch/backends/cuda/apply_rotary_emb.py # lmdeploy/pytorch/backends/default/apply_rotary_emb.py # lmdeploy/pytorch/config.py # lmdeploy/pytorch/kernels/cuda/apply_rotary_pos_emb.py # lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py # tests/pytorch/engine/test_cache_engine.py # tests/pytorch/kernel/test_apply_rotary.py # tests/pytorch/kernel/test_sparse_index_topk.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
LMDEPLOY_DISABLE_DSA_INDEXER_FUSION=1keeps the original path.Acc Benchmarks
Academic
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 of163,840tokens, 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
Acceptance of MTP
Aggregate results
Short aggregate per-position acceptance rates:
Long per-position acceptance rates:
Efficiency Benchmark
Throughput
Latency
DP8+EP8 config: max batch size 256, cache 0.6
Index Fusion Timeline
Full-model GLM-5.2 FP8, TP8 steady decode (representative rank-0
forward_cudagraph):The full model has 21 indexer layers, so the fusion removes exactly 13 CUDA-graph child kernels per layer (
21 × 13 = 273).Validation
Assistance
Assisted with Codex + GPT-5.6-Sol xHigh, reviewed manually