short context length composite command
This commit is contained in:
eusonice
@@ -0,0 +1,17 @@
|
||||
mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
|
||||
A1,1,8,8192,149.2685345000007,64,0.001756190618994587,656,0.013780533230866485,16.0625,8.0,57.75,2.0
|
||||
A1,1,8,16384,282.2252665000025,64,0.0018576933472394743,1360,0.014545118695104373,32.0625,8.0,57.75,4.0
|
||||
A1,1,8,32768,548.138730500006,64,0.0019129755692392889,2768,0.014961540835691614,64.0625,8.0,57.75,8.0
|
||||
A1,1,8,65536,1080.073058500367,64,0.0019416760593127277,5584,0.01517767698303756,128.0625,8.0,57.75,16.0
|
||||
A2,2,4,8192,142.16883275000066,32,0.0036877843748065785,656,0.028881154333033524,16.03125,8.0,24.75,4.0
|
||||
A2,2,4,16384,277.0102967500018,32,0.0037853322143696724,1360,0.029609007665885423,32.03125,8.0,24.75,8.0
|
||||
A2,2,4,32768,546.8006247504073,32,0.0038353138330044275,2768,0.02998167752182655,64.03125,8.0,24.75,16.0
|
||||
A2,2,4,65536,1087.0962007518285,32,0.003858263874988177,5584,0.03015188534127058,128.03125,8.0,24.75,32.0
|
||||
A4,4,2,8192,142.21417550000066,16,0.007373217165681769,656,0.05768763888097753,16.015625,8.0,8.25,8.0
|
||||
A4,4,2,16384,280.8788035001758,16,0.007466394665122732,1360,0.058373931374247456,32.015625,8.0,8.25,16.0
|
||||
A4,4,2,32768,558.5655195010398,16,0.007509063580845671,2768,0.058686042828569145,64.015625,8.0,8.25,32.0
|
||||
A4,4,2,65536,1114.319951501857,16,0.007528006645388841,5584,0.05882332081702009,128.015625,8.0,8.25,64.0
|
||||
B,8,1,8192,147.93485600006025,8,0.014176185766516818,656,0.11085960701508589,16.0078125,8.0,0.0,16.0
|
||||
B,8,1,16384,294.3171420004266,8,0.014250967413897551,1360,0.1113900460475132,32.0078125,8.0,0.0,32.0
|
||||
B,8,1,32768,587.2722140010417,8,0.01428401991446495,2768,0.11162115018757507,64.0078125,8.0,0.0,64.0
|
||||
B,8,1,65536,1173.1823580017838,8,0.014300603725891686,5584,0.11173710472707321,128.0078125,8.0,0.0,128.0
|
||||
|
@@ -0,0 +1,17 @@
|
||||
mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
|
||||
A1,1,8,8192,1103.359903999566,64,0.0019006962210579498,5248,0.001870649814732448,16.0625,64.0,114688.0,2.0
|
||||
A1,1,8,16384,2182.457823968016,64,0.0019218259129421755,10880,0.0018841143021603984,32.0625,64.0,229376.0,4.0
|
||||
A1,1,8,32768,4343.2822238495155,64,0.0019313983221991736,22144,0.0018898150239762977,64.0625,64.0,458752.0,8.0
|
||||
A1,1,8,65536,8666.440783650422,64,0.001935883071136998,44672,0.0018923570136126979,128.0625,64.0,917504.0,16.0
|
||||
A2,2,4,8192,1097.2535885006619,32,0.003822547535007557,2624,0.0037621203004135906,16.0625,64.0,49152.0,4.0
|
||||
A2,2,4,16384,2179.78425743804,32,0.003848366172648072,5440,0.003772850442394648,32.0625,64.0,98304.0,8.0
|
||||
A2,2,4,32768,4344.834905351148,32,0.003861416225357031,11072,0.003778279349528763,64.0625,64.0,196608.0,16.0
|
||||
A2,2,4,65536,8674.946611121853,32,0.003867969856656271,22336,0.003781003096658628,128.0625,64.0,393216.0,32.0
|
||||
A4,4,2,8192,1103.675383502932,16,0.007600611670231015,1312,0.007480460399321815,16.0625,64.0,16384.0,8.0
|
||||
A4,4,2,16384,2197.0616234262543,16,0.00763620638634706,2720,0.007486362614786297,32.0625,64.0,32768.0,16.0
|
||||
A4,4,2,32768,4383.834103368879,16,0.007654129058897454,5536,0.007489334501679554,64.0625,64.0,65536.0,32.0
|
||||
A4,4,2,65536,8757.379062883949,16,0.007663121981838932,11168,0.007490825682997995,128.0625,64.0,131072.0,64.0
|
||||
B,8,1,8192,1125.9457520018113,8,0.014900554462921552,656,0.01466500492643045,16.0625,64.0,0.0,16.0
|
||||
B,8,1,16384,2242.543191943168,8,0.014962669223294344,1360,0.014669059716747552,32.0625,64.0,0.0,32.0
|
||||
B,8,1,32768,4475.650071825993,8,0.014994216018466134,2768,0.014671388277951352,64.0625,64.0,0.0,64.0
|
||||
B,8,1,65536,8942.79183160376,8,0.015008481750139946,5584,0.014671033662702533,128.0625,64.0,0.0,128.0
|
||||
|
@@ -0,0 +1,191 @@
|
||||
"""GQA decode kernel — composite (second-level) variant.
|
||||
|
||||
Identical mapping to ``_gqa_attention_decode_short.py`` (sequence-shard
|
||||
+ chain reduce + FA2 head fusion, unified A1/A2/A4/B). Per-tile fusion
|
||||
follows the ADR-0065 two-composite pattern (cf. decode_opt2):
|
||||
|
||||
tile 0 primitives establish (m, ℓ, O)
|
||||
tile k>0 #1 Q·Kᵀ composite → scores (pinned primary-out)
|
||||
#2 softmax_merge prologue → P (pinned, auto-binds to GEMM a)
|
||||
+ P·V composite, ``out=O_local``
|
||||
+ add epilogue folding the result into O_local
|
||||
|
||||
The softmax_merge recipe (tl_recipes.py) folds max/exp/sum into the
|
||||
composite, eliminating the GEMM/MATH engine bubble that a flat
|
||||
primitives chain incurs. The intra-group chain reduce after the tile
|
||||
loop is unchanged.
|
||||
|
||||
Shard addressing, layouts, and caller contract are identical to the
|
||||
first-level decode kernel (ADR-0011 D-VA1).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024
|
||||
|
||||
|
||||
def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
|
||||
h_q: int, h_kv: int, S_kv: int) -> None:
|
||||
"""Validate composite-decode config — caller-side, sim-cost 0.
|
||||
|
||||
Mirrors first-level decode ``_validate_config``.
|
||||
"""
|
||||
if kv_per_cube not in (1, 2, 4, 8):
|
||||
raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
|
||||
if T_q != 1:
|
||||
raise ValueError(f"decode requires T_q == 1; got {T_q}")
|
||||
if P != 8:
|
||||
raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
|
||||
if h_q % h_kv != 0:
|
||||
raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
|
||||
if C != h_kv // kv_per_cube:
|
||||
raise ValueError(
|
||||
f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
|
||||
f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
|
||||
)
|
||||
group_size = P // kv_per_cube
|
||||
if S_kv % (group_size * TILE_S_KV) != 0:
|
||||
raise ValueError(
|
||||
f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
|
||||
f"({group_size * TILE_S_KV}); each PE's sequence shard must "
|
||||
f"be a whole number of tiles"
|
||||
)
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
"""Used only by the intra-group chain reduce below (recipe handles
|
||||
per-tile fold internally)."""
|
||||
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 gqa_attention_decode_short_composite_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,
|
||||
kv_per_cube: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Composite-GEMM decode — same mapping as first-level + tl.composite.
|
||||
|
||||
Caller must invoke ``_validate_config(...)`` first.
|
||||
"""
|
||||
group_size = P // kv_per_cube
|
||||
S_local = S_kv // group_size
|
||||
n_tiles_per_pe = S_local // TILE_S_KV
|
||||
|
||||
G = h_q // h_kv
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
pe_in_group = pe_id % group_size
|
||||
group_id_in_cube = pe_id // group_size
|
||||
|
||||
KV_ROW_BYTES = d_head * 2
|
||||
K_TILE_BYTES = d_head * TILE_S_KV * 2
|
||||
Q_ROW_BYTES = G * d_head * 2
|
||||
K_HEAD_BYTES = S_kv * d_head * 2
|
||||
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
|
||||
k_shard_base = (k_ptr
|
||||
+ cube_id * 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
|
||||
+ group_id_in_cube * V_HEAD_BYTES
|
||||
+ pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)
|
||||
|
||||
Q = tl.load(q_base, shape=(G, d_head), dtype="f16")
|
||||
|
||||
# ── Tile 0: establish running (m, ℓ, O) with primitives ──────────
|
||||
# (Reference: decode_opt2 — running state is set up with tl.dot/MATH
|
||||
# primitives, not composite. Recipe-driven composite enters in tile 1+.)
|
||||
K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.load(v_shard_base, shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
scores = tl.dot(Q, K_T)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
exp_scores = tl.exp(centered)
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
O_local = tl.dot(exp_scores, V)
|
||||
|
||||
# ── Tiles 1..n_tiles_per_pe-1: two composites per tile ──────────
|
||||
# #1 Q·Kᵀ composite → scores (pinned primary-out, fed into #2).
|
||||
# #2 softmax_merge prologue + P·V GEMM + add epilogue, all in one
|
||||
# composite: updates (m, ℓ) in place, computes P, runs P·V with
|
||||
# pinned auto-bind, and folds the result into O_local.
|
||||
for tile_idx in range(1, n_tiles_per_pe):
|
||||
with tl.scratch_scope():
|
||||
K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
|
||||
tl.composite(
|
||||
prologue=[{"op": "softmax_merge", "s": scores_t,
|
||||
"m": m_local, "l": l_local, "O": O_local}],
|
||||
op="gemm", b=V_t, out=O_local,
|
||||
epilogue=[{"op": "add", "other": O_local}],
|
||||
)
|
||||
|
||||
# ── Chain reduce ──
|
||||
group_cols = min(4, group_size)
|
||||
group_rows = (group_size + group_cols - 1) // group_cols
|
||||
pe_col_in_group = pe_in_group % group_cols
|
||||
pe_row_in_group = pe_in_group // group_cols
|
||||
|
||||
if group_cols > 1:
|
||||
if pe_col_in_group < group_cols - 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_in_group > 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_in_group == 0 and group_rows > 1:
|
||||
if pe_row_in_group < group_rows - 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_in_group > 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)
|
||||
|
||||
if pe_in_group == 0:
|
||||
O_final = O_local / l_local
|
||||
o_base = (o_ptr
|
||||
+ cube_id * kv_per_cube * Q_ROW_BYTES
|
||||
+ group_id_in_cube * Q_ROW_BYTES)
|
||||
tl.store(o_base, O_final)
|
||||
@@ -0,0 +1,216 @@
|
||||
"""GQA prefill kernel — composite (second-level) variant.
|
||||
|
||||
Identical mapping to ``_gqa_attention_prefill_short.py`` (Q-tile split +
|
||||
FA2 head fusion + IPCQ KV broadcast, unified A1/A2/A4/B).
|
||||
|
||||
Composite hybrid (option c1): Q·Kᵀ uses ``tl.composite(op="gemm")`` so
|
||||
GEMM and any future prologue/epilogue can overlap; P·V stays a plain
|
||||
``tl.dot``. The reason: prefill broadcasts V over IPCQ, so on non-root
|
||||
PEs V is a recv'd TCM slot (unpinned per ADR-0065 D4). Feeding it to a
|
||||
composite would trigger the scheduler's "stream from HBM" branch and
|
||||
PhysAddr-decode the slot address (PE scratch, bit-61 set) → PageFault.
|
||||
``tl.dot`` reads TCM in place, preserving the broadcast model.
|
||||
|
||||
This means prefill cannot host the ``softmax_merge`` prologue (which
|
||||
attaches to the P·V composite). Full second-level fusion would require
|
||||
either: (i) pinning recv handles (sim/API change, ADR-0065 supplement),
|
||||
or (ii) skipping the broadcast model — both out of scope here.
|
||||
|
||||
Shard addressing, layouts, and the caller contract are identical to
|
||||
the first-level kernel (see its docstring for the full Q/K/V/O layout
|
||||
and ADR-0011 D-VA1 offset computation).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024
|
||||
|
||||
|
||||
def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
|
||||
h_q: int, h_kv: int, S_kv: int) -> None:
|
||||
"""Validate composite-prefill config — caller-side, sim-cost 0.
|
||||
|
||||
Mirrors first-level prefill ``_validate_config``.
|
||||
"""
|
||||
if kv_per_cube not in (1, 2, 4, 8):
|
||||
raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
|
||||
if P != 8:
|
||||
raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}")
|
||||
if h_q % h_kv != 0:
|
||||
raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
|
||||
if C != h_kv // kv_per_cube:
|
||||
raise ValueError(
|
||||
f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
|
||||
f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
|
||||
)
|
||||
group_size = P // kv_per_cube
|
||||
if T_q < group_size:
|
||||
raise ValueError(
|
||||
f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
|
||||
f"has q-tile work for the broadcast chain"
|
||||
)
|
||||
if S_kv % TILE_S_KV != 0:
|
||||
raise ValueError(
|
||||
f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
|
||||
)
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
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 gqa_attention_prefill_short_composite_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,
|
||||
kv_per_cube: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Composite-GEMM prefill — same mapping as first-level + tl.composite.
|
||||
|
||||
Caller must invoke ``_validate_config(...)`` first.
|
||||
"""
|
||||
group_size = P // kv_per_cube
|
||||
|
||||
G = h_q // h_kv
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
pe_in_group = pe_id % group_size
|
||||
group_id_in_cube = pe_id // group_size
|
||||
|
||||
q_start = (T_q * pe_in_group) // group_size
|
||||
q_end = (T_q * (pe_in_group + 1)) // group_size
|
||||
T_q_pe = q_end - q_start
|
||||
|
||||
n_tiles = S_kv // TILE_S_KV
|
||||
|
||||
Q_ROW_BYTES = G * d_head * 2
|
||||
KV_ROW_BYTES = d_head * 2
|
||||
K_TILE_BYTES = d_head * TILE_S_KV * 2
|
||||
K_HEAD_BYTES = n_tiles * K_TILE_BYTES
|
||||
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 * T_q * Q_ROW_BYTES
|
||||
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
|
||||
k_head_shard_base = (k_ptr
|
||||
+ cube_id * kv_per_cube * K_HEAD_BYTES
|
||||
+ group_id_in_cube * K_HEAD_BYTES)
|
||||
v_head_shard_base = (v_ptr
|
||||
+ cube_id * kv_per_cube * V_HEAD_BYTES
|
||||
+ group_id_in_cube * V_HEAD_BYTES)
|
||||
|
||||
Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16")
|
||||
|
||||
group_cols = min(4, group_size)
|
||||
group_rows = (group_size + group_cols - 1) // group_cols
|
||||
pe_col_in_group = pe_in_group % group_cols
|
||||
pe_row_in_group = pe_in_group // group_cols
|
||||
|
||||
# ── Tile 0 — broadcast + persistent (m, ℓ, O) ──
|
||||
if pe_in_group == 0:
|
||||
K_T = tl.load(k_head_shard_base,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.load(v_head_shard_base,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
if group_rows > 1:
|
||||
tl.send(dir="intra_S", src=K_T)
|
||||
tl.send(dir="intra_S", src=V)
|
||||
elif pe_col_in_group == 0 and pe_row_in_group > 0:
|
||||
K_T = tl.recv(dir="intra_N",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.recv(dir="intra_N",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
else:
|
||||
K_T = tl.recv(dir="intra_W",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.recv(dir="intra_W",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
|
||||
# Q·Kᵀ is a composite GEMM (pinned load operands). P·V stays a
|
||||
# primitive tl.dot: in the broadcast topology, V on non-root PEs is
|
||||
# a recv'd TCM slot (unpinned), which the composite scheduler would
|
||||
# try to stream from HBM. tl.dot reads TCM in-place, preserving
|
||||
# the broadcast model. softmax_merge prologue requires a composite
|
||||
# P·V to host it — so prefill keeps Q·Kᵀ-only composite (no fused
|
||||
# softmax epilogue); see kernel docstring (c1).
|
||||
scores = tl.composite(op="gemm", a=Q, b=K_T)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
exp_scores = tl.exp(centered)
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
O_local = tl.dot(exp_scores, V)
|
||||
|
||||
# ── Tiles 1..n_tiles-1 ──
|
||||
for tile_idx in range(1, n_tiles):
|
||||
with tl.scratch_scope():
|
||||
if pe_in_group == 0:
|
||||
K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
if group_rows > 1:
|
||||
tl.send(dir="intra_S", src=K_T_t)
|
||||
tl.send(dir="intra_S", src=V_t)
|
||||
elif pe_col_in_group == 0 and pe_row_in_group > 0:
|
||||
K_T_t = tl.recv(dir="intra_N",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.recv(dir="intra_N",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
else:
|
||||
K_T_t = tl.recv(dir="intra_W",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.recv(dir="intra_W",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
# Q·Kᵀ composite; P·V stays primitive (see comment above).
|
||||
scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
|
||||
m_tile = tl.max(scores_t, axis=-1)
|
||||
centered_t = scores_t - m_tile
|
||||
exp_scores_t = tl.exp(centered_t)
|
||||
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)
|
||||
|
||||
O_final = O_local / l_local
|
||||
o_base = (o_ptr
|
||||
+ cube_id * kv_per_cube * T_q * Q_ROW_BYTES
|
||||
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
|
||||
tl.store(o_base, O_final)
|
||||
Reference in New Issue
Block a user