[Attention] Enable masked MHA for GLM-5 head dimensions (#53785)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Matthew Bonanni
co-authored by
OpenAI Codex
OpenAI Codex
parent
3a9bfc209f
commit
de9250ac9e
@@ -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
|
||||
|
||||
@@ -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"
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:],
|
||||
)
|
||||
)
|
||||
|
||||
@@ -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"),
|
||||
[
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user