ROCm / ROCm/FastFlowLM

KV Cache Quantization: Q2/Q3/Q4 MHA kernels exist in libmha.so, how to enable?

Open
#449 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
C++
Stars
1.9k
Forks
152
Avg merge
4h 14m
Merged PRs (30d)
11

Description

Summary

While exploring the FastFlowLM codebase to evaluate contributing a TurboQuant (KV cache compression) implementation, I discovered that libmha.so already contains NPU attention kernels for quantized KV cache:

MHA::Impl::_gen_mha_seq_d128_q2(...)
MHA::Impl::_gen_mha_seq_d128_q3(...)
MHA::Impl::_gen_mha_seq_d128_q4(...)
MHA::Impl::_gen_mha_seq_d256_q2(...)
MHA::Impl::_gen_mha_seq_d256_q4(...)

These suggest the NPU can natively consume 2/3/4-bit quantized KV cache during attention computation. However, there doesn't appear to be any user-facing flag or configuration to enable this — inference
currently uses full BF16 KV cache.

Questions

  1. Are these Q2/Q3/Q4 MHA kernels functional today? Is there an internal flag or config parameter to enable quantized KV cache?

  2. What quantization format do these kernels expect? Specifically:

    • Simple uniform/symmetric quantization (min/max scaling)?
    • Block-wise quantization with per-block scale factors?
    • Something else (e.g. Lloyd-Max optimal quantization as in TurboQuant)?
  3. Is there a plan to expose KV cache quantization as a user option? (e.g. --kv-quant q3)

Context: TurboQuant contribution

I'm interested in contributing a TurboQuant implementation for KV cache compression in FastFlowLM. TurboQuant (Google Research, ICLR 2026) uses random rotation
(Walsh-Hadamard transform) + Lloyd-Max optimal scalar quantization to compress KV cache to 3 bits with near-zero quality loss (~1% perplexity increase, 0.995 cosine similarity).

Benefits for NPU inference:

  • ~5x KV cache memory reduction (BF16 → TQ3)
  • Longer context windows on memory-constrained devices
  • Potentially faster attention due to reduced memory bandwidth
  • Training-free, no model fine-tuning required

The encode/decode path (WHT rotation + scalar quantization) is lightweight enough to run on CPU with AVX2, and the Q3 MHA kernel could consume the result directly on the NPU.

If the Q3 kernel format is compatible (or can be adapted), I'd like to contribute:

  • turboquant.hpp — CPU-side encode/decode (WHT + Lloyd-Max quantization)
  • Integration into fill_kv_cache() to compress K/V before storage
  • A --kv-quant CLI flag to enable it

Environment

  • Ryzen AI 9 HX 370 (XDNA2, 8 columns)
  • Fedora 43, kernel 6.19.9
  • FastFlowLM v0.9.37 (built from source)
  • XRT/xdna-driver 2.23.0

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by tracing fill_kv_cache() and the listed MHA Q2/Q3/Q4 symbols in libmha.so to determine whether an existing enablement path or expected quantization format is exposed. Done means establishing kernel support and format requirements, then defining or implementing a --kv-quant integration with the proposed turboquant.hpp path.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
ai-infra-agents, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.