[https://nvbugs/6445494][infra] Upgrade Triton to 3.7.1 for Torch 2.12.1 compatibility - #16303
[https://nvbugs/6445494][infra] Upgrade Triton to 3.7.1 for Torch 2.12.1 compatibility#16303HandongLi-01 wants to merge 4 commits into
Conversation
|
/bot run |
HandongLi-01
changed the title
|
PR_Github #58911 [ run ] triggered by Bot. Commit: |
|
PR_Github #58911 [ run ] completed with state |
|
/bot run |
|
PR_Github #58918 [ run ] triggered by Bot. Commit: |
|
/bot run |
|
PR_Github #58939 [ run ] triggered by Bot. Commit: |
|
PR_Github #58918 [ run ] completed with state |
|
PR_Github #58939 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59078 [ run ] triggered by Bot. Commit: |
|
PR_Github #59078 [ run ] completed with state
|
bb6d6c4 to
fe98fc6
Compare
|
/bot run |
|
PR_Github #59653 [ run ] triggered by Bot. Commit: |
|
PR_Github #59653 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59675 [ run ] triggered by Bot. Commit: |
|
PR_Github #59675 [ run ] completed with state
|
392e437 to
6d76449
Compare
|
/bot run |
6d76449 to
62e992e
Compare
|
/bot run |
HandongLi-01
changed the title
|
PR_Github #60281 [ run ] triggered by Bot. Commit: |
|
PR_Github #60281 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60847 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61199 [ run ] triggered by Bot. Commit: |
|
PR_Github #61199 [ run ] completed with state
|
f5d35b2 to
92dc585
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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_fused_moe.py (1)
2593-2633: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd a test-list entry for
test_fused_moe_triton_mxfp4The parametrization already covers unaligned sizes (INTERMEDIATE_SIZE=720and360), but there’s no matching entry undertests/integration/test_lists/test-db/ortests/integration/test_lists/qa/.🤖 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_fused_moe.py` around lines 2593 - 2633, The test-list configuration is missing an entry for test_fused_moe_triton_mxfp4. Add this test to the appropriate test-db and QA lists under the existing fused-MoE test-list conventions, preserving its parameterization for INTERMEDIATE_SIZE values 720 and 360.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_fused_moe.py`:
- Around line 2593-2633: The test-list configuration is missing an entry for
test_fused_moe_triton_mxfp4. Add this test to the appropriate test-db and QA
lists under the existing fused-MoE test-list conventions, preserving its
parameterization for INTERMEDIATE_SIZE values 720 and 360.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e7b5f975-334a-4f4e-bfb3-1821adcf6cda
📒 Files selected for processing (6)
requirements.txttensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.pytensorrt_llm/_torch/modules/triton_linear.pytests/unittest/_torch/modules/test_fused_moe.pytests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py
🚧 Files skipped from review as they are similar to previous changes (5)
- requirements.txt
- tensorrt_llm/_torch/modules/triton_linear.py
- tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py
- tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py
- tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py
c17dc45 to
817d2d6
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@docker/common/install_pytorch.sh`:
- Line 8: Update the PyPI install command in install_pytorch.sh to pin
torchvision to version 0.27.1 instead of leaving it unversioned. Preserve the
existing torch version variable and all other installation arguments.
In `@triton_kernels/matmul_details/opt_flags.py`:
- Around line 241-271: Update the bf16 exclusion in the conditional heuristic
after compute_block_k to compare lhs_dtype against Triton’s BF16 DataType
constant rather than torch.bfloat16, preserving the intended behavior of
skipping the swap for bf16 inputs.
🪄 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: 597c47dd-843e-46b6-8c3b-f6504280aaf9
⛔ Files ignored due to path filters (1)
security_scanning/poetry.lockis excluded by!**/*.lock
📒 Files selected for processing (66)
ATTRIBUTIONS-Python.mddocker/common/install_pytorch.shdocs/source/installation/installation-guide.mdjenkins/L0_Test.groovyrequirements.txtscripts/vendor_triton_kernels.pysecurity_scanning/pyproject.tomltensorrt_llm/_torch/auto_deploy/custom_ops/README.mdtensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.pytensorrt_llm/_torch/modules/triton_linear.pytensorrt_llm/tools/profiler/nsys_profile_tools/trtllm_engine_model.jsontests/unittest/_torch/modules/fused_moe/test_triton_mxfp4_swizzle.pytests/unittest/_torch/modules/test_fused_moe.pytests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.pytriton_kernels/README.mdtriton_kernels/VERSIONtriton_kernels/__init__.pytriton_kernels/compaction.pytriton_kernels/compaction_details/_masked_compaction.pytriton_kernels/distributed.pytriton_kernels/distributed_details/__init__.pytriton_kernels/distributed_details/mesh.pytriton_kernels/matmul.pytriton_kernels/matmul_details/__init__.pytriton_kernels/matmul_details/_common.pytriton_kernels/matmul_details/_matmul.pytriton_kernels/matmul_details/_p_matmul.pytriton_kernels/matmul_details/opt_flags.pytriton_kernels/matmul_details/opt_flags_details/__init__.pytriton_kernels/matmul_details/opt_flags_details/opt_flags_amd.pytriton_kernels/matmul_details/opt_flags_details/opt_flags_nvidia.pytriton_kernels/matmul_ogs.pytriton_kernels/meta.pytriton_kernels/numerics.pytriton_kernels/numerics_details/__init__.pytriton_kernels/numerics_details/flexpoint.pytriton_kernels/numerics_details/mxfp.pytriton_kernels/numerics_details/mxfp_details/_downcast_to_mxfp.pytriton_kernels/numerics_details/mxfp_details/_upcast_from_mxfp.pytriton_kernels/proton_opts.pytriton_kernels/reduce.pytriton_kernels/roofline.pytriton_kernels/specialize.pytriton_kernels/swiglu.pytriton_kernels/swiglu_details/_swiglu.pytriton_kernels/target_info.pytriton_kernels/tensor.pytriton_kernels/tensor_details/bitmatrix.pytriton_kernels/tensor_details/bitmatrix_details/sum_bitmatrix_rows.pytriton_kernels/tensor_details/dtype.pytriton_kernels/tensor_details/layout.pytriton_kernels/tensor_details/layout_details/base.pytriton_kernels/tensor_details/layout_details/blackwell_scale.pytriton_kernels/tensor_details/layout_details/blackwell_value.pytriton_kernels/tensor_details/layout_details/cdna4_scale.pytriton_kernels/tensor_details/layout_details/hopper_scale.pytriton_kernels/tensor_details/layout_details/hopper_value.pytriton_kernels/tensor_details/layout_details/strided.pytriton_kernels/tensor_details/layout_details/torch_utils.pytriton_kernels/tensor_details/ragged_tensor.pytriton_kernels/testing.pytriton_kernels/topk.pytriton_kernels/topk_details/__init__.pytriton_kernels/topk_details/_topk_backward.pytriton_kernels/topk_details/_topk_forward.py
💤 Files with no reviewable changes (1)
- triton_kernels/matmul_ogs.py
🚧 Files skipped from review as they are similar to previous changes (43)
- triton_kernels/VERSION
- docs/source/installation/installation-guide.md
- triton_kernels/topk_details/_topk_backward.py
- triton_kernels/swiglu.py
- triton_kernels/numerics_details/init.py
- triton_kernels/tensor_details/bitmatrix_details/sum_bitmatrix_rows.py
- triton_kernels/distributed_details/init.py
- triton_kernels/compaction_details/_masked_compaction.py
- triton_kernels/matmul_details/init.py
- triton_kernels/README.md
- triton_kernels/tensor_details/bitmatrix.py
- ATTRIBUTIONS-Python.md
- triton_kernels/init.py
- triton_kernels/matmul_details/opt_flags_details/init.py
- triton_kernels/meta.py
- triton_kernels/numerics_details/mxfp_details/_upcast_from_mxfp.py
- triton_kernels/numerics_details/flexpoint.py
- tensorrt_llm/_torch/auto_deploy/custom_ops/README.md
- triton_kernels/tensor_details/dtype.py
- security_scanning/pyproject.toml
- triton_kernels/target_info.py
- triton_kernels/numerics.py
- triton_kernels/specialize.py
- tests/unittest/_torch/modules/test_fused_moe.py
- triton_kernels/tensor_details/layout.py
- triton_kernels/topk_details/_topk_forward.py
- triton_kernels/tensor_details/layout_details/base.py
- triton_kernels/compaction.py
- tests/unittest/_torch/modules/fused_moe/test_triton_mxfp4_swizzle.py
- triton_kernels/tensor_details/layout_details/torch_utils.py
- requirements.txt
- tensorrt_llm/tools/profiler/nsys_profile_tools/trtllm_engine_model.json
- triton_kernels/matmul_details/opt_flags_details/opt_flags_amd.py
- triton_kernels/roofline.py
- jenkins/L0_Test.groovy
- triton_kernels/tensor_details/layout_details/strided.py
- triton_kernels/distributed.py
- triton_kernels/tensor_details/ragged_tensor.py
- triton_kernels/tensor_details/layout_details/hopper_scale.py
- triton_kernels/numerics_details/mxfp.py
- triton_kernels/matmul_details/opt_flags_details/opt_flags_nvidia.py
- triton_kernels/tensor_details/layout_details/blackwell_scale.py
- tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py
| # and closest to the version specified in | ||
| # https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-05.html#rel-26-05 | ||
| TORCH_VERSION="2.11.0" | ||
| TORCH_VERSION="2.12.1" |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== file outline ==\n'
ast-grep outline docker/common/install_pytorch.sh --view expanded || true
printf '\n== relevant lines ==\n'
cat -n docker/common/install_pytorch.sh | sed -n '1,180p'Repository: NVIDIA/TensorRT-LLM
Length of output: 3652
🏁 Script executed:
set -euo pipefail
printf '\n== file outline ==\n'
ast-grep outline docker/common/install_pytorch.sh --view expanded || true
printf '\n== relevant lines ==\n'
cat -n docker/common/install_pytorch.sh | sed -n '1,180p'Repository: NVIDIA/TensorRT-LLM
Length of output: 3652
🌐 Web query:
PyTorch torchvision compatibility matrix torch 2.12.1 torchvision version 2025
💡 Result:
For PyTorch version 2.12.1, the compatible version of torchvision is 0.27.1 [1][2]. According to the official PyTorch installation documentation, you should use the following command to install these compatible versions [2]: pip install torch==2.12.1 torchvision==0.27.1 The torchvision 0.27.x series is designed to support PyTorch 2.12 [3][4][5]. Torchvision 0.27.1 was released on June 17, 2026, to accompany the PyTorch 2.12.1 release [1][5]. This version supports Python versions 3.10 through 3.14 [3][4][5].
Citations:
- 1: https://dev-discuss.pytorch.org/t/pytorch-2-12-1-general-availability/3402
- 2: https://pytorch.org/get-started/previous-versions/
- 3: https://github.com/PyTorch/vision
- 4: https://github.com/pytorch/vision?tab=readme-ov-file
- 5: https://pypi.org/project/torchvision/
Pin torchvision in the PyPI install path. The source-build flow already resolves the matching vision branch, but pip3 install torch==${TORCH_VERSION} torchvision ... can drift as the index changes. Use torchvision==0.27.1 here to keep the image reproducible.
🤖 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 `@docker/common/install_pytorch.sh` at line 8, Update the PyPI install command
in install_pytorch.sh to pin torchvision to version 0.27.1 instead of leaving it
unversioned. Preserve the existing torch version variable and all other
installation arguments.
Source: MCP tools
| a_mx_scale_layout = None if not isinstance(precision_config.a_mx_scale, Tensor) else precision_config.a_mx_scale.storage.layout | ||
| b_mx_scale_layout = None if not isinstance(precision_config.b_mx_scale, Tensor) else precision_config.b_mx_scale.storage.layout | ||
|
|
||
| def _is_layout_strided(layout: Layout | None) -> bool: | ||
| return layout is None or isinstance(layout, StridedLayout) | ||
|
|
||
| requires_persistent = (not _is_layout_strided(a_mx_scale_layout) or not _is_layout_strided(b_mx_scale_layout)) and target_info.has_native_mxfp() | ||
| if constraints.get("is_persistent", None) is not None: | ||
| is_persistent = constraints["is_persistent"] | ||
| elif requires_persistent: | ||
| assert supports_persistent, "persistent kernel required but not supported" | ||
| is_persistent = True | ||
| else: | ||
| has_simple_epilogue = precision_config.max_num_imprecise_acc is None | ||
| is_persistent = supports_persistent and has_simple_epilogue and (tiles_per_sm >= 2.0 or lhs_dtype.itemsize <= 1) and out_dtype.itemsize < 4 | ||
| is_persistent = supports_persistent and has_simple_epilogue and (tiles_per_sm >= 2.0 or lhs_dtype.bitwidth <= 8) and out_dtype.bitwidth < 32 | ||
| # TMA is slower for batched matmuls with small m/n/k. | ||
| if m * n * k < 131072: | ||
| is_persistent = False | ||
| if isinstance(b_mx_scale_layout, HopperMXScaleLayout): | ||
| # TODO: persistent kernel is currently slower than non-persistent | ||
| is_persistent = False | ||
|
|
||
| # adjust block_n based on is_persistent signal | ||
| block_n = block_n_tma if is_persistent else block_n | ||
| # adjust block_m based on is_persistent signal | ||
| if is_persistent and opt_flags_nvidia.is_x_scale_swizzled(precision_config): | ||
| # a mx scale has been swizzled to BlackwellActMXScaleLayout, enforce block_m=128 to align with swizzling layout | ||
| block_m = 128 | ||
| # block k | ||
| block_k = opt_flags_nvidia.compute_block_k(m, k, is_persistent, lhs_dtype, rhs_dtype, precision_config, has_y_acc_in) | ||
| if block_n == 256 and block_k == 128 and block_m <= 64 and is_persistent and rhs_dtype == FP4 and k >= 4096 and tokens_per_expt > 1 and lhs_dtype != torch.bfloat16: | ||
| if block_n == 256 and block_k == 128 and block_m <= 64 and is_persistent and rhs_dtype == FP4 and k >= 4096 and slice_size > 1 and lhs_dtype != torch.bfloat16: |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Does triton_kernels DType (FloatType/IntegerType) define __eq__ that compares against torch.dtype, in triton-lang/triton triton_kernels v3.7.1?
💡 Result:
In the triton-lang/triton repository, specifically within the triton_kernels package (as observed in the current development codebase), the FloatType and IntegerType classes defined in triton_kernels/tensor_details/dtype.py [1] do not define an eq method that compares against torch.dtype. Instead, the triton_kernels package relies on explicit conversion functions—such as dtype_to_torch_dtype and torch_dtype_to_dtype—to map between Triton's internal tensor data types and PyTorch's torch.dtype [2]. While some test utilities in triton_kernels may implement helper classes that store a torch_dtype attribute for comparison purposes [3], the core FloatType and IntegerType dataclasses do not natively support equality comparison with torch.dtype objects [1]. Comparisons between these types and torch.dtype typically require explicit conversion or mapping [2].
Citations:
- 1: https://github.com/triton-lang/triton/blob/9499e261/python/triton_kernels/triton_kernels/tensor_details/dtype.py
- 2: https://github.com/triton-lang/triton/blob/9499e261/python/triton_kernels/triton_kernels/tensor.py
- 3: https://github.com/triton-lang/triton/blob/9499e261/python/triton_kernels/tests/test_matmul.py
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## Files of interest\n'
git ls-files triton_kernels/matmul_details/opt_flags.py triton_kernels/tensor_details/dtype.py triton_kernels/tensor.py 2>/dev/null || true
printf '\n## Outline: opt_flags.py\n'
ast-grep outline triton_kernels/matmul_details/opt_flags.py --view expanded || true
printf '\n## Outline: dtype.py\n'
ast-grep outline triton_kernels/tensor_details/dtype.py --view expanded || true
printf '\n## Relevant snippets\n'
sed -n '200,290p' triton_kernels/matmul_details/opt_flags.py
printf '\n--- dtype.py ---\n'
sed -n '1,260p' triton_kernels/tensor_details/dtype.py
printf '\n--- tensor.py (conversion helpers) ---\n'
sed -n '1,220p' triton_kernels/tensor.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 17501
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
p = Path("triton_kernels/matmul_details/opt_flags.py")
text = p.read_text()
for needle in ["lhs_dtype != torch.bfloat16", "rhs_dtype == FP4", "a.dtype", "b.dtype"]:
print(f"{needle}: {needle in text}")
print("\nContext around comparison:")
m = re.search(r'.{0,120}lhs_dtype != torch\.bfloat16.{0,120}', text, re.S)
print(m.group(0) if m else "not found")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 540
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
for path in ["triton_kernels/tensor_details/dtype.py", "triton_kernels/tensor.py"]:
p = Path(path)
if not p.exists():
print(f"{path}: missing")
continue
text = p.read_text()
for pattern in [
r'class\s+\w+.*?:',
r'def __eq__',
r'dtype_to_torch_dtype',
r'torch_dtype_to_dtype',
r'torch\.dtype',
]:
print(f"\n[{path}] pattern={pattern}")
for m in re.finditer(pattern, text):
start = max(0, m.start() - 80)
end = min(len(text), m.end() + 180)
print(text[start:end].splitlines()[0:8])
break
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 1933
Compare against BF16 here. lhs_dtype is a Triton DataType, so lhs_dtype != torch.bfloat16 is always true and the bf16 exclusion never triggers. Use the Triton BF16 constant (or convert to a torch.dtype first) so the swap heuristic actually skips bf16 inputs.
🤖 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 `@triton_kernels/matmul_details/opt_flags.py` around lines 241 - 271, Update
the bf16 exclusion in the conditional heuristic after compute_block_k to compare
lhs_dtype against Triton’s BF16 DataType constant rather than torch.bfloat16,
preserving the intended behavior of skipping the swap for bf16 inputs.
Source: Learnings
|
/bot run |
|
PR_Github #61556 [ run ] triggered by Bot. Commit: |
|
PR_Github #61556 [ run ] completed with state
|
Signed-off-by: handongl <15242213+HandongLi-01@users.noreply.github.com>
Signed-off-by: handongl <15242213+HandongLi-01@users.noreply.github.com>
Signed-off-by: handongl <15242213+HandongLi-01@users.noreply.github.com>
817d2d6 to
2f42fa1
Compare
|
/bot run |
|
PR_Github #61778 [ run ] triggered by Bot. Commit: |
HandongLi-01
changed the title
|
PR_Github #61778 [ run ] completed with state
|
Signed-off-by: handongl <15242213+HandongLi-01@users.noreply.github.com>
|
/bot run --stage-list "Build-Docker-Images" |
|
PR_Github #61835 [ run ] triggered by Bot. Commit: |
|
PR_Github #61835 [ run ] completed with state
|
6 participants
Description
Upgrade the public PyTorch stack to Torch 2.12.1 and Triton 3.7.1. Torch 2.12.1 requires
triton==3.7.1on Linux, replacing the previous Torch 2.11.0/Triton 3.6.0 combination. Triton 3.7.1 also includes the missing async-read proxy fence fix for Blackwell.This PR migrates TensorRT-LLM to the Triton 3.7 APIs and re-vendors
triton_kernelsfromv3.7.1.Triton API migration
triton_kernels.matmul_ogswithtriton_kernels.matmul.matmul_ogs()calls tomatmul()and migrate the renameda,b,c, andc_acc_inarguments.RoutingData,GatherIndx, andScatterIndxwrappers withTritonRoutingData,RaggedTensorMetadata, and plain gather/scatter tensors.triton_kernels.reduce, including masking invalid expert-parallel routing slots.PrecisionConfigfields toa_mx_scale,b_mx_scale, andc_mx_scale.StridedLayoutfallbacks.Dependency and CI updates
>=2.12.1,<=2.13.0a0.3.7.1.triton_kernelsfrom the Tritonv3.7.1tag.security_scanning/pyproject.tomlandpoetry.lockwith Poetry 2.4.1.nccl4py>=0.3.1,<0.4dependency introduced onmain.Test Coverage
Relevant coverage includes:
tests/unittest/others/test_triton_kernels_vendoring.pytests/unittest/_torch/modules/test_fused_moe.pytests/unittest/_torch/modules/test_triton_linear.pytests/unittest/_torch/modules/fused_moe/test_triton_mxfp4_swizzle.pytests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.pytests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_triton_moe.pyLocal validation completed:
PR Checklist
Dev Engineer Review
triton/triton-kernelsto 3.7.1 and updated vendored provenance metadata acrosstriton_kernels/*(includingtriton_kernels/VERSION,triton_kernels/README.md, and numerous header “Source” comments).matmul_ogs/routing-precision API to the newtriton_kernels.matmulflex-precision API:triton_kernels/matmul_ogs.py.tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.pyto usematmuland the new routing contract:TritonRoutingData.matmul_ogsfinal/top-k behavior with explicit masked reduction viacombine_expert_outputs()after GEMM2 output reshaping.tensorrt_llm/_torch/modules/triton_linear.pyand MXFP4 MoE fused op (tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py) to usematmuland updated precision/routing fields (notablyb_mx_scale,FlexCtx, and MXFP4 routing-data wiring).triton_kernels/distributed_details/mesh.py) and extended Tritontopk/reduce paths to use it.tensorrt_llm/tools/profiler/nsys_profile_tools/trtllm_engine_model.jsonMoE GEMM kernel matching regexes for additional fused/MoE variants.triton_kernels/reduce.pyAPI to acceptunpadded_batch_size.docker/common/install_pytorch.sh,docs/source/installation/installation-guide.md,requirements.txt,security_scanning/pyproject.toml.pin_torchparameter logic injenkins/L0_Test.groovyand adjusted “Extra PyTorch CUDA 13.2 install” pins accordingly.QA Engineer Review
Test changes
Touched unit/auto-deploy test code under
tests/:tests/unittest/_torch/modules/fused_moe/test_triton_mxfp4_swizzle.pyHopperMX*Layoutinstances intoconvert_layoutandconvert_layout_expert_chunked.tests/unittest/_torch/modules/test_fused_moe.pytest_fused_moe_triton_mxfp4MXFP4 FP32 conversion helpers to satisfy Triton 3.7.x 32-alignment requirements (pad/slice logic added for quant axis handling).tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.pytest_mxfp4_value_layout_uses_strided_layout_on_blackwelltest_mxfp4_value_layout_skips_conversion_on_blackwelltest_mxfp4_value_layout_skips_conversion_for_strided_fallbacktest_mxfp4_scale_layout_skips_conversion_on_blackwelltest_mxfp4_scale_layout_skips_conversion_for_strided_fallbacktest_mxfp4_scale_layout_keeps_hopper_layoutTest list / CI entry coverage
tests/integration/test_lists/,test-db/,qa/, orwaives.txtfor this PR’s diff.