gqa short-context: exclude KV deploy from wall; correct HBM peak; batch scaffold

Metric fixes (test harnesses, deploy artifact + wrong constant):
- wall = max(t_end) - min(t_start): exclude the one-time KV-cache deploy
  from the measured decode/prefill step (it was 92-99% of wall, mapping-
  invariant, and masked the real per-mapping separation).
- PEAK_PE_HBM_BPS 128 -> 256 B/ns (8 channels/PE x 32 GB/s); the old value
  was half the modeled per-PE HBM BW, so hbm_bw_util read >1.0 once wall
  was corrected. All six short-context sweep CSVs regenerated/repatched.

Result: the four KV mappings now separate along a {64,32,16,8}-active-PE
ladder (decode 8-kv/1-kv = 4.8x at 8K, 7.4x at 64K), not "modest/tied" as
before; decode is bandwidth-bound at 46-76% of the 256 GB/s per-PE ceiling.

Report (S5.2 rewrite):
- Replace the tied-wall / 8x-per-PE-util claims (both deploy artifacts)
  with the corrected separation and a density trade-off (per-CUBE KV
  footprint, Fig 16 top panel switched to a wall-invariant metric).
- Add a projected latency-vs-batched-throughput analysis (marked
  projected, not measured): dense mappings win throughput, 1-kv wins
  latency; converges at long context.
- Regenerate Fig 15/16/17/18; fix plot script hardcoded ROOT path.

Batch experiment (Part 2, cube_base):
- Add backward-compatible cube_base=0 scalar to the decode kernel so a
  batched user placed at DPPolicy.cube_start addresses 0-based shards.
  Default preserves single-user behavior (32 decode tests pass unchanged).
- New batch harness (skipped): concurrent B>=2 launches hit a sim routing
  issue (sip0.cube0.pe8); single-user path verified. Concurrency fix next.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-07-22 16:25:04 -07:00
parent cb6d21f145
commit 711a9a257f
23 changed files with 403 additions and 187 deletions
@@ -94,6 +94,7 @@ def gqa_attention_decode_short_kernel(
C: int,
P: int,
kv_per_cube: int,
cube_base: int = 0,
*,
tl,
) -> None:
@@ -106,6 +107,9 @@ def gqa_attention_decode_short_kernel(
G = h_q // h_kv
pe_id = tl.program_id(axis=0)
cube_id = tl.program_id(axis=1)
# User-local cube index: a batched user placed at cube_start=cube_base
# addresses its own shards from a 0-based cube index (default 0 = single user).
cube_local = cube_id - cube_base
pe_in_group = pe_id % group_size
group_id_in_cube = pe_id // group_size
@@ -116,13 +120,13 @@ def gqa_attention_decode_short_kernel(
V_HEAD_BYTES = S_kv * KV_ROW_BYTES
# Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
q_base = q_ptr + cube_id * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
q_base = q_ptr + cube_local * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
k_shard_base = (k_ptr
+ cube_id * kv_per_cube * K_HEAD_BYTES
+ cube_local * kv_per_cube * K_HEAD_BYTES
+ group_id_in_cube * K_HEAD_BYTES
+ pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
v_shard_base = (v_ptr
+ cube_id * kv_per_cube * V_HEAD_BYTES
+ cube_local * kv_per_cube * V_HEAD_BYTES
+ group_id_in_cube * V_HEAD_BYTES
+ pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)