Repository navigation
[Model][GLM-5.3-Flash] Warm up the DCP top-k merge kernels for the kpool indexer - #60032
Merged
NickLucche merged 4 commits intoOct 8, 2026
Merged
Conversation
Contributor
CI selector (shadow): 5 test steps (25 jobs) instead of 34 (47 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (5)
Would skip (today's rules run them) (32)
Would add (today's rules do not run them) (3)
AMD mirrors: would skip (32)
AMD mirrors: would add (0)none 1 changed files · base |
…ool indexer Under DCP the kpool indexer merges per-rank top-k with the CuteDSL/Triton DCP merge kernels, but only SparseAttnIndexer registered them for JIT warmup, so GLM-5.x compiled PackDCPTopkCandidatesKernel during inference. Register them from SparseAttnIndexerKpool as well. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
NickLucche
force-pushed
the
glm-flash-kpool-dcp-merge-warmup
branch
from
October 5, 2026 14:45
f75c346 to
984782e
Compare
Cover the warmup registration added to SparseAttnIndexerKpool.__init__: both DCP top-k merge kernels are registered under DCP with JIT warmup enabled, and neither is registered without DCP or with warmup disabled. Generated with Mistral Vibe. Co-Authored-By: Mistral Vibe <vibe@mistral.ai> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
NickLucche
marked this pull request as ready for review
October 5, 2026 17:50
NickLucche
requested review from
AndreasKaratzas,
DarkLight1337 and
ywang96
as code owners
October 5, 2026 17:50
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Member
Author
LopezCastroRoberto
requested changes
Oct 6, 2026
LopezCastroRoberto
left a comment
LopezCastroRoberto
left a comment
Contributor
There was a problem hiding this comment.
Thanks for catching this. I think we can drop the test and simplify the condition to current_platform.is_cuda() and has_cutedsl(). The registry already handles enable_jit_warmup, and DCP is guaranteed by the surrounding branch.
Address review feedback: drop the unit test and reduce the warmup condition to `current_platform.is_cuda() and has_cutedsl()`. The JIT warmup registry only compiles registered kernels when `enable_jit_warmup` is on (kernel_warmup gates registry.warmup()), and DCP > 1 is guaranteed by the surrounding branch, so the extra checks are redundant. Generated with Mistral Vibe. Co-Authored-By: Mistral Vibe <vibe@mistral.ai> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Member
Author
|
Thanks for reviewing @LopezCastroRoberto , addressed! |
LopezCastroRoberto
approved these changes
Oct 8, 2026
LopezCastroRoberto
left a comment
LopezCastroRoberto
left a comment
Contributor
There was a problem hiding this comment.
Gotcha! Sorry for the block @NickLucche. LGTM, thanks!
Contributor
|
/ci run |
|
✅ Triggered Buildkite CI #93544 for commit |
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
3 participants
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.
Purpose
Follow-up to #59211 (now merged; this branch is rebased on main).
Under DCP, the GLM-5.x kpool indexer merges the per-rank top-k with
_merge_dcp_topk_global, which runsPackDCPTopkCandidatesKernelandStableTopKFromGatheredCandidatesKernel. OnlySparseAttnIndexer.__init__registers these kernels for JIT warmup, so with the kpool indexer they were compiled on the first long-context request:This PR registers them from
SparseAttnIndexerKpool.__init__under the same conditions: DCP > 1, JIT warmup enabled, CUDA, and CuteDSL available. #59211 already makes the warmup keys use pool-level top-k and interleave for kpool models.Test Plan
zai-org/GLM-5.3-Flashon 8×H100 withvllm serve, bf16 KV cache, no MTP. TP8 + DCP8 with--cp-kv-cache-interleave-size 4, compared with TP8. GSM8K (1319 questions) at 5, 25, 50 and 100 shots, i.e. prompts from about 650 to about 17k tokens. Checked the server log for JIT monitor warnings.Test Result
PackDCPTopkCandidatesKernelno longer compiles during inference under DCP8. Without this change, it compiled during inference in the same setup. The other JIT warnings in the log (fused_moe_kernel, kpool compress kernels) also appear without DCP and are unrelated.The E2E runs used the same local SM90 FlashMLA bf16 DCP change as #59211, which Hopper needs for DCP at head_dim 512.
Duplicate check
I searched open PRs for kpool /
SparseAttnIndexerKpool/PackDCPTopkCandidateswarmup and found none. #54345 deals with prefill metadata JIT, not the DCP top-k merge kernels.AI assistance
This PR was written with AI assistance (Claude). I reviewed the change and ran the tests above.
🤖 Generated with Claude Code