Skip to content

[https://nvbugs/6275959][fix] In runJITWarmupGridIfRequested, std::min the warmup-only maxBatchSize against… #15221

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tensorrt-cicd wants to merge 2 commits into
base: main
Choose a base branch
from
+16 −1
Open

[https://nvbugs/6275959][fix] In runJITWarmupGridIfRequested, std::min the warmup-only maxBatchSize against… #15221

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9c89928
[nvbugs/6275959][fix] Skip TRTLLM-Gen FMHA JIT warmup for oversized e…
tensorrt-cicd Jun 17, 2026
5d3ccfe
[nvbugs/6275959][chore] Remove stale waiver after fix
trtllm-agent Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions tensorrt_llm/_torch/pyexecutor/model_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1681,6 +1681,22 @@ def _run_attention_warmup(self,
if not issubclass(self.attn_backend.Metadata, TrtllmAttentionMetadata):
return

# The C++ TRTLLM-Gen FMHA JIT warmup enumerates a (batchSize x seqLenKv) cartesian
# grid sized by engine maxima. PR #15305 densified the candidate lists, so for
# long-context configs (e.g., max_num_requests=2048, max_seq_len=131072 in
# disagg_config_ctxtp2_gentp2_llama31_8b_ucx.yaml), the grid produces thousands of
# NVRTC compilations that exceed the 600s server-start timeout. Skip the warmup
# whenever the product would blow the budget; any kernel not pre-warmed JIT-compiles
# lazily on first request, which is correct (just slower for that one request).
# The threshold matches the pre-PR #15305 effective grid size.
max_warmup_workload = self.batch_size * self.max_seq_len
if max_warmup_workload > 256 * 16384:
logger.info(
f"Skipping TRTLLM-Gen FMHA JIT warmup: engine config "
f"(max_batch_size={self.batch_size}, max_seq_len={self.max_seq_len}) "
f"would produce too many warmup grid points")
return
Comment on lines +1684 to +1698

@coderabbitai coderabbitai Bot Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift

Replace the all-or-nothing skip with capped FMHA warmup.

This branch returns before trtllm_gen_fmha_jit_warmup() runs. It leaves the TRTLLM-Gen FMHA grid unwarmed and moves NVRTC compilation into live request handling. That conflicts with the PR objective to clamp warmup dimensions to batch size 256 and sequence length 16384 while preserving runtime maxima.

Pass capped dimensions to the C++ warmup-grid path instead of skipping the entire warmup. Keep self.batch_size and self.max_seq_len unchanged for serving. Derive the limit from the actual candidate-grid budget and replace the inline 256 * 16384 values with named constants.

🤖 Prompt for 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.

In `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 1491 - 1505,
Replace the early return in the FMHA warmup flow with capped dimensions for the
C++ warmup-grid invocation, limiting batch size to 256 and sequence length to
16384 while preserving self.batch_size and self.max_seq_len for serving. Define
named constants for these caps and derive the candidate-grid workload limit from
their product, then pass the capped values to trtllm_gen_fmha_jit_warmup().


@contextlib.contextmanager
def trtllm_gen_fmha_jit_warmup():
previous = self._trtllm_gen_jit_warmup
Expand Down
1 change: 0 additions & 1 deletion tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,6 @@ full:H100/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mt
full:H100/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v2_kv_cache-True-True-cutlass-auto] SKIP (https://nvbugs/6529792)
full:H100/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.9-fp8kv=False] SKIP (https://nvbugs/6523809)
full:H100/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] SKIP (https://nvbugs/6273850)
full:H100/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959)
full:H100/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_32b_fp8_stress] SKIP (https://nvbugs/6312828)
full:H100_PCIe/unittest/llmapi/test_llm_pytorch.py::test_llama_7b_multi_lora_evict_and_reload_lora_gpu_cache SKIP (https://nvbugs/5682551)
full:H20/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=False] SKIP (https://nvbugs/6345827)
Expand Down
Loading