FastMQA — a Multi-Query Attention CUDA kernel

A small FlashAttention-style CUDA kernel for Multi-Query Attention (MQA) forward, exposed to PyTorch as a C++ extension, with a correctness suite and an honest benchmark against PyTorch baselines.

Status of this README (rewritten 2026-08-20). Earlier versions of this page claimed "production-ready" implementations, "validated" 98-99.8 % memory reductions with "512x concurrent user scaling", FlashAttention-3 / speculative-decoding / quantization feature lists, and Tesla-T4-verified results. An audit found that the committed "benchmark results" were simulated (a hard-coded 1.8x divisor applied to CPU timings, labelled "FastMQA CUDA (simulated)"), that the test files imported modules that do not exist in the repository, that one script fabricated a development timeline, and that the CUDA kernel itself kept a float scores[2048] array per thread (8 KB of register/local-memory spill), which made it 17-35x slower than even a naive PyTorch implementation (per-configuration ratios in the committed benchmark data). Everything below is limited to what the code in this tree does and what has been measured on hardware; the git history preserves what was removed.

What MQA is (and what is arithmetic vs. what is measured)

Multi-Query Attention (Shazeer 2019, arXiv:1911.02150) gives every one of the H query heads the same single K and V head. The KV cache for a layer therefore shrinks from 2*B*H*S*D to 2*B*S*D elements — H times smaller by construction. For the 32-head configurations below that is a 96.9 % reduction, and it is arithmetic that follows from the definition, not a benchmark result. What actually needs measuring is the speed and accuracy of the attention computation itself; that is what this repository measures.

What this is

• kernels/mqa_kernel.cu — one templated fp32 kernel implementing out = softmax(Q Kᵀ / √D) V with K/V broadcast over heads (non-causal, no mask). FlashAttention-style streaming: K/V pass through shared-memory tiles, each query row keeps a running max / running denominator / unnormalised accumulator (online softmax), and no thread ever materialises a full row of scores. One thread block per 16-32 query rows of one (batch, head); within a block each warp owns 2-4 rows and every shared-memory K/V element loaded is reused across all of a warp's rows. head_dim ∈ {32, 64, 128} (template), any seq_len. Launched on the current PyTorch stream, no device-wide synchronisation, launch-error checked.
• kernels/mqa_extension.cpp — argument validation (TORCH_CHECK) and pybind11 binding: fastmqa_cuda.forward(Q, K, V) with Q [B,H,S,D], K/V [B,1,S,D] → [B,H,S,D], fp32, forward-only.
• fastmqa.py — an MQAttention module (H query heads, one K/V head) that routes its attention core through the kernel when possible and through F.scaled_dot_product_attention with expanded K/V otherwise (masks, gradients, other dtypes, other head dims).
• tests/test_kernel.py — pytest suite: comparison against a float64 reference across tile-tail shapes (S = 1, 33, 127, 500, ...), large-magnitude softmax stability, uniform-attention identity, non-contiguous inputs, error cases, determinism.
• benchmarks/bench_mqa.py — CUDA-event benchmark against PyTorch baselines, including (optionally) the repository's own previous kernel, rebuilt from git history, so the before/after is measured rather than remembered.

Install

Requirements: a CUDA GPU, a CUDA-enabled PyTorch (torch>=2.0), and the matching CUDA toolkit with nvcc on PATH.

git clone https://github.com/JonSnow1807/FastMQA.git
cd FastMQA
pip install --no-build-isolation -e .
python -m pytest tests -q

Usage

import torch, fastmqa_cuda

B, H, S, D = 4, 32, 1024, 64
Q = torch.randn(B, H, S, D, device="cuda")
K = torch.randn(B, 1, S, D, device="cuda")   # single K/V head (MQA)
V = torch.randn(B, 1, S, D, device="cuda")
out = fastmqa_cuda.forward(Q, K, V)          # == softmax(Q K^T / sqrt(D)) V

# or as a layer (falls back to SDPA when the kernel does not apply):
from fastmqa import MQAttention
layer = MQAttention(hidden_dim=2048, num_heads=32).cuda().eval()
y = layer(torch.randn(2, 512, 2048, device="cuda"))

Measured performance

One GPU, one dtype: NVIDIA A100-SXM4-40GB, fp32, PyTorch 2.13.0+cu129 / CUDA 12.9. CUDA events, median per-call ms of 5 reps x 30 iters after 10 warmup calls, run from a clean checkout of the benchmarked commit; full data, provenance and the exact command in benchmarks/results/. "manual fp32" materialises the [B,H,S,S] score matrix with expanded K/V — what a straightforward PyTorch implementation (and this repo's pre-rewrite Python code) does. "old kernel" is this repository's previous CUDA kernel (commit 0f40e01), rebuilt from git history and timed on the same GPU. It has no S = 4096 entry: its fixed scores[2048] array is written unguarded up to index seq_len−1, so beyond 2048 the launch is undefined behaviour and the benchmark refuses to run it rather than time corrupted memory.

B x H x S x D	this kernel	manual fp32	SDPA fp32	SDPA fp16*	old kernel
4 x 32 x 512 x 64	1.15 ms	1.21 ms	0.43 ms	0.06 ms	23.1 ms
4 x 32 x 1024 x 64	4.04 ms	4.60 ms	1.60 ms	0.22 ms	82.3 ms
4 x 32 x 2048 x 64	16.0 ms	17.7 ms	6.32 ms	0.82 ms	306.0 ms
4 x 32 x 1024 x 128	10.3 ms	5.60 ms	2.55 ms	0.36 ms	177.5 ms
1 x 32 x 2048 x 128	10.3 ms	5.40 ms	2.64 ms	0.37 ms	188.2 ms
2 x 32 x 4096 x 128	81.5 ms	42.4 ms	20.0 ms	2.80 ms	UB, not run

* fp16 runs FlashAttention on tensor cores — a different precision, listed for context only.

Reading this honestly:

• vs. the kernel it replaces: 17-20x faster on every shape (and it no longer has a hard-coded maximum sequence length). This is the before/after of removing the per-thread float scores[2048] spill in favour of shared-memory tiling + online softmax.
• vs. a naive PyTorch implementation: faster at every head_dim 64 config (1.05-1.14x — the materialised S x S score matrix hurts the baseline as sequences grow), slower at head_dim 128 (~0.5x: the larger shared-memory tiles halve this kernel's rows-per-warp reuse and its occupancy).
• vs. F.scaled_dot_product_attention in fp32: 0.25-0.40x. PyTorch's memory-efficient backend is a CUTLASS kernel whose measured throughput (20.1-27.5 TFLOPS across these configs) exceeds the A100's fp32 CUDA-core roofline (~19.5 TFLOPS), i.e. it uses tensor cores; this kernel uses plain fp32 FMAs and reaches 38-44 % of that roofline at head_dim 64 (S >= 1024) and ~34 % at head_dim 128. Closing the gap would mean tensor-core MMA tiles, not more tuning of this design.
• The old README's "1.8x speedup" figure came from a benchmark that multiplied CPU timings by a constant; no measured configuration supports it and it does not reappear here.

Correctness

Max |error| against a float64 reference is ~8e-7 across the benchmark shapes — the same order as PyTorch's own fp32 SDPA on the identical problem (the test suite asserts the kernel stays within a small multiple of SDPA's error). The online softmax is exact up to fp32 rounding, not an approximation.

Limitations

• Forward only; fp32 only; non-causal; no attention mask (masked calls in MQAttention use the SDPA fallback).
• head_dim ∈ {32, 64, 128}.
• Measured on one GPU (A100-SXM4-40GB); other architectures compile (TORCH_CUDA_ARCH_LIST) but are unmeasured.
• No KV-cache decode path: the extension requires Q and K/V to share one seq_len (square, prefill-shaped attention), so a single-token query against a long cache is rejected by the shape checks — it does not run at all. A decode-optimised kernel would need a different interface and launch geometry.

Repository layout

kernels/            mqa_kernel.cu, mqa_kernel.h, mqa_extension.cpp
fastmqa.py          MQAttention module + kernel/SDPA routing
tests/              pytest suite (needs a GPU + the built extension)
benchmarks/         bench_mqa.py + committed results (JSON/Markdown)

License

MIT — see LICENSE.