[Correctness][Quant] per_channel_cast_fused mishandles expand tails

[GenTang]

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
• PyTorch reference implementation
• Existing test

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:

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

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:

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:

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.