kernbench2/docs/adr-proposed/ADR-0061-prog-tl-broadcast.md

ADR-0061: tl.broadcast / tl.repeat — data-faithful GQA head reuse

Status

Proposed

Supporting ADR for ADR-0060 (AHBM GQA Fused Attention).

⚠ Re-evaluated during design iteration — this is now OPTIONAL, not a GQA blocker. The original framing (below) held that GQA h_q > h_kv is blocked by the MemoryStore byte-conservation check on a broadcast view. That is true only of the baseline's head-packing hack (_view(K, (h_q·d, S_kv))). The correct GQA design (ADR-0060 §5.2) processes one KV head at a time and folds the G group rows into the matmul M dimension — using only byte-conserving reshapes — so it needs no broadcast op at all, and np.matmul already broadcasts a shared operand in Phase-2 data mode. This ADR therefore drops to a convenience primitive: it cleans up additive-mask construction across the G·T_q rows and serves general kernels. Lowest priority of the three supporting ADRs. The rest of this document is retained as the design for that convenience op.

Context

The GQA reuse lever

Grouped-Query Attention amortises K/V loads: G = H_q / H_kv query heads share one KV head. For Llama3-70B G = 8. In a kernel this means a KV tile of shape [S_kv, d] (for one KV head) must be matmul'd against the G query rows of its group. Loading K/V once and reusing it across the group is the decode-time efficiency win (ADR-0060 §5.2).

Why it does not work today

The existing mesh kernels reshape tensors with a metadata-only helper _view (src/kernbench/benches/_attention_mesh_kv.py:25-36, _attention_mesh_mlo.py:25-36):

def _view(handle, new_shape):
    return TensorHandle(..., shape=new_shape, nbytes=handle.nbytes, ...)
    #                          ^^^ new shape          ^^^ OLD nbytes (unchanged)

This rewrites shape but keeps the original nbytes. It is purely symbolic. It survives Phase 1 (timing) because Phase 1 never reads tensor bytes. It fails Phase 2 (enable_data=True) because the data executor materialises every op through MemoryStore.read, which enforces byte conservation (src/kernbench/sim_engine/memory_store.py:67-73):

if shape is not None and arr.shape != shape:
    if arr.nbytes != np.prod(shape) * arr.dtype.itemsize:
        raise ValueError(f"Shape mismatch: stored {arr.shape} ({arr.nbytes}B) "
                         f"vs requested {shape} ({...}B)")
    arr = arr.reshape(shape)

A reshape conserves bytes; a broadcast does not ([S_kv, 1, d] → [S_kv, G, d] multiplies the element count by G). So the moment a kernel tries to view a h_kv-headed K as if it had h_q heads, the stored array has 1/G of the requested bytes and read raises.

The current test suite documents this as a deliberate limitation (tests/attention/test_milestone_gqa_llama70b.py:137-142):

"v1 uses h_q == h_kv == 1 to avoid … GQA broadcast view (which is symbolic and does not survive MemoryStore's nbytes check under simulator data execution). Real GQA (h_q > h_kv) is deferred."

There is no broadcast / repeat / expand primitive in the tl API today (src/kernbench/triton_emu/tl_context.py exposes trans, zeros, full, arange as metadata-only ops; none replicate data).

Decision

Add a data-faithful broadcast op to the tl API and the data executor. Unlike _view, it updates nbytes and, in Phase 2, actually replicates the underlying ndarray. The latency it charges models the on-chip register/SRAM fan-out, not an HBM re-load (the KV tile is read from HBM once; broadcast is an on-chip replication).

D1. tl surface

def broadcast(self, x: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle:
    """Replicate x along size-1 axes to new_shape (numpy broadcast rules).

    new_shape must be broadcast-compatible with x.shape: every axis is
    either equal to x's, or x's is 1 (the axis being expanded). nbytes of
    the result reflect new_shape (unlike the metadata-only _view).
    """

tl.repeat(x, axis, n) is offered as a thin convenience wrapper over broadcast for the common "insert a group axis and expand it to G" pattern, but broadcast is the primitive.

D2. Command + op_log

Emit a new BroadcastCmd (sibling of MathCmd):

@dataclass(frozen=True)
class BroadcastCmd:
    src: TensorHandle
    out: TensorHandle          # out.nbytes reflects new_shape
    new_shape: tuple[int, ...]
    data_op: bool = True

op_log entry: op_kind="math", op_name="broadcast" (it runs on the MATH/vector engine, like the other element-wise ops). The output handle is allocated from PE scratch via the existing _make_compute_out (tl_context.py:148-156), so out.nbytes is computed from new_shape and the byte-conservation check passes in Phase 2.

D3. Phase 2 data execution

Add _execute_broadcast to the data executor (src/kernbench/sim_engine/data_executor.py):

def _execute_broadcast(self, op):
    p = op.params
    src = self._resolve_read(p["src_space"], p["src_addr"],
                             p["src_shape"], p["dtype"], op.t_start)
    result = np.broadcast_to(src, p["new_shape"]).copy()   # .copy() → writable, owns bytes
    self.store.write(p["dst_space"], p["dst_addr"], result)

.copy() is required: np.broadcast_to returns a zero-stride view, and MemoryStore stores references — a later consumer's nbytes check must see a real prod(new_shape)-element array.

D4. Latency model

Broadcast latency = vector-engine cost of writing prod(new_shape) elements (reuse pe_math._compute_ns(num_elements), src/kernbench/components/builtin/pe_math.py:46-51). This models the on-chip replication into the register file / SRAM that real GQA hardware performs. It is not an HBM transaction — no dma_read is logged, consistent with "K/V tile loaded once" (ADR-0060 §5.2).

Alternatives

A1. Indexed per-head dot (no broadcast op)

Loop the G query heads, each doing its own tl.dot(Q_g, Kᵀ) against the same (re-ref'd, not re-loaded) KV tile:

for g in range(G):
    s_g = tl.dot(Q[g], K_T)     # G separate GEMMs, KV reused via tl.ref

• Pro: no new primitive; works in Phase 2 today.
• Con: G× more GEMM commands (op_log inflation), no batched-GEMM modelling, and the per-call overhead is charged G times. It also obscures the "single batched matmul over the group" the hardware does. Rejected as the primary path; acceptable as a fallback if the batched 2D dot's shape handling proves awkward.

A2. Keep _view, relax the nbytes check

Make MemoryStore.read tolerate broadcast (stride-0) shapes. Rejected: it silently breaks byte conservation for all readers and would mask genuine shape bugs elsewhere. The check is a correctness invariant (ADR-0052), not an obstacle to route around.

A3. Materialise the broadcast at deploy time

Store K/V already expanded to h_q heads in HBM. Rejected: defeats the entire point of GQA (it G×'s KV-cache HBM footprint and bandwidth).

Consequences

Positive

• Unblocks real GQA (h_q > h_kv) in Phase 2 — the headline blocker for ADR-0060 and the deferred milestone work.
• A general, reusable primitive (any kernel needing numpy-style broadcast benefits, not just attention).
• Byte conservation stays intact; no weakening of MemoryStore.

Negative

• One new command type + one data-executor handler + one tl method.
• Phase 2 .copy() allocates the expanded array in the store; for very large groups this is real memory. Bounded by tile size in the tiled design (ADR-0060), so acceptable.

Test Requirements

1. Phase 1 timing: tl.broadcast([S,1,d] → [S,G,d]) logs one math/broadcast op with num_elements == S*G*d; no dma_read.
2. Phase 2 data: broadcasting a known array yields the numpy broadcast_to(...).copy() result; a downstream tl.dot consuming it passes the MemoryStore nbytes check (the exact case that fails today).
3. GQA end-to-end: re-run a reduced milestone-gqa-llama70b panel with h_q=8, h_kv=1 under enable_data=True; it must complete (today it raises ValueError: Shape mismatch).
4. Broadcast-incompatible shape (e.g. axis size 3 → 8) raises a clear error at emit time, not deep in Phase 2.