refactor: add per-model runtime-adapter contract + DeepSeek/GPT-OSS adapters

batchgen/contracts/__init__.py

@@ -0,0 +1,7 @@
+"""Model-support contracts the runtime core depends on.
+
+Generic runtime code (batchgen_worker.py / decode.py / prefill.py / wrappers)
+must stay model-agnostic; per-model behavior plugs in behind these contracts.
+See batchgen_design/model_architecture_spec.md (section 2.1) and
+batchgen_design/core_model_purity_audit.md.
+"""

batchgen/contracts/runtime_adapter.py

@@ -0,0 +1,90 @@
+"""Per-model runtime-behavior contract (modularization milestone M2).
+
+Mirrors `batchgen/cuda_graph/adapter.py`: a frozen `RuntimeState` is the *sole*
+coupling point between the generic runtime (`batchgen_worker.py`, `decode.py`,
+`prefill.py`) and per-model behavior. The core must not branch on a model name
+/ `model_type`; instead it calls the model's `ModelRuntimeAdapter`.
+
+This absorbs the behavioral leaks catalogued in
+`batchgen_design/core_model_purity_audit.md` (the `if "<model>" in model_type`
+branches for attention-backend config, position-id computation, and per-token
+KV byte sizing).
+
+**Phase A (this commit):** land the contract only. The generic files do not call
+it yet; per-model adapters + a dual-gated migration follow in Phase B/C
+(see `batchgen_design/blackwell/..` style Phase A/B/C from the cuda-graph work).
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional
+
+import torch
+
+
+class RuntimePhase(str, Enum):
+    PREFILL = "prefill"
+    DECODE = "decode"
+
+
+@dataclass(frozen=True)
+class RuntimeState:
+    """Snapshot the generic runtime passes to a `ModelRuntimeAdapter`.
+
+    This is the ONLY coupling surface: adapters read these fields and nothing
+    else (never reach into worker/model internals). Frozen so the invariant is
+    enforceable.
+    """
+
+    phase: RuntimePhase
+    attention_mask: Optional[torch.Tensor]
+    max_input_length: int
+    token_idx: int
+    device: Optional[torch.device] = None
+
+
+class ModelRuntimeAdapter(ABC):
+    """Per-model runtime behaviors that previously leaked into the core as
+    ``if "<model>" in model_type`` branches.
+
+    Subclasses live in `batchgen/models/<org>/<model>/` and are returned by the
+    model's initializer (``get_runtime_behavior_adapter()``). Only
+    `past_kv_byte_size` is abstract — the other two have model-agnostic defaults
+    that the common (GQA / non-MLA, no flash-attn toggle) case can use as-is.
+    """
+
+    def __init__(self, model_config: Any):
+        self.model_config = model_config
+
+    # --- attention backend (was decode.py:230, prefill.py:89/261) -----------
+    def configure_attention_backend(self, model: Any, *, phase: RuntimePhase) -> None:
+        """Default: no-op. Override to toggle e.g. ``_use_flash_attention_2``."""
+        return
+
+    # --- position ids (was decode.py:303-310 / 436-443) --------------------
+    def compute_position_ids(self, state: RuntimeState) -> torch.Tensor:
+        """Default: full ids in prefill, last-token id in decode.
+
+        MLA models (e.g. DeepSeek) override to return full ids in decode too.
+        """
+        from batchgen.utils import create_position_ids_from_attention_mask
+
+        pos = create_position_ids_from_attention_mask(state.attention_mask)
+        if state.phase == RuntimePhase.PREFILL:
+            return pos
+        return pos[:, -1].unsqueeze(-1)
+
+    # --- per-token KV byte size (was decode.py:376-411) --------------------
+    @abstractmethod
+    def past_kv_byte_size(self, state: RuntimeState) -> int:
+        """Bytes of one token's KV for this model (model-specific cache layout).
+
+        Computed from ``state.max_input_length + state.token_idx`` and the
+        adapter's ``model_config`` dimensions. There is no universal default
+        (the legacy code raised for unlisted models), so each model implements it.
+        """
+
+
+__all__ = ["RuntimePhase", "RuntimeState", "ModelRuntimeAdapter"]

batchgen/models/deepseek/deepseekv3/deepseekv3_initializer.py

@@ -250,6 +250,13 @@ def _parse_model_config(self):
         model_config.compressed_kv_dim = 576
         return model_config
 
+    def get_runtime_behavior_adapter(self):
+        """Per-model runtime behaviors (M2). See batchgen/contracts/runtime_adapter.py."""
+        from batchgen.models.deepseek.deepseekv3.runtime_adapter import (
+            DeepseekV3RuntimeAdapter,
+        )
+        return DeepseekV3RuntimeAdapter(self.model_config)
+
     def Init(self, weights_storage):
         try:
             torch.cuda.set_device(self.local_rank)

batchgen/models/deepseek/deepseekv3/runtime_adapter.py

@@ -0,0 +1,37 @@
+"""DeepSeek-V3 (MLA) runtime-behavior adapter (M2 phase B).
+
+Moves the `if "deepseek" in model_type` branches out of the generic runtime
+(decode.py / prefill.py) into this per-model adapter. See
+batchgen/contracts/runtime_adapter.py and
+batchgen_design/core_model_purity_audit.md.
+"""
+from __future__ import annotations
+
+import torch
+
+from batchgen.contracts.runtime_adapter import (
+    ModelRuntimeAdapter,
+    RuntimePhase,
+    RuntimeState,
+)
+
+
+class DeepseekV3RuntimeAdapter(ModelRuntimeAdapter):
+    def configure_attention_backend(self, model, *, phase: RuntimePhase) -> None:
+        # decode.py:230 set True in decode; prefill.py:89/261 set False in prefill.
+        model.model._use_flash_attention_2 = (phase == RuntimePhase.DECODE)
+
+    def compute_position_ids(self, state: RuntimeState) -> torch.Tensor:
+        # DeepSeek uses FULL position ids in both phases (decode.py:303-310 deepseek
+        # branch and :436-439), unlike the GQA default which slices the last token.
+        from batchgen.utils import create_position_ids_from_attention_mask
+
+        return create_position_ids_from_attention_mask(state.attention_mask)
+
+    def past_kv_byte_size(self, state: RuntimeState) -> int:
+        # decode.py:388-391 — MLA compressed KV; +1 token avoids a torch.cat in the
+        # attention forward.
+        return (
+            (state.max_input_length + state.token_idx + 1)
+            * self.model_config.compressed_kv_dim
+        )

batchgen/models/openai/gpt_oss_120b/gpt_oss_initializer.py

@@ -274,6 +274,13 @@ def _create_parameter_server(self) -> GptOss_Parameter_Server:
         logging.info("Parameter server initialized")
         return parameter_server
 
+    def get_runtime_behavior_adapter(self):
+        """Per-model runtime behaviors (M2). See batchgen/contracts/runtime_adapter.py."""
+        from batchgen.models.openai.gpt_oss_120b.runtime_adapter import (
+            GptOssRuntimeAdapter,
+        )
+        return GptOssRuntimeAdapter(self.model_config)
+
     def Init(self, weights_storage=None) -> Tuple:
         """Initialize the core engine and load weights.
 

batchgen/models/openai/gpt_oss_120b/runtime_adapter.py

@@ -0,0 +1,20 @@
+"""GPT-OSS (GQA) runtime-behavior adapter (M2 phase B).
+
+GPT-OSS uses the base defaults for position ids (last-token in decode) and
+attention backend (no flash-attn toggle); only the per-token KV byte size
+differs. See batchgen/contracts/runtime_adapter.py.
+"""
+from __future__ import annotations
+
+from batchgen.contracts.runtime_adapter import ModelRuntimeAdapter, RuntimeState
+
+
+class GptOssRuntimeAdapter(ModelRuntimeAdapter):
+    def past_kv_byte_size(self, state: RuntimeState) -> int:
+        # decode.py:400-407 — GQA: num_key_value_heads * head_dim, separate K and V.
+        return (
+            (state.max_input_length + state.token_idx + 1)
+            * self.model_config.num_key_value_heads
+            * self.model_config.head_dim
+            * 2  # K + V
+        )

tests/test_runtime_adapter_contract.py

@@ -0,0 +1,61 @@
+"""M2 Phase A: contract test for the per-model runtime adapter.
+
+GPU-free. Mirrors tests/cuda_graph_contract/test_adapter_contract.py.
+"""
+import pytest
+import torch
+
+from batchgen.contracts.runtime_adapter import (
+    ModelRuntimeAdapter,
+    RuntimePhase,
+    RuntimeState,
+)
+
+
+def test_abstract_method_set():
+    # Only past_kv_byte_size is mandatory; the rest have model-agnostic defaults.
+    assert ModelRuntimeAdapter.__abstractmethods__ == frozenset({"past_kv_byte_size"})
+
+
+def test_runtime_state_is_frozen():
+    s = RuntimeState(
+        phase=RuntimePhase.DECODE, attention_mask=None, max_input_length=8, token_idx=0
+    )
+    with pytest.raises(Exception):
+        s.token_idx = 1  # frozen dataclass -> FrozenInstanceError
+
+
+def test_phase_values():
+    assert RuntimePhase.PREFILL == "prefill"
+    assert RuntimePhase.DECODE == "decode"
+
+
+class _DummyAdapter(ModelRuntimeAdapter):
+    def past_kv_byte_size(self, state: RuntimeState) -> int:
+        return (state.max_input_length + state.token_idx) * 4
+
+
+def test_default_position_ids_shapes():
+    mask = torch.ones(2, 5, dtype=torch.long)
+    a = _DummyAdapter(model_config=None)
+    pre = a.compute_position_ids(
+        RuntimeState(RuntimePhase.PREFILL, mask, max_input_length=5, token_idx=0)
+    )
+    dec = a.compute_position_ids(
+        RuntimeState(RuntimePhase.DECODE, mask, max_input_length=5, token_idx=0)
+    )
+    assert tuple(pre.shape) == (2, 5)   # prefill: full
+    assert tuple(dec.shape) == (2, 1)   # decode: last token only
+
+
+def test_default_attention_backend_is_noop():
+    a = _DummyAdapter(model_config=None)
+    a.configure_attention_backend(object(), phase=RuntimePhase.DECODE)  # must not raise
+
+
+def test_past_kv_byte_size_uses_state():
+    a = _DummyAdapter(model_config=None)
+    n = a.past_kv_byte_size(
+        RuntimeState(RuntimePhase.DECODE, None, max_input_length=10, token_idx=3)
+    )
+    assert n == (10 + 3) * 4

tests/test_runtime_adapter_models.py

@@ -0,0 +1,55 @@
+"""M2 phase B: per-model runtime adapters reproduce the exact legacy inline
+formulas (decode.py / prefill.py) they replace. GPU-free.
+"""
+import torch
+
+from batchgen.contracts.runtime_adapter import RuntimePhase, RuntimeState
+from batchgen.models.deepseek.deepseekv3.runtime_adapter import DeepseekV3RuntimeAdapter
+from batchgen.models.openai.gpt_oss_120b.runtime_adapter import GptOssRuntimeAdapter
+
+
+class _Cfg:
+    def __init__(self, **kw):
+        self.__dict__.update(kw)
+
+
+class _Model:
+    class model:
+        _use_flash_attention_2 = None
+
+
+# ---- DeepSeek-V3 (MLA): all three behaviors -------------------------------
+def test_deepseek_past_kv_byte_size_matches_legacy():
+    a = DeepseekV3RuntimeAdapter(_Cfg(compressed_kv_dim=576))
+    s = RuntimeState(RuntimePhase.DECODE, None, max_input_length=100, token_idx=5)
+    assert a.past_kv_byte_size(s) == (100 + 5 + 1) * 576  # decode.py:388-391
+
+
+def test_deepseek_position_ids_are_full():
+    a = DeepseekV3RuntimeAdapter(_Cfg())
+    mask = torch.ones(2, 5, dtype=torch.long)
+    out = a.compute_position_ids(RuntimeState(RuntimePhase.DECODE, mask, 5, 0))
+    assert tuple(out.shape) == (2, 5)  # full, not last-token (2,1)
+
+
+def test_deepseek_flash_attention_toggle():
+    a = DeepseekV3RuntimeAdapter(_Cfg())
+    m = _Model()
+    a.configure_attention_backend(m, phase=RuntimePhase.DECODE)
+    assert m.model._use_flash_attention_2 is True   # decode.py:230
+    a.configure_attention_backend(m, phase=RuntimePhase.PREFILL)
+    assert m.model._use_flash_attention_2 is False  # prefill.py:89/261
+
+
+# ---- GPT-OSS (GQA): KV byte size only -------------------------------------
+def test_gpt_oss_past_kv_byte_size_matches_legacy():
+    a = GptOssRuntimeAdapter(_Cfg(num_key_value_heads=8, head_dim=64))
+    s = RuntimeState(RuntimePhase.DECODE, None, max_input_length=100, token_idx=5)
+    assert a.past_kv_byte_size(s) == (100 + 5 + 1) * 8 * 64 * 2  # decode.py:400-407
+
+
+def test_gpt_oss_position_ids_default_last_token():
+    a = GptOssRuntimeAdapter(_Cfg())
+    mask = torch.ones(2, 5, dtype=torch.long)
+    out = a.compute_position_ids(RuntimeState(RuntimePhase.DECODE, mask, 5, 0))
+    assert tuple(out.shape) == (2, 1)  # GQA default: last token only