Repository navigation
[CPU] Add KDA backend for GLM5Next - #57496
kunkunblueberry wants to merge 6 commits into
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
e69a889 to
953d226
Compare
8a07bf8 to
27b821f
Compare
|
Native recurrent milestone for #57346:
This is the first native step; packed/speculative native state routing and chunk-specific native API remain follow-up work. |
|
Validation update: real _C_AVX512 build loaded torch.ops._C.glm5next_kda_recurrent; native parity/speculative suite passed 21 tests. CPU binding smoke reports CpuPlatform, device_type=cpu, dispatch_key=CPU, is_cpu=True, and both common KDA symbols resolve to vllm.models.glm5next.cpu.kda. A production-shaped Glm5NextLinearAttention._forward decode smoke invoked vllm.models.glm5next.cpu.kda.fused_recurrent_kda and produced finite output. Focused CPU backend suite passed 19 tests with 2 expected skips under the server's torchcodec shim. Local diff-check and clang-format checks passed. |
c84b20b to
1867161
Compare
|
The branch history was rebuilt from current upstream/main into one signed-off feature commit (186716118a) to resolve the DCO failure from the original unsigned planning commit. The old local history is retained only as cpu-kda-glm5next-old/cpu-kda-glm5next-history; the fork branch was force-updated with --force-with-lease. |
|
Additional runtime fix: CPU layer construction now skips the CUDA-only Triton scatter warmup registration, which would otherwise fail before forward. The real _C_AVX512 focused suite remains green at 26 passed / 14 warnings after this change. Packed-prefill and prefill-to-decode layer tests exercise CPU causal convolution, CPU state helpers, native KDA, and cache handoff. CPU speculative model execution remains explicitly fail-closed because the existing CPU causal-conv primitive does not implement accepted-token rollback. |
Add the CPU GLM5Next KDA backend, AVX-512 recurrent operator registration, common-layer dispatch, CPU state helpers, and correctness/runtime coverage for recurrent, packed prefill, state slots, and prefill-to-decode execution. Refs vllm-project#57346 Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: kunkunblueberry <1833921874@qq.com>
The CPU state scatter helper is eager PyTorch and does not expose Triton warmup registration. Avoid invoking that CUDA-only setup during CPU layer construction. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: kunkunblueberry <1833921874@qq.com>
Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: kunkunblueberry <1833921874@qq.com>
Use a concrete sequence-bounds callable type, make validated optional values explicit to mypy, and mark mutually exclusive platform imports as intentional rebinding. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: kunkunblueberry <1833921874@qq.com>
6d8ac72 to
df95cdb
Compare
Self-review / Implementation notesThis review explains the design decisions and validation boundary for this PR. Scope and designThe initial audit showed that CPU support required more than changing the two KDA imports. The common layer also reaches causal-convolution and state gather/scatter boundaries that have GPU-oriented implementations. The change is therefore kept at the GLM5Next KDA execution boundary while preserving the existing projection, metadata, and routing code. I intentionally did not add Semantics preservedThe CPU backend preserves:
Reuse and native boundaryOrdinary CPU convolution reuses the existing CPU causal-convolution implementation. GLM5Next speculative decoding has an additional accepted-token rolling-window contract, so only that feature-specific wrapper is added here; generic causal-convolution APIs are unchanged. The Python backend is the explicit semantic fallback. When the supported dense-state conditions are met and ValidationThe independent reference implementation remains in the test module and shares no backend helper code. Coverage includes recurrent parity, FP32/BF16 native parity, packed prefill, sequence isolation, shuffled and NULL slots, prefill-to-decode handoff, speculative routing, CPU convolution rebinding, rolling convolution, and production-shaped With the real Recorded native parity errors were approximately Remaining boundaryThis PR makes no performance claim. AMX, additional precision-specific kernels, chunk-level parallelization, NUMA/batch benchmarks, and FP16/FP8 kernels are follow-ups. Full GLM-5.3-Flash checkpoint generation still depends on #57345 and a real checkpoint validation run. The parts I would particularly appreciate review on are the CPU boundary in This is an AI-assisted contribution; the human submitter will review every changed line and defend the design and validation results. |
Summary
This PR adds a correctness-first CPU execution path for GLM-5.3-Flash / GLM5Next KDA for #57346.
The path is deliberately feature-local:
It does not implement #57345 (CPU sparse MLA / KeyPool), Kimi K3 support, a generic GDN refactor, or a new GLM CPU model hierarchy. The two GLM5Next CPU PRs are intentionally separate so that KDA and sparse MLA can be reviewed and tested independently.
Self-review / Implementation notes
Scope and design
The initial source audit showed that CPU support required more than changing the two KDA imports. The common layer also reaches causal-convolution and state gather/scatter boundaries that have GPU-oriented implementations. I therefore kept the change at the GLM5Next KDA execution boundary and preserved the existing common projection, metadata, and routing code.
I did not add
glm5next/cpu/model.pyor a new backend registry. The hardware-specific difference is limited to KDA execution, so a model-level fork would add review and maintenance cost without solving a concrete blocker.The related CPU GDN implementation was used as a reference for stateful CPU integration, but it was not treated as a drop-in KDA implementation. KDA differs in safe-gate semantics, beta handling, V-major state layout, and speculative state routing.
KDA semantics preserved
The CPU backend keeps the existing GLM KDA contract:
A_logand per-head/key-dimensiong_bias;[slots, heads, V, K];chunk_kda_with_fused_gatecurrently evaluates each packed sequence as a token-wise recurrence. This is mathematically equivalent to the chunk state transition; chunk-level parallelization is an optimization and remains outside this correctness-first change.Reuse and native boundary
Ordinary CPU convolution continues to use the existing CPU causal-convolution implementation. GLM5Next speculative decoding has an additional accepted-token rolling-window contract, so only that feature-specific wrapper is implemented here; generic causal-convolution APIs are unchanged.
The Python backend is the explicit fallback and semantic contract. When the supported dense-state conditions are met and
torch.ops._C.glm5next_kda_recurrentis registered, the recurrent path can use the native AVX-512 operator. Unsupported configurations do not silently route to Triton or NVIDIA KDA.Validation
The independent reference implementation remains in the test module and shares no backend helper code. The focused suite covers:
Glm5NextLinearAttentionprefill, decode, and speculative paths.With the real
_C_AVX512extension loaded on the CPU server, the previously validated result is:The native parity checks recorded approximately
2.98e-7FP32 output error,5.96e-8FP32 final-state error,1.35e-3BF16 output error, and1.19e-7BF16 final-state error. Targeted Ruff, Ruff format, Typos, SPDX, lazy-import, forbidden-import, torch.cuda, and C++ clang-format hooks pass on the final files; manual Python 3.12 mypy also passes.The final detached verification worktree without the prebuilt CPU extension ran 23 tests, with two failures at the unregistered causal-convolution operator/native spy boundary. Those failures are environment/build limitations of that worktree, not a replacement for the extension-backed result above.
Deliberate non-goals and remaining boundary
This PR makes no performance claim. AMX, additional precision-specific vector kernels, chunk-level parallelization, NUMA/batch benchmarks, and FP16/FP8 kernel work are follow-ups. Full GLM-5.3-Flash checkpoint generation still depends on #57345 and a real checkpoint validation run.
Please review in particular the CPU boundary in
common/kda.py, state-slot semantics, the native operator guard/fallback, and whether any existing common helper should be reused without broadening this PR.Duplicate-work check
The issue and open-PR searches found no separate open implementation of #57346 or a GLM5Next CPU KDA backend. #57687 is the separate sparse MLA / KeyPool effort.
This is an AI-assisted contribution. The human submitter will review every changed line and defend the design and validation results.
Refs #57346