SIGN IN SIGN UP

Correctly log dense KV cache size from previously allocated array nbytes (#627)

Tl;DR: Gemma 4 E2B and E4B mix the 256-wide sliding window KV with
512-wide global KV. The manually calculated logger assumed every unique
layer as 512-wide, so it overstated E4B by 12/7 and E2B by 5/3.
Allocation was already right. The log now dynamically sums the allocated
arrays' nbytes. No allocation changes made.

## Summary
Gemma 4 unique cache layers are not one width. Sliding layers store K/V
at `head_dim=256`; global/full layers store `global_head_dim=512`. E4B
YOCO keeps 24 unique layers (20 sliding + 4 full). E2B keeps 15 (12 +
3). `_allocate_dense_caches` already builds each layer from
`kv_heads_per_layer` / `head_dim_per_layer`.

`_log_dense_cache` ignored those lists and multiplied `num_layers ×
blocks × 16 × num_kv_heads × self.head_dim × 2 × dtype`. `self.head_dim`
is `max(256, 512) = 512` from `resolve_max_head_dim` (kernel dispatch,
not sizing). Every unique layer was billed at 4096 bytes/token. Real E4B
is `20×2048 + 4×4096` per token (12/7). Real E2B is 5/3. The planner's
`get_cache_block_size_bytes()` / `_kv_layer_size_sum()` already matched
allocation, so serving size was fine; only the startup `KV cache: … MB`
line was wrong.

The logger now sums `key_caches` and `value_caches` `.nbytes` after
allocate, so it follows the arrays instead of a second formula. The
`dtype` argument is gone because nbytes already includes element size.
`_log_layout_cache` still uses `layout.total_bytes`: those per-layer
entries are views over shared slots, and summing nbytes would
double-count. Default 12B/26B/31B take that layout path and were not on
the bad line.

## Tests
- `test_heterogeneous_log_matches_allocation` — E4B-shaped 24 layers, 3
blocks: log 2,752,512 B, not the old 4,718,592 B product
- `test_uniform_log_matches_allocation` — uniform models still match the
old product
- `test_layout_log_reports_physical_bytes` — layout path still logs
`layout.total_bytes`, not the per-layer nbytes sum

No model or kernel change; no performance claim. CONTRIBUTING
golden-token and benchmark checks do not apply.

Checked and fixed with AI assistance (GPT-5.6 Sol, Opus 5, Grok 4.6). I
first raised the issue and reviewed the patch and take responsibility
for it. Fixes #628.

---------

Signed-off-by: Alexander Tu <awt2138@columbia.edu>
A
AstroWebDev committed
f4651f771255948ac7d72e5e8c3d60f0dd630ba5
Parent: 1c171b8
Committed by GitHub <noreply@github.com> on 8/19/2026, 10:02:56 AM