From 1b5063b9b93fc16726d0744f37492458668807ad Mon Sep 17 00:00:00 2001 From: Andrewxu313 Date: Tue, 9 Jun 2026 12:26:16 +0000 Subject: [PATCH 1/2] refactor: add per-model runtime-adapter contract (M2 phase A) Add batchgen/contracts/runtime_adapter.py: frozen RuntimeState (sole worker<->model coupling point) + ModelRuntimeAdapter ABC for the behavioral leaks audited in core_model_purity_audit.md (attention-backend config, position-id computation, per-token KV byte size). Defaults cover the common GQA/non-MLA case; past_kv_byte_size is abstract. Additive only -- generic files do not call it yet (phase A). Mirrors the cuda_graph adapter pattern. Adds GPU-free contract test (6 cases). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Andrewxu313 --- batchgen/contracts/__init__.py | 7 ++ batchgen/contracts/runtime_adapter.py | 90 ++++++++++++++++++++++++++ tests/test_runtime_adapter_contract.py | 61 +++++++++++++++++ 3 files changed, 158 insertions(+) create mode 100644 batchgen/contracts/__init__.py create mode 100644 batchgen/contracts/runtime_adapter.py create mode 100644 tests/test_runtime_adapter_contract.py diff --git a/batchgen/contracts/__init__.py b/batchgen/contracts/__init__.py new file mode 100644 index 00000000..f6f62c11 --- /dev/null +++ b/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. +""" diff --git a/batchgen/contracts/runtime_adapter.py b/batchgen/contracts/runtime_adapter.py new file mode 100644 index 00000000..1f043934 --- /dev/null +++ b/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 "" 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 "" in model_type`` branches. + + Subclasses live in `batchgen/models///` 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"] diff --git a/tests/test_runtime_adapter_contract.py b/tests/test_runtime_adapter_contract.py new file mode 100644 index 00000000..9f1bedbd --- /dev/null +++ b/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 From d4ad47e2930eb08b9c2d84c96151e6b1a5b333b2 Mon Sep 17 00:00:00 2001 From: Andrewxu313 Date: Wed, 10 Jun 2026 06:38:45 +0000 Subject: [PATCH 2/2] refactor: add DeepSeek-V3 and GPT-OSS runtime adapters (M2 phase B) Per-model ModelRuntimeAdapter implementations for the two models on the generic decode path, reproducing the exact legacy inline formulas they will replace: DeepSeek-V3 (full decode position_ids, compressed_kv_dim KV byte size, flash-attn toggle) and GPT-OSS (GQA KV byte size; base defaults otherwise). Exposed via each initializer's get_runtime_behavior_adapter(). Additive only -- decode.py/prefill.py do not call them yet. Equivalence unit tests prove adapter outputs match the original decode.py formulas. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Andrewxu313 --- .../deepseekv3/deepseekv3_initializer.py | 7 +++ .../deepseek/deepseekv3/runtime_adapter.py | 37 +++++++++++++ .../gpt_oss_120b/gpt_oss_initializer.py | 7 +++ .../openai/gpt_oss_120b/runtime_adapter.py | 20 +++++++ tests/test_runtime_adapter_models.py | 55 +++++++++++++++++++ 5 files changed, 126 insertions(+) create mode 100644 batchgen/models/deepseek/deepseekv3/runtime_adapter.py create mode 100644 batchgen/models/openai/gpt_oss_120b/runtime_adapter.py create mode 100644 tests/test_runtime_adapter_models.py diff --git a/batchgen/models/deepseek/deepseekv3/deepseekv3_initializer.py b/batchgen/models/deepseek/deepseekv3/deepseekv3_initializer.py index 8a58b897..3445bf25 100644 --- a/batchgen/models/deepseek/deepseekv3/deepseekv3_initializer.py +++ b/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) diff --git a/batchgen/models/deepseek/deepseekv3/runtime_adapter.py b/batchgen/models/deepseek/deepseekv3/runtime_adapter.py new file mode 100644 index 00000000..ce5e37a0 --- /dev/null +++ b/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 + ) diff --git a/batchgen/models/openai/gpt_oss_120b/gpt_oss_initializer.py b/batchgen/models/openai/gpt_oss_120b/gpt_oss_initializer.py index e30f416f..b780f104 100644 --- a/batchgen/models/openai/gpt_oss_120b/gpt_oss_initializer.py +++ b/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. diff --git a/batchgen/models/openai/gpt_oss_120b/runtime_adapter.py b/batchgen/models/openai/gpt_oss_120b/runtime_adapter.py new file mode 100644 index 00000000..dae63311 --- /dev/null +++ b/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 + ) diff --git a/tests/test_runtime_adapter_models.py b/tests/test_runtime_adapter_models.py new file mode 100644 index 00000000..c286986e --- /dev/null +++ b/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