Skip to content

fix(metal): fix/silence all kernel build warnings #29

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
asher merged 1 commit into from
Aug 1, 2026
Merged

fix(metal): fix/silence all kernel build warnings #29

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 67 additions & 66 deletions metal/mlx/backend/metal/kernels/kq_quantized.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1667,12 +1667,12 @@ template <typename T, int group_size, int bits, bool batched>
device T* y,
const constant int& in_vec_size,
const constant int& out_vec_size,
const constant int& x_batch_ndims,
const constant int* x_shape,
const constant int64_t* x_strides,
const constant int& w_batch_ndims,
const constant int* w_shape,
const constant int64_t* w_strides,
const constant int& /* x_batch_ndims */,
const constant int* /* x_shape */,
const constant int64_t* /* x_strides */,
const constant int& /* w_batch_ndims */,
const constant int* /* w_shape */,
const constant int64_t* /* w_strides */,
const constant int64_t* /* s_strides */,
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
Expand Down Expand Up @@ -1736,12 +1736,12 @@ template <typename T, int group_size, int bits, bool batched>
const device T* bias,
const constant int& in_vec_size,
const constant int& out_vec_size,
const constant int& x_batch_ndims,
const constant int* x_shape,
const constant int64_t* x_strides,
const constant int& w_batch_ndims,
const constant int* w_shape,
const constant int64_t* w_strides,
const constant int& /* x_batch_ndims */,
const constant int* /* x_shape */,
const constant int64_t* /* x_strides */,
const constant int& /* w_batch_ndims */,
const constant int* /* w_shape */,
const constant int64_t* /* w_strides */,
const constant int64_t* /* s_strides */,
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
Expand Down Expand Up @@ -1832,12 +1832,12 @@ template <typename T, int group_size, int bits, bool batched>
device T* y,
const constant int& in_vec_size,
const constant int& out_vec_size,
const constant int& x_batch_ndims,
const constant int* x_shape,
const constant int64_t* x_strides,
const constant int& w_batch_ndims,
const constant int* w_shape,
const constant int64_t* w_strides,
const constant int& /* x_batch_ndims */,
const constant int* /* x_shape */,
const constant int64_t* /* x_strides */,
const constant int& /* w_batch_ndims */,
const constant int* /* w_shape */,
const constant int64_t* /* w_strides */,
const constant int64_t* /* s_strides */,
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
Expand Down Expand Up @@ -1895,12 +1895,12 @@ template <typename T, int group_size, int bits, bool batched>
const device T* bias,
const constant int& in_vec_size,
const constant int& out_vec_size,
const constant int& x_batch_ndims,
const constant int* x_shape,
const constant int64_t* x_strides,
const constant int& w_batch_ndims,
const constant int* w_shape,
const constant int64_t* w_strides,
const constant int& /* x_batch_ndims */,
const constant int* /* x_shape */,
const constant int64_t* /* x_strides */,
const constant int& /* w_batch_ndims */,
const constant int* /* w_shape */,
const constant int64_t* /* w_strides */,
const constant int64_t* /* s_strides */,
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
Expand Down Expand Up @@ -2606,12 +2606,12 @@ template <typename T, int group_size, int bits, bool batched>
device T* y,
const constant int& in_vec_size,
const constant int& out_vec_size,
const constant int& x_batch_ndims,
const constant int* x_shape,
const constant int64_t* x_strides,
const constant int& w_batch_ndims,
const constant int* w_shape,
const constant int64_t* w_strides,
const constant int& /* x_batch_ndims */,
const constant int* /* x_shape */,
const constant int64_t* /* x_strides */,
const constant int& /* w_batch_ndims */,
const constant int* /* w_shape */,
const constant int64_t* /* w_strides */,
const constant int64_t* /* s_strides */,
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
Expand Down Expand Up @@ -2682,12 +2682,12 @@ template <typename T, int group_size, int bits, bool batched>
device T* y,
const constant int& in_vec_size,
const constant int& out_vec_size,
const constant int& x_batch_ndims,
const constant int* x_shape,
const constant int64_t* x_strides,
const constant int& w_batch_ndims,
const constant int* w_shape,
const constant int64_t* w_strides,
const constant int& /* x_batch_ndims */,
const constant int* /* x_shape */,
const constant int64_t* /* x_strides */,
const constant int& /* w_batch_ndims */,
const constant int* /* w_shape */,
const constant int64_t* /* w_strides */,
const constant int64_t* /* s_strides */,
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
Expand Down Expand Up @@ -3220,36 +3220,37 @@ METAL_FUNC void kq_gather_qmm_rhs_impl(
// rows per expert fragment each row tile into per-expert segments that each
// pay a full-tile mma K-loop, so utilization is ~1/segments-per-tile. The
// dispatch picks the largest BM not much above the batch's rows-per-expert.
#define KQ_DEFINE_GATHER_QMM_RHS(CODEC, LOADER) \
template <typename T, int group_size, int bits, bool aligned_N, int BM> \
[[kernel]] void kq_##CODEC##_gather_qmm_rhs( \
const device T* x [[buffer(0)]], \
const device uint8_t* w [[buffer(1)]], \
const device uint8_t* scales [[buffer(2)]], \
const device uint32_t* indices [[buffer(3)]], \
device T* y [[buffer(4)]], \
const constant int& M [[buffer(5)]], \
const constant int& N [[buffer(6)]], \
const constant int& K [[buffer(7)]], \
uint3 tid [[threadgroup_position_in_grid]], \
uint simd_gid [[simdgroup_index_in_threadgroup]], \
uint simd_lid [[thread_index_in_simdgroup]]) { \
constexpr int BK = 32, BN = 64; \
constexpr int BK_padded = (BK + 16 / sizeof(T)); \
using LoaderW = LOADER< \
T, \
BN, \
BK, \
BK_padded, \
/*reduction_dim=*/1, \
/*tgp_size=*/2 * 2 * SIMD_SIZE>; \
static_assert( \
group_size == LoaderW::weights_per_block, \
#CODEC " gather_qmm_rhs requires group_size == block size"); \
threadgroup T Xs[BM * BK_padded]; \
threadgroup T Ws[BN * BK_padded]; \
kq_gather_qmm_rhs_impl<T, LoaderW, aligned_N, BM, BK, BN>( \
x, w, indices, y, Xs, Ws, M, N, K, tid, simd_gid, simd_lid); \
#define KQ_DEFINE_GATHER_QMM_RHS(CODEC, LOADER) \
template <typename T, int group_size, int bits, bool aligned_N, int BM> \
[[kernel]] void kq_##CODEC##_gather_qmm_rhs( \
const device T* x [[buffer(0)]], \
const device uint8_t* w [[buffer(1)]], \
const device uint8_t* scales [[buffer(2)]], \
const device uint32_t* indices [[buffer(3)]], \
device T* y [[buffer(4)]], \
const constant int& M [[buffer(5)]], \
const constant int& N [[buffer(6)]], \
const constant int& K [[buffer(7)]], \
uint3 tid [[threadgroup_position_in_grid]], \
uint simd_gid [[simdgroup_index_in_threadgroup]], \
uint simd_lid [[thread_index_in_simdgroup]]) { \
(void)scales; /* wire-format blocks are self-scaled; no separate scales */ \
constexpr int BK = 32, BN = 64; \
constexpr int BK_padded = (BK + 16 / sizeof(T)); \
using LoaderW = LOADER< \
T, \
BN, \
BK, \
BK_padded, \
/*reduction_dim=*/1, \
/*tgp_size=*/2 * 2 * SIMD_SIZE>; \
static_assert( \
group_size == LoaderW::weights_per_block, \
#CODEC " gather_qmm_rhs requires group_size == block size"); \
threadgroup T Xs[BM * BK_padded]; \
threadgroup T Ws[BN * BK_padded]; \
kq_gather_qmm_rhs_impl<T, LoaderW, aligned_N, BM, BK, BN>( \
x, w, indices, y, Xs, Ws, M, N, K, tid, simd_gid, simd_lid); \
}

KQ_DEFINE_GATHER_QMM_RHS(q8_0, KqQ8_0BlockLoader)
Expand Down
48 changes: 24 additions & 24 deletions metal/mlx/backend/metal/kernels/kq_quantized_fp.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -459,12 +459,12 @@ template <typename T, int group_size, int bits, bool batched>
device T* y,
const constant int& in_vec_size,
const constant int& out_vec_size,
const constant int& x_batch_ndims,
const constant int* x_shape,
const constant int64_t* x_strides,
const constant int& w_batch_ndims,
const constant int* w_shape,
const constant int64_t* w_strides,
const constant int& /* x_batch_ndims */,
const constant int* /* x_shape */,
const constant int64_t* /* x_strides */,
const constant int& /* w_batch_ndims */,
const constant int* /* w_shape */,
const constant int64_t* /* w_strides */,
const constant int64_t* /* s_strides */,
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
Expand Down Expand Up @@ -519,12 +519,12 @@ template <typename T, int group_size, int bits, bool batched>
device T* y,
const constant int& in_vec_size,
const constant int& out_vec_size,
const constant int& x_batch_ndims,
const constant int* x_shape,
const constant int64_t* x_strides,
const constant int& w_batch_ndims,
const constant int* w_shape,
const constant int64_t* w_strides,
const constant int& /* x_batch_ndims */,
const constant int* /* x_shape */,
const constant int64_t* /* x_strides */,
const constant int& /* w_batch_ndims */,
const constant int* /* w_shape */,
const constant int64_t* /* w_strides */,
const constant int64_t* /* s_strides */,
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
Expand Down Expand Up @@ -963,12 +963,12 @@ template <typename T, int group_size, int bits, bool batched>
device T* y,
const constant int& in_vec_size,
const constant int& out_vec_size,
const constant int& x_batch_ndims,
const constant int* x_shape,
const constant int64_t* x_strides,
const constant int& w_batch_ndims,
const constant int* w_shape,
const constant int64_t* w_strides,
const constant int& /* x_batch_ndims */,
const constant int* /* x_shape */,
const constant int64_t* /* x_strides */,
const constant int& /* w_batch_ndims */,
const constant int* /* w_shape */,
const constant int64_t* /* w_strides */,
const constant int64_t* /* s_strides */,
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
Expand Down Expand Up @@ -1023,12 +1023,12 @@ template <typename T, int group_size, int bits, bool batched>
device T* y,
const constant int& in_vec_size,
const constant int& out_vec_size,
const constant int& x_batch_ndims,
const constant int* x_shape,
const constant int64_t* x_strides,
const constant int& w_batch_ndims,
const constant int* w_shape,
const constant int64_t* w_strides,
const constant int& /* x_batch_ndims */,
const constant int* /* x_shape */,
const constant int64_t* /* x_strides */,
const constant int& /* w_batch_ndims */,
const constant int* /* w_shape */,
const constant int64_t* /* w_strides */,
const constant int64_t* /* s_strides */,
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
Expand Down
Loading