[greg-kwasniewski1]

@coderabbitai summary

Description

invokeMxFP8Quantization gives each CTA one padded row. With the swizzled 128x4 scale-factor layout, a few rows
therefore launch PadUp(m, 128) = 128 CTAs of up to 512 threads. At M = 6, K = 5120, six CTAs convert data (two passes
over the row each) and 122 CTAs only write zero padding scale factors, one byte per store.

This PR adds quantize_mxfp8_small_m_swizzled<T> in quantization.cu and dispatches to it when all of these hold:

• SM 100 or SM 103;
• SF_VEC_SIZE == SF_OUTPUT_VEC_SIZE == 32, swizzled layout, b == 1;
• 1 <= m <= 32, n == padded_n, n % 128 == 0, n > 2048;
• 16-byte aligned input and scale-factor pointers.

Every other shape, layout, scale-vector size and architecture keeps quantize_with_block_size. SM107 is not touched,
so the SM107 path that #20623 adds in the same function is not affected; the two changes merge without conflicts.

How the kernel works:

• Grid (ceil(n / 1024), m), 128 threads per CTA; each thread converts one 8-element vector.
• Values and scale factors come from cvt_warp_fp16_to_mxfp8, the function that quantize_with_block_size calls, so
  the output is bitwise identical, including the exp2f_rcp handling of tiny blocks.
• A 128-column atom of the swizzled layout is 32 slots of 16 bytes; slot r holds the K32 scales of rows r, r + 32,
  r + 64 and r + 96. With m <= 32, slot r < m is [four scales of row r, 12 zero bytes] and every other slot is zero.
  The kernel writes each atom with 16-byte stores and covers every byte of the scale-factor buffer.
• The PDL contract does not change: cudaGridDependencySynchronize() comes before the first global access, and
  __threadfence(); __syncthreads(); come before cudaTriggerProgrammaticLaunchCompletion().

Up to 2048 columns both kernels run at the launch latency on B200 and the existing kernel is not slower, so the guard
starts above 2048 columns.

Performance (B200)

Per-launch time in µs, main → this PR, bf16, PDL on. Each value is the median of 9 rounds of a CUDA graph with 50
back-to-back launches. Main's and this PR's quantization.cu are compiled with the same nvcc and release flags
(-O3 -DNDEBUG, sm_100f) into one binary, and the rounds alternate between them.

M \ N	2048	4096	5120	6144	7168	8192	16384
1	1.68 (not dispatched)	2.05 → 1.76 (1.17x)	2.38 → 1.76 (1.35x)	2.34 → 1.77 (1.32x)	2.38 → 1.77 (1.35x)	2.54 → 1.76 (1.44x)	3.69 → 1.93 (1.91x)
2	1.68 (not dispatched)	2.05 → 1.68 (1.22x)	2.34 → 1.68 (1.39x)	2.38 → 1.88 (1.26x)	2.40 → 1.89 (1.27x)	2.58 → 1.89 (1.37x)	3.69 → 1.93 (1.91x)
4	1.69 (not dispatched)	2.05 → 1.80 (1.14x)	2.34 → 1.76 (1.32x)	2.38 → 1.80 (1.32x)	2.43 → 1.80 (1.35x)	2.62 → 1.97 (1.33x)	3.71 → 2.01 (1.84x)
6	1.68 (not dispatched)	2.05 → 1.80 (1.14x)	2.34 → 1.80 (1.29x)	2.38 → 1.81 (1.32x)	2.54 → 1.81 (1.40x)	2.62 → 1.80 (1.46x)	3.73 → 2.01 (1.86x)
8	1.73 (not dispatched)	2.05 → 1.80 (1.14x)	2.34 → 1.80 (1.30x)	2.38 → 1.80 (1.32x)	2.42 → 1.80 (1.34x)	2.62 → 1.80 (1.45x)	3.73 → 1.85 (2.02x)
16	1.72 (not dispatched)	2.10 → 1.73 (1.21x)	2.38 → 1.77 (1.35x)	2.38 → 1.76 (1.35x)	2.46 → 1.76 (1.40x)	2.63 → 1.77 (1.48x)	3.77 → 2.10 (1.79x)
32	1.73 (not dispatched)	2.13 → 1.77 (1.21x)	2.38 → 1.97 (1.21x)	2.46 → 2.09 (1.18x)	2.58 → 2.09 (1.23x)	2.67 → 2.10 (1.27x)	3.81 → 2.42 (1.58x)

Summary over all measured shapes:

	Dispatched (72 shapes: M ≤ 32, N ∈ {2176, 3584, 4096, 5120, 6144, 7168, 8192, 16384})	Not dispatched (71 shapes)
bf16, PDL on	1.05x – 2.02x, median 1.29x	0.98x – 1.02x
fp16, PDL on	1.09x – 2.09x, median 1.36x	0.98x – 1.02x
bf16, PDL off	1.04x – 1.92x, median 1.27x	0.98x – 1.02x
fp16, PDL off	1.08x – 2.00x, median 1.33x	0.98x – 1.02x

The not-dispatched shapes (M ∈ {33, 64, 128, 1024}, N ≤ 2048, N = 2880) run the same code as main; their spread is the
run-to-run noise.

Kernel duration from nsys (bf16, PDL off, median over all launches):

M × N	main	this PR
1 × 8192	2.50 µs	1.66 µs
6 × 5120	2.27 µs	1.73 µs
6 × 6144	2.27 µs	1.73 µs
16 × 8192	2.53 µs	1.70 µs
32 × 5120	2.27 µs	1.89 µs

Setup: one B200 (148 SMs), driver 615.71.10, nvcc 13.2, default clocks (not locked; maximum SM clock 1965 MHz).

Test Coverage

New tests in tests/unittest/_torch/kernels/parallel/test_fp8_quantize.py. The existing
unittest/_torch/kernels/parallel entries of l0_b200.yml and l0_gb300.yml run them.

• test_mxfp8_quantize_swizzled_matches_linear_layout: for m ∈ {1, 2, 6, 12, 31, 32, 33, 128},
  k ∈ {1024, 2048, 2176, 2880, 5120, 7168, 16384}, fp16 and bf16, the swizzled output equals the linear output with
  block_scale_interleave applied, byte for byte over the full value and scale-factor buffers, padding included. The
  inputs contain zero blocks, a block below the smallest UE8M0 scale, the dtype maximum, subnormals and negative zeros.
  Before the call, blocks filled with 0xFF are released to the caching allocator, so that an unwritten byte is likely
  to show.
• test_mxfp8_quantize_swizzled_cuda_graph_replays: CUDA-graph replays with new inputs give the bytes of eager calls,
  for dispatched and fallback shapes.

Existing tests on the same paths: test_mxfp8_quantize_torch_device, test_mxfp8_quantize_alignment_torch_device,
test_mxfp8_quantize_zero_rows, test_mxfp8_quantize_tiny_block_is_finite, and test_mxfp8_mxfp8_gemm.py for the
consumer GEMM.

Local validation on B200:

• A standalone binary that links main's and this PR's quantization.cu compared every byte of the value and
  scale-factor buffers for M ∈ {1, 2, 4, 6, 8, 12, 16, 24, 31, 32, 33, 64, 128, 1024} ×
  N ∈ {1024, 1152, 1536, 2048, 2880, 3584, 4096, 5120, 6144, 7168, 8192, 16384} × {fp16, bf16} × 3 seeds, with the
  outputs pre-filled with 0xFF and with 0x00: 3528 comparisons, 0 differing bytes.
• compute-sanitizer memcheck, racecheck and synccheck on the same binary: 0 errors.
• The two new tests pass against the unchanged kernel of the TensorRT-LLM 1.2.1 wheel (118 passed).
• quantization.cu compiles for sm_80, sm_90a, sm_100f, sm_103f and sm_120f with the same warnings as main. ptxas
  reports 30 registers, no stack frame and no spills for the new kernel.

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.

🤖 Generated with Claude Code

[greg-kwasniewski1]

/bot run

[tensorrt-cicd]

PR_Github #82237 [ run ] triggered by Bot. Commit: 3445c1d Link to invocation

[tensorrt-cicd]

PR_Github #82237 [ run ] completed with state FAILURE. Commit: 3445c1d
/LLM/main/L0_MergeRequest_PR pipeline #67866 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