diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite.png b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite.png index 7078001..459cc90 100644 Binary files a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite.png and b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite.png differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite_breakdown.png b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite_breakdown.png new file mode 100644 index 0000000..9e9bac0 Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite_breakdown.png differ diff --git a/scripts/paper/paper_plot_gqa_decode_composite_breakdown.py b/scripts/paper/paper_plot_gqa_decode_composite_breakdown.py new file mode 100644 index 0000000..2581117 --- /dev/null +++ b/scripts/paper/paper_plot_gqa_decode_composite_breakdown.py @@ -0,0 +1,131 @@ +"""Latency breakdown bar chart for the Case-6 composite-command decode study. + +Reads sweep_decode_composite.json and writes a single-figure stacked bar +chart comparing three variants — primitive hand-tiled (16×16×16), +composite GEMM, composite + softmax_merge — at the S_kv = 131 072 point: + + bottom stack: PE_CPU dispatch time (from pe_cpu_dispatch_cycles, ns at + 1 GHz — ADR-0064 Rev2 D3) + top stack: engine time (latency_ns − dispatch cycles) — DMA / GEMM / + MATH / IPCQ work the engine flushes on the critical path + +The dispatch/engine split is a first-order breakdown; in reality the +two paths overlap partially. The note on the plot calls this out. + +Run (after the composite bench sweep): + python scripts/paper/paper_plot_gqa_decode_composite_breakdown.py +""" +from __future__ import annotations + +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_FIG_DIR = ( + _REPO_ROOT / "src" / "kernbench" / "benches" + / "1H_milestone_output" / "gqa" / "long_ctx" +) +_SWEEP_JSON = _FIG_DIR / "sweep_decode_composite.json" +_PAPER_FIG_DIR = ( + _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures" +) + +# The three variants to compare (coarse `primitive` intentionally dropped). +_ORDER = ("primitive_tiled", "composite", "composite_extended") +_LABELS = { + "primitive_tiled": "primitive hand-tiled\n(16×16×16)", + "composite": "composite\nGEMM", + "composite_extended": "composite +\nsoftmax_merge", +} +_S_KV_TARGET = 131_072 + + +def main() -> None: + sweep = json.loads(_SWEEP_JSON.read_text()) + rows = {(r["variant"], r["S_kv"]): r for r in sweep["rows"]} + + dispatch_us = [] + engine_us = [] + totals_us = [] + for v in _ORDER: + r = rows[(v, _S_KV_TARGET)] + disp_ns = r["pe_cpu_dispatch_cycles"] + total_ns = r["latency_ns"] + eng_ns = max(0.0, total_ns - disp_ns) + dispatch_us.append(disp_ns / 1e3) + engine_us.append(eng_ns / 1e3) + totals_us.append(total_ns / 1e3) + + xs = list(range(len(_ORDER))) + labels = [_LABELS[v] for v in _ORDER] + + fig, ax = plt.subplots(figsize=(7.5, 5.0)) + + bars_eng = ax.bar( + xs, engine_us, + color="#4f8a4f", edgecolor="#2a4a2a", label="engine (DMA + GEMM + MATH + IPCQ)", + ) + bars_disp = ax.bar( + xs, dispatch_us, bottom=engine_us, + color="#c0504d", edgecolor="#5a2624", label="PE_CPU dispatch", + ) + + # Segment value labels — engine at mid, dispatch at mid of its stack. + for i, (eng, disp, tot) in enumerate(zip(engine_us, dispatch_us, totals_us)): + ax.text(i, eng / 2, f"{eng:.1f} µs", + ha="center", va="center", fontsize=9, color="white") + if disp / max(totals_us) > 0.04: # only label if visible + ax.text(i, eng + disp / 2, f"{disp:.1f} µs", + ha="center", va="center", fontsize=9, color="white") + else: + ax.annotate(f"{disp:.2f} µs", + xy=(i, tot), xytext=(0, 6), + textcoords="offset points", + ha="center", va="bottom", fontsize=8, + color="#5a2624") + offset_pts = 14 if disp / max(totals_us) > 0.04 else 20 + ax.annotate(f"total {tot:.1f}", + xy=(i, tot), xytext=(0, offset_pts), + textcoords="offset points", + ha="center", va="bottom", fontsize=9, + color="#333", fontweight="bold") + + ax.set_xticks(xs) + ax.set_xticklabels(labels, fontsize=10) + ax.set_ylabel("time (µs)") + ax.set_title( + f"Case-6 decode latency breakdown at $S_{{kv}}=${_S_KV_TARGET // 1024}K\n" + "(engine dominates all three — primitive hand-tiled adds " + f"{dispatch_us[0]:.0f} µs PE_CPU dispatch overhead)", + fontsize=11, + ) + ax.legend(loc="upper right", fontsize=9) + ax.grid(True, axis="y", ls=":", alpha=0.5) + ax.set_ylim(0, max(totals_us) * 1.15) + + fig.text( + 0.5, 0.02, + "First-order breakdown: dispatch and engine paths overlap partially on the real " + "critical path;\ntreat the split as an upper bound on dispatch's contribution.", + ha="center", fontsize=8, color="#666", + ) + fig.tight_layout(rect=(0, 0.06, 1, 1)) + + out = _FIG_DIR / "gqa_decode_long_ctx_composite_breakdown.png" + fig.savefig(out, dpi=150) + plt.close(fig) + print(f"wrote {out}") + + if _PAPER_FIG_DIR.is_dir(): + dst = _PAPER_FIG_DIR / out.name + dst.write_bytes(out.read_bytes()) + print(f"copied {dst}") + + +if __name__ == "__main__": + main() diff --git a/scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py b/scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py index cdf0684..2f60151 100644 --- a/scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py +++ b/scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py @@ -45,11 +45,12 @@ _N_RANKS = 64 # C·P for the Case-6 64-way split. # variant key → (display label, colour, marker) _VARIANT_STYLE = { - "primitive": ("primitive (tl.dot, hand-tiled)", "#c0504d", "o"), - "composite": ("composite GEMM", "#3b6ea5", "s"), - "composite_extended": ("composite + softmax_merge", "#4f8a4f", "^"), + "primitive_tiled": ("primitive hand-tiled (16×16×16)", "#8b5a2b", "D"), + "primitive": ("primitive (tl.dot, hand-tiled)", "#c0504d", "o"), + "composite": ("composite GEMM", "#3b6ea5", "s"), + "composite_extended": ("composite + softmax_merge", "#4f8a4f", "^"), } -_ORDER = ("primitive", "composite", "composite_extended") +_ORDER = ("primitive_tiled", "composite", "composite_extended") def _load() -> dict: diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite.png index 7078001..459cc90 100644 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite.png and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite_breakdown.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite_breakdown.png new file mode 100644 index 0000000..9e9bac0 Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite_breakdown.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json index f4060d6..fb28454 100644 --- a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json +++ b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json @@ -1,6 +1,7 @@ { "version": 2, "variants": [ + "primitive_tiled", "primitive", "composite", "composite_extended" @@ -20,6 +21,20 @@ 131072 ], "rows": [ + { + "variant": "primitive_tiled", + "S_kv": 8192, + "C": 8, + "P": 8, + "T_q": 1, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "pe_cpu_cmd_count": 523, + "pe_cpu_dispatch_cycles": 5011, + "latency_ns": 34727.00300000079, + "latency_derived": true + }, { "variant": "primitive", "S_kv": 8192, @@ -59,6 +74,20 @@ "pe_cpu_dispatch_cycles": 1032, "latency_ns": 30483.359500000362 }, + { + "variant": "primitive_tiled", + "S_kv": 65536, + "C": 8, + "P": 8, + "T_q": 1, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "pe_cpu_cmd_count": 3603, + "pe_cpu_dispatch_cycles": 34467, + "latency_ns": 265116.37900000165, + "latency_derived": true + }, { "variant": "primitive", "S_kv": 65536, @@ -98,6 +127,20 @@ "pe_cpu_dispatch_cycles": 1032, "latency_ns": 231211.1740000015 }, + { + "variant": "primitive_tiled", + "S_kv": 131072, + "C": 8, + "P": 8, + "T_q": 1, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "pe_cpu_cmd_count": 7133, + "pe_cpu_dispatch_cycles": 68228, + "latency_ns": 528202.5190000018, + "latency_derived": true + }, { "variant": "primitive", "S_kv": 131072, @@ -137,6 +180,19 @@ "pe_cpu_dispatch_cycles": 1032, "latency_ns": 460552.9805000025 }, + { + "variant": "primitive_tiled", + "S_kv": 262144, + "C": 8, + "P": 8, + "T_q": 1, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "pe_cpu_cmd_count": 14193, + "pe_cpu_dispatch_cycles": 135750, + "latency_ns": null + }, { "variant": "primitive", "S_kv": 262144, @@ -176,6 +232,19 @@ "pe_cpu_dispatch_cycles": 1032, "latency_ns": null }, + { + "variant": "primitive_tiled", + "S_kv": 524288, + "C": 8, + "P": 8, + "T_q": 1, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "pe_cpu_cmd_count": 28313, + "pe_cpu_dispatch_cycles": 270794, + "latency_ns": null + }, { "variant": "primitive", "S_kv": 524288, @@ -215,6 +284,19 @@ "pe_cpu_dispatch_cycles": 1032, "latency_ns": null }, + { + "variant": "primitive_tiled", + "S_kv": 1048576, + "C": 8, + "P": 8, + "T_q": 1, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "pe_cpu_cmd_count": 56553, + "pe_cpu_dispatch_cycles": 540882, + "latency_ns": null + }, { "variant": "primitive", "S_kv": 1048576, diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py new file mode 100644 index 0000000..26099c8 --- /dev/null +++ b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py @@ -0,0 +1,200 @@ +"""GQA decode kernel — Case 6, **primitive-TILED streamed** (16×16×16 MAC + per-block DMA). + +Same Case-6 placement and (m, ℓ, O) reduce as the primitive baseline +(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the difference is +that each local-attention matmul is *hand-blocked into 16×16×16 GEMMs* +(mac=16), **with each block re-fetching its operand slices from HBM** +(no operand cache) and each (mi, ni) output tile accumulated by a +**deferred sum outside the K-inner loop**. + +Per K-inner block: two ``tl.load`` DMAs (Q and K slice for Q·Kᵀ, or one +V slice for P·V) → one ``tl.dot`` GEMM. After the K-inner loop completes +for a given (mi, ni), the K/16 partial 16×16 results are summed via +sequential ``+`` (MathCmd add) into a single per-tile accumulator; the +sum is discarded (the shared M×N output handle stays zero-initialised, +which softmax reads under the zero-input decode-bench convention). + +This models a strict-streaming architecture (no HBM-operand cache) — the +worst-case dispatch endpoint: every 16³ compute block pays the ADR-0064 +D8 single-op FIXED overhead on DMA (per operand slice) AND on GEMM AND +on the per-tile add chain, exposing the full PE_CPU dispatch pressure +that composite forms (ADR-0065) absorb into PE_SCHEDULER. + +FLOPs are conserved (each 16³ GemmCmd carries the TFLOPS-model compute +of its block; the blocks sum to the full matmul); end-to-end compute +time is unchanged vs the coarse primitive — only the PE_CPU command +count and its dispatch cycles grow. Inputs are zero (decode bench +convention), so the blocked accumulation and the softmax-consumed +shared handle are identically zero — the value flow is a placeholder; +this study measures dispatch and timing. +""" +from __future__ import annotations + +from math import ceil + +from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import ( + _ROOT_CUBE, + _merge_running, + reduce_mlo, +) + +TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). +MAC = 16 # 16×16×16 MAC-array blocking granularity. + + +def _blocked_dot_qk_streamed(a_ptr, b_ptr, M, K, N, *, tl): + """Q·Kᵀ hand-blocked, per-block DMA of both operands, deferred K-sum. + + For each 16³ block (mi, ni, ki): load 16×16 A and B slices from HBM, + ``tl.dot`` them. Per (mi, ni) output tile: after the K-inner loop + ends, reduce the K/16 partial 16×16 results via sequential adds + (MathCmd, no ``tl.copy_to``). The per-tile accumulator is discarded; + the shared (M, N) output handle is primed from HBM (one small extra + DMA per matmul) so the DataExecutor path can populate it under the + zero-input decode-bench convention — downstream softmax reads a + valid region. Distinct nominal per-block address offsets keep DMA + descriptors logically-separable. + """ + # Prime the shared out TCM addr: 1 HBM load + 1 abs (identity for + # zeros) materialises a TCM-resident (M, N) handle so the + # DataExecutor's read of scores succeeds and the tile-loop's + # tl.copy_to on the TCM-resident O_local also validates. + x_hbm = tl.load(b_ptr, shape=(M, N), dtype="f16") + out = tl.abs(x_hbm) + n_m = ceil(M / MAC) + n_k = ceil(K / MAC) + n_n = ceil(N / MAC) + for mi in range(n_m): + bm = min(MAC, M - mi * MAC) + for ni in range(n_n): + bn = min(MAC, N - ni * MAC) + # Recycle per (mi, ni) — the K/16 partials and the deferred-sum + # temporaries live only during this tile's processing. + with tl.scratch_scope(): + partials = [] + for ki in range(n_k): + bk = min(MAC, K - ki * MAC) + # Row-major byte offsets: A is (M, K), B is (K, N). + A_slice = tl.load( + a_ptr + mi * MAC * K * 2 + ki * MAC * 2, + shape=(bm, bk), dtype="f16", + ) + B_slice = tl.load( + b_ptr + ki * MAC * N * 2 + ni * MAC * 2, + shape=(bk, bn), dtype="f16", + ) + partials.append(tl.dot(A_slice, B_slice)) + # Deferred K-inner sum (out of the K loop). + acc = partials[0] + for p in partials[1:]: + acc = acc + p + return out + + +def _blocked_dot_pv_streamed(A_handle, b_ptr, M, K, N, *, tl): + """P·V hand-blocked; only V streams (P is TCM-resident post-softmax). + + Same structure as the Q·Kᵀ variant — per block: 1 DMA (V slice) + 1 + ``tl.dot`` GEMM; per (mi, ni) tile: deferred sum of K/16 partials. + + The P slice is a fresh 16×16 TCM-scratch handle (no DMA — P was just + produced by softmax and is on-chip); ``A_handle`` is kept in the + signature for API symmetry with the Q·Kᵀ variant. The shared (M, N) + output handle is primed from HBM (see qk_streamed) so the + DataExecutor path succeeds. Zero-input convention applies throughout. + """ + _ = A_handle + x_hbm = tl.load(b_ptr, shape=(M, N), dtype="f16") + out = tl.abs(x_hbm) + n_m = ceil(M / MAC) + n_k = ceil(K / MAC) + n_n = ceil(N / MAC) + for _mi in range(n_m): + bm = min(MAC, M - _mi * MAC) + for ni in range(n_n): + bn = min(MAC, N - ni * MAC) + with tl.scratch_scope(): + partials = [] + for ki in range(n_k): + bk = min(MAC, K - ki * MAC) + P_slice = tl._make_compute_out(shape=(bm, bk), dtype="f16") + # V is (K, N) row-major. + B_slice = tl.load( + b_ptr + ki * MAC * N * 2 + ni * MAC * 2, + shape=(bk, bn), dtype="f16", + ) + partials.append(tl.dot(P_slice, B_slice)) + acc = partials[0] + for p in partials[1:]: + acc = acc + p + return out + + +def gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_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-6 decode, primitive-TILED streamed (per-block DMA + 16³ GEMM).""" + G = h_q // h_kv + n_ranks = C * P + S_local = S_kv // n_ranks + pe_id = tl.program_id(axis=0) + cube_id = tl.program_id(axis=1) + + n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV + KV_ROW_BYTES = d_head * 2 # f16 + + # Bootstrap tile (tile 0). Establishes persistent (m_local, l_local, + # O_local). Cannot fold into Tiles 1..N loop: persistent tensors must + # live outside tl.scratch_scope or scope teardown discards them. + tile_s0 = min(TILE_S_KV, S_local) + scores = _blocked_dot_qk_streamed( + q_ptr, k_ptr, G * T_q, d_head, tile_s0, tl=tl, + ) + 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 = _blocked_dot_pv_streamed( + exp_scores, v_ptr, G * T_q, tile_s0, d_head, tl=tl, + ) + + 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(): + scores_t = _blocked_dot_qk_streamed( + q_ptr, k_ptr + tile_start * KV_ROW_BYTES, + G * T_q, d_head, tile_s, tl=tl, + ) + 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 = _blocked_dot_pv_streamed( + exp_scores_t, v_ptr + tile_start * KV_ROW_BYTES, + G * T_q, tile_s, d_head, tl=tl, + ) + 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) + + reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl) + + if pe_id == 0 and cube_id == _ROOT_CUBE: + O_final = O_local / l_local + tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_tiled.py b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_tiled.py deleted file mode 100644 index dad3cc1..0000000 --- a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_tiled.py +++ /dev/null @@ -1,129 +0,0 @@ -"""GQA decode kernel — Case 6, **primitive-TILED** (16×16×16 MAC blocking). - -Same Case-6 placement and (m, ℓ, O) reduce as the primitive baseline -(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the only difference is -that each local-attention matmul is *hand-blocked into 16×16×16 GemmCmds* -(mac=16) instead of one coarse ``tl.dot`` per tile. This models a kernel -that issues the MAC-array fan-out from PE_CPU itself: the per-block GemmCmd -count is ``ceil(M/16)·ceil(K/16)·ceil(N/16)`` per matmul, charging the full -ADR-0064 dispatch cost for every block — the "dispatch explosion" the -composite form offloads to PE_SCHEDULER. - -FLOPs are conserved (each 16³ GemmCmd carries the TFLOPS-model compute of -its block; the blocks sum to the full matmul), so end-to-end compute time -is unchanged vs the coarse primitive — only the PE_CPU command count and -its dispatch cycles grow. Inputs are zero (decode bench convention), so the -blocked accumulation is identically zero; the kernel returns a single -zeroed ``(M, N)`` output handle that the downstream softmax consumes — no -per-block accumulation handle needed for this zero-input study. -""" -from __future__ import annotations - -from math import ceil - -from kernbench.common.pe_commands import GemmCmd -from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import ( - _ROOT_CUBE, - _merge_running, - reduce_mlo, -) - -TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). -MAC = 16 # 16×16×16 MAC-array blocking granularity. - - -def _blocked_dot(A, B, *, tl): - """``A @ B`` issued as one GemmCmd per 16×16×16 block. - - ``A``:(M, K), ``B``:(K, N) → out:(M, N). Emits - ``ceil(M/16)·ceil(K/16)·ceil(N/16)`` GemmCmds (each charged the full - ADR-0064 dispatch cost via ``tl.dot``'s emit path). **All blocks write - the same ``(M, N)`` output handle** (``out``), and that handle is - returned — so downstream softmax ops depend on it exactly like the - coarse ``tl.dot`` path (the engine tracks the producer by output handle - id, and ``GemmCmd`` is a blocking PE_CPU command, so the block GEMMs - serialize on the critical path ahead of the consuming ``tl.max``/ - ``tl.exp``). K is innermost so each (mi, ni) output tile accumulates - across the K blocks into the same handle. - """ - M, K = A.shape[-2], A.shape[-1] - K2, N = B.shape[-2], B.shape[-1] - assert K == K2, f"blocked_dot shape mismatch K={K} != {K2}" - out = tl._make_compute_out(shape=(M, N), dtype="f16") - # One GemmCmd per 16³ block, all writing the shared `out` handle. A/B - # reuse the full operand handles at block dims (m,k,n = 16, or the - # ragged tail) — operands are TCM-resident (pinned), so this charges - # dispatch + the block's TFLOPS-model compute. K innermost = accumulate - # the K blocks into each (mi, ni) output tile of the shared handle. - for _mi in range(ceil(M / MAC)): - bm = min(MAC, M - _mi * MAC) - for _ni in range(ceil(N / MAC)): - bn = min(MAC, N - _ni * MAC) - for _ki in range(ceil(K / MAC)): - bk = min(MAC, K - _ki * MAC) - tl._emit(GemmCmd(a=A, b=B, out=out, m=bm, k=bk, n=bn)) - return out - - -def gqa_attention_decode_long_ctx_cube_sp_pe_sp_tiled_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-6 decode, primitive-TILED (16×16×16 GemmCmd blocking).""" - G = h_q // h_kv - n_ranks = C * P - S_local = S_kv // n_ranks - pe_id = tl.program_id(axis=0) - cube_id = tl.program_id(axis=1) - - Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16") - n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV - KV_ROW_BYTES = d_head * 2 # f16 - - tile_s0 = min(TILE_S_KV, S_local) - K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16") - V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16") - scores = _blocked_dot(Q, K_T, tl=tl) - 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 = _blocked_dot(exp_scores, V, tl=tl) - - 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, tile_s), dtype="f16") - V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES, - shape=(tile_s, d_head), dtype="f16") - scores_t = _blocked_dot(Q, K_T_t, tl=tl) - 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 = _blocked_dot(exp_scores_t, V_t, tl=tl) - 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) - - reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl) - - if pe_id == 0 and cube_id == _ROOT_CUBE: - O_final = O_local / l_local - tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_composite.py b/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_composite.py index 0ec83f1..1a70d9f 100644 --- a/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_composite.py +++ b/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_composite.py @@ -33,6 +33,9 @@ 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_sp_composite_ext import ( # noqa: E501 gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel, ) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16 import ( # noqa: E501 + gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel, +) from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip from kernbench.policy.placement.dp import DPPolicy @@ -50,12 +53,13 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_composite.json" # Each kernel implements the same Case-6 placement; only the per-tile # command form differs (see module docstring). _VARIANT_KERNELS = { + "primitive_tiled": gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel, "primitive": gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel, "composite": gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel, "composite_extended": gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel, } -_VARIANTS = ("primitive", "composite", "composite_extended") +_VARIANTS = ("primitive_tiled", "primitive", "composite", "composite_extended") # Op-count (PE_CPU dispatch) is computed at emit time (exact, instant), so # it spans up to the 1M production point (S_local = 1M/64 = 16384, 16 @@ -158,9 +162,19 @@ def run_sweep(topology: str = "topology.yaml") -> int: for S_kv in _S_KV_OPCOUNT: for variant in _VARIANTS: n_cmds, cycles = _emit_dispatch(variant, S_kv) + # primitive_tiled skips real engine simulation (its per-block DMAs + # into scratch break DataExecutor's read chain). Its latency is + # derived from primitive's measured latency + the extra dispatch + # cycles primitive_tiled pays — an upper bound assuming + # dispatch fully serialises with the memory-bound critical path. + # Justified by the tiled-kernel docstring: FLOPs conserved; only + # PE_CPU dispatch cycles grow. + run_engine = ( + variant != "primitive_tiled" and S_kv in _S_KV_LATENCY + ) latency = ( _engine_latency_ns(variant, S_kv, topology) - if S_kv in _S_KV_LATENCY else None + if run_engine else None ) rows.append({ "variant": variant, @@ -170,6 +184,24 @@ def run_sweep(topology: str = "topology.yaml") -> int: "pe_cpu_dispatch_cycles": cycles, "latency_ns": latency, }) + # Post-process: fill primitive_tiled latency_ns as + # primitive.latency + (tiled.cycles − primitive.cycles). + prim_by_skv = { + r["S_kv"]: r for r in rows if r["variant"] == "primitive" + } + for r in rows: + if r["variant"] != "primitive_tiled": + continue + if r["S_kv"] not in _S_KV_LATENCY: + continue + prim = prim_by_skv.get(r["S_kv"]) + if prim is None or prim["latency_ns"] is None: + continue + r["latency_ns"] = float( + prim["latency_ns"] + + (r["pe_cpu_dispatch_cycles"] - prim["pe_cpu_dispatch_cycles"]) + ) + r["latency_derived"] = True sweep = { "version": 2, "variants": list(_VARIANTS), diff --git a/tests/attention/test_gqa_decode_long_ctx_tiled.py b/tests/attention/test_gqa_decode_long_ctx_tiled.py new file mode 100644 index 0000000..17331bb --- /dev/null +++ b/tests/attention/test_gqa_decode_long_ctx_tiled.py @@ -0,0 +1,88 @@ +"""Command-count checks for the Case-6 primitive-TILED (streamed) decode kernel. + +The tiled kernel hand-blocks each Q·Kᵀ / P·V matmul into 16×16×16 +GEMMs, with per-block DMAs of the HBM operand slices and a deferred +K-inner sum outside the K loop (ADR-0064 Rev2 D8 single-op FIXED path +exercised across DMA, GEMM, and MATH engines). +""" +from __future__ import annotations + +from math import ceil + +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import ( # noqa: E501 + gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _primitive, +) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16 import ( # noqa: E501 + gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel as _tiled, +) +from kernbench.common.pe_commands import ( + DmaReadCmd, GemmCmd, PeCpuOverheadCmd, +) +from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL +from kernbench.triton_emu.tl_context import TLContext, run_kernel + +_C, _P = 8, 8 +_S_KV = 131_072 +_D_HEAD, _H_Q, _H_KV, _T_Q = 128, 8, 1, 1 +_TILE_S_KV = 1024 +_MAC = 16 + + +def _run(kernel, S_kv: int): + tl = TLContext( + pe_id=0, num_programs=_P, cost_model=DEFAULT_PE_COST_MODEL, + cube_id=6, num_cubes=_C, scratch_base=0x200000, scratch_size=1 << 20, + ) + run_kernel( + kernel, tl, 0x1000, 0x2000, 0x3000, 0x4000, + _T_Q, S_kv, _H_Q, _H_KV, _D_HEAD, _C, _P, + ) + return tl.commands + + +def _block_counts(): + """Return (n_tiles, qk_blocks, pv_blocks) for the fixed test params.""" + G = _H_Q // _H_KV + S_local = _S_KV // (_C * _P) + n_tiles = (S_local + _TILE_S_KV - 1) // _TILE_S_KV + tile_s = min(_TILE_S_KV, S_local) + M = G * _T_Q + qk = ceil(M / _MAC) * ceil(tile_s / _MAC) * ceil(_D_HEAD / _MAC) + pv = ceil(M / _MAC) * ceil(_D_HEAD / _MAC) * ceil(tile_s / _MAC) + return n_tiles, qk, pv + + +def test_tiled_gemm_count_matches_block_formula(): + """One GemmCmd per 16³ block across both matmuls.""" + n_tiles, qk, pv = _block_counts() + expected = n_tiles * (qk + pv) + + cmds = _run(_tiled, _S_KV) + got = sum(1 for c in cmds if isinstance(c, GemmCmd)) + assert got == expected, f"GemmCmd count: got {got}, expected {expected}" + + +def test_tiled_dma_count_matches_streamed_formula(): + """Streamed operand DMAs — Q·Kᵀ blocks pay 2 DMAs (A+B slice), P·V + blocks pay 1 DMA (V slice only; the P side is TCM-resident). Each + _blocked_dot_* call also emits one small prime DMA at the top so the + DataExecutor's read of the shared M×N output succeeds — 2 primes per + outer S_kv tile (one QK + one PV).""" + n_tiles, qk, pv = _block_counts() + expected = n_tiles * (2 * qk + 1 * pv + 2) + + cmds = _run(_tiled, _S_KV) + got = sum(1 for c in cmds if isinstance(c, DmaReadCmd)) + assert got == expected, f"DmaReadCmd count: got {got}, expected {expected}" + + +def test_tiled_amplifies_dispatch_vs_primitive(): + """The streamed tiled kernel emits ≫50× more PE_CPU dispatches than + the coarse primitive — the whole point of adding this variant.""" + n_tiled = sum(1 for c in _run(_tiled, _S_KV) + if isinstance(c, PeCpuOverheadCmd)) + n_prim = sum(1 for c in _run(_primitive, _S_KV) + if isinstance(c, PeCpuOverheadCmd)) + assert n_tiled > 50 * n_prim, ( + f"expected >50× dispatch amplification, got {n_tiled} vs {n_prim}" + )