Repository: vllm-project/vllm-spyre
Pull request #114: [CB] 🔥 remove cb env vars
Author: login: prashantgupta24
Created: 2025-04-25T19:14:12Z | Merged: 2025-05-03T17:10:20Z

Fix https://github.com/vllm-project/vllm-spyre/issues/94


Diff:
diff --git a/examples/offline_inference_spyre_cb_test.py b/examples/offline_inference_spyre_cb_test.py
index 13d4a84fa..bea5b1764 100644
--- a/examples/offline_inference_spyre_cb_test.py
+++ b/examples/offline_inference_spyre_cb_test.py
@@ -16,9 +16,6 @@
 os.environ['VLLM_SPYRE_USE_CB'] = '1'
 os.environ['VLLM_USE_V1'] = '1'
 
-os.environ['VLLM_SPYRE_MAX_CONTEXT_LENGTH'] = '2048'
-os.environ['VLLM_SPYRE_MAX_BATCH_SIZE'] = str(max_num_seqs)
-
 # Sample prompts.
 template = (
     "Below is an instruction that describes a task. Write a response that "
@@ -63,7 +60,8 @@
 llm = LLM(model="/models/llama-194m",
           tokenizer="/models/llama-194m",
           max_model_len=2048,
-          block_size=2048)
+          block_size=2048,
+          max_num_seqs=max_num_seqs)
 
 # Generate texts from the prompts. The output is a list of RequestOutput objects
 # that contain the prompt, generated text, and other information.
diff --git a/vllm_spyre/envs.py b/vllm_spyre/envs.py
index 3b07259aa..9fa4a3066 100644
--- a/vllm_spyre/envs.py
+++ b/vllm_spyre/envs.py
@@ -7,8 +7,6 @@
     VLLM_SPYRE_WARMUP_NEW_TOKENS: Optional[list[int]] = None
     VLLM_SPYRE_WARMUP_BATCH_SIZES: Optional[list[int]] = None
     VLLM_SPYRE_USE_CB: bool = False
-    VLLM_SPYRE_MAX_BATCH_SIZE: int = 0
-    VLLM_SPYRE_MAX_CONTEXT_LENGTH: int = 0
     VLLM_SPYRE_RM_PADDED_BLOCKS: bool = False
     VLLM_SPYRE_PERF_METRIC_LOGGING_ENABLED: int = 0
     VLLM_SPYRE_PERF_METRIC_LOGGING_DIR: str = "/tmp"
@@ -51,14 +49,6 @@
     "VLLM_SPYRE_USE_CB":
     lambda: bool(int(os.getenv("VLLM_SPYRE_USE_CB", "0"))),
 
-    # Maximal supported batch size
-    "VLLM_SPYRE_MAX_BATCH_SIZE":
-    lambda: int(os.getenv("VLLM_SPYRE_MAX_BATCH_SIZE", "0")),
-
-    # Maximal supported context length
-    "VLLM_SPYRE_MAX_CONTEXT_LENGTH":
-    lambda: int(os.getenv("VLLM_SPYRE_MAX_CONTEXT_LENGTH", "0")),
-
     # If set, remove redundant (left) padded blocks
     "VLLM_SPYRE_RM_PADDED_BLOCKS":
     lambda: bool(int(os.getenv("VLLM_SPYRE_RM_PADDED_BLOCKS", "0"))),
diff --git a/vllm_spyre/model_executor/model_loader/spyre.py b/vllm_spyre/model_executor/model_loader/spyre.py
index 2436811f0..4658a4451 100644
--- a/vllm_spyre/model_executor/model_loader/spyre.py
+++ b/vllm_spyre/model_executor/model_loader/spyre.py
@@ -8,7 +8,7 @@
 import torch.nn as nn
 from fms.models import get_model
 from transformers import PretrainedConfig
-from vllm.config import ModelConfig, ParallelConfig
+from vllm.config import ModelConfig, ParallelConfig, SchedulerConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
@@ -40,6 +40,7 @@ def __init__(
         self,
         model_config: ModelConfig,
         parallel_config: ParallelConfig,
+        scheduler_config: SchedulerConfig,
         max_prompt_length: int,
         max_decode_length: int,
     ) -> None:
@@ -60,11 +61,13 @@ def __init__(
         # FMS Model
         if envs_spyre.VLLM_SPYRE_USE_CB:
             self.model = ContinuousBatchingFmsModel(model_config,
-                                                    parallel_config)
+                                                    parallel_config,
+                                                    scheduler_config)
         else:
             self.model = StaticBatchingFmsModel(
                 model_config,
                 parallel_config,
+                scheduler_config,
                 max_prompt_length,
                 max_decode_length,
             )
@@ -263,12 +266,16 @@ def load_weights(
 
 class ContinuousBatchingFmsModel(FmsModelBase):
 
-    def __init__(self, model_config: ModelConfig,
-                 parallel_config: ParallelConfig) -> None:
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        parallel_config: ParallelConfig,
+        scheduler_config: SchedulerConfig,
+    ) -> None:
 
         BLOCK_SIZE = 64
-        max_batch = envs_spyre.VLLM_SPYRE_MAX_BATCH_SIZE
-        max_model_len = envs_spyre.VLLM_SPYRE_MAX_CONTEXT_LENGTH
+        max_batch = scheduler_config.max_num_seqs
+        max_model_len = scheduler_config.max_model_len
 
         # edge case: prompt fills model length: can produce 1 token with prefill
         max_prompt_length = max_model_len
@@ -348,6 +355,7 @@ def __init__(
         self,
         model_config: ModelConfig,
         parallel_config: ParallelConfig,
+        _: SchedulerConfig,
         max_prompt_length: int,
         max_decode_length: int,
     ) -> None:
diff --git a/vllm_spyre/platform.py b/vllm_spyre/platform.py
index 9ea819ac7..be05a7e8f 100644
--- a/vllm_spyre/platform.py
+++ b/vllm_spyre/platform.py
@@ -1,4 +1,5 @@
 import operator
+import os
 from typing import TYPE_CHECKING, Optional, Union
 
 import torch
@@ -28,6 +29,7 @@ class SpyrePlatform(Platform):
     device_type: str = "cpu"
     supported_quantization: list[str] = ["gptq"]
     _warmup_shapes: Optional[tuple[dict[str, int], ...]] = None
+    _config: VllmConfig = None
 
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
@@ -42,43 +44,24 @@ def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        cls._config = vllm_config
         parallel_config = vllm_config.parallel_config
         scheduler_config = vllm_config.scheduler_config
         model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
 
         if scheduler_config.is_multi_step:
             raise NotImplementedError
 
-        # continuous batching related checks
-        if envs_spyre.VLLM_SPYRE_USE_CB and not envs.VLLM_USE_V1:
-            raise NotImplementedError(
-                "Continuous batching is only implemented for vLLM V1")
-
         # Near future TODO: vLLM will have an api to check whether v0 or v1 is
         # used that isn't just checking the environment variable
 
         if parallel_config.worker_cls == "auto":
-            if envs.VLLM_USE_V1:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.v1.worker.spyre_worker.SpyreWorker"
-            else:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.worker.spyre_worker.SpyreWorker"
-
-        if envs.VLLM_USE_V1:
-            # As of 0.7.3 the scheduler for V1 isn't actually pluggable like
-            # this yet
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                scheduler_config.scheduler_cls = \
-                    "vllm_spyre.v1.core.scheduler.ContinuousBatchingSpyreScheduler"
-            else:
-                scheduler_config.scheduler_cls = \
-                    "vllm_spyre.v1.core.scheduler.StaticBatchingSpyreScheduler"
-        else:
-            scheduler_config.scheduler_cls = \
-                "vllm_spyre.core.scheduler.SpyreScheduler"
+            parallel_config.worker_cls = (
+                f'vllm_spyre{".v1" if envs.VLLM_USE_V1 else ""}'\
+                    '.worker.spyre_worker.SpyreWorker')
 
-        if not envs_spyre.VLLM_SPYRE_USE_CB:
+        if not envs_spyre.VLLM_SPYRE_USE_CB:  # no CB
             # Override --max-num-seqs to the biggest warmup batch size
             # And override --max-model-len to the biggest warmup sequence
             cls._warmup_shapes = None
@@ -86,39 +69,41 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
             max_batch_size = 0
             max_seq_len = 0
             for shape in spyre_warmup_shapes:
-                max_batch_size = max(max_batch_size, shape['batch_size'])
+                max_batch_size = max(max_batch_size, shape["batch_size"])
                 max_seq_len = max(max_seq_len,
-                                  shape['prompt_length'] + shape['new_tokens'])
-
-        if envs.VLLM_USE_V1:
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                # For continuous batching we use max_num_seqs to control
-                # the max batch size respecting AIU Spyre KV cache size
-                scheduler_config.max_num_seqs =\
-                    envs_spyre.VLLM_SPYRE_MAX_BATCH_SIZE
-                # ToDo: this function check_and_update_config is called twice:
-                # 1st time scheduler_config.max_num_seqs is what user sets
-                # 2nd time scheduler_config.max_num_seqs is 128
-            else:
-                # The v0 scheduler will run out of blocks if this is overridden
-                scheduler_config.max_num_seqs = max_batch_size
-
-        cache_config = vllm_config.cache_config
-
-        if cache_config and model_config:
-            # Cache and model config aren't set in the individual worker procs
-            # These are set in the main engine process
-
-            # To disable any paged attention ops in the base scheduler, we both:
-            # - Set the block size (in tokens) to the maximum sequence length
-            #       so that the scheduler thinks an entire sequence will fit in
-            #       one single block.
-            # - Set the number of blocks to the maximum number of sequences, so
-            #       the scheduler always thinks there's a block available
-            if not envs_spyre.VLLM_SPYRE_USE_CB:
+                                  shape["prompt_length"] + shape["new_tokens"])
+            if model_config is not None:
                 model_config.max_model_len = max_seq_len
-            cache_config.block_size = model_config.max_model_len
 
+            if envs.VLLM_USE_V1:  # No CB with V1
+
+                scheduler_config.scheduler_cls = (
+                    "vllm_spyre.v1.core.scheduler."\
+                        "StaticBatchingSpyreScheduler")
+                scheduler_config.max_num_seqs = max_batch_size
+            else:  # No CB with V0
+                scheduler_config.scheduler_cls = (
+                    "vllm_spyre.core.scheduler.SpyreScheduler")
+        else:  # CB related checks
+            if not envs.VLLM_USE_V1:  # CB with V0
+                raise NotImplementedError(
+                    "Continuous batching is only implemented for vLLM V1")
+            else:  # CB with V1
+                # As of 0.7.3 the scheduler for V1 isn't actually pluggable like
+                # this yet
+                scheduler_config.scheduler_cls = "vllm_spyre.v1.core."\
+                        "scheduler.ContinuousBatchingSpyreScheduler"
+
+        # Cache and model config aren't set in the individual worker procs
+        # These are set in the main engine process
+
+        # To disable any paged attention ops in the base scheduler, we both:
+        # - Set the block size (in tokens) to the maximum sequence length
+        #       so that the scheduler thinks an entire sequence will fit in
+        #       one single block.
+        # - Set the number of blocks to the maximum number of sequences, so
+        #       the scheduler always thinks there's a block available
+        if cache_config is not None:
             if envs.VLLM_USE_V1:
                 # The V1 scheduler actually needs 2 blocks for each sequence...
                 cache_config.num_gpu_blocks_override = \
@@ -127,12 +112,30 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
                 cache_config.num_gpu_blocks_override = \
                     scheduler_config.max_num_seqs
 
-            logger.info(
-                "Overriding configurations based on warmup shapes. "
-                "max_model_len=%d, max_num_seqs=%d, block_size=%d, "
-                "num_gpu_blocks_override=%d", model_config.max_model_len,
-                scheduler_config.max_num_seqs, cache_config.block_size,
-                cache_config.num_gpu_blocks_override)
+            cache_config.block_size = model_config.max_model_len
+
+        logger.info(
+            "Overriding configurations based on warmup shapes. "
+            "max_model_len=%d, max_num_seqs=%d, block_size=%d, "
+            "num_gpu_blocks_override=%d",
+            model_config.max_model_len,
+            scheduler_config.max_num_seqs,
+            cache_config.block_size,
+            cache_config.num_gpu_blocks_override,
+        )
+
+        # set env vars for torch_sendnn to consume
+        os.environ["VLLM_DT_MAX_CONTEXT_LEN"] = str(
+            vllm_config.model_config.max_model_len)
+        os.environ["VLLM_DT_MAX_BATCH_SIZE"] = str(
+            vllm_config.scheduler_config.max_num_seqs)
+
+    @classmethod
+    def use_all_gather(cls) -> bool:
+        """
+        Whether to use allgather in LogitsProcessor to gather the logits.
+        """
+        return True
 
     @classmethod
     def is_pin_memory_available(cls) -> bool:
@@ -167,10 +170,9 @@ def get_warmup_shapes(cls, scheduler_config) -> tuple[dict[str, int], ...]:
                     "The lists in VLLM_SPYRE_WARMUP_PROMPT_LENS and "
                     "VLLM_SPYRE_WARMUP_NEW_TOKENS must have equal length")
 
-        if not envs_spyre.VLLM_SPYRE_USE_CB:
-            logger.info("VLLM_SPYRE_WARMUP_PROMPT_LENS = %s", wup_prompt_lens)
-            logger.info("VLLM_SPYRE_WARMUP_NEW_TOKENS = %s", wup_new_tokens)
-            logger.info("VLLM_SPYRE_WARMUP_BATCH_SIZES = %s", wup_batch_sizes)
+        logger.info("VLLM_SPYRE_WARMUP_PROMPT_LENS = %s", wup_prompt_lens)
+        logger.info("VLLM_SPYRE_WARMUP_NEW_TOKENS = %s", wup_new_tokens)
+        logger.info("VLLM_SPYRE_WARMUP_BATCH_SIZES = %s", wup_batch_sizes)
 
         cls._warmup_shapes = tuple(
             sorted([{
@@ -225,13 +227,13 @@ def validate_request(
             d = 64  # hardcoded AIU Spyre block size
             prompt_padding_len = ((n + d - 1) // d) * d
             if (prompt_padding_len + max_tokens
-                    > envs_spyre.VLLM_SPYRE_MAX_CONTEXT_LENGTH):
+                    > cls._config.scheduler_config.max_model_len):
                 raise ValueError(
                     "Could not add request: prompt length is "
                     f"{prompt_len} tokens, which gets padded to "
                     f"{prompt_padding_len} tokens, maximum number of output "
                     f"tokens is {max_tokens} tokens, but max model context "
-                    f"length is {envs_spyre.VLLM_SPYRE_MAX_CONTEXT_LENGTH}.")
+                    f"length is {cls._config.scheduler_config.max_model_len}.")
         else:
             # For non-continuous batching, check if the request matches a warmup
             # shape
diff --git a/vllm_spyre/v1/core/scheduler.py b/vllm_spyre/v1/core/scheduler.py
index 1cf7a1659..c4985d3f8 100644
--- a/vllm_spyre/v1/core/scheduler.py
+++ b/vllm_spyre/v1/core/scheduler.py
@@ -9,7 +9,6 @@
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request
 
-import vllm_spyre.envs as envs_spyre
 from vllm_spyre.platform import SpyrePlatform
 from vllm_spyre.v1.worker.spyre_model_runner import CBSpyreModelRunnerOutput
 
@@ -210,7 +209,7 @@ def schedule(self) -> "SchedulerOutput":
 
     def can_schedule(self, request) -> bool:
         max_prompt_batch_size = 1
-        max_context_len = envs_spyre.VLLM_SPYRE_MAX_CONTEXT_LENGTH
+        max_context_len = self.scheduler_config.max_model_len
 
         # running and waiting queues are both empty -> start new batch
         start_new_batch = len(self.running) + len(self.waiting) == 0
diff --git a/vllm_spyre/v1/worker/spyre_model_runner.py b/vllm_spyre/v1/worker/spyre_model_runner.py
index a1b02f09e..6e611d882 100644
--- a/vllm_spyre/v1/worker/spyre_model_runner.py
+++ b/vllm_spyre/v1/worker/spyre_model_runner.py
@@ -107,6 +107,7 @@ def load_model(self, prompt_lens: Iterable[int],
         self.model = SpyreCausalLM(
             self.model_config,
             parallel_config=self.parallel_config,
+            scheduler_config=self.scheduler_config,
             max_prompt_length=max_pad_length,
             max_decode_length=max_decode_length,
         )
@@ -569,8 +570,8 @@ def __init__(
         super().__init__(vllm_config=vllm_config,
                          is_driver_worker=is_driver_worker)
 
-        max_batch_size = envs_spyre.VLLM_SPYRE_MAX_BATCH_SIZE
-        max_model_len = envs_spyre.VLLM_SPYRE_MAX_CONTEXT_LENGTH
+        max_batch_size = vllm_config.scheduler_config.max_num_seqs
+        max_model_len = vllm_config.scheduler_config.max_model_len
 
         self.BLOCK_SIZE = 64
         NUM_BLOCKS = max_batch_size * max_model_len // self.BLOCK_SIZE  # 64
diff --git a/vllm_spyre/worker/spyre_model_runner.py b/vllm_spyre/worker/spyre_model_runner.py
index 692900f60..32f49bd4b 100644
--- a/vllm_spyre/worker/spyre_model_runner.py
+++ b/vllm_spyre/worker/spyre_model_runner.py
@@ -111,10 +111,13 @@ def load_model(self, prompt_lens: Iterable[int],
                    num_decode_tokens: Iterable[int]) -> None:
         max_pad_length = max(prompt_lens)
         max_decode_length = max(num_decode_tokens)
-        self.model = SpyreCausalLM(self.model_config,
-                                   parallel_config=self.parallel_config,
-                                   max_prompt_length=max_pad_length,
-                                   max_decode_length=max_decode_length)
+        self.model = SpyreCausalLM(
+            self.model_config,
+            parallel_config=self.parallel_config,
+            scheduler_config=self.scheduler_config,
+            max_prompt_length=max_pad_length,
+            max_decode_length=max_decode_length,
+        )
 
     @property
     def vocab_size(self) -> int:


Reviews:
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-25T19:29:21Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-25T19:31:10Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-28T02:28:28Z
user: login: wallashss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-28T14:33:58Z
user: login: joerunde
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-28T14:46:22Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T20:04:46Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T22:44:29Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T23:08:16Z
user: login: yannicks1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-30T23:45:11Z
user: login: joerunde
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-01T14:44:59Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-01T15:02:03Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-01T22:20:38Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-01T22:26:56Z
user: login: yannicks1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-02T14:48:39Z
user: login: yannicks1
body: Also, the code currently breaks here: https://github.com/vllm-project/vllm-spyre/blob/c0070a3d4b6bd90f6629c0f7b78211c7c87cc7c3/vllm_spyre/v1/core/scheduler.py#L213 
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-05-02T14:56:59Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-02T17:39:10Z
user: login: prashantgupta24
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-02T17:45:31Z
user: login: yannicks1
body: LGTM, great work!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-03T15:48:29Z

Review comments:
diff_hunk: @@ -77,26 +78,28 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         # Override --max-num-seqs to the biggest warmup batch size
         # And override --max-model-len to the biggest warmup sequence
         cls._warmup_shapes = None
-        spyre_warmup_shapes = cls.get_warmup_shapes(scheduler_config)
+        spyre_warmup_shapes = \
+            cls.get_warmup_shapes(scheduler_config)
         max_batch_size = 0
         max_seq_len = 0
         for shape in spyre_warmup_shapes:
             max_batch_size = max(max_batch_size, shape['batch_size'])
             max_seq_len = max(max_seq_len,
                               shape['prompt_length'] + shape['new_tokens'])
 
-        if envs.VLLM_USE_V1:
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                # For continuous batching we use max_num_seqs to control
-                # the max batch size respecting AIU Spyre KV cache size
-                scheduler_config.max_num_seqs =\
-                    envs_spyre.VLLM_SPYRE_MAX_BATCH_SIZE
-                # ToDo: this function check_and_update_config is called twice:
-                # 1st time scheduler_config.max_num_seqs is what user sets
-                # 2nd time scheduler_config.max_num_seqs is 128
-            else:
-                # The v0 scheduler will run out of blocks if this is overridden
-                scheduler_config.max_num_seqs = max_batch_size
+        # if envs.VLLM_USE_V1:
+        #     if envs_spyre.VLLM_SPYRE_USE_CB:
+        #         # For continuous batching we use max_num_seqs to control
+        #         # the max batch size respecting AIU Spyre KV cache size
+        #         scheduler_config.max_num_seqs =\
+        #             envs_spyre.VLLM_SPYRE_MAX_BATCH_SIZE
+        #         # ToDo: this function check_and_update_config is called twice:
+        #         # 1st time scheduler_config.max_num_seqs is what user sets
+        #         # 2nd time scheduler_config.max_num_seqs is 128
+        #     else:
+        #         # The v0 scheduler will run out of
+        #           blocks if this is overridden
+        #         scheduler_config.max_num_seqs = max_batch_size
path: vllm_spyre/platform.py
user: login: prashantgupta24
body: not 100% sure we need this or not?
```
scheduler_config.max_num_seqs = max_batch_size
```
created_at: 2025-04-25T19:29:21Z
author_association: COLLABORATOR
diff_hunk: @@ -133,6 +136,13 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
                 scheduler_config.max_num_seqs, cache_config.block_size,
                 cache_config.num_gpu_blocks_override)
 
+    @classmethod
+    def use_all_gather(cls) -> bool:
+        """
+        Whether to use allgather in LogitsProcessor to gather the logits.
+        """
+        return True
+
path: vllm_spyre/platform.py
user: login: prashantgupta24
body: This was the reason that there were two calls to `check_and_update_config` in `platform.py ` - since this function did not exist, the `interface.py`'s `use_all_gather` was being called which defaulted all the config values
created_at: 2025-04-25T19:31:10Z
author_association: COLLABORATOR
diff_hunk: @@ -133,6 +136,13 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
                 scheduler_config.max_num_seqs, cache_config.block_size,
                 cache_config.num_gpu_blocks_override)
 
+    @classmethod
+    def use_all_gather(cls) -> bool:
+        """
+        Whether to use allgather in LogitsProcessor to gather the logits.
+        """
+        return True
+
path: vllm_spyre/platform.py
user: login: prashantgupta24
body: Do we need to add a condition here like in `interaface.py`?
created_at: 2025-04-28T02:28:28Z
author_association: COLLABORATOR
diff_hunk: @@ -133,6 +136,13 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
                 scheduler_config.max_num_seqs, cache_config.block_size,
                 cache_config.num_gpu_blocks_override)
 
+    @classmethod
+    def use_all_gather(cls) -> bool:
+        """
+        Whether to use allgather in LogitsProcessor to gather the logits.
+        """
+        return True
+
path: vllm_spyre/platform.py
user: login: wallashss
body: I guess not, IIUC, when we use V1, it should use `use_all_gather` 
created_at: 2025-04-28T14:33:58Z
author_association: COLLABORATOR
diff_hunk: @@ -77,26 +78,28 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         # Override --max-num-seqs to the biggest warmup batch size
         # And override --max-model-len to the biggest warmup sequence
         cls._warmup_shapes = None
-        spyre_warmup_shapes = cls.get_warmup_shapes(scheduler_config)
+        spyre_warmup_shapes = \
+            cls.get_warmup_shapes(scheduler_config)
         max_batch_size = 0
         max_seq_len = 0
         for shape in spyre_warmup_shapes:
             max_batch_size = max(max_batch_size, shape['batch_size'])
             max_seq_len = max(max_seq_len,
                               shape['prompt_length'] + shape['new_tokens'])
 
-        if envs.VLLM_USE_V1:
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                # For continuous batching we use max_num_seqs to control
-                # the max batch size respecting AIU Spyre KV cache size
-                scheduler_config.max_num_seqs =\
-                    envs_spyre.VLLM_SPYRE_MAX_BATCH_SIZE
-                # ToDo: this function check_and_update_config is called twice:
-                # 1st time scheduler_config.max_num_seqs is what user sets
-                # 2nd time scheduler_config.max_num_seqs is 128
-            else:
-                # The v0 scheduler will run out of blocks if this is overridden
-                scheduler_config.max_num_seqs = max_batch_size
+        # if envs.VLLM_USE_V1:
+        #     if envs_spyre.VLLM_SPYRE_USE_CB:
+        #         # For continuous batching we use max_num_seqs to control
+        #         # the max batch size respecting AIU Spyre KV cache size
+        #         scheduler_config.max_num_seqs =\
+        #             envs_spyre.VLLM_SPYRE_MAX_BATCH_SIZE
+        #         # ToDo: this function check_and_update_config is called twice:
+        #         # 1st time scheduler_config.max_num_seqs is what user sets
+        #         # 2nd time scheduler_config.max_num_seqs is 128
+        #     else:
+        #         # The v0 scheduler will run out of
+        #           blocks if this is overridden
+        #         scheduler_config.max_num_seqs = max_batch_size
path: vllm_spyre/platform.py
user: login: joerunde
body: For the static batching use case the max batch size is still set with `VLLM_SPYRE_WARMUP_BATCH_SIZES`, so we override it to make sure the scheduling config is consistent. (e.g. setting both `--max-num-seqs=1` and `VLLM_SPYRE_WARMUP_BATCH_SIZES=2,4,6` wouldn't work out too well.)

The model and cache configs are also updated in a similar fashion below
created_at: 2025-04-28T14:46:22Z
author_association: COLLABORATOR
diff_hunk: @@ -16,9 +16,6 @@
 os.environ['VLLM_SPYRE_USE_CB'] = '1'
 os.environ['VLLM_USE_V1'] = '1'
 
-os.environ['VLLM_SPYRE_MAX_CONTEXT_LENGTH'] = '2048'
-os.environ['VLLM_SPYRE_MAX_BATCH_SIZE'] = str(max_num_seqs)
path: examples/offline_inference_spyre_cb_test.py
user: login: prashantgupta24
body: Do we want to leave these in since pytorch implementation might need them?
created_at: 2025-04-30T20:04:45Z
author_association: COLLABORATOR
diff_hunk: @@ -43,6 +44,7 @@ def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        cls._config = vllm_config
path: vllm_spyre/platform.py
user: login: prashantgupta24
body: Unfortunately using `get_current_vllm_config` doesn't work here, it always gets the default values :( That's why I am doing it like this
created_at: 2025-04-30T22:44:29Z
author_association: COLLABORATOR
diff_hunk: @@ -43,6 +44,7 @@ def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        cls._config = vllm_config
path: vllm_spyre/platform.py
user: login: prashantgupta24
body: Should I change it at the API layer itself? I can imagine validating requests might need config params in the future
created_at: 2025-04-30T23:08:15Z
author_association: COLLABORATOR
diff_hunk: @@ -16,9 +16,6 @@
 os.environ['VLLM_SPYRE_USE_CB'] = '1'
 os.environ['VLLM_USE_V1'] = '1'
 
-os.environ['VLLM_SPYRE_MAX_CONTEXT_LENGTH'] = '2048'
-os.environ['VLLM_SPYRE_MAX_BATCH_SIZE'] = str(max_num_seqs)
path: examples/offline_inference_spyre_cb_test.py
user: login: yannicks1
body: yes, let's keep setting them here for the time being just to be on the safe side. cc @joerunde 
created_at: 2025-04-30T23:45:11Z
author_association: COLLABORATOR
diff_hunk: @@ -43,6 +44,7 @@ def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        cls._config = vllm_config
path: vllm_spyre/platform.py
user: login: joerunde
body: I think `get_current_vllm_config` will only work post-engine-bootup, and `check_and_update_config` is called during engine creation. I'm curious, does it also not work if you call `get_current_vllm_config` inside of `validate_request`?

I'm fine with keeping a class cache around like this though
created_at: 2025-05-01T14:44:59Z
author_association: COLLABORATOR
diff_hunk: @@ -43,6 +44,7 @@ def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        cls._config = vllm_config
path: vllm_spyre/platform.py
user: login: prashantgupta24
body: > does it also not work if you call `get_current_vllm_config` inside of `validate_request`?

Nope :( 


created_at: 2025-05-01T15:02:03Z
author_association: COLLABORATOR
diff_hunk: @@ -16,9 +16,6 @@
 os.environ['VLLM_SPYRE_USE_CB'] = '1'
 os.environ['VLLM_USE_V1'] = '1'
 
-os.environ['VLLM_SPYRE_MAX_CONTEXT_LENGTH'] = '2048'
-os.environ['VLLM_SPYRE_MAX_BATCH_SIZE'] = str(max_num_seqs)
path: examples/offline_inference_spyre_cb_test.py
user: login: prashantgupta24
body: Removing them from here since we set them in the `check_and_update_config` function now. As Joe said, inside `check_and_update_config` set them back in `os.environ` to the values that are in `vllm_config`. That way the single source of truth is the CLI args, and we control setting the environment up for `torch_sendnn` to work properly.
created_at: 2025-05-01T22:20:38Z
author_association: COLLABORATOR
diff_hunk: @@ -168,10 +172,9 @@ def get_warmup_shapes(cls, scheduler_config) -> tuple[dict[str, int], ...]:
                     "The lists in VLLM_SPYRE_WARMUP_PROMPT_LENS and "
                     "VLLM_SPYRE_WARMUP_NEW_TOKENS must have equal length")
 
-        if not envs_spyre.VLLM_SPYRE_USE_CB:
path: vllm_spyre/platform.py
user: login: prashantgupta24
body: duplicated check, the function is called before checking `if not envs_spyre.VLLM_SPYRE_USE_CB` 
created_at: 2025-05-01T22:26:56Z
author_association: COLLABORATOR
diff_hunk: @@ -43,97 +45,97 @@ def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        cls._config = vllm_config
         parallel_config = vllm_config.parallel_config
         scheduler_config = vllm_config.scheduler_config
         model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
 
         if scheduler_config.is_multi_step:
             raise NotImplementedError
 
-        # continuous batching related checks
-        if envs_spyre.VLLM_SPYRE_USE_CB and not envs.VLLM_USE_V1:
-            raise NotImplementedError(
-                "Continuous batching is only implemented for vLLM V1")
-
         # Near future TODO: vLLM will have an api to check whether v0 or v1 is
         # used that isn't just checking the environment variable
 
         if parallel_config.worker_cls == "auto":
-            if envs.VLLM_USE_V1:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.v1.worker.spyre_worker.SpyreWorker"
-            else:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.worker.spyre_worker.SpyreWorker"
+            parallel_config.worker_cls = (
+                f'vllm_spyre{".v1" if envs.VLLM_USE_V1 else ""}'\
path: vllm_spyre/platform.py
user: login: yannicks1
body: nice hack:)
created_at: 2025-05-02T14:45:21Z
author_association: COLLABORATOR
diff_hunk: @@ -43,97 +45,97 @@ def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        cls._config = vllm_config
         parallel_config = vllm_config.parallel_config
         scheduler_config = vllm_config.scheduler_config
         model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
 
         if scheduler_config.is_multi_step:
             raise NotImplementedError
 
-        # continuous batching related checks
-        if envs_spyre.VLLM_SPYRE_USE_CB and not envs.VLLM_USE_V1:
-            raise NotImplementedError(
-                "Continuous batching is only implemented for vLLM V1")
-
         # Near future TODO: vLLM will have an api to check whether v0 or v1 is
         # used that isn't just checking the environment variable
 
         if parallel_config.worker_cls == "auto":
-            if envs.VLLM_USE_V1:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.v1.worker.spyre_worker.SpyreWorker"
-            else:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.worker.spyre_worker.SpyreWorker"
+            parallel_config.worker_cls = (
+                f'vllm_spyre{".v1" if envs.VLLM_USE_V1 else ""}'\
+                    '.worker.spyre_worker.SpyreWorker')
 
-        if envs.VLLM_USE_V1:
-            # As of 0.7.3 the scheduler for V1 isn't actually pluggable like
-            # this yet
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                scheduler_config.scheduler_cls = \
-                    "vllm_spyre.v1.core.scheduler.ContinuousBatchingSpyreScheduler"
-            else:
-                scheduler_config.scheduler_cls = \
-                    "vllm_spyre.v1.core.scheduler.StaticBatchingSpyreScheduler"
-        else:
-            scheduler_config.scheduler_cls = \
-                "vllm_spyre.core.scheduler.SpyreScheduler"
-
-        if not envs_spyre.VLLM_SPYRE_USE_CB:
+        # continuous batching related checks
path: vllm_spyre/platform.py
user: login: yannicks1
body: a bit confusing here as the lines after `# continuous batching related checks` are the static batching checks
created_at: 2025-05-02T14:45:55Z
author_association: COLLABORATOR
diff_hunk: @@ -43,97 +45,97 @@ def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        cls._config = vllm_config
         parallel_config = vllm_config.parallel_config
         scheduler_config = vllm_config.scheduler_config
         model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
 
         if scheduler_config.is_multi_step:
             raise NotImplementedError
 
-        # continuous batching related checks
-        if envs_spyre.VLLM_SPYRE_USE_CB and not envs.VLLM_USE_V1:
-            raise NotImplementedError(
-                "Continuous batching is only implemented for vLLM V1")
-
         # Near future TODO: vLLM will have an api to check whether v0 or v1 is
         # used that isn't just checking the environment variable
 
         if parallel_config.worker_cls == "auto":
-            if envs.VLLM_USE_V1:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.v1.worker.spyre_worker.SpyreWorker"
-            else:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.worker.spyre_worker.SpyreWorker"
+            parallel_config.worker_cls = (
+                f'vllm_spyre{".v1" if envs.VLLM_USE_V1 else ""}'\
+                    '.worker.spyre_worker.SpyreWorker')
 
-        if envs.VLLM_USE_V1:
-            # As of 0.7.3 the scheduler for V1 isn't actually pluggable like
-            # this yet
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                scheduler_config.scheduler_cls = \
-                    "vllm_spyre.v1.core.scheduler.ContinuousBatchingSpyreScheduler"
-            else:
-                scheduler_config.scheduler_cls = \
-                    "vllm_spyre.v1.core.scheduler.StaticBatchingSpyreScheduler"
-        else:
-            scheduler_config.scheduler_cls = \
-                "vllm_spyre.core.scheduler.SpyreScheduler"
-
-        if not envs_spyre.VLLM_SPYRE_USE_CB:
+        # continuous batching related checks
+        if not envs_spyre.VLLM_SPYRE_USE_CB:  # no CB
             # Override --max-num-seqs to the biggest warmup batch size
             # And override --max-model-len to the biggest warmup sequence
             cls._warmup_shapes = None
             spyre_warmup_shapes = cls.get_warmup_shapes(scheduler_config)
             max_batch_size = 0
             max_seq_len = 0
             for shape in spyre_warmup_shapes:
-                max_batch_size = max(max_batch_size, shape['batch_size'])
+                max_batch_size = max(max_batch_size, shape["batch_size"])
                 max_seq_len = max(max_seq_len,
-                                  shape['prompt_length'] + shape['new_tokens'])
+                                  shape["prompt_length"] + shape["new_tokens"])
+            model_config.max_model_len = max_seq_len
 
-        if envs.VLLM_USE_V1:
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                # For continuous batching we use max_num_seqs to control
-                # the max batch size respecting AIU Spyre KV cache size
-                scheduler_config.max_num_seqs =\
-                    envs_spyre.VLLM_SPYRE_MAX_BATCH_SIZE
-                # ToDo: this function check_and_update_config is called twice:
-                # 1st time scheduler_config.max_num_seqs is what user sets
-                # 2nd time scheduler_config.max_num_seqs is 128
-            else:
-                # The v0 scheduler will run out of blocks if this is overridden
-                scheduler_config.max_num_seqs = max_batch_size
+            if envs.VLLM_USE_V1:  # No CB with V1
 
-        cache_config = vllm_config.cache_config
+                scheduler_config.scheduler_cls = (
+                    "vllm_spyre.v1.core.scheduler."\
+                        "StaticBatchingSpyreScheduler")
+                scheduler_config.max_num_seqs = max_batch_size
+            else:  # No CB with V0
+                scheduler_config.scheduler_cls = (
+                    "vllm_spyre.core.scheduler.SpyreScheduler")
+        else:  # CB
+            if not envs.VLLM_USE_V1:  # CB with V0
+                raise NotImplementedError(
+                    "Continuous batching is only implemented for vLLM V1")
+            else:  # CB with V1
+                # As of 0.7.3 the scheduler for V1 isn't actually pluggable like
+                # this yet
+                scheduler_config.scheduler_cls = "vllm_spyre.v1.core."\
+                        "scheduler.ContinuousBatchingSpyreScheduler"
+
+        # Cache and model config aren't set in the individual worker procs
+        # These are set in the main engine process
+
+        # To disable any paged attention ops in the base scheduler, we both:
+        # - Set the block size (in tokens) to the maximum sequence length
+        #       so that the scheduler thinks an entire sequence will fit in
+        #       one single block.
+        # - Set the number of blocks to the maximum number of sequences, so
+        #       the scheduler always thinks there's a block available
+        if envs.VLLM_USE_V1:
+            # The V1 scheduler actually needs 2 blocks for each sequence...
+            cache_config.num_gpu_blocks_override = \
+                scheduler_config.max_num_seqs * 2
+        else:
+            cache_config.num_gpu_blocks_override = \
+                scheduler_config.max_num_seqs
+
+        cache_config.block_size = model_config.max_model_len
+
+        logger.info(
+            "Overriding configurations based on warmup shapes. "
+            "max_model_len=%d, max_num_seqs=%d, block_size=%d, "
+            "num_gpu_blocks_override=%d",
+            model_config.max_model_len,
+            scheduler_config.max_num_seqs,
+            cache_config.block_size,
+            cache_config.num_gpu_blocks_override,
+        )
+
+        # set env vars for torch_sendnn to consume
+        os.environ["VLLM_DT_MAX_CONTEXT_LEN"] = str(
+            vllm_config.model_config.max_model_len)
+        os.environ["VLLM_DT_MAX_BATCH_SIZE"] = str(
+            vllm_config.scheduler_config.max_num_seqs)
 
-        if cache_config and model_config:
path: vllm_spyre/platform.py
user: login: yannicks1
body: this check is removed, is this save? 
created_at: 2025-05-02T14:46:29Z
author_association: COLLABORATOR
diff_hunk: @@ -43,97 +45,97 @@ def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        cls._config = vllm_config
         parallel_config = vllm_config.parallel_config
         scheduler_config = vllm_config.scheduler_config
         model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
 
         if scheduler_config.is_multi_step:
             raise NotImplementedError
 
-        # continuous batching related checks
-        if envs_spyre.VLLM_SPYRE_USE_CB and not envs.VLLM_USE_V1:
-            raise NotImplementedError(
-                "Continuous batching is only implemented for vLLM V1")
-
         # Near future TODO: vLLM will have an api to check whether v0 or v1 is
         # used that isn't just checking the environment variable
 
         if parallel_config.worker_cls == "auto":
-            if envs.VLLM_USE_V1:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.v1.worker.spyre_worker.SpyreWorker"
-            else:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.worker.spyre_worker.SpyreWorker"
+            parallel_config.worker_cls = (
+                f'vllm_spyre{".v1" if envs.VLLM_USE_V1 else ""}'\
+                    '.worker.spyre_worker.SpyreWorker')
 
-        if envs.VLLM_USE_V1:
-            # As of 0.7.3 the scheduler for V1 isn't actually pluggable like
-            # this yet
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                scheduler_config.scheduler_cls = \
-                    "vllm_spyre.v1.core.scheduler.ContinuousBatchingSpyreScheduler"
-            else:
-                scheduler_config.scheduler_cls = \
-                    "vllm_spyre.v1.core.scheduler.StaticBatchingSpyreScheduler"
-        else:
-            scheduler_config.scheduler_cls = \
-                "vllm_spyre.core.scheduler.SpyreScheduler"
-
-        if not envs_spyre.VLLM_SPYRE_USE_CB:
+        # continuous batching related checks
+        if not envs_spyre.VLLM_SPYRE_USE_CB:  # no CB
             # Override --max-num-seqs to the biggest warmup batch size
             # And override --max-model-len to the biggest warmup sequence
             cls._warmup_shapes = None
             spyre_warmup_shapes = cls.get_warmup_shapes(scheduler_config)
             max_batch_size = 0
             max_seq_len = 0
             for shape in spyre_warmup_shapes:
-                max_batch_size = max(max_batch_size, shape['batch_size'])
+                max_batch_size = max(max_batch_size, shape["batch_size"])
                 max_seq_len = max(max_seq_len,
-                                  shape['prompt_length'] + shape['new_tokens'])
+                                  shape["prompt_length"] + shape["new_tokens"])
+            model_config.max_model_len = max_seq_len
 
-        if envs.VLLM_USE_V1:
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                # For continuous batching we use max_num_seqs to control
-                # the max batch size respecting AIU Spyre KV cache size
-                scheduler_config.max_num_seqs =\
-                    envs_spyre.VLLM_SPYRE_MAX_BATCH_SIZE
-                # ToDo: this function check_and_update_config is called twice:
-                # 1st time scheduler_config.max_num_seqs is what user sets
-                # 2nd time scheduler_config.max_num_seqs is 128
-            else:
-                # The v0 scheduler will run out of blocks if this is overridden
-                scheduler_config.max_num_seqs = max_batch_size
+            if envs.VLLM_USE_V1:  # No CB with V1
 
-        cache_config = vllm_config.cache_config
+                scheduler_config.scheduler_cls = (
+                    "vllm_spyre.v1.core.scheduler."\
+                        "StaticBatchingSpyreScheduler")
+                scheduler_config.max_num_seqs = max_batch_size
+            else:  # No CB with V0
+                scheduler_config.scheduler_cls = (
+                    "vllm_spyre.core.scheduler.SpyreScheduler")
+        else:  # CB
+            if not envs.VLLM_USE_V1:  # CB with V0
+                raise NotImplementedError(
+                    "Continuous batching is only implemented for vLLM V1")
+            else:  # CB with V1
+                # As of 0.7.3 the scheduler for V1 isn't actually pluggable like
+                # this yet
+                scheduler_config.scheduler_cls = "vllm_spyre.v1.core."\
+                        "scheduler.ContinuousBatchingSpyreScheduler"
+
+        # Cache and model config aren't set in the individual worker procs
+        # These are set in the main engine process
+
+        # To disable any paged attention ops in the base scheduler, we both:
+        # - Set the block size (in tokens) to the maximum sequence length
+        #       so that the scheduler thinks an entire sequence will fit in
+        #       one single block.
+        # - Set the number of blocks to the maximum number of sequences, so
+        #       the scheduler always thinks there's a block available
+        if envs.VLLM_USE_V1:
+            # The V1 scheduler actually needs 2 blocks for each sequence...
+            cache_config.num_gpu_blocks_override = \
+                scheduler_config.max_num_seqs * 2
+        else:
+            cache_config.num_gpu_blocks_override = \
+                scheduler_config.max_num_seqs
+
+        cache_config.block_size = model_config.max_model_len
+
+        logger.info(
+            "Overriding configurations based on warmup shapes. "
+            "max_model_len=%d, max_num_seqs=%d, block_size=%d, "
+            "num_gpu_blocks_override=%d",
+            model_config.max_model_len,
+            scheduler_config.max_num_seqs,
+            cache_config.block_size,
+            cache_config.num_gpu_blocks_override,
+        )
+
+        # set env vars for torch_sendnn to consume
+        os.environ["VLLM_DT_MAX_CONTEXT_LEN"] = str(
+            vllm_config.model_config.max_model_len)
+        os.environ["VLLM_DT_MAX_BATCH_SIZE"] = str(
+            vllm_config.scheduler_config.max_num_seqs)
 
-        if cache_config and model_config:
path: vllm_spyre/platform.py
user: login: prashantgupta24
body: Ever since I added the `use_all_gather` I haven't seen them None. But I can always add the conditions back to be safe 🤷 
created_at: 2025-05-02T17:39:10Z
author_association: COLLABORATOR
diff_hunk: @@ -43,97 +45,97 @@ def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        cls._config = vllm_config
         parallel_config = vllm_config.parallel_config
         scheduler_config = vllm_config.scheduler_config
         model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
 
         if scheduler_config.is_multi_step:
             raise NotImplementedError
 
-        # continuous batching related checks
-        if envs_spyre.VLLM_SPYRE_USE_CB and not envs.VLLM_USE_V1:
-            raise NotImplementedError(
-                "Continuous batching is only implemented for vLLM V1")
-
         # Near future TODO: vLLM will have an api to check whether v0 or v1 is
         # used that isn't just checking the environment variable
 
         if parallel_config.worker_cls == "auto":
-            if envs.VLLM_USE_V1:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.v1.worker.spyre_worker.SpyreWorker"
-            else:
-                parallel_config.worker_cls = \
-                    "vllm_spyre.worker.spyre_worker.SpyreWorker"
+            parallel_config.worker_cls = (
+                f'vllm_spyre{".v1" if envs.VLLM_USE_V1 else ""}'\
+                    '.worker.spyre_worker.SpyreWorker')
 
-        if envs.VLLM_USE_V1:
-            # As of 0.7.3 the scheduler for V1 isn't actually pluggable like
-            # this yet
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                scheduler_config.scheduler_cls = \
-                    "vllm_spyre.v1.core.scheduler.ContinuousBatchingSpyreScheduler"
-            else:
-                scheduler_config.scheduler_cls = \
-                    "vllm_spyre.v1.core.scheduler.StaticBatchingSpyreScheduler"
-        else:
-            scheduler_config.scheduler_cls = \
-                "vllm_spyre.core.scheduler.SpyreScheduler"
-
-        if not envs_spyre.VLLM_SPYRE_USE_CB:
+        # continuous batching related checks
+        if not envs_spyre.VLLM_SPYRE_USE_CB:  # no CB
             # Override --max-num-seqs to the biggest warmup batch size
             # And override --max-model-len to the biggest warmup sequence
             cls._warmup_shapes = None
             spyre_warmup_shapes = cls.get_warmup_shapes(scheduler_config)
             max_batch_size = 0
             max_seq_len = 0
             for shape in spyre_warmup_shapes:
-                max_batch_size = max(max_batch_size, shape['batch_size'])
+                max_batch_size = max(max_batch_size, shape["batch_size"])
                 max_seq_len = max(max_seq_len,
-                                  shape['prompt_length'] + shape['new_tokens'])
+                                  shape["prompt_length"] + shape["new_tokens"])
+            model_config.max_model_len = max_seq_len
 
-        if envs.VLLM_USE_V1:
-            if envs_spyre.VLLM_SPYRE_USE_CB:
-                # For continuous batching we use max_num_seqs to control
-                # the max batch size respecting AIU Spyre KV cache size
-                scheduler_config.max_num_seqs =\
-                    envs_spyre.VLLM_SPYRE_MAX_BATCH_SIZE
-                # ToDo: this function check_and_update_config is called twice:
-                # 1st time scheduler_config.max_num_seqs is what user sets
-                # 2nd time scheduler_config.max_num_seqs is 128
-            else:
-                # The v0 scheduler will run out of blocks if this is overridden
-                scheduler_config.max_num_seqs = max_batch_size
+            if envs.VLLM_USE_V1:  # No CB with V1
 
-        cache_config = vllm_config.cache_config
+                scheduler_config.scheduler_cls = (
+                    "vllm_spyre.v1.core.scheduler."\
+                        "StaticBatchingSpyreScheduler")
+                scheduler_config.max_num_seqs = max_batch_size
+            else:  # No CB with V0
+                scheduler_config.scheduler_cls = (
+                    "vllm_spyre.core.scheduler.SpyreScheduler")
+        else:  # CB
+            if not envs.VLLM_USE_V1:  # CB with V0
+                raise NotImplementedError(
+                    "Continuous batching is only implemented for vLLM V1")
+            else:  # CB with V1
+                # As of 0.7.3 the scheduler for V1 isn't actually pluggable like
+                # this yet
+                scheduler_config.scheduler_cls = "vllm_spyre.v1.core."\
+                        "scheduler.ContinuousBatchingSpyreScheduler"
+
+        # Cache and model config aren't set in the individual worker procs
+        # These are set in the main engine process
+
+        # To disable any paged attention ops in the base scheduler, we both:
+        # - Set the block size (in tokens) to the maximum sequence length
+        #       so that the scheduler thinks an entire sequence will fit in
+        #       one single block.
+        # - Set the number of blocks to the maximum number of sequences, so
+        #       the scheduler always thinks there's a block available
+        if envs.VLLM_USE_V1:
+            # The V1 scheduler actually needs 2 blocks for each sequence...
+            cache_config.num_gpu_blocks_override = \
+                scheduler_config.max_num_seqs * 2
+        else:
+            cache_config.num_gpu_blocks_override = \
+                scheduler_config.max_num_seqs
+
+        cache_config.block_size = model_config.max_model_len
+
+        logger.info(
+            "Overriding configurations based on warmup shapes. "
+            "max_model_len=%d, max_num_seqs=%d, block_size=%d, "
+            "num_gpu_blocks_override=%d",
+            model_config.max_model_len,
+            scheduler_config.max_num_seqs,
+            cache_config.block_size,
+            cache_config.num_gpu_blocks_override,
+        )
+
+        # set env vars for torch_sendnn to consume
+        os.environ["VLLM_DT_MAX_CONTEXT_LEN"] = str(
+            vllm_config.model_config.max_model_len)
+        os.environ["VLLM_DT_MAX_BATCH_SIZE"] = str(
+            vllm_config.scheduler_config.max_num_seqs)
 
-        if cache_config and model_config:
path: vllm_spyre/platform.py
user: login: prashantgupta24
body: Added the checks 👍 
created_at: 2025-05-02T17:45:31Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-25T19:14:21Z
body: 👋 Hi! Thank you for contributing to vLLM support on Spyre.
 Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, first install the linting requirements, then run `format.sh` and commit the changes. This can be done with `uv` directly: 
```
uv sync --frozen --group lint
```
 Or this can be done with `pip`: 
```
uv pip compile --group lint > requirements-lint.txt
pip install -r requirements-lint.txt
bash format.sh
```
Now you are good to go 🚀
author_association: NONE
user: login: joerunde
created_at: 2025-04-28T14:50:41Z
body: @prashantgupta24 There's an if check in here for `if cache_config and model_config:`, previously I had seen that on some invocations of `check_and_update_config` those were `None`. Can you check if that's still the case after fixing this issue with `use_all_gather`?

Also cc @yannicks1, I think we may want to actually leave the environment variables in place for a bit as it sounds like they're currently used by the torch stack as well
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-04-30T19:17:15Z
body: > @prashantgupta24 There's an if check in here for `if cache_config and model_config:`, previously I had seen that on some invocations of `check_and_update_config` those were `None`. Can you check if that's still the case after fixing this issue with `use_all_gather`?
> 
> Also cc @yannicks1, I think we may want to actually leave the environment variables in place for a bit as it sounds like they're currently used by the torch stack as well

@joerunde I see both not none when I put print statements, so that should solve it
```
cache config :  <vllm.config.CacheConfig object at 0x7f039597f550>
model config :  <vllm.config.ModelConfig object at 0x7f0395eca410>
```
author_association: COLLABORATOR
user: login: joerunde
created_at: 2025-05-01T14:40:20Z
body: > @joerunde I see both not none when I put print statements, so that should solve it

@prashantgupta24 would you mind removing the checks on cache and model configs then, and migrating the logic from that section to where it logically belongs above?
author_association: COLLABORATOR
user: login: joerunde
created_at: 2025-05-01T14:52:10Z
body: I think we should go one of two directions with this:

1. Delete the environment variables from `envs.py` completely, and inside `check_and_update_config` set them back in `os.environ` to the values that are in vllm_config. That way the single source of truth is the CLI args, and we control setting the environment up for `torch_sendnn` to work properly.
2. Keep the environment variables in `envs.py`, and override the `vllm_config` with their values inside `check_and_update_config`. That way the environment variables remain as the source of truth, but the rest of the code can still reference the config object instead of the environment variables.

I'd prefer going with option (1) so that our continuous batching is configured the same way as upstream vllm when running on gpus.
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-05-01T14:59:50Z
body: > I think we should go one of two directions with this:
> 
> 1. Delete the environment variables from `envs.py` completely, and inside `check_and_update_config` set them back in `os.environ` to the values that are in vllm_config. That way the single source of truth is the CLI args, and we control setting the environment up for `torch_sendnn` to work properly.
> 2. Keep the environment variables in `envs.py`, and override the `vllm_config` with their values inside `check_and_update_config`. That way the environment variables remain as the source of truth, but the rest of the code can still reference the config object instead of the environment variables.
> 
> I'd prefer going with option (1) so that our continuous batching is configured the same way as upstream vllm when running on gpus.

Yep option 1 is where I was going but then you mentioned that torch might still need them so I thought I can't do that. 

> Control setting the environment up for `torch_sendnn` to work properly

I'm not sure I fully understand how we would do that, can you elaborate? Is it by `set them back in os.environ to the values that are in vllm_config`? Basically torch just needs the 2 env vars? Can it not read `vllm_config`?
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-05-01T15:11:15Z
body: > > @joerunde I see both not none when I put print statements, so that should solve it
> 
> @prashantgupta24 would you mind removing the checks on cache and model configs then, and migrating the logic from that section to where it logically belongs above?

Let me clean up the whole logic block then, we use `VLLM_USE_V1` and `VLLM_SPYRE_USE_CB` far too many times within the function
author_association: COLLABORATOR
user: login: joerunde
created_at: 2025-05-01T15:18:38Z
body: > Yep option 1 is where I was going but then you mentioned that torch might still need them so I thought I can't do that.

> > Control setting the environment up for torch_sendnn to work properly

> I'm not sure I fully understand how we would do that, can you elaborate? Is it by set them back in os.environ to the values that are in vllm_config? Basically torch just needs the 2 env vars? Can it not read vllm_config?

well, there's no api currently to pass a `vllm_config` down to `torch_sendnn`, my simplified view of things is that it's basically just a compiler backend for `torch`- it operates on tensors and model graphs. My understanding is that it needs these two environment variables to be set so that it knows what shapes it needs to compile.

So something like this would set those environment variables for torch to use, without requiring the end user to do it:
```
def check_and_update_config(cls, vllm_config):
    os.environ["VLLM_SPYRE_MAX_CONTEXT_LENGTH"] = vllm_config.model_config.max_model_len
    os.environ["VLLM_SPYRE_MAX_BATCH_SIZE"] = vllm_config.scheduler_config.max_num_seqs
```
but, we'd have to make sure that those get set for all vllm processes so it may not be _quite_ that simple
author_association: COLLABORATOR
user: login: yannicks1
created_at: 2025-05-02T12:59:58Z
body: Yes, option (1) is the way to go here! these env variables have to be set for the AIU compiler team, but should not be consumed in vllm-spyre code. Please have a look at [this](https://github.com/vllm-project/vllm-spyre/pull/136) PR, we have agreed on another name. Thus also please remove it from `offline_inference_spyre_cb_test.py` @joerunde @prashantgupta24 
author_association: COLLABORATOR
user: login: yannicks1
created_at: 2025-05-02T13:36:49Z
body: sorry I accidentally updated your branch @prashantgupta24 , thought I was on another PR... :/
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-05-02T14:11:09Z
body: > Yes, option (1) is the way to go here! these env variables have to be set for the AIU compiler team, but should not be consumed in vllm-spyre code. Please have a look at [this](https://github.com/vllm-project/vllm-spyre/pull/136) PR, we have agreed on another name. Thus also please remove it from `offline_inference_spyre_cb_test.py` @joerunde @prashantgupta24

Done!
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-05-02T14:11:17Z
body: > sorry I accidentally updated your branch @prashantgupta24 , thought I was on another PR... :/

No problem!
author_association: COLLABORATOR
user: login: prashantgupta24
created_at: 2025-05-02T17:46:04Z
body: > Also, the code currently breaks here:
> 
> https://github.com/vllm-project/vllm-spyre/blob/c0070a3d4b6bd90f6629c0f7b78211c7c87cc7c3/vllm_spyre/v1/core/scheduler.py#L213

Fixed!
author_association: COLLABORATOR