deepseek-ai / deepseek-ai/TileKernels
[Correctness][Quant] per_channel_cast_fused mishandles expand tails
- Dominant language
- Python
- Stars
- 1.8k
- Forks
- 162
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
`per_channel_cast_fused()` accepts an expanded output length aligned to 16 tokens, while the TileLang kernel always processes 128-token tiles.
When:
- `pos_to_token is not None`
- `num_tokens_out % 16 == 0`
- `num_tokens_out % 128 != 0`
the last kernel block contains rows outside `num_tokens_out`. The out-of-range `pos_to_token` loads can be legalized as **zero** by TileLang, so those invalid rows are interpreted as references to token 0 and incorrectly participate in the per-channel `amax` reduction. This produces silently incorrect output scales.
Relevant code:
- [Kernel implementation](https://github.com/deepseek-ai/TileKernels/blob/main/tile_kernels/quant/per_channel_cast_fused_kernel.py)
- [PyTorch reference implementation](https://github.com/deepseek-ai/TileKernels/blob/main/tile_kernels/torch/per_channel_cast_fused.py)
- [Existing test](https://github.com/deepseek-ai/TileKernels/blob/main/tests/quant/test_per_channel_cast_fused.py)
## Root cause
Incorrect outputs can be observed when two conditions are satisfied:
1. num_tokens_out is not divisible by 128
2. Token 0 is not present within the valid entries of the final partial 128-row tile.
As previously noted, the outer wrapper supports arbitrary expanded lengths aligned to 16, while the underlying kernel uses a fixed `TILE_M = 128`
The mapping load does not check whether the row belongs to the logical output:
```python
if with_expand:
tmp = T.alloc_var(T.int32)
if k_id < VEC_M:
tmp = pos_to_token[k_id + m_offset]
for i in T.serial(VEC_M):
pos_to_token_local[i] = T.shfl_sync(tmp, i)
```
For the trailing partial tile, indices `k_id + m_offset` may exceed or equal `num_tokens_out`. These out-of-bounds loads return zero due to TileLang’s built-in memory legalization pass. The invalid entries are misinterpreted as `pos_to_token == 0`, causing token 0 to be erroneously loaded and included in the `amax_local` computation, which corrupts the final output scale.
## Minimal reproduction
```python
import torch
import tile_kernels
from tile_kernels.torch.per_channel_cast_fused import per_channel_cast_fused as torch_ref
N, H = 128, 128
x = torch.randn((N, H), dtype=torch.bfloat16, device="cuda")
# Token 0 should not be gathered, but give it a large value so that
# accidental inclusion in amax is easy to observe.
x[0].fill_(8)
# Length 16 is accepted by the wrapper, but the kernel uses TILE_M=128.
pos = torch.arange(1, 17, dtype=torch.int32, device="cuda").repeat(1)
_, sf = tile_kernels.quant.per_channel_cast_fused(
x, "e4m3", num_per_tokens=128, pos_to_token=pos)
_, sf_ref = torch_ref(
x, num_per_tokens=128, num_per_channels=None,
round_sf=False, pos_to_token=pos)
# Expected to fail
torch.testing.assert_close(sf, sf_ref, rtol=0, atol=0)
```
Observed result:
```
AssertionError: Tensor-likes are not equal!
Mismatched elements: 128 / 128 (100.0%)
Greatest absolute difference: 0.015136719681322575 at index (0, 105)
Greatest relative difference: 5.564102649688721 at index (0, 105)
```
No particular GPU architecture appears to be required for this reproduction. A CUDA GPU supported by TileLang/TileKernels should be sufficient.
## Why existing tests did not catch this
The expand tests generate `pos_to_token` using:
```python
tile_kernels.moe.get_fused_mapping(
topk_idx,
num_experts,
0,
128,
)
```
Each expert segment is aligned to 128 entries, so the resulting `pos_to_token.size(0)` is always divisible by 128. Therefore, the tests never exercise lengths accepted by the public API such as 16, 32, ...
## Proposed fix
Initialize the mapping value to the padding sentinel and only load an in-range output row:
```python
tmp = T.alloc_var(T.int32)
tmp = -1
row = k_id + m_offset
if k_id < VEC_M and row < num_tokens_out:
tmp = pos_to_token[row]
```
Invalid tail rows remain `-1` and follow the existing zero-fill path, so they do not participate in the `amax` reduction.
Contributor guide
No contributing guide indexed for this repository
Research direction
Read tile_kernels/quant/per_channel_cast_fused_kernel.py alongside tile_kernels/torch/per_channel_cast_fused.py, then run tests/quant/test_per_channel_cast_fused.py and the minimal reproduction. Add regression coverage for an expanded length such as 16 and verify the kernel output scales match the PyTorch reference without invalid tail rows affecting amax.
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
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 82/100