@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / aa6138169f
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 18:15 UTC

[MLA][Attention] Add OOT MLA prefill backend registration mechanism (#43325)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed May 26, 2026 at 21:56 UTC aa6138169f80863c0f4438403c2f78f64a13907c
3 files changed +227 -5
tests/v1/attention/test_mla_prefill_registry.py new
+137
@@ -0,0 +1,137 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for MLA prefill backend registry."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
9 +from vllm.v1.attention.backends.mla.prefill.registry import (
10 + MLAPrefillBackendEnum,
11 + register_mla_prefill_backend,
12 +)
13 +
14 +
15 +class CustomMLAPrefillBackend(MLAPrefillBackend):
16 + """Mock custom MLA prefill backend for testing."""
17 +
18 + supported_dtypes = [torch.bfloat16, torch.float16]
19 + requires_r1_mla_dimensions = False
20 +
21 + @staticmethod
22 + def get_name() -> str:
23 + return "CUSTOM"
24 +
25 + def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
26 + raise NotImplementedError
27 +
28 + def run_prefill_context_chunk(self, chunk_idx, q, k, v):
29 + raise NotImplementedError
30 +
31 +
32 +@pytest.fixture(autouse=True)
33 +def cleanup_overrides():
34 + """Clear any overrides after each test."""
35 + yield
36 + for member in MLAPrefillBackendEnum:
37 + member.clear_override()
38 +
39 +
40 +def test_custom_is_not_alias_of_any_backend():
41 + all_backends = list(MLAPrefillBackendEnum)
42 +
43 + aliases = []
44 + for backend in all_backends:
45 + if backend.name != "CUSTOM" and backend is MLAPrefillBackendEnum.CUSTOM:
46 + aliases.append(backend.name)
47 +
48 + assert len(aliases) == 0, (
49 + f"BUG! CUSTOM is an alias of: {', '.join(aliases)}!\n"
50 + f"CUSTOM.value = {repr(MLAPrefillBackendEnum.CUSTOM.value)}\n"
51 + f"All MLA prefill backend values:\n"
52 + + "\n".join(f" {b.name}: {repr(b.value)}" for b in all_backends)
53 + )
54 +
55 + assert MLAPrefillBackendEnum.CUSTOM.name == "CUSTOM"
56 +
57 +
58 +def test_custom_unregistered_raises():
59 + with pytest.raises(ValueError, match="must be registered before use"):
60 + MLAPrefillBackendEnum.CUSTOM.get_path()
61 +
62 +
63 +def test_register_custom_backend_with_class_path():
64 + register_mla_prefill_backend(
65 + backend=MLAPrefillBackendEnum.CUSTOM,
66 + class_path=(
67 + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
68 + ),
69 + )
70 +
71 + assert MLAPrefillBackendEnum.CUSTOM.is_overridden()
72 +
73 + class_path = MLAPrefillBackendEnum.CUSTOM.get_path()
74 + assert class_path == (
75 + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
76 + )
77 +
78 + backend_cls = MLAPrefillBackendEnum.CUSTOM.get_class()
79 + assert backend_cls.get_name() == "CUSTOM"
80 +
81 +
82 +def test_register_custom_backend_as_decorator():
83 + @register_mla_prefill_backend(MLAPrefillBackendEnum.CUSTOM)
84 + class DecoratedPrefillBackend(MLAPrefillBackend):
85 + supported_dtypes = [torch.bfloat16]
86 + requires_r1_mla_dimensions = False
87 +
88 + @staticmethod
89 + def get_name() -> str:
90 + return "DECORATED"
91 +
92 + def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
93 + raise NotImplementedError
94 +
95 + def run_prefill_context_chunk(self, chunk_idx, q, k, v):
96 + raise NotImplementedError
97 +
98 + assert MLAPrefillBackendEnum.CUSTOM.is_overridden()
99 + assert "DecoratedPrefillBackend" in MLAPrefillBackendEnum.CUSTOM.get_path()
100 +
101 +
102 +def test_override_existing_backend():
103 + original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path()
104 +
105 + register_mla_prefill_backend(
106 + backend=MLAPrefillBackendEnum.FLASH_ATTN,
107 + class_path=(
108 + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
109 + ),
110 + )
111 +
112 + assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
113 + assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() != original_path
114 +
115 + backend_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()
116 + assert backend_cls.get_name() == "CUSTOM"
117 +
118 +
119 +def test_clear_override():
120 + original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path()
121 +
122 + register_mla_prefill_backend(
123 + backend=MLAPrefillBackendEnum.FLASH_ATTN,
124 + class_path=(
125 + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
126 + ),
127 + )
128 + assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
129 +
130 + MLAPrefillBackendEnum.FLASH_ATTN.clear_override()
131 + assert not MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
132 + assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() == original_path
133 +
134 +
135 +def test_unknown_backend_name_raises():
136 + with pytest.raises(ValueError, match="Unknown MLA prefill backend"):
137 + MLAPrefillBackendEnum["NONEXISTENT"]
vllm/v1/attention/backends/mla/prefill/__init__.py
+5 -1
@@ -1,11 +1,15 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
4 -from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
4 +from vllm.v1.attention.backends.mla.prefill.registry import (
5 + MLAPrefillBackendEnum,
6 + register_mla_prefill_backend,
7 +)
8 from vllm.v1.attention.backends.mla.prefill.selector import get_mla_prefill_backend
9
10 __all__ = [
11 "MLAPrefillBackend",
12 "MLAPrefillBackendEnum",
13 "get_mla_prefill_backend",
14 + "register_mla_prefill_backend",
15 ]
vllm/v1/attention/backends/mla/prefill/registry.py
+85 -4
@@ -3,9 +3,10 @@
3 """Registry for MLA prefill backends.
4
5 This module provides an enumeration of all available MLA prefill backends
6 -and utilities for loading them.
6 +and utilities for loading and registering them.
7 """
8
9 +from collections.abc import Callable
10 from enum import Enum, EnumMeta
11 from typing import TYPE_CHECKING
12
@@ -47,11 +48,91 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
48 "vllm.v1.attention.backends.mla.prefill.tokenspeed_mla."
49 "TokenspeedMLAPrefillBackend"
50 )
51 + # Placeholder for third-party/custom backends - must be registered before use
52 + # set to None to avoid alias with other backend, whose value is an empty string
53 + CUSTOM = None
54
55 def get_path(self) -> str:
52 - """Get the fully qualified class path for this backend."""
53 - return self.value
56 + """Get the class path for this backend (respects overrides).
57 +
58 + Returns:
59 + The fully qualified class path string
60 +
61 + Raises:
62 + ValueError: If Backend.CUSTOM is used without being registered
63 + """
64 + path = _MLA_PREFILL_OVERRIDES.get(self, self.value)
65 + if not path:
66 + raise ValueError(
67 + f"MLA prefill backend {self.name} must be registered before "
68 + f"use. Use register_mla_prefill_backend("
69 + f"MLAPrefillBackendEnum.{self.name}, "
70 + f"'your.module.YourClass')"
71 + )
72 + return path
73
74 def get_class(self) -> "type[MLAPrefillBackend]":
56 - """Lazy load and return the backend class."""
75 + """Get the backend class (respects overrides).
76 +
77 + Returns:
78 + The backend class
79 +
80 + Raises:
81 + ImportError: If the backend class cannot be imported
82 + ValueError: If CUSTOM is used without being registered
83 + """
84 return resolve_obj_by_qualname(self.get_path())
85 +
86 + def is_overridden(self) -> bool:
87 + """Check if this backend has been overridden."""
88 + return self in _MLA_PREFILL_OVERRIDES
89 +
90 + def clear_override(self) -> None:
91 + """Clear any override for this backend, reverting to the default."""
92 + _MLA_PREFILL_OVERRIDES.pop(self, None)
93 +
94 +
95 +_MLA_PREFILL_OVERRIDES: dict[MLAPrefillBackendEnum, str] = {}
96 +
97 +
98 +def register_mla_prefill_backend(
99 + backend: MLAPrefillBackendEnum,
100 + class_path: str | None = None,
101 +) -> Callable[[type], type]:
102 + """Register or override an MLA prefill backend implementation.
103 +
104 + Args:
105 + backend: The MLAPrefillBackendEnum member to register.
106 + class_path: Optional class path. If not provided and used as
107 + decorator, will be auto-generated from the class.
108 +
109 + Returns:
110 + Decorator function if class_path is None, otherwise a no-op.
111 +
112 + Examples:
113 + # Override an existing MLA prefill backend
114 + @register_mla_prefill_backend(MLAPrefillBackendEnum.FLASH_ATTN)
115 + class MyCustomFlashAttn(MLAPrefillBackend):
116 + ...
117 +
118 + # Register a custom third-party MLA prefill backend
119 + @register_mla_prefill_backend(MLAPrefillBackendEnum.CUSTOM)
120 + class MyCustomPrefillBackend(MLAPrefillBackend):
121 + ...
122 +
123 + # Direct registration
124 + register_mla_prefill_backend(
125 + MLAPrefillBackendEnum.CUSTOM,
126 + "my.module.MyCustomPrefillBackend"
127 + )
128 + """
129 +
130 + def decorator(cls: type) -> type:
131 + _MLA_PREFILL_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}"
132 + return cls
133 +
134 + if class_path is not None:
135 + _MLA_PREFILL_OVERRIDES[backend] = class_path
136 + return lambda x: x
137 +
138 + return decorator