diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index bc24a8a5a1f..dc4b5402cab 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -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=` (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=` 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 | diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index bf3b400d9d7..1ab4949c400 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -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 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index e65d1d60402..f91ea85779d 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -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, diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py new file mode 100644 index 00000000000..068eb43faf4 --- /dev/null +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -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 diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index d131c9cc01f..73ef8b91582 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -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=` (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=` argument (e.g., `FLASHMLA`, `TRITON_MLA`).", + "", ] ) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 18973f5d66d..b5dc7a5bf60 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -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." + ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 4afe2319570..82eecc8cd49 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -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 diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index c0d52620c08..2753326755f 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -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}" diff --git a/vllm/v1/attention/backends/mla/prefill/__init__.py b/vllm/v1/attention/backends/mla/prefill/__init__.py new file mode 100644 index 00000000000..ae5b7ae8259 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/__init__.py @@ -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", +] diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py new file mode 100644 index 00000000000..9c850a0b1d9 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -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 diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py new file mode 100644 index 00000000000..42d77b12d0d --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -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, + ) diff --git a/vllm/v1/attention/backends/mla/prefill/flashinfer.py b/vllm/v1/attention/backends/mla/prefill/flashinfer.py new file mode 100644 index 00000000000..a1107df304f --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/flashinfer.py @@ -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() diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py new file mode 100644 index 00000000000..3a3242f6036 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/registry.py @@ -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()) diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py new file mode 100644 index 00000000000..fdb8be6d65d --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -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}." + ) diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py new file mode 100644 index 00000000000..2ffe244ebe0 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -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() diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index 7aa8a646f41..c45a631008e 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -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],