A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 12:30 UTC
uncomment flex backend for batch invariant mode (#40842)
Signed-off-by: Angel Li <liangel@meta.com>
liangel-02 committed
Apr 29, 2026 at 00:05 UTC
7fd05e05aeb3664ca19346771dc559d93423acd4
2 files changed
+8
-9
tests/v1/determinism/utils.py
+1
@@ -26,6 +26,7 @@ TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
26
BACKENDS: list[str] = [
27
"FLASH_ATTN",
28
"TRITON_ATTN",
29
+ "FLEX_ATTENTION",
30
]
31
32
# FlashInfer temporarily disabled due to invariant CTA sizes.
vllm/v1/attention/backends/flex_attention.py
+7
-9
@@ -99,6 +99,10 @@ class FlexAttentionBackend(AttentionBackend):
99
"""FlexAttention supports both decoder and encoder-only attention."""
100
return attn_type in (AttentionType.DECODER, AttentionType.ENCODER_ONLY)
101
102
+ @classmethod
103
+ def supports_batch_invariance(cls) -> bool:
104
+ return True
105
+
106
@classmethod
107
def supports_mm_prefix(cls) -> bool:
108
"""FlexAttention supports full attention for image tokens."""
@@ -326,15 +330,9 @@ class BlockSparsityHint(NamedTuple):
330
331
332
def copy_to_persistent(dst, src):
329
- try:
330
- dst = dst.as_strided(src.shape, src.stride())
331
- except RuntimeError as e:
332
- raise RuntimeError(
333
- f"Fail to re-stride a persistent tensor of shape {dst.shape} "
334
- f"for a tensor of shape {src.shape}"
335
- ) from e
336
- dst.copy_(src)
337
- return dst
333
+ sliced = dst[tuple(slice(0, s) for s in src.shape)]
334
+ sliced.copy_(src)
335
+ return sliced
336
337
338
@dataclass