L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix for NAX overflow. (#3092)

Browse Source
This commit is contained in:
Awni Hannun
2026-02-02 18:54:01 -08:00
committed by GitHub
parent b56782be52
commit adcbb91a9e
3 changed files with 18 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h
+6 -2
View File
@@ -157,10 +157,14 @@ template <
const short tm = SM * (simd_group_id / WN);
const short tn = SN * (simd_group_id % WN);
const short sgp_sm = align_M ? SM : min(SM, short(params->M - (c_row + tm)));
const int sgp_sm_int =
align_M ? int(SM) : min(int(SM), params->M - (c_row + tm));
const short sgp_sm = short(sgp_sm_int);
const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM);
const short sgp_sn = align_N ? SN : min(SN, short(params->N - (c_col + tn)));
const int sgp_sn_int =
align_N ? int(SN) : min(int(SN), params->N - (c_col + tn));
const short sgp_sn = short(sgp_sn_int);
const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN);
A += transpose_a ? tm : (tm * params->lda);
mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h
+6 -2
View File
@@ -53,10 +53,14 @@ gather_mm_rhs_nax(
const short tm = SM * (simd_group_id / WN);
const short tn = SN * (simd_group_id % WN);
const short sgp_sm = align_M ? SM : min(SM, short(params->M - (c_row + tm)));
const int sgp_sm_int =
align_M ? int(SM) : min(int(SM), params->M - (c_row + tm));
const short sgp_sm = short(sgp_sm_int);
const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM);
const short sgp_sn = align_N ? SN : min(SN, short(params->N - (c_col + tn)));
const int sgp_sn_int =
align_N ? int(SN) : min(int(SN), params->N - (c_col + tn));
const short sgp_sn = short(sgp_sn_int);
const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN);
A += transpose_a ? tm : (tm * params->lda);
mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.h
+6 -2
View File
@@ -86,10 +86,14 @@ template <
const short tm = SM * (simd_group_id / WN);
const short tn = SN * (simd_group_id % WN);
const short sgp_sm = align_M ? SM : min(SM, short(params->M - (c_row + tm)));
const int sgp_sm_int =
align_M ? int(SM) : min(int(SM), params->M - (c_row + tm));
const short sgp_sm = short(sgp_sm_int);
const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM);
const short sgp_sn = align_N ? SN : min(SN, short(params->N - (c_col + tn)));
const int sgp_sn_int =
align_N ? int(SN) : min(int(SN), params->N - (c_col + tn));
const short sgp_sn = short(sgp_sn_int);
const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN);
A += transpose_a ? tm : (tm * params->lda);