[None][perf] Size the per-tensor FP8 dynamic-quant amax grid to the input by hyukn · Pull Request #16569 · NVIDIA/TensorRT-LLM

hyukn · 2026-07-18T23:34:13Z

Description

The per-tensor FP8 (e4m3) dynamic quantization path (invokeComputeScalesAndQuantizeMatrix and invokeComputeFP8QuantizeScale in cudaFp8Utils.cu) computes the amax with computeFP8QuantizeScale<PER_TENSOR>, launched with a fixed grid(1024) and reducing across blocks with a 1024-way atomicMax to a single scale address. On a small activation only a fraction of the 1024 blocks load any data — the rest still launch, block-reduce zeros, and contend on the single atomic — so the kernel has a data-independent floor (~6–8us) no matter how little it actually processes (~200x off the memory-bound time for a 256KB input). This is a property of the launch/reduction design, not the data, and it applies to any model that uses the dynamic per-tensor FP8 quant path.

This PR sizes the amax grid to the input via the existing scaleMatrixVecGridSize(numel) (as the PER_CHANNEL and the vectorized apply paths already do), and adds a vectorized PER_TENSOR amax kernel (computeFP8QuantizeScalePerTensorVec, 128-bit loads / 8 bf16 per thread) mirroring the existing vectorized apply kernel. The net change is cpp/tensorrt_llm/common/cudaFp8Utils.cu only.

This supersedes the earlier Triton-based approach in this PR (first commit, kept in history): fixing the in-tree kernel is uniformly ≥ a Triton replacement, needs no new dependency, and serves both prefill and decode with one kernel.

Test Coverage

Kernel microbench — per-tensor amax only, B200, under CUDA graph (us/call), original fixed-grid + 1024-way-atomic vs this fix (numel-sized grid + vectorized):

rows M (K=2048)	original	this PR	speedup
1	6.1	4.1	1.5x
64	8.2	4.1	2.0x
256	8.2	6.1	1.3x
1024	10.2	6.1	1.7x
65536	87.4	43.0	2.0x

The fix removes the fixed ~6–8us floor at skinny (small-M) shapes and, via the vectorized loads, is also ~2x faster at very large M — i.e. no large-batch (prefill) regression.

Correctness. The per-tensor amax is an exact reduction and the scale is computed with the same amax / 448, so the op output is bitwise-unchanged. Verified on B200 against a torch reference for M = 1 … 65536: the returned scale is bitwise-equal and the quantized bytes match the apply-only op at every shape. Existing FP8-quantization tests cover the op path.

PR Checklist

PR description clearly explains what and why.
Follows the TRT-LLM coding guidelines.
Numerically identical to the prior kernel (exact per-tensor amax); covered by existing FP8-quant tests.
API-compatible: internal kernel launch change only, no signature changes.
No new dependencies.
Commit is DCO signed-off.

coderabbitai · 2026-07-18T23:40:28Z

📝 Walkthrough

Walkthrough

Changes

FP8 Dynamic Quantization

Layer / File(s)	Summary
Triton scale computation `tensorrt_llm/_torch/modules/fp8_dynamic_quant_triton.py`	Adds two-stage device-side FP8 e4m3 maximum reduction and produces the clamped, bf16-rounded fp32 scale.
FP8 linear-path integration `tensorrt_llm/_torch/modules/fp8_dynamic_quant_triton.py`, `tensorrt_llm/_torch/modules/linear.py`	Adds the combined dynamic quantization helper and uses it in `FP8QDQLinearMethod.apply` when static input scaling is unavailable.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant FP8QDQLinearMethod
  participant fp8_dynamic_quantize
  participant TritonScaleReduction
  participant static_quantize_e4m3_per_tensor
  FP8QDQLinearMethod->>fp8_dynamic_quantize: quantize dynamic activation
  fp8_dynamic_quantize->>TritonScaleReduction: compute scale
  TritonScaleReduction-->>fp8_dynamic_quantize: return scale
  fp8_dynamic_quantize->>static_quantize_e4m3_per_tensor: quantize input with scale
  static_quantize_e4m3_per_tensor-->>FP8QDQLinearMethod: return quantized input and scale

Suggested reviewers: yihuilu512

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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 is specific and related to the FP8 dynamic-quantization performance change, though it is narrower than the full changeset.
Description check	✅ Passed	The description includes the required sections and explains the issue, solution, tests, and checklist items.

✨ 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: 3

🤖 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/fp8_dynamic_quant_triton.py`:
- Around line 81-85: Add Google-style Args and Returns sections to both public
helpers, fp8_dynamic_quant_scale and fp8_dynamic_quantize, documenting their
arguments and return values. Update fp8_dynamic_quantize’s signature to annotate
its return type as tuple[torch.Tensor, torch.Tensor], using Python 3.10 built-in
generic syntax.
- Around line 52-62: Update the FP8 dynamic-quantization path around
FP8QDQLinearMethod.apply and _amax_partial_kernel to handle strided 2-D inputs
before the pointer-linear reduction. Either make the reduction operate on a
contiguous representation of the logical input or pass and apply the input
strides when loading elements, ensuring views such as x[:, ::2] include exactly
their logical values and exclude skipped storage.
- Line 61: Update the accumulator reduction in the Triton quantization path near
acc and quantize_e4m3_per_tensor so NaN inputs receive the same handling as the
CUDA implementation; either configure the maximum operation to propagate NaNs
explicitly or add a regression parity test covering NaN-containing inputs before
claiming bitwise identity.

🪄 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: 2bd2fb03-9a92-42f7-bc59-2c77e76a2a96

📥 Commits

Reviewing files that changed from the base of the PR and between 634abb8 and f77c236.

📒 Files selected for processing (2)

tensorrt_llm/_torch/modules/fp8_dynamic_quant_triton.py
tensorrt_llm/_torch/modules/linear.py

hyukn · 2026-07-20T03:41:32Z

/bot run --disable-fail-fst

tensorrt-cicd · 2026-07-20T03:47:18Z

PR_Github #60254 Bot args parsing error: usage: /bot [-h]
{run,kill,skip,submit,reviewers,reuse-pipeline,reuse-review} ...
/bot: error: unrecognized arguments: --disable-fail-fst

Link to invocation

kaiyux

Should we move this file to tensorrt_llm/_torch/modules/fused_ops?

kaiyux · 2026-07-20T05:44:52Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T05:51:07Z

PR_Github #60278 [ run ] triggered by Bot. Commit: e460a9d Link to invocation

tensorrt-cicd · 2026-07-20T09:38:58Z

PR_Github #60278 [ run ] completed with state SUCCESS. Commit: e460a9d
/LLM/main/L0_MergeRequest_PR pipeline #48635 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

hyukn · 2026-07-20T09:49:25Z

Should we move this file to tensorrt_llm/_torch/modules/fused_ops?

Done.

hyukn · 2026-07-20T09:51:22Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T09:57:23Z

PR_Github #60347 [ run ] triggered by Bot. Commit: d25d8e0 Link to invocation

The dynamic-quantization path of FP8QDQLinearMethod computes the per-tensor amax with torch.ops.tensorrt_llm.quantize_e4m3_per_tensor, whose amax half (computeFP8QuantizeScale in cudaFp8Utils.cu) is a shared-memory block reduce that is SM-underutilized / sync-bound (~6us) at skinny decode shapes (M=64) and serializes ahead of the projection GEMM across the ~30 GDN layers. Replace ONLY the amax->scale computation with a fused Triton per-tensor kernel (1.84x vs the CUDA fused op under CUDA graph), keeping the same apply-only quantize op. The scale and quantized output are bitwise-identical to torch.ops.tensorrt_llm.quantize_e4m3_per_tensor, so this is a pure throughput change with no accuracy impact (GSM8K unchanged within noise). Measured on Qwen3.6-35B-A3B-NVFP4, B200, ISL 1024 / OSL 6144, 9-point concurrency sweep: +2.5% (concurrency 256) to +14.7% (concurrency 1) output throughput; +5.4-5.6% at the canonical high-concurrency operating point. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

hyukn · 2026-07-20T11:56:13Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T12:06:40Z

PR_Github #60363 [ run ] triggered by Bot. Commit: 0631b79 Link to invocation

tensorrt-cicd · 2026-07-20T12:08:12Z

PR_Github #60347 [ run ] completed with state ABORTED. Commit: d25d8e0

Link to invocation

tensorrt-cicd · 2026-07-20T15:57:52Z

PR_Github #60363 [ run ] completed with state FAILURE. Commit: 0631b79
/LLM/main/L0_MergeRequest_PR pipeline #48704 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

hyukn · 2026-07-20T20:42:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T20:51:10Z

PR_Github #60452 [ run ] triggered by Bot. Commit: 0631b79 Link to invocation

tensorrt-cicd · 2026-07-20T22:32:55Z

PR_Github #60452 [ run ] completed with state SUCCESS. Commit: 0631b79
/LLM/main/L0_MergeRequest_PR pipeline #48787 completed with status: 'SUCCESS'

CI Report

Link to invocation

JunyiXu-nv

LGTM from runtime side.

Supersede the Triton FP8 dynamic-quant amax (previous commit) with a fix to the in-tree CUDA kernel it was working around. The dynamic per-tensor path (invokeComputeScalesAndQuantizeMatrix / invokeComputeFP8QuantizeScale) launched computeFP8QuantizeScale<PER_TENSOR> with a FIXED grid(1024) and reduced across blocks with a 1024-way atomicMax to a single scale address. On a skinny decode activation ([64,2048]) only ~128 of the 1024 blocks load data; the rest launch, block-reduce zeros, and still contend on the one atomic — a data-independent ~6us floor. Fix: size the amax grid to the input via the existing scaleMatrixVecGridSize(), and add a vectorized PER_TENSOR amax kernel (128-bit loads) mirroring the existing vectorized apply kernel. This reverts the Python/Triton change (fp8_dynamic_quant.py, linear.py, its unit test and l0 entry); the previous commit is kept in history. Why the CUDA fix over the Triton replacement: - Per-tensor amax is exact (max reduction), so the op output is bitwise- identical to baseline (verified: scale + quantized bytes across M=1..65536; GSM8K 96.75 vs 96.25 baseline, within noise). - It is uniformly >= the Triton kernel and does NOT regress at large M (prefill), which the fixed-grid-128 Triton kernel did (crossover ~M=2048). - E2E (Qwen3.6-35B-A3B-NVFP4, B200, ISL1024/OSL6144, concurrency 1..256): mean +7.34% vs baseline, vs +4.87% for the Triton version. - One kernel for prefill and decode; no new Triton dependency. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

hyukn · 2026-07-21T12:56:30Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-21T13:02:49Z

PR_Github #60697 [ run ] triggered by Bot. Commit: 87fd271 Link to invocation

tensorrt-cicd · 2026-07-21T18:00:00Z

PR_Github #60697 [ run ] completed with state FAILURE. Commit: 87fd271
/LLM/main/L0_MergeRequest_PR pipeline #48988 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

hyukn · 2026-07-22T00:07:21Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-22T00:13:40Z

PR_Github #60826 [ run ] triggered by Bot. Commit: 87fd271 Link to invocation

tensorrt-cicd · 2026-07-22T02:42:27Z

PR_Github #60826 [ run ] completed with state SUCCESS. Commit: 87fd271
/LLM/main/L0_MergeRequest_PR pipeline #49102 completed with status: 'SUCCESS'

CI Report

Link to invocation

…nput (NVIDIA#16569) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

hyukn requested a review from a team as a code owner July 18, 2026 23:34

hyukn requested review from YihuiLu512 and zhaoyangwang-nvidia July 18, 2026 23:34

github-actions Bot assigned hyukn Jul 18, 2026

coderabbitai Bot reviewed Jul 18, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/modules/fused_ops/fp8_dynamic_quant.py Outdated

Comment thread tensorrt_llm/_torch/modules/fused_ops/fp8_dynamic_quant.py Outdated

Comment thread tensorrt_llm/_torch/modules/fp8_dynamic_quant_triton.py Outdated

hyukn force-pushed the perf/fp8-fused-dynamic-quant-amax branch from f77c236 to e460a9d Compare July 20, 2026 01:02

kaiyux reviewed Jul 20, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/modules/fused_ops/fp8_dynamic_quant.py Outdated

hyukn force-pushed the perf/fp8-fused-dynamic-quant-amax branch 2 times, most recently from c0e43f0 to 04b89ce Compare July 20, 2026 09:04

hyukn force-pushed the perf/fp8-fused-dynamic-quant-amax branch from 04b89ce to d25d8e0 Compare July 20, 2026 09:50

hyukn requested review from a team as code owners July 20, 2026 09:50

hyukn requested review from brnguyen2, mzweilz, niukuo, schetlur-nv and xinhe-nv July 20, 2026 09:50

hyukn force-pushed the perf/fp8-fused-dynamic-quant-amax branch from d25d8e0 to 0631b79 Compare July 20, 2026 10:59

mikeiovine requested review from mikeiovine and removed request for schetlur-nv July 20, 2026 15:25

kaiyux approved these changes Jul 21, 2026

View reviewed changes

xinhe-nv approved these changes Jul 21, 2026

View reviewed changes

JunyiXu-nv approved these changes Jul 21, 2026

View reviewed changes

hyukn changed the title ~~[None][perf] Fuse Triton FP8 dynamic-quant amax for dense projections~~ [None][perf] Size the per-tensor FP8 dynamic-quant amax grid to the input Jul 22, 2026

hyukn merged commit 09e5d0c into NVIDIA:main Jul 22, 2026
16 checks passed

yuanjingx87 pushed a commit to yuanjingx87/TensorRT-LLM that referenced this pull request Jul 26, 2026

[None][perf] Size the per-tensor FP8 dynamic-quant amax grid to the i…

e405ce9

…nput (NVIDIA#16569) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

Uh oh!

Conversation

hyukn commented Jul 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

Uh oh!

coderabbitai Bot commented Jul 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

hyukn commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

kaiyux left a comment

Choose a reason for hiding this comment

Uh oh!

kaiyux commented Jul 20, 2026

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

hyukn commented Jul 20, 2026

Uh oh!

hyukn commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

hyukn commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

hyukn commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

JunyiXu-nv left a comment

Choose a reason for hiding this comment

Uh oh!

hyukn commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

hyukn commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

hyukn commented Jul 18, 2026 •

edited

Loading

coderabbitai Bot commented Jul 18, 2026 •

edited

Loading