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