Reduce local KV memory in BF16 Flash attention

[Mikyx-1]

Closes #1016. Supersedes #1017, which was based on main.

The default Flash path on dev reserves full-context BF16 KV/K/V storage for local layers and also allocates compact tiled buffers it does not use. This change gives each owning Gemma layer separate BF16 buffers: local rings retain the attention window plus batch and SIMD padding, while global layers retain the configured context. The default Flash path no longer allocates unused tiled buffers.

Shared layers alias their source buffers. Ring growth preserves live flat and transposed history, snapshots remain independent, and cache recycling clears the owned buffers. API sessions and conversation bindings use the runtime-aware constructor. The legacy constructor and tiled, DeepSeek, and encoder/decoder layouts remain supported.

Validation

Against upstream/dev at 1658f88, all 13 A/B pairs on Gemma 3 270M and 1B produced identical decoded bytes and token IDs: 1,664 generated tokens per build. Runs used a shared baseline-calibrated matrix plan, six pinned workers on an i5-12400F, context 8,192, prefill batch 256, and 128 greedy output tokens.

Medians from three alternating pairs with a 3,367-token prompt:

Model	Cache capacity, MiB	Peak RSS, MiB	Prefill, tokens/s	Decode, tokens/s
270M	503.0 → 73.7	638.0 → 559.2	1512.0 → 1499.5	46.89 → 46.57
1B	705.5 → 101.4	1540.1 → 1430.7	259.4 → 259.6	14.14 → 14.10

The measured benefit is lower memory use; these results do not demonstrate a speedup. A slower single short-prompt 270M sample did not reproduce in three confirmation pairs (median decode 49.29 → 49.21 tokens/s).

• Seven regression tests passed in Release and under AddressSanitizer with leak detection. Coverage includes capacities, context caps, shared aliases, wraparound, growth, transposed history, snapshots, clearing, and 18 exact full-cache/ring kernel comparisons on AVX2 and EMU128.
• ASan instrumented the changed cache/attention units and tests; supporting libraries were reused without instrumentation.
• CLI, API server, single and batch benchmarks built; conversation binding compile check passed. Bazel dependencies were updated, but Bazel was not run.
• End-to-end model coverage is 270M and 1B; Gemma 4 shared allocation is tested by configuration. Vision and continuous multi-query inference were not benchmarked.

The commit contains production changes, regression tests, and their build wiring. Benchmark harnesses, reports, and generated artifacts are excluded.

[jan-wassenberg]

Thank you! It seems like this would make a big difference for longer sequences.

But unfortunately, as-is, this is adding too much tech debt :/ We'd have a third layout: kv/k/v, compact, per-layer - semi-invisible, based on whether we use the 3-arg ctor or not. And the abstraction leaks, user code has to remember to check ring.

Let's think about alternatives. Can we use the compact everywhere in flash attention?

[jan-wassenberg]

This is not constexpr, please use normal variable naming (snake_case).

[jan-wassenberg]

RuntimeConfig is the older/deprecated mechanism. Prefer to pass attention_impl directly, and ideally later we'll move it to InferenceArgs; but definitely let's avoid further dependency on RuntimeConfig.