From 3d778dc2a63794e09cb4fbcfe83c6afa0abe1254 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Wed, 7 Jan 2026 18:11:38 +0100 Subject: [PATCH 1/9] improve swizzling --- mlx/backend/cuda/quantized/qqmm_utils.cu | 193 +++++++++++++---------- 1 file changed, 110 insertions(+), 83 deletions(-) diff --git a/mlx/backend/cuda/quantized/qqmm_utils.cu b/mlx/backend/cuda/quantized/qqmm_utils.cu index ff19057b08..29f5ddd10e 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -43,95 +43,123 @@ namespace cg = cooperative_groups; // [252, 253, 254, 255], // [380, 381, 382, 383], // [508, 509, 510, 511]]]]], -__device__ size_t -scale_tiled_offset(size_t scale_index, size_t num_rows, size_t num_scale_cols) { - // Compute the tiled layout offset for scale factors used in tensor cores - // This function maps from a linear scale index to the tiled layout expected - // by tensor cores (and cublaslt). - // - // Input: linear scale index (e.g., for a matrix M x K with group_size, - // scale_index ranges from 0 to (M * K/group_size - 1)) - // - // The tiled layout organizes scales into tiles of 128 rows x 4 columns, - // where each tile is subdivided into 4 sub-blocks of 32 rows x 4 columns. - size_t row = scale_index / num_scale_cols; - size_t col = scale_index % num_scale_cols; - - constexpr size_t rows_per_tile = 128; - constexpr size_t rows_per_sub_block = 32; - constexpr size_t cols_per_sub_block = 4; - constexpr size_t sub_blocks_per_tile = 4; // Vertically stacked - - // Decompose row position - size_t tile_row = row / rows_per_tile; // Which tile row - size_t row_in_tile = row % rows_per_tile; // Row within tile - size_t sub_block_row = - row_in_tile / rows_per_sub_block; // Sub-block within tile - size_t row_in_sub_block = - row_in_tile % rows_per_sub_block; // Row in sub-block - - // Decompose column position - size_t col_tile = col / cols_per_sub_block; // Which column tile - size_t col_in_sub_block = col % cols_per_sub_block; // Column within sub-block - - // Compute tile index and offset within tile - size_t num_col_tiles = cuda::ceil_div(num_scale_cols, cols_per_sub_block); - size_t tile_idx = tile_row * num_col_tiles + col_tile; - - size_t offset_in_tile = - (row_in_sub_block * sub_blocks_per_tile * cols_per_sub_block) + - (sub_block_row * cols_per_sub_block) + col_in_sub_block; - - constexpr size_t tile_size = rows_per_tile * cols_per_sub_block; - return tile_idx * tile_size + offset_in_tile; + +inline std::tuple get_swizzle_launch_args( + size_t M_swizzled, + size_t K_swizzled, + int tile_rows = 128, + int tile_cols = 4, + int tiles_per_lane = 1) { + constexpr int lanes_per_block = 32 // 32 threads per warp + int lanes_per_block = + tiles_per_lane* lanes_per_block // tiles_per_lane = 1 if load int, 4 + // if int4 + const int tiles_per_block = lanes_per_block * tiles_per_lane; + const int warps_per_block = tile_rows / 4; // 128 / 4 = 32 + + const int num_tiles_k = K_swizzled / tile_cols; + const int num_tiles_m = M_swizzled / tile_rows; + + dim3 grid; + grid.x = cuda::ceil_div(num_tiles_k, tiles_per_block); + grid.y = num_tiles_m; + grid.z = 1; + + // Block is always (32, 32) = 1024 threads + dim3 block(lanes_per_block, warps_per_block, 1); + + return std::make_tuple(grid, block); } namespace cu { -__global__ void repack_scales( +__global__ void swizzle_scaling_factors_impl( const uint8_t* scales_linear, - uint8_t* scales_tiled, - size_t input_rows, - size_t input_cols, - size_t output_rows, - size_t output_cols) { - auto block_size = cg::this_thread_block().dim_threads(); + uint8_t* scales_swizzled, + const size_t M, + const size_t K, + const size_t M_swizzled, + const size_t K_swizzled) { + // M_swizzled and K_swizzled are dimensions of scales_tiled array + // (padded to full tiles 128x4 if M or K are not multiples of tile sizes) + + // Tile dimensions for scale factors + constexpr int tile_dim_row = 128; + constexpr int tile_dim_col = 4; + constexpr int tile_size = tile_dim_row * tile_dim_col; // 512 bytes + constexpr int num_tile_rows_per_thread = 4; // always 4 + constexpr int num_tiles_per_thread = 1; + constexpr int int4_per_tile = + tile_dim_row* tile_dim_col / 16 constexpr int lanes_per_block = 32; + constexpr int num_tiles_per_block = lanes_per_block * num_tiles_per_thread; + // Each thread loads 4 rows of 4 bytes x 1 column of scales (16 bytes -- 16 + // scales) thread (0, 0) loads scales at rows 0,32,64,96 of tile 0 thread (1, + // 0) loads rows 0,32,64,96 of of tile 1 therefore a warp loads: consecutive 4 + // bytes x 32 = 128 bytes 4 times with a stride 32 this can be improved + // further with processing 2/4 tiles within a warp: load 8/16 bytes but it + // will be a bit more complicated, because then we would need to shuffle + // registers so for now each thread process exactly 1 tile + + auto block_size = cg::this_thread_block().dim_threads(); // (32, 32, 1) auto block_idx = cg::this_thread_block().group_index(); auto idx_in_block = cg::this_thread_block().thread_index(); - auto tidx = block_idx.x * block_size.x + idx_in_block.x; - auto tidy = block_idx.y * block_size.y + idx_in_block.y; - - auto grid_dim_x = - cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x; - - size_t output_index = tidx + grid_dim_x * size_t(tidy); - size_t output_size = output_rows * output_cols; - - if (output_index >= output_size) { - return; + auto tidx = idx_in_block.x; // lane within warp (0, 31) + auto tidy = idx_in_block.y; // warp index (0, 31) + auto linear_tid = tidy * block_size.x + tidx; + + const int bid_x = block_idx.x; + const int bid_y = block_idx.y; + + // [M tile * 128 * bytes] + [K tile * 32] + const int K_int = K_swizzled / 4; + // incase of overflow cast to size_t + const size_t input_block_offset = + static_cast(bid_y) * tile_dim_row * K_int + + static_cast(bid_x) * num_tiles_per_block; + const int* input_block = + reinterpret_cast(scales_linear) + input_block_offset; + + const size_t output_offset = + static_cast(bid_y) * tile_dim_row * K_int + + static_cast(bid_x) * num_tiles_per_block * tile_size / 4; + + int* output_block = reinterpret_cast(scales_swizzled) + output_offset; + + const int num_tiles_k = K_swizzled / tile_dim_col; + const int grid_dim_x = cg::this_grid().dim_blocks().x; + const int grid_dim_y = cg::this_grid().dim_blocks().y; + + // each thread store 16 consecutive scales, but this store is strided within a + // warp so we will write first to the shared memory (each thread 4 int32 -- 4 + // rows -- 16 scales) then it will be coaleased store shared -> global within + // a thread block + extern __shared__ int4 strided_scales_thread[]; + // load + int thread_tile_rows[num_tile_rows_per_thread]; +#pragma unroll + for (int i = 0; i < num_tile_rows_per_thread; i++) { + const int thread_offset = + (i * block_size.x + tidy) * K_int + tidx * num_tiles_per_thread; + thread_tile_rows[i] = __ldg(input_block + thread_offset); } - - size_t tiled_offset = - scale_tiled_offset(output_index, output_rows, output_cols); - - size_t row = output_index / output_cols; - size_t col = output_index % output_cols; - - // Probably this can be done better with 2 separated paths for valid and - // padding - if (row < input_rows && col < input_cols) { - size_t input_index = row * input_cols + col; - scales_tiled[tiled_offset] = scales_linear[input_index]; - } else { - // Zero-fill padding region - scales_tiled[tiled_offset] = 0; + // todo add padding + boundaries + // write 4 ints to the shared memory + strided_scales_thread[tidx * tile_size / 16 + tidy] = + *reinterpret_cast(thread_tile_rows); + __syncthreads(); + + // shared -> global + __align__(16) int4* output_block_int4 = reinterpret_cast(output_block); + const int total_int4s = tiles_in_block * tile_size / 16; +#pragma unroll + for (int i = linear_tid; i < total_int4s; i += block_size.x * block_size.y) { + output_block_int4[i] = strided_scales_thread[i]; } } - } // namespace cu -void repack_scales( +void swizzle_scaling_factors( const array& scales, array& scales_tiled, cu::CommandEncoder& enc, @@ -140,24 +168,23 @@ void repack_scales( enc.set_output_array(scales_tiled); // Note: scales_tiled is padded to full tiles so if num_rows or num_cols - // are not multiples of tile sizes, the extra space is filled with zeros + // are not multiples of tile sizes size_t input_rows = scales.shape(-2); size_t input_cols = scales.shape(-1); size_t output_rows = scales_tiled.shape(-2); size_t output_cols = scales_tiled.shape(-1); - size_t output_size = output_rows * output_cols; - bool large = output_size > UINT_MAX; - auto [num_blocks, block_dims] = get_launch_args( - output_size, scales_tiled.shape(), scales_tiled.strides(), large); + uint32_t smem_bytes = 128 * 4 * 32; + auto [num_blocks, block_dims] = + get_swizzle_launch_args(output_rows, output_cols); enc.add_kernel_node( - cu::repack_scales, + cu::swizzle_scaling_factors_impl, num_blocks, block_dims, - 0, + smem_bytes, gpu_ptr(scales), gpu_ptr(scales_tiled), input_rows, From db08b8765c97d258fc40f309ad914a1fa6f73c4a Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Thu, 8 Jan 2026 01:52:29 +0100 Subject: [PATCH 2/9] added padding, fixed typos --- mlx/backend/cuda/quantized/qqmm_utils.cu | 78 +++++++++++++----------- 1 file changed, 44 insertions(+), 34 deletions(-) diff --git a/mlx/backend/cuda/quantized/qqmm_utils.cu b/mlx/backend/cuda/quantized/qqmm_utils.cu index 29f5ddd10e..96fa637026 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -50,11 +50,8 @@ inline std::tuple get_swizzle_launch_args( int tile_rows = 128, int tile_cols = 4, int tiles_per_lane = 1) { - constexpr int lanes_per_block = 32 // 32 threads per warp - int lanes_per_block = - tiles_per_lane* lanes_per_block // tiles_per_lane = 1 if load int, 4 - // if int4 - const int tiles_per_block = lanes_per_block * tiles_per_lane; + constexpr int lanes_per_block = 32; // 32 threads per warp + const int tiles_per_block = lanes_per_block * tiles_per_lane; const int warps_per_block = tile_rows / 4; // 128 / 4 = 32 const int num_tiles_k = K_swizzled / tile_cols; @@ -73,7 +70,7 @@ inline std::tuple get_swizzle_launch_args( namespace cu { -__global__ void swizzle_scaling_factors_impl( +__global__ void swizzle_scales( const uint8_t* scales_linear, uint8_t* scales_swizzled, const size_t M, @@ -89,17 +86,12 @@ __global__ void swizzle_scaling_factors_impl( constexpr int tile_size = tile_dim_row * tile_dim_col; // 512 bytes constexpr int num_tile_rows_per_thread = 4; // always 4 constexpr int num_tiles_per_thread = 1; - constexpr int int4_per_tile = - tile_dim_row* tile_dim_col / 16 constexpr int lanes_per_block = 32; + constexpr int lanes_per_block = 32; constexpr int num_tiles_per_block = lanes_per_block * num_tiles_per_thread; // Each thread loads 4 rows of 4 bytes x 1 column of scales (16 bytes -- 16 // scales) thread (0, 0) loads scales at rows 0,32,64,96 of tile 0 thread (1, // 0) loads rows 0,32,64,96 of of tile 1 therefore a warp loads: consecutive 4 - // bytes x 32 = 128 bytes 4 times with a stride 32 this can be improved - // further with processing 2/4 tiles within a warp: load 8/16 bytes but it - // will be a bit more complicated, because then we would need to shuffle - // registers so for now each thread process exactly 1 tile - + // bytes x 32 = 128 bytes 4 times with a stride 32 auto block_size = cg::this_thread_block().dim_threads(); // (32, 32, 1) auto block_idx = cg::this_thread_block().group_index(); auto idx_in_block = cg::this_thread_block().thread_index(); @@ -114,11 +106,11 @@ __global__ void swizzle_scaling_factors_impl( // [M tile * 128 * bytes] + [K tile * 32] const int K_int = K_swizzled / 4; // incase of overflow cast to size_t - const size_t input_block_offset = + const size_t block_offset = static_cast(bid_y) * tile_dim_row * K_int + static_cast(bid_x) * num_tiles_per_block; const int* input_block = - reinterpret_cast(scales_linear) + input_block_offset; + reinterpret_cast(scales_linear) + block_offset; const size_t output_offset = static_cast(bid_y) * tile_dim_row * K_int + @@ -130,43 +122,62 @@ __global__ void swizzle_scaling_factors_impl( const int grid_dim_x = cg::this_grid().dim_blocks().x; const int grid_dim_y = cg::this_grid().dim_blocks().y; - // each thread store 16 consecutive scales, but this store is strided within a - // warp so we will write first to the shared memory (each thread 4 int32 -- 4 - // rows -- 16 scales) then it will be coaleased store shared -> global within - // a thread block + bool pad_rows = (bid_y == grid_dim_y - 1) && + (M < M_swizzled); // if the last and is partial + bool pad_cols = (bid_x == grid_dim_x - 1) && + (K < K_swizzled); // if the last and is partial + + int num_tiles_per_block_ = num_tiles_per_block; + if (bid_x == grid_dim_x - 1) { + num_tiles_per_block_ = (K_int - 1) % num_tiles_per_block + 1; + } + bool valid_tile = threadIdx.x * num_tiles_per_thread < num_tiles_per_block_; + // Each thread loads 16 scales from 4 rows (stride 32) and packs them into + // int4. The store is strided within a warp (stride 32 int4s), so we first + // write to shared memory, then do a coalesced store from shared to global extern __shared__ int4 strided_scales_thread[]; // load int thread_tile_rows[num_tile_rows_per_thread]; + if (valid_tile) { #pragma unroll - for (int i = 0; i < num_tile_rows_per_thread; i++) { - const int thread_offset = - (i * block_size.x + tidy) * K_int + tidx * num_tiles_per_thread; - thread_tile_rows[i] = __ldg(input_block + thread_offset); + for (int i = 0; i < num_tile_rows_per_thread; i++) { + const int thread_offset = + (i * block_size.x + tidy) * K_int + tidx * num_tiles_per_thread; + thread_tile_rows[i] = __ldg(input_block + thread_offset); + if (pad_rows || pad_cols) { + // check bytes we need to pad + for (int j = 0; j < num_tile_rows_per_thread * sizeof(int); j++) { + const size_t element_idx = + (block_offset + thread_offset) * sizeof(int) + j; + if (element_idx / K_swizzled >= M || + (element_idx % K_swizzled) >= K) { + reinterpret_cast(&thread_tile_rows[i])[j] = 0; + } + } + } + // write 4 ints to the shared memory + } + strided_scales_thread[tidx * tile_size / 16 + tidy] = + *reinterpret_cast(thread_tile_rows); } - // todo add padding + boundaries - // write 4 ints to the shared memory - strided_scales_thread[tidx * tile_size / 16 + tidy] = - *reinterpret_cast(thread_tile_rows); __syncthreads(); // shared -> global - __align__(16) int4* output_block_int4 = reinterpret_cast(output_block); - const int total_int4s = tiles_in_block * tile_size / 16; + const int total_int4s = num_tiles_per_block_ * tile_size / 16; #pragma unroll for (int i = linear_tid; i < total_int4s; i += block_size.x * block_size.y) { - output_block_int4[i] = strided_scales_thread[i]; + reinterpret_cast(output_block)[i] = strided_scales_thread[i]; } } } // namespace cu -void swizzle_scaling_factors( +void swizzle_scales( const array& scales, array& scales_tiled, cu::CommandEncoder& enc, const Stream& s) { enc.set_input_array(scales); enc.set_output_array(scales_tiled); - // Note: scales_tiled is padded to full tiles so if num_rows or num_cols // are not multiples of tile sizes @@ -179,9 +190,8 @@ void swizzle_scaling_factors( uint32_t smem_bytes = 128 * 4 * 32; auto [num_blocks, block_dims] = get_swizzle_launch_args(output_rows, output_cols); - enc.add_kernel_node( - cu::swizzle_scaling_factors_impl, + cu::swizzle_scales, num_blocks, block_dims, smem_bytes, From 53182060ee541f66acdb28f2a22b95006564416e Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Thu, 8 Jan 2026 14:08:01 +0100 Subject: [PATCH 3/9] refactoring --- mlx/backend/cuda/quantized/qqmm.cpp | 8 ++++---- mlx/backend/cuda/quantized/qqmm_utils.cu | 11 +++++------ 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/mlx/backend/cuda/quantized/qqmm.cpp b/mlx/backend/cuda/quantized/qqmm.cpp index e7260da8c1..b57031d028 100644 --- a/mlx/backend/cuda/quantized/qqmm.cpp +++ b/mlx/backend/cuda/quantized/qqmm.cpp @@ -46,7 +46,7 @@ inline array ensure_row_contiguous_matrix( return x_copy; } -array pad_and_repack_scales( +array pad_and_swizzle_scales( const array& scale, cu::CommandEncoder& encoder, const Stream& s) { @@ -64,7 +64,7 @@ array pad_and_repack_scales( cu::malloc_async(pad_outer * pad_inner, encoder), Shape{pad_outer, pad_inner}, scale.dtype()); - repack_scales(scale, scale_tiled, encoder, s); + swizzle_scales(scale, scale_tiled, encoder, s); encoder.add_temporary(scale_tiled); return scale_tiled; @@ -176,8 +176,8 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { int K = K_packed * (32 / bits_); // Repack scales from linear to tiled layout for tensor cores - array scale_x = pad_and_repack_scales(scale_x_pre, encoder, s); - array scale_w = pad_and_repack_scales(scale_w_pre, encoder, s); + array scale_x = pad_and_swizzle_scales(scale_x_pre, encoder, s); + array scale_w = pad_and_swizzle_scales(scale_w_pre, encoder, s); bool x_transposed = false; bool w_transposed = true; // always transposed diff --git a/mlx/backend/cuda/quantized/qqmm_utils.cu b/mlx/backend/cuda/quantized/qqmm_utils.cu index 96fa637026..17250f5b58 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -64,8 +64,8 @@ inline std::tuple get_swizzle_launch_args( // Block is always (32, 32) = 1024 threads dim3 block(lanes_per_block, warps_per_block, 1); - - return std::make_tuple(grid, block); + int shared_mem_bytes = tile_rows * tile_cols * tiles_per_block; + return std::make_tuple(grid, block, shared_mem_bytes); } namespace cu { @@ -155,8 +155,8 @@ __global__ void swizzle_scales( } } } - // write 4 ints to the shared memory } + // write 4 ints to the shared memory strided_scales_thread[tidx * tile_size / 16 + tidy] = *reinterpret_cast(thread_tile_rows); } @@ -187,14 +187,13 @@ void swizzle_scales( size_t output_rows = scales_tiled.shape(-2); size_t output_cols = scales_tiled.shape(-1); - uint32_t smem_bytes = 128 * 4 * 32; - auto [num_blocks, block_dims] = + auto [num_blocks, block_dims, shared_mem_bytes] = get_swizzle_launch_args(output_rows, output_cols); enc.add_kernel_node( cu::swizzle_scales, num_blocks, block_dims, - smem_bytes, + shared_mem_bytes, gpu_ptr(scales), gpu_ptr(scales_tiled), input_rows, From d6f3a88ba641f651ed7280b345fd3bfb03bb376d Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Thu, 8 Jan 2026 14:18:16 +0100 Subject: [PATCH 4/9] Revert "refactoring" This reverts commit 53182060ee541f66acdb28f2a22b95006564416e. --- mlx/backend/cuda/quantized/qqmm.cpp | 8 ++++---- mlx/backend/cuda/quantized/qqmm_utils.cu | 11 ++++++----- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/mlx/backend/cuda/quantized/qqmm.cpp b/mlx/backend/cuda/quantized/qqmm.cpp index b57031d028..e7260da8c1 100644 --- a/mlx/backend/cuda/quantized/qqmm.cpp +++ b/mlx/backend/cuda/quantized/qqmm.cpp @@ -46,7 +46,7 @@ inline array ensure_row_contiguous_matrix( return x_copy; } -array pad_and_swizzle_scales( +array pad_and_repack_scales( const array& scale, cu::CommandEncoder& encoder, const Stream& s) { @@ -64,7 +64,7 @@ array pad_and_swizzle_scales( cu::malloc_async(pad_outer * pad_inner, encoder), Shape{pad_outer, pad_inner}, scale.dtype()); - swizzle_scales(scale, scale_tiled, encoder, s); + repack_scales(scale, scale_tiled, encoder, s); encoder.add_temporary(scale_tiled); return scale_tiled; @@ -176,8 +176,8 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { int K = K_packed * (32 / bits_); // Repack scales from linear to tiled layout for tensor cores - array scale_x = pad_and_swizzle_scales(scale_x_pre, encoder, s); - array scale_w = pad_and_swizzle_scales(scale_w_pre, encoder, s); + array scale_x = pad_and_repack_scales(scale_x_pre, encoder, s); + array scale_w = pad_and_repack_scales(scale_w_pre, encoder, s); bool x_transposed = false; bool w_transposed = true; // always transposed diff --git a/mlx/backend/cuda/quantized/qqmm_utils.cu b/mlx/backend/cuda/quantized/qqmm_utils.cu index 17250f5b58..96fa637026 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -64,8 +64,8 @@ inline std::tuple get_swizzle_launch_args( // Block is always (32, 32) = 1024 threads dim3 block(lanes_per_block, warps_per_block, 1); - int shared_mem_bytes = tile_rows * tile_cols * tiles_per_block; - return std::make_tuple(grid, block, shared_mem_bytes); + + return std::make_tuple(grid, block); } namespace cu { @@ -155,8 +155,8 @@ __global__ void swizzle_scales( } } } + // write 4 ints to the shared memory } - // write 4 ints to the shared memory strided_scales_thread[tidx * tile_size / 16 + tidy] = *reinterpret_cast(thread_tile_rows); } @@ -187,13 +187,14 @@ void swizzle_scales( size_t output_rows = scales_tiled.shape(-2); size_t output_cols = scales_tiled.shape(-1); - auto [num_blocks, block_dims, shared_mem_bytes] = + uint32_t smem_bytes = 128 * 4 * 32; + auto [num_blocks, block_dims] = get_swizzle_launch_args(output_rows, output_cols); enc.add_kernel_node( cu::swizzle_scales, num_blocks, block_dims, - shared_mem_bytes, + smem_bytes, gpu_ptr(scales), gpu_ptr(scales_tiled), input_rows, From f30c91f6802b59f4babf1b6e5e17c7ca902cff0d Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Thu, 8 Jan 2026 22:59:01 +0100 Subject: [PATCH 5/9] xor bank conflicts --- mlx/backend/cuda/quantized/qqmm.cpp | 8 +-- mlx/backend/cuda/quantized/qqmm_utils.cu | 91 +++++++++++------------- mlx/backend/cuda/quantized/qqmm_utils.h | 2 +- 3 files changed, 47 insertions(+), 54 deletions(-) diff --git a/mlx/backend/cuda/quantized/qqmm.cpp b/mlx/backend/cuda/quantized/qqmm.cpp index e7260da8c1..b57031d028 100644 --- a/mlx/backend/cuda/quantized/qqmm.cpp +++ b/mlx/backend/cuda/quantized/qqmm.cpp @@ -46,7 +46,7 @@ inline array ensure_row_contiguous_matrix( return x_copy; } -array pad_and_repack_scales( +array pad_and_swizzle_scales( const array& scale, cu::CommandEncoder& encoder, const Stream& s) { @@ -64,7 +64,7 @@ array pad_and_repack_scales( cu::malloc_async(pad_outer * pad_inner, encoder), Shape{pad_outer, pad_inner}, scale.dtype()); - repack_scales(scale, scale_tiled, encoder, s); + swizzle_scales(scale, scale_tiled, encoder, s); encoder.add_temporary(scale_tiled); return scale_tiled; @@ -176,8 +176,8 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { int K = K_packed * (32 / bits_); // Repack scales from linear to tiled layout for tensor cores - array scale_x = pad_and_repack_scales(scale_x_pre, encoder, s); - array scale_w = pad_and_repack_scales(scale_w_pre, encoder, s); + array scale_x = pad_and_swizzle_scales(scale_x_pre, encoder, s); + array scale_w = pad_and_swizzle_scales(scale_w_pre, encoder, s); bool x_transposed = false; bool w_transposed = true; // always transposed diff --git a/mlx/backend/cuda/quantized/qqmm_utils.cu b/mlx/backend/cuda/quantized/qqmm_utils.cu index 96fa637026..fd0ead3e45 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -44,13 +44,13 @@ namespace cg = cooperative_groups; // [380, 381, 382, 383], // [508, 509, 510, 511]]]]], -inline std::tuple get_swizzle_launch_args( +inline std::tuple get_swizzle_launch_args( size_t M_swizzled, size_t K_swizzled, int tile_rows = 128, int tile_cols = 4, int tiles_per_lane = 1) { - constexpr int lanes_per_block = 32; // 32 threads per warp + constexpr int lanes_per_block = 32; const int tiles_per_block = lanes_per_block * tiles_per_lane; const int warps_per_block = tile_rows / 4; // 128 / 4 = 32 @@ -61,11 +61,11 @@ inline std::tuple get_swizzle_launch_args( grid.x = cuda::ceil_div(num_tiles_k, tiles_per_block); grid.y = num_tiles_m; grid.z = 1; - + int shared_mem_bytes = tile_rows * tile_cols * tiles_per_block; // Block is always (32, 32) = 1024 threads dim3 block(lanes_per_block, warps_per_block, 1); - return std::make_tuple(grid, block); + return std::make_tuple(grid, block, shared_mem_bytes); } namespace cu { @@ -77,66 +77,51 @@ __global__ void swizzle_scales( const size_t K, const size_t M_swizzled, const size_t K_swizzled) { - // M_swizzled and K_swizzled are dimensions of scales_tiled array - // (padded to full tiles 128x4 if M or K are not multiples of tile sizes) - - // Tile dimensions for scale factors constexpr int tile_dim_row = 128; constexpr int tile_dim_col = 4; - constexpr int tile_size = tile_dim_row * tile_dim_col; // 512 bytes - constexpr int num_tile_rows_per_thread = 4; // always 4 + constexpr int tile_size = tile_dim_row * tile_dim_col; + constexpr int num_tile_rows_per_thread = 4; constexpr int num_tiles_per_thread = 1; constexpr int lanes_per_block = 32; - constexpr int num_tiles_per_block = lanes_per_block * num_tiles_per_thread; - // Each thread loads 4 rows of 4 bytes x 1 column of scales (16 bytes -- 16 - // scales) thread (0, 0) loads scales at rows 0,32,64,96 of tile 0 thread (1, - // 0) loads rows 0,32,64,96 of of tile 1 therefore a warp loads: consecutive 4 - // bytes x 32 = 128 bytes 4 times with a stride 32 - auto block_size = cg::this_thread_block().dim_threads(); // (32, 32, 1) + constexpr int max_tiles_per_block = lanes_per_block * num_tiles_per_thread; + + constexpr int tile_stride = tile_size / 16; // 32 int4s per tile + + auto block_size = cg::this_thread_block().dim_threads(); auto block_idx = cg::this_thread_block().group_index(); auto idx_in_block = cg::this_thread_block().thread_index(); - auto tidx = idx_in_block.x; // lane within warp (0, 31) - auto tidy = idx_in_block.y; // warp index (0, 31) + auto tidx = idx_in_block.x; + auto tidy = idx_in_block.y; auto linear_tid = tidy * block_size.x + tidx; const int bid_x = block_idx.x; const int bid_y = block_idx.y; - // [M tile * 128 * bytes] + [K tile * 32] const int K_int = K_swizzled / 4; - // incase of overflow cast to size_t const size_t block_offset = static_cast(bid_y) * tile_dim_row * K_int + - static_cast(bid_x) * num_tiles_per_block; + static_cast(bid_x) * max_tiles_per_block; const int* input_block = reinterpret_cast(scales_linear) + block_offset; const size_t output_offset = static_cast(bid_y) * tile_dim_row * K_int + - static_cast(bid_x) * num_tiles_per_block * tile_size / 4; - + static_cast(bid_x) * max_tiles_per_block * tile_size / 4; int* output_block = reinterpret_cast(scales_swizzled) + output_offset; - const int num_tiles_k = K_swizzled / tile_dim_col; const int grid_dim_x = cg::this_grid().dim_blocks().x; const int grid_dim_y = cg::this_grid().dim_blocks().y; - bool pad_rows = (bid_y == grid_dim_y - 1) && - (M < M_swizzled); // if the last and is partial - bool pad_cols = (bid_x == grid_dim_x - 1) && - (K < K_swizzled); // if the last and is partial + bool pad_rows = (bid_y == grid_dim_y - 1) && (M < M_swizzled); + bool pad_cols = (bid_x == grid_dim_x - 1) && (K < K_swizzled); + + int remaining = K_int - bid_x * max_tiles_per_block; + int tiles_in_block = std::min(remaining, max_tiles_per_block); + bool valid_tile = threadIdx.x * num_tiles_per_thread < tiles_in_block; + + __shared__ int4 strided_scales_thread[max_tiles_per_block * tile_stride]; - int num_tiles_per_block_ = num_tiles_per_block; - if (bid_x == grid_dim_x - 1) { - num_tiles_per_block_ = (K_int - 1) % num_tiles_per_block + 1; - } - bool valid_tile = threadIdx.x * num_tiles_per_thread < num_tiles_per_block_; - // Each thread loads 16 scales from 4 rows (stride 32) and packs them into - // int4. The store is strided within a warp (stride 32 int4s), so we first - // write to shared memory, then do a coalesced store from shared to global - extern __shared__ int4 strided_scales_thread[]; - // load int thread_tile_rows[num_tile_rows_per_thread]; if (valid_tile) { #pragma unroll @@ -145,7 +130,6 @@ __global__ void swizzle_scales( (i * block_size.x + tidy) * K_int + tidx * num_tiles_per_thread; thread_tile_rows[i] = __ldg(input_block + thread_offset); if (pad_rows || pad_cols) { - // check bytes we need to pad for (int j = 0; j < num_tile_rows_per_thread * sizeof(int); j++) { const size_t element_idx = (block_offset + thread_offset) * sizeof(int) + j; @@ -155,18 +139,29 @@ __global__ void swizzle_scales( } } } - // write 4 ints to the shared memory } - strided_scales_thread[tidx * tile_size / 16 + tidy] = + // store to shared with XOR swizzle to avoid 32 way bank conflict + int base_idx = tidx * tile_stride + tidy; + int xor_bits = (tidy >> 3) & 0x3; + int swizzled_idx = base_idx ^ xor_bits; + strided_scales_thread[swizzled_idx] = *reinterpret_cast(thread_tile_rows); } - __syncthreads(); - // shared -> global - const int total_int4s = num_tiles_per_block_ * tile_size / 16; + cg::thread_block block = cg::this_thread_block(); + cg::sync(block); + + // load from shared with XOR swizzle to avoid 4 way bank conflict + const int total_int4s = tiles_in_block * tile_stride; #pragma unroll for (int i = linear_tid; i < total_int4s; i += block_size.x * block_size.y) { - reinterpret_cast(output_block)[i] = strided_scales_thread[i]; + int tile_idx = i / tile_stride; + int row_idx = i % tile_stride; + int base_idx = tile_idx * tile_stride + row_idx; + int xor_bits = (row_idx >> 3) & 0x3; + int swizzled_idx = base_idx ^ xor_bits; + reinterpret_cast(output_block)[i] = + strided_scales_thread[swizzled_idx]; } } } // namespace cu @@ -180,21 +175,19 @@ void swizzle_scales( enc.set_output_array(scales_tiled); // Note: scales_tiled is padded to full tiles so if num_rows or num_cols // are not multiples of tile sizes - size_t input_rows = scales.shape(-2); size_t input_cols = scales.shape(-1); size_t output_rows = scales_tiled.shape(-2); size_t output_cols = scales_tiled.shape(-1); - uint32_t smem_bytes = 128 * 4 * 32; - auto [num_blocks, block_dims] = + auto [num_blocks, block_dims, shared_mem_bytes] = get_swizzle_launch_args(output_rows, output_cols); enc.add_kernel_node( cu::swizzle_scales, num_blocks, block_dims, - smem_bytes, + shared_mem_bytes, gpu_ptr(scales), gpu_ptr(scales_tiled), input_rows, diff --git a/mlx/backend/cuda/quantized/qqmm_utils.h b/mlx/backend/cuda/quantized/qqmm_utils.h index 126cc298b2..0a9a78f70c 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.h +++ b/mlx/backend/cuda/quantized/qqmm_utils.h @@ -21,7 +21,7 @@ inline std::pair get_padded_scale_dims(int num_rows, int num_cols) { return {padded_rows, padded_cols}; } -void repack_scales( +void swizzle_scales( const array& scales, array& scales_tiled, cu::CommandEncoder& enc, From 8d0a768ac0b633a73ec198203ae2eaa03af4a647 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Fri, 9 Jan 2026 00:44:02 +0100 Subject: [PATCH 6/9] fix cases when K is not multiple of 4 --- mlx/backend/cuda/quantized/qqmm.cpp | 2 - mlx/backend/cuda/quantized/qqmm_utils.cu | 87 +++++++++++++++++------- 2 files changed, 63 insertions(+), 26 deletions(-) diff --git a/mlx/backend/cuda/quantized/qqmm.cpp b/mlx/backend/cuda/quantized/qqmm.cpp index b57031d028..21e41cf0ed 100644 --- a/mlx/backend/cuda/quantized/qqmm.cpp +++ b/mlx/backend/cuda/quantized/qqmm.cpp @@ -69,9 +69,7 @@ array pad_and_swizzle_scales( encoder.add_temporary(scale_tiled); return scale_tiled; } -} // namespace -namespace { void qqmm_impl( cu::CommandEncoder& encoder, int M, diff --git a/mlx/backend/cuda/quantized/qqmm_utils.cu b/mlx/backend/cuda/quantized/qqmm_utils.cu index fd0ead3e45..4854e18bbf 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -87,6 +87,11 @@ __global__ void swizzle_scales( constexpr int tile_stride = tile_size / 16; // 32 int4s per tile + // Each thread loads 4 rows of 4 bytes x 1 column of scales (16 bytes -- 16 + // scales) thread (0, 0) loads scales at rows 0,32,64,96 of tile 0 thread (1, + // 0) loads rows 0,32,64,96 of of tile 1 therefore a warp loads: consecutive 4 + // bytes x 32 = 128 bytes + // Note: it is better to have K aligned to 4 auto block_size = cg::this_thread_block().dim_threads(); auto block_idx = cg::this_thread_block().group_index(); auto idx_in_block = cg::this_thread_block().thread_index(); @@ -99,11 +104,6 @@ __global__ void swizzle_scales( const int bid_y = block_idx.y; const int K_int = K_swizzled / 4; - const size_t block_offset = - static_cast(bid_y) * tile_dim_row * K_int + - static_cast(bid_x) * max_tiles_per_block; - const int* input_block = - reinterpret_cast(scales_linear) + block_offset; const size_t output_offset = static_cast(bid_y) * tile_dim_row * K_int + @@ -113,34 +113,74 @@ __global__ void swizzle_scales( const int grid_dim_x = cg::this_grid().dim_blocks().x; const int grid_dim_y = cg::this_grid().dim_blocks().y; - bool pad_rows = (bid_y == grid_dim_y - 1) && (M < M_swizzled); - bool pad_cols = (bid_x == grid_dim_x - 1) && (K < K_swizzled); - int remaining = K_int - bid_x * max_tiles_per_block; - int tiles_in_block = std::min(remaining, max_tiles_per_block); - bool valid_tile = threadIdx.x * num_tiles_per_thread < tiles_in_block; + int tiles_in_block = min(remaining, max_tiles_per_block); + bool valid_tile = tidx * num_tiles_per_thread < tiles_in_block; + // Each thread loads 16 scales from 4 rows (stride 32) and packs them into + // int4. The store is strided within a warp (stride 32 int4s), so we first + // write to shared memory, then do a coalesced store from shared to global __shared__ int4 strided_scales_thread[max_tiles_per_block * tile_stride]; - int thread_tile_rows[num_tile_rows_per_thread]; + // Initialize to zero for padding + int thread_tile_rows[num_tile_rows_per_thread] = {0, 0, 0, 0}; + if (valid_tile) { + const size_t col_base = + static_cast(bid_x) * max_tiles_per_block * tile_dim_col + + tidx * tile_dim_col; + + const bool aligned_k = (K % 4 == 0); + + if (aligned_k) { + // fast path: K is aligned, use vectorized loads with stride K/4 + const int K_stride = K / 4; + const size_t block_offset = + static_cast(bid_y) * tile_dim_row * K_stride + + static_cast(bid_x) * max_tiles_per_block; + const int* input_block = + reinterpret_cast(scales_linear) + block_offset; + +// load +#pragma unroll + for (int i = 0; i < num_tile_rows_per_thread; i++) { + const size_t row = + static_cast(bid_y) * tile_dim_row + i * block_size.x + tidy; + const int thread_offset = + (i * block_size.x + tidy) * K_stride + tidx * num_tiles_per_thread; + + if (row < M && col_base + tile_dim_col <= K) { + thread_tile_rows[i] = __ldg(input_block + thread_offset); + } else if (row < M) { +// partial tile at K boundary: load byte-by-byte +#pragma unroll + for (int c = 0; c < tile_dim_col; c++) { + if (col_base + c < K) { + reinterpret_cast(&thread_tile_rows[i])[c] = + scales_linear[row * K + col_base + c]; + } + } + } + } + } else { #pragma unroll - for (int i = 0; i < num_tile_rows_per_thread; i++) { - const int thread_offset = - (i * block_size.x + tidy) * K_int + tidx * num_tiles_per_thread; - thread_tile_rows[i] = __ldg(input_block + thread_offset); - if (pad_rows || pad_cols) { - for (int j = 0; j < num_tile_rows_per_thread * sizeof(int); j++) { - const size_t element_idx = - (block_offset + thread_offset) * sizeof(int) + j; - if (element_idx / K_swizzled >= M || - (element_idx % K_swizzled) >= K) { - reinterpret_cast(&thread_tile_rows[i])[j] = 0; + for (int i = 0; i < num_tile_rows_per_thread; i++) { + const size_t row = + static_cast(bid_y) * tile_dim_row + i * block_size.x + tidy; + if (row < M) { + const size_t row_start = row * K; +#pragma unroll + for (int c = 0; c < tile_dim_col; c++) { + if (col_base + c < K) { + reinterpret_cast(&thread_tile_rows[i])[c] = + scales_linear[row_start + col_base + c]; + } } } } } - // store to shared with XOR swizzle to avoid 32 way bank conflict + + // store to shared with XOR swizzle to avoid bank conflicts int base_idx = tidx * tile_stride + tidy; int xor_bits = (tidy >> 3) & 0x3; int swizzled_idx = base_idx ^ xor_bits; @@ -151,7 +191,6 @@ __global__ void swizzle_scales( cg::thread_block block = cg::this_thread_block(); cg::sync(block); - // load from shared with XOR swizzle to avoid 4 way bank conflict const int total_int4s = tiles_in_block * tile_stride; #pragma unroll for (int i = linear_tid; i < total_int4s; i += block_size.x * block_size.y) { From d99e57835e0329fd815ccd5e396b510cd13880e5 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Fri, 9 Jan 2026 02:54:13 +0100 Subject: [PATCH 7/9] drop repetitive comment --- mlx/backend/cuda/quantized/qqmm_utils.cu | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/mlx/backend/cuda/quantized/qqmm_utils.cu b/mlx/backend/cuda/quantized/qqmm_utils.cu index 4854e18bbf..16e915d1af 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -87,11 +87,11 @@ __global__ void swizzle_scales( constexpr int tile_stride = tile_size / 16; // 32 int4s per tile - // Each thread loads 4 rows of 4 bytes x 1 column of scales (16 bytes -- 16 - // scales) thread (0, 0) loads scales at rows 0,32,64,96 of tile 0 thread (1, - // 0) loads rows 0,32,64,96 of of tile 1 therefore a warp loads: consecutive 4 - // bytes x 32 = 128 bytes - // Note: it is better to have K aligned to 4 + // Each thread loads 16 scales from 4 rows (stride 32) and packs them into + // int4. For example: thread (0, 0) loads scales at rows 0,32,64,96 of tile 0, + // thread (1, 0) loads rows 0,32,64,96 of of tile 1, etc. + // The store is strided within a warp (stride 32 int4s), so we first + // write to shared memory, then do a coalesced store from shared to global auto block_size = cg::this_thread_block().dim_threads(); auto block_idx = cg::this_thread_block().group_index(); auto idx_in_block = cg::this_thread_block().thread_index(); @@ -117,9 +117,6 @@ __global__ void swizzle_scales( int tiles_in_block = min(remaining, max_tiles_per_block); bool valid_tile = tidx * num_tiles_per_thread < tiles_in_block; - // Each thread loads 16 scales from 4 rows (stride 32) and packs them into - // int4. The store is strided within a warp (stride 32 int4s), so we first - // write to shared memory, then do a coalesced store from shared to global __shared__ int4 strided_scales_thread[max_tiles_per_block * tile_stride]; // Initialize to zero for padding @@ -140,7 +137,6 @@ __global__ void swizzle_scales( static_cast(bid_x) * max_tiles_per_block; const int* input_block = reinterpret_cast(scales_linear) + block_offset; - // load #pragma unroll for (int i = 0; i < num_tile_rows_per_thread; i++) { @@ -148,7 +144,6 @@ __global__ void swizzle_scales( static_cast(bid_y) * tile_dim_row + i * block_size.x + tidy; const int thread_offset = (i * block_size.x + tidy) * K_stride + tidx * num_tiles_per_thread; - if (row < M && col_base + tile_dim_col <= K) { thread_tile_rows[i] = __ldg(input_block + thread_offset); } else if (row < M) { @@ -179,7 +174,6 @@ __global__ void swizzle_scales( } } } - // store to shared with XOR swizzle to avoid bank conflicts int base_idx = tidx * tile_stride + tidy; int xor_bits = (tidy >> 3) & 0x3; From 046e8b4dddaa35a5051f2a4740e5b872ad0cffa9 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Sat, 10 Jan 2026 15:14:11 +0100 Subject: [PATCH 8/9] drop shared memory bytes from launch kernel args --- mlx/backend/cuda/quantized/qqmm_utils.cu | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/mlx/backend/cuda/quantized/qqmm_utils.cu b/mlx/backend/cuda/quantized/qqmm_utils.cu index 16e915d1af..428d117a31 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -44,7 +44,7 @@ namespace cg = cooperative_groups; // [380, 381, 382, 383], // [508, 509, 510, 511]]]]], -inline std::tuple get_swizzle_launch_args( +inline std::tuple get_swizzle_launch_args( size_t M_swizzled, size_t K_swizzled, int tile_rows = 128, @@ -61,11 +61,10 @@ inline std::tuple get_swizzle_launch_args( grid.x = cuda::ceil_div(num_tiles_k, tiles_per_block); grid.y = num_tiles_m; grid.z = 1; - int shared_mem_bytes = tile_rows * tile_cols * tiles_per_block; // Block is always (32, 32) = 1024 threads dim3 block(lanes_per_block, warps_per_block, 1); - return std::make_tuple(grid, block, shared_mem_bytes); + return std::make_tuple(grid, block); } namespace cu { @@ -214,13 +213,13 @@ void swizzle_scales( size_t output_rows = scales_tiled.shape(-2); size_t output_cols = scales_tiled.shape(-1); - auto [num_blocks, block_dims, shared_mem_bytes] = + auto [num_blocks, block_dims] = get_swizzle_launch_args(output_rows, output_cols); enc.add_kernel_node( cu::swizzle_scales, num_blocks, block_dims, - shared_mem_bytes, + 0, gpu_ptr(scales), gpu_ptr(scales_tiled), input_rows, From a116413af19576179f803a4c25b30152b7aecf8e Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Sat, 10 Jan 2026 16:06:47 +0100 Subject: [PATCH 9/9] constexpr at the top of the file --- mlx/backend/cuda/quantized/qqmm_utils.cu | 54 +++++++++++------------- 1 file changed, 25 insertions(+), 29 deletions(-) diff --git a/mlx/backend/cuda/quantized/qqmm_utils.cu b/mlx/backend/cuda/quantized/qqmm_utils.cu index 428d117a31..c8764709b9 100644 --- a/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -10,6 +10,11 @@ namespace mlx::core { namespace cg = cooperative_groups; +constexpr int TILE_ROWS = 128; +constexpr int TILE_COLS = 4; +constexpr int TILES_PER_LANE = 1; +constexpr int LANES_PER_BLOCK = 32; + // To pass scales to tensor cores, they need to be repacked into a tiled layout // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout // Tiled layout for scale factors is very well described in CUTLASS @@ -46,23 +51,19 @@ namespace cg = cooperative_groups; inline std::tuple get_swizzle_launch_args( size_t M_swizzled, - size_t K_swizzled, - int tile_rows = 128, - int tile_cols = 4, - int tiles_per_lane = 1) { - constexpr int lanes_per_block = 32; - const int tiles_per_block = lanes_per_block * tiles_per_lane; - const int warps_per_block = tile_rows / 4; // 128 / 4 = 32 + size_t K_swizzled) { + constexpr int tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE; + constexpr int warps_per_block = TILE_ROWS / 4; // 128 / 4 = 32 - const int num_tiles_k = K_swizzled / tile_cols; - const int num_tiles_m = M_swizzled / tile_rows; + const int num_tiles_k = K_swizzled / TILE_COLS; + const int num_tiles_m = M_swizzled / TILE_ROWS; dim3 grid; grid.x = cuda::ceil_div(num_tiles_k, tiles_per_block); grid.y = num_tiles_m; grid.z = 1; // Block is always (32, 32) = 1024 threads - dim3 block(lanes_per_block, warps_per_block, 1); + dim3 block(LANES_PER_BLOCK, warps_per_block, 1); return std::make_tuple(grid, block); } @@ -76,13 +77,9 @@ __global__ void swizzle_scales( const size_t K, const size_t M_swizzled, const size_t K_swizzled) { - constexpr int tile_dim_row = 128; - constexpr int tile_dim_col = 4; - constexpr int tile_size = tile_dim_row * tile_dim_col; + constexpr int tile_size = TILE_ROWS * TILE_COLS; constexpr int num_tile_rows_per_thread = 4; - constexpr int num_tiles_per_thread = 1; - constexpr int lanes_per_block = 32; - constexpr int max_tiles_per_block = lanes_per_block * num_tiles_per_thread; + constexpr int max_tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE; constexpr int tile_stride = tile_size / 16; // 32 int4s per tile @@ -104,8 +101,7 @@ __global__ void swizzle_scales( const int K_int = K_swizzled / 4; - const size_t output_offset = - static_cast(bid_y) * tile_dim_row * K_int + + const size_t output_offset = static_cast(bid_y) * TILE_ROWS * K_int + static_cast(bid_x) * max_tiles_per_block * tile_size / 4; int* output_block = reinterpret_cast(scales_swizzled) + output_offset; @@ -114,17 +110,17 @@ __global__ void swizzle_scales( int remaining = K_int - bid_x * max_tiles_per_block; int tiles_in_block = min(remaining, max_tiles_per_block); - bool valid_tile = tidx * num_tiles_per_thread < tiles_in_block; + bool valid_tile = tidx * TILES_PER_LANE < tiles_in_block; __shared__ int4 strided_scales_thread[max_tiles_per_block * tile_stride]; // Initialize to zero for padding - int thread_tile_rows[num_tile_rows_per_thread] = {0, 0, 0, 0}; + int thread_tile_rows[num_tile_rows_per_thread] = {0}; if (valid_tile) { const size_t col_base = - static_cast(bid_x) * max_tiles_per_block * tile_dim_col + - tidx * tile_dim_col; + static_cast(bid_x) * max_tiles_per_block * TILE_COLS + + tidx * TILE_COLS; const bool aligned_k = (K % 4 == 0); @@ -132,7 +128,7 @@ __global__ void swizzle_scales( // fast path: K is aligned, use vectorized loads with stride K/4 const int K_stride = K / 4; const size_t block_offset = - static_cast(bid_y) * tile_dim_row * K_stride + + static_cast(bid_y) * TILE_ROWS * K_stride + static_cast(bid_x) * max_tiles_per_block; const int* input_block = reinterpret_cast(scales_linear) + block_offset; @@ -140,15 +136,15 @@ __global__ void swizzle_scales( #pragma unroll for (int i = 0; i < num_tile_rows_per_thread; i++) { const size_t row = - static_cast(bid_y) * tile_dim_row + i * block_size.x + tidy; + static_cast(bid_y) * TILE_ROWS + i * block_size.x + tidy; const int thread_offset = - (i * block_size.x + tidy) * K_stride + tidx * num_tiles_per_thread; - if (row < M && col_base + tile_dim_col <= K) { + (i * block_size.x + tidy) * K_stride + tidx * TILES_PER_LANE; + if (row < M && col_base + TILE_COLS <= K) { thread_tile_rows[i] = __ldg(input_block + thread_offset); } else if (row < M) { // partial tile at K boundary: load byte-by-byte #pragma unroll - for (int c = 0; c < tile_dim_col; c++) { + for (int c = 0; c < TILE_COLS; c++) { if (col_base + c < K) { reinterpret_cast(&thread_tile_rows[i])[c] = scales_linear[row * K + col_base + c]; @@ -160,11 +156,11 @@ __global__ void swizzle_scales( #pragma unroll for (int i = 0; i < num_tile_rows_per_thread; i++) { const size_t row = - static_cast(bid_y) * tile_dim_row + i * block_size.x + tidy; + static_cast(bid_y) * TILE_ROWS + i * block_size.x + tidy; if (row < M) { const size_t row_start = row * K; #pragma unroll - for (int c = 0; c < tile_dim_col; c++) { + for (int c = 0; c < TILE_COLS; c++) { if (col_base + c < K) { reinterpret_cast(&thread_tile_rows[i])[c] = scales_linear[row_start + col_base + c];