[None][feat] Enable Marlin NVFP4 on Ada Lovelace#16749
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:
WalkthroughMarlin NVFP4 support now includes Ada SM89 alongside Hopper SM90–SM99. CUDA compilation, kernel guards, Python backend selection, fused MoE checks, CI gating, documentation, and L40S accuracy coverage were updated. ChangesMarlin NVFP4 Ada support
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant NVFP4Runner
participant SMHelper
participant MarlinKernel
NVFP4Runner->>SMHelper: Check SM89-SM99 support
SMHelper->>MarlinKernel: Select Marlin NVFP4 execution
MarlinKernel->>NVFP4Runner: Return execution result
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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py (1)
55-78: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winAlign the MoE eligibility table with the shared SM89–SM99 contract.
_QUANT_SUPPORT_TABLEstill permits only{89, 90}, whileis_nvfp4_marlin_supported_sm()accepts SM89–SM99. On SM91–SM99, generic quantization eligibility can reject this backend beforecan_implement()runs. Use the same complete supported range and update the error text accordingly.🤖 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/modules/fused_moe/fused_moe_marlin.py` around lines 55 - 78, Update _QUANT_SUPPORT_TABLE’s NVFP4 sm_constraint to include the full SM89–SM99 range accepted by is_nvfp4_marlin_supported_sm(), and revise the can_implement() unsupported-architecture error text to describe SM89–SM99 instead of only Ada and Hopper.tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)
1222-1239: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocument
marlinas a valid backend.The implementation accepts
marlin, and this PR’s callers use it, but theValid backendsdocumentation still omits it. Add'marlin'to that list.🤖 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 1222 - 1239, Update the allowed_backends documentation in the FP4 operation’s docstring to include 'marlin' in the Valid backends list, alongside the existing backend names. Do not change the implementation or other backend behavior.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)
1156-1162: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse the shared Marlin SM predicate.
Use
is_nvfp4_marlin_supported_sm(sm_version)here instead of duplicating the numeric range. The fused MoE path and other checks already use the helper, so this keeps future architecture changes consistent.🤖 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 1156 - 1162, Update the Marlin selection condition in the backend-selection logic to call is_nvfp4_marlin_supported_sm(sm_version) instead of checking the hardcoded 89–99 range, while preserving the existing allowed_backends requirement and tactic assignment.
🤖 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/kernels/marlin/marlin_nvfp4.h`:
- Around line 28-31: Keep the shared isMarlinNvfp4SmSupported definition in
marlin_nvfp4.h and remove the identical per-file definitions from
marlin_nvfp4_gemm.cu, marlin_nvfp4_moe_gemm.cu, and marlin_repack.cu, leaving
their callers unchanged.
---
Outside diff comments:
In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`:
- Around line 1222-1239: Update the allowed_backends documentation in the FP4
operation’s docstring to include 'marlin' in the Valid backends list, alongside
the existing backend names. Do not change the implementation or other backend
behavior.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py`:
- Around line 55-78: Update _QUANT_SUPPORT_TABLE’s NVFP4 sm_constraint to
include the full SM89–SM99 range accepted by is_nvfp4_marlin_supported_sm(), and
revise the can_implement() unsupported-architecture error text to describe
SM89–SM99 instead of only Ada and Hopper.
---
Nitpick comments:
In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`:
- Around line 1156-1162: Update the Marlin selection condition in the
backend-selection logic to call is_nvfp4_marlin_supported_sm(sm_version) instead
of checking the hardcoded 89–99 range, while preserving the existing
allowed_backends requirement and tactic assignment.
🪄 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: c484fbb0-f31b-4e58-94d1-805000e70898
📒 Files selected for processing (18)
cpp/tensorrt_llm/kernels/CMakeLists.txtcpp/tensorrt_llm/kernels/marlin/marlin.cuhcpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.hcpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cucpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cucpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.hcpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.hcpp/tensorrt_llm/kernels/marlin/marlin_repack.cutensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.pytensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/utils.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/test-db/l0_l40s.ymltests/unittest/_torch/modules/moe/moe_test_utils.pytests/unittest/_torch/thop/parallel/test_fp4_linear.pytests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py
|
I noticed there is [perf] tag in the PR title but only added functionality test case. Do we need add perf coverage for it ? |
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
Marlin NVFP4 enablement is cleanly gated: SM guards widen 90→{89,90} only inside Marlin-specific files, Marlin stays explicit opt-in (AUTO/default NVFP4 never selects it), and the device macro/CMake keep Hopper(SM90) and Blackwell(SM100+) unchanged. LGTM.
9a9d74a to
8147490
Compare
rmeghwal-nv
changed the title
Good point. This PR is mainly architecture enablement with functional coverage, not a measured performance change. I updated the title from [perf] to [feat]. |
| self.allowed_backends) > 0, "No allowed backends available" | ||
| sm_version = get_sm_version() | ||
| if "marlin" in self.allowed_backends and 90 <= sm_version <= 99: | ||
| if "marlin" in self.allowed_backends and 89 <= sm_version <= 99: |
There was a problem hiding this comment.
I think we can utilize from tensorrt_llm._torch.utils import is_nvfp4_marlin_supported_sm and use this to check whether Marlin is supported based on the hardware compute capability in this file.
| @skip_pre_ada | ||
| @skip_post_hopper | ||
| @pytest.mark.skip_less_device_memory(32000) | ||
| def test_nvfp4_marlin_single_gpu(self): | ||
| with LLM( | ||
| f"{llm_models_root()}/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", | ||
| kv_cache_config=KvCacheConfig( | ||
| enable_block_reuse=False, | ||
| mamba_ssm_cache_dtype="float16", | ||
| free_gpu_memory_fraction=0.5, | ||
| ), | ||
| tensor_parallel_size=1, | ||
| moe_expert_parallel_size=1, | ||
| max_batch_size=8, | ||
| moe_config=MoeConfig(backend="MARLIN"), | ||
| nvfp4_gemm_config={"allowed_backends": ["marlin"]}, | ||
| ) as llm: | ||
| task = MMLU(self.MODEL_NAME) | ||
| task.evaluate(llm, | ||
| extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) | ||
|
|
There was a problem hiding this comment.
Do you think we can merge this with the test_nvfp4_marlin_multi_gpus above. We can add a parametrize indicating the number of GPUs, as the code looks similar in these 2 test functions.
There was a problem hiding this comment.
I have removed this test, and updated old test only ( changed skip code)
| # (E4M3 input, BF16 output, paged_kv, head_dim=128, sm_89) FMHA cubin. | ||
| - accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[fp8-1-attn_dp_off-flashinfer] | ||
| # PyTorch: force NVFP4 Marlin on Ada to cover the SM89-compiled Marlin kernels. | ||
| - accuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_nvfp4_marlin_single_gpu |
There was a problem hiding this comment.
Update the test name here if you decide to merge the test functions.
| @pytest.mark.skipif( | ||
| get_sm_version() < 90 or get_sm_version() >= 100, | ||
| reason="Marlin NVFP4 backend runs Hopper", | ||
| not (89 <= get_sm_version() < 100), |
There was a problem hiding this comment.
Similarly, we can use is_nvfp4_marlin_supported_sm to check the Marlin and SM match so no hard-coded SM number everywhere.
| quant_config = override_quant_config | ||
| layer_prefix = f"[layer_idx={layer_idx}] " if layer_idx is not None else "" | ||
| if moe_backend.upper() == "MARLIN": | ||
| # Marlin MoE is a Hopper-specific NVFP4 W4A16 backend. Require nvfp4 |
There was a problem hiding this comment.
Please help to update the MOE_DEVELOPER_GUIDE.md.
5f35908 to
ad72494
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/fused_moe/MOE_DEVELOPER_GUIDE.md`:
- Line 153: Align both Marlin NVFP4 hardware entries with
MarlinFusedMoE.can_implement by changing the advertised range from SM89-SM99 to
SM89-SM90. Update the entries at
tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md lines 153-153 and
206-206; no runtime predicate change is needed.
🪄 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: 2bcd77ac-6f53-45b7-8e52-27c083ea8b9e
📒 Files selected for processing (19)
cpp/tensorrt_llm/kernels/CMakeLists.txtcpp/tensorrt_llm/kernels/marlin/marlin.cuhcpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.hcpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cucpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cucpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.hcpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.hcpp/tensorrt_llm/kernels/marlin/marlin_repack.cutensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.pytensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/utils.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/test-db/l0_l40s.ymltests/unittest/_torch/modules/moe/moe_test_utils.pytests/unittest/_torch/thop/parallel/test_fp4_linear.pytests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py
🚧 Files skipped from review as they are similar to previous changes (18)
- cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h
- cpp/tensorrt_llm/kernels/marlin/marlin.cuh
- tensorrt_llm/_torch/modules/fused_moe/create_moe.py
- cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h
- cpp/tensorrt_llm/kernels/CMakeLists.txt
- cpp/tensorrt_llm/kernels/marlin/marlin_repack.cu
- tests/integration/test_lists/test-db/l0_l40s.yml
- tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py
- tests/integration/defs/accuracy/test_llm_api_pytorch.py
- cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu
- cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h
- cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu
- tests/unittest/_torch/modules/moe/moe_test_utils.py
- tensorrt_llm/_torch/modules/linear.py
- tensorrt_llm/_torch/utils.py
- tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py
- tests/unittest/_torch/thop/parallel/test_fp4_linear.py
- tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
| | `mega_moe/mega_moe_deepgemm.py` | `MegaMoEDeepGemm` | SM100/SM103 | W4A8_MXFP4_MXFP8 via DeepGEMM `fp8_fp4_mega_moe` fused dispatch+GEMM+act+GEMM+combine kernel; requires `hidden_size % 512 == 0` | `FUSED_COMM` | | ||
| | `mega_moe/mega_moe_cute_dsl.py` | `MegaMoECuteDsl` | SM100/SM103 | NVFP4 via ported CuteDSL `Sm100MegaMoEKernel` fused dispatch+FC1+act+FC2+combine kernel; requires CUDA 13 Cutlass DSL runtime (PR #14354) and NVSHMEM provider (hard gate); threads per-expert `fc31_alpha`/`fc2_alpha`/`fc1_norm_const` through the kernel ABI and supports SwiGLU clamp via `swiglu_limit`; default deepgemm graph (topk score folded before fc1-out quant, host `combine_output.sum(dim=1)`) | `FUSED_COMM` | | ||
| | `fused_moe_marlin.py` | `MarlinFusedMoE` | SM90 only | W4A16 NVFP4 on Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); supports attention-DP + EP via external comm (scheduler precomputes routing; dispatch payload is plain BF16, no activation scales); non-NVFP4 layers (e.g. unquantized MTP draft layers) fall back to Cutlass in `get_moe_cls`; no dynamic EPLB | `EXTERNAL_COMM` | | ||
| | `fused_moe_marlin.py` | `MarlinFusedMoE` | SM89-SM99 | W4A16 NVFP4 on Ada/Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); supports attention-DP + EP via external comm (scheduler precomputes routing; dispatch payload is plain BF16, no activation scales); non-NVFP4 layers (e.g. unquantized MTP draft layers) fall back to Cutlass in `get_moe_cls`; no dynamic EPLB | `EXTERNAL_COMM` | |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Keep both capability tables consistent with the runtime gate.
The documentation now advertises Marlin NVFP4 on SM89-SM99, while the supplied MarlinFusedMoE.can_implement implementation accepts only SM89 and SM90. Update both entries to SM89-SM90, or expand the runtime predicate before documenting SM91–SM99 support.
tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md#L153-L153: change the backend hardware column toSM89-SM90.tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md#L206-L206: change the Marlin NVFP4 hardware entry toSM89-SM90.
📍 Affects 1 file
tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md#L153-L153(this comment)tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md#L206-L206
🤖 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/modules/fused_moe/MOE_DEVELOPER_GUIDE.md` at line 153,
Align both Marlin NVFP4 hardware entries with MarlinFusedMoE.can_implement by
changing the advertised range from SM89-SM99 to SM89-SM90. Update the entries at
tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md lines 153-153 and
206-206; no runtime predicate change is needed.
|
/bot run --disable-fail-fast |
|
PR_Github #61607 [ run ] triggered by Bot. Commit: |
|
PR_Github #61607 [ run ] completed with state
|
|
@rmeghwal-nv Please re-run the |
…ion to Hopper Extend Marlin NVFP4 GEMM, MoE, and repack kernels from Hopper-only (SM90) to also target Ada (SM89), with shared SM gating in C++ and Python. Signed-off-by: Rameshwar Meghwal <rmeghwal@nvidia.com>
ad72494 to
209418b
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/modules/moe/moe_test_utils.py (1)
1465-1474: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd boundary coverage for the generalized SM gate.
Please test SM89, SM90/SM99, unsupported SM88/SM100, and non-NVFP4 MARLIN configurations; this logic controls whether the entire MARLIN CI matrix runs.
🤖 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/moe/moe_test_utils.py` around lines 1465 - 1474, Add focused tests for the MARLIN CI gating logic around the helper containing the “Rule 0a” checks: cover supported SM89, SM90, and SM99 values, reject unsupported SM88 and SM100 values, and verify non-NVFP4 quantization is rejected regardless of SM version. Use mocks or existing test hooks for get_sm_version and is_nvfp4_marlin_supported_sm, and assert the returned skip messages for each case.
🤖 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 `@tests/unittest/_torch/modules/moe/moe_test_utils.py`:
- Around line 1465-1474: Update the surrounding CI-rule comments and docstring
for the MARLIN backend to consistently describe NVFP4 support as SM89–SM99,
including replacing any wording that says Hopper-only or “Hopper (SM90)”. Leave
the existing gating logic unchanged.
---
Nitpick comments:
In `@tests/unittest/_torch/modules/moe/moe_test_utils.py`:
- Around line 1465-1474: Add focused tests for the MARLIN CI gating logic around
the helper containing the “Rule 0a” checks: cover supported SM89, SM90, and SM99
values, reject unsupported SM88 and SM100 values, and verify non-NVFP4
quantization is rejected regardless of SM version. Use mocks or existing test
hooks for get_sm_version and is_nvfp4_marlin_supported_sm, and assert the
returned skip messages for each case.
🪄 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: 80010470-7e96-4e46-bf2a-f02e324a2560
📒 Files selected for processing (19)
cpp/tensorrt_llm/kernels/CMakeLists.txtcpp/tensorrt_llm/kernels/marlin/marlin.cuhcpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.hcpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cucpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cucpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.hcpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.hcpp/tensorrt_llm/kernels/marlin/marlin_repack.cutensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.pytensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/utils.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/test-db/l0_l40s.ymltests/unittest/_torch/modules/moe/moe_test_utils.pytests/unittest/_torch/thop/parallel/test_fp4_linear.pytests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py
🚧 Files skipped from review as they are similar to previous changes (18)
- cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu
- cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h
- tests/integration/test_lists/test-db/l0_l40s.yml
- cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h
- cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h
- tensorrt_llm/_torch/modules/fused_moe/create_moe.py
- tensorrt_llm/_torch/modules/linear.py
- tensorrt_llm/_torch/utils.py
- tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py
- tests/unittest/_torch/thop/parallel/test_fp4_linear.py
- tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md
- cpp/tensorrt_llm/kernels/marlin/marlin.cuh
- tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py
- cpp/tensorrt_llm/kernels/marlin/marlin_repack.cu
- cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu
- cpp/tensorrt_llm/kernels/CMakeLists.txt
- tests/integration/defs/accuracy/test_llm_api_pytorch.py
- tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
| # --- Rule 0a: MARLIN backend only runs NVFP4 on Ada/Hopper (SM89/SM90) --- | ||
| if backend_type == MoeBackendType.MARLIN: | ||
| from tensorrt_llm._torch.utils import is_nvfp4_marlin_supported_sm | ||
| from tensorrt_llm._utils import get_sm_version | ||
|
|
||
| if quant_algo != QuantAlgo.NVFP4: | ||
| return f"[CI accel] MARLIN only tests NVFP4 in CI (got {quant_algo})" | ||
| sm_version = get_sm_version() | ||
| if sm_version != 90: | ||
| return f"[CI accel] MARLIN only runs on Hopper (SM90) in CI (got SM{sm_version})" | ||
| if not is_nvfp4_marlin_supported_sm(sm_version): | ||
| return f"[CI accel] MARLIN only runs on Ada/Hopper (SM89-99) in CI (got SM{sm_version})" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the surrounding CI-rule documentation.
The implementation accepts every SM version from 89 through 99, but the nearby comments/docstring still describe Hopper-only or SM89/SM90 support. Align them with the actual SM89-SM99 contract to prevent future gating mistakes.
Proposed wording
- # --- Rule 0a: MARLIN backend only runs NVFP4 on Ada/Hopper (SM89/SM90) ---
+ # --- Rule 0a: MARLIN backend only runs NVFP4 on Ada/Hopper (SM89-SM99) ---Also update the rule description above that still says “Hopper (SM90)”.
📝 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.
| # --- Rule 0a: MARLIN backend only runs NVFP4 on Ada/Hopper (SM89/SM90) --- | |
| if backend_type == MoeBackendType.MARLIN: | |
| from tensorrt_llm._torch.utils import is_nvfp4_marlin_supported_sm | |
| from tensorrt_llm._utils import get_sm_version | |
| if quant_algo != QuantAlgo.NVFP4: | |
| return f"[CI accel] MARLIN only tests NVFP4 in CI (got {quant_algo})" | |
| sm_version = get_sm_version() | |
| if sm_version != 90: | |
| return f"[CI accel] MARLIN only runs on Hopper (SM90) in CI (got SM{sm_version})" | |
| if not is_nvfp4_marlin_supported_sm(sm_version): | |
| return f"[CI accel] MARLIN only runs on Ada/Hopper (SM89-99) in CI (got SM{sm_version})" | |
| # --- Rule 0a: MARLIN backend only runs NVFP4 on Ada/Hopper (SM89-SM99) --- | |
| if backend_type == MoeBackendType.MARLIN: | |
| from tensorrt_llm._torch.utils import is_nvfp4_marlin_supported_sm | |
| from tensorrt_llm._utils import get_sm_version | |
| if quant_algo != QuantAlgo.NVFP4: | |
| return f"[CI accel] MARLIN only tests NVFP4 in CI (got {quant_algo})" | |
| sm_version = get_sm_version() | |
| if not is_nvfp4_marlin_supported_sm(sm_version): | |
| return f"[CI accel] MARLIN only runs on Ada/Hopper (SM89-99) in CI (got SM{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 `@tests/unittest/_torch/modules/moe/moe_test_utils.py` around lines 1465 -
1474, Update the surrounding CI-rule comments and docstring for the MARLIN
backend to consistently describe NVFP4 support as SM89–SM99, including replacing
any wording that says Hopper-only or “Hopper (SM90)”. Leave the existing gating
logic unchanged.
|
/bot run --disable-fail-fast |
|
PR_Github #61692 [ run ] triggered by Bot. Commit: |
6 participants
Description
Extend the existing opt-in Marlin NVFP4 path from Hopper-only coverage to also support Ada Lovelace / L40S (SM89), while preserving Hopper support.
This change compiles the Marlin NVFP4 kernels for SM89 and SM90, adds shared Ada/Hopper SM support gates in C++ and Python, and updates runtime checks for Marlin GEMM, fused MoE GEMM, and repack paths.
Marlin remains strictly opt-in. Users select it through
moe_config.backend: MARLINfor MoE layers andnvfp4_gemm_config.allowed_backends: ["marlin"]for dense / Mamba Linear layers. Default Blackwell paths and non-Marlin fallback backends remain unchanged.This PR also adds focused 1-GPU L40S coverage by reusing the existing Nemotron 3 Nano NVFP4 Marlin accuracy test with
tp_size=1.Test Coverage
Added / updated coverage:
tp_size=1:accuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_nvfp4_marlin_multi_gpus[tp_size=1]tp_size=1parametrization to the L40S pre-merge PyTorch lane:tests/integration/test_lists/test-db/l0_l40s.ymlValidation run locally:
build_wheel_targetsintensorrt_llm/devel:devel-fresh-l40s.cuobjdumpfor:marlin_nvfp4_gemm.cu.omarlin_nvfp4_moe_gemm.cu.omarlin_repack.cu.obuild/tensorrt_llm-1.3.0rc21-cp312-cp312-linux_x86_64.whltests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.pymax_batch_size=32.max_batch_size=128.87.5on 8 samples.python3 scripts/check_test_list.py --validate --test-lists-dir tests/integration/test_lists --test-base-dir tests/integration/defsOK: 2028 unique test entries validated.git diff --check.Notes:
tp_size=1because the currentl0_l40s.ymllane issystem_gpu_count: gte: 1, lte: 1.tp_size=8.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.Dev Engineer Review
MARLIN_NVFP4_DEVICE_SUPPORTED) and shared runtime helper (marlin_nvfp4::isMarlinNvfp4SmSupported(int sm)).is_nvfp4_marlin_supported_sm()and updatedis_nvfp4_marlin_enabled()gating.marlin_srcOBJECT library configuration).moe_config.backend: MARLINand NVFP4 gemmallowed_backends: ["marlin"]); broadened only the supported SM range under those explicit opt-in conditions.QA Engineer Review
tests/integration/defs/accuracy/test_llm_api_pytorch.pyTestNemotronV3Nano.test_nvfp4_marlin_multi_gpus: broadened skip gating to Ada, adjusted memory skip threshold, added MPI-world-size-aware skipping.tests/unittest/_torch/thop/parallel/test_fp4_linear.pytest_fp4_linear_marlin: updated skip logic to use the shared Ada/Hopper SM support predicate.tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.pyTestNvfp4MarlinGemm(decorators/skip gating): switched to Ada-or-Hopper-aware skip rule.tests/unittest/_torch/modules/moe/moe_test_utils.pytests/integration/test_lists/test-db/l0_l40s.ymlaccuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_nvfp4_marlin_multi_gpus[tp_size=1]under the L40S single-GPU pre-merge PyTorch list.test-db/l0_l40s.ymlvia the addedTestNemotronV3Nano::test_nvfp4_marlin_multi_gpus[tp_size=1]entry.