Multi-tensor swizzle kernels fail with "too many resources requested for launch" (missing __launch_bounds__)
Maintainers usually reply within 2 days
Assessment
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Newbie friendliness
- 25/100
Research direction
The four multi-tensor kernels live in transformer_engine/common/swizzle/swizzle.cu and are launched from launch_multi_tensor_swizzle_scaling_factors (the line in the report's stack trace); #2076 added launch_bounds(TB_DIM * TB_DIM) to the single-tensor kernels, so compare those declarations first. Done means the row, row and col variants stay within 64 registers at 1024 threads and the MultiTensorSwizzleTestSuite cases pass when run via tests/cpp/build/operator/test_operator --gtest_filter='MultiTensorSwizzleTestSuite'. Note PR #3622 is already open against this issue, so coordinate before starting.
Written by the indexing model from the issue text.
Description
Describe the bug
The multi-tensor swizzle kernels in transformer_engine/common/swizzle/swizzle.cu (multi_tensor_swizzle_{row,col}_scaling_kernel and multi_tensor_unswizzle_{row,col}_scaling_kernel) are launched with 1024 threads per block (dim3 block_size(TB_DIM, TB_DIM)), but unlike every other swizzle kernel they have no __launch_bounds__(TB_DIM * TB_DIM). Without it, ptxas may give a thread more than 64 registers. 1024 threads × more than 64 registers exceeds the 64K registers available to a block, so the launch fails with:
CUDA Error: too many resources requested for launch
#2076 added __launch_bounds__ to the single-tensor swizzle kernels; the multi-tensor kernels added a week earlier in #2019 were not included.
Whether it fails depends on the GPU arch and CUDA version, because the register count does. Measured with cuobjdump --dump-resource-usage on swizzle.cu compiled with the TE build flags (registers per thread; > 64 cannot launch with 1024 threads):
| Kernel | CUDA 12.8 sm_90 | CUDA 12.8 sm_100 | CUDA 12.8 sm_120 | CUDA 13.4 sm_90 | CUDA 13.4 sm_100 | CUDA 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 |
| other multi-tensor (un)swizzle variants | ≤ 56 | ≤ 40 | ≤ 48 | ≤ 55 | ≤ 40 | ≤ 48 |
So on sm_120 it fails with both CUDA versions, and with CUDA 12.8 (the minimum for Blackwell) it should also fail on sm_90 and sm_100.
Steps/Code to reproduce bug
RTX 5090 (CC 12.0), TE built with NVTE_CUDA_ARCHS=120, CUDA 12.8:
tests/cpp/build/operator/test_operator --gtest_filter='*MultiTensorSwizzleTestSuite*'
3 cases fail, all routed to multi_tensor_swizzle_row_scaling_kernel<int4>:
[ FAILED ] OperatorTest/MultiTensorSwizzleTestSuite.TestMultiTensorSwizzle/n2_M128_K1024_row
[ FAILED ] OperatorTest/MultiTensorSwizzleTestSuite.TestMultiTensorSwizzle/n3_M256_K4096_row
[ FAILED ] OperatorTest/MultiTensorSwizzleTestSuite.TestMultiTensorSwizzle/n2_M128_K8192_row
C++ exception with description "transformer_engine/common/swizzle/swizzle.cu:1360 in function launch_multi_tensor_swizzle_scaling_factors: CUDA Error: too many resources requested for launch" thrown in the test body.
(n2_M128_K1024_row is meant to cover the narrow-K kernel, which needs 128 KB of shared memory. On a 99 KiB GPU the dispatcher correctly falls back to the regular multi-tensor kernel, which then fails to launch.)
The row<int2> and col<int4> variants are not reached by the current test shapes, e.g. {2, 128, 4352, true} (row, vec_load_size = 2) and {2, 512, 4096, false} (col, vec_load_size = 4) would cover them.
Expected behavior
The multi-tensor swizzle succeeds for all shapes, like the single-tensor path.
Proposed fix
Add __launch_bounds__(TB_DIM * TB_DIM) to the four multi-tensor kernels, matching the other swizzle kernels, and add test shapes that reach the row<int2> and col<int4> variants. I'm happy to open a PR.
Environment overview
- Environment location: Docker on vast.ai
- Method of Transformer Engine install: from source (
mainat 5759fa0f) - Docker image: Ubuntu 24.04 with
/venv/main
Environment details
- OS version: Ubuntu 24.04
- PyTorch version: 2.11.0+cu128
- Python version: 3.12
- Transformer Engine version: 2.21.0.dev0+5759fa0f
- CUDA version: 12.8 (register table also with 13.4)
- CUDNN version: 9.19
Device details
- GPU model: NVIDIA GeForce RTX 5090 (CC 12.0), driver 580.95.05
- Dominant language
- Python
- Stars
- 3.6k
- Forks
- 844
- Avg merge
- 4d 42m
- Merged PRs (30d)
- 49
Getting set up
- No Dockerfile or Docker Compose file
- Has a pull request template
- Read the contributing guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
More from NVIDIA/TransformerEngine
-
[Bug] Backend selection picks FA3 for training with head_dim_qk=192 / v_head_dim=128, but FA3 backward cannot run itPossibly taken @yuweih205 claimed this 29 days ago. Openattention
Difficulty 2/5 1-3 hours Newbie friendliness 85/100
NVIDIA/TransformerEngine#3481 · 4 comments ·
Maintainers usually reply within 2 days
-
bug
Difficulty 2/5 1-3 hours Newbie friendliness 68/100
NVIDIA/TransformerEngine#2189 · 7 comments · 5 reactions ·
Maintainers usually reply within 2 days
-
[BUG] Grouped MXFP8 quantization is not concurrency safe with multiple streamsPossibly taken @kainzhong claimed this today. Openbug
Difficulty 4/5 3-5 days Newbie friendliness 25/100
NVIDIA/TransformerEngine#3630 ·
Maintainers usually reply within 2 days
-
[bug] NVFP4 + `torch.compile`: errors with 3D inputPossibly taken @pggPL claimed this today. Openbug
Difficulty 2/5 1-3 hours Newbie friendliness 25/100
NVIDIA/TransformerEngine#3626 ·
Maintainers usually reply within 2 days
-
[PyTorch] Avoid selecting FA4 for deterministic training on SM120Possibly taken A pull request linked to this issue is open or already merged. Openbug
Difficulty 3/5 1-2 days Newbie friendliness 35/100
NVIDIA/TransformerEngine#3594 ·
Maintainers usually reply within 2 days
All issues in NVIDIA/TransformerEngine
Similar issues
-
Device Details tables: FS/SF columns contradict each other (nfet_01v8 Vt row, pfet_01v8 Idsat row)Open
Difficulty 2/5 1-3 hours Newbie friendliness 75/100
google/skywater-pdk#450 ·
-
Drained trajectory arrays are overwritten when the sequence buffer is reusedPossibly taken @sylvesterkaczmarek claimed this today. Open
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
google-deepmind/bsuite#56 ·
-
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
LearningCircuit/local-deep-research#7206 ·
Maintainers usually reply within 1 day
-
Difficulty 2/5 1-3 hours Newbie friendliness 68/100
chingu-voyages/V62-tier3-team-33#285 ·
Maintainers usually reply within 1 day
-
Proxy drops log notifications from backends that don't send FastMCP's msg/extra dictPossibly taken @asasemahmed claimed this today. Openbug server
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
Maintainers usually reply within 1 day