Repository navigation
[Performance] Optimize GLM-5.3-Flash Triton KPool and gated norm - #17542
Conversation
github-actions
Bot
added
documentation
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces significant performance optimizations for the GLM-5.3-Flash Triton KPool and gated RMS norm implementations. By fusing operations, improving memory reuse, and correcting the head-wise scoring logic, the changes streamline the indexing and normalization paths. These improvements reduce overall model latency and improve throughput, particularly for large-batch scenarios, while maintaining correctness across various context lengths and cache configurations. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Ops][Feature] Optimize GLM5 Next KPool lightning indexer with fused Triton kernels and buffer reuseSuggested PR Summary:
### What this PR does / why we need it?
This PR optimizes the GLM5 Next KPool lightning indexer by fusing paged scoring, ReLU application, and token-index expansion into Triton kernels. It introduces direct writing to the SFA output buffer (`output_buffer`) and supports packing the causal tail (`pack_tail=True`) directly within the Triton expansion kernel. This eliminates the need for separate host-dispatched operations, manual tail scattering, padding-mask fills, and buffer copies. Additionally, it optimizes the KPool tail compression kernel by checking token validity early, and adds an option to skip saving statistics in the gated layer norm/RMS norm kernels when they are not needed.
### Does this PR introduce _any_ user-facing change?
No. These are internal performance optimizations and kernel refactorings for the GLM5-Next model backend on Ascend NPUs.
### How was this patch tested?
The changes were validated using the existing and updated end-to-end and unit tests:
- `tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_pool_key_indexer_triton.py` (updated to test `pack_tail` and strided output buffers)
- `tests/ut/models/test_glm5next_kpool_model_backend.py` (updated to verify mock arguments for `output_buffer` and `pack_tail`)
AuroraEmiya
left a comment
AuroraEmiya
left a comment
There was a problem hiding this comment.
One inline performance concern:
AuroraEmiya
left a comment
AuroraEmiya
left a comment
There was a problem hiding this comment.
One inline portability/validation concern:
|
/rerun Rerun (failed jobs only):
|
1 similar comment
|
/rerun Rerun (failed jobs only):
|
Skip incomplete or invalid pools before request lookup and compression. Preserve the separate tail update launch so history is read before it is overwritten. Size gated norm tiles to small batches and omit unused RMS inference rstd output. A3 validation passed 18 existing compressor cases and bitwise gated norm comparisons across model and shared API branches. Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com>
Restore per-head ReLU scoring with bounded scratch, query reuse and packed cache reads. Fuse pool expansion, causal tail, graph padding and writes into the SFA buffer. Keep context lengths dynamic in Triton and select A3 tile sizes from measured batch behavior. Validate against the independent CPU oracle, 27 existing operator cases, graph replay, nonfinite inputs, large cache addressing, cold-cache specialization checks and paired A3 microbenchmarks. Reuse existing tests without adding nightly cases. Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com>
Retain the existing FP32 head-weighted query and dot-product order while keeping fused expansion and paged-cache optimizations. Keep per-head ReLU scoring outside this performance change and use 128-pool tiles for the FP32 reduction. Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com>
Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com>
Use a stable paged-cache path for FULL graph capture and replay so padded token rows cannot imply query reuse. Preserve the existing packing heuristic after live-token trimming outside FULL graphs and cover the model dispatch contract. Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com>
Keep operator docs focused on behavior, interfaces and constraints. Remove validation sections and measurement-specific wording. Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com>
Bundle pinned baseline and candidate operators for A3/A5 checks. Compare only live indexer score rows and separately assert padding sentinels, with CPU regression coverage for the 49000-cell false mismatch. Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com>
06b816f to
81922df
Compare
Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com>
AuroraEmiya
left a comment
AuroraEmiya
left a comment
There was a problem hiding this comment.
LGTM after fix magic number : )
Replace masked request-end literals with the int32 maximum and explain their ordering requirement. Pass the same value as a Triton compile-time argument in tail compression and pool scoring. Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com>
…m-project#17542) ### What this PR does / why we need it? Reduce GLM-5.3-Flash Triton operator overhead: - Fuse KPool score sanitization and index expansion into the final SFA buffer; reuse queries and cached keys where applicable. - Keep FULL graphs on paged reads so padding cannot trigger cache packing. - Skip invalid and incomplete pools before compression; avoid unused gated-norm statistics and tune small batches. Refs vllm-project#15665 ### Does this PR introduce _any_ user-facing change? Lower inference latency with mainline scoring semantics preserved. No new configuration flags or CANN changes. ### How was this patch tested? - A3 at `4a73c393d`: 27 existing operator cases passed; 13 indexer shapes and sampled gated-norm branches matched mainline bitwise. - Full-model results on a compatible runtime overlay (W8A8, TP4/DP2/EP8, MTP off, 3500 input/128 output tokens): TPOT improved 7.81% at C1 and 7.08% at C8. Full GSM8K (8-shot, 10240-token limit): 1274/1319 baseline → 1280/1319 (97.04%). These measurements precede the FULL-graph dispatch fix. - Graph-dispatch fix: 12 CPU backend tests passed with unavailable vLLM/NPU imports stubbed. Full-model graph/performance rechecks remain pending. - A3 and A5 operator correctness checks passed (A5 result reported by the user). - vLLM main: vllm-project/vllm@ced6857 --------- Signed-off-by: Li Jiahang <216526138+lijiahang226@users.noreply.github.com>
合并 feat/glm5-kpool-ascendc-indexer 44 commits + 队友线 2 commits: - glm5_kpool_split AscendC 内核转正(split 达成: 64k TTFT -37.5% / 192k -59% vs triton) - glm5_kpool_indexer: dispatcher PR#17542 契约补全 + padded rows mask + absolute positions - 上游 backport: PR#17115/vllm-project#17613/vllm-project#17542/vllm-project#17524(P/D FullAttn TP parallel / mooncake) - Restore causal_conv1d_update_npu Triton kernel (dropped by vllm-project#14620) - memcache: filter cacheable KV groups + rewarm SSD on load - ascend_store: port vllm-project#16945 empty RoPE view regression tests 原始历史: backup/glm5.3-flash-0.29-pre-squash (bd6c295)
- glm5_kpool_split AscendC 内核(64k TTFT -37.5% / 192k -59% vs triton)+ dispatcher PR#17542 契约 - 上游 backport: PR#17115/vllm-project#17613/vllm-project#17542/vllm-project#17524 - causal_conv1d_update_npu Triton 恢复 + memcache rewarm + vllm-project#16945 回归测试 squash 自 47 commits; 原始历史: backup/glm5.3-flash-0.29-pre-squash
3 participants
What this PR does / why we need it?
Reduce GLM-5.3-Flash Triton operator overhead:
Refs #15665
Does this PR introduce any user-facing change?
Lower inference latency with mainline scoring semantics preserved. No new configuration flags or CANN changes.
How was this patch tested?
A3 at
4a73c393d: 27 existing operator cases passed; 13 indexer shapes and sampled gated-norm branches matched mainline bitwise.Full-model results on a compatible runtime overlay (W8A8, TP4/DP2/EP8, MTP off, 3500 input/128 output tokens): TPOT improved 7.81% at C1 and 7.08% at C8. Full GSM8K (8-shot, 10240-token limit): 1274/1319 baseline → 1280/1319 (97.04%). These measurements precede the FULL-graph dispatch fix.
Graph-dispatch fix: 12 CPU backend tests passed with unavailable vLLM/NPU imports stubbed. Full-model graph/performance rechecks remain pending.
A3 and A5 operator correctness checks passed (A5 result reported by the user).
vLLM main: vllm-project/vllm@ced6857