[None][fix] Warm up flashinfer sampling module during engine warmup - #17286
[None][fix] Warm up flashinfer sampling module during engine warmup#17286qiaoxj07 wants to merge 4 commits into
Conversation
flashinfer ships its sampling kernels as source and JIT-builds them with nvcc on first use. Today the only thing that exercises the non-greedy sampler during warmup is the advanced-sampling CUDA graph capture pass, so with cuda_graph_config=None nothing triggers that build until the first non-greedy request is served. The build then runs inside the executor loop and takes ~80-105s with a cold JIT cache (the cache lives under $HOME inside the container, so it is cold on every job). While one rank compiles it never launches its MoE all-to-all dispatch, so peer ranks GPU-spin on its completion flag and block in tp_gather. If the build crosses the 300s hang-detector threshold the detector fires and aborts the job. Observed on a 16-node GB300 disaggregated DeepSeek-V4 run: the hang detector caught a CTX rank with sample_from_logits_op -> top_k_mask_logits_op -> get_sampling_module -> build_and_load -> run_ninja on its stack, while its peers sat in tp_gather. Build the module during warmup instead, where the hang detector is not armed and no peer is waiting on this rank. get_sampling_module() covers every sampling kernel the sampler uses and is cached, so this is a no-op once built. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughChangesFlashInfer Sampling Warmup
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 1189-1192: Move the warmup_sampling_module() call in warmup()
above the non-HELIX context-parallel early-return guard, ensuring FlashInfer
initializes before any return path. Preserve the existing guard and the call’s
current warmup behavior for all engine configurations.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: df1f93d0-0f97-4508-b240-2653e7c4fd60
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py
|
/bot run --disable-fail-fast |
|
PR_Github #63957 [ run ] triggered by Bot. Commit: |
…turns warmup() returns early when no KV cache manager is allocated, and again for non-HELIX context-parallel engines. Both returns preceded the flashinfer warmup call, so those configurations still JIT-built the sampling kernels during serving -- the exact failure this change fixes. Neither guard has any bearing on whether the sampling kernels are needed, and the warmup reads no engine state, so hoist it to the top of warmup() where it runs on every path. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>
|
Addressed the CodeRabbit finding in 3546e88. The call is now hoisted to the top of Cluster validationRunning now on a 4-node GB300 disagg config that reproduces the preconditions (ctx Early signal is what the fix predicts: the For reference, this is the stack the hang detector caught on an unpatched run, which is what motivated the change: |
|
/bot run --disable-fail-fast |
|
PR_Github #63960 [ run ] triggered by Bot. Commit: |
|
PR_Github #63957 [ run ] completed with state |
|
PR_Github #63961 [ run ] triggered by Bot. Commit: |
|
PR_Github #63960 [ run ] completed with state |
Cluster validation completeTwo 4-node GB300 disagg runs with this patch bind-mounted over the container's modules, on a config that reproduces the preconditions (ctx Mid-serving stalls
~156s of multi-minute mid-serving stalls collapse to ~20s, and both patched runs served 2250+ iterations with zero hang-detector events. The unpatched run died at iteration 9. The torch.compile ruled outPatched run B also carried CaveatThis is not a controlled A/B. The validation runs are 4-node |
The prewarm runs ahead of every guard in warmup(), so a broken JIT toolchain would abort startup even for greedy-only deployments that never reach these kernels. Log and continue instead, matching the neighbouring pre-JIT warmups; deployments that do sample non-greedily still surface the same error at the real call site. The catch is broad because the build can fail as a subprocess error, a missing-cache error, or a shared-object load error. The narrow catch used by the DeepGEMM prewarm guards against signature drift, which does not apply here -- get_sampling_module() takes no arguments. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>
|
Good catch — applied in 25611d3. You're right that hoisting the call above the guards changed the blast radius: before, a broken JIT toolchain only surfaced at the first non-greedy request (and never at all for greedy-only deployments); after, it aborted startup for everyone. The prewarm is an optimization, so it should never be able to do that. Now wrapped in On the exception breadth: I went with a broad catch rather than matching the DeepGEMM prewarm's narrow |
Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #63998 [ run ] triggered by Bot. Commit: |
|
PR_Github #63961 [ run ] completed with state |
|
PR_Github #63998 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64053 [ run ] triggered by Bot. Commit: |
|
PR_Github #64053 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64075 [ run ] triggered by Bot. Commit: |
|
PR_Github #64075 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64118 [ run ] triggered by Bot. Commit: |
5 participants
Description
flashinfer ships its sampling kernels as source and JIT-builds them with
nvccon first use. The only thing that exercises the non-greedy sampler during warmup today is the advanced-sampling CUDA graph capture pass (_run_capture_pass(force_non_greedy=True)), so withcuda_graph_config: nullnothing triggers that build until the first non-greedy request is actually served.The build then runs inside the executor loop and takes ~80–105s with a cold JIT cache. The cache lives under
$HOMEinside the container (/root/.cache/flashinfer/...), so it is cold on every job. While one rank compiles, it never launches its MoE all-to-all dispatch, so peer ranks GPU-spin on its completion flag and block intp_gather. If the build crosses the 300s hang-detector threshold, the detector fires and aborts the whole job.This was observed on a 16-node GB300 disaggregated DeepSeek-V4 run. The hang detector caught a CTX rank with this stack:
while its peers sat in
_flush_pending_transfer_responses -> _enqueue_responses -> tp_gather.Generation workers were unaffected in the same runs because they set
cuda_graph_config.batch_sizes, which makes the advanced-sampling capture pass run and build the module during warmup.Fix
Build the module during engine warmup instead, where the hang detector is not armed and no peer is waiting on this rank.
get_sampling_module()covers every sampling kernel the sampler uses and is@functools.cached upstream, so this is a no-op once built.Not gated on spec-decoding: the plain
TorchSamplernon-greedy path reaches the same kernels, so gating would leave that case exposed. Cost is one build at startup that would otherwise be paid on the first non-greedy request anyway.Test Coverage
Verified the no-flashinfer branch returns cleanly before any
flashinfer.*attribute access (the module binds a placeholder that raisesImportErroron attribute access when flashinfer is absent).Verified idempotency: repeated calls build once.
Cluster validation on two 4-node GB300 disagg runs (ctx
cuda_graph_config: null+ MTP), patch bind-mounted over the container's modules:Hang detectedeventsMPI_AbortThe
ninja: Entering .../cached_ops/samplingblock now appears during warmup, ahead of the first warmup forward, and never during serving. Patched run B additionally carriedTORCH_LOGS=recompiles: 111 recompiles were logged, all in the DeepSeek-V4 sparse-attention helpers, and none inside any stall window — so Dynamo/Inductor is not a contributor.Not a controlled A/B: validation runs are 4-node
ctx1 + gen1@ concurrency 224 vs. 16-nodectx4 + gen1@ 1418 for the failures. See the validation comment for the full breakdown.PR Checklist
[JIRA/NVBUG/None][type] descriptionDev Engineer Review
warmup_sampling_module()to eagerly initialize cached FlashInfer sampling kernels.PyTorchModelEngine.warmup(), before early-return guards and other warmup stages.QA Engineer Review
No test changes.