[None][perf] Optimize MiniMax-M3 MXFP8 GEMMs - #17238
Conversation
d350653 to
4f71320
Compare
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
4f71320 to
3c1a752
Compare
WalkthroughMXFP8 GEMM now supports thread-safe tactic caching, autotuned Torch operators, configurable native or FlashInfer dispatch, and decode-graph capture integration. The executor coordinates separate warmup paths, with tests covering backend selection, tactic caching, graph capture, and numerical correctness. ChangesMXFP8 autotuning
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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: 2
🧹 Nitpick comments (6)
tests/unittest/_torch/modules/test_mxfp8_linear.py (2)
354-357: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the local
Linearvariable so it does not shadow theflashinfermodule.Line 355 imports
flashinferinto a local name. Line 375 rebinds the same local name to aLinearinstance. The availability check and the module reference are then unreachable in the rest of the test. Use a distinct name, for exampleflashinfer_linear.Also applies to: 375-381
🤖 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 `@tests/unittest/_torch/modules/test_mxfp8_linear.py` around lines 354 - 357, Rename the local Linear instance created later in the test to a distinct name such as flashinfer_linear, preserving flashinfer as the imported module name from the availability check. Update all subsequent references to that Linear instance within the affected test.
88-88: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPatch the module reference instead of global
torch.ops.
monkeypatch.setattr(linear_module.torch, "ops", ...)mutates the realtorchmodule, becauselinear_module.torchis the globaltorchmodule object. Every module in the process sees the fakeopsnamespace for the duration of the test.monkeypatchrestores it, but an unrelatedtorch.ops.*call inside the code under test fails with a confusingAttributeError.Replace
torchon the module under test instead, so onlylinear_modulesees the fake.♻️ Proposed scoping fix
- monkeypatch.setattr(linear_module.torch, "ops", SimpleNamespace(trtllm=fake_trtllm_ops)) + monkeypatch.setattr( + linear_module, "torch", SimpleNamespace(ops=SimpleNamespace(trtllm=fake_trtllm_ops)) + )Add any other
torchattributes thatMXFP8LinearMethod.applyneeds, for exampleonesandfloat32.🤖 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 `@tests/unittest/_torch/modules/test_mxfp8_linear.py` at line 88, Update the test’s monkeypatch to replace the module-level torch reference on linear_module, rather than assigning to linear_module.torch.ops and mutating global torch.ops. Build the fake torch object with the attributes MXFP8LinearMethod.apply requires, including the fake ops namespace and any needed symbols such as ones and float32.tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py (1)
79-82: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDo not hard-code the tactic count.
expected_tactics = 20 if getSMVersion() == 100 else 10couples this pre-merge test to the exact CUTLASS configuration list. Bothl0_b200.ymlandl0_b300.ymlrun this test on every PR, so any added or removed configuration fails CI for an unrelated change. Assert that at least one configuration exists and iterate the reported count instead.♻️ Proposed fix
- expected_tactics = 20 if getSMVersion() == 100 else 10 - assert runner.get_num_configs() == expected_tactics + num_configs = runner.get_num_configs() + assert num_configs > 0 - for tactic in [-1, *range(expected_tactics)]: + for tactic in [-1, *range(num_configs)]:🤖 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 `@tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py` around lines 79 - 82, Replace the hard-coded expected_tactics calculation with a check that runner.get_num_configs() is greater than zero, then iterate tactics using the count returned by runner.get_num_configs() while preserving the existing -1 tactic case.tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py (1)
244-244: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNarrow the environment patch.
patch.dict("os.environ", {}, clear=True)removes every environment variable for the duration of the block. The test only needsTRTLLM_MXFP8_GEMM_BACKENDto be absent. Clearing everything can change unrelated behavior inside the warmup path, for example logging or autotuner cache-path handling.♻️ Proposed narrower patch
- patch.dict("os.environ", {}, clear=True), + patch.dict("os.environ", {}),Then delete only the one variable inside the block:
os.environ.pop("TRTLLM_MXFP8_GEMM_BACKEND", None)🤖 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 `@tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py` at line 244, In the warmup test’s environment patch, stop clearing the entire environment and preserve all unrelated variables. Within the relevant block, remove only TRTLLM_MXFP8_GEMM_BACKEND using the existing test setup around the warmup invocation.tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)
605-620: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate
unique_idand consider narrowing the tactic-cache sync.Two points:
unique_idhas no return annotation. The coding guidelines require an annotation on every function.sync_tactic_cachere-scans every cache entry for the op on each autotuned call while tuning mode is active. Each MXFP8 linear layer repeats the same full scan and re-registers already registered tactics. The work is redundant during warmup.Track the last synchronized cache size, or sync once after warmup instead of per call.
As per coding guidelines: "Annotate every function, use `None` for procedures".♻️ Proposed annotation fix
- def unique_id(self): + def unique_id(self) -> tuple[torch.dtype, int]: return (self.output_dtype, self.sm_version)🤖 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/custom_ops/torch_custom_ops.py` around lines 605 - 620, Annotate the `unique_id` function with its concrete return type. In `sync_tactic_cache`, avoid rescanning and re-registering the entire profiling cache on every autotuned call by tracking the last synchronized cache size or performing synchronization once after warmup, while preserving synchronization for newly added cache entries.Source: Coding guidelines
cpp/tensorrt_llm/thop/mxfp8Gemm.cpp (1)
285-289: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a message to the bounds check.
TORCH_CHECK(configIdx >= 0 && configIdx < getNumConfigs());produces an error without context. A message makes an out-of-range tactic index from the Python autotuner easy to diagnose.♻️ Proposed diagnostic message
- TORCH_CHECK(configIdx >= 0 && configIdx < getNumConfigs()); + TORCH_CHECK(configIdx >= 0 && configIdx < getNumConfigs(), "MXFP8 config index ", configIdx, + " is out of range [0, ", getNumConfigs(), ").");🤖 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/mxfp8Gemm.cpp` around lines 285 - 289, Update the TORCH_CHECK in getConfig to include a descriptive message for an invalid configIdx, identifying the out-of-range tactic index and the valid configuration range so Python autotuner failures are diagnosable.
🤖 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/modules/linear.py`:
- Around line 3120-3123: In tensorrt_llm/_torch/modules/linear.py:3120-3123,
remove the backend check from needs_native_autotune, gate operator selection in
apply() on the FlashInfer dispatch decision, and consistently clear
use_native_autotuner and _native_autotuned in disable_flashinfer_auto(). In
tensorrt_llm/_torch/pyexecutor/model_engine.py:1569-1582, collect
native_mxfp8_methods before calling quant_method.enable_flashinfer_auto() so
native tuning remains eligible.
In `@tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py`:
- Around line 319-341: Add the missing _release_megamoe_profiling_scratch mock
or callable to the engine SimpleNamespace in
test_flashinfer_mxfp8_autotunes_before_graph_capture, matching the stub’s other
engine methods so both _run_autotuner_warmup calls can complete without
AttributeError.
---
Nitpick comments:
In `@cpp/tensorrt_llm/thop/mxfp8Gemm.cpp`:
- Around line 285-289: Update the TORCH_CHECK in getConfig to include a
descriptive message for an invalid configIdx, identifying the out-of-range
tactic index and the valid configuration range so Python autotuner failures are
diagnosable.
In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`:
- Around line 605-620: Annotate the `unique_id` function with its concrete
return type. In `sync_tactic_cache`, avoid rescanning and re-registering the
entire profiling cache on every autotuned call by tracking the last synchronized
cache size or performing synchronization once after warmup, while preserving
synchronization for newly added cache entries.
In `@tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py`:
- Line 244: In the warmup test’s environment patch, stop clearing the entire
environment and preserve all unrelated variables. Within the relevant block,
remove only TRTLLM_MXFP8_GEMM_BACKEND using the existing test setup around the
warmup invocation.
In `@tests/unittest/_torch/modules/test_mxfp8_linear.py`:
- Around line 354-357: Rename the local Linear instance created later in the
test to a distinct name such as flashinfer_linear, preserving flashinfer as the
imported module name from the availability check. Update all subsequent
references to that Linear instance within the affected test.
- Line 88: Update the test’s monkeypatch to replace the module-level torch
reference on linear_module, rather than assigning to linear_module.torch.ops and
mutating global torch.ops. Build the fake torch object with the attributes
MXFP8LinearMethod.apply requires, including the fake ops namespace and any
needed symbols such as ones and float32.
In `@tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py`:
- Around line 79-82: Replace the hard-coded expected_tactics calculation with a
check that runner.get_num_configs() is greater than zero, then iterate tactics
using the count returned by runner.get_num_configs() while preserving the
existing -1 tactic case.
🪄 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: 2cb98156-18c1-4b0a-b787-853c692a480a
📒 Files selected for processing (10)
cpp/tensorrt_llm/thop/mxfp8Gemm.cpptensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_b300.ymltests/unittest/_torch/executor/test_pytorch_model_engine_warmup.pytests/unittest/_torch/modules/test_mxfp8_linear.pytests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py
|
Could you address these correctness edges before merge?
These cases affect runtime behavior, not only initialization performance. |
peihu-nv
changed the title
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The main issue is that the PR's two optimizations appear to be mutually exclusive within a single process (detailed inline at model_engine.py:1577): enabling the FlashInfer decode-graph path flips backend to "auto", which disqualifies the layer from native large-M autotuning. The description says the changes "together" cover both ends of the workload, but in aggregated MiniMax-M3 serving with decode graphs — the default configuration — context/prefill GEMMs stay on the untuned default CUTLASS config. The two perf measurements were separate disagg arms (CTX-only and GEN-only processes), so they don't demonstrate the combined behavior. If the exclusivity is intentional, please say so in the description and docstring; if not, the backend == "trtllm" condition in needs_native_autotune looks like the culprit, since auto mode still runs the native op for everything except captured decode graphs.
Two smaller asks:
- This is a nontrivial perf feature; it should carry a TRTLLM JIRA ticket rather than
[None]. TRTLLM_MXFP8_GEMM_BACKENDis a new user-facing env var documented only in a docstring — consider adding it to the docs where other TRTLLM_* knobs live.
The test coverage is genuinely good, especially the GPU graph-replay parity test for the FlashInfer scale-layout contract.
| quant_method.enable_flashinfer_auto() | ||
| if quant_method.needs_flashinfer_autotune: | ||
| flashinfer_mxfp8_methods.append(quant_method) | ||
| if enable_trtllm_autotuner and quant_method.needs_native_autotune: |
There was a problem hiding this comment.
Ordering bug (or undocumented design decision): enable_flashinfer_auto() on line 1574 flips backend to "auto", and needs_native_autotune requires backend == "trtllm" (linear.py:3121-3123) — so any layer that gets the FlashInfer decode path is silently excluded from native large-M autotuning. In auto mode, eager/context execution still calls the plain mxfp8_mxfp8_gemm, whose tactic cache is now never populated, so large-M context GEMMs run the untuned default config. For aggregated MiniMax-M3 serving with decode graphs, the PR's second optimization never engages. If the intent is that the native op is still the eager path in auto mode, dropping the backend == "trtllm" condition (or checking needs_native_autotune before calling enable_flashinfer_auto()) would let both apply.
There was a problem hiding this comment.
Good catch, thanks! Fixed in c031804: native and FlashInfer tuning are now independent and run in separate warmup passes. Auto mode keeps the tuned native path for eager/context GEMMs and uses FlashInfer only for decode graphs.
| return act.new_empty((act.size(0), weight.size(0)), dtype=output_dtype) | ||
|
|
||
|
|
||
| _MXFP8_LARGE_M_BUCKETS = (8192, 16384, 32768) |
There was a problem hiding this comment.
These bucket bands duplicate kMxfp8LargeMMin/kMxfp8M16kMin/etc. in cpp/tensorrt_llm/thop/mxfp8Gemm.cpp:47-53. If either side drifts, the failure is silent: Python tunes and registers tactics at bucket M while C++ maps runtime M to a different key, so every lookup misses and falls back to the default config — no error, just the perf win quietly disappearing. Consider exposing the bucket mapping from the C++ runner (single source), or at minimum a unit test that round-trips register_tactic/get_cached_tactic across the band boundaries to pin the two implementations together. The thresholds themselves (6553/13106/19659) also deserve a one-line derivation comment.
There was a problem hiding this comment.
Thanks, agreed. I added a GB300 boundary test that registers tactics through C++ and verifies the Python and C++ mappings at every bucket boundary. I also documented where the thresholds came from.
| global_scale = torch.ones([1], | ||
| dtype=torch.float32, | ||
| device=input.device) | ||
| gemm = (torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned |
There was a problem hiding this comment.
needs_native_autotune defaults True for every MXFP8LinearMethod, and only _run_autotuner_warmup ever resolves it (mark/disable). Any path that skips that warmup — helix CP returns early at model_engine.py:1208, or Linear used outside the PyExecutor engine — leaves it True forever, so every apply() routes through mxfp8_mxfp8_gemm_autotuned and pays an AutoTuner.choose_one cache-miss lookup per GEMM call for the life of the process. Consider defaulting use_native_autotuner off and having the engine opt in, so non-engine users keep the direct op call.
There was a problem hiding this comment.
Good point. Fixed in daed63f: native autotuning now defaults off, and PyTorchModelEngine explicitly enables it during startup warmup. Standalone and Helix paths therefore use the direct native GEMM.
| flashinfer_autotune_context = ( | ||
| flashinfer_mxfp8_autotune() if self.cuda_graph_runner.is_warmup_only | ||
| and flashinfer_methods else contextlib.nullcontext()) | ||
| with flashinfer_autotune_context, flashinfer_mxfp8_decode_graph_capture( |
There was a problem hiding this comment.
Graphs can also be captured lazily during serving (needs_capture path in forward, model_engine.py:6538) — outside this contextvar. A decode graph key not captured during warmup (e.g. a spec-decode on/off variant) bakes the native GEMM instead of FlashInfer, so per-shape decode perf silently depends on which keys happened to capture during warmup. Worth either setting the capture contextvar in the lazy-capture path too, or noting the limitation here.
There was a problem hiding this comment.
Thanks for flagging this. I double-checked the current graph-runner behavior: runtime shapes are first padded to an existing graph key. If the resulting key is still missing, _capture_allowed is false during serving, so it falls back to eager execution instead of capturing a native-GEMM graph. Therefore, no change is needed in this PR.
| - unittest/_torch/modules/test_rotary_embedding.py | ||
| - unittest/_torch/modules/mamba | ||
| - unittest/_torch/modules/tests_lora_modules | ||
| - unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py |
There was a problem hiding this comment.
This same pre-merge section already contains unittest/_torch/thop/parallel TIMEOUT (90) (line 143), which picks up the new file automatically — this explicit entry double-registers it. Same duplication in l0_b300.yml (directory at line 23, explicit entry at line 36). The test_mxfp8_linear.py additions are fine since modules tests are enumerated individually.
There was a problem hiding this comment.
Good catch, thanks. I removed the duplicate explicit entries from both B200 and B300 lists; the existing directory entries already include this test.
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/modules/test_mxfp8_linear.py (1)
298-322: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winReset
synced_cache_keyswith the profiling cache lifecycle.MXFP8GemmRunner.synced_cache_keysis shared across instances, butAutoTuner.clear_cache()does not clear it. A reprofiled cache entry with the samecache_keycan therefore skipregister_tactic().id(tuner.profiling_cache)is also reusable after object destruction. Key synchronization state by cache lifetime or clear it when the profiling cache is cleared.Test coverage:
test_mxfp8_auto_fallback_does_not_rearm_native_autotuningcovers fallback state, andtest_mxfp8_native_autotuner_syncs_profilescovers repeated synchronization for one live cache. Add a regression test for cache clearing and cache replacement. The test file is listed intest-db/l0_b200.ymlandtest-db/l0_b300.yml. Coverage verdict: insufficient.🤖 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 `@tests/unittest/_torch/modules/test_mxfp8_linear.py` around lines 298 - 322, Update MXFP8GemmRunner.sync_tactic_cache and the synced_cache_keys lifecycle so clearing or replacing AutoTuner.profiling_cache cannot suppress register_tactic for a reprofiled cache_key; avoid relying solely on reusable profiling-cache object IDs, and reset or scope synchronization state to the cache lifetime. Extend test_mxfp8_native_autotuner_syncs_profiles with regression coverage for cache clearing and cache replacement while preserving the existing repeated-sync behavior.
🧹 Nitpick comments (5)
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (2)
576-576: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the required modern type annotations.
Replace
List[...]with built-inlist[...]. Type**kwargsasobjectbecause the override accepts arbitrary keyword values.Proposed annotation update
-def _mxfp8_scale_infer_shape(input_shapes: List[List[int]]) -> int: +def _mxfp8_scale_infer_shape(input_shapes: list[list[int]]) -> int: - def get_valid_tactics(self, inputs: List[torch.Tensor], - profile: OptimizationProfile, **kwargs) -> List[int]: + def get_valid_tactics(self, inputs: list[torch.Tensor], + profile: OptimizationProfile, + **kwargs: object) -> list[int]:Also applies to: 607-608
🤖 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/custom_ops/torch_custom_ops.py` at line 576, Update _mxfp8_scale_infer_shape and the additionally referenced definitions to use modern built-in list[...] annotations instead of List[...]. Annotate any **kwargs parameters in the affected override as object while preserving the existing method signatures and behavior.Source: Coding guidelines
583-585: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the mutable runner caches with
ClassVar.
MXFP8GemmRunner.runner_dictandsynced_cache_keysare shared by instances ofMXFP8GemmRunner, not process-wide. Combine each annotation with its initializer. Apply the same change to the other mutable runner-cache declarations in this file; Ruff reports 11RUF012violations.🤖 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/custom_ops/torch_custom_ops.py` around lines 583 - 585, Update the mutable runner-cache declarations throughout this file, including MXFP8GemmRunner.runner_dict and synced_cache_keys, to combine each ClassVar annotation with its initializer. Apply the same ClassVar treatment to all 11 runner-cache declarations reported by Ruff RUF012, preserving their existing types and initial values.Source: Linters/SAST tools
tests/unittest/_torch/modules/test_mxfp8_linear.py (1)
88-93: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake the fake
torchdelegate unknown attributes to the real module.
fake_torchexposes onlyops,ones, andfloat32. Any othertorch.*attribute reached byMXFP8LinearMethod.applyraisesAttributeError, which reads as a test bug rather than a production change.onesandfloat32were already added for that reason. A delegating proxy keeps the operator stubs while letting every other attribute resolve to the realtorch.♻️ Proposed refactor
- fake_torch = SimpleNamespace( - ops=SimpleNamespace(trtllm=fake_trtllm_ops), - ones=torch.ones, - float32=torch.float32, - ) - monkeypatch.setattr(linear_module, "torch", fake_torch) + class _FakeTorch: + ops = SimpleNamespace(trtllm=fake_trtllm_ops) + + def __getattr__(self, name): + return getattr(torch, name) + + monkeypatch.setattr(linear_module, "torch", _FakeTorch())🤖 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 `@tests/unittest/_torch/modules/test_mxfp8_linear.py` around lines 88 - 93, Update the fake_torch setup in the MXFP8LinearMethod test so unknown attributes delegate to the real torch module while preserving the custom ops, ones, and float32 overrides. Use a proxy mechanism such as attribute fallback to torch, then continue monkeypatching linear_module.torch with that delegating fake.tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py (2)
481-555: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the forced-FlashInfer rank mismatch.
This test covers the automatic (
auto) path, where a rank mismatch logs a warning and falls back._run_autotuner_warmuphas a second branch: if any method hasbackend == "flashinfer"(set throughTRTLLM_MXFP8_GEMM_BACKEND=flashinfer), the mismatch must raiseRuntimeErrorinstead of falling back silently. No test asserts that branch, so a regression that downgrades the error to a warning would pass CI.Add a variant that sets
TRTLLM_MXFP8_GEMM_BACKEND=flashinferand assertsRuntimeErrorwithassertRaises.🤖 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 `@tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py` around lines 481 - 555, Add a test variant for _run_autotuner_warmup that sets TRTLLM_MXFP8_GEMM_BACKEND to flashinfer while retaining the simulated TP rank mismatch, and assert the call raises RuntimeError. Verify the forced FlashInfer path fails before warmup rather than falling back, using the existing engine, tuner, and dist setup from test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup.
220-576: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winComplete MXFP8 warmup coverage and registration.
The five added tests run through
unittest/_torch/executorinl0_b300.yml, butl0_b200.ymldoes not include this directory. Add the module to the B200 list if B200 is an intended target. Add tests for the forced-FlashInferRuntimeErrorpath and pipeline-parallel rank divergence.Test coverage summary: The MXFP8 linear and GEMM tests are listed in both B200 and B300. Warmup coverage remains insufficient.
🤖 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 `@tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py` around lines 220 - 576, Add the warmup test module containing test_flashinfer_mxfp8_respects_disabled_global_autotuner and related MXFP8 tests to the B200 test list when B200 is an intended target. Extend the warmup coverage with tests for the forced-FlashInfer RuntimeError path and pipeline-parallel rank divergence, following the existing PyTorchModelEngine._run_autotuner_warmup test patterns.Source: Path instructions
🤖 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.
Outside diff comments:
In `@tests/unittest/_torch/modules/test_mxfp8_linear.py`:
- Around line 298-322: Update MXFP8GemmRunner.sync_tactic_cache and the
synced_cache_keys lifecycle so clearing or replacing AutoTuner.profiling_cache
cannot suppress register_tactic for a reprofiled cache_key; avoid relying solely
on reusable profiling-cache object IDs, and reset or scope synchronization state
to the cache lifetime. Extend test_mxfp8_native_autotuner_syncs_profiles with
regression coverage for cache clearing and cache replacement while preserving
the existing repeated-sync behavior.
---
Nitpick comments:
In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`:
- Line 576: Update _mxfp8_scale_infer_shape and the additionally referenced
definitions to use modern built-in list[...] annotations instead of List[...].
Annotate any **kwargs parameters in the affected override as object while
preserving the existing method signatures and behavior.
- Around line 583-585: Update the mutable runner-cache declarations throughout
this file, including MXFP8GemmRunner.runner_dict and synced_cache_keys, to
combine each ClassVar annotation with its initializer. Apply the same ClassVar
treatment to all 11 runner-cache declarations reported by Ruff RUF012,
preserving their existing types and initial values.
In `@tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py`:
- Around line 481-555: Add a test variant for _run_autotuner_warmup that sets
TRTLLM_MXFP8_GEMM_BACKEND to flashinfer while retaining the simulated TP rank
mismatch, and assert the call raises RuntimeError. Verify the forced FlashInfer
path fails before warmup rather than falling back, using the existing engine,
tuner, and dist setup from
test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup.
- Around line 220-576: Add the warmup test module containing
test_flashinfer_mxfp8_respects_disabled_global_autotuner and related MXFP8 tests
to the B200 test list when B200 is an intended target. Extend the warmup
coverage with tests for the forced-FlashInfer RuntimeError path and
pipeline-parallel rank divergence, following the existing
PyTorchModelEngine._run_autotuner_warmup test patterns.
In `@tests/unittest/_torch/modules/test_mxfp8_linear.py`:
- Around line 88-93: Update the fake_torch setup in the MXFP8LinearMethod test
so unknown attributes delegate to the real torch module while preserving the
custom ops, ones, and float32 overrides. Use a proxy mechanism such as attribute
fallback to torch, then continue monkeypatching linear_module.torch with that
delegating fake.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 51483533-5edf-4903-b893-e8ea35bcd9d2
📒 Files selected for processing (7)
cpp/tensorrt_llm/thop/mxfp8Gemm.cpptensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/executor/test_pytorch_model_engine_warmup.pytests/unittest/_torch/modules/test_mxfp8_linear.pytests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py
- cpp/tensorrt_llm/thop/mxfp8Gemm.cpp
- tensorrt_llm/_torch/modules/linear.py
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py (1)
170-172: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd the procedure return annotation.
Add
-> Nonetotest_mxfp8_native_tactic_cache_large_m_bucket_boundaries. The Python guidelines require annotations on every function.Proposed fix
def test_mxfp8_native_tactic_cache_large_m_bucket_boundaries( lower_bound: int, upper_bound: int, bucket: int -): +) -> None:As per coding guidelines, annotate every function.
🤖 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 `@tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py` around lines 170 - 172, Add the required `-> None` return annotation to the `test_mxfp8_native_tactic_cache_large_m_bucket_boundaries` test function signature, leaving its parameters and implementation unchanged.Source: Coding guidelines
🤖 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 `@docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md`:
- Around line 108-112: Update the `flashinfer` entry in the execution-mode
descriptions so it states that eager, context/prefill, and piecewise execution
remain on the native TensorRT LLM GEMM path, while FlashInfer applies only to
eligible decode CUDA-graph GEMMs. Leave the `trtllm` and `auto` descriptions
unchanged.
---
Nitpick comments:
In `@tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py`:
- Around line 170-172: Add the required `-> None` return annotation to the
`test_mxfp8_native_tactic_cache_large_m_bucket_boundaries` test function
signature, leaving its parameters and implementation unchanged.
🪄 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: d40f8b20-d73c-4b22-9457-a78e277c449c
📒 Files selected for processing (9)
docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.mdtensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_b300.ymltests/unittest/_torch/executor/test_pytorch_model_engine_warmup.pytests/unittest/_torch/modules/test_mxfp8_linear.pytests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py
💤 Files with no reviewable changes (2)
- tests/integration/test_lists/test-db/l0_b200.yml
- tests/integration/test_lists/test-db/l0_b300.yml
🚧 Files skipped from review as they are similar to previous changes (4)
- tensorrt_llm/_torch/pyexecutor/model_engine.py
- tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
- tensorrt_llm/_torch/modules/linear.py
- tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
| * `trtllm` uses the native TensorRT LLM GEMM for eager execution and CUDA graphs. | ||
| * `flashinfer` forces FlashInfer for both eager execution and CUDA graphs; it | ||
| requires the pinned `flashinfer-python` package and Blackwell MXFP8 support. | ||
| * `auto` keeps eager execution on the native GEMM and uses FlashInfer in | ||
| captured decode CUDA graphs after startup tuning. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Correct the flashinfer behavior description.
Line 109 states that flashinfer uses FlashInfer for eager GEMMs. The stated runtime contract keeps eager, context/prefill, and piecewise execution on the native TensorRT LLM GEMM path. Document FlashInfer as applying only to eligible decode CUDA-graph GEMMs.
Proposed documentation fix
-* `flashinfer` forces FlashInfer for both eager execution and CUDA graphs; it
- requires the pinned `flashinfer-python` package and Blackwell MXFP8 support.
+* `flashinfer` forces FlashInfer for eligible decode CUDA-graph GEMMs. Eager,
+ context/prefill, and piecewise graph execution use the native GEMM. It
+ requires the pinned `flashinfer-python` package and Blackwell MXFP8 support.As per PR objectives, FlashInfer dispatch is limited to eligible decode CUDA-graph GEMMs.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| * `trtllm` uses the native TensorRT LLM GEMM for eager execution and CUDA graphs. | |
| * `flashinfer` forces FlashInfer for both eager execution and CUDA graphs; it | |
| requires the pinned `flashinfer-python` package and Blackwell MXFP8 support. | |
| * `auto` keeps eager execution on the native GEMM and uses FlashInfer in | |
| captured decode CUDA graphs after startup tuning. | |
| * `trtllm` uses the native TensorRT LLM GEMM for eager execution and CUDA graphs. | |
| * `flashinfer` forces FlashInfer for eligible decode CUDA-graph GEMMs. Eager, | |
| context/prefill, and piecewise graph execution use the native GEMM. It | |
| requires the pinned `flashinfer-python` package and Blackwell MXFP8 support. | |
| * `auto` keeps eager execution on the native GEMM and uses FlashInfer in | |
| captured decode CUDA graphs after startup tuning. |
🤖 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 `@docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md`
around lines 108 - 112, Update the `flashinfer` entry in the execution-mode
descriptions so it states that eager, context/prefill, and piecewise execution
remain on the native TensorRT LLM GEMM path, while FlashInfer applies only to
eligible decode CUDA-graph GEMMs. Leave the `trtllm` and `auto` descriptions
unchanged.
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
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. |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cpp/tensorrt_llm/thop/mxfp8Gemm.cpp (1)
221-237: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd Doxygen documentation for the new C++ interfaces.
Document
mxfp8_mxfp8_gemm,MXFP8GemmRunner, and its public methods with Doxygen comments. Include tensor shapes, supported dtypes, tactic sentinel values, cache scope, and error conditions.The coding guidelines require Doxygen documentation for new interfaces.
Also applies to: 243-294
🤖 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/mxfp8Gemm.cpp` around lines 221 - 237, Add Doxygen comments for mxfp8_mxfp8_gemm, MXFP8GemmRunner, and each public method. Document tensor shapes, supported input/output dtypes, tactic sentinel values, cache scope, and the error conditions each interface can raise, while keeping the existing declarations and behavior unchanged.Source: Coding guidelines
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)
584-612: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new Python interfaces and tensor contracts.
Add Google-style docstrings to
MXFP8GemmRunnerandmxfp8_mxfp8_gemm_autotuned. Document input tensor dimensions, required dtypes, scale layouts, output shape, tactic semantics, and raised exceptions.The coding guidelines require docstrings for externally usable interfaces and dimensions for public Tensor-like arguments.
Also applies to: 640-676
🤖 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/custom_ops/torch_custom_ops.py` around lines 584 - 612, Add Google-style docstrings to the public `MXFP8GemmRunner` class and `mxfp8_mxfp8_gemm_autotuned` function. Document tensor dimensions, required dtypes, scale layouts, output shape, tactic-selection semantics, and exceptions raised, including dimension details for every public Tensor-like argument; preserve the existing behavior and interfaces.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@cpp/tensorrt_llm/thop/mxfp8Gemm.cpp`:
- Around line 221-237: Add Doxygen comments for mxfp8_mxfp8_gemm,
MXFP8GemmRunner, and each public method. Document tensor shapes, supported
input/output dtypes, tactic sentinel values, cache scope, and the error
conditions each interface can raise, while keeping the existing declarations and
behavior unchanged.
In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`:
- Around line 584-612: Add Google-style docstrings to the public
`MXFP8GemmRunner` class and `mxfp8_mxfp8_gemm_autotuned` function. Document
tensor dimensions, required dtypes, scale layouts, output shape,
tactic-selection semantics, and exceptions raised, including dimension details
for every public Tensor-like argument; preserve the existing behavior and
interfaces.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6a40e78c-6864-407e-b72c-cbf43f4b4985
📒 Files selected for processing (11)
cpp/tensorrt_llm/thop/mxfp8Gemm.cppdocs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.mdtensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_b300.ymltests/unittest/_torch/executor/test_pytorch_model_engine_warmup.pytests/unittest/_torch/modules/test_mxfp8_linear.pytests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py
🚧 Files skipped from review as they are similar to previous changes (9)
- tests/integration/test_lists/test-db/l0_b200.yml
- tests/unittest/_torch/thop/parallel/test_mxfp8_mxfp8_gemm.py
- tensorrt_llm/_torch/models/modeling_minimaxm3.py
- docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md
- tensorrt_llm/_torch/pyexecutor/model_engine.py
- tests/integration/test_lists/test-db/l0_b300.yml
- tests/unittest/_torch/modules/test_mxfp8_linear.py
- tensorrt_llm/_torch/modules/linear.py
- tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
|
/bot run |
|
PR_Github #64214 [ run ] triggered by Bot. Commit: |
|
PR_Github #64214 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
6 participants
Summary
MXFP8GemmRunner.Dev Engineer Review
enable_autotuner=Falsedisables FlashInfer and native MXFP8 warmup.QA Engineer Review
Modified test-list files:
tests/integration/test_lists/test-db/l0_b200.ymlunittest/_torch/modules/test_mxfp8_linear.py.tests/integration/test_lists/test-db/l0_b300.ymlunittest/_torch/modules/test_mxfp8_linear.py.Added or updated test coverage:
Description
This PR combines two complementary MiniMax-M3 MXFP8 GEMM optimizations:
flashinfer.mm_mxfp8. Eager execution, context/prefill execution, and piecewise CUDA graphs retain the native TensorRT-LLM path.PyTorchModelEngineexplicitly opts MXFP8 layers into native tuning during standard startup autotuning. The engine profiles the compiled CUTLASS tactic portfolio and caches the best tactic per shape. Standalone modules and engine paths that skip autotuner warmup, including Helix CP, remain on the direct native GEMM instead of retaining a serving-time Python autotuner lookup.Native and FlashInfer tuning use separate startup forwards. FlashInfer availability is synchronized across tensor-parallel ranks, and a missing warmup batch safely falls back to the plain native GEMM. The native tactic cache uses stable 8K, 16K, and 32K context buckets, with Python/C++ boundary consistency covered by round-trip tests.
For advanced debugging and performance experiments, the MiniMax-M3 deployment guide now documents
TRTLLM_MXFP8_GEMM_BACKEND={trtllm,flashinfer,auto}. Normal deployments should leave the variable unset and use automatic selection.Together these changes select a better MXFP8 implementation at both ends of the workload: FlashInfer for the small-M captured decode path and shape-specific native CUTLASS tactics for large-M execution.
This main-branch port consolidates the implementations originally merged into
feat/m3_with_msain #16695 and #16816.Performance
FlashInfer decode path
Matched GB200 MiniMax-M3 1P1D disaggregated serving:
Native large-M autotuning
Matched GB300 MiniMax-M3 2P8D disaggregated serving (random 8K/1K, concurrency 256, 2,560 requests; CTX image changed only):
Test Coverage
unittest/_torch/thop/parallelsuite entries without duplicate registration.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions).
If PR introduces API changes, an appropriate PR label is added—either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.