Rebase Arctic Inference onto vLLM v0.26.0

arctic_inference/server/worker.py

@@ -138,14 +138,26 @@ async def initialize(self, engine_kwargs: dict[str, Any], extra_env: dict[str, s
         import vllm.plugins
         vllm.plugins.load_general_plugins()
 
-        if not arctic_inference_effective_enabled():
-            engine_kwargs.pop("forest_cascade_attn_configs", None)
-
-        # from vllm.v1.engine.async_llm import AsyncLLM
-        # from vllm.engine.arg_utils import AsyncEngineArgs
-
         from vllm.v1.engine.async_llm import AsyncLLM
-        from arctic_inference.vllm.args import ArcticAsyncEngineArgs
+
+        # If we reach here with Arctic enabled, load_general_plugins() above has
+        # already applied the Arctic patches: the plugin raises on an
+        # enabled-but-version-mismatch (unless ARCTIC_INFERENCE_SKIP_VERSION_CHECK
+        # forced it through), so an enabled process that got this far is patched.
+        # Arctic engine args are therefore valid exactly when Arctic is enabled;
+        # otherwise fall back to vanilla AsyncEngineArgs and drop Arctic-only
+        # kwargs that vanilla vLLM would reject. Source of the kwarg list:
+        # ArcticArgs in arctic_inference/vllm/args.py.
+        use_arctic = arctic_inference_effective_enabled()
+        if use_arctic:
+            from arctic_inference.vllm.args import \
+                ArcticAsyncEngineArgs as _EngineArgs
+        else:
+            from vllm.engine.arg_utils import AsyncEngineArgs as _EngineArgs
+            for _k in ("ulysses_sequence_parallel_size", "enable_shift_parallel",
+                       "shift_parallel_threshold", "forest_cascade_attn_configs",
+                       "fp32_lm_head"):
+                engine_kwargs.pop(_k, None)
 
         # invariance setup
         #print(f"{os.environ.get('VLLM_BATCH_INVARIANT')=}")
@@ -162,7 +174,7 @@ async def initialize(self, engine_kwargs: dict[str, Any], extra_env: dict[str, s
             "arctic_inference.server.weight_sync.WeightSyncExtension",
         )
 
-        engine_args = ArcticAsyncEngineArgs(**engine_kwargs)
+        engine_args = _EngineArgs(**engine_kwargs)
         logger.info("Worker %d engine_args.enable_sleep_mode=%s",
                      os.getpid(), getattr(engine_args, 'enable_sleep_mode', 'N/A'))
         vllm_config = engine_args.create_engine_config()

arctic_inference/utils.py

@@ -13,17 +13,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import re
-from importlib.metadata import requires
+# The exact vLLM version this branch's patches target. This is the single source
+# of truth for the plugin's version check and is intentionally decoupled from the
+# install extras: the `vllm` extra is left unpinned so users can resolve vLLM
+# against their own torch. Arctic patches are only applied when the *installed*
+# vLLM matches this version; on any other version the plugin skips patching and
+# vLLM runs unmodified (no acceleration).
+VLLM_PATCH_VERSION = "0.26.0"
 
 
 def get_compatible_vllm_version():
-    reqs = requires("arctic_inference")
-    for req in reqs:
-        match = re.match("vllm==(.*); extra == \"vllm\"", req)
-        if match is not None:
-            return match.groups()[0]
-
+    """The exact vLLM version the plugin's patches are written against."""
+    return VLLM_PATCH_VERSION
+
+
+def plugin_version_compatible() -> bool:
+    """True iff the installed vLLM exactly matches the plugin's supported pin."""
+    import vllm
+    want = get_compatible_vllm_version()
+    return want is not None and vllm.__version__ == want
+
 
 # For debugging
 def print0(*args, **kwargs):

arctic_inference/vllm/args.py

@@ -20,13 +20,41 @@
 from dataclasses import dataclass, fields
 
 from vllm.config import ParallelConfig
+from vllm.config.compilation import CUDAGraphMode
 from vllm.config.utils import is_init_field
 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
+from vllm.logger import init_logger
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from arctic_inference.patching import ArcticPatch
 from arctic_inference.vllm.config import ArcticParallelConfig
 
+logger = init_logger(__name__)
+
+
+def _maybe_force_piecewise_for_fca(vllm_config) -> None:
+    """Force cudagraph_mode down to PIECEWISE when Forest Cascade Attention is
+    configured.
+
+    FCA only fires under piecewise cudagraphs -- its per-batch gate requires
+    ``not use_full_cuda_graph`` -- so any full-cudagraph mode would silently
+    disable FCA. With shift parallelism it would also dispatch full graphs the
+    shift model never captured. Downgrading here (base + shift, at engine-config
+    creation) keeps the behavior correct and crash-free. An explicit NONE or
+    PIECEWISE mode is preserved; only full modes are downgraded.
+    """
+    if getattr(vllm_config, "_forest_cascade_attn_config", None) is None:
+        return
+    cc = vllm_config.compilation_config
+    cg = cc.cudagraph_mode
+    if cg is not None and cg.has_full_cudagraphs():
+        logger.warning(
+            "Forest Cascade Attention is configured; forcing cudagraph_mode "
+            "%s -> PIECEWISE so FCA can fire (FCA is disabled under full "
+            "cudagraphs).", cg,
+        )
+        cc.cudagraph_mode = CUDAGraphMode.PIECEWISE
+
 
 @dataclass
 class ArcticArgs:
@@ -38,6 +66,27 @@ class ArcticArgs:
     fp32_lm_head: bool = False
 
 
+def _ensure_arctic_fields(engine_args) -> None:
+    """Backfill ArcticArgs defaults onto a plain ``EngineArgs`` instance.
+
+    v0.26 plugin-load timing (runtime-only drift): both 0.24 and 0.26 load
+    general plugins from ``EngineArgs.__post_init__`` -> the *first*
+    ``EngineArgs(...)`` (e.g. ``LLM.__init__`` building it directly) is
+    constructed *before* Arctic's patches are applied, so ``EngineArgsPatch.
+    __new__`` never upgrades it to ``ArcticEngineArgs`` and it lacks the
+    ArcticArgs fields. By the time the (now-patched) ``create_engine_config``
+    runs on that instance, reading ``self.ulysses_sequence_parallel_size``
+    raises ``AttributeError``. On this path the Arctic args can only ever be
+    their defaults (a base ``EngineArgs.__init__`` would reject them as unknown
+    kwargs), so backfilling the dataclass defaults is exact, not a papering-over
+    shim. The CLI path (``from_cli_args`` -> ``ArcticEngineArgs``) already has
+    them and is unaffected.
+    """
+    for f in fields(ArcticArgs):
+        if not hasattr(engine_args, f.name):
+            setattr(engine_args, f.name, f.default)
+
+
 @dataclass
 class ArcticEngineArgs(EngineArgs, ArcticArgs):
     pass
@@ -139,6 +188,7 @@ def from_cli_args(cls, args: argparse.Namespace):
         return EngineArgsPatch._orig_from_cli_args.__func__(cls, args)
 
     def create_engine_config(self, *args, **kwargs):
+        _ensure_arctic_fields(self)
         if (self.ulysses_sequence_parallel_size > 1 and
                 self.distributed_executor_backend is None):
             self.distributed_executor_backend = "mp"
@@ -175,6 +225,8 @@ def create_engine_config(self, *args, **kwargs):
         else:
             vllm_config._forest_cascade_attn_config = None
 
+        _maybe_force_piecewise_for_fca(vllm_config)
+
         return vllm_config
 
 
@@ -196,6 +248,7 @@ def __post_init__(self):
         self._orig_post_init()
 
     def create_engine_config(self, *args, **kwargs):
+        _ensure_arctic_fields(self)
         if (self.ulysses_sequence_parallel_size > 1 and
                 self.distributed_executor_backend is None):
             self.distributed_executor_backend = "mp"
@@ -231,4 +284,6 @@ def create_engine_config(self, *args, **kwargs):
         else:
             vllm_config._forest_cascade_attn_config = None
 
+        _maybe_force_piecewise_for_fca(vllm_config)
+
         return vllm_config

arctic_inference/vllm/attention/flash_attn_forest_cascade.py

@@ -14,16 +14,20 @@
 import numpy as np
 import torch
 
+import vllm.envs as envs
 from vllm.model_executor.layers.attention import Attention
+from vllm.utils.torch_utils import (
+    canonicalize_singleton_dim_strides,
+    is_quantized_kv_cache,
+)
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionImpl,
     AttentionType,
     MultipleOf,
-    is_quantized_kv_cache,
 )
 from vllm.v1.attention.backends.fa_utils import (
-    flash_attn_supports_fp8,
+    flash_attn_supports_quant_query_input,
     get_flash_attn_version,
     is_flash_attn_varlen_func_available,
 )
@@ -45,9 +49,7 @@
 from vllm.config.cache import CacheDType
 from vllm.distributed.parallel_state import get_dcp_group
 from vllm.logger import init_logger
-from vllm.model_executor.layers.batch_invariant import (
-    vllm_is_batch_invariant,
-)
+from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backend import (
@@ -119,21 +121,29 @@ def get_kv_cache_shape(
     ) -> tuple[int, ...]:
         if block_size % 16 != 0:
             raise ValueError("Block size must be a multiple of 16.")
-        return (2, num_blocks, block_size, num_kv_heads, head_size)
+        # v0.26: K and V are packed into the content dim: logical (B, H, N, 2*D).
+        return (num_blocks, num_kv_heads, block_size, 2 * head_size)
 
     @staticmethod
     def get_kv_cache_stride_order(
         include_num_layers_dimension: bool = False,
     ) -> tuple[int, ...]:
+        # `stride_order` indicates the permutation that gets us from
+        # `get_kv_cache_shape` (logical (B, H, N, 2*D)) to the actual memory
+        # layout we want.
         cache_layout = get_kv_cache_layout()
         if cache_layout == "NHD" and include_num_layers_dimension:
-            return (2, 0, 1, 3, 4, 5)
+            # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size)
+            return (1, 0, 3, 2, 4)
         elif cache_layout == "NHD":
-            stride_order = (0, 1, 2, 3, 4)
+            # (num_blocks, block_size, num_kv_heads, 2*head_size)
+            stride_order = (0, 2, 1, 3)
         elif cache_layout == "HND" and include_num_layers_dimension:
-            return (2, 4, 0, 1, 3, 5)
+            # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size)
+            return (1, 2, 0, 3, 4)
         elif cache_layout == "HND":
-            stride_order = (0, 1, 3, 2, 4)
+            # (num_blocks, num_kv_heads, block_size, 2*head_size)
+            stride_order = (0, 1, 2, 3)
         else:
             raise ValueError(f"Unknown cache layout format {cache_layout}.")
         return stride_order
@@ -155,9 +165,14 @@ def supports_kv_cache_dtype(
     ) -> bool:
         if kv_cache_dtype is None:
             return True
-        if kv_cache_dtype.startswith("fp8"):
-            return flash_attn_supports_fp8()
-        return kv_cache_dtype in ["auto"]
+        if kv_cache_dtype in ("fp8", "fp8_e4m3"):
+            if current_platform.is_xpu():
+                return True
+            return (
+                get_flash_attn_version() == 3
+                and current_platform.is_device_capability_family(90)
+            )
+        return kv_cache_dtype in ["auto", "float16", "bfloat16"]
 
     @classmethod
     def supports_sink(cls) -> bool:
@@ -454,7 +469,7 @@ def build(
         ):
             max_num_splits = self.max_num_splits
 
-        if vllm_is_batch_invariant():
+        if envs.VLLM_BATCH_INVARIANT:
             max_num_splits = 1
 
         def schedule(
@@ -871,15 +886,8 @@ def __init__(
 
         self.attn_type = attn_type
         self.vllm_flash_attn_version = get_flash_attn_version()
-        self.batch_invariant_enabled = vllm_is_batch_invariant()
-
-        if (
-            is_quantized_kv_cache(self.kv_cache_dtype)
-            and not flash_attn_supports_fp8()
-        ):
-            raise NotImplementedError(
-                "FlashAttention does not support fp8 kv-cache on this device."
-            )
+        # Cache the batch invariant result for use in forward passes
+        self.batch_invariant_enabled = envs.VLLM_BATCH_INVARIANT
 
         self.sinks = sinks
         if self.sinks is not None:
@@ -891,7 +899,7 @@ def __init__(
                 "heads in the layer"
             )
 
-        self.supports_quant_query_input = True
+        self.supports_quant_query_input = flash_attn_supports_quant_query_input()
 
     def forward(
         self,
@@ -912,7 +920,7 @@ def forward(
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
             kv_cache: shape =
-                [2, num_blocks, block_size, num_kv_heads, head_size]
+                [num_blocks, 2, block_size, num_kv_heads, head_size]
             attn_metadata: Metadata for attention.
         Returns:
             shape = [num_tokens, num_heads * head_size]
@@ -945,8 +953,18 @@ def forward(
                 layer,
             )
 
-        # For decoder and cross-attention, use KV cache as before
-        key_cache, value_cache = kv_cache.unbind(0)
+        # For decoder and cross-attention, use KV cache as before.
+        # v0.26: KV cache is packed as [num_blocks, num_kv_heads, block_size,
+        # 2 * head_size]; key/value are split out of the last dim (was a
+        # separate dim-1 `unbind(1)` in v0.24). The resulting views keep the
+        # same logical shape [num_blocks, block_size, num_kv_heads, head_size]
+        # as before, so the downstream cascade paths are unaffected.
+        # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D))
+        key_cache, value_cache = kv_cache.transpose(1, 2).split(
+            self.head_size, dim=-1
+        )
+        key_cache = canonicalize_singleton_dim_strides(key_cache)
+        value_cache = canonicalize_singleton_dim_strides(value_cache)
 
         if (
             self.kv_sharing_target_layer_name is None
@@ -1718,7 +1736,7 @@ def forest_cascade_attention(
     suffix_descale_shape = (B, num_kv_heads)
 
     effective_num_splits = (
-        1 if vllm_is_batch_invariant() else max_num_splits
+        1 if envs.VLLM_BATCH_INVARIANT else max_num_splits
     )
 
     # Prefix attention: groups as sequences, causal=False.
@@ -1881,7 +1899,7 @@ def cascade_attention(
             else None
         ),
         s_aux=s_aux,
-        num_splits=1 if vllm_is_batch_invariant() else max_num_splits,
+        num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits,
     )
 
     descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2])
@@ -1918,7 +1936,7 @@ def cascade_attention(
             if v_descale is not None
             else None
         ),
-        num_splits=1 if vllm_is_batch_invariant() else max_num_splits,
+        num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits,
     )
 
     # Merge prefix and suffix outputs, and store the result in output.