From 1f8760f4d3eb7519a495f8d84c7a6b054add8739 Mon Sep 17 00:00:00 2001 From: booth-algo Date: Wed, 3 Jun 2026 02:05:28 +0100 Subject: [PATCH] fix(attention): correct causal mask for multi-tile decoder attention (seq_len > mlen) The native per-head flash-attention kernel (_flash_attention_mha) added the same static (mlen, mlen) triangular causal mask to every QK score tile. That is correct only when seq_len <= mlen, where the score matrix is a single tile. When seq_len > mlen the seq x seq score matrix tiles into a grid and the mask must depend on tile position: strictly-future key tiles (k_idx > q_idx) need to be fully masked but the triangular mask left their lower triangle open (leaking future keys), and strictly-past key tiles (k_idx < q_idx) need no mask but the triangular mask wrongly masked their upper triangle (dropping valid past keys). Only the diagonal tile was correct, producing a partial numerical divergence (~58% allclose) on decoder runs at seq_len > mlen. The kernel now derives each tile's global causal geometry: strictly-future key tiles contribute exp(-inf)=0 and are skipped entirely, the triangular mask is applied only on the diagonal tile, and strictly-past tiles are left unmasked. A NotImplementedError guards the unsupported kv_seq_len != seq_len (KV-cache query offset) case instead of silently miscomputing. seq_len <= mlen remains the single-tile special case and is byte-identical; non-causal callers (causal_mask=None, e.g. vision) are unaffected. This regime was never exercised before: the sub-64 decoder fixes (#54/#55) validated seq_len=4 < mlen (single tile), the default mlen=256 keeps seq_len=64 < mlen (single tile), and vision is non-causal. Verified by isolation (native_64x64x16 seq=64 PASS vs seq=128 was-FAIL-now-PASS, identical config) and by the full SmolVLM2 1L decoder sweep at mlen=16/32 (all blen) going from allclose FAIL to PASS. Adds test_mha_causal_skips_future_tiles_and_masks_only_diagonal asserting future tiles are skipped and the mask lands on diagonal tiles only. --- aten/plena/program_attention.py | 34 ++++++++++++++++++++++- aten/tests/test_plena_compiler.py | 45 +++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 1 deletion(-) diff --git a/aten/plena/program_attention.py b/aten/plena/program_attention.py index 932d190..d1df9fa 100644 --- a/aten/plena/program_attention.py +++ b/aten/plena/program_attention.py @@ -198,6 +198,12 @@ def _flash_attention_mha( # b*seq_len*mlen), so its per-batch base offset is seq_len*mlen. o_batch_stride = seq_len * mlen + # Position of the first query relative to the keys: queries are the LAST + # seq_len of the kv_seq_len key positions, so query row r of block q_idx is + # global position q_offset + q_idx*mlen + r. For prefill (kv_seq_len == + # seq_len) this is 0. + q_offset = kv_seq_len - seq_len + for batch_idx in range(batch_size): if batch_size == 1 and total_q_rows == seq_len: Q_batch = Q @@ -231,6 +237,32 @@ def _flash_attention_mha( for k_idx in range(num_k_blocks): block_cols = min(mlen, kv_seq_len - k_idx * mlen) + needs_triangular_mask = False + if causal_mask is not None: + # Causal geometry across tiles. Query rows of block q_idx are + # global positions q_offset + q_idx*mlen + [0, block_rows); + # key cols of block k_idx are k_idx*mlen + [0, block_cols). A + # key block entirely in the strict future of every query row + # contributes nothing (exp(-inf)=0) and is skipped; one + # entirely in the past is fully visible (no mask); only a + # straddling block needs the triangular mask. The static + # (mlen, mlen) mask encodes a zero-diagonal triangle, which is + # exactly right when the straddle sits on the q_idx == k_idx + # diagonal (q_offset == 0, i.e. prefill). seq_len <= mlen is + # the single-block special case of this and is unchanged. + key_first = k_idx * mlen + query_first = q_offset + q_idx * mlen + query_last = query_first + block_rows - 1 + if key_first > query_last: + continue # whole key block is in the strict future + needs_triangular_mask = key_first + block_cols - 1 > query_first + if needs_triangular_mask and query_first != key_first: + raise NotImplementedError( + "Causal mask across tiles with kv_seq_len != seq_len " + "(non-zero query offset) is unsupported: the static " + "(mlen, mlen) mask only encodes the zero diagonal " + f"(q_offset={q_offset}, q_idx={q_idx}, k_idx={k_idx})." + ) physical_k_idx = batch_k_block_base + k_idx self.vram_sub_projection_T_to( Q_batch, @@ -244,7 +276,7 @@ def _flash_attention_mha( valid_col_mask = valid_col_masks.get(block_cols) if valid_col_mask is not None: self.vram_add(S_block, valid_col_mask, num_rows=block_rows) - if causal_mask is not None: + if needs_triangular_mask: self.vram_add(S_block, causal_mask) softmax_valid_cols = None if valid_col_mask is not None else block_cols self.online_softmax_block(S_block, scale, rows=block_rows, valid_cols=softmax_valid_cols) diff --git a/aten/tests/test_plena_compiler.py b/aten/tests/test_plena_compiler.py index 2a0d3e2..22bf3ab 100644 --- a/aten/tests/test_plena_compiler.py +++ b/aten/tests/test_plena_compiler.py @@ -688,6 +688,51 @@ def test_mha_accepts_batch_slabs(): print(" PASS test_mha_accepts_batch_slabs") +def test_mha_causal_skips_future_tiles_and_masks_only_diagonal(): + """Causal MHA across multiple seq tiles (seq_len > mlen). + + Regression for the seq>mlen causal bug: the static (mlen, mlen) triangular + mask used to be added to *every* score tile, which leaked future keys on + above-diagonal tiles and dropped valid past keys on below-diagonal tiles + (correct only for the single-tile seq_len <= mlen case). The kernel must now + mask only the diagonal tile, skip strictly-future key tiles entirely, and + leave strictly-past tiles unmasked. + """ + from compiler.aten.plena import PlenaCompiler + + prog = PlenaCompiler(mlen=64, blen=4) + q_input = prog.input("Q", shape=(128, 64), physical_shape=(128, 64), prestaged_vram_addr=0) + k_input = prog.input("K", shape=(128, 64), physical_shape=(128, 64)) + v_input = prog.input("V", shape=(128, 64), physical_shape=(128, 64)) + q = prog.load_batch(q_input, name="Q") + mask_input = prog.input("causal_mask", shape=(64, 64)) + mask = prog.load_batch(mask_input, name="CAUSAL_MASK") + + o = prog.flash_attention( + q, + k_input, + v_input, + scale=1.0 / 8.0, + causal_mask=mask, + batch_size=1, + seq_len=128, + kv_seq_len=128, + ) + asm = prog.compile() + + assert o.shape == (128, 64) + # 2 query tiles x 2 key tiles. Causal keeps (q0,k0), (q1,k0), (q1,k1) and + # drops the strictly-future (q0,k1): key tile 0 is consumed by both query + # tiles, key tile 1 only by the second query tile. + assert asm.count("Compute PV = P @ V[k_idx=0]") == 2 + assert asm.count("Compute PV = P @ V[k_idx=1]") == 1 + # The triangular mask is added on the two diagonal tiles only, never on the + # below-diagonal fully-visible tile (pre-fix this was 4). + assert asm.count("+= CAUSAL_MASK") == 2 + + print(" PASS test_mha_causal_skips_future_tiles_and_masks_only_diagonal") + + def test_compile_native_hf_decoder_golden_vs_hf(): """Golden (MXFP8+BF16) should closely match HF float32 at native dims.""" from compiler.aten.plena_frontend import compile_native_hf_decoder