google / google/gemma.cpp

Default flash attention retains redundant BF16 KV projection history

オープン
#1,029 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る
主要言語
C++
スター
7k
フォーク
660
平均マージ
20時間 43分
マージ済み PR(30日)
33

説明

## Summary

Default flash attention retains two BF16 representations of KV history.
The sequence-major projection buffer remains allocated and populated after
its keys and values have been transposed into the buffers attention reads.
This increases resident memory as context length grows.

## Affected path

- Backend: default `--attention_impl flash`.
- Confirmed with Gemma 3 270M, 1B, and 4B text inference.
- Measured baseline: `ffc1abc05abdf11d875d25d36c8859553ddf2641`.
- The retained-history path is also present on current `dev` (`b68def1`).
- T5 and DeepSeek use their legacy buffers differently.

## What happens

1. `ComputeQKV` writes BF16 projections into sequence-major `kv_cache`.
2. It applies normalization and positional encoding with BF16 rounding.
3. It transposes the results into `k_cache` and `v_cache`.
4. `FlashAttention` consumes those transposed buffers.
5. The original projection buffer still retains every layer and position.

References:
- [Projection and transpose path](https://github.com/google/gemma.cpp/blob/ffc1abc05abdf11d875d25d36c8859553ddf2641/gemma/attention.cc#L208-L322).
- [Full-sequence cache allocations](https://github.com/google/gemma.cpp/blob/ffc1abc05abdf11d875d25d36c8859553ddf2641/gemma/kv_cache.cc#L275-L297).

## Reproduction and observed cost

Run default-flash Gemma 3 270M with a 32,736-token text prompt,
sequence capacity 32,768, prefill batch 4,096, and 16 decode tokens.
Inspect the allocated cache extents and peak process RSS during inference.

The projection buffer occupies 578 MiB including row padding.
The transposed K/V buffers occupy another 576 MiB.
Measured peak process RSS is 1,778.59 MiB on this workload.
This is active retained data, not merely unused virtual address space.

Environment: Linux, Intel i5-12400F, six pinned threads,
Release AVX2/Haswell build, no oneDNN, and approximately 15.5 GiB RAM.

## Expected behavior

Projection intermediates should not retain a second complete KV history
after attention's persistent representation has been produced.
Required context, existing BF16 rounding, and model behavior must be preserved.

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

まず gemma/attention.cc の208-322行目と gemma/kv_cache.cc の275-297行目を読み、projection、transpose、cache allocation を追跡してください。指定された prompt と sequence settings で default-flash Gemma 3 workload を再現し、その後 cache extents と peak RSS を調べてください。context length、BF16 rounding、model behavior を変更せずに冗長な KV history が保持されなくなれば完了です。

索引モデルが issue の本文から書いたものです。

評価

技術スタック
cpp
領域
machine-learning, performance
issue の種類
バグ
難易度
4/5
見積もり時間
3〜5日
活発さ
活発
明瞭さ
明確に書かれている
初心者へのやさしさ
68/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。