fix(attention): correct causal mask for multi-tile decoder attention (seq_len > mlen)

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)

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