Expose Flex attention causal/decode masks as static methods (#26514)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-28 01:03:32 -07:00
committed by GitHub
co-authored by Claude Opus 4.7
parent e5f5d84780
commit b429a30428
@@ -70,10 +70,12 @@ class TorchFlexAttnBackend(AttentionBackend):
)
)
def _causal_mask(self, b, h, q_idx, kv_idx):
@staticmethod
def _causal_mask(b, h, q_idx, kv_idx):
return q_idx >= kv_idx
def _decode_mask(self, b, h, q_idx, kv_idx):
@staticmethod
def _decode_mask(b, h, q_idx, kv_idx):
return q_idx <= kv_idx
def _run_flex_forward_extend(