From de9250ac9e9b249133fff14e15d9248a1ebbcdb8 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 27 Aug 2026 17:17:12 -0400 Subject: [PATCH] [Attention] Enable masked MHA for GLM-5 head dimensions (#53785) Signed-off-by: Matthew Bonanni Co-authored-by: OpenAI Codex --- benchmarks/attention_benchmarks/benchmark.py | 38 ++++- .../mla_sparse_masked_mha_vs_mqa_glm5.yaml | 59 +++++++ benchmarks/attention_benchmarks/mla_runner.py | 32 ++++ cmake/external_projects/vllm_flash_attn.cmake | 2 +- .../layers/test_mla_short_prefill_indexer.py | 82 ++++++++- tests/v1/attention/test_mla_backends.py | 119 +++++++++++++ .../v1/attention/test_sparse_mla_backends.py | 47 ++++- .../layers/attention/mla_attention.py | 160 ++++++++++++++---- .../layers/attention/sparse_mla_attention.py | 25 ++- vllm/models/deepseek_v32/attention.py | 37 ++-- 10 files changed, 527 insertions(+), 74 deletions(-) create mode 100644 benchmarks/attention_benchmarks/configs/mla_sparse_masked_mha_vs_mqa_glm5.yaml diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 8f01be94d2..9c1ea1add7 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -1143,12 +1143,26 @@ def main(): ] else: variants = available_variants + model_sweep = getattr(args, "model_parameter_sweep", None) + if model_sweep is None: + num_q_heads_values = [args.num_q_heads] + elif model_sweep.param_name == "num_q_heads": + num_q_heads_values = model_sweep.values + else: + raise ValueError( + "mha_vs_mqa mode only supports a num_q_heads model parameter sweep" + ) + benchmark_specs = [ + (spec, num_q_heads) + for spec in args.batch_specs + for num_q_heads in num_q_heads_values + ] masked_mha_max_seq_len = getattr( args, "sparse_mla_masked_mha_max_seq_len", None ) formatter = ResultsFormatter(console) total = 0 - for spec in args.batch_specs: + for spec, _ in benchmark_specs: q_len = max(request.q_len for request in parse_batch_spec(spec)) for variant_label, _, _ in variants: if ( @@ -1164,7 +1178,7 @@ def main(): total += len(backends) with tqdm(total=total, desc="Benchmarking") as pbar: - for spec in args.batch_specs: + for spec, num_q_heads in benchmark_specs: q_len = max(request.q_len for request in parse_batch_spec(spec)) for backend in backends: for variant_label, force_mqa, mha_mode in variants: @@ -1178,12 +1192,17 @@ def main(): and q_len > masked_mha_max_seq_len ): continue + result_backend = f"{backend}_{variant_label}" + if model_sweep is not None: + result_backend = model_sweep.get_label( + result_backend, num_q_heads + ) config = BenchmarkConfig( - backend=f"{backend}_{variant_label}", + backend=result_backend, batch_spec=spec, num_layers=args.num_layers, head_dim=args.head_dim, - num_q_heads=args.num_q_heads, + num_q_heads=num_q_heads, num_kv_heads=args.num_kv_heads, block_size=args.block_size, device=args.device, @@ -1211,7 +1230,7 @@ def main(): # run_mla_benchmark needs the real backend name from mla_runner import run_mla_benchmark as run_mla - run_label = f"{backend}_{variant_label} {spec}" + run_label = f"{backend}_{variant_label} {num_q_heads}h {spec}" pbar.set_postfix_str(run_label) try: @@ -1248,7 +1267,14 @@ def main(): # Display results with variant labels as separate "backends" console.print("\n[bold green]MHA vs MQA Results:[/]") - variant_backends = [f"{b}_{v}" for b in backends for v, _, _ in variants] + variant_backends = [] + for num_q_heads in num_q_heads_values: + for backend in backends: + for variant, _, _ in variants: + label = f"{backend}_{variant}" + if model_sweep is not None: + label = model_sweep.get_label(label, num_q_heads) + variant_backends.append(label) formatter.print_table(all_results, variant_backends) # Handle model parameter sweep mode diff --git a/benchmarks/attention_benchmarks/configs/mla_sparse_masked_mha_vs_mqa_glm5.yaml b/benchmarks/attention_benchmarks/configs/mla_sparse_masked_mha_vs_mqa_glm5.yaml new file mode 100644 index 0000000000..682792cbe8 --- /dev/null +++ b/benchmarks/attention_benchmarks/configs/mla_sparse_masked_mha_vs_mqa_glm5.yaml @@ -0,0 +1,59 @@ +mode: mha_vs_mqa +model: + name: "glm-5.x" + num_layers: 78 + num_q_heads: 64 + num_kv_heads: 1 + head_dim: 576 + kv_lora_rank: 512 + qk_nope_head_dim: 192 + qk_rope_head_dim: 64 + v_head_dim: 256 + block_size: 128 + max_model_len: 131072 +model_parameter_sweep: + param_name: "num_q_heads" + values: [64, 32, 16, 8] + label_format: "{backend}_{value}h" +batch_specs: + - "q256s2048" + - "q512s2048" + - "q1024s2048" + - "q1536s2048" + - "q2048" + - "q256s4096" + - "q512s4096" + - "q1024s4096" + - "q1536s4096" + - "q2048s4096" + - "q4096" + - "q256s8192" + - "q512s8192" + - "q1024s8192" + - "q1536s8192" + - "q2048s8192" + - "q4096s8192" + - "q8192" + - "q256s16384" + - "q512s16384" + - "q1024s16384" + - "q1536s16384" + - "q2048s16384" + - "q4096s16384" + - "q8192s16384" + - "q16384" + - "q256s32768" + - "q512s32768" + - "q1024s32768" + - "q1536s32768" + - "q2048s32768" + - "q4096s32768" + - "q8192s32768" + - "q16384s32768" + - "q32768" +backends: ["FLASHMLA_SPARSE", "FLASHINFER_MLA_SPARSE"] +device: "cuda:0" +profile_memory: false +sparse_mla_masked_mha_max_seq_len: 131072 +sparse_mla_mha_variants: ["masked_mha", "mqa"] +sparse_mla_topk_pattern: "random" diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index 6d56fb11ac..790b2f8944 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -313,6 +313,9 @@ _BACKEND_PROPERTIES = { "FLASHMLA_SPARSE": { "query_format": "concat", # Single concatenated tensor (vs tuple) }, + "FLASHINFER_MLA_SPARSE": { + "query_format": "concat", # Single concatenated tensor (vs tuple) + }, } @@ -1184,6 +1187,35 @@ def _run_mla_benchmark_batched( ) # Run each benchmark with the shared impl for config, threshold, num_splits in configs_with_params: + force_masked_mha = ( + is_sparse and getattr(config, "sparse_mla_mha_mode", "auto") == "masked" + ) + if force_masked_mha: + impl.masked_mha_available = True + from vllm.model_executor.layers.attention.sparse_mla_attention import ( + GLOBAL_TOPK_MASK_MAX_BYTES, + _topk_mask_shape, + ) + + requests = parse_batch_spec(config.batch_spec) + max_query_len = max(request.q_len for request in requests) + mask_shape = _topk_mask_shape( + len(requests), max_query_len, max_query_len + ) + mask_words = max( + GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize, + mask_shape[0] * mask_shape[1] * mask_shape[2], + ) + if ( + builder_instance.topk_mask_workspace is None + or builder_instance.topk_mask_workspace.numel() != mask_words + ): + builder_instance.topk_mask_workspace = torch.zeros( + mask_words, + dtype=torch.int32, + device=device, + ) + # Set threshold for this benchmark (FlashAttn/FlashMLA only) original_threshold = None effective_threshold = threshold diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 2170c5c33c..2758788606 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG 4eb29359a51b4fe67c8c0f0c5e931236971aa492 + GIT_TAG 06bdd47c0d0383daf6a2ff0c418faff9c6da16e5 GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index e798d8c7f5..8bd97a3ab4 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -8,6 +8,7 @@ import torch import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer from vllm.config import CUDAGraphMode +from vllm.models.deepseek_v32 import attention as deepseek_v32_attention from vllm.models.deepseek_v32.attention import DeepseekV32Attention from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata @@ -49,7 +50,6 @@ def make_mla_metadata(*, use_dense_mha: bool = True, num_decode_tokens: int = 0) [ "short", "threshold_mismatch", - "mqa_only_layer", "force_mqa", "mla_decode", "capture", @@ -141,11 +141,7 @@ def test_short_prefill_updates_k_cache_before_scoring_decision( topk_indices = torch.full((7, 2048), 17, dtype=torch.int32) def run_indexer(): - if batch_kind == "mqa_only_layer": - assert not DeepseekV32Attention.supports_dense_mha_prefill - dense_mha_layer = "" - else: - dense_mha_layer = MLA_LAYER + assert DeepseekV32Attention.supports_dense_mha_prefill return sparse_indexer.sparse_attn_indexer( hidden_states, INDEXER_LAYER, @@ -163,7 +159,7 @@ def test_short_prefill_updates_k_cache_before_scoring_decision( topk_indices, False, False, - dense_mha_layer, + MLA_LAYER, ) if should_skip: @@ -222,3 +218,75 @@ def test_skipped_k_cache_insert_accepts_no_k( assert result is topk_indices assert torch.all(topk_indices == -1) + + +@pytest.mark.parametrize("fp8_query", [False, True]) +def test_deepseek_v32_dispatches_selected_mha( + monkeypatch: pytest.MonkeyPatch, + fp8_query: bool, +) -> None: + attn_metadata = SimpleNamespace(num_actual_tokens=2) + kv_cache = torch.empty(1) + monkeypatch.setattr( + deepseek_v32_attention, + "get_attention_context", + lambda _: (attn_metadata, None, kv_cache, None), + ) + + observed = {} + + def record_forward_impl(*args): + observed["args"] = args + + layer = SimpleNamespace( + indexer=None, + skip_topk=False, + layer_name=MLA_LAYER, + use_pcp=False, + _fp8_query=fp8_query, + _use_sparse_mha=lambda _: True, + rotary_emb=lambda _positions, q: (q + 1, None), + forward_impl=record_forward_impl, + ) + q_nope = torch.randn(2, 1, 2) + q_pe = torch.randn(2, 1, 2) + mqa_q = torch.randn(2, 1, 2) + mha_q = torch.cat((q_nope, q_pe + 1 if fp8_query else mqa_q), dim=-1) + kv_c = torch.empty(2, 2) + k_pe = torch.empty(2, 2) + output = torch.empty(2, 2) + + DeepseekV32Attention._sparse_indexer_and_attn( + layer, + torch.arange(2), + torch.empty(2, 2), + q_nope, + q_pe, + None, + None, + None, + kv_c, + k_pe, + torch.empty(2, 1, 2), + mqa_q, + output, + ) + + expected_args = ( + mha_q, + kv_c, + kv_cache, + attn_metadata, + output, + ) + actual_args = observed["args"] + torch.testing.assert_close(actual_args[0], mha_q) + assert actual_args[2].shape == (2, 1, 2) + assert actual_args[2].data_ptr() == k_pe.data_ptr() + assert all( + actual is expected + for actual, expected in zip( + actual_args[1:2] + actual_args[3:], + expected_args[1:], + ) + ) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index f054d1a7e0..455060413f 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.attention.mla_attention import ( MLAAttention, QueryLenSupport, _DecodeConcatQuantFP8, + _use_masked_mha, build_mla_chunked_context_metadata, ) from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape @@ -68,6 +69,124 @@ if current_platform.is_rocm(): DEVICE_TYPE = current_platform.device_type +@pytest.mark.parametrize( + ("tensor_parallel_size", "query_len", "expected"), + [ + (1, 8192, True), + (1, 9216, False), + (2, 20480, True), + (2, 24576, False), + (4, 48 * 1024, True), + (4, 56 * 1024, False), + (8, 112 * 1024, True), + (8, 128 * 1024, False), + ], +) +def test_glm5_masked_mha_pure_prefill_routing( + tensor_parallel_size, query_len, expected +): + assert ( + _use_masked_mha( + backend_name="FLASHMLA_SPARSE", + tensor_parallel_size=tensor_parallel_size, + qk_head_dim=256, + v_head_dim=256, + query_len=query_len, + seq_len=query_len, + has_context=False, + ) + is expected + ) + + +@pytest.mark.parametrize( + ("tensor_parallel_size", "query_len", "seq_len", "expected"), + [ + (1, 8192, 8192, False), + (2, 4096, 6144, True), + (2, 4096, 8192, False), + (2, 8192, 10240, True), + (4, 16384, 24576, True), + (4, 16384, 32768, False), + (4, 32768, 34 * 1024, True), + (4, 32768, 36 * 1024, False), + (8, 32768, 49152, True), + (8, 32768, 65536, False), + ], +) +def test_glm5_masked_mha_context_routing( + tensor_parallel_size, query_len, seq_len, expected +): + assert ( + _use_masked_mha( + backend_name="FLASHMLA_SPARSE", + tensor_parallel_size=tensor_parallel_size, + qk_head_dim=256, + v_head_dim=256, + query_len=query_len, + seq_len=seq_len, + has_context=True, + ) + is expected + ) + + +@pytest.mark.parametrize( + ("tensor_parallel_size", "query_len", "seq_len", "has_context", "expected"), + [ + (4, 36 * 1024, 36 * 1024, False, True), + (4, 40 * 1024, 40 * 1024, False, False), + (4, 4 * 1024, 12 * 1024, True, True), + (4, 4 * 1024, 16 * 1024, True, False), + (4, 24 * 1024, 28 * 1024, True, True), + (4, 32 * 1024, 36 * 1024, True, False), + (8, 64 * 1024, 64 * 1024, False, True), + (8, 68 * 1024, 68 * 1024, False, False), + (8, 4 * 1024, 20 * 1024, True, True), + (8, 16 * 1024, 32 * 1024, True, True), + (8, 48 * 1024, 64 * 1024, True, True), + (8, 56 * 1024, 72 * 1024, True, True), + (8, 63 * 1024, 79 * 1024, True, False), + ], +) +def test_glm5_flashinfer_masked_mha_routing( + tensor_parallel_size, query_len, seq_len, has_context, expected +): + assert ( + _use_masked_mha( + backend_name="FLASHINFER_MLA_SPARSE", + tensor_parallel_size=tensor_parallel_size, + qk_head_dim=256, + v_head_dim=256, + query_len=query_len, + seq_len=seq_len, + has_context=has_context, + ) + is expected + ) + + +def test_masked_mha_routing_is_dimension_specific(): + assert _use_masked_mha( + backend_name="FLASHMLA_SPARSE", + tensor_parallel_size=1, + qk_head_dim=192, + v_head_dim=128, + query_len=1536, + seq_len=2048, + has_context=True, + ) + assert not _use_masked_mha( + backend_name="FLASHMLA_SPARSE", + tensor_parallel_size=1, + qk_head_dim=128, + v_head_dim=128, + query_len=1536, + seq_len=2048, + has_context=True, + ) + + @pytest.mark.parametrize( ("cache_dtype", "expected_quant_mode"), [ diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index b65fe24243..cb1c6faaf7 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -21,6 +21,7 @@ from tests.v1.attention.utils import ( ) from vllm import _custom_ops as ops from vllm.config import set_current_vllm_config +from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( GLOBAL_TOPK_MASK_MAX_BYTES, _masked_mha_workspace_fits, @@ -884,6 +885,35 @@ def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expect ) +@pytest.mark.parametrize( + ("backend_name", "tensor_parallel_size", "query_len"), + [ + ("FLASHMLA_SPARSE", 4, 48 * 1024), + ("FLASHMLA_SPARSE", 8, 112 * 1024), + ("FLASHINFER_MLA_SPARSE", 4, 36 * 1024), + ("FLASHINFER_MLA_SPARSE", 8, 64 * 1024), + ], +) +def test_masked_mha_workspace_guards_long_routing_policy( + backend_name, tensor_parallel_size, query_len +): + assert _use_masked_mha( + backend_name=backend_name, + tensor_parallel_size=tensor_parallel_size, + qk_head_dim=256, + v_head_dim=256, + query_len=query_len, + seq_len=query_len, + has_context=False, + ) + assert not _masked_mha_workspace_fits( + batch_size=1, + max_query_len=query_len, + max_context_chunk_seq_len=0, + workspace_numel=GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize, + ) + + def test_masked_mha_workspace_fits_accounts_for_batch_and_context(): """Request count and context chunk length are independent multipliers.""" base = dict(batch_size=2, max_query_len=2048, max_context_chunk_seq_len=2048) @@ -912,14 +942,25 @@ PREFILL_BATCH_SPECS = { ) @pytest.mark.parametrize("batch_name", list(PREFILL_BATCH_SPECS.keys())) @pytest.mark.parametrize("kv_cache_dtype", ["auto"]) +@pytest.mark.parametrize( + ("num_heads", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"), + [ + pytest.param(128, 128, 64, 128, id="deepseek_hd192_v128"), + pytest.param(64, 192, 64, 256, id="glm5_hd256_v256"), + ], +) def test_sparse_backend_prefill_correctness( default_vllm_config, dist_init, batch_name, kv_cache_dtype, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, workspace_init, ): - """Test single-pass dense and masked MHA for sparse MLA prefill.""" + """Test dense and masked MHA across supported sparse MLA dimensions.""" backend_cls = FlashMLASparseBackend batch_spec = PREFILL_BATCH_SPECS[batch_name] @@ -927,11 +968,7 @@ def test_sparse_backend_prefill_correctness( dtype = torch.bfloat16 block_size = 64 - num_heads = 128 kv_lora_rank = 512 - qk_nope_head_dim = 128 - qk_rope_head_dim = 64 - v_head_dim = 128 head_size = kv_lora_rank + qk_rope_head_dim masked_mha = batch_name.startswith("masked_mha") topk_tokens = 200 if masked_mha else 512 diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 6e2902a3c1..fad2854861 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -827,25 +827,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): use_mha = True if self.impl.is_sparse and num_mha_tokens > 0: - prefill = getattr(attn_metadata, "prefill", None) - use_dense_mha = getattr(prefill, "use_dense_mha", False) - prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined] - use_masked_mha = ( - self.prefill_backend is not None - and self.impl.masked_mha_available # type: ignore[attr-defined] - and self.impl.dcp_world_size <= 1 - and prefill is not None - and _use_masked_mha( - backend_name=self.attn_backend.get_name(), - tensor_parallel_size=self._vllm_config.parallel_config.tensor_parallel_size, - query_len=prefill.max_query_len, - seq_len=prefill_max_seq_len, - ) - and self.impl.masked_mha_workspace_fits(prefill) # type: ignore[attr-defined] - ) - use_mha = (use_dense_mha or use_masked_mha) and not ( - self._vllm_config.attention_config.sparse_mla_force_mqa - ) + use_mha = self._use_sparse_mha(attn_metadata) if not use_mha: num_mqa_tokens = q.size(0) num_mha_tokens = 0 @@ -1063,6 +1045,31 @@ class MLAAttention(nn.Module, AttentionLayerBase): output_padded[num_actual_toks:].zero_() return output_padded + def _use_sparse_mha(self, attn_metadata: "MLACommonMetadata") -> bool: + prefill = attn_metadata.prefill + if prefill is None: + return False + use_masked_mha = ( + self.prefill_backend is not None + and self.impl.masked_mha_available # type: ignore[attr-defined] + and self.impl.dcp_world_size <= 1 + and _use_masked_mha( + backend_name=self.attn_backend.get_name(), + tensor_parallel_size=( + self._vllm_config.parallel_config.tensor_parallel_size + ), + qk_head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + query_len=prefill.max_query_len, + seq_len=attn_metadata.prefill_max_seq_len, # type: ignore[attr-defined] + has_context=prefill.chunked_context is not None, + ) + and self.impl.masked_mha_workspace_fits(prefill) # type: ignore[attr-defined] + ) + return (prefill.use_dense_mha or use_masked_mha) and not ( + self._vllm_config.attention_config.sparse_mla_force_mqa + ) + def process_weights_after_loading(self, act_dtype: torch.dtype): # Let per-backend impls do their own weight packing first (no-op # unless overridden), mirroring Attention.process_weights_after_loading. @@ -1619,36 +1626,121 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims: ) -_DSV32_MASKED_MHA_THRESHOLDS: dict[str, dict[int, tuple[int | None, ...]]] = { - "FLASHMLA_SPARSE": { - 1: (1536, 4096, None, None, None), - 2: (512, 1024, 4096, None, None), - 4: (512, 1024, 1536, 8192, None), - 8: (512, 512, 1024, 2048, 16384), +_MaskedMHARanges = tuple[tuple[int, int], ...] + +_MASKED_MHA_THRESHOLDS: dict[ + tuple[int, int], + dict[str, dict[int, tuple[_MaskedMHARanges, _MaskedMHARanges | None]]], +] = { + (192, 128): { + "FLASHMLA_SPARSE": { + 1: (((2048, 1536), (4096, 4096)), None), + 2: (((2048, 512), (4096, 1024), (8192, 4096)), None), + 4: ( + ((2048, 512), (4096, 1024), (8192, 1536), (16384, 8192)), + None, + ), + 8: ( + ( + (2048, 512), + (4096, 512), + (8192, 1024), + (16384, 2048), + (32768, 16384), + ), + None, + ), + }, + "FLASHINFER_MLA_SPARSE": { + 8: ( + ( + (2048, 512), + (4096, 1024), + (8192, 1024), + (16384, 2048), + (32768, 32768), + ), + None, + ), + }, }, - "FLASHINFER_MLA_SPARSE": { - 8: (512, 1024, 1024, 2048, 32768), + (256, 256): { + "FLASHMLA_SPARSE": { + 1: (((8 * 1024, 0),), ()), + 2: ( + ((20 * 1024, 0),), + ((6 * 1024, 4 * 1024), (10 * 1024, 8 * 1024)), + ), + 4: ( + ((48 * 1024, 0),), + ( + (8 * 1024, 4 * 1024), + (16 * 1024, 8 * 1024), + (24 * 1024, 16 * 1024), + (34 * 1024, 32 * 1024), + ), + ), + 8: ( + ((112 * 1024, 0),), + ( + (8 * 1024, 4 * 1024), + (16 * 1024, 8 * 1024), + (32 * 1024, 16 * 1024), + (48 * 1024, 32 * 1024), + ), + ), + }, + "FLASHINFER_MLA_SPARSE": { + 4: ( + ((36 * 1024, 0),), + ( + (12 * 1024, 4 * 1024), + (16 * 1024, 8 * 1024), + (24 * 1024, 16 * 1024), + (28 * 1024, 24 * 1024), + ), + ), + 8: ( + ((64 * 1024, 0),), + ( + (20 * 1024, 4 * 1024), + (24 * 1024, 8 * 1024), + (32 * 1024, 16 * 1024), + (48 * 1024, 32 * 1024), + (56 * 1024, 40 * 1024), + (64 * 1024, 48 * 1024), + (72 * 1024, 56 * 1024), + ), + ), + }, }, } -_DSV32_SEQ_LEN_BUCKETS = (2048, 4096, 8192, 16384, 32768) def _use_masked_mha( *, backend_name: str, tensor_parallel_size: int, + qk_head_dim: int, + v_head_dim: int, query_len: int, seq_len: int, + has_context: bool, ) -> bool: - thresholds = _DSV32_MASKED_MHA_THRESHOLDS.get(backend_name, {}).get( - tensor_parallel_size + thresholds = ( + _MASKED_MHA_THRESHOLDS.get((qk_head_dim, v_head_dim), {}) + .get(backend_name, {}) + .get(tensor_parallel_size) ) if thresholds is None: return False - for bucket_idx, bucket_seq_len in enumerate(_DSV32_SEQ_LEN_BUCKETS): - if seq_len <= bucket_seq_len: - min_query_len = thresholds[bucket_idx] - return min_query_len is not None and query_len >= min_query_len + pure_prefill_ranges, context_ranges = thresholds + ranges = pure_prefill_ranges + if has_context and context_ranges is not None: + ranges = context_ranges + for max_seq_len, min_query_len in ranges: + if seq_len <= max_seq_len: + return query_len >= min_query_len return False diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 1fe267f3c8..ab6946132f 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -88,16 +88,27 @@ def _is_masked_mha_available( """Check if masked MHA can ever fire for this model configuration.""" if not current_platform.is_device_capability_family(100): return False - if ( - num_heads_total != 128 - or kv_lora_rank != 512 - or qk_nope_head_dim != 128 - or qk_rope_head_dim != 64 - or v_head_dim != 128 + model_dims = ( + num_heads_total, + kv_lora_rank, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + ) + if model_dims not in ( + (128, 512, 128, 64, 128), + (64, 512, 192, 64, 256), ): return False qk_head_dim = qk_nope_head_dim + qk_rope_head_dim - fa_version = get_flash_attn_version(head_size=qk_head_dim, head_size_v=v_head_dim) + if qk_head_dim == 256 and v_head_dim == 256: + # This path uses contiguous K/V, so it does not need the paged-KV + # features that keep head-dim 256 disabled in the general FA selector. + fa_version = get_flash_attn_version() + else: + fa_version = get_flash_attn_version( + head_size=qk_head_dim, head_size_v=v_head_dim + ) return fa_version == 4 and not is_quantized_kv_cache(kv_cache_dtype) diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 87e7ec24e3..e0dd5bfb5e 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -119,8 +119,6 @@ class DeepseekV32Attention(MLAAttention): indexer: "DeepseekV32Indexer | None" indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer - supports_dense_mha_prefill = False - def __init__( self, vllm_config: VllmConfig, @@ -303,15 +301,6 @@ class DeepseekV32Attention(MLAAttention): device=hidden_states.device, ) forward_context = get_forward_context() - attn_metadata_raw = forward_context.attn_metadata - if isinstance(attn_metadata_raw, dict): - attn_metadata = attn_metadata_raw.get(self.layer_name) - elif isinstance(attn_metadata_raw, list): - attn_metadata = attn_metadata_raw[0].get(self.layer_name) - else: - attn_metadata = attn_metadata_raw - attn_metadata = cast("MLACommonMetadata | None", attn_metadata) - slot_mapping = forward_context.slot_mapping assert isinstance(slot_mapping, dict) mla_slot = slot_mapping.get(self.layer_name) @@ -337,7 +326,7 @@ class DeepseekV32Attention(MLAAttention): indexer_softmax_scale = 0.0 indexer_n_head_scale = 0.0 - if attn_metadata is None or self.use_pcp: + if forward_context.attn_metadata is None or self.use_pcp: mla_kv_cache = None mla_k_scale = None indexer_k_cache = None @@ -346,8 +335,8 @@ class DeepseekV32Attention(MLAAttention): mla_kv_cache = self.kv_cache mla_k_scale = self._k_scale - kv_c_out = torch.empty_like(kv_c) if self.use_pcp else None - k_pe_out = torch.empty_like(k_pe) if self.use_pcp else None + kv_c_out = torch.empty_like(kv_c) + k_pe_out = torch.empty_like(k_pe) q_c = fused_norm_rope( positions, q_c, @@ -403,7 +392,10 @@ class DeepseekV32Attention(MLAAttention): ) self._sparse_indexer_and_attn( + positions, q_c, + q_nope, + q_pe, index_q_fp8, index_k_out, index_weights_out, @@ -418,7 +410,10 @@ class DeepseekV32Attention(MLAAttention): @eager_break_during_capture def _sparse_indexer_and_attn( self, + positions: torch.Tensor, q_c: torch.Tensor, + q_nope: torch.Tensor, + q_pe: torch.Tensor, index_q_fp8: torch.Tensor | None, index_k: torch.Tensor | None, index_weights_out: torch.Tensor | None, @@ -498,6 +493,20 @@ class DeepseekV32Attention(MLAAttention): output.zero_() return + if self._use_sparse_mha(attn_metadata): + assert kv_c is not None and k_pe is not None + mha_q_pe = self.rotary_emb(positions, q_pe)[0] if self._fp8_query else mqa_q + mha_q = torch.cat((q_nope, mha_q_pe), dim=-1) + self.forward_impl( + mha_q, + kv_c, + k_pe.unsqueeze(1), + kv_cache, + attn_metadata, + output, + ) + return + if self._fp8_kv_needs_view: kv_cache = kv_cache.view(torch.float8_e4m3fn) if self._fp8_query: