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 15, 06:00 UTC
Fixes IMA for TP w/ flex-attention (#19712)
Signed-off-by: drisspg <drisspguessous@gmail.com>
Driss Guessous committed
Jun 16, 2025 at 21:01 UTC
ddfed314f9c3f2f09b207bc4c4c1b6c35af53e04
2 files changed
+2
-10
tests/kernels/test_flex_attention.py
-2
@@ -51,7 +51,6 @@ def test_flex_attention_vs_default_backend(monkeypatch):
51
with monkeypatch.context() as m:
52
m.setenv("VLLM_USE_V1", "1")
53
m.setenv("VLLM_ATTENTION_BACKEND", "FLEX_ATTENTION")
54
- m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
54
55
set_seed(seed)
56
@@ -66,7 +65,6 @@ def test_flex_attention_vs_default_backend(monkeypatch):
65
# Run with default backend
66
with monkeypatch.context() as m:
67
m.setenv("VLLM_USE_V1", "1")
69
- m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
68
set_seed(seed)
69
llm_default = LLM(
70
model_name,
vllm/v1/attention/backends/flex_attention.py
+2
-8
@@ -13,7 +13,6 @@ from torch.nn.attention.flex_attention import (BlockMask, _mask_mod_signature,
13
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
14
AttentionMetadata, AttentionType,
15
is_quantized_kv_cache)
16
-from vllm.distributed import get_tensor_model_parallel_world_size
16
from vllm.logger import init_logger
17
from vllm.platforms import current_platform
18
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
@@ -237,17 +236,13 @@ class FlexAttentionMetadata:
236
237
def build_block_mask(self) -> BlockMask:
238
assert self.mask_mod is not None
240
- # FIXME: With TP>1, create_block_mask_compiled will raise
241
- # CUDA error: an illegal memory access was encountered
242
- create_block_mask_fn = (create_block_mask_compiled
243
- if get_tensor_model_parallel_world_size() == 1
244
- else create_block_mask)
245
- return create_block_mask_fn(
239
+ return create_block_mask_compiled(
240
self.mask_mod,
241
None,
242
None,
243
self.num_actual_tokens,
244
self.total_cache_tokens,
245
+ device=self.block_table.device,
246
)
247
248
def __post_init__(self):
@@ -429,7 +424,6 @@ class FlexAttentionImpl(AttentionImpl):
424
shape = [num_tokens, num_heads * head_size]
425
"""
426
assert output is not None, "Output tensor must be provided."
432
-
427
if output_scale is not None:
428
raise NotImplementedError(
429
"fused output quantization is not yet supported"