Back Home

GitHub Repo

Transformers Community Proposes Native GQA Optimization for CPUs, with Local Tests Showing About 20% Faster Decoding

A Transformers technical discussion on October 12–13 proposed using native GQA for CPU inference with attention masks, avoiding K/V tensor copies. The authors measured 1.21× faster decoding, while a second machine showed only 1.06×; the proposal has not been merged.

Andrew Wippler from Lancaster, USA · CC BY 2.0 · Image source
zh-Hant

On October 12, the Transformers community proposed optimizing the CPU attention path, and a project collaborator responded the next day. The proposal concerns the use_gqa_in_sdpa check: the authors say native grouped-query attention (GQA) is currently enabled for CPU calls only when there is no attention mask. Batches with padding and many generation requests still go through repeat_kv, which copies K/V to match the full number of query heads. Technical proposal

The technical benefit is less data movement before attention computation. GQA lets multiple query heads share fewer K/V heads; if the framework expands the tensors first, that shared structure still incurs extra copying. Hugging Face documentation says sdpa uses PyTorch’s attention implementation and that engineers can select a backend with attn_implementation. Choosing SDPA alone therefore does not establish whether those internal copies are avoided. Attention backend documentation

The authors tested a Llama configuration matching SmolLM2-135M’s shape: 9 query heads, 3 K/V heads, bf16, batch size 2, a 2,048-token left-padded prefix, and 4 CPU threads. On PyTorch 2.14.1, median decoding-step time fell from 71.5 to 59.3 milliseconds, or about 1.21× faster; a second machine showed a 1.06× speedup. The authors reported identical logits and greedy-generation tokens, and said the relevant tests passed. These are still proposer-run, local validations. Test setup and results

The proposal says CPU kernels in PyTorch 2.9 and later can accept native GQA with a mask, and limits the change to cases where K/V do not require gradients, execution is not compiled or traced, and the K/V head dimensions match. A collaborator has asked why compiled and traced execution is excluded and questioned whether the accumulating conditions will be maintainable. The issue remains open, so this should not be treated as an improvement in a released version. Engineers should track compatibility testing and merge status, then measure performance on their own CPUs and with their own sequence lengths and batch configurations. The reported decoding-step results also cannot be directly extrapolated to end-to-end generation latency. Collaborator response

Sources

  1. CPU SDPA could use native GQA with a mask on torch >= 2.9
  2. Transformers Attention backends