ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

paper(gqa): two-regime "use of composite commands" — decode + compute-bound prefill

Browse Source 
Complete the composite-command study across both attention regimes and
write up the result, resolving when the composite command helps latency.

Decode (memory-bound, T_q=1, M=8): the command form is latency-neutral —
the kernel is bound by KV streaming and the MAC array is near-idle, so the
composite's only benefit here is host-issue offload (O(n_tiles) -> O(1)
PE_CPU commands). Figure: gqa_decode_long_ctx_composite (latency flat,
command count saturates).

Prefill (compute-bound, large M=G*T_q tile-filling): add three command-form
variants of a single-rank FlashAttention prefill kernel
(_gqa_prefill_compute_bound: primitive / composite / composite_extended).
Here the composite's scheduler-internal per-tile DMA<->compute pipeline
keeps the MAC array fed while the primitive's blocking tl.dot starves it,
so composite wins on MAC utilization (68% flat -> 83%) and wall-clock
(19% faster at ctx=1024), with the margin growing with context (deeper
P.V reduction = more tiles to pipeline) — the compute-bound mirror of the
GEMM result in section 3. Figure: gqa_prefill_compute_bound (latency +
MAC util). Sweep wired as GQA_1H_SWEEPS=prefill_cb; tests cover structure,
e2e, and composite<primitive in the compute-bound regime.

The synthesis: composite has two benefits set by roofline position —
host-issue offload (always) and MAC-array feeding (compute-bound only).
Decode exercises the first, prefill both.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook Kang
2026-06-23 17:04:26 -07:00
parent cd6f0ed91d
commit 24c705419e
13 changed files with 970 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_composite.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 136 KiB

src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_compute_bound.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 145 KiB

src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json
+258
View File
@@ -0,0 +1,258 @@
{
"version": 2,
"variants": [
"primitive",
"composite",
"composite_extended"
],
"s_kv_opcount": [
8192,
65536,
131072,
262144,
524288,
1048576
],
"s_kv_latency": [
8192,
32768,
65536,
131072
],
"rows": [
{
"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
},
{
"variant": "composite",
"S_kv": 8192,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 94,
"pe_cpu_dispatch_cycles": 978,
"latency_ns": 30566.1840000007
},
{
"variant": "composite_extended",
"S_kv": 8192,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 98,
"pe_cpu_dispatch_cycles": 1032,
"latency_ns": 30483.359500000362
},
{
"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
},
{
"variant": "composite",
"S_kv": 65536,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 94,
"pe_cpu_dispatch_cycles": 978,
"latency_ns": 231270.18400000152
},
{
"variant": "composite_extended",
"S_kv": 65536,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 98,
"pe_cpu_dispatch_cycles": 1032,
"latency_ns": 231211.1740000015
},
{
"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
},
{
"variant": "composite",
"S_kv": 131072,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 94,
"pe_cpu_dispatch_cycles": 978,
"latency_ns": 460651.3170000027
},
{
"variant": "composite_extended",
"S_kv": 131072,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 98,
"pe_cpu_dispatch_cycles": 1032,
"latency_ns": 460552.9805000025
},
{
"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,
"latency_ns": null
},
{
"variant": "composite",
"S_kv": 262144,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 94,
"pe_cpu_dispatch_cycles": 978,
"latency_ns": null
},
{
"variant": "composite_extended",
"S_kv": 262144,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 98,
"pe_cpu_dispatch_cycles": 1032,
"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,
"latency_ns": null
},
{
"variant": "composite",
"S_kv": 524288,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 94,
"pe_cpu_dispatch_cycles": 978,
"latency_ns": null
},
{
"variant": "composite_extended",
"S_kv": 524288,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 98,
"pe_cpu_dispatch_cycles": 1032,
"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,
"latency_ns": null
},
{
"variant": "composite",
"S_kv": 1048576,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 94,
"pe_cpu_dispatch_cycles": 978,
"latency_ns": null
},
{
"variant": "composite_extended",
"S_kv": 1048576,
"C": 8,
"P": 8,
"T_q": 1,
"d_head": 128,
"h_q": 8,
"h_kv": 1,
"pe_cpu_cmd_count": 98,
"pe_cpu_dispatch_cycles": 1032,
"latency_ns": null
}
]
}
src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_prefill_compute_bound.json
+105
View File
@@ -0,0 +1,105 @@
{
"version": 1,
"variants": [
"primitive",
"composite",
"composite_extended"
],
"ctx_points": [
256,
512,
1024
],
"rows": [
{
"variant": "primitive",
"ctx_len": 256,
"M": 2048,
"latency_ns": 48857.361999999725,
"gemm_busy_ns": 33554.43200000003,
"dma_busy_ns": 20880.000000000004,
"achieved_tflops": 5.4942683151825005,
"mac_util": 0.6867835393978126
},
{
"variant": "composite",
"ctx_len": 256,
"M": 2048,
"latency_ns": 49887.505999999594,
"gemm_busy_ns": 34531.32799999314,
"dma_busy_ns": 1095773.440000072,
"achieved_tflops": 5.380815308746887,
"mac_util": 0.6726019135933609
},
{
"variant": "composite_extended",
"ctx_len": 256,
"M": 2048,
"latency_ns": 50280.91399999841,
"gemm_busy_ns": 129096.19199996963,
"dma_busy_ns": 633563.5200000411,
"achieved_tflops": 5.338714725830331,
"mac_util": 0.6673393407287914
},
{
"variant": "primitive",
"ctx_len": 512,
"M": 4096,
"latency_ns": 197496.81800000605,
"gemm_busy_ns": 134217.72800000012,
"dma_busy_ns": 66880.0,
"achieved_tflops": 5.436755056985105,
"mac_util": 0.6795943821231382
},
{
"variant": "composite",
"ctx_len": 512,
"M": 4096,
"latency_ns": 177341.23400000148,
"gemm_busy_ns": 141168.63999995415,
"dma_busy_ns": 8567063.039998509,
"achieved_tflops": 6.054665346469796,
"mac_util": 0.7568331683087245
},
{
"variant": "composite_extended",
"ctx_len": 512,
"M": 4096,
"latency_ns": 174702.7699999963,
"gemm_busy_ns": 1076198.3999996716,
"dma_busy_ns": 6594394.87999886,
"achieved_tflops": 6.146106464139193,
"mac_util": 0.7682633080173992
},
{
"variant": "primitive",
"ctx_len": 1024,
"M": 8192,
"latency_ns": 794090.4820000405,
"gemm_busy_ns": 536870.9120000004,
"dma_busy_ns": 234240.0,
"achieved_tflops": 5.4086623544239645,
"mac_util": 0.6760827943029956
},
{
"variant": "composite",
"ctx_len": 1024,
"M": 8192,
"latency_ns": 668110.7059999453,
"gemm_busy_ns": 873011.1999923651,
"dma_busy_ns": 67794150.3999821,
"achieved_tflops": 6.428526376585786,
"mac_util": 0.8035657970732233
},
{
"variant": "composite_extended",
"ctx_len": 1024,
"M": 8192,
"latency_ns": 646937.2019999702,
"gemm_busy_ns": 8870907.903995967,
"dma_busy_ns": 59655820.79998447,
"achieved_tflops": 6.638924586068552,
"mac_util": 0.829865573258569
}
]
}
src/kernbench/benches/gqa_helpers/long_ctx/_gqa_prefill_compute_bound.py
+161
View File
@@ -0,0 +1,161 @@
"""Compute-bound prefill attention — 3 command-form variants (single-rank).
Companion to the memory-bound decode study
(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp*``). Prefill processes a
block of T_q query positions at once, so the score / context GEMMs have a
large M = G·T_q and high arithmetic intensity (~M flops/byte) — the
workload is **compute-bound** (above the roofline ridge), unlike T_q=1
decode (M=8, memory-bound). This is the regime where the composite
command's value shows: it streams DMA↔compute per HW tile to keep the MAC
array fed, while the primitive kernel serializes load→dot and starves it.
Single-rank (C=P=1): no cross-device reduce — the focus is the per-PE
GEMM-issue mechanism. FlashAttention 2-D tiling (Q-block × S_kv-tile,
online softmax) bounds the TCM scratch.
Three forms, differing only in how each Q-block's local attention is
issued (placement/softmax identical):
primitive tl.dot per S_kv-tile + primitive online merge.
composite one coarse composite GEMM per Q-block over the whole
S_kv (K, V as HBM refs → PE_SCHEDULER tiles/streams).
composite_extended Q·Kᵀ composite + softmax_merge recipe composite.
"""
from __future__ import annotations
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import _merge_running
M_BLOCK = 128 # query rows per FlashAttention Q-block (tile-filling).
TILE_S_KV = 256 # primitive per-tile S_kv width.
_SEED_S_KV = 64 # composite_extended recipe seed slice.
def _qblock_bounds(T_q: int, h_q: int, h_kv: int):
G = h_q // h_kv
M_total = G * T_q
n_qblocks = (M_total + M_BLOCK - 1) // M_BLOCK
return G, M_total, n_qblocks
# ── primitive: hand-tiled monolithic dots + online merge ─────────────
def gqa_prefill_primitive_kernel(
q_ptr, k_ptr, v_ptr, o_ptr,
T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
) -> None:
G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
ROW = d_head * 2 # f16
n_tiles = (S_kv + TILE_S_KV - 1) // TILE_S_KV
for qb in range(n_qblocks):
# Per-Q-block scratch_scope: this block's running (m,ℓ,O) plus its
# tile-0 transients are freed before the next block, so scratch
# stays O(one block) rather than accumulating across all blocks.
with tl.scratch_scope():
m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
q_off = qb * M_BLOCK * ROW
Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
tile_s0 = min(TILE_S_KV, S_kv)
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 = tl.dot(Q, K_T)
m_local = tl.max(scores, axis=-1)
exp_s = tl.exp(scores - m_local)
l_local = tl.sum(exp_s, axis=-1)
O_local = tl.dot(exp_s, V)
for ti in range(1, n_tiles):
start = ti * TILE_S_KV
tile_s = min(TILE_S_KV, S_kv - start)
with tl.scratch_scope():
K_T_t = tl.load(k_ptr + start * ROW,
shape=(d_head, tile_s), dtype="f16")
V_t = tl.load(v_ptr + start * ROW,
shape=(tile_s, d_head), dtype="f16")
scores_t = tl.dot(Q, K_T_t)
m_tile = tl.max(scores_t, axis=-1)
exp_t = tl.exp(scores_t - m_tile)
l_tile = tl.sum(exp_t, axis=-1)
O_tile = tl.dot(exp_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
tl.store(o_ptr + q_off, O_final)
# ── composite: one coarse composite GEMM per Q-block ─────────────────
def gqa_prefill_composite_kernel(
q_ptr, k_ptr, v_ptr, o_ptr,
T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
) -> None:
G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
ROW = d_head * 2
for qb in range(n_qblocks):
with tl.scratch_scope():
m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
q_off = qb * M_BLOCK * ROW
Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
K_T = tl.ref(k_ptr, shape=(d_head, S_kv), dtype="f16")
V = tl.ref(v_ptr, shape=(S_kv, d_head), dtype="f16")
scores = tl.composite(op="gemm", a=Q, b=K_T) # Q·Kᵀ, one command
m_local = tl.max(scores, axis=-1)
exp_s = tl.exp(scores - m_local)
l_local = tl.sum(exp_s, axis=-1)
O_local = tl.zeros((m_blk, d_head), dtype="f16")
tl.composite(op="gemm", a=exp_s, b=V, out=O_local) # P·V, one command
O_final = O_local / l_local
tl.store(o_ptr + q_off, O_final)
# ── composite_extended: Q·Kᵀ composite + softmax_merge recipe ────────
def gqa_prefill_composite_ext_kernel(
q_ptr, k_ptr, v_ptr, o_ptr,
T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
) -> None:
G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
ROW = d_head * 2
for qb in range(n_qblocks):
with tl.scratch_scope():
m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
q_off = qb * M_BLOCK * ROW
Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
seed = min(_SEED_S_KV, S_kv)
K_T0 = tl.load(k_ptr, shape=(d_head, seed), dtype="f16")
V0 = tl.load(v_ptr, shape=(seed, d_head), dtype="f16")
scores0 = tl.dot(Q, K_T0)
m_local = tl.max(scores0, axis=-1)
exp0 = tl.exp(scores0 - m_local)
l_local = tl.sum(exp0, axis=-1)
O_local = tl.dot(exp0, V0)
rest = S_kv - seed
if rest > 0:
K_T1 = tl.ref(k_ptr + seed * ROW,
shape=(d_head, rest), dtype="f16")
V1 = tl.ref(v_ptr + seed * ROW,
shape=(rest, d_head), dtype="f16")
scores1 = tl.composite(op="gemm", a=Q, b=K_T1)
tl.composite(
prologue=[{"op": "softmax_merge", "s": scores1,
"m": m_local, "l": l_local, "O": O_local}],
op="gemm", b=V1, out=O_local,
epilogue=[{"op": "add", "other": O_local}],
)
O_final = O_local / l_local
tl.store(o_ptr + q_off, O_final)
src/kernbench/benches/gqa_helpers/long_ctx/gqa_prefill_compute_bound.py
+142
View File
@@ -0,0 +1,142 @@
"""milestone-1h-gqa: compute-bound prefill composite-command study.
Three command-form variants of a single-rank compute-bound prefill
attention kernel (``_gqa_prefill_compute_bound``), swept over context
length S_kv = T_q. Unlike the memory-bound decode study, prefill has a
large M = G·T_q, so the score / context GEMMs are compute-bound — the
regime where the composite command keeps the MAC array fed (DMA↔compute
pipelining) and the hand-tiled primitive starves it on the serial
load→dot path.
Records per (variant, context) the end-to-end latency, the GEMM-engine
busy time, and the MAC occupancy (gemm_busy / e2e) so the comparative
plot can show composite winning on both latency and utilization.
Runs in data mode (engine latency). Gated via the umbrella
``GQA_1H_SWEEPS=prefill_cb``.
"""
from __future__ import annotations
import json
from pathlib import Path
from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
gqa_prefill_composite_ext_kernel,
gqa_prefill_composite_kernel,
gqa_prefill_primitive_kernel,
)
from kernbench.policy.placement.dp import DPPolicy
_OUTPUT_DIR = (
Path(__file__).resolve().parents[2]
/ "1H_milestone_output" / "gqa" / "long_ctx"
)
_SWEEP_JSON = _OUTPUT_DIR / "sweep_prefill_compute_bound.json"
_VARIANT_KERNELS = {
"primitive": gqa_prefill_primitive_kernel,
"composite": gqa_prefill_composite_kernel,
"composite_extended": gqa_prefill_composite_ext_kernel,
}
_VARIANTS = ("primitive", "composite", "composite_extended")
# Context length S_kv = T_q (prefill processes T_q tokens against S_kv=T_q
# keys). M = G·T_q = 8·T_q is tile-filling/compute-bound at every point.
_CTX_POINTS = (256, 512, 1024)
_H_Q, _H_KV, _D_HEAD = 8, 1, 128
_PEAK_TFLOPS = 8.0 # per-PE f16 GEMM peak (topology.yaml pe_gemm.peak_tflops_f16)
def _run_panel_fn(variant: str, ctx_len: int):
kernel = _VARIANT_KERNELS[variant]
panel = f"prefill_cb_{variant}_c{ctx_len}"
def _bench_fn(ctx):
dp = DPPolicy(cube="replicate", pe="replicate",
num_cubes=1, num_pes=1)
q = ctx.zeros((ctx_len, _H_Q * _D_HEAD),
dtype="f16", dp=dp, name=f"{panel}_q")
k = ctx.zeros((ctx_len, _H_KV * _D_HEAD),
dtype="f16", dp=dp, name=f"{panel}_k")
v = ctx.zeros((ctx_len, _H_KV * _D_HEAD),
dtype="f16", dp=dp, name=f"{panel}_v")
o = ctx.empty((ctx_len, _H_Q * _D_HEAD),
dtype="f16", dp=dp, name=f"{panel}_o")
ctx.launch(panel, kernel, q, k, v, o,
ctx_len, ctx_len, _H_Q, _H_KV, _D_HEAD, 1, 1,
_auto_dim_remap=False)
return _bench_fn
def _end_to_end_ns(op_log) -> float:
if not op_log:
return 0.0
return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
def _engine_busy_ns(op_log, suffix: str) -> float:
return sum(r.t_end - r.t_start
for r in op_log if r.component_id.endswith("." + suffix))
def _run_panel(variant: str, ctx_len: int, topology: str) -> dict:
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology
topo = resolve_topology(topology)
result = run_bench(
topology=topo, bench_fn=_run_panel_fn(variant, ctx_len),
device=resolve_device(None),
engine_factory=lambda t, d: GraphEngine(
getattr(t, "topology_obj", t), enable_data=True,
),
)
if not result.completion.ok:
raise RuntimeError(
f"gqa-prefill-cb {variant}@{ctx_len} failed: {result.completion}"
)
op_log = result.engine.op_log
e2e = _end_to_end_ns(op_log)
gemm = _engine_busy_ns(op_log, "pe_gemm")
dma = _engine_busy_ns(op_log, "pe_dma")
# Useful attention flops (Q·Kᵀ + P·V), single rank.
G = _H_Q // _H_KV
M = G * ctx_len
useful_flops = 4.0 * M * _D_HEAD * ctx_len
# MAC utilization = achieved / peak. ``achieved_tflops`` uses wall-clock
# (useful_flops / e2e) so it is bounded by peak even when the composite
# path overlaps many tile GEMMs (gemm_busy is a sum over overlapping ops
# and is kept only for diagnostics).
return {
"variant": variant,
"ctx_len": ctx_len,
"M": M,
"latency_ns": e2e,
"gemm_busy_ns": gemm,
"dma_busy_ns": dma,
"achieved_tflops": (useful_flops / e2e / 1e3) if e2e > 0 else 0.0,
"mac_util": (useful_flops / e2e / 1e3 / _PEAK_TFLOPS) if e2e > 0 else 0.0,
}
def run_sweep(topology: str = "topology.yaml") -> int:
"""Drive all (variant, context) prefill panels; write sweep.json."""
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
rows = [
_run_panel(variant, ctx_len, topology)
for ctx_len in _CTX_POINTS
for variant in _VARIANTS
]
sweep = {
"version": 1,
"variants": list(_VARIANTS),
"ctx_points": list(_CTX_POINTS),
"rows": rows,
}
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
print(f" gqa-prefill-cb: {len(rows)} rows -> {_SWEEP_JSON}")
return len(rows)
src/kernbench/benches/milestone_1h_gqa.py
+4
View File
@@ -28,6 +28,9 @@ from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (
run_sweep as _run_composite_sweep,
)
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
run_sweep as _run_prefill_cb_sweep,
)
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases import (
run_sweep as _run_prefill_sweep,
)
@@ -62,6 +65,7 @@ def run(torch) -> None:
"prefill": _run_prefill_sweep,
"decode": _run_decode_sweep,
"composite": _run_composite_sweep,
"prefill_cb": _run_prefill_cb_sweep,
}
unknown = [s for s in sweeps if s not in runners]
if unknown: