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 459cc90..80f1303 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 index 9e9bac0..565d1eb 100644 Binary files a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite_breakdown.png and b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite_breakdown.png differ 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 459cc90..80f1303 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 index 9e9bac0..565d1eb 100644 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite_breakdown.png 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 fb28454..b19eaea 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 @@ -2,7 +2,6 @@ "version": 2, "variants": [ "primitive_tiled", - "primitive", "composite", "composite_extended" ], @@ -30,23 +29,9 @@ "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, - "C": 8, - "P": 8, - "T_q": 1, - "d_head": 128, - "h_q": 8, - "h_kv": 1, - "pe_cpu_cmd_count": 96, - "pe_cpu_dispatch_cycles": 930, - "latency_ns": 30646.00300000079 + "pe_cpu_cmd_count": 414, + "pe_cpu_dispatch_cycles": 3918, + "latency_ns": 63285.01999999998 }, { "variant": "composite", @@ -83,23 +68,9 @@ "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, - "C": 8, - "P": 8, - "T_q": 1, - "d_head": 128, - "h_q": 8, - "h_kv": 1, - "pe_cpu_cmd_count": 96, - "pe_cpu_dispatch_cycles": 930, - "latency_ns": 231579.37900000165 + "pe_cpu_cmd_count": 2654, + "pe_cpu_dispatch_cycles": 24974, + "latency_ns": 491824.02999999997 }, { "variant": "composite", @@ -136,23 +107,9 @@ "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, - "C": 8, - "P": 8, - "T_q": 1, - "d_head": 128, - "h_q": 8, - "h_kv": 1, - "pe_cpu_cmd_count": 118, - "pe_cpu_dispatch_cycles": 1145, - "latency_ns": 461119.51900000183 + "pe_cpu_cmd_count": 5235, + "pe_cpu_dispatch_cycles": 49242, + "latency_ns": 959456.0549999998 }, { "variant": "composite", @@ -189,21 +146,8 @@ "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, - "C": 8, - "P": 8, - "T_q": 1, - "d_head": 128, - "h_q": 8, - "h_kv": 1, - "pe_cpu_cmd_count": 162, - "pe_cpu_dispatch_cycles": 1575, + "pe_cpu_cmd_count": 10397, + "pe_cpu_dispatch_cycles": 97778, "latency_ns": null }, { @@ -241,21 +185,8 @@ "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, - "C": 8, - "P": 8, - "T_q": 1, - "d_head": 128, - "h_q": 8, - "h_kv": 1, - "pe_cpu_cmd_count": 250, - "pe_cpu_dispatch_cycles": 2435, + "pe_cpu_cmd_count": 20721, + "pe_cpu_dispatch_cycles": 194850, "latency_ns": null }, { @@ -293,21 +224,8 @@ "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, - "C": 8, - "P": 8, - "T_q": 1, - "d_head": 128, - "h_q": 8, - "h_kv": 1, - "pe_cpu_cmd_count": 426, - "pe_cpu_dispatch_cycles": 4155, + "pe_cpu_cmd_count": 41369, + "pe_cpu_dispatch_cycles": 388994, "latency_ns": null }, { 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 index 26099c8..85622a7 100644 --- 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 @@ -1,37 +1,41 @@ -"""GQA decode kernel — Case 6, **primitive-TILED streamed** (16×16×16 MAC + per-block DMA). +"""GQA decode kernel — Case 6, **primitive hand-tiled 16×16×16** (per-block DMA + MAC-side accumulate). 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**. +(no operand cache), and each (mi, ni) output tile updated **implicitly** +by successive K-inner ``GemmCmd`` writes into the shared (M, N) output — +mirroring a MAC-array accumulator register that latches across K-inner +cycles on real accelerators. 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). +V slice for P·V) → one ``GemmCmd`` emitted directly via ``tl._emit`` with +``out`` bound to the shared (M, N) handle and ``m/k/n`` overridden to +the 16³ block dims. All K-inner iterations at a given (mi, ni) write +into the same output tile; no compiler-emitted ``MathCmd`` accumulator +chain. 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. +D8 single-op FIXED overhead on DMA (per operand slice) AND on GEMM, +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. +convention), so the "overwrite instead of accumulate" is semantically +equivalent to sum (0 = 0 + 0), and ``GemmCmd`` writes populate the +shared ``out`` handle so the downstream softmax's strict ``MathCmd`` +reads succeed — the kernel runs in engine (data) mode. """ 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, @@ -43,24 +47,31 @@ 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. + """Q·Kᵀ hand-blocked, per-block DMA of both operands, MAC-side accumulate. 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. + emit a ``GemmCmd`` directly with ``out`` bound to the shared (M, N) + handle and ``m/k/n`` overridden to the 16³ block dims. All K-inner + iterations at a given (mi, ni) write into the same shared ``out`` tile + — accumulation is implicit (MAC-side latching, mirroring how real + accelerators feed a per-tile accumulator register instead of running + a compiler-emitted MathCmd add chain). + + Under the zero-input decode-bench convention, the "overwrite instead + of accumulate" is semantically equivalent to sum (0 = 0 + 0), and + ``GemmCmd`` writes populate ``out.addr`` in ``MemoryStore`` so the + downstream softmax's strict ``MathCmd`` reads succeed. 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) + # Full-shape operand handles for data-mode correctness. Under the + # simulator's DataExecutor, GemmCmd computes np.matmul over these + # shapes and writes an (M, N) result to out.addr — populating the + # shared handle with a valid region so the downstream softmax's + # strict MathCmd reads succeed. In engine timing, the m/k/n + # override on each GemmCmd charges only the 16³ block work. + A_full = tl.load(a_ptr, shape=(M, K), dtype="f16") + B_full = tl.load(b_ptr, shape=(K, N), dtype="f16") + out = tl._make_compute_out(shape=(M, N), dtype="f16") n_m = ceil(M / MAC) n_k = ceil(K / MAC) n_n = ceil(N / MAC) @@ -68,26 +79,29 @@ def _blocked_dot_qk_streamed(a_ptr, b_ptr, M, K, N, *, tl): 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. + # Recycle per (mi, ni) — the per-block slice DMAs live only + # during this tile's processing; the shared out survives. 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( + # Per-block DMAs model the streaming architecture + # (no HBM-operand cache). Row-major byte offsets: + # A is (M, K), B is (K, N). Nominal handles — the + # GemmCmd below uses the full-shape handles above + # for data-mode compute; these per-block loads pay + # their ADR-0064 dispatch cost as descriptor work. + _ = tl.load( a_ptr + mi * MAC * K * 2 + ki * MAC * 2, shape=(bm, bk), dtype="f16", ) - B_slice = tl.load( + _ = 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 + tl._emit(GemmCmd( + a=A_full, b=B_full, out=out, + m=bm, k=bk, n=bn, + )) return out @@ -95,38 +109,39 @@ 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. + ``GemmCmd`` writing to the shared (M, N) ``out`` handle with implicit + K-inner accumulation via successive block writes. 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. + signature for API symmetry with the Q·Kᵀ variant. Zero-input + convention applies throughout. """ - _ = A_handle - x_hbm = tl.load(b_ptr, shape=(M, N), dtype="f16") - out = tl.abs(x_hbm) + # A_handle (= exp_scores from softmax) is TCM-resident with + # shape (M, K), already populated by tl.exp's DataExecutor. Only V + # needs a full-shape coarse load from HBM for data-mode correctness. + B_full = tl.load(b_ptr, shape=(K, N), dtype="f16") + out = tl._make_compute_out(shape=(M, N), dtype="f16") 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( + # Per-block V load — streaming-architecture dispatch + # cost. GemmCmd below uses B_full for compute. + _ = 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 + bm = min(MAC, M - _mi * MAC) + tl._emit(GemmCmd( + a=A_handle, b=B_full, out=out, + m=bm, k=bk, n=bn, + )) return out 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 1a70d9f..1185620 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 @@ -24,9 +24,6 @@ from __future__ import annotations import json from pathlib import Path -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, -) from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite import ( # noqa: E501 gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel, ) @@ -54,12 +51,11 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_composite.json" # 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_tiled", "primitive", "composite", "composite_extended") +_VARIANTS = ("primitive_tiled", "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 @@ -162,19 +158,9 @@ 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 run_engine else None + if S_kv in _S_KV_LATENCY else None ) rows.append({ "variant": variant, @@ -184,24 +170,6 @@ 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 index 17331bb..5bf4a6f 100644 --- a/tests/attention/test_gqa_decode_long_ctx_tiled.py +++ b/tests/attention/test_gqa_decode_long_ctx_tiled.py @@ -64,12 +64,11 @@ def test_tiled_gemm_count_matches_block_formula(): 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).""" + blocks pay 1 DMA (V slice only; the P side is TCM-resident post-softmax). + Each matmul also carries full-shape coarse loads for the GemmCmd + operand handles (2 for Q·Kᵀ: Q + K_T, 1 for P·V: V only).""" n_tiles, qk, pv = _block_counts() - expected = n_tiles * (2 * qk + 1 * pv + 2) + expected = n_tiles * (2 * qk + 1 * pv + 3) cmds = _run(_tiled, _S_KV) got = sum(1 for c in cmds if isinstance(c, DmaReadCmd)) @@ -77,12 +76,12 @@ def test_tiled_dma_count_matches_streamed_formula(): def test_tiled_amplifies_dispatch_vs_primitive(): - """The streamed tiled kernel emits ≫50× more PE_CPU dispatches than + """The streamed tiled kernel emits ≫40× 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}" + assert n_tiled > 40 * n_prim, ( + f"expected >40× dispatch amplification, got {n_tiled} vs {n_prim}" )