[tdophung]

Description

common::reduce_dbias adds up the FP32 dbias workspace (one partial row per CTA of the quantize kernel) into the
dbias vector. It serves:

• the MXFP8 quantize with fused dbias, on both the CUDA and the CuTeDSL backends (quantize_mxfp8.cuh,
  quantize_mxfp8_cutedsl.cuh);
• the FP8 tensor-scaling quantize with fused dbias (quantize_fp8.cuh).

The grouped reduction (grouped_reduce_dbias) and the cast-transpose dbias reductions are separate kernels and do not
change.

The old kernel ran 256 threads per CTA with four columns each (two for FP32 output). A 7168-column BF16 workspace
therefore launched 7 CTAs, and each thread waited on one load per row. This PR uses one thread per column, 64
threads per CTA:

• On sm_80 and newer, each thread copies its column into shared memory with cp.async, 32 rows per stage and 4
  stages deep, and adds each stage once it has landed.
• Below sm_80, each thread loads its column directly.

Each column is still added in row order from 0 and rounded once, so dbias is bit-identical to before. The
cols % nvec check is dropped: any column count works.

The change is on by default for every caller above; there is no option. It was split out of #3761.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• transformer_engine/common/cast/core/grouped_layout.cuh: the new reduce_dbias_kernel and its launch.
• tests/cpp/util/test_reduce_dbias.cu (new): checks reduce_dbias bit for bit against a host sum in row order.
  • FP32, BF16 and FP16 output.
  • Data spans 2^-30..2^30 with mixed signs, so any other order or grouping of the adds changes the result.
  • Rows 1-2000; columns 1-7170, including counts that are not multiples of 4.

Performance

One GB300 GPU, CuTeDSL backend on, this PR against its base main b9e8e57, on the same node:

Op (BF16 in, MXFP8 E4M3 out)	Shape	main (us)	this PR (us)	Speedup (round 1 / round 2)	Kernels, both arms
cast + dbias (bgrad_quantize)	4096x7168	28.83	25.15	1.15x (1.157 / 1.136)	CuTeDSL general kernel + reduce_dbias
cast + dbias (bgrad_quantize)	4096x36864	100.58	95.25	1.06x (1.054 / 1.058)	same

• Bits: dbias, the FP8 data and the scales are bit-identical to main's (SHA-256). In both arms they are bitwise
  equal to the CUDA backend's.
• Not timed here: the CUDA backend and the FP8 tensor-scaling quantize call the same reduce_dbias.

Measurement details:

• Hardware: one NVIDIA GB300, exclusive node, driver 580.173.10.
• Software:
  • Container nvcr.io/nvidia/pytorch:26.10-py3, build 477939455, index digest
    sha256:6d347d649d5dd108b4652f1b9cea6a3c859692ac408617a9108616da40e3da1f.
  • PyTorch 2.15.0a0+875d815, CUDA 13.4.3, cuDNN 9.27.0.42, NCCL 2.32.3.
  • nvidia-cutlass-dsl 4.8.0, apache-tvm-ffi 0.1.14.post1, cuda-core 1.2.2.
  • TE 2.22.0.dev0 built from source at each arm's commit, with
    NVTE_FRAMEWORK=pytorch NVTE_WITH_CUTEDSL=1 NVTE_CUDA_ARCHS=100 (nvcc 13.4).
• Backend and environment: NVTE_ENABLE_CUTEDSL_BACKEND=1 NVTE_ACTIVATION_TABLE=1 in both arms; the tables do
  not apply to a plain cast, and main ignores the variable. The kernels each arm launched were checked with
  torch.profiler.
• Call: one tex.bgrad_quantize call with an MXFP8Quantizer (MXFP8 block scaling): BF16 input from randn,
  E4M3 output, rowwise and columnwise, GEMM-swizzled scales.
• Timing:
  • The window holds the quantize kernel and reduce_dbias, between two CUDA events.
  • Before each call, the L2 is flushed by reading a buffer twice its size, and a GPU sleep longer than the host
    launch is enqueued. Launch overhead therefore stays outside the window.
  • 20 warm-up calls run first; the first one compiles the CuTeDSL kernel.
  • Each run reports the median of 500 calls.
• Order: main, PR, PR, main on one GPU. The table gives the mean of each arm's two medians, and the speedup per
  round.
• Command:
  NVTE_ENABLE_CUTEDSL_BACKEND=1 NVTE_ACTIVATION_TABLE=1 python bench_quantize.py --op cast_dbias --m 4096 --n 7168 --in-dtype bf16 --fp8 e4m3 --direction both --swizzle 1 --reps 500 --warmup 20 --flush read
  (--n 36864 for the second row). The script is not in TE's benchmarks/.

Testing

All on one GB300 (SM 10.3); the C++ tests were built for sm_103a.

Suite	CuTeDSL backend off	CuTeDSL backend on, NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN=ERROR
test_util --gtest_filter='*ReduceDbias*' (new)	143 passed (backend-independent)
test_operator FP8 dbias, *CastDBias*:*CastDBiasDGelu*	168 passed	168 passed
test_operator MXFP8 cast and dbias, *FusedCastMXFP8*:*SwizzledScalesFusedCastMXFP8*	688 passed, 488 skipped	688 passed, 488 skipped
compute-sanitizer racecheck and synccheck, ReduceDbias 512x65 and 1000x1028	0 hazards, 0 errors

• Lint: cpplint 1.6.0 and clang-format 18.1.6 are clean.
• Not run here: the minimum-version builds (GCC 11, sm_70, where the plain-loop path compiles), SM 10.0, and
  JAX.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes