@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 82ae4164ee

[2/N][Attention] Enable masked MHA for sparse MLA prefills (#48770)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Matthew Bonanni committed Jul 31, 2026 at 09:01 UTC 82ae4164ee016d4daecd2033c26f5c0827984a80

11 files changed +855 -26

benchmarks/attention_benchmarks/benchmark.py

+17 -1

@@ -617,7 +617,7 @@ def main():

617 "--sparse-mla-mha-variants",

618 nargs="+",

619 default=None,

620 - choices=["dense_mha", "mqa"],

620 + choices=["dense_mha", "masked_mha", "mqa"],

621 help="Sparse MLA variants to run in mha_vs_mqa mode. Defaults to both.",

622 )

623

@@ -748,6 +748,9 @@ def main():

748 args.sparse_mla_mha_variants = yaml_config.get(

749 "sparse_mla_mha_variants", args.sparse_mla_mha_variants

750 )

751 + args.sparse_mla_masked_mha_max_seq_len = yaml_config.get(

752 + "sparse_mla_masked_mha_max_seq_len", None

753 + )

754

755 # Parameter sweep configuration

756 if "parameter_sweep" in yaml_config:

@@ -1119,6 +1122,7 @@ def main():

1122 console.print(f"Prefill backend: {prefill_backend}")

1123 available_variants = [

1124 ("dense_mha", False, "dense"),

1125 + ("masked_mha", False, "masked"),

1126 ("mqa", True, "auto"),

1127 ]

1128 requested_variants = getattr(args, "sparse_mla_mha_variants", None)

@@ -1139,6 +1143,9 @@ def main():

1143 ]

1144 else:

1145 variants = available_variants

1146 + masked_mha_max_seq_len = getattr(

1147 + args, "sparse_mla_masked_mha_max_seq_len", None

1148 + )

1149 formatter = ResultsFormatter(console)

1150 total = 0

1151 for spec in args.batch_specs:

@@ -1148,6 +1155,10 @@ def main():

1155 variant_label == "dense_mha"

1156 and dense_mha_max_seq_len is not None

1157 and q_len > dense_mha_max_seq_len

1158 + ) or (

1159 + variant_label == "masked_mha"

1160 + and masked_mha_max_seq_len is not None

1161 + and q_len > masked_mha_max_seq_len

1162 ):

1163 continue

1164 total += len(backends)

@@ -1161,6 +1172,10 @@ def main():

1172 variant_label == "dense_mha"

1173 and dense_mha_max_seq_len is not None

1174 and q_len > dense_mha_max_seq_len

1175 + ) or (

1176 + variant_label == "masked_mha"

1177 + and masked_mha_max_seq_len is not None

1178 + and q_len > masked_mha_max_seq_len

1179 ):

1180 continue

1181 config = BenchmarkConfig(

@@ -1190,6 +1205,7 @@ def main():

1205 sparse_mla_dense_mha_max_seq_len=dense_mha_max_seq_len,

1206 sparse_mla_topk_pattern=sparse_mla_topk_pattern,

1207 prefill_backend=prefill_backend,

1208 + sparse_mla_masked_mha_max_seq_len=masked_mha_max_seq_len,

1209 )

1210

1211 # run_mla_benchmark needs the real backend name

benchmarks/attention_benchmarks/common.py

+2 -1

@@ -299,8 +299,9 @@ class BenchmarkConfig:

299 num_kv_splits: int | None = None # CUTLASS MLA

300 reorder_batch_threshold: int | None = None # FlashAttn MLA, FlashMLA

301 sparse_mla_force_mqa: bool = False # Force MQA path for sparse MLA

302 - sparse_mla_mha_mode: str = "auto" # "auto" or "dense"

302 + sparse_mla_mha_mode: str = "auto" # "auto", "dense", or "masked"

303 sparse_mla_dense_mha_max_seq_len: int | None = None

304 + sparse_mla_masked_mha_max_seq_len: int | None = None

305 sparse_mla_topk_pattern: str = "random" # "random", "prefix", "sliding_window"

306 num_splits: int | None = None # FlashAttention split-K (0=auto, 1=disabled)

307

benchmarks/attention_benchmarks/configs/mla_sparse_masked_mha_vs_mqa.yaml new

+123

@@ -0,0 +1,123 @@

1 +# Sparse MLA benchmark: forward_mha vs forward_mqa crossover plot

2 +#

3 +# Dense pure-prefill sweep to generate a smooth speedup curve. The 100 points

4 +# run from q32 through q262144, approximately even on the log2 scale, and include

5 +# routing boundaries plus extra samples around the 6k-8k crossover region.

6 +# Dense MHA is only benchmarked through q32768 because it dominates runtime above

7 +# that range.

8 +#

9 +# Usage:

10 +# python benchmark.py --config configs/mla_sparse_mha_vs_mqa.yaml

11 +

12 +mode: mha_vs_mqa

13 +

14 +model:

15 + name: "deepseek-v3"

16 + num_layers: 60

17 + num_q_heads: 16

18 + num_kv_heads: 1

19 + head_dim: 576

20 + kv_lora_rank: 512

21 + qk_nope_head_dim: 128

22 + qk_rope_head_dim: 64

23 + v_head_dim: 128

24 + block_size: 128

25 + max_model_len: 262144

26 +

27 +batch_specs:

28 + - "q32"

29 + - "q36"

30 + - "q40"

31 + - "q45"

32 + - "q50"

33 + - "q56"

34 + - "q62"

35 + - "q70"

36 + - "q78"

37 + - "q87"

38 + - "q97"

39 + - "q109"

40 + - "q122"

41 + - "q136"

42 + - "q152"

43 + - "q170"

44 + - "q190"

45 + - "q212"

46 + - "q237"

47 + - "q255"

48 + - "q256"

49 + - "q257"

50 + - "q265"

51 + - "q296"

52 + - "q331"

53 + - "q370"

54 + - "q413"

55 + - "q462"

56 + - "q516"

57 + - "q577"

58 + - "q645"

59 + - "q721"

60 + - "q806"

61 + - "q901"

62 + - "q1007"

63 + - "q1125"

64 + - "q1257"

65 + - "q1405"

66 + - "q1571"

67 + - "q1756"

68 + - "q1962"

69 + - "q2048"

70 + - "q2049"

71 + - "q2193"

72 + - "q2451"

73 + - "q2740"

74 + - "q3062"

75 + - "q3422"

76 + - "q3825"

77 + - "q4275"

78 + - "q4778"

79 + - "q5340"

80 + - "q5854"

81 + - "q5969"

82 + - "q6000"

83 + - "q6144"

84 + - "q6367"

85 + - "q6671"

86 + - "q6925"

87 + - "q7168"

88 + - "q7456"

89 + - "q7532"

90 + - "q8192"

91 + - "q8193"

92 + - "q8333"

93 + - "q9314"

94 + - "q10410"

95 + - "q11635"

96 + - "q13004"

97 + - "q14534"

98 + - "q16244"

99 + - "q18156"

100 + - "q20292"

101 + - "q22680"

102 + - "q25349"

103 + - "q28332"

104 + - "q31665"

105 + - "q32768"

106 + - "q32769"

107 + - "q35391"

108 + - "q39556"

109 + - "q44210"

110 + - "q49413"

111 + - "q55227"

112 + - "q61726"

113 + - "q65536"

114 +

115 +backends:

116 + - FLASHMLA_SPARSE

117 +

118 +device: "cuda:0"

119 +profile_memory: true

120 +sparse_mla_dense_mha_max_seq_len: 25349

121 +sparse_mla_masked_mha_max_seq_len: 8192

122 +sparse_mla_mha_variants: ["dense_mha", "masked_mha", "mqa"]

123 +sparse_mla_topk_pattern: "random"

benchmarks/attention_benchmarks/mla_runner.py

@@ -931,10 +931,17 @@ def _run_single_benchmark(

931 prefill_quant_op = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)

932

933 fused_output = output_scale is not None and fuse_quant_op

934 + mha_mode = getattr(config, "sparse_mla_mha_mode", "auto")

935

936 # Build forward function (runs a single decode/prefill pass)

937 def forward_fn():

938 results = []

939 + impl._sparse_mla_force_masked_mha = ( # type: ignore[attr-defined]

940 + mha_mode == "masked"

941 + )

942 + impl._sparse_mla_force_dense_mha = ( # type: ignore[attr-defined]

943 + mha_mode == "dense"

944 + )

945 if has_decode:

946 results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer))

947 if has_prefill:

cmake/external_projects/vllm_flash_attn.cmake

+1 -1

@@ -39,7 +39,7 @@ else()

39 FetchContent_Declare(

40 vllm-flash-attn

41 GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git

42 - GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32

42 + GIT_TAG ec4a882d70b19235c824848db944d84a6cb8f9ee

43 GIT_PROGRESS TRUE

44 # Don't share the vllm-flash-attn build between build types

45 BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

tests/v1/attention/test_sparse_mla_backends.py

+26 -6

@@ -862,6 +862,8 @@ def test_split_prefill_chunks(seq_lens, max_buf, expected):

862 PREFILL_BATCH_SPECS = {

863 "short_dense_mha": BatchSpec(seq_lens=[64, 128], query_lens=[64, 128]),

864 "short_context_dense_mha": BatchSpec(seq_lens=[128, 160], query_lens=[64, 32]),

865 + "masked_mha": BatchSpec(seq_lens=[256], query_lens=[256]),

866 + "masked_mha_chunked_context": BatchSpec(seq_lens=[448, 384], query_lens=[256, 256]),

867 }

868

869

@@ -878,7 +880,7 @@ def test_sparse_backend_prefill_correctness(

880 kv_cache_dtype,

881 workspace_init,

882 ):

881 - """Test single-pass FA4 dense forward_mha for sparse MLA prefill."""

883 + """Test single-pass dense and masked MHA for sparse MLA prefill."""

884 backend_cls = FlashMLASparseBackend

885 batch_spec = PREFILL_BATCH_SPECS[batch_name]

886

@@ -892,7 +894,8 @@ def test_sparse_backend_prefill_correctness(

894 qk_rope_head_dim = 64

895 v_head_dim = 128

896 head_size = kv_lora_rank + qk_rope_head_dim

895 - topk_tokens = 512

897 + masked_mha = batch_name.startswith("masked_mha")

898 + topk_tokens = 200 if masked_mha else 512

899

900 max_seqlen = max(batch_spec.seq_lens)

901 total_cache_tokens = sum(batch_spec.seq_lens)

@@ -918,6 +921,7 @@ def test_sparse_backend_prefill_correctness(

921 model_type="deepseek_v2",

922 )

923 model_config.dtype = dtype

924 + model_config.model_arch_config.total_num_attention_heads = num_heads

925 model_config.get_num_attention_heads = MethodType(

926 lambda self, parallel_config: num_heads, model_config

927 )

@@ -942,13 +946,14 @@ def test_sparse_backend_prefill_correctness(

946

947 # Compute dense reference outputs.

948 total_query_tokens = sum(query_lens)

945 - sparse_indices = torch.zeros(

946 - total_query_tokens, topk_tokens, dtype=torch.int32, device=device

949 + sparse_indices = torch.full(

950 + (total_query_tokens, topk_tokens), -1, dtype=torch.int32, device=device

951 )

952

953 all_q, all_kv_c_new, all_k_pe_new = [], [], []

954 kv_c_contexts, k_pe_contexts = [], []

955 reference_outputs = []

956 + global_token_idx = 0

957

958 for i in range(batch_spec.batch_size):

959 s_len = seq_lens[i]

@@ -981,8 +986,15 @@ def test_sparse_backend_prefill_correctness(

986 for j in range(q_len):

987 attend_end = ctx_len + j + 1

988 q_tok = q_mha[j : j + 1] # (1, H, D_qk)

984 - k_attend = k_all[:attend_end] # (N, H, D_qk)

985 - v_attend = v_all[:attend_end] # (N, H, D_v)

989 + if masked_mha:

990 + actual_topk = min(topk_tokens, attend_end)

991 + attend_indices = torch.randperm(attend_end, device=device)[:actual_topk]

992 + sparse_indices[global_token_idx, :actual_topk] = attend_indices

993 + k_attend = k_all[attend_indices]

994 + v_attend = v_all[attend_indices]

995 + else:

996 + k_attend = k_all[:attend_end] # (N, H, D_qk)

997 + v_attend = v_all[:attend_end] # (N, H, D_v)

998

999 q_sdpa = q_tok.unsqueeze(0).transpose(1, 2).float()

1000 k_sdpa = k_attend.unsqueeze(0).transpose(1, 2).float()

@@ -993,6 +1005,7 @@ def test_sparse_backend_prefill_correctness(

1005 )

1006 out = out.transpose(1, 2).squeeze(0) # (1, H, D_v)

1007 reference_outputs.append(out.to(dtype).flatten(start_dim=-2))

1008 + global_token_idx += 1

1009

1010 all_q.append(q_mha)

1011 all_kv_c_new.append(kv_c_full[ctx_len:])

@@ -1047,6 +1060,13 @@ def test_sparse_backend_prefill_correctness(

1060

1061 builder_cls = backend_cls.get_builder_cls()

1062 builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)

1063 + if batch_name == "masked_mha_chunked_context":

1064 + builder.chunked_prefill_workspace_size = block_size * batch_spec.batch_size

1065 + builder.chunked_prefill_workspace = torch.empty(

1066 + (builder.chunked_prefill_workspace_size, head_size),

1067 + dtype=dtype,

1068 + device=device,

1069 + )

1070 # Drive the queries through the dense-MHA prefill path directly (the routing

1071 # threshold would otherwise classify these short queries as MQA decodes).

1072 builder.reorder_batch_threshold = 1

tests/v1/attention/test_sparse_mla_mask.py new

+29

@@ -0,0 +1,29 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +import pytest

5 +import torch

6 +

7 +pytest.importorskip("cutlass")

8 +

9 +from vllm.model_executor.layers.attention.sparse_mla_attention import (

10 + _build_topk_mask,

11 +)

12 +

13 +

14 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")

15 +def test_build_topk_mask_single_request_matches_generic_path() -> None:

16 + topk = torch.tensor(

17 + [[0, 31, 32, 63, -1], [1, 64, 127, -1, -1]],

18 + dtype=torch.int32,

19 + device="cuda",

20 + )

21 +

22 + num_words = (128 + 31) // 32

23 + single_out = torch.zeros(1, 2, num_words, dtype=torch.int32, device="cuda")

24 + generic_out = torch.zeros(2, 1, num_words, dtype=torch.int32, device="cuda")

25 + single_req = _build_topk_mask([topk], [2], 2, 128, single_out)

26 + generic = _build_topk_mask([topk[:1], topk[1:]], [1, 1], 1, 128, generic_out)

27 +

28 + torch.testing.assert_close(single_req[0, 0], generic[0, 0])

29 + torch.testing.assert_close(single_req[0, 1], generic[1, 0])

vllm/model_executor/layers/attention/mla_attention.py

+62 -5

@@ -772,8 +772,25 @@ class MLAAttention(nn.Module, AttentionLayerBase):

772 num_mha_tokens = q.size(0) - num_mqa_tokens

773

774 if self.impl.is_sparse and num_mha_tokens > 0:

775 - prefill_metadata = getattr(attn_metadata, "prefill", None)

776 - if not getattr(prefill_metadata, "use_dense_mha", False):

775 + prefill = getattr(attn_metadata, "prefill", None)

776 + use_dense_mha = getattr(prefill, "use_dense_mha", False)

777 + prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined]

778 + use_masked_mha = (

779 + self.prefill_backend is not None

780 + and self.impl.masked_mha_available # type: ignore[attr-defined]

781 + and self.impl.dcp_world_size <= 1

782 + and prefill is not None

783 + and _use_masked_mha(

784 + backend_name=self.attn_backend.get_name(),

785 + tensor_parallel_size=self._vllm_config.parallel_config.tensor_parallel_size,

786 + query_len=prefill.max_query_len,

787 + seq_len=prefill_max_seq_len,

788 + )

789 + )

790 + use_mha = (use_dense_mha or use_masked_mha) and not (

791 + self._vllm_config.attention_config.sparse_mla_force_mqa

792 + )

793 + if not use_mha:

794 num_mqa_tokens = q.size(0)

795 num_mha_tokens = 0

796

@@ -781,6 +798,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):

798 quant_key is not None

799 and self.prefill_backend is not None

800 and self.prefill_backend.supports_quant_output(quant_key)

801 + and (

802 + not self.impl.is_sparse

803 + or attn_metadata.prefill_max_seq_len # type: ignore[attr-defined]

804 + <= attn_metadata.topk_tokens # type: ignore[attr-defined]

805 + )

806 and attn_metadata is not None

807 and attn_metadata.prefill is not None

808 and attn_metadata.prefill.chunked_context is None

@@ -1387,6 +1409,7 @@ class MLACommonPrefillMetadata:

1409 seq_tot: list[int]

1410 max_seq_lens: list[int]

1411 seq_lens: torch.Tensor

1412 + context_lens: torch.Tensor

1413 workspace: torch.Tensor

1414 token_to_seq: torch.Tensor

1415 chunk_total_token: list[int]

@@ -1408,10 +1431,9 @@ class MLACommonPrefillMetadata:

1431 q_data_type: torch.dtype | None = None

1432 output_dtype: torch.dtype | None = None

1433 prefill_backend: MLAPrefillBackend | None = None

1411 - # Whether the prefill suffix is routed through dense MHA.

1412 - # Indexer scoring may be skipped only for a pure-prefill batch,

1413 - # since decode tokens still consume top-k indices.

1434 + query_lens_cpu: torch.Tensor | None = None

1435 use_dense_mha: bool = False

1436 + topk_mask_workspace: torch.Tensor | None = None

1437

1438

1439 @dataclass

@@ -1508,6 +1530,39 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims:

1530 )

1531

1532

1533 +_DSV32_MASKED_MHA_THRESHOLDS: dict[str, dict[int, tuple[int | None, ...]]] = {

1534 + "FLASHMLA_SPARSE": {

1535 + 1: (1536, 4096, None, None, None),

1536 + 2: (512, 1024, 4096, None, None),

1537 + 4: (512, 1024, 1536, 8192, None),

1538 + 8: (512, 512, 1024, 2048, 16384),

1539 + },

1540 + "FLASHINFER_MLA_SPARSE": {

1541 + 8: (512, 1024, 1024, 2048, 32768),

1542 + },

1543 +}

1544 +_DSV32_SEQ_LEN_BUCKETS = (2048, 4096, 8192, 16384, 32768)

1545 +

1546 +

1547 +def _use_masked_mha(

1548 + *,

1549 + backend_name: str,

1550 + tensor_parallel_size: int,

1551 + query_len: int,

1552 + seq_len: int,

1553 +) -> bool:

1554 + thresholds = _DSV32_MASKED_MHA_THRESHOLDS.get(backend_name, {}).get(

1555 + tensor_parallel_size

1556 + )

1557 + if thresholds is None:

1558 + return False

1559 + for bucket_idx, bucket_seq_len in enumerate(_DSV32_SEQ_LEN_BUCKETS):

1560 + if seq_len <= bucket_seq_len:

1561 + min_query_len = thresholds[bucket_idx]

1562 + return min_query_len is not None and query_len >= min_query_len

1563 + return False

1564 +

1565 +

1566 @functools.cache

1567 def backend_supports_prefill_query_quantization() -> bool:

1568 """Check if the selected MLA prefill backend supports query quantization.

@@ -1682,6 +1737,7 @@ def build_mla_chunked_context_metadata(

1737 seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),

1738 max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

1739 seq_lens=chunk_seq_lens,

1740 + context_lens=context_lens_cpu.to(device, non_blocking=True),

1741 token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),

1742 chunk_total_token=chunk_total_token.tolist(),

1743 workspace=chunked_prefill_workspace,

@@ -1705,6 +1761,7 @@ def build_mla_chunked_context_metadata(

1761 seq_tot=chunk_seq_lens.sum(dim=1).tolist(),

1762 max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

1763 seq_lens=chunk_seq_lens,

1764 + context_lens=context_lens_cpu.to(device, non_blocking=True),

1765 token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),

1766 chunk_total_token=chunk_total_token,

1767 workspace=chunked_prefill_workspace,

vllm/model_executor/layers/attention/sparse_mla_attention.py

+519 -12

@@ -1,29 +1,34 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -"""Shared forward_mha implementation and metadata builder for sparse MLA backends."""

3 +"""Shared MHA implementation and metadata builder for sparse MLA backends."""

5 from shutil import which

6 -from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar

6 +from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast

8 import numpy as np

9 import torch

11 -from vllm.distributed import get_dcp_group

11 +from vllm import _custom_ops as ops

12 +from vllm.distributed import (

13 + get_dcp_group,

14 + get_tensor_model_parallel_world_size,

15 +)

16 from vllm.logger import init_logger

17 from vllm.model_executor.layers.attention.mla_attention import (

18 MLACommonBaseImpl,

19 + MLACommonMetadata,

20 MLACommonPrefillMetadata,

21 build_mla_chunked_context_metadata,

22 get_mla_dims,

23 )

24 from vllm.platforms import current_platform

25 +from vllm.triton_utils import tl, triton

26 from vllm.utils.flashinfer import has_flashinfer

21 -from vllm.utils.torch_utils import np_to_pinned_tensor

22 -from vllm.v1.attention.backend import (

23 - AttentionMetadata,

24 - AttentionMetadataBuilder,

25 -)

27 +from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor

28 +from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder

29 +from vllm.v1.attention.backends.fa_utils import get_flash_attn_version

30 from vllm.v1.attention.backends.utils import split_decodes_and_prefills

31 +from vllm.v1.attention.ops.merge_attn_states import merge_attn_states

33 if TYPE_CHECKING:

34 from vllm.config import VllmConfig

@@ -35,6 +40,32 @@ logger = init_logger(__name__)

41 T = TypeVar("T", bound=AttentionMetadata)

43 +GLOBAL_TOPK_MASK_MAX_BYTES = 64 * 1024 * 1024 # 64 MiB

44 +

45 +

46 +def _is_masked_mha_available(

47 + num_heads_total: int,

48 + kv_lora_rank: int,

49 + qk_nope_head_dim: int,

50 + qk_rope_head_dim: int,

51 + v_head_dim: int,

52 + kv_cache_dtype: str,

53 +) -> bool:

54 + """Check if masked MHA can ever fire for this model configuration."""

55 + if not current_platform.is_device_capability_family(100):

56 + return False

57 + if (

58 + num_heads_total != 128

59 + or kv_lora_rank != 512

60 + or qk_nope_head_dim != 128

61 + or qk_rope_head_dim != 64

62 + or v_head_dim != 128

63 + ):

64 + return False

65 + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim

66 + fa_version = get_flash_attn_version(head_size=qk_head_dim, head_size_v=v_head_dim)

67 + return fa_version == 4 and not is_quantized_kv_cache(kv_cache_dtype)

68 +

70 class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):

71 metadata_cls: type[T]

@@ -86,6 +117,20 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):

117 dtype=self.model_config.dtype,

118 device=device,

119 )

120 + self.topk_mask_workspace: torch.Tensor | None = None

121 + if _is_masked_mha_available(

122 + self.model_config.model_arch_config.total_num_attention_heads,

123 + self.mla_dims.kv_lora_rank,

124 + self.mla_dims.qk_nope_head_dim,

125 + self.mla_dims.qk_rope_head_dim,

126 + self.mla_dims.v_head_dim,

127 + vllm_config.cache_config.cache_dtype,

128 + ):

129 + self.topk_mask_workspace = torch.zeros(

130 + GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize,

131 + dtype=torch.int32,

132 + device=device,

133 + )

134 layer_prefill_backend = vllm_config.compilation_config.static_forward_context[

135 layer_names[0]

136 ].prefill_backend

@@ -194,6 +239,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):

239 block_table=common_attn_metadata.block_table_tensor[num_decodes:, ...],

240 query_start_loc=prefill_query_start_loc,

241 max_query_len=prefill_max_query_len,

242 + query_lens_cpu=prefill_query_lens_cpu,

243 chunked_context=self._build_chunked_context_fields(

244 common_attn_metadata,

245 num_decodes,

@@ -207,6 +253,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):

253 prefill_max_seq_len <= self.topk_tokens

254 and not self.vllm_config.attention_config.sparse_mla_force_mqa

255 ),

256 + topk_mask_workspace=self.topk_mask_workspace,

257 )

258 self._prefill_backend.prepare_metadata(prefill)

259

@@ -260,9 +307,139 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):

307 )

308

309

310 +@triton.jit

311 +def _scatter_topk_kernel(

312 + mask_ptr,

313 + topk_ptr,

314 + cu_q_lens_ptr,

315 + num_words: tl.constexpr,

316 + mask_row_stride: tl.constexpr,

317 + num_topk: tl.constexpr,

318 + topk_stride: tl.constexpr,

319 + max_q_len: tl.constexpr,

320 + BLOCK_TOPK: tl.constexpr,

321 + BLOCK_WORDS: tl.constexpr,

322 +):

323 + row_idx = tl.program_id(0)

324 + row_ptr = mask_ptr + row_idx * mask_row_stride

325 + word_offsets = tl.arange(0, BLOCK_WORDS)

326 + tl.store(

327 + row_ptr + word_offsets,

328 + tl.zeros([BLOCK_WORDS], dtype=tl.int32),

329 + mask=word_offsets < num_words,

330 + )

331 +

332 + req_idx = row_idx // max_q_len

333 + q_local = row_idx % max_q_len

334 + q_start = tl.load(cu_q_lens_ptr + req_idx)

335 + q_len = tl.load(cu_q_lens_ptr + req_idx + 1) - q_start

336 + if q_local < q_len:

337 + src_row = q_start + q_local

338 + offsets = tl.arange(0, BLOCK_TOPK)

339 + in_range = offsets < num_topk

340 + indices = tl.load(

341 + topk_ptr + src_row * topk_stride + offsets,

342 + mask=in_range,

343 + other=-1,

344 + )

345 + valid = in_range & (indices >= 0)

346 + word_indices = indices >> 5

347 + bits = (1 << (indices & 31)).to(tl.int32)

348 + tl.atomic_or(row_ptr + word_indices, bits, mask=valid)

349 +

350 +

351 +@triton.jit

352 +def _scatter_topk_single_req_kernel(

353 + mask_ptr,

354 + topk_ptr,

355 + num_words: tl.constexpr,

356 + mask_row_stride: tl.constexpr,

357 + num_topk: tl.constexpr,

358 + topk_stride: tl.constexpr,

359 + total_q: tl.constexpr,

360 + BLOCK_TOPK: tl.constexpr,

361 + BLOCK_WORDS: tl.constexpr,

362 +):

363 + row_idx = tl.program_id(0)

364 + row_ptr = mask_ptr + row_idx * mask_row_stride

365 + word_offsets = tl.arange(0, BLOCK_WORDS)

366 + tl.store(

367 + row_ptr + word_offsets,

368 + tl.zeros([BLOCK_WORDS], dtype=tl.int32),

369 + mask=word_offsets < num_words,

370 + )

371 + if row_idx < total_q:

372 + offsets = tl.arange(0, BLOCK_TOPK)

373 + in_range = offsets < num_topk

374 + indices = tl.load(

375 + topk_ptr + row_idx * topk_stride + offsets,

376 + mask=in_range,

377 + other=-1,

378 + )

379 + valid = in_range & (indices >= 0)

380 + word_indices = indices >> 5

381 + bits = (1 << (indices & 31)).to(tl.int32)

382 + tl.atomic_or(row_ptr + word_indices, bits, mask=valid)

383 +

384 +

385 +def _build_topk_mask(

386 + topk_indices_per_req: list[torch.Tensor],

387 + q_lens: list[int],

388 + max_q_len: int,

389 + max_seq_len: int,

390 + out: torch.Tensor,

391 +) -> torch.Tensor:

392 + """Build a bit-packed top-k mask into ``out[:B, :max_Q, :num_words]``."""

393 + batch_size = len(q_lens)

394 + num_words = (max_seq_len + 31) // 32

395 + total_rows = batch_size * max_q_len

396 + if total_rows == 0:

397 + return out[:batch_size, :max_q_len, :num_words]

398 +

399 + total_q = sum(q_lens)

400 + mask_row_stride = out.stride(-2)

401 + block_words = triton.next_power_of_2(num_words)

402 +

403 + if batch_size == 1:

404 + topk_packed = topk_indices_per_req[0]

405 + num_topk = topk_packed.shape[1]

406 + _scatter_topk_single_req_kernel[(max_q_len,)](

407 + out,

408 + topk_packed,

409 + num_words=num_words,

410 + mask_row_stride=mask_row_stride,

411 + num_topk=num_topk,

412 + topk_stride=topk_packed.stride(0),

413 + total_q=total_q,

414 + BLOCK_TOPK=triton.next_power_of_2(num_topk),

415 + BLOCK_WORDS=block_words,

416 + )

417 + return out[:1, :max_q_len, :num_words]

418 +

419 + topk_packed = torch.cat(topk_indices_per_req, dim=0)

420 + num_topk = topk_packed.shape[1]

421 + q_lens_tensor = np_to_pinned_tensor(np.asarray(q_lens, dtype=np.int32)).to(

422 + out.device, non_blocking=True

423 + )

424 + cu_q_lens = out.new_zeros(batch_size + 1)

425 + torch.cumsum(q_lens_tensor, dim=0, out=cu_q_lens[1:])

426 + _scatter_topk_kernel[(total_rows,)](

427 + out,

428 + topk_packed,

429 + cu_q_lens,

430 + num_words=num_words,

431 + mask_row_stride=mask_row_stride,

432 + num_topk=num_topk,

433 + topk_stride=topk_packed.stride(0),

434 + max_q_len=max_q_len,

435 + BLOCK_TOPK=triton.next_power_of_2(num_topk),

436 + BLOCK_WORDS=block_words,

437 + )

438 + return out[:batch_size, :max_q_len, :num_words]

439 +

440 +

441 class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):

264 - """Sparse MLA base: shared dense-MHA prefill (from MLACommonBaseImpl) plus the

265 - sparse top-k MQA decode path. Subclasses implement forward_mqa."""

442 + """Sparse MLA base with dense and masked-MHA prefill paths."""

443

444 is_sparse = True

445

@@ -278,7 +455,6 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):

455 logits_soft_cap: float | None,

456 attn_type: str,

457 kv_sharing_target_layer_name: str | None,

281 - # MLA-specific

458 q_lora_rank: int | None,

459 kv_lora_rank: int,

460 qk_nope_head_dim: int,

@@ -312,7 +488,6 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):

488 if indexer is not None

489 else topk_indices_buffer

490 )

315 -

491 self._use_flashinfer_concat_mla_k = (

492 has_flashinfer()

493 and which("ninja") is not None

@@ -320,3 +495,335 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):

495 and (self.qk_nope_head_dim == 128)

496 and (self.qk_rope_head_dim == 64)

497 )

498 + self.masked_mha_available = _is_masked_mha_available(

499 + num_heads_total=num_heads * get_tensor_model_parallel_world_size(),

500 + kv_lora_rank=kv_lora_rank,

501 + qk_nope_head_dim=qk_nope_head_dim,

502 + qk_rope_head_dim=qk_rope_head_dim,

503 + v_head_dim=v_head_dim,

504 + kv_cache_dtype=kv_cache_dtype,

505 + )

506 +

507 + @staticmethod

508 + def _slice_topk_per_req(

509 + topk_all: torch.Tensor,

510 + q_lens: list[int],

511 + ) -> list[torch.Tensor]:

512 + topk_per_req = []

513 + offset = 0

514 + for q_len in q_lens:

515 + topk_per_req.append(topk_all[offset : offset + q_len])

516 + offset += q_len

517 + return topk_per_req

518 +

519 + @staticmethod

520 + def _remap_topk_to_ranges(

521 + topk_per_req: list[torch.Tensor],

522 + range_starts: list[int] | torch.Tensor,

523 + range_lens: list[int],

524 + ) -> list[torch.Tensor]:

525 + remapped = []

526 + for topk, start, length in zip(topk_per_req, range_starts, range_lens):

527 + valid = (topk >= start) & (topk < start + length)

528 + remapped.append(torch.where(valid, topk - start, -1))

529 + return remapped

530 +

531 + def _project_kv(

532 + self, kv_c_normed: torch.Tensor, k_pe: torch.Tensor

533 + ) -> tuple[torch.Tensor, torch.Tensor]:

534 + kv_nope = self.kv_b_proj(kv_c_normed)[0].view(

535 + -1,

536 + self.num_heads,

537 + self.qk_nope_head_dim + self.v_head_dim,

538 + )

539 + k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

540 + return self._concat_k_nope_k_pe(k_nope, k_pe), v

541 +

542 + @staticmethod

543 + def _try_build_global_mask(

544 + topk_per_req: list[torch.Tensor],

545 + q_lens: list[int],

546 + max_query_len: int,

547 + max_seq_len: int,

548 + topk_mask_workspace: torch.Tensor,

549 + ) -> torch.Tensor | None:

550 + """Build a full-sequence top-k mask if it fits within the budget.

551 +

552 + When the mask fits, it is reused across the suffix and all context

553 + chunks, avoiding per-chunk mask rebuilds. Returns None when the

554 + mask is too large, signalling the caller to fall back to per-chunk

555 + index remapping.

556 + """

557 + batch_size = len(q_lens)

558 + tile_m = 128 if max_query_len <= 128 else 256

559 + padded_q_len = triton.cdiv(max_query_len, tile_m) * tile_m

560 + num_words_padded = (max_seq_len + 31) // 32 + 1

561 + needed = batch_size * padded_q_len * num_words_padded

562 + if needed * torch.int32.itemsize > GLOBAL_TOPK_MASK_MAX_BYTES:

563 + return None

564 +

565 + mask = topk_mask_workspace[:needed].view(

566 + batch_size, padded_q_len, num_words_padded

567 + )

568 + _build_topk_mask(

569 + topk_per_req,

570 + q_lens,

571 + padded_q_len,

572 + max_seq_len,

573 + mask,

574 + )

575 + return mask

576 +

577 + def _run_masked_mha(

578 + self,

579 + q: torch.Tensor,

580 + k: torch.Tensor,

581 + v: torch.Tensor,

582 + cu_seqlens_q: torch.Tensor,

583 + cu_seqlens_k: torch.Tensor,

584 + max_seqlen_q: int,

585 + max_seqlen_k: int,

586 + topk_per_req: list[torch.Tensor],

587 + q_lens: list[int],

588 + causal: bool,

589 + return_softmax_lse: bool = False,

590 + dense_mask: torch.Tensor | None = None,

591 + key_starts: torch.Tensor | None = None,

592 + topk_mask_workspace: torch.Tensor | None = None,

593 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

594 + from vllm.model_executor.layers.attention.sparse_mla_mask import (

595 + dense_mask_mod,

596 + offset_dense_mask_mod,

597 + )

598 + from vllm.vllm_flash_attn import flash_attn_varlen_func

599 +

600 + tile_m = 128 if max_seqlen_q <= 128 else 256

601 + padded_q_len = triton.cdiv(max_seqlen_q, tile_m) * tile_m

602 + if dense_mask is None:

603 + batch_size = len(q_lens)

604 + num_words = (max_seqlen_k + 31) // 32

605 + assert topk_mask_workspace is not None

606 + workspace_3d = topk_mask_workspace[

607 + : batch_size * padded_q_len * num_words

608 + ].view(batch_size, padded_q_len, num_words)

609 + dense_mask = _build_topk_mask(

610 + topk_per_req,

611 + q_lens,

612 + padded_q_len,

613 + max_seqlen_k,

614 + workspace_3d,

615 + )

616 + if key_starts is not None:

617 + dense_mask[:, 0, -1].copy_(key_starts)

618 + kwargs = {

619 + "q": q,

620 + "k": k,

621 + "v": v,

622 + "cu_seqlens_q": cu_seqlens_q,

623 + "cu_seqlens_k": cu_seqlens_k,

624 + "max_seqlen_q": max_seqlen_q,

625 + "max_seqlen_k": max_seqlen_k,

626 + "softmax_scale": self.scale,

627 + "return_softmax_lse": return_softmax_lse,

628 + "fa_version": 4,

629 + "mask_mod": dense_mask_mod if key_starts is None else offset_dense_mask_mod,

630 + "aux_tensors": [dense_mask],

631 + "aux_tensor_leading_dims": [2],

632 + "causal": causal,

633 + }

634 +

635 + return flash_attn_varlen_func(**kwargs)

636 +

637 + def _compute_context_mha(

638 + self,

639 + q: torch.Tensor,

640 + kv_c_and_k_pe_cache: torch.Tensor,

641 + prefill_metadata: MLACommonPrefillMetadata,

642 + k_scale: torch.Tensor,

643 + q_lens: list[int],

644 + topk_per_req: list[torch.Tensor],

645 + dense_mask: torch.Tensor | None = None,

646 + ) -> tuple[torch.Tensor, torch.Tensor]:

647 + if self.dcp_world_size > 1:

648 + raise NotImplementedError(

649 + "Masked MHA with context does not yet support decode context "

650 + "parallelism"

651 + )

652 +

653 + chunked_context = prefill_metadata.chunked_context

654 + assert chunked_context is not None

655 + use_global_mask = dense_mask is not None

656 + output: torch.Tensor | None = None

657 + output_lse: torch.Tensor | None = None

658 + workspace = chunked_context.workspace

659 +

660 + for i, toks in enumerate(chunked_context.seq_tot):

661 + if toks == 0:

662 + continue

663 + ops.gather_and_maybe_dequant_cache(

664 + src_cache=kv_c_and_k_pe_cache,

665 + dst=workspace,

666 + block_table=prefill_metadata.block_table,

667 + cu_seq_lens=chunked_context.cu_seq_lens[i],

668 + token_to_seq=chunked_context.token_to_seq[i],

669 + num_tokens=chunked_context.chunk_total_token[i],

670 + kv_cache_dtype=self.kv_cache_dtype,

671 + scale=k_scale,

672 + seq_starts=chunked_context.starts[i],

673 + )

674 +

675 + chunk_kv_c = workspace[:toks, : self.kv_lora_rank]

676 + chunk_k_pe = workspace[:toks, self.kv_lora_rank :].unsqueeze(1)

677 + k, v = self._project_kv(chunk_kv_c, chunk_k_pe)

678 + chunk_lens = chunked_context.seq_lens[i].tolist()

679 + chunk_topk = (

680 + topk_per_req

681 + if use_global_mask

682 + else self._remap_topk_to_ranges(

683 + topk_per_req,

684 + chunked_context.starts[i],

685 + chunk_lens,

686 + )

687 + )

688 + attn_out, lse = self._run_masked_mha(

689 + q=q,

690 + k=k,

691 + v=v,

692 + cu_seqlens_q=prefill_metadata.query_start_loc,

693 + cu_seqlens_k=chunked_context.cu_seq_lens[i],

694 + max_seqlen_q=prefill_metadata.max_query_len,

695 + max_seqlen_k=chunked_context.max_seq_lens[i],

696 + topk_per_req=chunk_topk,

697 + q_lens=q_lens,

698 + causal=False,

699 + return_softmax_lse=True,

700 + dense_mask=dense_mask,

701 + key_starts=(chunked_context.starts[i] if use_global_mask else None),

702 + topk_mask_workspace=prefill_metadata.topk_mask_workspace,

703 + )

704 +

705 + if output is None:

706 + output = attn_out

707 + output_lse = lse

708 + else:

709 + assert output_lse is not None

710 + merge_attn_states(

711 + output=output,

712 + output_lse=output_lse,

713 + prefix_output=output,

714 + prefix_lse=output_lse,

715 + suffix_output=attn_out,

716 + suffix_lse=lse,

717 + )

718 +

719 + assert output is not None and output_lse is not None

720 + return output, output_lse

721 +

722 + def forward_mha( # type: ignore[override]

723 + self,

724 + q: torch.Tensor,

725 + kv_c_normed: torch.Tensor,

726 + k_pe: torch.Tensor,

727 + kv_c_and_k_pe_cache: torch.Tensor,

728 + attn_metadata: T,

729 + k_scale: torch.Tensor,

730 + output: torch.Tensor,

731 + output_scale: torch.Tensor | None = None,

732 + ) -> None:

733 + prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined]

734 + topk_tokens = attn_metadata.topk_tokens # type: ignore[attr-defined]

735 + force_dense = getattr(self, "_sparse_mla_force_dense_mha", False)

736 + force_masked = getattr(self, "_sparse_mla_force_masked_mha", False)

737 + if force_dense or (prefill_max_seq_len <= topk_tokens and not force_masked):

738 + return super().forward_mha(

739 + q,

740 + kv_c_normed,

741 + k_pe,

742 + kv_c_and_k_pe_cache,

743 + cast(MLACommonMetadata, attn_metadata),

744 + k_scale,

745 + output,

746 + output_scale,

747 + )

748 +

749 + assert output_scale is None

750 + assert self.masked_mha_available

751 + prefill_metadata = attn_metadata.prefill # type: ignore[attr-defined]

752 + assert prefill_metadata is not None

753 + assert prefill_metadata.query_lens_cpu is not None

754 + assert self.topk_indices_buffer is not None

755 +

756 + q_lens = prefill_metadata.query_lens_cpu.tolist()

757 + num_decode_tokens = attn_metadata.num_decode_tokens # type: ignore[attr-defined]

758 + topk_all = self.topk_indices_buffer[

759 + num_decode_tokens : num_decode_tokens + q.shape[0]

760 + ]

761 + topk_per_req = self._slice_topk_per_req(topk_all, q_lens)

762 +

763 + k, v = self._project_kv(kv_c_normed, k_pe)

764 + chunked_context = prefill_metadata.chunked_context

765 + if chunked_context is None:

766 + attn_out = self._run_masked_mha(

767 + q=q,

768 + k=k,

769 + v=v,

770 + cu_seqlens_q=prefill_metadata.query_start_loc,

771 + cu_seqlens_k=prefill_metadata.query_start_loc,

772 + max_seqlen_q=prefill_metadata.max_query_len,

773 + max_seqlen_k=prefill_metadata.max_query_len,

774 + topk_per_req=topk_per_req,

775 + q_lens=q_lens,

776 + causal=True,

777 + topk_mask_workspace=prefill_metadata.topk_mask_workspace,

778 + )

779 + assert isinstance(attn_out, torch.Tensor)

780 + output.copy_(attn_out[..., : self.v_head_dim].flatten(start_dim=-2))

781 + return

782 +

783 + context_lens = chunked_context.seq_lens.sum(dim=0).tolist()

784 + dense_mask = self._try_build_global_mask(

785 + topk_per_req,

786 + q_lens,

787 + prefill_metadata.max_query_len,

788 + prefill_max_seq_len,

789 + prefill_metadata.topk_mask_workspace,

790 + )

791 + if dense_mask is not None:

792 + suffix_topk = topk_per_req

793 + else:

794 + suffix_topk = self._remap_topk_to_ranges(topk_per_req, context_lens, q_lens)

795 + suffix_output, suffix_lse = self._run_masked_mha(

796 + q=q,

797 + k=k,

798 + v=v,

799 + cu_seqlens_q=prefill_metadata.query_start_loc,

800 + cu_seqlens_k=prefill_metadata.query_start_loc,

801 + max_seqlen_q=prefill_metadata.max_query_len,

802 + max_seqlen_k=prefill_metadata.max_query_len,

803 + topk_per_req=suffix_topk,

804 + q_lens=q_lens,

805 + causal=True,

806 + return_softmax_lse=True,

807 + dense_mask=dense_mask,

808 + key_starts=(

809 + chunked_context.context_lens if dense_mask is not None else None

810 + ),

811 + topk_mask_workspace=prefill_metadata.topk_mask_workspace,

812 + )

813 + context_output, context_lse = self._compute_context_mha(

814 + q=q,

815 + kv_c_and_k_pe_cache=kv_c_and_k_pe_cache,

816 + prefill_metadata=prefill_metadata,

817 + k_scale=k_scale,

818 + q_lens=q_lens,

819 + topk_per_req=topk_per_req,

820 + dense_mask=dense_mask,

821 + )

822 + merge_attn_states(

823 + output=output.view(-1, self.num_heads, self.v_head_dim),

824 + prefix_output=context_output[..., : self.v_head_dim],

825 + prefix_lse=context_lse,

826 + suffix_output=suffix_output[..., : self.v_head_dim],

827 + suffix_lse=suffix_lse,

828 + prefill_tokens_with_context=chunked_context.prefill_tokens_with_context,

829 + )

vllm/model_executor/layers/attention/sparse_mla_mask.py new

+56

@@ -0,0 +1,56 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +import cutlass

5 +import cutlass.cute as cute

6 +

7 +from vllm.vllm_flash_attn.cute import utils

8 +

9 +

10 +@cute.jit

11 +def dense_mask_mod(

12 + batch: cute.TensorSSA,

13 + head: cute.TensorSSA,

14 + q_idx: cute.TensorSSA,

15 + kv_idx: cute.TensorSSA,

16 + seqlen_info,

17 + aux_tensors: list,

18 +) -> cute.TensorSSA:

19 + dense_mask = aux_tensors[0]

20 + batch_idx = utils.ssa_to_scalar(batch)

21 + q_idx = utils.ssa_to_scalar(q_idx)

22 + kv_idx = utils.ssa_to_scalar(kv_idx)

23 + word_idx = kv_idx >> 5

24 + bit_idx = cutlass.Uint32(kv_idx & 31)

25 + word = dense_mask[batch_idx, q_idx, word_idx]

26 + result = cute.make_rmem_tensor(1, dtype=cutlass.Uint32)

27 + result[0] = utils.shr_u32(cutlass.Uint32(word), bit_idx)

28 + return result.load()

29 +

30 +

31 +dense_mask_mod.__vec_size__ = 32

32 +

33 +

34 +@cute.jit

35 +def offset_dense_mask_mod(

36 + batch: cute.TensorSSA,

37 + head: cute.TensorSSA,

38 + q_idx: cute.TensorSSA,

39 + kv_idx: cute.TensorSSA,

40 + seqlen_info,

41 + aux_tensors: list,

42 +) -> cute.TensorSSA:

43 + dense_mask = aux_tensors[0]

44 + batch_idx = utils.ssa_to_scalar(batch)

45 + q_idx = utils.ssa_to_scalar(q_idx)

46 + key_start = dense_mask[batch_idx, 0, dense_mask.shape[2] - 1]

47 + kv_idx = utils.ssa_to_scalar(kv_idx) + key_start

48 + word_idx = kv_idx >> 5

49 + bit_idx = cutlass.Uint32(kv_idx & 31)

50 + word = dense_mask[batch_idx, q_idx, word_idx]

51 + result = cute.make_rmem_tensor(1, dtype=cutlass.Uint32)

52 + result[0] = utils.shr_u32(cutlass.Uint32(word), bit_idx)

53 + return result.load()

54 +

55 +

56 +offset_dense_mask_mod.__vec_size__ = 32

vllm/vllm_flash_attn/flash_attn_interface.py

+13

@@ -210,7 +210,9 @@ def flash_attn_varlen_func(

210 cp_tot_seqused_k=None,

211 # FA4 only

212 mask_mod=None,

213 + block_sparse_tensors=None,

214 aux_tensors=None,

215 + aux_tensor_leading_dims=None,

216 dynamic_causal: "torch.Tensor | None" = None,

217 ):

218 """dropout_p should be set to 0.0 during evaluation

@@ -385,6 +387,15 @@ def flash_attn_varlen_func(

387 )

388 elif fa_version == 4:

389 assert alibi_slopes is None, "Alibi is not supported in FA4"

390 + if block_sparse_tensors is not None:

391 + assert block_sparse_tensors.full_block_cnt is not None, (

392 + "FA4 block_sparse_tensors must materialize empty full_block_cnt "

393 + "instead of passing None"

394 + )

395 + assert block_sparse_tensors.full_block_idx is not None, (

396 + "FA4 block_sparse_tensors must materialize empty full_block_idx "

397 + "instead of passing None"

398 + )

399

400 from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd

401

@@ -409,7 +420,9 @@ def flash_attn_varlen_func(

420 out=out,

421 learnable_sink=s_aux,

422 mask_mod=mask_mod,

423 + block_sparse_tensors=block_sparse_tensors,

424 aux_tensors=aux_tensors,

425 + aux_tensor_leading_dims=aux_tensor_leading_dims,

426 output_scale=output_scale,

427 )

428 else: