From 94a4957572edc0c84d338caf4057ae556455ddaf Mon Sep 17 00:00:00 2001 From: booth-algo Date: Mon, 1 Jun 2026 15:41:00 +0100 Subject: [PATCH 1/2] Support batch_size>1 in non-packed MHA decoder (delegate to kernel) The non-packed decoder attention previously raised NotImplementedError for batch_size>1. Rather than re-implement a per-batch loop in the decoder (which NaN-ed under hand-written column-block-major slab slicing), span the per-head Q/K/V/O across all batches and let the flash-attention kernel's own batch loop handle it. head_dim<=mlen works now; head_dim>mlen still raises gracefully via the kernel's guard (sub-64 lands separately). Also make the batch>1 output comparison truthful. The compact golden holds only the active rows (batch_size*seq_len), which live rpb-strided in VRAM (batch b at physical row b*rows_per_batch). Set num_batches to the active-row count and carry rows_per_batch/active_seq_per_batch so the verifier extracts the strided active rows instead of comparing all physical (padding) rows. batch=2 @64/64/16 seq=4: 97.37% allclose PASS (previously a 29% false-negative); batch=1 unchanged at 99.96%. --- aten/plena_frontend.py | 100 +++++++++++++++++++++++++++++++++-------- 1 file changed, 82 insertions(+), 18 deletions(-) diff --git a/aten/plena_frontend.py b/aten/plena_frontend.py index da97d91..f091176 100644 --- a/aten/plena_frontend.py +++ b/aten/plena_frontend.py @@ -1056,9 +1056,15 @@ def _emit_attention_block( checkpoint_recorder: StageCheckpointRecorder | None = None, active_seq_len: int | None = None, active_hidden: int | None = None, + batch_size: int = 1, + rows_per_batch: int | None = None, + active_seq_len_per_batch: int | None = None, ): active_seq_len = active_seq_len or seq_len active_hidden = active_hidden or current.shape[1] + active_seq_len_per_batch = active_seq_len_per_batch or seq_len + rows_per_batch = rows_per_batch or max(prog.mlen, seq_len) + total_physical_rows = batch_size * rows_per_batch if batch_size > 1 else max(prog.mlen, seq_len) if checkpoint_recorder is not None: checkpoint_recorder.record( prog, @@ -1080,7 +1086,13 @@ def _emit_attention_block( semantic="attention RMS-normalized input", ) - Q = _linear_projection(prog, current, layer_inputs.w_q, f"Q_{layer_idx}") + Q = _linear_projection( + prog, + current, + layer_inputs.w_q, + f"Q_{layer_idx}", + physical_shape=(total_physical_rows, total_q_dim) if batch_size > 1 else None, + ) q_full_addr = prog.get_vram_addr(Q.name) if checkpoint_recorder is not None: checkpoint_recorder.record( @@ -1092,7 +1104,13 @@ def _emit_attention_block( semantic="Q projection before per-head RoPE", ) - O_full = prog.alloc(f"O_full_{layer_idx}", seq_len, total_q_dim, strict=False) + O_full = prog.alloc( + f"O_full_{layer_idx}", + current.shape[0] if batch_size > 1 else seq_len, + total_q_dim, + strict=False, + physical_shape=(total_physical_rows, total_q_dim) if batch_size > 1 else None, + ) o_full_addr = prog.get_vram_addr(O_full.name) prog.vram_fill_zero(O_full) # Each head occupies head_dim//mlen col-blocks of O_full, so the per-head span @@ -1110,18 +1128,26 @@ def _emit_attention_block( rope_inputs, layer_idx, num_kv_heads, + physical_rows=total_physical_rows if batch_size > 1 else None, checkpoint_recorder=checkpoint_recorder, active_seq_len=active_seq_len, active_head_dim=head_dim, ) rope_matrix, cos_var, sin_var = rope_inputs + q_h_phys = (total_physical_rows, head_dim) if batch_size > 1 else None for h in range(num_heads): kv_h = h // ratio K_stored, V_stored = kv_stored[kv_h] q_h_addr = q_full_addr + h * head_stride - Q_h = prog.alloc_at(f"Q_h{h}_{layer_idx}", seq_len, head_dim, q_h_addr) + Q_h = prog.alloc_at( + f"Q_h{h}_{layer_idx}", + current.shape[0] if batch_size > 1 else seq_len, + head_dim, + q_h_addr, + physical_shape=q_h_phys, + ) _apply_rope_projection( prog, Q_h, @@ -1131,6 +1157,8 @@ def _emit_attention_block( f"Q_rot_{layer_idx}_h{h}", ) + # Delegate the per-batch loop to the kernel (proven for head_dim <= mlen; + # the kernel's own guard blocks head_dim > mlen until the kernel fix lands). O_h = ops.flash_attention( prog, Q_h, @@ -1138,20 +1166,47 @@ def _emit_attention_block( V_stored, scale, causal_mask=causal_mask, - batch_size=1, - seq_len=active_seq_len, - kv_seq_len=active_seq_len, + batch_size=batch_size, + seq_len=active_seq_len_per_batch, + kv_seq_len=active_seq_len_per_batch, ) - o_h_dest_addr = o_full_addr + h * head_stride - _copy_into_vram_view( - prog, - O_h, - f"O_dest_h{h}_{layer_idx}", - seq_len, - head_dim, - o_h_dest_addr, - ) + if batch_size == 1: + o_h_dest_addr = o_full_addr + h * head_stride + _copy_into_vram_view( + prog, + O_h, + f"O_dest_h{h}_{layer_idx}", + seq_len, + head_dim, + o_h_dest_addr, + ) + else: + # The kernel packs batches at active_seq stride; O_full needs them at the + # decoder's rows_per_batch stride (so O_proj feeds the rpb-strided + # residual). Remap each batch's active rows into its O_full slab. + o_h_addr = prog.get_vram_addr(O_h.name) + o_h_batch_stride = active_seq_len_per_batch * prog.mlen + o_full_batch_stride = rows_per_batch * prog.mlen + for b in range(batch_size): + O_hb = prog.alloc_at( + f"O_src_h{h}_b{b}_{layer_idx}", + active_seq_len_per_batch, + head_dim, + o_h_addr + b * o_h_batch_stride, + physical_shape=O_h.physical_shape, + ) + # O_full is already zero-filled, so add directly into batch b's slab. + # (Do NOT use _copy_into_vram_view: its vram_fill_zero uses the + # full-height physical_shape and would clobber adjacent heads/batches.) + O_dst = prog.alloc_at( + f"O_dest_h{h}_b{b}_{layer_idx}", + active_seq_len_per_batch, + head_dim, + o_full_addr + h * head_stride + b * o_full_batch_stride, + physical_shape=(total_physical_rows, head_dim), + ) + prog.vram_add(O_dst, O_hb) _free_named_tensors(prog, ("O", "S", "PV")) if checkpoint_recorder is not None: @@ -3333,8 +3388,6 @@ def _verbose(message: str = ""): active_seq_len_per_batch=seq_len, ) else: - if batch_size != 1: - raise NotImplementedError("native non-packed MHA decoder lowering does not yet support batch_size > 1") current_after_attn = _emit_attention_block( prog, current, @@ -3353,6 +3406,9 @@ def _verbose(message: str = ""): checkpoints, checkpoint_rows, hidden, + batch_size=batch_size, + rows_per_batch=rows_per_batch, + active_seq_len_per_batch=seq_len, ) current = _emit_ffn_block( @@ -3423,6 +3479,12 @@ def _verbose(message: str = ""): out_features = padded_hidden comparison_rows = compile_seq_rows if batch_size > 1 else seq_len + # The compact golden holds only the active rows (batch_size*seq_len). For batch>1 + # the active rows are rpb-strided (batch b lives at physical row b*rows_per_batch), + # so num_batches counts active rows and the reader extracts them via + # rows_per_batch/active_seq_per_batch instead of reading contiguous rows. For + # batch==1 this equals seq_len and the reader falls back to contiguous (unchanged). + active_comparison_rows = batch_size * seq_len # Output is column-block-major: each batch's `out_features` span ceil(out_features/mlen) # column blocks, and consecutive col-blocks of a batch are `physical_rows` rows apart @@ -3435,11 +3497,13 @@ def _verbose(message: str = ""): comparison_params = { "start_row_idx": o_vram_addr // mlen, "num_rows": _num_col_blocks * _physical_rows if out_features > mlen else comparison_rows, - "num_batches": comparison_rows, + "num_batches": active_comparison_rows, "elements_per_batch": out_features, "row_dim": mlen, "physical_rows": _physical_rows, "use_stride_mode": out_features > mlen, + "rows_per_batch": rows_per_batch if batch_size > 1 else None, + "active_seq_per_batch": seq_len if batch_size > 1 else None, } info = { From 0f55571086973e3a65e62a3bf0c251f6161fb96c Mon Sep 17 00:00:00 2001 From: booth-algo Date: Mon, 1 Jun 2026 16:40:19 +0100 Subject: [PATCH 2/2] Support head_dim > mlen in batched MHA kernel (true sub-64 multi-batch) The non-packed flash-attention kernel guarded head_dim > mlen with batch_size > 1 because the per-batch Q/O views truncated physical_shape to one batch, which is only valid for a single head_dim col-block. Tensors are column-block-major (col-block cb at cb*R*mlen, row r at r*mlen), so for head_dim > mlen a batch's rows are interleaved across head_dim/mlen col-blocks and the truncated views mis-strided every col-block past the first. Fix: preserve the full physical_shape in Q_batch/O_batch so the col-block stride stays R*mlen, and set the per-batch Q base offset to q_rows_per_batch*mlen (the flat distance to batch b's first row within col-block 0) instead of q_rows_per_batch*physical_shape[1] (which over-skips by head_dim/mlen col-blocks). The two expressions coincide at head_dim == mlen, so the existing head_dim <= mlen batched path is unchanged. o_batch_stride (seq_len*mlen) already matched the decoder's O_h read and is kept. Verified: 16/16/4 batch=2 seq=4 (head_dim 64 > mlen 16, 4 col-blocks) = 99.07% allclose PASS, no NaN -- previously raised NotImplementedError. 64/64/16 batch=2 regression = 97.37%, byte-identical to before (no-op at head_dim == mlen). batch=1 unaffected (single-batch fast path). --- aten/plena/program_attention.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/aten/plena/program_attention.py b/aten/plena/program_attention.py index d84f1ed..932d190 100644 --- a/aten/plena/program_attention.py +++ b/aten/plena/program_attention.py @@ -159,8 +159,6 @@ def _flash_attention_mha( raise ValueError(f"K physical rows per batch {k_rows_per_batch} must be multiple of MLEN={mlen}") if batch_size > 1 and v_rows_per_batch % mlen != 0: raise ValueError(f"V physical rows per batch {v_rows_per_batch} must be multiple of MLEN={mlen}") - if head_dim > mlen and batch_size > 1: - raise NotImplementedError("Batched MHA currently supports head_dim <= MLEN.") if scale is None: scale = 1.0 / math.sqrt(head_dim) @@ -189,7 +187,15 @@ def _flash_attention_mha( q_base = self.get_vram_addr(Q.name) o_base = self.get_vram_addr(O.name) - q_batch_stride = q_rows_per_batch * Q.physical_shape[1] + # Tensors are column-block-major: col-block cb of a tensor with physical height R + # starts at cb*R*mlen, and row r within a col-block is at r*mlen. So advancing one + # batch (q_rows_per_batch rows) within col-block 0 skips q_rows_per_batch*mlen flat + # elements — NOT q_rows_per_batch*physical_shape[1], which over-skips by head_dim/mlen + # col-blocks when head_dim > mlen. The two expressions coincide at head_dim == mlen + # (physical_shape[1] == mlen), so the head_dim <= mlen path is unchanged. + q_batch_stride = q_rows_per_batch * mlen + # O packs batches contiguously by seq_len rows (the decoder reads O_h batch b at + # b*seq_len*mlen), so its per-batch base offset is seq_len*mlen. o_batch_stride = seq_len * mlen for batch_idx in range(batch_size): @@ -202,14 +208,20 @@ def _flash_attention_mha( seq_len, head_dim, q_base + batch_idx * q_batch_stride, - physical_shape=(q_rows_per_batch, Q.physical_shape[1]), + # Preserve the full height so the col-block stride stays R*mlen; the + # per-batch base offset then lands batch b correctly inside every + # col-block. Truncating to q_rows_per_batch only works for a single + # col-block (head_dim <= mlen) and mis-strides higher col-blocks. + physical_shape=Q.physical_shape, ) O_batch = self.alloc_at( f"_mha_O_b{batch_idx}", seq_len, head_dim, o_base + batch_idx * o_batch_stride, - physical_shape=(max(mlen, seq_len), O.physical_shape[1]), + # Same reasoning as Q_batch: keep O's full height so writes to higher + # col-blocks (head_dim > mlen) land at R*mlen strides. + physical_shape=O.physical_shape, ) batch_k_block_base = batch_idx * k_row_blocks_per_batch