Repository navigation
metal: add gemv_wide for fp16/bf16 matmuls of a few rows - #3888
Merged
Merged
Conversation
fp16/bf16 products of a few rows against a large transposed weight matrix (x @ w.T with M = 2..15) route to GEMM tiles that pad M to 64 rows and leave much of the memory bandwidth unused. gemv_wide streams each block of weight rows once and applies it to up to five input vectors held in registers; chunks beyond the grid loop inside one threadgroup, so vocab-width outputs also stream the matrix once. It covers Matmul, AddMM, and GatherMM, batched or sliced, on the M3 generation and later. Speedups are per-kernel GPU time against the stock routing, with uncached weights. The shapes come from a qwen3.6 deployment (the layers a quantized model keeps in bf16). in_proj_a/b, [M, 2048] x [2048, 32]: M M3 Ultra M4 Pro M5 Max 2 3.0x 2.3x 6.0x 4 2.3x 1.8x 4.7x 8 2.3x 1.7x 4.8x lm_head, [M, 2048] x [2048, 248320]: M M3 Ultra M4 Pro M5 Max 2 1.7x 2.3x 1.4x 4 1.7x 2.2x 1.4x 8 1.8x 2.1x 1.3x
1 task done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
fp16/bf16 products of a few rows against a large transposed weight matrix
(
x @ w.Twith M = 2..15) route to GEMM tiles that pad M to 64 rowsand leave much of the memory bandwidth unused.
gemv_widestreams each block of weight rows once and applies it to upto five input vectors held in registers; chunks beyond the grid loop
inside one threadgroup, so vocab-width outputs also stream the matrix
once. It covers
Matmul,AddMM, andGatherMM, batched or sliced, onthe M3 generation and later.
Speedups are per-kernel GPU time against the stock routing, with
uncached weights. The shapes come from a qwen3.6 deployment (the
layers a quantized model keeps in bf16).
in_proj_a/b,[M, 2048] x [2048, 32]:lm_head,[M, 2048] x [2048, 248320]: