Skip to content

[Metal] Add fast route for dequantizing fp4 to half in fp_qmm_t - #4709

Merged
zcbenz merged 3 commits into
ml-explore:mainfrom
dhiltgen:qmm-nvfp4-half-tile
Oct 21, 2026
Merged

zcbenz merged 3 commits into
ml-explore:mainfrom
dhiltgen:qmm-nvfp4-half-tile

Conversation

@dhiltgen

Copy link
Copy Markdown
Contributor

This change speeds up prompt processing for fp4 models on older generations. Instead of dequantizing to f32, applying group scale, then converting to bf16, decode to f16 tiles and scale with one op, and leverage bigger tiles for long prompts. For nvfp4 there is no precision loss. Same for mxfp4 fp16 based models. For mxfp4 bf16 based models to avoid precision loss, those are left on the prior code path.

mlx_lm.benchmark, 2048-token prompt, mlx-community Qwen3.8-27B-nvfp4 prompt tps
(generation tps unchanged).

GPU main prompt tps this PR prompt tps
M1 Pro 54.4 76.5
M2 Pro 69.7 96.8
M2 Ultra 275.0 378.1
M3 Ultra 451.8 492.0
M4 Pro 132.1 144.4
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: co-developed with a coding agent

@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Oct 19, 2026
int BK = 32,
int BN = 32>
int BN = 32,
bool half_tile = false>

@zcbenz zcbenz Oct 19, 2026 •

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than having the caller pass the half_tile parameter, can we just deduce inside fp_qmm_t?

constexpr half_tile =
    (metal::is_same_v<T, float16> && bits == ... && group_size == ...) ||
    (metal::is_same_v<T, bfloat16> && bits == ... && group_size == ...);

@zcbenz zcbenz removed the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Oct 19, 2026
@dhiltgen
dhiltgen force-pushed the qmm-nvfp4-half-tile branch from 12f0150 to e075da1 Compare October 19, 2026 15:56
dhiltgen added a commit to dhiltgen/ollama that referenced this pull request Oct 20, 2026
Half-tile fp4 quantized matmul, which speeds up 4-bit prefill. Applied
to the pinned MLX until the pin includes it; drop this commit then.
dhiltgen added a commit to dhiltgen/ollama that referenced this pull request Oct 20, 2026
Quantized matmul tiled for 6-16 rows, the shape of speculative verify,
rebased onto ml-explore/mlx#4709. Applied to the pinned MLX until the
pin includes it; drop this commit then.
@zcbenz
zcbenz force-pushed the qmm-nvfp4-half-tile branch from e075da1 to ff91940 Compare October 21, 2026 05:05
@zcbenz zcbenz changed the title [Metal] Use half tile for fp4 qmm [Metal] Add fast route for dequantizing fp4 to half in fp_qmm_t Oct 21, 2026
@zcbenz
zcbenz merged commit 52a6c23 into ml-explore:main Oct 21, 2026
30 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants