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