Feat/precomputed noise pool

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:

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())

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",
+]