[Common] Add launch bounds to multi-tensor swizzle kernels

[ravimajeti]

Description

The multi-tensor swizzle and unswizzle kernels in common/swizzle/swizzle.cu launch with TB_DIM x TB_DIM (1024) threads per block. Every other swizzle kernel has __launch_bounds__(TB_DIM * TB_DIM), added in #2076, but these four don't. Without the bound, ptxas doesn't cap registers at 64 per thread, and some variants use up to 99. 1024 threads x more than 64 registers is more than a block can have, so the launch fails with too many resources requested for launch.

Which variants fail depends on the target arch and the CUDA version, because the register count does. Measured with cuobjdump --dump-resource-usage (registers per thread; more than 64 can't launch with 1024 threads):

Kernel	12.8 sm_90	12.8 sm_100	12.8 sm_120	13.4 sm_90	13.4 sm_100	13.4 sm_120
multi_tensor_swizzle_row_scaling_kernel<int4>	89	89	96	61	50	56
multi_tensor_swizzle_row_scaling_kernel<int2>	32	61	71	32	61	70
multi_tensor_swizzle_col_scaling_kernel<int4>	89	99	99	64	56	61

With this PR, every multi-tensor (un)swizzle variant is at most 64 registers on all six configurations, with no spills (STACK:0, LOCAL:0).

Fixes #3621

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

• common/swizzle/swizzle.cu: add __launch_bounds__(TB_DIM* TB_DIM) to multi_tensor_swizzle_{row,col}_scaling_kernel and multi_tensor_unswizzle_{row,col}_scaling_kernel, matching the other swizzle kernels.
• tests/cpp/operator/test_multi_swizzle.cu: add two shapes to multi_tensor_test_cases, which is shared by the swizzle, unswizzle and roundtrip suites. Before this, no test reached the row<int2> or col<int4> variants:
  • {2, 128, 4352, true}: rowwise num_tiles_k = 34, which selects vec_load_size = 2.
  • {2, 512, 4096, false}: columnwise num_tiles_k = 4, which selects vec_load_size = 4. num_tiles_m = 32, so it isn't the narrow-M path.

Testing

RTX 5090 (SM 12.0), CUDA 12.8, NVTE_FRAMEWORK=pytorch NVTE_CUDA_ARCHS=120, main @ 5759fa0:

Check	main + new tests	This PR
test_operator --gtest_filter='*wizzle*'	497 passed, 5 failed (too many resources requested for launch)	502 passed
tests/pytorch/test_weight_swizzle_in_layers.py	36 passed, 16 skipped	36 passed, 16 skipped

The 5 failures on main are n2_M128_K1024_row, n3_M256_K4096_row and n2_M128_K8192_row, which already exist, plus the two new shapes. n2_M128_K1024_row is meant to cover the narrow-K kernel. It needs 128 KB of shared memory, so on a GPU with 99 KiB the dispatcher falls back to the regular multi-tensor kernel, which then fails to launch.

The register numbers for sm_90 and sm_100 are from compile-only builds of swizzle.cu with the TE build flags. I don't have an H100 or B200, so I haven't run the tests there. Based on the register counts, a CUDA 12.8 build on those GPUs should fail the existing K4096/K8192 rowwise cases.

pre-commit passes on both changed files.

Performance

RTX 5090, CUDA 12.8. I benchmarked two copies of libtransformer_engine.so (main and this PR) with the same binary that calls nvte_multi_tensor_{,un}swizzle_scaling_factors, alternating between them for 3 rounds each. GPU time is 200 calls captured in a CUDA graph and replayed, divided by 200, taking the median of 7 repeats. Run-to-run spread is ≤ 0.01 µs.

Variant	Tensors × (M, K)	main GPU µs	This PR GPU µs	Change
row<int>	8 × (4096, 4224)	4.95	4.98	+0.6%
row<int>	32 × (512, 1408)	2.34	2.37	+1.2%
row<int2>	8 × (4096, 4352)	launch error	4.02
row<int4>	8 × (4096, 7168)	launch error	5.30
col<int>	8 × (4224, 4096)	4.70	4.83	+2.7%
col<int2>	8 × (768, 7168)	3.09	3.05	−1.3%
col<int2>	32 × (1792, 4096)	7.48	7.41	−1.0%
col<int4>	8 × (4096, 7168)	launch error	5.68
unswizzle row<int>	8 × (4096, 4224)	5.35	5.35	0.0%
unswizzle row<int4>	8 × (4096, 7168)	5.68	5.60	−1.4%
unswizzle col<int2>	8 × (768, 7168)	3.17	3.19	+0.7%
unswizzle col<int4>	8 × (4096, 7168)	5.87	5.87	0.0%

For variants that already launched, GPU time changes by −1.4% to +2.7% (at most 0.13 µs), in both directions. With host launch overhead included (back-to-back calls without a graph), each call takes 4–8 µs and is unchanged. The scales stay in L2 across repeated calls, so these are L2-resident timings. Clocks were not locked.

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

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 5/5

[Medium risk] Adds launch bounds to GPU kernel functions.

The PR appears safe to merge; no actionable regression was identified.

Summary

The PR adds launch bounds to four multi-tensor scale swizzle kernels and adds C++ cases targeting previously uncovered vector variants.

• The bounds match the kernels’ 1024-thread launch configuration.
• The added shapes select the intended regular-kernel variants.

_{Reviews (1) · Last reviewed commit: "[Common] Add launch bounds to multi-tens..."}

[ptrendx]

/te-ci core