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). This ADR removes the single hard blocker that keeps the existing attention kernels at h_q == h_kv == 1: there is no way to expand a KV head to the G query heads that share it without violating the MemoryStore byte-conservation check.

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.