Introducing Sim mode

sendnn_inference/envs.py

@@ -171,6 +171,21 @@ def clear_env_cache():
     "SENDNN_INFERENCE_MM_DEVICE": lambda: parse_mm_device(
         os.getenv("SENDNN_INFERENCE_MM_DEVICE", "auto")
     ),
+    # Sim-mode: replace the real Spyre forward with a no-op model and substitute
+    # virtual durations into request_metrics.jsonl. Lets us exercise scheduler /
+    # runner / batching logic on a laptop without AIU hardware. Use with
+    # DYNAMO_BACKEND=eager.
+    "SENDNN_INFERENCE_SIM_MODE": lambda: bool(int(os.getenv("SENDNN_INFERENCE_SIM_MODE", "0"))),
+    # Virtual duration (ms) charged for each prefill forward step in sim mode.
+    # Operator-supplied; default 0 means no virtual time accumulates.
+    "SENDNN_INFERENCE_SIM_PREFILL_MS": lambda: float(
+        os.getenv("SENDNN_INFERENCE_SIM_PREFILL_MS", "0")
+    ),
+    # Virtual duration (ms) charged for each decode forward step in sim mode.
+    # Operator-supplied; default 0 means no virtual time accumulates.
+    "SENDNN_INFERENCE_SIM_DECODE_MS": lambda: float(
+        os.getenv("SENDNN_INFERENCE_SIM_DECODE_MS", "0")
+    ),
 }
 # --8<-- [end:env-vars-definition]
 

sendnn_inference/platform.py

@@ -242,6 +242,18 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         if not is_decoder and not is_pooling:
             raise ValueError("Only the 'generate' and 'pooling' runners are supported")
 
+        if envs_spyre.SENDNN_INFERENCE_SIM_MODE:
+            if parallel_config.tensor_parallel_size != 1:
+                raise ValueError(
+                    "SENDNN_INFERENCE_SIM_MODE only supports tensor_parallel_size=1, "
+                    f"got {parallel_config.tensor_parallel_size}."
+                )
+            if envs_spyre.SENDNN_INFERENCE_DYNAMO_BACKEND != "eager":
+                raise ValueError(
+                    "SENDNN_INFERENCE_SIM_MODE requires SENDNN_INFERENCE_DYNAMO_BACKEND=eager, "
+                    f"got '{envs_spyre.SENDNN_INFERENCE_DYNAMO_BACKEND}'."
+                )
+
         if parallel_config.worker_cls == "auto":
             parallel_config.worker_cls = "sendnn_inference.v1.worker.spyre_worker.SpyreWorker"
 

sendnn_inference/v1/metrics/stats_logger.py

@@ -121,6 +121,12 @@ def record(
             self.iso_format
         )[:-3]
 
+        if envs_spyre.SENDNN_INFERENCE_SIM_MODE:
+            # In sim mode, virtual timings are emitted directly from the engine
+            # process to sim_metrics.jsonl. The wall-clock fields here are ~0
+            # and would be misleading; skip writing.
+            return
+
         records_to_write: list[str] = []
         for r in iteration_stats.finished_requests:
             # Calculate some estimates to add to the engine stats

sendnn_inference/v1/sim.py

@@ -0,0 +1,280 @@
+"""Sim-mode plumbing: no-op model + virtual-clock state.
+
+Activated by SENDNN_INFERENCE_SIM_MODE=1. The runner instantiates
+``MockSpyreCausalLM`` instead of ``SpyreCausalLM`` (no FMS load, no
+torch.compile, no senlib) and feeds each forward step into ``SimState``,
+which advances a virtual clock by ``SIM_PREFILL_MS`` or ``SIM_DECODE_MS``
+and accumulates per-request timing. When a request finishes, the runner
+calls ``finalize_and_write`` which appends a JSONL line of virtual stats
+to ``<perf_dir>/sim_metrics.jsonl``.
+
+A separate output file (rather than substituting into vLLM's
+request_metrics.jsonl) avoids the AsyncLLM process boundary: the
+FileStatLogger that emits request_metrics.jsonl runs in a different
+process from the runner, so it cannot see SimState. Sim mode disables
+that logger so only sim_metrics.jsonl is written.
+
+Token timestamps and ITL: each forward step advances the global virtual
+clock. We record, per request, the end-time of every prefill step and
+every decode step it participates in. The first sampled token is produced
+by the *last* prefill chunk; subsequent tokens come from each decode
+step. This gives a per-token virtual timeline and a meaningful ITL — the
+gap between two consecutive decode tokens widens whenever an intervening
+prefill of another request happens.
+"""
+
+import json
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from threading import Lock
+from types import SimpleNamespace
+
+import torch
+from vllm.config import VllmConfig
+from vllm.forward_context import get_forward_context
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
+from transformers import AutoTokenizer
+
+import sendnn_inference.envs as envs_spyre
+from sendnn_inference.model_executor.model_loader.spyre import SpyreAttentionMetadata
+
+
+# ---------------------------------------------------------------------------
+# Mock model
+# ---------------------------------------------------------------------------
+
+
+class MockSpyreCausalLM:
+    """No-op stand-in for SpyreCausalLM.
+
+    Returns dummy logits without running any real forward pass. Also used
+    by unit tests that exercise scheduler/runner logic without a real model.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+    ) -> None:
+        self.sampler = Sampler()
+
+        # boolean tensor of length batch size with indices:
+        # True for unfinished sequences and
+        # False for finished or padded sequences
+        self.indices = None
+
+        # number of right pads (relevant for continuous batching only)
+        self.n_pads_right = 0
+
+        self.vocab_size = vllm_config.model_config.get_vocab_size()
+
+        # ChunkedPrefillModelRunner.vocab_size reads .fms_model.config.src_vocab_size
+        # and .is_multimodal directly; provide minimal shims so warmup works.
+        self.is_multimodal = False
+        self.fms_model = SimpleNamespace(config=SimpleNamespace(src_vocab_size=self.vocab_size))
+
+        # These variables are here for future test scenarios to use
+        self.last_input_ids: torch.Tensor | None = None
+        self.last_positions: torch.Tensor | None = None
+        self.last_masks: torch.Tensor | None = None
+        self.last_is_prompt: bool | None = None
+        self.last_attn_metadata: SpyreAttentionMetadata | None = None
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            vllm_config.model_config.model, revision=vllm_config.model_config.revision
+        )
+        self.a_token = self.tokenizer.encode("a", add_special_tokens=False)[0]
+
+    def get_maybe_mm_embeddings(self, *args, **kwargs):
+        # This model is not multimodal
+        return None
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+    def forward(
+        self,
+        input_ids_or_embeds: torch.Tensor,
+        positions: torch.Tensor,
+        masks: torch.Tensor,
+        is_prompt: bool,
+    ) -> torch.Tensor:
+        # These variables are here for future test scenarios to use;
+        # NOTE: for now, we always use input IDs since this isn't multimodal.
+        self.last_input_ids = input_ids_or_embeds
+        self.last_positions = positions
+        self.last_masks = masks
+        self.last_is_prompt = is_prompt
+
+        forward_context = get_forward_context()
+
+        assert isinstance(forward_context.attn_metadata, SpyreAttentionMetadata)
+        self.last_attn_metadata = forward_context.attn_metadata
+
+        batch_size = input_ids_or_embeds.shape[0]
+
+        # make the logits predictable
+        logits = torch.zeros(
+            (batch_size, self.vocab_size), dtype=torch.float32, device=input_ids_or_embeds.device
+        )
+        logits[:, self.a_token] = 1
+        return logits
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> SamplerOutput | None:
+        next_tokens = self.sampler(logits, sampling_metadata)
+        return next_tokens
+
+    def set_past_key_value_states(self, num_blocks) -> None:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# Virtual-clock state
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _RequestSimRecord:
+    virtual_arrival: float
+    last_prefill_end: float | None = None
+    decode_step_ends: list[float] = field(default_factory=list)
+    virtual_completion: float | None = None
+    num_prefill_chunks: int = 0
+
+
+class SimState:
+    def __init__(self) -> None:
+        self.virtual_clock_seconds: float = 0.0
+        self._records: dict[str, _RequestSimRecord] = {}
+        self._lock = Lock()
+        self._fp = None
+
+    def _ensure_file(self):
+        if self._fp is not None:
+            return
+        out_dir = Path(envs_spyre.SENDNN_INFERENCE_PERF_METRIC_LOGGING_DIR)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        path = out_dir / "sim_metrics.jsonl"
+        if path.exists():
+            path.unlink()
+        self._fp = path.open("a", buffering=1)
+
+    def has_record(self, req_id: str) -> bool:
+        with self._lock:
+            return req_id in self._records
+
+    def record_step(
+        self,
+        is_prompt: bool,
+        prefill_ms: float,
+        decode_ms: float,
+        scheduler_output: SchedulerOutput,
+    ) -> None:
+        step_seconds = (prefill_ms if is_prompt else decode_ms) / 1000.0
+        end_t = self.virtual_clock_seconds + step_seconds
+        new_req_ids = [r.req_id for r in scheduler_output.scheduled_new_reqs]
+        cached_req_ids = list(scheduler_output.scheduled_cached_reqs.req_ids)
+
+        with self._lock:
+            for rid in new_req_ids:
+                if rid not in self._records:
+                    self._records[rid] = _RequestSimRecord(
+                        virtual_arrival=self.virtual_clock_seconds
+                    )
+
+            for rid in new_req_ids + cached_req_ids:
+                rec = self._records.get(rid)
+                if rec is None:
+                    rec = _RequestSimRecord(virtual_arrival=self.virtual_clock_seconds)
+                    self._records[rid] = rec
+                if is_prompt:
+                    rec.num_prefill_chunks += 1
+                    rec.last_prefill_end = end_t
+                else:
+                    rec.decode_step_ends.append(end_t)
+                rec.virtual_completion = end_t
+
+            self.virtual_clock_seconds = end_t
+
+    def finalize_and_write(
+        self,
+        req_id: str,
+        num_prompt_tokens: int,
+    ) -> None:
+        prefill_ms = envs_spyre.SENDNN_INFERENCE_SIM_PREFILL_MS
+        with self._lock:
+            rec = self._records.pop(req_id, None)
+        if rec is None:
+            return
+
+        # Token emit times (absolute virtual seconds): the first comes from
+        # the last prefill chunk; each subsequent from a decode step.
+        token_emit_times: list[float] = []
+        if rec.last_prefill_end is not None:
+            token_emit_times.append(rec.last_prefill_end)
+        token_emit_times.extend(rec.decode_step_ends)
+        num_generation_tokens = len(token_emit_times)
+
+        if num_generation_tokens == 0:
+            # Request never produced a token (e.g., immediate cancel). Skip.
+            return
+
+        first_token_t = token_emit_times[0]
+        last_token_t = token_emit_times[-1]
+        ttft = first_token_t - rec.virtual_arrival
+        decode_time = last_token_t - first_token_t  # bench convention
+        prefill_time = rec.num_prefill_chunks * prefill_ms / 1000.0
+
+        # ITLs between successive emitted tokens (size = num_generation_tokens - 1)
+        itls = [
+            token_emit_times[i] - token_emit_times[i - 1] for i in range(1, num_generation_tokens)
+        ]
+
+        completion = rec.virtual_completion if rec.virtual_completion is not None else last_token_t
+        e2e_latency = completion - rec.virtual_arrival
+        # In sim mode the scheduler picks a request immediately when it arrives,
+        # so there is no front-of-queue wait; report 0 for bench parity.
+        queued_time = 0.0
+        # Inference time: bench defines it as last_token_ts - scheduled_ts.
+        # We approximate scheduled_ts as virtual_arrival.
+        inference_time = last_token_t - rec.virtual_arrival
+        mean_tpot = decode_time / max(num_generation_tokens - 1, 1)
+
+        record = {
+            "timestamp": datetime.now().strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3],
+            "request_id": req_id,
+            "num_prompt_tokens": num_prompt_tokens,
+            "num_generation_tokens": num_generation_tokens,
+            "num_prefill_chunks": rec.num_prefill_chunks,
+            "num_decode_steps": len(rec.decode_step_ends),
+            "virtual_arrival_seconds": rec.virtual_arrival,
+            "virtual_completion_seconds": completion,
+            "e2e_latency_seconds": e2e_latency,
+            "queued_time_seconds": queued_time,
+            "prefill_time_seconds": prefill_time,
+            "inference_time_seconds": inference_time,
+            "decode_time_seconds": decode_time,
+            "time_to_first_token_seconds": ttft,
+            "mean_time_per_output_token_seconds": mean_tpot,
+            "inter_token_latencies_seconds": itls,
+        }
+        with self._lock:
+            self._ensure_file()
+            assert self._fp is not None
+            self._fp.write(json.dumps(record) + "\n")
+
+
+_sim_state: SimState | None = None
+
+
+def get_sim_state() -> SimState:
+    global _sim_state
+    if _sim_state is None:
+        _sim_state = SimState()
+    return _sim_state

sendnn_inference/v1/worker/spyre_model_runner.py

@@ -723,11 +723,24 @@ def __init__(
         # Initialize performance metric logger for tracking embedding times
         self.perf_logger = create_perf_metric_logger(rank=rank)
 
+        self._sim_state = None
+        self._mock_causal_lm = None
+        if envs_spyre.SENDNN_INFERENCE_SIM_MODE:
+            from sendnn_inference.v1.sim import MockSpyreCausalLM, get_sim_state
+
+            self._sim_state = get_sim_state()
+            self._mock_causal_lm = MockSpyreCausalLM
+
     def load_model(self) -> None:
-        self._model = SpyreCausalLM(
-            vllm_config=self.vllm_config,
-            rank=self.rank,
-        )
+        if envs_spyre.SENDNN_INFERENCE_SIM_MODE:
+            logger.info("SENDNN_INFERENCE_SIM_MODE=1: loading MockSpyreCausalLM (no-op forward)")
+            assert self._mock_causal_lm is not None
+            self._model = self._mock_causal_lm(vllm_config=self.vllm_config)  # ty: ignore[invalid-assignment]
+        else:
+            self._model = SpyreCausalLM(
+                vllm_config=self.vllm_config,
+                rank=self.rank,
+            )
 
     @property
     def vocab_size(self) -> int:
@@ -1461,6 +1474,15 @@ def _update_batch(self, scheduler_output: SchedulerOutput):
 
         if scheduler_output.finished_req_ids:
             for req_id in scheduler_output.finished_req_ids:
+                if self._sim_state is not None:
+                    finished_state = self.requests.get(req_id)
+                    num_prompt_tokens = (
+                        len(finished_state.prompt_token_ids) if finished_state is not None else 0
+                    )
+                    self._sim_state.finalize_and_write(
+                        req_id=req_id,
+                        num_prompt_tokens=num_prompt_tokens,
+                    )
                 self.input_batch.remove_request(req_id)
                 # TODO: Processing multiple removals at once can break alignment
                 # of logitprocs. Refactor so that we can batch removals to the
@@ -1558,6 +1580,14 @@ def execute_model(
                 is_prompt=model_input.is_prompt,
             )
 
+        if self._sim_state is not None:
+            self._sim_state.record_step(
+                is_prompt=model_input.is_prompt,
+                prefill_ms=envs_spyre.SENDNN_INFERENCE_SIM_PREFILL_MS,
+                decode_ms=envs_spyre.SENDNN_INFERENCE_SIM_DECODE_MS,
+                scheduler_output=scheduler_output,
+            )
+
         # If the prompt is being prefilled we don't have to sample
         # and generate a new token.
         if is_prefill and self.check_incomplete_prefill(scheduler_output):