refactor: split fp8 into fp8_ops adapter and fp8 policy module

- fp8_ops is the only module touching the pybind (kernel interface)
- fp8.py keeps scaling state, delayed amax and aten::linear dispatch
- remove circular imports between old fp8_ops/fp8_state/fp8_dispatch

.github/workflows/release.yml

@@ -54,6 +54,9 @@ jobs:
       - name: Build wheel (with CUDA kernels)
         run: |
           CSRC_KERNELS=true pip wheel . --no-deps --no-build-isolation -w dist/
+          for f in dist/*.whl; do
+            mv "$f" "dist/$(basename "$f" .whl)+${{ matrix.cuda_tag }}.whl"
+          done
 
       - uses: actions/upload-artifact@v4
         with:

Dockerfile

@@ -57,8 +57,13 @@ COPY docs/ ./docs/
 COPY pyproject.toml .
 COPY README.md .
 
-# Create non-root user
-RUN useradd -m astrai && chown -R astrai:astrai /app
+# Create non-root user matching the host uid/gid (passed via build args)
+ARG USER_UID=1000
+ARG USER_GID=1000
+RUN groupadd -g "${USER_GID}" astrai \
+    && useradd -m -u "${USER_UID}" -g astrai astrai \
+    && chown -R astrai:astrai /app
+ENV HOME=/home/astrai
 USER astrai
 
 ENV PYTHONUNBUFFERED=1 \

astrai/__init__.py

@@ -1,9 +1,6 @@
-__version__ = "1.3.12"
+__version__ = "1.3.13"
 __author__ = "ViperEkura"
 
-import logging
-import os
-
 from astrai.config import (
     AutoRegressiveLMConfig,
     BaseModelConfig,
@@ -28,6 +25,7 @@ from astrai.inference import (
     run_server,
     sample,
 )
+from astrai.logging import setup_logging
 from astrai.model import (
     AutoModel,
     AutoRegressiveLM,
@@ -55,30 +53,6 @@ from astrai.trainer import (
     Trainer,
 )
 
-
-def setup_logging(level: str = "INFO"):
-    """Attach a handler to the ``astrai`` logger (only, not root).
-
-    Call once per process, e.g. at the top of CLI scripts.
-    Set ``ASTR_LOG_LEVEL`` to override the default ``INFO``.
-    """
-    _logger = logging.getLogger("astrai")
-    if _logger.handlers:
-        return
-    _level = getattr(
-        logging, os.environ.get("ASTR_LOG_LEVEL", level).upper(), logging.INFO
-    )
-    _logger.setLevel(_level)
-    _handler = logging.StreamHandler()
-    _handler.setFormatter(
-        logging.Formatter(
-            "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
-            datefmt="%Y-%m-%d %H:%M:%S",
-        )
-    )
-    _logger.addHandler(_handler)
-
-
 __all__ = [
     "AutoRegressiveLM",
     "AutoRegressiveLMConfig",
@@ -122,3 +96,5 @@ __all__ = [
     "setup_logging",
     "spawn_parallel_fn",
 ]
+
+setup_logging()

astrai/config/train_config.py

@@ -48,6 +48,7 @@ class TrainConfig(BaseConfig):
         random_seed (int): Random seed. Defaults to 3407.
         num_workers (int): Number of workers for dataloader. Defaults to 0.
         prefetch_factor (Optional[int]): Prefetch factor for dataloader. Defaults to None.
+        persistent_workers (bool): Keep DataLoader workers alive between epochs. Defaults to False.
         pin_memory (bool): Pin memory for dataloader. Defaults to False.
         collate_fn (Optional[Callable[[List[Any]], Any]]): Collate function for dataloader (e.g. dpo_collate_fn). Defaults to None.
         nprocs (int): Number of processes for distributed training. Defaults to 1.
@@ -98,6 +99,7 @@ class TrainConfig(BaseConfig):
     random_seed: int = 3407
     num_workers: int = 0
     prefetch_factor: Optional[int] = None
+    persistent_workers: bool = False
     pin_memory: bool = False
     collate_fn: Optional[Callable[[List[Any]], Any]] = None
 

astrai/dataset/__init__.py

@@ -14,6 +14,7 @@ from astrai.dataset.storage import (
     Streamable,
     detect_format,
 )
+from astrai.dataset.streaming import StreamingSeqDataset
 from astrai.serialization import (
     load_bin,
     save_bin,
@@ -34,4 +35,5 @@ __all__ = [
     "save_bin",
     "load_bin",
     "RDSampler",
+    "StreamingSeqDataset",
 ]

astrai/dataset/dataset.py

@@ -25,20 +25,50 @@ function (pure ``record -> Dict[str, Tensor]``) is forwarded to
 
 from abc import ABC, abstractmethod
 from functools import partial
+from pathlib import Path
 from typing import Callable, Dict, List, Optional
 
 import torch
 from torch import Tensor
 from torch.utils.data import Dataset
 
+from astrai.config.preprocess_config import PipelineConfig
 from astrai.dataset.storage import (
     Store,
     StoreFactory,
     detect_format,
 )
 from astrai.factory import BaseFactory
+from astrai.preprocessing.transform import TokenizeTransform
 from astrai.tokenize import AutoTokenizer
 
+_DEFAULT_MESSAGES_CONFIG = {
+    "version": 1,
+    "input": {"sections": [{"field": "messages", "action": "$role", "template": True}]},
+    "mask": {"system": "mask", "user": "mask", "assistant": "train"},
+    "mask_default": "mask",
+    "output": {"position_ids_mode": "doc_reset"},
+}
+
+
+def _build_jsonl_transform(
+    path: str, tokenizer_path: Optional[str] = None
+) -> Optional["TokenizeTransform"]:
+    """Auto-build a TokenizeTransform for JSONL eager loading.
+
+    Reads ``dataset_config.json`` from the data dir if present, or
+    falls back to the built-in chatml SFT config when *tokenizer_path*
+    is provided.
+    """
+    root = Path(path)
+    config_path = root / "dataset_config.json" if root.is_dir() else None
+    if config_path is not None and config_path.exists():
+        return TokenizeTransform.from_config_file(str(config_path))
+    if tokenizer_path:
+        config = PipelineConfig.from_dict(_DEFAULT_MESSAGES_CONFIG)
+        return TokenizeTransform(config, tokenizer_path)
+    return None
+
 
 def dpo_tokenize(
     record: dict,
@@ -349,16 +379,18 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
         )
         if processor is not None:
             store.load(load_path, processor=processor, **kwargs)
+        elif storage_type == "jsonl":
+            transform = _build_jsonl_transform(load_path, tokenizer_path)
+            if transform is None:
+                raise FileNotFoundError(
+                    f"JSONL dataset config not found. Expected "
+                    f"dataset_config.json alongside *.jsonl files, pass "
+                    f"tokenizer_path= for the built-in messages config, or "
+                    f"use processor= for lazy on-the-fly tokenisation."
+                )
+            store.load(load_path, transform=transform, **kwargs)
         else:
-            load_kwargs = dict(kwargs)
-            if (
-                tokenizer_path is not None
-                and storage_type == "jsonl"
-                and train_type in ("seq", "sft")
-                and "tokenizer_path" not in load_kwargs
-            ):
-                load_kwargs["tokenizer_path"] = tokenizer_path
-            store.load(load_path, **load_kwargs)
+            store.load(load_path, **kwargs)
 
         return cls.create(train_type, store=store)
 

astrai/dataset/storage.py

@@ -55,9 +55,7 @@ from typing import Callable, Dict, List, Optional, Tuple, Union
 import torch
 from torch import Tensor
 
-from astrai.config.preprocess_config import PipelineConfig
 from astrai.factory import BaseFactory
-from astrai.preprocessing.transform import TokenizeTransform
 from astrai.serialization import (
     load_bin,
     load_bin_offsets,
@@ -536,19 +534,8 @@ class JsonlStore(Store, Streamable, Recordable):
       ``len(store)`` returns ``num_records``; stream primitives raise.
     """
 
-    CONFIG_NAME = "dataset_config.json"
     segments_are_records = True
 
-    _DEFAULT_MESSAGES_CONFIG = {
-        "version": 1,
-        "input": {
-            "sections": [{"field": "messages", "action": "$role", "template": True}]
-        },
-        "mask": {"system": "mask", "user": "mask", "assistant": "train"},
-        "mask_default": "mask",
-        "output": {"position_ids_mode": "doc_reset"},
-    }
-
     def __init__(
         self,
         window_size: int = 0,
@@ -569,22 +556,10 @@ class JsonlStore(Store, Streamable, Recordable):
             return
 
         if transform is None:
-            root = Path(path)
-            config_path = root / self.CONFIG_NAME if root.is_dir() else None
-            if config_path is not None and config_path.exists():
-                transform = TokenizeTransform.from_config_file(str(config_path))
-            else:
-                tokenizer_path = kwargs.get("tokenizer_path")
-                if not tokenizer_path:
-                    raise FileNotFoundError(
-                        f"JSONL dataset config not found. Expected "
-                        f"{self.CONFIG_NAME} alongside *.jsonl files, pass an "
-                        f"explicit transform, pass processor= for lazy "
-                        f"on-the-fly tokenisation, or pass tokenizer_path= to "
-                        f"use the built-in messages config."
-                    )
-                config = PipelineConfig.from_dict(self._DEFAULT_MESSAGES_CONFIG)
-                transform = TokenizeTransform(config, tokenizer_path)
+            raise ValueError(
+                "JsonlStore eager mode requires transform=. "
+                "Use DatasetFactory.load() which auto-constructs it."
+            )
 
         transformed = transform.apply(records)
         self._normalize(transformed)

astrai/dataset/streaming.py

@@ -0,0 +1,122 @@
+"""Streaming IterableDataset for pre-training with shard-level shuffle.
+
+Unlike the map-style datasets, the streaming dataset yields windows
+sequentially through each data shard — no random access, no sampler.
+Each DataLoader worker independently streams its assigned shard subset,
+giving better OS page-cache locality for large-scale (TB+) datasets.
+
+Key properties:
+- Implements ``torch.utils.data.IterableDataset``.
+- ``__len__`` returns total window count so ``compute_total_steps`` works.
+- Shard-level shuffle with deterministic seed (reproducible across runs).
+- Distributed: each rank gets a disjoint subset of shards.
+- Multi-worker: each worker within a rank gets a disjoint subset.
+"""
+
+import random
+from typing import Iterator, Optional
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+from torch.utils.data import IterableDataset
+
+from astrai.dataset.storage import Store
+
+
+def _resolve_rank_and_world_size() -> tuple[int, int]:
+    if dist.is_available() and dist.is_initialized():
+        return dist.get_rank(), dist.get_world_size()
+    return 0, 1
+
+
+def _total_windows(token_count, window_size, stride):
+    if token_count <= window_size:
+        return 0
+    return (token_count - 1 - window_size) // stride + 1
+
+
+class StreamingSeqDataset(IterableDataset):
+    """Streaming next-token prediction dataset.
+
+    Yields ``{"input_ids": [L], "target_ids": [L]}`` dicts by sliding a
+    window sequentially through each data shard.  Shards are shuffled
+    deterministically.  Distributed and multi-worker DataLoader modes are
+    supported: each consumer gets a disjoint shard subset.
+
+    Args:
+        store: Already-loaded Store with a ``"sequence"`` key.
+        window_size: Context length per sample.
+        stride: Step between consecutive windows (default: window_size).
+        shuffle: Shuffle shard order.
+        seed: Base seed for deterministic shard shuffle.
+    """
+
+    def __init__(
+        self,
+        store: Store,
+        window_size: int,
+        stride: Optional[int] = None,
+        shuffle: bool = True,
+        seed: int = 42,
+        rank: Optional[int] = None,
+        world_size: Optional[int] = None,
+    ):
+        super().__init__()
+        if window_size <= 0:
+            raise ValueError("window_size must be positive")
+        self.store = store
+        self.window_size = window_size
+        self.stride = stride if stride is not None else window_size
+        self.shuffle = shuffle
+        self.seed = seed
+        self._rank, self._world_size = (
+            rank,
+            world_size if rank is not None else _resolve_rank_and_world_size(),
+        )
+
+        if "sequence" not in store.keys:
+            raise KeyError(
+                f"Store is missing required key 'sequence'; "
+                f"available keys: {sorted(store.keys)}"
+            )
+
+    @property
+    def num_samples(self) -> int:
+        return _total_windows(self.store.token_count, self.window_size, self.stride)
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def __iter__(self) -> Iterator[dict[str, Tensor]]:
+        segments = self.store._data["sequence"]
+        n_shards = len(segments)
+
+        indices = list(range(n_shards))
+        if self.shuffle:
+            rng = random.Random(self.seed)
+            rng.shuffle(indices)
+
+        worker_info = torch.utils.data.get_worker_info()
+        if worker_info is None:
+            num_consumers = self._world_size
+            consumer_id = self._rank
+        else:
+            num_consumers = self._world_size * worker_info.num_workers
+            consumer_id = self._rank * worker_info.num_workers + worker_info.id
+
+        my_shards = [
+            i for idx, i in enumerate(indices) if idx % num_consumers == consumer_id
+        ]
+
+        for shard_idx in my_shards:
+            segment = segments[shard_idx]
+            seq_len = segment.shape[0]
+            for begin in range(0, seq_len - self.window_size, self.stride):
+                end = begin + self.window_size
+                yield {
+                    "input_ids": torch.as_tensor(segment[begin:end], dtype=torch.long),
+                    "target_ids": torch.as_tensor(
+                        segment[begin + 1 : end + 1], dtype=torch.long
+                    ),
+                }

astrai/extension/attention_backend.py

@@ -34,6 +34,7 @@ import enum
 import functools
 import importlib
 import os
+import threading
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Optional, Union
@@ -50,10 +51,15 @@ from astrai.extension.loader import is_available
 from astrai.factory import BaseFactory
 
 if TYPE_CHECKING:
-    from astrai.inference.core.cache import KVCache
+    from astrai.inference.cache import KVCache
 
-_current_backend: contextvars.ContextVar["AttentionBackend"] = contextvars.ContextVar(
-    "attn_backend"
+
+_default_backend: Optional["AttentionBackend"] = None
+_default_backend_lock = threading.Lock()
+_env_backend_name: Optional[str] = None
+_env_backend: Optional["AttentionBackend"] = None
+_current_backend: contextvars.ContextVar[Optional["AttentionBackend"]] = (
+    contextvars.ContextVar("attn_backend", default=None)
 )
 
 
@@ -100,9 +106,6 @@ class ATTN_BACKEND(enum.Enum):
     FLASH = "flash"
 
 
-_default_backend: Optional["AttentionBackend"] = None
-
-
 def _priority_backends() -> list["AttentionBackend"]:
     """Available backends in priority order: cuda -> flash -> torch."""
     backends: list[AttentionBackend] = []
@@ -135,46 +138,83 @@ def _backend_supports(
     if isinstance(backend, FlashAttnBackend):
         if not flash_attn_available():
             return False
+        if q.dtype not in (torch.float16, torch.bfloat16):
+            return False
         if q.size(1) == 1 and kv_cache is not None:
             return True
-        return not (attn_mask is not None and not is_causal)
+        if attn_mask is None or is_causal:
+            return True
+        return attn_mask.dim() == 4
     return True
 
 
 def _resolve_default_backend() -> "AttentionBackend":
     """Pick the highest-priority available backend (cuda -> flash -> torch).
 
-    Set ``ASTR_BACKEND`` to override: ``ASTR_BACKEND=cuda``, ``torch_native``,
-    or ``flash``.  The value is the registered name (same as the
-    ``ATTN_BACKEND`` enum value).
-
     Resolved lazily on first ``get_backend()`` and cached.  Per-call
     capability fallback happens in ``attention()``, so the default is
     safe for training and fp32 models.
     """
-    forced = os.environ.get("ASTR_BACKEND", "").strip().lower()
-    if forced:
-        try:
-            return AttentionBackendFactory.create(forced)
-        except (ValueError, RuntimeError):
-            pass
     return _priority_backends()[0]
 
 
-def get_backend() -> "AttentionBackend":
-    """Return the active backend for the current thread/context.
+def _environment_backend() -> Optional["AttentionBackend"]:
+    """Resolve the process-wide ``ASTR_BACKEND`` override, if configured."""
+    global _env_backend, _env_backend_name
+    name = os.environ.get("ASTR_BACKEND", "").strip().lower()
+    if not name:
+        return None
+    if name != _env_backend_name:
+        with _default_backend_lock:
+            if name != _env_backend_name:
+                try:
+                    _env_backend = AttentionBackendFactory.create(name)
+                except (ValueError, RuntimeError):
+                    _env_backend = None
+                _env_backend_name = name
+    return _env_backend
 
-    Falls back to the highest-priority available backend (cuda -> flash ->
-    torch_native) when no backend has been activated via ``with``.  Set
-    ``ASTR_BACKEND`` to override the default.
+
+def _resolve_backend(
+    backend: Optional[Union[str, ATTN_BACKEND, "AttentionBackend", type]] = None,
+) -> "AttentionBackend":
+    """Resolve a backend configuration, defaulting to the process policy."""
+    if backend is not None:
+        if isinstance(backend, ATTN_BACKEND):
+            return AttentionBackendFactory.create(backend.value)
+        if isinstance(backend, str):
+            return AttentionBackendFactory.create(backend)
+        if isinstance(backend, type) and issubclass(backend, AttentionBackend):
+            return backend()
+        if isinstance(backend, AttentionBackend):
+            return backend
+        raise TypeError(
+            f"expected a registered name, ATTN_BACKEND, AttentionBackend type, "
+            f"or instance, got {type(backend).__name__}"
+        )
+
+    global _default_backend
+    if _default_backend is None:
+        with _default_backend_lock:
+            if _default_backend is None:
+                _default_backend = _resolve_default_backend()
+    return _default_backend
+
+
+def get_backend(
+    use_default: bool = True,
+) -> Optional["AttentionBackend"]:
+    """Return the context override, optionally falling back to the process default.
+
+    ``ASTR_BACKEND`` is a process-wide override and takes precedence over the
+    context value. Pass ``use_default=False`` at request submission to retain
+    only an environment override or the caller's :func:`attn_backend` value.
     """
-    try:
-        return _current_backend.get()
-    except LookupError:
-        global _default_backend
-        if _default_backend is None:
-            _default_backend = _resolve_default_backend()
-        return _default_backend
+    return (
+        _environment_backend()
+        or _current_backend.get()
+        or (_resolve_backend() if use_default else None)
+    )
 
 
 @contextmanager
@@ -193,20 +233,7 @@ def attn_backend(backend: Union[str, ATTN_BACKEND, "AttentionBackend", type]):
         with attn_backend(TorchNativeBackend()):
             ...
     """
-    if isinstance(backend, ATTN_BACKEND):
-        instance = AttentionBackendFactory.create(backend.value)
-    elif isinstance(backend, str):
-        instance = AttentionBackendFactory.create(backend)
-    elif isinstance(backend, type) and issubclass(backend, AttentionBackend):
-        instance = backend()
-    elif isinstance(backend, AttentionBackend):
-        instance = backend
-    else:
-        raise TypeError(
-            f"expected a registered name, ATTN_BACKEND, AttentionBackend type, "
-            f"or instance, "
-            f"got {type(backend).__name__}"
-        )
+    instance = _resolve_backend(backend)
     token = _current_backend.set(instance)
     try:
         yield instance
@@ -277,9 +304,15 @@ def attention(
     """
     backend = get_backend()
     if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal):
-        # The active backend cannot run this call (e.g. CUDA on a training /
-        # fp32 / unsupported-head_dim input) — fall back to the highest-
-        # priority backend that can, ending at torch SDPA.
+        explicit = get_backend(use_default=False)
+        if explicit is not None:
+            raise RuntimeError(
+                f"Explicitly-set backend {type(backend).__name__} cannot "
+                f"handle this attention call (shape={q.shape}, "
+                f"dtype={q.dtype}, kv_cache={'none' if kv_cache is None else 'present'}, "
+                f"attn_mask={'none' if attn_mask is None else 'present'}). "
+                f"Remove the attn_backend() context or switch to a compatible backend."
+            )
         for candidate in _priority_backends():
             if isinstance(candidate, type(backend)):
                 continue
@@ -512,7 +545,6 @@ class CudaBackend(AttentionBackend):
             kv_cache.req_to_token,
             kv_cache.req_pool_indices,
             kv_indptr,
-            kv_cache.max_len,
             is_causal=True,
             o_part_buf=kv_cache.decode_o_part,
             ml_part_buf=kv_cache.decode_ml_part,
@@ -558,7 +590,6 @@ class CudaBackend(AttentionBackend):
             kv_indptr,
             qo_indptr,
             attn_mask,
-            q_len,
             is_causal=is_causal,
         )
         return out.reshape(b, q_len, q.size(2), q.size(3)).flatten(2)
@@ -626,7 +657,7 @@ class FlashAttnBackend(AttentionBackend):
             k = repeat_kv(k, n_rep)
             v = repeat_kv(v, n_rep)
 
-        if attn_mask is not None and not is_causal:
+        if attn_mask is not None and not is_causal and attn_mask.dim() != 4:
             raise ValueError(
                 "FlashAttnBackend does not support a custom attention mask; "
                 "use a causal mask or select TorchNativeBackend."
@@ -638,7 +669,10 @@ class FlashAttnBackend(AttentionBackend):
                 "Install with `pip install flash-attn`."
             )
         out = fa.flash_attn_func(
-            q.contiguous(), k.contiguous(), v.contiguous(), causal=is_causal
+            q.contiguous(),
+            k.contiguous(),
+            v.contiguous(),
+            causal=is_causal or (attn_mask is not None and attn_mask.dim() == 4),
         )
         return out.contiguous().flatten(2)
 

astrai/extension/attention_ops.py

@@ -97,7 +97,6 @@ def attn_paged_decode(
     req_to_token: torch.Tensor,
     req_pool_indices: torch.Tensor,
     kv_indptr: torch.Tensor,
-    max_seq_len: int,
     mask: Optional[torch.Tensor] = None,
     is_causal: bool = False,
     o_part_buf: Optional[torch.Tensor] = None,
@@ -117,8 +116,7 @@ def attn_paged_decode(
         req_to_token: [num_reqs, max_context_len] (int64) — token -> slot
         req_pool_indices: [batch] (int64) — rows into req_to_token
         kv_indptr: [batch+1] (int32) — prefix sum of per-request seq_lens
-        max_seq_len: max per-request seq_len (Python int, for split computation)
-        mask: 2D [batch, max_seq_len] (bool, True=keep) or None
+        mask: 2D [batch, max_context_len] (bool, True=keep) or None
         is_causal: apply causal mask
         o_part_buf: pre-allocated split-KV o partial buffer (workflow bypass)
         ml_part_buf: pre-allocated split-KV m/l buffer (workflow bypass)
@@ -136,7 +134,6 @@ def attn_paged_decode(
         req_to_token,
         req_pool_indices,
         kv_indptr,
-        max_seq_len,
         mask=mask,
         causal_offset=causal_offset,
         o_part_buf=o_part_buf,
@@ -154,7 +151,6 @@ def attn_paged_prefill(
     kv_indptr: torch.Tensor,
     qo_indptr: torch.Tensor,
     mask: Optional[torch.Tensor] = None,
-    max_q_len: int = 0,
     is_causal: bool = False,
 ) -> torch.Tensor:
     """SGLang-style paged prefill (ragged batch, flat KV pool).
@@ -172,7 +168,6 @@ def attn_paged_prefill(
         kv_indptr: [batch+1] (int32) — prefix sum of per-request kv_lens
         qo_indptr: [batch+1] (int32) — prefix sum of per-request q_lens
         mask: 4D [batch, 1, q_len, kv_len] (bool, True=keep) or None
-        max_q_len: max per-request q_len (Python int, for grid computation)
         is_causal: apply causal mask
 
     Returns:
@@ -189,6 +184,5 @@ def attn_paged_prefill(
         kv_indptr,
         qo_indptr,
         mask,
-        max_q_len,
         causal_offset=causal_offset,
     )

astrai/extension/fp8.py

@@ -0,0 +1,284 @@
+"""FP8 training: scaling state and aten::linear dispatch.
+
+Layered (see also ``fp8_ops.py`` for the CUDA interface adapter):
+
+1. Kernel interface: "fp8_ops" — the only module touching the pybind.
+2. Training state (this module): per-tensor scales, amax history, delayed
+   scaling, and the ``fp8_autocast`` context (TE-style, like
+   ``torch.autocast``).
+3. aten::linear integration (this module): registers the CUDA impl and the
+   M/N alignment guard.
+
+Usage::
+
+    from astrai.extension.fp8 import fp8_autocast
+
+    with fp8_autocast(enabled=True):
+        logits = model(input_ids)
+        loss.backward()
+
+Importing this module registers the aten::linear CUDA implementation.
+"""
+
+from contextlib import contextmanager
+
+import torch
+from torch.library import Library
+
+from astrai.extension.fp8_ops import (
+    linear_backward_scaled,
+    linear_forward_scaled,
+)
+
+E4M3_MAX = 448.0
+
+
+# ---------------------------------------------------------------------------
+# Layer 2: training state (scales, amax history, delayed scaling, autocast)
+# ---------------------------------------------------------------------------
+
+
+class FP8TensorMeta:
+    """Scales + amax state for one weight tensor and its paired activations.
+
+    - weight: delayed scale from a 16-step amax history window (TE style)
+    - x/g: delayed one step, reuse the quantize kernel's free atomic amax
+    """
+
+    __slots__ = (
+        "scale",
+        "scale_inv",
+        "amax_history",
+        "idx",
+        "x_scale",
+        "x_scale_inv",
+        "g_scale",
+        "g_scale_inv",
+    )
+
+    def __init__(self, device: torch.device, update_interval: int):
+        self.scale = torch.ones(1, device=device, dtype=torch.float32)
+        self.scale_inv = torch.ones(1, device=device, dtype=torch.float32)
+        self.amax_history = torch.ones(
+            update_interval, device=device, dtype=torch.float32
+        )
+        self.idx = 0
+        self.x_scale = torch.ones(1, device=device, dtype=torch.float32)
+        self.x_scale_inv = torch.ones(1, device=device, dtype=torch.float32)
+        self.g_scale = torch.ones(1, device=device, dtype=torch.float32)
+        self.g_scale_inv = torch.ones(1, device=device, dtype=torch.float32)
+
+    def record(self, amax: torch.Tensor) -> None:
+        """Push the latest amax into the ring buffer (device-side copy, no sync)."""
+        self.amax_history[self.idx] = amax.reshape(())
+        self.idx = (self.idx + 1) % self.amax_history.numel()
+
+    def refresh(self) -> None:
+        """Recompute scale from the amax history window (delayed scaling)."""
+        amax = self.amax_history.max()
+        if amax > 0:
+            self.scale.copy_(amax / E4M3_MAX)
+            self.scale_inv.copy_(E4M3_MAX / amax)
+
+
+class FP8State:
+    """Global fp8 training state, TE-style."""
+
+    def __init__(self, update_interval: int = 16):
+        self.enabled = False
+        self.update_interval = update_interval
+        self.step_count = 0
+        self._metas: dict[tuple, FP8TensorMeta] = {}
+        self._last_device: torch.device | None = None
+
+    def _get_device(self, t: torch.Tensor) -> torch.device:
+        if self._last_device is None:
+            self._last_device = t.device
+        return t.device
+
+    def get_weight_meta(self, w: torch.Tensor) -> FP8TensorMeta:
+        key = (w.data_ptr(), w.shape, w.dtype)
+        meta = self._metas.get(key)
+        if meta is None:
+            meta = FP8TensorMeta(self._get_device(w), self.update_interval)
+            self._metas[key] = meta
+        return meta
+
+    def step(self) -> None:
+        """Advance the counter and refresh all weight scales every N steps."""
+        self.step_count += 1
+        if self.step_count % self.update_interval == 0:
+            for meta in self._metas.values():
+                meta.refresh()
+
+    def reset(self) -> None:
+        self.enabled = False
+        self.step_count = 0
+        self._metas.clear()
+        self._last_device = None
+
+
+# Global singleton: autograd backward runs on the engine worker threads, so
+# thread-local state would lose the fp8 flag during loss.backward(). The GIL
+# protects Python-side mutation; the CUDA kernels take their own mutex.
+_state = FP8State()
+
+
+def fp8_state() -> FP8State:
+    return _state
+
+
+@contextmanager
+def fp8_autocast(enabled: bool = True, update_interval: int = 16):
+    """Autocast-style context: fp8 linear dispatch on this thread.
+
+    Usage::
+
+        with fp8_autocast(enabled=True):
+            logits = model(input_ids)   # aten::linear -> fp8 path
+            loss.backward()
+
+    The scale-update counter advances once per ``enter`` (one training step),
+    refreshing weight scales from their amax history every ``update_interval``.
+    """
+    state = fp8_state()
+    prev_enabled = state.enabled
+    prev_interval = state.update_interval
+    state.enabled = enabled
+    state.update_interval = update_interval
+    try:
+        if enabled:
+            state.step()
+        yield
+    finally:
+        state.enabled = prev_enabled
+        state.update_interval = prev_interval
+
+
+def _update_delayed_scale(scale, scale_inv, amax) -> None:
+    """scale = amax / 448 for the *next* call (device-side, no sync)."""
+    amax_f = amax.reshape(()).to(torch.float32).clamp_min(1e-12)
+    scale.copy_(amax_f / E4M3_MAX)
+    scale_inv.copy_(E4M3_MAX / amax_f)
+
+
+def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None):
+    """TE-style scaled fp8 linear forward (called from the aten::linear impl).
+
+    x uses the delayed scale of its paired weight meta (amax from the previous
+    forward of this linear); the quantize kernel emits the current amax for the
+    next step. No extra abs/max reduce.
+    """
+    if bias is None:
+        bias = torch.empty(0, device=x.device, dtype=x.dtype)
+    state = fp8_state()
+    meta = state.get_weight_meta(w)
+    amax_x = torch.empty(1, device=x.device, dtype=torch.float32)
+    amax_w = torch.empty(1, device=x.device, dtype=torch.float32)
+    out = linear_forward_scaled(
+        x,
+        w,
+        bias,
+        meta.x_scale,
+        meta.scale,
+        meta.x_scale_inv,
+        meta.scale_inv,
+        amax_x,
+        amax_w,
+    )
+    meta.record(amax_w)
+    _update_delayed_scale(meta.x_scale, meta.x_scale_inv, amax_x)
+    return out
+
+
+def fp8_linear_backward(g, x, w, masks):
+    """TE-style scaled fp8 linear backward (called from aten::linear_backward)."""
+    state = fp8_state()
+    meta = state.get_weight_meta(w)
+    amax_g = torch.empty(1, device=g.device, dtype=torch.float32)
+    out = linear_backward_scaled(
+        g,
+        x,
+        w,
+        masks,
+        meta.g_scale,
+        meta.scale,
+        meta.x_scale,
+        meta.g_scale_inv,
+        meta.scale_inv,
+        meta.x_scale_inv,
+        amax_g,
+    )
+    _update_delayed_scale(meta.g_scale, meta.g_scale_inv, amax_g)
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Layer 3: aten::linear integration
+# ---------------------------------------------------------------------------
+
+
+def fp8_linear_enable(enabled: bool = True) -> None:
+    """Toggle fp8 dispatch for aten::linear (global; backward runs on engine
+    worker threads, so a thread-local flag would be lost during backward)."""
+    fp8_state().enabled = enabled
+
+
+def fp8_linear_enabled() -> bool:
+    return fp8_state().enabled
+
+
+def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool:
+    """cuBLASLt fp8 requires M % 16 == 0 and N % 16 == 0 (K is padded)."""
+    m = x.numel() // x.size(-1)
+    return m % 16 == 0 and w.size(0) % 16 == 0
+
+
+def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
+    if (
+        fp8_linear_enabled()
+        and x.dtype == torch.bfloat16
+        and w.dtype == torch.bfloat16
+        and _fp8_supported(x, w)
+    ):
+        return fp8_linear_forward(x, w, bias)
+    return torch.ops.aten.linear.default.redispatch(
+        torch._C.DispatchKeySet(torch._C.DispatchKey.CompositeImplicitAutograd),
+        x,
+        w,
+        bias,
+    )
+
+
+def _linear_backward_cuda_impl(input_tensor, grad_output, weight, output_mask):
+    if (
+        fp8_linear_enabled()
+        and weight.dtype == torch.bfloat16
+        and _fp8_supported(grad_output, weight)
+    ):
+        return fp8_linear_backward(grad_output, input_tensor, weight, list(output_mask))
+    compute_dtype = weight.dtype
+    grad = grad_output.to(compute_dtype)
+    grad_2d = grad.reshape(-1, weight.size(0))
+    input_2d = input_tensor.reshape(-1, input_tensor.size(-1)).to(compute_dtype)
+    grad_input = (
+        torch.mm(grad_2d, weight)
+        if output_mask[0]
+        else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype)
+    )
+    grad_weight = (
+        torch.mm(grad_2d.t(), input_2d)
+        if output_mask[1]
+        else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype)
+    )
+    grad_bias = (
+        grad.sum(dim=0)
+        if output_mask[2]
+        else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype)
+    )
+    return grad_input.reshape_as(input_tensor), grad_weight, grad_bias
+
+
+_lib = Library("aten", "IMPL", "CUDA")
+_lib.impl("linear", _linear_cuda_impl)
+_lib.impl("linear_backward", _linear_backward_cuda_impl)

astrai/extension/fp8_ops.py

@@ -0,0 +1,73 @@
+"""FP8 CUDA kernel interface adapter (the only module touching the pybind.
+
+Isolates the ``fp8_mm`` CUDA extension behind stable Python functions:
+- availability / dtype checks and clear errors
+- torch.library ``custom::fp8_mm`` registration (meta + CPU fallback)
+- quantize-in-GEMM primitives used by ``fp8.py`` training state
+
+Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
+this module is stateless.
+"""
+
+import torch
+from torch.library import custom_op
+
+from astrai.extension.loader import get_module, is_available
+
+
+def _mod():
+    if not is_available("fp8_mm"):
+        raise RuntimeError(
+            "CUDA kernel 'fp8_mm' is not available. Build with CSRC_KERNELS=true."
+        )
+    return get_module("fp8_mm")
+
+
+@custom_op("custom::fp8_mm", mutates_args=())
+def fp8_mm(
+    a: torch.Tensor, b: torch.Tensor, sx: torch.Tensor, sw: torch.Tensor
+) -> torch.Tensor:
+    """FP8 e4m3 GEMM: a[M,K] x b[N,K] -> bf16[M,N] (pre-scaled inputs)."""
+
+
+@fp8_mm.register_fake
+def _fp8_mm_fake(a, b, sx, sw):
+    return torch.empty((a.size(0), b.size(1)), device=a.device, dtype=torch.bfloat16)
+
+
+@fp8_mm.register_kernel("cuda")
+def _fp8_mm_cuda(a, b, sx, sw):
+    return _mod().fp8_mm(a, b)
+
+
+@fp8_mm.register_kernel("cpu")
+def _fp8_mm_cpu(a, b, sx, sw):
+    return torch.mm(a.float(), b.float().t()).to(torch.bfloat16)
+
+
+def linear_forward_scaled(x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w):
+    """Quantize x/w with per-tensor scales + cuBLASLt GEMM + bias -> bf16.
+
+    x/w: [..., K] / [N, K] bf16; sx/sw: f32 scale tensors (device scalars);
+    sx_inv/sw_inv: 1/scale; amax_x/amax_w: f32 buffers receiving max-abs.
+    """
+    if not (x.dtype == torch.bfloat16 and w.dtype == torch.bfloat16):
+        raise TypeError(f"fp8 forward requires bf16 inputs, got {x.dtype}/{w.dtype}")
+    return _mod().fp8_linear_forward_scaled(
+        x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w
+    )
+
+
+def linear_backward_scaled(g, x, w, masks, sg, sw, sx, sg_inv, sw_inv, sx_inv, amax_g):
+    """dX = g @ W, dW = g^T @ X, dB = sum(g) with per-tensor scales."""
+    if not (
+        g.dtype == torch.bfloat16
+        and x.dtype == torch.bfloat16
+        and w.dtype == torch.bfloat16
+    ):
+        raise TypeError(
+            f"fp8 backward requires bf16 inputs, got {g.dtype}/{x.dtype}/{w.dtype}"
+        )
+    return _mod().fp8_linear_backward_scaled(
+        g, x, w, masks, sg, sw, sx, sg_inv, sw_inv, sx_inv, amax_g
+    )

astrai/extension/loader.py

@@ -17,6 +17,7 @@ KERNEL_NAMES = [
     "attn_paged_decode",
     "attn_paged_prefill",
     "rotary_emb",
+    "fp8_mm",
 ]
 
 _available: dict[str, bool] = {}

astrai/inference/__init__.py

@@ -1,15 +1,29 @@
 """Inference module for continuous batching.
 
-Layers:
-  - core/:        Core inference loop (cache, executor, scheduler, task)
-  - api/:         HTTP orchestration (ProtocolHandler, server)
-  - protocols/:   Response builders (OpenAI, Anthropic)
-  - transport/:   SSE transport utilities
-  - engine.py:    Facade (InferenceEngine)
-  - sample.py:    Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy, FrequencyPenaltyStrategy)
+Subpackages:
+  - cache/:     KV cache (buffers, strategies, pool)
+  - runtime/:   Execution + sampling (executor, CUDA graph, sampling strategies)
+  - task/:      Request lifecycle + performance metrics
+  - network/:   HTTP protocol handling (server, protocol, OpenAI/Anthropic builders)
+
+Modules:
+  - scheduler.py:  Continuous batching loop
+  - workspace.py:  Pre-allocated GPU buffers
+  - engine.py:     Facade (InferenceEngine)
 """
 
-from astrai.inference.api import (
+from astrai.inference.cache import (
+    Allocator,
+    KVCache,
+    KVStorage,
+    PagePool,
+    RadixCache,
+    ReqToTokenPool,
+    TaskCacheManager,
+    page_hash,
+)
+from astrai.inference.engine import InferenceEngine
+from astrai.inference.network import (
     AnthropicMessage,
     BaseToolParser,
     ChatCompletionRequest,
@@ -25,25 +39,10 @@ from astrai.inference.api import (
     get_app,
     run_server,
 )
-from astrai.inference.api.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
-from astrai.inference.core import (
-    STOP,
-    Allocator,
-    Executor,
-    InferenceScheduler,
-    KVCache,
-    KVStorage,
-    PagePool,
-    RadixCache,
-    ReqToTokenPool,
-    Task,
-    TaskManager,
-    TaskStatus,
-    page_hash,
-)
-from astrai.inference.engine import InferenceEngine
-from astrai.inference.sample import (
+from astrai.inference.network.anthropic import AnthropicResponseBuilder
+from astrai.inference.network.openai import OpenAIResponseBuilder
+from astrai.inference.runtime.executor import Executor
+from astrai.inference.runtime.sample import (
     BaseSamplingStrategy,
     FrequencyPenaltyStrategy,
     SamplingPipeline,
@@ -52,6 +51,8 @@ from astrai.inference.sample import (
     TopPStrategy,
     sample,
 )
+from astrai.inference.scheduler import InferenceScheduler
+from astrai.inference.task import STOP, Task, TaskManager, TaskStatus
 
 __all__ = [
     "InferenceEngine",
@@ -67,6 +68,7 @@ __all__ = [
     "PagePool",
     "RadixCache",
     "ReqToTokenPool",
+    "TaskCacheManager",
     "page_hash",
     "sample",
     "BaseSamplingStrategy",

astrai/inference/cache/__init__.py

@@ -0,0 +1,27 @@
+"""KV cache subsystem: buffers, strategies, pool management."""
+
+from astrai.inference.cache.buffer import KVCache, KVStorage, ReqToTokenPool
+from astrai.inference.cache.pool import PagePool, TaskCacheManager, page_hash
+from astrai.inference.cache.strategy import (
+    AllocationStrategy,
+    Allocator,
+    ContiguousStrategy,
+    PagedStrategy,
+    RadixCache,
+    TaskCacheState,
+)
+
+__all__ = [
+    "KVCache",
+    "KVStorage",
+    "ReqToTokenPool",
+    "Allocator",
+    "RadixCache",
+    "TaskCacheState",
+    "AllocationStrategy",
+    "ContiguousStrategy",
+    "PagedStrategy",
+    "PagePool",
+    "TaskCacheManager",
+    "page_hash",
+]

astrai/inference/cache/buffer.py

@@ -0,0 +1,104 @@
+"""Physical KV cache buffers.
+
+Layer 1 — ``KVStorage``:   flat token-level K/V GPU buffers [n_layers, size, n_kv_heads, head_dim]
+Layer 2 — ``ReqToTokenPool``: index table [req_idx, pos] → physical token slot
+Layer 3 — ``KVCache``:       pure dataclass passed to the model for direct buffer access
+
+These classes have no knowledge of tasks, allocation policies, or scheduling.
+They are the "dumb" physical storage layer.
+"""
+
+import threading
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+from torch import Tensor
+
+
+class ReqToTokenPool:
+    """Maps [req_idx, pos] → physical token slot in KV storage.
+
+    Each row is one request; each column is a sequence position.  The value
+    at [req_idx, pos] is the flat index into the KV storage buffers.
+    """
+
+    def __init__(self, size: int, max_context_len: int, device: torch.device):
+        self.size = size
+        self.max_context_len = max_context_len
+        self.req_to_token = torch.zeros(
+            (size, max_context_len), dtype=torch.long, device=device
+        )
+        self.free_slots = list(range(size))
+        self._lock = threading.Lock()
+
+    def alloc(self, num_reqs: int) -> Optional[List[int]]:
+        with self._lock:
+            if num_reqs > len(self.free_slots):
+                return None
+            slots = self.free_slots[:num_reqs]
+            self.free_slots = self.free_slots[num_reqs:]
+            return slots
+
+    def free(self, req_indices: List[int]):
+        with self._lock:
+            self.free_slots.extend(req_indices)
+
+    def write(self, indices, values):
+        self.req_to_token[indices] = values
+
+
+class KVStorage:
+    """Token-level KV cache storage.
+
+    Buffers: ``[n_layers, size, n_kv_heads, head_dim]``.  Each token occupies
+    one slot indexed by ``ReqToTokenPool``.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        n_layers: int,
+        n_kv_heads: int,
+        head_dim: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        self.size = size
+        self.k_buffer = torch.empty(
+            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
+        )
+        self.v_buffer = torch.empty(
+            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
+        )
+
+    def get_key_buffer(self, layer_id: int) -> Tensor:
+        return self.k_buffer[layer_id]
+
+    def get_value_buffer(self, layer_id: int) -> Tensor:
+        return self.v_buffer[layer_id]
+
+    def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None:
+        self.k_buffer[layer_id, loc] = k
+        self.v_buffer[layer_id, loc] = v
+
+
+@dataclass
+class KVCache:
+    """Pure data struct passed to model for KV cache I/O.
+
+    The attention layer does raw buffer indexing — no methods, no abstraction.
+    """
+
+    k_buffer: Tensor
+    v_buffer: Tensor
+    req_to_token: Tensor
+    req_pool_indices: Tensor
+    seq_lens: Tensor
+    out_cache_loc: Tensor
+    max_len: int = 0
+    kv_indptr: Optional[Tensor] = None
+    qo_indptr: Optional[Tensor] = None
+    decode_o_part: Optional[Tensor] = None
+    decode_ml_part: Optional[Tensor] = None
+    decode_out: Optional[Tensor] = None

astrai/inference/cache/pool.py

@@ -0,0 +1,351 @@
+"""KV cache orchestration: PagePool + TaskCacheManager.
+
+PagePool owns the physical buffers (``KVStorage`` + ``ReqToTokenPool``)
+and wires them to an allocation strategy.  It assembles the ``KVCache``
+dataclass passed to the model forward.
+
+TaskCacheManager owns the ``task_id`` → ``TaskCacheState`` mapping and
+delegates physical slot allocation to the strategy, and KV bind to the pool.
+
+See ``cache_buffer.py`` for the raw buffer primitives and ``cache_strategy.py``
+for the allocation policies.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import torch
+
+from astrai.inference.cache.buffer import KVCache, KVStorage, ReqToTokenPool
+from astrai.inference.cache.strategy import (
+    AllocationStrategy,
+    Allocator,
+    ContiguousStrategy,
+    PagedStrategy,
+    RadixCache,
+    TaskCacheState,
+)
+from astrai.inference.workspace import InferenceWorkspace
+
+# Re-export everything so existing ``from astrai.inference.cache import ...``
+# continues to work unchanged after the file split.
+__all__ = [
+    "KVCache",
+    "KVStorage",
+    "ReqToTokenPool",
+    "Allocator",
+    "RadixCache",
+    "AllocationStrategy",
+    "ContiguousStrategy",
+    "PagedStrategy",
+    "PagePool",
+    "TaskCacheManager",
+    "TaskCacheState",
+    "page_hash",
+]
+
+# ---- helpers ----
+
+
+def page_hash(
+    token_ids: List[int], page_idx: int, page_size: int, parent_hash: int = 0
+) -> int:
+    start = page_idx * page_size
+    end = min(start + page_size, len(token_ids))
+    h = parent_hash
+    for i in range(start, end):
+        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
+    return h
+
+
+def _is_steady_increment(
+    prev_sig: Optional[tuple],
+    prev_vals: Optional[List[int]],
+    cur_sig: tuple,
+    cur_vals: List[int],
+) -> bool:
+    return (
+        prev_sig is not None
+        and prev_vals is not None
+        and prev_sig == cur_sig
+        and len(prev_vals) == len(cur_vals)
+        and all(c == p + 1 for c, p in zip(cur_vals, prev_vals))
+    )
+
+
+# ---- task-scoped bind state ----
+@dataclass
+class _BindState:
+    """Cached bind metadata for steady-state decode increment detection."""
+
+    sig: tuple
+    seq_lens: List[int]
+
+
+# ---- pool + manager ----
+
+
+class PagePool:
+    """Physical KV cache: buffers + req-to-token table + allocation strategy + bind.
+
+    Does not know about tasks — task lifecycle is managed by
+    :class:`TaskCacheManager`, which holds a reference to this pool.
+    """
+
+    def __init__(
+        self,
+        n_layers: int,
+        n_kv_heads: int,
+        head_dim: int,
+        max_batch_size: int,
+        max_seq_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        page_size: int = 1,
+        n_tokens: Optional[int] = None,
+    ):
+        self.page_size = page_size
+        self.max_batch_size = max_batch_size
+        self.max_seq_len = max_seq_len
+        self.device = device
+        self.dtype = dtype
+        self.n_layers = n_layers
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = head_dim
+
+        self.contiguous = n_tokens is None
+        self.n_tokens = max_batch_size * max_seq_len if self.contiguous else n_tokens
+
+        self._storage = KVStorage(
+            self.n_tokens, n_layers, n_kv_heads, head_dim, device, dtype
+        )
+        self._req_pool = ReqToTokenPool(max_batch_size, max_seq_len, device)
+
+        if self.contiguous:
+            for i in range(max_batch_size):
+                self._req_pool.req_to_token[i] = torch.arange(
+                    i * max_seq_len, (i + 1) * max_seq_len, device=device
+                )
+            self._strategy: AllocationStrategy = ContiguousStrategy()
+        else:
+            n_pages = self.n_tokens // page_size
+            alloc = Allocator(n_pages)
+            prefix = RadixCache(page_size) if page_size > 1 else None
+            if prefix is not None:
+                alloc.on_evict = prefix.evict
+            self._strategy = PagedStrategy(
+                alloc, prefix, page_size, self._req_pool, device
+            )
+
+    @property
+    def strategy(self) -> AllocationStrategy:
+        return self._strategy
+
+    @property
+    def req_pool(self) -> ReqToTokenPool:
+        return self._req_pool
+
+    def bind_tasks(
+        self,
+        req_indices: List[int],
+        seq_lens: List[int],
+        workspace: InferenceWorkspace,
+        device: Optional[torch.device] = None,
+        start_pos: Optional[int] = None,
+        incremental: bool = False,
+    ) -> KVCache:
+        """Assemble the ``KVCache`` metadata for a batch of tasks.
+
+        Args:
+            req_indices: request slot indices (from ``ReqToTokenPool``).
+            seq_lens:    current sequence length per task.
+            workspace:   pre-allocated fixed-shape buffers (CUDA-graph safe).
+            start_pos:   if set, produce **prefill** cache (full q_len range).
+                         If ``None``, produce **decode** cache (last position).
+            incremental: if ``True``, reuse workspace state from previous step
+                         by incrementing counters in-place (decode hot path).
+
+        Returns:
+            ``KVCache`` dataclass with the correct output shapes for the
+            attention backend (prefill: ``[B, q_len]``, decode: ``[B, 1]``).
+        """
+        if device is None:
+            device = workspace.device
+        b = len(req_indices)
+
+        rpi_buf = workspace.req_pool_indices
+        sl_buf = workspace.seq_lens
+        kvp_buf = workspace.kv_indptr
+        inc_buf = workspace.inc
+        ocl_buf = workspace.out_cache_loc
+
+        if incremental:
+            sl_buf[:b] += 1
+            kvp_buf[: b + 1] += inc_buf[: b + 1]
+        else:
+            rpi_buf[:b].copy_(
+                torch.tensor(req_indices, dtype=torch.long, device=device)
+            )
+            sl_buf[:b].copy_(torch.tensor(seq_lens, dtype=torch.long, device=device))
+            kvp_buf[: b + 1].zero_()
+            kvp_buf[1 : b + 1] = sl_buf[:b].cumsum(0).to(torch.int32)
+
+        req_pool_indices = rpi_buf[:b]
+        seq_lens_t = sl_buf[:b]
+        kv_indptr = kvp_buf[: b + 1]
+
+        if start_pos is not None:
+            # ---- prefill: out_cache_loc covers prefix range [start_pos:seq_len] ----
+            seq_len = seq_lens[0]
+            out_cache_loc = self._req_pool.req_to_token[
+                req_pool_indices, start_pos:seq_len
+            ]
+            q_len = seq_len - start_pos
+            workspace.qo_indptr[: b + 1].copy_(
+                torch.arange(b + 1, dtype=torch.int32, device=device) * q_len
+            )
+            qo_indptr = workspace.qo_indptr[: b + 1]
+            decode_o_part = decode_ml_part = decode_out = None
+        else:
+            # ---- decode: out_cache_loc is a single column (last position) ----
+            write_pos = seq_lens_t - 1
+            loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
+            ocl_buf[:b].copy_(loc)
+            out_cache_loc = ocl_buf[:b]
+            qo_indptr = None
+            decode_o_part = getattr(workspace, "decode_o_part", None)
+            decode_ml_part = getattr(workspace, "decode_ml_part", None)
+            decode_out = getattr(workspace, "decode_out", None)
+
+        return KVCache(
+            k_buffer=self._storage.k_buffer,
+            v_buffer=self._storage.v_buffer,
+            req_to_token=self._req_pool.req_to_token,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens_t,
+            out_cache_loc=out_cache_loc,
+            max_len=max(seq_lens),
+            kv_indptr=kv_indptr,
+            qo_indptr=qo_indptr,
+            decode_o_part=decode_o_part,
+            decode_ml_part=decode_ml_part,
+            decode_out=decode_out,
+        )
+
+
+class TaskCacheManager:
+    """Task ↔ KV slot lifecycle manager.
+
+    Sole owner of ``task_id → TaskCacheState``.  Delegates physical slot
+    allocation to the strategy (via ``pool.strategy``) and KV bind to
+    ``pool.bind_tasks()``.
+
+    Usage::
+
+        pool = PagePool(...)
+        mgr = TaskCacheManager(pool)
+        mgr.task_alloc("req_1", [101, 202, 303])
+        ...
+        kv = mgr.bind(["req_1"], workspace)
+    """
+
+    def __init__(self, pool: PagePool):
+        self._pool = pool
+        self._strategy = pool.strategy
+        self._req_pool = pool.req_pool
+        self._max_seq_len = pool.max_seq_len
+        self._states: Dict[str, TaskCacheState] = {}
+        self._bind_state: Optional[_BindState] = None
+        self._bind_was_steady = False
+
+    # -- public task lifecycle --
+
+    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
+        self._bind_state = None
+        req_slots = self._req_pool.alloc(1)
+        if req_slots is None:
+            return False
+        state = TaskCacheState(req_idx=req_slots[0])
+        self._states[task_id] = state
+        if not self._strategy.alloc(state, prompt_ids):
+            self._rollback(state, task_id)
+            return False
+        self._strategy.write_indices(state, prompt_ids)
+        state.length = len(prompt_ids)
+        return True
+
+    def task_free(self, task_id: str):
+        self._bind_state = None
+        state = self._states.pop(task_id, None)
+        if state is None:
+            return
+        self._strategy.free(state)
+        self._req_pool.free([state.req_idx])
+
+    def task_extend(self, task_id: str, pos: int) -> bool:
+        state = self._states.get(task_id)
+        if state is None or pos >= self._max_seq_len:
+            return False
+        if not self._strategy.extend(state, pos):
+            return False
+        state.length = pos + 1
+        return True
+
+    def task_cached(self, task_id: str) -> int:
+        state = self._states.get(task_id)
+        return state.cached if state is not None else 0
+
+    def task_record_hashes(
+        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
+    ):
+        state = self._states.get(task_id)
+        if state is not None:
+            self._strategy.record_hashes(state, prompt_ids, start_logical_page)
+
+    @staticmethod
+    def task_cacheable_ids(task_id: str, prompt_ids: List[int], output_ids: List[int]):
+        return list(prompt_ids) + list(output_ids[:-1])
+
+    # -- bind (assemble KVCache for the model forward) --
+
+    def bind(
+        self,
+        task_ids: List[str],
+        workspace: InferenceWorkspace,
+        device: Optional[torch.device] = None,
+        start_pos: Optional[int] = None,
+    ) -> KVCache:
+        """Build ``KVCache`` for an ordered list of task IDs."""
+        states = [self._states[tid] for tid in task_ids]
+        req_indices = [s.req_idx for s in states]
+        seq_lens = [s.length for s in states]
+        sig = tuple(req_indices)
+
+        prev = self._bind_state
+        incremental = (
+            start_pos is None
+            and prev is not None
+            and _is_steady_increment(prev.sig, prev.seq_lens, sig, seq_lens)
+        )
+        self._bind_state = _BindState(sig, list(seq_lens))
+        self._bind_was_steady = incremental
+
+        return self._pool.bind_tasks(
+            req_indices,
+            seq_lens,
+            workspace,
+            device=device,
+            start_pos=start_pos,
+            incremental=incremental,
+        )
+
+    @property
+    def bind_was_steady(self) -> bool:
+        return self._bind_was_steady
+
+    # -- internals --
+
+    def _rollback(self, state: TaskCacheState, task_id: str):
+        self._strategy.free(state)
+        self._req_pool.free([state.req_idx])
+        self._states.pop(task_id, None)

astrai/inference/cache/strategy.py

@@ -0,0 +1,320 @@
+"""KV cache allocation layer.
+
+Encapsulates the physical slot allocation policy, isolated from GPU buffers
+and task lifecycle management.
+
+- ``TaskCacheState``: data contract between strategy and manager (per-task slot state)
+- ``Allocator``:       bitmask-based page allocator with LRU eviction
+- ``RadixCache``:      page-granular prefix index (exact token match)
+- ``AllocationStrategy``: ABC for physical slot allocation
+- ``ContiguousStrategy``: statically partitioned, no dynamic allocation
+- ``PagedStrategy``:    dynamic paged allocation from a shared pool
+"""
+
+import threading
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Optional, OrderedDict
+
+import torch
+
+from astrai.inference.cache.buffer import ReqToTokenPool
+
+# ---- data contract: per-task slot state ----
+
+
+@dataclass
+class TaskCacheState:
+    """Per-task cache allocation state.
+
+    Co-locates all task-owned cache metadata so the alloc/free/extend
+    lifecycle is atomic.  Owned by ``TaskCacheManager``, consumed by
+    every ``AllocationStrategy`` method.
+    """
+
+    req_idx: int
+    length: int = 0
+    cached: int = 0
+    pages: List[int] = field(default_factory=list)
+
+
+# ---- allocation primitives ----
+
+
+class Allocator:
+    """Bitmask-based page allocator with ref-counting and LRU eviction."""
+
+    def __init__(self, n_pages: int):
+        self._free_mask = (1 << n_pages) - 1
+        self._refs: List[int] = [0] * n_pages
+        self._lru: OrderedDict[int, None] = OrderedDict()
+        self.on_evict: Optional[Callable[[int], None]] = None
+        self._lock = threading.Lock()
+
+    def alloc(self) -> int:
+        with self._lock:
+            if self._free_mask:
+                lsb = self._free_mask & -self._free_mask
+                idx = lsb.bit_length() - 1
+                self._free_mask ^= lsb
+                self._refs[idx] = 1
+                return idx
+            if self._lru:
+                idx, _ = self._lru.popitem(last=False)
+                if self.on_evict:
+                    self.on_evict(idx)
+                self._refs[idx] = 1
+                self._free_mask &= ~(1 << idx)
+                return idx
+            return -1
+
+    def free(self, idx: int, keep_cached: bool = False):
+        with self._lock:
+            self._refs[idx] -= 1
+            if self._refs[idx] == 0:
+                if keep_cached:
+                    self._lru[idx] = None
+                else:
+                    self._free_mask |= 1 << idx
+
+    def inc_ref(self, idx: int):
+        with self._lock:
+            self._refs[idx] += 1
+            self._lru.pop(idx, None)
+
+    def ref_count(self, idx: int) -> int:
+        with self._lock:
+            return self._refs[idx]
+
+    def touch(self, idx: int):
+        with self._lock:
+            if idx in self._lru:
+                self._lru.move_to_end(idx)
+
+
+class RadixNode:
+    """A page-aligned edge in the CPU-side prefix radix trie."""
+
+    __slots__ = ("parent", "children", "page_idx", "tokens", "lock_ref")
+
+    def __init__(self, parent=None, tokens=(), page_idx=None):
+        self.parent = parent
+        self.children: Dict[tuple, "RadixNode"] = {}
+        self.page_idx = page_idx
+        self.tokens = tuple(tokens)
+        self.lock_ref = 0
+
+
+class RadixCache:
+    """Page-granular radix prefix index with exact token matching."""
+
+    def __init__(self, page_size: int):
+        self._page_size = page_size
+        self._root = RadixNode()
+        self._page_to_node: Dict[int, RadixNode] = {}
+        self._lock = threading.Lock()
+
+    def evict(self, idx: int):
+        with self._lock:
+            node = self._page_to_node.pop(idx, None)
+            if node is None:
+                return
+            node.page_idx = None
+            parent = node.parent
+            if parent is not None:
+                parent.children.pop(node.tokens, None)
+
+    def has_page(self, idx: int) -> bool:
+        with self._lock:
+            return idx in self._page_to_node
+
+    def lookup(self, token_ids: List[int]) -> List[int]:
+        with self._lock:
+            full_pages = len(token_ids) // self._page_size
+            hits: List[int] = []
+            node = self._root
+            for i in range(full_pages):
+                start = i * self._page_size
+                page_tokens = tuple(token_ids[start : start + self._page_size])
+                child = node.children.get(page_tokens)
+                if child is None or child.page_idx is None:
+                    break
+                hits.append(child.page_idx)
+                node = child
+            return hits
+
+    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
+        with self._lock:
+            full_pages = len(token_ids) // self._page_size
+            if logical_page_idx >= full_pages:
+                return
+            old = self._page_to_node.pop(page_idx, None)
+            if old is not None and old.parent is not None:
+                old.parent.children.pop(old.tokens, None)
+
+            node = self._root
+            for i in range(logical_page_idx + 1):
+                start = i * self._page_size
+                page_tokens = tuple(token_ids[start : start + self._page_size])
+                child = node.children.get(page_tokens)
+                if child is None:
+                    child = RadixNode(node, page_tokens)
+                    node.children[page_tokens] = child
+                node = child
+            if node.page_idx is not None and node.page_idx != page_idx:
+                replaced = node.page_idx
+                self._page_to_node.pop(replaced, None)
+            node.page_idx = page_idx
+            self._page_to_node[page_idx] = node
+
+    def release(self, pages: List[int]) -> None:
+        with self._lock:
+            for page_idx in pages:
+                node = self._page_to_node.get(page_idx)
+                if node is not None and node.lock_ref:
+                    node.lock_ref -= 1
+
+
+class AllocationStrategy(ABC):
+    """Physical slot allocation policy.
+
+    Subclasses implement the actual allocation semantics.  This ABC declares
+    the contract; there are no default implementations.
+    """
+
+    @abstractmethod
+    def alloc(self, state: TaskCacheState, prompt_ids: List[int]) -> bool: ...
+
+    @abstractmethod
+    def free(self, state: TaskCacheState) -> None: ...
+
+    @abstractmethod
+    def extend(self, state: TaskCacheState, pos: int) -> bool: ...
+
+    @abstractmethod
+    def write_indices(self, state: TaskCacheState, prompt_ids: List[int]) -> None: ...
+
+    @abstractmethod
+    def record_hashes(
+        self,
+        state: TaskCacheState,
+        prompt_ids: List[int],
+        start: int,
+    ) -> None: ...
+
+
+class ContiguousStrategy(AllocationStrategy):
+    """Static contiguous allocation: slots are pre-assigned at pool init.
+
+    No dynamic allocation or prefix caching.  All operations are no-ops
+    because ``ReqToTokenPool`` is pre-filled with contiguous ranges.
+    """
+
+    def alloc(self, state: TaskCacheState, prompt_ids: List[int]) -> bool:
+        return True
+
+    def free(self, state: TaskCacheState) -> None:
+        pass
+
+    def extend(self, state: TaskCacheState, pos: int) -> bool:
+        return True
+
+    def write_indices(self, state: TaskCacheState, prompt_ids: List[int]) -> None:
+        pass
+
+    def record_hashes(
+        self,
+        state: TaskCacheState,
+        prompt_ids: List[int],
+        start: int,
+    ) -> None:
+        pass
+
+
+class PagedStrategy(AllocationStrategy):
+    """Dynamic paged allocation from a shared bitmask pool.
+
+    ``page_size`` is a parameter, not a separate strategy: at ``page_size=1``
+    each allocated page *is* one token slot (``page * 1 + 0``), and prefix
+    caching is simply disabled (``prefix=None``).  The unified page formula
+    ``pages[page_idx] * page_size + offset`` holds for both.
+    """
+
+    def __init__(
+        self,
+        alloc: Allocator,
+        prefix: Optional[RadixCache],
+        page_size: int,
+        req_pool: ReqToTokenPool,
+        device,
+    ):
+        self._alloc = alloc
+        self._prefix = prefix
+        self._page_size = page_size
+        self._req_pool = req_pool
+        self._device = device
+
+    def alloc(self, state: TaskCacheState, prompt_ids: List[int]) -> bool:
+        if self._prefix is not None:
+            hits = self._prefix.lookup(prompt_ids)
+            state.cached = len(hits) * self._page_size
+            for p in hits:
+                self._alloc.inc_ref(p)
+            state.pages = list(hits)
+
+        remaining = len(prompt_ids) - state.cached
+        if remaining <= 0:
+            return True
+        n_new = (remaining + self._page_size - 1) // self._page_size
+        for _ in range(n_new):
+            p = self._alloc.alloc()
+            if p < 0:
+                return False
+            state.pages.append(p)
+        return True
+
+    def free(self, state: TaskCacheState) -> None:
+        if self._prefix is not None:
+            for p in state.pages:
+                keep = self._prefix.has_page(p)
+                self._alloc.free(p, keep_cached=keep)
+                if not keep:
+                    self._prefix.evict(p)
+        else:
+            for p in state.pages:
+                self._alloc.free(p)
+
+    def extend(self, state: TaskCacheState, pos: int) -> bool:
+        page_idx = pos // self._page_size
+        if page_idx >= len(state.pages):
+            p = self._alloc.alloc()
+            if p < 0:
+                return False
+            state.pages.append(p)
+        offset = pos % self._page_size
+        self._req_pool.req_to_token[state.req_idx, pos] = (
+            state.pages[page_idx] * self._page_size + offset
+        )
+        return True
+
+    def write_indices(self, state: TaskCacheState, prompt_ids: List[int]) -> None:
+        total = len(prompt_ids)
+        for pos in range(total):
+            page_idx = pos // self._page_size
+            offset = pos % self._page_size
+            if page_idx < len(state.pages):
+                self._req_pool.req_to_token[state.req_idx, pos] = (
+                    state.pages[page_idx] * self._page_size + offset
+                )
+
+    def record_hashes(
+        self,
+        state: TaskCacheState,
+        prompt_ids: List[int],
+        start: int,
+    ) -> None:
+        if self._prefix is None:
+            return
+        full = len(prompt_ids) // self._page_size
+        for i in range(start, min(full, len(state.pages))):
+            self._prefix.record(state.pages[i], prompt_ids, i)

astrai/inference/core/__init__.py

@@ -1,30 +0,0 @@
-"""Inference core: cache, executor, scheduler, task management."""
-
-from astrai.inference.core.cache import (
-    Allocator,
-    KVCache,
-    KVStorage,
-    PagePool,
-    RadixCache,
-    ReqToTokenPool,
-    page_hash,
-)
-from astrai.inference.core.executor import Executor
-from astrai.inference.core.scheduler import InferenceScheduler
-from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
-
-__all__ = [
-    "Allocator",
-    "KVCache",
-    "KVStorage",
-    "PagePool",
-    "RadixCache",
-    "ReqToTokenPool",
-    "page_hash",
-    "Executor",
-    "InferenceScheduler",
-    "STOP",
-    "Task",
-    "TaskManager",
-    "TaskStatus",
-]

astrai/inference/core/cache.py

@@ -1,638 +0,0 @@
-"""KV cache architecture: three-layer separation (SGLang-inspired).
-
-Layer 1 — KVStorage:      flat token-level K/V buffers [n_layers, size, H, D]
-Layer 2 — ReqToTokenPool:  index table [req_idx, pos] → physical token slot
-Layer 3 — Allocator:       slot/page allocation with ref-counting and LRU
-
-PagePool orchestrates all three plus RadixCache (prefix addressing).
-KVCache is a pure dataclass passed to the model for direct buffer access.
-
-Two modes:
-  - contiguous (default): pre-allocated per-request blocks, no dynamic alloc
-  - paged: shared pool with on-demand allocation, prefix caching support
-"""
-
-import threading
-from collections import OrderedDict
-from dataclasses import dataclass
-from typing import Callable, Dict, List, Optional
-
-import torch
-from torch import Tensor
-
-from astrai.inference.core.workspace import InferenceWorkspace
-
-
-def page_hash(
-    token_ids: List[int], page_idx: int, page_size: int, parent_hash: int = 0
-) -> int:
-    start = page_idx * page_size
-    end = min(start + page_size, len(token_ids))
-    h = parent_hash
-    for i in range(start, end):
-        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
-    return h
-
-
-class Allocator:
-    """Bitmask-based page allocator with ref-counting and LRU eviction."""
-
-    def __init__(self, n_pages: int):
-        self._free_mask = (1 << n_pages) - 1
-        self._refs: List[int] = [0] * n_pages
-        self._lru: OrderedDict[int, None] = OrderedDict()
-        self.on_evict: Optional[Callable[[int], None]] = None
-        self._lock = threading.Lock()
-
-    def alloc(self) -> int:
-        with self._lock:
-            if self._free_mask:
-                lsb = self._free_mask & -self._free_mask
-                idx = lsb.bit_length() - 1
-                self._free_mask ^= lsb
-                self._refs[idx] = 1
-                return idx
-            if self._lru:
-                idx, _ = self._lru.popitem(last=False)
-                if self.on_evict:
-                    self.on_evict(idx)
-                self._refs[idx] = 1
-                self._free_mask &= ~(1 << idx)
-                return idx
-            return -1
-
-    def free(self, idx: int, keep_cached: bool = False):
-        with self._lock:
-            self._refs[idx] -= 1
-            if self._refs[idx] == 0:
-                if keep_cached:
-                    self._lru[idx] = None
-                else:
-                    self._free_mask |= 1 << idx
-
-    def inc_ref(self, idx: int):
-        with self._lock:
-            self._refs[idx] += 1
-            self._lru.pop(idx, None)
-
-    def ref_count(self, idx: int) -> int:
-        with self._lock:
-            return self._refs[idx]
-
-    def touch(self, idx: int):
-        with self._lock:
-            if idx in self._lru:
-                self._lru.move_to_end(idx)
-
-
-class RadixNode:
-    """A page-aligned edge in the CPU-side prefix radix."""
-
-    __slots__ = ("parent", "children", "page_idx", "tokens", "lock_ref")
-
-    def __init__(self, parent=None, tokens=(), page_idx=None):
-        self.parent = parent
-        self.children: Dict[tuple, "RadixNode"] = {}
-        self.page_idx = page_idx
-        self.tokens = tuple(tokens)
-        self.lock_ref = 0
-
-
-class RadixCache:
-    """Page-granular radix prefix index with exact token matching."""
-
-    def __init__(self, page_size: int):
-        self._page_size = page_size
-        self._root = RadixNode()
-        self._page_to_node: Dict[int, RadixNode] = {}
-        # Retained as an introspection-compatible map; matching never relies on
-        # this lossy value.
-        self._page_to_hash: Dict[int, int] = {}
-        self._lock = threading.Lock()
-
-    def evict(self, idx: int):
-        with self._lock:
-            node = self._page_to_node.pop(idx, None)
-            self._page_to_hash.pop(idx, None)
-            if node is None:
-                return
-            node.page_idx = None
-            parent = node.parent
-            if parent is not None:
-                parent.children.pop(node.tokens, None)
-
-    def has_page(self, idx: int) -> bool:
-        with self._lock:
-            return idx in self._page_to_node
-
-    def lookup(self, token_ids: List[int]) -> List[int]:
-        with self._lock:
-            full_pages = len(token_ids) // self._page_size
-            hits: List[int] = []
-            node = self._root
-            for i in range(full_pages):
-                start = i * self._page_size
-                page_tokens = tuple(token_ids[start : start + self._page_size])
-                child = node.children.get(page_tokens)
-                if child is None or child.page_idx is None:
-                    break
-                hits.append(child.page_idx)
-                node = child
-            return hits
-
-    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
-        with self._lock:
-            full_pages = len(token_ids) // self._page_size
-            if logical_page_idx >= full_pages:
-                return
-            old = self._page_to_node.pop(page_idx, None)
-            self._page_to_hash.pop(page_idx, None)
-            if old is not None and old.parent is not None:
-                old.parent.children.pop(old.tokens, None)
-
-            node = self._root
-            for i in range(logical_page_idx + 1):
-                start = i * self._page_size
-                page_tokens = tuple(token_ids[start : start + self._page_size])
-                child = node.children.get(page_tokens)
-                if child is None:
-                    child = RadixNode(node, page_tokens)
-                    node.children[page_tokens] = child
-                node = child
-            if node.page_idx is not None and node.page_idx != page_idx:
-                replaced = node.page_idx
-                self._page_to_node.pop(replaced, None)
-                self._page_to_hash.pop(replaced, None)
-            node.page_idx = page_idx
-            self._page_to_node[page_idx] = node
-            self._page_to_hash[page_idx] = page_hash(
-                token_ids, logical_page_idx, self._page_size
-            )
-
-    def release(self, pages: List[int]) -> None:
-        with self._lock:
-            for page_idx in pages:
-                node = self._page_to_node.get(page_idx)
-                if node is not None and node.lock_ref:
-                    node.lock_ref -= 1
-
-
-class ReqToTokenPool:
-    """Maps [req_idx, pos] -> physical token slot in KV storage.
-
-    Each row is one request; each column is a sequence position. The value
-    at [req_idx, pos] is the flat index into the KV storage buffers.
-    """
-
-    def __init__(self, size: int, max_context_len: int, device: torch.device):
-        self.size = size
-        self.max_context_len = max_context_len
-        self.req_to_token = torch.zeros(
-            (size, max_context_len), dtype=torch.long, device=device
-        )
-        self.free_slots = list(range(size))
-        self._lock = threading.Lock()
-
-    def alloc(self, num_reqs: int) -> Optional[List[int]]:
-        with self._lock:
-            if num_reqs > len(self.free_slots):
-                return None
-            slots = self.free_slots[:num_reqs]
-            self.free_slots = self.free_slots[num_reqs:]
-            return slots
-
-    def free(self, req_indices: List[int]):
-        with self._lock:
-            self.free_slots.extend(req_indices)
-
-    def write(self, indices, values):
-        self.req_to_token[indices] = values
-
-
-class KVStorage:
-    """Token-level KV cache storage.
-
-    Buffers: [n_layers, size, n_kv_heads, head_dim]. Each token occupies
-    one slot indexed by ReqToTokenPool.
-    """
-
-    def __init__(
-        self,
-        size: int,
-        n_layers: int,
-        n_kv_heads: int,
-        head_dim: int,
-        device: torch.device,
-        dtype: torch.dtype,
-    ):
-        self.size = size
-        self.k_buffer = torch.empty(
-            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
-        )
-        self.v_buffer = torch.empty(
-            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
-        )
-
-    def get_key_buffer(self, layer_id: int) -> Tensor:
-        return self.k_buffer[layer_id]
-
-    def get_value_buffer(self, layer_id: int) -> Tensor:
-        return self.v_buffer[layer_id]
-
-    def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None:
-        self.k_buffer[layer_id, loc] = k
-        self.v_buffer[layer_id, loc] = v
-
-
-@dataclass
-class KVCache:
-    """Pure data struct passed to model for KV cache I/O.
-
-    The attention layer does raw buffer indexing — no methods, no abstraction.
-
-    Attributes:
-        k_buffer: [n_layers, size, n_kv_heads, head_dim]
-        v_buffer: [n_layers, size, n_kv_heads, head_dim]
-        req_to_token: [num_reqs, max_ctx_len] — index table
-        req_pool_indices: [batch_size] — row indices into req_to_token
-        seq_lens: [batch_size] — per-request total sequence lengths
-        out_cache_loc: [batch, new_seq_len] or [batch, 1] — write indices
-        max_len: max(seq_lens) as Python int — avoids GPU sync in decode
-        kv_indptr: [batch+1] int32 — prefix sum of seq_lens, precomputed once
-            per step so the attention backend avoids rebuilding it per layer.
-        qo_indptr: [batch+1] int32 — prefill qo prefix-sum (None in decode)
-        decode_o_part: split-KV o partial workspace (mirrors FlashInfer)
-        decode_ml_part: split-KV m/l partial workspace (mirrors FlashInfer)
-        decode_out: pre-allocated decode output buffer (graph-safe)
-    """
-
-    k_buffer: Tensor
-    v_buffer: Tensor
-    req_to_token: Tensor
-    req_pool_indices: Tensor
-    seq_lens: Tensor
-    out_cache_loc: Tensor
-    max_len: int = 0
-    kv_indptr: Optional[Tensor] = None
-    qo_indptr: Optional[Tensor] = None
-    decode_o_part: Optional[Tensor] = None
-    decode_ml_part: Optional[Tensor] = None
-    decode_out: Optional[Tensor] = None
-
-
-class PagePool:
-    """Top-level KV cache manager.
-
-    Combines KVStorage + ReqToTokenPool + Allocator + RadixCache.
-
-    Args:
-        n_layers: Number of transformer layers.
-        n_kv_heads: Number of KV attention heads.
-        head_dim: Dimension per head.
-        max_batch_size: Maximum concurrent requests.
-        max_seq_len: Maximum sequence length per request.
-        device, dtype: Tensor device and dtype.
-        page_size: Page size for paged mode (1 = token-level).
-        n_tokens: Total token slots for paged mode. None = contiguous mode
-            (pre-allocates max_batch_size * max_seq_len).
-    """
-
-    def __init__(
-        self,
-        n_layers: int,
-        n_kv_heads: int,
-        head_dim: int,
-        max_batch_size: int,
-        max_seq_len: int,
-        device: torch.device,
-        dtype: torch.dtype,
-        page_size: int = 1,
-        n_tokens: Optional[int] = None,
-    ):
-        self.page_size = page_size
-        self.max_batch_size = max_batch_size
-        self.max_seq_len = max_seq_len
-        self.device = device
-        self.dtype = dtype
-        self.n_layers = n_layers
-        self.n_kv_heads = n_kv_heads
-        self.head_dim = head_dim
-
-        self.contiguous = n_tokens is None
-        if self.contiguous:
-            self.n_tokens = max_batch_size * max_seq_len
-        else:
-            self.n_tokens = n_tokens
-
-        self._storage = KVStorage(
-            self.n_tokens, n_layers, n_kv_heads, head_dim, device, dtype
-        )
-        self._req_pool = ReqToTokenPool(max_batch_size, max_seq_len, device)
-
-        if self.contiguous:
-            for i in range(max_batch_size):
-                self._req_pool.req_to_token[i] = torch.arange(
-                    i * max_seq_len, (i + 1) * max_seq_len, device=device
-                )
-            self._alloc: Optional[Allocator] = None
-            self._prefix: Optional[RadixCache] = None
-        else:
-            n_pages = self.n_tokens // page_size
-            self._alloc = Allocator(n_pages)
-            self._prefix = RadixCache(page_size) if page_size > 1 else None
-            if self._prefix is not None:
-                self._alloc.on_evict = self._prefix.evict
-
-        self._task_req: Dict[str, int] = {}
-        self._task_len: Dict[int, int] = {}
-        self._task_cached: Dict[str, int] = {}
-        self._task_slots: Dict[str, List[int]] = {}
-        self._task_pages: Dict[str, List[int]] = {}
-        self._lock = threading.Lock()
-
-        # Steady-state decode validation state: the ordered task set and its
-        # Python seq_lens mirror.  When the same set advances every sequence
-        # by exactly one token per step, bind_tasks updates the stable
-        # buffers in-place (+=1 / +=inc) instead of re-cumsumming.  Any
-        # task-set change is a miss and rebuilds.
-        self._bind_sig: Optional[tuple] = None
-        self._bind_seq_lens: Optional[List[int]] = None
-
-    # ---- task lifecycle ----
-
-    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
-        req_slots = self._req_pool.alloc(1)
-        if req_slots is None:
-            return False
-        req_idx = req_slots[0]
-        self._task_req[task_id] = req_idx
-
-        if self.contiguous:
-            self._task_len[req_idx] = len(prompt_ids)
-            self._task_cached[task_id] = 0
-            return True
-
-        n_tokens_needed = len(prompt_ids)
-        cached = 0
-
-        if self._prefix is not None:
-            hits = self._prefix.lookup(prompt_ids)
-            cached = len(hits) * self.page_size
-            for p in hits:
-                self._alloc.inc_ref(p)
-            self._task_pages[task_id] = list(hits)
-            self._task_slots[task_id] = []
-        else:
-            self._task_pages[task_id] = []
-            self._task_slots[task_id] = []
-
-        remaining = n_tokens_needed - cached
-        if remaining > 0:
-            if self.page_size == 1:
-                slots = self._alloc_tokens(remaining)
-                if slots is None:
-                    for p in self._task_pages[task_id]:
-                        self._alloc.free(p)
-                    self._req_pool.free([req_idx])
-                    del self._task_req[task_id]
-                    return False
-                self._task_slots[task_id] = slots
-            else:
-                n_new_pages = (remaining + self.page_size - 1) // self.page_size
-                new_pages = []
-                for _ in range(n_new_pages):
-                    p = self._alloc.alloc()
-                    if p < 0:
-                        for hp in self._task_pages[task_id]:
-                            self._alloc.free(hp)
-                        for np_ in new_pages:
-                            self._alloc.free(np_)
-                        self._req_pool.free([req_idx])
-                        del self._task_req[task_id]
-                        return False
-                    new_pages.append(p)
-                self._task_pages[task_id].extend(new_pages)
-
-        self._write_req_to_token(task_id, prompt_ids, cached)
-        self._task_len[req_idx] = len(prompt_ids)
-        self._task_cached[task_id] = cached
-        return True
-
-    def task_free(self, task_id: str):
-        req_idx = self._task_req.pop(task_id, None)
-        if req_idx is None:
-            return
-        self._task_len.pop(req_idx, None)
-        self._task_cached.pop(task_id, None)
-
-        if not self.contiguous:
-            if self._prefix is not None:
-                for p in self._task_pages.get(task_id, []):
-                    keep = self._prefix.has_page(p)
-                    self._alloc.free(p, keep_cached=keep)
-                    if not keep:
-                        self._prefix.evict(p)
-            else:
-                for p in self._task_pages.get(task_id, []):
-                    self._alloc.free(p)
-            self._task_pages.pop(task_id, None)
-            self._task_slots.pop(task_id, None)
-
-        self._req_pool.free([req_idx])
-
-    def task_extend(self, task_id: str, pos: int) -> bool:
-        req_idx = self._task_req.get(task_id)
-        if req_idx is None or pos >= self.max_seq_len:
-            return False
-
-        # Paged mode must also claim a physical slot for the new token;
-        # contiguous mode's block is pre-allocated so this is a no-op.
-        if not self.contiguous and not self._extend_slot(task_id, req_idx, pos):
-            return False
-
-        self._task_len[req_idx] = pos + 1
-        return True
-
-    def _extend_slot(self, task_id: str, req_idx: int, pos: int) -> bool:
-        """Allocate the physical slot for one extended token (paged mode)."""
-        if self.page_size == 1:
-            slots = self._alloc_tokens(1)
-            if slots is None:
-                return False
-            self._task_slots.setdefault(task_id, []).extend(slots)
-            self._req_pool.req_to_token[req_idx, pos] = slots[0]
-            return True
-
-        page_idx = pos // self.page_size
-        existing = self._task_pages.get(task_id, [])
-        if page_idx >= len(existing):
-            p = self._alloc.alloc()
-            if p < 0:
-                return False
-            existing.append(p)
-            self._task_pages[task_id] = existing
-        page_offset = pos % self.page_size
-        page = existing[page_idx]
-        token_slot = page * self.page_size + page_offset
-        self._req_pool.req_to_token[req_idx, pos] = token_slot
-        return True
-
-    def task_cached(self, task_id: str) -> int:
-        return self._task_cached.get(task_id, 0)
-
-    def task_record_hashes(
-        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
-    ):
-        if self._prefix is None or self.contiguous:
-            return
-        pages = self._task_pages.get(task_id, [])
-        full_pages = len(prompt_ids) // self.page_size
-        for i in range(start_logical_page, min(full_pages, len(pages))):
-            self._prefix.record(pages[i], prompt_ids, i)
-
-    def task_cacheable_ids(
-        self, task_id: str, prompt_ids: List[int], output_ids: List[int]
-    ):
-        """Return the sequence whose KV entries are already materialized.
-
-        The first sampled output is produced by prompt prefill, and the last
-        sampled output has not been decoded into KV yet.  Therefore the cache
-        can safely retain the prompt plus every output except the last one.
-        """
-        return list(prompt_ids) + list(output_ids[:-1])
-
-    # ---- bind for forward ----
-
-    def bind_tasks(
-        self,
-        task_ids: List[str],
-        workspace: InferenceWorkspace,
-        device: Optional[torch.device] = None,
-        start_pos: Optional[int] = None,
-    ) -> KVCache:
-        if device is None:
-            device = workspace.device
-        req_indices = [self._task_req[tid] for tid in task_ids]
-        # Per-request lengths come from the pool's own tracking (task_alloc
-        # sets len(prompt_ids); task_extend sets pos+1), so callers need not
-        # pass them.
-        seq_lens = [self._task_len[req_idx] for req_idx in req_indices]
-        b = len(task_ids)
-        sig = tuple(task_ids)
-
-        # Write into the caller's workspace buffers (fixed addresses, sized
-        # to max_batch/max_seq at init) — the sole owner of the per-step
-        # KV bind tensors.
-        rpi_buf = workspace.req_pool_indices
-        sl_buf = workspace.seq_lens
-        kvp_buf = workspace.kv_indptr
-        inc_buf = workspace.inc
-        ocl_buf = workspace.out_cache_loc
-
-        incremental = (
-            start_pos is None
-            and self._bind_sig is not None
-            and self._bind_sig == sig
-            and self._bind_seq_lens is not None
-            and len(self._bind_seq_lens) == b
-            and all(s == p + 1 for s, p in zip(seq_lens, self._bind_seq_lens))
-        )
-        if incremental:
-            # Steady-state decode: advance the stable buffers in-place.
-            # Normal-mode buffers keep ``+=`` legal regardless of whether
-            # this runs inside ``torch.inference_mode()``.
-            sl_buf[:b] += 1
-            kvp_buf[: b + 1] += inc_buf[: b + 1]
-            req_pool_indices = rpi_buf[:b]
-            seq_lens_t = sl_buf[:b]
-            kv_indptr = kvp_buf[: b + 1]
-        else:
-            # Cold path: fill the stable buffers from fresh host tensors.
-            rpi_buf[:b].copy_(
-                torch.tensor(req_indices, dtype=torch.long, device=device)
-            )
-            sl_buf[:b].copy_(torch.tensor(seq_lens, dtype=torch.long, device=device))
-            kvp_buf[: b + 1].zero_()
-            kvp_buf[1 : b + 1] = sl_buf[:b].cumsum(0).to(torch.int32)
-            req_pool_indices = rpi_buf[:b]
-            seq_lens_t = sl_buf[:b]
-            kv_indptr = kvp_buf[: b + 1]
-        self._bind_sig = sig
-        self._bind_seq_lens = list(seq_lens)
-
-        if start_pos is not None:
-            seq_len = seq_lens[0]
-            out_cache_loc = self._req_pool.req_to_token[
-                req_pool_indices, start_pos:seq_len
-            ]
-            # Ragged query segmentation for the prefill kernel, computed once
-            # (was rebuilt per layer in CudaBackend.fwd_prefill).
-            q_len = seq_len - start_pos
-            workspace.qo_indptr[: b + 1].copy_(
-                torch.arange(b + 1, dtype=torch.int32, device=device) * q_len
-            )
-            qo_indptr = workspace.qo_indptr[: b + 1]
-            decode_o_part, decode_ml_part = None, None
-            decode_out = None
-        else:
-            write_pos = seq_lens_t - 1
-            loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
-            ocl_buf[:b].copy_(loc)
-            out_cache_loc = ocl_buf[:b]
-            qo_indptr = None
-            decode_o_part = getattr(workspace, "decode_o_part", None)
-            decode_ml_part = getattr(workspace, "decode_ml_part", None)
-            decode_out = getattr(workspace, "decode_out", None)
-
-        return KVCache(
-            k_buffer=self._storage.k_buffer,
-            v_buffer=self._storage.v_buffer,
-            req_to_token=self._req_pool.req_to_token,
-            req_pool_indices=req_pool_indices,
-            seq_lens=seq_lens_t,
-            out_cache_loc=out_cache_loc,
-            max_len=max(seq_lens),
-            kv_indptr=kv_indptr,
-            qo_indptr=qo_indptr,
-            decode_o_part=decode_o_part,
-            decode_ml_part=decode_ml_part,
-            decode_out=decode_out,
-        )
-
-    # ---- internals ----
-
-    def _alloc_tokens(self, n: int) -> Optional[List[int]]:
-        if self.page_size != 1:
-            raise RuntimeError("_alloc_tokens is for page_size=1 only")
-        slots = []
-        for _ in range(n):
-            p = self._alloc.alloc()
-            if p < 0:
-                for s in slots:
-                    self._alloc.free(s)
-                return None
-            slots.append(p)
-        return slots
-
-    def _write_req_to_token(self, task_id: str, prompt_ids: List[int], cached: int):
-        req_idx = self._task_req[task_id]
-        total = len(prompt_ids)
-
-        if self.contiguous:
-            return
-
-        if self.page_size == 1:
-            slots = self._task_slots.get(task_id, [])
-            all_slots = slots[: total - cached]
-            if all_slots:
-                self._req_pool.req_to_token[req_idx, cached:total] = torch.tensor(
-                    all_slots, dtype=torch.long, device=self.device
-                )
-        else:
-            pages = self._task_pages.get(task_id, [])
-            for pos in range(cached, total):
-                page_idx = pos // self.page_size
-                page_offset = pos % self.page_size
-                if page_idx < len(pages):
-                    token_slot = pages[page_idx] * self.page_size + page_offset
-                    self._req_pool.req_to_token[req_idx, pos] = token_slot

astrai/inference/engine.py

@@ -8,9 +8,10 @@ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple,
 import torch
 import torch.nn as nn
 
-from astrai.inference.core.cache import PagePool
-from astrai.inference.core.scheduler import InferenceScheduler
-from astrai.inference.core.task import STOP
+from astrai.extension import ATTN_BACKEND, AttentionBackend, get_backend
+from astrai.inference.cache import PagePool
+from astrai.inference.scheduler import InferenceScheduler
+from astrai.inference.task import STOP
 from astrai.tokenize import AutoTokenizer
 
 
@@ -74,6 +75,8 @@ class InferenceEngine:
         max_batch_size: int = 1,
         max_seq_len: Optional[int] = None,
         cache: Optional[PagePool] = None,
+        enable_cuda_graph: bool = True,
+        backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None,
     ):
         self.model = model
         self.tokenizer = tokenizer
@@ -83,6 +86,8 @@ class InferenceEngine:
             max_batch_size=max_batch_size,
             max_seq_len=max_seq_len,
             cache=cache,
+            enable_cuda_graph=enable_cuda_graph,
+            backend=backend,
         )
 
         self.scheduler.start()
@@ -172,6 +177,7 @@ class InferenceEngine:
         rep_window: int,
     ) -> Union[Generator, str, List[str]]:
         n = len(prompts)
+        request_backend = get_backend(use_default=False)
         result = GenerateResult(count=n)
         task_ids = [
             self.scheduler.add_task(
@@ -182,6 +188,7 @@ class InferenceEngine:
                 top_k=top_k,
                 frequency_penalty=frequency_penalty,
                 rep_window=rep_window,
+                backend=request_backend,
                 stream_callback=lambda token, idx=i: result.append(token, idx),
             )
             for i, p in enumerate(prompts)
@@ -204,27 +211,31 @@ class InferenceEngine:
 
         def gen():
             nonlocal remaining
-            try:
-                while remaining > 0:
-                    items = result.pop_all()
-                    for idx, token in items:
-                        if token is STOP:
-                            if not finished[idx]:
-                                finished[idx] = True
-                                remaining -= 1
-                        else:
-                            yield (idx, token) if is_batch else token
-                    if remaining > 0:
-                        result.wait(timeout=0.05)
-            finally:
-                for tid in task_ids:
-                    self.scheduler.remove_task(tid)
+            while remaining > 0:
+                items = result.pop_all()
+                for idx, token in items:
+                    if token is STOP:
+                        if not finished[idx]:
+                            finished[idx] = True
+                            remaining -= 1
+                    else:
+                        yield (idx, token) if is_batch else token
+                if remaining > 0:
+                    result.wait(timeout=0.05)
 
         return gen()
 
     def get_stats(self) -> Dict[str, Any]:
         return self.scheduler.get_stats()
 
+    @property
+    def backend_name(self) -> str:
+        return self.scheduler.backend_name
+
+    @property
+    def cuda_graph_enabled(self) -> bool:
+        return self.scheduler.cuda_graph_enabled
+
     def shutdown(self):
         self.scheduler.stop()
         if torch.cuda.is_available():

astrai/inference/metrics.py

@@ -0,0 +1,223 @@
+"""Unified per-task perf/stats: timing records, context-manager scopes, aggregate reporting."""
+
+import time
+from collections import deque
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any, Deque, Dict, Generator, List, Literal, Optional
+
+
+@dataclass
+class TaskTiming:
+    """Timestamp snapshots and computed metrics for one generation task.
+
+    Created by :class:`MetricsCollector` at task-registration time;
+    updated via ``record`` / ``mark_finished``.
+    """
+
+    task_id: str
+    arrival_time: float
+    prefill_start_time: Optional[float] = None
+    first_token_time: Optional[float] = None
+    finish_time: Optional[float] = None
+    input_tokens: int = 0
+    output_tokens: int = 0
+    _decode_steps: int = 0
+    _decode_total_s: float = 0.0
+
+    # derived metrics
+
+    @property
+    def queue_wait_ms(self) -> Optional[float]:
+        if self.prefill_start_time is not None:
+            return (self.prefill_start_time - self.arrival_time) * 1000
+        return None
+
+    @property
+    def ttft_ms(self) -> Optional[float]:
+        if self.first_token_time is not None:
+            return (self.first_token_time - self.arrival_time) * 1000
+        return None
+
+    @property
+    def prefill_tps(self) -> Optional[float]:
+        if self.prefill_start_time is not None and self.first_token_time is not None:
+            d = self.first_token_time - self.prefill_start_time
+            if d > 0 and self.input_tokens > 0:
+                return self.input_tokens / d
+        return None
+
+    @property
+    def decode_tps(self) -> Optional[float]:
+        if self.first_token_time is not None and self.finish_time is not None:
+            d = self.finish_time - self.first_token_time
+            dt = self.output_tokens - 1
+            if dt > 0 and d > 0:
+                return dt / d
+        return None
+
+    @property
+    def decode_avg_ms(self) -> Optional[float]:
+        if self._decode_steps > 0 and self._decode_total_s > 0:
+            return (self._decode_total_s / self._decode_steps) * 1000
+        return None
+
+    @property
+    def e2e_latency_ms(self) -> Optional[float]:
+        if self.finish_time is not None:
+            return (self.finish_time - self.arrival_time) * 1000
+        return None
+
+    @property
+    def total_tps(self) -> Optional[float]:
+        if self.finish_time is not None:
+            total = self.input_tokens + self.output_tokens
+            d = self.finish_time - self.arrival_time
+            if total > 0 and d > 0:
+                return total / d
+        return None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "task_id": self.task_id,
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "queue_wait_ms": (
+                round(self.queue_wait_ms, 2) if self.queue_wait_ms is not None else None
+            ),
+            "ttft_ms": (round(self.ttft_ms, 2) if self.ttft_ms is not None else None),
+            "prefill_tps": (
+                round(self.prefill_tps, 2) if self.prefill_tps is not None else None
+            ),
+            "decode_tps": (
+                round(self.decode_tps, 2) if self.decode_tps is not None else None
+            ),
+            "decode_avg_ms": (
+                round(self.decode_avg_ms, 2) if self.decode_avg_ms is not None else None
+            ),
+            "total_tps": (
+                round(self.total_tps, 2) if self.total_tps is not None else None
+            ),
+            "e2e_latency_ms": (
+                round(self.e2e_latency_ms, 2)
+                if self.e2e_latency_ms is not None
+                else None
+            ),
+        }
+
+
+class MetricsCollector:
+    """Single-owner perf/stats hub for all generation tasks.
+
+    Usage::
+
+        metrics = MetricsCollector()
+        metrics.register(task_id, arrival_time)
+
+        with metrics.record(task_ids, "prefill"):
+            run_prefill(...)
+
+        metrics.mark_finished(task_id, input_tokens, output_tokens)
+
+        stats = metrics.get_stats()
+    """
+
+    def __init__(self, max_recent: int = 128):
+        self._timings: Dict[str, TaskTiming] = {}
+        self._completed: Deque[TaskTiming] = deque(maxlen=max_recent)
+
+        self._ttft_ms_sum = 0.0
+        self._ttft_ms_count = 0
+        self._decode_tps_sum = 0.0
+        self._decode_tps_count = 0
+        self._e2e_ms_sum = 0.0
+        self._e2e_ms_count = 0
+
+    def register(self, task_id: str):
+        """Create a timing record for a newly-created task."""
+        self._timings[task_id] = TaskTiming(task_id=task_id, arrival_time=time.time())
+
+    def mark_finished(self, task_id: str, input_tokens: int, output_tokens: int):
+        """Close timing for a finished/aborted task and move it to completed."""
+        timing = self._timings.pop(task_id, None)
+        if timing is None:
+            return
+        timing.finish_time = time.time()
+        timing.input_tokens = input_tokens
+        timing.output_tokens = output_tokens
+        self._completed.append(timing)
+        self._accumulate(timing)
+
+    def clear(self):
+        """Reset all state (e.g. on engine shutdown)."""
+        self._timings.clear()
+        self._completed.clear()
+        self._ttft_ms_sum = 0.0
+        self._ttft_ms_count = 0
+        self._decode_tps_sum = 0.0
+        self._decode_tps_count = 0
+        self._e2e_ms_sum = 0.0
+        self._e2e_ms_count = 0
+
+    # timing scopes
+
+    @contextmanager
+    def record(
+        self, task_ids: List[str], phase: Literal["prefill", "decode"]
+    ) -> Generator[None, None, None]:
+        tic = time.time()
+        yield
+        toc = time.time()
+        dt = toc - tic
+        for tid in task_ids:
+            t = self._timings.get(tid)
+            if t is None:
+                continue
+            if phase == "prefill":
+                t.prefill_start_time = tic
+                t.first_token_time = toc
+            elif phase == "decode":
+                t._decode_steps += 1
+                t._decode_total_s += dt
+
+    # access
+
+    def get_timing(self, task_id: str) -> Optional[TaskTiming]:
+        """Return the timing record for *task_id* (active or completed)."""
+        if task_id in self._timings:
+            return self._timings[task_id]
+        for t in self._completed:
+            if t.task_id == task_id:
+                return t
+        return None
+
+    # aggregate stats
+
+    def get_stats(self) -> Dict[str, Any]:
+        stats: Dict[str, Any] = {}
+        if self._ttft_ms_count > 0:
+            stats["avg_ttft_ms"] = round(self._ttft_ms_sum / self._ttft_ms_count, 2)
+        if self._decode_tps_count > 0:
+            stats["avg_decode_tps"] = round(
+                self._decode_tps_sum / self._decode_tps_count, 2
+            )
+        if self._e2e_ms_count > 0:
+            stats["avg_e2e_latency_ms"] = round(
+                self._e2e_ms_sum / self._e2e_ms_count, 2
+            )
+        if self._completed:
+            stats["recent_tasks"] = [t.to_dict() for t in self._completed]
+        return stats
+
+    # internal
+
+    def _accumulate(self, t: TaskTiming):
+        if t.ttft_ms is not None:
+            self._ttft_ms_sum += t.ttft_ms
+            self._ttft_ms_count += 1
+        if t.decode_tps is not None:
+            self._decode_tps_sum += t.decode_tps
+            self._decode_tps_count += 1
+        if t.e2e_latency_ms is not None:
+            self._e2e_ms_sum += t.e2e_latency_ms
+            self._e2e_ms_count += 1

astrai/inference/network/__init__.py

@@ -4,8 +4,7 @@
 lazy singleton FastAPI instance.
 """
 
-from astrai.inference.api.protocol import GenContext, ProtocolHandler, StopChecker
-from astrai.inference.api.server import (
+from astrai.inference.network.app import (
     AnthropicMessage,
     ChatCompletionRequest,
     ChatMessage,
@@ -15,7 +14,8 @@ from astrai.inference.api.server import (
     get_app,
     run_server,
 )
-from astrai.inference.api.tool_parser import (
+from astrai.inference.network.protocol import GenContext, ProtocolHandler, StopChecker
+from astrai.inference.network.tool_parser import (
     BaseToolParser,
     SimpleJsonToolParser,
     ToolParserFactory,

astrai/inference/network/anthropic.py

@@ -6,13 +6,13 @@ from typing import Any, Dict, List, Tuple, Union
 
 from pydantic import BaseModel
 
-from astrai.inference.api.protocol import (
+from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.protocol import (
     GenContext,
     ResponseBuilder,
     StopInfo,
     sse_event,
 )
-from astrai.inference.engine import InferenceEngine
 
 
 def _extract_text(content: Union[str, List[Dict[str, Any]]]) -> str:

astrai/inference/network/app.py

@@ -18,10 +18,10 @@ import uvicorn
 from fastapi import APIRouter, FastAPI, HTTPException
 from pydantic import BaseModel, Field
 
-from astrai.inference.api.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
-from astrai.inference.api.protocol import ProtocolHandler
 from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.anthropic import AnthropicResponseBuilder
+from astrai.inference.network.openai import OpenAIResponseBuilder
+from astrai.inference.network.protocol import ProtocolHandler
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 

astrai/inference/network/openai.py

@@ -7,14 +7,14 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 from pydantic import BaseModel
 
-from astrai.inference.api.protocol import (
+from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.protocol import (
     GenContext,
     ResponseBuilder,
     StopInfo,
     sse_event,
 )
-from astrai.inference.api.tool_parser import BaseToolParser, ToolParserFactory
-from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.tool_parser import BaseToolParser, ToolParserFactory
 
 logger = logging.getLogger(__name__)
 

astrai/inference/network/protocol.py

@@ -181,12 +181,10 @@ class ProtocolHandler:
         self, agen: AsyncGenerator, ctx: GenContext, stop_sequences: List[str]
     ) -> Dict[str, Any]:
         checker = StopChecker(stop_sequences)
-        chunks: List[str] = []
         body = ""
         matched = None
 
         async for token in agen:
-            chunks.append(token)
             body += token
 
             matched = checker.check(body)
@@ -195,6 +193,5 @@ class ProtocolHandler:
 
             ctx.completion_tokens += 1
 
-        content = "".join(chunks)
         stop = StopInfo(matched=matched, body=body)
-        return self.builder.format_response(ctx, content, stop)
+        return self.builder.format_response(ctx, body, stop)

astrai/inference/runtime/__init__.py

@@ -0,0 +1,25 @@
+"""Execution primitives: forward passes, CUDA graphs, and sampling."""
+
+from astrai.inference.runtime.executor import Executor
+from astrai.inference.runtime.graph import CudaGraphContext
+from astrai.inference.runtime.sample import (
+    BaseSamplingStrategy,
+    FrequencyPenaltyStrategy,
+    SamplingPipeline,
+    TemperatureStrategy,
+    TopKStrategy,
+    TopPStrategy,
+    sample,
+)
+
+__all__ = [
+    "Executor",
+    "CudaGraphContext",
+    "BaseSamplingStrategy",
+    "FrequencyPenaltyStrategy",
+    "SamplingPipeline",
+    "TemperatureStrategy",
+    "TopKStrategy",
+    "TopPStrategy",
+    "sample",
+]

astrai/inference/runtime/executor.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import time
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -9,33 +8,41 @@ import torch
 from torch import Tensor
 
 from astrai.extension.attention_backend import (
-    ATTN_BACKEND,
     CudaBackend,
-    attn_backend,
     get_backend,
 )
-from astrai.inference.core.cache import PagePool
-from astrai.inference.core.graph import CudaGraphContext
-from astrai.inference.core.task import Task
-from astrai.inference.core.workspace import InferenceWorkspace
-from astrai.inference.sample import sample
+from astrai.inference.cache import PagePool, TaskCacheManager
+from astrai.inference.runtime.graph import CudaGraphContext
+from astrai.inference.runtime.sample import sample
+from astrai.inference.task import Task
+from astrai.inference.workspace import InferenceWorkspace
 from astrai.model.automodel import AutoModel
-from astrai.tokenize.tokenizer import AutoTokenizer
 
 logger = logging.getLogger(__name__)
-_TIMED = os.environ.get("ASTRAI_TIMED", "") == "1"
 
 
 @contextmanager
 def timed(label: str, log: Optional[logging.Logger] = None):
-    """Wall-clock debug timer, enabled via ``ASTRAI_TIMED=1``."""
-    if not _TIMED:
+    """GPU-precise timer via CUDA events; falls back to perf_counter on CPU."""
+    log = log or logger
+    if not log.isEnabledFor(logging.DEBUG):
         yield
         return
-    tic = time.perf_counter()
+    use_cuda = torch.cuda.is_available()
+    if use_cuda:
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+    else:
+        tic = time.perf_counter()
     yield
-    elapsed_ms = (time.perf_counter() - tic) * 1000
-    (log or logger).info("%s %.1fms", label, elapsed_ms)
+    if use_cuda:
+        end.record()
+        torch.cuda.synchronize()
+        elapsed_ms = start.elapsed_time(end)
+    else:
+        elapsed_ms = (time.perf_counter() - tic) * 1000
+    log.debug("%s %.2fms", label, elapsed_ms)
 
 
 @dataclass
@@ -54,6 +61,19 @@ class SamplingBatchInfo:
     has_freq: bool  # any frequency_penalty != 0 (avoids per-step GPU .any())
 
 
+@dataclass
+class DecodeSteadyState:
+    """Cached decode metadata for the steady-state case.
+
+    When the same ordered task set decodes one token per step, sampling
+    params and task signature are reused; only positions advance by 1.
+    """
+
+    task_sig: tuple
+    positions: list[int]
+    sampling_info: SamplingBatchInfo
+
+
 def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
     pin = str(device).startswith("cuda")
     freq_penalties = torch.tensor(
@@ -77,12 +97,37 @@ def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
 def _warmup_cuda_graphs(
     model: AutoModel,
     pool: PagePool,
+    task_cache: TaskCacheManager,
     ws: InferenceWorkspace,
-    gctx: "CudaGraphContext",
+    gctx: CudaGraphContext,
     max_batch_size: int,
     prompt_len: int = 1,
     device: Optional[str] = None,
 ):
+    dev = device or next(model.parameters()).device
+
+    # Prefill warmup: cuBLAS auto-tunes for the actual prompt-length tensor
+    # shapes on first call (F.linear is the dominant cost).  This also warms
+    # up the CUDA context (driver init) and compiles the graph-capture trace
+    # that follows.  Custom .so kernels do NOT need this — they are pre-built.
+    warmup_len = 64
+    tid = "_warmup_prefill"
+    if task_cache.task_alloc(tid, list(range(warmup_len))):
+        with (
+            torch.inference_mode(),
+            timed("warmup prefill", logger),
+        ):
+            kv = task_cache.bind([tid], ws, start_pos=0)
+            ids_in = torch.arange(warmup_len, device=dev).unsqueeze(0)
+            pos_in = ids_in
+            model(
+                ids_in,
+                input_mask=pos_in.unsqueeze(-1) >= torch.arange(warmup_len, device=dev),
+                kv_cache=kv,
+                position_ids=pos_in,
+            )
+        task_cache.task_free(tid)
+
     batch_sizes = [1]
     n = 2
     while n <= max_batch_size:
@@ -91,47 +136,29 @@ def _warmup_cuda_graphs(
     if max_batch_size not in batch_sizes:
         batch_sizes.append(max_batch_size)
 
-    dev = device or next(model.parameters()).device
-
     for b in batch_sizes:
-        task_ids = [f"_gr_{b}_{i}" for i in range(b)]
+        task_ids = [f"_warmup_decode_{b}_{i}" for i in range(b)]
         prompt_tokens = [list(range(prompt_len)) for _ in range(b)]
         alloc_ok = True
         for tid, pt in zip(task_ids, prompt_tokens):
-            if not pool.task_alloc(tid, pt):
+            if not task_cache.task_alloc(tid, pt):
                 alloc_ok = False
                 break
         if not alloc_ok:
             for tid in task_ids:
-                pool.task_free(tid)
+                task_cache.task_free(tid)
             continue
 
         with (
             torch.inference_mode(),
-            attn_backend(ATTN_BACKEND.CUDA),
-            timed(f"warmup prefill b={b}", logger),
-        ):
-            kv_cache = pool.bind_tasks(task_ids, ws, start_pos=0)
-            ids_in = torch.tensor(prompt_tokens, dtype=torch.long, device=dev)
-            pos_in = torch.arange(prompt_len, device=dev).unsqueeze(0).expand(b, -1)
-            model(
-                ids_in,
-                input_mask=pos_in.unsqueeze(-1) >= torch.arange(prompt_len, device=dev),
-                kv_cache=kv_cache,
-                position_ids=pos_in,
-            )
-
-        with (
-            torch.inference_mode(),
-            attn_backend(ATTN_BACKEND.CUDA),
             timed(f"warmup decode b={b}", logger),
         ):
             for step in range(2):
-                seq_pos = prompt_len + step
+                seq_pos = step
                 ws.position_ids[:b] = seq_pos
                 for tid in task_ids:
-                    pool.task_extend(tid, seq_pos)
-                kv = pool.bind_tasks(task_ids, ws)
+                    task_cache.task_extend(tid, seq_pos)
+                kv = task_cache.bind(task_ids, ws)
                 input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)
                 ids_buf = ws.fill_input_ids([step] * b)
                 gctx.forward(
@@ -144,7 +171,7 @@ def _warmup_cuda_graphs(
                 )
 
         for tid in task_ids:
-            pool.task_free(tid)
+            task_cache.task_free(tid)
         torch.cuda.synchronize()
 
 
@@ -154,22 +181,22 @@ class Executor:
     def __init__(
         self,
         model: AutoModel,
-        tokenizer: AutoTokenizer,
         kv_cache: PagePool,
+        task_cache: TaskCacheManager,
         device: Optional[str] = None,
         dtype: Optional[torch.dtype] = None,
+        enable_cuda_graph: bool = True,
     ):
         self.model = model
-        self.tokenizer = tokenizer
         self.kv_cache = kv_cache
+        self.task_cache = task_cache
         self.device = device or next(model.parameters()).device
         self.dtype = dtype or next(model.parameters()).dtype
 
-        # Per-step decode cache for the steady-state case where the same
-        # ordered task set decodes one token per step.  Sampling params are
-        # constant across steps; position_ids grows by exactly 1.  Single-slot:
-        # any task-set change is a cache miss.
-        self._decode_cache: Optional[tuple] = None
+        # Per-step decode cache for the steady-state case (same ordered
+        # task set decodes one token per step).  Sampling params stay
+        # constant; only positions advance.
+        self._decode_cache: Optional[DecodeSteadyState] = None
 
         # Pre-allocated fixed-shape buffers for the decode hot path
         # (input_ids, decode mask, KV bind metadata).  Eagerly sized at init
@@ -178,8 +205,10 @@ class Executor:
         config = model.config
         max_q_heads = config.num_attention_heads
         head_dim = config.hidden_size // config.num_attention_heads
-        self._head_dim = head_dim
-        self._graph_supported = CudaBackend.supports(head_dim=head_dim)
+        backend = get_backend()
+        self._graph_supported = backend.supports_graph() and CudaBackend.supports(
+            head_dim=head_dim
+        )
         self._workspace = InferenceWorkspace(
             max_batch_size=kv_cache.max_batch_size,
             max_seq_len=kv_cache.max_seq_len,
@@ -193,7 +222,8 @@ class Executor:
         # Enabled at init-time via _warmup_cuda_graphs for CudaBackend
         # on supported head_dims; left disabled otherwise.
         self._graph_ctx = CudaGraphContext()
-        self._try_enable_cuda_graph()
+        if enable_cuda_graph:
+            self._try_enable_cuda_graph()
 
     def _try_enable_cuda_graph(self):
         if not self._graph_supported:
@@ -203,12 +233,17 @@ class Executor:
         _warmup_cuda_graphs(
             self.model,
             self.kv_cache,
+            self.task_cache,
             self._workspace,
             self._graph_ctx,
             max_batch_size=self.kv_cache.max_batch_size,
             device=self.device,
         )
 
+    @property
+    def cuda_graph_enabled(self) -> bool:
+        return self._graph_ctx.enabled and self._graph_supported
+
     def _sample_logits(
         self,
         logits: Tensor,
@@ -296,7 +331,7 @@ class Executor:
                 input_ids,
                 input_mask=input_mask,
                 position_ids=position_ids,
-                kv_cache=self.kv_cache.bind_tasks(
+                kv_cache=self.task_cache.bind(
                     task_ids,
                     self._workspace,
                     start_pos=start_pos,
@@ -338,28 +373,27 @@ class Executor:
         task_ids = [t.task_id for t in tasks]
         cur_positions = [t.next_pos for t in tasks]
 
-        sig = tuple(task_ids)
-        cached = self._decode_cache
-        if (
-            cached is not None
-            and cached[0] == sig
-            and cur_positions == [p + 1 for p in cached[1]]
-        ):
-            info = cached[2]
+        kv_cache = self.task_cache.bind(task_ids, ws)
+
+        task_sig = tuple(task_ids)
+        reuse_decode_state = (
+            self.task_cache.bind_was_steady
+            and self._decode_cache is not None
+            and self._decode_cache.task_sig == task_sig
+        )
+        if reuse_decode_state:
+            info = self._decode_cache.sampling_info
             ws.position_ids[:b] += 1
-            self._decode_cache = (sig, cur_positions, info)
         else:
             info = _build_sampling_batch_info(tasks, self.device)
             ws.position_ids[:b].copy_(
                 torch.tensor(cur_positions, dtype=torch.long, device=self.device)
             )
-            self._decode_cache = (sig, cur_positions, info)
+        self._decode_cache = DecodeSteadyState(task_sig, cur_positions, info)
 
         total_len = max(cur_positions) + 1
         input_mask = ws.decode_mask(ws.position_ids[:b], total_len)
 
-        kv_cache = self.kv_cache.bind_tasks(task_ids, ws)
-
         # ---- forward (graph replay or live run + capture) ----
 
         use_graph = (

astrai/inference/runtime/sample.py

@@ -363,20 +363,6 @@ def sample(
         ``True`` — a ``(token_ids, chosen_logprobs)`` tuple where
         ``chosen_logprobs`` has shape ``[batch]``.
     """
-    greedy = (
-        bool((temperature == 0).all())
-        if isinstance(temperature, Tensor)
-        else temperature == 0
-    )
-
-    if greedy:
-        tokens = logits.argmax(dim=-1)
-        if not return_logprobs:
-            return tokens
-        log_probs = torch.log_softmax(logits.float(), dim=-1)
-        chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
-        return tokens, chosen
-
     has_freq = (
         (isinstance(frequency_penalty, Tensor) and (frequency_penalty != 0).any())
         if isinstance(frequency_penalty, Tensor)

astrai/inference/scheduler.py

@@ -1,13 +1,21 @@
 import logging
 import threading
 import uuid
-from typing import Any, Dict, List, Optional, Tuple
+from contextlib import nullcontext
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 
-from astrai.inference.core.cache import PagePool
-from astrai.inference.core.executor import Executor
-from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
+from astrai.extension import (
+    ATTN_BACKEND,
+    AttentionBackend,
+    attn_backend,
+    get_backend,
+)
+from astrai.inference.cache import PagePool, TaskCacheManager
+from astrai.inference.metrics import MetricsCollector
+from astrai.inference.runtime.executor import Executor
+from astrai.inference.task import STOP, Task, TaskManager, TaskStatus
 from astrai.model.automodel import AutoModel
 from astrai.tokenize.tokenizer import AutoTokenizer
 
@@ -26,6 +34,8 @@ class InferenceScheduler:
         device: Optional[str] = None,
         dtype: Optional[torch.dtype] = None,
         cache: Optional[PagePool] = None,
+        enable_cuda_graph: bool = True,
+        backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None,
     ):
         config = model.config
 
@@ -56,19 +66,42 @@ class InferenceScheduler:
                 dtype=self.dtype,
             )
 
+        self._metrics = MetricsCollector()
+
+        self._task_cache = TaskCacheManager(self._cache)
+
         self._task_mgr = TaskManager(
             tokenizer=tokenizer,
             max_batch_size=max_batch_size,
             max_seq_len=self.max_seq_len,
+            metrics=self._metrics,
         )
 
-        self._executor = Executor(
-            model=model,
-            tokenizer=tokenizer,
-            kv_cache=self._cache,
-            device=self.device,
-            dtype=self.dtype,
-        )
+        if backend is None:
+            self._backend = None
+            default_backend = get_backend()
+            self._backend_name = type(default_backend).__name__
+            with attn_backend(default_backend):
+                self._executor = Executor(
+                    model=model,
+                    kv_cache=self._cache,
+                    task_cache=self._task_cache,
+                    device=self.device,
+                    dtype=self.dtype,
+                    enable_cuda_graph=enable_cuda_graph,
+                )
+        else:
+            with attn_backend(backend):
+                self._backend = get_backend()
+                self._backend_name = type(self._backend).__name__
+                self._executor = Executor(
+                    model=model,
+                    kv_cache=self._cache,
+                    task_cache=self._task_cache,
+                    device=self.device,
+                    dtype=self.dtype,
+                    enable_cuda_graph=enable_cuda_graph,
+                )
 
         self._stop_event = threading.Event()
         self._loop_thread: Optional[threading.Thread] = None
@@ -78,11 +111,31 @@ class InferenceScheduler:
 
     def remove_task(self, task_id: str):
         for task in self._task_mgr.remove_task(task_id):
-            self._cache.task_free(task.task_id)
+            self._task_cache.task_free(task.task_id)
 
     def get_stats(self) -> Dict[str, Any]:
         return self._task_mgr.get_stats()
 
+    @property
+    def backend_name(self) -> str:
+        return self._backend_name
+
+    @property
+    def cuda_graph_enabled(self) -> bool:
+        return self._executor.cuda_graph_enabled
+
+    def _backend_context(self):
+        if self._backend is None:
+            return nullcontext()
+        return attn_backend(self._backend)
+
+    @staticmethod
+    def _task_backend_groups(tasks: List[Task]):
+        groups = {}
+        for task in tasks:
+            groups.setdefault(task.backend, (task.backend, []))[1].append(task)
+        return groups.values()
+
     def _step(
         self, tasks: List[Task], return_logprobs: bool = False
     ) -> Tuple[List[Task], List[Task]]:
@@ -106,32 +159,45 @@ class InferenceScheduler:
             already appended to ``output_ids``) and tasks that hit the
             sequence cap and were marked ``ABORTED``.
         """
-        cache = self._cache
-
-        to_prefill = [t for t in tasks if t.output_tokens == 0 and t.prompt_ids]
+        to_prefill = [t for t in tasks if not t.prefill_done and t.prompt_ids]
         prefilled_ids = set()
         produced: List[Task] = []
         if to_prefill:
             for t in to_prefill:
                 t.input_tokens = len(t.prompt_ids)
 
-            groups: Dict[Tuple[int, int], List[Task]] = {}
+            groups: Dict[Tuple[int, int, Optional[AttentionBackend]], List[Task]] = {}
             for t in to_prefill:
-                start_pos = min(cache.task_cached(t.task_id), len(t.prompt_ids) - 1)
-                groups.setdefault((len(t.prompt_ids), start_pos), []).append(t)
-
-            for (prompt_len, start_pos), group in groups.items():
-                prefilled, step_out = self._executor.execute_prefill(
-                    group, prompt_len, start_pos, return_logprobs=return_logprobs
+                start_pos = min(
+                    self._task_cache.task_cached(t.task_id), len(t.prompt_ids) - 1
                 )
+                groups.setdefault((len(t.prompt_ids), start_pos, t.backend), []).append(
+                    t
+                )
+
+            for (prompt_len, start_pos, _), group in groups.items():
+                backend = group[0].backend
+                backend_context = (
+                    attn_backend(backend) if backend is not None else nullcontext()
+                )
+                with (
+                    backend_context,
+                    self._metrics.record([t.task_id for t in group], "prefill"),
+                ):
+                    prefilled, step_out = self._executor.execute_prefill(
+                        group, prompt_len, start_pos, return_logprobs=return_logprobs
+                    )
+
                 for t, out in zip(prefilled, step_out):
                     t.output_ids.append(out[0] if return_logprobs else out)
                     t.output_tokens += 1
+                    t.mark_prefill_done()
                     prefilled_ids.add(t.task_id)
                     produced.append(t)
-                start_logical_page = start_pos // getattr(cache, "page_size", 64)
+
+                start_logical_page = start_pos // self._cache.page_size
                 for t in group:
-                    cache.task_record_hashes(
+                    self._task_cache.task_record_hashes(
                         t.task_id, t.prompt_ids, start_logical_page
                     )
 
@@ -140,79 +206,86 @@ class InferenceScheduler:
         for t in tasks:
             if t.task_id in prefilled_ids:
                 continue
-            if cache.task_extend(t.task_id, t.next_pos):
+            if self._task_cache.task_extend(t.task_id, t.next_pos):
                 decoded.append(t)
             else:
                 t.status = TaskStatus.ABORTED
                 aborted.append(t)
 
-        if decoded:
-            step_out = self._executor.execute_decode(
-                decoded, return_logprobs=return_logprobs
+        for backend, group in self._task_backend_groups(decoded):
+            backend_context = (
+                attn_backend(backend) if backend is not None else nullcontext()
             )
-            for t, out in zip(decoded, step_out):
+            with (
+                backend_context,
+                self._metrics.record([t.task_id for t in group], "decode"),
+            ):
+                step_out = self._executor.execute_decode(
+                    group, return_logprobs=return_logprobs
+                )
+            for t, out in zip(group, step_out):
                 t.output_ids.append(out[0] if return_logprobs else out)
                 t.output_tokens += 1
+                t.advance_kv()
                 produced.append(t)
 
         return produced, aborted
 
     def _run_generation_loop(self):
         stop_ids = self._task_mgr.tokenizer.stop_ids
-        cache = self._cache
         try:
-            while not self._stop_event.is_set():
-                finished = self._task_mgr.remove_finished_tasks(stop_ids)
-                for task in finished:
-                    if task.status == TaskStatus.FINISHED:
-                        cache.task_record_hashes(
-                            task.task_id,
-                            cache.task_cacheable_ids(
-                                task.task_id, task.prompt_ids, task.output_ids
-                            ),
-                        )
-                    cache.task_free(task.task_id)
+            with self._backend_context():
+                while not self._stop_event.is_set():
+                    finished = self._task_mgr.remove_finished_tasks(stop_ids)
+                    for task in finished:
+                        if task.status == TaskStatus.FINISHED:
+                            self._task_cache.task_record_hashes(
+                                task.task_id,
+                                self._task_cache.task_cacheable_ids(
+                                    task.task_id, task.prompt_ids, task.output_ids
+                                ),
+                            )
+                        self._task_cache.task_free(task.task_id)
 
-                active = self._task_mgr.get_active_tasks()
-                available = self._task_mgr.max_batch_size - len(active)
-                if available > 0:
-                    candidates = self._task_mgr.pull_candidates(available)
-                    failed = []
-                    for task in candidates:
-                        if cache.task_alloc(task.task_id, task.prompt_ids):
-                            self._task_mgr.activate(task)
-                        else:
-                            failed.append(task)
-                    if failed:
-                        self._task_mgr.return_to_waiting(failed)
+                    active = self._task_mgr.get_active_tasks()
+                    available = self._task_mgr.max_batch_size - len(active)
+                    if available > 0:
+                        candidates = self._task_mgr.pull_candidates(available)
+                        failed = []
+                        for task in candidates:
+                            if self._task_cache.task_alloc(
+                                task.task_id, task.prompt_ids
+                            ):
+                                self._task_mgr.activate(task)
+                            else:
+                                failed.append(task)
+                        if failed:
+                            self._task_mgr.return_to_waiting(failed)
 
-                if not self._task_mgr.has_work():
-                    self._task_mgr.wait_for_tasks(timeout=1.0)
-                    continue
+                    if not self._task_mgr.has_work():
+                        self._task_mgr.wait_for_tasks(timeout=1.0)
+                        continue
 
-                active = self._task_mgr.get_active_tasks()
+                    active = self._task_mgr.get_active_tasks()
 
-                decoded, aborted = self._step(active)
+                    decoded, aborted = self._step(active)
 
-                for t in aborted:
-                    self._task_mgr.invoke_callback(t.task_id, STOP)
-
-                for t in decoded:
-                    new_text = t.decode_new_token(self._task_mgr.tokenizer)
-                    if new_text:
-                        self._task_mgr.invoke_callback(t.task_id, new_text)
-                    if t.is_finished(stop_ids):
-                        remaining = t.flush_remaining(self._task_mgr.tokenizer)
-                        if remaining:
-                            self._task_mgr.invoke_callback(t.task_id, remaining)
+                    for t in aborted:
                         self._task_mgr.invoke_callback(t.task_id, STOP)
 
+                    for t in decoded:
+                        new_text = t.decode_new_token(self._task_mgr.tokenizer)
+                        if new_text:
+                            self._task_mgr.invoke_callback(t.task_id, new_text)
+                        if t.is_finished(stop_ids):
+                            self._task_mgr.invoke_callback(t.task_id, STOP)
+
         except Exception as e:
             self._stop_event.set()
             logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
             for task in self._task_mgr.get_active_tasks():
                 self._task_mgr.invoke_callback(task.task_id, STOP)
-                cache.task_free(task.task_id)
+                self._task_cache.task_free(task.task_id)
             for task in self._task_mgr.get_waiting_tasks():
                 self._task_mgr.invoke_callback(task.task_id, STOP)
             self._task_mgr.clear_queues()
@@ -233,10 +306,10 @@ class InferenceScheduler:
         self._loop_thread = None
         for task in self._task_mgr.get_active_tasks():
             self._task_mgr.invoke_callback(task.task_id, STOP)
-            self._cache.task_free(task.task_id)
+            self._task_cache.task_free(task.task_id)
         for task in self._task_mgr.get_waiting_tasks():
             self._task_mgr.invoke_callback(task.task_id, STOP)
-            self._cache.task_free(task.task_id)
+            self._task_cache.task_free(task.task_id)
         self._task_mgr.clear_queues()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -275,8 +348,8 @@ class InferenceScheduler:
             ``List[Tuple[List[int], List[float]]]``.
         """
         stop_ids = self._task_mgr.tokenizer.stop_ids
-        cache = self._cache
         seq_cap = self.max_seq_len
+        request_backend = get_backend(use_default=False)
 
         tasks: List[Task] = []
         for ids in prompt_ids_list:
@@ -300,23 +373,29 @@ class InferenceScheduler:
                 top_k=top_k,
                 frequency_penalty=frequency_penalty,
                 rep_window=rep_window,
+                backend=request_backend,
             )
-            if not cache.task_alloc(task.task_id, task.prompt_ids):
+            if not self._task_cache.task_alloc(task.task_id, task.prompt_ids):
                 tasks.append(None)
                 continue
             task.input_tokens = len(task.prompt_ids)
+            self._metrics.register(task.task_id)
             tasks.append(task)
 
         try:
             live = [t for t in tasks if t is not None]
 
-            while live:
-                decoded, _ = self._step(live, return_logprobs=return_logprobs)
-                live = [t for t in decoded if not t.is_finished(stop_ids)]
+            with self._backend_context():
+                while live:
+                    decoded, _ = self._step(live, return_logprobs=return_logprobs)
+                    live = [t for t in decoded if not t.is_finished(stop_ids)]
         finally:
             for t in tasks:
                 if t is not None:
-                    cache.task_free(t.task_id)
+                    self._metrics.mark_finished(
+                        t.task_id, t.input_tokens, t.output_tokens
+                    )
+                    self._task_cache.task_free(t.task_id)
 
         results: List[Any] = []
         for t in tasks:

astrai/inference/task.py

@@ -1,16 +1,17 @@
-import logging
 import threading
 import time
 import uuid
 from collections import deque
 from enum import Enum
-from typing import Any, Callable, Deque, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, List, Optional
 
 from tokenizers.decoders import DecodeStream
 
+from astrai.inference.metrics import MetricsCollector
 from astrai.tokenize.tokenizer import AutoTokenizer
 
-logger = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from astrai.extension import AttentionBackend
 
 STOP = object()
 
@@ -64,6 +65,7 @@ class Task:
         top_k: int = 50,
         frequency_penalty: float = 0.0,
         rep_window: int = 64,
+        backend: Optional["AttentionBackend"] = None,
     ):
         self.task_id = task_id
         self.prompt_ids = prompt_ids
@@ -73,16 +75,25 @@ class Task:
         self.top_k = top_k
         self.frequency_penalty = frequency_penalty
         self.rep_window = rep_window
+        self.backend = backend
 
         self.status = TaskStatus.PENDING
         self.output_ids: List[int] = []
         self.output_logprobs: List[float] = []
         self.input_tokens: int = 0
         self.output_tokens: int = 0
-        self.arrival_time = time.time()
-        self.finish_time: Optional[float] = None
+        self._kv_len: int = 0
         self._decoder: Optional[StreamDecoder] = None
 
+    def mark_prefill_done(self):
+        """Prompt KV is materialized by prefill; first output sampled but
+        not yet written to KV."""
+        self._kv_len = self.input_tokens
+
+    def advance_kv(self):
+        """One more position written to KV (after a decode forward)."""
+        self._kv_len += 1
+
     def decode_new_token(self, tokenizer: AutoTokenizer) -> str:
         """Decode the last appended output token, buffering incomplete
         multi-byte sequences across calls.
@@ -93,20 +104,15 @@ class Task:
             self._decoder = StreamDecoder(tokenizer)
         return self._decoder.push(self.output_ids[-1])
 
-    def flush_remaining(self, tokenizer: AutoTokenizer) -> str:
-        """Emit any text still buffered in the decoder.
-
-        With the Rust-native DecodeStream, the stream is always in a
-        correct state — any completed text was already emitted by the
-        last ``push``.  A trailing incomplete multi-byte sequence has no
-        valid text to emit, so this is a no-op.
-        """
-        return ""
-
     @property
     def next_pos(self) -> int:
-        # The first output is sampled from prefill and enters KV on the next step.
-        return self.input_tokens + max(0, len(self.output_ids) - 1)
+        """KV position where the next decode step will write."""
+        return self._kv_len
+
+    @property
+    def prefill_done(self) -> bool:
+        """True when all prompt KV entries are materialized."""
+        return self._kv_len >= self.input_tokens > 0
 
     def is_finished(self, stop_ids: List[int]) -> bool:
         if self.max_tokens is not None and self.output_tokens >= self.max_tokens:
@@ -124,6 +130,7 @@ class TaskManager:
         tokenizer: AutoTokenizer,
         max_batch_size: int = 16,
         max_seq_len: int = 8192,
+        metrics: Optional["MetricsCollector"] = None,
     ):
         self.tokenizer = tokenizer
         self.max_batch_size = max_batch_size
@@ -139,6 +146,8 @@ class TaskManager:
         self._total_tasks = 0
         self._total_tokens = 0
 
+        self._metrics = metrics
+
     def add_task(
         self,
         prompt: str,
@@ -148,6 +157,7 @@ class TaskManager:
         top_k: int = 50,
         frequency_penalty: float = 0.0,
         rep_window: int = 64,
+        backend: Optional["AttentionBackend"] = None,
         stream_callback: Optional[Callable[[str], None]] = None,
     ) -> str:
         task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
@@ -155,11 +165,6 @@ class TaskManager:
         if len(prompt_ids) > self.max_seq_len:
             prompt_ids = prompt_ids[-self.max_seq_len :]
 
-        if len(prompt_ids) > self.max_seq_len:
-            if stream_callback:
-                stream_callback(STOP)
-            return task_id
-
         if max_tokens is None:
             max_tokens = self.max_seq_len - len(prompt_ids)
         else:
@@ -174,6 +179,7 @@ class TaskManager:
             top_k=top_k,
             frequency_penalty=frequency_penalty,
             rep_window=rep_window,
+            backend=backend,
         )
 
         with self._lock:
@@ -182,6 +188,9 @@ class TaskManager:
             if stream_callback:
                 self._callbacks[task_id] = stream_callback
 
+        if self._metrics is not None:
+            self._metrics.register(task_id)
+
         self._task_event.set()
         return task_id
 
@@ -201,26 +210,33 @@ class TaskManager:
             cb(token)
 
     def get_stats(self) -> Dict[str, Any]:
-        return {
+        stats: Dict[str, Any] = {
             "total_tasks": self._total_tasks,
             "total_tokens": self._total_tokens,
             "active_tasks": len(self.active_tasks),
             "waiting_queue": len(self.waiting_queue),
         }
+        if self._metrics is not None:
+            stats.update(self._metrics.get_stats())
+        return stats
 
     def remove_finished_tasks(self, stop_ids: List[int]) -> List[Task]:
         with self._lock:
             finished = []
             for task in self.active_tasks:
                 if task.status == TaskStatus.ABORTED:
-                    task.finish_time = time.time()
                     finished.append(task)
                 elif task.is_finished(stop_ids):
                     task.status = TaskStatus.FINISHED
-                    task.finish_time = time.time()
                     finished.append(task)
                     self._total_tokens += task.output_tokens
 
+            if self._metrics is not None:
+                for task in finished:
+                    self._metrics.mark_finished(
+                        task.task_id, task.input_tokens, task.output_tokens
+                    )
+
             self.active_tasks = [
                 t
                 for t in self.active_tasks

astrai/logging.py

@@ -0,0 +1,27 @@
+import logging
+import os
+
+
+def setup_logging(level: str = "INFO"):
+    """Attach a StreamHandler to the ``astrai`` logger (idempotent).
+
+    Call once per process at the top of CLI scripts.
+    Set ``ASTR_LOG_LEVEL`` env var to override the default level.
+
+    Level names: ``DEBUG``, ``INFO``, ``WARNING``, ``ERROR``, ``CRITICAL``.
+    ``DEBUG`` enables per-step prefill/decode timing logs
+    (:func:`astrai.inference.runtime.executor.timed`).
+    """
+    logger = logging.getLogger("astrai")
+    if logger.handlers:
+        return
+    level_name = os.environ.get("ASTR_LOG_LEVEL", level).upper()
+    logger.setLevel(getattr(logging, level_name, logging.INFO))
+    handler = logging.StreamHandler()
+    handler.setFormatter(
+        logging.Formatter(
+            "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+    )
+    logger.addHandler(handler)

astrai/model/components/attention.py

@@ -8,7 +8,7 @@ from torch import Tensor
 from astrai.extension import attention
 from astrai.extension.rotary_backend import apply_rotary_emb
 from astrai.factory import BaseFactory
-from astrai.inference.core.cache import KVCache
+from astrai.inference.cache import KVCache
 from astrai.model.components.linear import Linear
 from astrai.model.components.norm import RMSNorm
 

astrai/model/components/decoder_block.py

@@ -4,7 +4,7 @@ from typing import Optional, TypedDict
 import torch.nn as nn
 from torch import Tensor
 
-from astrai.inference.core.cache import KVCache
+from astrai.inference.cache import KVCache
 from astrai.model.components.attention import AttnFactory
 from astrai.model.components.mlp import FFNFactory, RouterStats
 from astrai.model.components.norm import RMSNorm

astrai/model/transformer.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 from torch import Tensor
 
 from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.inference.core.cache import KVCache
+from astrai.inference.cache import KVCache
 from astrai.model.automodel import AutoModel, ModelFactory
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding

astrai/trainer/metric_util.py

@@ -1,3 +1,4 @@
+import math
 from typing import Dict
 
 import torch
@@ -27,6 +28,8 @@ class GradSNRTracker:
 
     SNR = E[g]^2 / Var(g) = E[g]^2 / (E[g^2] - E[g]^2)
 
+    The reported value is the power ratio in decibels: ``10 * log10(SNR)``.
+
     The tracker accumulates per-parameter EMA moments across optimizer steps.
     Call ``update`` after backward (before ``optimizer.step``) and read
     ``snr`` to get the aggregate SNR across all parameters.
@@ -64,7 +67,8 @@ class GradSNRTracker:
             noise = (v - m.pow(2)).clamp(min=0).sum().item()
             total_signal += signal
             total_noise += noise
-        return total_signal / (total_noise + self.eps)
+        snr = total_signal / (total_noise + self.eps)
+        return 10.0 * math.log10(max(snr, self.eps))
 
 
 def ctx_get_loss(ctx):

astrai/trainer/rollout.py

@@ -6,7 +6,7 @@ Provides:
 - :class:`BaseRewardModel` — pluggable reward interface
 - :class:`RolloutGenerator` — KV-cache-backed generation of grouped
   responses + decoding (no reward); delegates the generation loop to
-  :class:`~astrai.inference.core.scheduler.InferenceScheduler.run_batch`
+  :class:`~astrai.inference.scheduler.InferenceScheduler.run_batch`
   so rollout and the production inference server share one code path
 - :class:`RolloutRunner` — orchestrates generation + scoring with a
   step-driven cache; its ``__call__`` returns ``(RolloutResult, is_fresh)``
@@ -20,7 +20,7 @@ from typing import Dict, List, Optional, Tuple
 import torch
 from torch import Tensor
 
-from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.scheduler import InferenceScheduler
 
 
 @dataclass(kw_only=True)
@@ -101,7 +101,7 @@ class RolloutGenerator:
     """Pure generation + decoding for a group of responses per prompt.
 
     Delegates the prefill/decode loop to
-    :meth:`~astrai.inference.core.scheduler.InferenceScheduler.run_batch`,
+    :meth:`~astrai.inference.scheduler.InferenceScheduler.run_batch`,
     which uses a real KV cache (no O(n²) recompute).  Has no dependency
     on any reward model; can be reused in isolation for offline
     generation, qualitative sampling, or eval pipelines.

astrai/trainer/train_context.py

@@ -10,7 +10,7 @@ from torch.utils.data import DataLoader, random_split
 
 from astrai.config.train_config import TrainConfig
 from astrai.dataset import RDSampler
-from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.scheduler import InferenceScheduler
 from astrai.model.components.lora import inject_lora
 from astrai.parallel.executor import BaseExecutor, ExecutorFactory, create_ref_model
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
@@ -66,6 +66,15 @@ class TrainContext:
         )
 
 
+@dataclass
+class _PreloadedState:
+    model_config: dict = field(default_factory=dict)
+    state_dict: Optional[dict] = None
+    epoch: int = 0
+    consumed_samples: int = 0
+    checkpoint: Optional[Checkpoint] = None
+
+
 class TrainContextBuilder:
     def __init__(
         self,
@@ -81,213 +90,231 @@ class TrainContextBuilder:
         return self
 
     def build(self) -> TrainContext:
-        cfg = self.config
-        device = get_current_device()
+        # Resolve persisted state.
+        preloaded_state = self._load_preloaded_state()
 
-        executor = ExecutorFactory.create(
+        # Build the core training components and restore their persisted state.
+        executor = self._create_executor()
+        context = self._create_context(preloaded_state, executor)
+        self._prepare_model(context, executor, preloaded_state)
+        self._restore_optimizer_state(context)
+
+        # Resolve datasets.
+        train_dataset, val_dataset = self._get_datasets()
+        self._create_dataloaders(context, train_dataset, val_dataset)
+
+        # Strategies depend on the prepared model; online rollout depends on both.
+        strategy_kwargs = self._create_strategy(context, executor)
+        self._configure_rollout(context, strategy_kwargs)
+
+        return context
+
+    def _create_executor(self) -> BaseExecutor:
+        cfg = self.config
+        return ExecutorFactory.create(
             cfg.parallel_mode,
             grad_accum_steps=cfg.grad_accum_steps,
             **cfg.executor_kwargs,
         )
 
-        model_config = {}
+    def _load_preloaded_state(self) -> _PreloadedState:
+        cfg = self.config
+        state = _PreloadedState(
+            epoch=cfg.start_epoch,
+            consumed_samples=cfg.start_samples * get_world_size(),
+        )
         if self._param_path:
             config_path = Path(self._param_path) / "config.json"
             if config_path.exists():
-                model_config = load_json(config_path)
-
-        preloaded_state_dict = None
-        preloaded_epoch = cfg.start_epoch
-        preloaded_consumed = cfg.start_samples * get_world_size()
-        preloaded_checkpoint = None
-        if self._param_path:
+                state.model_config = load_json(config_path)
             checkpoint = Checkpoint.load_any(self._param_path)
             if checkpoint is not None:
-                preloaded_state_dict = checkpoint.state_dict
-                if checkpoint.config:
-                    model_config = checkpoint.config
+                state.state_dict = checkpoint.state_dict
+                state.model_config = checkpoint.config or state.model_config
                 if self._resume:
-                    preloaded_epoch = checkpoint.epoch
+                    state.epoch = checkpoint.epoch
                     per_step = (
                         cfg.batch_per_device * get_world_size() * cfg.grad_accum_steps
                     )
-                    preloaded_consumed = (
-                        checkpoint.consumed_samples // per_step
-                    ) * per_step
-                    preloaded_checkpoint = checkpoint
+                    state.consumed_samples = (
+                        checkpoint.consumed_samples // per_step * per_step
+                    )
+                    state.checkpoint = checkpoint
+        if not state.model_config and hasattr(cfg.model_fn(), "config"):
+            state.model_config = cfg.model_fn().config.to_dict()
+        return state
 
-        if not model_config and hasattr(cfg.model_fn(), "config"):
-            model_config = cfg.model_fn().config.to_dict()
+    def _create_context(
+        self, state: _PreloadedState, executor: BaseExecutor
+    ) -> TrainContext:
+        return TrainContext(
+            world_size=get_world_size(),
+            rank=get_rank(),
+            config=self.config,
+            model_config=state.model_config,
+            executor=executor,
+            epoch=state.epoch,
+            consumed_samples=state.consumed_samples,
+            checkpoint=state.checkpoint,
+        )
 
-        def _before_wrap(m):
-            m = m.to(device=device)
+    def _prepare_model(
+        self, context: TrainContext, executor: BaseExecutor, state: _PreloadedState
+    ) -> None:
+        cfg = self.config
+        device = get_current_device()
+
+        def before_wrap(model):
+            model = model.to(device=device)
             if cfg.lora is not None:
                 inject_lora(
-                    m,
+                    model,
                     r=cfg.lora.r,
                     alpha=cfg.lora.alpha,
                     target_modules=set(cfg.lora.target_modules),
                 )
-            if preloaded_state_dict is not None:
-                m.load_state_dict(preloaded_state_dict, strict=False)
-            return m
+            if state.state_dict is not None:
+                model.load_state_dict(state.state_dict, strict=False)
+            return model
 
-        def _after_wrap(m):
+        def after_wrap(model):
             if cfg.compile_mode is not None:
                 logger.info("torch.compile enabled (mode=%s)", cfg.compile_mode)
-                m = torch.compile(m, mode=cfg.compile_mode)
-            return m
-
-        context = TrainContext(
-            world_size=get_world_size(),
-            rank=get_rank(),
-            config=cfg,
-            model_config=model_config,
-            executor=executor,
-            epoch=preloaded_epoch,
-            consumed_samples=preloaded_consumed,
-            checkpoint=preloaded_checkpoint,
-        )
+                model = torch.compile(model, mode=cfg.compile_mode)
+            return model
 
         context.model, context.optimizer, context.scheduler = executor.prepare(
             cfg.model_fn,
             cfg.optimizer_fn,
             cfg.scheduler_fn,
-            before_wrap=_before_wrap,
-            after_wrap=_after_wrap,
+            before_wrap=before_wrap,
+            after_wrap=after_wrap,
         )
 
-        train_dataset = cfg.dataset
-        val_dataset = cfg.val_dataset
+    def _get_datasets(self):
+        cfg = self.config
+        if cfg.val_dataset is not None or cfg.val_split is None:
+            return cfg.dataset, cfg.val_dataset
+        n_val = max(1, int(len(cfg.dataset) * cfg.val_split))
+        generator = torch.Generator().manual_seed(cfg.random_seed)
+        return random_split(
+            cfg.dataset, [len(cfg.dataset) - n_val, n_val], generator=generator
+        )
 
-        if val_dataset is None and cfg.val_split is not None:
-            n_total = len(cfg.dataset)
-            n_val = max(1, int(n_total * cfg.val_split))
-            n_train = n_total - n_val
-            generator = torch.Generator().manual_seed(cfg.random_seed)
-            train_dataset, val_dataset = random_split(
-                cfg.dataset, [n_train, n_val], generator=generator
+    def _create_dataloaders(
+        self, context: TrainContext, train_dataset, val_dataset
+    ) -> None:
+        cfg = self.config
+        sampler_offset = context.consumed_samples // context.world_size
+        if self._resume and sampler_offset > 0:
+            samples_per_replica = (
+                len(train_dataset) + context.world_size - 1
+            ) // context.world_size
+            if samples_per_replica > 0:
+                context.epoch = sampler_offset // samples_per_replica
+        context.dataloader = self._create_dataloader(
+            train_dataset, context.epoch, sampler_offset
+        )
+        if val_dataset is not None:
+            context.val_dataloader = self._create_dataloader(
+                val_dataset, 0, 0, shuffle=False
             )
 
-        sampler_offset = context.consumed_samples // context.world_size
-
-        if self._resume and sampler_offset > 0:
-            offset = context.world_size - 1
-            num_samples_per_replica = (
-                len(train_dataset) + offset
-            ) // context.world_size
-            if num_samples_per_replica > 0:
-                context.epoch = sampler_offset // num_samples_per_replica
-
+    def _create_dataloader(
+        self, dataset, epoch: int, start_iter: int, shuffle: bool = True
+    ):
+        cfg = self.config
         sampler = RDSampler(
-            data_source=train_dataset,
-            start_epoch=context.epoch,
-            start_iter=sampler_offset,
+            dataset,
+            start_epoch=epoch,
+            start_iter=start_iter,
             seed=cfg.random_seed,
+            shuffle=shuffle,
         )
-        context.dataloader = DataLoader(
-            train_dataset,
+        loader_kwargs = dict(
+            dataset=dataset,
             batch_size=cfg.batch_per_device,
             sampler=sampler,
             num_workers=cfg.num_workers,
             pin_memory=cfg.pin_memory,
-            prefetch_factor=cfg.prefetch_factor,
             collate_fn=cfg.collate_fn,
         )
-
-        if val_dataset is not None:
-            val_sampler = RDSampler(
-                data_source=val_dataset,
-                start_epoch=0,
-                start_iter=0,
-                seed=cfg.random_seed,
-                shuffle=False,
-            )
-            context.val_dataloader = DataLoader(
-                val_dataset,
-                batch_size=cfg.batch_per_device,
-                sampler=val_sampler,
-                num_workers=cfg.num_workers,
-                pin_memory=cfg.pin_memory,
-                prefetch_factor=cfg.prefetch_factor,
-                collate_fn=cfg.collate_fn,
-            )
-
-        if context.checkpoint and context.checkpoint.extra:
-            extra = context.checkpoint.extra
-            for name in ("optimizer", "scheduler"):
-                if name in extra:
-                    obj = getattr(context, name, None)
-                    if obj is not None:
-                        obj.load_state_dict(extra[name])
-
-        strategy_kwargs = dict(cfg.extra_kwargs)
-        strategy_kwargs.setdefault("moe_aux_loss_coef", cfg.moe_aux_loss_coef)
-
-        needs_ref = cfg.strategy in (
-            "dpo",
-            "grpo",
-            "online_grpo",
-            "online_dpo",
+        # PyTorch rejects prefetch_factor/persistent_workers when workers=0.
+        if cfg.num_workers > 0:
+            loader_kwargs["persistent_workers"] = cfg.persistent_workers
+            if cfg.prefetch_factor is not None:
+                loader_kwargs["prefetch_factor"] = cfg.prefetch_factor
+        return DataLoader(
+            **loader_kwargs,
         )
-        needs_old = cfg.strategy in ("grpo", "online_grpo")
 
-        if needs_ref:
-            strategy_kwargs["ref_model"] = create_ref_model(
-                cfg.model_fn, executor=executor, model=context.model, device=device
+    def _restore_optimizer_state(self, context: TrainContext) -> None:
+        if context.checkpoint and context.checkpoint.extra:
+            for name in ("optimizer", "scheduler"):
+                if (
+                    name in context.checkpoint.extra
+                    and getattr(context, name, None) is not None
+                ):
+                    getattr(context, name).load_state_dict(
+                        context.checkpoint.extra[name]
+                    )
+
+    def _create_strategy(self, context: TrainContext, executor: BaseExecutor) -> dict:
+        cfg = self.config
+        kwargs = dict(cfg.extra_kwargs)
+        kwargs.setdefault("moe_aux_loss_coef", cfg.moe_aux_loss_coef)
+        if cfg.strategy in ("dpo", "grpo", "online_grpo", "online_dpo"):
+            kwargs["ref_model"] = create_ref_model(
+                cfg.model_fn,
+                executor=executor,
+                model=context.model,
+                device=get_current_device(),
             )
-
-        if needs_old:
-            strategy_kwargs["old_model"] = create_ref_model(
-                cfg.model_fn, executor=executor, model=context.model, device=device
+        if cfg.strategy in ("grpo", "online_grpo"):
+            kwargs["old_model"] = create_ref_model(
+                cfg.model_fn,
+                executor=executor,
+                model=context.model,
+                device=get_current_device(),
             )
-
         context.strategy = StrategyFactory.create(
             cfg.strategy,
             model=context.model,
-            device=device,
+            device=get_current_device(),
             executor=executor,
-            **strategy_kwargs,
+            **kwargs,
         )
+        return kwargs
 
-        # Enable online rollout when the train_type is an ``online_*`` variant.
-        is_online = cfg.strategy.startswith("online_")
-        if is_online:
-            if not context.strategy.supports_online():
-                raise ValueError(
-                    f"Strategy '{cfg.strategy}' does not support online rollout"
-                )
-            if cfg.reward_model_fn is None:
-                raise ValueError("reward_model_fn is required for online RL strategies")
-
-            tokenizer = AutoTokenizer.from_pretrained(self._param_path)
-            reward_model = cfg.reward_model_fn()
-
-            group_size = strategy_kwargs.get("group_size", 1)
-            rollout_batch_size = group_size * max(1, cfg.batch_per_device)
-            max_seq_len = getattr(context.model.config, "max_position_embeddings", None)
-
-            scheduler = InferenceScheduler(
-                model=context.model,
-                tokenizer=tokenizer,
-                max_batch_size=rollout_batch_size,
-                max_seq_len=max_seq_len,
+    def _configure_rollout(self, context: TrainContext, strategy_kwargs: dict) -> None:
+        cfg = self.config
+        if not cfg.strategy.startswith("online_"):
+            return
+        if not context.strategy.supports_online():
+            raise ValueError(
+                f"Strategy '{cfg.strategy}' does not support online rollout"
             )
-
-            generator = RolloutGenerator(
-                scheduler=scheduler,
-                tokenizer=tokenizer,
-                max_tokens=cfg.rollout_max_tokens,
-                group_size=group_size,
-                temperature=cfg.rollout_temperature,
-                top_k=cfg.rollout_top_k,
-                top_p=cfg.rollout_top_p,
-            )
-            runner = RolloutRunner(
+        tokenizer = AutoTokenizer.from_pretrained(self._param_path)
+        group_size = strategy_kwargs.get("group_size", 1)
+        scheduler = InferenceScheduler(
+            model=context.model,
+            tokenizer=tokenizer,
+            max_batch_size=group_size * max(1, cfg.batch_per_device),
+            max_seq_len=getattr(context.model.config, "max_position_embeddings", None),
+        )
+        generator = RolloutGenerator(
+            scheduler=scheduler,
+            tokenizer=tokenizer,
+            max_tokens=cfg.rollout_max_tokens,
+            group_size=group_size,
+            temperature=cfg.rollout_temperature,
+            top_k=cfg.rollout_top_k,
+            top_p=cfg.rollout_top_p,
+        )
+        context.strategy.set_rollout_runner(
+            RolloutRunner(
                 generator=generator,
-                reward_model=reward_model,
+                reward_model=cfg.reward_model_fn(),
                 rollout_interval=cfg.rollout_interval,
             )
-            context.strategy.set_rollout_runner(runner)
-
-        return context
+        )

csrc/CMakeLists.txt

@@ -48,7 +48,7 @@ set(TORCH_LIBS
 
 set(CMAKE_CUDA_ARCHITECTURES "${ASTRAI_CUDA_ARCH}")
 
-set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb)
+set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb fp8_mm)
 
 foreach(name ${KERNELS})
     add_library(${name} MODULE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${name}.cu")
@@ -61,6 +61,9 @@ foreach(name ${KERNELS})
         "${PYTHON_INCLUDE_DIR}")
 
     target_link_libraries(${name} PRIVATE ${TORCH_LIBS})
+    if(${name} STREQUAL "fp8_mm")
+        target_link_libraries(${name} PRIVATE CUDA::cublasLt)
+    endif()
     target_link_options(${name} PRIVATE "-Wl,-rpath,${TORCH_LIB_DIR}")
 
     target_compile_options(${name} PRIVATE

csrc/kernels/attn_common.h

@@ -18,49 +18,52 @@ enum TensorLayout : int {
 // drift out of sync.
 template<typename T, typename AT = float>
 struct AttentionParams {
-    // ---- shared across all paths ----
+    // Shape
     int batch;
     int q_head;
     int kv_head;
     int head_dim;
-    int use_mask;
-    int causal_offset;   // -1 = non-causal; >=0 = absolute position of first Q token
-    int num_splits;
+    int q_len;   // Per-request in contiguous mode; total_q in paged mode.
+    int kv_len;  // Contiguous mode; paged mode uses kv_indptr.
+
+    // Attention behavior
     float scale;
+    // -1 = non-causal; >=0 = absolute position of first Q token
+    int causal_offset;
+    int use_mask;
 
-    // Q strides (element offsets for each dim — layout-agnostic)
-    int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
-
-    // Mask: 2D [batch, kv_len], 3D [batch, q_len, kv_len],
-    // or 4D [batch, n_heads, q_len, kv_len] (head dim broadcasts when stride=0)
-    int mask_b_stride;   // batch stride
-    int mask_h_stride;   // head stride (0 = broadcast across heads)
-    int mask_q_stride;   // q stride (0 = all q rows share)
+    // pointers
+    const T* __restrict__ q_ptr;
+    const T* __restrict__ k_ptr;
+    const T* __restrict__ v_ptr;
+    T* __restrict__ o_ptr;
     const bool* __restrict__ mask;
 
-    const T* __restrict__ q;
-    T* __restrict__ o;
+    // strides
+    int q_b_stride;
+    int q_h_stride;
+    int q_l_stride;
+    int q_d_stride;
+
+    int kv_b_stride;
+    int kv_h_stride;
+    int kv_l_stride;
+    int kv_d_stride;
+
+    int mask_b_stride;
+    int mask_h_stride;
+    int mask_l_stride;
+
+    // Paged K/V addressing
+    const int64_t* __restrict__ req_to_token;     // [num_reqs, max_context_len]
+    const int64_t* __restrict__ req_pool_indices; // [batch]
+    const int* __restrict__ kv_indptr;             // [batch + 1]
+    const int* __restrict__ qo_indptr;             // [batch + 1] or nullptr for decode
+    int max_context_len; // req_to_token stride (dim 1)
+
+    // Decode split-KV workspace
+    int num_splits;
     AT* __restrict__ o_part;
     AT* __restrict__ ml_part;
 
-    // ---- contiguous K/V mode ----
-    int q_len;
-    int kv_len;
-    int kv_stride_b, kv_stride_h, kv_stride_l, kv_stride_d;
-    const T* __restrict__ k;
-    const T* __restrict__ v;
-
-    // ---- paged (SGLang flat pool) mode ----
-    const T* __restrict__ k_cache;
-    const T* __restrict__ v_cache;
-
-    // Indexing
-    const int64_t* __restrict__ req_to_token;      // [num_reqs, max_context_len]
-    const int64_t* __restrict__ req_pool_indices;   // [batch]
-    const int* __restrict__ kv_indptr;               // [batch+1]
-    const int* __restrict__ qo_indptr;               // [batch+1] or nullptr (decode)
-    int max_context_len;  // req_to_token stride (dim 1)
-    int max_seq_len;      // max per-request seq_len (host-side, for split computation)
-    int total_q;          // total Q tokens across all requests (host-side, for grid)
-    int max_q_len;        // max per-request q_len (host-side, for prefill grid)
 };

csrc/kernels/attn_decode.cu

@@ -22,15 +22,22 @@ torch::Tensor attn_decode(
 
     auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
     auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
-    p.o = (bf16*)O_view.data_ptr();
+    p.o_ptr = (bf16*)O_view.data_ptr();
 
     if (o_part_buf.has_value() && ml_part_buf.has_value()
         && o_part_buf->defined() && ml_part_buf->defined()) {
         TORCH_CHECK(o_part_buf->scalar_type() == torch::kFloat32, "o_part_buf must be f32");
         TORCH_CHECK(ml_part_buf->scalar_type() == torch::kFloat32, "ml_part_buf must be f32");
         int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim;
+        int64_t ml_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * 2;
         TORCH_CHECK(o_part_buf->numel() >= o_needed,
                      "o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel());
+        TORCH_CHECK(ml_part_buf->numel() >= ml_needed,
+                     "ml_part_buf too small: need ", ml_needed, " got ", ml_part_buf->numel());
+        TORCH_CHECK(o_part_buf->is_cuda() && ml_part_buf->is_cuda(),
+                     "split buffers must be CUDA tensors");
+        TORCH_CHECK(o_part_buf->is_contiguous() && ml_part_buf->is_contiguous(),
+                     "split buffers must be contiguous");
         p.o_part = (float*)o_part_buf->data_ptr();
         p.ml_part = (float*)ml_part_buf->data_ptr();
     } else {

csrc/kernels/attn_decode_split_kv.cuh

@@ -27,9 +27,9 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
     // Q: [batch, q_head, q_len=1, head_dim] — stride-based
     float q_reg[8];
     int q_off = KV::q_decode_base(p, batch, q_head)
-              + lane * hd_per_thread * p.q_stride_d;
+              + lane * hd_per_thread * p.q_d_stride;
     for (int i = 0; i < hd_per_thread; i++)
-        q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
+        q_reg[i] = __bfloat162float(p.q_ptr[q_off + i * p.q_d_stride]);
 
     int mask_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
 
@@ -138,6 +138,6 @@ __global__ void attn_decode_combine_kernel(AttentionParams<bf16> p) {
     }
 
     float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
-    int o_off = KV::q_decode_base(p, batch, q_head) + d * p.q_stride_d;
-    p.o[o_off] = __float2bfloat16(acc * inv);
+    int o_off = KV::q_decode_base(p, batch, q_head) + d * p.q_d_stride;
+    p.o_ptr[o_off] = __float2bfloat16(acc * inv);
 }

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -48,7 +48,7 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
     const int qrb = gid + 8;
     const bool va = qra < G, vb = qrb < G;
     unsigned Qa[Traits::KD][4];
-    load_q_mma_frags<Traits::KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
+    load_q_mma_frags<Traits::KD>(p.q_ptr + q_base, p.q_h_stride, p.q_d_stride,
                                   qra, qrb, va, vb, tid4, Qa);
 
     float Oacc[Traits::DN8][4];
@@ -107,10 +107,11 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
         int maxc = IsCausal ? KV::decode_attend_len(p, batch) : seq_len;
         mma_softmax_tile<Traits, HasMask>(kv0, maxc, maxc,
                                            0, 0,
-                                           p.mask_b_stride, 0, 0,
-                                           batch, 0,
-                                           p.mask,
-                                           Sacc, Oacc, m0, m1, l0, l1, lane);
+                                            p.mask_b_stride, p.mask_h_stride, p.mask_l_stride,
+                                            batch, q_head0 + gid, q_head0 + gid + 8,
+                                            p.mask,
+                                            va, vb,
+                                            Sacc, Oacc, m0, m1, l0, l1, lane);
 
         mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
     };
@@ -121,7 +122,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
             load_tile(ti_begin + i, i);
 
         for (int it = 0; it < ntiles; it++) {
-            cp_async_wait_group<STAGES - 1>();
+            if (it + 1 == ntiles)
+                cp_async_wait_group<0>();
+            else
+                cp_async_wait_group<STAGES - 1>();
             __syncwarp();
             process_tile(it, it & (STAGES - 1));
             __syncwarp();

csrc/kernels/attn_dispatchers.cuh

@@ -66,9 +66,9 @@ struct PrefillLauncherMMA {
         constexpr int WARPS = 4;
         constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
         using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
-        int q_len = KV::host_q_len(p);
-        dim3 grid((q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS),
-                  p.q_head, p.batch);
+        constexpr int ROWS = Traits::BR * WARPS;
+        dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head,
+                  KV::kPaged ? 1 : p.batch);
         dim3 block(Traits::NUM_THREADS);
         attn_prefill_split_q_mma_kernel<Traits, KV, IsCausal, HasMask>
             <<<grid, block, 0, stream>>>(p);
@@ -80,9 +80,9 @@ template <typename KV>
 struct PrefillLauncherScalar {
     template <int HEAD_DIM, bool IsCausal, bool HasMask>
     static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
-        constexpr int G = 8, ROWS = 32, P_BC = 32;
-        int q_len = KV::host_q_len(p);
-        dim3 grid((q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
+        constexpr int G = (HEAD_DIM == 32) ? 4 : 8, ROWS = 32, P_BC = 32;
+        dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head,
+                  KV::kPaged ? 1 : p.batch);
         dim3 block(G, ROWS);
         attn_prefill_split_q_kernel_t<HEAD_DIM, KV, G, ROWS, P_BC, IsCausal, HasMask>
             <<<grid, block, 0, stream>>>(p);
@@ -133,7 +133,7 @@ static inline void dispatch_paged_prefill(AttentionParams<bf16>& p, cudaStream_t
 template <typename KV>
 struct DecodeLauncherMMA {
     template <int HEAD_DIM, bool IsCausal, bool HasMask>
-    static void launch(AttentionParams<bf16>& p, int group_size, cudaStream_t stream) {
+    static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
         int G = p.q_head / p.kv_head;
         constexpr int MAX_G = 16;
         int num_passes = (G + MAX_G - 1) / MAX_G;
@@ -153,14 +153,19 @@ struct DecodeLauncherMMA {
 template <typename KV>
 struct DecodeLauncherScalar {
     template <int HEAD_DIM, bool IsCausal, bool HasMask>
-    static void launch(AttentionParams<bf16>& p, int group_size, cudaStream_t stream) {
+    static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
         int kv_len = KV::host_kv_len(p);
         int chunks_total = (kv_len + DC_CHUNK - 1) / DC_CHUNK;
         p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
         size_t smem = 2 * DC_CHUNK * p.head_dim * sizeof(bf16);
+        int group_size = p.q_head / p.kv_head;
         int g = min(group_size, 32);  // cap at 32 to respect 1024-thread limit
         dim3 grid(p.batch * p.kv_head, 1, p.num_splits);
         dim3 block(32, g);
+        cudaFuncSetAttribute(
+            attn_decode_split_kv_kernel<HEAD_DIM, KV, IsCausal, HasMask>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize,
+            smem);
         attn_decode_split_kv_kernel<HEAD_DIM, KV, IsCausal, HasMask>
             <<<grid, block, smem, stream>>>(p);
     }
@@ -170,16 +175,15 @@ template <int HEAD_DIM>
 static inline void dispatch_decode(AttentionParams<bf16>& p, cudaStream_t stream) {
     bool is_causal = (p.causal_offset >= 0);
     bool has_mask = (p.use_mask && p.mask);
-    int group_size = p.q_head / p.kv_head;
 
 #ifndef ASTRAI_NO_MMA
     DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                          DecodeLauncherMMA<ContigKV>::template launch,
-                         HEAD_DIM, p, group_size, stream);
+                         HEAD_DIM, p, stream);
 #else
     DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                          DecodeLauncherScalar<ContigKV>::template launch,
-                         HEAD_DIM, p, group_size, stream);
+                         HEAD_DIM, p, stream);
 #endif
 
     attn_decode_combine_kernel<ContigKV><<<p.batch * p.q_head, p.head_dim, 0, stream>>>(p);
@@ -189,16 +193,15 @@ template <int HEAD_DIM>
 static inline void dispatch_paged_decode(AttentionParams<bf16>& p, cudaStream_t stream) {
     bool is_causal = (p.causal_offset >= 0);
     bool has_mask = (p.use_mask && p.mask);
-    int group_size = p.q_head / p.kv_head;
 
 #ifndef ASTRAI_NO_MMA
     DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                          DecodeLauncherMMA<PagedKV>::template launch,
-                         HEAD_DIM, p, group_size, stream);
+                         HEAD_DIM, p, stream);
 #else
     DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                          DecodeLauncherScalar<PagedKV>::template launch,
-                         HEAD_DIM, p, group_size, stream);
+                         HEAD_DIM, p, stream);
 #endif
 
     attn_decode_combine_kernel<PagedKV><<<p.batch * p.q_head, p.head_dim, 0, stream>>>(p);

csrc/kernels/attn_entry_utils.cuh

@@ -42,10 +42,10 @@ inline void extract_q_dims_and_strides(torch::Tensor& q, int64_t layout, P& p) {
     p.q_head = (int)q.size(1);
     p.q_len = (int)q.size(2);
     p.head_dim = (int)q.size(3);
-    p.q_stride_b = (int)q.stride(0);
-    p.q_stride_h = (int)q.stride(1);
-    p.q_stride_l = (int)q.stride(2);
-    p.q_stride_d = (int)q.stride(3);
+    p.q_b_stride = (int)q.stride(0);
+    p.q_h_stride = (int)q.stride(1);
+    p.q_l_stride = (int)q.stride(2);
+    p.q_d_stride = (int)q.stride(3);
 }
 
 // ---- Shared mask packing ----
@@ -63,17 +63,17 @@ inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
         if (m.dim() == 2) {
             p.mask_b_stride = (int)m.stride(0);
             p.mask_h_stride = 0;
-            p.mask_q_stride = 0;
+            p.mask_l_stride = 0;
         } else if (m.dim() == 3) {
             TORCH_CHECK(m.size(1) == 1 || m.size(1) == p.q_len, "mask q_len mismatch");
             p.mask_b_stride = (int)m.stride(0);
             p.mask_h_stride = 0;
-            p.mask_q_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
+            p.mask_l_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
         } else if (m.dim() == 4) {
             TORCH_CHECK(m.size(2) == 1 || m.size(2) == p.q_len, "mask q_len mismatch");
             p.mask_b_stride = (int)m.stride(0);
             p.mask_h_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
-            p.mask_q_stride = (m.size(2) == 1) ? 0 : (int)m.stride(2);
+            p.mask_l_stride = (m.size(2) == 1) ? 0 : (int)m.stride(2);
         } else {
             TORCH_CHECK(false, "mask must be 2D, 3D, or 4D");
         }
@@ -82,7 +82,7 @@ inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
         p.mask = nullptr;
         p.mask_b_stride = 0;
         p.mask_h_stride = 0;
-        p.mask_q_stride = 0;
+        p.mask_l_stride = 0;
     }
 }
 
@@ -106,28 +106,31 @@ inline void attn_pack_params(
     TORCH_CHECK(v.dtype() == torch::kBFloat16);
     TORCH_CHECK(k.sizes() == v.sizes(), "K and V must have identical shapes");
     TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "Q/K/V must be 4D");
-
     extract_q_dims_and_strides(q, layout, p);
 
     if (layout == BLHD) k = k.transpose(1, 2), v = v.transpose(1, 2);
 
     p.kv_head = (int)k.size(1);
     p.kv_len = (int)k.size(2);
+    TORCH_CHECK(p.q_head % p.kv_head == 0,
+                "q_head must be divisible by kv_head");
     TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q");
+    TORCH_CHECK(q.stride(3) == 1 && k.stride(3) == 1 && v.stride(3) == 1,
+                "Q/K/V head_dim must be contiguous");
 
-    p.kv_stride_b = (int)k.stride(0);
-    p.kv_stride_h = (int)k.stride(1);
-    p.kv_stride_l = (int)k.stride(2);
-    p.kv_stride_d = (int)k.stride(3);
+    p.kv_b_stride = (int)k.stride(0);
+    p.kv_h_stride = (int)k.stride(1);
+    p.kv_l_stride = (int)k.stride(2);
+    p.kv_d_stride = (int)k.stride(3);
 
     p.causal_offset = (int)causal_offset;
     p.use_mask = mask.has_value() ? 1 : 0;
     p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
 
-    p.q = (const T*)q.data_ptr();
-    p.k = (const T*)k.data_ptr();
-    p.v = (const T*)v.data_ptr();
-    p.o = nullptr;
+    p.q_ptr = (const T*)q.data_ptr();
+    p.k_ptr = (const T*)k.data_ptr();
+    p.v_ptr = (const T*)v.data_ptr();
+    p.o_ptr = nullptr;
     p.o_part = nullptr;
     p.ml_part = nullptr;
 
@@ -145,7 +148,6 @@ inline void attn_pack_paged_decode_params(
     torch::Tensor req_to_token,
     torch::Tensor req_pool_indices,
     torch::Tensor kv_indptr,
-    int64_t max_seq_len,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
@@ -170,24 +172,23 @@ inline void attn_pack_paged_decode_params(
     p.head_dim = (int)q.size(2);
     p.kv_head = (int)k_cache.size(1);
     TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch");
+    TORCH_CHECK(q.stride(2) == 1 && k_cache.stride(2) == 1 && v_cache.stride(2) == 1,
+                "Q/K/V head_dim must be contiguous");
     TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
     TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head");
 
-    p.q_stride_l = (int)q.stride(0);
-    p.q_stride_h = (int)q.stride(1);
-    p.q_stride_d = (int)q.stride(2);
+    p.q_l_stride = (int)q.stride(0);
+    p.q_h_stride = (int)q.stride(1);
+    p.q_d_stride = (int)q.stride(2);
 
-    p.k_cache = (const T*)k_cache.data_ptr();
-    p.v_cache = (const T*)v_cache.data_ptr();
-    p.q = (const T*)q.data_ptr();
+    p.k_ptr = (const T*)k_cache.data_ptr();
+    p.v_ptr = (const T*)v_cache.data_ptr();
+    p.q_ptr = (const T*)q.data_ptr();
     p.req_to_token = req_to_token.data_ptr<int64_t>();
     p.req_pool_indices = req_pool_indices.data_ptr<int64_t>();
     p.kv_indptr = kv_indptr.data_ptr<int>();
     p.qo_indptr = nullptr;
     p.max_context_len = (int)req_to_token.size(1);
-    p.max_seq_len = (int)max_seq_len;
-    p.total_q = p.batch;  // decode: 1 Q token per request
-    p.max_q_len = 1;
 
     p.causal_offset = (int)causal_offset;
     p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
@@ -199,16 +200,16 @@ inline void attn_pack_paged_decode_params(
         TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
         p.mask_b_stride = (int)m.stride(0);
         p.mask_h_stride = 0;
-        p.mask_q_stride = 0;
+        p.mask_l_stride = 0;
         p.mask = m.data_ptr<bool>();
     } else {
         p.mask = nullptr;
         p.mask_b_stride = 0;
         p.mask_h_stride = 0;
-        p.mask_q_stride = 0;
+        p.mask_l_stride = 0;
     }
 
-    p.o = nullptr;
+    p.o_ptr = nullptr;
     p.o_part = nullptr;
     p.ml_part = nullptr;
 }
@@ -226,7 +227,6 @@ inline void attn_pack_paged_prefill_params(
     torch::Tensor kv_indptr,
     torch::Tensor qo_indptr,
     c10::optional<torch::Tensor> mask,
-    int64_t max_q_len,
     int64_t causal_offset,
     double scale,
     AttentionParams<T>& p
@@ -249,31 +249,29 @@ inline void attn_pack_paged_prefill_params(
 
     p.q_head = (int)q.size(1);
     p.head_dim = (int)q.size(2);
+    p.q_len = (int)q.size(0);
     p.kv_head = (int)k_cache.size(1);
     p.batch = (int)req_pool_indices.size(0);
     TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch");
+    TORCH_CHECK(q.stride(2) == 1 && k_cache.stride(2) == 1 && v_cache.stride(2) == 1,
+                "Q/K/V head_dim must be contiguous");
     TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
     TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head");
     TORCH_CHECK(kv_indptr.size(0) == p.batch + 1, "kv_indptr must be [batch+1]");
     TORCH_CHECK(qo_indptr.size(0) == p.batch + 1, "qo_indptr must be [batch+1]");
 
-    p.q_stride_l = (int)q.stride(0);
-    p.q_stride_h = (int)q.stride(1);
-    p.q_stride_d = (int)q.stride(2);
+    p.q_l_stride = (int)q.stride(0);
+    p.q_h_stride = (int)q.stride(1);
+    p.q_d_stride = (int)q.stride(2);
 
-    p.k_cache = (const T*)k_cache.data_ptr();
-    p.v_cache = (const T*)v_cache.data_ptr();
-    p.q = (const T*)q.data_ptr();
+    p.k_ptr = (const T*)k_cache.data_ptr();
+    p.v_ptr = (const T*)v_cache.data_ptr();
+    p.q_ptr = (const T*)q.data_ptr();
     p.req_to_token = req_to_token.data_ptr<int64_t>();
     p.req_pool_indices = req_pool_indices.data_ptr<int64_t>();
     p.kv_indptr = kv_indptr.data_ptr<int>();
     p.qo_indptr = qo_indptr.data_ptr<int>();
     p.max_context_len = (int)req_to_token.size(1);
-    p.total_q = (int)q.size(0);  // prefill: flattened Q across all requests
-    p.max_q_len = (int)max_q_len;
-    // max_seq_len is unused by the prefill path (decode uses it for split
-    // computation); fill with max_q_len only to keep the POD struct defined.
-    p.max_seq_len = p.max_q_len;
 
     p.causal_offset = (int)causal_offset;
     p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
@@ -285,14 +283,14 @@ inline void attn_pack_paged_prefill_params(
             TORCH_CHECK(m.size(1) <= p.max_context_len, "mask kv_len mismatch");
             p.mask_b_stride = (int)m.stride(0);
             p.mask_h_stride = 0;
-            p.mask_q_stride = 0;
+            p.mask_l_stride = 0;
         } else if (m.dim() == 4) {
             TORCH_CHECK(m.size(1) == 1 || m.size(1) == p.q_head, "mask head mismatch");
-            TORCH_CHECK(m.size(2) == 1 || m.size(2) == p.max_q_len, "mask q_len mismatch");
+            TORCH_CHECK(m.size(2) > 0 && m.size(2) <= p.q_len, "mask q_len mismatch");
             TORCH_CHECK(m.size(3) <= p.max_context_len, "mask kv_len mismatch");
             p.mask_b_stride = (int)m.stride(0);
             p.mask_h_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
-            p.mask_q_stride = (m.size(2) == 1) ? 0 : (int)m.stride(2);
+            p.mask_l_stride = (m.size(2) == 1) ? 0 : (int)m.stride(2);
         } else {
             TORCH_CHECK(false, "mask must be 2D or 4D");
         }
@@ -301,11 +299,11 @@ inline void attn_pack_paged_prefill_params(
         p.mask = nullptr;
         p.mask_b_stride = 0;
         p.mask_h_stride = 0;
-        p.mask_q_stride = 0;
+        p.mask_l_stride = 0;
     }
     p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
 
-    p.o = nullptr;
+    p.o_ptr = nullptr;
     p.o_part = nullptr;
     p.ml_part = nullptr;
 }

csrc/kernels/attn_kv_source.cuh

@@ -13,11 +13,11 @@
 //
 //   ContigKV:  K/V are dense [batch, kv_head, kv_len, head_dim] tensors.
 //              Params fields used: k, v, kv_stride_*, kv_len, q_len,
-//              q_stride_b, causal_offset.
+//              q_b_stride, causal_offset.
 //   PagedKV:   K/V live in a flat pool [size, kv_head, head_dim] indexed via
 //              req_to_token.  Params fields used: k_cache, v_cache,
 //              req_to_token, req_pool_indices, kv_indptr, qo_indptr,
-//              max_context_len, q_stride_l.
+//              max_context_len, q_l_stride.
 //
 // Addressing state that is constant across a whole kernel invocation for one
 // (batch, kv_head) pair is captured once by make_ctx<HEAD_DIM>() and passed
@@ -32,7 +32,7 @@ using bf16 = __nv_bfloat16;
 
 // Hoisted per-(batch, kv_head) addressing context.
 struct KVContext {
-    int kv_base;          // contig: batch*kv_stride_b + kv_head*kv_stride_h
+    int kv_base;          // contig: batch*kv_b_stride + kv_head*kv_h_stride
     int64_t req_idx;      // paged: req_pool_indices[batch]
     int64_t rtt_stride;   // paged: max_context_len
     int64_t pool_stride;  // paged: kv_head * HEAD_DIM
@@ -57,22 +57,30 @@ struct ContigKV {
     static constexpr bool kPaged = false;
 
     // host-side length hooks (grid + split computation in the launchers)
-    HOST_DEV_FORCEINLINE int host_q_len(const AttentionParams<bf16>& p) {
-        return p.q_len;
+    HOST_DEV_FORCEINLINE int host_q_blocks(const AttentionParams<bf16>& p, int rows) {
+        return (p.q_len + rows - 1) / rows;
+    }
+    template <int ROWS>
+    HOST_DEV_FORCEINLINE bool map_q_tile(const AttentionParams<bf16>&,
+                                         int flat_tile, int grid_batch,
+                                         int& batch, int& q_tile) {
+        batch = grid_batch;
+        q_tile = flat_tile;
+        return true;
     }
     HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
         return p.kv_len;
     }
 
-    // prefill: element offset of the request's Q rows (kernel adds qrow*q_stride_l)
+    // prefill: element offset of the request's Q rows (kernel adds qrow*q_l_stride)
     HOST_DEV_FORCEINLINE int q_base(
         const AttentionParams<bf16>& p, int batch, int q_head) {
-        return batch * p.q_stride_b + q_head * p.q_stride_h;
+        return batch * p.q_b_stride + q_head * p.q_h_stride;
     }
     // decode: same offset (q_len == 1, so there is no row stride component)
     HOST_DEV_FORCEINLINE int q_decode_base(
         const AttentionParams<bf16>& p, int batch, int q_head) {
-        return batch * p.q_stride_b + q_head * p.q_stride_h;
+        return batch * p.q_b_stride + q_head * p.q_h_stride;
     }
 
     HOST_DEV_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int batch) {
@@ -93,13 +101,13 @@ struct ContigKV {
     HOST_DEV_FORCEINLINE KVContext make_ctx(
         const AttentionParams<bf16>& p, int batch, int kv_head) {
         KVContext c = {};
-        c.kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
+        c.kv_base = batch * p.kv_b_stride + kv_head * p.kv_h_stride;
         return c;
     }
     HOST_DEV_FORCEINLINE KVAddr kv_addr(
         const AttentionParams<bf16>& p, const KVContext& c, int kc, int d, bool valid) {
-        const int g_off = c.kv_base + kc * p.kv_stride_l + d * p.kv_stride_d;
-        return {&p.k[g_off], &p.v[g_off], valid};
+        const int g_off = c.kv_base + kc * p.kv_l_stride + d * p.kv_d_stride;
+        return {&p.k_ptr[g_off], &p.v_ptr[g_off], valid};
     }
 };
 
@@ -107,8 +115,26 @@ struct ContigKV {
 struct PagedKV {
     static constexpr bool kPaged = true;
 
-    HOST_DEV_FORCEINLINE int host_q_len(const AttentionParams<bf16>& p) {
-        return p.max_q_len;
+    HOST_DEV_FORCEINLINE int host_q_blocks(const AttentionParams<bf16>& p, int rows) {
+        // sum(ceil(q_len[b] / rows)) <= ceil(total_q / rows) + batch - 1.
+        return (p.q_len + rows - 1) / rows + p.batch - 1;
+    }
+    template <int ROWS>
+    HOST_DEV_FORCEINLINE bool map_q_tile(const AttentionParams<bf16>& p,
+                                         int flat_tile, int,
+                                         int& batch, int& q_tile) {
+        int tile_base = 0;
+        for (int b = 0; b < p.batch; ++b) {
+            int len = p.qo_indptr[b + 1] - p.qo_indptr[b];
+            int tiles = (len + ROWS - 1) / ROWS;
+            if (flat_tile < tile_base + tiles) {
+                batch = b;
+                q_tile = flat_tile - tile_base;
+                return true;
+            }
+            tile_base += tiles;
+        }
+        return false;
     }
     HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
         return p.max_context_len;
@@ -117,12 +143,12 @@ struct PagedKV {
     // prefill: Q rows start at qo_indptr[batch] (ragged batch base)
     HOST_DEV_FORCEINLINE int q_base(
         const AttentionParams<bf16>& p, int batch, int q_head) {
-        return p.qo_indptr[batch] * p.q_stride_l + q_head * p.q_stride_h;
+        return p.qo_indptr[batch] * p.q_l_stride + q_head * p.q_h_stride;
     }
     // decode: Q is [batch, q_head, head_dim], so batch is the outer row
     HOST_DEV_FORCEINLINE int q_decode_base(
         const AttentionParams<bf16>& p, int batch, int q_head) {
-        return batch * p.q_stride_l + q_head * p.q_stride_h;
+        return batch * p.q_l_stride + q_head * p.q_h_stride;
     }
 
     HOST_DEV_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int batch) {
@@ -154,6 +180,33 @@ struct PagedKV {
         const int64_t slot = valid ? p.req_to_token[c.req_idx * c.rtt_stride + kc] : 0;
         const bool ok = valid && (slot >= 0);
         const int64_t gmem_off = slot * c.pool_stride + c.head_off + d;
-        return {&p.k_cache[gmem_off], &p.v_cache[gmem_off], ok};
+        return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], ok};
     }
 };
+
+// ---- Q-block mapping ----
+// Contiguous grids map directly to (batch, q_tile). Paged grids flatten the
+// ragged Q tiles, so one thread resolves the request and broadcasts it.
+template <int ROWS, typename KV>
+__device__ __forceinline__ bool map_q_block(
+    const AttentionParams<bf16>& p, int& batch, int& q_tile) {
+    if constexpr (!KV::kPaged) {
+        batch = blockIdx.z;
+        q_tile = blockIdx.x;
+        return true;
+    } else {
+        __shared__ int mapped_batch;
+        __shared__ int mapped_q_tile;
+
+        if ((threadIdx.x | threadIdx.y) == 0) {
+            mapped_batch = -1;
+            KV::template map_q_tile<ROWS>(
+                p, blockIdx.x, blockIdx.z, mapped_batch, mapped_q_tile);
+        }
+        __syncthreads();
+
+        batch = mapped_batch;
+        q_tile = mapped_q_tile;
+        return batch >= 0;
+    }
+}

csrc/kernels/attn_mma_utils.cuh

@@ -133,8 +133,8 @@ __device__ __forceinline__ void cp_async_wait_group() {
 // ---------------------------------------------------------------------------
 // Q-load: load query rows directly from global memory into mma A-operand
 // register layout. One call replaces ~15 duplicated lines in each MMA kernel.
-// stride_row is p.q_stride_h for decode (q_len=1, G heads) or
-//              p.q_stride_l for prefill (multi-q rows).
+// stride_row is p.q_h_stride for decode (q_len=1, G heads) or
+//              p.q_l_stride for prefill (multi-q rows).
 // ---------------------------------------------------------------------------
 template <int KD>
 __device__ inline void load_q_mma_frags(
@@ -198,9 +198,10 @@ __device__ inline void mma_softmax_tile(
     int kv0,
     int maxc0, int maxc1,
     int qrow0, int qrow1,
-    int mask_b_stride, int mask_h_stride, int mask_q_stride,
-    int mask_batch, int mask_head,
+    int mask_b_stride, int mask_h_stride, int mask_l_stride,
+    int mask_batch, int mask_head0, int mask_head1,
     const bool* __restrict__ mask,
+    bool valid0, bool valid1,
     float Sacc[Traits::NC8][4],
     float Oacc[Traits::DN8][4],
     float& m0, float& m1,
@@ -210,16 +211,16 @@ __device__ inline void mma_softmax_tile(
     int tid4 = lane & 3;
 
     float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
-    int mask_base0 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow0 * mask_q_stride;
-    int mask_base1 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow1 * mask_q_stride;
+    int mask_base0 = mask_batch * mask_b_stride + mask_head0 * mask_h_stride + qrow0 * mask_l_stride;
+    int mask_base1 = mask_batch * mask_b_stride + mask_head1 * mask_h_stride + qrow1 * mask_l_stride;
     #pragma unroll
     for (int n8 = 0; n8 < Traits::NC8; n8++) {
         int cc = kv0 + n8 * 8 + 2 * tid4;
         int c1 = cc + 1;
-        bool b0 = (cc >= maxc0) || (HasMask && !mask[mask_base0 + cc]);
-        bool b1 = (c1 >= maxc0) || (HasMask && !mask[mask_base0 + c1]);
-        bool b2 = (cc >= maxc1) || (HasMask && !mask[mask_base1 + cc]);
-        bool b3 = (c1 >= maxc1) || (HasMask && !mask[mask_base1 + c1]);
+        bool b0 = !valid0 || (cc >= maxc0) || (HasMask && !mask[mask_base0 + cc]);
+        bool b1 = !valid0 || (c1 >= maxc0) || (HasMask && !mask[mask_base0 + c1]);
+        bool b2 = !valid1 || (cc >= maxc1) || (HasMask && !mask[mask_base1 + cc]);
+        bool b3 = !valid1 || (c1 >= maxc1) || (HasMask && !mask[mask_base1 + c1]);
         float s0 = b0 ? -FLT_MAX : Sacc[n8][0];
         float s1 = b1 ? -FLT_MAX : Sacc[n8][1];
         float s2 = b2 ? -FLT_MAX : Sacc[n8][2];

csrc/kernels/attn_paged_decode.cu

@@ -8,7 +8,6 @@ torch::Tensor attn_paged_decode(
     torch::Tensor req_to_token,
     torch::Tensor req_pool_indices,
     torch::Tensor kv_indptr,
-    int64_t max_seq_len,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
@@ -22,29 +21,38 @@ torch::Tensor attn_paged_decode(
     AttentionParams<bf16> p;
     attn_pack_paged_decode_params(q, k_cache, v_cache,
                                    req_to_token, req_pool_indices, kv_indptr,
-                                   max_seq_len, mask, causal_offset, scale, p);
+                                   mask, causal_offset, scale, p);
 
     torch::Tensor O;
     if (out_buf.has_value() && out_buf->defined()) {
         TORCH_CHECK(out_buf->dtype() == q.dtype(), "out_buf dtype must match q");
+        TORCH_CHECK(out_buf->is_cuda() && out_buf->is_contiguous(),
+                    "out_buf must be a contiguous CUDA tensor");
         TORCH_CHECK(out_buf->size(0) >= q.size(0), "out_buf batch too small");
-        TORCH_CHECK(out_buf->size(1) >= q.size(1), "out_buf heads too small");
-        TORCH_CHECK(out_buf->size(2) >= q.size(2), "out_buf head_dim too small");
-        O = out_buf.value().slice(0, 0, q.size(0))
-                            .slice(1, 0, q.size(1))
-                            .slice(2, 0, q.size(2));
+        TORCH_CHECK(out_buf->size(1) == q.size(1), "out_buf heads must match q");
+        TORCH_CHECK(out_buf->size(2) == q.size(2), "out_buf head_dim must match q");
+        TORCH_CHECK(q.is_contiguous(),
+                    "q must be contiguous when out_buf is provided");
+        O = out_buf.value().slice(0, 0, q.size(0));
     } else {
         O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
     }
-    p.o = (bf16*)O.data_ptr();
+    p.o_ptr = (bf16*)O.data_ptr();
 
     if (o_part_buf.has_value() && ml_part_buf.has_value()
         && o_part_buf->defined() && ml_part_buf->defined()) {
         TORCH_CHECK(o_part_buf->scalar_type() == torch::kFloat32, "o_part_buf must be f32");
         TORCH_CHECK(ml_part_buf->scalar_type() == torch::kFloat32, "ml_part_buf must be f32");
         int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim;
+        int64_t ml_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * 2;
         TORCH_CHECK(o_part_buf->numel() >= o_needed,
                      "o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel());
+        TORCH_CHECK(ml_part_buf->numel() >= ml_needed,
+                     "ml_part_buf too small: need ", ml_needed, " got ", ml_part_buf->numel());
+        TORCH_CHECK(o_part_buf->is_cuda() && ml_part_buf->is_cuda(),
+                    "split buffers must be CUDA tensors");
+        TORCH_CHECK(o_part_buf->is_contiguous() && ml_part_buf->is_contiguous(),
+                    "split buffers must be contiguous");
         p.o_part = (float*)o_part_buf->data_ptr();
         p.ml_part = (float*)ml_part_buf->data_ptr();
     } else {
@@ -63,7 +71,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("req_to_token"),
         py::arg("req_pool_indices"),
         py::arg("kv_indptr"),
-        py::arg("max_seq_len"),
         py::arg("mask") = py::none(),
         py::arg("causal_offset") = -1,
         py::arg("scale") = 0.0,

csrc/kernels/attn_paged_prefill.cu

@@ -10,7 +10,6 @@ torch::Tensor attn_paged_prefill(
     torch::Tensor kv_indptr,
     torch::Tensor qo_indptr,
     c10::optional<torch::Tensor> mask,
-    int64_t max_q_len,
     int64_t causal_offset,
     double scale
 ) {
@@ -21,10 +20,10 @@ torch::Tensor attn_paged_prefill(
     attn_pack_paged_prefill_params(q, k_cache, v_cache,
                                     req_to_token, req_pool_indices,
                                     kv_indptr, qo_indptr, mask,
-                                    max_q_len, causal_offset, scale, p);
+                                    causal_offset, scale, p);
 
     auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
-    p.o = (bf16*)O.data_ptr();
+    p.o_ptr = (bf16*)O.data_ptr();
 
     DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_prefill, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
@@ -41,7 +40,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("kv_indptr"),
         py::arg("qo_indptr"),
         py::arg("mask") = py::none(),
-        py::arg("max_q_len"),
         py::arg("causal_offset") = -1,
         py::arg("scale") = 0.0,
         "SGLang-style paged prefill: flat KV pool + ragged batch.");

csrc/kernels/attn_prefill.cu

@@ -19,7 +19,7 @@ torch::Tensor attn_prefill(
 
     auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
     auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
-    p.o = (bf16*)O_view.data_ptr();
+    p.o_ptr = (bf16*)O_view.data_ptr();
 
     DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());

csrc/kernels/attn_prefill_split_q.cuh

@@ -36,9 +36,11 @@ template <int HEAD_DIM, typename KV, int G, int ROWS, int P_BC, bool IsCausal, b
 __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     constexpr int DPT = HEAD_DIM / G;
 
-    int q_tile = blockIdx.x;
+    int batch, q_tile;
+    if (!map_q_block<ROWS, KV>(p, batch, q_tile))
+        return;
+
     int q_head = blockIdx.y;
-    int batch  = blockIdx.z;
     int gpos   = threadIdx.x;  // 0..G-1  (which d-chunk)
     int row    = threadIdx.y;  // 0..ROWS-1
     int q_row  = q_tile * ROWS + row;
@@ -57,10 +59,10 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     const int q_base = KV::q_base(p, batch, q_head);
     float qreg[DPT];
     if (q_row < q_len) {
-        int q_off = q_base + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d;
+        int q_off = q_base + q_row * p.q_l_stride + gpos * DPT * p.q_d_stride;
 #pragma unroll
         for (int i = 0; i < DPT; i++)
-            qreg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
+            qreg[i] = __bfloat162float(p.q_ptr[q_off + i * p.q_d_stride]);
     }
 
     float m = -FLT_MAX, l = 0.0f;
@@ -105,7 +107,7 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
             }
         }
 
-        int mask_row_base = mask_batch_base + q_row * p.mask_q_stride;
+        int mask_row_base = mask_batch_base + q_row * p.mask_l_stride;
         for (int s = 0; s < lim; s++) {
             const bf16* kr = sK + s * HEAD_DIM + gpos * DPT;
             float part = 0.0f;
@@ -145,10 +147,10 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     }
 
     if (q_row < q_len) {
-        int o_off = q_base + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d;
+        int o_off = q_base + q_row * p.q_l_stride + gpos * DPT * p.q_d_stride;
         float rl = (l > 1e-20f) ? (1.0f / l) : 0.0f;
 #pragma unroll
         for (int i = 0; i < DPT; i++)
-            p.o[o_off + i * p.q_stride_d] = __float2bfloat16(acc[i] * rl);
+            p.o_ptr[o_off + i * p.q_d_stride] = __float2bfloat16(acc[i] * rl);
     }
 }

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -24,9 +24,11 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     const int tid4 = lane & 3;   // 0..3
 
     const int q_head = blockIdx.y;
-    const int batch = blockIdx.z;
+    int batch, q_tile;
+    if (!map_q_block<Traits::BR * Traits::WARPS, KV>(p, batch, q_tile))
+        return;
     const int kv_head = q_head / (p.q_head / p.kv_head);
-    const int qrow0 = (blockIdx.x * Traits::WARPS + warp) * Traits::BR;
+    const int qrow0 = (q_tile * Traits::WARPS + warp) * Traits::BR;
 
     // Per-request dims (from KV policy — paged reads kv_indptr/qo_indptr).
     const int seq_len    = KV::kv_len(p, batch);
@@ -45,7 +47,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     const int qrb = qrow0 + gid + 8;
     const bool va = qra < q_len, vb = qrb < q_len;
     unsigned Qa[Traits::KD][4];
-    load_q_mma_frags<Traits::KD>(p.q + q_base, p.q_stride_l, p.q_stride_d,
+    load_q_mma_frags<Traits::KD>(p.q_ptr + q_base, p.q_l_stride, p.q_d_stride,
                                   qra, qrb, va, vb, tid4, Qa);
 
     float Oacc[Traits::DN8][4];
@@ -61,7 +63,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     // Causal tile-skip bounds (dead code when IsCausal == false)
     const int max_kv = qrow0 + Traits::BR - 1 + causal_off;
     const int block_max_kv =
-        blockIdx.x * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1
+        q_tile * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1
         + causal_off;
 
     int t_end = tiles - 1;
@@ -122,9 +124,10 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
                                  : seq_len;
             mma_softmax_tile<Traits, HasMask>(kv0, maxc0, maxc1,
                                                qr0, qr1,
-                                               p.mask_b_stride, p.mask_h_stride, p.mask_q_stride,
-                                               batch, q_head,
+                                               p.mask_b_stride, p.mask_h_stride, p.mask_l_stride,
+                                               batch, q_head, q_head,
                                                p.mask,
+                                               va, vb,
                                                Sacc, Oacc, m0, m1, l0, l1, lane);
 
             mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
@@ -142,13 +145,13 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
             __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0,
                                                       Oacc[dn8][1] * rl0);
             *reinterpret_cast<__nv_bfloat162*>(
-                &p.o[o_base + qr0 * p.q_stride_l + d * p.q_stride_d]) = v;
+                &p.o_ptr[o_base + qr0 * p.q_l_stride + d * p.q_d_stride]) = v;
         }
         if (qr1 < q_len) {
             __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
                                                       Oacc[dn8][3] * rl1);
             *reinterpret_cast<__nv_bfloat162*>(
-                &p.o[o_base + qr1 * p.q_stride_l + d * p.q_stride_d]) = v;
+                &p.o_ptr[o_base + qr1 * p.q_l_stride + d * p.q_d_stride]) = v;
         }
     }
 }

csrc/kernels/fp8_mm.cu

@@ -0,0 +1,441 @@
+// FP8 e4m3 matrix multiply via cuBLASLt (sm89 TN layout).
+//
+// cuBLASLt exposes fp8 kernels only for op(A)=T, op(B)=N on Ada; we exploit
+// the identity: row-major a[M,K] == A^T as col-major [K,M] (zero copy), and
+// row-major wT[N,K] == B as col-major [K,N] (zero copy). The col-major
+// result D[M,N] is C^T in row-major terms, so we transpose the output once.
+//
+// Inputs arrive pre-scaled fp8 e4m3 tensors; output is unscaled fp32.
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cublasLt.h>
+#include <cuda_fp8.h>
+#include <cstdint>
+#include <mutex>
+#include <unordered_map>
+
+static std::recursive_mutex g_mutex;
+
+static cublasLtHandle_t g_handle = nullptr;
+static cublasLtMatmulDesc_t g_desc = nullptr;
+static cublasLtMatrixLayout_t g_layout_a = nullptr;
+static cublasLtMatrixLayout_t g_layout_b = nullptr;
+static cublasLtMatrixLayout_t g_layout_c = nullptr;
+static cublasLtMatmulPreference_t g_pref = nullptr;
+static void* g_workspace = nullptr;
+static size_t g_ws_size = 0;
+
+struct ShapeKey {
+    int64_t m;
+    int64_t k;
+    int64_t n;
+    bool operator==(const ShapeKey& other) const {
+        return m == other.m && k == other.k && n == other.n;
+    }
+};
+
+struct ShapeKeyHash {
+    size_t operator()(const ShapeKey& s) const {
+        size_t h = std::hash<int64_t>()(s.m);
+        h ^= std::hash<int64_t>()(s.k) + 0x9e3779b9 + (h << 6) + (h >> 2);
+        h ^= std::hash<int64_t>()(s.n) + 0x9e3779b9 + (h << 6) + (h >> 2);
+        return h;
+    }
+};
+
+using AlgoCache = std::unordered_map<ShapeKey, cublasLtMatmulAlgo_t, ShapeKeyHash>;
+
+static void create_matmul_config(cublasLtMatmulDesc_t* desc,
+                                 cublasLtMatrixLayout_t* layout_a,
+                                 cublasLtMatrixLayout_t* layout_b,
+                                 cublasLtMatrixLayout_t* layout_c) {
+    cublasOperation_t ta = CUBLAS_OP_T, tb = CUBLAS_OP_N;
+    TORCH_CHECK(cublasLtMatmulDescCreate(desc, CUBLAS_COMPUTE_32F, CUDA_R_32F) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatmulDescSetAttribute(
+                    *desc, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatmulDescSetAttribute(
+                    *desc, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatrixLayoutCreate(layout_a, CUDA_R_8F_E4M3, 1, 1, 1) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatrixLayoutCreate(layout_b, CUDA_R_8F_E4M3, 1, 1, 1) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatrixLayoutCreate(layout_c, CUDA_R_16BF, 1, 1, 1) ==
+                CUBLAS_STATUS_SUCCESS);
+}
+
+static void ensure_cublas_lt() {
+    std::lock_guard<std::recursive_mutex> lock(g_mutex);
+    if (g_handle) {
+        return;
+    }
+    TORCH_CHECK(cublasLtCreate(&g_handle) == CUBLAS_STATUS_SUCCESS);
+    create_matmul_config(&g_desc, &g_layout_a, &g_layout_b, &g_layout_c);
+    TORCH_CHECK(cublasLtMatmulPreferenceCreate(&g_pref) == CUBLAS_STATUS_SUCCESS);
+    size_t ws = 16 * 1024 * 1024;
+    TORCH_CHECK(cublasLtMatmulPreferenceSetAttribute(
+                   g_pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws, sizeof(ws)) ==
+               CUBLAS_STATUS_SUCCESS);
+}
+
+static cublasStatus_t get_algo_cached(int64_t m, int64_t k, int64_t n,
+                                      AlgoCache* cache,
+                                      cublasLtMatmulAlgo_t* algo);
+
+static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor out,
+                          int64_t m, int64_t k, int64_t n,
+                          const float* a_scale, const float* b_scale,
+                          cudaStream_t stream);
+
+static const float k_scale_one = 1.0f;
+
+static void set_layout(cublasLtMatrixLayout_t layout, int64_t rows, int64_t cols,
+                       int64_t ld) {
+    TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_ROWS,
+                                                 &rows, sizeof(rows)) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_COLS,
+                                                 &cols, sizeof(cols)) ==
+                CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_LD,
+                                                 &ld, sizeof(ld)) ==
+                CUBLAS_STATUS_SUCCESS);
+}
+
+torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b) {
+    TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn, "a must be float8_e4m3fn");
+    TORCH_CHECK(b.scalar_type() == torch::kFloat8_e4m3fn, "b must be float8_e4m3fn");
+    TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "2D tensors required");
+    const at::cuda::OptionalCUDAGuard guard(a.device());
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    auto a_c = a.contiguous();
+    auto b_c = b.contiguous();
+    int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0);
+    TORCH_CHECK(b_c.size(1) == k, "inner dim mismatch");
+
+    auto buf = torch::empty({m, n}, a_c.options().dtype(torch::kBFloat16));
+    ensure_cublas_lt();
+    fp8_gemm_into(a_c, b_c, buf, m, k, n, &k_scale_one, &k_scale_one,
+                  stream.stream());
+    return buf;
+}
+
+
+// ---------------------------------------------------------------------------
+// Quantize: bf16 * scale_inv -> fp8, one atomicMax amax per kernel call.
+// amax_ptr must be zeroed before launch; float-bits atomicMax works because
+// |v| >= 0 has a monotonic IEEE bit pattern.
+// ---------------------------------------------------------------------------
+
+template <typename T8>
+__device__ __forceinline__ T8 cast_fp8(float v);
+
+template <>
+__device__ __forceinline__ __nv_fp8_e4m3 cast_fp8<__nv_fp8_e4m3>(float v) {
+    return __nv_fp8_e4m3(v);
+}
+
+template <>
+__device__ __forceinline__ __nv_fp8_e5m2 cast_fp8<__nv_fp8_e5m2>(float v) {
+    return __nv_fp8_e5m2(v);
+}
+
+template <typename T8>
+__global__ void quantize_kernel(const __nv_bfloat16* __restrict__ src,
+                                const float* __restrict__ scale_inv,
+                                T8* __restrict__ dst,
+                                float* __restrict__ amax_ptr, int64_t n) {
+    int64_t i = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
+    float amax = 0.f;
+    if (i < n) {
+        float v = __bfloat162float(src[i]) * *scale_inv;
+        dst[i] = cast_fp8<T8>(v);
+        amax = fabsf(v);
+    }
+    for (int off = 16; off; off >>= 1)
+        amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off));
+    __shared__ float sm[8];
+    if ((threadIdx.x & 31) == 0) sm[threadIdx.x >> 5] = amax;
+    __syncthreads();
+    if (threadIdx.x == 0) {
+        float m = 0.f;
+        for (int w = 0; w < blockDim.x / 32; ++w) m = fmaxf(m, sm[w]);
+        atomicMax(reinterpret_cast<unsigned*>(amax_ptr), __float_as_uint(m));
+    }
+}
+
+// Same but with a transpose (rows x cols bf16 row-major -> fp8 [cols, rows]).
+template <typename T8>
+__global__ void transpose_quantize_kernel(
+    const __nv_bfloat16* __restrict__ src, const float* __restrict__ scale_inv,
+    T8* __restrict__ dst, float* __restrict__ amax_ptr, int64_t rows,
+    int64_t cols) {
+    __shared__ T8 tile[32][33];
+    int64_t x = blockIdx.x * 32 + threadIdx.x;
+    int64_t y = blockIdx.y * 32 + threadIdx.y;
+    float amax = 0.f;
+    for (int j = 0; j < 32; j += 8) {
+        if (x < cols && y + j < rows) {
+            float v = __bfloat162float(src[(y + j) * cols + x]) * *scale_inv;
+            tile[threadIdx.y + j][threadIdx.x] = cast_fp8<T8>(v);
+            amax = fmaxf(amax, fabsf(v));
+        }
+    }
+    __syncthreads();
+
+    x = blockIdx.y * 32 + threadIdx.x;
+    y = blockIdx.x * 32 + threadIdx.y;
+    for (int j = 0; j < 32; j += 8) {
+        if (x < rows && y + j < cols) {
+            dst[(y + j) * rows + x] = tile[threadIdx.x][threadIdx.y + j];
+        }
+    }
+    for (int off = 16; off; off >>= 1)
+        amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off));
+    __shared__ float sm[8];
+    if ((threadIdx.x & 31) == 0) sm[threadIdx.x >> 5] = amax;
+    __syncthreads();
+    if (threadIdx.x == 0) {
+        float m = 0.f;
+        for (int w = 0; w < blockDim.x / 32; ++w) m = fmaxf(m, sm[w]);
+        atomicMax(reinterpret_cast<unsigned*>(amax_ptr), __float_as_uint(m));
+    }
+}
+
+__global__ void bias_add_bf16_kernel(
+    __nv_bfloat16* __restrict__ dst, const __nv_bfloat16* __restrict__ bias,
+    int64_t total, int64_t n) {
+    // GEMM and output use the same row-major [M,N] layout.
+    int64_t idx = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
+    if (idx >= total) return;
+    float v = __bfloat162float(dst[idx]);
+    dst[idx] = __float2bfloat16(v + __bfloat162float(bias[idx % n]));
+}
+
+static cublasStatus_t get_algo_cached(int64_t m, int64_t k, int64_t n,
+                                      AlgoCache* cache,
+                                      cublasLtMatmulAlgo_t* algo) {
+    std::lock_guard<std::recursive_mutex> lock(g_mutex);
+    ShapeKey key{m, k, n};
+    auto it = cache->find(key);
+    if (it != cache->end()) {
+        *algo = it->second;
+        return CUBLAS_STATUS_SUCCESS;
+    }
+    cublasLtMatmulHeuristicResult_t heur;
+    int returned = 0;
+    cublasStatus_t st = cublasLtMatmulAlgoGetHeuristic(
+        g_handle, g_desc, g_layout_a, g_layout_b, g_layout_c, g_layout_c, g_pref, 1,
+        &heur, &returned);
+    if (st != CUBLAS_STATUS_SUCCESS || returned == 0)
+        return CUBLAS_STATUS_NOT_SUPPORTED;
+    if (heur.workspaceSize > g_ws_size) {
+        if (g_workspace) cudaFree(g_workspace);
+        TORCH_CHECK(cudaMalloc(&g_workspace, heur.workspaceSize) == cudaSuccess);
+        g_ws_size = heur.workspaceSize;
+    }
+    cache->emplace(key, heur.algo);
+    *algo = heur.algo;
+    return CUBLAS_STATUS_SUCCESS;
+}
+
+static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor out,
+                          int64_t m, int64_t k, int64_t n,
+                          const float* a_scale, const float* b_scale,
+                          cudaStream_t stream) {
+    std::lock_guard<std::recursive_mutex> lock(g_mutex);
+    set_layout(g_layout_a, k, n, k);  // param A = rhs (op=T -> [N,K])
+    set_layout(g_layout_b, k, m, k);  // param B = lhs (op=N -> [K,M])
+    set_layout(g_layout_c, n, m, n);  // col-major [N,M] == row-major [M,N]
+    // Per-tensor FP32 scales applied inside the GEMM:
+    // D = alpha * A_SCALE * B_SCALE * A * B (alpha = 1).
+    TORCH_CHECK(cublasLtMatmulDescSetAttribute(
+                    g_desc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &a_scale,
+                    sizeof(a_scale)) == CUBLAS_STATUS_SUCCESS);
+    TORCH_CHECK(cublasLtMatmulDescSetAttribute(
+                    g_desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_scale,
+                    sizeof(b_scale)) == CUBLAS_STATUS_SUCCESS);
+    float alpha = 1.0f, beta = 0.0f;
+    static AlgoCache cache;
+    cublasLtMatmulAlgo_t algo;
+    cublasStatus_t st = get_algo_cached(m, k, n, &cache, &algo);
+    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
+                "cublasLtMatmulAlgoGetHeuristic failed: ", cublasLtGetStatusName(st));
+    st = cublasLtMatmul(g_handle, g_desc, &alpha, rhs.data_ptr(), g_layout_a,
+                        lhs.data_ptr(), g_layout_b, &beta, out.data_ptr(), g_layout_c,
+                        out.data_ptr(), g_layout_c, &algo, g_workspace, g_ws_size,
+                        stream);
+    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
+                "cublasLtMatmul failed: ", cublasLtGetStatusName(st));
+}
+
+// ---------------------------------------------------------------------------
+// Scaled FP8 linear forward: quantize x/w with per-tensor scales -> cublasLt
+// GEMM (scales applied inside) -> bias in-place -> bf16 [..., N].
+// sx/sw: f32 scale tensors (device scalars); sx_inv/sw_inv: 1/scale.
+// amax_x/amax_w: f32 buffers receiving max-abs of the quantized tensors.
+// ---------------------------------------------------------------------------
+
+torch::Tensor fp8_linear_forward_scaled(torch::Tensor x, torch::Tensor w,
+                                        torch::Tensor bias, torch::Tensor sx,
+                                        torch::Tensor sw, torch::Tensor sx_inv,
+                                        torch::Tensor sw_inv,
+                                        torch::Tensor amax_x,
+                                        torch::Tensor amax_w) {
+    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(x.dtype() == torch::kBFloat16 && w.dtype() == torch::kBFloat16,
+                "x and w must be bf16");
+    const at::cuda::OptionalCUDAGuard guard(x.device());
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    auto x_c = x.reshape({-1, w.size(1)}).contiguous();
+    auto w_c = w.contiguous();
+    int64_t m = x_c.size(0), k = x_c.size(1), n = w_c.size(0);
+    TORCH_CHECK(w_c.size(1) == k, "inner dim mismatch");
+    ensure_cublas_lt();
+
+    const float* sx_ptr = sx.data_ptr<float>();
+    const float* sw_ptr = sw.data_ptr<float>();
+    const float* sxi_ptr = sx_inv.data_ptr<float>();
+    const float* swi_ptr = sw_inv.data_ptr<float>();
+    float* amax_x_ptr = amax_x.data_ptr<float>();
+    float* amax_w_ptr = amax_w.data_ptr<float>();
+    C10_CUDA_CHECK(cudaMemsetAsync(amax_x_ptr, 0, sizeof(float), stream.stream()));
+    C10_CUDA_CHECK(cudaMemsetAsync(amax_w_ptr, 0, sizeof(float), stream.stream()));
+
+    auto x8 = torch::empty({m, k}, x_c.options().dtype(torch::kFloat8_e4m3fn));
+    auto w8 = torch::empty({n, k}, w_c.options().dtype(torch::kFloat8_e4m3fn));
+    int64_t block = 256;
+    quantize_kernel<__nv_fp8_e4m3>
+        <<<(unsigned)((m * k + block - 1) / block), block, 0, stream.stream()>>>(
+            reinterpret_cast<const __nv_bfloat16*>(x_c.data_ptr()), sxi_ptr,
+            reinterpret_cast<__nv_fp8_e4m3*>(x8.data_ptr()), amax_x_ptr, m * k);
+    quantize_kernel<__nv_fp8_e4m3>
+        <<<(unsigned)((n * k + block - 1) / block), block, 0, stream.stream()>>>(
+            reinterpret_cast<const __nv_bfloat16*>(w_c.data_ptr()), swi_ptr,
+            reinterpret_cast<__nv_fp8_e4m3*>(w8.data_ptr()), amax_w_ptr, n * k);
+    C10_CUDA_CHECK(cudaGetLastError());
+
+    auto out = torch::empty({m, n}, x_c.options());
+    fp8_gemm_into(x8, w8, out, m, k, n, sw_ptr, sx_ptr, stream.stream());
+
+    if (bias.defined() && bias.numel() > 0) {
+        TORCH_CHECK(bias.scalar_type() == torch::kBFloat16 && bias.numel() == n,
+                    "bias must be bf16 with shape [N]");
+        bias_add_bf16_kernel<<<(unsigned)((m * n + block - 1) / block), block, 0, stream>>>(
+            reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
+            reinterpret_cast<const __nv_bfloat16*>(bias.data_ptr()), m * n, n);
+        C10_CUDA_CHECK(cudaGetLastError());
+    }
+
+    std::vector<int64_t> shape(x.sizes().begin(), x.sizes().end() - 1);
+    shape.push_back(n);
+    return out.reshape(shape);
+}
+
+// ---------------------------------------------------------------------------
+// Scaled FP8 linear backward: dX = g @ W, dW = g^T @ X, dB = sum(g).
+// Scales: g uses sg (immediate), w/x reuse the forward scales.
+// ---------------------------------------------------------------------------
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward_scaled(
+    torch::Tensor g, torch::Tensor x, torch::Tensor w,
+    std::vector<int64_t> masks, torch::Tensor sg, torch::Tensor sw,
+    torch::Tensor sx, torch::Tensor sg_inv, torch::Tensor sw_inv,
+    torch::Tensor sx_inv, torch::Tensor amax_g) {
+    const at::cuda::OptionalCUDAGuard guard(g.device());
+    TORCH_CHECK(g.dtype() == torch::kBFloat16 && x.dtype() == torch::kBFloat16 &&
+                    w.dtype() == torch::kBFloat16,
+                "g, x, and w must be bf16");
+    auto stream = at::cuda::getCurrentCUDAStream();
+    auto g_c = g.reshape({-1, w.size(0)}).contiguous();
+    auto x_c = x.reshape({-1, x.size(-1)}).contiguous();
+    auto w_c = w.contiguous();
+    int64_t m = g_c.size(0);
+    int64_t n = w.size(0);
+    int64_t k = w.size(1);
+    TORCH_CHECK(x_c.size(0) == m && x_c.size(1) == k && g_c.size(1) == n,
+                "backward shape mismatch");
+
+    auto grad_input = torch::empty_like(x);
+    auto grad_weight = torch::empty_like(w);
+    auto grad_bias = torch::empty({0}, g_c.options().dtype(g.dtype()));
+    ensure_cublas_lt();
+
+    const float* sg_ptr = sg.data_ptr<float>();
+    const float* sw_ptr = sw.data_ptr<float>();
+    const float* sx_ptr = sx.data_ptr<float>();
+    const float* sgi_ptr = sg_inv.data_ptr<float>();
+    const float* swi_ptr = sw_inv.data_ptr<float>();
+    const float* sxi_ptr = sx_inv.data_ptr<float>();
+    float* amax_g_ptr = amax_g.data_ptr<float>();
+    C10_CUDA_CHECK(cudaMemsetAsync(amax_g_ptr, 0, sizeof(float), stream.stream()));
+
+    auto fp8_options = g_c.options().dtype(torch::kFloat8_e4m3fn);
+    auto g8 = torch::empty({m, n}, fp8_options);
+    auto gt8 = masks[1] ? torch::empty({n, m}, fp8_options) : torch::Tensor();
+    auto wt8 = masks[0] ? torch::empty({k, n}, fp8_options) : torch::Tensor();
+    auto xt8 = masks[1] ? torch::empty({k, m}, fp8_options) : torch::Tensor();
+
+    int64_t block = 256;
+    quantize_kernel<__nv_fp8_e4m3>
+        <<<(unsigned)((m * n + block - 1) / block), block, 0, stream.stream()>>>(
+            reinterpret_cast<const __nv_bfloat16*>(g_c.data_ptr()), sgi_ptr,
+            reinterpret_cast<__nv_fp8_e4m3*>(g8.data_ptr()), amax_g_ptr, m * n);
+    dim3 threads(32, 8);
+    if (masks[0]) {
+        dim3 blocks((k + 31) / 32, (n + 31) / 32);
+        transpose_quantize_kernel<__nv_fp8_e4m3>
+            <<<blocks, threads, 0, stream.stream()>>>(
+                reinterpret_cast<const __nv_bfloat16*>(w_c.data_ptr()), swi_ptr,
+                reinterpret_cast<__nv_fp8_e4m3*>(wt8.data_ptr()), amax_g_ptr,
+                n, k);
+        fp8_gemm_into(g8, wt8, grad_input.reshape({m, k}), m, n, k, sg_ptr,
+                      sw_ptr, stream.stream());
+    }
+    if (masks[1]) {
+        dim3 g_blocks((n + 31) / 32, (m + 31) / 32);
+        dim3 x_blocks((k + 31) / 32, (m + 31) / 32);
+        transpose_quantize_kernel<__nv_fp8_e4m3>
+            <<<g_blocks, threads, 0, stream.stream()>>>(
+                reinterpret_cast<const __nv_bfloat16*>(g_c.data_ptr()), sgi_ptr,
+                reinterpret_cast<__nv_fp8_e4m3*>(gt8.data_ptr()), amax_g_ptr,
+                m, n);
+        transpose_quantize_kernel<__nv_fp8_e4m3>
+            <<<x_blocks, threads, 0, stream.stream()>>>(
+                reinterpret_cast<const __nv_bfloat16*>(x_c.data_ptr()), sxi_ptr,
+                reinterpret_cast<__nv_fp8_e4m3*>(xt8.data_ptr()), amax_g_ptr,
+                m, k);
+        fp8_gemm_into(gt8, xt8, grad_weight, n, m, k, sg_ptr, sx_ptr,
+                      stream.stream());
+    }
+    C10_CUDA_CHECK(cudaGetLastError());
+    if (masks[2]) {
+        grad_bias = g_c.sum(0).to(g.dtype());
+    }
+    return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>(
+        grad_input, grad_weight, grad_bias);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"),
+          "FP8 e4m3 GEMM: a[M,K] x b[N,K] -> bf16[M,N] (pre-scaled inputs)");
+    m.def("fp8_linear_forward_scaled", &fp8_linear_forward_scaled,
+          py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("sx"),
+          py::arg("sw"), py::arg("sx_inv"), py::arg("sw_inv"),
+          py::arg("amax_x"), py::arg("amax_w"),
+          "Scaled FP8 linear forward: quantize with per-tensor scales + "
+          "cublasLt GEMM (scales applied inside) + bias -> bf16");
+    m.def("fp8_linear_backward_scaled", &fp8_linear_backward_scaled,
+          py::arg("g"), py::arg("x"), py::arg("w"), py::arg("masks"),
+          py::arg("sg"), py::arg("sw"), py::arg("sx"), py::arg("sg_inv"),
+          py::arg("sw_inv"), py::arg("sx_inv"), py::arg("amax_g"),
+          "Scaled FP8 linear backward: dX = g*sw @ W, dW = (g*sx)^T @ X, "
+          "dB = sum(g)");
+}

csrc/tests/attn_paged_test.cu

@@ -9,6 +9,9 @@
 #include "test_utils.cuh"
 #include "../kernels/attn_dispatchers.cuh"
 
+struct PagedDecodeDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_paged_decode<H>(p, 0); } };
+struct PagedPrefillDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_paged_prefill<H>(p, 0); } };
+
 // ---- CPU reference: paged decode with variable seq_lens ----
 // Q: [B, Hq, D], K/V pool: [pool_size, Hkv, D]
 // req_to_token: [num_reqs, max_ctx_len], req_pool_indices: [B]
@@ -65,7 +68,7 @@ static void cpu_paged_prefill_ref(
     const float* Q, const float* K_pool, const float* V_pool,
     const int64_t* req_to_token, const int64_t* req_pool_indices,
     const int* kv_indptr, const int* qo_indptr,
-    const bool* mask, int mask_q_stride, int mask_kv_stride,
+    const bool* mask, int mask_l_stride, int mask_kv_stride,
     int B, int Hq, int Hkv, int D, int max_ctx_len, int causal,
     float* O)
 {
@@ -84,7 +87,7 @@ static void cpu_paged_prefill_ref(
                 float accum[256] = {0.0f};
                 int lim = causal ? min(seq_len, causal_off + qi + 1) : seq_len;
                 for (int kj = 0; kj < lim; kj++) {
-                    if (mask && !mask[b * mask_q_stride * mask_kv_stride
+                    if (mask && !mask[b * mask_l_stride * mask_kv_stride
                                      + qi * mask_kv_stride + kj]) continue;
                     int64_t slot = req_to_token[req_idx * max_ctx_len + kj];
                     float dot = 0.0f;
@@ -127,14 +130,15 @@ inline void print_paged_row(const char* cfg, float max_err, bool pass) {
 // ======================================================================
 template <int HEAD_DIM>
 static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
-                            int causal, int seed) {
+                            int causal, int seed, int context_capacity = 0,
+                            int fixed_seq_len = 0) {
     // Variable seq_lens per request
     srand(seed);
     std::vector<int> seq_lens(B);
     for (int b = 0; b < B; b++)
-        seq_lens[b] = 8 + rand() % (max_seq - 8);
+        seq_lens[b] = fixed_seq_len ? fixed_seq_len : 8 + rand() % (max_seq - 8);
     int max_sl = *std::max_element(seq_lens.begin(), seq_lens.end());
-    int max_ctx = max_sl + 16;
+    int max_ctx = context_capacity ? context_capacity : max_sl + 16;
 
     int pool_size = B * max_ctx;
     int num_reqs = B + 4;
@@ -214,19 +218,19 @@ static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
     // Kernel launch
     AttentionParams<bf16> p;
     p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
-    p.head_dim = HEAD_DIM; p.total_q = B;
-    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
-    p.max_context_len = max_ctx; p.max_seq_len = max_sl;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
     p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
     p.mask = nullptr; p.mask_b_stride = 0;
-    p.mask_h_stride = 0; p.mask_q_stride = 0;
+    p.mask_h_stride = 0; p.mask_l_stride = 0;
     p.scale = 1.0f / sqrtf((float)HEAD_DIM);
-    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
     p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
     p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
-    p.o = d_o; p.o_part = d_op; p.ml_part = d_ml;
+    p.o_ptr = d_o; p.o_part = d_op; p.ml_part = d_ml;
 
-    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p, 0); });
+    dispatch_by_head_dim(HEAD_DIM, PagedDecodeDispatch{p});
     cudaDeviceSynchronize();
 
     bf16* h_o_bf = (bf16*)malloc(sz_q);
@@ -234,7 +238,7 @@ static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
     float* h_o_got = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
     for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
 
-    const float atol = 0.02f, rtol = 0.02f;
+    const float atol = 0.01f, rtol = 0.01f;
     bool pass = true;
     float max_err = 0.0f;
     for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
@@ -349,19 +353,19 @@ static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
 
     AttentionParams<bf16> p;
     p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
-    p.head_dim = HEAD_DIM; p.total_q = B;
-    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
-    p.max_context_len = max_ctx; p.max_seq_len = max_sl;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
     p.causal_offset = -1; p.use_mask = 1;
     p.mask = d_mask; p.mask_b_stride = max_sl;
-    p.mask_h_stride = 0; p.mask_q_stride = 0;
+    p.mask_h_stride = 0; p.mask_l_stride = 0;
     p.scale = 1.0f / sqrtf((float)HEAD_DIM);
-    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
     p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
     p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
-    p.o = d_o; p.o_part = d_op; p.ml_part = d_ml;
+    p.o_ptr = d_o; p.o_part = d_op; p.ml_part = d_ml;
 
-    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p, 0); });
+    dispatch_by_head_dim(HEAD_DIM, PagedDecodeDispatch{p});
     cudaDeviceSynchronize();
 
     bf16* h_o_bf = (bf16*)malloc(sz_q);
@@ -369,7 +373,7 @@ static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
     float* h_o_got = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
     for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
 
-    const float atol = 0.02f, rtol = 0.02f;
+    const float atol = 0.01f, rtol = 0.01f;
     bool pass = true;
     float max_err = 0.0f;
     for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
@@ -482,22 +486,20 @@ static int run_prefill_test(int B, int Hq, int Hkv,
     // Kernel launch
     AttentionParams<bf16> p;
     p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
-    p.head_dim = HEAD_DIM; p.total_q = total_q;
-    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
-    p.max_context_len = max_ctx; p.max_seq_len = max_sl;
-    int max_ql = 0;
-    for (int b = 0; b < B; b++) max_ql = max(max_ql, q_lens[b]);
-    p.max_q_len = max_ql;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
+    p.q_len = total_q;
     p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
     p.mask = nullptr; p.mask_b_stride = 0;
-    p.mask_h_stride = 0; p.mask_q_stride = 0;
+    p.mask_h_stride = 0; p.mask_l_stride = 0;
     p.scale = 1.0f / sqrtf((float)HEAD_DIM);
-    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
     p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
     p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
-    p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
+    p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr;
 
-    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p, 0); });
+    dispatch_by_head_dim(HEAD_DIM, PagedPrefillDispatch{p});
     cudaDeviceSynchronize();
 
     bf16* h_o_bf = (bf16*)malloc(sz_q);
@@ -505,7 +507,7 @@ static int run_prefill_test(int B, int Hq, int Hkv,
     float* h_o_got = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
     for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
 
-    const float atol = 0.02f, rtol = 0.02f;
+    const float atol = 0.01f, rtol = 0.01f;
     bool pass = true;
     float max_err = 0.0f;
     for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) {
@@ -619,20 +621,20 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
 
     AttentionParams<bf16> p;
     p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
-    p.head_dim = HEAD_DIM; p.total_q = total_q;
-    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
-    p.max_context_len = max_ctx; p.max_seq_len = q_len;
-    p.max_q_len = q_len;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
+    p.q_len = B * q_len;
     p.causal_offset = -1; p.use_mask = 1;
     p.mask = d_mask; p.mask_b_stride = q_len * q_len;
-    p.mask_h_stride = 0; p.mask_q_stride = q_len;
+    p.mask_h_stride = 0; p.mask_l_stride = q_len;
     p.scale = 1.0f / sqrtf((float)HEAD_DIM);
-    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
     p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
     p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
-    p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
+    p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr;
 
-    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p, 0); });
+    dispatch_by_head_dim(HEAD_DIM, PagedPrefillDispatch{p});
     cudaDeviceSynchronize();
 
     bf16* h_o_bf = (bf16*)malloc(sz_q);
@@ -640,7 +642,7 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
     float* h_o_got = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
     for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
 
-    const float atol = 0.02f, rtol = 0.02f;
+    const float atol = 0.01f, rtol = 0.01f;
     bool pass = true;
     float max_err = 0.0f;
     for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) {
@@ -709,19 +711,19 @@ static void bench_decode(int B, int Hq, int Hkv, int seq_len) {
 
     AttentionParams<bf16> p;
     p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
-    p.head_dim = HEAD_DIM; p.total_q = B;
-    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
-    p.max_context_len = max_ctx; p.max_seq_len = seq_len;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
     p.causal_offset = 0; p.use_mask = 0;
     p.mask = nullptr; p.mask_b_stride = 0;
     p.scale = 1.0f / sqrtf((float)HEAD_DIM);
-    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
     p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
     p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
-    p.o = d_o; p.o_part = d_op; p.ml_part = d_ml;
+    p.o_ptr = d_o; p.o_part = d_op; p.ml_part = d_ml;
 
     auto launch = [&]() {
-        dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p, 0); });
+        dispatch_by_head_dim(HEAD_DIM, PagedDecodeDispatch{p});
     };
     // Decode: q_len=1, query is the last token → attends to all [0, seq_len).
     // FLOPs = 2 * (QK^T + PV) = 4 * B * Hq * seq_len * D.
@@ -786,20 +788,20 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau
 
     AttentionParams<bf16> p;
     p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
-    p.head_dim = HEAD_DIM; p.total_q = total_q;
-    p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
-    p.max_context_len = max_ctx; p.max_seq_len = kv_len;
-    p.total_q = total_q; p.max_q_len = q_len;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
+    p.q_len = B * q_len;
     p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
     p.mask = nullptr; p.mask_b_stride = 0;
     p.scale = 1.0f / sqrtf((float)HEAD_DIM);
-    p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
     p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
     p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
-    p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
+    p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr;
 
     auto launch = [&]() {
-        dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p, 0); });
+        dispatch_by_head_dim(HEAD_DIM, PagedPrefillDispatch{p});
     };
     // FLOPs = 2 * (QK^T + PV) = 4 * effective_qk_pairs * Hq * D.
     // Non-causal: effective = q_len * kv_len.
@@ -844,6 +846,9 @@ int main() {
     fail += run_decode_test<256>(1, 2, 1, 256, 0, 9);
     fail += run_decode_test<128>(16, 32, 4, 2048, 0, 10);
     fail += run_decode_test<128>(32, 32, 4, 1024, 0, 11);
+    // Production keeps a fixed 32768-wide request table.  This forces 32
+    // splits, so seq_len > 512 gives each split multiple cp.async tiles.
+    fail += run_decode_test<64>(1, 24, 4, 1100, 0, 12, 32768, 1100);
 
     // Decode with 2D mask (regression: mixed seq_lens + HasMask)
     fail += run_decode_mask_test<128>(2, 8, 2, 256, 30);

csrc/tests/attn_test.cu

@@ -9,6 +9,9 @@ nvcc -I csrc -arch=sm_89 -O3 \
 #include "test_utils.cuh"
 #include "../kernels/attn_dispatchers.cuh"
 
+struct DecodeDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_decode<H>(p, 0); } };
+struct PrefillDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_prefill<H>(p, 0); } };
+
 // Split-K scratch (torch-free)
 struct DecodeScratch {
     float* o_part = nullptr;
@@ -30,8 +33,6 @@ static void free_scratch(DecodeScratch& sc) {
 // ======================================================================
 
 static int run_decode_test(int B, int Hq, int Hk, int sl, int D, int causal) {
-    int gs = Hq / Hk;
-
     size_t nQ = B*Hq*1*D, nKV = B*Hk*sl*D;
     float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
     for (size_t i=0;i<nQ;i++) hQ[i]=randf();
@@ -60,14 +61,14 @@ static int run_decode_test(int B, int Hq, int Hk, int sl, int D, int causal) {
     p.use_mask=0; p.causal_offset=causal?0:-1;
     p.scale=1.0f/sqrtf((float)D);
     set_default_strides(p);
-    p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
+    p.q_ptr=dQ; p.k_ptr=dK; p.v_ptr=dV; p.mask=nullptr; p.o_ptr=dO;
 
     DecodeScratch sc;
     setup_scratch(p, sc);
     p.o_part = sc.o_part; p.ml_part = sc.ml_part;
 
     double t0=now_ms();
-    dispatch_by_head_dim(D, [&]<int H>() { dispatch_decode<H>(p, 0); });
+    dispatch_by_head_dim(D, DecodeDispatch{p});
     cudaDeviceSynchronize();
     (void)t0;
     cudaError_t err=cudaGetLastError();
@@ -140,13 +141,13 @@ static void bench_decode() {
         p.head_dim = D; p.use_mask = 0; p.causal_offset = -1;
         p.scale = 1.0f / sqrtf((float)D);
         set_default_strides(p);
-        p.q = dQ; p.k = dK; p.v = dV; p.mask = nullptr; p.o = dO;
+        p.q_ptr = dQ; p.k_ptr = dK; p.v_ptr = dV; p.mask = nullptr; p.o_ptr = dO;
 
         DecodeScratch sc;
         setup_scratch(p, sc);
         p.o_part = sc.o_part; p.ml_part = sc.ml_part;
 
-        auto launch = [&]() { dispatch_by_head_dim(D, [&]<int H>() { dispatch_decode<H>(p, 0); }); };
+        auto launch = [&]() { dispatch_by_head_dim(D, DecodeDispatch{p}); };
         double flops = 4.0 * B * Hq * (double)sl * D;
         BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops);
 
@@ -187,10 +188,10 @@ static int run_prefill_test(int B, int Hq, int Hk, int ql, int kl, int D, int ca
     p.use_mask=0; p.causal_offset=causal?0:-1;
     set_default_strides(p);
     p.scale=1.0f/sqrtf((float)D);
-    p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
+    p.q_ptr=dQ; p.k_ptr=dK; p.v_ptr=dV; p.mask=nullptr; p.o_ptr=dO;
 
     double t0=now_ms();
-    dispatch_by_head_dim(D, [&]<int H>() { dispatch_prefill<H>(p, 0); });
+    dispatch_by_head_dim(D, PrefillDispatch{p});
     cudaDeviceSynchronize();
     (void)t0;
     cudaError_t err=cudaGetLastError();
@@ -261,9 +262,9 @@ static void bench_prefill() {
         p.use_mask=0; p.causal_offset=causal?0:-1;
         set_default_strides(p);
         p.scale=1.0f/sqrtf((float)D);
-        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
+        p.q_ptr=dQ; p.k_ptr=dK; p.v_ptr=dV; p.mask=nullptr; p.o_ptr=dO;
 
-        auto launch = [&]() { dispatch_by_head_dim(D, [&]<int H>() { dispatch_prefill<H>(p, 0); }); };
+        auto launch = [&]() { dispatch_by_head_dim(D, PrefillDispatch{p}); };
         for (int i=0;i<WARMUP;i++) launch();
         cudaDeviceSynchronize();
         cudaError_t err=cudaGetLastError();
@@ -322,6 +323,7 @@ int main() {
     // ---- PREFILL ----
     {
         const int configs[][7] = {
+            {1,2,1,64,128,32,0},     // scalar fallback D=32
             {1,2,1,64,128,64,0},     // tiny: B,Hq,Hk,q,kv,D,causal
             {1,32,4,512,512,128,0},  // standard
             {1,32,4,128,256,128,0},  // medium

csrc/tests/test_utils.cuh

@@ -107,29 +107,29 @@ void dispatch_by_head_dim(int head_dim, Fn&& fn) {
 // Set default strides for contiguous b h l d layout on AttentionParams.
 template<typename P>
 inline void set_default_strides(P& p) {
-    p.q_stride_b  = p.q_head * p.q_len * p.head_dim;
-    p.q_stride_h  = p.q_len * p.head_dim;
-    p.q_stride_l  = p.head_dim;
-    p.q_stride_d  = 1;
-    p.kv_stride_b = p.kv_head * p.kv_len * p.head_dim;
-    p.kv_stride_h = p.kv_len * p.head_dim;
-    p.kv_stride_l = p.head_dim;
-    p.kv_stride_d = 1;
+    p.q_b_stride  = p.q_head * p.q_len * p.head_dim;
+    p.q_h_stride  = p.q_len * p.head_dim;
+    p.q_l_stride  = p.head_dim;
+    p.q_d_stride  = 1;
+    p.kv_b_stride = p.kv_head * p.kv_len * p.head_dim;
+    p.kv_h_stride = p.kv_len * p.head_dim;
+    p.kv_l_stride = p.head_dim;
+    p.kv_d_stride = 1;
     p.mask_b_stride = p.kv_len;
     p.mask_h_stride = 0;
-    p.mask_q_stride = 0;
+    p.mask_l_stride = 0;
 }
 
 // Set default Q strides for a paged decode params struct.
 template<typename P>
 inline void set_default_paged_strides(P& p) {
-    p.q_stride_b  = p.q_head * p.q_len * p.head_dim;
-    p.q_stride_h  = p.q_len * p.head_dim;
-    p.q_stride_l  = p.head_dim;
-    p.q_stride_d  = 1;
+    p.q_b_stride  = p.q_head * p.q_len * p.head_dim;
+    p.q_h_stride  = p.q_len * p.head_dim;
+    p.q_l_stride  = p.head_dim;
+    p.q_d_stride  = 1;
     p.mask_b_stride = p.kv_len;
     p.mask_h_stride = 0;
-    p.mask_q_stride = 0;
+    p.mask_l_stride = 0;
 }
 
 // Generic CPU reference for multi-query / grouped-query attention.

docker-compose.yml

@@ -5,7 +5,9 @@ services:
       dockerfile: Dockerfile
       args:
         CUDA_TAG: ${CUDA_TAG:-cu128}
-    user: "${UID:-1000}:${GID:-1000}"
+        USER_UID: ${ASTRAI_UID:-1000}
+        USER_GID: ${ASTRAI_GID:-1000}
+    user: "${ASTRAI_UID:-1000}:${ASTRAI_GID:-1000}"
     ports:
       - "8000:8000"
     volumes:
@@ -33,7 +35,9 @@ services:
       dockerfile: Dockerfile
       args:
         CUDA_TAG: ${CUDA_TAG:-cu128}
-    user: "${UID:-1000}:${GID:-1000}"
+        USER_UID: ${ASTRAI_UID:-1000}
+        USER_GID: ${ASTRAI_GID:-1000}
+    user: "${ASTRAI_UID:-1000}:${ASTRAI_GID:-1000}"
     ports:
       - "8000:8000"
     volumes:
@@ -46,3 +50,44 @@ services:
       retries: 3
       start_period: 120s
     restart: unless-stopped
+
+  trainer:
+    profiles: [train]
+    build:
+      context: .
+      dockerfile: Dockerfile
+      args:
+        CUDA_TAG: ${CUDA_TAG:-cu128}
+        USER_UID: ${ASTRAI_UID:-1000}
+        USER_GID: ${ASTRAI_GID:-1000}
+    init: true
+    user: "${ASTRAI_UID:-1000}:${ASTRAI_GID:-1000}"
+    volumes:
+      - ${TRAIN_DATA_DIR:-./data}:/data:ro
+      - ${TRAIN_MODEL_DIR:-./params}:/models/base:ro
+      - ${TRAIN_CHECKPOINT_DIR:-./checkpoints}:/checkpoints
+    environment:
+      - TRAIN_JOB_NAME=${TRAIN_JOB_NAME:-astrai-train}
+      - TRAIN_CONFIG=${TRAIN_CONFIG:-}
+      - BASE_MODEL=${BASE_MODEL:-/models/base}
+      - CHECKPOINT_ROOT=/checkpoints
+      - TRAIN_GPU_COUNT=${TRAIN_GPU_COUNT:-all}
+      - CUDA_VISIBLE_DEVICES
+      - NCCL_P2P_DISABLE
+      - NCCL_NET_GDR_LEVEL
+    entrypoint: ["bash", "/app/scripts/docker/train-entrypoint.sh"]
+    ipc: ${TRAIN_IPC_MODE:-host}
+    stop_grace_period: ${TRAIN_STOP_GRACE_PERIOD:-10m}
+    restart: "no"
+    logging:
+      driver: json-file
+      options:
+        max-size: ${TRAIN_LOG_MAX_SIZE:-100m}
+        max-file: ${TRAIN_LOG_MAX_FILES:-5}
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]

docs/README-zh-CN.md

@@ -187,7 +187,7 @@ docker run --gpus all -it astrai:latest
 
 # 运行推理服务
 docker run --gpus all -p 8000:8000 astrai:latest \
-  python -m scripts.tools.server --port 8000 --device cuda
+  python scripts/tools/server.py --port 8000 --device cuda
 
 # 挂载数据卷
 docker run --gpus all -v /path/to/data:/data -it astrai:latest

docs/developer/architecture.md

@@ -315,6 +315,7 @@ classDiagram
             <<TypedDict>>
             +Tensor hidden_states
             +Optional[Tensor] aux_loss
+            +Optional[RouterStats] router_stats
         }
 
         class GQA {
@@ -361,6 +362,7 @@ classDiagram
             <<TypedDict>>
             +Tensor hidden_states
             +Optional[Tensor] aux_loss
+            +Optional[RouterStats] router_stats
         }
 
         class DeepSeekMoE {
@@ -807,7 +809,6 @@ classDiagram
             +AutoTokenizer tokenizer
             +InferenceScheduler scheduler
             +generate(prompt, stream, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window) Union[Generator, str, List[str]]
-            +generate_with_request(request) Union[Generator, str, List[str]]
             +generate_async(prompt, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window) AsyncGenerator
             +get_stats() Dict
             +shutdown()
@@ -915,6 +916,9 @@ classDiagram
             +int max_len
             +Optional[Tensor] kv_indptr
             +Optional[Tensor] qo_indptr
+            +Optional[Tensor] decode_o_part
+            +Optional[Tensor] decode_ml_part
+            +Optional[Tensor] decode_out
         }
 
         class PagePool {
@@ -980,17 +984,6 @@ classDiagram
             +get_stats() Dict
         }
 
-        class GenerationRequest {
-            +List[Dict] messages
-            +int top_k
-            +float top_p
-            +float temperature
-            +Optional[int] max_tokens
-            +float frequency_penalty
-            +int rep_window
-            +bool stream
-        }
-
         class BaseSamplingStrategy {
             <<abstract>>
             +apply(logits, filter_value, input_ids, input_mask) Tensor
@@ -1407,7 +1400,7 @@ classDiagram
     CheckpointCallback ..> Checkpoint : creates
     PagePool ..> KVCache : binds
     PagePool ..> InferenceWorkspace : fills
-    InferenceEngine ..> GenerationRequest : uses
+    InferenceEngine ..> GenerateResult : uses
     InferenceEngine ..> GenerateResult : creates
     OpenAIResponseBuilder ..> ChatCompletionRequest : receives
     AnthropicResponseBuilder ..> MessagesRequest : receives
@@ -1443,7 +1436,7 @@ classDiagram
 | **astrai.model** | ModelFactory, AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, LoRAConfig, LoRALinear, RotaryEmbedding, Embedding | Neural network model |
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
 | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory, RawRollout, RolloutResult, BaseRewardModel, RolloutGenerator, RolloutRunner | Training workflow |
-| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
+| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
 | **astrai.extension** | AttentionBackend, TorchNativeBackend, CudaBackend, attn_backend, ATTN_BACKEND, attn_decode, attn_prefill, attn_paged_decode, attn_paged_prefill, rotary_emb, apply_rotary_emb, rotary_backend, is_available | CUDA attention + rotary kernels, backend abstraction, auto-dispatch |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, LaunchStrategy, TorchrunStrategy, LocalStrategy, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler | Distributed parallel & gradient accumulation |
 | **astrai.factory** | BaseFactory | Component registration |
@@ -1462,7 +1455,7 @@ classDiagram
 | **Observer** | `TrainCallback`, callback implementations | Training process monitoring |
 | **Context** | `TrainContext` | Unified training state bag |
 | **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
-| **Strategy (Attention)** | `AttentionBackend`, `TorchNativeBackend`, `CudaBackend` | Attention computation backend switching via context manager |
+| **Strategy (Attention)** | `AttentionBackend`, `CudaBackend`, `FlashAttnBackend`, `TorchNativeBackend` | Attention computation backend switching via context manager |
 | **Auto-dispatch (Rotary)** | `apply_rotary_emb`, `rotary_backend.py`, `rotary_ops.py` | Rotary embedding CUDA kernel auto-dispatch with torch fallback |
 | **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor`, `FSDPExecutor` | Gradient accumulation & model distribution |
 | **Storage** | `Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support |
@@ -1475,7 +1468,7 @@ classDiagram
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` for loss, `BaseExecutor` for gradient accumulation + model distribution
 3. **Strategy Selection**: `StrategyFactory` creates strategy by `train_type`
 4. **Executor Selection**: `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)` → `NoneExecutor` / `DDPExecutor` / `FSDPExecutor`
-5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (`TorchNativeBackend` default, `CudaBackend` for CUDA kernels). Rotary embedding auto-dispatches to CUDA kernel when available (inference mode), else torch complex multiply (training).
+5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (cuda > flash > torch priority; `ASTR_BACKEND` env var overrides default; `TorchNativeBackend` fallback). Rotary embedding auto-dispatches to CUDA kernel when available, else torch complex multiply.
 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
 7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (`MmapStore`/`JsonlStore`) loads data with explicit `_length` and multi-segment `_data`
 8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata; `CheckpointCallback` performs rank-0 training saves, with extra state saved as `{key}.pt`

docs/developer/cuda_kernels.md

@@ -86,8 +86,12 @@ Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 mod
 `astrai/extension/attention_backend.py` provides the backend abstraction:
 
 - **`AttentionBackend`** (ABC): `fwd_decode` / `fwd_prefill` abstract methods, `forward` dispatches by q_len
-- **`TorchNativeBackend`**: SDPA with indirect KV cache gather (default)
-- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`)
+- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`). Default on GPU.
+- **`FlashAttnBackend`**: Optional flash-attn dispatch with `flash_attn_with_kvcache` fast path.
+- **`TorchNativeBackend`**: SDPA with indirect KV cache gather (always-available fallback)
+
+Default priority: cuda > flash > torch. Set ``ASTR_BACKEND=cuda|torch_native|flash``
+to override the default.
 
 Select a backend via context manager (mirrors `torch.nn.attention.sdpa_kernel`):
 
@@ -98,7 +102,7 @@ with attn_backend(ATTN_BACKEND.CUDA):
     engine.generate("hello")
 ```
 
-`CudaBackend` falls back to `TorchNativeBackend` when a kernel is not available.
+`CudaBackend` falls back to `FlashAttnBackend` (when flash-attn is installed and supports the input dtype) or `TorchNativeBackend` otherwise.
 
 ### Rotary Backend
 

docs/developer/dataflow.md

@@ -72,7 +72,7 @@ metadata, but the preprocessing `BinWriter` currently does not write offsets.
 
 - If `load_path` is a file: `.jsonl` selects `"jsonl"`; other suffixes raise `ValueError`.
 - If `load_path` is a directory: any recursive `*.bin` plus a `meta.json` selects `"bin"`; otherwise any recursive `*.jsonl` selects `"jsonl"`.
-- Detection does not require `dataset_config.json`; configuration is selected later when `JsonlStore.load()` chooses a transform.
+- Detection does not require `dataset_config.json`; configuration is selected later when `DatasetFactory.load()` constructs a transform via `_build_jsonl_transform()` and passes it to `JsonlStore.load()`.
 
 ### Store Backends
 
@@ -89,12 +89,17 @@ access methods.
 **MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`. `segments_are_records=False` — bin segments are contiguous streams; record access is driven by `_offsets` (written when `save_bin(..., record_keys=...)` was used at preprocessing time).
 
 **JsonlStore**: Reads a `.jsonl` file or the sorted top-level `*.jsonl` files in
-a directory. Eager transform selection uses the first available route:
+a directory. Eager transform selection is owned by
+`DatasetFactory._build_jsonl_transform()` (called from `DatasetFactory.load()`) —
+the factory picks the first available route:
 
-1. An explicit `transform=` argument.
+1. An explicit `transform=` argument passed through `store.load()`.
 2. `dataset_config.json` in the JSONL directory. It follows `PipelineConfig` and may add `tokenizer_path`; when omitted, the config directory is used.
 3. The built-in `messages` transform when `tokenizer_path=` is supplied. It masks system/user turns, trains assistant turns, and emits document-reset position IDs.
 
+`JsonlStore.load()` requires `transform=` to be passed explicitly for eager mode
+(raises `ValueError` if missing).
+
 Only DPO gets an automatic lazy route from `DatasetFactory`: raw JSONL plus
 `tokenizer_path` installs `dpo_processor` and tokenizes each record in
 `fetch_record`. GRPO does not currently have an automatic lazy processor.

docs/developer/docker-training.md

@@ -0,0 +1,57 @@
+# Containerized Training Deployment
+
+Rules for running AstrAI distributed training in containers, distilled from real deployment failures. Read before touching `Dockerfile`, `docker-compose.yml`, `scripts/train.sh`, `train-entrypoint.sh`. AGENTS.md mirrors this locally; this file is the committed version.
+
+## Architecture
+
+```
+scripts/train.sh           host-side CLI: env loading, preflight, compose wrapper, lifecycle
+  └── docker-compose.yml   GPU passthrough, mounts, in-container env vars, entrypoint
+        └── train-entrypoint.sh   GPU-count resolution, parallel-mode selection, auto-resume
+              └── train.py --config /run/astrai/train.yaml
+```
+
+| Layer | Responsible for | NOT responsible for |
+|-------|-----------------|---------------------|
+| `train.sh` | host paths, `.env.train`, preflight, lifecycle | training args, GPU selection, parallel mode |
+| compose | GPU passthrough, mounts, in-container env (NCCL) | training args (beyond `TRAIN_*` forwarding) |
+| entrypoint | `--ckpt_dir/--nprocs/--parallel_mode/--param_path`, resume | hyperparameters (YAML/CLI) |
+| `train.yaml` | hyperparameters (`_merge_yaml_into_kwargs`, CLI wins) | container paths, process count |
+
+## Path Conventions
+
+| Host var | Container | Perm | Purpose |
+|---|---|---|---|
+| `TRAIN_DATA_DIR` | `/data` | ro | dataset (`data_root_path` must be `/data`) |
+| `TRAIN_MODEL_DIR` | `/models/base` | ro | base model (`config.json` + `model.safetensors`) |
+| `TRAIN_CHECKPOINT_DIR` | `/checkpoints` | rw | checkpoint root, per-`TRAIN_JOB_NAME` subdirs |
+| `TRAIN_CONFIG_FILE` | `/run/astrai/train.yaml` | ro | training YAML (mounted only on `start`) |
+| code | `/app` | image | **not a mount**; rebuild image for code changes |
+
+## Hard Rules
+
+1. **Filter GPUs once**: compose passes the full physical set (`count: all`); `CUDA_VISIBLE_DEVICES` filters inside by physical index. Never `count: N` + physical indices (double filter leaves 1 card → `device_id out of range`).
+2. **In-container UID = host UID**: Dockerfile builds the user via `USER_UID/USER_GID` args; `train.sh` injects `ASTRAI_UID/GID` (bash `UID` is readonly). compose `user:` alone does not create the /etc/passwd entry — torch's `getpass.getuser()` then dies with `uid not found`.
+3. **In-container env vars are explicit**: `.env.train` (`--env-file`) is only compose's interpolation dictionary — never reaches the container. A var arrives only via a value-less `environment` entry (`- VAR`, read from the calling process env).
+4. **NCCL hang workaround** (this host): `NCCL_P2P_DISABLE=1` + `NCCL_NET_GDR_LEVEL=0` must be in-container.
+5. **Checkpoint complete =** `meta.json + config.json + model.safetensors + optimizer.pt + scheduler.pt`; `start` auto-resumes the latest complete one.
+6. **tqdm is silent without a TTY**: add `disable=False` in `astrai/trainer/train_callback.py`; `metric.jsonl` (per step) works as progress evidence regardless.
+
+## Operations
+
+```bash
+bash scripts/train.sh init        # first run: dirs + .env.train (edit per machine)
+bash scripts/train.sh preflight   # validate Docker/paths/GPU/model/YAML/compose
+bash scripts/train.sh start       # build + start in background (auto-resume)
+bash scripts/train.sh start --foreground -- --dry-run   # print plan only
+bash scripts/train.sh logs | status | stop | restart
+bash scripts/train.sh clean --keep 5   # prune old checkpoints (--force to delete)
+```
+
+## Files
+
+- `docker-compose.yml` — trainer service: `count: all`, `ASTRAI_UID/GID` build args + `user:`, env whitelist, mounts
+- `Dockerfile` — production stage builds user from `USER_UID/USER_GID`; `ENV HOME=/home/astrai`; `USER astrai`
+- `scripts/train.sh` — `load_env` filters `UID=` lines (readonly var); `compose()` injects `ASTRAI_UID/GID`
+- `scripts/docker/train-entrypoint.sh` — GPU-count resolution, parallel mode, resume
+- `.env.train`, `train.yaml` — host-specific; templates from `scripts/train.sh init`; scientific-notation floats (`2e-5`) parse correctly since train.py uses the YAML 1.2 float schema

docs/developer/internals.md

@@ -178,8 +178,10 @@ Three-layer separation (SGLang-inspired):
 
 Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/attention_backend.py`):
 
-- **`TorchNativeBackend`** (default): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
-- **`CudaBackend`**: decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool). Falls back to `TorchNativeBackend` when kernel unavailable.
+- **`CudaBackend`** (default): decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool). Falls back to `FlashAttnBackend` when dtype unsupported.
+- **`FlashAttnBackend`**: optional flash-attn dispatch with `flash_attn_with_kvcache` fast path for contiguous cache; falls back to KV gather + `flash_attn_func`.
+- **`TorchNativeBackend`** (always-available fallback): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
+- Default priority: cuda > flash > torch. Set `ASTR_BACKEND=cuda|torch_native|flash` to override.
 
 Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.py`, which auto-dispatches to the fused CUDA kernel (`rotary_emb.cu`) during inference or torch complex multiply during training (for autograd compatibility). Both attention backends share the same rotary dispatch.
 

docs/get-started.md

@@ -36,7 +36,7 @@ pip install -e .
 # pip install -e ".[dev]"
 ```
 
-> **CUDA kernels** are opt-in. They are not built by default. When built, they can be activated via `with attn_backend(ATTN_BACKEND.CUDA):` for accelerated decode/prefill, and the fused rotary embedding kernel is auto-dispatched when available. You can skip them for normal usage.
+> **CUDA kernels** are opt-in at build time (`CSRC_KERNELS=true`). Once built, `CudaBackend` is the default attention backend on GPU (cuda > flash > torch priority). Override via `ASTR_BACKEND` env var or `attn_backend()` context manager. Fused rotary embedding kernel is auto-dispatched when available. Skip for CPU-only usage.
 
 ## 2. Download Model Weights
 
@@ -232,7 +232,7 @@ docker build -t astrai:latest .
 
 # Run inference server with GPU
 docker run --gpus all -p 8000:8000 astrai:latest \
-  python -m scripts.tools.server --port 8000 --device cuda
+  python scripts/tools/server.py --port 8000 --device cuda
 
 # Docker Compose (GPU)
 docker compose up -d

docs/guides/evaluation.md

@@ -214,6 +214,7 @@ python scripts/eval/evaluate_ifd.py \
 | `--sentinel_text` | `\n` | Prefix for unconditional pass (`""` → bos/pad fallback) |
 | `--per_token` | False | Include per-token IFD breakdown |
 | `--max_samples` | None | Random subsample per file |
+| `--append_eos` / `--no-append_eos` | `True` | Append (or skip) EOS token to instruction/response |
 
 **How it works**: Two forward passes per batch — (1) conditional: packed BFD sequence with context + response, (2) unconditional: response prefixed with a sentinel. IFD = mean_conditional_loss / mean_unconditional_loss. IFD > 1 means the instruction makes the response harder to predict (higher quality data).
 

docs/guides/inference.md

@@ -65,10 +65,14 @@ Attention computation (cache I/O + SDPA/kernel dispatch) is decoupled from the m
 
 ```
 AttentionBackend (ABC)
-  ├── TorchNativeBackend   SDPA + indirect KV cache gather (default)
-  └── CudaBackend          CUDA kernel dispatch (attn_paged_decode, attn_paged_prefill)
+  ├── CudaBackend          CUDA kernel dispatch (default on GPU)
+  ├── FlashAttnBackend     Optional flash-attn dispatch (fallback)
+  └── TorchNativeBackend   SDPA + indirect KV cache gather (always-available fallback)
 ```
 
+Default priority: cuda > flash > torch.  Set ``ASTR_BACKEND=cuda|torch_native|flash``
+to override.
+
 Select via context manager (mirrors `torch.nn.attention.sdpa_kernel`):
 
 ```python
@@ -82,7 +86,7 @@ with attn_backend(ATTN_BACKEND.CUDA):
 
 `CudaBackend` prefill path: writes K/V, then calls `attn_paged_prefill` — a ragged-batch (paged) prefill kernel that reads K/V directly from the flat pool via `req_to_token`, addressing each request's `q_len`/`kv_len` through `qo_indptr` and `kv_indptr`. No explicit K/V gather needed.
 
-Fallback: `CudaBackend` delegates to `TorchNativeBackend` when a CUDA kernel is not available.
+Fallback: when `CudaBackend` cannot handle an input (wrong dtype or head_dim), `FlashAttnBackend` is tried next (if installed), then `TorchNativeBackend`.
 
 ### Rotary Embedding Backend
 
@@ -143,7 +147,6 @@ Adding a protocol = one builder file, no handler subclassing needed.
 ```
 InferenceEngine
   ├── generate(prompt, stream, ...) → str | List[str] | Generator
-  ├── generate_with_request(req)    → same
   ├── generate_async(prompt, ...)   → AsyncGenerator
   ├── get_stats()                   → Dict
   └── shutdown()
@@ -230,19 +233,6 @@ The HTTP protocols and direct engine API have distinct request models and defaul
 | `stream` | Optional[bool] | False | Stream output |
 | `stop_sequences` | Optional[List[str]] | None | Stop sequences |
 
-**Engine** (`GenerationRequest`):
-
-| Param | Type | Default | Description |
-|-------|------|---------|-------------|
-| `messages` | List[Dict[str, str]] | required | Messages to format before generation |
-| `top_k` | int | 50 | Top-k count; 0 disables filtering |
-| `top_p` | float | 1.0 | Nucleus threshold |
-| `temperature` | float | 1.0 | Sampling temperature; 0 enables greedy decoding |
-| `max_tokens` | Optional[int] | None | Max generation length |
-| `frequency_penalty` | float | 0.0 | Frequency penalty (-2.0 to 2.0) |
-| `rep_window` | int | 64 | Recent-token window used by the frequency penalty |
-| `stream` | bool | False | Stream output |
-
 ### SSE Streaming Format
 
 **OpenAI** (`/v1/chat/completions`, `stream=true`):

scripts/demo/stream_chat.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 import torch
 
-from astrai.inference import InferenceEngine
+from astrai import InferenceEngine
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 

scripts/docker.sh

@@ -1,257 +1,95 @@
 #!/bin/bash
-
-# AstrAI Docker Script
-# Build and manage Docker images
-
 set -e
 
-# Colors
-RED='\033[0;31m'
-GREEN='\033[0;32m'
-YELLOW='\033[1;33m'
-BLUE='\033[0;34m'
-NC='\033[0m' # No Color
+IMAGE_NAME="${ASTRAI_IMAGE:-astrai}"
+IMAGE_TAG="${ASTRAI_TAG:-latest}"
+PORT="8000"
+GPU=true
+RUN_ARGS=()
 
-# Default values
-IMAGE_NAME="astrai"
-IMAGE_TAG="latest"
-REGISTRY=""
-CONTAINER_ID=""
+usage() {
+    cat <<EOF
+Usage: $0 <command> [options]
 
-# Print colored messages
-print_info() {
-    echo -e "${BLUE}[INFO]${NC} $1"
+Commands:
+  build              Build the image
+  run [--] [ARGS]    Run a container; ARGS after -- are passed to the container
+
+Options:
+  --gpu              Enable GPU support (default)
+  --no-gpu           Disable GPU support
+  --port PORT        Host port for run (default: 8000)
+  -h, --help         Show this help
+
+Environment:
+  ASTRAI_IMAGE       Image name (default: astrai)
+  ASTRAI_TAG         Image tag (default: latest)
+
+Examples:
+  $0 build
+  $0 run
+  $0 run --port 8080 -- python -m scripts.tools.server --port 8000 --device cuda
+EOF
 }
 
-print_success() {
-    echo -e "${GREEN}[SUCCESS]${NC} $1"
-}
-
-print_error() {
-    echo -e "${RED}[ERROR]${NC} $1"
-}
-
-print_warning() {
-    echo -e "${YELLOW}[WARNING]${NC} $1"
-}
-
-# Check if Docker is installed
-check_docker() {
-    if ! command -v docker &> /dev/null; then
-        print_error "Docker is not installed"
-        exit 1
-    fi
-    print_success "Docker version: $(docker --version)"
-}
-
-# Build Docker image
 build_image() {
-    local dockerfile="${1:-Dockerfile}"
-    local context="${2:-.}"
-
-    if [ ! -f "$dockerfile" ]; then
-        print_error "Dockerfile not found: $dockerfile"
-        exit 1
-    fi
-
-    print_info "Building Docker image: ${IMAGE_NAME}:${IMAGE_TAG}"
-    docker build -t "${IMAGE_NAME}:${IMAGE_TAG}" -f "$dockerfile" "$context"
-    print_success "Image built successfully"
+    docker build -t "${IMAGE_NAME}:${IMAGE_TAG}" .
 }
 
-# Run container
 run_container() {
-    local port="${1:-8000}"
-    local gpu="${2:-false}"
-
-    print_info "Running container on port $port..."
-
-    if [ "$gpu" = true ]; then
-        docker run --gpus all -p "${port}:8000" "${IMAGE_NAME}:${IMAGE_TAG}"
-    else
-        docker run -p "${port}:8000" "${IMAGE_NAME}:${IMAGE_TAG}"
-    fi
+    local gpu_args=()
+    [ "$GPU" = true ] && gpu_args=(--gpus all)
+    docker run "${gpu_args[@]}" -p "${PORT}:8000" "${IMAGE_NAME}:${IMAGE_TAG}" "$@"
 }
 
-# Push image to registry
-push_image() {
-    if [ -z "$REGISTRY" ]; then
-        print_error "Registry not set. Use --registry option"
-        exit 1
-    fi
-
-    local full_tag="${REGISTRY}/${IMAGE_NAME}:${IMAGE_TAG}"
-    print_info "Tagging image: ${full_tag}"
-    docker tag "${IMAGE_NAME}:${IMAGE_TAG}" "$full_tag"
-
-    print_info "Pushing image to registry..."
-    docker push "$full_tag"
-    print_success "Image pushed successfully"
-}
-
-# Remove image
-remove_image() {
-    print_info "Removing image: ${IMAGE_NAME}:${IMAGE_TAG}"
-    docker rmi "${IMAGE_NAME}:${IMAGE_TAG}" 2>/dev/null || print_warning "Image not found"
-    print_success "Image removed"
-}
-
-# Show image info
-show_info() {
-    print_info "Image information:"
-    docker images "${IMAGE_NAME}"
-}
-
-# Show logs
-show_logs() {
-    local container_id="$1"
-    if [ -z "$container_id" ]; then
-        print_error "Container ID required"
-        exit 1
-    fi
-    docker logs "$container_id"
-}
-
-# Main function
 main() {
-    echo "========================================"
-    echo "      AstrAI Docker Management"
-    echo "========================================"
-    echo ""
+    local command=""
 
-    COMMAND=""
-    DOCKERFILE="Dockerfile"
-    CONTEXT="."
-    PORT="8000"
-    GPU=false
-
-    # Parse arguments
     while [[ $# -gt 0 ]]; do
-        case $1 in
-            build)
-                COMMAND="build"
+        case "$1" in
+            build|run)
+                command="$1"
                 shift
                 ;;
-            run)
-                COMMAND="run"
-                shift
-                ;;
-            push)
-                COMMAND="push"
-                shift
-                ;;
-            remove|rm)
-                COMMAND="remove"
-                shift
-                ;;
-            info)
-                COMMAND="info"
-                shift
-                ;;
-            logs)
-                COMMAND="logs"
-                shift
-                ;;
-            --image)
-                IMAGE_NAME="$2"
-                shift 2
-                ;;
-            --tag)
-                IMAGE_TAG="$2"
-                shift 2
-                ;;
-            --registry)
-                REGISTRY="$2"
-                shift 2
-                ;;
-            --dockerfile)
-                DOCKERFILE="$2"
-                shift 2
-                ;;
-            --context)
-                CONTEXT="$2"
-                shift 2
-                ;;
-            --port)
-                PORT="$2"
-                shift 2
-                ;;
-            --container)
-                CONTAINER_ID="$2"
-                shift 2
-                ;;
             --gpu)
                 GPU=true
                 shift
                 ;;
-            --help)
-                echo "Usage: $0 <command> [options]"
-                echo ""
-                echo "Commands:"
-                echo "  build     Build Docker image"
-                echo "  run       Run container"
-                echo "  push      Push image to registry"
-                echo "  remove    Remove image"
-                echo "  info      Show image information"
-                echo "  logs      Show container logs"
-                echo ""
-                echo "Options:"
-                echo "  --image NAME       Image name (default: astrai)"
-                echo "  --tag TAG          Image tag (default: latest)"
-                echo "  --registry URL     Registry URL for push"
-                echo "  --dockerfile FILE  Dockerfile path (default: Dockerfile)"
-                echo "  --context PATH     Build context (default: .)"
-                echo "  --port PORT        Port for run (default: 8000)"
-                echo "  --container ID     Container ID for logs"
-                echo "  --gpu              Enable GPU support"
-                echo "  --help             Show this help message"
-                echo ""
-                echo "Examples:"
-                echo "  $0 build"
-                echo "  $0 build --tag v1.0.0"
-                echo "  $0 run --port 8080"
-                echo "  $0 run --gpu"
-                echo "  $0 logs --container abc123"
-                echo "  $0 push --registry ghcr.io/username"
+            --no-gpu)
+                GPU=false
+                shift
+                ;;
+            --port)
+                PORT="$2"
+                shift 2
+                ;;
+            --)
+                shift
+                RUN_ARGS=("$@")
+                break
+                ;;
+            -h|--help)
+                usage
                 exit 0
                 ;;
             *)
-                if [ -z "$COMMAND" ]; then
-                    print_error "Unknown command: $1"
-                    exit 1
-                fi
-                shift
+                echo "Unknown option: $1" >&2
+                usage >&2
+                exit 1
                 ;;
         esac
     done
 
-    check_docker
-
-    case "$COMMAND" in
+    case "$command" in
         build)
-            build_image "$DOCKERFILE" "$CONTEXT"
+            build_image
             ;;
         run)
-            run_container "$PORT" "$GPU"
-            ;;
-        push)
-            push_image
-            ;;
-        remove)
-            remove_image
-            ;;
-        info)
-            show_info
-            ;;
-        logs)
-            show_logs "$CONTAINER_ID"
-            ;;
-        "")
-            print_error "No command specified. Use --help for usage"
-            exit 1
+            run_container "${RUN_ARGS[@]}"
             ;;
         *)
-            print_error "Unknown command: $COMMAND"
+            echo "No command specified. Use --help for usage" >&2
+            usage >&2
             exit 1
             ;;
     esac

scripts/docker/lib/train-common.sh

@@ -0,0 +1,66 @@
+#!/usr/bin/env bash
+
+log_info() {
+    printf '[INFO] %s\n' "$*"
+}
+
+log_warn() {
+    printf '[WARN] %s\n' "$*" >&2
+}
+
+die() {
+    printf '[ERROR] %s\n' "$*" >&2
+    exit 1
+}
+
+require_command() {
+    command -v "$1" >/dev/null 2>&1 || die "Required command not found: $1"
+}
+
+validate_job_name() {
+    [[ "$1" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] ||
+        die "Invalid TRAIN_JOB_NAME '$1'; use letters, numbers, dot, underscore, or dash"
+}
+
+checkpoint_is_complete() {
+    local checkpoint="$1"
+    local file
+
+    [[ -d "${checkpoint}" ]] || return 1
+
+    for file in meta.json config.json model.safetensors optimizer.pt scheduler.pt; do
+        [[ -s "${checkpoint}/${file}" ]] || return 1
+    done
+
+    return 0
+}
+
+checkpoint_coordinates() {
+    local name
+
+    name="$(basename "$1")"
+    [[ "${name}" =~ ^epoch_([0-9]+)_step_([0-9]+)$ ]] || return 1
+    printf '%d %d\n' "$((10#${BASH_REMATCH[1]}))" "$((10#${BASH_REMATCH[2]}))"
+}
+
+list_complete_checkpoints() {
+    local checkpoint_dir="$1"
+    local checkpoint coordinates epoch step
+
+    for checkpoint in "${checkpoint_dir}"/epoch_*_step_*; do
+        [[ -d "${checkpoint}" ]] || continue
+        coordinates="$(checkpoint_coordinates "${checkpoint}")" || continue
+        checkpoint_is_complete "${checkpoint}" || continue
+        read -r epoch step <<<"${coordinates}"
+        printf '%012d %012d %s\n' "${epoch}" "${step}" "${checkpoint}"
+    done | sort -n -k1,1 -k2,2
+}
+
+find_latest_checkpoint() {
+    local checkpoint_dir="$1"
+    local latest
+
+    latest="$(list_complete_checkpoints "${checkpoint_dir}" | tail -n 1)"
+    [[ -n "${latest}" ]] || return 1
+    printf '%s\n' "${latest#* * }"
+}

scripts/docker/train-entrypoint.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+source "${SCRIPT_DIR}/lib/train-common.sh"
+
+TRAIN_JOB_NAME="${TRAIN_JOB_NAME:?TRAIN_JOB_NAME is required}"
+CHECKPOINT_ROOT="${CHECKPOINT_ROOT:-/checkpoints}"
+CHECKPOINT_DIR="${CHECKPOINT_ROOT}/${TRAIN_JOB_NAME}"
+BASE_MODEL="${BASE_MODEL:-/models/base}"
+TRAIN_CONFIG="${TRAIN_CONFIG:-}"
+TRAIN_GPU_COUNT="${TRAIN_GPU_COUNT:-all}"
+
+validate_job_name "${TRAIN_JOB_NAME}"
+if [[ "${TRAIN_GPU_COUNT}" == "all" ]]; then
+    TRAIN_GPU_COUNT="$(python -c 'import torch; print(torch.cuda.device_count())')"
+fi
+[[ "${TRAIN_GPU_COUNT}" =~ ^[1-9][0-9]*$ ]] || die "No visible GPU found"
+if [[ -n "${TRAIN_CONFIG}" ]]; then
+    [[ -f "${TRAIN_CONFIG}" ]] || die "Training config not found: ${TRAIN_CONFIG}"
+fi
+[[ -r /data ]] || die "Training data directory is not readable: /data"
+
+mkdir -p "${CHECKPOINT_DIR}"
+[[ -w "${CHECKPOINT_DIR}" ]] || die "Checkpoint directory is not writable: ${CHECKPOINT_DIR}"
+
+latest_checkpoint="$(find_latest_checkpoint "${CHECKPOINT_DIR}" || true)"
+
+train_args=(
+    python scripts/tools/train.py
+    --ckpt_dir "${CHECKPOINT_DIR}"
+    --nprocs "${TRAIN_GPU_COUNT}"
+)
+
+if [[ -n "${TRAIN_CONFIG}" ]]; then
+    train_args+=(--config "${TRAIN_CONFIG}")
+fi
+
+if (( TRAIN_GPU_COUNT > 1 )); then
+    train_args+=(--parallel_mode ddp)
+else
+    train_args+=(--parallel_mode none)
+fi
+
+if [[ -n "${latest_checkpoint}" ]]; then
+    log_info "Resuming ${TRAIN_JOB_NAME} from ${latest_checkpoint}"
+    train_args+=(--param_path "${latest_checkpoint}" --resume)
+else
+    [[ -s "${BASE_MODEL}/config.json" ]] || die "Base model config not found: ${BASE_MODEL}/config.json"
+    [[ -s "${BASE_MODEL}/model.safetensors" ]] || die "Base model weights not found: ${BASE_MODEL}/model.safetensors"
+    log_info "Starting ${TRAIN_JOB_NAME} from ${BASE_MODEL}"
+    train_args+=(--param_path "${BASE_MODEL}")
+fi
+
+log_info "GPUs=${TRAIN_GPU_COUNT}, checkpoints=${CHECKPOINT_DIR}"
+
+# Replace the shell so the container init forwards SIGTERM to the trainer.
+exec "${train_args[@]}" "$@"

scripts/tools/benchmark.py

@@ -1,17 +1,19 @@
 import json
+import time
 from pathlib import Path
 from typing import Optional, Union
 
 import click
 import torch
 
-from astrai import setup_logging
 from astrai.config import BaseModelConfig, ConfigFactory
 from astrai.extension import ATTN_BACKEND, AttentionBackendFactory, attn_backend
-from astrai.inference.core.cache import PagePool
-from astrai.inference.core.graph import CudaGraphContext
-from astrai.inference.core.workspace import InferenceWorkspace
+from astrai.inference.cache import PagePool, TaskCacheManager
+from astrai.inference.engine import InferenceEngine
+from astrai.inference.runtime.graph import CudaGraphContext
+from astrai.inference.workspace import InferenceWorkspace
 from astrai.model import AutoModel, AutoRegressiveLM
+from astrai.tokenize import AutoTokenizer
 
 _DTYPES = ["bfloat16", "float16", "float32"]
 _CACHES = ["contiguous", "paged"]
@@ -59,6 +61,7 @@ class GenerationBenchmark:
         cache_type: str = "contiguous",
         backend: Union[str, ATTN_BACKEND] = ATTN_BACKEND.CUDA,
         cuda_graph: bool = False,
+        tokenizer: Optional[AutoTokenizer] = None,
     ):
         self.device = device
         self.dtype = dtype
@@ -67,8 +70,18 @@ class GenerationBenchmark:
         self.config = config
         self.backend = backend
         self.cuda_graph = cuda_graph
+        self.tokenizer = tokenizer
 
     def _make_pool(self, batch_size: int, max_seq_len: int) -> PagePool:
+        if self.cache_type == "contiguous":
+            n_tokens = None
+        elif self.cache_type == "paged":
+            # Keep the total token capacity equal to contiguous mode while
+            # routing allocation through the shared paged pool.
+            n_tokens = batch_size * max_seq_len
+        else:
+            raise ValueError(f"unsupported cache type: {self.cache_type}")
+
         return PagePool(
             n_layers=self.config.num_hidden_layers,
             n_kv_heads=self.config.num_key_value_heads,
@@ -78,7 +91,7 @@ class GenerationBenchmark:
             device=self.device,
             dtype=self.dtype,
             page_size=1,
-            n_tokens=None,
+            n_tokens=n_tokens,
         )
 
     @staticmethod
@@ -92,9 +105,14 @@ class GenerationBenchmark:
             dtype=pool.dtype,
         )
 
+    @staticmethod
+    def _make_task_cache(pool: PagePool) -> TaskCacheManager:
+        return TaskCacheManager(pool)
+
     def _run_prefill(
         self,
         pool: PagePool,
+        task_cache: TaskCacheManager,
         batch_size: int,
         prompt_len: int,
         workspace: InferenceWorkspace,
@@ -113,9 +131,9 @@ class GenerationBenchmark:
 
         task_ids = [f"bench_{i}" for i in range(batch_size)]
         for tid in task_ids:
-            pool.task_alloc(tid, list(range(prompt_len)))
+            task_cache.task_alloc(tid, list(range(prompt_len)))
 
-        kv_cache = pool.bind_tasks(task_ids, workspace, self.device, start_pos=0)
+        kv_cache = task_cache.bind(task_ids, workspace, self.device, start_pos=0)
         with torch.inference_mode(), attn_backend(self.backend):
             self.model(
                 input_ids,
@@ -129,6 +147,7 @@ class GenerationBenchmark:
     def _run_decode_step(
         self,
         pool: PagePool,
+        task_cache: TaskCacheManager,
         task_ids: list,
         seq_len: int,
         workspace: InferenceWorkspace,
@@ -142,11 +161,11 @@ class GenerationBenchmark:
         )
         total_len = seq_len + 1
         for tid in task_ids:
-            pool.task_extend(tid, seq_len)
+            task_cache.task_extend(tid, seq_len)
         input_mask = position_ids[:, :, None] >= torch.arange(
             total_len, device=self.device
         )
-        kv_cache = pool.bind_tasks(task_ids, workspace, self.device)
+        kv_cache = task_cache.bind(task_ids, workspace, self.device)
         with torch.inference_mode(), attn_backend(self.backend):
             self.model(
                 input_ids,
@@ -161,13 +180,12 @@ class GenerationBenchmark:
         prompt_length: int = 512,
         num_trials: int = 5,
     ) -> BenchmarkResult:
-        import time
-
         pool = self._make_pool(batch_size, prompt_length)
         workspace = self._make_workspace(pool, self.config)
+        task_cache = self._make_task_cache(pool)
         task_ids = [f"bench_prefill_{i}" for i in range(batch_size)]
         for tid in task_ids:
-            pool.task_alloc(tid, list(range(prompt_length)))
+            task_cache.task_alloc(tid, list(range(prompt_length)))
 
         input_ids = torch.randint(
             0, self.config.vocab_size, (batch_size, prompt_length), device=self.device
@@ -180,7 +198,7 @@ class GenerationBenchmark:
         input_mask = position_ids.unsqueeze(-1) >= torch.arange(
             prompt_length, device=self.device
         )
-        kv_cache = pool.bind_tasks(task_ids, workspace, self.device, start_pos=0)
+        kv_cache = task_cache.bind(task_ids, workspace, self.device, start_pos=0)
 
         for _ in range(3):
             with torch.inference_mode(), attn_backend(self.backend):
@@ -221,12 +239,60 @@ class GenerationBenchmark:
         gen_length: int = 128,
         num_trials: int = 5,
     ) -> BenchmarkResult:
-        if self.cuda_graph and self.backend == "cuda":
-            return self._run_graph_decode_benchmark(
-                batch_size, prompt_length, gen_length, num_trials
-            )
-        return self._run_plain_decode_benchmark(
-            batch_size, prompt_length, gen_length, num_trials
+        if self.tokenizer is None:
+            raise ValueError("Engine decode benchmark requires a tokenizer")
+
+        # Use the real engine so scheduler, executor, sampling, and graph
+        # warmup/replay are included in the measured generation path.
+        phrase = "Benchmark the language model with a realistic generation prompt. "
+        prompt_ids = self.tokenizer.encode(
+            (phrase * (prompt_length // 10 + 2)).strip()
+        )[:prompt_length]
+        prompt = self.tokenizer.decode(prompt_ids, skip_special_tokens=False)
+        prompt_tokens = len(self.tokenizer.encode(prompt))
+        max_seq_len = prompt_tokens + gen_length
+        pool = self._make_pool(batch_size, max_seq_len)
+        engine = InferenceEngine(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            max_batch_size=batch_size,
+            max_seq_len=max_seq_len,
+            cache=pool,
+            enable_cuda_graph=self.cuda_graph,
+            backend=self.backend,
+        )
+        prompts = [prompt] * batch_size
+
+        try:
+            # Capture graphs and populate the allocator before timing. The
+            # first request also includes model/scheduler startup effects.
+            engine.generate(prompts, max_tokens=gen_length, temperature=0.0)
+            if self.device.startswith("cuda"):
+                torch.cuda.synchronize()
+
+            t0 = time.perf_counter()
+            for _ in range(num_trials):
+                engine.generate(prompts, max_tokens=gen_length, temperature=0.0)
+            if self.device.startswith("cuda"):
+                torch.cuda.synchronize()
+            elapsed = time.perf_counter() - t0
+        finally:
+            engine.shutdown()
+
+        tokens = batch_size * gen_length * num_trials
+        return BenchmarkResult(
+            name="decode",
+            batch_size=batch_size,
+            seq_len=gen_length,
+            tokens_per_second=tokens / elapsed,
+            latency_ms=elapsed / (gen_length * num_trials) * 1000,
+            metadata={
+                "benchmark_type": "engine_decode",
+                "num_trials": num_trials,
+                "prompt_length": prompt_tokens,
+                "backend": engine.backend_name,
+                "cuda_graph": engine.cuda_graph_enabled,
+            },
         )
 
     def _run_graph_decode_benchmark(
@@ -236,12 +302,13 @@ class GenerationBenchmark:
         gen_length: int,
         num_trials: int,
     ) -> BenchmarkResult:
-        import time
-
         max_seq_len = prompt_length + 5 + gen_length * num_trials
         pool = self._make_pool(batch_size, max_seq_len)
         workspace = self._make_workspace(pool, self.config)
-        task_ids = self._run_prefill(pool, batch_size, prompt_length, workspace)
+        task_cache = self._make_task_cache(pool)
+        task_ids = self._run_prefill(
+            pool, task_cache, batch_size, prompt_length, workspace
+        )
 
         b = batch_size
         input_ids_buf = torch.zeros(b, 1, dtype=torch.long, device=self.device)
@@ -257,8 +324,8 @@ class GenerationBenchmark:
             )
             position_ids_buf[:] = seq_len
             for tid in task_ids:
-                pool.task_extend(tid, seq_len)
-            kv_cache = pool.bind_tasks(task_ids, workspace, self.device)
+                task_cache.task_extend(tid, seq_len)
+            kv_cache = task_cache.bind(task_ids, workspace, self.device)
 
             input_mask = torch.ge(
                 position_ids_buf[:, None],
@@ -309,20 +376,25 @@ class GenerationBenchmark:
         gen_length: int,
         num_trials: int,
     ) -> BenchmarkResult:
-        import time
-
         max_seq_len = prompt_length + 5 + gen_length * num_trials
         pool = self._make_pool(batch_size, max_seq_len)
         workspace = self._make_workspace(pool, self.config)
-        task_ids = self._run_prefill(pool, batch_size, prompt_length, workspace)
+        task_cache = self._make_task_cache(pool)
+        task_ids = self._run_prefill(
+            pool, task_cache, batch_size, prompt_length, workspace
+        )
 
         for i in range(5):
-            self._run_decode_step(pool, task_ids, prompt_length + i, workspace)
+            self._run_decode_step(
+                pool, task_cache, task_ids, prompt_length + i, workspace
+            )
         torch.cuda.synchronize()
 
         t0 = time.perf_counter()
         for i in range(gen_length * num_trials):
-            self._run_decode_step(pool, task_ids, prompt_length + 5 + i, workspace)
+            self._run_decode_step(
+                pool, task_cache, task_ids, prompt_length + 5 + i, workspace
+            )
         torch.cuda.synchronize()
         elapsed = time.perf_counter() - t0
         tokens = batch_size * gen_length * num_trials
@@ -378,9 +450,9 @@ def print_benchmark_result(result: BenchmarkResult) -> None:
 @click.option("--prefill_only", is_flag=True, help="Prefill benchmark only.")
 @click.option("--decode_only", is_flag=True, help="Decode benchmark only.")
 @click.option(
-    "--cuda-graph",
-    is_flag=True,
-    help="Enable CUDA graph capture for decode (cuda backend only).",
+    "--cuda-graph/--no-cuda-graph",
+    default=True,
+    help="Enable or disable CUDA graph capture for engine decode.",
 )
 @click.option(
     "--ckpt",
@@ -439,6 +511,8 @@ def benchmark_command(
             f"({sum(p.numel() for p in model.parameters()) / 1e9:.2f}B params)"
         )
 
+    tokenizer = AutoTokenizer.from_pretrained(ckpt or Path("params"))
+
     model.to(device=device, dtype=dtype_map[dtype])
     model.eval()
 
@@ -453,6 +527,7 @@ def benchmark_command(
             cache_type=cache,
             backend=name,
             cuda_graph=cuda_graph,
+            tokenizer=tokenizer,
         )
 
         click.secho(
@@ -478,5 +553,4 @@ def benchmark_command(
 
 
 if __name__ == "__main__":
-    setup_logging()
     benchmark_command()

scripts/tools/generate.py

@@ -6,7 +6,6 @@ import click
 import torch
 from tqdm import tqdm
 
-from astrai import setup_logging
 from astrai.inference import InferenceEngine
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
@@ -157,5 +156,4 @@ def generate_command(**kwargs):
 
 
 if __name__ == "__main__":
-    setup_logging()
     generate_command()

scripts/tools/preprocess.py

@@ -2,7 +2,6 @@
 
 import click
 
-from astrai import setup_logging
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.preprocessing.pipeline import Pipeline
 
@@ -48,5 +47,4 @@ def preprocess_command(inputs, output_dir, pipeline_config, tokenizer_path, batc
 
 
 if __name__ == "__main__":
-    setup_logging()
     preprocess_command()

scripts/tools/server.py

@@ -3,7 +3,6 @@ from pathlib import Path
 import click
 import torch
 
-from astrai import setup_logging
 from astrai.inference import run_server
 
 _DTYPES = ["bfloat16", "float16", "float32"]
@@ -65,5 +64,4 @@ def server_command(
 
 
 if __name__ == "__main__":
-    setup_logging()
     server_command()

scripts/tools/train.py

@@ -1,14 +1,15 @@
 import os
+import re
 from collections import OrderedDict
 from collections.abc import Callable
 from functools import partial
 
 import click
 import torch
+import yaml
 from click.core import ParameterSource
 from torch import optim
 
-from astrai import setup_logging
 from astrai.config import AutoRegressiveLMConfig, TrainConfig
 from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn
 from astrai.model import AutoRegressiveLM
@@ -49,13 +50,30 @@ def opt(*param_decls, group: str, **kwargs):
     return click.option(*param_decls, **kwargs)
 
 
+_YAML_FLOAT_PATTERN = re.compile(
+    r"""^(?:[-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+]?[0-9]+)?
+    |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
+    |[-+]?\.(?:inf|Inf|INF)
+    |\.(?:nan|NaN|NAN))$""",
+    re.X,
+)
+
+
+def _enable_yaml12_floats() -> None:
+    """PyYAML implements YAML 1.1, where ``2e-5`` parses as a string; switch its
+    float resolver to the YAML 1.2 core schema so scientific notation works."""
+    yaml.SafeLoader.add_implicit_resolver(
+        "tag:yaml.org,2002:float", _YAML_FLOAT_PATTERN, list("-+0123456789.")
+    )
+
+
 def _merge_yaml_into_kwargs(
     config_path: str,
     passed_kwargs: dict,
     explicit_keys: set[str] | None = None,
 ) -> dict:
     """Merge Click defaults, YAML values, then explicit CLI values."""
-    import yaml
+    _enable_yaml12_floats()
 
     with open(config_path) as f:
         cfg = yaml.safe_load(f) or {}
@@ -252,6 +270,12 @@ _START_METHODS = ["spawn", "fork", "forkserver"]
     group="Data Loading",
     help="Pin memory.",
 )
+@opt(
+    "--persistent_workers/--no-persistent_workers",
+    default=True,
+    group="Data Loading",
+    help="Keep DataLoader workers alive between epochs.",
+)
 @opt(
     "--window_size",
     type=int,
@@ -607,6 +631,7 @@ def train(
     random_seed: int,
     num_workers: int,
     pin_memory: bool,
+    persistent_workers: bool,
     gradient_checkpointing: bool,
     window_size: int,
     stride: int,
@@ -798,6 +823,7 @@ def train(
         random_seed=random_seed,
         num_workers=num_workers,
         pin_memory=pin_memory,
+        persistent_workers=persistent_workers,
         nprocs=nprocs,
         backend=backend,
         master_addr=master_addr,
@@ -828,5 +854,4 @@ def train(
 
 
 if __name__ == "__main__":
-    setup_logging()
     train_command()

scripts/train.sh

@@ -0,0 +1,327 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+source "${ROOT_DIR}/scripts/docker/lib/train-common.sh"
+
+ENV_FILE="${TRAIN_ENV_FILE:-${ROOT_DIR}/.env.train}"
+COMPOSE_BASE=(
+    docker compose
+    --project-directory "${ROOT_DIR}"
+    --file "${ROOT_DIR}/docker-compose.yml"
+    --profile train
+)
+
+usage() {
+    cat <<'EOF'
+Usage: scripts/train.sh <command> [options]
+
+Commands:
+  init                 Create local directories and .env.train
+  preflight            Validate Docker, paths, GPU settings, and Compose
+  build                Build the trainer image
+  start [--foreground] [-- ARGS...] Start or resume training
+  stop                 Gracefully stop and checkpoint training
+  restart              Stop, then start training
+  logs                 Follow trainer logs
+  status               Show container and latest checkpoint status
+  latest               Print the latest complete checkpoint path
+  list                 List all complete checkpoints
+  clean [--keep N]     Preview old checkpoint removal
+  clean --force        Remove old checkpoints after previewing
+
+Environment:
+  TRAIN_ENV_FILE       Env file path (default: .env.train)
+  TRAIN_CONFIG_FILE    Optional host YAML mounted only when the job starts
+
+Training arguments come from an externally mounted TRAIN_CONFIG or ARGS passed
+after --. The image does not contain experiment configuration.
+EOF
+}
+
+load_env() {
+    if [[ -f "${ENV_FILE}" ]]; then
+        set -a
+        # UID/GID are readonly in bash; compose gets them via ASTRAI_UID/GID in compose()
+        # shellcheck disable=SC1090
+        source <(grep -v -E '^[[:space:]]*(UID|GID)=' "${ENV_FILE}")
+        set +a
+    fi
+
+    TRAIN_JOB_NAME="${TRAIN_JOB_NAME:-astrai-train}"
+    TRAIN_DATA_DIR="${TRAIN_DATA_DIR:-./data}"
+    TRAIN_MODEL_DIR="${TRAIN_MODEL_DIR:-./params}"
+    TRAIN_CHECKPOINT_DIR="${TRAIN_CHECKPOINT_DIR:-./checkpoints}"
+    TRAIN_GPU_COUNT="${TRAIN_GPU_COUNT:-all}"
+    TRAIN_STOP_TIMEOUT="${TRAIN_STOP_TIMEOUT:-600}"
+
+    validate_job_name "${TRAIN_JOB_NAME}"
+}
+
+resolve_path() {
+    if [[ "$1" = /* ]]; then
+        printf '%s\n' "$1"
+    else
+        printf '%s/%s\n' "${ROOT_DIR}" "${1#./}"
+    fi
+}
+
+checkpoint_dir() {
+    printf '%s/%s\n' "$(resolve_path "${TRAIN_CHECKPOINT_DIR}")" "${TRAIN_JOB_NAME}"
+}
+
+compose() {
+    local -a command=("${COMPOSE_BASE[@]}")
+
+    if [[ -f "${ENV_FILE}" ]]; then
+        command+=(--env-file "${ENV_FILE}")
+    fi
+
+    # Inject the host user into compose so container processes share the
+    # checkpoint directory ownership (bash UID/GID are readonly).
+    ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${command[@]}" "$@"
+}
+
+init_environment() {
+    local data_dir model_dir checkpoints_dir
+
+    data_dir="$(resolve_path "${TRAIN_DATA_DIR}")"
+    model_dir="$(resolve_path "${TRAIN_MODEL_DIR}")"
+    checkpoints_dir="$(resolve_path "${TRAIN_CHECKPOINT_DIR}")"
+    mkdir -p "${data_dir}" "${model_dir}" "${checkpoints_dir}"
+
+    if [[ ! -f "${ENV_FILE}" ]]; then
+        cat >"${ENV_FILE}" <<'EOF'
+TRAIN_JOB_NAME=astrai-train
+TRAIN_DATA_DIR=./data
+TRAIN_MODEL_DIR=./params
+TRAIN_CHECKPOINT_DIR=./checkpoints
+TRAIN_CONFIG_FILE=
+TRAIN_GPU_COUNT=all
+# CUDA_VISIBLE_DEVICES=0,1
+CUDA_TAG=cu128
+TRAIN_IPC_MODE=host
+TRAIN_STOP_GRACE_PERIOD=10m
+TRAIN_STOP_TIMEOUT=600
+CHECKPOINT_KEEP_LAST=5
+EOF
+        log_info "Created ${ENV_FILE}"
+    else
+        log_info "Keeping existing ${ENV_FILE}"
+    fi
+    log_info "Data: ${data_dir}"
+    log_info "Model: ${model_dir}"
+    log_info "Checkpoints: ${checkpoints_dir}"
+}
+
+preflight() {
+    local data_dir model_dir checkpoints_dir config_file latest visible_count
+
+    require_command docker
+    docker info >/dev/null 2>&1 || die "Docker daemon is unavailable"
+    [[ "${TRAIN_GPU_COUNT}" == "all" || "${TRAIN_GPU_COUNT}" =~ ^[1-9][0-9]*$ ]] ||
+        die "TRAIN_GPU_COUNT must be 'all' or a positive integer"
+
+    data_dir="$(resolve_path "${TRAIN_DATA_DIR}")"
+    model_dir="$(resolve_path "${TRAIN_MODEL_DIR}")"
+    checkpoints_dir="$(resolve_path "${TRAIN_CHECKPOINT_DIR}")"
+    [[ -d "${data_dir}" ]] || die "Training data directory not found: ${data_dir}"
+    mkdir -p "${checkpoints_dir}/${TRAIN_JOB_NAME}"
+    [[ -w "${checkpoints_dir}/${TRAIN_JOB_NAME}" ]] || die "Checkpoint directory is not writable"
+
+    if [[ -n "${TRAIN_CONFIG_FILE:-}" ]]; then
+        config_file="$(resolve_path "${TRAIN_CONFIG_FILE}")"
+        [[ -f "${config_file}" ]] || die "Training config not found: ${config_file}"
+    fi
+
+    latest="$(find_latest_checkpoint "${checkpoints_dir}/${TRAIN_JOB_NAME}" || true)"
+    if [[ -z "${latest}" ]]; then
+        [[ -s "${model_dir}/config.json" ]] || die "Model config not found: ${model_dir}/config.json"
+        [[ -s "${model_dir}/model.safetensors" ]] || die "Model weights not found: ${model_dir}/model.safetensors"
+    else
+        log_info "Resume candidate: ${latest}"
+    fi
+
+    if [[ -n "${CUDA_VISIBLE_DEVICES:-}" && "${TRAIN_GPU_COUNT}" != "all" ]]; then
+        IFS=',' read -r -a visible_gpus <<<"${CUDA_VISIBLE_DEVICES}"
+        visible_count="${#visible_gpus[@]}"
+        (( visible_count == TRAIN_GPU_COUNT )) ||
+            die "TRAIN_GPU_COUNT=${TRAIN_GPU_COUNT}, but CUDA_VISIBLE_DEVICES exposes ${visible_count} GPU(s)"
+    fi
+
+    compose config --quiet
+    log_info "Preflight passed for ${TRAIN_JOB_NAME} (GPU request: ${TRAIN_GPU_COUNT})"
+}
+
+start_training() {
+    local foreground="$1"
+    local config_file container running
+    local -a run_options=()
+    shift
+
+    preflight
+    if [[ -n "${TRAIN_CONFIG_FILE:-}" ]]; then
+        config_file="$(resolve_path "${TRAIN_CONFIG_FILE}")"
+        run_options+=(
+            --volume "${config_file}:/run/astrai/train.yaml:ro"
+            --env TRAIN_CONFIG=/run/astrai/train.yaml
+        )
+    elif [[ -z "${TRAIN_CONFIG:-}" && $# -eq 0 ]]; then
+        die "Set TRAIN_CONFIG_FILE or pass complete trainer arguments after --"
+    fi
+
+    container="astrai-trainer-${TRAIN_JOB_NAME}"
+    running="$(docker inspect --format '{{.State.Running}}' "${container}" 2>/dev/null || true)"
+    [[ "${running}" != "true" ]] || die "Trainer is already running: ${container}"
+    docker rm "${container}" >/dev/null 2>&1 || true
+    if [[ "${foreground}" == "true" ]]; then
+        compose run --build --rm "${run_options[@]}" trainer "$@"
+    else
+        compose run -d --build --name "${container}" \
+            "${run_options[@]}" trainer "$@"
+        log_info "Training started; run scripts/train.sh logs to follow it"
+    fi
+}
+
+stop_training() {
+    log_info "Stopping trainer with ${TRAIN_STOP_TIMEOUT}s grace period"
+    docker stop --timeout "${TRAIN_STOP_TIMEOUT}" "astrai-trainer-${TRAIN_JOB_NAME}" >/dev/null 2>&1 ||
+        log_warn "Trainer container is not running"
+}
+
+restart_training() {
+    local container="astrai-trainer-${TRAIN_JOB_NAME}"
+
+    docker inspect "${container}" >/dev/null 2>&1 ||
+        die "Trainer container not found; use start with a config or CLI arguments first"
+    log_info "Restarting trainer with ${TRAIN_STOP_TIMEOUT}s grace period"
+    docker restart --timeout "${TRAIN_STOP_TIMEOUT}" "${container}" >/dev/null
+}
+
+show_status() {
+    local latest
+
+    docker ps -a --filter "name=^/astrai-trainer-${TRAIN_JOB_NAME}$"
+    latest="$(find_latest_checkpoint "$(checkpoint_dir)" || true)"
+    if [[ -n "${latest}" ]]; then
+        log_info "Latest checkpoint: ${latest}"
+    else
+        log_info "No complete checkpoint found for ${TRAIN_JOB_NAME}"
+    fi
+}
+
+clean_checkpoints() {
+    local keep="$1" force="$2" dir count remove_count index path
+    local -a checkpoints=()
+
+    [[ "${keep}" =~ ^[1-9][0-9]*$ ]] || die "--keep must be a positive integer"
+    dir="$(checkpoint_dir)"
+    while IFS= read -r line; do
+        [[ -n "${line}" ]] && checkpoints+=("${line#* * }")
+    done < <(list_complete_checkpoints "${dir}")
+
+    count="${#checkpoints[@]}"
+    remove_count=$((count - keep))
+    if (( remove_count <= 0 )); then
+        log_info "Nothing to clean; ${count} complete checkpoint(s), keeping ${keep}"
+        return
+    fi
+
+    for ((index = 0; index < remove_count; index++)); do
+        path="${checkpoints[index]}"
+        if [[ "${force}" == "true" ]]; then
+            rm -rf -- "${path}"
+            log_info "Removed ${path}"
+        else
+            printf 'Would remove %s\n' "${path}"
+        fi
+    done
+    [[ "${force}" == "true" ]] || log_warn "Preview only; add --force to delete"
+}
+
+main() {
+    local command="${1:-}" foreground=false keep="${CHECKPOINT_KEEP_LAST:-5}" force=false
+    local -a train_args=()
+    [[ -n "${command}" ]] || { usage; exit 1; }
+    shift || true
+    load_env
+
+    case "${command}" in
+        init)
+            init_environment
+            ;;
+        preflight)
+            preflight
+            ;;
+        build)
+            preflight
+            compose build trainer
+            ;;
+        start)
+            while [[ $# -gt 0 ]]; do
+                case "$1" in
+                    --foreground)
+                        foreground=true
+                        shift
+                        ;;
+                    --)
+                        shift
+                        train_args=("$@")
+                        break
+                        ;;
+                    *)
+                        die "Unknown start option: $1 (put trainer arguments after --)"
+                        ;;
+                esac
+            done
+            start_training "${foreground}" "${train_args[@]}"
+            ;;
+        stop)
+            stop_training
+            ;;
+        restart)
+            restart_training
+            ;;
+        logs)
+            docker logs -f --tail "${TRAIN_LOG_TAIL:-200}" "astrai-trainer-${TRAIN_JOB_NAME}"
+            ;;
+        status)
+            show_status
+            ;;
+        latest)
+            find_latest_checkpoint "$(checkpoint_dir)" || die "No complete checkpoint found"
+            ;;
+        list)
+            list_complete_checkpoints "$(checkpoint_dir)" | while read -r _epoch _step path; do
+                printf '%s\n' "${path}"
+            done
+            ;;
+        clean)
+            while [[ $# -gt 0 ]]; do
+                case "$1" in
+                    --keep)
+                        [[ $# -ge 2 ]] || die "--keep requires a value"
+                        keep="$2"
+                        shift 2
+                        ;;
+                    --force)
+                        force=true
+                        shift
+                        ;;
+                    *)
+                        die "Unknown clean option: $1"
+                        ;;
+                esac
+            done
+            clean_checkpoints "${keep}" "${force}"
+            ;;
+        help|-h|--help)
+            usage
+            ;;
+        *)
+            die "Unknown command: ${command}"
+            ;;
+    esac
+}
+
+main "$@"

tests/data/test_dataset.py

@@ -9,6 +9,7 @@ import torch
 from astrai.dataset.dataset import (
     DatasetFactory,
     GRPODataset,
+    _build_jsonl_transform,
     dpo_tokenize,
     grpo_collate_fn,
 )
@@ -525,7 +526,7 @@ def test_json_store_seq(base_test_env):
     )
 
     store = StoreFactory.create("jsonl")
-    store.load(data_dir)
+    store.load(data_dir, transform=_build_jsonl_transform(data_dir))
     assert len(store) > 0
     assert "sequence" in store.keys
 
@@ -580,7 +581,7 @@ def test_json_store_no_tokenizer_path(base_test_env):
         json.dump(config, f, ensure_ascii=False, indent=2)
 
     store = StoreFactory.create("jsonl")
-    store.load(data_dir)
+    store.load(data_dir, transform=_build_jsonl_transform(data_dir))
     assert len(store) > 0
     assert "sequence" in store.keys
     assert "loss_mask" in store.keys
@@ -597,7 +598,7 @@ def test_jsonl_store_seq(base_test_env):
     )
 
     store = StoreFactory.create("jsonl")
-    store.load(data_dir)
+    store.load(data_dir, transform=_build_jsonl_transform(data_dir))
     assert len(store) > 0
     assert "sequence" in store.keys
 
@@ -638,7 +639,7 @@ def test_jsonl_store_sft(base_test_env):
     )
 
     store = StoreFactory.create("jsonl")
-    store.load(data_dir)
+    store.load(data_dir, transform=_build_jsonl_transform(data_dir))
     assert "sequence" in store.keys
     assert "loss_mask" in store.keys
     assert "position_ids" in store.keys
@@ -1085,7 +1086,7 @@ def test_jsonl_store_eager_len_returns_token_count(base_test_env):
     )
 
     store = JsonlStore()
-    store.load(data_dir)
+    store.load(data_dir, transform=_build_jsonl_transform(data_dir))
 
     assert store.num_records == 2
     assert len(store.keys) > 0

tests/extension/test_backend.py

@@ -15,8 +15,8 @@ from astrai.extension import (
 )
 
 
-def test_default_backend_is_torch_native():
-    """Default is the highest-priority available backend (flash > cuda > torch)."""
+def test_default_backend_resolves_to_available():
+    """Default backend is the first available in cuda > flash > torch order."""
     from astrai.extension.attention_backend import (
         CudaBackend,
         FlashAttnBackend,
@@ -26,7 +26,6 @@ def test_default_backend_is_torch_native():
 
     backend = get_backend()
     assert isinstance(backend, (CudaBackend, FlashAttnBackend, TorchNativeBackend))
-    assert isinstance(backend, type(_resolve_default_backend()))
 
 
 def test_attn_backend_context_with_enum():
@@ -43,6 +42,20 @@ def test_attn_backend_context_with_registered_name():
     assert get_backend() is default
 
 
+def test_backend_can_read_only_context_selection():
+    assert get_backend(use_default=False) is None
+    with attn_backend("cuda") as backend:
+        assert get_backend(use_default=False) is backend
+    assert get_backend(use_default=False) is None
+
+
+def test_environment_backend_overrides_context(monkeypatch):
+    monkeypatch.setenv("ASTR_BACKEND", "torch_native")
+    with attn_backend("cuda"):
+        assert type(get_backend()).__name__ == "TorchNativeBackend"
+        assert type(get_backend(use_default=False)).__name__ == "TorchNativeBackend"
+
+
 def test_attention_backend_factory_lists_builtin_backends():
     assert AttentionBackendFactory.list_registered() == [
         "cuda",

tests/extension/test_backend_equivalence.py

@@ -7,11 +7,15 @@ seq_lens with padding mask), and end-to-end scheduler.run_batch.
 import torch
 
 from astrai.extension import ATTN_BACKEND, attn_backend
-from astrai.inference.core.cache import PagePool
-from astrai.inference.core.workspace import InferenceWorkspace
+from astrai.inference.cache import PagePool, TaskCacheManager
+from astrai.inference.workspace import InferenceWorkspace
 from tests.extension.conftest import D, skip_no_kernel
 
 
+def _mk_task_cache(pool: PagePool) -> TaskCacheManager:
+    return TaskCacheManager(pool)
+
+
 def _ws(pool: PagePool) -> InferenceWorkspace:
     return InferenceWorkspace(
         pool.max_batch_size,
@@ -25,27 +29,26 @@ def _ws(pool: PagePool) -> InferenceWorkspace:
 
 @skip_no_kernel
 def test_training_forward_matches_torch(cuda_model):
-    """Training forward (kv_cache=None) should produce identical logits.
+    """Training forward (kv_cache=None) uses torch-native SDPA.
 
-    CudaBackend is now safe as a default: for training (``kv_cache=None``)
-    or non-bf16 inputs it falls back to torch SDPA.  Verify the fallback
-    path matches the torch-native forward exactly.
+    CudaBackend does not support training (requires kv_cache).
+    Torch-native backend must match default (which falls back to torch).
     """
 
     model, _ = cuda_model
     input_ids = torch.randint(0, 1000, (2, 16), device="cuda")
 
     with torch.no_grad():
-        out_torch = model(input_ids)
+        out_default = model(input_ids)
 
-    with attn_backend(ATTN_BACKEND.CUDA):
+    with attn_backend(ATTN_BACKEND.TORCH_NATIVE):
         with torch.no_grad():
-            out_cuda = model(input_ids)
+            out_torch = model(input_ids)
 
     torch.testing.assert_close(
-        out_cuda["logits"], out_torch["logits"], atol=1e-6, rtol=1e-6
+        out_torch["logits"], out_default["logits"], atol=1e-6, rtol=1e-6
     )
-    assert out_torch["logits"].shape[0] == 2
+    assert out_default["logits"].shape[0] == 2
 
 
 @skip_no_kernel
@@ -75,20 +78,21 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
         dtype=torch.bfloat16,
     )
 
+    task_cache = _mk_task_cache(cache)
     ws = _ws(cache)
-    cache.task_alloc("t1", prompt_ids[0])
-    cache.task_alloc("t2", prompt_ids[1])
-    kv1 = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
+    task_cache.task_alloc("t1", prompt_ids[0])
+    task_cache.task_alloc("t2", prompt_ids[1])
+    kv1 = task_cache.bind(["t1", "t2"], ws, start_pos=0)
     with torch.inference_mode():
         out_torch = model(
             input_ids, input_mask=input_mask, kv_cache=kv1, position_ids=position_ids
         )
 
-    cache.task_free("t1")
-    cache.task_free("t2")
-    cache.task_alloc("t1", prompt_ids[0])
-    cache.task_alloc("t2", prompt_ids[1])
-    kv2 = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
+    task_cache.task_free("t1")
+    task_cache.task_free("t2")
+    task_cache.task_alloc("t1", prompt_ids[0])
+    task_cache.task_alloc("t2", prompt_ids[1])
+    kv2 = task_cache.bind(["t1", "t2"], ws, start_pos=0)
     with attn_backend(ATTN_BACKEND.CUDA):
         with torch.inference_mode():
             out_cuda = model(
@@ -139,10 +143,11 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
         input_mask[i, : len(p)] = True
         position_ids[i, : len(p)] = torch.arange(len(p), device=device)
 
+    task_cache = _mk_task_cache(cache)
     ws = _ws(cache)
-    cache.task_alloc("t1", prompt_ids[0])
-    cache.task_alloc("t2", prompt_ids[1])
-    kv = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
+    task_cache.task_alloc("t1", prompt_ids[0])
+    task_cache.task_alloc("t2", prompt_ids[1])
+    kv = task_cache.bind(["t1", "t2"], ws, start_pos=0)
     with torch.inference_mode():
         model(input_ids, input_mask=input_mask, kv_cache=kv, position_ids=position_ids)
 
@@ -152,15 +157,15 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     total_len = 9
     dec_mask = dec_pos[:, None, None] >= torch.arange(total_len, device=device)
 
-    cache.task_extend("t1", 8)
-    cache.task_extend("t2", 6)
-    kv_t = cache.bind_tasks(["t1", "t2"], ws)
+    task_cache.task_extend("t1", 8)
+    task_cache.task_extend("t2", 6)
+    kv_t = task_cache.bind(["t1", "t2"], ws)
     with torch.inference_mode():
         out_torch = model(
             dec_ids, input_mask=dec_mask, kv_cache=kv_t, position_ids=dec_pos
         )
 
-    kv_c = cache.bind_tasks(["t1", "t2"], ws)
+    kv_c = task_cache.bind(["t1", "t2"], ws)
     with attn_backend(ATTN_BACKEND.CUDA):
         with torch.inference_mode():
             out_cuda = model(
@@ -174,7 +179,7 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
 @skip_no_kernel
 def test_run_batch_cuda_matches_torch_greedy(cuda_model):
     """Greedy decode (temperature=0) should produce identical tokens."""
-    from astrai.inference.core.scheduler import InferenceScheduler
+    from astrai.inference.scheduler import InferenceScheduler
     from tests.helpers import FakeTokenizer
 
     model, _ = cuda_model

tests/inference/test_cache.py

@@ -8,9 +8,10 @@ from astrai.inference import (
     PagePool,
     RadixCache,
     ReqToTokenPool,
+    TaskCacheManager,
     page_hash,
 )
-from astrai.inference.core.workspace import InferenceWorkspace
+from astrai.inference.workspace import InferenceWorkspace
 
 
 def _ws(pool: PagePool) -> InferenceWorkspace:
@@ -25,6 +26,10 @@ def _ws(pool: PagePool) -> InferenceWorkspace:
     )
 
 
+def _make_task_cache(pool: PagePool) -> TaskCacheManager:
+    return TaskCacheManager(pool)
+
+
 # ---- page_hash ----
 
 
@@ -116,9 +121,9 @@ def test_prefix_cache_ignores_partial_last_page():
 def test_prefix_cache_on_evict_clears_mappings():
     prefix = RadixCache(64)
     prefix.record(0, list(range(64)), 0)
-    assert 0 in prefix._page_to_hash
+    assert prefix.has_page(0)
     prefix.evict(0)
-    assert 0 not in prefix._page_to_hash
+    assert not prefix.has_page(0)
 
 
 def test_prefix_cache_has_page():
@@ -162,7 +167,8 @@ def test_prefix_cache_does_not_record_partial_page():
 
 def test_page_pool_task_cacheable_ids_excludes_unmaterialized_tail():
     pool = _make_paged_pool_ps64()
-    assert pool.task_cacheable_ids("missing", [1, 2], [3, 4]) == [1, 2, 3]
+    task_cache = _make_task_cache(pool)
+    assert task_cache.task_cacheable_ids("missing", [1, 2], [3, 4]) == [1, 2, 3]
 
 
 # ---- ReqToTokenPool ----
@@ -243,31 +249,35 @@ def _make_contiguous_pool(**kwargs):
 
 def test_page_pool_contiguous_task_alloc_free():
     pool = _make_contiguous_pool()
-    assert pool.task_alloc("t1", [1, 2, 3])
-    assert "t1" in pool._task_req
-    pool.task_free("t1")
-    assert "t1" not in pool._task_req
+    task_cache = _make_task_cache(pool)
+    assert task_cache.task_alloc("t1", [1, 2, 3])
+    assert "t1" in task_cache._states
+    task_cache.task_free("t1")
+    assert "t1" not in task_cache._states
 
 
 def test_page_pool_contiguous_task_extend():
     pool = _make_contiguous_pool()
-    pool.task_alloc("t1", [1, 2, 3])
-    assert pool.task_extend("t1", 3)
-    assert pool.task_extend("t1", 63)
-    assert not pool.task_extend("t1", 64)
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", [1, 2, 3])
+    assert task_cache.task_extend("t1", 3)
+    assert task_cache.task_extend("t1", 63)
+    assert not task_cache.task_extend("t1", 64)
 
 
 def test_page_pool_contiguous_task_cached():
     pool = _make_contiguous_pool()
-    pool.task_alloc("t1", [1, 2, 3])
-    assert pool.task_cached("t1") == 0
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", [1, 2, 3])
+    assert task_cache.task_cached("t1") == 0
 
 
 def test_page_pool_contiguous_bind_tasks_prefill():
     pool = _make_contiguous_pool()
-    pool.task_alloc("t1", list(range(10)))
-    pool.task_alloc("t2", list(range(10)))
-    kv = pool.bind_tasks(["t1", "t2"], _ws(pool), start_pos=0)
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(10)))
+    task_cache.task_alloc("t2", list(range(10)))
+    kv = task_cache.bind(["t1", "t2"], _ws(pool), start_pos=0)
     assert kv.out_cache_loc.shape == (2, 10)
     assert kv.seq_lens.tolist() == [10, 10]
     assert kv.req_pool_indices.shape == (2,)
@@ -275,12 +285,13 @@ def test_page_pool_contiguous_bind_tasks_prefill():
 
 def test_page_pool_contiguous_bind_tasks_decode():
     pool = _make_contiguous_pool()
-    pool.task_alloc("t1", list(range(10)))
-    pool.task_alloc("t2", list(range(8)))
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(10)))
+    task_cache.task_alloc("t2", list(range(8)))
     # Simulate one decode extension so seq_lens advance to 11 and 9.
-    assert pool.task_extend("t1", 10)
-    assert pool.task_extend("t2", 8)
-    kv = pool.bind_tasks(["t1", "t2"], _ws(pool))
+    assert task_cache.task_extend("t1", 10)
+    assert task_cache.task_extend("t2", 8)
+    kv = task_cache.bind(["t1", "t2"], _ws(pool))
     assert kv.out_cache_loc.shape == (2, 1)
     assert kv.seq_lens.tolist() == [11, 9]
 
@@ -288,9 +299,10 @@ def test_page_pool_contiguous_bind_tasks_decode():
 def test_page_pool_contiguous_bind_roundtrip():
     """Write KV via bind_tasks, then gather via req_to_token indexing."""
     pool = _make_contiguous_pool(n_layers=1, n_kv_heads=2, head_dim=4)
-    pool.task_alloc("t1", list(range(4)))
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(4)))
 
-    kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
+    kv = task_cache.bind(["t1"], _ws(pool), start_pos=0)
     k = torch.randn(1, 4, 2, 4)
     v = torch.randn(1, 4, 2, 4)
     kv.k_buffer[0, kv.out_cache_loc] = k
@@ -324,35 +336,38 @@ def _make_paged_pool(**kwargs):
 
 def test_page_pool_paged_task_alloc():
     pool = _make_paged_pool()
-    assert pool.task_alloc("t1", list(range(10)))
-    req_idx = pool._task_req["t1"]
-    slots = pool._task_slots["t1"]
-    assert len(slots) == 10
-    assert pool._req_pool.req_to_token[req_idx, 0].item() == slots[0]
+    task_cache = _make_task_cache(pool)
+    assert task_cache.task_alloc("t1", list(range(10)))
+    state = task_cache._states["t1"]
+    assert len(state.pages) == 10
+    assert pool.req_pool.req_to_token[state.req_idx, 0].item() == state.pages[0]
 
 
 def test_page_pool_paged_task_extend():
     pool = _make_paged_pool()
-    pool.task_alloc("t1", list(range(4)))
-    assert pool.task_extend("t1", 4)
-    req_idx = pool._task_req["t1"]
-    slot = pool._req_pool.req_to_token[req_idx, 4].item()
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(4)))
+    assert task_cache.task_extend("t1", 4)
+    req_idx = task_cache._states["t1"].req_idx
+    slot = pool.req_pool.req_to_token[req_idx, 4].item()
     assert slot >= 0
 
 
 def test_page_pool_paged_task_free_releases_slots():
     pool = _make_paged_pool(n_tokens=16)
-    pool.task_alloc("t1", list(range(8)))
-    pool.task_free("t1")
-    assert "t1" not in pool._task_req
-    assert len(pool._req_pool.free_slots) == 4
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(8)))
+    task_cache.task_free("t1")
+    assert "t1" not in task_cache._states
+    assert len(pool.req_pool.free_slots) == 4
 
 
 def test_page_pool_paged_bind_roundtrip():
     pool = _make_paged_pool(n_layers=1, n_kv_heads=2, head_dim=4)
-    pool.task_alloc("t1", list(range(4)))
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(4)))
 
-    kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
+    kv = task_cache.bind(["t1"], _ws(pool), start_pos=0)
     k = torch.randn(1, 4, 2, 4)
     v = torch.randn(1, 4, 2, 4)
     kv.k_buffer[0, kv.out_cache_loc] = k
@@ -384,26 +399,53 @@ def _make_paged_pool_ps64(**kwargs):
 
 def test_page_pool_paged_ps64_task_alloc():
     pool = _make_paged_pool_ps64()
+    task_cache = _make_task_cache(pool)
     prompt = list(range(200))
-    assert pool.task_alloc("t1", prompt)
-    assert pool.task_cached("t1") == 0
+    assert task_cache.task_alloc("t1", prompt)
+    assert task_cache.task_cached("t1") == 0
     n_pages = (200 + 63) // 64
-    assert len(pool._task_pages["t1"]) == n_pages
+    assert len(task_cache._states["t1"].pages) == n_pages
 
 
 def test_page_pool_paged_ps64_task_extend_crosses_page():
     pool = _make_paged_pool_ps64()
-    pool.task_alloc("t1", list(range(64)))
-    assert pool.task_extend("t1", 64)
-    assert len(pool._task_pages["t1"]) >= 2
+    task_cache = _make_task_cache(pool)
+    task_cache.task_alloc("t1", list(range(64)))
+    assert task_cache.task_extend("t1", 64)
+    assert len(task_cache._states["t1"].pages) >= 2
+
+
+def test_page_pool_prefix_hit_populates_request_mapping():
+    pool = _make_paged_pool_ps64(page_size=2, max_seq_len=8, n_tokens=16)
+    task_cache = _make_task_cache(pool)
+    prompt = [11, 12, 13, 14]
+
+    assert task_cache.task_alloc("first", prompt)
+    task_cache.task_record_hashes("first", prompt)
+    task_cache.task_free("first")
+
+    assert task_cache.task_alloc("second", prompt)
+    second_state = task_cache._states["second"]
+    expected = [
+        page * pool.page_size + offset
+        for page in second_state.pages
+        for offset in range(pool.page_size)
+    ]
+
+    assert second_state.cached == len(prompt)
+    assert (
+        pool.req_pool.req_to_token[second_state.req_idx, : len(prompt)].tolist()
+        == expected
+    )
 
 
 def test_page_pool_paged_ps64_bind_roundtrip():
     pool = _make_paged_pool_ps64(n_layers=1, n_kv_heads=2, head_dim=4)
+    task_cache = _make_task_cache(pool)
     prompt = list(range(128))
-    pool.task_alloc("t1", prompt)
+    task_cache.task_alloc("t1", prompt)
 
-    kv = pool.bind_tasks(["t1"], _ws(pool), start_pos=0)
+    kv = task_cache.bind(["t1"], _ws(pool), start_pos=0)
     k = torch.randn(1, 128, 2, 4)
     v = torch.randn(1, 128, 2, 4)
     kv.k_buffer[0, kv.out_cache_loc] = k

tests/inference/test_engine.py

@@ -3,6 +3,7 @@
 import threading
 from unittest.mock import MagicMock, patch
 
+from astrai.extension import TorchNativeBackend, attn_backend
 from astrai.inference import STOP
 from astrai.inference.engine import GenerateResult, InferenceEngine
 
@@ -199,3 +200,40 @@ def test_engine_generate_zero_max_tokens_stream_is_empty():
         eng = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1)
         assert list(eng.generate("hello", stream=True, max_tokens=0)) == []
         instance.add_task.assert_not_called()
+
+
+def test_engine_passes_backend_to_scheduler():
+    mock_model = MagicMock()
+    mock_tokenizer = MagicMock()
+
+    with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
+        InferenceEngine(
+            mock_model,
+            mock_tokenizer,
+            max_batch_size=1,
+            backend="torch_native",
+        )
+
+    assert MockSched.call_args.kwargs["backend"] == "torch_native"
+
+
+def test_generate_captures_calling_backend_context():
+    mock_model = MagicMock()
+    mock_tokenizer = MagicMock()
+    captured = []
+
+    with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
+        instance = MockSched.return_value
+
+        def fake_add(prompt, **kwargs):
+            captured.append(kwargs["backend"])
+            kwargs["stream_callback"](STOP)
+            return "task"
+
+        instance.add_task.side_effect = fake_add
+        engine = InferenceEngine(mock_model, mock_tokenizer)
+        with attn_backend("torch_native"):
+            assert engine.generate("hello") == ""
+
+    assert len(captured) == 1
+    assert isinstance(captured[0], TorchNativeBackend)

tests/inference/test_protocol.py

@@ -5,9 +5,9 @@ from unittest.mock import MagicMock
 
 import pytest
 
-from astrai.inference.api.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
-from astrai.inference.api.protocol import GenContext, StopChecker, StopInfo
+from astrai.inference.network.anthropic import AnthropicResponseBuilder
+from astrai.inference.network.openai import OpenAIResponseBuilder
+from astrai.inference.network.protocol import GenContext, StopChecker, StopInfo
 
 
 def _make_ctx(**kwargs):

tests/inference/test_sample.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from astrai.inference.sample import (
+from astrai.inference.runtime.sample import (
     FrequencyPenaltyStrategy,
     SamplingPipeline,
     TemperatureStrategy,
@@ -268,7 +268,7 @@ def test_sample_return_logprobs_matches_manual_computation():
     logits = torch.randn(2, 30)
     tokens, logprobs = sample(logits, temperature=0.7, top_p=0.95, return_logprobs=True)
     # Recompute with the same pipeline
-    from astrai.inference.sample import (
+    from astrai.inference.runtime.sample import (
         SamplingPipeline,
         TemperatureStrategy,
         TopPStrategy,

tests/inference/test_scheduler.py

@@ -1,12 +1,17 @@
 """Tests for scheduler concurrency."""
 
 import threading
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
 
+from astrai.extension import CudaBackend, TorchNativeBackend, get_backend
 from astrai.inference import InferenceScheduler
+from astrai.inference.metrics import MetricsCollector
+from astrai.inference.runtime.executor import DecodeSteadyState, Executor
+from astrai.inference.task import Task
 from astrai.model.transformer import AutoRegressiveLM
 from tests.helpers import FakeTokenizer, make_rollout_config
 
@@ -38,8 +43,8 @@ def test_scheduler_concurrent_add_task(mock_model_and_tokenizer):
     """Test concurrent add_task operations."""
     mock_model, mock_tokenizer = mock_model_and_tokenizer
 
-    with patch("astrai.inference.core.scheduler.AutoModel"):
-        with patch("astrai.inference.core.scheduler.AutoTokenizer"):
+    with patch("astrai.inference.scheduler.AutoModel"):
+        with patch("astrai.inference.scheduler.AutoTokenizer"):
             scheduler = InferenceScheduler(
                 model=mock_model,
                 tokenizer=mock_tokenizer,
@@ -73,12 +78,75 @@ def test_scheduler_concurrent_add_task(mock_model_and_tokenizer):
     assert len(results["task_ids"]) == 50
 
 
+def test_generation_loop_activates_backend_in_worker_thread():
+    scheduler = object.__new__(InferenceScheduler)
+    scheduler._backend = TorchNativeBackend()
+    scheduler._stop_event = threading.Event()
+    scheduler._task_cache = MagicMock()
+
+    observed = []
+    task_mgr = MagicMock()
+    task_mgr.tokenizer.stop_ids = [0]
+    task_mgr.remove_finished_tasks.return_value = []
+    task_mgr.get_active_tasks.return_value = []
+    task_mgr.max_batch_size = 1
+    task_mgr.pull_candidates.return_value = []
+    task_mgr.has_work.return_value = False
+
+    def observe_backend(*args, **kwargs):
+        observed.append(type(get_backend()))
+        scheduler._stop_event.set()
+
+    task_mgr.wait_for_tasks.side_effect = observe_backend
+    scheduler._task_mgr = task_mgr
+
+    thread = threading.Thread(target=scheduler._run_generation_loop)
+    thread.start()
+    thread.join(timeout=5)
+
+    assert not thread.is_alive()
+    assert observed == [TorchNativeBackend]
+
+
+def test_step_splits_decode_batch_by_request_backend():
+    scheduler = object.__new__(InferenceScheduler)
+    scheduler._task_cache = MagicMock()
+    scheduler._task_cache.task_extend.return_value = True
+    scheduler._metrics = MetricsCollector()
+    scheduler._executor = MagicMock()
+
+    observed = []
+
+    def execute(tasks, **kwargs):
+        observed.append((type(get_backend()), [task.task_id for task in tasks]))
+        return [1] * len(tasks)
+
+    scheduler._executor.execute_decode.side_effect = execute
+
+    torch_task = Task("torch", [1], backend=TorchNativeBackend())
+    cuda_task = Task("cuda", [1], backend=CudaBackend())
+    for task in (torch_task, cuda_task):
+        task.input_tokens = 1
+        task.output_ids = [1]
+        task.mark_prefill_done()
+        scheduler._metrics.register(task.task_id)
+
+    produced, aborted = scheduler._step([torch_task, cuda_task])
+
+    assert aborted == []
+    assert produced == [torch_task, cuda_task]
+    assert observed == [
+        (TorchNativeBackend, ["torch"]),
+        (CudaBackend, ["cuda"]),
+    ]
+
+
 def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
     """Test concurrent add and remove task operations."""
     mock_model, mock_tokenizer = mock_model_and_tokenizer
 
-    with patch("astrai.inference.core.scheduler.AutoModel"):
-        with patch("astrai.inference.core.scheduler.AutoTokenizer"):
+    with patch("astrai.inference.scheduler.AutoModel"):
+        with patch("astrai.inference.scheduler.AutoTokenizer"):
             scheduler = InferenceScheduler(
                 model=mock_model,
                 tokenizer=mock_tokenizer,
@@ -126,8 +194,8 @@ def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
     """Test concurrent get_stats operations."""
     mock_model, mock_tokenizer = mock_model_and_tokenizer
 
-    with patch("astrai.inference.core.scheduler.AutoModel"):
-        with patch("astrai.inference.core.scheduler.AutoTokenizer"):
+    with patch("astrai.inference.scheduler.AutoModel"):
+        with patch("astrai.inference.scheduler.AutoTokenizer"):
             scheduler = InferenceScheduler(
                 model=mock_model,
                 tokenizer=mock_tokenizer,
@@ -180,7 +248,7 @@ def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
 def _make_real_scheduler(device):
     """Build a scheduler backed by a tiny real model for run_batch tests."""
     cfg = make_rollout_config(max_position_embeddings=64)
-    model = AutoRegressiveLM(cfg).to(device=device).eval()
+    model = AutoRegressiveLM(cfg).to(device=device, dtype=torch.bfloat16).eval()
     tokenizer = FakeTokenizer()
     scheduler = InferenceScheduler(
         model=model,
@@ -302,3 +370,45 @@ def test_run_batch_too_long_prompt_skipped(device):
         assert len(results[1]) <= 2
     finally:
         scheduler.stop()
+
+
+def test_decode_does_not_reuse_previous_batch_state():
+    executor = object.__new__(Executor)
+    executor.device = torch.device("cpu")
+    executor.task_cache = MagicMock()
+    executor.task_cache.bind_was_steady = True
+    executor.task_cache.bind.return_value = MagicMock()
+    executor._graph_supported = False
+    executor._graph_ctx = SimpleNamespace(enabled=False)
+
+    workspace = MagicMock()
+    workspace.position_ids = torch.tensor([2], dtype=torch.long)
+    workspace.fill_input_ids.return_value = torch.tensor([7], dtype=torch.long)
+    workspace.decode_mask.return_value = torch.ones(1, 1, 9, dtype=torch.bool)
+    executor._workspace = workspace
+    executor.model = MagicMock(
+        return_value={"logits": torch.zeros(1, 1, 10, dtype=torch.float32)}
+    )
+
+    old_info = object()
+    new_info = object()
+    executor._decode_cache = DecodeSteadyState(("old",), [2], old_info)
+    executor._sample_logits = MagicMock(return_value=[3])
+
+    task = Task("new", list(range(8)), temperature=0)
+    task.input_tokens = 8
+    task.output_ids = [7]
+    task.mark_prefill_done()
+
+    with patch(
+        "astrai.inference.runtime.executor._build_sampling_batch_info",
+        return_value=new_info,
+    ):
+        assert executor.execute_decode([task]) == [3]
+
+    assert workspace.position_ids.tolist() == [8]
+    assert executor._decode_cache.task_sig == ("new",)
+    executor._sample_logits.assert_called_once()
+    args, kwargs = executor._sample_logits.call_args
+    assert args[1:] == ([task], False)
+    assert kwargs["info"] is new_info

tests/inference/test_task.py

@@ -20,11 +20,12 @@ def test_task_default_status_is_pending():
 def test_task_next_pos():
     task = Task("id1", [1, 2, 3])
     task.input_tokens = 5
+    task.mark_prefill_done()
     assert task.next_pos == 5
-    task.output_ids.append(4)
-    assert task.next_pos == 5
-    task.output_ids.append(5)
+    task.advance_kv()
     assert task.next_pos == 6
+    task.advance_kv()
+    assert task.next_pos == 7
 
 
 def test_task_is_finished_max_tokens():

tests/inference/test_tool_parser.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from astrai.inference.api.tool_parser import (
+from astrai.inference.network.tool_parser import (
     _TOOL_CALL_HEAD_RE,
     BaseToolParser,
     SimpleJsonToolParser,

tests/trainer/test_metric_util.py

@@ -0,0 +1,17 @@
+import pytest
+import torch
+
+from astrai.trainer.metric_util import GradSNRTracker
+
+
+def test_grad_snr_is_reported_in_decibels():
+    model = torch.nn.Linear(1, 1, bias=False)
+    tracker = GradSNRTracker(beta=0.5, eps=1e-8)
+
+    model.weight.grad = torch.tensor([[1.0]])
+    tracker.update(model)
+    model.weight.grad = torch.tensor([[3.0]])
+    tracker.update(model)
+
+    # E[g]^2 / Var(g) = 4 / 1 = 4, which is 6.0206 dB.
+    assert tracker.snr == pytest.approx(10.0 * torch.log10(torch.tensor(4.0)).item())

tests/trainer/test_rollout.py

@@ -3,7 +3,7 @@
 import pytest
 import torch
 
-from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.scheduler import InferenceScheduler
 from astrai.trainer.rollout import (
     BaseRewardModel,
     RawRollout,