feat(sampling): Add async ring buffer and Spyre-optimized samplers

sendnn_inference/model_executor/model_loader/spyre.py

@@ -23,6 +23,7 @@
 import sendnn_inference.multimodal as spyre_mm
 import sendnn_inference.utils as utils_spyre
 from sendnn_inference.platform import SpyrePlatform
+from sendnn_inference.v1.sample.spyre_sampler import SpyreSampler
 
 try:
     import backends.dynamo_tracer  # ty: ignore[unresolved-import] # noqa
@@ -112,8 +113,17 @@ def __init__(
         rank: int,
     ) -> None:
         super().__init__()
-
-        self.sampler = Sampler()
+        vllm_config_compatible = SpyreSampler.is_vllm_config_compatible(vllm_config)
+        if not vllm_config_compatible:
+            logger.warning(
+                "The provided vllm_config is not compatible with SpyreSampler. "
+                "Falling back to default Sampler with reduced performance on Spyre platform."
+            )
+            self.sampler = Sampler()
+        else:
+            # SpyreSampler is a vLLM Sampler subclass that uses top-k/top-p sampling
+            # implementations optimized for Spyre platform.
+            self.sampler = SpyreSampler(vllm_config=vllm_config)
 
         # boolean tensor of length batch size with indices:
         # True for unfinished sequences and

sendnn_inference/v1/sample/async_ring_buffer.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the sendnn-inference project
+
+import contextlib
+import queue
+import threading
+from abc import ABC, abstractmethod
+from collections.abc import Generator
+
+import torch
+
+
+class AsyncRingBuffer(ABC):
+    """Pre-generates data rows on a background thread via a ring buffer.
+
+    Maintains a contiguous ``(S, V)`` tensor (``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.
+
+    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
+
+        # buffer allocation
+        self._buf = torch.empty(self._S, self._V, dtype=torch.float32)
+
+        # 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, daemon=True)
+        self._thread.start()
+
+    @abstractmethod
+    def _refill_slice(self, start: int, end: int) -> None:
+        """Fill ``self._buf[start:end]`` with fresh values in-place."""
+        ...
+
+    @property
+    def vocab_size(self) -> int:
+        return self._V
+
+    @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 consumer to fill up at least n many values ahead
+        with self._cond:
+            self._cond.wait_for(lambda: self._tail >= end)
+
+        # get view (zero-copy)
+        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:
+        while True:
+            req = self._refill_q.get()
+
+            # handle termination signal
+            if req is None:
+                break
+
+            # refill buffer
+            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()
+
+    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 ``exponential_().log_()``."""
+
+    def _refill_slice(self, start: int, end: int) -> None:
+        self._buf[start:end].exponential_().log_()
+
+
+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._buf[start + i].fill_(self._total_generated)
+            self._total_generated += 1

sendnn_inference/v1/sample/spyre_sampler.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the sendnn-inference project
+
+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
+
+from sendnn_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,
+    ):
+        """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.
+
+        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")
+
+        super().__init__(logprobs_mode=logprobs_mode, use_fp64_gumbel=False)
+
+        # 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")
+
+        # 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,
+        )
+
+    @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()