Skip to content

Metal flash attention requires K.type == V.type — blocks asymmetric-precision KV cache schemes #25900

Description

@rajveer43

Prerequisites

  • I am running the latest code. Mention the version if possible as well.
  • I carefully followed the README.md.
  • I searched using keywords relevant to my issue to make sure that I am creating a new issue that is not already open (or closed).
  • I reviewed the Discussions, and have a new and useful enhancement to share.

Feature Description

Context

I maintain VeloxQuant-MLX, a KV cache compression library for MLX/Apple Silicon. Several of the compression schemes we've implemented and measured (e.g. an adaptation of RaBitQ — 1-bit keys + 4-bit values) rely on asymmetric key/value precision: keys and values are quantized at different bit-widths because their error tolerance differs (keys feed the attention scores, values feed the output — the failure modes aren't symmetric).

I noticed ggml-metal-device.m's support check for GGML_OP_FLASH_ATTN_EXT requires src[1]->type == src[2]->type (K and V must share one GGML type). That means any asymmetric K/V precision scheme currently can't use Metal flash attention at all — it would have to fall back to unfused attention, which is a meaningfully different performance path on Apple Silicon.

I saw this come up implicitly in #20977 (several proposed schemes there also use asymmetric K/V precision) but didn't see it discussed as a standalone constraint, so raising it separately.

Is this a known/deliberate constraint?

Wanted to check before assuming it's an oversight — is K.type == V.type required by something structural in the Metal FA kernel (e.g. a fused load path that assumes one dequant function for both), or is it a simplification that could be relaxed for a subset of type pairs the way CUDA appears to allow?

Offer

Separately from this: given the discussion on #20977 about KLD/PPL evidence being the actual bar for new KV quant work (rather than a new ggml_type), I'd be glad to run the KLD/PPL comparison protocol from CONTRIBUTING.md — small GGUF-converted model, current rotated baseline (post-#21038) as the comparison point — using our RVQ approach, and post the numbers on #20977 rather than opening a separate feature request. Flagging here first since the Metal constraint above would affect whether that's even testable on Apple Silicon specifically.

Environment

  • llama.cpp: master (referencing commit around 571d0d5 / 2026-07-18)
  • Relevant file: ggml/src/ggml-metal/ggml-metal-device.m, GGML_OP_FLASH_ATTN_EXT support check (~line 1228)

Motivation

Possible Implementation

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions