gqa(decode): add d_head-TP kernels + measurement runner + figure generators
Two new long-ctx decode attention kernels for the d_head-TP sharding
variants the 6-case chart predicts:
· _gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead.py
Cube-SP × PE-TP-d_head (Case 4 in chart). Per cube holds
S_kv/C tokens of full d_head; per PE holds same tokens but
only d_head/P dims. Partial Q·Kᵀ scores reduced intra-cube
before softmax; outer (m,ℓ,O) merge two-phase (intra+inter).
· _gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp.py
Cube-TP-d_head × PE-SP (Case 5). Per cube holds full S_kv
with only d_head/C dims; per PE holds S_kv/P of those dims.
Partial scores reduced inter-cube (UCIe) before softmax.
Sweep dispatch (gqa_decode_long_ctx_4cases.py) extended with two
new panels so the milestone-1h-gqa sweep covers all 6 cases.
Smoke test scripts/verify_case4_dhead_tp.py runs Cases 4/5/6 at
S_kv=2K to validate the kernels load and execute end-to-end.
Plus the figure-generation toolchain that produced the committed
PNGs in the prior commit (dd3337f):
· paper_plot_gqa_4cases_summary.py - 3-panel summary +
2-panel (analytical / paired-measured) chart generator.
_plot_comm now takes mode="analytical" | "paired".
· paper_plot_gqa_kv_sharding_diagram.py - 6-case 2-D KV-tensor
diagram + companion comparison-table PNG.
· measure_gqa_decode_placement_comm.py - runs all 6 kernels at
S_kv=8K, sums actual IPCQ-copy bytes from engine.op_log,
scales partial-score AR ×128 to S_kv=1M, writes
gqa_3cases_measured_comm.json (committed).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,57 @@
|
||||
{
|
||||
"S_kv_measured": 8192,
|
||||
"S_kv_headline": 1048576,
|
||||
"n_layers": 80,
|
||||
"num_pes": 64,
|
||||
"wo_ffn_per_token_bytes": 1310720,
|
||||
"cases": {
|
||||
"1": {
|
||||
"label": "Case 1 (Cube-Repl x PE-repl)",
|
||||
"total_ipcq_bytes_one_layer": 0,
|
||||
"per_pe_partial_score_bytes_one_layer": 0,
|
||||
"per_pe_mlo_bytes_one_layer": 0,
|
||||
"per_pe_attn_bytes_per_token_at_1M": 0,
|
||||
"per_pe_total_bytes_per_token_at_1M": 1310720
|
||||
},
|
||||
"2": {
|
||||
"label": "Case 2 (Cube-SP x PE-repl)",
|
||||
"total_ipcq_bytes_one_layer": 14560,
|
||||
"per_pe_partial_score_bytes_one_layer": 0,
|
||||
"per_pe_mlo_bytes_one_layer": 227,
|
||||
"per_pe_attn_bytes_per_token_at_1M": 18160,
|
||||
"per_pe_total_bytes_per_token_at_1M": 1328880
|
||||
},
|
||||
"3": {
|
||||
"label": "Case 3 (Cube-Repl x PE-SP)",
|
||||
"total_ipcq_bytes_one_layer": 116480,
|
||||
"per_pe_partial_score_bytes_one_layer": 0,
|
||||
"per_pe_mlo_bytes_one_layer": 1820,
|
||||
"per_pe_attn_bytes_per_token_at_1M": 145600,
|
||||
"per_pe_total_bytes_per_token_at_1M": 1456320
|
||||
},
|
||||
"4": {
|
||||
"label": "Case 4 (Cube-SP x PE-TP d_head)",
|
||||
"total_ipcq_bytes_one_layer": 1851136,
|
||||
"per_pe_partial_score_bytes_one_layer": 16384,
|
||||
"per_pe_mlo_bytes_one_layer": 12540,
|
||||
"per_pe_attn_bytes_per_token_at_1M": 168775360,
|
||||
"per_pe_total_bytes_per_token_at_1M": 170086080
|
||||
},
|
||||
"5": {
|
||||
"label": "Case 5 (Cube-TP d_head x PE-SP)",
|
||||
"total_ipcq_bytes_one_layer": 1851136,
|
||||
"per_pe_partial_score_bytes_one_layer": 16384,
|
||||
"per_pe_mlo_bytes_one_layer": 12540,
|
||||
"per_pe_attn_bytes_per_token_at_1M": 168775360,
|
||||
"per_pe_total_bytes_per_token_at_1M": 170086080
|
||||
},
|
||||
"6": {
|
||||
"label": "Case 6 (Cube-SP x PE-SP) [*]",
|
||||
"total_ipcq_bytes_one_layer": 131040,
|
||||
"per_pe_partial_score_bytes_one_layer": 0,
|
||||
"per_pe_mlo_bytes_one_layer": 2047,
|
||||
"per_pe_attn_bytes_per_token_at_1M": 163760,
|
||||
"per_pe_total_bytes_per_token_at_1M": 1474480
|
||||
}
|
||||
}
|
||||
}
|
||||
+363
@@ -0,0 +1,363 @@
|
||||
"""GQA decode kernel — Case 4 (Cube-SP × PE-TP on d_head).
|
||||
|
||||
Per the per-PE memory/comm summary chart (paper_plot_gqa_4cases_summary.py):
|
||||
this is the d_head-TP variant — KV is sharded 64-way (cube=row_wise on
|
||||
S_kv, pe=column_wise on d_head), so per-PE KV memory matches Case 6 ★
|
||||
(640 MB / PE at 1M context across 80 layers), but the attention
|
||||
algorithm needs an extra mid-layer partial-score AllReduce because
|
||||
each PE only owns d_head/P dimensions of K/V.
|
||||
|
||||
Algorithm (single layer, decode T_q=1, B=1):
|
||||
1. Each PE loads its own d_head/P slice of Q, K, V (Q replicated
|
||||
in HBM; PE-load via column-wise DPPolicy carves out the slice).
|
||||
2. Per S_kv tile, each PE computes a PARTIAL score tile:
|
||||
partial_score = Q_shard · K_shard^T shape: (G·T_q, S_local)
|
||||
3. AR-SUM these partial scores INTRA-CUBE (across the P PEs of a
|
||||
cube) so every PE has the full Q · K^T for the cube's S_local
|
||||
tokens. Online-softmax merge folds tiles into a running state.
|
||||
4. After all tiles, each PE has (m_i, ℓ_i, O_i_shard) for cube i,
|
||||
where O_i_shard is the cube's partial attention output on this
|
||||
PE's d_head/P slice.
|
||||
5. INTER-CUBE online-softmax merge of (m, ℓ, O_shard) across C
|
||||
cubes (lrab center-root reduce). After this, every PE has the
|
||||
full-sequence attention output for its d_head/P slice.
|
||||
6. Final normalise + store: the P PEs of the root cube each write
|
||||
their (G·T_q, d_head/P) slice of the output.
|
||||
|
||||
Tensor layout:
|
||||
Q : (T_q, h_q · d_head) cube=replicate, pe=column_wise on d_head
|
||||
→ per-PE shape (T_q, h_q · d_head/P), byte-reshaped to
|
||||
(G·T_q, d_head/P) for the per-PE GEMM input.
|
||||
K : (S_kv, h_kv · d_head) cube=row_wise on S_kv, pe=column_wise
|
||||
on d_head → per-PE shape (S_local, d_head/P), byte-reshaped to
|
||||
(d_head/P, S_local) for the K^T-form GEMM.
|
||||
V : same layout as K.
|
||||
O : (T_q, h_q · d_head) cube=replicate, pe=column_wise on d_head
|
||||
→ each writer PE owns its d_head/P slice.
|
||||
|
||||
Topology / SFR:
|
||||
- Requires ``configure_sfr_intercube_multisip`` (provides disjoint
|
||||
``intra_*`` and ``E/W/N/S`` namespaces).
|
||||
- Intra-CUBE PEs in a 2×4 grid; AR-SUM of partial scores uses a
|
||||
reduce-to-PE0-then-broadcast pattern over intra_E/W + intra_N/S.
|
||||
- Inter-CUBE lrab center-root reduce on (m, ℓ, O_shard) over the
|
||||
4×2 cube sub-mesh (root cube = 6).
|
||||
|
||||
Naming note: this kernel exemplifies the chart's "PE-TP" label
|
||||
meaning **d_head TP** (NOT the batch-axis TP of the original slide-17
|
||||
Case 1). The chart-side Case label is "Case 4 Cube-SP × PE-TP".
|
||||
|
||||
ADR-0065 P5 scope note: this kernel is meant for op_log / latency
|
||||
measurement (matches the milestone-1h-gqa sweep mode). Numeric
|
||||
data-mode parity is a separate follow-up.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
|
||||
|
||||
# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh).
|
||||
_SUB_W = 4
|
||||
_SUB_H = 2
|
||||
_ROOT_COL = _SUB_W // 2 # 2
|
||||
_ROOT_ROW = _SUB_H // 2 # 1
|
||||
_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL # 6
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples.
|
||||
|
||||
O is the per-PE d_head-slice partial; merge is element-wise so it
|
||||
composes naturally even with O being a d_head shard.
|
||||
"""
|
||||
m_new = tl.maximum(m_local, m_other)
|
||||
scale_old = tl.exp(m_local - m_new)
|
||||
scale_new = tl.exp(m_other - m_new)
|
||||
l_new = l_local * scale_old + l_other * scale_new
|
||||
O_new = O_local * scale_old + O_other * scale_new
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def _intra_cube_allreduce_sum(tensor, *, tl,
|
||||
pe_col: int, pe_row: int,
|
||||
pe_cols_used: int, pe_rows_used: int):
|
||||
"""Sum-AllReduce a tensor across the 8 PEs of one cube (2×4 grid).
|
||||
|
||||
Reduce-to-PE0 (row chain west + col bridge north) followed by
|
||||
broadcast-from-PE0 (col bridge south + row chain east). Each phase
|
||||
operates in place on ``tensor`` via ``tl.copy_to``.
|
||||
|
||||
Used here to combine PARTIAL Q·K^T score tiles across the P=8 PEs
|
||||
of a cube — each PE has only d_head/P dims, so each contributes a
|
||||
partial dot product; the full score = elementwise sum across PEs.
|
||||
"""
|
||||
# ── Phase 1a: row reduce (west) — accumulate at pe_col=0 ────────
|
||||
if pe_cols_used > 1:
|
||||
if pe_col == pe_cols_used - 1:
|
||||
tl.send(dir="intra_W", src=tensor)
|
||||
elif 0 < pe_col < pe_cols_used - 1:
|
||||
with tl.scratch_scope():
|
||||
other = tl.recv(dir="intra_E", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, tensor + other)
|
||||
tl.send(dir="intra_W", src=tensor)
|
||||
elif pe_col == 0:
|
||||
with tl.scratch_scope():
|
||||
other = tl.recv(dir="intra_E", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, tensor + other)
|
||||
|
||||
# ── Phase 1b: col reduce on pe_col=0 (north) — accumulate at row=0 ─
|
||||
if pe_col == 0 and pe_rows_used > 1:
|
||||
if pe_row == pe_rows_used - 1:
|
||||
tl.send(dir="intra_N", src=tensor)
|
||||
elif pe_row == 0:
|
||||
with tl.scratch_scope():
|
||||
other = tl.recv(dir="intra_S", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, tensor + other)
|
||||
|
||||
# ── Phase 2a: col broadcast (south) — row=0 → row=1 on pe_col=0 ─
|
||||
if pe_col == 0 and pe_rows_used > 1:
|
||||
if pe_row == 0:
|
||||
tl.send(dir="intra_S", src=tensor)
|
||||
elif pe_row > 0:
|
||||
with tl.scratch_scope():
|
||||
received = tl.recv(dir="intra_N", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, received)
|
||||
|
||||
# ── Phase 2b: row broadcast (east) — pe_col=0 → all on each row ─
|
||||
if pe_cols_used > 1:
|
||||
if pe_col == 0:
|
||||
tl.send(dir="intra_E", src=tensor)
|
||||
elif 0 < pe_col < pe_cols_used - 1:
|
||||
with tl.scratch_scope():
|
||||
received = tl.recv(dir="intra_W", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, received)
|
||||
tl.send(dir="intra_E", src=tensor)
|
||||
elif pe_col == pe_cols_used - 1:
|
||||
with tl.scratch_scope():
|
||||
received = tl.recv(dir="intra_W", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, received)
|
||||
|
||||
|
||||
def gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
P: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Case-4 decode: Cube-SP × PE-TP-on-d_head, lrab center-root at cube 6."""
|
||||
G = h_q // h_kv
|
||||
if d_head % P != 0:
|
||||
raise ValueError(
|
||||
f"Case 4 requires d_head={d_head} divisible by P={P}"
|
||||
)
|
||||
if S_kv % C != 0:
|
||||
raise ValueError(
|
||||
f"Case 4 requires S_kv={S_kv} divisible by C={C}"
|
||||
)
|
||||
|
||||
d_head_local = d_head // P # 128 / 8 = 16 in the headline setup
|
||||
S_local = S_kv // C
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
|
||||
PE_GRID_COLS = 4
|
||||
pe_col = pe_id % PE_GRID_COLS
|
||||
pe_row = pe_id // PE_GRID_COLS
|
||||
pe_cols_used = min(PE_GRID_COLS, P)
|
||||
pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
|
||||
|
||||
KV_ROW_BYTES = d_head_local * 2 # bytes per K/V row at this PE's d_head slice
|
||||
|
||||
# ── Load this PE's d_head slice of Q (byte-conserving reshape) ──
|
||||
# Q in HBM is (T_q, h_q · d_head); the column_wise pe shard gives
|
||||
# each PE (T_q, h_q · d_head/P). Reshape to (G·T_q, d_head/P) for the
|
||||
# GEMM: byte-conserving since G = h_q (h_kv = 1).
|
||||
Q = tl.load(q_ptr, shape=(G * T_q, d_head_local), dtype="f16")
|
||||
|
||||
# ── Local attention (S_kv-axis tile sweep) with intra-cube AR sum ─
|
||||
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
|
||||
|
||||
# Tile 0: establish persistent (m, ℓ, O_shard). O_shard is the
|
||||
# per-PE d_head/P slice of the partial attention output.
|
||||
tile_s0 = min(TILE_S_KV, S_local)
|
||||
K_T = tl.load(k_ptr, shape=(d_head_local, tile_s0), dtype="f16")
|
||||
V = tl.load(v_ptr, shape=(tile_s0, d_head_local), dtype="f16")
|
||||
partial_scores0 = tl.dot(Q, K_T)
|
||||
# AR-SUM partial scores intra-cube → every PE in the cube has the
|
||||
# full Q · K^T_tile0 scores.
|
||||
_intra_cube_allreduce_sum(
|
||||
partial_scores0, tl=tl,
|
||||
pe_col=pe_col, pe_row=pe_row,
|
||||
pe_cols_used=pe_cols_used, pe_rows_used=pe_rows_used,
|
||||
)
|
||||
m_local = tl.max(partial_scores0, axis=-1)
|
||||
exp_scores = tl.exp(partial_scores0 - m_local)
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
# Partial output sharded on d_head: (G·T_q, d_head/P) per PE.
|
||||
O_local = tl.dot(exp_scores, V)
|
||||
|
||||
for tile_idx in range(1, n_tiles):
|
||||
tile_start = tile_idx * TILE_S_KV
|
||||
tile_s = min(TILE_S_KV, S_local - tile_start)
|
||||
with tl.scratch_scope():
|
||||
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(d_head_local, tile_s), dtype="f16")
|
||||
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(tile_s, d_head_local), dtype="f16")
|
||||
partial_scores_t = tl.dot(Q, K_T_t)
|
||||
_intra_cube_allreduce_sum(
|
||||
partial_scores_t, tl=tl,
|
||||
pe_col=pe_col, pe_row=pe_row,
|
||||
pe_cols_used=pe_cols_used, pe_rows_used=pe_rows_used,
|
||||
)
|
||||
m_tile = tl.max(partial_scores_t, axis=-1)
|
||||
exp_scores_t = tl.exp(partial_scores_t - m_tile)
|
||||
l_tile = tl.sum(exp_scores_t, axis=-1)
|
||||
O_tile = tl.dot(exp_scores_t, V_t)
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
|
||||
# ── Inter-CUBE lrab-adapted center-root reduce on (m, ℓ, O_shard) ─
|
||||
# Every PE participates (each contributes its own d_head/P slice
|
||||
# of O_shard); the merge is element-wise so d_head-sharding is
|
||||
# preserved through the AR.
|
||||
row = cube_id // _SUB_W
|
||||
col = cube_id % _SUB_W
|
||||
|
||||
# Phase 1: row reduce — converge at col == _ROOT_COL.
|
||||
if col == 0:
|
||||
tl.send(dir="E", src=m_local)
|
||||
tl.send(dir="E", src=l_local)
|
||||
tl.send(dir="E", src=O_local)
|
||||
elif 0 < col < _ROOT_COL:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
tl.send(dir="E", src=m_local)
|
||||
tl.send(dir="E", src=l_local)
|
||||
tl.send(dir="E", src=O_local)
|
||||
elif col == _ROOT_COL:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
elif _ROOT_COL < col < _SUB_W - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
tl.send(dir="W", src=m_local)
|
||||
tl.send(dir="W", src=l_local)
|
||||
tl.send(dir="W", src=O_local)
|
||||
elif col == _SUB_W - 1:
|
||||
tl.send(dir="W", src=m_local)
|
||||
tl.send(dir="W", src=l_local)
|
||||
tl.send(dir="W", src=O_local)
|
||||
|
||||
# Phase 2: col reduce on col == _ROOT_COL — converge at row == _ROOT_ROW.
|
||||
if col == _ROOT_COL:
|
||||
if row == 0:
|
||||
tl.send(dir="S", src=m_local)
|
||||
tl.send(dir="S", src=l_local)
|
||||
tl.send(dir="S", src=O_local)
|
||||
elif 0 < row < _ROOT_ROW:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
tl.send(dir="S", src=m_local)
|
||||
tl.send(dir="S", src=l_local)
|
||||
tl.send(dir="S", src=O_local)
|
||||
elif row == _ROOT_ROW:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if _SUB_H - 1 > _ROOT_ROW:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
elif _ROOT_ROW < row < _SUB_H - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
tl.send(dir="N", src=m_local)
|
||||
tl.send(dir="N", src=l_local)
|
||||
tl.send(dir="N", src=O_local)
|
||||
elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW:
|
||||
tl.send(dir="N", src=m_local)
|
||||
tl.send(dir="N", src=l_local)
|
||||
tl.send(dir="N", src=O_local)
|
||||
|
||||
# ── Final normalise + store (root cube: each P-PE writes its d_head slice) ─
|
||||
if cube_id == _ROOT_CUBE:
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr, O_final)
|
||||
+265
@@ -0,0 +1,265 @@
|
||||
"""GQA decode kernel — Case 5 (Cube-TP on d_head × PE-SP on S_kv).
|
||||
|
||||
Per the per-PE memory/comm summary chart (paper_plot_gqa_4cases_summary.py):
|
||||
this is the second d_head-TP variant — KV is sharded 64-way (cube=
|
||||
column_wise on d_head, pe=row_wise on S_kv). Per-PE KV memory matches
|
||||
Cases 4 and 6 ★ (640 MB / PE at 1M context across 80 layers), but the
|
||||
attention algorithm pays the partial-score AR over INTER-cube fabric
|
||||
(UCIe, slower) instead of Case 4's intra-cube fabric (NoC).
|
||||
|
||||
Symmetric mirror of Case 4 — intra and inter-cube AR roles SWAPPED:
|
||||
|
||||
Step Case 4 (Cube-SP × PE-TP/d_head) Case 5 (THIS — Cube-TP/d_head × PE-SP)
|
||||
─────────────── ──────────────────────────────── ──────────────────────────────────────
|
||||
partial Q·K^T AR-SUM intra-cube (NoC) AR-SUM inter-cube (UCIe)
|
||||
online-softmax inter-cube lrab on (m, ℓ, O) intra-cube merge on (m, ℓ, O)
|
||||
(every PE participates; each (online-softmax reduce-to-PE0)
|
||||
PE owns its d_head/P slice)
|
||||
writer root cube's P PEs (each writes each cube's PE 0 writes its
|
||||
its d_head/P slice) d_head/C slice (C writers total)
|
||||
|
||||
Algorithm:
|
||||
1. Each PE loads its (d_head/C × S_kv/P) slice of K and V; full-d_head
|
||||
Q from HBM but only its d_head/C slice via column_wise DPPolicy.
|
||||
2. Per S_kv tile (per PE — already a strict subset of the cube's
|
||||
S_kv), compute partial_score = Q_shard · K_shard^T. Shape:
|
||||
(G·T_q, S_kv/P). AR-SUM these PARTIAL scores INTER-CUBE across
|
||||
the C cubes (same PE rank, different d_head slices) so every PE
|
||||
in cube c, rank p has the full Q · K^T_p scores for THIS PE's
|
||||
S_kv/P token range.
|
||||
3. Local softmax + P · V_shard → partial output (G·T_q, d_head/C)
|
||||
for this PE's S_kv/P slice and this cube's d_head/C slice.
|
||||
4. INTRA-CUBE online-softmax merge of (m, ℓ, O) reduce-to-PE0
|
||||
across the cube's P PEs (different S_kv slices, same d_head slice).
|
||||
5. PE 0 of each cube normalises and writes the cube's d_head/C slice.
|
||||
|
||||
Tensor layout:
|
||||
Q : (T_q, h_q · d_head) cube=column_wise on d_head, pe=replicate.
|
||||
Per PE shape (T_q, h_q · d_head/C); reshape to (G·T_q, d_head/C).
|
||||
K : (S_kv, h_kv · d_head) cube=column_wise on d_head, pe=row_wise on S_kv.
|
||||
Per PE shape (S_kv/P, d_head/C); reshape to (d_head/C, S_kv/P) for K^T.
|
||||
V : same layout as K.
|
||||
O : (T_q, h_q · d_head) cube=column_wise on d_head, pe=replicate.
|
||||
Per-cube writer (PE 0) owns its d_head/C slice.
|
||||
|
||||
Topology / SFR:
|
||||
- Requires ``configure_sfr_intercube_multisip`` (disjoint intra_*
|
||||
and E/W/N/S namespaces).
|
||||
- INTRA-CUBE PEs arranged 2×4 (used for online-softmax merge).
|
||||
- INTER-CUBE 4×2 cube sub-mesh (used for partial-score AR sum).
|
||||
|
||||
ADR-0065 P5 scope note: op_log / latency measurement only; numeric
|
||||
data-mode parity is a separate follow-up (matches the milestone-1h-gqa
|
||||
sweep mode).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024 # per-tile S_kv width within a PE's own S_kv/P shard.
|
||||
|
||||
# Cube sub-mesh geometry (4×2). For Case 5 we pick cube 0 as the AR-sum
|
||||
# root (simpler than lrab; partial-score AR doesn't need geometric centering
|
||||
# because the reduction is element-wise sum, not online-softmax merge).
|
||||
_SUB_W = 4
|
||||
_SUB_H = 2
|
||||
_AR_ROOT_CUBE = 0
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
|
||||
m_new = tl.maximum(m_local, m_other)
|
||||
scale_old = tl.exp(m_local - m_new)
|
||||
scale_new = tl.exp(m_other - m_new)
|
||||
l_new = l_local * scale_old + l_other * scale_new
|
||||
O_new = O_local * scale_old + O_other * scale_new
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def _inter_cube_allreduce_sum(tensor, *, tl, cube_id: int,
|
||||
sub_w: int = _SUB_W, sub_h: int = _SUB_H):
|
||||
"""Sum-AllReduce a tensor INTER-CUBE across the 4×2 cube sub-mesh.
|
||||
|
||||
Reduce-to-cube-0 (row chain west + col bridge north) followed by
|
||||
broadcast-from-cube-0 (col bridge south + row chain east). The
|
||||
payload is the partial-score tile — each cube has a different
|
||||
d_head slice's contribution; element-wise sum recovers the full
|
||||
Q·K^T scores for this PE rank's S_kv/P tokens.
|
||||
|
||||
All P PEs of every cube participate independently (per-PE-rank
|
||||
inter-cube fabric lanes).
|
||||
"""
|
||||
row = cube_id // sub_w
|
||||
col = cube_id % sub_w
|
||||
|
||||
# ── Phase 1a: row reduce (west) — accumulate at col=0 ────────────
|
||||
if sub_w > 1:
|
||||
if col == sub_w - 1:
|
||||
tl.send(dir="W", src=tensor)
|
||||
elif 0 < col < sub_w - 1:
|
||||
with tl.scratch_scope():
|
||||
other = tl.recv(dir="E", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, tensor + other)
|
||||
tl.send(dir="W", src=tensor)
|
||||
elif col == 0:
|
||||
with tl.scratch_scope():
|
||||
other = tl.recv(dir="E", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, tensor + other)
|
||||
|
||||
# ── Phase 1b: col reduce on col=0 (north) — accumulate at row=0 ──
|
||||
if col == 0 and sub_h > 1:
|
||||
if row == sub_h - 1:
|
||||
tl.send(dir="N", src=tensor)
|
||||
elif row == 0:
|
||||
with tl.scratch_scope():
|
||||
other = tl.recv(dir="S", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, tensor + other)
|
||||
|
||||
# ── Phase 2a: col broadcast (south) — row=0 → row=1 on col=0 ─────
|
||||
if col == 0 and sub_h > 1:
|
||||
if row == 0:
|
||||
tl.send(dir="S", src=tensor)
|
||||
elif row > 0:
|
||||
with tl.scratch_scope():
|
||||
received = tl.recv(dir="N", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, received)
|
||||
|
||||
# ── Phase 2b: row broadcast (east) — col=0 → all on each row ─────
|
||||
if sub_w > 1:
|
||||
if col == 0:
|
||||
tl.send(dir="E", src=tensor)
|
||||
elif 0 < col < sub_w - 1:
|
||||
with tl.scratch_scope():
|
||||
received = tl.recv(dir="W", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, received)
|
||||
tl.send(dir="E", src=tensor)
|
||||
elif col == sub_w - 1:
|
||||
with tl.scratch_scope():
|
||||
received = tl.recv(dir="W", shape=tensor.shape, dtype="f16")
|
||||
tl.copy_to(tensor, received)
|
||||
|
||||
|
||||
def gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
P: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Case-5 decode: Cube-TP-on-d_head × PE-SP-on-S_kv."""
|
||||
G = h_q // h_kv
|
||||
if d_head % C != 0:
|
||||
raise ValueError(
|
||||
f"Case 5 requires d_head={d_head} divisible by C={C}"
|
||||
)
|
||||
if S_kv % P != 0:
|
||||
raise ValueError(
|
||||
f"Case 5 requires S_kv={S_kv} divisible by P={P}"
|
||||
)
|
||||
|
||||
d_head_local = d_head // C # 128 / 8 = 16 in the headline setup
|
||||
S_local_pe = S_kv // P # each PE owns S_kv/P tokens
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
|
||||
PE_GRID_COLS = 4
|
||||
pe_col = pe_id % PE_GRID_COLS
|
||||
pe_row = pe_id // PE_GRID_COLS
|
||||
pe_cols_used = min(PE_GRID_COLS, P)
|
||||
pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
|
||||
|
||||
KV_ROW_BYTES = d_head_local * 2 # bytes per K/V row at this PE's d_head slice
|
||||
|
||||
# ── Load this PE's d_head slice of Q ────────────────────────────
|
||||
# Q in HBM is (T_q, h_q · d_head); column_wise cube shard gives each
|
||||
# cube (T_q, h_q · d_head/C); pe=replicate within cube → every PE
|
||||
# in the cube has the same data. Reshape to (G·T_q, d_head/C).
|
||||
Q = tl.load(q_ptr, shape=(G * T_q, d_head_local), dtype="f16")
|
||||
|
||||
# ── Local attention (tile sweep on this PE's S_kv/P shard) ──────
|
||||
n_tiles = (S_local_pe + TILE_S_KV - 1) // TILE_S_KV
|
||||
|
||||
# Tile 0: bootstrap (m, ℓ, O_shard). The partial Q·K^T is across
|
||||
# d_head — must AR inter-cube before softmax.
|
||||
tile_s0 = min(TILE_S_KV, S_local_pe)
|
||||
K_T = tl.load(k_ptr, shape=(d_head_local, tile_s0), dtype="f16")
|
||||
V = tl.load(v_ptr, shape=(tile_s0, d_head_local), dtype="f16")
|
||||
partial_scores0 = tl.dot(Q, K_T)
|
||||
_inter_cube_allreduce_sum(partial_scores0, tl=tl, cube_id=cube_id)
|
||||
m_local = tl.max(partial_scores0, axis=-1)
|
||||
exp_scores = tl.exp(partial_scores0 - m_local)
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
# Partial output (G·T_q, d_head/C) — sharded on d_head only; PE
|
||||
# still holds the full d_head/C slice for its S_kv/P tokens.
|
||||
O_local = tl.dot(exp_scores, V)
|
||||
|
||||
for tile_idx in range(1, n_tiles):
|
||||
tile_start = tile_idx * TILE_S_KV
|
||||
tile_s = min(TILE_S_KV, S_local_pe - tile_start)
|
||||
with tl.scratch_scope():
|
||||
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(d_head_local, tile_s), dtype="f16")
|
||||
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(tile_s, d_head_local), dtype="f16")
|
||||
partial_scores_t = tl.dot(Q, K_T_t)
|
||||
_inter_cube_allreduce_sum(partial_scores_t, tl=tl,
|
||||
cube_id=cube_id)
|
||||
m_tile = tl.max(partial_scores_t, axis=-1)
|
||||
exp_scores_t = tl.exp(partial_scores_t - m_tile)
|
||||
l_tile = tl.sum(exp_scores_t, axis=-1)
|
||||
O_tile = tl.dot(exp_scores_t, V_t)
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
|
||||
# ── INTRA-CUBE online-softmax reduce-to-PE0 (different S_kv slices) ─
|
||||
# Each cube ends with (m, ℓ, O) for its d_head/C slice at PE 0.
|
||||
if pe_cols_used > 1:
|
||||
if pe_col < pe_cols_used - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if pe_col > 0:
|
||||
tl.send(dir="intra_W", src=m_local)
|
||||
tl.send(dir="intra_W", src=l_local)
|
||||
tl.send(dir="intra_W", src=O_local)
|
||||
|
||||
if pe_col == 0 and pe_rows_used > 1:
|
||||
if pe_row < pe_rows_used - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if pe_row > 0:
|
||||
tl.send(dir="intra_N", src=m_local)
|
||||
tl.send(dir="intra_N", src=l_local)
|
||||
tl.send(dir="intra_N", src=O_local)
|
||||
|
||||
# ── Final normalise + store (one writer per cube: PE 0; each
|
||||
# cube owns a disjoint d_head/C slice of the output) ─────────
|
||||
if pe_id == 0:
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr, O_final)
|
||||
@@ -38,6 +38,12 @@ from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp import (
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
|
||||
gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import (
|
||||
_ccl_cfg,
|
||||
_summarize_op_log,
|
||||
@@ -58,10 +64,12 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep_decode.json"
|
||||
|
||||
|
||||
_PANELS = (
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp", # Case 4 ★ optimal
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp", # Case 2 (no comm; 8× memory)
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp", # Case 3 (intra-CUBE AR only; 8× memory)
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp", # Case 1 (inter-CUBE lrab only; PE-TP B=1 waste)
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp", # Case 6 ★ optimal (S_kv-SP both axes; was "Case 4" in slide-17 numbering)
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp", # Case 1 (no sharding; 40 GB / PE @ 1M)
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp", # Case 3 (intra-cube AR on (m,ℓ,O); 5 GB / PE)
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp", # Case 2 (inter-cube AR on (m,ℓ,O); 5 GB / PE)
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead", # Case 4 (Cube-SP × PE-TP-on-d_head — d_head-TP partial-score AR intra-cube)
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp", # Case 5 (Cube-TP-on-d_head × PE-SP — d_head-TP partial-score AR inter-cube)
|
||||
)
|
||||
|
||||
# Each entry: (kind, panel-specific params).
|
||||
@@ -102,6 +110,23 @@ _PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
|
||||
"T_q": 1, "S_kv": 8_192,
|
||||
"d_head": 128, "h_q": 8, "h_kv": 1,
|
||||
}),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead": ("decode_long_ctx_cube_sp_pe_tp_dhead", {
|
||||
# Case 4 (new numbering): Cube-SP on S_kv × PE-TP on d_head.
|
||||
# Partial-score AR happens intra-cube (NoC); (m,ℓ,O) merge happens
|
||||
# inter-cube (UCIe). Memory-equal to Case 6; comm dominated by
|
||||
# partial-score AR which scales with S_kv.
|
||||
"C": 8, "P": 8,
|
||||
"T_q": 1, "S_kv": 8_192,
|
||||
"d_head": 128, "h_q": 8, "h_kv": 1,
|
||||
}),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp": ("decode_long_ctx_cube_tp_dhead_pe_sp", {
|
||||
# Case 5 (new numbering): Cube-TP on d_head × PE-SP on S_kv.
|
||||
# Mirror of Case 4: partial-score AR happens INTER-cube (UCIe);
|
||||
# (m,ℓ,O) merge happens INTRA-cube (NoC). Memory-equal to Case 6.
|
||||
"C": 8, "P": 8,
|
||||
"T_q": 1, "S_kv": 8_192,
|
||||
"d_head": 128, "h_q": 8, "h_kv": 1,
|
||||
}),
|
||||
}
|
||||
|
||||
|
||||
@@ -238,6 +263,74 @@ def _run_decode_panel_long_ctx_cube_sp_pe_sp(
|
||||
)
|
||||
|
||||
|
||||
def _run_decode_panel_long_ctx_cube_sp_pe_tp_dhead(
|
||||
ctx, *, panel: str, C: int, P: int,
|
||||
T_q: int, S_kv: int,
|
||||
d_head: int, h_q: int, h_kv: int,
|
||||
) -> None:
|
||||
"""Case 4 (new) runner: Cube-SP × PE-TP-on-d_head.
|
||||
|
||||
Q is cube-replicated, PE-sharded column-wise (d_head/P per PE).
|
||||
K, V are cube=row_wise on S_kv, pe=column_wise on d_head — each PE
|
||||
owns (S_kv/C, d_head/P). The kernel does an intra-cube AR-SUM of
|
||||
partial scores then an inter-cube online-softmax merge of
|
||||
(m, ℓ, O_shard).
|
||||
"""
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_q = DPPolicy(cube="replicate", pe="column_wise",
|
||||
num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="column_wise",
|
||||
num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((T_q, h_q * d_head),
|
||||
dtype="f16", dp=dp_q, name=f"{panel}_q")
|
||||
k = ctx.zeros((S_kv, h_kv * d_head),
|
||||
dtype="f16", dp=dp_kv, name=f"{panel}_k")
|
||||
v = ctx.zeros((S_kv, h_kv * d_head),
|
||||
dtype="f16", dp=dp_kv, name=f"{panel}_v")
|
||||
o = ctx.empty((T_q, h_q * d_head),
|
||||
dtype="f16", dp=dp_q, name=f"{panel}_o")
|
||||
ctx.launch(
|
||||
panel, gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_q, h_kv, d_head, C, P,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
|
||||
def _run_decode_panel_long_ctx_cube_tp_dhead_pe_sp(
|
||||
ctx, *, panel: str, C: int, P: int,
|
||||
T_q: int, S_kv: int,
|
||||
d_head: int, h_q: int, h_kv: int,
|
||||
) -> None:
|
||||
"""Case 5 (new) runner: Cube-TP-on-d_head × PE-SP-on-S_kv.
|
||||
|
||||
Q is cube-sharded column_wise (d_head/C per cube), pe-replicated
|
||||
within cube. K, V are cube=column_wise on d_head, pe=row_wise on
|
||||
S_kv — each PE owns (S_kv/P, d_head/C). The kernel does an
|
||||
inter-cube AR-SUM of partial scores then an intra-cube online-
|
||||
softmax merge of (m, ℓ, O) reduce-to-PE0.
|
||||
"""
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_q = DPPolicy(cube="column_wise", pe="replicate",
|
||||
num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="column_wise", pe="row_wise",
|
||||
num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((T_q, h_q * d_head),
|
||||
dtype="f16", dp=dp_q, name=f"{panel}_q")
|
||||
k = ctx.zeros((S_kv, h_kv * d_head),
|
||||
dtype="f16", dp=dp_kv, name=f"{panel}_k")
|
||||
v = ctx.zeros((S_kv, h_kv * d_head),
|
||||
dtype="f16", dp=dp_kv, name=f"{panel}_v")
|
||||
o = ctx.empty((T_q, h_q * d_head),
|
||||
dtype="f16", dp=dp_q, name=f"{panel}_o")
|
||||
ctx.launch(
|
||||
panel, gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_q, h_kv, d_head, C, P,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
|
||||
def _make_bench_fn(panel: str):
|
||||
kind, params = _PANEL_DISPATCH[panel]
|
||||
|
||||
@@ -250,6 +343,10 @@ def _make_bench_fn(panel: str):
|
||||
_run_decode_panel_long_ctx_cube_repl_pe_sp(ctx, panel=panel, **params)
|
||||
elif kind == "decode_long_ctx_cube_sp_pe_tp":
|
||||
_run_decode_panel_long_ctx_cube_sp_pe_tp(ctx, panel=panel, **params)
|
||||
elif kind == "decode_long_ctx_cube_sp_pe_tp_dhead":
|
||||
_run_decode_panel_long_ctx_cube_sp_pe_tp_dhead(ctx, panel=panel, **params)
|
||||
elif kind == "decode_long_ctx_cube_tp_dhead_pe_sp":
|
||||
_run_decode_panel_long_ctx_cube_tp_dhead_pe_sp(ctx, panel=panel, **params)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"milestone-gqa-decode-long-ctx-4cases panel {panel!r} has "
|
||||
|
||||
Reference in New Issue
Block a user