deepseek-ai / deepseek-ai/TileKernels

[Correctness][Quant] per_channel_cast_fused mishandles expand tails

Open
#26 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.