karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[Attention] Abstract the MLA prefill backends and eliminate cuDNN (#32623)

Browse Source 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Matthew Bonanni
2026-05-01 13:36:20 -04:00
committed by GitHub
co-authored by Michael Goin Lucas Wilkinson mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 51295793a2
commit f3fef12350
16 changed files with 1629 additions and 708 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/attention_backends.md
+12 -8
View File
@@ -192,21 +192,25 @@ MLA uses separate backends for prefill and decode phases.
### Prefill Backends
The prefill backend is selected at runtime based on hardware and
configuration.
To explicitly select a prefill backend, use
`-ac.mla_prefill_backend=<BACKEND>` (e.g., `FLASH_ATTN`, `FLASHINFER`).
Otherwise, the prefill backend is selected automatically at runtime based on
hardware and configuration.
| Backend | Description | Compute Cap. | Enable | Disable | Notes |
| ------- | ----------- | ------------ | ------ | ------- | ----- |
| TRT-LLM Ragged‡ | TensorRT-LLM ragged attention | 10.x | Default on SM100 | `-ac.use_trtllm_ragged_deepseek_prefill=0` | DeepSeek R1 dims only |
| FlashInfer | FlashInfer CUTLASS backend | 10.x | `-ac.disable_flashinfer_prefill=0` | `-ac.disable_flashinfer_prefill=1` | DeepSeek R1 dims only |
| cuDNN | cuDNN-based attention | 10.x | `-ac.use_cudnn_prefill=1` | `-ac.use_cudnn_prefill=0` | |
| FlashAttention | FlashAttention varlen (FA2/FA3) | Any | Default fallback | Use other backends | FA3 on SM90, FA2 otherwise |
| Backend | Description | Dtypes | Compute Cap. | Notes |
| ------- | ----------- | ------ | ------------ | ----- |
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise |
| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | DeepSeek R1 dims only |
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
> **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
> On other GPUs, FlashAttention is used as the default.
### Decode Backends
MLA decode backends are selected using the standard
`-ac.backend=<BACKEND>` argument (e.g., `FLASHMLA`, `TRITON_MLA`).
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
tests/engine/test_arg_utils.py
-3
View File
@@ -333,8 +333,6 @@ def test_attention_config():
"true",
"--attention-config.flash_attn_max_num_splits_for_cuda_graph",
"16",
"--attention-config.use_cudnn_prefill",
"true",
"--attention-config.use_trtllm_ragged_deepseek_prefill",
"true",
"--attention-config.use_trtllm_attention",
@@ -352,7 +350,6 @@ def test_attention_config():
assert engine_args.attention_config.flash_attn_version == 3
assert engine_args.attention_config.use_prefill_decode_attention is True
assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 16
assert engine_args.attention_config.use_cudnn_prefill is True
assert engine_args.attention_config.use_trtllm_ragged_deepseek_prefill is True
assert engine_args.attention_config.use_trtllm_attention is True
assert engine_args.attention_config.disable_flashinfer_prefill is True
tests/v1/attention/test_mla_backends.py
+1
View File
@@ -672,6 +672,7 @@ def run_attention_backend(
def test_backend_correctness(
default_vllm_config,
dist_init,
workspace_init,
batch_spec_name: str,
model: str,
tensor_parallel_size: int,
tests/v1/attention/test_mla_prefill_selector.py
+304
View File
@@ -0,0 +1,304 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for MLA prefill backend selector."""
from unittest.mock import MagicMock, patch
import pytest
import torch
from vllm.config import AttentionConfig, ModelConfig, VllmConfig
from vllm.platforms.interface import DeviceCapability
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
from vllm.v1.attention.backends.mla.prefill.selector import (
MLAPrefillSelectorConfig,
_auto_select_mla_prefill_backend,
get_mla_prefill_backend,
is_deepseek_r1_mla_compatible,
)
@pytest.fixture(autouse=True)
def clear_cache():
"""Clear lru cache to ensure each test case runs without caching."""
_auto_select_mla_prefill_backend.cache_clear()
def _make_mock_model_config(
qk_nope_head_dim: int = 128,
qk_rope_head_dim: int = 64,
v_head_dim: int = 128,
dtype: torch.dtype = torch.bfloat16,
) -> ModelConfig:
mock_config = MagicMock(spec=ModelConfig)
mock_config.dtype = dtype
mock_config.hf_text_config = MagicMock()
mock_config.hf_text_config.qk_nope_head_dim = qk_nope_head_dim
mock_config.hf_text_config.qk_rope_head_dim = qk_rope_head_dim
mock_config.hf_text_config.v_head_dim = v_head_dim
return mock_config
def _make_vllm_config(
model_config: ModelConfig | None = None,
mla_prefill_backend: MLAPrefillBackendEnum | None = None,
) -> VllmConfig:
if model_config is None:
model_config = _make_mock_model_config()
attention_config = AttentionConfig(mla_prefill_backend=mla_prefill_backend)
mock_vllm_config = MagicMock(spec=VllmConfig)
mock_vllm_config.model_config = model_config
mock_vllm_config.attention_config = attention_config
return mock_vllm_config
class TestGetMLAPrefillBackend:
"""Tests for get_mla_prefill_backend (public API)."""
def test_no_device_capability_returns_flash_attn(self):
vllm_config = _make_vllm_config()
with patch("vllm.platforms.current_platform") as mock_platform:
mock_platform.get_device_capability.return_value = None
backend = get_mla_prefill_backend(vllm_config)
assert backend.get_name() == "FLASH_ATTN"
def test_explicit_flash_attn_selection(self):
try:
flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()
except ImportError:
pytest.skip("FLASH_ATTN backend not available")
return
vllm_config = _make_vllm_config(
mla_prefill_backend=MLAPrefillBackendEnum.FLASH_ATTN,
)
with patch("vllm.platforms.current_platform") as mock_platform:
mock_platform.get_device_capability.return_value = DeviceCapability(
major=9, minor=0
)
with patch.object(
flash_attn_cls,
"validate_configuration",
return_value=[],
):
backend = get_mla_prefill_backend(vllm_config)
assert backend.get_name() == "FLASH_ATTN"
def test_explicit_backend_invalid_raises_error(self):
vllm_config = _make_vllm_config(
mla_prefill_backend=MLAPrefillBackendEnum.FLASHINFER,
)
with patch("vllm.platforms.current_platform") as mock_platform:
mock_platform.get_device_capability.return_value = DeviceCapability(
major=9, minor=0
)
with pytest.raises(ValueError, match="is not valid"):
get_mla_prefill_backend(vllm_config)
def test_explicit_backend_import_error_raises(self):
vllm_config = _make_vllm_config(
mla_prefill_backend=MLAPrefillBackendEnum.TRTLLM_RAGGED,
)
with patch("vllm.platforms.current_platform") as mock_platform:
mock_platform.get_device_capability.return_value = DeviceCapability(
major=10, minor=0
)
with (
patch.object(
MLAPrefillBackendEnum.TRTLLM_RAGGED,
"get_class",
side_effect=ImportError("trtllm not installed"),
),
pytest.raises(ValueError, match="is not valid"),
):
get_mla_prefill_backend(vllm_config)
def test_auto_selection_on_hopper(self):
try:
flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()
except ImportError:
pytest.skip("FLASH_ATTN backend not available")
return
vllm_config = _make_vllm_config()
with patch("vllm.platforms.current_platform") as mock_platform:
mock_platform.get_device_capability.return_value = DeviceCapability(
major=9, minor=0
)
with patch.object(
flash_attn_cls,
"validate_configuration",
return_value=[],
):
backend = get_mla_prefill_backend(vllm_config)
assert backend.get_name() == "FLASH_ATTN"
class TestAutoSelectMLAPrefillBackend:
"""Tests for fallback and error paths in auto-selection."""
def test_blackwell_falls_back_to_trtllm(self):
vllm_config = _make_vllm_config()
capability = DeviceCapability(major=10, minor=0)
selector_config = MLAPrefillSelectorConfig(
dtype=torch.bfloat16,
is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),
)
try:
trtllm_cls = MLAPrefillBackendEnum.TRTLLM_RAGGED.get_class()
except ImportError:
pytest.skip("TRTLLM_RAGGED backend not available")
return
with (
patch.object(
MLAPrefillBackendEnum.FLASH_ATTN,
"get_class",
side_effect=ImportError("FLASH_ATTN not available"),
),
patch.object(trtllm_cls, "validate_configuration", return_value=[]),
):
backend = _auto_select_mla_prefill_backend(
capability,
selector_config,
)
assert backend.get_name() == "TRTLLM_RAGGED"
def test_all_fail_raises_error(self):
vllm_config = _make_vllm_config()
capability = DeviceCapability(major=10, minor=0)
selector_config = MLAPrefillSelectorConfig(
dtype=torch.bfloat16,
is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),
)
def mock_get_class(backend_enum): # noqa: ARG001
cls = MagicMock()
cls.validate_configuration.return_value = ["not available"]
return cls
with patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class):
_auto_select_mla_prefill_backend.cache_clear()
with pytest.raises(ValueError, match="No valid MLA"):
_auto_select_mla_prefill_backend(
capability,
selector_config,
)
class TestBackendValidation:
"""Tests for backend validation logic."""
def test_r1_dimension_requirement(self):
try:
from vllm.v1.attention.backends.mla.prefill.flashinfer import (
FlashInferPrefillBackend,
)
except ImportError:
pytest.skip("FlashInfer prefill backend not available")
return
assert FlashInferPrefillBackend.requires_r1_mla_dimensions is True
vllm_config = _make_vllm_config(
model_config=_make_mock_model_config(
qk_nope_head_dim=128,
qk_rope_head_dim=64,
v_head_dim=128,
)
)
capability = DeviceCapability(major=10, minor=0)
selector_config = MLAPrefillSelectorConfig(
dtype=torch.bfloat16,
is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),
)
with patch.object(FlashInferPrefillBackend, "is_available", return_value=True):
invalid_reasons = FlashInferPrefillBackend.validate_configuration(
capability,
selector_config,
)
assert len(invalid_reasons) == 0
vllm_config_invalid = _make_vllm_config(
model_config=_make_mock_model_config(
qk_nope_head_dim=64,
qk_rope_head_dim=64,
v_head_dim=128,
)
)
selector_config_invalid = MLAPrefillSelectorConfig(
dtype=torch.bfloat16,
is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config_invalid),
)
with patch.object(FlashInferPrefillBackend, "is_available", return_value=True):
invalid_reasons = FlashInferPrefillBackend.validate_configuration(
capability,
selector_config_invalid,
)
assert len(invalid_reasons) == 1
assert "DeepSeek R1 MLA dimensions" in invalid_reasons[0]
class TestMLAPrefillBackendParsing:
"""Tests for string-based mla_prefill_backend parsing from CLI args."""
def test_valid_string_parses_to_enum(self):
config = AttentionConfig(
mla_prefill_backend="FLASH_ATTN", # type: ignore[arg-type]
)
assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN
def test_invalid_string_raises_error(self):
with pytest.raises(ValueError, match="Unknown MLA prefill backend"):
AttentionConfig(
mla_prefill_backend="NONEXISTENT", # type: ignore[arg-type]
)
class TestDeprecatedFlagMigration:
"""Tests for _migrate_deprecated_mla_prefill_flags in AttentionConfig."""
def test_no_deprecated_flags_leaves_backend_none(self):
config = AttentionConfig()
assert config.mla_prefill_backend is None
def test_use_trtllm_ragged_migrates_to_trtllm_ragged(self):
config = AttentionConfig(use_trtllm_ragged_deepseek_prefill=True)
assert config.mla_prefill_backend == MLAPrefillBackendEnum.TRTLLM_RAGGED
def test_disable_flashinfer_prefill_migrates_to_flash_attn(self):
config = AttentionConfig(disable_flashinfer_prefill=True)
assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN
def test_explicit_backend_ignores_deprecated_flags(self):
config = AttentionConfig(
mla_prefill_backend=MLAPrefillBackendEnum.FLASH_ATTN,
use_cudnn_prefill=True,
)
assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN
def test_cudnn_raises_error(self):
match = "cuDNN MLA prefill backend has been removed"
with pytest.raises(ValueError, match=match):
AttentionConfig(use_cudnn_prefill=True)
def test_trtllm_takes_priority_over_disable_flashinfer(self):
config = AttentionConfig(
use_trtllm_ragged_deepseek_prefill=True,
disable_flashinfer_prefill=True,
)
assert config.mla_prefill_backend == MLAPrefillBackendEnum.TRTLLM_RAGGED
tools/pre_commit/generate_attention_backend_docs.py
+253 -88
View File
@@ -30,7 +30,6 @@ REPO_ROOT = Path(__file__).parent.parent.parent
RELEVANT_PATTERNS = [
"vllm/v1/attention/backends/*.py",
"vllm/v1/attention/backends/**/*.py",
"vllm/v1/attention/backends/fa_utils.py",
"vllm/model_executor/layers/attention/mla_attention.py",
"vllm/platforms/cuda.py",
"tools/pre_commit/generate_attention_backend_docs.py",
@@ -68,6 +67,11 @@ def is_relevant_file(filepath: str) -> bool:
return any(fnmatch.fnmatch(path_str, pattern) for pattern in RELEVANT_PATTERNS)
MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill"
MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py"
MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py"
# ---------------------------------------------------------------------------
# AST utility helpers
# ---------------------------------------------------------------------------
@@ -293,6 +297,242 @@ def get_file_from_class_path(class_path: str) -> Path | None:
return py_file if py_file.exists() else None
def parse_mla_prefill_registry() -> dict[str, str]:
"""Parse MLAPrefillBackendEnum from the prefill registry.
Returns:
A dict mapping backend names to their class paths.
"""
if not MLA_PREFILL_REGISTRY_FILE.exists():
return {}
try:
tree = ast.parse(MLA_PREFILL_REGISTRY_FILE.read_text())
except Exception:
return {}
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef) and node.name == "MLAPrefillBackendEnum":
return _extract_enum_values(node)
return {}
def parse_mla_prefill_priorities() -> dict[str, list[str]]:
"""Parse MLA prefill backend priorities from selector.py.
Returns:
A dict with keys like 'blackwell' and 'default' containing
lists of backend enum names in priority order.
"""
if not MLA_PREFILL_SELECTOR_FILE.exists():
return {}
try:
tree = ast.parse(MLA_PREFILL_SELECTOR_FILE.read_text())
except Exception:
return {}
priorities: dict[str, list[str]] = {}
for node in ast.walk(tree):
if not isinstance(node, ast.FunctionDef):
continue
if node.name != "_get_mla_prefill_backend_priorities":
continue
# Look for if statements checking device_capability.major
for stmt in ast.walk(node):
if not isinstance(stmt, ast.If):
continue
# Check if it's a capability.major == 10 check (Blackwell)
is_blackwell = (
isinstance(stmt.test, ast.Compare)
and isinstance(stmt.test.left, ast.Attribute)
and stmt.test.left.attr == "major"
and stmt.test.comparators
and isinstance(stmt.test.comparators[0], ast.Constant)
and stmt.test.comparators[0].value == 10
)
# Extract backends from return statements
for body_stmt in stmt.body:
if isinstance(body_stmt, ast.Return) and isinstance(
body_stmt.value, ast.List
):
backends = []
for elt in body_stmt.value.elts:
if isinstance(elt, ast.Attribute):
backends.append(elt.attr)
if is_blackwell:
priorities["blackwell"] = backends
else:
priorities["default"] = backends
# Extract from else branch
for else_stmt in stmt.orelse:
if isinstance(else_stmt, ast.Return) and isinstance(
else_stmt.value, ast.List
):
backends = []
for elt in else_stmt.value.elts:
if isinstance(elt, ast.Attribute):
backends.append(elt.attr)
priorities["default"] = backends
return priorities
def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None:
"""Parse a single MLA prefill backend file to extract its properties.
Args:
class_path: The fully qualified class path.
Returns:
A dict with backend properties, or None if parsing fails.
"""
file_path = get_file_from_class_path(class_path)
if file_path is None:
return None
try:
tree = ast.parse(file_path.read_text())
except Exception:
return None
class_name = class_path.rsplit(".", 1)[1]
class_node = find_class_in_ast(tree, class_name)
if class_node is None:
return None
info: dict[str, Any] = {
"compute_capability": "Any",
"requires_r1_dims": False,
"dtypes": "fp16, bf16", # Default from base class
}
# Parse class variables
for item in class_node.body:
if isinstance(item, ast.Assign):
for target in item.targets:
if (
isinstance(target, ast.Name)
and target.id == "requires_r1_mla_dimensions"
and isinstance(item.value, ast.Constant)
):
info["requires_r1_dims"] = item.value.value
# Parse supported_dtypes class variable
if (
isinstance(item, ast.AnnAssign)
and isinstance(item.target, ast.Name)
and item.target.id == "supported_dtypes"
and isinstance(item.value, ast.List)
):
dtype_map = {"float16": "fp16", "bfloat16": "bf16", "float32": "fp32"}
dtypes = []
for elt in item.value.elts:
if isinstance(elt, ast.Attribute):
dtypes.append(dtype_map.get(elt.attr, elt.attr))
if dtypes:
info["dtypes"] = ", ".join(dtypes)
# Parse get_name static method
get_name_method = find_method(class_node, "get_name")
if get_name_method:
for n in ast.walk(get_name_method):
if isinstance(n, ast.Return) and isinstance(n.value, ast.Constant):
info["name"] = n.value.value
# Parse supports_compute_capability classmethod
cc_method = find_method(class_node, "supports_compute_capability")
if cc_method:
for n in ast.walk(cc_method):
# Look for capability.major == 10 style checks
if (
isinstance(n, ast.Compare)
and isinstance(n.left, ast.Attribute)
and n.left.attr == "major"
and n.comparators
and isinstance(n.comparators[0], ast.Constant)
):
major = n.comparators[0].value
info["compute_capability"] = f"{major}.x"
return info
def parse_mla_prefill_backends() -> list[dict[str, Any]]:
"""Parse MLA prefill backend options from the prefill registry.
MLA uses different backends for prefill vs decode. The decode backends are
registered in the main registry, but prefill backends have their own
registry at vllm/v1/attention/backends/mla/prefill/registry.py.
Returns a list of prefill backend info dicts with their requirements.
"""
registry = parse_mla_prefill_registry()
priorities = parse_mla_prefill_priorities()
if not registry:
return []
# Get the priority order (Blackwell order shows all backends)
priority_order = priorities.get("blackwell", list(registry.keys()))
prefill_backends: list[dict[str, Any]] = []
# Backend-specific metadata that can't be easily parsed from code
backend_metadata = {
"TRTLLM_RAGGED": {
"description": "TensorRT-LLM ragged attention",
},
"FLASHINFER": {
"description": "FlashInfer CUTLASS backend",
},
"FLASH_ATTN": {
"description": "FlashAttention varlen (FA2/FA3/FA4)",
},
}
for backend_name in priority_order:
if backend_name not in registry:
continue
class_path = registry[backend_name]
backend_info = parse_mla_prefill_backend_file(class_path)
if backend_info is None:
continue
metadata = backend_metadata.get(backend_name, {})
display_name = backend_info.get("name", backend_name)
# Add marker for default Blackwell backend
marker = ""
if backend_name == priority_order[0] and priorities.get("blackwell"):
marker = "‡"
notes = ""
if backend_info.get("requires_r1_dims"):
notes = "DeepSeek R1 dims only"
elif backend_name == "FLASH_ATTN":
notes = "FA4 on SM100+, FA3 on SM90, FA2 otherwise"
prefill_backends.append(
{
"name": display_name,
"marker": marker,
"description": metadata.get("description", ""),
"dtypes": backend_info.get("dtypes", "fp16, bf16"),
"compute_capability": backend_info.get("compute_capability", "Any"),
"notes": notes,
}
)
return prefill_backends
# ---------------------------------------------------------------------------
# Backend feature extraction from AST
# ---------------------------------------------------------------------------
@@ -807,86 +1047,6 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
}
def parse_mla_prefill_backends() -> list[dict[str, Any]]:
"""Parse MLA prefill backend options from mla_attention.py.
MLA uses different backends for prefill vs decode. The decode backends are
registered in the registry, but prefill backends are selected at runtime
based on conditions in MLACommonImpl.__init__.
Returns a list of prefill backend info dicts with their requirements.
"""
if not MLA_ATTENTION_FILE.exists():
return []
try:
tree = ast.parse(MLA_ATTENTION_FILE.read_text())
except Exception:
return []
# Find compute capability requirements by parsing use_* functions
trtllm_cc = _find_cc_in_function(tree, "use_trtllm_ragged_deepseek_prefill")
flashinfer_cc = _find_cc_in_function(tree, "use_flashinfer_prefill")
cudnn_cc = _find_cc_in_function(tree, "use_cudnn_prefill")
# Build prefill backend list based on what we found
# Order matches the priority in MLACommonImpl.__init__
prefill_backends: list[dict[str, Any]] = []
# TRT-LLM Ragged (highest priority if available)
if trtllm_cc:
prefill_backends.append(
{
"name": "TRT-LLM Ragged‡",
"description": "TensorRT-LLM ragged attention",
"compute_capability": trtllm_cc,
"enable": "Default on SM100",
"disable": "`-ac.use_trtllm_ragged_deepseek_prefill=0`",
"notes": "DeepSeek R1 dims only",
}
)
# FlashInfer prefill
if flashinfer_cc:
prefill_backends.append(
{
"name": "FlashInfer",
"description": "FlashInfer CUTLASS backend",
"compute_capability": flashinfer_cc,
"enable": "`-ac.disable_flashinfer_prefill=0`",
"disable": "`-ac.disable_flashinfer_prefill=1`",
"notes": "DeepSeek R1 dims only",
}
)
# cuDNN prefill
if cudnn_cc:
prefill_backends.append(
{
"name": "cuDNN",
"description": "cuDNN-based attention",
"compute_capability": cudnn_cc,
"enable": "`-ac.use_cudnn_prefill=1`",
"disable": "`-ac.use_cudnn_prefill=0`",
"notes": "",
}
)
# FlashAttention is always available as fallback
prefill_backends.append(
{
"name": "FlashAttention",
"description": "FlashAttention varlen (FA2/FA3)",
"compute_capability": "Any",
"enable": "Default fallback",
"disable": "Use other backends",
"notes": "FA3 on SM90, FA2 otherwise",
}
)
return prefill_backends
# ---------------------------------------------------------------------------
# Backend variant expansion (FA2/FA3/FA4, FlashInfer native/TRTLLM)
# ---------------------------------------------------------------------------
@@ -1415,20 +1575,22 @@ def generate_mla_section(
"",
"### Prefill Backends",
"",
"The prefill backend is selected at runtime based on hardware and",
"configuration.",
"To explicitly select a prefill backend, use",
"`-ac.mla_prefill_backend=<BACKEND>` (e.g., `FLASH_ATTN`, `FLASHINFER`).",
"Otherwise, the prefill backend is selected automatically at runtime based on",
"hardware and configuration.",
"",
"| Backend | Description | Compute Cap. | Enable | Disable | Notes |",
"| ------- | ----------- | ------------ | ------ | ------- | ----- |",
"| Backend | Description | Dtypes | Compute Cap. | Notes |",
"| ------- | ----------- | ------ | ------------ | ----- |",
]
for backend in prefill_backends:
row = "| {} | {} | {} | {} | {} | {} |".format(
row = "| `{}`{} | {} | {} | {} | {} |".format(
backend["name"],
backend.get("marker", ""),
backend["description"],
backend.get("dtypes", "fp16, bf16"),
backend["compute_capability"],
backend["enable"],
backend["disable"],
backend.get("notes", ""),
)
lines.append(row.replace(" ", " "))
@@ -1441,6 +1603,9 @@ def generate_mla_section(
"",
"### Decode Backends",
"",
"MLA decode backends are selected using the standard",
"`-ac.backend=<BACKEND>` argument (e.g., `FLASHMLA`, `TRITON_MLA`).",
"",
]
)
vllm/config/attention.py
+57 -2
View File
@@ -6,8 +6,12 @@ from typing import Any, Literal
from pydantic import field_validator
from vllm.config.utils import config
from vllm.logger import init_logger
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
from vllm.v1.attention.backends.registry import AttentionBackendEnum
logger = init_logger(__name__)
@config
class AttentionConfig:
@@ -33,7 +37,7 @@ class AttentionConfig:
and buffers can be pre-allocated to avoid inflating the memory estimate."""
use_cudnn_prefill: bool = False
"""Whether to use cudnn prefill."""
"""Deprecated: cuDNN prefill backend has been removed."""
use_trtllm_ragged_deepseek_prefill: bool = False
"""Whether to use TRTLLM ragged deepseek prefill."""
@@ -42,12 +46,18 @@ class AttentionConfig:
"""If set to True/False, use or don't use the TRTLLM attention backend
in flashinfer. If None, auto-detect the attention backend in flashinfer."""
disable_flashinfer_prefill: bool = True
disable_flashinfer_prefill: bool | None = None
"""Whether to disable flashinfer prefill."""
disable_flashinfer_q_quantization: bool = False
"""If set, when using fp8 kv, do not quantize Q to fp8."""
mla_prefill_backend: MLAPrefillBackendEnum | None = None
"""MLA prefill backend to use. If None, will be selected automatically.
Valid options: FLASH_ATTN (FA3/FA4), FLASHINFER, TRTLLM_RAGGED.
This option supersedes use_trtllm_ragged_deepseek_prefill
and disable_flashinfer_prefill which are deprecated."""
use_prefill_query_quantization: bool = False
"""If set, quantize query for attention in prefill."""
@@ -84,3 +94,48 @@ class AttentionConfig:
return None
return AttentionBackendEnum[value.upper()]
return value
@field_validator("mla_prefill_backend", mode="before")
@classmethod
def validate_mla_prefill_backend_before(cls, value: Any) -> Any:
"""Enable parsing of the `mla_prefill_backend` enum type from string."""
if isinstance(value, str):
return MLAPrefillBackendEnum[value.upper()]
return value
def __post_init__(self) -> None:
self._migrate_deprecated_mla_prefill_flags()
def _migrate_deprecated_mla_prefill_flags(self) -> None:
"""Migrate deprecated MLA prefill flags to mla_prefill_backend."""
# If the new option is already set, it takes precedence
if self.mla_prefill_backend is not None:
return
# Check for deprecated flags and migrate them.
# Only the first flag encountered sets the backend.
if self.use_cudnn_prefill:
raise ValueError(
"The cuDNN MLA prefill backend has been removed. "
"Use --attention-config.mla_prefill_backend=FLASH_ATTN or "
"FLASHINFER or TRTLLM_RAGGED instead."
)
if self.use_trtllm_ragged_deepseek_prefill:
if self.mla_prefill_backend is None:
self.mla_prefill_backend = MLAPrefillBackendEnum.TRTLLM_RAGGED
logger.warning_once(
"use_trtllm_ragged_deepseek_prefill is deprecated and "
"will be removed in v0.22. Use "
"--attention-config.mla_prefill_backend=TRTLLM_RAGGED "
"instead."
)
if self.disable_flashinfer_prefill:
if self.mla_prefill_backend is None:
self.mla_prefill_backend = MLAPrefillBackendEnum.FLASH_ATTN
logger.warning_once(
"disable_flashinfer_prefill is deprecated and will be removed "
"in v0.22. Use --attention-config.mla_prefill_backend="
"FLASH_ATTN instead."
)
vllm/model_executor/layers/attention/mla_attention.py
+58 -595
View File
@@ -189,12 +189,9 @@ return curr_o @ W_O
import functools
from abc import abstractmethod
from dataclasses import dataclass, field
from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast
if TYPE_CHECKING:
from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
from typing import ClassVar, Generic, TypeVar, cast
import torch
import torch.nn as nn
@@ -242,7 +239,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kNvfp4Dynamic,
)
from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer, has_nvidia_artifactory
from vllm.utils.flashinfer import has_flashinfer
from vllm.utils.math_utils import cdiv, round_down
from vllm.utils.torch_utils import (
LayerNameType,
@@ -262,11 +259,9 @@ from vllm.v1.attention.backend import (
MLAAttentionImpl,
SparseMLAAttentionImpl,
)
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
from vllm.v1.attention.backends.mla.prefill import MLAPrefillBackend
from vllm.v1.attention.backends.utils import (
get_dcp_local_seq_lens,
get_per_layer_parameters,
infer_global_hyperparameters,
split_decodes_and_prefills,
)
from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
@@ -1123,33 +1118,6 @@ class QueryLenSupport(Enum):
VARLEN = "varlen"
try:
from vllm.vllm_flash_attn import ( # type: ignore[attr-defined]
flash_attn_varlen_func,
)
is_vllm_fa = True
except ImportError:
is_vllm_fa = False
flash_attn_varlen_func = None # type: ignore[assignment]
# On ROCm, vllm_flash_attn is not available, try upstream flash_attn instead.
# On CUDA, vllm_flash_attn should always be available (built with vLLM),
# so we don't attempt the fallback there.
if current_platform.is_rocm():
try:
from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]
except ImportError:
logger.debug(
"flash_attn not available on ROCm; "
"MLA models using TRITON_MLA will require flash_attn. "
"AITER_MLA backends use aiter kernels instead."
)
elif current_platform.is_xpu():
from vllm._xpu_ops import xpu_ops
flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func # type: ignore[no-redef,attr-defined,assignment]
def dynamic_per_batched_tensor_quant(
x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn
):
@@ -1161,9 +1129,6 @@ def dynamic_per_batched_tensor_quant(
return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()
logger = init_logger(__name__)
@CustomOp.register(
"mla_decode_concat_quant_fp8",
dynamic_arg_dims={"decode_ql_nope": 0, "decode_q_pe": 0},
@@ -1197,9 +1162,6 @@ class _DecodeConcatQuantFP8(QuantFP8):
forward_hip = _make_forward(QuantFP8.forward_hip) # type: ignore[arg-type]
CUDNN_WORKSPACE_SIZE = 12800
class MLACommonBackend(AttentionBackend):
@staticmethod
def get_name() -> str:
@@ -1268,26 +1230,9 @@ class MLACommonPrefillMetadata:
query_start_loc: torch.Tensor
max_query_len: int
chunked_context: ChunkedContextMetadata | None = None
query_seq_lens: torch.Tensor | None = None
workspace_buffer: torch.Tensor | None = None
q_data_type: torch.dtype | None = None
output_dtype: torch.dtype | None = None
@dataclass
class FlashInferPrefillMetadata(MLACommonPrefillMetadata):
prefill_main: "BatchPrefillWithRaggedKVCacheWrapper | None" = None
prefill_chunks: "list[BatchPrefillWithRaggedKVCacheWrapper]" = field(
default_factory=list
)
@dataclass
class CudnnPrefillMetadata(MLACommonPrefillMetadata):
class ChunkedContextMetadata(MLACommonPrefillMetadata.ChunkedContextMetadata):
seq_lens: torch.Tensor
cudnn_workspace: torch.Tensor | None = None
prefill_backend: MLAPrefillBackend | None = None
@dataclass
@@ -1333,13 +1278,8 @@ class MLACommonMetadata(AttentionMetadata, Generic[D]):
# The dimension of the attention heads
head_dim: int | None = None
prefill: MLACommonPrefillMetadata | None = None
decode: D | None = None
prefill: (
MLACommonPrefillMetadata
| FlashInferPrefillMetadata
| CudnnPrefillMetadata
| None
) = None
def __post_init__(self):
if self.head_dim is not None and not MLACommonBackend.supports_head_size(
@@ -1352,64 +1292,6 @@ M = TypeVar("M", bound=MLACommonMetadata)
A = TypeVar("A", bound=AttentionMetadata)
def is_deepseek_r1_mla_compatible(vllm_config: VllmConfig) -> bool:
# Check if model has DeepSeek R1 compatible MLA dimensions:
# qk_nope_head_dim = 128, qk_rope_head_dim = 64, v_head_dim = 128
# which results in query/key head dim = 192.
if vllm_config.model_config is None:
return False
hf_text_config = vllm_config.model_config.hf_text_config
qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
qk_rope_head_dim = getattr(hf_text_config, "qk_rope_head_dim", 1)
v_head_dim = getattr(hf_text_config, "v_head_dim", 1)
return qk_nope_head_dim == 128 and qk_rope_head_dim == 64 and v_head_dim == 128
@functools.cache
def use_flashinfer_prefill() -> bool:
from vllm.config import get_current_vllm_config
vllm_config = get_current_vllm_config()
if not (
not vllm_config.attention_config.disable_flashinfer_prefill
and has_flashinfer()
and not vllm_config.attention_config.use_cudnn_prefill
and current_platform.is_device_capability_family(100)
):
return False
return is_deepseek_r1_mla_compatible(vllm_config)
@functools.cache
def use_cudnn_prefill() -> bool:
from vllm.config import get_current_vllm_config
vllm_config = get_current_vllm_config()
return (
has_flashinfer()
and vllm_config.attention_config.use_cudnn_prefill
and current_platform.is_device_capability_family(100)
and has_nvidia_artifactory()
)
@functools.cache
def use_trtllm_ragged_deepseek_prefill() -> bool:
"""Check if TRT-LLM ragged DeepSeek prefill should be used."""
from vllm.config import get_current_vllm_config
vllm_config = get_current_vllm_config()
if not (
has_flashinfer()
and vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill
and current_platform.is_device_capability_family(100)
):
return False
return is_deepseek_r1_mla_compatible(vllm_config)
@dataclass
class MLADims:
q_lora_rank: int | None
@@ -1447,15 +1329,14 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims:
@functools.cache
def backend_supports_prefill_query_quantization() -> bool:
"""Check if the selected MLA backend supports prefill query quantization.
"""Check if the selected MLA prefill backend supports query quantization.
Currently supported backends:
- FlashInfer prefill
- TRT-LLM ragged DeepSeek prefill
- FlashInfer
- TRT-LLM Ragged
Not supported:
- cuDNN Prefill
- FlashAttention
- FlashAttention (FA3/FA4)
- Non-GB200 devices (FP8 prefill requires device capability 100)
"""
# FP8 prefill query quantization requires GB200 (device capability 100)
@@ -1463,7 +1344,15 @@ def backend_supports_prefill_query_quantization() -> bool:
if not current_platform.is_device_capability_family(100):
return False
return use_flashinfer_prefill() or use_trtllm_ragged_deepseek_prefill()
from vllm.config import get_current_vllm_config
from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
vllm_config = get_current_vllm_config()
backend_cls = get_mla_prefill_backend(vllm_config)
return backend_cls.get_name() in (
"FLASHINFER",
"TRTLLM_RAGGED",
)
class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
@@ -1574,7 +1463,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
metadata_cls if metadata_cls is not None else MLACommonMetadata
)
self.kv_cache_spec = kv_cache_spec
scheduler_config = vllm_config.scheduler_config
self.model_config = vllm_config.model_config
parallel_config = vllm_config.parallel_config
self.compilation_config = vllm_config.compilation_config
@@ -1634,139 +1522,32 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
device=device,
)
self._use_cudnn_prefill = use_cudnn_prefill()
self._use_fi_prefill = use_flashinfer_prefill()
self._use_trtllm_ragged_prefill = use_trtllm_ragged_deepseek_prefill()
self.prefill_metadata_cls = (
FlashInferPrefillMetadata
if self._use_fi_prefill
else CudnnPrefillMetadata
if self._use_cudnn_prefill
else MLACommonPrefillMetadata
from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
prefill_backend_cls = get_mla_prefill_backend(vllm_config)
self._prefill_backend = prefill_backend_cls(
num_heads=self.num_heads,
scale=self.model_config.get_head_size() ** -0.5,
kv_lora_rank=self.mla_dims.kv_lora_rank,
qk_nope_head_dim=self.mla_dims.qk_nope_head_dim,
qk_rope_head_dim=self.mla_dims.qk_rope_head_dim,
v_head_dim=self.mla_dims.v_head_dim,
vllm_config=vllm_config,
device=device,
layer_names=layer_names,
)
if self._use_fi_prefill:
self._workspace_buffer = torch.empty(
envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
dtype=torch.uint8,
device=device,
)
self._fi_prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None
self._fi_prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []
self._global_hyperparameters = infer_global_hyperparameters(
get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl) # type: ignore[type-abstract]
)
if self._use_trtllm_ragged_prefill:
self._workspace_buffer = torch.empty(
envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
dtype=torch.uint8,
device=device,
)
if self._use_cudnn_prefill:
self.cudnn_workspace = torch.empty(
CUDNN_WORKSPACE_SIZE * scheduler_config.max_num_seqs,
dtype=torch.int8,
device=device,
)
supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY
self._init_reorder_batch_threshold(
self.reorder_batch_threshold, supports_spec_decode, supports_dcp_with_varlen
)
# Validate consistency between query_len_support and reorder_batch_threshold
if self.query_len_support == QueryLenSupport.SINGLE_ONLY:
assert self.reorder_batch_threshold == 1, (
f"reorder_batch_threshold must be 1 when query_len_support is "
f"SINGLE_ONLY, got {self.reorder_batch_threshold}"
)
def _build_fi_prefill_wrappers(self, prefill: FlashInferPrefillMetadata):
qo_indptr = prefill.query_start_loc
has_context = False
if prefill.chunked_context is not None:
chunked_context = prefill.chunked_context
has_context = True
if self._fi_prefill_main is None:
from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
self._fi_prefill_main = BatchPrefillWithRaggedKVCacheWrapper(
self._workspace_buffer, "NHD", backend="cutlass"
)
if has_context:
num_chunks = chunked_context.cu_seq_lens.shape[0]
# Allocate more prefill chunk wrappers if needed
if len(self._fi_prefill_chunks) < num_chunks:
from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
for _ in range(len(self._fi_prefill_chunks), num_chunks):
self._fi_prefill_chunks.append(
BatchPrefillWithRaggedKVCacheWrapper(
self._workspace_buffer, "NHD", backend="cutlass"
)
)
assert num_chunks <= len(self._fi_prefill_chunks)
# In MLA, the non-latent num_qo_heads == num_kv_heads
num_qo_heads = self.num_heads
num_kv_heads = num_qo_heads
# Sanity: Verify that num_kv_heads == 1 since it is latent space
assert self.kv_cache_spec.num_kv_heads == 1
# Get non-latent head_dim_qk and head_dim_vo
head_dim_qk = self.mla_dims.qk_nope_head_dim + self.mla_dims.qk_rope_head_dim
head_dim_vo = self.mla_dims.v_head_dim
# For main run, qo_indptr == kv_indptr
kv_indptr = qo_indptr.clone()
# Prepare main prefill
self._fi_prefill_main.plan(
qo_indptr=qo_indptr,
kv_indptr=kv_indptr,
num_qo_heads=num_qo_heads,
num_kv_heads=num_kv_heads,
head_dim_qk=head_dim_qk,
head_dim_vo=head_dim_vo,
causal=True, # This is main run
sm_scale=self._global_hyperparameters.sm_scale,
window_left=self._global_hyperparameters.window_left,
logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
q_data_type=self.q_data_type,
o_data_type=prefill.output_dtype,
)
# Prepare context prefills
if has_context:
for i in range(num_chunks):
kv_indptr_chunk = chunked_context.cu_seq_lens[i]
self._fi_prefill_chunks[i].plan(
qo_indptr=qo_indptr,
kv_indptr=kv_indptr_chunk,
num_qo_heads=num_qo_heads,
num_kv_heads=num_kv_heads,
head_dim_qk=head_dim_qk,
head_dim_vo=head_dim_vo,
causal=False, # This is context run
sm_scale=self._global_hyperparameters.sm_scale,
window_left=self._global_hyperparameters.window_left,
logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
q_data_type=self.q_data_type,
o_data_type=prefill.output_dtype,
)
prefill.prefill_main = self._fi_prefill_main
prefill.prefill_chunks = self._fi_prefill_chunks
def _build_decode(
self,
block_table_tensor: torch.Tensor,
@@ -1972,18 +1753,14 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
dtype=torch.int32,
)
chunked_context_metadata_cls = (
CudnnPrefillMetadata.ChunkedContextMetadata
if self._use_cudnn_prefill
else MLACommonPrefillMetadata.ChunkedContextMetadata
)
prefill_tokens_with_context = None
if num_prefills_with_context_cpu > 0:
prefill_tokens_with_context = prefill_query_start_loc_cpu[
num_prefills_with_context_cpu
].item()
_ChunkedMetadata = MLACommonPrefillMetadata.ChunkedContextMetadata
if self.dcp_world_size > 1:
chunked_context_metadata = chunked_context_metadata_cls(
chunked_context_metadata = _ChunkedMetadata(
cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
starts=local_chunk_starts.to(device, non_blocking=True),
seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),
@@ -2004,7 +1781,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
prefill_tokens_with_context=prefill_tokens_with_context,
)
else:
chunked_context_metadata = chunked_context_metadata_cls(
chunked_context_metadata = _ChunkedMetadata(
cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
starts=chunk_starts.to(device, non_blocking=True),
seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
@@ -2018,35 +1795,22 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
prefill_tokens_with_context=prefill_tokens_with_context,
)
if self._use_cudnn_prefill:
chunked_context_metadata.seq_lens = chunk_seq_lens
assert (
max(chunked_context_metadata.max_seq_lens)
<= self.chunked_prefill_workspace_size
)
prefill_metadata = self.prefill_metadata_cls(
prefill_metadata = MLACommonPrefillMetadata(
block_table=block_table_tensor[reqs_start:, ...],
query_start_loc=prefill_query_start_loc,
max_query_len=max_query_len,
chunked_context=chunked_context_metadata,
output_dtype=self.model_config.dtype,
q_data_type=self.q_data_type,
prefill_backend=self._prefill_backend,
)
if self._use_cudnn_prefill:
assert isinstance(prefill_metadata, CudnnPrefillMetadata)
prefill_metadata.query_seq_lens = (
prefill_query_start_loc[1:] - prefill_query_start_loc[:-1]
)
prefill_metadata.cudnn_workspace = self.cudnn_workspace
if self._use_trtllm_ragged_prefill:
prefill_metadata.query_seq_lens = (
prefill_query_start_loc[1:] - prefill_query_start_loc[:-1]
)
prefill_metadata.workspace_buffer = self._workspace_buffer
self._prefill_backend.prepare_metadata(prefill_metadata)
decode_metadata = None
if num_decodes > 0:
@@ -2091,10 +1855,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
decode=decode_metadata,
)
if self._use_fi_prefill and num_prefills > 0:
assert isinstance(attn_metadata.prefill, FlashInferPrefillMetadata)
self._build_fi_prefill_wrappers(attn_metadata.prefill)
return attn_metadata # type: ignore[return-value]
@@ -2240,308 +2000,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
and (self.qk_rope_head_dim == 64)
)
if use_trtllm_ragged_deepseek_prefill():
logger.info_once("Using TRT-LLM ragged DeepSeek prefill for MLA")
self._run_prefill_context_chunk = (
self._run_prefill_context_chunk_trtllm_ragged
)
self._run_prefill_new_tokens = self._run_prefill_new_tokens_trtllm_ragged
self._pad_v = False
elif use_flashinfer_prefill():
logger.info_once("Using FlashInfer prefill for MLA")
self._run_prefill_context_chunk = self._run_prefill_context_chunk_fi
self._run_prefill_new_tokens = self._run_prefill_new_tokens_fi
self._pad_v = False
elif use_cudnn_prefill():
logger.info_once("Using CUDNN prefill for MLA")
self._run_prefill_context_chunk = self._run_prefill_context_chunk_cudnn
self._run_prefill_new_tokens = self._run_prefill_new_tokens_cudnn
self._pad_v = False
else: # Use FlashAttention
if flash_attn_varlen_func is None:
raise RuntimeError(
"MLA attention requires FlashAttention but it is not "
"available. Please install flash_attn or use "
"--attention-backend ROCM_AITER_MLA."
)
logger.info_once("Using FlashAttention prefill for MLA")
self._run_prefill_context_chunk = self._run_prefill_context_chunk_fa
self._run_prefill_new_tokens = self._run_prefill_new_tokens_fa
# Handle the differences between the flash_attn_varlen from
# flash_attn and the one from vllm_flash_attn. The former is used on
# RoCM and the latter has an additional parameter to control
# FA2 vs FA3
self.flash_attn_varlen_func = flash_attn_varlen_func
self.vllm_flash_attn_version = get_flash_attn_version(
head_size=self.qk_head_dim
)
if self.vllm_flash_attn_version is not None:
self.flash_attn_varlen_func = functools.partial(
flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version
)
# For MLA the v head dim is smaller than qk head dim so we pad out
# v with 0s to match the qk head dim for attention backends that do
# not support different headdims.
# FA3 on Hopper (SM90) and FA4 natively handle diff headdims.
device_capability = current_platform.get_device_capability()
self._pad_v = self.vllm_flash_attn_version is None or not (
(
self.vllm_flash_attn_version == 3
and device_capability is not None
and device_capability[0] == 9
)
or self.vllm_flash_attn_version == 4
)
self.dcp_world_size: int = -1
self.cp_kv_cache_interleave_size: int = (
get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
)
def _flash_attn_varlen_diff_headdims(
self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs
):
maybe_padded_v = v
if self._pad_v:
maybe_padded_v = torch.nn.functional.pad(
v, [0, q.shape[-1] - v.shape[-1]], value=0
)
if is_vllm_fa:
kwargs["return_softmax_lse"] = return_softmax_lse
else:
# ROCm leverages the upstream flash_attn, which takes a parameter
# called "return_attn_probs" instead of return_softmax_lse
kwargs["return_attn_probs"] = return_softmax_lse
if envs.VLLM_BATCH_INVARIANT:
kwargs["num_splits"] = 1
attn_out = self.flash_attn_varlen_func(
q=q,
k=k,
v=maybe_padded_v,
softmax_scale=softmax_scale,
**kwargs,
)
# Unpack the output if there is multiple results
lse = None
if isinstance(attn_out, tuple):
attn_out, lse = attn_out[0], attn_out[1]
# Remain consistent with old `flash_attn_varlen_func` where there
# is only one output tensor if `return_softmax_lse` is False.
if return_softmax_lse:
return attn_out, lse
return attn_out
def _run_prefill_new_tokens_fa(
self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
):
return self._flash_attn_varlen_diff_headdims(
q=q,
k=k,
v=v,
cu_seqlens_q=prefill.query_start_loc,
cu_seqlens_k=prefill.query_start_loc,
max_seqlen_q=prefill.max_query_len,
max_seqlen_k=prefill.max_query_len,
softmax_scale=self.scale,
causal=True,
return_softmax_lse=return_softmax_lse,
)
def _run_prefill_new_tokens_fi(
self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
):
assert isinstance(prefill, FlashInferPrefillMetadata)
assert prefill.prefill_main is not None
ret = prefill.prefill_main.run(
q=q,
k=k,
v=v,
return_lse=return_softmax_lse,
)
if isinstance(ret, tuple):
return ret[0], ret[1].transpose(0, 1).contiguous()
return ret
def _run_prefill_new_tokens_cudnn(
self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
):
assert isinstance(prefill, CudnnPrefillMetadata)
assert prefill.query_seq_lens is not None
from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
output, lse = cudnn_batch_prefill_with_kv_cache(
q=q,
k_cache=k,
v_cache=v,
scale=self.scale,
workspace_buffer=prefill.cudnn_workspace,
max_token_per_sequence=prefill.max_query_len,
max_sequence_kv=prefill.max_query_len,
actual_seq_lens_q=prefill.query_seq_lens.view(-1, 1, 1, 1),
actual_seq_lens_kv=prefill.query_seq_lens.view(-1, 1, 1, 1),
causal=True,
# Do not support False for now
return_lse=True,
# Indicates actual_seq_lens are on GPU or CPU.
is_cuda_graph_compatible=True,
)
if return_softmax_lse:
return output, lse
return output
def _run_prefill_context_chunk_fa(
self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
):
assert prefill.chunked_context is not None
return self._flash_attn_varlen_diff_headdims(
q=q,
k=k,
v=v,
cu_seqlens_q=prefill.query_start_loc,
cu_seqlens_k=prefill.chunked_context.cu_seq_lens[chunk_idx],
max_seqlen_q=prefill.max_query_len,
max_seqlen_k=prefill.chunked_context.max_seq_lens[chunk_idx],
softmax_scale=self.scale,
causal=False, # Context is unmasked
return_softmax_lse=True,
)
def _run_prefill_context_chunk_fi(
self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
):
assert isinstance(prefill, FlashInferPrefillMetadata)
attn_out, lse = prefill.prefill_chunks[chunk_idx].run(
q=q,
k=k,
v=v,
return_lse=True,
)
# Convert from (q_len, num_heads) to (num_heads, q_len)
return attn_out, lse.transpose(0, 1).contiguous()
def _run_prefill_context_chunk_cudnn(
self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
):
assert isinstance(prefill, CudnnPrefillMetadata)
assert prefill.chunked_context is not None
assert prefill.chunked_context.seq_lens[chunk_idx] is not None
assert prefill.query_seq_lens is not None
from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
return cudnn_batch_prefill_with_kv_cache(
q=q,
k_cache=k,
v_cache=v,
scale=self.scale,
workspace_buffer=prefill.cudnn_workspace,
max_token_per_sequence=prefill.max_query_len,
max_sequence_kv=prefill.chunked_context.max_seq_lens[chunk_idx],
actual_seq_lens_q=prefill.query_seq_lens.view(-1, 1, 1, 1),
actual_seq_lens_kv=prefill.chunked_context.seq_lens[chunk_idx].view(
-1, 1, 1, 1
),
causal=False,
return_lse=True,
# Indicates actual_seq_lens are on GPU or CPU.
is_cuda_graph_compatible=True,
)
def _run_prefill_new_tokens_trtllm_ragged(
self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
):
"""TRT-LLM ragged attention for new tokens (causal)."""
from flashinfer.prefill import trtllm_ragged_attention_deepseek
assert prefill.query_seq_lens is not None
assert prefill.workspace_buffer is not None
# allocate BF16 / FP16 output tensor for TRT-LLM ragged attention
out = torch.empty(
q.shape[0],
q.shape[1],
v.shape[2],
device=q.device,
dtype=prefill.output_dtype,
)
ret = trtllm_ragged_attention_deepseek(
query=q,
key=k,
value=v,
workspace_buffer=prefill.workspace_buffer,
seq_lens=prefill.query_seq_lens,
max_q_len=prefill.max_query_len,
max_kv_len=prefill.max_query_len,
bmm1_scale=self.scale,
bmm2_scale=1.0,
o_sf_scale=1.0,
batch_size=prefill.query_seq_lens.shape[0],
window_left=-1,
cum_seq_lens_q=prefill.query_start_loc,
cum_seq_lens_kv=prefill.query_start_loc,
enable_pdl=False,
is_causal=True,
return_lse=return_softmax_lse,
out=out,
)
if isinstance(ret, tuple):
# Convert from (q_len, num_heads) to (num_heads, q_len)
return ret[0], ret[1].transpose(0, 1).contiguous()
return ret
def _run_prefill_context_chunk_trtllm_ragged(
self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
):
"""TRT-LLM ragged attention for context chunks (non-causal)."""
from flashinfer.prefill import trtllm_ragged_attention_deepseek
assert prefill.chunked_context is not None
assert prefill.chunked_context.seq_lens[chunk_idx] is not None
assert prefill.workspace_buffer is not None
out = torch.empty(
q.shape[0],
q.shape[1],
v.shape[2],
device=q.device,
dtype=prefill.output_dtype,
)
attn_out, lse = trtllm_ragged_attention_deepseek(
query=q,
key=k,
value=v,
workspace_buffer=prefill.workspace_buffer,
seq_lens=prefill.chunked_context.seq_lens[chunk_idx],
max_q_len=prefill.max_query_len,
max_kv_len=prefill.chunked_context.max_seq_lens[chunk_idx],
bmm1_scale=self.scale,
bmm2_scale=1.0,
o_sf_scale=1.0,
batch_size=prefill.chunked_context.seq_lens[chunk_idx].shape[0],
window_left=-1,
cum_seq_lens_q=prefill.query_start_loc,
cum_seq_lens_kv=prefill.chunked_context.cu_seq_lens[chunk_idx],
enable_pdl=False,
is_causal=False,
return_lse=True,
out=out,
)
# Convert from (q_len, num_heads) to (num_heads, q_len)
return attn_out, lse.transpose(0, 1).contiguous()
def _concat_k_nope_k_pe(
self, k_nope: torch.Tensor, k_pe: torch.Tensor
) -> torch.Tensor:
@@ -2582,6 +2046,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
):
assert attn_metadata.prefill is not None
prefill_metadata = attn_metadata.prefill
assert prefill_metadata.prefill_backend is not None
assert prefill_metadata.chunked_context is not None
use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
@@ -2649,12 +2114,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
k = self._concat_k_nope_k_pe(k_nope, k_pe)
attn_output, attn_softmax_lse = self._run_prefill_context_chunk(
prefill=prefill_metadata,
chunk_idx=i,
q=q,
k=k,
v=v,
attn_output, attn_softmax_lse = (
prefill_metadata.prefill_backend.run_prefill_context_chunk(
chunk_idx=i,
q=q,
k=k,
v=v,
)
)
if output is None:
@@ -2687,6 +2153,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
assert k_scale is None, "DCP not support scaled kvcache now."
assert attn_metadata.prefill is not None
prefill_metadata = attn_metadata.prefill
assert prefill_metadata.prefill_backend is not None
assert prefill_metadata.chunked_context is not None
assert prefill_metadata.chunked_context.padded_local_chunk_seq_lens is not None
assert prefill_metadata.chunked_context.local_context_lens_allranks is not None
@@ -2753,12 +2220,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
k = self._concat_k_nope_k_pe(k_nope, k_pe)
attn_output, attn_softmax_lse = self._run_prefill_context_chunk(
prefill=prefill_metadata,
chunk_idx=i,
q=q,
k=k,
v=v,
attn_output, attn_softmax_lse = (
prefill_metadata.prefill_backend.run_prefill_context_chunk(
chunk_idx=i,
q=q,
k=k,
v=v,
)
)
if output is None:
@@ -2790,11 +2258,11 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
k_scale: torch.Tensor,
output: torch.Tensor,
) -> None:
# TODO (zyongye): Prefill function here
assert attn_metadata.prefill is not None
assert self.dcp_world_size != -1
prefill_metadata = attn_metadata.prefill
assert prefill_metadata.prefill_backend is not None
use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
# Convert q to FP8 if FP8 prefill attention is enabled
@@ -2813,8 +2281,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
k = k.to(prefill_metadata.q_data_type)
v = v.to(prefill_metadata.q_data_type)
output_prefill = self._run_prefill_new_tokens(
prefill=prefill_metadata,
output_prefill = prefill_metadata.prefill_backend.run_prefill_new_tokens(
q=q,
k=k,
v=v,
@@ -2839,11 +2306,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
q, kv_c_and_k_pe_cache, attn_metadata, k_scale
)
# unpad if necessary
if self._pad_v:
context_output = context_output[..., : v.shape[-1]]
suffix_output = suffix_output[..., : v.shape[-1]]
output = output.view(-1, self.num_heads, self.v_head_dim)
merge_attn_states(
output=output,
@@ -2854,7 +2316,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
prefill_tokens_with_context=prefill_metadata.chunked_context.prefill_tokens_with_context,
)
else:
output_prefill = output_prefill[..., : v.shape[-1]].flatten(start_dim=-2)
assert isinstance(output_prefill, torch.Tensor)
output_prefill = output_prefill.flatten(start_dim=-2)
output.copy_(output_prefill)
@abstractmethod
vllm/platforms/interface.py
+3
View File
@@ -86,6 +86,9 @@ class DeviceCapability(NamedTuple):
return NotImplemented
return (self.major, self.minor) > (other.major, other.minor)
def __hash__(self) -> int:
return hash((self.major, self.minor))
def as_version_str(self) -> str:
return f"{self.major}.{self.minor}"
vllm/v1/attention/backends/mla/prefill/__init__.py
+11
View File
@@ -0,0 +1,11 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
from vllm.v1.attention.backends.mla.prefill.selector import get_mla_prefill_backend
__all__ = [
"MLAPrefillBackend",
"MLAPrefillBackendEnum",
"get_mla_prefill_backend",
]
vllm/v1/attention/backends/mla/prefill/base.py
+125
View File
@@ -0,0 +1,125 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Abstract base class for MLA prefill backends."""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, ClassVar
import torch
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.model_executor.layers.attention.mla_attention import (
MLACommonPrefillMetadata,
)
from vllm.platforms.interface import DeviceCapability
from vllm.v1.attention.backends.mla.prefill.selector import (
MLAPrefillSelectorConfig,
)
class MLAPrefillBackend(ABC):
"""Abstract base class for MLA prefill backends."""
supported_dtypes: ClassVar[list[torch.dtype]] = [
torch.float16,
torch.bfloat16,
]
requires_r1_mla_dimensions: ClassVar[bool] = False
@staticmethod
@abstractmethod
def get_name() -> str:
raise NotImplementedError
@classmethod
def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:
return True
@classmethod
def supports_dtype(cls, dtype: torch.dtype) -> bool:
return dtype in cls.supported_dtypes
@classmethod
def is_available(cls) -> bool:
return True
@classmethod
def validate_configuration(
cls,
device_capability: "DeviceCapability",
selector_config: "MLAPrefillSelectorConfig",
) -> list[str]:
invalid_reasons: list[str] = []
if not cls.supports_compute_capability(device_capability):
invalid_reasons.append(
f"compute capability {device_capability.major}."
f"{device_capability.minor} not supported"
)
if not cls.supports_dtype(selector_config.dtype):
invalid_reasons.append(f"dtype {selector_config.dtype} not supported")
if not cls.is_available():
invalid_reasons.append("required dependencies not available")
if cls.requires_r1_mla_dimensions and not selector_config.is_r1_compatible:
invalid_reasons.append(
"model does not have DeepSeek R1 MLA dimensions "
"(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128)"
)
return invalid_reasons
def __init__(
self,
num_heads: int,
scale: float,
kv_lora_rank: int,
qk_nope_head_dim: int,
qk_rope_head_dim: int,
v_head_dim: int,
vllm_config: "VllmConfig",
device: torch.device,
layer_names: list[str] | None = None,
) -> None:
self.num_heads = num_heads
self.scale = scale
self.kv_lora_rank = kv_lora_rank
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_rope_head_dim = qk_rope_head_dim
self.v_head_dim = v_head_dim
self.vllm_config = vllm_config
self.device = device
self.layer_names = layer_names
def prepare_metadata( # noqa: B027
self,
prefill_metadata: "MLACommonPrefillMetadata",
) -> None:
"""Prepare backend-specific metadata before the forward pass.
Called by the metadata builder after constructing the prefill metadata.
"""
self._prefill_metadata = prefill_metadata
@abstractmethod
def run_prefill_new_tokens(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
return_softmax_lse: bool,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
raise NotImplementedError
@abstractmethod
def run_prefill_context_chunk(
self,
chunk_idx: int,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
raise NotImplementedError
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+180
View File
@@ -0,0 +1,180 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""FlashAttention backend for MLA prefill."""
import functools
from typing import TYPE_CHECKING
import torch
import vllm.envs as envs
from vllm.platforms import current_platform
from vllm.v1.attention.backends.fa_utils import (
get_flash_attn_version,
is_flash_attn_varlen_func_available,
)
from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
if TYPE_CHECKING:
from vllm.config import VllmConfig
if is_flash_attn_varlen_func_available():
from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
else:
flash_attn_varlen_func = None # type: ignore[assignment]
class FlashAttnPrefillBackend(MLAPrefillBackend):
"""FlashAttention backend for MLA prefill."""
@staticmethod
def get_name() -> str:
return "FLASH_ATTN"
@classmethod
def is_available(cls) -> bool:
return is_flash_attn_varlen_func_available()
def __init__(
self,
num_heads: int,
scale: float,
kv_lora_rank: int,
qk_nope_head_dim: int,
qk_rope_head_dim: int,
v_head_dim: int,
vllm_config: "VllmConfig",
device: torch.device,
layer_names: list[str] | None = None,
) -> None:
super().__init__(
num_heads=num_heads,
scale=scale,
kv_lora_rank=kv_lora_rank,
qk_nope_head_dim=qk_nope_head_dim,
qk_rope_head_dim=qk_rope_head_dim,
v_head_dim=v_head_dim,
vllm_config=vllm_config,
device=device,
layer_names=layer_names,
)
# Handle the differences between the flash_attn_varlen from
# flash_attn and the one from vllm_flash_attn
assert flash_attn_varlen_func is not None, (
"FlashAttnPrefillBackend requires flash_attn_varlen_func. "
"Ensure FlashAttnPrefillBackend.is_available() is checked first."
)
qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
self.flash_attn_varlen_func = flash_attn_varlen_func
self.vllm_flash_attn_version = get_flash_attn_version(head_size=qk_head_dim)
if self.vllm_flash_attn_version is not None:
self.flash_attn_varlen_func = functools.partial(
flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version
)
# Determine if we need to pad V
# For MLA the v head dim is smaller than qk head dim so we pad out
# v with 0s to match the qk head dim for attention backends that do
# not support different headdims.
# FA3 on Hopper (SM90) and FA4 natively handle diff headdims.
device_capability = current_platform.get_device_capability()
self.requires_v_padding = self.vllm_flash_attn_version is None or not (
(
self.vllm_flash_attn_version == 3
and device_capability is not None
and device_capability[0] == 9
)
or self.vllm_flash_attn_version == 4
)
# Track whether we're using vllm's FA or upstream (for ROCm)
self._is_vllm_fa = current_platform.is_cuda() or current_platform.is_xpu()
def _flash_attn_varlen_diff_headdims(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
return_softmax_lse: bool = False,
softmax_scale: float | None = None,
**kwargs,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
maybe_padded_v = v
if self.requires_v_padding:
maybe_padded_v = torch.nn.functional.pad(
v, [0, q.shape[-1] - v.shape[-1]], value=0
)
if self._is_vllm_fa:
kwargs["return_softmax_lse"] = return_softmax_lse
else:
# ROCm leverages the upstream flash_attn, which takes a parameter
# called "return_attn_probs" instead of return_softmax_lse
kwargs["return_attn_probs"] = return_softmax_lse
if envs.VLLM_BATCH_INVARIANT:
kwargs["num_splits"] = 1
attn_out = self.flash_attn_varlen_func(
q=q,
k=k,
v=maybe_padded_v,
softmax_scale=softmax_scale,
**kwargs,
)
# Unpack the output if there are multiple results
lse = None
if isinstance(attn_out, tuple):
attn_out, lse = attn_out[0], attn_out[1]
# Unpad output back to v_head_dim if we padded V
if self.requires_v_padding:
attn_out = attn_out[..., : v.shape[-1]]
# Remain consistent with old `flash_attn_varlen_func` where there
# is only one output tensor if `return_softmax_lse` is False.
if return_softmax_lse:
return attn_out, lse
return attn_out
def run_prefill_new_tokens(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
return_softmax_lse: bool,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
return self._flash_attn_varlen_diff_headdims(
q=q,
k=k,
v=v,
cu_seqlens_q=self._prefill_metadata.query_start_loc,
cu_seqlens_k=self._prefill_metadata.query_start_loc,
max_seqlen_q=self._prefill_metadata.max_query_len,
max_seqlen_k=self._prefill_metadata.max_query_len,
softmax_scale=self.scale,
causal=True,
return_softmax_lse=return_softmax_lse,
)
def run_prefill_context_chunk(
self,
chunk_idx: int,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
assert self._prefill_metadata.chunked_context is not None
return self._flash_attn_varlen_diff_headdims(
q=q,
k=k,
v=v,
cu_seqlens_q=self._prefill_metadata.query_start_loc,
cu_seqlens_k=self._prefill_metadata.chunked_context.cu_seq_lens[chunk_idx],
max_seqlen_q=self._prefill_metadata.max_query_len,
max_seqlen_k=self._prefill_metadata.chunked_context.max_seq_lens[chunk_idx],
softmax_scale=self.scale,
causal=False, # Context is unmasked
return_softmax_lse=True,
)
vllm/v1/attention/backends/mla/prefill/flashinfer.py
+211
View File
@@ -0,0 +1,211 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""FlashInfer backend for MLA prefill."""
from typing import TYPE_CHECKING
import torch
import vllm.envs as envs
from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
from vllm.v1.attention.backends.utils import (
get_per_layer_parameters,
infer_global_hyperparameters,
)
from vllm.v1.worker.workspace import current_workspace_manager
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.model_executor.layers.attention.mla_attention import (
MLACommonPrefillMetadata,
)
from vllm.platforms.interface import DeviceCapability
try:
from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
except ImportError:
BatchPrefillWithRaggedKVCacheWrapper = object # type: ignore[misc,assignment]
_DEFAULT_NUM_CHUNKS = 32
class FlashInferPrefillBackend(MLAPrefillBackend):
"""FlashInfer backend for MLA prefill."""
requires_r1_mla_dimensions = True
@staticmethod
def get_name() -> str:
return "FLASHINFER"
@classmethod
def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:
return device_capability.major == 10
@classmethod
def is_available(cls) -> bool:
try:
from flashinfer import (
BatchPrefillWithRaggedKVCacheWrapper, # noqa: F401
)
return True
except ImportError:
return False
def __init__(
self,
num_heads: int,
scale: float,
kv_lora_rank: int,
qk_nope_head_dim: int,
qk_rope_head_dim: int,
v_head_dim: int,
vllm_config: "VllmConfig",
device: torch.device,
layer_names: list[str] | None = None,
) -> None:
super().__init__(
num_heads=num_heads,
scale=scale,
kv_lora_rank=kv_lora_rank,
qk_nope_head_dim=qk_nope_head_dim,
qk_rope_head_dim=qk_rope_head_dim,
v_head_dim=v_head_dim,
vllm_config=vllm_config,
device=device,
layer_names=layer_names,
)
self._prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None
self._prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []
if layer_names is None:
raise ValueError(
"FlashInferPrefillBackend requires layer_names to "
"initialize global hyperparameters."
)
from vllm.model_executor.layers.attention.mla_attention import (
MLACommonImpl,
)
self._global_hyperparameters = infer_global_hyperparameters(
get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl) # type: ignore[type-abstract]
)
def _ensure_chunks(
self,
num_chunks: int,
workspace_buffer: torch.Tensor,
) -> None:
if len(self._prefill_chunks) < num_chunks:
for _ in range(len(self._prefill_chunks), num_chunks):
self._prefill_chunks.append(
BatchPrefillWithRaggedKVCacheWrapper(
workspace_buffer, "NHD", backend="cutlass"
)
)
def prepare_metadata(
self,
prefill_metadata: "MLACommonPrefillMetadata",
) -> None:
qo_indptr = prefill_metadata.query_start_loc
has_context = prefill_metadata.chunked_context is not None
(workspace_buffer,) = current_workspace_manager().get_simultaneous(
((envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), torch.uint8),
)
if self._prefill_main is None:
self._prefill_main = BatchPrefillWithRaggedKVCacheWrapper(
workspace_buffer, "NHD", backend="cutlass"
)
self._ensure_chunks(_DEFAULT_NUM_CHUNKS, workspace_buffer)
if has_context:
chunked_context = prefill_metadata.chunked_context
assert chunked_context is not None
num_chunks = chunked_context.cu_seq_lens.shape[0]
self._ensure_chunks(num_chunks, workspace_buffer)
num_qo_heads = self.num_heads
num_kv_heads = num_qo_heads
head_dim_qk = self.qk_nope_head_dim + self.qk_rope_head_dim
head_dim_vo = self.v_head_dim
kv_indptr = qo_indptr.clone()
assert self._prefill_main is not None
self._prefill_main.plan(
qo_indptr=qo_indptr,
kv_indptr=kv_indptr,
num_qo_heads=num_qo_heads,
num_kv_heads=num_kv_heads,
head_dim_qk=head_dim_qk,
head_dim_vo=head_dim_vo,
causal=True,
sm_scale=self._global_hyperparameters.sm_scale,
window_left=self._global_hyperparameters.window_left,
logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
q_data_type=prefill_metadata.q_data_type,
o_data_type=prefill_metadata.output_dtype,
)
if has_context:
chunked_context = prefill_metadata.chunked_context
assert chunked_context is not None
for i in range(num_chunks):
kv_indptr_chunk = chunked_context.cu_seq_lens[i]
self._prefill_chunks[i].plan(
qo_indptr=qo_indptr,
kv_indptr=kv_indptr_chunk,
num_qo_heads=num_qo_heads,
num_kv_heads=num_kv_heads,
head_dim_qk=head_dim_qk,
head_dim_vo=head_dim_vo,
causal=False,
sm_scale=self._global_hyperparameters.sm_scale,
window_left=self._global_hyperparameters.window_left,
logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
q_data_type=prefill_metadata.q_data_type,
o_data_type=prefill_metadata.output_dtype,
)
def run_prefill_new_tokens(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
return_softmax_lse: bool,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
assert self._prefill_main is not None
ret = self._prefill_main.run(
q=q,
k=k,
v=v,
return_lse=return_softmax_lse,
)
if isinstance(ret, tuple):
# Convert from (q_len, num_heads) to (num_heads, q_len)
return ret[0], ret[1].transpose(0, 1).contiguous()
return ret
def run_prefill_context_chunk(
self,
chunk_idx: int,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
attn_out, lse = self._prefill_chunks[chunk_idx].run(
q=q,
k=k,
v=v,
return_lse=True,
)
# Convert from (q_len, num_heads) to (num_heads, q_len)
return attn_out, lse.transpose(0, 1).contiguous()
vllm/v1/attention/backends/mla/prefill/registry.py
+53
View File
@@ -0,0 +1,53 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Registry for MLA prefill backends.
This module provides an enumeration of all available MLA prefill backends
and utilities for loading them.
"""
from enum import Enum, EnumMeta
from typing import TYPE_CHECKING
from vllm.utils.import_utils import resolve_obj_by_qualname
if TYPE_CHECKING:
from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
class _MLAPrefillBackendEnumMeta(EnumMeta):
"""Metaclass for MLAPrefillBackendEnum to provide better error messages."""
def __getitem__(cls, name: str):
try:
return super().__getitem__(name)
except KeyError:
members = cls.__members__.keys()
valid_backends = ", ".join(members)
raise ValueError(
f"Unknown MLA prefill backend: '{name}'. "
f"Valid options are: {valid_backends}"
) from None
class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
"""Enumeration of all supported MLA prefill backends."""
FLASH_ATTN = (
"vllm.v1.attention.backends.mla.prefill.flash_attn.FlashAttnPrefillBackend"
)
FLASHINFER = (
"vllm.v1.attention.backends.mla.prefill.flashinfer.FlashInferPrefillBackend"
)
TRTLLM_RAGGED = (
"vllm.v1.attention.backends.mla.prefill.trtllm_ragged."
"TrtllmRaggedPrefillBackend"
)
def get_path(self) -> str:
"""Get the fully qualified class path for this backend."""
return self.value
def get_class(self) -> "type[MLAPrefillBackend]":
"""Lazy load and return the backend class."""
return resolve_obj_by_qualname(self.get_path())
vllm/v1/attention/backends/mla/prefill/selector.py
+183
View File
@@ -0,0 +1,183 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Selector for MLA prefill backends.
This module provides functions for selecting the appropriate MLA prefill
backend based on device capabilities and configuration.
"""
from functools import cache
from typing import TYPE_CHECKING, NamedTuple
import torch
from vllm.logger import init_logger
from vllm.platforms.interface import DeviceCapability
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
logger = init_logger(__name__)
class MLAPrefillSelectorConfig(NamedTuple):
"""Hashable configuration for MLA prefill backend selection.
This is analogous to AttentionSelectorConfig and contains model-specific
configuration needed to select an MLA prefill backend, extracted from
VllmConfig into a hashable form for caching.
"""
dtype: torch.dtype
is_r1_compatible: bool
def is_deepseek_r1_mla_compatible(vllm_config: "VllmConfig") -> bool:
"""Check if model has DeepSeek R1 compatible MLA dimensions.
DeepSeek R1 MLA dimensions are:
- qk_nope_head_dim = 128
- qk_rope_head_dim = 64
- v_head_dim = 128
"""
if vllm_config.model_config is None:
return False
hf_text_config = vllm_config.model_config.hf_text_config
qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
qk_rope_head_dim = getattr(hf_text_config, "qk_rope_head_dim", 1)
v_head_dim = getattr(hf_text_config, "v_head_dim", 1)
return qk_nope_head_dim == 128 and qk_rope_head_dim == 64 and v_head_dim == 128
def _get_mla_prefill_backend_priorities(
device_capability: DeviceCapability,
) -> list[MLAPrefillBackendEnum]:
"""Get MLA prefill backend priorities based on device capability.
Args:
device_capability: The device's compute capability.
Returns:
List of backends in priority order (highest priority first).
"""
if device_capability.major == 10: # Blackwell
return [
MLAPrefillBackendEnum.FLASH_ATTN,
MLAPrefillBackendEnum.TRTLLM_RAGGED,
MLAPrefillBackendEnum.FLASHINFER,
]
else: # Hopper (SM90) and older
return [
MLAPrefillBackendEnum.FLASH_ATTN,
]
def get_mla_prefill_backend(
vllm_config: "VllmConfig",
) -> "type[MLAPrefillBackend]":
"""Select the MLA prefill backend based on configuration and device.
This function first checks for explicit user preferences via
mla_prefill_backend in AttentionConfig, then falls back to automatic
priority-based selection.
Args:
vllm_config: The vLLM configuration.
Returns:
The selected prefill backend class.
"""
from vllm.platforms import current_platform
device_capability = current_platform.get_device_capability()
if device_capability is None:
logger.info_once(
"Device capability not available, using FlashAttention MLA prefill backend."
)
return MLAPrefillBackendEnum.FLASH_ATTN.get_class()
attention_config = vllm_config.attention_config
selector_config = MLAPrefillSelectorConfig(
dtype=vllm_config.model_config.dtype,
is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),
)
if attention_config.mla_prefill_backend is not None:
selected_backend = attention_config.mla_prefill_backend
backend_cls: type[MLAPrefillBackend] | None = None
try:
backend_cls = selected_backend.get_class()
invalid_reasons = backend_cls.validate_configuration(
device_capability, selector_config
)
except ImportError:
invalid_reasons = ["ImportError"]
if invalid_reasons:
raise ValueError(
f"Selected MLA prefill backend {selected_backend.name} "
f"is not valid for this configuration. "
f"Reason: {invalid_reasons}"
)
assert backend_cls is not None
logger.info("Using %s MLA prefill backend.", selected_backend.name)
return backend_cls
return _auto_select_mla_prefill_backend(
device_capability,
selector_config,
)
@cache
def _auto_select_mla_prefill_backend(
device_capability: DeviceCapability,
selector_config: MLAPrefillSelectorConfig,
) -> "type[MLAPrefillBackend]":
"""Auto-select the best available MLA prefill backend.
Args:
device_capability: The device's compute capability.
selector_config: Hashable configuration for backend selection.
Returns:
The selected prefill backend class.
"""
priorities = _get_mla_prefill_backend_priorities(device_capability)
all_invalid_reasons: dict[str, list[str]] = {}
for backend_enum in priorities:
backend_cls: type[MLAPrefillBackend] | None = None
try:
backend_cls = backend_enum.get_class()
invalid_reasons = backend_cls.validate_configuration(
device_capability, selector_config
)
except ImportError:
invalid_reasons = ["ImportError"]
if not invalid_reasons:
assert backend_cls is not None
logger.info_once("Using %s MLA prefill backend.", backend_enum.name)
return backend_cls
all_invalid_reasons[backend_enum.name] = invalid_reasons
reasons_str = (
"{"
+ ", ".join(
f"{name}: [{', '.join(reasons)}]"
for name, reasons in all_invalid_reasons.items()
)
+ "}"
)
config_str = repr(selector_config)
logger.debug_once(
"Some MLA prefill backends are not valid with %s. Reasons: %s.",
config_str,
reasons_str,
)
raise ValueError(
f"No valid MLA prefill backend found with {config_str}. Reasons: {reasons_str}."
)
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
+178
View File
@@ -0,0 +1,178 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""TRT-LLM Ragged backend for MLA prefill."""
from typing import TYPE_CHECKING
import torch
import vllm.envs as envs
from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
from vllm.v1.worker.workspace import current_workspace_manager
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.model_executor.layers.attention.mla_attention import (
MLACommonPrefillMetadata,
)
from vllm.platforms.interface import DeviceCapability
class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
"""TRT-LLM Ragged backend for MLA prefill."""
requires_r1_mla_dimensions = True
@staticmethod
def get_name() -> str:
return "TRTLLM_RAGGED"
@classmethod
def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:
return device_capability.major == 10
@classmethod
def is_available(cls) -> bool:
try:
from flashinfer.prefill import (
trtllm_ragged_attention_deepseek, # noqa: F401
)
return True
except ImportError:
return False
def __init__(
self,
num_heads: int,
scale: float,
kv_lora_rank: int,
qk_nope_head_dim: int,
qk_rope_head_dim: int,
v_head_dim: int,
vllm_config: "VllmConfig",
device: torch.device,
layer_names: list[str] | None = None,
) -> None:
super().__init__(
num_heads=num_heads,
scale=scale,
kv_lora_rank=kv_lora_rank,
qk_nope_head_dim=qk_nope_head_dim,
qk_rope_head_dim=qk_rope_head_dim,
v_head_dim=v_head_dim,
vllm_config=vllm_config,
device=device,
layer_names=layer_names,
)
def _get_workspace_buffer(self) -> torch.Tensor:
(workspace_buffer,) = current_workspace_manager().get_simultaneous(
(
(envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,),
torch.uint8,
),
)
return workspace_buffer
def prepare_metadata(
self,
prefill_metadata: "MLACommonPrefillMetadata",
) -> None:
super().prepare_metadata(prefill_metadata)
self._query_seq_lens = (
prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]
)
def run_prefill_new_tokens(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
return_softmax_lse: bool,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
from flashinfer.prefill import trtllm_ragged_attention_deepseek
workspace_buffer = self._get_workspace_buffer()
out = torch.empty(
q.shape[0],
q.shape[1],
v.shape[2],
device=q.device,
dtype=self._prefill_metadata.output_dtype,
)
ret = trtllm_ragged_attention_deepseek(
query=q,
key=k,
value=v,
workspace_buffer=workspace_buffer,
seq_lens=self._query_seq_lens,
max_q_len=self._prefill_metadata.max_query_len,
max_kv_len=self._prefill_metadata.max_query_len,
bmm1_scale=self.scale,
bmm2_scale=1.0,
o_sf_scale=1.0,
batch_size=self._query_seq_lens.shape[0],
window_left=-1,
cum_seq_lens_q=self._prefill_metadata.query_start_loc,
cum_seq_lens_kv=self._prefill_metadata.query_start_loc,
enable_pdl=False,
is_causal=True,
return_lse=return_softmax_lse,
out=out,
)
if isinstance(ret, tuple):
# Convert from (q_len, num_heads) to (num_heads, q_len)
return ret[0], ret[1].transpose(0, 1).contiguous()
return ret
def run_prefill_context_chunk(
self,
chunk_idx: int,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
from flashinfer.prefill import trtllm_ragged_attention_deepseek
assert self._prefill_metadata.chunked_context is not None
assert self._prefill_metadata.chunked_context.seq_lens[chunk_idx] is not None
workspace_buffer = self._get_workspace_buffer()
out = torch.empty(
q.shape[0],
q.shape[1],
v.shape[2],
device=q.device,
dtype=self._prefill_metadata.output_dtype,
)
attn_out, lse = trtllm_ragged_attention_deepseek(
query=q,
key=k,
value=v,
workspace_buffer=workspace_buffer,
seq_lens=self._prefill_metadata.chunked_context.seq_lens[chunk_idx],
max_q_len=self._prefill_metadata.max_query_len,
max_kv_len=self._prefill_metadata.chunked_context.max_seq_lens[chunk_idx],
bmm1_scale=self.scale,
bmm2_scale=1.0,
o_sf_scale=1.0,
batch_size=self._prefill_metadata.chunked_context.seq_lens[chunk_idx].shape[
0
],
window_left=-1,
cum_seq_lens_q=self._prefill_metadata.query_start_loc,
cum_seq_lens_kv=self._prefill_metadata.chunked_context.cu_seq_lens[
chunk_idx
],
enable_pdl=False,
is_causal=False,
return_lse=True,
out=out,
)
# Convert from (q_len, num_heads) to (num_heads, q_len)
return attn_out, lse.transpose(0, 1).contiguous()
vllm/v1/attention/backends/mla/triton_mla.py
-12
View File
@@ -123,18 +123,6 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
self._sm_count = current_platform.num_compute_units()
def _flash_attn_varlen_diff_headdims(
self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs
):
return super()._flash_attn_varlen_diff_headdims(
q,
k,
v,
return_softmax_lse=return_softmax_lse,
softmax_scale=softmax_scale,
**kwargs,
)
def forward_mqa(
self,
q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],