Skip to content

Support batch_size>1 in native MHA decoder, incl. true sub-64 (head_dim>mlen) #55

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
booth-algo merged 2 commits into from
Jun 2, 2026
Merged

Support batch_size>1 in native MHA decoder, incl. true sub-64 (head_dim>mlen) #55

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
94a4957
Support batch_size>1 in non-packed MHA decoder (delegate to kernel)
booth-algo Jun 1, 2026
0f55571
Support head_dim > mlen in batched MHA kernel (true sub-64 multi-batch)
booth-algo Jun 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 17 additions & 5 deletions aten/plena/program_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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):
Expand All @@ -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
Expand Down
100 changes: 82 additions & 18 deletions aten/plena_frontend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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(
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -1131,27 +1157,56 @@ 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,
K_stored,
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:
Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand All @@ -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 = {
Expand Down
Loading