[https://nvbugs/6432948][fix] Add per-engine flag _trtllm_gen_jit_warmup_done; skip… - #16205
[https://nvbugs/6432948][fix] Add per-engine flag _trtllm_gen_jit_warmup_done; skip…#16205trtllm-agent wants to merge 3 commits into
_trtllm_gen_jit_warmup_done; skip…#16205Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughFP8 block-scale MoE tactic selection now excludes undersized tiles for fused and unfused workloads. Fallback selection validates configurations, skips undersized tiles, prioritizes the heuristic tile, and errors when no valid configuration exists. ChangesFP8 MoE tactic selection
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
98dcf31 to
39b1dd5
Compare
| perf/test_perf.py::test_perf[whisper_large_v3-bench-float16-input_output_len:128,20] SKIP | ||
| perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp4_v2_grace_blackwell-r1_fp4_v2_dep4_mtp1_1k8k] SKIP (https://nvbugs/6422339) | ||
| perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp8_blackwell-r1_fp8_tp8_mtp3_1k1k] SKIP (https://nvbugs/6445332) | ||
| perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp8_blackwell-r1_fp8_tp8_mtp3_8k1k] SKIP (https://nvbugs/6432948) |
There was a problem hiding this comment.
The code change looks like a good improvement, but I'm not seeing how it addresses the root cause of the flaky failure. The error was cudaStreamDestroy(stream): an illegal memory access was encountered, not an out of memory issue.
If you want to move forward with this as a fix, please run the test several times to ensure it passes consistently.
39b1dd5 to
2995242
Compare
2995242 to
4cb6611
Compare
4cb6611 to
acb9a9d
Compare
f6ceed3 to
f00fd2b
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 `@cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp`:
- Line 423: Update the tile-selection flow around computeSelectedTileN to first
filter the supported tile list to values at least fusedMinTileN, pass that
eligible list into the heuristic, and explicitly reject fusedMinTileN values
above the largest supported tile. Ensure the later tileN threshold check only
evaluates tiles from the filtered list, while preserving valid selection of tile
64 and larger.
- Around line 482-483: Add braces around the body of the conditional in the
tileN candidate logic, preserving the existing push_back behavior when t !=
tileN.
- Around line 485-489: Replace the independent TLLM_MOE_FUSED_MIN_TILEN parsing
in both selection paths, including fusedMinTileNFallback, with one shared C++17
std::from_chars-based parser. Parse the environment variable once, reject
malformed, non-positive, and otherwise unsupported thresholds, and use the
validated fallback consistently so tileN 8/16 cubins cannot be selected by
invalid input.
🪄 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: 4f9de668-f431-4985-88fc-dab84dc98e63
📒 Files selected for processing (1)
cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp
| static int const fusedMinTileNFallback = []() | ||
| { | ||
| char const* env = std::getenv("TLLM_MOE_FUSED_MIN_TILEN"); | ||
| return env != nullptr ? std::atoi(env) : 32; | ||
| }(); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Verify the declared C++ language standard before using std::from_chars.
rg -n -C 2 'CXX_STANDARD|cxx_std_|CMAKE_CXX_STANDARD' \
-g 'CMakeLists.txt' -g '*.cmake' -g 'meson.build' -g 'BUILD' -g 'BUILD.bazel' .
# Trace all threshold parsing sites that must use the same checked accessor.
rg -n -C 4 'TLLM_MOE_FUSED_MIN_TILEN|std::atoi' cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cppRepository: NVIDIA/TensorRT-LLM
Length of output: 6687
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- target source ---'
sed -n '380,435p' cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp
sed -n '455,515p' cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp
printf '%s\n' '--- declarations and includes relevant to parsing/selection ---'
sed -n '1,90p' cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp
rg -n -C 3 'computeSelectedTileN|fusedMinTileN|TLLM_MOE_FUSED_MIN_TILEN|atoi|from_chars|charconv' \
cpp/tensorrt_llm/thop cpp/tests
printf '%s\n' '--- repository guidance ---'
if [ -f CODING_GUIDELINES.md ]; then
sed -n '1,220p' CODING_GUIDELINES.md
fiRepository: NVIDIA/TensorRT-LLM
Length of output: 30379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ctypes
# Model the exact conversion used by std::atoi for representative environment
# values, without executing repository code.
samples = ["", "abc", "32abc", "-1", "0", "8", "16", "32", "64", "999999999999999999999"]
for value in samples:
encoded = value.encode()
result = ctypes.CDLL(None).atoi(encoded)
print(f"{value!r} -> {result}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 282
Parse TLLM_MOE_FUSED_MIN_TILEN once and reject invalid values.
std::atoi returns 0 for malformed input such as abc. This disables the small-tile exclusion and can select the known unsafe tileN 8/16 cubins. Both selection paths parse the variable independently. Use one shared checked parser and reject malformed or unsupported thresholds. std::from_chars is available under the project’s C++17 standard.
🤖 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 `@cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp` around lines 485 - 489, Replace
the independent TLLM_MOE_FUSED_MIN_TILEN parsing in both selection paths,
including fusedMinTileNFallback, with one shared C++17 std::from_chars-based
parser. Parse the environment variable once, reject malformed, non-positive, and
otherwise unsupported thresholds, and use the validated fallback consistently so
tileN 8/16 cubins cannot be selected by invalid input.
Source: Coding guidelines
0e4f44b to
7c525aa
Compare
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The widening itself looks right: the fault is in the shared small-tile cubins, and gating on num_fused_shared_experts > 0 meant the default (unfused) configuration never got the protection.
Two things about the PR as posted:
The title and description describe a different change. They're about _trtllm_gen_jit_warmup_done and skipping the Phase-2 FMHA JIT warmup; the diff is the MoE tactic-selection exclusion. The commit message on the head commit is accurate and detailed — please copy that into the PR body so reviewers and git log archaeologists see the same story.
Perf validation. Excluding tileN 8/16 for every caller means small-batch decode (1-8 tokens/expert) now runs on tileN 32 instead of a tile matched to the batch. The commit message documents crash-freedom across shapes but no throughput/latency numbers. Since this affects the default DeepSeek-R1 FP8 path, a decode-side ITL comparison before/after would be worth having on record — if it's a real cost, that strengthens the case for chasing the kernel-side fix rather than living with the WAR.
Please document TLLM_MOE_FUSED_MIN_TILEN before this merges. It is user-visible, it now changes kernel selection for every FP8 block-scale MoE caller rather than just the fused path, and it is currently documented nowhere outside this source file. Add it to wherever the debug/tuning env vars are listed, including the default (32), what 0 does, and a note that it exists only as a WAR for the small-tile IMA. An undocumented global knob that silently changes tactic selection is how the next debugging session loses a day.
| if (t != tileN) | ||
| tileN_candidates.push_back(t); | ||
| } | ||
| static int const fusedMinTileNFallback = []() |
There was a problem hiding this comment.
The threshold is now read twice — once here and once at line 414 — as two independent function-local statics parsing the same env var. They agree today, but nothing enforces that: a future change to the bound or the variable name in one place leaves the other site admitting tileN 8/16, and this fallback is the path warmup takes, so the regression would show as a rare illegal memory access rather than a test failure.
Hoist it into a single file-scope accessor (e.g. static int moeMinTileN() with the cached lambda inside) and call it from both sites.
| // mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty | ||
| // passing-config list, so the exclusion has to happen at tactic-selection time. | ||
| // TLLM_MOE_FUSED_MIN_TILEN overrides the threshold (0 disables) for A/B experiments. | ||
| static int const fusedMinTileN = []() |
There was a problem hiding this comment.
TLLM_MOE_FUSED_MIN_TILEN / fusedMinTileN no longer have anything to do with the fused path — the exclusion is now global. The name will mislead the next person debugging a tile-selection issue on an unfused model into thinking the knob doesn't apply to them.
Rename the locals to minTileN; for the env var, either rename it (and keep the old name as a deprecated alias if anyone is using it in experiments) or add a comment stating the name is historical.
| tileN = std::clamp(nextPowerOfTwo(avg_tokens_per_expert), mSupportedTileN.front(), mSupportedTileN.back()); | ||
|
|
||
| if (num_fused_shared_experts.value_or(0) > 0) | ||
| // getDefaultValidConfigIndex only pairs the per-GEMM "default" indices without |
There was a problem hiding this comment.
This lands a second, unrelated behavior change: every non-fused call previously went through getDefaultValidConfigIndex and now takes getValidConfigIndices(...).front(). That's a change in the selected kernel config for all default-path (non-autotuned) FP8 block-scale MoE traffic, not just an exclusion of the faulty tiles.
The comment's justification also doesn't match Runner::getDefaultValidConfigIndex (runner.cu:781): indexGemm1/indexGemm2 each come from the respective GEMM's own valid-index list for this problem size, and the pair must be present in mPassingConfigs or it throws — so the returned pair does appear to be validated against the problem size. If there's a case where it isn't, please state it concretely; otherwise scope this PR to the tileN exclusion and keep getDefaultValidConfigIndex for tileN >= 32.
| config = mRunners.at(tileN)->getDefaultValidConfigIndex( | ||
| total_experts_per_token, hidden_size, intermediate_size, num_total_local_experts, num_tokens); | ||
| } | ||
| TLLM_CHECK_WITH_INFO(config != -1, "No valid TRTLLM-Gen config found for FP8 block-scale MoE."); |
There was a problem hiding this comment.
Non-fused callers could not previously fail here — getDefaultValidConfigIndex either returned a config or threw with its own message. Now a shape with no validated config for any tileN >= 32 aborts the run. The commit message says the tactic list is never emptied for the models checked (DeepSeek-R1 EP 1/4/8, Qwen3-235B/30B), but that's an empirical claim over a sample.
Include the problem dimensions in the message (num_tokens, hidden_size, intermediate_size, experts_per_token, local_experts, and the min-tileN in force) so a report from an unchecked model is actionable without a repro.
871454d to
01878c9
Compare
…ase-2 warmup create_py_executor's two-phase KV-cache-estimation flow instantiates PyExecutor twice against the same model_engine, so warmup - and therefore _run_attention_warmup - runs twice. The TRTLLM-Gen FMHA JIT kernel cache is process-global, so the second grid enumeration compiles nothing new. For DeepSeek-R1 FP8 TP=8 MTP3 with max_num_tokens=12288 on B200, running the grid a second time - after autotuner exploration and CUDA-graph capture have consumed most of GPU memory - occasionally triggers an illegal memory access asynchronously reported at torch.cuda.synchronize() in _run_attention_warmup. Add a per-engine _trtllm_gen_jit_warmup_done flag: first call runs the full grid and sets the flag; subsequent calls short-circuit with a logger.info. Correctness is preserved because the JIT cache from Phase-1 already covers every shape Phase-2 could request; any kernel not yet compiled would JIT-compile lazily on first request anyway. Verified: 1 passed in 230.48s on B200 tp8 (previously EXIT_CODE=1 at 1071s with CUDA IMA in _run_attention_warmup during Phase-2 restart). Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
01878c9 to
86e4d94
Compare
7 participants
Summary
_trtllm_gen_jit_warmup_done; skip _run_attention_warmup's JIT warmup body on the second and subsequent invocations with an informational log — Phase-1 semantics unchanged, redundant Phase-2 grid enumeration eliminated.Test plan
Links
Dev Engineer Review
fp8BlockScaleMoe.cppto exclude tiles below the configured minimum for all tactic-selection callers.QA Engineer Review
No test changes.