diff --git a/docs/user_guide/configuration.md b/docs/user_guide/configuration.md index 3f9e1e862..3afe7ba4a 100644 --- a/docs/user_guide/configuration.md +++ b/docs/user_guide/configuration.md @@ -28,6 +28,44 @@ llm = LLM( See the [Examples](../examples/offline_inference/torch_spyre_inference.md) page for more usage patterns. +## Host sampler (async noise + log-space Gumbel) + +Spyre replaces vLLM's default host sampler with a Spyre-optimized path +ported from [sendnn-inference#1046](https://github.com/torch-spyre/sendnn-inference/pull/1046) +(Holtz), in three stages: + +1. **Async noise ring buffer** — Exp(1) log-noise is filled on a background + thread; the decode loop borrows zero-copy rows instead of calling + `exponential_()` on the critical path. +2. **TP rank-0 sampling** — when tensor-parallel and logits are on CPU, only + rank 0 samples and broadcasts token ids (and logprobs) to the other ranks. +3. **Log-space Gumbel** — sample as `argmax(logits - log_noise)` instead of + `argmax(probs / noise)`, which removes softmax from the hot path while + preserving token order / distribution. + +If `vllm_config` lacks `max_num_seqs` or vocab size, the runner falls back to +vLLM's default `Sampler` (same as sendnn). Sampling still runs on the **CPU**. + +### Configuration (`spyre_inference.envs`) + +All host-sampler knobs live in `spyre_inference/envs.py` (vLLM-style lazy +module with defaults, docs, optional `enable_envs_cache()`, and `is_set()` for +override detection). Prefer `import spyre_inference.envs as envs` over +scattered `os.environ.get` calls. + +| Variable | Default | Meaning | +|---|---|---| +| `SPYRE_USE_SPYRE_SAMPLER` | `1` | Set to `0` to force upstream `Sampler` | +| `SPYRE_ASYNC_NOISE_SCALE` | `4` | Ring depth = scale × `max_num_seqs` (must be ≥ 2) | + +Do **not** inject sampler timing into the production path. Measure with the +[Kineto / Spyre profiler](kineto_profiling.md) instead. + +```bash +SPYRE_ASYNC_NOISE_SCALE=8 \ + python examples/offline_inference/torch_spyre_inference.py +``` + ## pyproject.toml Reference The `pyproject.toml` includes several key build configurations: diff --git a/spyre_inference/envs.py b/spyre_inference/envs.py new file mode 100644 index 000000000..66dd34ff6 --- /dev/null +++ b/spyre_inference/envs.py @@ -0,0 +1,106 @@ +# Copyright 2026 The Spyre-Inference Authors. +# +# 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. + +"""Spyre-inference environment variables. + +Central place for config levers: documentation, defaults, lazy evaluation, +and optional caching after service init. Prefer:: + + import spyre_inference.envs as envs + + scale = envs.SPYRE_ASYNC_NOISE_SCALE + +over scattered ``os.environ.get`` calls. + +Do **not** add production-path timing toggles here (e.g. a hypothetical +``SPYRE_SAMPLER_TIMING``). Use the Spyre / Kineto profiler instead — see +``docs/user_guide/kineto_profiling.md``. +""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + # Help type checkers resolve lazy attributes from environment_variables. + SPYRE_USE_SPYRE_SAMPLER: bool = True + SPYRE_ASYNC_NOISE_SCALE: int = 4 + +# Populated by ``enable_envs_cache()``; ``None`` means uncached / lazy. +_env_cache: dict[str, Any] | None = None + + +def _async_noise_scale() -> int: + raw = os.getenv("SPYRE_ASYNC_NOISE_SCALE", "4") + scale = int(raw) + if scale < 2: + raise ValueError( + f"SPYRE_ASYNC_NOISE_SCALE must be >= 2 (got {scale}); " + "the async ring buffer needs at least one full batch ahead of the consumer." + ) + return scale + + +environment_variables: dict[str, Callable[[], Any]] = { + # Host sampler from sendnn-inference#1046 (async log-noise ring buffer, + # TP rank-0 sample + broadcast, log-space Gumbel). On by default when + # vllm_config is compatible; set to 0 to force upstream Sampler. + "SPYRE_USE_SPYRE_SAMPLER": lambda: os.getenv("SPYRE_USE_SPYRE_SAMPLER", "1") == "1", + # Depth of the host-side async Exp(1) log-noise ring buffer: + # rows = scale * max_num_seqs. Must be >= 2. + "SPYRE_ASYNC_NOISE_SCALE": _async_noise_scale, +} + + +def __getattr__(name: str) -> Any: + """Lazy attribute access into ``environment_variables``. + + After ``enable_envs_cache()``, values are served from ``_env_cache`` (do + not change env after service init if cache is enabled). + """ + if name not in environment_variables: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + if _env_cache is not None: + return _env_cache[name] + return environment_variables[name]() + + +def enable_envs_cache() -> None: + """Cache env lookups after service initialization.""" + global _env_cache + if _env_cache is not None: + return + _env_cache = {key: getter() for key, getter in environment_variables.items()} + + +def disable_envs_cache() -> None: + """Clear the env cache (for tests that mutate ``os.environ``).""" + global _env_cache + _env_cache = None + + +def is_set(name: str) -> bool: + """Return True if ``name`` is present in the process environment. + + Distinguishes an explicit override from the documented default. + """ + if name in environment_variables: + return name in os.environ + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return list(environment_variables.keys()) diff --git a/spyre_inference/v1/sample/__init__.py b/spyre_inference/v1/sample/__init__.py new file mode 100644 index 000000000..0d609bbbb --- /dev/null +++ b/spyre_inference/v1/sample/__init__.py @@ -0,0 +1,68 @@ +# Copyright 2026 The Spyre-Inference Authors. +# +# 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. + +"""Host sampler path for Spyre (async noise ring buffer + log-space Gumbel). + +Port of sendnn-inference#1046 (Holtz): async Exp(1) log-noise, TP rank-0 +sample + broadcast, and log-space Gumbel. + +Config levers live in ``spyre_inference.envs`` (``SPYRE_USE_SPYRE_SAMPLER``, +``SPYRE_ASYNC_NOISE_SCALE``). +""" + +from __future__ import annotations + +import warnings + +from vllm.config import VllmConfig +from vllm.v1.sample.sampler import Sampler + +import spyre_inference.envs as envs +from spyre_inference.v1.sample.async_ring_buffer import ( + AsyncExponential_RingBuffer, + AsyncRingBuffer, +) +from spyre_inference.v1.sample.spyre_sampler import SpyreSampler +from spyre_inference.v1.sample.spyre_topk_topp_sampler import SpyreTopKTopPSampler + + +def build_spyre_sampler(vllm_config: VllmConfig) -> Sampler: + """Build Holtz SpyreSampler, or fall back to upstream Sampler. + + Falls back when ``SPYRE_USE_SPYRE_SAMPLER=0`` or when ``vllm_config`` lacks + ``max_num_seqs`` / vocab size (same as sendnn ``SpyreCausalLM``). + """ + logprobs_mode = vllm_config.model_config.logprobs_mode + if not envs.SPYRE_USE_SPYRE_SAMPLER: + return Sampler(logprobs_mode=logprobs_mode) + if not SpyreSampler.is_vllm_config_compatible(vllm_config): + warnings.warn( + "The provided vllm_config is not compatible with SpyreSampler. " + "Falling back to default Sampler with reduced performance on Spyre platform.", + stacklevel=2, + ) + return Sampler(logprobs_mode=logprobs_mode) + return SpyreSampler( + vllm_config=vllm_config, + logprobs_mode=logprobs_mode, + ) + + +__all__ = [ + "AsyncExponential_RingBuffer", + "AsyncRingBuffer", + "SpyreSampler", + "SpyreTopKTopPSampler", + "build_spyre_sampler", +] diff --git a/spyre_inference/v1/sample/async_ring_buffer.py b/spyre_inference/v1/sample/async_ring_buffer.py new file mode 100644 index 000000000..a754eba4b --- /dev/null +++ b/spyre_inference/v1/sample/async_ring_buffer.py @@ -0,0 +1,211 @@ +# Copyright 2026 The Spyre-Inference Authors. +# +# 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. + +import contextlib +import queue +import threading +from abc import ABC, abstractmethod +from collections.abc import Generator + +import numpy as np +import torch + + +class AsyncRingBuffer(ABC): + """Pre-generates data rows on a background thread via a ring buffer. + + Maintains a contiguous ``(S, V)`` buffer (``S = scale * max_batch_size``) + and two shared counters: + + * ``_read_pos`` — next row index the consumer will read from. + * ``_tail`` — upper bound (in unwrapped space) up to which the consumer + may read without stalling. + + On init the buffer is fully filled and ``_tail = S``. The consumer + advances ``_read_pos`` after each call; when it approaches the end of the + buffer it wraps back to 0. Each consumed segment is enqueued for the + background thread to refill, which increments ``_tail`` once done. + + Storage is a NumPy array shared with a Torch CPU view. The producer + thread must not call Torch ops: under the Spyre plugin, background-thread + Torch tensor mutations can abort the producer, which then deadlocks + ``borrow_rows`` after the first wrap. + + Args: + vocab_size: Number of columns ``V``. + max_batch_size: Maximum rows per :meth:`get_rows` call ``B``. + scale: Buffer depth multiplier; ``S = scale * B``. Must be >= 2 so + there is always at least one full batch of pre-filled rows ahead + of the consumer. + """ + + def __init__( + self, + vocab_size: int, + max_batch_size: int, + scale: int = 4, + ) -> None: + assert scale >= 2, "scale must be >= 2" + self._V = vocab_size + self._B = max_batch_size + self._S = scale * max_batch_size + + # NumPy backing store; Torch view shares the same memory for zero-copy + # borrows. Producer refills via NumPy only (see class docstring). + self._np = np.empty((self._S, self._V), dtype=np.float32) + self._buf = torch.from_numpy(self._np) + + self._error: BaseException | None = None + + # first-time buffer initialization + self._refill_slice(0, self._S) + self._tail: int = self._S + self._read_pos: int = 0 + + # _tail and _read_pos are guarded by _cond. + self._cond = threading.Condition(threading.Lock()) + + # Refill requests: (start, end, wrap) + self._refill_q: queue.Queue[tuple[int, int, bool] | None] = queue.Queue() + + self._thread = threading.Thread( + target=self._produce, + name="async-ring-buffer", + daemon=True, + ) + self._thread.start() + + @abstractmethod + def _refill_slice(self, start: int, end: int) -> None: + """Fill ``self._np[start:end]`` with fresh values in-place (NumPy only).""" + ... + + @property + def vocab_size(self) -> int: + return self._V + + def _raise_if_producer_failed(self) -> None: + if self._error is not None: + raise RuntimeError("async ring buffer producer failed") from self._error + if not self._thread.is_alive() and self._error is None: + # Thread exited without recording an error (e.g. unexpected break). + raise RuntimeError("async ring buffer producer thread is not alive") + + @contextlib.contextmanager + def borrow_rows(self, n: int) -> Generator[torch.Tensor, None, None]: + """Context manager that yields a zero-copy ``(n, V)`` view. + + The backing rows are released for refill automatically when the + ``with`` block exits, even if an exception is raised. The view + must not be used after the block. + + Args: + n: Number of rows to borrow. Must satisfy ``1 <= n <= B``. + + Raises: + ValueError: If ``n`` is outside the valid range. + + Example:: + + with buf.borrow_rows(batch_size) as noise: + tokens = probs.div(noise).argmax(dim=-1) + """ + if n > self._B or n < 1: + raise ValueError(f"n (got {n}) must satisfy 1 <= n <= {self._B} (max_batch_size)") + + start = self._read_pos + end = start + n + + # wait for the producer to fill up at least n many values ahead + with self._cond: + while self._tail < end: + self._raise_if_producer_failed() + self._cond.wait(timeout=1.0) + self._raise_if_producer_failed() + + # get view (zero-copy into the shared Torch/NumPy buffer) + view = self._buf[start:end] + + wrap: bool = end > self._S - self._B + if wrap: + with self._cond: + self._tail -= self._S + + self._read_pos = 0 + else: + self._read_pos = end + + try: + # yield view to outside consumer + yield view + finally: + # issue refill request once view has been consumed and returned + self._refill_q.put((start, end, wrap)) + + def _produce(self) -> None: + try: + while True: + req = self._refill_q.get() + + # handle termination signal + if req is None: + break + + # refill buffer (NumPy only — see class docstring) + start, end, wrap = req + self._refill_slice(start, end) + + increment = (self._S - start) if wrap else (end - start) + with self._cond: + self._tail += increment + self._cond.notify_all() + except BaseException as exc: + self._error = exc + with self._cond: + self._cond.notify_all() + raise + + def shutdown(self) -> None: + """Signal the background thread to stop and wait for it to exit.""" + self._refill_q.put(None) + self._thread.join() + + +class AsyncExponential_RingBuffer(AsyncRingBuffer): + """Ring buffer that pre-generates exponential log noise via Exp(1) then log.""" + + def _refill_slice(self, start: int, end: int) -> None: + # Match torch.Tensor.exponential_() default (rate=1) then log_(). + n = end - start + out = self._np[start:end] + out[:] = np.random.exponential(scale=1.0, size=(n, self._V)) + np.log(out, out=out) + + +class _AsyncCounterRingBuffer(AsyncRingBuffer): + """Ring buffer that fills each row with the cumulative row index. + + Used in tests to verify that consumers receive the correct rows in order + without repeating any. + """ + + def __init__(self, vocab_size: int, max_batch_size: int, scale: int = 4) -> None: + self._total_generated: int = 0 + super().__init__(vocab_size, max_batch_size, scale) + + def _refill_slice(self, start: int, end: int) -> None: + n = end - start + for i in range(n): + self._np[start + i, :] = self._total_generated + self._total_generated += 1 diff --git a/spyre_inference/v1/sample/spyre_sampler.py b/spyre_inference/v1/sample/spyre_sampler.py new file mode 100644 index 000000000..c4071f467 --- /dev/null +++ b/spyre_inference/v1/sample/spyre_sampler.py @@ -0,0 +1,188 @@ +# Copyright 2026 The Spyre-Inference Authors. +# +# 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. + +import torch +from vllm.config import VllmConfig +from vllm.config.model import LogprobsMode +from vllm.distributed import get_tp_group +from vllm.distributed.parallel_state import _TP +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler + +import spyre_inference.envs as envs +from spyre_inference.v1.sample.spyre_topk_topp_sampler import SpyreTopKTopPSampler + + +class SpyreSampler(Sampler): + """A vLLM Sampler subclass that uses top-k/top-p sampling implementations optimized for Spyre + platform. + """ + + def __init__( + self, + vllm_config: VllmConfig, + logprobs_mode: LogprobsMode = "raw_logprobs", + use_fp64_gumbel: bool = False, + noise_scale: int | None = None, + ): + """Initialize the SpyreSampler with Spyre-optimized sampling components. + + Initializes the parent Sampler and replaces the default top-k/top-p sampler + with a Spyre-specific implementation. Configuration parameters are extracted + from vllm_config to ensure the sampler is properly tuned for the target + hardware and model. + + Args: + vllm_config: The VLLMConfig instance containing model configuration, + vocabulary size, and concurrency settings needed to initialize + the Spyre sampler. + logprobs_mode: See vllm.v1.sample.sampler.Sampler for details. + use_fp64_gumbel: See vllm.v1.sample.sampler.Sampler for details. + This parameter is not supported by SpyreSampler. + Defaults to False. + noise_scale: Async ring-buffer depth multiplier. ``None`` reads + ``envs.SPYRE_ASYNC_NOISE_SCALE`` (default 4). + + Raises: + ValueError: If use_fp64_gumbel is True, as SpyreSampler does + not support 64-bit Gumbel noise computation. + ValueError: If vllm_config does not provide max_num_seqs or vocab_size, + which are required for SpyreSampler initialization. + """ + if use_fp64_gumbel: + raise ValueError("SpyreSampler does not support use_fp64_gumbel=True") + + # Current spyre-inference vLLM Sampler only takes logprobs_mode + # (sendnn passes use_fp64_gumbel=False as well). + super().__init__(logprobs_mode=logprobs_mode) + + # read concurrency and vocab size from vllm_config + max_concurrency = SpyreSampler._try_get_concurrency(vllm_config) + if max_concurrency is None: + raise ValueError("SpyreSampler requires vllm_config to specify max_num_seqs") + vocab_size = SpyreSampler._try_get_vocab_size(vllm_config) + if vocab_size is None: + raise ValueError("SpyreSampler requires vllm_config to specify vocab_size") + + if noise_scale is None: + noise_scale = envs.SPYRE_ASYNC_NOISE_SCALE + + # override topk_topp_sampler with spyre-specific topk-topp-sampler + self.topk_topp_sampler: SpyreTopKTopPSampler = SpyreTopKTopPSampler( + max_batch_size=max_concurrency, + vocab_size=vocab_size, + logprobs_mode=logprobs_mode, + noise_scale=noise_scale, + ) + + @staticmethod + def is_vllm_config_compatible(vllm_config: VllmConfig) -> bool: + """Check if the provided VllmConfig provides all necessary parameters for SpyreSampler + initialization. + """ + has_concurrency = SpyreSampler._try_get_concurrency(vllm_config) is not None + has_vocab_size = SpyreSampler._try_get_vocab_size(vllm_config) is not None + return has_concurrency and has_vocab_size + + @staticmethod + def _try_get_concurrency(vllm_config: VllmConfig) -> int | None: + """Try to extract the max_num_seqs parameter from the VllmConfig. + + Returns: + The max_num_seqs value if present, otherwise None. + """ + return getattr(vllm_config.scheduler_config, "max_num_seqs", None) + + @staticmethod + def _try_get_vocab_size(vllm_config: VllmConfig) -> int | None: + """Try to extract the vocab_size parameter from the VllmConfig. + + Returns: + The vocab_size value if present, otherwise None. + """ + if hasattr(vllm_config, "model_config") and hasattr(vllm_config.model_config, "hf_config"): + hf_cfg = vllm_config.model_config.hf_config + if hasattr(hf_cfg, "vocab_size"): + # convention: HuggingFace model configs have a vocab_size attribute + return hf_cfg.vocab_size + elif hasattr(hf_cfg, "text_config") and hasattr(hf_cfg.text_config, "vocab_size"): + # fallback: some multi-modal HuggingFace model configs have a text_config + # with a vocab_size attribute + return hf_cfg.text_config.vocab_size + return None + + def forward( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + predict_bonus_token: bool = False, + logprobs_mode_override: LogprobsMode | None = None, + ) -> SamplerOutput: + if logits.device.type == "cpu" and _TP is not None: + # if the sampler runs on CPU and is distributed across tensor parallel ranks, + # use an optimized path on CPU that avoids redundant computations across ranks + return self.forward_cpu_tp( + logits, sampling_metadata, predict_bonus_token, logprobs_mode_override + ) + else: + # if the sampler does not run on CPU, fall back to the base class implementation + return super().forward( + logits, sampling_metadata, predict_bonus_token, logprobs_mode_override + ) + + def forward_cpu_tp( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + predict_bonus_token: bool = False, + logprobs_mode_override: LogprobsMode | None = None, + ) -> SamplerOutput: + """Overrides the upstream sampler to run only on the first TP-rank and broadcast + results to other TP ranks. + + This is a correctness fix, because independent sampling across ranks would diverge + the computation across ranks over time. Further, this fix improves performance if + the sampler runs on CPU by avoiding redundant computations. + """ + + tp_group = get_tp_group() + if tp_group.is_first_rank: + sampler_output = super().forward( + logits, sampling_metadata, predict_bonus_token, logprobs_mode_override + ) + else: + # Allocate placeholder; will be filled by the broadcast below. + num_reqs = logits.shape[0] + sampler_output = SamplerOutput( + sampled_token_ids=torch.empty( + (num_reqs, 1), dtype=torch.int32, device=logits.device + ), + logprobs_tensors=None, + ) + + # Broadcast sampled token ids from TP rank 0 to all other TP ranks so + # that every rank feeds identical tokens into the next forward pass. + tp_group.broadcast(sampler_output.sampled_token_ids, src=0) + + # Broadcast the logprobs_tensors (broadcast_object handles None) and + # update sampler outputs + logprobs_tensors = tp_group.broadcast_object(sampler_output.logprobs_tensors, src=0) + sampler_output.logprobs_tensors = logprobs_tensors + + return sampler_output + + def shutdown(self) -> None: + """Shutdown the sampler and clean up resources.""" + self.topk_topp_sampler.shutdown() diff --git a/spyre_inference/v1/sample/spyre_topk_topp_sampler.py b/spyre_inference/v1/sample/spyre_topk_topp_sampler.py new file mode 100644 index 000000000..cf7ee530a --- /dev/null +++ b/spyre_inference/v1/sample/spyre_topk_topp_sampler.py @@ -0,0 +1,106 @@ +# Copyright 2026 The Spyre-Inference Authors. +# +# 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. + +import warnings + +import torch +from vllm.config.model import LogprobsMode +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p, TopKTopPSampler + +from spyre_inference.v1.sample.async_ring_buffer import AsyncExponential_RingBuffer + + +class SpyreTopKTopPSampler(TopKTopPSampler): + """Top-k/top-p sampler optimized for Spyre hardware via asynchronous noise pre-sampling. + + This removes CPU-bound noise generation from the latency-critical sampling path by + pre-drawing noise into a ring buffer that the decoder can consume via zero-copy views + during token selection. The buffer is pre-allocated based on vocab_size and multiples + of max_batch_size to support zero-copy access patterns. + """ + + def __init__( + self, + vocab_size: int, + max_batch_size: int, + logprobs_mode: LogprobsMode = "raw_logprobs", + noise_scale: int = 4, + ): + """Initialize the SpyreTopKTopPSampler with a asynchronous exponential + noise ring buffer. + + Args: + vocab_size: The size of the vocabulary (number of possible tokens). + Used to allocate noise buffer rows of appropriate size. + max_batch_size: The maximum batch size that will be processed. + Determines the total capacity of the pre-allocated noise buffer. + logprobs_mode: See vllm.v1.sample.ops.topk_topp_sampler for details. + noise_scale: Ring-buffer depth multiplier (``rows = scale * max_batch_size``). + Must be >= 2. Typically from ``envs.SPYRE_ASYNC_NOISE_SCALE``. + """ + super().__init__(logprobs_mode=logprobs_mode) + + self._noise_buffer = AsyncExponential_RingBuffer( + vocab_size=vocab_size, + max_batch_size=max_batch_size, + scale=noise_scale, + ) + # Always use the native path (async log-noise Gumbel); do not dispatch + # to CUDA/flashinfer variants from the base class. (spyre-inference + # adaptation; sendnn relies on the CPU dispatch in TopKTopPSampler.) + self.forward = self.forward_native + + def forward_native( + self, + logits: torch.Tensor, + generators: dict[int, torch.Generator], + k: torch.Tensor | None, + p: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Apply top-k/top-p filtering and sample tokens using pre-drawn noise.""" + + if generators: + warnings.warn( + "Generators are not supported by SpyreTopKTopPSampler. Falling back to base class.", + stacklevel=2, + ) + return super().forward_native(logits, generators, k, p) + + logits = apply_top_k_top_p(logits, k, p) + logits_to_return = None + if self.logprobs_mode == "processed_logits": + logits_to_return = logits + elif self.logprobs_mode == "processed_logprobs": + logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32) + + with self._noise_buffer.borrow_rows(n=logits.shape[0]) as log_noise: + sample_result = SpyreTopKTopPSampler._sample_with_predrawn_log_noise(logits, log_noise) + + return sample_result, logits_to_return + + def shutdown(self) -> None: + """Shutdown the sampler and clean up resources.""" + self._noise_buffer.shutdown() + + @staticmethod + def _sample_with_predrawn_noise(probs: torch.Tensor, noise: torch.Tensor) -> torch.Tensor: + """Sample using pre-drawn exponential noise (no exponential_() call).""" + return probs.div(noise).argmax(dim=-1).view(-1) + + @staticmethod + def _sample_with_predrawn_log_noise( + logits: torch.Tensor, log_noise: torch.Tensor + ) -> torch.Tensor: + """Sample using pre-drawn exponential log noise (no exponential_() call).""" + return (logits - log_noise).argmax(dim=-1).view(-1) diff --git a/spyre_inference/v1/worker/spyre_model_runner.py b/spyre_inference/v1/worker/spyre_model_runner.py index 1b35c95ce..c6f3d1c27 100644 --- a/spyre_inference/v1/worker/spyre_model_runner.py +++ b/spyre_inference/v1/worker/spyre_model_runner.py @@ -81,6 +81,7 @@ copy_pooler_output_to_cpu, select_rows, ) +from spyre_inference.v1.sample import build_spyre_sampler logger = init_logger(__name__) @@ -363,6 +364,14 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # _make_buffer (overridden below) places float .gpu tensors on Spyre # regardless of self.device. + # Host sampler (sendnn-inference#1046): async Exp(1) log-noise ring + # buffer + log-space Gumbel + TP rank-0 sample/broadcast. Falls back + # to upstream Sampler when vllm_config lacks max_num_seqs / vocab. + if hasattr(self, "sampler"): + self.sampler = build_spyre_sampler(vllm_config) + if getattr(self, "rejection_sampler", None) is not None: + self.rejection_sampler.sampler = self.sampler + # Disable GPU-specific features (same as CPUModelRunner) self.use_cuda_graph = False self.cascade_attn_enabled = False diff --git a/tests/test_async_ring_buffer.py b/tests/test_async_ring_buffer.py new file mode 100644 index 000000000..b24db8305 --- /dev/null +++ b/tests/test_async_ring_buffer.py @@ -0,0 +1,154 @@ +# Copyright 2026 The Spyre-Inference Authors. +# +# 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. + +import pytest + +from spyre_inference.v1.sample.async_ring_buffer import ( + AsyncExponential_RingBuffer, + AsyncRingBuffer, + _AsyncCounterRingBuffer, +) + + +class TestAsyncRingBuffer: + """Tests for AsyncRingBuffer, AsyncExponentialRingBuffer, and + _AsyncCounterRingBuffer.""" + + def test_abc_cannot_instantiate(self): + """AsyncRingBuffer is abstract and must not be instantiated directly.""" + with pytest.raises(TypeError): + AsyncRingBuffer(vocab_size=10, max_batch_size=4) # type: ignore[abstract] + + def test_counter_no_duplicates(self): + """Every row index must be yielded at most once across all borrows. + + The counter buffer does NOT guarantee strict global sequentiality: + rows near the wrap boundary can be skipped. The invariant it *does* + guarantee is that no row index is ever handed to the consumer twice. + """ + V, B, scale = 3, 4, 4 + buf = _AsyncCounterRingBuffer(vocab_size=V, max_batch_size=B, scale=scale) + total_steps = 20 + seen: list[int] = [] + try: + for _ in range(total_steps): + with buf.borrow_rows(1) as rows: + assert rows.shape == (1, V) + # All columns of a row share the same counter value. + val = int(rows[0, 0].item()) + seen.append(val) + finally: + buf.shutdown() + + assert len(seen) == len(set(seen)), f"duplicate row indices returned: {seen}" + + def test_counter_variable_batch_sizes(self): + """No row index must appear twice across borrows of varying size.""" + V, B, scale = 2, 4, 4 + buf = _AsyncCounterRingBuffer(vocab_size=V, max_batch_size=B, scale=scale) + batch_sizes = [1, 2, 3, 4, 1, 4, 2, 3] + seen: list[int] = [] + try: + for b in batch_sizes: + with buf.borrow_rows(b) as rows: + assert rows.shape == (b, V) + for i in range(b): + seen.append(int(rows[i, 0].item())) + finally: + buf.shutdown() + + assert len(seen) == len(set(seen)), f"duplicate row indices returned: {seen}" + + def test_counter_wrap_around(self): + """No row index is repeated when the buffer wraps multiple times.""" + V, B, scale = 2, 2, 4 # S = 8 + buf = _AsyncCounterRingBuffer(vocab_size=V, max_batch_size=B, scale=scale) + n_steps = 5 * scale + seen: list[int] = [] + try: + for _ in range(n_steps): + with buf.borrow_rows(B) as rows: + for i in range(B): + seen.append(int(rows[i, 0].item())) + finally: + buf.shutdown() + + assert len(seen) == len(set(seen)), f"duplicate row indices returned: {seen}" + + def test_exponential_shape(self): + """AsyncExponentialRingBuffer returns correctly shaped tensors.""" + V, B = 16, 4 + buf = AsyncExponential_RingBuffer(vocab_size=V, max_batch_size=B) + try: + for b in [1, 2, B]: + with buf.borrow_rows(b) as rows: + assert rows.shape == (b, V) + finally: + buf.shutdown() + + def test_borrow_is_zero_copy(self): + """borrow_rows must yield a view into the backing buffer, not a copy.""" + V, B = 8, 4 + buf = AsyncExponential_RingBuffer(vocab_size=V, max_batch_size=B) + try: + with buf.borrow_rows(B) as rows: + assert rows.untyped_storage().data_ptr() == buf._buf.untyped_storage().data_ptr() + finally: + buf.shutdown() + + def test_borrow_out_of_bounds(self): + """borrow_rows must raise a ValueError when n is outside the valid range.""" + V, B = 8, 4 + buf = AsyncExponential_RingBuffer(vocab_size=V, max_batch_size=B) + try: + with pytest.raises(ValueError, match="n.*must satisfy"), buf.borrow_rows(B + 1): + pass + finally: + buf.shutdown() + + def test_release_on_exception(self): + """Release must occur even when the consumer body raises.""" + V, B = 4, 2 + buf = _AsyncCounterRingBuffer(vocab_size=V, max_batch_size=B) + try: + with pytest.raises(RuntimeError, match="intentional") as _, buf.borrow_rows(B): + raise RuntimeError("intentional") + finally: + # If release happened, a second borrow must succeed. + with buf.borrow_rows(B) as rows: + assert rows.shape == (B, V) + buf.shutdown() + + def test_stop_joins_thread(self): + """stop() must cause the background thread to finish.""" + buf = AsyncExponential_RingBuffer(vocab_size=4, max_batch_size=2) + assert buf._thread.is_alive() + buf.shutdown() + assert not buf._thread.is_alive() + + @pytest.mark.parametrize("scale", [2, 3, 4, 8]) + def test_scale_invariance(self, scale: int): + """Different scale values must all produce no duplicate row indices.""" + V, B = 4, 3 + buf = _AsyncCounterRingBuffer(vocab_size=V, max_batch_size=B, scale=scale) + seen: list[int] = [] + try: + for _ in range(scale * 3): + with buf.borrow_rows(B) as rows: + for i in range(B): + seen.append(int(rows[i, 0].item())) + finally: + buf.shutdown() + + assert len(seen) == len(set(seen)), f"duplicate row indices returned: {seen}" diff --git a/tests/test_envs.py b/tests/test_envs.py new file mode 100644 index 000000000..4aa87c334 --- /dev/null +++ b/tests/test_envs.py @@ -0,0 +1,82 @@ +# Copyright 2026 The Spyre-Inference Authors. +# +# 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. + +"""Unit tests for ``spyre_inference.envs`` (no torch/vLLM required).""" + +from __future__ import annotations + +import importlib + +import pytest + + +@pytest.fixture(autouse=True) +def _fresh_envs(monkeypatch: pytest.MonkeyPatch): + """Reload envs with a clean cache for each test.""" + import spyre_inference.envs as envs + + envs.disable_envs_cache() + monkeypatch.delenv("SPYRE_USE_SPYRE_SAMPLER", raising=False) + monkeypatch.delenv("SPYRE_ASYNC_NOISE_SCALE", raising=False) + importlib.reload(envs) + yield envs + envs.disable_envs_cache() + + +def test_use_spyre_sampler_default_on(monkeypatch: pytest.MonkeyPatch): + import spyre_inference.envs as envs + + monkeypatch.delenv("SPYRE_USE_SPYRE_SAMPLER", raising=False) + assert envs.SPYRE_USE_SPYRE_SAMPLER is True + assert envs.is_set("SPYRE_USE_SPYRE_SAMPLER") is False + + +def test_use_spyre_sampler_override_off(monkeypatch: pytest.MonkeyPatch): + import spyre_inference.envs as envs + + monkeypatch.setenv("SPYRE_USE_SPYRE_SAMPLER", "0") + assert envs.SPYRE_USE_SPYRE_SAMPLER is False + assert envs.is_set("SPYRE_USE_SPYRE_SAMPLER") is True + + +def test_async_noise_scale_default(monkeypatch: pytest.MonkeyPatch): + import spyre_inference.envs as envs + + monkeypatch.delenv("SPYRE_ASYNC_NOISE_SCALE", raising=False) + assert envs.SPYRE_ASYNC_NOISE_SCALE == 4 + + +def test_async_noise_scale_override(monkeypatch: pytest.MonkeyPatch): + import spyre_inference.envs as envs + + monkeypatch.setenv("SPYRE_ASYNC_NOISE_SCALE", "8") + assert envs.SPYRE_ASYNC_NOISE_SCALE == 8 + + +def test_async_noise_scale_rejects_below_two(monkeypatch: pytest.MonkeyPatch): + import spyre_inference.envs as envs + + monkeypatch.setenv("SPYRE_ASYNC_NOISE_SCALE", "1") + with pytest.raises(ValueError, match="must be >= 2"): + _ = envs.SPYRE_ASYNC_NOISE_SCALE + + +def test_envs_cache_freezes_value(monkeypatch: pytest.MonkeyPatch): + import spyre_inference.envs as envs + + monkeypatch.setenv("SPYRE_ASYNC_NOISE_SCALE", "6") + envs.enable_envs_cache() + assert envs.SPYRE_ASYNC_NOISE_SCALE == 6 + monkeypatch.setenv("SPYRE_ASYNC_NOISE_SCALE", "10") + assert envs.SPYRE_ASYNC_NOISE_SCALE == 6 diff --git a/tests/test_spyre_sampler.py b/tests/test_spyre_sampler.py new file mode 100644 index 000000000..57abdbac4 --- /dev/null +++ b/tests/test_spyre_sampler.py @@ -0,0 +1,93 @@ +# Copyright 2026 The Spyre-Inference Authors. +# +# 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. + +from types import SimpleNamespace + +import pytest + +from spyre_inference.v1.sample.spyre_sampler import SpyreSampler + + +def _make_vllm_config(max_concurrency=1, vocab_size=128, use_text_config=False): + if use_text_config: + hf_config = SimpleNamespace(text_config=SimpleNamespace(vocab_size=vocab_size)) + else: + hf_config = SimpleNamespace(vocab_size=vocab_size) + return SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_seqs=max_concurrency), + model_config=SimpleNamespace(hf_config=hf_config), + ) + + +@pytest.fixture( + params=[ + _make_vllm_config(use_text_config=False), + _make_vllm_config(use_text_config=True), + ], + ids=["hf_config_vocab_size", "text_config_vocab_size"], +) +def valid_vllm_config(request): + """Return valid vLLM config variants expected to initialize SpyreSampler.""" + return request.param + + +@pytest.fixture( + params=[ + SimpleNamespace( + scheduler_config=SimpleNamespace(), + model_config=SimpleNamespace(hf_config=SimpleNamespace(vocab_size=128)), + ), + SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_seqs=1), + model_config=SimpleNamespace(hf_config=SimpleNamespace()), + ), + SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_seqs=1), + model_config=SimpleNamespace(hf_config=SimpleNamespace(text_config=SimpleNamespace())), + ), + ], + ids=["missing_concurrency", "missing_vocab_size", "missing_nested_vocab_size"], +) +def invalid_vllm_config(request): + """Return incomplete vLLM config objects that should fail validation.""" + return request.param + + +class TestSpyreSampler: + """Test suite for SpyreSampler.""" + + def test_initialization_rejects_fp64_gumbel(self, valid_vllm_config): + """Test that SpyreSampler raises ValueError with use_fp64_gumbel=True.""" + with pytest.raises(ValueError, match="SpyreSampler does not support use_fp64_gumbel=True"): + SpyreSampler(vllm_config=valid_vllm_config, use_fp64_gumbel=True) + + def test_initialization_accepts_supported_vllm_config_variants(self, valid_vllm_config): + """SpyreSampler should initialize when the required vLLM config fields are present.""" + sampler = SpyreSampler(vllm_config=valid_vllm_config) + + assert sampler.topk_topp_sampler is not None + assert sampler.topk_topp_sampler._noise_buffer is not None + assert SpyreSampler.is_vllm_config_compatible(valid_vllm_config) is True + + def test_initialization_rejects_incomplete_vllm_config(self, invalid_vllm_config): + """SpyreSampler should require both concurrency and vocabulary size metadata.""" + with pytest.raises(ValueError): + SpyreSampler(vllm_config=invalid_vllm_config) + + assert SpyreSampler.is_vllm_config_compatible(invalid_vllm_config) is False + + def test_is_vllm_config_compatible(self, valid_vllm_config, invalid_vllm_config): + """Compatibility checks should only pass when both required values are present.""" + assert SpyreSampler.is_vllm_config_compatible(valid_vllm_config) is True + assert SpyreSampler.is_vllm_config_compatible(invalid_vllm_config) is False diff --git a/tests/test_spyre_topk_topp_sampler.py b/tests/test_spyre_topk_topp_sampler.py new file mode 100644 index 000000000..907cd2b7c --- /dev/null +++ b/tests/test_spyre_topk_topp_sampler.py @@ -0,0 +1,87 @@ +# Copyright 2026 The Spyre-Inference Authors. +# +# 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. + +import torch + +from spyre_inference.v1.sample.spyre_topk_topp_sampler import SpyreTopKTopPSampler + + +class TestSpyreTopKTopPSampler: + """Test suite for SpyreTopKTopPSampler.""" + + def test_initialization_with_valid_params(self): + """Test that SpyreTopKTopPSampler initializes successfully with valid parameters.""" + vocab_size = 1000 + max_batch_size = 32 + + sampler = SpyreTopKTopPSampler( + vocab_size=vocab_size, + max_batch_size=max_batch_size, + logprobs_mode="raw_logprobs", + ) + + assert sampler is not None + assert sampler._noise_buffer is not None + sampler.shutdown() + + def test_forward_returns_valid_samples(self): + """Test that forward pass returns valid sampled token indices.""" + vocab_size = 100 + max_batch_size = 8 + batch_size = 4 + + sampler = SpyreTopKTopPSampler( + vocab_size=vocab_size, + max_batch_size=max_batch_size, + ) + + # Create dummy logits + logits = torch.randn(batch_size, vocab_size) + + # Forward pass without top-k/top-p constraints + samples, logprobs = sampler.forward( + logits=logits, + generators={}, + k=None, + p=None, + ) + + # Verify output shapes and types + assert samples.shape == (batch_size,), ( + f"Expected shape ({batch_size},), got {samples.shape}" + ) + assert samples.dtype == torch.long, f"Expected dtype torch.long, got {samples.dtype}" + assert logprobs is None, "Expected logprobs to be None with raw_logprobs mode" + + # Verify sampled tokens are within vocab range + assert (samples >= 0).all() and (samples < vocab_size).all(), ( + "Sampled tokens should be within vocabulary range" + ) + + sampler.shutdown() + + def test_gumble_max_trick(self): + """Test that sampling with log noise yields same tokens as regular sampling.""" + batch_size = 32 + vocab_size = 10000 + + logits = torch.randn(batch_size, vocab_size) + probs = logits.softmax(dim=-1, dtype=logits.dtype) + noise = torch.empty_like(logits).exponential_() + log_noise = torch.log(noise) + + expected_sampled_ids = SpyreTopKTopPSampler._sample_with_predrawn_noise(probs, noise) + gumble_sampled_ids = SpyreTopKTopPSampler._sample_with_predrawn_log_noise(logits, log_noise) + + assert torch.all(expected_sampled_ids == gumble_sampled_ids)