@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cf56cf78b4

Fork 0

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 14, 05:30 UTC

[V1] Add sliding window support to Flex Attention backend (#24089)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Isotr0py committed Sep 21, 2025 at 13:08 UTC cf56cf78b47e5f9b6a81ce0d50a94f9291922315

2 files changed +229 -69

tests/v1/attention/test_attention_backends.py

+157 -51

@@ -1,15 +1,20 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Tests for v1 attention backends without GPUModelRunner dependency."""

4 +from functools import partial

5 +from typing import Optional, Union

7 import pytest

8 import torch

9 +from torch.nn.attention.flex_attention import create_block_mask, flex_attention

11 from tests.v1.attention.utils import (BatchSpec, _Backend,

12 create_common_attn_metadata,

13 create_standard_kv_cache_spec,

14 create_vllm_config,

15 get_attention_backend)

16 +from vllm.config import ModelConfig

17 +from vllm.platforms import current_platform

18 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv, is_torch_equal_or_newer

19 from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,

20 set_kv_cache_layout)

@@ -183,13 +188,19 @@ class MockAttentionLayer:

188 self._v_scale_float = 1.0

189

190

186 -def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,

187 - layer_names: list[str], vllm_config,

188 - device: torch.device,

189 - common_attn_metadata: CommonAttentionMetadata,

190 - query: torch.Tensor, key: torch.Tensor,

191 - value: torch.Tensor,

192 - kv_cache: torch.Tensor) -> torch.Tensor:

191 +def run_attention_backend(

192 + backend: _Backend,

193 + kv_cache_spec: FullAttentionSpec,

194 + layer_names: list[str],

195 + vllm_config,

196 + device: torch.device,

197 + common_attn_metadata: CommonAttentionMetadata,

198 + query: torch.Tensor,

199 + key: torch.Tensor,

200 + value: torch.Tensor,

201 + kv_cache: torch.Tensor,

202 + sliding_window: Optional[int] = None,

203 +) -> torch.Tensor:

204 """Run attention computation using the specified backend's AttentionImpl."""

205

206 # Handle special case for FLEX_ATTENTION_SLOW

@@ -253,7 +264,7 @@ def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,

264 scale=scale,

265 num_kv_heads=num_kv_heads,

266 alibi_slopes=None,

256 - sliding_window=None,

267 + sliding_window=sliding_window,

268 kv_cache_dtype="auto",

269 )

270

@@ -275,13 +286,16 @@ def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,

286 return output

287

288

278 -@pytest.mark.parametrize("batch_spec_name", [

279 - "small_decode", "small_prefill", "mixed_small", "medium_decode",

280 - "medium_prefill", "mixed_medium", "large_decode", "large_prefill",

281 - "single_decode", "single_prefill"

282 -])

283 -@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])

284 -def test_backend_correctness(batch_spec_name: str, model: str):

289 +def _test_backend_correctness(

290 + batch_spec: BatchSpec,

291 + model: str,

292 + backend_to_test: list[Union[_Backend, str]],

293 + mask_mod,

294 + *,

295 + block_size: int = 16,

296 + atol: float = 1e-2,

297 + rtol: float = 1e-2,

298 +):

299 """

300 Test that all backends produce similar outputs to a reference implementation

301 using torch.nn.functional.scaled_dot_product_attention.

@@ -297,9 +311,10 @@ def test_backend_correctness(batch_spec_name: str, model: str):

311 simulated paged KV cache.

312 5. Comparing the vLLM backend's output to the ground-truth SDPA output.

313 """

300 - batch_spec = BATCH_SPECS[batch_spec_name]

314 + current_platform.seed_everything(42)

315 vllm_config = create_vllm_config(model_name=model,

316 max_model_len=max(batch_spec.seq_lens),

317 + block_size=block_size,

318 num_gpu_blocks=8192)

319 device = torch.device("cuda:0")

320

@@ -314,6 +329,7 @@ def test_backend_correctness(batch_spec_name: str, model: str):

329 num_kv_heads = vllm_config.model_config.get_num_kv_heads(

330 vllm_config.parallel_config)

331 head_size = vllm_config.model_config.get_head_size()

332 + sliding_window = vllm_config.model_config.get_sliding_window()

333 dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)

334 block_size = vllm_config.cache_config.block_size

335 scale = 1.0 / (head_size**0.5)

@@ -361,22 +377,21 @@ def test_backend_correctness(batch_spec_name: str, model: str):

377 # Create causal mask: query token i attends to positions 0 to

378 # (context_len + i)

379 kv_len = s_len

364 - offset = context_len

365 - attn_mask = torch.full((q_len, kv_len),

366 - float('-inf'),

367 - device=device,

368 - dtype=dtype)

369 - for i in range(q_len):

370 - attn_mask[i, :offset + i + 1] = 0.0

371 -

372 - sdpa_out_i = torch.nn.functional.scaled_dot_product_attention(

373 - q_sdpa_in,

374 - k_sdpa_in,

375 - v_sdpa_in,

376 - attn_mask=attn_mask,

377 - scale=scale,

378 - enable_gqa=True)

379 - # Convert back to (L, H, D)

380 +

381 + final_mask_mod = partial(mask_mod, context_len=context_len)

382 + block_mask = create_block_mask(final_mask_mod,

383 + B=None,

384 + H=None,

385 + Q_LEN=q_len,

386 + KV_LEN=kv_len,

387 + device=device)

388 + sdpa_out_i = flex_attention(q_sdpa_in,

389 + k_sdpa_in,

390 + v_sdpa_in,

391 + block_mask=block_mask,

392 + scale=scale,

393 + enable_gqa=True)

394 +

395 all_sdpa_outputs.append(sdpa_out_i.transpose(1, 2).squeeze(0))

396

397 # Inputs for vLLM backends are just the new tokens

@@ -412,7 +427,7 @@ def test_backend_correctness(batch_spec_name: str, model: str):

427 # 4. Run vLLM backends and compare

428 # Note: flex_attention has known Triton kernel compatibility issues

429 # with test infrastructures

415 - for backend_name in BACKENDS_TO_TEST:

430 + for backend_name in backend_to_test:

431 # FlashAttentionm + FlexAttention:

432 # [2, num_blocks, block_size, num_kv_heads, head_size]

433 # FlashInfer:

@@ -427,12 +442,19 @@ def test_backend_correctness(batch_spec_name: str, model: str):

442 2, 3).contiguous().transpose(2, 3)

443 set_kv_cache_layout("HND")

444

430 - backend_output = run_attention_backend(backend_name, kv_cache_spec,

431 - ["placeholder"], vllm_config,

432 - device, common_attn_metadata,

433 - query_vllm, key_vllm,

434 - value_vllm,

435 - kv_cache_for_backend)

445 + backend_output = run_attention_backend(

446 + backend_name,

447 + kv_cache_spec,

448 + ["placeholder"],

449 + vllm_config,

450 + device,

451 + common_attn_metadata,

452 + query_vllm,

453 + key_vllm,

454 + value_vllm,

455 + kv_cache_for_backend,

456 + sliding_window=sliding_window,

457 + )

458

459 # Check shape and dtype consistency

460 assert backend_output.shape == sdpa_output.shape, (

@@ -446,18 +468,102 @@ def test_backend_correctness(batch_spec_name: str, model: str):

468 f"[{backend_name}] produced non-finite values")

469

470 # Check numerical similarity

449 - rtol = 1e-2

450 - atol = 5e-3

451 -

452 - max_diff = torch.max(torch.abs(backend_output - sdpa_output)).item()

453 - max_rel_diff = torch.max(

454 - torch.abs(backend_output - sdpa_output) /

455 - torch.abs(sdpa_output)).item()

456 - all_close = torch.allclose(backend_output,

471 + def error_msg(msg: str, backend_name: str):

472 + return (f"[{backend_name}] output differs from SDPA baseline. "

473 + f"{msg}")

474 +

475 + torch.testing.assert_close(backend_output,

476 sdpa_output,

477 rtol=rtol,

459 - atol=atol)

478 + atol=atol,

479 + msg=partial(error_msg,

480 + backend_name=backend_name))

481

461 - assert all_close, (

462 - f"[{backend_name}] output differs from SDPA baseline. "

463 - f"Max diff: {max_diff:.6f}, max rel diff: {max_rel_diff:.6f})")

\ No newline at end of file

482 +

483 +@pytest.mark.parametrize("batch_spec_name", [

484 + "small_decode", "small_prefill", "mixed_small", "medium_decode",

485 + "medium_prefill", "mixed_medium", "large_decode", "large_prefill",

486 + "single_decode", "single_prefill"

487 +])

488 +@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])

489 +def test_causal_backend_correctness(batch_spec_name: str, model: str):

490 + """Test backend's correctness with causal attention."""

491 +

492 + def causal_mask_mod(

493 + b: torch.Tensor,

494 + h: torch.Tensor,

495 + q_idx: torch.Tensor,

496 + kv_idx: torch.Tensor,

497 + *,

498 + context_len: int,

499 + ):

500 + return (q_idx + context_len) >= kv_idx

501 +

502 + batch_spec = BATCH_SPECS[batch_spec_name]

503 + LARGE_BLOCK_BACKENDS = ([_Backend.FLEX_ATTENTION]

504 + if is_torch_equal_or_newer("2.9.0.dev0") else [])

505 + SMALL_BLOCK_BACKENDS = [

506 + x for x in BACKENDS_TO_TEST if x not in LARGE_BLOCK_BACKENDS

507 + ]

508 + _test_backend_correctness(batch_spec, model, SMALL_BLOCK_BACKENDS,

509 + causal_mask_mod)

510 +

511 + # Fast FlexAttention needs to run with block_size=128

512 + if LARGE_BLOCK_BACKENDS:

513 + _test_backend_correctness(batch_spec,

514 + model,

515 + LARGE_BLOCK_BACKENDS,

516 + causal_mask_mod,

517 + block_size=128)

518 +

519 +

520 +SLIDING_WINDOW_BACKENDS_TO_TEST = [

521 + _Backend.FLASH_ATTN_VLLM_V1, _Backend.FLEX_ATTENTION,

522 + _Backend.TRITON_ATTN_VLLM_V1, "FLEX_ATTENTION_SLOW"

523 +]

524 +

525 +

526 +@pytest.mark.parametrize("batch_spec_name", [

527 + "small_decode", "small_prefill", "mixed_medium", "large_decode",

528 + "large_prefill"

529 +])

530 +@pytest.mark.parametrize("model", ["microsoft/Phi-tiny-MoE-instruct"])

531 +def test_sliding_window_backend_correctness(batch_spec_name: str, model: str):

532 + """Test backend's correctness with sliding window attention."""

533 +

534 + def sliding_window_mask_mod(

535 + b: torch.Tensor,

536 + h: torch.Tensor,

537 + q_idx: torch.Tensor,

538 + kv_idx: torch.Tensor,

539 + *,

540 + context_len: int,

541 + sliding_window: int,

542 + ):

543 + causal_mask = q_idx + context_len >= kv_idx

544 + window_mask = q_idx + context_len - kv_idx < sliding_window

545 + return causal_mask & window_mask

546 +

547 + batch_spec = BATCH_SPECS[batch_spec_name]

548 + model_config = ModelConfig(model=model,

549 + max_model_len=max(batch_spec.seq_lens))

550 + sliding_window = model_config.get_sliding_window()

551 + sliding_window_mask_mod_fn = partial(sliding_window_mask_mod,

552 + sliding_window=sliding_window)

553 +

554 + LARGE_BLOCK_BACKENDS = ([_Backend.FLEX_ATTENTION]

555 + if is_torch_equal_or_newer("2.9.0.dev0") else [])

556 + SMALL_BLOCK_BACKENDS = [

557 + x for x in SLIDING_WINDOW_BACKENDS_TO_TEST

558 + if x not in LARGE_BLOCK_BACKENDS

559 + ]

560 + _test_backend_correctness(batch_spec, model, SMALL_BLOCK_BACKENDS,

561 + sliding_window_mask_mod_fn)

562 +

563 + # Fast FlexAttention needs to run with block_size=128

564 + if LARGE_BLOCK_BACKENDS:

565 + _test_backend_correctness(batch_spec,

566 + model,

567 + LARGE_BLOCK_BACKENDS,

568 + sliding_window_mask_mod_fn,

569 + block_size=128)

vllm/v1/attention/backends/flex_attention.py

+72 -18

@@ -9,7 +9,7 @@ import torch

9 import torch._dynamo.decorators

10 import torch.nn.functional as F

11 from torch.nn.attention.flex_attention import (BlockMask, _mask_mod_signature,

12 - _score_mod_signature,

12 + _score_mod_signature, and_masks,

13 create_block_mask,

14 flex_attention)

@@ -292,6 +292,7 @@ class FlexAttentionMetadata:

292 q_block_size: int = 16

293 kv_block_size: int = 16

294 transformed_score_mod: Optional[_score_mod_signature] = None

295 + sliding_window: Optional[int] = None

296

297 def _convert_physical_to_logical(

298 self,

@@ -380,6 +381,53 @@ class FlexAttentionMetadata:

381

382 return final_mask_mod

383

384 + def get_sliding_window_mask_mod(self) -> _mask_mod_signature:

385 + """Creates the sliding window mask_mod function for FlexAttention.

386 +

387 + Note that the sliding window mask here is bidirectional, we need

388 + to mask it with the bidirectional/causal mask for encoder/decoder.

389 + """

390 +

391 + if self.sliding_window is None:

392 + raise ValueError(

393 + "sliding_window must be set for sliding window attention")

394 +

395 + def sliding_window_mask_mod(b: torch.Tensor, h: torch.Tensor,

396 + q_idx: torch.Tensor, kv_idx: torch.Tensor):

397 + return torch.abs(q_idx - kv_idx) < self.sliding_window

398 +

399 + def final_mask_mod(

400 + b: torch.Tensor,

401 + h: torch.Tensor,

402 + q_idx: torch.Tensor,

403 + physical_kv_idx: torch.Tensor,

404 + ) -> torch.Tensor:

405 + (is_valid, logical_q_idx,

406 + logical_kv_idx) = self._convert_physical_to_logical(

407 + self.doc_ids, q_idx, physical_kv_idx)

408 + return torch.where(

409 + is_valid,

410 + sliding_window_mask_mod(b, h, logical_q_idx, logical_kv_idx),

411 + False,

412 + )

413 +

414 + return final_mask_mod if self.causal else sliding_window_mask_mod

415 +

416 + def get_mask_mod(self):

417 + # Stage-1: initialize the base mask_mod

418 + # (causal mask for decoder or bidirectional mask for encoder)

419 + if self.causal:

420 + mask_mod = self.get_causal_mask_mod()

421 + else:

422 + mask_mod = self.get_bidirectional_mask_mod()

423 + # stage-2: add external mask_mod for special attention during

424 + # forwarding runtime to create the combined mask_mod.

425 + if self.sliding_window is not None:

426 + # Add sliding window mask for sliding window attention

427 + sliding_window_mask_mod = self.get_sliding_window_mask_mod()

428 + mask_mod = and_masks(mask_mod, sliding_window_mask_mod)

429 + return mask_mod

430 +

431 def get_transformed_score_mod(self) -> Optional[_score_mod_signature]:

432 """Creates the transformed score_mod function for FlexAttention.

433

@@ -472,12 +520,9 @@ class FlexAttentionMetadata:

520 return BlockMask.from_kv_blocks(**block_mask_kwargs)

521

522 def build_block_mask(self) -> BlockMask:

475 - if self.causal:

476 - mask_mod = self.get_causal_mask_mod()

477 - kv_len = self.total_cache_tokens

478 - else:

479 - mask_mod = self.get_bidirectional_mask_mod()

480 - kv_len = self.num_actual_tokens

523 + mask_mod = self.get_mask_mod()

524 + kv_len = (self.total_cache_tokens

525 + if self.causal else self.num_actual_tokens)

526 return create_block_mask_compiled(

527 mask_mod,

528 None,

@@ -498,11 +543,7 @@ class FlexAttentionMetadata:

543 self.doc_ids = _offsets_to_doc_ids_tensor(self.query_start_loc)

544 self.num_blocks = self.total_cache_tokens // self.block_size

545

501 - if self.causal:

502 - self.mask_mod = self.get_causal_mask_mod()

503 - else:

504 - self.mask_mod = self.get_bidirectional_mask_mod()

505 -

546 + self.mask_mod = self.get_mask_mod()

547 self.transformed_score_mod = self.get_transformed_score_mod()

548

549 if self.direct_build and self.causal:

@@ -607,7 +648,7 @@ class FlexAttentionMetadataBuilder(

648

649

650 class FlexAttentionImpl(AttentionImpl):

610 - sliding_window: Optional[tuple[int, int]]

651 + sliding_window: Optional[int]

652 alibi_slopes: Optional[torch.Tensor]

653 logits_soft_cap: Optional[float]

654

@@ -641,11 +682,9 @@ class FlexAttentionImpl(AttentionImpl):

682 "FlexAttention does not support alibi slopes yet.")

683 else:

684 self.alibi_slopes = None

644 - if sliding_window is not None:

645 - raise NotImplementedError(

646 - "FlexAttention does not support sliding window yet.")

647 - else:

648 - self.sliding_window = (-1, -1)

685 +

686 + self.sliding_window = sliding_window

687 +

688 self.kv_cache_dtype = kv_cache_dtype

689 self.logits_soft_cap = logits_soft_cap

690 if self.logits_soft_cap is not None:

@@ -712,6 +751,21 @@ class FlexAttentionImpl(AttentionImpl):

751

752 num_actual_tokens = attn_metadata.num_actual_tokens

753

754 + if attn_metadata.sliding_window != self.sliding_window:

755 + attn_metadata.sliding_window = self.sliding_window

756 + if attn_metadata.direct_build:

757 + # TODO: Support skipping the computation of sliding window

758 + # in direct block mask building code path.

759 + logger.warning_once(

760 + "Using direct block mask building with sliding window, "

761 + "which is suboptimal now. Performance may be degraded.")

762 + # update mask mod in attention metadata

763 + attn_metadata.mask_mod = attn_metadata.get_mask_mod()

764 + attn_metadata.block_mask = (

765 + attn_metadata._build_block_mask_direct())

766 + else:

767 + attn_metadata.block_mask = attn_metadata.build_block_mask()

768 +

769 if not attn_metadata.causal:

770 assert self.attn_type == AttentionType.ENCODER_ONLY

771