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:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user