From 46ee0e9068e7e1cce29050d87079bbb8ba7b4fd7 Mon Sep 17 00:00:00 2001 From: CCYeh Date: Mon, 12 Jan 2026 15:16:05 +0100 Subject: [PATCH] Fix `grid_dim_x` calculations (#2980) --- mlx/backend/cuda/quantized/affine_quantize.cu | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/mlx/backend/cuda/quantized/affine_quantize.cu b/mlx/backend/cuda/quantized/affine_quantize.cu index da4b3413..4892bc7d 100644 --- a/mlx/backend/cuda/quantized/affine_quantize.cu +++ b/mlx/backend/cuda/quantized/affine_quantize.cu @@ -23,8 +23,7 @@ affine_quantize(const T* w, uint8_t* out, T* scales, T* biases, size_t size) { 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; + auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; constexpr float eps = 1e-7; constexpr int simd_size = WARP_SIZE; constexpr float n_bins = (1 << bits) - 1; @@ -141,8 +140,7 @@ __global__ void affine_dequantize( 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; + auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; constexpr int pack_factor = get_pack_factor(); constexpr int bytes_per_pack = get_bytes_per_pack();