From bff98f4f938e0e2c9ae33cf801530ee0f358f9be Mon Sep 17 00:00:00 2001 From: Gen TANG Date: Tue, 28 Jul 2026 18:50:38 +0800 Subject: [PATCH 1/3] add test case to cover expand tails issue --- tests/quant/test_per_channel_cast_fused.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/quant/test_per_channel_cast_fused.py b/tests/quant/test_per_channel_cast_fused.py index 5acf0c1..eaff846 100644 --- a/tests/quant/test_per_channel_cast_fused.py +++ b/tests/quant/test_per_channel_cast_fused.py @@ -13,7 +13,7 @@ os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0' -def generate_test_data(params): +def generate_test_data(params, alignment=128): num_send_tokens = params['num_send_tokens'] num_topk = params['num_topk'] num_experts = params['num_experts'] @@ -28,7 +28,7 @@ def generate_test_data(params): topk_idx = generate_topk_idx(params) num_tokens = topk_idx.shape[0] _, pos_to_token, _, token_topk_to_pos, _, _, _, _ = ( - tile_kernels.moe.get_fused_mapping(topk_idx, num_experts, 0, 128) + tile_kernels.moe.get_fused_mapping(topk_idx, num_experts, 0, alignment) ) x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') x = tile_kernels.moe.expand_to_fused(x, token_topk_to_pos, pos_to_token) @@ -74,8 +74,10 @@ def generate_test_params(is_benchmark: bool) -> list[dict]: @pytest.mark.parametrize('params', generate_test_params(is_benchmark=False), ids=make_param_id) -def test_per_channel_cast_fused(params): - _, _, _, func, func_ref = generate_test_data(params) +@pytest.mark.parametrize('alignment', [16, 128]) +# use alignment=16 to cover expand tails issues. +def test_per_channel_cast_fused(params, alignment): + _, _, _, func, func_ref = generate_test_data(params, alignment) x_fp8, x_fp8_sf = func() x_fp8_ref, x_fp8_sf_ref = func_ref() From 8655a64c7af4e68c7c1e173fcec9933d6c644064 Mon Sep 17 00:00:00 2001 From: Gen TANG Date: Tue, 28 Jul 2026 21:25:32 +0800 Subject: [PATCH 2/3] fix expand tails bug --- tile_kernels/quant/per_channel_cast_fused_kernel.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tile_kernels/quant/per_channel_cast_fused_kernel.py b/tile_kernels/quant/per_channel_cast_fused_kernel.py index cce18cf..de86a69 100644 --- a/tile_kernels/quant/per_channel_cast_fused_kernel.py +++ b/tile_kernels/quant/per_channel_cast_fused_kernel.py @@ -71,8 +71,12 @@ def per_channel_cast_fused_kernel( T.assume(num_tokens_out % 128 == 0 or (with_expand and num_tokens_out % 16 == 0)) if with_expand: tmp = T.alloc_var(T.int32) - if k_id < VEC_M: - tmp = pos_to_token[k_id + m_offset] + # Invalid tail rows remain -1 and follow the existing zero-fill path, + # so they do not participate in the amax reduction. + tmp = -1 + row = k_id + m_offset + if k_id < VEC_M and row < num_tokens_out: + tmp = pos_to_token[row] for i in T.serial(VEC_M): pos_to_token_local[i] = T.shfl_sync(tmp, i) From bd36724d1215510e280c28b6f0deffcb2d5d5057 Mon Sep 17 00:00:00 2001 From: Gen TANG Date: Tue, 28 Jul 2026 22:22:59 +0800 Subject: [PATCH 3/3] add store bounds guard --- tile_kernels/quant/per_channel_cast_fused_kernel.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tile_kernels/quant/per_channel_cast_fused_kernel.py b/tile_kernels/quant/per_channel_cast_fused_kernel.py index de86a69..91cb31f 100644 --- a/tile_kernels/quant/per_channel_cast_fused_kernel.py +++ b/tile_kernels/quant/per_channel_cast_fused_kernel.py @@ -132,8 +132,10 @@ def per_channel_cast_fused_kernel( out_local[j] = in_local[j] * sf_invs_local[i] * amax_local[j] else: out_local[j] = in_local[j] * amax_local[j] - for j in T.vectorized(VEC_K): - out[i + m_offset, j + k_offset] = out_local[j] + out_row = i + m_offset + if out_row < num_tokens_out: + for j in T.vectorized(VEC_K): + out[out_row, j + k_offset] = out_local[j] return per_channel_cast_fused_kernel