perf(glm5next): select KDA backend per server - #495
Conversation
GLM KDA decode retains the shape plan and creates a fresh B12X binding for each call from vLLM workspace scratch. This prevents a binding from outliving recyclable scratch or KV-cache state while leaving non-B12X and prefill paths unchanged. Validation: tests/models/test_glm5next_model.py (26 passed); ruff check and format check. Assisted-by: OpenAI Codex <codex@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org>
Use the packed Triton KDA path for plain GLM-5.3 decode and the B12X KDA path for speculative decode by default. Plain decode avoids B12X metadata staging and validation overhead, while speculative decode retains B12X checkpoint selection and rollback handling. Expose glm53_kda_decode_backend=auto|b12x|triton through additional_config. The default auto mode preserves B12X initialization for speculative requests; explicit values remain available for qualification and diagnosis. Unknown values fail during model construction. Validation: six targeted GLM model tests; Ruff check and format; direct BF16 kernel comparison against the B12X reference for batch 1 over 64 sequential steps and batch 16; TP4 W4A4 decode benchmarks on physical GPUs 4-7. Signed-off-by: Martin Vit <martin@voipmonitor.org>
Select one KDA implementation for the lifetime of every recurrent cache. Auto mode uses Triton when speculative decoding is disabled and B12X when an MTP or DFlash configuration is present. Explicit b12x and triton selections remain unchanged. Both implementations satisfy the BF16 operation reference, but switching their rounding behavior between plain and speculative steps changes subsequent draft acceptance. A server-level choice preserves stable recurrent-state behavior and retains the qualified speculative rollback path. Compatibility: the additional_config interface remains glm53_kda_decode_backend=auto|b12x|triton. Validation: nine targeted GLM model tests, Ruff check and format, and git diff validation. Signed-off-by: Martin Vit <martin@voipmonitor.org>
|
Important Review skippedNo new commits to review since the last review. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthroughGLM-5.3 KDA now validates backend configuration and selects B12X decoding for supported cases. B12X decoding stores a reusable plan, acquires workspace scratch buffers per call, supports direct tensor binding, and validates token capacity. ChangesGLM-5.3 KDA decode
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: 🟡 Moderate · up to The change selects a server-stable KDA backend and uses per-invocation scratch, but some mutable execution buffers remain shared and rollback after interrupted or failed recurrent-state updates is not established. Overlapping requests or failures could mix request data or corrupt later decoding, so the PR is not merge-ready until those execution and recovery guarantees are explicitly ensured or accepted. Sequence Diagram(s)sequenceDiagram
participant Attention as Glm5NextLinearAttention
participant Workspace as current_workspace_manager
participant API as B12X KDA API
Attention->>Workspace: get_simultaneous(plan.shapes_and_dtypes())
Workspace-->>Attention: scratch buffers
Attention->>API: bind_kda(plan, caller tensors or scratch buffers)
API-->>Attention: binding
Attention->>API: run_kda(binding)
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Use the existing B12X binding contract for speculative verified-prefix tensors while preserving the established staged path for ordinary decode. Supply live token count from an independent caller-owned scalar so B12X can validate packed metadata end to end. Co-authored-by: OpenAI Codex <codex@openai.com>
|
@coderabbitai review |
|
|
Closed as not planned. The qualified GLM-5.3-Flash runtime keeps target KDA ownership in B12X instead of selecting Triton KDA for non-speculative serving. The B12X-only policy avoids maintaining two target-state implementations with different checkpoint and rollback contracts. This pull request is not part of the qualified merge checklist in issue #590. |
2 participants
Resulting behavior
Status: implemented and qualified for GLM-5.3 tensor-parallel serving at tensor parallel size 4.
vLLM creates a B12X KDA binding from caller-owned scratch for each invocation. No binding outlives recyclable workspace scratch or scheduler-owned recurrent state.
GLM-5.3 selects one KDA implementation for the lifetime of a server:
glm53_kda_decode_backend=triton|b12xprovides an explicit qualification and diagnostic override.Keeping one implementation for the recurrent cache lifetime prevents differences in BF16 operation order from changing later speculative acceptance.
Source contract
local-inference-lab/vllm:dev/jovian-judgementatc79f35ca00e8e93e0943a0d79b85b22b18aac939.da60b74f2a6aadbb0dcb53a97590b159fae96431.plan.bind(scratch=...)on every invocation. vLLM does not cache a B12X workspace, arena, or binding.Validation
tests/models/test_glm5next_model.py -k 'kda or b12x': 15 passed, 24 deselected.git diff --check: passed.Duplicate-work check
local-inference-lab/vllm#490integrates the proposed publicrun_kda_single_tokenAPI fromlocal-inference-lab/b12x#249for ordinary non-speculative decode. This pull request does not add that API. It provides caller-scratch binding lifetime and server-stable backend selection, including speculative transaction semantics. The direct B12X single-token optimization can be implemented later as an internal B12X route without changing this server-lifetime contract.Open upstream vLLM pull requests were searched for GLM KDA, B12X, and Triton backend selection. No open upstream pull request implements this contract.
Review disclosure
OpenAI Codex assisted with implementation, kernel comparison, tests, profiling, benchmarking, and pull-request preparation. Human review of every changed line and the recurrent-state contract is required before merge.
Summary by CodeRabbit