From 82fbd40133768c1989d9bcfc3087ae20bac5d7e0 Mon Sep 17 00:00:00 2001 From: JayceSu98 Date: Fri, 12 Jun 2026 06:30:17 +0000 Subject: [PATCH] [BugFix][Quant] Guard scale tails and FP32 absmax The full H100/CuTeDSL correctness run covers token and hidden sizes that do not land exactly on the kernel tile, scaling-factor block, or E5M6 packing granularity. Those edge tiles exposed two independent quantization issues. First, cast-back and lossless/per-token cast kernels could load or store scaling factors for padded token/channel blocks. The same edge condition also applied to packed E5M6 output rows and columns, where the final tile may not contain a full logical token row or hidden group. Second, per-token scale generation used in_config.dtype as the absmax fragment dtype. That dtype describes the input tensor storage, not the reduction accumulator. For BF16 and narrow quantization paths, keeping absmax in the input dtype can round or overflow the scale path before the final output cast, producing byte mismatches against the PyTorch reference. Guard scale loads/stores and packed E5M6 stores with the real token/channel extents, keep invalid scale lanes zero-initialized, and use FP32 fragments for per-token absmax reductions before computing the reciprocal scale. Co-authored-by: dingsg --- tile_kernels/quant/cast_back_kernel.py | 9 +++++++-- .../quant/per_block_cast_lossless_kernel.py | 18 ++++++++++++------ tile_kernels/quant/per_token_cast_kernel.py | 8 +++++--- .../quant/per_token_cast_to_e5m6_kernel.py | 12 ++++++++---- 4 files changed, 32 insertions(+), 15 deletions(-) diff --git a/tile_kernels/quant/cast_back_kernel.py b/tile_kernels/quant/cast_back_kernel.py index 69d6a6f..36c47af 100644 --- a/tile_kernels/quant/cast_back_kernel.py +++ b/tile_kernels/quant/cast_back_kernel.py @@ -60,11 +60,16 @@ def cast_back_kernel( out_fragment = T.alloc_fragment((TILE_M, TILE_K), out_dtype) T.copy(x[pid_token * TILE_M, pid_hidden * TILE_K], x_shared, disable_tma=True) + num_sf_token_blocks = T.ceildiv(num_tokens, num_per_tokens) + num_sf_channel_blocks = T.ceildiv(hidden, num_per_channels) for i, j in T.Parallel(T.ceildiv(TILE_M, num_per_tokens), T.ceildiv(TILE_K, num_per_channels)): token_index = pid_token * TILE_M // num_per_tokens + i channel_index = pid_hidden * TILE_K // num_per_channels + j - sf = load_sf(x_sf, token_index, channel_index, in_config) - sf_shared[i, j] = transform_sf(sf, in_config) + if token_index < num_sf_token_blocks and channel_index < num_sf_channel_blocks: + sf = load_sf(x_sf, token_index, channel_index, in_config) + sf_shared[i, j] = transform_sf(sf, in_config) + else: + sf_shared[i, j] = 0 for i, j in T.Parallel(TILE_M, TILE_K): out_fragment[i, j] = x_shared[i, j] * sf_shared[i // num_per_tokens, j // num_per_channels] diff --git a/tile_kernels/quant/per_block_cast_lossless_kernel.py b/tile_kernels/quant/per_block_cast_lossless_kernel.py index 9645aef..84f9c70 100644 --- a/tile_kernels/quant/per_block_cast_lossless_kernel.py +++ b/tile_kernels/quant/per_block_cast_lossless_kernel.py @@ -88,10 +88,13 @@ def per_block_cast_lossless_kernel( # Load scaling factor of x to fragment T.fill(x_sf_fragment, 0) + num_in_sf_blocks_m = T.ceildiv(num_tokens, in_config.sf_block[0]) + num_in_sf_blocks_k = T.ceildiv(hidden, in_config.sf_block[1]) for i, j in T.Parallel(num_in_sf_per_block_m, num_in_sf_per_block_k): m_idx = pid_token * block_m // in_config.sf_block[0] + i k_idx = pid_hidden * block_k // in_config.sf_block[1] + j - x_sf_fragment[i, j] = load_sf(x_sf, m_idx, k_idx, in_config) + if m_idx < num_in_sf_blocks_m and k_idx < num_in_sf_blocks_k: + x_sf_fragment[i, j] = load_sf(x_sf, m_idx, k_idx, in_config) # Alloc fragments x_sf_uint32_fragment = T.alloc_fragment((num_in_sf_per_block_m, num_in_sf_per_block_k), T.uint32) @@ -137,14 +140,17 @@ def per_block_cast_lossless_kernel( x_out_fragment[i, j] = T.cast(T.float32(x_in_shared[i, j]) * sf, out_config.dtype) # Store scaling factor back to global memory + num_out_sf_blocks_m = T.ceildiv(num_tokens, out_config.sf_block[0]) + num_out_sf_blocks_k = T.ceildiv(hidden, out_config.sf_block[1]) for i, j in T.Parallel(num_out_sf_per_block_m, num_out_sf_per_block_k): sf_m_idx = pid_token * num_out_sf_per_block_m + i sf_k_idx = pid_hidden * num_out_sf_per_block_k + j - if out_config.use_packed_ue8m0: - sf = T.uint8(out_sf_uint32_fragment[i, j]) - else: - sf = transform_sf_to_fp32(out_sf_uint32_fragment[i, j]) - store_sf(out_sf, sf, sf_m_idx, sf_k_idx, out_config) + if sf_m_idx < num_out_sf_blocks_m and sf_k_idx < num_out_sf_blocks_k: + if out_config.use_packed_ue8m0: + sf = T.uint8(out_sf_uint32_fragment[i, j]) + else: + sf = transform_sf_to_fp32(out_sf_uint32_fragment[i, j]) + store_sf(out_sf, sf, sf_m_idx, sf_k_idx, out_config) T.copy(x_out_fragment, out[pid_token * block_m: (pid_token + 1) * block_m, pid_hidden * block_k: (pid_hidden + 1) * block_k]) diff --git a/tile_kernels/quant/per_token_cast_kernel.py b/tile_kernels/quant/per_token_cast_kernel.py index 1678648..7234b7f 100644 --- a/tile_kernels/quant/per_token_cast_kernel.py +++ b/tile_kernels/quant/per_token_cast_kernel.py @@ -115,7 +115,8 @@ def per_token_cast_kernel( # Store SF m_idx = pid_token * block_m + i k_idx = pid_hidden * num_groups + j - store_sf(out_sf, sf, m_idx, k_idx, out_config) + if m_idx < num_tokens: + store_sf(out_sf, sf, m_idx, k_idx, out_config) sf_inv_fragment[i, j] = sf_inv # Store casted values @@ -131,7 +132,7 @@ def per_token_cast_kernel( sf = load_sf(out_sf, pid_token * block_m + i, pid_hidden * num_groups + j, out_config) sf_inv_fragment[i, j] = 1 / sf else: - amax_fragment = T.alloc_fragment((block_m, num_groups), in_config.dtype) + amax_fragment = T.alloc_fragment((block_m, num_groups), T.float32) x_fragment_reshaped = T.reshape(x_fragment, [block_m, num_groups, num_per_channels]) # Reduce SF T.reduce_absmax(x_fragment_reshaped, amax_fragment, dim=2) @@ -142,7 +143,8 @@ def per_token_cast_kernel( # Store SF m_idx = pid_token * block_m + i k_idx = pid_hidden * num_groups + j - store_sf(out_sf, sf, m_idx, k_idx, out_config) + if m_idx < num_tokens: + store_sf(out_sf, sf, m_idx, k_idx, out_config) sf_inv_fragment[i, j] = sf_inv # Store casted values diff --git a/tile_kernels/quant/per_token_cast_to_e5m6_kernel.py b/tile_kernels/quant/per_token_cast_to_e5m6_kernel.py index 0956c13..01be08a 100644 --- a/tile_kernels/quant/per_token_cast_to_e5m6_kernel.py +++ b/tile_kernels/quant/per_token_cast_to_e5m6_kernel.py @@ -119,7 +119,7 @@ def per_token_cast_to_e5m6_kernel( # Copy input into registers T.copy(x[pid_token * block_m, pid_hidden * block_k], x_fragment) - amax_fragment = T.alloc_fragment((block_m, num_groups), in_config.dtype) + amax_fragment = T.alloc_fragment((block_m, num_groups), T.float32) x_fragment_reshaped = T.reshape(x_fragment, [block_m, num_groups, num_per_channels]) # Reduce SF T.reduce_absmax(x_fragment_reshaped, amax_fragment, dim=2) @@ -130,7 +130,8 @@ def per_token_cast_to_e5m6_kernel( # Store SF m_idx = pid_token * block_m + i k_idx = pid_hidden * num_groups + j - store_sf(out_sf, sf, m_idx, k_idx, out_config) + if m_idx < num_tokens: + store_sf(out_sf, sf, m_idx, k_idx, out_config) sf_inv_fragment[i, j] = sf_inv T.annotate_layout({ @@ -150,8 +151,11 @@ def per_token_cast_to_e5m6_kernel( for j in T.serial(8): in_local[j] = out_fragment[x, y * 8 + j] float_to_e5m6(in_local, out_local) - for j in T.serial(3): - out[pid_token * block_m + x, pid_hidden * (block_k // 8 * 3) + y * 3 + j] = out_local[j] + m_idx = pid_token * block_m + x + k_idx = pid_hidden * block_k + y * 8 + if m_idx < num_tokens and k_idx < hidden: + for j in T.serial(3): + out[m_idx, (k_idx // 8) * 3 + j] = out_local[j] return per_token_cast_to_e5m6_kernel