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:
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(
|
||||
|
||||
Reference in New Issue
Block a user