pytorch / pytorch/pytorch

[PT2] Device-TMA reduction stores lose tma_min_block_sizes

Open
#196,557 0 comments 0 reactions 0 assignees View on GitHub
bot-triaged module: inductor oncall: pt2 release triage triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

## Describe the bug
`store_reduction()` bypasses `codegen_block_ptr()` and fails to record an in-kernel device-TMA descriptor. The resulting kernel can select XBLOCK=1 for an fp32 descriptor store, which Triton rejects because 1 * 4 < 16 bytes.

## Minimal reproducer
CUDA_VISIBLE_DEVICES=0 python3 repro_tma_store_block_shape_lt_16_bytes.py

## Expected behavior
The reduction compiles and matches eager while preserving TMA minimum block
size constraints.

## Error logs
File "/usr/local/lib/python3.12/dist-packages/triton/compiler/compiler.py", line 80, in make_ir
return ast_to_ttir(self.fn, self, context=context, options=options, codegen_fns=codegen_fns,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch._inductor.exc.InductorError: CompilationError: at 24:4:
for r0_offset in tl.range(0, r0_numel, R0_BLOCK):
r0_index = r0_offset + r0_base
r0_mask = r0_index < r0_numel
roffset = r0_offset
rindex = r0_index
r0_1 = r0_index
tmp0 = tl.load(in_ptr0 + (24576*((r0_1 + 8192*x0) // 12288) + (((r0_1 + 8192*x0) % 12288))), r0_mask & xmask, eviction_policy='evict_last', other=0.0)
tmp1 = tl.broadcast_to(tmp0, [XBLOCK, R0_BLOCK])
tmp3 = _tmp2 + tmp1
_tmp2 = tl.where(r0_mask & xmask, tmp3, _tmp2)
tmp2 = tl.sum(_tmp2, 1)[:, None]
tl.make_tensor_descriptor(out_ptr0, shape=[3], strides=[1], block_shape=[XBLOCK]).store([xoffset], tl.reshape(tl.broadcast_to(tmp2, [XBLOCK, 1]), [XBLOCK]).to(tl.float32))
^
Descriptor block shape must have at least 16 bytes in the last dimension, but got 1 * 4 = 4 bytes

## Root cause and proposed fix
`store_reduction()` bypasses `codegen_block_ptr()`. Mark the directly formatted
`TensorDescriptorOptions` descriptor as device TMA so codegen retains
`tma_min_block_sizes`; make the repaired padded-output case a passing regression.

[repro_tma_store_block_shape_lt_16_bytes.py](https://github.com/user-attachments/files/32044671/repro_tma_store_block_shape_lt_16_bytes.py)

### Versions

Collecting environment information...
PyTorch version: 2.15.0a0+gita7a0e15 (PyTorch main at commit a7a0e15894ee2c285d5065c3672291a19b0a8091 (built from source).)
Is debug build: False
CUDA used to build PyTorch: 13.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.3 LTS (x86_64) (docker container, host kernel 5.10.134)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0
CMake version: version 3.31.6
Libc version: glibc-2.39

Python version: 3.12.3 (64-bit runtime)
Python platform: Linux-5.10.134-16.3.al8.x86_64-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 13.1.115
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA H20 (sm_90, 96 GiB)
Nvidia driver version: 535.183.06
cuDNN version: 9.19.0

Versions of relevant libraries:
[pip3] numpy==2.1.0
[pip3] pytorch-triton==3.6.0+git9844da95.nv26.2
[pip3] torch==2.15.0a0+git40e238c
[pip3] triton_kernels==1.0.0+git9844da95.nv26.2

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo

Contributor guide

Open the contributing guide

Research direction

Start by tracing store_reduction() and its direct TensorDescriptorOptions formatting, then compare that path with codegen_block_ptr(). Use repro_tma_store_block_shape_lt_16_bytes.py to reproduce the descriptor block-size failure. Done means the padded-output reduction compiles, matches eager behavior, preserves TMA minimum block sizes, and has a passing regression.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.