ByteDance-Seed / ByteDance-Seed/Triton-distributed

gqa_fwd_batch_decode: every non-power-of-two head dim fails to compile — the rope-split Q load indexes the wrong axis

Open
#200 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.5k
Forks
172
PR merge metrics
No merged PRs in 30d

Description

## Summary

`gqa_fwd_batch_decode` and `gqa_fwd_batch_decode_persistent` split the head dim into a power-of-two part plus a remainder (`BLOCK_DPE = q_head_dim - 2**int(log2(q_head_dim))`) precisely so that non-power-of-two head dims are supported — 96 and 192 are the usual MLA/DeepSeek shapes. That branch cannot run: the remainder ("rope split") Q load builds its offsets with `offs_dpe[:, None]` while the mask on the next line says `mask_dpe[None, :]`, so the two disagree and Triton rejects the shapes with `Cannot make_shape_compatible: incompatible dimensions at index 0`. Any head dim that is not a power of two raises at compile time, which makes the whole rope-split path dead.

The ordinary Q load six lines above puts the head-dim offsets on the last axis, and only the K-cache load below is legitimately `[:, None]` because K is loaded transposed. The same line appears in `kernel_gqa_fwd_batch_decode_split_kv_persistent` and in the mega-kernel copy at `python/triton_dist/mega_triton_kernel/kernels/flash_decode.py`.

https://github.com/ByteDance-Seed/Triton-distributed/blob/8260bc34398c2b8f36dc840fd22f741ca9294584/python/triton_dist/kernels/nvidia/flash_decode.py#L187-L191

## Reproduction

```python
import torch
from triton_dist.kernels.nvidia import gqa_fwd_batch_decode

NUM_BLOCKS, KV_LENS = 4096, [463, 18, 1320]

def run(head_size, num_query_heads=32, num_kv_heads=8, block_size=1):
torch.manual_seed(0)
q = torch.randn(len(KV_LENS), num_query_heads, head_size, dtype=torch.float16, device="cuda")
kv = torch.randn(NUM_BLOCKS, 2, block_size, num_kv_heads, head_size, dtype=torch.float16, device="cuda")
tables = torch.randint(0, NUM_BLOCKS, (len(KV_LENS), max(KV_LENS)), dtype=torch.int32, device="cuda")
ws = torch.zeros([len(KV_LENS) * num_query_heads * 32], dtype=torch.int32, device="cuda")
gqa_fwd_batch_decode(q, kv[:, 0].contiguous(), kv[:, 1].contiguous(), ws, [1] * len(KV_LENS),
torch.tensor(KV_LENS, dtype=torch.int32, device="cuda"), tables, head_size**-0.5, 0.0)

for head_size in (128, 256, 96, 192):
try:
run(head_size)
print(f"head_size={head_size:3d} ok")
except Exception as exc:
print(f"head_size={head_size:3d} {type(exc).__name__}: {str(exc).strip().splitlines()[-1]}")
```

```
head_size=128 ok
head_size=256 ok
head_size= 96 CompilationError: ValueError('Cannot make_shape_compatible: incompatible dimensions at index 0: 16 and 32')
head_size=192 CompilationError: ValueError('Cannot make_shape_compatible: incompatible dimensions at index 0: 16 and 64')
```

`gqa_fwd_batch_decode_persistent` fails identically at the same two head dims. `16` is `BLOCK_H`; `32` and `64` are `BLOCK_DPE`.

## Environment

- Repo commit tested: `8260bc34398c2b8f36dc840fd22f741ca9294584` (main)
- GPU: NVIDIA B200 (sm_100), driver 595.71.05, CUDA 13.2
- torch 2.12.1+cu130, the repo's vendored Triton 3.4.0, Python 3.11

Contributor guide

Open the contributing guide

Research direction

Start with the rope-split Q load around lines 187-191 of python/triton_dist/kernels/nvidia/flash_decode.py, then inspect the matching line in kernel_gqa_fwd_batch_decode_split_kv_persistent and python/triton_dist/mega_triton_kernel/kernels/flash_decode.py. Run the supplied reproduction for head sizes 128, 256, 96, and 192; done means the non-power-of-two cases compile and run for both batch-decode paths.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
74/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.