[TRTLLM-15316][feat] sm107 gemm + quant

[farazkh80]

Description

SM107 support for the CUTLASS GEMM dispatch paths and FP8 block-scale quantization kernels.

• SM107 is an SM100-family member and runs the same cutlass::arch::Sm100 kernels, so this is mostly predicate widening: sm == 100 || sm == 103 checks become isSM100Family() (a strict superset — SM100/SM103 select exactly the same paths as before, only SM107 is newly reachable). The build-system half landed in [None][infra] Recognize SM107 (Rubin) in build config and arch detection #17336.
• MXFP8 quantization entry points are templated on SF_VEC_SIZE/SF_OUTPUT_VEC_SIZE; defaults reproduce the previous 32-wide behavior exactly.
• One intentional behavior change: fp8_quantize_1x128_packed_ue8m0 gains use_r128c4_layout=True, switching its default output to the standard R128c4 scale layout. Its only in-tree consumer (deep_gemm) is updated in the same commit to request the legacy layout explicitly, so that path is unchanged.
• No Python model-layer wire-up is included — the SM107 paths become reachable in a follow-up, so this PR is inert on its own.

Test Coverage

test_fp8_quantize.py validates the new R128c4 default (test_fp8_quantize_1x128_packed_ue8m0_r128c4*) and keeps the legacy deep_gemm contract under test_..._legacy_layout. Dispatch changes are covered by the existing FP4 / FP8-rowwise / fpA_intB / FP8 block-scale GEMM suites. Requesting B200 and H100 stages — Hopper shares the fp8BlockScalingGemm dispatch and quantization.* templates, though no Hopper behavior should change.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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

• Expanded SM100-family dispatch and kernel generation to include SM107.
• Preserved existing SM100, SM103, SM120, and SM121 behavior.
• Added SM107 support across CUTLASS, FP4, FP8, FP8 rowwise, FP8 block-scale, and fpA-intB GEMM paths.
• Added SM107 FP4 candidate validation and SM100-family architecture normalization.
• Restricted SM107 generated-kernel selection to Sm100f.
• Parameterized MXFP8 quantization by SF_VEC_SIZE and SF_OUTPUT_VEC_SIZE.
• Preserved 32-wide quantization defaults and added BF16 (128, 32) support.
• Added replicated scale-factor output and optional R128c4 layout support.
• Preserved the legacy fp8_quantize_1x128_packed_ue8m0 layout as the default.
• Updated the C++ API, Torch schema, fake implementation, and Python callers consistently.
• Updated CUDA architecture configuration for SM107.
• Review should verify SM107 compilation, candidate selection, scale-factor bounds, error handling, and legacy-layout compatibility.

QA Engineer Review

Modified or added test functions in tests/unittest/_torch/thop/parallel/test_fp8_quantize.py:

• Renamed and updated test_fp8_quantize_1x128_packed_ue8m0_matches_legacy to test_fp8_quantize_1x128_packed_ue8m0_r128c4.
• Removed test_fp8_quantize_1x128_packed_ue8m0_padded_rows_are_zero.
• Added test_fp8_quantize_1x128_packed_ue8m0_r128c4_zero_blocks.
• Added test_fp8_quantize_1x128_packed_ue8m0_r128c4_padding_is_zero.
• Added test_fp8_quantize_1x128_packed_ue8m0_legacy_layout.
• Updated FP8 block-scale GEMM skip conditions to include SM107.

Coverage in tests/integration/test_lists/ is not established from the available change summary.

Verdict: needs follow-up.

[farazkh80]

/bot run

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 74ba743b-6e2d-4231-91a6-51feed2aa6ac

📥 Commits

Reviewing files that changed from the base of the PR and between 69b8195 and b684ebb.

📒 Files selected for processing (1)

• tests/unittest/_torch/thop/parallel/test_fp8_quantize.py

🚧 Files skipped from review as they are similar to previous changes (1)

• tests/unittest/_torch/thop/parallel/test_fp8_quantize.py

---

Walkthrough

The change broadens SM100-family CUDA build and dispatch handling, including SM107. It adds configurable MXFP8 scale-vector sizes and replication, plus optional R128c4 output while preserving the legacy packed layout by default.

Changes

Blackwell architecture support

Layer / File(s)	Summary
Architecture build and generation cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt, cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py	Shared architecture lists include SM100-family targets, including SM107. Kernel generation covers architectures 100 through 109.
GEMM dispatch, preprocessing, and compatibility cpp/tensorrt_llm/kernels/cutlass_kernels/*, cpp/tensorrt_llm/thop/fp8BlockScalingGemm.cpp, cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/KernelRunner.cpp, tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py	FP4, FP8, FP-A/int-B, rowwise, block-scaling, preprocessing, and TRT-LLM-Gen compatibility paths recognize the SM100 family and SM107.
Heuristic architecture normalization cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp	SM107 FP4 candidates require TMA epilogue schedules. Other unknown SM100-family versions normalize to SM100.

MXFP8 scale-factor layouts

Layer / File(s)	Summary
Configurable MXFP8 quantization kernel cpp/tensorrt_llm/kernels/quantization.*	MXFP8 quantization accepts input and output scale-vector sizes. Scale groups support power-of-two widths through 32 lanes, with replication-aware addressing and stores.
Selectable quantization operator layouts cpp/tensorrt_llm/thop/fp8Quantize.cpp	fp8_quantize_1x128_packed_ue8m0 supports R128c4 and legacy packed layouts. Legacy packed output remains the default.
Custom operation integration and validation tensorrt_llm/_torch/custom_ops/*, tests/unittest/_torch/thop/parallel/test_fp8_quantize.py	The fake operator models both layouts. The fused helper selects the legacy layout explicitly. Tests cover R128c4 output, padding, zero blocks, and compatibility.

Estimated code review effort: 4 (Complex) | ~45 minutes

Merge Risk: ⚪ Minimal · up to b684e

The current change is merge-ready after normal checks and review; no actionable merge-blocking risk remains.

Sequence Diagram(s)

sequenceDiagram
  participant TorchCaller
  participant fp8_quantize_1x128_packed_ue8m0
  participant invokeMxFP8Quantization
  participant MXFP8Kernel
  TorchCaller->>fp8_quantize_1x128_packed_ue8m0: pass use_r128c4_layout
  fp8_quantize_1x128_packed_ue8m0->>invokeMxFP8Quantization: select scale-vector configuration
  invokeMxFP8Quantization->>MXFP8Kernel: launch MXFP8 quantization
  MXFP8Kernel-->>fp8_quantize_1x128_packed_ue8m0: return FP8 values and selected scale layout

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 47.37% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the SM107 GEMM and quantization feature and follows the required ticket and type format.
Description check	✅ Passed	The description explains the objectives, behavior changes, compatibility impact, test coverage, and includes the required checklist sections.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

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 `@cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py`:
- Around line 1008-1011: Update the SM100 branch in the architecture-generation
loop to pass the has_arch predicate to generate_sm100_operations instead of the
integer sm, preserving the existing enabled-architecture check and break
behavior.

In `@cpp/tensorrt_llm/kernels/quantization.cu`:
- Around line 418-423: Add the missing explicit instantiation for
invokeMxFP8Quantization<__nv_bfloat16, 128, 128> alongside the existing BF16
specializations, ensuring the header-exposed default resolves correctly;
alternatively, remove the unsupported default if that specialization is
intentionally unavailable.

In `@cpp/tensorrt_llm/thop/fp8Quantize.cpp`:
- Line 274: Update the test calls to fp8_quantize_1x128_packed_ue8m0 at both
referenced cases so they explicitly pass use_r128c4_layout=False, preserving the
existing assertions for the legacy packed layout.

In `@tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py`:
- Around line 778-788: Add the return type annotation `tuple[torch.Tensor,
torch.Tensor]` to the fake operator function signature, preserving its existing
parameters and return behavior.

🪄 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: 0a0aa5a8-72ee-4a7a-a6c0-3eaba60748e4

📥 Commits

Reviewing files that changed from the base of the PR and between 0e2d598 and edf6d3c.

📒 Files selected for processing (14)

• cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt
• cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
• cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.cpp
• cpp/tensorrt_llm/kernels/cutlass_kernels/fp4_gemm/fp4_gemm_template.h
• cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm_template.h
• cpp/tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h
• cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py
• cpp/tensorrt_llm/kernels/quantization.cu
• cpp/tensorrt_llm/kernels/quantization.cuh
• cpp/tensorrt_llm/kernels/quantization.h
• cpp/tensorrt_llm/thop/fp8BlockScalingGemm.cpp
• cpp/tensorrt_llm/thop/fp8Quantize.cpp
• tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

[coderabbitai]

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.

[mingyangHao]

• ### [P1] Preserve the existing default output contract

This op previously accepted a single argument, returned the legacy strided int32 [m,
packed_sf_k] scale tensor, and supported K % 16 == 0.

Defaulting use_r128c4_layout to True silently changes the output dtype, shape, and
layout to a flat uint8 R128c4 buffer, while also tightening the requirement to K %
128 == 0. Updating the only in-tree caller to pass False does not protect downstream
users calling the registered torch.ops.trtllm op directly.

Please keep use_r128c4_layout=False as the default and let the future Rubin consumer
opt in explicitly, or register a separate R128c4 op. If the breaking default change
is intentional, this PR should be marked api-breaking and include migration guidance.

———

[P1] Add target-architecture coverage for the newly enabled SM107 paths

The PR description says the existing GEMM suites cover the dispatch changes, but the
test that directly exercises fp8_block_scaling_gemm still excludes SM107:

TensorRT-LLM/tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py

Lines 66 to 109 in 14cbc87

	@pytest.mark.skipif(
	getSMVersion() != 100 and getSMVersion() != 89 and getSMVersion() != 120,
	reason="The test is for Blackwell and Ada only. Current SM is %d." %
	getSMVersion(),
	)
	@pytest.mark.parametrize(
	"k, n",
	[(7168, 2112), (1536, 24576), (512, 32768), (16384, 7168), (7168, 4096),
	(2048, 7168), (1024, 1024)],
	)
	@pytest.mark.parametrize(
	"m",
	[7, 64, 128, 4096],
	)
	@pytest.mark.parametrize(
	"dtype",
	[torch.bfloat16],
	)
	def test_fp8_block_scale_gemm(dtype, m, k, n):
	torch.random.manual_seed(0)
	a = torch.randn((m, k), device='cuda', dtype=dtype) / k
	b = torch.randn((n, k), device='cuda', dtype=dtype) / k
	if getSMVersion() == 120:
	act_a_fp8, act_a_sf = fp8_utils.per_token_quant_and_transform(a)
	act_b_fp8, act_b_sf = per_block_cast_to_fp8_e8m0(b)
	act_b_sf = fp8_utils.transform_sf_into_required_layout(
	act_b_sf,
	mn=act_b_fp8.shape[0],
	k=act_b_fp8.shape[1],
	recipe=(1, 128, 128),
	is_sfa=False)
	else:
	act_a_fp8, act_a_sf = torch.ops.trtllm.fp8_quantize_1x128(a)
	act_b_fp8, act_b_sf = per_block_cast_to_fp8(b)
	output_expected = a @ b.t()
	output = torch.ops.trtllm.fp8_block_scaling_gemm(act_a_fp8, act_b_fp8,
	act_a_sf, act_b_sf)
	diff = calc_diff(output, output_expected)
	assert diff < 1e-3
	torch.testing.assert_close(output, output_expected, atol=1e-3, rtol=1e-3)

The relevant FP4 linear and MoE GEMM tests also remain restricted to SM100/SM103:

TensorRT-LLM/tests/unittest/_torch/thop/parallel/test_fp4_linear.py

Lines 100 to 112 in 14cbc87

	# compare
	torch.cuda.synchronize()
	torch.testing.assert_close(output, output_ref)
	@pytest.mark.skipif(sys.version_info < (3, 12),
	reason="cutlass-dsl 4.1.0 requires Python 3.12+")
	@pytest.mark.skipif(
	get_sm_version() not in [100, 103],
	reason="This test is only supported in sm100 and sm103 architecture",
	)
	@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE,
	reason="cutlass-dsl is not available")

TensorRT-LLM/tests/unittest/_torch/thop/parallel/test_moe_densegemm.py

Lines 75 to 85 in 14cbc87

	return ref
	@pytest.mark.skipif(
	get_sm_version() not in (100, 103),
	reason="This test is only supported on SM 100 and SM 103 GPUs",
	)
	@pytest.mark.parametrize("num_expert", [1, 4, 8, 32])
	@pytest.mark.parametrize("weight_per_expert", [256, 512, 1024, 2816])
	@pytest.mark.parametrize("num_tokens", [1, 8, 127, 256])

The new quantization tests validate the producer layout, but they do not validate
SM107 cubin compatibility, tactic selection, or the producer-to-GEMM-consumer
contract.

Before merging, please provide:

• A CUDA 13.4 build targeting SM107.

• Rubin correctness/smoke coverage for FP8 block-scale GEMM, FP4 dense/MoE GEMM, FP8
  rowwise GEMM, and fpA-intB GEMM.

• An R128c4 quantization-to-GEMM integration test.

B200 and H100 stages cannot validate the newly enabled SM107 runtime paths.

———

[P2] Make the SM107 TRT-LLM Gen compatibility rule explicit

This PR allows SM107 to enter TrtllmGenGemmRunner, but isSMCompatible(107, kernelSM)
currently falls through to return true, accepting every kernel SM rather than only
family-compatible Sm100f kernels:

TensorRT-LLM/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/KernelRunner.cpp

Lines 41 to 56 in 14cbc87

	constexpr bool isSMCompatible(int gpuSM, SmVersion kernelSM)
	{
	if (gpuSM == 103)
	{
	return kernelSM == SmVersion::Sm103a || kernelSM == SmVersion::Sm100f;
	}
	else if (gpuSM == 100)
	{
	return kernelSM == SmVersion::Sm100a || kernelSM == SmVersion::Sm100f;
	}
	else if (gpuSM == 90)
	{
	return kernelSM == SmVersion::Sm90a;
	}
	return true;
	}

The current generated metadata happens to contain only Sm100f candidates for this
exact DeepSeek FP8 configuration, so the present snapshot should work. However,
correctness relies on that incidental property of the kernel bundle; a future bundle
containing matching Sm100a or Sm103a candidates could select an incompatible cubin.

Please add an explicit SM107 rule that accepts only Sm100f kernels before enabling
this dispatch path.

[zhaoyangwang-nvidia]

Approving on behalf of trt-llm-runtime-devs

[farazkh80]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #69761 [ run ] triggered by Bot. Commit: 27c8293 Link to invocation

[tensorrt-cicd]

PR_Github #69761 [ run ] completed with state FAILURE. Commit: 27c8293
/LLM/main/L0_MergeRequest_PR pipeline #57058 completed with status: 'UNSTABLE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[farazkh80]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #69818 [ run ] triggered by Bot. Commit: 27c8293 Link to invocation

[tensorrt-cicd]

PR_Github #69818 [ run ] completed with state SUCCESS. Commit: 27c8293
/LLM/main/L0_MergeRequest_PR pipeline #57113 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[farazkh80]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #69976 [ run ] triggered by Bot. Commit: be67661 Link to invocation

[tensorrt-cicd]

PR_Github #69976 [ run ] completed with state FAILURE. Commit: be67661
/LLM/main/L0_MergeRequest_PR pipeline #57258 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[farazkh80]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #70050 [ run ] triggered by Bot. Commit: be67661 Link to invocation

[tensorrt-cicd]

PR_Github #70050 [ run ] completed with state SUCCESS. Commit: be67661
/LLM/main/L0_MergeRequest_PR pipeline #57325 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[farazkh80]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #70154 [ run ] triggered by Bot. Commit: be67661 Link to invocation

[tensorrt-cicd]

PR_Github #70154 [ run ] completed with state SUCCESS. Commit: be67661
/LLM/main/L0_MergeRequest_PR pipeline #57419 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[farazkh80]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #70159 [ run ] triggered by Bot. Commit: be67661 Link to invocation

[tensorrt-cicd]

PR_Github #70159 [ run ] completed with state SUCCESS. Commit: be67661
/LLM/main/L0_MergeRequest_PR pipeline #57422 completed with status: 'SUCCESS'

CI Report

Link to invocation