From b429a30428b03783a542c2d34e6e20301c56575c Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Thu, 28 May 2026 01:03:32 -0700 Subject: [PATCH] Expose Flex attention causal/decode masks as static methods (#26514) Co-authored-by: Claude Opus 4.7 (1M context) --- python/sglang/srt/layers/attention/torch_flex_backend.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/torch_flex_backend.py b/python/sglang/srt/layers/attention/torch_flex_backend.py index 1af8508cb..1c83d397f 100644 --- a/python/sglang/srt/layers/attention/torch_flex_backend.py +++ b/python/sglang/srt/layers/attention/torch_flex_backend.py @@ -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(