Skip to content
Open
Changes from all commits
Commits
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
81 changes: 40 additions & 41 deletions cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -400,10 +400,16 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder
int64_t const totalExpertsPerToken = topK + numFusedSharedExpert.value_or(0);
int64_t const numTotalLocalExperts = numLocalExperts + numFusedSharedExpert.value_or(0);
// WAR: the small-tile (tileN 8/16) dynB TRTLLM-Gen batched-GEMM cubins flakily hit an
// illegal memory access (garbage TMA-descriptor pointer, MMU fault in the gemm2 K-loop)
// when shared experts are fused into the grouped GEMM (num_fused_shared_experts > 0);
// illegal memory access (garbage TMA-descriptor pointer, MMU fault in the gemm2 K-loop);
// tileN >= 32 is unaffected (10/10 clean vs minutes-to-crash baseline on B300 TP=4).
// Restrict the fused path to tileN >= 32 until the kernel-side fix lands (nvbug TBD).
// Originally scoped to the fused shared-expert path, but the defect is in the shared
// small-tile cubins and not caused by expert fusion: DeepSeek-R1 FP8 TP=8 (unfused)
// faults identically during warmup, where the 1/2/8-token shapes are the only ones that
// can select tileN 8/16 (12288 tokens gets tileN 64/128 and always passes). Excluding
// the small tiles for every caller is safe because every FP8 block-scale MoE shape also
// offers a tileN >= 32 tactic, so the returned list is never emptied. The tiles stay in
// 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 = []()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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.

{
Expand All @@ -414,7 +420,7 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder
std::vector<std::vector<int64_t>> tactics;
for (auto& [tileN, runner] : mRunners)
{
if (numFusedSharedExpert.value_or(0) > 0 && tileN < fusedMinTileN)
if (tileN < fusedMinTileN)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I think the IMA issues we met only happens when numFusedSharedExpert large than 0. Please explain why apply it for all the fp8 tactic.

Comment thread
coderabbitai[bot] marked this conversation as resolved.
{
continue;
}
Expand Down Expand Up @@ -461,49 +467,42 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder
= static_cast<float>(num_tokens * total_experts_per_token) / num_total_local_experts;
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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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.

// re-validating them against the actual problem size, which can return a config
// whose kernel is absent (illegal memory access at launch). Pick an
// explicitly-validated config instead -- the same set the autotuner draws from --
// searching the heuristic tileN first. Small warmup batches clamp the heuristic to
// mSupportedTileN.front(), so this fallback is the path that reaches the defective
// small-tile cubins; it needs the same tileN >= 32 exclusion as getValidConfigs
// (see the WAR comment there).
config = -1;
std::vector<int32_t> tileN_candidates{static_cast<int32_t>(tileN)};
for (auto t : mSupportedTileN)
{
// getDefaultValidConfigIndex only pairs the per-GEMM "default" indices without
// re-validating them against the actual problem size. For the inflated fused
// expert/topK counts that can return a config whose kernel is absent (illegal
// memory access at launch). Pick an explicitly-validated config instead -- the
// same set the autotuner draws from -- searching the heuristic tileN first.
config = -1;
std::vector<int32_t> tileN_candidates{static_cast<int32_t>(tileN)};
for (auto t : mSupportedTileN)
if (t != tileN)
tileN_candidates.push_back(t);
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}
static int const fusedMinTileNFallback = []()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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.

{
char const* env = std::getenv("TLLM_MOE_FUSED_MIN_TILEN");
return env != nullptr ? std::atoi(env) : 32;
}();
Comment on lines +485 to +489

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.

🩺 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.cpp

Repository: 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
fi

Repository: 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}")
PY

Repository: 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

for (auto t : tileN_candidates)
{
if (t < fusedMinTileNFallback)
{
if (t != tileN)
tileN_candidates.push_back(t);
continue;
}
// Same small-tile exclusion as getValidConfigs (see the WAR comment there).
static int const fusedMinTileNFallback = []()
{
char const* env = std::getenv("TLLM_MOE_FUSED_MIN_TILEN");
return env != nullptr ? std::atoi(env) : 32;
}();
for (auto t : tileN_candidates)
auto valid = mRunners.at(t)->getValidConfigIndices(
total_experts_per_token, hidden_size, intermediate_size, num_total_local_experts, num_tokens);
if (!valid.empty())
{
if (t < fusedMinTileNFallback)
{
continue;
}
auto valid = mRunners.at(t)->getValidConfigIndices(
total_experts_per_token, hidden_size, intermediate_size, num_total_local_experts, num_tokens);
if (!valid.empty())
{
tileN = t;
config = valid.front();
break;
}
tileN = t;
config = valid.front();
break;
}
TLLM_CHECK_WITH_INFO(
config != -1, "No valid TRTLLM-Gen config found for fused shared-expert FP8 block-scale MoE.");
}
else
{
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.");

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

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.

}

return run_fp8_block_scale_moe(routing_logits, routing_bias, hidden_states, hidden_states_scale, gemm1_weights,
Expand Down
Loading