Conversation
| uint32_t blocks_y = (N + rows_per_block - 1) / rows_per_block; | ||
| const uint32_t* mat_ptr = gpu_ptr<uint32_t>(mat); | ||
| const T* vec_ptr = gpu_ptr<T>(vec); | ||
| int n = 1; |
There was a problem hiding this comment.
this block is just a refactoring if I am reading it correctly?
There was a problem hiding this comment.
Correct, no logic change, since my change adds a second path that needs the same computation I refactored to DRY it out.
| bool can_use_qmm_sm80 = supports(supports_qmm_sm80); | ||
| bool can_use_qmm_naive = supports(supports_qmm_naive); | ||
| bool can_use_fp_gather_qmv = | ||
| !global_scale && supports(supports_fp_gather_qmv); |
There was a problem hiding this comment.
I there a reason why fp_gather_qmv does not support global scale? I think it might be useful if weights are in nvfp4
There was a problem hiding this comment.
Yes, agreed. I was trying to avoid a hard dependency between these 2 PRs, so I have that optimization on another branch and was planning to PR it once these are in.
| if (!w.flags().row_contiguous || !scales.flags().row_contiguous) { | ||
| return false; | ||
| } | ||
| // Four packed words per thread beat the qmm kernels from K = 1024 on. |
There was a problem hiding this comment.
I am not sure, I fully understand this line.. Could please provide more details?
There was a problem hiding this comment.
I'll clean up the comment. Basically if the shape fills the warp it wins over the qmm kernels. I did test gemma with a K of 704 and it did see an ~8% prompt speedup in the new path, but it doesn't fill the warp so it felt messy to set the K limit to a partial. K 512 underperforms prompt by ~1%.
| int bits, | ||
| int group_size, | ||
| bool use_mx_scale> | ||
| __global__ void fp_qmv_gather( |
There was a problem hiding this comment.
Maybe I am missing something, why not just fp_gather_qmv?
| !global_scale && supports(supports_fp_gather_qmv); | ||
| bool can_use_qmv = supports(supports_qmv); | ||
|
|
||
| auto call_fp_gather_qmv = [&]() { |
There was a problem hiding this comment.
I dont think that separate lambda is needed here.
I think we can fold it in call_qmv similar to QuantizedMatmul:
bool can_use_fp_gather_qmv = !global_scale && supports(supports_fp_gather_qmv);
bool can_use_qmv = supports(supports_qmv) || can_use_fp_gather_qmv;
auto call_qmv = [&]() {
out.set_data(cu::malloc_async(out.nbytes(), encoder));
if (can_use_fp_gather_qmv) {
fp_gather_qmv(x, w, scales, lhs_indices, rhs_indices, out, bits_, group_size_, encoder);
} else {
gather_qmv(x, w, scales, biases, lhs_indices, rhs_indices, out, bits_, group_size_, mode_, encoder);
}
};
|
Thank you for the pull request! Great improvements! I left some comments and questions. |
nastya236
added
the
await response
d6e2103 to
d56fa41
Compare
2 participants
This speeds up gather_qmm on CUDA for MoE models.
mlx-lm benchmark examples with gemma4 26b and Nemotron 3 Nano 30B (p2048/g128)
Generation speed unchanged.
Note: this carries a few lines from #4507 but is otherwise independent of the global scale work so these can merge in either order.