Skip to content

Reduce default-flash RAM by reusing BF16 projection scratch - #1030

Open
Mikyx-1 wants to merge 1 commit into
google:devfrom
Mikyx-1:fix/flash-kv-projection-scratch
Open

Reduce default-flash RAM by reusing BF16 projection scratch#1030
Mikyx-1 wants to merge 1 commit into
google:devfrom
Mikyx-1:fix/flash-kv-projection-scratch

Conversation

@Mikyx-1

@Mikyx-1 Mikyx-1 commented Sep 12, 2026

Copy link
Copy Markdown

Default flash attention retains a full sequence-major projection history
after transposing it into the K/V buffers consumed by attention.
This change reuses BF16 projection scratch and reduces peak process RAM
while preserving the projection rounding and attention computation.

Fixes #1029

Implementation

  • Reuse one layer of batch-sized BF16 projection scratch per query.
  • Size it for the widest layer; grow it for larger batches or layer widths.
  • Convert older-constructor caches when GemmaAttention first uses them.
  • Preserve sequence capacity and the transposed K/V history.
  • Copy the actual K/V buffers when cloning the new scratch layout.
  • Skip unused compact-buffer allocations for the default-flash runtime path.
  • Preserve normalization, RoPE positions, and both BF16 rounding boundaries.

Measured RAM effect

Peak RSS includes weights, KV caches, and inference working buffers.
These percentages describe the whole process, not model-weight compression.

Gemma 3 model Prompt tokens Before After Reduction
270M 32,736 1,778.59 MiB 1,205.35 MiB 32.2%
1B 16,352 2,369.89 MiB 1,957.65 MiB 17.4%
4B 4,064 5,754.25 MiB 5,216.35 MiB 9.3%

For 270M, the projection allocation falls from 578.00 to 4.25 MiB.
Combined projection plus K/V allocations fall from 1,154.00 to 580.25 MiB.
Dormant virtual allocations are not counted as physical RAM savings.

Measured speed effect

Model Prefill seconds, before → after Decode tokens/s, before → after
270M 39.4237 → 39.3970 37.4104 → 37.8506
1B 68.7325 → 68.5894 12.1040 → 11.9374
4B 74.3011 → 74.3155 2.9791 → 3.0151

Observed decode changes: +1.2% for 270M, -1.4% for 1B, +1.2% for 4B.
Prefill is essentially unchanged. The small measured speedups are tentative:
16-token decode intervals and few repeats do not establish statistical significance.

Measurement setup

  • Baseline ffc1abc; measured fix caa2e13; Linux / Intel i5-12400F.
  • Release AVX2/Haswell, six pinned threads, no oneDNN or spinning.
  • 270M: capacity 32,768, batch 4,096, three stock-autotuning repetitions.
  • 1B: capacity 16,384, batch 4,096; 4B: capacity 4,096, batch 1,024.
  • 1B/4B: two repeats with fixed matmul candidates and forced continuation.
  • 4B uses the existing padded checkpoint with --map 1 in both versions.
  • Timings are medians; all runs are serial with a RAM/swap watchdog.
  • Model snapshots show no swap; system swap grew 162.25 MiB in the 4B baseline.

Validation and scope

  • Five focused cache/attention tests passed on AVX2.
  • Coverage includes two queries, both constructor paths, heterogeneous widths,
    scratch growth, retained sequence capacity, cache copying, and T5 storage.
  • Fixed-matmul verification: all 32 full-logit hashes match for each model.
  • Each version also matches itself across two repetitions.
  • That covers 8,388,608 logit values per model under identical continuation.
  • T5, DeepSeek, and alternative tiled-backend inference storage are unchanged.
  • No quantization, local-window rings, context truncation, or weight changes.
  • The existing tiled test could not compile due to missing mock/span symbols.
  • Only three implementation files and two test files are included.
  • Benchmark artifacts, diagnostic harnesses, and reports are not committed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant