[https://nvbugs/6424956][fix] Support large FP8 quantization grids

[lfr-0531]

Dev Engineer Review

• Updated FP8 quantization kernel launches to distribute FP8 row blocks across grid.y and grid.z, preventing failures when max_num_tokens pushes grid.y beyond CUDA’s 65,535 dimension limit.
• Adjusted kernel indexing to linearize blockIdx.z and blockIdx.y into the FP8 row-block index (m_idx), while preserving the prior mapping behavior for launches within the previous limit (e.g., grid.z == 1).
• Added makeQuantizeGrid(numPackedSfK, mBlocks) to construct the bounded 3D grid rather than hardcoding dim3(numPackedSfK, m_blocks, 1).
• Preserved existing semantics/mapping for smaller launches; no public API, dependency, or configuration changes were introduced.
• GPU execution remains unverified in CI due to the reported B200 NVLink Fabric state and CUDA error 802; build/registration/binary inspection/pre-commit checks passed.

QA Engineer Review

• Extended the SM100-only fp8_quantize_1x128_packed_ue8m0_matches_legacy parameterization with a boundary case m=262141, k=128 (id="grid-y-overflow") to exercise the grid-dimension limit fix.
• The boundary regression asserts bit-identical FP8 output and exact packed UE8M0 scale equality between the fused FP8 quant+pack path and the legacy quantize-and-pack path.
• No changes were made to integration test list files (tests/integration/test_lists/test-db, tests/integration/test_lists/qa, or tests/integration/test_lists/*.txt).
• Verdict: needs follow-up.

Description

NVBug 6424956 reports a cudaLaunchKernelEx invalid-argument failure when
max_num_tokens is 64K. The fused 1x128 FP8 quantization kernels mapped all
row blocks to grid.y; after token expansion, an aligned row count of 262144
requires 65536 row blocks, exceeding CUDA's 65535 limit for grid.y.

This change distributes row blocks across grid.y and grid.z and linearizes
the two indices in the kernel. Launches within the existing limit retain
grid.z == 1 and the same row mapping. A boundary regression case compares
the fused path with the legacy quantize-and-pack path.

There are no API or dependency changes. The expected functional impact is
limited to enabling large FP8 quantization launches that previously failed.
GPU execution remains unverified because the available B200 node's NVLink
Fabric stayed in In Progress state and CUDA initialization returned error
802; the PR is therefore submitted as a draft.

Related: NVBug 6424956, TRTLLM-14125.

Test Coverage

• Built tensorrt_llm, th_common, and Python bindings for SM100.
• Linked and installed the built libraries into the worktree package.
• Loaded the new libth_common.so and confirmed the target op is
  registered and all mapped TensorRT-LLM libraries come from this
  worktree.
• Confirmed the compiled kernel reads SR_CTAID.Z with cuobjdump.
• Passed all staged-file pre-commit hooks.
• Run
  test_fp8_quantize.py::test_fp8_quantize_1x128_packed_ue8m0_matches_legacy[grid-y-overflow]
  on B200. This is blocked by the node-level NVLink Fabric/CUDA error 802
  described above.

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.

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62201 [ run ] triggered by Bot. Commit: 4cd6a03 Link to invocation

[coderabbitai]

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: 47082978-e0a9-4122-b98f-e9776b16763f

📥 Commits

Reviewing files that changed from the base of the PR and between 4cd6a03 and 4a39aeb.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu

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

• cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu

---

Walkthrough

FP8 packed quantization now distributes row blocks across CUDA grid.y and grid.z, combines both dimensions when indexing kernel rows, updates two launch paths, and adds an SM100 test case covering m=262141, k=128.

Changes

FP8 quantization grid scaling

Layer / File(s)	Summary
Partition row blocks across grid dimensions cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu	Adds makeQuantizeGrid with a 65535 limit for grid.y/grid.z, and updates kernel indexing to combine both dimensions.
Use the grid mapping in launches and tests cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu, tests/unittest/_torch/thop/parallel/test_fp8_quantize.py	Updates both quantization launch functions and adds the grid-y-overflow SM100 test case.

Estimated code review effort: 2 (Simple) | ~10 minutes

Sequence Diagram(s)

sequenceDiagram
  participant QuantizationLaunch
  participant makeQuantizeGrid
  participant CUDAKernel
  QuantizationLaunch->>makeQuantizeGrid: Pass num_packed_sf_k and m_blocks
  makeQuantizeGrid-->>QuantizationLaunch: Return 3D CUDA grid
  QuantizationLaunch->>CUDAKernel: Launch with grid.y and grid.z
  CUDAKernel->>CUDAKernel: Combine blockIdx.z and blockIdx.y into mBlockIdx

Loading

Suggested reviewers: zongfeijing

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise, specific, and matches the main change: fixing large FP8 quantization grids.
Description check	✅ Passed	The description follows the template and covers the issue, solution, test coverage, and checklist.
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.

✨ 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: 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 `@tests/unittest/_torch/thop/parallel/test_fp8_quantize.py`:
- Line 400: Add an equivalent pytest parameter for the grid-y overflow boundary
case to the test coverage invoking launch_fp8_quantize_1x128_cutedsl_bf16_e4m3,
matching the existing packed-path inputs and expected behavior while retaining
the descriptive grid-y-overflow identifier.

🪄 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: 4e98a3c0-1adb-400c-8297-a713a6aee6ab

📥 Commits

Reviewing files that changed from the base of the PR and between f3d4c85 and 4cd6a03.

📒 Files selected for processing (2)

• cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu
• tests/unittest/_torch/thop/parallel/test_fp8_quantize.py

[tensorrt-cicd]

PR_Github #62201 [ run ] completed with state FAILURE. Commit: 4cd6a03
/LLM/main/L0_MergeRequest_PR pipeline #50371 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

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62350 [ run ] triggered by Bot. Commit: 4cd6a03 Link to invocation

[tensorrt-cicd]

PR_Github #62350 [ run ] completed with state SUCCESS. Commit: 4cd6a03
/LLM/main/L0_MergeRequest_PR pipeline #50517 completed with status: 'SUCCESS'

CI Report

Link to invocation

[jiaganc]

Reproduce script passed. LGTM.

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62493 [ run ] triggered by Bot. Commit: 4a39aeb Link to invocation

[tensorrt-cicd]

PR_Github #62493 [ run ] completed with state SUCCESS. Commit: 4a39aeb
/LLM/main/L0_MergeRequest_PR pipeline #50640 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

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62544 [ run ] triggered by Bot. Commit: 4a39aeb Link to invocation

[tensorrt-cicd]

PR_Github #62544 [ run ] completed with state SUCCESS. Commit: 4a39aeb
/LLM/main/L0_MergeRequest_PR pipeline #50685 completed with status: 'SUCCESS'

CI Report

Link to invocation

[mingyangHao]

LGTM

[zongfeijing]

LGTM