From fa49953a012f69a5c2dfb3fa5b4fdb85b5480d30 Mon Sep 17 00:00:00 2001 From: ye-wang_snow Date: Wed, 22 Apr 2026 22:41:28 +0000 Subject: [PATCH 1/2] runtime rope --- arctic_inference/vllm/model_runner.py | 134 +++ arctic_inference/vllm/patches.py | 9 + arctic_inference/vllm/rope/__init__.py | 85 ++ arctic_inference/vllm/rope/multi_cache_ntk.py | 443 +++++++++ arctic_inference/vllm/rope/patches.py | 334 +++++++ tests/unit_tests/test_multi_cache_rope.py | 938 ++++++++++++++++++ 6 files changed, 1943 insertions(+) create mode 100644 arctic_inference/vllm/rope/__init__.py create mode 100644 arctic_inference/vllm/rope/multi_cache_ntk.py create mode 100644 arctic_inference/vllm/rope/patches.py create mode 100644 tests/unit_tests/test_multi_cache_rope.py diff --git a/arctic_inference/vllm/model_runner.py b/arctic_inference/vllm/model_runner.py index 2cff488ef..171697f15 100644 --- a/arctic_inference/vllm/model_runner.py +++ b/arctic_inference/vllm/model_runner.py @@ -113,6 +113,7 @@ class GPUModelRunnerPatch(ArcticPatch[GPUModelRunner]): _orig_bookkeeping_sync = GPUModelRunner._bookkeeping_sync _orig_sample_tokens = GPUModelRunner.sample_tokens _orig_initialize_kv_cache = GPUModelRunner.initialize_kv_cache + _orig_model_forward = GPUModelRunner._model_forward # _orig_pad_for_sequence_parallelism = GPUModelRunner._pad_for_sequence_parallelism def __init__( @@ -605,6 +606,139 @@ def _sample( sampler_output.sampled_token_ids) return sampler_output + # ------------------------------------------------------------------ + # Multi-cache dynamic NTK RoPE plumbing + # + # When any rotary layer in ``self.model`` is an instance of + # :class:`MultiCacheDynamicNTKRotaryEmbedding` (installed via + # ``rope_type="multi_cache_ntk"`` in the HF config, or auto-promoted + # via ``ARCTIC_INFERENCE_MULTI_CACHE_ROPE=1``), we must refresh its + # per-token offset buffer on every real forward. + # + # We write directly into the module's ``runtime_bucket_offsets`` + # buffer *before* invoking the model. That design has two important + # properties: + # + # 1. CUDA graph safe. The rotary forward has no Python-level + # control flow on tensor values and reads offsets from the + # buffer via a fixed-shape slice. Graph capture records the + # load; graph replay picks up whatever was most recently + # written into the buffer. No graph breaks, no re-capture. + # 2. No host sync. The per-token seq-len tensor is built on the + # CPU side from ``self.seq_lens`` + ``self.query_start_loc`` (both + # maintained by vLLM anyway) and pushed to GPU non-blocking. The + # seq_len -> offset translation happens on-device inside + # :meth:`MultiCacheDynamicNTKRotaryEmbedding.update_runtime_seq_lens`. + # + # We cache the list of multi-cache rotary modules per model identity + # so shift-model swaps don't leak into each other. + # ------------------------------------------------------------------ + + def _build_rope_seq_lens_per_token_gpu( + self, num_tokens_padded: int + ) -> Optional[torch.Tensor]: + """Return a GPU int32 tensor of per-token seq-lens. + + Shape is ``[num_tokens_padded]``. Padding tokens (if any) are + assigned a seq-len of ``1`` so the bucket router picks factor-1 + (the unscaled cache). Padding tokens are masked out by + attention so the exact offset is irrelevant, but we want the + routing to land in a numerically safe bucket. + """ + num_reqs = getattr(getattr(self, "input_batch", None), "num_reqs", 0) + if not num_reqs: + return None + + seq_lens_cpu = self.seq_lens.np[:num_reqs] + # query_start_loc stores cumulative token counts per request in + # slots [0..num_reqs]. diff() gives per-request scheduled tokens. + qsl = self.query_start_loc.np[: num_reqs + 1] + num_scheduled_per_req = np.diff(qsl).astype(np.int64, copy=False) + num_tokens_unpadded = int(num_scheduled_per_req.sum()) + if num_tokens_unpadded <= 0: + return None + + per_token_cpu = np.empty(num_tokens_padded, dtype=np.int32) + per_token_cpu[:num_tokens_unpadded] = np.repeat( + seq_lens_cpu.astype(np.int32, copy=False), + num_scheduled_per_req, + ) + if num_tokens_padded > num_tokens_unpadded: + # Padding tokens: seq_len=1 routes to factor-1 (unscaled). + per_token_cpu[num_tokens_unpadded:] = 1 + + cpu_tensor = torch.from_numpy(per_token_cpu) + return cpu_tensor.to( + self.device, dtype=torch.int32, non_blocking=True, + ) + + def _runtime_rope_modules(self) -> list: + """Return the list of :class:`MultiCacheDynamicNTKRotaryEmbedding` + instances inside ``self.model``. + + Cached per model identity so shift-model swaps don't return a + stale list. The list is empty when the model uses no multi-cache + rope, which is the fast path we check first in ``_model_forward``. + """ + model = getattr(self, "model", None) + if model is None: + return [] + cache = getattr(self, "_arctic_runtime_rope_modules", None) + cache_id = getattr(self, "_arctic_runtime_rope_modules_id", None) + if cache is not None and cache_id == id(model): + return cache + found: list = [] + try: + from arctic_inference.vllm.rope import ( + MultiCacheDynamicNTKRotaryEmbedding, + ) + + for module in model.modules(): + if isinstance(module, MultiCacheDynamicNTKRotaryEmbedding): + found.append(module) + except Exception: + found = [] + self._arctic_runtime_rope_modules = found + self._arctic_runtime_rope_modules_id = id(model) + return found + + def _model_forward( + self, + input_ids: Optional[torch.Tensor] = None, + positions: Optional[torch.Tensor] = None, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + **model_kwargs, + ): + rope_modules = self._runtime_rope_modules() + if rope_modules: + # Write per-token seq-lens (which the module translates to + # offsets) into each multi-cache rope module's buffer + # *before* the model is invoked. The in-place write is + # what makes this CUDA-graph safe: captured graphs read + # the buffer's storage, and each replay sees the latest + # values. If we can't materialise a seq-lens tensor (e.g. + # empty batch) we leave the buffer alone -- the rotary will + # then use whatever was written previously (or the init + # value of 0, which is the unscaled factor-1 offset). + num_tokens_padded = ( + int(positions.shape[-1]) if positions is not None else 0 + ) + gpu_seq_lens = self._build_rope_seq_lens_per_token_gpu( + num_tokens_padded, + ) + if gpu_seq_lens is not None: + for rope_mod in rope_modules: + rope_mod.update_runtime_seq_lens(gpu_seq_lens) + + return self._orig_model_forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + ) + @torch.inference_mode() def execute_model( self, diff --git a/arctic_inference/vllm/patches.py b/arctic_inference/vllm/patches.py index 0d13baed7..8e5d54960 100644 --- a/arctic_inference/vllm/patches.py +++ b/arctic_inference/vllm/patches.py @@ -27,6 +27,7 @@ SpeculativeConfigPatch, VllmConfigPatch, MLPSpeculatorConfigPatch) +from arctic_inference.vllm.rope import apply_rope_runtime_patches from arctic_inference.vllm.stats import (SpecDecodingStatsPatch, SpecDecodingLoggingPatch) from arctic_inference.vllm.structured_output import XgrammarBackendPatch @@ -298,5 +299,13 @@ def apply_arctic_patches(): XgrammarBackendPatch.apply_patch() MLPSpeculatorConfigPatch.apply_patch() + # Multi-cache dynamic NTK RoPE (per-factor static caches concatenated; + # per-token bucket routing picks the factor whose cache best covers + # each request's seq_len). Installs a wrapper around vLLM's get_rope + # so that rope_type="multi_cache_ntk" is dispatched to + # arctic_inference.vllm.rope.MultiCacheDynamicNTKRotaryEmbedding. + # Must run before any model loads. + apply_rope_runtime_patches() + # Main optimization patches. apply_shift_parallel_patches() diff --git a/arctic_inference/vllm/rope/__init__.py b/arctic_inference/vllm/rope/__init__.py new file mode 100644 index 000000000..bb832b3a7 --- /dev/null +++ b/arctic_inference/vllm/rope/__init__.py @@ -0,0 +1,85 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Multi-cache dynamic NTK RoPE support for ArcticInference. + +This subpackage adds a new ``rope_type="multi_cache_ntk"`` that +implements a *bucketed* dynamic-NTK scheme: one static-rope cache per +configured factor, concatenated into a single ``cos_sin_cache``. At +forward time, each token is routed (per its request's seq_len) to the +factor whose cache best covers that length, and the existing CUDA rope +kernel is re-used unchanged -- we simply add a per-token offset tensor +to ``positions`` before the kernel indexes the unified cache. + +Why this shape +-------------- +* Within any given bucket, the behavior is bit-identical to vLLM's + static :class:`~vllm.model_executor.layers.rotary_embedding.DynamicNTKScalingRotaryEmbedding` + at that bucket's factor. This matches the training distribution for + models that were fine-tuned with static rope at some factor. +* No new CUDA kernel. The forward reuses ``ops.rotary_embedding`` with + ``positions + per_token_offset``. Cost is a single elementwise add. +* CUDA-graph safe. The per-token offset tensor lives in a registered + buffer; the model runner writes it in place each forward. Captured + graphs see fresh values on replay with no re-capture. + +Design summary +-------------- +* :class:`MultiCacheDynamicNTKRotaryEmbedding` is a drop-in + :class:`~vllm.model_executor.layers.rotary_embedding.base.RotaryEmbedding` + subclass. It owns a registered GPU buffer ``runtime_bucket_offsets`` + sized for ``max_num_batched_tokens`` and a concatenated + ``cos_sin_cache`` covering every factor. +* The model runner + (:class:`arctic_inference.vllm.model_runner.GPUModelRunnerPatch`) + computes per-token seq-lens and calls + :meth:`MultiCacheDynamicNTKRotaryEmbedding.update_runtime_seq_lens` + before invoking the model; that method does the seq_len -> offset + translation in-place on the GPU buffer. + +Public API +---------- +* :class:`MultiCacheDynamicNTKRotaryEmbedding` +* :func:`apply_rope_runtime_patches` +""" + +from arctic_inference.vllm.rope.multi_cache_ntk import ( + DEFAULT_FACTORS, + MultiCacheDynamicNTKRotaryEmbedding, +) + +__all__ = [ + "DEFAULT_FACTORS", + "MultiCacheDynamicNTKRotaryEmbedding", + "apply_rope_runtime_patches", +] + + +def apply_rope_runtime_patches() -> None: + """Install the multi-cache RoPE patches. + + This must be called *before* any model loads so that + :func:`vllm.model_executor.layers.rotary_embedding.get_rope` knows + how to dispatch the new ``multi_cache_ntk`` rope type. The + companion :class:`GPUModelRunnerPatch` hook is installed elsewhere + (in :mod:`arctic_inference.vllm.model_runner`) because it needs to + run after CUDA is safely importable. + """ + from arctic_inference.vllm.rope.patches import ( + _install_apply_dict_overrides_patch, + _install_get_rope_patch, + ) + + _install_apply_dict_overrides_patch() + _install_get_rope_patch() diff --git a/arctic_inference/vllm/rope/multi_cache_ntk.py b/arctic_inference/vllm/rope/multi_cache_ntk.py new file mode 100644 index 000000000..eb62de5a0 --- /dev/null +++ b/arctic_inference/vllm/rope/multi_cache_ntk.py @@ -0,0 +1,443 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Multi-cache dynamic NTK RoPE for ArcticInference. + +This module implements :class:`MultiCacheDynamicNTKRotaryEmbedding`: a +RoPE layer that precomputes *multiple* cos/sin caches (one per static +scaling factor) and routes each token to the cache whose factor best +covers its current sequence length. Within any given bucket, the +layer is bit-identical to vLLM's +:class:`~vllm.model_executor.layers.rotary_embedding.DynamicNTKScalingRotaryEmbedding` +built with that factor. The module is therefore best understood as +"a seq_len-routed ensemble of static dynamic-NTK configurations". + +Why this shape (vs. an on-the-fly per-token formula)? + +* **Training distribution.** Models are typically fine-tuned with a + *single* static rope factor. Matching that exact cache for requests + within that factor's regime preserves the training distribution. + The continuous (per-token) HF formula replaces the static eff_base + with a seq_len-dependent one, which the model has never seen. +* **Performance.** The forward reuses vLLM's existing fused rope CUDA + kernel (``ops.rotary_embedding``) unchanged -- we just add a + per-token offset tensor to ``positions`` before the kernel looks up + ``cos_sin_cache``. No new kernels required. This is the same trick + :class:`DeepseekScalingRotaryEmbedding` uses for YaRN offsets. + +Cache layout +------------ +The caches are concatenated along axis 0, exactly like +:class:`LinearScalingRotaryEmbedding`:: + + cos_sin_cache = [ cache_F1 (size F1 * m) , + cache_F2 (size F2 * m) , + ... + cache_Fk (size Fk * m) ] + +For a request at position ``p`` that the router assigns to factor +``F_i``, the kernel indexes ``cos_sin_cache[offsets[i] + p]``. +``offsets[i]`` is the cumulative size of caches strictly before +``cache_F_i``. + +Routing +------- +``factor = clamp(ceil(seq_len / m), 1, max_factor)``: pick the smallest +factor whose cache covers the request's seq_len. Routing is +per-*token*, so mixed batches (e.g. a short decode request alongside a +long prefill request) rotate each token under its own factor's cache. + +CUDA-graph safety +----------------- +* The per-token offset tensor lives in a registered GPU buffer on the + module (``runtime_bucket_offsets``). +* :meth:`update_runtime_seq_lens` writes into that buffer **in place** + each forward. The model runner calls this before invoking the + model. +* The forward reads ``self.runtime_bucket_offsets[:num_tokens]`` with + a fixed-shape slice; captured graphs record a load from this + buffer's storage, and each replay picks up whatever was most + recently written. +""" + +from __future__ import annotations + +import math +from typing import Optional, Sequence + +import torch + +from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding + + +_DEFAULT_BUFFER_SIZE = 2048 + +#: Default factor set (fixed integers 1..6, each bucket one ``m`` wide). +DEFAULT_FACTORS: tuple[float, ...] = (1.0, 2.0, 3.0, 4.0, 5.0, 6.0) + + +class MultiCacheDynamicNTKRotaryEmbedding(RotaryEmbedding): + """RoPE with N concatenated static dynamic-NTK caches + per-token routing. + + Args: + head_size: Attention head dimension. + rotary_dim: Number of dimensions to rotate (``<= head_size``). + max_position_embeddings: The **original** max position (``m``) + the model was trained on. Each factor ``F``'s cache spans + ``ceil(F * m)`` positions. + base: ``rope_theta`` from the HF config. + is_neox_style: NeoX (halves) vs GPT-J (interleaved) rotation. + dtype: Cache dtype (``cos_sin_cache`` is cast to this). + factors: Sorted list of scaling factors. Defaults to + ``DEFAULT_FACTORS`` = ``(1.0, 2.0, 3.0, 4.0, 5.0, 6.0)``. + Each factor must be ``>= 1.0``. + max_num_batched_tokens: Upper bound on ``num_tokens`` in any + single forward. Determines the size of + ``runtime_bucket_offsets``. The model runner writes actual + per-token seq-lens (via :meth:`update_runtime_seq_lens`, + which converts them to offsets) before each forward. + """ + + DEFAULT_FACTORS = DEFAULT_FACTORS + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: float, + is_neox_style: bool, + dtype: torch.dtype, + factors: Optional[Sequence[float]] = None, + max_num_batched_tokens: Optional[int] = None, + ) -> None: + raw_factors = ( + list(factors) if factors is not None else list(DEFAULT_FACTORS) + ) + # Dedupe + sort; validate monotone >= 1.0. + normalized = sorted({float(f) for f in raw_factors}) + if not normalized or normalized[0] < 1.0: + raise ValueError( + "factors must be a non-empty list of values >= 1.0, got " + f"{normalized}" + ) + self.factors: list[float] = normalized + self._original_max_position_embeddings = int(max_position_embeddings) + + # Per-factor cache sizes and cumulative offsets. + self._per_factor_max_len: list[int] = [ + int(math.ceil(F * self._original_max_position_embeddings)) + for F in self.factors + ] + cumulative = 0 + per_factor_offsets: list[int] = [] + for mx in self._per_factor_max_len: + per_factor_offsets.append(cumulative) + cumulative += mx + self._per_factor_offsets: list[int] = per_factor_offsets + self._total_cache_len: int = cumulative + + # Runtime buffer sizing. + if max_num_batched_tokens is None: + buf_size = _DEFAULT_BUFFER_SIZE + else: + buf_size = int(max_num_batched_tokens) + self._runtime_buffer_size = max(1, buf_size) + + # Parent __init__ calls ``_compute_cos_sin_cache`` and registers + # ``cos_sin_cache``. Our override concatenates per-factor + # caches; it relies on ``self.factors`` et al already being set. + super().__init__( + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=max_position_embeddings, + base=base, + is_neox_style=is_neox_style, + dtype=dtype, + ) + + # Per-token offset buffer. Initialised to 0 = factor-1 cache, + # which is the unscaled base. This keeps profile / dummy runs + # that bypass the model runner in a numerically safe regime. + self.register_buffer( + "runtime_bucket_offsets", + torch.zeros(self._runtime_buffer_size, dtype=torch.long), + persistent=False, + ) + + # Lookup table used by ``update_runtime_seq_lens`` to translate + # a bucket index -> offset into the concatenated cache. + # Registered as a buffer so it rides with the module to the + # correct device (via ``.to(device)``). + self.register_buffer( + "_factor_offsets_tensor", + torch.tensor(self._per_factor_offsets, dtype=torch.long), + persistent=False, + ) + + # ------------------------------------------------------------------ + # Cache construction (per-factor dynamic-NTK caches, concatenated) + # ------------------------------------------------------------------ + def _compute_cos_sin_cache(self) -> torch.Tensor: + """Build ``cos_sin_cache`` as the concatenation of per-factor caches. + + For each factor ``F``: + + 1. ``max_len = ceil(F * m)`` positions (what vLLM's static + ``DynamicNTKScalingRotaryEmbedding`` would allocate). + 2. ``eff_base = base * ((F * max_len / m) - (F - 1)) ** (d/(d-2))`` + (the dynamic NTK formula evaluated at the cache's own max + seq_len, matching the static class's behavior). + 3. ``inv_freq = 1 / eff_base ** (arange(0, d, 2)/d)``. + 4. ``cache_F = concat(cos(t*inv_freq), sin(t*inv_freq))`` for + ``t in [0, max_len)``. + + Concatenating all ``cache_F`` along dim 0 yields the unified + ``cos_sin_cache``. Per-factor starts are recorded in + ``self._per_factor_offsets``. + """ + m = self._original_max_position_embeddings + d = self.rotary_dim + caches: list[torch.Tensor] = [] + for F in self.factors: + max_len = int(math.ceil(F * m)) + eff_base = self.base * ( + (F * max_len / m) - (F - 1.0) + ) ** (d / (d - 2)) + inv_freq = 1.0 / ( + eff_base + ** (torch.arange(0, d, 2, dtype=torch.float) / d) + ) + t = torch.arange(max_len, dtype=torch.float) + freqs = torch.einsum("i,j -> ij", t, inv_freq) + cos = freqs.cos() + sin = freqs.sin() + caches.append(torch.cat((cos, sin), dim=-1)) + return torch.cat(caches, dim=0) + + # ------------------------------------------------------------------ + # Introspection + # ------------------------------------------------------------------ + def factor_offset(self, factor: float) -> int: + """Return the offset into the concatenated cache for ``factor``.""" + try: + idx = self.factors.index(float(factor)) + except ValueError as e: + raise ValueError( + f"factor {factor!r} not in configured factors " + f"{self.factors!r}" + ) from e + return self._per_factor_offsets[idx] + + def bucket_for_seq_len(self, seq_len: int) -> int: + """Scalar routing function: return the 0-indexed bucket for a seq_len. + + ``bucket = clamp(ceil(seq_len / m), 1, len(factors)) - 1``. + Matches the vectorized routing in + :meth:`update_runtime_seq_lens`. + """ + m = self._original_max_position_embeddings + k = len(self.factors) + if seq_len <= 0: + return 0 + bucket_one_indexed = min(k, max(1, math.ceil(seq_len / m))) + return bucket_one_indexed - 1 + + # ------------------------------------------------------------------ + # Public state update (called by the model runner each forward) + # ------------------------------------------------------------------ + def update_runtime_seq_lens( + self, seq_lens: torch.Tensor, non_blocking: bool = True + ) -> None: + """Convert per-token seq_lens to offsets, in place. + + ``seq_lens`` is shape ``[num_tokens]`` (on any device, any + integer dtype). For each token ``i``: + + .. code-block:: python + + bucket[i] = clamp(ceil(seq_lens[i] / m), 1, k) - 1 + offset[i] = factor_offsets[bucket[i]] + + The computed ``offset`` is written into + ``self.runtime_bucket_offsets`` **in place**. Preserving the + buffer's storage is what makes captured CUDA graphs see fresh + values on replay. + """ + n = int(seq_lens.shape[0]) + if n > self._runtime_buffer_size: + raise ValueError( + f"MultiCacheDynamicNTKRotaryEmbedding: received {n} tokens " + f"but runtime buffer was sized for {self._runtime_buffer_size}. " + "Increase max_num_batched_tokens." + ) + if n == 0: + return + + m = self._original_max_position_embeddings + max_bucket = len(self.factors) + + # Move to the buffer's device/dtype in one shot. ``copy_`` on + # the destination would work, but we need the ceil_div + + # clamp arithmetic in int64 on-device anyway. + sl = seq_lens.to( + device=self.runtime_bucket_offsets.device, + dtype=torch.long, + non_blocking=non_blocking, + ) + # Integer ceil_div for positive x: (x + m - 1) // m. + bucket_idx = torch.clamp( + (sl + m - 1) // m, + min=1, + max=max_bucket, + ) - 1 + offsets = self._factor_offsets_tensor[bucket_idx] + self.runtime_bucket_offsets[:n].copy_(offsets, non_blocking=non_blocking) + + # ------------------------------------------------------------------ + # Forward methods: delegate to vLLM's fused kernel with + # ``positions + runtime_bucket_offsets`` as the effective positions. + # No per-token formula recompute; just a tiny add + the same CUDA + # rope kernel the rest of vLLM uses. + # ------------------------------------------------------------------ + def _effective_positions(self, positions: torch.Tensor) -> torch.Tensor: + num_tokens = positions.shape[0] + # Use the bucket-offset slice matching the batch size. The + # slice is a view, not a copy; no allocation. + offsets = self.runtime_bucket_offsets[:num_tokens] + return torch.add(positions, offsets) + + def forward_native( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """PyTorch-native forward (also used on CPU).""" + effective = self._effective_positions(positions) + return self.forward_static( + effective, + query, + key, + self.head_size, + self.rotary_dim, + self.cos_sin_cache, + self.is_neox_style, + ) + + def forward_cuda( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Fused CUDA rope: same kernel as static rope, but with + ``positions + per-token offset``. This is the fast path.""" + if self.use_flashinfer: + # Mirror DeepseekScalingRotaryEmbedding.forward_cuda: pass + # the combined positions to flashinfer. (Our base class + # sets use_flashinfer = False by default, so this branch + # is inactive unless explicitly enabled.) + effective = self._effective_positions(positions) + torch.ops.vllm.flashinfer_rotary_embedding( + effective, + query, + key, + self.head_size, + self.cos_sin_cache, + self.is_neox_style, + ) + return query, key + + from vllm import _custom_ops as ops + + effective = self._effective_positions(positions) + self._match_cos_sin_cache_dtype(query) + # In-place kernel: writes into query / key. + ops.rotary_embedding( + effective, + query, + key, + self.head_size, + self.cos_sin_cache, + self.is_neox_style, + ) + return query, key + + def forward_hip( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Same as forward_cuda for our purposes (fused kernel handles + both). The ROCm Triton rope kernel doesn't take offsets, so + route through the plain CUDA path.""" + return self.forward_cuda(positions, query, key) + + def forward_xpu( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + if key is None: + # XPU kernel doesn't support key=None; fall back to native. + return self.forward_native(positions, query, key) + from vllm._ipex_ops import ipex_ops as ops + + effective = self._effective_positions(positions) + self._match_cos_sin_cache_dtype(query) + ops.rotary_embedding( + effective, + query, + key, + self.head_size, + self.cos_sin_cache, + self.is_neox_style, + ) + return query, key + + def forward_cpu( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + from vllm import _custom_ops as ops + + effective = self._effective_positions(positions) + self._match_cos_sin_cache_dtype(query) + ops.rotary_embedding( + effective, + query, + key, + self.head_size, + self.cos_sin_cache, + self.is_neox_style, + ) + return query, key + + # ------------------------------------------------------------------ + # Diagnostics + # ------------------------------------------------------------------ + def extra_repr(self) -> str: + parent = super().extra_repr() + return ( + f"{parent}, factors={self.factors}, " + f"orig_max_position={self._original_max_position_embeddings}, " + f"runtime_buffer_size={self._runtime_buffer_size}, " + f"total_cache_len={self._total_cache_len}, " + f"multi_cache=True" + ) diff --git a/arctic_inference/vllm/rope/patches.py b/arctic_inference/vllm/rope/patches.py new file mode 100644 index 000000000..fa307bec9 --- /dev/null +++ b/arctic_inference/vllm/rope/patches.py @@ -0,0 +1,334 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Patches to wire multi-cache dynamic NTK RoPE into vLLM's get_rope dispatch. + +We cannot add a new branch to +:func:`vllm.model_executor.layers.rotary_embedding.get_rope` directly +because that function lives upstream. Instead we wrap it: our wrapper +handles ``rope_type`` values that the multi-cache subsystem owns and +delegates everything else to the original. The wrapper reuses vLLM's +existing :data:`_ROPE_DICT` memoization so repeated calls with the +same key return a singleton. +""" + +from __future__ import annotations + +import os +from typing import Any + +import torch + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +#: rope_type values we dispatch to :class:`MultiCacheDynamicNTKRotaryEmbedding`. +_MULTI_CACHE_ROPE_TYPES = frozenset( + [ + "multi_cache_ntk", + # Short alias, reads as "the multi-cache variant of dynamic". + "dynamic_multi_cache", + ] +) + + +def _build_multi_cache_ntk( + head_size: int, + rotary_dim: int, + max_position: int, + base: float, + is_neox_style: bool, + rope_parameters: dict[str, Any], + dtype: torch.dtype, +): + """Build a :class:`MultiCacheDynamicNTKRotaryEmbedding` from rope params. + + Supported ``rope_parameters`` keys (all optional): + + * ``factors``: list of scaling factors (e.g. ``[1, 2, 3, 4, 5, 6]``). + Defaults to + :data:`arctic_inference.vllm.rope.multi_cache_ntk.DEFAULT_FACTORS`. + * ``original_max_position_embeddings``: the trained context (``m``). + Each factor ``F`` gets a cache spanning ``ceil(F * m)`` positions. + Defaults to ``max_position`` when omitted. + + The trailing ``factor`` (singular) key from legacy dynamic configs is + folded into ``factors`` as ``[factor]`` when present and ``factors`` + is not explicitly set. This keeps backward compat with configs + written for the old static ``dynamic`` rope. + + ``max_num_batched_tokens`` is looked up from the active + :class:`vllm.config.VllmConfig` to size the runtime offset buffer. + The module falls back to a conservative default when no config is + active (typical in unit tests). + """ + from arctic_inference.vllm.rope.multi_cache_ntk import ( + DEFAULT_FACTORS, + MultiCacheDynamicNTKRotaryEmbedding, + ) + + orig_max = rope_parameters.get( + "original_max_position_embeddings", max_position + ) + + # Resolve the factor list: explicit ``factors`` wins; otherwise + # fall back to a single-value list from ``factor``; otherwise + # defaults. This ordering is deliberate so that configs that + # already set ``factors`` are not second-guessed. + if "factors" in rope_parameters: + raw_factors = list(rope_parameters["factors"]) + elif "factor" in rope_parameters: + raw_factors = [float(rope_parameters["factor"])] + else: + raw_factors = list(DEFAULT_FACTORS) + + max_num_batched_tokens: int | None = None + try: + from vllm.config import get_current_vllm_config_or_none + + cfg = get_current_vllm_config_or_none() + if ( + cfg is not None + and getattr(cfg, "scheduler_config", None) is not None + ): + max_num_batched_tokens = int( + cfg.scheduler_config.max_num_batched_tokens + ) + except Exception: # pragma: no cover - defensive + max_num_batched_tokens = None + + return MultiCacheDynamicNTKRotaryEmbedding( + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=int(orig_max), + base=base, + is_neox_style=is_neox_style, + dtype=dtype, + factors=raw_factors, + max_num_batched_tokens=max_num_batched_tokens, + ) + + +def _install_apply_dict_overrides_patch() -> None: + """Re-run ``patch_rope_parameters`` after dict-valued hf_overrides apply. + + In vLLM 0.14.1, :meth:`ModelConfig.__post_init__` splits its + ``hf_overrides`` into *flat* values (passed into ``get_config``, which + runs ``patch_rope_parameters``) and *dict* values (applied via + ``_apply_dict_overrides`` *after* the HF config has already been + loaded and patched). ``rope_scaling`` is a dict, so it always lands + in the second bucket. The consequence is that + ``config.rope_parameters`` never picks up the override, so the + scaling factor in ``_get_and_verify_max_len`` stays at ``1.0`` and + ``derived_max_model_len`` stays at the unscaled ``max_position_embeddings``. + + We wrap ``_apply_dict_overrides`` to invoke ``patch_rope_parameters`` + on the target config (and its text sub-config) after the dict values + land, which re-materializes ``rope_parameters`` from the just-written + ``rope_scaling``. Idempotent; a second install is a no-op. + """ + from vllm.config.model import ModelConfig + from vllm.transformers_utils.config import patch_rope_parameters + + orig = ModelConfig._apply_dict_overrides + if getattr(orig, "_arctic_wrapped", False): + return + + def _apply_dict_overrides_arctic(self, config, overrides): + orig(self, config, overrides) + if "rope_scaling" not in overrides: + return + patch_rope_parameters(config) + try: + text_cfg = config.get_text_config() + except Exception: # pragma: no cover - non-text configs + text_cfg = None + if text_cfg is not None and text_cfg is not config: + patch_rope_parameters(text_cfg) + + _apply_dict_overrides_arctic._arctic_wrapped = True # type: ignore[attr-defined] + _apply_dict_overrides_arctic._arctic_original = orig # type: ignore[attr-defined] + ModelConfig._apply_dict_overrides = _apply_dict_overrides_arctic # type: ignore[assignment] + + +def _install_get_rope_patch() -> None: + """Install the ``get_rope`` wrapper. + + Idempotent: installing twice is a no-op. The wrapper is attached as + ``_arctic_wrapped`` so we can detect it on subsequent imports. + """ + import vllm.model_executor.layers.rotary_embedding as rope_mod + + orig = rope_mod.get_rope + if getattr(orig, "_arctic_wrapped", False): + return + + def get_rope_arctic( + head_size: int, + max_position: int, + is_neox_style: bool = True, + rope_parameters: dict[str, Any] | None = None, + dtype: torch.dtype | None = None, + dual_chunk_attention_config: dict[str, Any] | None = None, + ): + # Optionally promote legacy dynamic rope configs to multi-cache. + rope_parameters = maybe_promote_rope_parameters(rope_parameters) + + scaling_type = None + if rope_parameters is not None: + scaling_type = rope_parameters.get("rope_type") + + if scaling_type not in _MULTI_CACHE_ROPE_TYPES: + return orig( + head_size=head_size, + max_position=max_position, + is_neox_style=is_neox_style, + rope_parameters=rope_parameters, + dtype=dtype, + dual_chunk_attention_config=dual_chunk_attention_config, + ) + + if dual_chunk_attention_config is not None: + raise ValueError( + "multi_cache_ntk rope is incompatible with " + "dual_chunk_attention_config" + ) + + if dtype is None: + dtype = torch.get_default_dtype() + assert rope_parameters is not None + + base = rope_parameters.get("rope_theta", 10000) + partial_rotary_factor = rope_parameters.get( + "partial_rotary_factor", 1.0, + ) + if partial_rotary_factor <= 0.0 or partial_rotary_factor > 1.0: + raise ValueError( + f"{partial_rotary_factor=} must be between 0.0 and 1.0" + ) + rotary_dim = int(head_size * partial_rotary_factor) + + # Reuse vLLM's cache so repeated calls for the same layer family + # get the same singleton module. + rope_parameters_tuple = { + k: tuple(v) if isinstance(v, list) else v + for k, v in rope_parameters.items() + } + rope_parameters_args = tuple(rope_parameters_tuple.items()) + key = ( + head_size, + rotary_dim, + max_position, + is_neox_style, + rope_parameters_args, + None, # dual_chunk_attention_args + dtype, + ) + cache = rope_mod._ROPE_DICT + if key in cache: + return cache[key] + + module = _build_multi_cache_ntk( + head_size=head_size, + rotary_dim=rotary_dim, + max_position=max_position, + base=base, + is_neox_style=is_neox_style, + rope_parameters=rope_parameters, + dtype=dtype, + ) + cache[key] = module + logger.info( + "ArcticInference installed MultiCacheDynamicNTKRotaryEmbedding " + "(head_size=%d, rotary_dim=%d, factors=%s, orig_max=%s, " + "max_position=%s, runtime_buffer_size=%s, total_cache_len=%s)", + head_size, + rotary_dim, + module.factors, + rope_parameters.get( + "original_max_position_embeddings", max_position + ), + max_position, + module._runtime_buffer_size, + module._total_cache_len, + ) + return module + + get_rope_arctic._arctic_wrapped = True # type: ignore[attr-defined] + get_rope_arctic._arctic_original = orig # type: ignore[attr-defined] + rope_mod.get_rope = get_rope_arctic + + # Also patch the re-export path used by some model implementations + # that import ``get_rope`` from the layers package root. + try: + import vllm.model_executor.layers as layers_mod + + if hasattr(layers_mod, "get_rope"): + layers_mod.get_rope = get_rope_arctic + except Exception: # pragma: no cover - defensive + pass + + +# -------------------------------------------------------------------------- +# Config promotion helper +# -------------------------------------------------------------------------- + +_ARCTIC_MULTI_CACHE_ROPE_ENV = "ARCTIC_INFERENCE_MULTI_CACHE_ROPE" + + +def multi_cache_rope_enabled() -> bool: + """Return True if ArcticInference should promote dynamic rope to multi-cache. + + Controlled by the env var ``ARCTIC_INFERENCE_MULTI_CACHE_ROPE``: + ``"1"`` enables, anything else (or unset) disables. + """ + return os.getenv(_ARCTIC_MULTI_CACHE_ROPE_ENV, "0") == "1" + + +def maybe_promote_rope_parameters( + rope_parameters: dict[str, Any] | None, +) -> dict[str, Any] | None: + """Rewrite ``rope_parameters["rope_type"]`` to the multi-cache variant. + + Only fires when :func:`multi_cache_rope_enabled` returns True and the + incoming ``rope_type`` is the static-dynamic alias. The rewrite is + a shallow copy so the caller's dict is not mutated. Users who want + to opt-in per-model can set ``rope_type="multi_cache_ntk"`` directly + in the HF config. + + If the incoming config specifies a single ``factor`` but no + ``factors``, the promoter leaves it alone and + :func:`_build_multi_cache_ntk` will fold it into a single-bucket + multi-cache (degenerate but valid). To activate the full default + factor set, the caller should set ``factors`` explicitly. The + legacy "alpha" variant is not promoted because it uses a different + base formula. + """ + if rope_parameters is None or not multi_cache_rope_enabled(): + return rope_parameters + rope_type = rope_parameters.get("rope_type") + if rope_type in ("dynamic",): + if "alpha" in rope_parameters and "factor" not in rope_parameters: + # The alpha variant has no multi-cache analogue; leave it. + return rope_parameters + promoted = dict(rope_parameters) + promoted["rope_type"] = "multi_cache_ntk" + logger.info( + "ArcticInference: promoting rope_type='dynamic' to " + "'multi_cache_ntk' (factors=%s)", + promoted.get("factors", promoted.get("factor", "default")), + ) + return promoted + return rope_parameters diff --git a/tests/unit_tests/test_multi_cache_rope.py b/tests/unit_tests/test_multi_cache_rope.py new file mode 100644 index 000000000..69c99e6ac --- /dev/null +++ b/tests/unit_tests/test_multi_cache_rope.py @@ -0,0 +1,938 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Numerical and wiring tests for MultiCacheDynamicNTKRotaryEmbedding. + +These tests run on the GPU because that is the deployment target: the +forward re-uses vLLM's fused ``ops.rotary_embedding`` CUDA kernel, and +we want the tests to exercise the same device, dtype, and kernel paths +that production will. Tests that depend on ``get_rope`` wiring are +device-independent and stay on CPU. + +The ``forward_cuda`` path is a single elementwise add followed by the +same CUDA rope kernel the rest of vLLM uses. That's what makes it +safe under CUDA graphs. A few of the tests below additionally +exercise ``torch.compile(fullgraph=True)`` as a static proxy for +"no graph breaks", plus a real CUDA-graph capture/replay test for the +end-to-end safety property. +""" + +from __future__ import annotations + +import math +from typing import Tuple + +import pytest +import torch + +pytest.importorskip("vllm") + + +if not torch.cuda.is_available(): + pytest.skip( + "CUDA is required for multi-cache RoPE tests (they run on GPU " + "to match the production code path).", + allow_module_level=True, + ) + + +@pytest.fixture(autouse=True) +def _set_default_dtype(): + prev = torch.get_default_dtype() + torch.set_default_dtype(torch.float32) + yield + torch.set_default_dtype(prev) + + +@pytest.fixture(autouse=True) +def _default_vllm_config(): + """Provide a default ``VllmConfig`` for every test. + + ``RotaryEmbedding`` is a ``CustomOp`` whose ``__init__`` reads + ``get_current_vllm_config()`` to pick the platform forward. Outside + a real engine that context is unset, so direct construction in tests + raises ``AssertionError: Current vLLM config is not set.``. This + mirrors the ``default_vllm_config`` fixture in vLLM's own test suite. + """ + from vllm.config import VllmConfig, set_current_vllm_config + + with set_current_vllm_config(VllmConfig()): + yield + + +@pytest.fixture +def device() -> torch.device: + return torch.device("cuda:0") + + +# -------------------------------------------------------------------------- +# Helpers +# -------------------------------------------------------------------------- + + +def _make_multi_cache( + *, + head_size: int = 64, + rotary_dim: int = 64, + max_position_embeddings: int = 2048, + base: float = 10000.0, + factors=None, + is_neox_style: bool = True, + dtype: torch.dtype = torch.float32, + max_num_batched_tokens: int | None = None, + device: torch.device | None = None, +): + from arctic_inference.vllm.rope.multi_cache_ntk import ( + MultiCacheDynamicNTKRotaryEmbedding, + ) + + mod = MultiCacheDynamicNTKRotaryEmbedding( + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=max_position_embeddings, + base=base, + is_neox_style=is_neox_style, + dtype=dtype, + factors=factors, + max_num_batched_tokens=max_num_batched_tokens, + ) + if device is not None: + mod = mod.to(device) + return mod + + +def _make_static_dynamic_ntk( + *, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: float, + is_neox_style: bool, + scaling_factor: float, + dtype: torch.dtype, + device: torch.device, +): + """Build vLLM's static :class:`DynamicNTKScalingRotaryEmbedding` as a + reference. Within any bucket our multi-cache must match this class + bit-for-bit at the corresponding factor.""" + from vllm.model_executor.layers.rotary_embedding.dynamic_ntk_scaling_rope import ( + DynamicNTKScalingRotaryEmbedding, + ) + + mod = DynamicNTKScalingRotaryEmbedding( + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=max_position_embeddings, + base=base, + is_neox_style=is_neox_style, + scaling_factor=scaling_factor, + dtype=dtype, + ) + return mod.to(device) + + +def _make_inputs( + *, + num_tokens: int, + num_heads: int = 4, + num_kv_heads: int = 2, + head_size: int = 64, + seq_len_start: int = 0, + dtype: torch.dtype = torch.float32, + seed: int = 0, + device: torch.device | None = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + gen_device = device or torch.device("cpu") + gen = torch.Generator(device=gen_device).manual_seed(seed) + positions = torch.arange( + seq_len_start, seq_len_start + num_tokens, + dtype=torch.long, device=gen_device, + ) + query = torch.randn( + num_tokens, num_heads * head_size, dtype=dtype, + device=gen_device, generator=gen, + ) + key = torch.randn( + num_tokens, num_kv_heads * head_size, dtype=dtype, + device=gen_device, generator=gen, + ) + return positions, query, key + + +# -------------------------------------------------------------------------- +# Cache construction +# -------------------------------------------------------------------------- + + +def test_per_factor_caches_match_vllm_static_dynamic_ntk(device): + """For each configured factor F, the slice of the concatenated cache + belonging to F must be bit-identical to + :class:`DynamicNTKScalingRotaryEmbedding(factor=F)`'s cache. + """ + head_size = 64 + rotary_dim = 64 + max_pos = 2048 + base = 10000.0 + factors = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + + multi = _make_multi_cache( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, base=base, + factors=factors, device=device, + ) + multi_cache = multi.cos_sin_cache + + for F in factors: + static = _make_static_dynamic_ntk( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, base=base, + is_neox_style=True, scaling_factor=F, + dtype=torch.float32, device=device, + ) + offset = multi.factor_offset(F) + length = int(math.ceil(F * max_pos)) + multi_slice = multi_cache[offset: offset + length] + assert multi_slice.shape == static.cos_sin_cache.shape, ( + f"cache shape mismatch for F={F}: " + f"{tuple(multi_slice.shape)} vs {tuple(static.cos_sin_cache.shape)}" + ) + torch.testing.assert_close( + multi_slice, static.cos_sin_cache, rtol=0, atol=0, + ) + + +def test_cache_layout_offsets_are_cumulative(device): + """``_per_factor_offsets`` must equal the prefix sums of the + per-factor cache sizes, and the total cache length must equal the + concatenated cache's length.""" + max_pos = 2048 + factors = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + multi = _make_multi_cache( + max_position_embeddings=max_pos, factors=factors, device=device, + ) + + expected_sizes = [int(math.ceil(F * max_pos)) for F in factors] + expected_offsets = [0] + for sz in expected_sizes[:-1]: + expected_offsets.append(expected_offsets[-1] + sz) + + assert multi._per_factor_offsets == expected_offsets + assert multi._total_cache_len == sum(expected_sizes) + assert multi.cos_sin_cache.shape[0] == sum(expected_sizes) + + +def test_factors_are_sorted_and_deduped(device): + """Construction must sort + dedupe the factor list. Downstream + routing assumes strictly increasing unique factors.""" + multi = _make_multi_cache( + factors=[4.0, 2.0, 4.0, 1.0, 3.0], device=device, + ) + assert multi.factors == [1.0, 2.0, 3.0, 4.0] + + +def test_invalid_factors_raise(): + """Factors < 1.0 don't make sense (there's no factor smaller than + the unscaled baseline) and must be rejected.""" + with pytest.raises(ValueError, match=">=.*1.0"): + _make_multi_cache(factors=[0.5, 1.0]) + with pytest.raises(ValueError, match=">=.*1.0"): + _make_multi_cache(factors=[]) + + +# -------------------------------------------------------------------------- +# Routing +# -------------------------------------------------------------------------- + + +def test_bucket_for_seq_len_picks_smallest_covering_factor(device): + """Scalar routing function: + ``bucket = clamp(ceil(seq_len / m), 1, k) - 1``.""" + m = 2048 + multi = _make_multi_cache( + max_position_embeddings=m, + factors=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + device=device, + ) + assert multi.bucket_for_seq_len(0) == 0 # degenerate safe-default + assert multi.bucket_for_seq_len(1) == 0 + assert multi.bucket_for_seq_len(m) == 0 # exactly m fits factor 1 + assert multi.bucket_for_seq_len(m + 1) == 1 + assert multi.bucket_for_seq_len(2 * m) == 1 + assert multi.bucket_for_seq_len(2 * m + 1) == 2 + assert multi.bucket_for_seq_len(6 * m) == 5 + assert multi.bucket_for_seq_len(6 * m + 1) == 5 # clamped to last + assert multi.bucket_for_seq_len(100 * m) == 5 + + +def test_update_runtime_seq_lens_translates_to_offsets(device): + """The vectorized update must produce the same offsets the scalar + routing function returns.""" + m = 2048 + factors = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + multi = _make_multi_cache( + max_position_embeddings=m, factors=factors, + max_num_batched_tokens=64, device=device, + ) + + # Hand-picked seq_lens that exercise every bucket, including the + # boundary + overflow cases. + seq_lens_cpu = [ + 0, 1, 1024, m, + m + 1, m * 2, + m * 2 + 1, m * 3, + m * 3 + 1, m * 4, + m * 4 + 1, m * 5, + m * 5 + 1, m * 6, m * 6 + 1, 10_000_000, + ] + seq_lens = torch.tensor(seq_lens_cpu, dtype=torch.int32, device=device) + multi.update_runtime_seq_lens(seq_lens) + + expected = [ + multi.factor_offset(factors[multi.bucket_for_seq_len(int(s))]) + for s in seq_lens_cpu + ] + got = multi.runtime_bucket_offsets[: len(seq_lens_cpu)].tolist() + assert got == expected + + +def test_update_runtime_seq_lens_is_in_place(device): + """The update must preserve the buffer's storage so CUDA graph + replay sees new values without a re-capture.""" + multi = _make_multi_cache( + max_num_batched_tokens=256, device=device, + ) + storage_ptr_before = multi.runtime_bucket_offsets.data_ptr() + buffer_id_before = id(multi.runtime_bucket_offsets) + + multi.update_runtime_seq_lens( + torch.arange(1, 33, dtype=torch.int32, device=device), + ) + + assert multi.runtime_bucket_offsets.data_ptr() == storage_ptr_before + assert id(multi.runtime_bucket_offsets) == buffer_id_before + + +def test_update_runtime_seq_lens_rejects_oversized_input(device): + """Guard against silent truncation: too many tokens is an error.""" + multi = _make_multi_cache(max_num_batched_tokens=32, device=device) + with pytest.raises(ValueError, match="runtime buffer was sized"): + multi.update_runtime_seq_lens( + torch.ones(64, dtype=torch.int32, device=device), + ) + + +def test_update_runtime_seq_lens_handles_empty_batch(device): + """``seq_lens`` of shape [0] is a no-op (doesn't touch the buffer).""" + multi = _make_multi_cache( + max_num_batched_tokens=32, device=device, + ) + before = multi.runtime_bucket_offsets.clone() + multi.update_runtime_seq_lens( + torch.empty(0, dtype=torch.int32, device=device), + ) + torch.testing.assert_close(multi.runtime_bucket_offsets, before) + + +# -------------------------------------------------------------------------- +# Forward correctness vs static DynamicNTK reference +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("factor", [1.0, 2.0, 4.0, 6.0]) +def test_forward_single_bucket_matches_static_rope(device, factor): + """All-tokens-same-seq_len batch: the output must match vLLM's + :class:`DynamicNTKScalingRotaryEmbedding(factor=factor)` exactly. + + This is the core invariant: "within a bucket, we *are* static rope".""" + head_size, rotary_dim = 64, 64 + max_pos = 2048 + base = 10000.0 + factors = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + num_tokens = 32 + # Pick a seq_len inside the bucket range. For factor=F, the bucket + # covers ((F-1)*m, F*m]. We aim mid-range: (F - 0.5) * m. + seq_len = max(1, int((factor - 0.5) * max_pos)) + + multi = _make_multi_cache( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, base=base, + factors=factors, max_num_batched_tokens=num_tokens, + device=device, + ) + static = _make_static_dynamic_ntk( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, base=base, + is_neox_style=True, scaling_factor=factor, + dtype=torch.float32, device=device, + ) + + positions, q, k = _make_inputs( + num_tokens=num_tokens, head_size=head_size, + seq_len_start=max(0, seq_len - num_tokens), device=device, + ) + + multi.update_runtime_seq_lens( + torch.full( + (num_tokens,), seq_len, dtype=torch.int32, device=device, + ), + ) + got_q, got_k = multi.forward_native(positions, q.clone(), k.clone()) + ref_q, ref_k = static.forward_native(positions, q.clone(), k.clone()) + + torch.testing.assert_close(got_q, ref_q, rtol=0, atol=0) + torch.testing.assert_close(got_k, ref_k, rtol=0, atol=0) + + +def test_forward_mixed_batch_routes_each_request_correctly(device): + """A batch mixing two requests with different seq_lens (so different + buckets) must rotate each request's slice of the output under its + own bucket's static rope cache.""" + head_size, rotary_dim = 64, 64 + max_pos = 2048 + base = 10000.0 + factors = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + num_a, num_b = 17, 11 + seq_len_a = 1800 # bucket 0 (factor 1) + seq_len_b = 7500 # bucket 3 (factor 4) + + multi = _make_multi_cache( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, base=base, + factors=factors, max_num_batched_tokens=num_a + num_b, + device=device, + ) + + pos_a, q_a, k_a = _make_inputs( + num_tokens=num_a, head_size=head_size, + seq_len_start=0, seed=0, device=device, + ) + pos_b, q_b, k_b = _make_inputs( + num_tokens=num_b, head_size=head_size, + seq_len_start=seq_len_b - num_b, seed=42, device=device, + ) + + positions = torch.cat([pos_a, pos_b], dim=0) + query = torch.cat([q_a, q_b], dim=0) + key = torch.cat([k_a, k_b], dim=0) + per_token = torch.cat([ + torch.full((num_a,), seq_len_a, dtype=torch.int32, device=device), + torch.full((num_b,), seq_len_b, dtype=torch.int32, device=device), + ]) + + multi.update_runtime_seq_lens(per_token) + got_q, got_k = multi.forward_native( + positions, query.clone(), key.clone(), + ) + + # Build per-bucket reference with the appropriate static factor. + static_a = _make_static_dynamic_ntk( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, base=base, + is_neox_style=True, + scaling_factor=factors[multi.bucket_for_seq_len(seq_len_a)], + dtype=torch.float32, device=device, + ) + static_b = _make_static_dynamic_ntk( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, base=base, + is_neox_style=True, + scaling_factor=factors[multi.bucket_for_seq_len(seq_len_b)], + dtype=torch.float32, device=device, + ) + ref_q_a, ref_k_a = static_a.forward_native( + pos_a, q_a.clone(), k_a.clone(), + ) + ref_q_b, ref_k_b = static_b.forward_native( + pos_b, q_b.clone(), k_b.clone(), + ) + + torch.testing.assert_close(got_q[:num_a], ref_q_a, rtol=0, atol=0) + torch.testing.assert_close(got_k[:num_a], ref_k_a, rtol=0, atol=0) + torch.testing.assert_close(got_q[num_a:], ref_q_b, rtol=0, atol=0) + torch.testing.assert_close(got_k[num_a:], ref_k_b, rtol=0, atol=0) + + +def test_forward_cuda_matches_forward_native(device): + """The fused CUDA path and the PyTorch path must produce the same + output; that's the contract that lets us use the kernel with no + further adaptation. Use a mid-range seq_len and a non-trivial + bucket to exercise a real offset.""" + head_size, rotary_dim = 64, 64 + max_pos = 2048 + num_tokens = 32 + seq_len = 5000 # bucket 2 (factor 3) + + multi = _make_multi_cache( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, + max_num_batched_tokens=num_tokens, device=device, + ) + positions, q, k = _make_inputs( + num_tokens=num_tokens, head_size=head_size, + seq_len_start=seq_len - num_tokens, device=device, + ) + + multi.update_runtime_seq_lens( + torch.full( + (num_tokens,), seq_len, dtype=torch.int32, device=device, + ), + ) + + q_native = q.clone() + k_native = k.clone() + got_q_native, got_k_native = multi.forward_native( + positions, q_native, k_native, + ) + + q_cuda = q.clone() + k_cuda = k.clone() + # forward_cuda writes in place, so capture post-call tensors. + got_q_cuda, got_k_cuda = multi.forward_cuda( + positions, q_cuda, k_cuda, + ) + + torch.testing.assert_close( + got_q_native, got_q_cuda, rtol=1e-4, atol=1e-5, + ) + torch.testing.assert_close( + got_k_native, got_k_cuda, rtol=1e-4, atol=1e-5, + ) + + +# -------------------------------------------------------------------------- +# Buffer sizing + init state +# -------------------------------------------------------------------------- + + +def test_runtime_buffer_sized_by_max_num_batched_tokens(): + multi = _make_multi_cache( + max_position_embeddings=2048, + max_num_batched_tokens=4096, + ) + assert multi.runtime_bucket_offsets.shape == (4096,) + + +def test_runtime_buffer_has_conservative_default_when_nothing_specified(): + multi = _make_multi_cache( + max_position_embeddings=2048, + max_num_batched_tokens=None, + ) + # Conservative default >= 1 and matches the module-private constant. + from arctic_inference.vllm.rope import multi_cache_ntk as mcn + + assert multi.runtime_bucket_offsets.shape[0] == mcn._DEFAULT_BUFFER_SIZE + + +def test_runtime_buffer_initialized_to_factor_one_offset(): + """Before any update, all offsets should be 0 -- i.e. all tokens + route to the unscaled factor-1 cache. This is the safe regime for + profile / dummy runs that bypass the model runner.""" + multi = _make_multi_cache( + max_position_embeddings=2048, factors=[1.0, 2.0, 4.0], + max_num_batched_tokens=128, + ) + assert (multi.runtime_bucket_offsets == 0).all() + + +# -------------------------------------------------------------------------- +# Graph-compatibility smoke tests +# -------------------------------------------------------------------------- + + +def test_forward_compiles_without_graph_breaks(device): + """``torch.compile(fullgraph=True)`` traces ``forward_native`` without + hitting Python-level control flow on tensor values. This is a + static proxy for "safe to capture in a CUDA graph".""" + pytest.importorskip("torch._dynamo") + + head_size, rotary_dim = 64, 64 + max_pos = 2048 + num_tokens = 16 + + multi = _make_multi_cache( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, + max_num_batched_tokens=num_tokens, device=device, + ) + multi.update_runtime_seq_lens( + torch.full( + (num_tokens,), 5000, dtype=torch.int32, device=device, + ), + ) + + positions, q, k = _make_inputs( + num_tokens=num_tokens, head_size=head_size, + seq_len_start=4000, device=device, + ) + + import torch._dynamo as dynamo + + dynamo.reset() + try: + compiled = torch.compile( + multi.forward_native, + backend="eager", + fullgraph=True, + dynamic=False, + ) + eager_q, eager_k = multi.forward_native( + positions, q.clone(), k.clone(), + ) + compiled_q, compiled_k = compiled( + positions, q.clone(), k.clone(), + ) + finally: + dynamo.reset() + + torch.testing.assert_close( + eager_q, compiled_q, rtol=1e-5, atol=1e-5, + ) + torch.testing.assert_close( + eager_k, compiled_k, rtol=1e-5, atol=1e-5, + ) + + +def test_compiled_forward_sees_fresh_buffer_values_between_calls(device): + """After compilation, updating the buffer must still change the + output on the next call. If the buffer read were constant-folded + into the graph, the second call would reuse the first call's + values -- which is exactly the silent correctness bug we care + about for CUDA graph replay.""" + pytest.importorskip("torch._dynamo") + + head_size, rotary_dim = 64, 64 + max_pos = 2048 + num_tokens = 16 + + multi = _make_multi_cache( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, + max_num_batched_tokens=num_tokens, device=device, + ) + positions, q, k = _make_inputs( + num_tokens=num_tokens, head_size=head_size, + seq_len_start=0, device=device, + ) + + import torch._dynamo as dynamo + + dynamo.reset() + try: + compiled = torch.compile( + multi.forward_native, + backend="eager", + fullgraph=True, + dynamic=False, + ) + + # First call: all tokens route to bucket 0 (factor 1). + multi.update_runtime_seq_lens( + torch.full( + (num_tokens,), 1000, dtype=torch.int32, device=device, + ), + ) + out_a_q, out_a_k = compiled(positions, q.clone(), k.clone()) + + # Second call: all tokens route to bucket 3 (factor 4). + multi.update_runtime_seq_lens( + torch.full( + (num_tokens,), 7000, dtype=torch.int32, device=device, + ), + ) + out_b_q, out_b_k = compiled(positions, q.clone(), k.clone()) + finally: + dynamo.reset() + + assert not torch.allclose(out_a_q, out_b_q) + assert not torch.allclose(out_a_k, out_b_k) + + +def test_forward_survives_cuda_graph_capture_and_replay(device): + """End-to-end validation of the CUDA-graph safety property. + + We capture the rotary forward into a CUDA graph while the runtime + buffer routes tokens to bucket A, then update the buffer so they + route to bucket B and replay. Replay must produce the same output + as running the forward eagerly under bucket B, proving the graph + reads live values from the buffer rather than constant-folding. + """ + head_size, rotary_dim = 64, 64 + max_pos = 2048 + num_tokens = 16 + + multi = _make_multi_cache( + head_size=head_size, rotary_dim=rotary_dim, + max_position_embeddings=max_pos, + max_num_batched_tokens=num_tokens, device=device, + ) + positions, q, k = _make_inputs( + num_tokens=num_tokens, head_size=head_size, + seq_len_start=0, device=device, + ) + + # Static tensors backing the CUDA graph. + q_buf = q.clone() + k_buf = k.clone() + out_q = torch.empty_like(q_buf) + out_k = torch.empty_like(k_buf) + + # Warmup required by CUDA graph capture. + multi.update_runtime_seq_lens( + torch.full((num_tokens,), 1000, dtype=torch.int32, device=device), + ) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(2): + multi.forward_native( + positions, q_buf.clone(), k_buf.clone(), + ) + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + q_in = q_buf.clone() + k_in = k_buf.clone() + captured_q, captured_k = multi.forward_native( + positions, q_in, k_in, + ) + out_q.copy_(captured_q) + out_k.copy_(captured_k) + + # Replay after routing to a different bucket. The graph must pick + # up the new offsets via the live buffer read. + multi.update_runtime_seq_lens( + torch.full((num_tokens,), 7000, dtype=torch.int32, device=device), + ) + graph.replay() + torch.cuda.synchronize() + + # Reference: eager forward under the new bucket. + ref_q, ref_k = multi.forward_native( + positions, q_buf.clone(), k_buf.clone(), + ) + + torch.testing.assert_close(out_q, ref_q, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(out_k, ref_k, rtol=1e-5, atol=1e-5) + + +# -------------------------------------------------------------------------- +# ``get_rope`` wiring +# -------------------------------------------------------------------------- + + +def test_get_rope_dispatches_multi_cache_type(): + from arctic_inference.vllm.rope import ( + MultiCacheDynamicNTKRotaryEmbedding, + apply_rope_runtime_patches, + ) + + apply_rope_runtime_patches() + from vllm.model_executor.layers.rotary_embedding import ( + get_rope as patched_get_rope, + ) + + mod = patched_get_rope( + head_size=64, + max_position=16384, + is_neox_style=True, + rope_parameters={ + "rope_type": "multi_cache_ntk", + "factors": [1.0, 2.0, 4.0], + "rope_theta": 10000.0, + "original_max_position_embeddings": 2048, + }, + dtype=torch.float32, + ) + assert isinstance(mod, MultiCacheDynamicNTKRotaryEmbedding) + assert mod.factors == [1.0, 2.0, 4.0] + assert mod._original_max_position_embeddings == 2048 + + +def test_get_rope_uses_default_factors_when_factors_not_specified(): + from arctic_inference.vllm.rope import ( + DEFAULT_FACTORS, + MultiCacheDynamicNTKRotaryEmbedding, + apply_rope_runtime_patches, + ) + + apply_rope_runtime_patches() + from vllm.model_executor.layers.rotary_embedding import ( + get_rope as patched_get_rope, + ) + + mod = patched_get_rope( + head_size=64, + max_position=16384, + is_neox_style=True, + rope_parameters={ + "rope_type": "multi_cache_ntk", + "rope_theta": 10000.0, + "original_max_position_embeddings": 2048, + }, + dtype=torch.float32, + ) + assert isinstance(mod, MultiCacheDynamicNTKRotaryEmbedding) + assert mod.factors == list(DEFAULT_FACTORS) + + +def test_get_rope_folds_legacy_factor_into_single_bucket(): + """Backward compat: ``rope_type=multi_cache_ntk`` with a scalar + ``factor`` (and no ``factors``) should produce a single-bucket + multi-cache at that factor. This keeps configs written for the old + static dynamic rope from silently changing behavior when the + ``rope_type`` is changed.""" + from arctic_inference.vllm.rope import ( + MultiCacheDynamicNTKRotaryEmbedding, + apply_rope_runtime_patches, + ) + + apply_rope_runtime_patches() + from vllm.model_executor.layers.rotary_embedding import ( + get_rope as patched_get_rope, + ) + + mod = patched_get_rope( + head_size=64, + max_position=16384, + is_neox_style=True, + rope_parameters={ + "rope_type": "multi_cache_ntk", + "factor": 4.0, + "rope_theta": 10000.0, + "original_max_position_embeddings": 2048, + }, + dtype=torch.float32, + ) + assert isinstance(mod, MultiCacheDynamicNTKRotaryEmbedding) + assert mod.factors == [4.0] + + +def test_get_rope_promotion_via_env(monkeypatch): + from arctic_inference.vllm.rope import ( + DEFAULT_FACTORS, + MultiCacheDynamicNTKRotaryEmbedding, + apply_rope_runtime_patches, + ) + + apply_rope_runtime_patches() + monkeypatch.setenv("ARCTIC_INFERENCE_MULTI_CACHE_ROPE", "1") + + from vllm.model_executor.layers.rotary_embedding import ( + get_rope as patched_get_rope, + ) + + mod = patched_get_rope( + head_size=64, + max_position=16384, + is_neox_style=True, + rope_parameters={ + "rope_type": "dynamic", + "factor": 4.0, + "rope_theta": 10000.0, + }, + dtype=torch.float32, + ) + # Env-promoted: gets a single-bucket multi-cache using the ``factor`` + # from the legacy config (see ``_build_multi_cache_ntk`` semantics). + assert isinstance(mod, MultiCacheDynamicNTKRotaryEmbedding) + assert mod.factors == [4.0] + # No ``factors`` key means the default list is NOT pulled in when + # a legacy ``factor`` is present; explicit opt-in required. + assert mod.factors != list(DEFAULT_FACTORS) + + +def test_get_rope_default_still_works_after_patching(): + """Non-multi-cache rope types must continue to work unchanged after + our wrapper is installed.""" + from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding + from arctic_inference.vllm.rope import apply_rope_runtime_patches + + apply_rope_runtime_patches() + from vllm.model_executor.layers.rotary_embedding import ( + get_rope as patched_get_rope, + ) + + mod = patched_get_rope( + head_size=64, + max_position=2048, + is_neox_style=True, + rope_parameters=None, + dtype=torch.float32, + ) + assert isinstance(mod, RotaryEmbedding) + + +def test_promotion_skips_alpha_variant(monkeypatch): + """The legacy alpha variant uses a different base formula, so it + must NOT be promoted to ``multi_cache_ntk`` even with the env flag + set.""" + from arctic_inference.vllm.rope import ( + MultiCacheDynamicNTKRotaryEmbedding, + apply_rope_runtime_patches, + ) + + apply_rope_runtime_patches() + monkeypatch.setenv("ARCTIC_INFERENCE_MULTI_CACHE_ROPE", "1") + from vllm.model_executor.layers.rotary_embedding import ( + get_rope as patched_get_rope, + ) + + mod = patched_get_rope( + head_size=64, + max_position=16384, + is_neox_style=True, + rope_parameters={ + "rope_type": "dynamic", + "alpha": 1.0, + "rope_theta": 10000.0, + }, + dtype=torch.float32, + ) + assert not isinstance(mod, MultiCacheDynamicNTKRotaryEmbedding) + + +def test_get_rope_threads_max_num_batched_tokens_from_vllm_config(): + """When a vLLM config is active the rope module should size its + runtime buffer to match ``scheduler_config.max_num_batched_tokens``. + """ + from arctic_inference.vllm.rope import ( + MultiCacheDynamicNTKRotaryEmbedding, + apply_rope_runtime_patches, + ) + from vllm.config import VllmConfig, set_current_vllm_config + + apply_rope_runtime_patches() + from vllm.model_executor.layers.rotary_embedding import ( + get_rope as patched_get_rope, + ) + + cfg = VllmConfig() + cfg.scheduler_config.max_num_batched_tokens = 7777 + with set_current_vllm_config(cfg): + mod = patched_get_rope( + head_size=64, + max_position=32768, + is_neox_style=True, + rope_parameters={ + "rope_type": "multi_cache_ntk", + "factors": [1.0, 2.0, 4.0], + "rope_theta": 10000.0, + "original_max_position_embeddings": 2048, + }, + dtype=torch.float32, + ) + assert isinstance(mod, MultiCacheDynamicNTKRotaryEmbedding) + assert mod.runtime_bucket_offsets.shape[0] == 7777 From ad27a89e0fb88aa17fa9f189d220d8152da03225 Mon Sep 17 00:00:00 2001 From: ye-wang_snow Date: Thu, 23 Apr 2026 06:58:52 +0000 Subject: [PATCH 2/2] multi-cache rope: uniform factor across request lifetime --- arctic_inference/vllm/model_runner.py | 80 +++++- tests/unit_tests/test_multi_cache_rope.py | 320 ++++++++++++++++++++++ 2 files changed, 398 insertions(+), 2 deletions(-) diff --git a/arctic_inference/vllm/model_runner.py b/arctic_inference/vllm/model_runner.py index 171697f15..4bdfe4d75 100644 --- a/arctic_inference/vllm/model_runner.py +++ b/arctic_inference/vllm/model_runner.py @@ -644,12 +644,88 @@ def _build_rope_seq_lens_per_token_gpu( (the unscaled cache). Padding tokens are masked out by attention so the exact offset is irrelevant, but we want the routing to land in a numerically safe bucket. + + Routing pinning: the effective seq-len is + ``max(projected_max, num_computed + num_scheduled)``, where + ``projected_max = num_prompt_tokens + sampling_params.max_tokens`` + (falls through to ``num_prompt_tokens`` for pooling requests + that have no sampling params). This keeps the rotary basis + uniform across both prefill AND the entire decode for a given + request: + + * **Chunked prefill** -- using only ``self.seq_lens`` + (``num_computed + num_scheduled`` -- the end-of-current-chunk + position) would leave early chunks of a long prompt routed + to a smaller factor while later chunks route to a larger + factor, producing a KV cache that holds keys rotated under + multiple rope bases for a single request. Pinning to the + full prompt length guarantees every chunk sees the same + basis. + * **Decode-boundary crossing** -- a request with prompt=30k + and ``max_tokens=15k`` reaches total length 45k by the end + of generation, which crosses F=1 into F=2. Without + projecting ``max_tokens`` into the routing, prefill keys + would be stored under F=1 and the late-decode queries / + keys would flip to F=2, recreating the same basis mismatch + at the decode boundary that chunked prefill creates at a + chunk boundary. Projecting ``max_tokens`` from the start + pins the request to the factor that covers its final length + for the request's entire lifetime. + + During decode, ``raw_seq_lens`` may still exceed + ``projected_max`` (e.g. the request generated slightly past + the projection due to spec decode or reused slots). The + ``max`` fall-through keeps routing monotone in that edge + case. """ - num_reqs = getattr(getattr(self, "input_batch", None), "num_reqs", 0) + input_batch = getattr(self, "input_batch", None) + num_reqs = getattr(input_batch, "num_reqs", 0) if not num_reqs: return None - seq_lens_cpu = self.seq_lens.np[:num_reqs] + raw_seq_lens = self.seq_lens.np[:num_reqs] + num_prompt_tokens = getattr(input_batch, "num_prompt_tokens", None) + if num_prompt_tokens is not None: + # Per-request projection = prompt + max_tokens. For pooling + # requests (no sampling_params) or missing attributes, the + # projection collapses to prompt-only, preserving the + # prefill-only chunked-prefill guarantee. + projected_max = num_prompt_tokens[:num_reqs].astype( + np.int64, copy=True, + ) + req_ids = getattr(input_batch, "req_ids", None) + requests = getattr(self, "requests", None) + if req_ids is not None and requests is not None: + for i in range(num_reqs): + req_id = req_ids[i] + if req_id is None: + continue + req_state = requests.get(req_id) + if req_state is None: + continue + sampling_params = getattr( + req_state, "sampling_params", None, + ) + if sampling_params is None: + continue + max_tokens = getattr(sampling_params, "max_tokens", None) + if max_tokens is None: + continue + projected_max[i] = ( + int(num_prompt_tokens[i]) + int(max_tokens) + ) + seq_lens_cpu = np.maximum( + projected_max, + raw_seq_lens.astype(np.int64, copy=False), + ) + else: + # Defensive fallback: if input_batch doesn't expose the + # prompt-length array (older vLLM shapes), fall back to the + # raw seq_lens. This preserves the old behaviour and the + # unit tests for the multi-cache math still pass; the + # chunked-prefill guarantee just weakens to "best-effort". + seq_lens_cpu = raw_seq_lens + # query_start_loc stores cumulative token counts per request in # slots [0..num_reqs]. diff() gives per-request scheduled tokens. qsl = self.query_start_loc.np[: num_reqs + 1] diff --git a/tests/unit_tests/test_multi_cache_rope.py b/tests/unit_tests/test_multi_cache_rope.py index 69c99e6ac..5bf60c88d 100644 --- a/tests/unit_tests/test_multi_cache_rope.py +++ b/tests/unit_tests/test_multi_cache_rope.py @@ -344,6 +344,326 @@ def test_update_runtime_seq_lens_handles_empty_batch(device): torch.testing.assert_close(multi.runtime_bucket_offsets, before) +# -------------------------------------------------------------------------- +# Model-runner hook: chunked-prefill seq-len derivation +# -------------------------------------------------------------------------- + + +def _call_build_rope_seq_lens_per_token_gpu( + *, + device: torch.device, + num_prompt_tokens: list[int] | None, + seq_lens: list[int], + num_scheduled_per_req: list[int], + num_tokens_padded: int, + max_tokens_per_req: list[int | None] | None = None, + expose_req_ids: bool = True, + expose_requests: bool = True, +) -> torch.Tensor: + """Invoke the patched helper on a minimal fake-``self`` object. + + The method reads ``input_batch.num_reqs``, + ``input_batch.num_prompt_tokens``, ``input_batch.req_ids``, + ``self.requests``, ``self.seq_lens.np``, + ``self.query_start_loc.np``, and ``self.device``. We build + ``SimpleNamespace`` objects with exactly those attributes and call + the method unbound. This keeps the test narrow and avoids booting + a real vLLM engine. + + Args: + max_tokens_per_req: if provided, each non-``None`` entry stocks + a fake ``CachedRequestState.sampling_params.max_tokens`` + under ``self.requests[req_id]`` so the projection code + path is exercised. ``None`` entries simulate pooling + requests (no ``sampling_params``). + expose_req_ids / expose_requests: when ``False``, omit the + corresponding attribute from the fake namespaces to + exercise defensive fallback branches. + """ + from types import SimpleNamespace + + import numpy as np + + from arctic_inference.vllm.model_runner import GPUModelRunnerPatch + + num_reqs = len(seq_lens) + assert num_reqs == len(num_scheduled_per_req) + if max_tokens_per_req is not None: + assert len(max_tokens_per_req) == num_reqs + + qsl = np.zeros(num_reqs + 1, dtype=np.int32) + qsl[1:] = np.cumsum(num_scheduled_per_req) + + input_batch = SimpleNamespace(num_reqs=num_reqs) + if num_prompt_tokens is not None: + assert len(num_prompt_tokens) == num_reqs + input_batch.num_prompt_tokens = np.asarray( + num_prompt_tokens, dtype=np.int32, + ) + req_ids = [f"req_{i}" for i in range(num_reqs)] + if expose_req_ids: + input_batch.req_ids = req_ids + + requests: dict | None = {} + if max_tokens_per_req is not None and expose_requests: + for i, rid in enumerate(req_ids): + max_toks = max_tokens_per_req[i] + if max_toks is None: + requests[rid] = SimpleNamespace(sampling_params=None) + else: + requests[rid] = SimpleNamespace( + sampling_params=SimpleNamespace(max_tokens=int(max_toks)) + ) + if not expose_requests: + requests = None + + fake_self = SimpleNamespace( + input_batch=input_batch, + seq_lens=SimpleNamespace(np=np.asarray(seq_lens, dtype=np.int32)), + query_start_loc=SimpleNamespace(np=qsl), + device=device, + ) + if requests is not None: + fake_self.requests = requests + return GPUModelRunnerPatch._build_rope_seq_lens_per_token_gpu( + fake_self, num_tokens_padded, + ) + + +def test_chunked_prefill_pins_seq_len_to_full_prompt_length(device): + """Under chunked prefill, all chunks of a long request must see the + request's full prompt length, not the end-of-chunk position. + + This is the regression guard for the KV-cache basis pollution bug: + ``vllm.v1.worker.gpu_model_runner`` sets ``seq_lens[i] = + num_computed_tokens[i] + num_scheduled_tokens[i]`` -- the end of + the current chunk, not the total prompt length. If we route by + that value, a 130k-token prompt processed over multiple 8k chunks + would see seq_len=8k on the first chunk (→ F=1) and only eventually + cross into F=4 on a later chunk, leaving the KV cache holding keys + rotated under inconsistent rope bases for a single request. + + Concretely we simulate two requests sharing one forward: + * req0: short prompt (4000 tokens), all in this chunk. + * req1: long prompt (130000 tokens), only 4192 tokens scheduled + this chunk (the rest of max_num_batched_tokens=8192). + ``seq_lens`` reported by the scheduler is ``[4000, 4192]`` -- the + end-of-chunk positions. Routing by that would pin req1 to F=1. + After the fix, req1's per-token seq-lens must be 130000 across + all 4192 of its tokens in this chunk, so it routes to F=4. + """ + num_prompt_tokens = [4000, 130_000] + num_scheduled_per_req = [4000, 4192] + seq_lens = [4000, 4192] # end-of-chunk positions from vLLM scheduler + num_tokens_unpadded = sum(num_scheduled_per_req) + num_tokens_padded = num_tokens_unpadded + + out = _call_build_rope_seq_lens_per_token_gpu( + device=device, + num_prompt_tokens=num_prompt_tokens, + seq_lens=seq_lens, + num_scheduled_per_req=num_scheduled_per_req, + num_tokens_padded=num_tokens_padded, + ) + + assert out is not None + assert out.shape == (num_tokens_padded,) + expected = torch.tensor( + [4000] * 4000 + [130_000] * 4192, + dtype=torch.int32, device=device, + ) + torch.testing.assert_close(out, expected) + + +def test_decode_seq_len_exceeds_prompt_and_keeps_growing(device): + """During decode, ``seq_lens`` > ``num_prompt_tokens`` and the + effective seq-len we route on should track decode growth. The + ``max(num_prompt_tokens, seq_lens)`` form must fall through to + ``seq_lens`` in that regime so factor upgrades during long decode + still happen.""" + num_prompt_tokens = [40_000] + seq_lens = [41_005] # prompt done + 1005 decode tokens generated + out = _call_build_rope_seq_lens_per_token_gpu( + device=device, + num_prompt_tokens=num_prompt_tokens, + seq_lens=seq_lens, + num_scheduled_per_req=[1], + num_tokens_padded=1, + ) + assert out is not None + # Must be the larger of the two (41005), not pinned to 40000. + assert int(out.item()) == 41_005 + + +def test_fallback_path_when_num_prompt_tokens_missing(device): + """If the input_batch doesn't expose ``num_prompt_tokens`` (older + vLLM shapes), the helper must still work using ``seq_lens`` only. + This preserves the pre-fix behaviour for environments where the + new field isn't available.""" + out = _call_build_rope_seq_lens_per_token_gpu( + device=device, + num_prompt_tokens=None, + seq_lens=[8192], + num_scheduled_per_req=[8192], + num_tokens_padded=8192, + ) + assert out is not None + expected = torch.full( + (8192,), 8192, dtype=torch.int32, device=device, + ) + torch.testing.assert_close(out, expected) + + +# -------------------------------------------------------------------------- +# Routing must project max_tokens so the rope basis is uniform across +# prefill + decode, not just within prefill. +# -------------------------------------------------------------------------- + + +def test_projected_seq_len_pins_factor_to_prompt_plus_max_tokens(device): + """A 30k-prompt request with ``max_tokens=15k`` can reach total + length 45k during decode, crossing F=1 (max_len=40960) into F=2. + + Without max_tokens projection, the first decode step (seq_len = + 30001) would route to F=1, so prefill + early decode keys are + stored under F=1. Late-decode queries/keys at seq_len=45000 flip + to F=2 and attend back to F=1 keys → basis mismatch and attention + corruption, exactly analogous to the chunked-prefill bug. + + With projection, even at prefill time we already route to F=2 + (projected_max = 30k + 15k = 45k), so the entire request's + keys/queries share a single F=2 basis. + """ + # Decode step: prefill finished, 1 decode token scheduled. + out = _call_build_rope_seq_lens_per_token_gpu( + device=device, + num_prompt_tokens=[30_000], + seq_lens=[30_001], + num_scheduled_per_req=[1], + num_tokens_padded=1, + max_tokens_per_req=[15_000], + ) + assert out is not None + # Projected max dominates over seq_lens=30001 and over + # num_prompt_tokens=30000, routing the decode step to the factor + # that covers the request's final length (45000 → F=2). + assert int(out.item()) == 45_000 + + +def test_projected_seq_len_pins_prefill_to_same_factor_as_decode(device): + """Prefill of the same 30k-prompt + 15k-max_tokens request must + also route to F=2, not F=1. This is the uniformity guarantee: + every token of the request's lifetime (prefill chunks + decode + steps) picks the same factor.""" + # Prefill chunk: first 8192 tokens of the 30k prompt. + out = _call_build_rope_seq_lens_per_token_gpu( + device=device, + num_prompt_tokens=[30_000], + seq_lens=[8192], # end-of-chunk position + num_scheduled_per_req=[8192], + num_tokens_padded=8192, + max_tokens_per_req=[15_000], + ) + assert out is not None + # All tokens of the prefill chunk must see the projected 45k, + # not the end-of-chunk 8192 and not just the prompt length 30000. + assert torch.all(out == 45_000).item() + + +def test_projected_seq_len_leaves_short_requests_alone(device): + """A short request (prompt=5k, max_tokens=2k → projected=7k) must + stay in F=1. Projection should not artificially escalate factors + for requests that never cross a boundary.""" + out = _call_build_rope_seq_lens_per_token_gpu( + device=device, + num_prompt_tokens=[5_000], + seq_lens=[5_000], + num_scheduled_per_req=[5_000], + num_tokens_padded=5_000, + max_tokens_per_req=[2_000], + ) + assert out is not None + assert torch.all(out == 7_000).item() + + +def test_projected_seq_len_handles_pooling_request_without_sampling_params( + device, +): + """Pooling requests have ``sampling_params=None`` (they don't + generate). Projection must fall through to prompt-only for them, + leaving the pre-projection chunked-prefill behaviour intact.""" + out = _call_build_rope_seq_lens_per_token_gpu( + device=device, + num_prompt_tokens=[130_000], + seq_lens=[8192], # chunked-prefill chunk position + num_scheduled_per_req=[8192], + num_tokens_padded=8192, + max_tokens_per_req=[None], # no sampling_params + ) + assert out is not None + # No max_tokens → projection = prompt = 130k. All tokens should + # see 130k (not 8192, the chunk position). + assert torch.all(out == 130_000).item() + + +def test_projected_seq_len_falls_back_when_requests_dict_absent(device): + """If ``self.requests`` isn't attached to the runner (defensive + path), the helper must still work and emit prompt-length routing + rather than crashing on attribute access.""" + out = _call_build_rope_seq_lens_per_token_gpu( + device=device, + num_prompt_tokens=[30_000], + seq_lens=[8192], + num_scheduled_per_req=[8192], + num_tokens_padded=8192, + max_tokens_per_req=[15_000], + expose_requests=False, + ) + assert out is not None + # Without the requests dict we can't read max_tokens, so we fall + # back to prompt-length routing (30k). + assert torch.all(out == 30_000).item() + + +def test_projected_seq_len_mixed_batch_routes_each_request_independently( + device, +): + """Two requests in the same forward with different projected + lengths must each get their own per-token routing signal, not a + batch-wide value.""" + # req_0: short prompt + small max_tokens → stays in F=1 regime (7k). + # req_1: medium prompt + large max_tokens → crosses into F=2 (95k). + out = _call_build_rope_seq_lens_per_token_gpu( + device=device, + num_prompt_tokens=[5_000, 80_000], + seq_lens=[5_000, 8192], # req_1 is in chunked prefill + num_scheduled_per_req=[5_000, 8192], + num_tokens_padded=5_000 + 8192, + max_tokens_per_req=[2_000, 15_000], + ) + assert out is not None + assert torch.all(out[:5_000] == 7_000).item() # 5k+2k projection + assert torch.all(out[5_000:] == 95_000).item() # 80k+15k projection + + +def test_projected_seq_len_respects_late_decode_past_projection(device): + """Edge: if ``seq_lens`` has somehow grown past + ``prompt + max_tokens`` (e.g. reused batch slot, spec decode + overshoot), the max fall-through must still pick the larger + value so routing is monotone non-decreasing through the request's + lifetime.""" + out = _call_build_rope_seq_lens_per_token_gpu( + device=device, + num_prompt_tokens=[30_000], + seq_lens=[50_000], # past projection + num_scheduled_per_req=[1], + num_tokens_padded=1, + max_tokens_per_req=[15_000], # projection = 45_000 + ) + assert out is not None + assert int(out.item()) == 50_000 + + # -------------------------------------------------------------------------- # Forward correctness vs static DynamicNTK reference # --------------------------------------------------------------------------