refactor: simplify attention mask handling

2026-07-20 20:36:16 +08:00
parent a6e920fdb0
commit d7ac66fb73
6 changed files with 34 additions and 38 deletions
@@ -26,12 +26,14 @@ class DecoderBlock(nn.Module):
        rotary_emb: Tensor,
        attention_mask: Optional[Tensor] = None,
        paged_cache: Optional[CacheView] = None,
+        is_causal: bool = False,
    ) -> Tensor:
        attn_output = self.attention(
            self.input_norm(x),
            rotary_emb,
            attention_mask,
            paged_cache,
+            is_causal,
        )
        x = attn_output + x
        x = self.mlp(self.post_attention_norm(x)) + x