[None][perf] Use FlashInfer MXFP8 GEMM for MiniMax-M3 decode - #17237
[None][perf] Use FlashInfer MXFP8 GEMM for MiniMax-M3 decode#17237peihu-nv wants to merge 4 commits into
Conversation
|
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:
WalkthroughMXFP8 linear execution now supports FlashInfer, TensorRT-LLM, and automatic backend selection. Executor warmup performs FlashInfer autotuning when enabled. CUDA graph generation uses the FlashInfer decode-capture context. MiniMax-M3 opts into this behavior. ChangesFlashInfer MXFP8 execution
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ModelEngine
participant MXFP8LinearMethod
participant FlashInfer
ModelEngine->>MXFP8LinearMethod: discover MXFP8 methods
ModelEngine->>FlashInfer: enter autotuning context
ModelEngine->>MXFP8LinearMethod: run warmup forward
MXFP8LinearMethod->>FlashInfer: autotune FlashInfer GEMM
ModelEngine->>FlashInfer: capture decode graph
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: 3
🧹 Nitpick comments (4)
tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py (2)
247-247: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNarrow the environment patch.
patch.dict("os.environ", {}, clear=True)removes every environment variable for the whole block. The code under test only readsTRTLLM_MXFP8_GEMM_BACKENDandTLLM_AUTOTUNER_CACHE_PATH. Clearing everything also removes variables that CUDA, PyTorch, and MPI read, which can make the test order-dependent or environment-dependent.Remove only the keys the test controls.
♻️ Proposed narrower patch
- patch.dict("os.environ", {}, clear=True), + patch.dict( + "os.environ", + {}, + clear=False, + ),Then drop the two keys explicitly inside the block, for example with
os.environ.pop("TRTLLM_MXFP8_GEMM_BACKEND", None)and
os.environ.pop("TLLM_AUTOTUNER_CACHE_PATH", None);patch.dictrestores both on exit.🤖 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 247, Update the environment patch in the affected test to preserve unrelated environment variables instead of using clear=True. Within the patch block, explicitly remove only TRTLLM_MXFP8_GEMM_BACKEND and TLLM_AUTOTUNER_CACHE_PATH so patch.dict restores those controlled keys on exit.
219-293: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd tests for the unexecuted FlashInfer warmup branches.
Test coverage is insufficient. Add cases for
ran_forward is False, forced-flashinferfailure, theTRTLLM_MXFP8_GEMM_BACKENDoverride, and disabled CUDA graphs. The test file is not listed undertests/integration/test_lists/test-db/ortests/integration/test_lists/qa/.Changed tests: added
TestWarmupCleanup.test_flashinfer_mxfp8_autotunes_before_graph_capture; modified or removed 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` around lines 219 - 293, Expand TestWarmupCleanup coverage around PyTorchModelEngine._run_autotuner_warmup with separate cases for ran_forward=False, forced FlashInfer failure, the TRTLLM_MXFP8_GEMM_BACKEND override, and disabled CUDA graphs; assert each branch’s backend, autotune, and warmup behavior. Register the test file in the appropriate tests/integration/test_lists/test-db/ or qa/ list.Source: Path instructions
tests/unittest/_torch/modules/test_mxfp8_linear.py (2)
67-75: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
linear_module.torchis the globaltorchmodule, so this replacestorch.opsprocess-wide.
linear_module.torchis not a module-local alias; it is the same object astorch. Line 74 therefore swaps the wholetorch.opsnamespace for aSimpleNamespacethat exposes onlymxfp8_quantizeandmxfp8_mxfp8_gemm.monkeypatchrestores it at teardown, so there is no cross-test leak today. The current code paths under test reach no other op, so the tests pass.The risk is future breakage: if
MXFP8LinearMethod.applyever calls anothertorch.ops.*entry, these tests fail with a confusingAttributeErrorrather than a meaningful assertion.Patching only the
trtllmnamespace is narrower. Note thattorch.ops.trtllmmay be absent when the compiled library is not loaded, which is presumably why the broader swap was chosen; if so, a short comment recording that constraint would help the next reader.🤖 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 67 - 75, Update _mock_mxfp8_ops to patch only the torch.ops.trtllm namespace instead of replacing the process-wide torch.ops object, while preserving the mock mxfp8 operations and restoration behavior. Account for torch.ops.trtllm being absent when the compiled library is unavailable, and add a brief comment documenting that constraint if the fallback replacement remains necessary.
112-141: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest coverage summary (QA review) and one missing dispatch branch.
Changed test functions in this file:
- Added:
_mock_mxfp8_ops(helper),test_mxfp8_flashinfer_call_contract,test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer,test_mxfp8_rejects_unknown_backend.- Modified:
test_mxfp8_dispatch_returns_mxfp8_method(now isolatesTRTLLM_MXFP8_GEMM_BACKENDand asserts the default backend).- Removed: none.
Test list registration: I cannot confirm from the supplied context whether this file appears under
tests/integration/test_lists/test-db/for CI ortests/integration/test_lists/qa/for manual QA. The new tests need no GPU, while the pre-existing tests in the same file call.cuda(), so the file may need entries in more than one list.What the new tests cover:
- Default backend resolution and rejection of an unknown backend value.
- Forced
flashinferdispatch: argument order, zero-copy weight transpose, scale forwarding, and exact kwargs.- Auto mode: eager execution stays on the native GEMM, decode graph capture routes to
mm_mxfp8, and leaving the capture scope restores the native path.Coverage gap in the dispatch predicate:
MXFP8LinearMethod.applyreaches FlashInfer through two independent conditions (tensorrt_llm/_torch/modules/linear.pylines 3196-3200). This test exercises only_FLASHINFER_MXFP8_DECODE_GRAPH_CAPTURE_ACTIVEcombined with_flashinfer_autotuned. The_FLASHINFER_MXFP8_AUTOTUNE_ACTIVEbranch is untested, and so is the case where decode capture is active butmark_flashinfer_autotuned()was never called.Verdict: needs follow-up. Add the two missing predicate cases.
As per path instructions for
tests/**: "Always produce a test coverage summary, even if no issues are found" and the summary must state "Whether each changed test is listed in the appropriate test list files under tests/integration/test_lists/".💚 Proposed additional assertions for the untested predicate branches
method = MXFP8LinearMethod() assert method.enable_flashinfer_auto() assert method.apply(module, activation, bias=None) is native_output native_gemm.assert_called_once() mm_mxfp8.assert_not_called() + # Decode capture without a completed autotune must stay on the native path. + with flashinfer_mxfp8_decode_graph_capture(): + assert method.apply(module, activation, bias=None) is native_output + mm_mxfp8.assert_not_called() + method.mark_flashinfer_autotuned() with flashinfer_mxfp8_decode_graph_capture(): assert method.apply(module, activation, bias=None) is flashinfer_output mm_mxfp8.assert_called_once()Cover the autotune-active branch in a separate test, because entering
flashinfer_mxfp8_autotune()also importsflashinfer.autotuner:def test_mxfp8_auto_dispatches_flashinfer_while_autotuning(monkeypatch): monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) monkeypatch.setattr(linear_module, "_mxfp8_cutlass_op_available", lambda: True) flashinfer_output = torch.empty((2, 3), dtype=torch.bfloat16) mm_mxfp8 = Mock(return_value=flashinfer_output) monkeypatch.setitem(sys.modules, "flashinfer", SimpleNamespace(mm_mxfp8=mm_mxfp8)) _mock_mxfp8_ops(monkeypatch) module = SimpleNamespace( weight=torch.empty((3, 4), dtype=torch.float8_e4m3fn), weight_scale=torch.empty(512, dtype=torch.uint8), dtype=torch.bfloat16, ) activation = torch.randn((2, 4), dtype=torch.bfloat16) method = MXFP8LinearMethod() assert method.enable_flashinfer_auto() # Autotuning alone must route to FlashInfer, without mark_flashinfer_autotuned(). token = linear_module._FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.set(True) try: assert method.apply(module, activation, bias=None) is flashinfer_output finally: linear_module._FLASHINFER_MXFP8_AUTOTUNE_ACTIVE.reset(token) mm_mxfp8.assert_called_once()🤖 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 112 - 141, Add coverage for both missing branches in MXFP8LinearMethod.apply: verify flashinfer dispatch when _FLASHINFER_MXFP8_AUTOTUNE_ACTIVE is set without mark_flashinfer_autotuned(), and verify decode-capture dispatch when that method has not been called. Add focused tests using the existing _mock_mxfp8_ops setup and ensure each temporary state is reset after the assertion; include the requested test coverage summary and state whether the changed tests are listed in the appropriate integration test-list files.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.
Inline comments:
In `@tensorrt_llm/_torch/modules/linear.py`:
- Around line 3044-3054: The flashinfer_mxfp8_autotune context manager imports
autotune from an unsupported module. Load autotune from flashinfer.autotune
within _load_flashinfer, retain its unavailable state for auto-mode fallback to
trtllm, and update flashinfer_mxfp8_autotune to use the loaded API without
importing flashinfer.autotuner.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 1597-1602: Add generation-shaped warmup forwards before CUDA graph
capture, covering every configured capture shape derived from batch_size and
draft_len. Ensure these warmups run while flashinfer_autotune_context is active
so FlashInfer tunes each GEMM profile’s M, N, and K, while preserving the
existing context-shape warmup and capture flow.
In `@tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py`:
- Around line 252-284: Add the missing _release_megamoe_profiling_scratch
callable to the engine SimpleNamespace stub used by
test_flashinfer_mxfp8_autotunes_before_graph_capture, so _run_autotuner_warmup
can invoke it unconditionally without raising AttributeError.
---
Nitpick comments:
In `@tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py`:
- Line 247: Update the environment patch in the affected test to preserve
unrelated environment variables instead of using clear=True. Within the patch
block, explicitly remove only TRTLLM_MXFP8_GEMM_BACKEND and
TLLM_AUTOTUNER_CACHE_PATH so patch.dict restores those controlled keys on exit.
- Around line 219-293: Expand TestWarmupCleanup coverage around
PyTorchModelEngine._run_autotuner_warmup with separate cases for
ran_forward=False, forced FlashInfer failure, the TRTLLM_MXFP8_GEMM_BACKEND
override, and disabled CUDA graphs; assert each branch’s backend, autotune, and
warmup behavior. Register the test file in the appropriate
tests/integration/test_lists/test-db/ or qa/ list.
In `@tests/unittest/_torch/modules/test_mxfp8_linear.py`:
- Around line 67-75: Update _mock_mxfp8_ops to patch only the torch.ops.trtllm
namespace instead of replacing the process-wide torch.ops object, while
preserving the mock mxfp8 operations and restoration behavior. Account for
torch.ops.trtllm being absent when the compiled library is unavailable, and add
a brief comment documenting that constraint if the fallback replacement remains
necessary.
- Around line 112-141: Add coverage for both missing branches in
MXFP8LinearMethod.apply: verify flashinfer dispatch when
_FLASHINFER_MXFP8_AUTOTUNE_ACTIVE is set without mark_flashinfer_autotuned(),
and verify decode-capture dispatch when that method has not been called. Add
focused tests using the existing _mock_mxfp8_ops setup and ensure each temporary
state is reset after the assertion; include the requested test coverage summary
and state whether the changed tests are listed in the appropriate integration
test-list files.
🪄 Autofix (Beta)
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: 672fcd03-9107-4807-b1c6-b774b210541b
📒 Files selected for processing (5)
tensorrt_llm/_torch/models/modeling_minimaxm3.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.py
ac57e15 to
befb30b
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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 3044-3064: Add the required return annotations to the changed
symbols: annotate both context managers, flashinfer_mxfp8_autotune and
flashinfer_mxfp8_decode_graph_capture, with Iterator[None], and annotate
MXFP8LinearMethod.__init__ with -> None. Ensure the necessary Iterator import is
available.
In `@tests/unittest/_torch/modules/test_mxfp8_linear.py`:
- Around line 252-255: Move the monkeypatch.delenv call for
TRTLLM_MXFP8_GEMM_BACKEND to before the reference Linear construction and its
MXFP8LinearMethod initialization, while preserving the existing flashinfer
method setup and assertion afterward.
- Around line 264-268: Update the test setup around the graph-capture flow using
the visible _flashinfer_mxfp8 method so it is wrapped or mocked before capture,
then assert its call count increases during
flashinfer_mxfp8_decode_graph_capture and graph replay. Keep the existing output
comparison, but make the test explicitly verify dispatch reaches FlashInfer
rather than only matching native results.
- Around line 56-65: Add tests/unittest/_torch/modules/test_mxfp8_linear.py to
the appropriate CI test lists, including l0_b200.yml and l0_b300.yml, so all
listed MXFP8 tests run in CI. Add a QA test-list entry only if these tests
require manual QA.
🪄 Autofix (Beta)
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: b02b9133-e36b-4345-be13-0168f8abc7cc
📒 Files selected for processing (5)
tensorrt_llm/_torch/models/modeling_minimaxm3.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.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/_torch/models/modeling_minimaxm3.py
- tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
- tensorrt_llm/_torch/pyexecutor/model_engine.py
| def test_mxfp8_dispatch_returns_mxfp8_method(monkeypatch): | ||
| """get_quant_method must dispatch QuantAlgo.MXFP8 to MXFP8LinearMethod. | ||
|
|
||
| This is a pure dispatch check; no CUDA required. | ||
| """ | ||
| monkeypatch.delenv("TRTLLM_MXFP8_GEMM_BACKEND", raising=False) | ||
| qc = QuantConfig(quant_algo=QuantAlgo.MXFP8, group_size=32) | ||
| method = get_quant_method(qc) | ||
| assert isinstance(method, MXFP8LinearMethod) | ||
| assert method.backend == "trtllm" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
for list_dir in tests/integration/test_lists/test-db tests/integration/test_lists/qa; do
echo "=== ${list_dir} ==="
rg -n --glob '*.yml' \
-e 'test_mxfp8_dispatch_returns_mxfp8_method' \
-e 'test_mxfp8_flashinfer_call_contract' \
-e 'test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer' \
-e 'test_mxfp8_rejects_unknown_backend' \
-e 'test_mxfp8_flashinfer_decode_graph_matches_native' \
"${list_dir}" || true
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 242
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== changed test file ==="
git diff --unified=3 -- tests/unittest/_torch/modules/test_mxfp8_linear.py
echo "=== test-list files ==="
find tests/integration/test_lists/test-db tests/integration/test_lists/qa \
-maxdepth 2 -type f \( -name '*.yml' -o -name '*.yaml' -o -name '*.json' -o -name '*.txt' \) -print 2>/dev/null | sort
echo "=== test file outline and relevant definitions ==="
ast-grep outline tests/unittest/_torch/modules/test_mxfp8_linear.py
rg -n '^\s*def test_|^\s*async def test_' tests/unittest/_torch/modules/test_mxfp8_linear.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 5990
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== changed test file ==="
git diff --unified=3 -- tests/unittest/_torch/modules/test_mxfp8_linear.py
echo "=== test-list files ==="
find tests/integration/test_lists/test-db tests/integration/test_lists/qa \
-maxdepth 2 -type f \( -name '*.yml' -o -name '*.yaml' -o -name '*.json' -o -name '*.txt' \) -print 2>/dev/null | sort
echo "=== test definitions ==="
rg -n '^\s*(async )?def test_' tests/unittest/_torch/modules/test_mxfp8_linear.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 5324
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== MXFP8 and test-file references in test lists ==="
rg -n -i --glob '*.{yml,yaml,txt,json}' \
-e 'mxfp8' \
-e 'test_mxfp8_linear' \
-e 'tests/unittest/_torch/modules' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa || true
echo "=== representative list schemas and entries ==="
for file in \
tests/integration/test_lists/test-db/l0_sanity_check.yml \
tests/integration/test_lists/test-db/l0_h100.yml \
tests/integration/test_lists/qa/llm_function_core.txt \
tests/integration/test_lists/qa/llm_config_database.yml; do
if [ -f "$file" ]; then
echo "--- $file ---"
sed -n '1,100p' "$file"
fi
done
echo "=== all repository references to the test file and changed test names ==="
rg -n \
-e 'test_mxfp8_linear\.py' \
-e 'test_mxfp8_dispatch_returns_mxfp8_method' \
-e 'test_mxfp8_flashinfer_call_contract' \
-e 'test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer' \
-e 'test_mxfp8_rejects_unknown_backend' \
-e 'test_mxfp8_flashinfer_decode_graph_matches_native' \
. || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 31191
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== test-list references ==="
rg -n -i \
-e 'mxfp8' \
-e 'test_mxfp8_linear' \
-e 'tests/unittest/_torch/modules' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa || true
echo "=== repository references ==="
rg -n \
-e 'test_mxfp8_linear\.py' \
-e 'test_mxfp8_dispatch_returns_mxfp8_method' \
-e 'test_mxfp8_flashinfer_call_contract' \
-e 'test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer' \
-e 'test_mxfp8_rejects_unknown_backend' \
-e 'test_mxfp8_flashinfer_decode_graph_matches_native' \
. || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 5990
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== list entries that could include the test by directory ==="
rg -n --glob '*.{yml,yaml,txt,json}' \
-e '(^|[[:space:]/])-?[[:space:]]*unittest/_torch($|/|[[:space:]])' \
-e '(^|[[:space:]/])-?[[:space:]]*unittest/_torch/modules($|/|[[:space:]])' \
-e 'unittest/_torch/modules/test_mxfp8_linear\.py' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa || true
echo "=== relevant test-list membership summary ==="
python3 - <<'PY'
from pathlib import Path
root = Path("tests/integration/test_lists")
needles = (
"unittest/_torch/modules/test_mxfp8_linear.py",
"unittest/_torch/modules",
"unittest/_torch",
)
for path in sorted(root.glob("**/*")):
if path.suffix not in {".yml", ".yaml", ".txt", ".json"}:
continue
text = path.read_text(errors="replace")
hits = [line.strip() for line in text.splitlines()
if any(needle in line for needle in needles)]
if hits:
print(path)
for line in hits:
print(f" {line}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n --glob '*.{yml,yaml,txt,json}' \
-e 'unittest/_torch/modules/test_mxfp8_linear\.py' \
-e 'unittest/_torch/modules([[:space:]]|$|/)' \
-e 'unittest/_torch([[:space:]]|$|/)' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
target = Path("tests/unittest/_torch/modules/test_mxfp8_linear.py")
target_selector = "unittest/_torch/modules/test_mxfp8_linear.py"
list_root = Path("tests/integration/test_lists")
print("=== covering entries for target test file ===")
found = False
for path in sorted(list_root.glob("**/*")):
if path.suffix not in {".yml", ".yaml", ".txt", ".json"}:
continue
for line_number, line in enumerate(
path.read_text(errors="replace").splitlines(), 1
):
match = re.search(r"unittest/_torch/modules(?:/[^ \t#'\"]*)?", line)
if not match:
continue
entry = match.group(0).rstrip(",")
entry_path = Path(entry.split("::", 1)[0])
if entry_path == Path(target_selector) or (
entry_path != Path(".")
and target_selector.startswith(str(entry_path).rstrip("/") + "/")
):
print(f"{path}:{line_number}: {line.strip()}")
found = True
print("COVERED" if found else "NOT_COVERED")
print("=== exact selector references ===")
names = [
"test_mxfp8_dispatch_returns_mxfp8_method",
"test_mxfp8_flashinfer_call_contract",
"test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer",
"test_mxfp8_rejects_unknown_backend",
"test_mxfp8_flashinfer_decode_graph_matches_native",
]
for name in names:
locations = []
for path in sorted(list_root.glob("**/*")):
if path.suffix not in {".yml", ".yaml", ".txt", ".json"}:
continue
for line_number, line in enumerate(
path.read_text(errors="replace").splitlines(), 1
):
if name in line:
locations.append(f"{path}:{line_number}")
print(f"{name}: {locations or 'NO_REFERENCE'}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 540
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
target = "unittest/_torch/modules/test_mxfp8_linear.py"
root = Path("tests/integration/test_lists")
covering = []
for path in sorted(root.glob("**/*")):
if path.suffix not in {".yml", ".yaml", ".txt", ".json"}:
continue
for number, line in enumerate(path.read_text(errors="replace").splitlines(), 1):
match = re.search(r"unittest/_torch/modules(?:/[^ \t#'\"]*)?", line)
if match:
entry = match.group(0).split("::", 1)[0].rstrip("/")
if target == entry or target.startswith(entry + "/"):
covering.append(f"{path}:{number}: {line.strip()}")
print("\n".join(covering) if covering else "NOT_COVERED")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 169
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== MXFP8 test markers and skip conditions ==="
sed -n '1,280p' tests/unittest/_torch/modules/test_mxfp8_linear.py | \
rg -n -C 3 'pytest|skip|cuda|flashinfer|cudagraph|parametrize|backend'
echo "=== nearby module-test list patterns ==="
for file in \
tests/integration/test_lists/test-db/l0_h100.yml \
tests/integration/test_lists/test-db/l0_b200.yml \
tests/integration/test_lists/test-db/l0_b300.yml \
tests/integration/test_lists/qa/llm_function_core.txt; do
echo "--- $file ---"
rg -n -C 2 'test_triton_linear|test_fused_activation_quant|unittest/_torch/modules' "$file" | head -80
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 18200
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,280p' tests/unittest/_torch/modules/test_mxfp8_linear.py | \
rg -n -C 3 'pytest|skip|cuda|flashinfer|cudagraph|parametrize|backend'
for file in \
tests/integration/test_lists/test-db/l0_h100.yml \
tests/integration/test_lists/test-db/l0_b200.yml \
tests/integration/test_lists/test-db/l0_b300.yml; do
echo "--- $file ---"
rg -n -C 2 'test_triton_linear|test_fused_activation_quant|unittest/_torch/modules' "$file" | head -80
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 18050
Add the MXFP8 tests to the test lists.
Changed tests: test_mxfp8_dispatch_returns_mxfp8_method, test_mxfp8_flashinfer_call_contract, test_mxfp8_auto_keeps_eager_native_and_captures_flashinfer, test_mxfp8_rejects_unknown_backend, and test_mxfp8_flashinfer_decode_graph_matches_native. None is listed in CI or QA test lists. Add the file to the appropriate CI lists, such as l0_b200.yml and l0_b300.yml. Add a QA entry if manual QA requires these tests. 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 56 - 65, Add
tests/unittest/_torch/modules/test_mxfp8_linear.py to the appropriate CI test
lists, including l0_b200.yml and l0_b300.yml, so all listed MXFP8 tests run in
CI. Add a QA test-list entry only if these tests require manual QA.
Source: Path instructions
befb30b to
d20592c
Compare
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
d20592c to
554b88e
Compare
|
/bot run |
|
PR_Github #63837 [ run ] triggered by Bot. Commit: |
|
PR_Github #63837 [ run ] completed with state
|
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
3 participants
Dev Engineer Review
mm_mxfp8support for MiniMax-M3 decode CUDA graphs.trtllm,flashinfer, andautothroughTRTLLM_MXFP8_GEMM_BACKEND.QA Engineer Review
test_mxfp8_dispatch_returns_mxfp8_methodtest_mxfp8_flashinfer_call_contracttest_mxfp8_auto_keeps_eager_native_and_captures_flashinfertest_mxfp8_rejects_unknown_backendtest_mxfp8_flashinfer_decode_graph_matches_nativetest_pytorch_model_engine_warmup.pytests/unittest/_torch/modules/test_mxfp8_linear.pyis included in:tests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_b300.ymlDescription
MiniMax-M3 decode CUDA graphs execute 240 MXFP8 linear GEMMs per replay. This change uses autotuned FlashInfer
mm_mxfp8for captured decode GEMMs while reusing TensorRT-LLM's existing quantized activations, weights, and scales without repacking.The automatic path is limited to MiniMax-M3 decode CUDA-graph capture. Eager execution, CTX/prefill, and piecewise prefill graphs retain the native TensorRT-LLM backend.
TRTLLM_MXFP8_GEMM_BACKEND=flashinferis an explicit override and intentionally applies outside decode capture; the native fallback is preserved.Matched GB200 1P1D disaggregated serving (TP4/EP4, attention DP, FP8 KV cache, 8K/1K, concurrency 64): +10.38% total/output throughput and -9.78% median TPOT.
Test Coverage
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.