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

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 6, 02:00 UTC

[FlexAttention] allow custom mask mod (#37692)

Signed-off-by: Angel Li <liangel@meta.com>

liangel-02 committed Mar 24, 2026 at 16:03 UTC 8c47fdfdb17c7bab93c70722b7adf71d87deddb9

2 files changed +121 -16

tests/kernels/test_flex_attention.py

+51

@@ -14,6 +14,7 @@ from tests.v1.attention.utils import (

14 create_vllm_config,

15 )

16 from vllm.v1.attention.backends.flex_attention import (

17 + BlockSparsityHint,

18 FlexAttentionMetadataBuilder,

19 physical_to_logical_mapping,

20 )

@@ -223,5 +224,55 @@ def test_physical_to_logical_mapping_handles_reused_blocks():

224 assert out2[0, 2].item() == 1

225

226

227 +@pytest.mark.skipif(

228 + not torch.cuda.is_available() or TORCH_VERSION < DIRECT_BUILD_VERSION,

229 + reason="CUDA not available or PyTorch version < 2.9",

230 +)

231 +def test_block_sparsity_hint_prunes_blocks():

232 + """Test that BlockSparsityHint prunes KV blocks from the direct build path.

233 +

234 + Uses a hint that only keeps the diagonal (q_block == kv_block) to verify

235 + that off-diagonal blocks are excluded from the resulting BlockMask.

236 + """

237 + device = torch.device("cuda")

238 +

239 + vllm_config = create_vllm_config(

240 + model_name="facebook/opt-125m",

241 + block_size=16,

242 + max_model_len=1024,

243 + )

244 + kv_cache_spec = create_standard_kv_cache_spec(vllm_config)

245 +

246 + batch_spec = BatchSpec(

247 + seq_lens=[256],

248 + query_lens=[256],

249 + name="test_sparsity_hint",

250 + )

251 +

252 + common_attn_metadata = create_common_attn_metadata(

253 + batch_spec, vllm_config.cache_config.block_size, device

254 + )

255 +

256 + builder = FlexAttentionMetadataBuilder(kv_cache_spec, [], vllm_config, device)

257 +

258 + metadata_no_hint = builder.build(

259 + common_prefix_len=0, common_attn_metadata=common_attn_metadata

260 + )

261 + metadata_no_hint.block_mask = metadata_no_hint._build_block_mask_direct()

262 + assert metadata_no_hint.block_mask.kv_num_blocks.max().item() > 1

263 +

264 + def diagonal_hint(q_block_idx, kv_block_idx, block_size):

265 + return q_block_idx == kv_block_idx

266 +

267 + metadata_with_hint = builder.build(

268 + common_prefix_len=0, common_attn_metadata=common_attn_metadata

269 + )

270 + metadata_with_hint.block_sparsity_hint = BlockSparsityHint(

271 + hint_fn=diagonal_hint,

272 + )

273 + metadata_with_hint.block_mask = metadata_with_hint._build_block_mask_direct()

274 + assert metadata_with_hint.block_mask.kv_num_blocks.max().item() <= 1

275 +

276 +

277 if __name__ == "__main__":

278 pytest.main([__file__])

vllm/v1/attention/backends/flex_attention.py

+70 -16

@@ -3,9 +3,10 @@

3 """Attention layer with FlexAttention."""

5 import math

6 +from collections.abc import Callable

7 from dataclasses import dataclass

8 from functools import cached_property

8 -from typing import ClassVar

9 +from typing import ClassVar, NamedTuple

11 import torch

12 import torch._dynamo.decorators

@@ -294,6 +295,27 @@ def causal_mask_mod(

295 return q_idx >= kv_idx

296

297

298 +# Type alias for the block sparsity hint callable signature.

299 +_block_sparsity_hint_signature = Callable[

300 + [torch.Tensor, torch.Tensor, int], torch.Tensor

301 +]

302 +

303 +

304 +class BlockSparsityHint(NamedTuple):

305 + """This prunes KV blocks from the BlockMask before the flex_attention kernel

306 + is invoked, so that blocks that are fully masked never get loaded.

307 + Use this with custom mask_mods that are sparse to avoid

308 + the kernel iterating over all KV blocks unnecessarily.

309 +

310 + Attributes:

311 + hint_fn: (q_block_idx [num_tokens, 1], kv_block_idx [1, num_kv_blocks],

312 + block_size int) -> bool Tensor [num_tokens, num_kv_blocks].

313 + Returns True for block pairs that may contain non-masked elements.

314 + """

315 +

316 + hint_fn: _block_sparsity_hint_signature

317 +

318 +

319 @dataclass

320 class FlexAttentionMetadata:

321 causal: bool

@@ -335,6 +357,7 @@ class FlexAttentionMetadata:

357 transformed_score_mod: _score_mod_signature | None = None

358 sliding_window: int | None = None

359 mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None

360 + block_sparsity_hint: BlockSparsityHint | None = None

361

362 @cached_property

363 def logical_block_ids(self):

@@ -378,7 +401,7 @@ class FlexAttentionMetadata:

401

402 return is_valid, logical_q_idx, logical_kv_idx

403

381 - def get_causal_mask_mod(self) -> _mask_mod_signature:

404 + def get_paged_mask_mod(self) -> _mask_mod_signature:

405 """Creates the mask_mod function for FlexAttention.

406

407 This function creates the combined mask mod function that handles:

@@ -504,8 +527,9 @@ class FlexAttentionMetadata:

527 def get_mask_mod(self):

528 # Stage-1: initialize the base mask_mod

529 # (causal mask for decoder or bidirectional mask for encoder)

507 - if self.causal:

508 - mask_mod = self.get_causal_mask_mod()

530 + has_custom_mask = self.logical_mask_mod is not causal_mask_mod

531 + if self.causal or has_custom_mask:

532 + mask_mod = self.get_paged_mask_mod()

533 else:

534 mask_mod = self.get_bidirectional_mask_mod()

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

@@ -591,7 +615,9 @@ class FlexAttentionMetadata:

615 self.doc_ids, : cdiv(self.max_seq_len, self.block_size)

616 ]

617

594 - if self.sliding_window and self.causal:

618 + custom_hint = self.block_sparsity_hint is not None

619 +

620 + if self.sliding_window or custom_hint:

621 device = used_pages.device

622 assert self.doc_ids is not None

623 token_indices = torch.arange(

@@ -602,10 +628,24 @@ class FlexAttentionMetadata:

628 - self.query_start_loc[self.doc_ids]

629 + self.decode_offset[self.doc_ids]

630 )

605 - min_kv_idx = torch.clamp(logical_q_idx - (self.sliding_window - 1), min=0)

606 - min_block_idx = min_kv_idx // self.block_size

607 - sliding_mask = self.logical_block_ids >= min_block_idx[:, None]

608 - used_pages.masked_fill_(~sliding_mask, 0)

631 +

632 + if self.sliding_window:

633 + assert self.sliding_window is not None

634 + min_kv_idx = torch.clamp(

635 + logical_q_idx - (self.sliding_window - 1), min=0

636 + )

637 + min_block_idx = min_kv_idx // self.block_size

638 + sliding_mask = self.logical_block_ids >= min_block_idx[:, None]

639 + used_pages.masked_fill_(~sliding_mask, 0)

640 + if custom_hint:

641 + assert self.block_sparsity_hint is not None

642 + q_block_idx = logical_q_idx // self.block_size

643 + hint_mask = self.block_sparsity_hint.hint_fn(

644 + q_block_idx[:, None],

645 + self.logical_block_ids[None, :],

646 + self.block_size,

647 + )

648 + used_pages.masked_fill_(~hint_mask, 0)

649

650 used_pages_padded = pad_to_multiple(

651 used_pages, multiple=self.q_block_size, dim=0

@@ -660,11 +700,6 @@ class FlexAttentionMetadata:

700 self.mask_mod = self.get_mask_mod()

701 self.transformed_score_mod = self.get_transformed_score_mod()

702

663 - if self.direct_build and self.causal:

664 - self.block_mask = self._build_block_mask_direct()

665 - else:

666 - self.block_mask = self.build_block_mask()

667 -

703

704 class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadata]):

705 def __init__(

@@ -770,6 +805,8 @@ class FlexAttentionImpl(AttentionImpl):

805 alibi_slopes: torch.Tensor | None

806 logits_soft_cap: float | None

807 mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None

808 + logical_mask_mod: _mask_mod_signature | None = None

809 + block_sparsity_hint: BlockSparsityHint | None = None

810

811 def __init__(

812 self,

@@ -907,8 +944,25 @@ class FlexAttentionImpl(AttentionImpl):

944 attn_metadata.mask_mod = attn_metadata.get_mask_mod()

945 needs_rebuild_block_mask = True

946

910 - if needs_rebuild_block_mask:

911 - if attn_metadata.direct_build and attn_metadata.causal:

947 + layer_mask_mod = getattr(layer, "logical_mask_mod", None)

948 + if (

949 + layer_mask_mod is not None

950 + and attn_metadata.logical_mask_mod is not layer_mask_mod

951 + ):

952 + attn_metadata.logical_mask_mod = layer_mask_mod

953 + attn_metadata.mask_mod = attn_metadata.get_mask_mod()

954 + needs_rebuild_block_mask = True

955 +

956 + layer_hint = getattr(layer, "block_sparsity_hint", None)

957 + if (

958 + layer_hint is not None

959 + and attn_metadata.block_sparsity_hint is not layer_hint

960 + ):

961 + attn_metadata.block_sparsity_hint = layer_hint

962 + needs_rebuild_block_mask = True

963 +

964 + if needs_rebuild_block_mask or attn_metadata.block_mask is None:

965 + if attn_metadata.direct_build:

966 attn_metadata.block_mask = attn_metadata._build_block_mask_direct()

967 else:

968 attn_metadata.block_mask = attn_metadata.build_block_mask()