fix: size KV pool from prompt/gen args in benchmark

- Drop hardcoded CACHE_MAX_SEQ=2048 which overflowed at long prompts
- Size prefill pool to prompt_length and decode pool to prompt+5+gen*num_trials
- Unblocks decode/prefill benchmark at prompt 4096+ (was KV cache index OOB)

CONTRIBUTING.md

@@ -20,9 +20,6 @@ Run the following checks **in order** — CI will reject if any fail.
 ruff format .
 ```
 
-> **Note**: `ruff format` may rename parameters (e.g. `mask` → `attn_mask`).  
-> Always review the diff after formatting.
-
 ### 2. Import sorting
 
 ```bash
@@ -44,7 +41,7 @@ python -u -m pytest tests/ -v
 
 > Failed tests may leave orphan tempdirs under `%TEMP%`. Clean them manually if needed.
 
-### 4. (Optional) Full pre-commit check
+### 4. (Optional) Full pre-commit check script
 
 If you have Git Bash available:
 
@@ -52,12 +49,17 @@ If you have Git Bash available:
 bash scripts/pre_commit.sh
 ```
 
-This runs format check, import sort check, and tests in one go.
+The script installs development dependencies by default, then runs the format
+check, import sort check, and tests. If dependencies are already installed, use:
+
+```bash
+bash scripts/pre_commit.sh --skip-deps
+```
 
 ## Commit Style
 
 ```
-fix/feat/chore/docs/refactor/perf/test/style/ci/build/revert : short description (~50 chars)
+type: short description (~50 chars)
 
 - bullet point body (each ~60 chars)
 ```
@@ -73,7 +75,7 @@ fix/feat/chore/docs/refactor/perf/test/style/ci/build/revert : short description
 |---------|-------|-----|
 | `ruff check --select I` fails | Wrong import order | `ruff check . --select I --fix .` then `ruff format .` |
 | `ruff format` changed many files | Not formatted before commit | Review diff carefully before staging |
-| Pre-commit hook rejects | Tests or lint failed | Fix individually, do not `--no-verify` |
+| Pre-commit check script fails | Dependency install, tests, or lint failed | Fix the failing step; use `--skip-deps` only when dependencies are already installed |
 | Tests fail with tempdir left | Test crash | Clean `%TEMP%` manually |
 
 ## Submitting Changes

README.md

@@ -56,6 +56,8 @@ End-to-end walkthrough in 5 steps:
 
 **1. Install**
 
+AstrAI requires Python 3.12+ and pins PyTorch exactly to `2.11.0`. Training, `scripts/tools/generate.py`, generation evaluations, and the generation demos require CUDA; CPU support is limited to components with an explicit CPU device path, such as the HTTP server and direct-scoring evaluations.
+
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
@@ -132,7 +134,7 @@ Check out the demos in the `scripts/demo/` folder:
 # Download model weights (required before running demos)
 python scripts/demo/download.py                      # model → params/
 
-# Interactive streaming chat (multi-turn, maintains history)
+# Single-turn interactive streaming prompt loop (no conversation history)
 python scripts/demo/stream_chat.py
 # Type your message after >>, type !exit to quit
 
@@ -183,7 +185,7 @@ docker run --gpus all -v /path/to/data:/data -it astrai:latest
 # Docker Compose (GPU, default)
 docker compose up -d
 
-# Docker Compose (CPU only)
+# Docker Compose CPU server profile (CUDA-only generation scripts/demos are unavailable)
 docker compose --profile cpu up -d
 ```
 
@@ -256,4 +258,4 @@ This project is licensed under the [GPL-3.0 License](LICENSE).
 
 <div align="center">
   <em>A lightweight Transformer framework designed for both high performance and ease of use.</em>
-</div>
+</div>

astrai/config/model_config.py

@@ -63,6 +63,11 @@ class AutoRegressiveLMConfig(BaseModelConfig):
         n_shared_experts (Optional[int]): Number of shared experts, MoE only. Defaults to None.
         n_activated_experts (Optional[int]): Number of activated experts per token, MoE only. Defaults to None.
         topk_method (Optional[str]): Top-k routing method, MoE only. Defaults to None.
+        moe_intermediate_size (Optional[int]): Expert hidden dim, defaults to intermediate_size if None. MoE only.
+        shared_expert_intermediate_size (Optional[int]): Shared expert hidden dim, defaults to intermediate_size if None. MoE only.
+        norm_topk_prob (bool): Normalize top-k routing probabilities. Defaults to True.
+        decoder_sparse_step (int): Frequency of MoE layers, 1=every layer. Defaults to 1.
+        mlp_only_layers (Optional[list[int]]): Layer indices using dense MLP instead of MoE. Defaults to None.
     """
 
     vocab_size: Optional[int] = None
@@ -87,6 +92,11 @@ class AutoRegressiveLMConfig(BaseModelConfig):
     n_shared_experts: Optional[int] = None
     n_activated_experts: Optional[int] = None
     topk_method: Optional[str] = None
+    moe_intermediate_size: Optional[int] = None
+    shared_expert_intermediate_size: Optional[int] = None
+    norm_topk_prob: bool = True
+    decoder_sparse_step: int = 1
+    mlp_only_layers: Optional[list[int]] = None
 
     @field_validator("attn_type")
     def _validate_attn_type(cls, v: str) -> str:
@@ -102,6 +112,12 @@ class AutoRegressiveLMConfig(BaseModelConfig):
             raise ValueError(f"ffn_type must be one of {sorted(_FFN_TYPES)}, got {v!r}")
         return v
 
+    @field_validator("decoder_sparse_step")
+    def _validate_decoder_sparse_step(cls, v: int) -> int:
+        if v < 1:
+            raise ValueError(f"decoder_sparse_step must be at least 1, got {v}")
+        return v
+
 
 @dataclass
 @ConfigFactory.register("embedding")

astrai/config/train_config.py

@@ -61,6 +61,7 @@ class TrainConfig(BaseConfig):
         val_split (Optional[float]): Ratio to split from training dataset for validation, e.g. 0.05. Defaults to None.
         val_step (int): Number of optimizer steps between validation runs. Defaults to 1000.
         neftune_alpha (float): NEFTune noise alpha, 0=disabled, typical: 5.0. Defaults to 0.0.
+        moe_aux_loss_coef (float): Weight applied to the MoE load-balancing loss. Defaults to 0.01.
         rollout_interval (int): Number of optimizer steps between online rollouts. Defaults to 512.
         rollout_temperature (float): Sampling temperature for online rollout. Defaults to 0.7.
         rollout_top_k (int): Top-k filtering for online rollout, 0=disable. Defaults to 0.
@@ -112,6 +113,7 @@ class TrainConfig(BaseConfig):
     val_split: Optional[float] = None
     val_step: int = 1000
     neftune_alpha: float = 0.0
+    moe_aux_loss_coef: float = 0.01
 
     rollout_interval: int = 512
     rollout_temperature: float = 0.7
@@ -187,7 +189,9 @@ class TrainConfig(BaseConfig):
             raise ValueError(f"rollout_top_p must be in (0, 1], got {v}")
         return v
 
-    @field_validator("rollout_top_k", "num_workers", "neftune_alpha")
+    @field_validator(
+        "rollout_top_k", "num_workers", "neftune_alpha", "moe_aux_loss_coef"
+    )
     def _validate_non_negative(cls, v):
         if v < 0:
             raise ValueError(f"must be non-negative, got {v}")

astrai/model/__init__.py

@@ -9,7 +9,7 @@ from astrai.model.components.lora import (
     merge_lora,
     save_lora,
 )
-from astrai.model.components.mlp import MLP
+from astrai.model.components.mlp import MLP, DeepSeekMoE
 from astrai.model.components.norm import RMSNorm
 from astrai.model.encoder import EmbeddingEncoder
 from astrai.model.transformer import AutoRegressiveLM
@@ -19,6 +19,7 @@ __all__ = [
     "Linear",
     "RMSNorm",
     "MLP",
+    "DeepSeekMoE",
     "GQA",
     "DecoderBlock",
     # Models

astrai/model/components/__init__.py

@@ -3,7 +3,7 @@ from astrai.model.components.attention import GQA, MLA
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
 from astrai.model.components.linear import Linear
-from astrai.model.components.mlp import MLP
+from astrai.model.components.mlp import MLP, DeepSeekMoE
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import (
     RotaryEmbedding,
@@ -14,6 +14,7 @@ __all__ = [
     "Linear",
     "RMSNorm",
     "MLP",
+    "DeepSeekMoE",
     "Embedding",
     "GQA",
     "MLA",

astrai/model/components/decoder_block.py

@@ -1,5 +1,5 @@
 from dataclasses import asdict
-from typing import Optional
+from typing import Optional, TypedDict
 
 import torch.nn as nn
 from torch import Tensor
@@ -10,6 +10,11 @@ from astrai.model.components.mlp import FFNFactory
 from astrai.model.components.norm import RMSNorm
 
 
+class DecoderOutput(TypedDict):
+    hidden_states: Tensor
+    aux_loss: Optional[Tensor]
+
+
 class DecoderBlock(nn.Module):
     def __init__(self, config, layer_id: int):
         super().__init__()
@@ -26,7 +31,20 @@ class DecoderBlock(nn.Module):
         self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
         self.input_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
         self.post_attention_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.mlp = FFNFactory.create(config.ffn_type, **cfg)
+        ffn_type = self._resolve_ffn_type(config, layer_id)
+        self.mlp = FFNFactory.create(ffn_type, **cfg)
+
+    @staticmethod
+    def _resolve_ffn_type(config, layer_id: int) -> str:
+        if config.ffn_type != "moe":
+            return config.ffn_type
+        mlp_only = config.mlp_only_layers or []
+        if layer_id in mlp_only:
+            return "mlp"
+        if config.decoder_sparse_step > 1:
+            if (layer_id + 1) % config.decoder_sparse_step != 0:
+                return "mlp"
+        return "moe"
 
     def forward(
         self,
@@ -35,7 +53,7 @@ class DecoderBlock(nn.Module):
         attention_mask: Optional[Tensor] = None,
         kv_cache: Optional[KVCache] = None,
         is_causal: bool = False,
-    ) -> Tensor:
+    ) -> DecoderOutput:
         attn_output = self.attention(
             self.input_norm(x),
             rotary_emb,
@@ -44,6 +62,8 @@ class DecoderBlock(nn.Module):
             is_causal,
         )
         x = attn_output + x
-        x = self.mlp(self.post_attention_norm(x)) + x
+        normalized = self.post_attention_norm(x)
+        mlp_output = self.mlp(normalized)
+        x = mlp_output["hidden_states"] + x
 
-        return x
+        return {"hidden_states": x, "aux_loss": mlp_output["aux_loss"]}

astrai/model/components/mlp.py

@@ -1,3 +1,5 @@
+from typing import Optional, TypedDict
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -11,6 +13,16 @@ class FFNFactory(BaseFactory[nn.Module]):
     pass
 
 
+class FFNOutput(TypedDict):
+    hidden_states: Tensor
+    aux_loss: Optional[Tensor]
+
+
+class RoutedOutput(TypedDict):
+    hidden_states: Tensor
+    aux_loss: Optional[Tensor]
+
+
 @FFNFactory.register("mlp")
 class MLP(nn.Module):
     def __init__(self, dim: int, dim_ffn: int, down_init_std: float = 0.02):
@@ -19,10 +31,10 @@ class MLP(nn.Module):
         self.gate = Linear(dim, dim_ffn)
         self.down = Linear(dim_ffn, dim, init_std=down_init_std)
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, x: Tensor) -> FFNOutput:
         gated = self.up(x) * F.silu(self.gate(x))
         out = self.down(gated)
-        return out
+        return {"hidden_states": out, "aux_loss": None}
 
 
 @FFNFactory.register("moe")
@@ -36,6 +48,9 @@ class DeepSeekMoE(nn.Module):
         n_activated_experts: int = 2,
         topk_method: str = "greedy",
         n_layers: int = 1,
+        moe_intermediate_size: Optional[int] = None,
+        shared_expert_intermediate_size: Optional[int] = None,
+        norm_topk_prob: bool = True,
     ):
         super().__init__()
         self.dim = dim
@@ -43,6 +58,16 @@ class DeepSeekMoE(nn.Module):
         self.n_shared_experts = n_shared_experts
         self.n_activated_experts = n_activated_experts
         self.topk_method = topk_method
+        self.norm_topk_prob = norm_topk_prob
+
+        expert_dim_ffn = (
+            moe_intermediate_size if moe_intermediate_size is not None else dim_ffn
+        )
+        shared_dim_ffn = (
+            shared_expert_intermediate_size
+            if shared_expert_intermediate_size is not None
+            else dim_ffn
+        )
 
         self.router = Linear(dim, n_routed_experts, bias=False)
         moe_scale = 1 / max(n_shared_experts, 1) + 1 / n_activated_experts
@@ -50,33 +75,37 @@ class DeepSeekMoE(nn.Module):
 
         self.shared_experts = nn.ModuleList(
             [
-                MLP(dim, dim_ffn, down_init_std=down_init_std)
+                MLP(dim, shared_dim_ffn, down_init_std=down_init_std)
                 for _ in range(n_shared_experts)
             ]
         )
         self.routed_experts = nn.ModuleList(
             [
-                MLP(dim, dim_ffn, down_init_std=down_init_std)
+                MLP(dim, expert_dim_ffn, down_init_std=down_init_std)
                 for _ in range(n_routed_experts)
             ]
         )
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, x: Tensor) -> FFNOutput:
+        include_aux_loss = self.training and torch.is_grad_enabled()
         bsz, seq_len, dim = x.shape
         x_flat = x.view(-1, dim)
 
         shared_out = self._shared_forward(x_flat)
-        routed_out = self._routed_forward(x_flat)
+        routed_output = self._routed_forward(x_flat, include_aux_loss)
 
-        out = (shared_out + routed_out).view(bsz, seq_len, dim)
-        return out
+        out = (shared_out + routed_output["hidden_states"]).view(bsz, seq_len, dim)
+        return {"hidden_states": out, "aux_loss": routed_output["aux_loss"]}
 
     def _shared_forward(self, x: Tensor) -> Tensor:
         if self.n_shared_experts == 0:
             return torch.zeros_like(x)
-        return sum(e(x) for e in self.shared_experts) / self.n_shared_experts
+        return (
+            sum(e(x)["hidden_states"] for e in self.shared_experts)
+            / self.n_shared_experts
+        )
 
-    def _routed_forward(self, x: Tensor) -> Tensor:
+    def _routed_forward(self, x: Tensor, include_aux_loss: bool) -> RoutedOutput:
         N, D = x.shape
         K = self.n_activated_experts
 
@@ -84,7 +113,17 @@ class DeepSeekMoE(nn.Module):
         router_probs = torch.softmax(router_logits.float(), dim=-1).to(x.dtype)
 
         topk_weights, topk_indices = torch.topk(router_probs, K, dim=-1)
-        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+        if self.norm_topk_prob:
+            topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+        aux_loss = None
+        if include_aux_loss:
+            expert_load = F.one_hot(
+                topk_indices, num_classes=self.n_routed_experts
+            ).float()
+            expert_load = expert_load.mean(dim=(0, 1))
+            router_prob = router_probs.float().mean(dim=0)
+            aux_loss = self.n_routed_experts * (expert_load * router_prob).sum()
 
         output = torch.zeros(N, D, device=x.device, dtype=x.dtype)
         for expert_idx in range(self.n_routed_experts):
@@ -92,9 +131,11 @@ class DeepSeekMoE(nn.Module):
             token_idx, k_idx = expert_mask.nonzero(as_tuple=True)
             if token_idx.numel() == 0:
                 continue
+            expert = self.routed_experts[expert_idx]
             expert_input = x[token_idx]
-            expert_output = self.routed_experts[expert_idx](expert_input)
+            expert_output = expert(expert_input)["hidden_states"]
+
             weights = topk_weights[token_idx, k_idx].unsqueeze(-1)
             output.index_add_(0, token_idx, expert_output * weights)
 
-        return output
+        return {"hidden_states": output, "aux_loss": aux_loss}

astrai/model/encoder.py

@@ -70,7 +70,7 @@ class EmbeddingEncoder(AutoModel):
         attn_mask = process_attention_mask(input_mask)
 
         for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask)
+            x = layer(x, rotary_emb, attn_mask)["hidden_states"]
 
         hidden_states = self.norm(x)
 

astrai/model/transformer.py

@@ -113,10 +113,23 @@ class AutoRegressiveLM(AutoModel):
         attn_mask = process_attention_mask(input_mask)
         use_sdpa_causal_mask = attn_mask is None
 
+        aux_losses = []
         for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, kv_cache, use_sdpa_causal_mask)
+            layer_output = layer(
+                x,
+                rotary_emb,
+                attn_mask,
+                kv_cache,
+                use_sdpa_causal_mask,
+            )
+            x = layer_output["hidden_states"]
+            if layer_output["aux_loss"] is not None:
+                aux_losses.append(layer_output["aux_loss"])
 
         hidden_states = self.norm(x)
         logits = self.lm_head(hidden_states)
 
-        return {"logits": logits, "hidden_states": hidden_states}
+        output = {"logits": logits, "hidden_states": hidden_states}
+        if aux_losses:
+            output["aux_loss"] = torch.stack(aux_losses).mean()
+        return output

astrai/trainer/strategy.py

@@ -1,7 +1,7 @@
 """Training strategy implementations with factory pattern."""
 
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, Union
+from typing import Callable, Dict, Optional, TypedDict, Union
 
 import torch
 import torch.nn as nn
@@ -13,6 +13,16 @@ from astrai.parallel.executor import broadcast_state_dict
 from astrai.trainer.rollout import RolloutResult
 
 
+class LossOutput(TypedDict):
+    loss: Tensor
+    metrics: Dict[str, float]
+
+
+class LogprobsOutput(TypedDict):
+    logprobs: Tensor
+    aux_loss: Optional[Tensor]
+
+
 def move_to_device(batch: Dict[str, Tensor], device: str) -> Dict[str, Tensor]:
     """Move batch tensors to specified device with non-blocking transfer."""
     return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
@@ -24,7 +34,7 @@ def get_logprobs(
     attn_mask: Tensor,
     loss_mask: Tensor,
     reduction: str,
-) -> Tensor:
+) -> LogprobsOutput:
     """Compute token-wise log probabilities from model outputs.
 
     Args:
@@ -46,10 +56,11 @@ def get_logprobs(
     shifted_input_ids = input_ids[:, 1:]
     shifted_loss_mask = loss_mask[:, 1:]
 
-    logits = model(
+    outputs = model(
         input_ids[:, :-1],
         attn_mask[:, :, :-1, :-1] if attn_mask.dim() == 4 else attn_mask[:, :-1],
-    )["logits"]
+    )
+    logits = outputs["logits"]
     log_probs = torch.log_softmax(logits.float(), dim=-1)
 
     token_logprobs = torch.gather(
@@ -57,13 +68,14 @@ def get_logprobs(
     ).squeeze(-1)
 
     if reduction == "mean":
-        return (token_logprobs * shifted_loss_mask).sum(dim=-1) / shifted_loss_mask.sum(
+        logprobs = (token_logprobs * shifted_loss_mask).sum(
             dim=-1
-        ).clamp(min=1.0)
+        ) / shifted_loss_mask.sum(dim=-1).clamp(min=1.0)
     elif reduction == "sum":
-        return (token_logprobs * shifted_loss_mask).sum(dim=-1)
+        logprobs = (token_logprobs * shifted_loss_mask).sum(dim=-1)
     else:
-        return token_logprobs * shifted_loss_mask
+        logprobs = token_logprobs * shifted_loss_mask
+    return {"logprobs": logprobs, "aux_loss": outputs.get("aux_loss")}
 
 
 def make_doc_boundary_mask(position_ids: Tensor) -> Tensor:
@@ -102,6 +114,7 @@ class BaseStrategy(ABC):
         self.model = model
         self.device = device
         self.executor = kwargs.pop("executor", None)
+        self.moe_aux_loss_coef = kwargs.pop("moe_aux_loss_coef", 0.01)
         self.extra_kwargs = kwargs
         self._rollout_runner = None
 
@@ -117,6 +130,33 @@ class BaseStrategy(ABC):
         """
         raise NotImplementedError
 
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
+        return self._normalize_output(self.compute_loss(batch))
+
+    def _loss_output(
+        self,
+        task_loss: Tensor,
+        metrics: Dict[str, Tensor],
+        aux_loss: Optional[Tensor] = None,
+    ) -> LossOutput:
+        total_loss = task_loss
+        if aux_loss is not None:
+            weighted_aux_loss = self.moe_aux_loss_coef * aux_loss
+            total_loss = total_loss + weighted_aux_loss
+            metrics["moe_aux_loss"] = aux_loss
+            metrics["moe_aux_loss_weighted"] = weighted_aux_loss
+        metrics["loss"] = total_loss
+        return {
+            "loss": total_loss,
+            "metrics": {name: value.detach().item() for name, value in metrics.items()},
+        }
+
+    @staticmethod
+    def _normalize_output(output: Union[LossOutput, Tensor]) -> LossOutput:
+        if isinstance(output, dict):
+            return output
+        return {"loss": output, "metrics": {"loss": output.detach().item()}}
+
     def supports_online(self) -> bool:
         """Whether this strategy can operate with a rollout runner.
 
@@ -153,17 +193,17 @@ class BaseStrategy(ABC):
         if self._rollout_runner is not None:
             self._rollout_runner.step()
 
-    def __call__(self, batch: Dict[str, Tensor]) -> Tensor:
+    def __call__(self, batch: Dict[str, Tensor]) -> LossOutput:
         """Run offline or online forward depending on runner injection."""
         if self._rollout_runner is None:
-            return self.compute_loss(batch)
+            return self.compute_loss_output(batch)
 
         result, is_fresh = self._rollout_runner(batch)
         if is_fresh:
             self._on_rollout_refresh()
 
         train_batch = self.prepare_from_rollout(result)
-        return self.compute_loss(train_batch)
+        return self.compute_loss_output(train_batch)
 
 
 class StrategyFactory(BaseFactory["BaseStrategy"]):
@@ -203,9 +243,13 @@ class SEQStrategy(BaseStrategy):
         self.label_smoothing = label_smoothing
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        return self.compute_loss_output(batch)["loss"]
+
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         input_ids, target_ids = batch["input_ids"], batch["target_ids"]
-        logits = self.model(input_ids=input_ids)["logits"]
+        outputs = self.model(input_ids=input_ids)
+        logits = outputs["logits"]
 
         loss = F.cross_entropy(
             input=logits.flatten(0, 1).float(),
@@ -213,7 +257,7 @@ class SEQStrategy(BaseStrategy):
             label_smoothing=self.label_smoothing,
         )
 
-        return loss
+        return self._loss_output(loss, {"task_loss": loss}, outputs.get("aux_loss"))
 
 
 @StrategyFactory.register("sft")
@@ -234,6 +278,9 @@ class SFTStrategy(BaseStrategy):
         self.label_smoothing = label_smoothing
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        return self.compute_loss_output(batch)["loss"]
+
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         input_ids, target_ids, position_ids, loss_mask = (
             batch["input_ids"],
@@ -245,9 +292,10 @@ class SFTStrategy(BaseStrategy):
         ignore_index = -100
         input_mask = make_doc_boundary_mask(position_ids)
         target_ids = target_ids.masked_fill(~loss_mask, ignore_index)
-        logits = self.model(
+        outputs = self.model(
             input_ids=input_ids, position_ids=position_ids, input_mask=input_mask
-        )["logits"]
+        )
+        logits = outputs["logits"]
 
         loss = F.cross_entropy(
             input=logits.flatten(0, 1).float(),
@@ -256,7 +304,7 @@ class SFTStrategy(BaseStrategy):
             label_smoothing=self.label_smoothing,
         )
 
-        return loss
+        return self._loss_output(loss, {"task_loss": loss}, outputs.get("aux_loss"))
 
 
 @StrategyFactory.register("dpo")
@@ -282,6 +330,9 @@ class DPOStrategy(BaseStrategy):
         self.reduction = reduction
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        return self.compute_loss_output(batch)["loss"]
+
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         chosen_ids, rejected_ids = batch["chosen"], batch["rejected"]
         chosen_mask, rejected_mask = batch["chosen_mask"], batch["rejected_mask"]
@@ -297,22 +348,25 @@ class DPOStrategy(BaseStrategy):
         )[None, None, :, :]  # [1, 1, S, S]
         full_mask = key_pad & causal  # [B*2, 1, S, S] — composed
 
-        log_pi = get_logprobs(
+        policy_output = get_logprobs(
             self.model,
             concat_ids,
             full_mask,
             concat_loss_mask,
             self.reduction,
         )
+        log_pi = policy_output["logprobs"]
+        aux_loss = policy_output["aux_loss"]
 
         with torch.no_grad():
-            log_ref = get_logprobs(
+            ref_output = get_logprobs(
                 self.ref_model,
                 concat_ids,
                 full_mask,
                 concat_loss_mask,
                 self.reduction,
             )
+            log_ref = ref_output["logprobs"]
 
         log_pi_chosen = log_pi[: chosen_ids.shape[0]]
         log_pi_rejected = log_pi[chosen_ids.shape[0] :]
@@ -325,7 +379,7 @@ class DPOStrategy(BaseStrategy):
         ratio_diff = pi_log_ratio - ref_log_ratio
         dpo_loss = -F.logsigmoid(self.beta * ratio_diff).mean()
 
-        return dpo_loss
+        return self._loss_output(dpo_loss, {"dpo_loss": dpo_loss}, aux_loss)
 
     def supports_online(self) -> bool:
         return True
@@ -398,6 +452,9 @@ class GRPOStrategy(BaseStrategy):
             self.old_model.load_state_dict(state_dict)
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        return self.compute_loss_output(batch)["loss"]
+
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         prompts = batch["prompts"]
         responses = batch["responses"]
@@ -438,16 +495,23 @@ class GRPOStrategy(BaseStrategy):
         # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len).
         # Response token logprobs occupy the last ``response_len`` positions
         # (the first response token is predicted from the last prompt token).
-        token_log_probs_policy = get_logprobs(
+        policy_output = get_logprobs(
             self.model, full_sequences, attn_mask, full_masks, "none"
-        )[:, prompt_len - 1 :]
+        )
+        token_log_probs_policy = policy_output["logprobs"]
+        aux_loss = policy_output["aux_loss"]
+        token_log_probs_policy = token_log_probs_policy[:, prompt_len - 1 :]
         with torch.no_grad():
-            token_log_probs_old = get_logprobs(
+            old_output = get_logprobs(
                 self.old_model, full_sequences, attn_mask, full_masks, "none"
-            )[:, prompt_len - 1 :]
-            token_log_probs_ref = get_logprobs(
+            )
+            token_log_probs_old = old_output["logprobs"]
+            token_log_probs_old = token_log_probs_old[:, prompt_len - 1 :]
+            ref_output = get_logprobs(
                 self.ref_model, full_sequences, attn_mask, full_masks, "none"
-            )[:, prompt_len - 1 :]
+            )
+            token_log_probs_ref = ref_output["logprobs"]
+            token_log_probs_ref = token_log_probs_ref[:, prompt_len - 1 :]
 
         # Reshape to [B, G, response_len]
         token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1)
@@ -480,9 +544,12 @@ class GRPOStrategy(BaseStrategy):
         kl_per_token = r - torch.log(r + eps) - 1.0
         kl_penalty = self.kl_coef * (kl_per_token * token_masks).sum() / token_count
 
-        total_loss = policy_loss + kl_penalty
-
-        return total_loss
+        task_loss = policy_loss + kl_penalty
+        return self._loss_output(
+            task_loss,
+            {"policy_loss": policy_loss, "kl_loss": kl_penalty},
+            aux_loss,
+        )
 
     def supports_online(self) -> bool:
         return True

astrai/trainer/train_callback.py

@@ -260,11 +260,28 @@ class MetricCallback(TrainCallback):
         }
 
     def _metrics(self, context: TrainContext, names):
-        return {
-            m: self._metric_funcs[m](context)
-            for m in names
-            if self._metric_funcs[m](context) is not None
-        }
+        metrics = dict(context.metrics)
+        for name in names:
+            metric_fn = self._metric_funcs.get(name)
+            if metric_fn is None:
+                continue
+            value = metric_fn(context)
+            if value is not None:
+                metrics[name] = value
+        selected = set(context.metrics) | set(names)
+        selected.discard("*")
+        result = {name: metrics[name] for name in selected if name in metrics}
+        if context.world_size > 1 and dist.is_initialized() and result:
+            metric_names = sorted(result)
+            values = torch.tensor(
+                [result[name] for name in metric_names],
+                dtype=torch.float32,
+                device=get_current_device(),
+            )
+            dist.all_reduce(values, op=dist.ReduceOp.SUM)
+            values /= context.world_size
+            result.update(zip(metric_names, values.tolist()))
+        return result
 
     @only_on_rank(0)
     def _append(self, event_type: str, context: TrainContext, **extra):
@@ -286,8 +303,8 @@ class MetricCallback(TrainCallback):
 
         with torch.no_grad():
             for batch in context.val_dataloader:
-                loss = context.strategy(batch)
-                total_loss += loss.item()
+                loss_output = context.strategy(batch)
+                total_loss += loss_output["loss"].item()
                 num_batches += 1
 
         if context.world_size > 1 and dist.is_initialized():

astrai/trainer/train_context.py

@@ -38,6 +38,7 @@ class TrainContext:
     epoch: int = field(default=0)
     consumed_samples: int = field(default=0)
     loss: float = field(default=0.0)
+    metrics: Dict[str, float] = field(default_factory=dict)
     grad_norm: Optional[float] = field(default=None)
     grad_snr_tracker: GradSNRTracker = field(default_factory=GradSNRTracker)
     val_dataloader: Optional[DataLoader] = field(default=None)
@@ -221,6 +222,7 @@ class TrainContextBuilder:
                         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",

astrai/trainer/trainer.py

@@ -82,9 +82,10 @@ class Trainer:
                         break
                     with executor.accumulate(context.model):
                         self._call_callbacks("on_batch_begin", context)
-                        loss = context.strategy(batch)
-                        context.loss = loss.item()
-                        stand_loss = loss / executor.grad_accum_steps
+                        loss_output = context.strategy(batch)
+                        context.loss = loss_output["loss"].item()
+                        context.metrics = loss_output["metrics"]
+                        stand_loss = loss_output["loss"] / executor.grad_accum_steps
                         executor.backward(stand_loss)
                         context.consumed_samples += (
                             context.config.batch_per_device * context.world_size

csrc/kernels/attn_decode.cu

@@ -10,6 +10,9 @@ torch::Tensor attn_decode(
     double scale,
     int64_t layout
 ) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
     AttentionParams<bf16> p;
     attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
     TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1");
@@ -20,7 +23,7 @@ torch::Tensor attn_decode(
     p.o = (bf16*)O_view.data_ptr();
 
     alloc_split_partials(p);
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p);
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
     return O;
 }

csrc/kernels/attn_dispatchers.cuh

@@ -16,29 +16,20 @@
 #include "attn_paged_prefill_split_q_mma.cuh"
 #endif
 
-// Cached SM count — cudaDeviceGetAttribute is a host-side call that was
-// invoked on every decode/paged-decode launch.  Cache per-device so multi-GPU
-// setups with heterogeneous GPUs still get the right count, while the common
-// single-GPU path hits the cache after the first call.
-inline int get_sm_count() {
-    int dev = 0;
-    cudaGetDevice(&dev);
-    static int cached_dev = -1;
-    static int cached_count = 0;
-    if (dev != cached_dev) {
-        cudaDeviceGetAttribute(&cached_count, cudaDevAttrMultiProcessorCount, dev);
-        cached_dev = dev;
-    }
-    return cached_count;
-}
-
 // Split-KV: compute number of splits to fill all SMs for small-batch decode.
 // Caps splits so each split processes at least `min_tiles_per_split` tiles,
 // avoiding excessive loop/prologue overhead when tiles are small.
+//
+// Target total grid blocks (`TARGET_BLOCKS`) rather than scaling splits by SM
+// count.  Decode blocks are single-warp (32 threads) and a SM hosts ~11 of
+// them, so the old `2*sm/base` cap badly undersplit at large batch (B=16 got
+// 3 splits, optimal ~8).  Measured (L20, grid search): bandwidth saturates
+// near 256-512 total blocks; 512 minimizes worst-case latency across the
+// B x kv grid; more is pure oversplit overhead.
+constexpr int DECODE_TARGET_BLOCKS = 512;
 inline int compute_num_splits(int base_blocks, int tiles_total,
                                int min_tiles_per_split = 1) {
-    int sm_count = get_sm_count();
-    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
+    int n = (DECODE_TARGET_BLOCKS + base_blocks - 1) / base_blocks;
     int max_by_work = tiles_total / min_tiles_per_split;
     return std::max(1, std::min(n, std::min(max_by_work, MAX_SPLITS)));
 }
@@ -66,33 +57,33 @@ inline int compute_num_splits(int base_blocks, int tiles_total,
 
 #ifndef ASTRAI_NO_MMA
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_prefill_mma(AttentionParams<bf16>& p) {
+static inline void launch_prefill_mma(AttentionParams<bf16>& p, cudaStream_t stream) {
     constexpr int WARPS = 4;
     constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
     using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
     dim3 grid((p.q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS), p.q_head, p.batch);
     dim3 block(Traits::NUM_THREADS);
-    attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block>>>(p);
+    attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block, 0, stream>>>(p);
 }
 #endif
 
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_prefill_scalar(AttentionParams<bf16>& p) {
+static inline void launch_prefill_scalar(AttentionParams<bf16>& p, cudaStream_t stream) {
     constexpr int G = 8, ROWS = 32, P_BC = 32;
     dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
     dim3 block(G, ROWS);
-    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC, IsCausal, HasMask><<<grid, block>>>(p);
+    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC, IsCausal, HasMask><<<grid, block, 0, stream>>>(p);
 }
 
 template <int HEAD_DIM>
-static inline void dispatch_prefill(AttentionParams<bf16>& p) {
+static inline void dispatch_prefill(AttentionParams<bf16>& p, cudaStream_t stream) {
     bool is_causal = (p.causal_offset >= 0);
     bool has_mask = (p.use_mask && p.mask);
 
 #ifndef ASTRAI_NO_MMA
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_mma, HEAD_DIM, p);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_mma, HEAD_DIM, p, stream);
 #else
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_scalar, HEAD_DIM, p);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_scalar, HEAD_DIM, p, stream);
 #endif
 }
 
@@ -106,44 +97,44 @@ static inline void dispatch_prefill(AttentionParams<bf16>& p) {
 // enabling STAGES=2 (double-buffer) within the 32KB smem budget — eliminates
 // the 176-byte spill that STAGES=1+BC=32 suffered.
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_decode_mma(AttentionParams<bf16>& p, int group_size) {
+static inline void launch_decode_mma(AttentionParams<bf16>& p, int group_size, 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;
     constexpr int BC = 16;
     int tiles_total = (p.kv_len + BC - 1) / BC;
-    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total, 2);
+    p.num_splits = compute_num_splits(p.batch * p.kv_head * num_passes, tiles_total, 2);
     constexpr int STAGES = 2;
     using Traits = KernelTraits<HEAD_DIM, BC, 1, STAGES>;
     dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits);
-    attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask><<<grid, 32>>>(p);
+    attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask><<<grid, 32, 0, stream>>>(p);
 }
 #endif
 
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_decode_scalar(AttentionParams<bf16>& p, int group_size) {
+static inline void launch_decode_scalar(AttentionParams<bf16>& p, int group_size, cudaStream_t stream) {
     int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
     p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
     size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
     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);
-    attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem>>>(p);
+    attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem, stream>>>(p);
 }
 
 template <int HEAD_DIM>
-static inline void dispatch_decode(AttentionParams<bf16>& p) {
+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, launch_decode_mma, HEAD_DIM, p, group_size);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_mma, HEAD_DIM, p, group_size, stream);
 #else
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_scalar, HEAD_DIM, p, group_size);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_scalar, HEAD_DIM, p, group_size, stream);
 #endif
 
-    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim, 0, stream>>>(p);
 }
 
 // ======================================================================
@@ -152,7 +143,7 @@ static inline void dispatch_decode(AttentionParams<bf16>& p) {
 
 #ifndef ASTRAI_NO_MMA
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int) {
+static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, cudaStream_t stream) {
     int G = p.q_head / p.kv_head;
     constexpr int MAX_G = 16;
     constexpr int BC = 16;
@@ -162,34 +153,34 @@ static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int) {
     constexpr int STAGES = 2;
     using Traits = KernelTraits<HEAD_DIM, BC, 1, STAGES>;
     dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits);
-    paged_attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask> <<<grid, 32>>>(p);
+    paged_attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask> <<<grid, 32, 0, stream>>>(p);
 }
 #endif
 
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_decode_scalar(PagedAttentionParams<bf16>& p, int group_size) {
+static inline void launch_paged_decode_scalar(PagedAttentionParams<bf16>& p, int group_size, cudaStream_t stream) {
     int chunks_total = (p.max_seq_len + PDC_CHUNK - 1) / PDC_CHUNK;
     p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
     size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
     int g = min(group_size, 32);
     dim3 grid(p.batch * p.kv_head, 1, p.num_splits);
     dim3 block(32, g);
-    paged_attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem>>>(p);
+    paged_attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem, stream>>>(p);
 }
 
 template <int HEAD_DIM>
-static inline void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
+static inline void dispatch_paged_decode(PagedAttentionParams<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, launch_paged_decode_mma, HEAD_DIM, p, 0);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_mma, HEAD_DIM, p, stream);
 #else
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_scalar, HEAD_DIM, p, group_size);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_scalar, HEAD_DIM, p, group_size, stream);
 #endif
 
-    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim, 0, stream>>>(p);
 }
 
 // ======================================================================
@@ -198,35 +189,35 @@ static inline void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
 
 #ifndef ASTRAI_NO_MMA
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_prefill_mma(PagedAttentionParams<bf16>& p) {
+static inline void launch_paged_prefill_mma(PagedAttentionParams<bf16>& p, cudaStream_t stream) {
     constexpr int WARPS = 4;
     constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
     using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
     int max_q_tiles = (p.max_q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS);
     dim3 grid(max_q_tiles, p.q_head, p.batch);
     dim3 block(Traits::NUM_THREADS);
-    paged_attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block>>>(p);
+    paged_attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block, 0, stream>>>(p);
 }
 #endif
 
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_prefill_scalar(PagedAttentionParams<bf16>& p) {
+static inline void launch_paged_prefill_scalar(PagedAttentionParams<bf16>& p, cudaStream_t stream) {
     constexpr int G = 8, ROWS = 32, P_BC = 32;
     int max_q_tiles = (p.max_q_len + ROWS - 1) / ROWS;
     dim3 grid(max_q_tiles, p.q_head, p.batch);
     dim3 block(G, ROWS);
     paged_attn_prefill_split_q_kernel<HEAD_DIM, G, ROWS, P_BC, IsCausal, HasMask>
-        <<<grid, block>>>(p);
+        <<<grid, block, 0, stream>>>(p);
 }
 
 template <int HEAD_DIM>
-static inline void dispatch_paged_prefill(PagedAttentionParams<bf16>& p) {
+static inline void dispatch_paged_prefill(PagedAttentionParams<bf16>& p, cudaStream_t stream) {
     bool is_causal = (p.causal_offset >= 0);
     bool has_mask = (p.use_mask && p.mask);
 
 #ifndef ASTRAI_NO_MMA
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_prefill_mma, HEAD_DIM, p);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_prefill_mma, HEAD_DIM, p, stream);
 #else
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_prefill_scalar, HEAD_DIM, p);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_prefill_scalar, HEAD_DIM, p, stream);
 #endif
 }

csrc/kernels/attn_entry_utils.cuh

@@ -8,14 +8,14 @@
 using bf16 = __nv_bfloat16;
 
 // Dispatch head_dim: shared macro — avoids C++20 lambda template syntax.
-// Usage: DISPATCH_HEAD_DIM(hd, fn, arg)
-//   Expands to: fn<32>(arg); fn<64>(arg); etc.
-#define DISPATCH_HEAD_DIM(hd, fn, arg) \
+// Usage: DISPATCH_HEAD_DIM(hd, fn, args...)
+//   Expands to: fn<32>(args...); fn<64>(args...); etc.
+#define DISPATCH_HEAD_DIM(hd, fn, ...) \
     switch (hd) { \
-        case 32:  fn<32>(arg); break; \
-        case 64:  fn<64>(arg); break; \
-        case 128: fn<128>(arg); break; \
-        case 256: fn<256>(arg); break; \
+        case 32:  fn<32>(__VA_ARGS__); break; \
+        case 64:  fn<64>(__VA_ARGS__); break; \
+        case 128: fn<128>(__VA_ARGS__); break; \
+        case 256: fn<256>(__VA_ARGS__); break; \
         default: \
             TORCH_CHECK(false, "unsupported head_dim ", hd, \
                          " (supported: 32, 64, 128, 256)"); \

csrc/kernels/attn_paged_decode.cu

@@ -13,6 +13,9 @@ torch::Tensor attn_paged_decode(
     int64_t causal_offset,
     double scale
 ) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
     PagedAttentionParams<bf16> p;
     attn_pack_paged_decode_params(q, k_cache, v_cache,
                                    req_to_token, req_pool_indices, kv_indptr,
@@ -22,7 +25,7 @@ torch::Tensor attn_paged_decode(
     p.o = (bf16*)O.data_ptr();
 
     alloc_split_partials(p);
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p);
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
     return O;
 }

csrc/kernels/attn_paged_prefill.cu

@@ -14,6 +14,9 @@ torch::Tensor attn_paged_prefill(
     int64_t causal_offset,
     double scale
 ) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
     PagedAttentionParams<bf16> p;
     attn_pack_paged_prefill_params(q, k_cache, v_cache,
                                     req_to_token, req_pool_indices,
@@ -23,7 +26,7 @@ torch::Tensor attn_paged_prefill(
     auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
     p.o = (bf16*)O.data_ptr();
 
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_prefill, p);
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_prefill, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
     return O;
 }

csrc/kernels/attn_prefill.cu

@@ -10,6 +10,9 @@ torch::Tensor attn_prefill(
     double scale,
     int64_t layout
 ) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
     AttentionParams<bf16> p;
     attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
     TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
@@ -18,7 +21,7 @@ torch::Tensor attn_prefill(
     auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
     p.o = (bf16*)O_view.data_ptr();
 
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p);
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
     return O;
 }

csrc/kernels/rotary_emb.cu

@@ -49,6 +49,7 @@ torch::Tensor rotary_emb(
     torch::Tensor freqs_cis
 ) {
     const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+    auto stream = at::cuda::getCurrentCUDAStream();
 
     TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
     TORCH_CHECK(freqs_cis.is_cuda(), "freqs_cis must be on CUDA");
@@ -77,7 +78,7 @@ torch::Tensor rotary_emb(
     int block = 256;
     int grid = std::min((total + block - 1) / block, 1024);
 
-    rotary_emb_kernel<<<grid, block>>>(
+    rotary_emb_kernel<<<grid, block, 0, stream>>>(
         reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
         freqs_cis.data_ptr<float>(),
         reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),

csrc/tests/attn_decode_test.cu

@@ -1,185 +0,0 @@
-/*
-Pure-C test — uses shared dispatcher.
-nvcc -I csrc -arch=sm_89 -O3 \
-    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
-    csrc/tests/attn_decode_test.cu -o test && ./test
-*/
-
-#include "test_utils.cuh"
-#include "../kernels/attn_dispatchers.cuh"
-
-// Split-K scratch (torch-free)
-struct DecodeScratch {
-    float* o_part = nullptr;
-    float* ml_part = nullptr;
-};
-
-static void setup_scratch(AttentionParams<bf16>& p, DecodeScratch& sc) {
-    int max_splits = 32;
-    cudaMalloc(&sc.o_part, (size_t)p.batch * p.q_head * max_splits * p.head_dim * sizeof(float));
-    cudaMalloc(&sc.ml_part, (size_t)p.batch * p.q_head * max_splits * 2 * sizeof(float));
-}
-
-static void free_scratch(DecodeScratch& sc) {
-    cudaFree(sc.o_part); cudaFree(sc.ml_part);
-}
-
-// Warmed-up, CUDA-event timed sweep over the production decode MMA path.
-static void bench() {
-    const int cfgs[][5] = {
-        {1, 32, 4, 512, 128},
-        {1, 32, 4, 1024, 128},
-        {1, 32, 4, 2048, 128},
-        {1, 32, 4, 4096, 128},
-        {16, 32, 4, 2048, 128},
-        {32, 32, 4, 1024, 128},
-    };
-    const int WARMUP = 10, ITERS = 100;
-    printf("\n===== DECODE BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
-    print_bench_header();
-
-    for (int ci = 0; ci < 6; ci++) {
-        int B = cfgs[ci][0], Hq = cfgs[ci][1], Hk = cfgs[ci][2];
-        int sl = cfgs[ci][3], D = cfgs[ci][4];
-        size_t nQ = (size_t)B * Hq * D;
-        size_t nKV = (size_t)B * Hk * sl * D;
-
-        bf16 *dQ, *dK, *dV, *dO;
-        cudaMalloc(&dQ, nQ*2); cudaMalloc(&dK, nKV*2);
-        cudaMalloc(&dV, nKV*2); cudaMalloc(&dO, nQ*2);
-        size_t big = nQ > nKV ? nQ : nKV; bf16* tmp = new bf16[big];
-        for (size_t i = 0; i < nQ; i++)  tmp[i] = f2bf(randf());
-        cudaMemcpy(dQ, tmp, nQ*2, cudaMemcpyHostToDevice);
-        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
-        cudaMemcpy(dK, tmp, nKV*2, cudaMemcpyHostToDevice);
-        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
-        cudaMemcpy(dV, tmp, nKV*2, cudaMemcpyHostToDevice);
-        delete[] tmp;
-
-        AttentionParams<bf16> p;
-        p.batch = B; p.q_head = Hq; p.kv_head = Hk; p.q_len = 1; p.kv_len = sl;
-        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;
-
-        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); }); };
-        double flops = 4.0 * B * Hq * (double)sl * D;
-        double bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
-        BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops, bytes);
-
-        char cfg[64];
-        snprintf(cfg, sizeof(cfg),
-                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
-                 B, Hq, Hk, 1, sl, D, 0);
-        print_bench_row(cfg, r);
-
-        cudaFree(dQ); cudaFree(dK); cudaFree(dV); cudaFree(dO);
-        free_scratch(sc);
-    }
-}
-
-static int run_test(int B, int Hq, int Hk, int sl, int D, int causal) {
-    int gs = Hq / Hk;
-    printf("=== B=%d Hq=%d Hk=%d seq=%d D=%d gs=%d causal=%d ===\n",
-           B,Hq,Hk,sl,D,gs,causal);
-
-    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();
-    for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
-
-    bool* hMask=new bool[B*sl];
-    for (int i=0;i<B*sl;i++) hMask[i]=true;
-
-    bf16 *dQ,*dK,*dV,*dO,*tmp;
-    bool* dMask;
-    cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
-    cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
-    cudaMalloc(&dMask,B*sl);
-
-    tmp=new bf16[max(nQ,nKV)];
-    for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
-    cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
-    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
-    cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
-    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
-    cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
-    cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
-
-    AttentionParams<bf16> p;
-    p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
-    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;
-
-    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); });
-    cudaDeviceSynchronize();
-    double kms=now_ms()-t0;
-    cudaError_t err=cudaGetLastError();
-    if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
-
-    bf16* hOut=new bf16[nQ];
-    cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
-
-    float* ref=new float[nQ];
-    cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, causal ? 0 : -1);
-
-    float max_abs_err=0, max_rel_err=0;
-    for (size_t i=0;i<nQ;i++){
-        float err=fabsf(bf2f(hOut[i])-ref[i]);
-        if(err>max_abs_err) max_abs_err=err;
-        float rel=err/fmaxf(fabsf(ref[i]), 1e-8f);
-        if(rel>max_rel_err) max_rel_err=rel;
-    }
-    const float atol=0.01f, rtol=0.01f;
-    bool pass=true;
-    for (size_t i=0;i<nQ;i++){
-        float err=fabsf(bf2f(hOut[i])-ref[i]);
-        if (err > atol + rtol * fabsf(ref[i])) { pass=false; break; }
-    }
-    printf("kernel: %.3f ms  max_abs_err: %.6e  max_rel_err: %.6e  %s\n\n",
-           kms, max_abs_err, max_rel_err, pass?"PASS":"FAIL");
-
-    cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
-    free_scratch(sc);
-    delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
-
-    return pass ? 0 : 1;
-}
-
-int main() {
-    const int configs[][6] = {
-        {1, 2, 1, 64, 32, 0},
-        {1, 32, 4, 512, 128, 0},
-        {1, 32, 4, 1024, 128, 0},
-        {1, 32, 4, 512, 128, 1},
-    };
-    int n_cfgs = sizeof(configs) / sizeof(configs[0]);
-    int fail = 0;
-
-    for (int ci = 0; ci < n_cfgs; ci++) {
-        int B = configs[ci][0], Hq = configs[ci][1], Hk = configs[ci][2];
-        int sl = configs[ci][3], D = configs[ci][4], causal = configs[ci][5];
-        fail += run_test(B, Hq, Hk, sl, D, causal);
-        if (fail) break;
-    }
-
-    if (fail) {
-        printf("FAILED\n");
-        return fail;
-    }
-    printf("All tests passed!\n");
-    bench();
-    return 0;
-}

csrc/tests/attn_paged_test.cu

@@ -1,6 +1,7 @@
 // Compile:
 //   nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
-//        --extra-device-vectorization csrc/tests/attn_paged_test.cu \
+//        --extra-device-vectorization -Xcompiler -fopenmp \
+//        csrc/tests/attn_paged_test.cu \
 //        -o /tmp/test_paged && /tmp/test_paged
 
 #include <cstring>
@@ -24,6 +25,7 @@ static void cpu_paged_decode_ref(
     for (int b = 0; b < B; b++) {
         int seq_len = kv_indptr[b + 1] - kv_indptr[b];
         int64_t req_idx = req_pool_indices[b];
+        #pragma omp parallel for schedule(dynamic)
         for (int h = 0; h < Hq; h++) {
             int kv_h = h / n_rep;
             float mv = -INFINITY, sv = 0.0f;
@@ -74,9 +76,10 @@ static void cpu_paged_prefill_ref(
         int q_len = qo_indptr[b + 1] - qo_indptr[b];
         int causal_off = seq_len - q_len;
         int64_t req_idx = req_pool_indices[b];
+        #pragma omp parallel for collapse(2) schedule(dynamic)
         for (int h = 0; h < Hq; h++) {
-            int kv_h = h / n_rep;
             for (int qi = 0; qi < q_len; qi++) {
+                int kv_h = h / n_rep;
                 float mv = -INFINITY, sv = 0.0f;
                 float accum[256] = {0.0f};
                 int lim = causal ? min(seq_len, causal_off + qi + 1) : seq_len;
@@ -106,6 +109,19 @@ static void cpu_paged_prefill_ref(
     }
 }
 
+// ---- paged validation table (kernel vs CPU ref, abs error only) ----
+inline void print_paged_header() {
+    printf("%-58s | %11s | %6s\n",
+           "config", "max_err", "result");
+    printf("----------------------------------------------------------------"
+           "--------------------------------\n");
+}
+
+inline void print_paged_row(const char* cfg, float max_err, bool pass) {
+    printf("%-58s | %11.3e | %s\n",
+           cfg, max_err, pass ? "PASS" : "FAIL");
+}
+
 // ======================================================================
 // DECODE TEST
 // ======================================================================
@@ -123,10 +139,9 @@ static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
     int pool_size = B * max_ctx;
     int num_reqs = B + 4;
 
-    printf("DECODE B=%d Hq=%d Hkv=%d D=%d seqs=[", B, Hq, Hkv, HEAD_DIM);
-    for (int b = 0; b < B; b++) printf("%d%s", seq_lens[b], b < B-1 ? "," : "");
-    printf("] causal=%d ... ", causal);
-    fflush(stdout);
+    char cfg[80];
+    snprintf(cfg, sizeof(cfg), "DECODE B=%d Hq=%d Hkv=%d D=%d max_sl=%d causal=%d",
+             B, Hq, Hkv, HEAD_DIM, max_sl, causal);
 
     size_t sz_q  = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
     size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
@@ -211,7 +226,7 @@ static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
     p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
     p.o = 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); });
+    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p, 0); });
     cudaDeviceSynchronize();
 
     bf16* h_o_bf = (bf16*)malloc(sz_q);
@@ -228,8 +243,7 @@ static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
         if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
     }
 
-    if (pass) printf("PASS (max_err=%.4e)\n", max_err);
-    else      printf("FAIL (max_err=%.4e)\n", max_err);
+    print_paged_row(cfg, max_err, pass);
 
     free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
     free(h_kvi); free(h_q_f); free(h_k_f); free(h_v_f);
@@ -254,8 +268,9 @@ static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
     int pool_size = B * max_ctx;
     int num_reqs = B + 4;
 
-    printf("DECODE-MASK B=%d Hq=%d Hkv=%d D=%d max_sl=%d ... ", B, Hq, Hkv, HEAD_DIM, max_sl);
-    fflush(stdout);
+    char cfg[80];
+    snprintf(cfg, sizeof(cfg), "DECODE-MASK B=%d Hq=%d Hkv=%d D=%d max_sl=%d",
+             B, Hq, Hkv, HEAD_DIM, max_sl);
 
     size_t sz_q  = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
     size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
@@ -346,7 +361,7 @@ static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
     p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
     p.o = 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); });
+    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p, 0); });
     cudaDeviceSynchronize();
 
     bf16* h_o_bf = (bf16*)malloc(sz_q);
@@ -363,8 +378,7 @@ static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
         if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
     }
 
-    if (pass) printf("PASS (max_err=%.4e)\n", max_err);
-    else      printf("FAIL (max_err=%.4e)\n", max_err);
+    print_paged_row(cfg, max_err, pass);
 
     free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
     free(h_kvi); free(h_mask); free(h_q_f); free(h_k_f); free(h_v_f);
@@ -393,12 +407,9 @@ static int run_prefill_test(int B, int Hq, int Hkv,
     int pool_size = B * max_ctx;
     int num_reqs = B + 4;
 
-    printf("PREFILL B=%d Hq=%d Hkv=%d D=%d q_lens=[", B, Hq, Hkv, HEAD_DIM);
-    for (int b = 0; b < B; b++) printf("%d%s", q_lens[b], b < B-1 ? "," : "");
-    printf("] kv_lens=[");
-    for (int b = 0; b < B; b++) printf("%d%s", kv_lens[b], b < B-1 ? "," : "");
-    printf("] causal=%d ... ", causal);
-    fflush(stdout);
+    char cfg[80];
+    snprintf(cfg, sizeof(cfg), "PREFILL B=%d Hq=%d Hkv=%d D=%d max_sl=%d causal=%d",
+             B, Hq, Hkv, HEAD_DIM, max_sl, causal);
 
     size_t sz_q  = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
     size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
@@ -486,7 +497,7 @@ static int run_prefill_test(int B, int Hq, int Hkv,
     p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
     p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
 
-    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p); });
+    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p, 0); });
     cudaDeviceSynchronize();
 
     bf16* h_o_bf = (bf16*)malloc(sz_q);
@@ -503,8 +514,7 @@ static int run_prefill_test(int B, int Hq, int Hkv,
         if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
     }
 
-    if (pass) printf("PASS (max_err=%.4e)\n", max_err);
-    else      printf("FAIL (max_err=%.4e)\n", max_err);
+    print_paged_row(cfg, max_err, pass);
 
     free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
     free(h_kvi); free(h_qoi); free(h_q_f); free(h_k_f); free(h_v_f);
@@ -527,7 +537,9 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
     int pool_size = B * max_ctx;
     int num_reqs = B + 4;
 
-    printf("PREFILL-MASK Hq=%d Hkv=%d D=%d q_len=%d ... ", Hq, Hkv, HEAD_DIM, q_len);
+    char cfg[80];
+    snprintf(cfg, sizeof(cfg), "PREFILL-MASK Hq=%d Hkv=%d D=%d q_len=%d",
+             Hq, Hkv, HEAD_DIM, q_len);
     fflush(stdout);
 
     size_t sz_q  = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
@@ -620,7 +632,7 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
     p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
     p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
 
-    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p); });
+    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p, 0); });
     cudaDeviceSynchronize();
 
     bf16* h_o_bf = (bf16*)malloc(sz_q);
@@ -637,8 +649,7 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
         if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
     }
 
-    if (pass) printf("PASS (max_err=%.4e)\n", max_err);
-    else      printf("FAIL (max_err=%.4e)\n", max_err);
+    print_paged_row(cfg, max_err, pass);
 
     free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
     free(h_kvi); free(h_qoi); free(h_mask); free(h_q_f); free(h_k_f); free(h_v_f);
@@ -710,15 +721,12 @@ static void bench_decode(int B, int Hq, int Hkv, int seq_len) {
     p.o = 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); });
+        dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p, 0); });
     };
     // 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.
     double flops = 4.0 * B * Hq * (double)seq_len * HEAD_DIM;
-    // HBM: K+V read (Q/O negligible for decode).
-    size_t nKV = (size_t)B * Hkv * seq_len * HEAD_DIM;
-    double bytes = 2.0 * nKV * sizeof(bf16);
-    BenchResult r = bench_kernel(launch, 10, 100, flops, bytes);
+    BenchResult r = bench_kernel(launch, 3, 10, flops);
 
     char cfg[64];
     snprintf(cfg, sizeof(cfg), "DEC B=%2d Hq=%2d Hk=%d kv=%4d D=%3d",
@@ -791,7 +799,7 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau
     p.o = 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); });
+        dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p, 0); });
     };
     // FLOPs = 2 * (QK^T + PV) = 4 * effective_qk_pairs * Hq * D.
     // Non-causal: effective = q_len * kv_len.
@@ -807,11 +815,7 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau
         eff_kv = (double)q_len * kv_len;
     }
     double flops = 4.0 * B * Hq * eff_kv * HEAD_DIM;
-    // HBM: Q read + K read + V read + O write.
-    size_t nKV = (size_t)B * Hkv * kv_len * HEAD_DIM;
-    size_t nQ  = (size_t)total_q * Hq * HEAD_DIM;
-    double bytes = (2.0 * nQ + 2.0 * nKV) * sizeof(bf16);
-    BenchResult r = bench_kernel(launch, 10, 100, flops, bytes);
+    BenchResult r = bench_kernel(launch, 3, 10, flops);
 
     char cfg[80];
     snprintf(cfg, sizeof(cfg), "PRE B=%d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d c=%d",
@@ -827,7 +831,8 @@ int main() {
     int fail = 0;
 
     // ===== DECODE TESTS =====
-    printf("=== Paged Decode Tests ===\n\n");
+    printf("=== Paged Decode Tests ===\n");
+    print_paged_header();
     fail += run_decode_test<128>(1, 32, 4, 512, 0, 1);
     fail += run_decode_test<128>(1, 32, 4, 1024, 0, 2);
     fail += run_decode_test<128>(4, 32, 4, 512, 0, 3);
@@ -848,7 +853,8 @@ int main() {
     if (fail) { printf("\nFAILED decode tests\n"); return fail; }
 
     // ===== PREFILL TESTS =====
-    printf("\n=== Paged Prefill Tests ===\n\n");
+    printf("\n=== Paged Prefill Tests ===\n");
+    print_paged_header();
     // Single request, pure prefill (q_len == kv_len)
     {
         std::vector<int> ql = {512};

csrc/tests/attn_prefill_test.cu

@@ -1,169 +0,0 @@
-/*
-Pure-C test — uses shared dispatcher.
-nvcc -I csrc -arch=sm_89 -O3 \
-    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
-    csrc/tests/attn_prefill_test.cu -o test && ./test
-*/
-
-#include "test_utils.cuh"
-#include "../kernels/attn_dispatchers.cuh"
-
-// Warmed-up, CUDA-event timed throughput sweep over the production MMA path.
-static void bench() {
-    const int cfgs[][7] = {
-        {1,32,4,512,512,128,0},
-        {1,32,4,1024,1024,128,0},
-        {1,32,4,2048,2048,128,0},
-        {1,32,4,2048,2048,128,1},
-        {4,32,4,2048,2048,128,1},
-        {1,32,4,4096,4096,128,1},
-    };
-    int n = sizeof(cfgs)/sizeof(cfgs[0]);
-    const int WARMUP = 10, ITERS = 50;
-    printf("\n===== PREFILL BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
-    printf("%-46s | %10s | %10s | %10s\n",
-           "config", "latency", "bandwidth", "throughput");
-    printf("---------------------------------------------------------------"
-           "----------------------------\n");
-
-    for (int ci = 0; ci < n; ci++) {
-        int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
-        int ql=cfgs[ci][3], kl=cfgs[ci][4], D=cfgs[ci][5], causal=cfgs[ci][6];
-        size_t nQ=(size_t)B*Hq*ql*D, nKV=(size_t)B*Hk*kl*D;
-
-        bf16 *dQ,*dK,*dV,*dO,*tmp;
-        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
-        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
-        size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
-        for (size_t i=0;i<nQ;i++)  tmp[i]=f2bf(randf());
-        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
-        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
-        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
-        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
-        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
-
-        AttentionParams<bf16> p;
-        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
-        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;
-
-        auto launch = [&]() { dispatch_by_head_dim(D, [&]<int H>() { dispatch_prefill<H>(p); }); };
-        for (int i=0;i<WARMUP;i++) launch();
-        cudaDeviceSynchronize();
-        cudaError_t err=cudaGetLastError();
-        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
-
-        cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
-        cudaEventRecord(s);
-        for (int i=0;i<ITERS;i++) launch();
-        cudaEventRecord(e); cudaEventSynchronize(e);
-        float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
-
-        double flops = 4.0*B*Hq*(double)ql*kl*D;
-        if (causal) flops *= 0.5;
-        double tflops = flops/(ms*1e-3)/1e12;
-        double bytes = 2.0 * (2.0*nQ + 2.0*nKV);
-        double gbps = bytes/(ms*1e-3)/1e9;
-
-        char cfg[64];
-        snprintf(cfg, sizeof(cfg),
-                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
-                 B,Hq,Hk,ql,kl,D,causal);
-        printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
-               cfg, ms, gbps, tflops);
-
-        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
-        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
-    }
-}
-
-static int run_test(int B, int Hq, int Hk, int ql, int kl, int D, int causal) {
-    printf("=== B=%d Hq=%d Hk=%d q=%d kv=%d D=%d causal=%d ===\n",
-           B,Hq,Hk,ql,kl,D,causal);
-
-    size_t nQ = B*Hq*ql*D, nKV = B*Hk*kl*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();
-    for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
-
-    bf16 *dQ,*dK,*dV,*dO,*tmp;
-    cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
-    cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
-    tmp=new bf16[max(nQ,nKV)];
-    for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
-    cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
-    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
-    cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
-    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
-    cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
-
-    AttentionParams<bf16> p;
-    p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
-    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;
-
-    double t0=now_ms();
-    dispatch_by_head_dim(D, [&]<int H>() { dispatch_prefill<H>(p); });
-    cudaDeviceSynchronize();
-    double kms=now_ms()-t0;
-    cudaError_t err=cudaGetLastError();
-    if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
-
-    bf16* hOut=new bf16[nQ];
-    cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
-
-    float* ref=new float[nQ];
-    cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal ? 0 : -1);
-
-    float max_abs_err=0, max_rel_err=0;
-    for (size_t i=0;i<nQ;i++) {
-        float err=fabsf(bf2f(hOut[i])-ref[i]);
-        if(err>max_abs_err) max_abs_err=err;
-        float rel=err/fmaxf(fabsf(ref[i]), 1e-8f);
-        if(rel>max_rel_err) max_rel_err=rel;
-    }
-    const float atol=0.01f, rtol=0.01f;
-    bool pass=true;
-    for (size_t i=0;i<nQ;i++) {
-        float err=fabsf(bf2f(hOut[i])-ref[i]);
-        if (err > atol + rtol * fabsf(ref[i])) { pass=false; break; }
-    }
-    printf("kernel: %.3f ms  max_abs_err: %.6e  max_rel_err: %.6e  %s\n\n",
-           kms, max_abs_err, max_rel_err, pass?"PASS":"FAIL");
-
-    cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
-    delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
-
-    return pass ? 0 : 1;
-}
-
-int main() {
-    const int configs[][7] = {
-        {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
-        {1,4,2,256,256,128,1},   // causal
-    };
-    int n_configs = sizeof(configs) / sizeof(configs[0]);
-    int fail = 0;
-
-    for (int ci = 0; ci < n_configs; ci++) {
-        int B=configs[ci][0], Hq=configs[ci][1], Hk=configs[ci][2];
-        int ql=configs[ci][3], kl=configs[ci][4], D=configs[ci][5];
-        int causal=configs[ci][6];
-        fail += run_test(B, Hq, Hk, ql, kl, D, causal);
-        if (fail) break;
-    }
-
-    if (fail) {
-        printf("FAILED\n");
-        return fail;
-    }
-    printf("All tests passed!\n");
-    bench();
-    return 0;
-}

csrc/tests/attn_test.cu

@@ -0,0 +1,346 @@
+/*
+Pure-C test — uses shared dispatcher.  Combines the decode (split-KV) and
+prefill (split-Q) correctness checks + benchmarks into one binary.
+nvcc -I csrc -arch=sm_89 -O3 \
+    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
+    -Xcompiler -fopenmp csrc/tests/attn_test.cu -o test && ./test
+*/
+
+#include "test_utils.cuh"
+#include "../kernels/attn_dispatchers.cuh"
+
+// Split-K scratch (torch-free)
+struct DecodeScratch {
+    float* o_part = nullptr;
+    float* ml_part = nullptr;
+};
+
+static void setup_scratch(AttentionParams<bf16>& p, DecodeScratch& sc) {
+    int max_splits = 32;
+    cudaMalloc(&sc.o_part, (size_t)p.batch * p.q_head * max_splits * p.head_dim * sizeof(float));
+    cudaMalloc(&sc.ml_part, (size_t)p.batch * p.q_head * max_splits * 2 * sizeof(float));
+}
+
+static void free_scratch(DecodeScratch& sc) {
+    cudaFree(sc.o_part); cudaFree(sc.ml_part);
+}
+
+// ======================================================================
+// DECODE
+// ======================================================================
+
+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();
+    for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
+
+    bool* hMask=new bool[B*sl];
+    for (int i=0;i<B*sl;i++) hMask[i]=true;
+
+    bf16 *dQ,*dK,*dV,*dO,*tmp;
+    bool* dMask;
+    cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+    cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+    cudaMalloc(&dMask,B*sl);
+
+    tmp=new bf16[max(nQ,nKV)];
+    for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
+    cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
+    cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
+    cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+    cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
+
+    AttentionParams<bf16> p;
+    p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
+    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;
+
+    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); });
+    cudaDeviceSynchronize();
+    (void)t0;
+    cudaError_t err=cudaGetLastError();
+    if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
+
+    bf16* hOut=new bf16[nQ];
+    cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
+
+    float* ref=new float[nQ];
+    cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, causal ? 0 : -1);
+
+    float max_abs_err=0, max_rel_err=0;
+    for (size_t i=0;i<nQ;i++){
+        float err=fabsf(bf2f(hOut[i])-ref[i]);
+        if(err>max_abs_err) max_abs_err=err;
+        float rel=err/fmaxf(fabsf(ref[i]), 1e-8f);
+        if(rel>max_rel_err) max_rel_err=rel;
+    }
+    const float atol=0.01f, rtol=0.01f;
+    bool pass=true;
+    for (size_t i=0;i<nQ;i++){
+        float err=fabsf(bf2f(hOut[i])-ref[i]);
+        if (err > atol + rtol * fabsf(ref[i])) { pass=false; break; }
+    }
+    char cfg[64];
+    snprintf(cfg, sizeof(cfg), "B=%2d Hq=%2d Hk=%d seq=%4d D=%3d causal=%d",
+             B, Hq, Hk, sl, D, causal);
+    print_test_row(cfg, max_abs_err, max_rel_err, pass);
+
+    cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
+    free_scratch(sc);
+    delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
+
+    return pass ? 0 : 1;
+}
+
+static void bench_decode() {
+    const int cfgs[][5] = {
+        {1, 32, 4, 512, 128},
+        {1, 32, 4, 1024, 128},
+        {1, 32, 4, 2048, 128},
+        {1, 32, 4, 4096, 128},
+        {16, 32, 4, 2048, 128},
+        {32, 32, 4, 1024, 128},
+    };
+    const int WARMUP = 3, ITERS = 10;
+    printf("\n===== DECODE BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
+    print_bench_header();
+
+    for (int ci = 0; ci < 6; ci++) {
+        int B = cfgs[ci][0], Hq = cfgs[ci][1], Hk = cfgs[ci][2];
+        int sl = cfgs[ci][3], D = cfgs[ci][4];
+        size_t nQ = (size_t)B * Hq * D;
+        size_t nKV = (size_t)B * Hk * sl * D;
+
+        bf16 *dQ, *dK, *dV, *dO;
+        cudaMalloc(&dQ, nQ*2); cudaMalloc(&dK, nKV*2);
+        cudaMalloc(&dV, nKV*2); cudaMalloc(&dO, nQ*2);
+        size_t big = nQ > nKV ? nQ : nKV; bf16* tmp = new bf16[big];
+        for (size_t i = 0; i < nQ; i++)  tmp[i] = f2bf(randf());
+        cudaMemcpy(dQ, tmp, nQ*2, cudaMemcpyHostToDevice);
+        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
+        cudaMemcpy(dK, tmp, nKV*2, cudaMemcpyHostToDevice);
+        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
+        cudaMemcpy(dV, tmp, nKV*2, cudaMemcpyHostToDevice);
+        delete[] tmp;
+
+        AttentionParams<bf16> p;
+        p.batch = B; p.q_head = Hq; p.kv_head = Hk; p.q_len = 1; p.kv_len = sl;
+        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;
+
+        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); }); };
+        double flops = 4.0 * B * Hq * (double)sl * D;
+        BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops);
+
+        char cfg[64];
+        snprintf(cfg, sizeof(cfg),
+                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
+                 B, Hq, Hk, 1, sl, D, 0);
+        print_bench_row(cfg, r);
+
+        cudaFree(dQ); cudaFree(dK); cudaFree(dV); cudaFree(dO);
+        free_scratch(sc);
+    }
+}
+
+// ======================================================================
+// PREFILL
+// ======================================================================
+
+static int run_prefill_test(int B, int Hq, int Hk, int ql, int kl, int D, int causal) {
+    size_t nQ = B*Hq*ql*D, nKV = B*Hk*kl*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();
+    for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
+
+    bf16 *dQ,*dK,*dV,*dO,*tmp;
+    cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+    cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+    tmp=new bf16[max(nQ,nKV)];
+    for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
+    cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
+    cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
+    cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+
+    AttentionParams<bf16> p;
+    p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
+    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;
+
+    double t0=now_ms();
+    dispatch_by_head_dim(D, [&]<int H>() { dispatch_prefill<H>(p, 0); });
+    cudaDeviceSynchronize();
+    (void)t0;
+    cudaError_t err=cudaGetLastError();
+    if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
+
+    bf16* hOut=new bf16[nQ];
+    cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
+
+    float* ref=new float[nQ];
+    cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal ? 0 : -1);
+
+    float max_abs_err=0, max_rel_err=0;
+    for (size_t i=0;i<nQ;i++) {
+        float err=fabsf(bf2f(hOut[i])-ref[i]);
+        if(err>max_abs_err) max_abs_err=err;
+        float rel=err/fmaxf(fabsf(ref[i]), 1e-8f);
+        if(rel>max_rel_err) max_rel_err=rel;
+    }
+    const float atol=0.01f, rtol=0.01f;
+    bool pass=true;
+    for (size_t i=0;i<nQ;i++) {
+        float err=fabsf(bf2f(hOut[i])-ref[i]);
+        if (err > atol + rtol * fabsf(ref[i])) { pass=false; break; }
+    }
+    char cfg[64];
+    snprintf(cfg, sizeof(cfg), "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
+             B, Hq, Hk, ql, kl, D, causal);
+    print_test_row(cfg, max_abs_err, max_rel_err, pass);
+
+    cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
+    delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
+
+    return pass ? 0 : 1;
+}
+
+static void bench_prefill() {
+    const int cfgs[][7] = {
+        {1,32,4,512,512,128,0},
+        {1,32,4,1024,1024,128,0},
+        {1,32,4,2048,2048,128,0},
+        {1,32,4,2048,2048,128,1},
+        {4,32,4,2048,2048,128,1},
+        {1,32,4,4096,4096,128,1},
+    };
+    int n = sizeof(cfgs)/sizeof(cfgs[0]);
+    const int WARMUP = 3, ITERS = 10;
+    printf("\n===== PREFILL BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
+    print_bench_header();
+
+    for (int ci = 0; ci < n; ci++) {
+        int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
+        int ql=cfgs[ci][3], kl=cfgs[ci][4], D=cfgs[ci][5], causal=cfgs[ci][6];
+        size_t nQ=(size_t)B*Hq*ql*D, nKV=(size_t)B*Hk*kl*D;
+
+        bf16 *dQ,*dK,*dV,*dO,*tmp;
+        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+        size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
+        for (size_t i=0;i<nQ;i++)  tmp[i]=f2bf(randf());
+        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+
+        AttentionParams<bf16> p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
+        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;
+
+        auto launch = [&]() { dispatch_by_head_dim(D, [&]<int H>() { dispatch_prefill<H>(p, 0); }); };
+        for (int i=0;i<WARMUP;i++) launch();
+        cudaDeviceSynchronize();
+        cudaError_t err=cudaGetLastError();
+        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
+
+        cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
+        cudaEventRecord(s);
+        for (int i=0;i<ITERS;i++) launch();
+        cudaEventRecord(e); cudaEventSynchronize(e);
+        float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
+
+        double flops = 4.0*B*Hq*(double)ql*kl*D;
+        if (causal) flops *= 0.5;
+        double tflops = flops/(ms*1e-3)/1e12;
+        BenchResult r{ms, tflops};
+
+        char cfg[64];
+        snprintf(cfg, sizeof(cfg),
+                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
+                 B,Hq,Hk,ql,kl,D,causal);
+        print_bench_row(cfg, r);
+
+        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
+        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
+    }
+}
+
+// ======================================================================
+// MAIN
+// ======================================================================
+
+int main() {
+    int fail = 0;
+
+    // ---- DECODE ----
+    {
+        const int configs[][6] = {
+            {1, 2, 1, 64, 32, 0},
+            {1, 32, 4, 512, 128, 0},
+            {1, 32, 4, 1024, 128, 0},
+            {1, 32, 4, 512, 128, 1},
+        };
+        int n_cfgs = sizeof(configs) / sizeof(configs[0]);
+        printf("=== DECODE TESTS ===\n");
+        print_test_header();
+        for (int ci = 0; ci < n_cfgs; ci++) {
+            int B = configs[ci][0], Hq = configs[ci][1], Hk = configs[ci][2];
+            int sl = configs[ci][3], D = configs[ci][4], causal = configs[ci][5];
+            fail += run_decode_test(B, Hq, Hk, sl, D, causal);
+            if (fail) break;
+        }
+        if (fail) { printf("FAILED decode tests\n"); return fail; }
+        bench_decode();
+    }
+
+    // ---- PREFILL ----
+    {
+        const int configs[][7] = {
+            {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
+            {1,4,2,256,256,128,1},   // causal
+        };
+        int n_configs = sizeof(configs) / sizeof(configs[0]);
+        printf("\n=== PREFILL TESTS ===\n");
+        print_test_header();
+        for (int ci = 0; ci < n_configs; ci++) {
+            int B=configs[ci][0], Hq=configs[ci][1], Hk=configs[ci][2];
+            int ql=configs[ci][3], kl=configs[ci][4], D=configs[ci][5];
+            int causal=configs[ci][6];
+            fail += run_prefill_test(B, Hq, Hk, ql, kl, D, causal);
+            if (fail) break;
+        }
+        if (fail) { printf("FAILED prefill tests\n"); return fail; }
+        bench_prefill();
+    }
+
+    printf("\nAll tests passed!\n");
+    return 0;
+}

csrc/tests/test_utils.cuh

@@ -29,19 +29,18 @@ inline double now_ms() {
 
 struct BenchResult {
     float ms;
-    double gbps;
     double tflops;
 };
 
 template <typename Fn>
 BenchResult bench_kernel(Fn launch, int warmup, int iters,
-                         double flops, double bytes) {
+                         double flops) {
     for (int i = 0; i < warmup; i++) launch();
     cudaDeviceSynchronize();
     cudaError_t err = cudaGetLastError();
     if (err != cudaSuccess) {
         printf("CUDA error before bench: %s\n", cudaGetErrorString(err));
-        return {0, 0, 0};
+        return {0, 0};
     }
 
     cudaEvent_t s, e;
@@ -52,19 +51,33 @@ BenchResult bench_kernel(Fn launch, int warmup, int iters,
     float ms = 0; cudaEventElapsedTime(&ms, s, e); ms /= iters;
     cudaEventDestroy(s); cudaEventDestroy(e);
 
-    return {ms, bytes / (ms * 1e-3) / 1e9, flops / (ms * 1e-3) / 1e12};
+    return {ms, flops / (ms * 1e-3) / 1e12};
 }
 
 inline void print_bench_header() {
-    printf("%-46s | %10s | %10s | %10s\n",
-           "config", "latency", "bandwidth", "throughput");
+    printf("%-46s | %10s | %10s\n",
+           "config", "latency", "TFLOP/s");
     printf("---------------------------------------------------------------"
            "----------------------------\n");
 }
 
 inline void print_bench_row(const char* cfg, const BenchResult& r) {
-    printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
-           cfg, r.ms, r.gbps, r.tflops);
+    printf("%-46s | %7.4f ms | %6.2f\n",
+           cfg, r.ms, r.tflops);
+}
+
+// ---- validation table (kernel vs CPU reference) ----
+inline void print_test_header() {
+    printf("%-46s | %11s | %11s | %6s\n",
+           "config", "max_abs_err", "max_rel_err", "result");
+    printf("----------------------------------------------------------------"
+           "----------------------------\n");
+}
+
+inline void print_test_row(const char* cfg, float max_abs_err,
+                           float max_rel_err, bool pass) {
+    printf("%-46s | %11.3e | %11.3e | %s\n",
+           cfg, max_abs_err, max_rel_err, pass ? "PASS" : "FAIL");
 }
 
 template <int... Ds>
@@ -135,9 +148,10 @@ static void cpu_attention_ref(
     float scale = 1.0f / sqrtf((float)D);
     int n_rep = Hq / Hk;
     for (int b = 0; b < B; b++) {
+        #pragma omp parallel for collapse(2) schedule(dynamic)
         for (int h = 0; h < Hq; h++) {
-            int kv_h = h / n_rep;
             for (int qi = 0; qi < q_len; qi++) {
+                int kv_h = h / n_rep;
                 float mv = -INFINITY, sv = 0.0f;
                 float accum[256] = {0.0f};
                 int lim = kv_len;

docs/README-zh-CN.md

@@ -62,6 +62,8 @@
 
 **1. 安装**
 
+AstrAI 需要 Python 3.12+，并精确固定 PyTorch 版本为 `2.11.0`。训练、`scripts/tools/generate.py`、生成式评估和生成演示需要 CUDA；CPU 支持仅适用于提供明确 CPU 设备路径的组件，例如 HTTP 服务和直接打分评估。
+
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
@@ -138,7 +140,7 @@ curl http://localhost:8000/v1/chat/completions \
 # 下载模型权重（运行演示前必需）
 python scripts/demo/download.py                      # model → params/
 
-# 交互式流式聊天（多轮对话，保持历史记录）
+# 单轮交互式流式提示循环（不保留对话历史）
 python scripts/demo/stream_chat.py
 # 在 >> 后输入消息，输入 !exit 退出
 
@@ -189,7 +191,7 @@ docker run --gpus all -v /path/to/data:/data -it astrai:latest
 # Docker Compose（GPU，默认）
 docker compose up -d
 
-# Docker Compose（仅 CPU）
+# Docker Compose CPU 服务配置（不支持仅限 CUDA 的生成脚本和演示）
 docker compose --profile cpu up -d
 ```
 
@@ -239,7 +241,7 @@ SSE 流式格式、错误码和统计端点详见[推理文档](guides/inference
 
 ### 贡献
 
-我们欢迎贡献！请参阅[贡献指南](../../CONTRIBUTING.md)了解详情。
+我们欢迎贡献！请参阅[贡献指南](../CONTRIBUTING.md)了解详情。
 
 1. Fork 本仓库。
 2. 创建功能分支。
@@ -256,10 +258,10 @@ SSE 流式格式、错误码和统计端点详见[推理文档](guides/inference
 
 ### 许可证
 
-本项目采用 [GPL-3.0 许可证](../../LICENSE)。
+本项目采用 [GPL-3.0 许可证](../LICENSE)。
 
 ---
 
 <div align="center">
   <em>专为高性能与易用性设计的轻量级 Transformer 框架。</em>
-</div>
+</div>

docs/developer/architecture.md

@@ -4,7 +4,7 @@
 
 - [Class Diagram](#class-diagram) — Full Mermaid class diagram across 10+ namespaces
 - [Module Overview](#module-overview) — Component inventory per module
-- [Design Patterns](#design-patterns) — 13 documented patterns with classes
+- [Design Patterns](#design-patterns) — 15 documented patterns with classes
 - [Core Relationships](#core-relationships) — 11 key inter-component relationships
 
 ## Class Diagram
@@ -49,6 +49,11 @@ classDiagram
             +Optional[int] n_shared_experts
             +Optional[int] n_activated_experts
             +Optional[str] topk_method
+            +Optional[int] moe_intermediate_size
+            +Optional[int] shared_expert_intermediate_size
+            +bool norm_topk_prob
+            +int decoder_sparse_step
+            +Optional[List[int]] mlp_only_layers
         }
 
         class EncoderConfig {
@@ -63,6 +68,7 @@ classDiagram
             +Optional[int] num_attention_heads
             +Optional[int] num_key_value_heads
             +Optional[bool] use_qk_norm
+            +Optional[bool] use_gated_attention
             +str ffn_type
             +Optional[dict] rope_scaling
             +Optional[str] pooling_type
@@ -114,22 +120,25 @@ classDiagram
             +Dataset dataset
             +Callable optimizer_fn
             +Callable scheduler_fn
+            +Optional[str] optimizer_name
+            +Dict[str, Any] optimizer_hyperparameters
             +int n_epoch
             +int batch_per_device
             +int grad_accum_steps
             +Optional[float] max_grad_norm
             +list gradient_checkpointing_modules
+            +Optional[str] compile_mode
             +int start_epoch
             +int start_samples
             +str ckpt_dir
             +int ckpt_interval
-            +str log_dir
             +List[str] metrics
             +Optional[LoRAConfig] lora
             +int random_seed
             +int num_workers
             +Optional[int] prefetch_factor
             +bool pin_memory
+            +Optional[Callable] collate_fn
             +int nprocs
             +str backend
             +str master_addr
@@ -140,6 +149,7 @@ classDiagram
             +Optional[float] val_split
             +int val_step
             +float neftune_alpha
+            +float moe_aux_loss_coef
             +str parallel_mode
             +int rollout_interval
             +float rollout_temperature
@@ -149,7 +159,6 @@ classDiagram
             +Optional[Callable] reward_model_fn
             +dict executor_kwargs
             +dict extra_kwargs
-            +validate()
         }
 
     }
@@ -205,10 +214,6 @@ classDiagram
             -_fetch_record_key(key, index) Tensor
         }
 
-        class H5Store {
-            +load(path)
-        }
-
         class MmapStore {
             +List _mmap_refs
             +load(path)
@@ -260,11 +265,15 @@ classDiagram
     }
 
     namespace model {
-        class AutoModel {
-            +BaseModelConfig config
+        class ModelFactory {
             +Dict _entries
             +register(name) decorator
             +get_component_class(name) Type
+        }
+
+        class AutoModel {
+            <<nn.Module>>
+            +BaseModelConfig config
             +from_pretrained(path, disable_random_init, strict) nn.Module
             +save_pretrained(save_directory)
             +to(*args, **kwargs) Self
@@ -299,7 +308,13 @@ classDiagram
             +RMSNorm input_norm
             +nn.Module mlp        # MLP or DeepSeekMoE via FFNFactory
             +RMSNorm post_attention_norm
-            +forward(x, rotary_emb, attention_mask, kv_cache) Tensor
+            +forward(x, rotary_emb, attention_mask, kv_cache, is_causal) DecoderOutput
+        }
+
+        class DecoderOutput {
+            <<TypedDict>>
+            +Tensor hidden_states
+            +Optional[Tensor] aux_loss
         }
 
         class GQA {
@@ -314,7 +329,7 @@ classDiagram
             +Linear q_proj, k_proj, v_proj, o_proj
             +Linear gate  # only if use_gated_attention
             +RMSNorm q_norm, k_norm  # only if use_qk_norm
-            +forward(x, rotary_emb, attn_mask, kv_cache) Tensor
+            +forward(x, rotary_emb, attn_mask, kv_cache, is_causal) Tensor
         }
 
         class MLA {
@@ -334,12 +349,18 @@ classDiagram
             +Linear gate  # only if use_gated_attention
             +RMSNorm kv_norm
             +RMSNorm q_norm, k_norm  # only if use_qk_norm
-            +forward(x, rotary_emb, attn_mask, kv_cache) Tensor
+            +forward(x, rotary_emb, attn_mask, kv_cache, is_causal) Tensor
         }
 
         class MLP {
             +Linear up, gate, down
-            +forward(x) Tensor
+            +forward(x) FFNOutput
+        }
+
+        class FFNOutput {
+            <<TypedDict>>
+            +Tensor hidden_states
+            +Optional[Tensor] aux_loss
         }
 
         class DeepSeekMoE {
@@ -351,7 +372,7 @@ classDiagram
             +Linear router
             +ModuleList shared_experts
             +ModuleList routed_experts
-            +forward(x) Tensor
+            +forward(x) FFNOutput
         }
 
         class AttnFactory {
@@ -380,9 +401,8 @@ classDiagram
             +int max_len
             +float base
             +Optional[Dict] rope_scaling
-            +Tensor cos_table
-            +Tensor sin_table
-            +forward(x, position_ids=None) Tuple[Tensor, Tensor]
+            +Tensor freqs_cis
+            +forward(x, position_ids=None) Tensor
         }
 
         class Embedding {
@@ -486,10 +506,6 @@ classDiagram
             +save(output_dir, domain, shard_idx, tensors)
         }
 
-        class H5Writer {
-            +save(output_dir, domain, shard_idx, tensors)
-        }
-
         class Pipeline {
             +PipelineConfig config
             +List[str] paths
@@ -559,7 +575,7 @@ classDiagram
         class Trainer {
             +TrainConfig train_config
             +List[TrainCallback] callbacks
-            +train(resume_dir)
+            +train(param_path=None, resume=False)
             -_get_default_callbacks() List[TrainCallback]
         }
 
@@ -576,13 +592,17 @@ classDiagram
             +int epoch
             +int consumed_samples
             +float loss
-            +float grad_norm
+            +Dict[str, float] metrics
+            +Optional[float] grad_norm
+            +GradSNRTracker grad_snr_tracker
             +DataLoader val_dataloader
-            +float val_loss
+            +Optional[float] val_loss
             +int world_size
             +int rank
             +dict kwargs
-            +optimizer_step() int
+            +stop_requested (property) bool
+            +optimizer_step (property) int
+            +request_stop()
         }
 
         class TrainContextBuilder {
@@ -594,11 +614,22 @@ classDiagram
         class BaseStrategy {
             +Callable model
             +Optional[BaseExecutor] executor
-            +Optional[Callable] model_fn
+            +float moe_aux_loss_coef
             +dict extra_kwargs
             +str device
-            +__call__(batch) Tensor
+            +__call__(batch) LossOutput
             +compute_loss(batch) Tensor
+            +compute_loss_output(batch) LossOutput
+            +supports_online() bool
+            +set_rollout_runner(runner)
+            +prepare_from_rollout(result) Dict
+            +on_optimizer_step()
+        }
+
+        class LossOutput {
+            <<TypedDict>>
+            +Tensor loss
+            +Dict[str, float] metrics
         }
 
         class StrategyFactory {
@@ -636,9 +667,12 @@ classDiagram
 
         class RawRollout {
             +Tensor prompts
+            +Tensor prompt_mask
             +Tensor responses
             +Tensor response_mask
             +Tensor logprobs_old
+            +List[str] prompt_texts
+            +List[List[str]] response_texts
         }
 
         class RolloutResult {
@@ -647,10 +681,18 @@ classDiagram
 
         class BaseRewardModel {
             <<abstract>>
-            +score(prompts, responses) Tensor
+            +score(List[str] prompts, List[List[str]] responses) Tensor
         }
 
         class RolloutGenerator {
+            +InferenceScheduler scheduler
+            +int max_tokens
+            +int group_size
+            +float temperature
+            +int top_k
+            +float top_p
+            +float frequency_penalty
+            +int rep_window
             +generate(batch) RawRollout
         }
 
@@ -740,7 +782,7 @@ classDiagram
         }
 
         class MetricCallback {
-            +Path log_dir
+            +Path ckpt_dir
             +int save_interval
             +List[str] metrics
             +int val_step
@@ -764,9 +806,9 @@ classDiagram
             +nn.Module model
             +AutoTokenizer tokenizer
             +InferenceScheduler scheduler
-            +generate(prompt, stream, max_tokens, temperature, top_p, top_k) Union[Generator, str, List[str]]
+            +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) AsyncGenerator
+            +generate_async(prompt, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window) AsyncGenerator
             +get_stats() Dict
             +shutdown()
         }
@@ -774,18 +816,18 @@ classDiagram
         class Executor {
             +AutoModel model
             +AutoTokenizer tokenizer
-            +KVCache page_cache
+            +PagePool kv_cache
             +Optional[str] device
             +Optional[torch.dtype] dtype
             +execute_prefill(tasks, prompt_len, start_pos)
-            +execute_decode(tasks) List[int]
+            +execute_decode(tasks, return_logprobs=False) Union[List[int], List[Tuple[int, float]]]
         }
 
         class InferenceScheduler {
-            +KVCache _page_cache
+            +PagePool _cache
             +Executor _executor
             +TaskManager _task_mgr
-            +bool _running
+            +Event _stop_event
             +Thread _loop_thread
             +int max_seq_len
             +str device
@@ -795,6 +837,7 @@ classDiagram
             +start()
             +stop()
             +get_stats() Dict
+            +run_batch(prompt_ids_list, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window, return_logprobs) Union[List[List[int]], List[Tuple[List[int], List[float]]]]
         }
 
         class Allocator {
@@ -816,16 +859,6 @@ classDiagram
             +record(page_idx, token_ids, logical_page_idx)
         }
 
-        class PagePool {
-            -Allocator _alloc
-            -PrefixCache _prefix
-            +alloc() int
-            +free(idx)
-            +inc_ref(idx)
-            +lookup(token_ids) List[int]
-            +record(page_idx, token_ids, logical_page_idx)
-        }
-
         class KVStorage {
             +int size
             +Tensor k_buffer
@@ -852,8 +885,7 @@ classDiagram
             +Tensor seq_lens
             +Tensor out_cache_loc
             +int max_len
-            +Optional[Tensor] page_table
-            +Optional[Tensor] decode_mask
+            +Optional[Tensor] kv_indptr
         }
 
         class PagePool {
@@ -878,6 +910,8 @@ classDiagram
         +float temperature
         +float top_p
         +int top_k
+        +float frequency_penalty
+        +int rep_window
         +TaskStatus status
         +List output_ids
         +int input_tokens
@@ -923,27 +957,29 @@ classDiagram
             +float top_p
             +float temperature
             +Optional[int] max_tokens
+            +float frequency_penalty
+            +int rep_window
             +bool stream
         }
 
         class BaseSamplingStrategy {
             <<abstract>>
-            +apply(logits, filter_value) Tensor
+            +apply(logits, filter_value, input_ids, input_mask) Tensor
         }
 
         class TemperatureStrategy {
             +float temperature
-            +apply(logits, filter_value) Tensor
+            +apply(logits, filter_value, input_ids, input_mask) Tensor
         }
 
         class TopKStrategy {
             +int top_k
-            +apply(logits, filter_value) Tensor
+            +apply(logits, filter_value, input_ids, input_mask) Tensor
         }
 
         class TopPStrategy {
             +float top_p
-            +apply(logits, filter_value) Tensor
+            +apply(logits, filter_value, input_ids, input_mask) Tensor
         }
 
         class FrequencyPenaltyStrategy {
@@ -953,8 +989,8 @@ classDiagram
 
         class SamplingPipeline {
             +List[BaseSamplingStrategy] strategies
-            +apply(logits, filter_value) Tensor
-            +sample(logits, filter_value) Tensor
+            +apply(logits, filter_value, input_ids, input_mask) Tensor
+            +sample(logits, filter_value, input_ids, input_mask, return_logprobs) Union[Tensor, Tuple[Tensor, Tensor]]
         }
 
         class StreamDecoder {
@@ -1029,7 +1065,7 @@ classDiagram
             <<abstract>>
             +prepare(request, engine) Tuple[str, GenContext, List[str]]
             +format_stream_start(ctx) List[str]
-            +format_chunk(token) List[str]
+            +format_chunk(token, **kwargs) List[str]
             +format_stream_end(ctx, stop) List[str]
             +format_response(ctx, content, stop) Dict
         }
@@ -1037,7 +1073,7 @@ classDiagram
         class OpenAIResponseBuilder {
             +prepare(request, engine) Tuple
             +format_stream_start(ctx) List[str]
-            +format_chunk(token) List[str]
+            +format_chunk(token, **kwargs) List[str]
             +format_stream_end(ctx, stop) List[str]
             +format_response(ctx, content, stop) Dict
         }
@@ -1045,7 +1081,7 @@ classDiagram
         class AnthropicResponseBuilder {
             +prepare(request, engine) Tuple
             +format_stream_start(ctx) List[str]
-            +format_chunk(token) List[str]
+            +format_chunk(token, **kwargs) List[str]
             +format_stream_end(ctx, stop) List[str]
             +format_response(ctx, content, stop) Dict
         }
@@ -1153,10 +1189,13 @@ classDiagram
 
         class BaseExecutor {
             +GradientState gradient_state
-            +prepare(model_fn, optimizer_fn, scheduler_fn, before_wrap) tuple
+            +prepare(model_fn, optimizer_fn, scheduler_fn, before_wrap, after_wrap) tuple
             +accumulate(model) context manager
             +backward(loss)
             +unwrap_model(model) dict
+            +checkpoint_context(model) context manager
+            +clip_grad_norm(model, max_norm) float
+            +use_distributed (property) bool
             +sync_gradients (property) bool
             +grad_accum_steps (property) int
         }
@@ -1173,7 +1212,8 @@ classDiagram
         class FSDPExecutor {
             -_prepare_model(model) nn.Module
             -_no_sync(model) context manager
-            +unwrap_model(model) dict
+            +unwrap_model(model) Optional[dict]
+            +clip_grad_norm(model, max_norm) float
         }
 
         class ExecutorFactory {
@@ -1182,33 +1222,6 @@ classDiagram
             +create(parallel_mode, **kwargs) BaseExecutor
         }
 
-        class ParallelModel {
-            +dist.ProcessGroup process_group
-            +int rank
-            +int world_size
-        }
-
-        class ColumnParallelLinear {
-            +int in_features
-            +int out_features
-            +int out_features_per_rank
-            +bool gather_results
-            +Parameter weight
-            +Optional[Parameter] bias
-            +forward(x) Tensor
-            +load_state_dict(state_dict)
-        }
-
-        class RowParallelLinear {
-            +int in_features
-            +int out_features
-            +int in_features_per_rank
-            +bool reduce_results
-            +Parameter weight
-            +Optional[Parameter] bias
-            +forward(x) Tensor
-            +load_state_dict(state_dict)
-        }
     }
 
     %% Relationships — UML notation: <|-- generalization, *-- composition, o-- aggregation, --> association, ..> dependency
@@ -1230,11 +1243,8 @@ classDiagram
     BaseDataset <|-- SFTDataset
     BaseDataset <|-- DPODataset
     BaseDataset <|-- GRPODataset
-    Store <|-- H5Store
     Store <|-- MmapStore
     Store <|-- JsonlStore
-    H5Store --|> Streamable
-    H5Store --|> Recordable
     MmapStore --|> Streamable
     MmapStore --|> Recordable
     JsonlStore --|> Streamable
@@ -1243,8 +1253,6 @@ classDiagram
     BaseSamplingStrategy <|-- TopKStrategy
     BaseSamplingStrategy <|-- TopPStrategy
     BaseSamplingStrategy <|-- FrequencyPenaltyStrategy
-    ParallelModel <|-- RowParallelLinear
-    ParallelModel <|-- ColumnParallelLinear
     AutoModel <|-- AutoRegressiveLM
     AutoModel <|-- EmbeddingEncoder
     BaseConfig <|-- BaseModelConfig
@@ -1255,7 +1263,7 @@ classDiagram
     BaseConfig <|-- PipelineConfig
     BaseModelConfig <|-- AutoRegressiveLMConfig
     BaseModelConfig <|-- EncoderConfig
-    BaseFactory <|-- AutoModel
+    BaseFactory <|-- ModelFactory
     BaseFactory <|-- AttnFactory
     BaseFactory <|-- FFNFactory
     BaseFactory <|-- DatasetFactory
@@ -1286,7 +1294,6 @@ classDiagram
     PositionIdStrategy <|-- DocResetPositionId
     PositionIdStrategy <|-- ContinuousPositionId
     StoreWriter <|-- BinWriter
-    StoreWriter <|-- H5Writer
     RawRollout <|-- RolloutResult
     LaunchStrategy <|-- TorchrunStrategy
     LaunchStrategy <|-- LocalStrategy
@@ -1317,8 +1324,6 @@ classDiagram
     %% --- Aggregation (weak ownership) ---
     AutoModel o-- BaseModelConfig
     AutoTokenizer o-- ChatTemplate
-    PagePool o-- Allocator
-    PagePool o-- PrefixCache
     Trainer o-- TrainCallback
     TrainContext o-- BaseStrategy
     TrainContext o-- BaseScheduler
@@ -1352,11 +1357,12 @@ classDiagram
     FFNFactory ..> DeepSeekMoE : creates
     DecoderBlock ..> AttnFactory : uses
     DecoderBlock ..> FFNFactory : uses
-    StoreFactory ..> H5Store : creates
     StoreFactory ..> MmapStore : creates
     StoreFactory ..> JsonlStore : creates
     ConfigFactory ..> AutoRegressiveLMConfig : creates
     ConfigFactory ..> EncoderConfig : creates
+    ModelFactory ..> AutoRegressiveLM : creates
+    ModelFactory ..> EmbeddingEncoder : creates
     ExecutorFactory ..> NoneExecutor : creates
     ExecutorFactory ..> DDPExecutor : creates
     ExecutorFactory ..> FSDPExecutor : creates
@@ -1399,10 +1405,10 @@ classDiagram
 | Module | Components | Description |
 |--------|------------|-------------|
 | **astrai.config** | BaseConfig, BaseModelConfig, AutoRegressiveLMConfig, EncoderConfig, ConfigFactory, TrainConfig, PipelineConfig, InputConfig, ProcessingConfig, OutputConfig | Configuration management (to_dict/from_dict, to_file/from_file) |
-| **astrai.preprocessing** | SectionRenderer, BaseMaskBuilder, MaskBuilderFactory, SectionedMaskBuilder, SingleOutputMaskBuilder, MultiOutputMaskBuilder, Pipeline, TokenizeTransform, PackingStrategy, PackingStrategyFactory, SimplePacking, BFDPacking, BFDSplitPacking, PositionIdStrategy, PositionIdStrategyFactory, NoPositionId, DocResetPositionId, ContinuousPositionId, StoreWriter, StoreWriterFactory, BinWriter, H5Writer | Declarative JSON-driven data preprocessing |
-| **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, Store, Streamable, Recordable, H5Store, MmapStore, JsonlSource, JsonlStore, StoreFactory, RDSampler, DatasetFactory | Dataset loading and management |
+| **astrai.preprocessing** | SectionRenderer, BaseMaskBuilder, MaskBuilderFactory, SectionedMaskBuilder, SingleOutputMaskBuilder, MultiOutputMaskBuilder, Pipeline, TokenizeTransform, PackingStrategy, PackingStrategyFactory, SimplePacking, BFDPacking, BFDSplitPacking, PositionIdStrategy, PositionIdStrategyFactory, NoPositionId, DocResetPositionId, ContinuousPositionId, StoreWriter, StoreWriterFactory, BinWriter | Declarative JSON-driven data preprocessing |
+| **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, Store, Streamable, Recordable, MmapStore, JsonlSource, JsonlStore, StoreFactory, RDSampler, DatasetFactory | Dataset loading and management |
 | **astrai.serialization** | Checkpoint | Model serialization |
-| **astrai.model** | AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, LoRAConfig, LoRALinear, RotaryEmbedding, Embedding | Neural network model |
+| **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, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, PrefixCache, 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 |
@@ -1415,7 +1421,7 @@ classDiagram
 
 | Pattern | Classes | Purpose |
 |---------|---------|---------|
-| **Factory** | `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory`, `StoreFactory`, `ConfigFactory`, `ExecutorFactory`, `MaskBuilderFactory`, `StoreWriterFactory`, `PackingStrategyFactory`, `PositionIdStrategyFactory`, `ToolParserFactory` | Decorator-based component creation |
+| **Factory** | `ModelFactory`, `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory`, `StoreFactory`, `ConfigFactory`, `ExecutorFactory`, `MaskBuilderFactory`, `StoreWriterFactory`, `PackingStrategyFactory`, `PositionIdStrategyFactory`, `ToolParserFactory` | Decorator-based component creation |
 | **Registry** | `BaseFactory` | Component registration |
 | **Strategy** | `SEQStrategy`, `SFTStrategy`, `DPOStrategy`, `GRPOStrategy` | Training strategy switching |
 | **Strategy (Sampling)** | `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `SamplingPipeline` | Composable logit transformations |
@@ -1427,9 +1433,9 @@ classDiagram
 | **Strategy (Attention)** | `AttentionBackend`, `TorchNativeBackend`, `CudaBackend` | 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`, `H5Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support |
+| **Storage** | `Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support |
 | **Producer-Consumer** | `InferenceScheduler`, `Task`, queues | Continuous batching |
-| **AutoModel Registry** | `AutoModel`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading |
+| **Model Registry** | `ModelFactory`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading |
 
 ## Core Relationships
 
@@ -1439,10 +1445,10 @@ classDiagram
 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).
 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
-7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (H5Store/MmapStore/JsonlStore) loads data with explicit `_length` and multi-segment `_data`
-8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata (rank-0 only), extra state saved as `{key}.pt`
+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`
 9. **Scheduler**: `SchedulerFactory` creates `CosineScheduler`/`SGDRScheduler`/`WSDScheduler`
 10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
 11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers
 
-> Document Update Time: 2026-07-31
+> Document Update Time: 2026-08-02

docs/developer/cuda_kernels.md

@@ -117,23 +117,22 @@ Each `csrc/tests/*.cu` file has the `nvcc` compile command in its header comment
 ```bash
 nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
      --ptxas-options=-O3,-v --extra-device-vectorization \
-     csrc/tests/attn_decode_test.cu -o /tmp/test && /tmp/test
+     -Xcompiler -fopenmp csrc/tests/attn_test.cu -o /tmp/test && /tmp/test
 ```
 
 Test files:
-- `attn_decode_test.cu` — basic decode kernel
-- `attn_paged_decode_test.cu` — paged decode kernel
-- `attn_prefill_test.cu` — prefill kernel
+- `attn_test.cu` — decode + prefill kernels (correctness tables + benchmarks)
+- `attn_paged_test.cu` — paged decode/prefill kernels
 
 ## Benchmarks
 
 Hardware: NVIDIA L20 (sm_89, 46 GB), CUDA 12.8, driver 570.86.
 
-Reproduce:
+Reproduce (decode + prefill in `attn_test.cu`, paged in `attn_paged_test.cu`):
 ```bash
 nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
      --ptxas-options=-O3,-v --extra-device-vectorization \
-     csrc/tests/attn_<name>_test.cu -o /tmp/test && /tmp/test
+     -Xcompiler -fopenmp csrc/tests/attn_test.cu -o /tmp/test && /tmp/test
 ```
 
 ## Known Optimization Targets
@@ -165,9 +164,8 @@ csrc/
 │   └── attn_warp_utils.cuh              # Warp-level utilities
 └── tests/
     ├── test_utils.cuh           # Shared test utilities
-    ├── attn_decode_test.cu      # Decode kernel test
-    ├── attn_paged_decode_test.cu # Paged decode test
-    └── attn_prefill_test.cu     # Prefill kernel test
+    ├── attn_test.cu             # Decode + prefill kernels
+    └── attn_paged_test.cu       # Paged decode/prefill kernels
 ```
 
 Compiled `.so` files are placed in `astrai/extension/lib/`, separate from Python source files.

docs/developer/dataflow.md

@@ -14,26 +14,30 @@ This document describes the data pipeline: from raw text to model input tensors.
 ## Overview
 
 ```
-JSONL Lines → Pipeline (mask builder) → Tokenized Tensors
-                                              ↓
-                                      .h5 or .bin storage
-                                              ↓
-                                      Store.load()
+JSON / JSONL Records → Pipeline (mask builder) → Tokenized Tensors
+                                               ↓
+                                           .bin storage
+                                               ↓
+                                       Store.load()
                                               ↓
                                       Store.fetch(begin, end, keys)
                                               ↓
-                                      BaseDataset.__getitem__(idx)
-                                              ↓
-                                      Sampler → DataLoader → Training / Inference
+                                       Dataset.__getitem__(idx)
+                                               ↓
+                                       RDSampler → DataLoader → Training
 ```
 
 ## Data Preparation
 
-Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or binary (`.bin` + `meta.json`) files with keyed tensor groups.
+The offline `Pipeline` accepts `.jsonl` records and `.json` files containing one
+object or a list of objects. It tokenizes them and writes binary shards (`.bin`
+plus `meta.json`) with keyed tensor groups. Binary is the only registered output
+writer; the pipeline cannot emit JSONL.
 
 ### Tokenization
 
-The `Pipeline` reads JSONL lines, applies the mask builder (see [Preprocessing](../guides/preprocessing.md)), and produces flat token sequences:
+The `Pipeline` reads JSON/JSONL records, applies the mask builder (see
+[Preprocessing](../guides/preprocessing.md)), and produces token sequences:
 
 ```python
 # Per JSONL line: messages → chat template → token IDs + loss mask
@@ -42,84 +46,114 @@ loss_mask = [0, 0, 0, 1, 1, 1, 1, 1, 1]        # 0=masked, 1=train
 # Stored as flat tensors, packed with other lines by packing strategy
 ```
 
-The output `meta.json` records the storage format, key names, dtype, total token count, and tensor shapes for each shard.
+For default single-output preprocessing, the stored keys are `sequence` and
+`position_ids`, plus `loss_mask` when masking is required. Packing is supported
+for single-output data with a `sequence` key. Shard flushing counts the primary
+flat sequence for each record: `sequence` in single-output mode, otherwise the
+first flat source output.
+
+The exact shard `meta.json` schema is a top-level mapping from key to tensor
+metadata. It does not contain a storage-format or total-token field:
+
+```json
+{
+  "sequence": {"shape": [123456], "dtype": "int32"},
+  "loss_mask": {"shape": [123456], "dtype": "bool"},
+  "position_ids": {"shape": [123456], "dtype": "int32"}
+}
+```
+
+Record-aware binary data may also include `"offsets": [0, ...]` inside a key's
+metadata, but the preprocessing `BinWriter` currently does not write offsets.
 
 ### Format Detection
 
 `detect_format(load_path)` inspects the path:
 
-- If `load_path` is a file: checks suffix — `.h5`/`.hdf5` → `"h5"`, `.jsonl` → `"jsonl"`, unknown suffix raises `ValueError`
-- If `load_path` is a directory: recursively globs for `*.h5`/`*.hdf5` files → `"h5"`, `*.bin` + `**/meta.json` → `"bin"`, or `*.jsonl` + `dataset_config.json` → `"jsonl"`
+- 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.
 
 ### Store Backends
 
 Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:
 
 ```
-StoreFactory.create("h5")    → H5Store
 StoreFactory.create("bin")   → MmapStore
 StoreFactory.create("jsonl") → JsonlStore
 ```
 
-All three inherit `Store` (base, owns `_data`/`_cum`/`_offsets`/`_normalize`) plus the `Streamable` and `Recordable` mixins, so every backend supports both `fetch(begin, end, keys)` (stream) and `fetch_record(index, keys)` (record) APIs.
-
-**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage. `segments_are_records=True` — each `data_i` dataset is one record.
+Both stores inherit `Store` and compose the `Streamable` and `Recordable`
+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**: On-the-fly tokenization of raw JSONL files at load time. Requires a `dataset_config.json` alongside the `.jsonl` files following the same `PipelineConfig` schema with an additional `tokenizer_path` field. Two modes: eager (default, applies `TokenizeTransform` to all records at load) and lazy (`processor=fn` given, defers tokenisation to `fetch_record` — used by DPO/GRPO).
+**JsonlStore**: Reads a `.jsonl` file or the sorted top-level `*.jsonl` files in
+a directory. Eager transform selection uses the first available route:
 
-All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based stream indexing) + `Store._offsets[Dict[str, List[int]]]` (per-record offsets for record-mode indexing). Nested keys (GRPO `responses`/`masks` as `List[List[Tensor]]`) are stored as-is and excluded from both bookkeepings — they are only accessed record-by-record.
+1. An explicit `transform=` argument.
+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.
+
+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.
+
+Eager-loaded stores normalize tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for stream indexing) + `Store._offsets[Dict[str, List[int]]]` (per-record offsets for record indexing). Nested JSONL keys such as GRPO `responses`/`masks` are kept as record values and excluded from stream bookkeeping. Lazy DPO instead retains raw records and processes them in `fetch_record`.
 
 ## Data Keys by Training Type
 
 | Type | Storage Keys | Access Mode |
 |------|-------------|-------------|
-| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) | stream (`fetch`) |
+| `seq` | `sequence`, `position_ids` by default (`SEQDataset` consumes only `sequence`) | stream (`fetch`) |
 | `sft` | `sequence`, `loss_mask`, `position_ids` | stream (`fetch`) |
 | `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` | record (`fetch_record`) |
 | `grpo` | `prompts`, `responses`, `masks`, `rewards` | record (`fetch_record`) |
 
+Offline `.bin` output from DPO/GRPO preprocessing is not currently loadable for
+training. DPO shards are written without record offsets, while GRPO response
+groups are flattened without preserving record/group boundaries. Supported raw
+routes are eager JSONL for SEQ/SFT and automatic lazy JSONL for DPO. GRPO
+requires a caller-built, already-loaded record store.
+
 ## Dataset Architecture
 
 ```
-DatasetFactory.load(
-    train_type, load_path=None, window_size=0, stride=None,
-    storage_type=None, tokenizer_path=None,
-    max_len=2048, store=None
-)
-  → BaseDataset.load(load_path, storage_type=None)
-    → detect_format(load_path)
-    → StoreFactory.create(storage_type)
-    → Store.load(load_path)
-      → _normalize(raw)  # base Store, shared by both backends
-        → Store._data[Dict[str, List[Tensor]]]
-          + _cum[Dict[str, List[int]]]   (stream mode)
-          + _offsets[Dict[str, List[int]]]  (record mode)
+DatasetFactory.load(...)
+  → detect_format(load_path)
+  → optionally build dpo_processor for raw JSONL
+  → StoreFactory.create(storage_type, window_size, stride)
+  → Store.load(load_path, transform=... or processor=...)
+  → DatasetFactory.create(train_type, store=store)
 
 Stream datasets (SEQ/SFT):
   BaseDataset.__getitem__(idx)
-    → get_index(idx) → [begin, end)
+    → Store.sample_window(idx) → [begin, end)
     → Store.fetch(begin, end, keys) → Tensor / Dict[str, Tensor]
 
-Record datasets (DPO/GRPO via RecordDataset):
-  RecordDataset.__getitem__(idx)
+Record datasets (DPO/GRPO):
+  DPODataset/GRPODataset.__getitem__(idx)
     → Store.fetch_record(idx, keys) → Tensor / Dict[str, Tensor]
 ```
 
-Class hierarchy: `BaseDataset` ← `SEQDataset` / `SFTDataset` (stream); `BaseDataset` ← `RecordDataset` ← `DPODataset` / `GRPODataset` (record).
+Class hierarchy: `BaseDataset` is the direct base of `SEQDataset`, `SFTDataset`,
+`DPODataset`, and `GRPODataset`. There is no `RecordDataset` class.
 
 `window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`, optional). Only meaningful for stream datasets — record datasets ignore both. `storage_type` defaults to `None` (auto-detect via `detect_format`).
 
-`tokenizer_path` triggers lazy on-the-fly tokenisation for record datasets on raw JSONL (DPO builds a `dpo_processor`; SEQ/SFT/pre-tokenised backends ignore it). `store` (pre-built `Store`) bypasses `load_path`/`storage_type`/`tokenizer_path` entirely — the caller controls Store construction.
+For raw JSONL, `tokenizer_path` builds the lazy processor only for DPO. For
+SEQ/SFT it is forwarded to `JsonlStore` so the built-in eager `messages`
+transform can be selected when no `dataset_config.json` exists. GRPO receives no
+automatic processor. A pre-built `store` bypasses path, format, tokenizer,
+window, and stride setup entirely.
 
 `Store.fetch(begin, end, keys)` (stream mode, on `Streamable`): accepts a single key (`str`) returning a `Tensor`, or a list of keys returning `Dict[str, Tensor]`. Internally uses `bisect` across multi-segment tensors. Raises `RuntimeError("Store not loaded")` if called before `load()`.
 
-`Store.fetch_record(index, keys)` (record mode, on `Recordable`): same key API. Uses `_offsets[key]` when present (bin layout with per-record offsets), otherwise indexes `_data[key]` directly (H5/JSONL where each segment is one record).
+`Store.fetch_record(index, keys)` (record mode, on `Recordable`): same key API. Uses `_offsets[key]` when present for binary record layouts; otherwise it indexes per-record JSONL tensors directly.
 
 ## Sampler
 
-`ResumableDistributedSampler` supports checkpoint-aware distributed sampling:
+`RDSampler` supports checkpoint-aware distributed sampling:
 
 - Tracks `start_epoch` / `start_iter` for resume
 - Shuffle via `torch.Generator(seed + epoch)`

docs/developer/internals.md

@@ -41,7 +41,12 @@ RoPE embeds position into Q/K vectors via complex rotation:
 
 $$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
 
-`RotaryEmbedding` pre-computes `cos_table` and `sin_table` (f32, `[max_len, dim/2]`). `forward()` returns a `(cos, sin)` tuple indexed by `position_ids`. `apply_rotary_emb` applies the rotation: during training it uses torch complex multiply (autograd-compatible); during inference it auto-dispatches to a fused CUDA kernel when available. The key property is that the dot product $q_i^T k_j$ depends only on the relative position $i - j$, not the absolute positions.
+`RotaryEmbedding` pre-computes a complex `freqs_cis` buffer. `forward()` returns
+a tensor indexed by `position_ids`. `apply_rotary_emb` applies the rotation:
+during training it uses torch complex multiply (autograd-compatible); during
+inference it auto-dispatches to a fused CUDA kernel when available. The key
+property is that the dot product $q_i^T k_j$ depends only on the relative
+position $i - j$, not the absolute positions.
 
 **Critical for inference**: RoPE is applied **before** KV cache write, not after. If applied after caching, position encoding drift occurs because cached K/V would have stale rotation factors.
 
@@ -51,13 +56,13 @@ $$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{
 
 Next-token cross-entropy with optional label smoothing:
 
-$$ L_{\text{PT}} = -\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta) $$
+$$ L_{\text{PT}} = -\frac{1}{T}\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta) $$
 
 ### SFT (Supervised Fine-Tuning)
 
 Masked cross-entropy (`ignore_index=-100`) over response tokens only:
 
-$$ L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta) $$
+$$ L_{\text{SFT}} = -\frac{1}{L}\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta) $$
 
 Prompt tokens are masked out via `loss_mask`; only response tokens contribute to the loss.
 
@@ -81,6 +86,14 @@ Where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the per-toke
 
 Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`.
 
+### MoE Load Balancing
+
+MoE layers add a differentiable load-balancing term based on mean router probabilities and top-k expert assignment frequency. The training objective is:
+
+$$ L = L_{\text{task}} + \lambda_{\text{MoE}} L_{\text{aux}} $$
+
+`TrainConfig.moe_aux_loss_coef` controls $\lambda_{\text{MoE}}$ (default `0.01`). The unweighted and weighted auxiliary losses are logged separately.
+
 ## Training Loop Internals
 
 Two-level loop: **epoch** → **batch**. Optimizer step fires every `grad_accum_steps` batches.
@@ -90,11 +103,12 @@ on_train_begin
   model.train()
   on_epoch_begin
     for batch in dataloader:
-      on_batch_begin
       with executor.accumulate(model):
-        loss = strategy.compute_loss(batch)
-        context.loss = loss.item()
-        stand_loss = loss / executor.grad_accum_steps
+        on_batch_begin
+        loss_output = strategy(batch)
+        context.loss = loss_output["loss"].item()
+        context.metrics = loss_output["metrics"]
+        stand_loss = loss_output["loss"] / executor.grad_accum_steps
         executor.backward(stand_loss)
         context.consumed_samples += (
             context.config.batch_per_device * context.world_size
@@ -104,6 +118,7 @@ on_train_begin
         if executor.sync_gradients:
           on_optimizer_step
           optimizer.step()
+          strategy.on_optimizer_step()
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()
@@ -112,21 +127,23 @@ on_train_end
 ```
 
 The loss is divided by `grad_accum_steps` before `backward()`, so accumulated gradients sum to the correct mean.
+Strategy metrics are detached and converted to Python `float` values before the
+`LossOutput` is returned; only `LossOutput.loss` remains a differentiable tensor.
 
 ## Callback Lifecycle
 
 | Hook | Fires | Default callback |
 |------|-------|-----------------|
-| `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
+| `on_train_begin` | Before training starts | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
 | `on_batch_begin` | Every batch | — |
-| `on_optimizer_step` | Every accumulation window | `GradientClippingCallback`, `MetricCallback`, `ProgressBarCallback` |
+| `on_optimizer_step` | Every accumulation window | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
 | `on_batch_end` | Every batch | `CheckpointCallback` |
 | `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
 | `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
-| `on_train_end` | Training ends (always via finally) | `CheckpointCallback`, `MetricCallback`, `GradientCheckpointingCallback` |
+| `on_train_end` | Training exits after `on_train_begin` completes (via `finally`) | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 
-Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm), `gradient_clipping` (always registered; computes grad norm, clips only when `max_grad_norm` is not `None`).
+Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm, rank-0), `gradient_clipping`. The gradient-clipping callback is always registered and always calls `executor.clip_grad_norm()` with the numeric `max_grad_norm` value.
 
 ## KV Cache Mathematics
 
@@ -151,7 +168,7 @@ Three-layer separation (SGLang-inspired):
 - **ReqToTokenPool**: Index table `[req_idx, pos] → physical token slot`, shared across all layers.
 - **Allocator + PrefixCache**: Paged-mode slot allocation with ref-counting, LRU eviction, and hash-based prefix sharing.
 
-`PagePool` orchestrates all three. In contiguous mode (default), `req_to_token` is a trivial linear mapping. In paged mode, slots are allocated on demand with prefix caching support. `bind_tasks()` returns a `KVCache` dataclass with precomputed `page_table` and `decode_mask` fields (computed once per decode step, shared across all layers). Attention layers access buffers directly — no methods, no abstraction.
+`PagePool` orchestrates all three. In contiguous mode (default), `req_to_token` is a trivial linear mapping. In paged mode, slots are allocated on demand with prefix caching support. `bind_tasks()` returns a `KVCache` dataclass with `kv_indptr`, a prefix-sum index over sequence lengths computed once per step and shared across layers. Attention layers access buffers directly — no methods, no abstraction.
 
 ### Attention Backend
 
@@ -230,4 +247,4 @@ total_steps          = (batches_per_replica // grad_accum_steps) * n_epoch
 
 This accounts for data-parallel sharding — each rank processes `1/nprocs` of the dataset.
 
-> Document Update Time: 2026-07-31
+> Document Update Time: 2026-08-02

docs/get-started.md

@@ -2,11 +2,23 @@
 
 This guide walks you through installing AstrAI, downloading a model, running inference, preprocessing data, and launching your first training job.
 
+## Contents
+
+- [Prerequisites](#prerequisites)
+- [1. Install](#1-install)
+- [2. Download Model Weights](#2-download-model-weights)
+- [3. Run Inference](#3-run-inference)
+- [4. Preprocess Data](#4-preprocess-data)
+- [5. Train](#5-train)
+- [6. Evaluate](#6-evaluate)
+- [7. Docker](#7-docker)
+- [Next Steps](#next-steps)
+
 ## Prerequisites
 
 - **Python 3.12+**
-- **PyTorch 2.11+** (CUDA 12.8 recommended for GPU support)
-- NVIDIA GPU with CUDA (optional but recommended; CPU works for inference)
+- **PyTorch 2.11.0** (the exact version pinned by AstrAI; CUDA 12.8 build recommended for GPU support)
+- NVIDIA GPU with CUDA for training, `scripts/tools/generate.py`, generation evaluations, and demos. The HTTP server and direct-scoring evaluations can run on CPU where their CLI exposes a CPU device.
 
 ## 1. Install
 
@@ -55,7 +67,7 @@ python scripts/demo/stream_chat.py
 # Type your message after >>, type !exit to quit
 ```
 
-This starts a multi-turn interactive chat session with streaming output.
+This starts a single-turn interactive prompt loop with streaming output. Each prompt is independent; conversation history is not retained.
 
 ### Start an HTTP Server
 
@@ -192,6 +204,13 @@ See [Training Guide](guides/training.md) for loss formulas and strategies. See [
 
 ## 6. Evaluate
 
+HumanEval and MMLU download their benchmark data through HuggingFace
+`datasets`, which is not part of the base install:
+
+```bash
+pip install datasets
+```
+
 ```bash
 # HumanEval (code generation, auto-downloads dataset)
 python scripts/eval/evaluate_humaneval.py --param_path ./params --num_samples 20

docs/guides/distributed.md

@@ -2,6 +2,18 @@
 
 AstrAI supports three parallel modes: **single GPU** (`none`), **Data Parallel** (`ddp`), and **Fully Sharded Data Parallel** (`fsdp`). This guide covers when to use each, how to launch multi-GPU training, and how gradient accumulation works.
 
+## Contents
+
+- [Quick Start](#quick-start)
+- [Parallel Modes](#parallel-modes)
+- [Gradient Accumulation](#gradient-accumulation)
+- [Process Launching](#process-launching)
+- [NCCL Troubleshooting](#nccl-troubleshooting)
+- [Checkpoint Saving](#checkpoint-saving)
+- [Total Steps Calculation](#total-steps-calculation)
+- [Real Examples](#real-examples)
+- [CLI Parameters](#cli-parameters)
+
 ## Quick Start
 
 ### Single GPU
@@ -21,9 +33,6 @@ python scripts/tools/train.py \
 
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
-export NCCL_P2P_DISABLE=1
-export NCCL_NET_GDR_LEVEL=0
-
 python scripts/tools/train.py \
     --train_type=sft \
     --param_path ./params \
@@ -38,9 +47,6 @@ python scripts/tools/train.py \
 
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
-export NCCL_P2P_DISABLE=1
-export NCCL_NET_GDR_LEVEL=0
-
 python scripts/tools/train.py \
     --train_type=sft \
     --param_path ./params \
@@ -110,7 +116,7 @@ AstrAI auto-detects the launch method:
 
 | Detection | Strategy | Use Case |
 |-----------|----------|----------|
-| `torchelastic` / `torchrun` env vars | `TorchrunStrategy` | External orchestrator (torchrun, SLURM, K8s) |
+| `torchelastic` / `torchrun` env vars | `TorchrunStrategy` | External orchestrator (`torchrun`, K8s) |
 | `RANK` + `WORLD_SIZE` env vars | `TorchrunStrategy` | External launch |
 | Neither | `LocalStrategy` | `python scripts/tools/train.py` (in-process spawn) |
 
@@ -126,23 +132,28 @@ For multi-node or SLURM environments:
 torchrun --nproc_per_node=4 scripts/tools/train.py \
     --train_type=sft \
     --parallel_mode=ddp \
+    --nprocs=4 \
     --param_path ./params \
     --data_root_path ./dataset \
     --batch_per_device=4
 ```
 
-When launched via torchrun, AstrAI reads `RANK`, `WORLD_SIZE`, `LOCAL_RANK` from the environment and uses `TorchrunStrategy`. The `--nprocs` flag is ignored (the orchestrator controls process count).
+When launched via `torchrun`, the launcher creates the worker processes. AstrAI reads `RANK`, `WORLD_SIZE`, and `LOCAL_RANK` from the environment and uses `TorchrunStrategy`; `--nprocs` does not control process creation in this mode.
 
-## NCCL Environment Variables
+The current training CLI still uses `--nprocs` when calculating scheduler `total_steps`. Set it to the global `WORLD_SIZE` so the step count reflects data-parallel sharding, including multi-node runs.
 
-For multi-GPU training, you **must** set these environment variables:
+Raw Slurm variables such as `SLURM_PROCID`, `SLURM_NTASKS`, and `SLURM_LOCALID` are not recognized automatically. Launch through `torchrun`, or map the scheduler's variables to `RANK`, `WORLD_SIZE`, `LOCAL_RANK`, `MASTER_ADDR`, and `MASTER_PORT` before starting AstrAI. The same requirement applies to launchers that expose only OpenMPI-specific variables.
+
+## NCCL Troubleshooting
+
+The following variables are troubleshooting options for hardware or network configurations where NCCL hangs or fails. They are not general requirements and can reduce performance by disabling peer-to-peer or GPUDirect RDMA paths:
 
 ```bash
 export NCCL_P2P_DISABLE=1
 export NCCL_NET_GDR_LEVEL=0
 ```
 
-These are required on certain hardware configurations (see `AGENTS.md`). Without them, NCCL may hang or crash during collective operations. These are set in the training shell scripts (`train-seq.sh`, `train-sft.sh`, `train-dpo.sh`) but not in Python code — you must export them before launching.
+Apply them only after confirming the relevant NCCL transport is the source of the failure. AstrAI does not set them in Python.
 
 ## Checkpoint Saving
 
@@ -176,9 +187,6 @@ This ensures the LR schedule is correctly scaled regardless of the number of GPU
 
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
-export NCCL_P2P_DISABLE=1
-export NCCL_NET_GDR_LEVEL=0
-
 python scripts/tools/train.py \
     --train_type=seq \
     --param_path ./params \
@@ -240,7 +248,7 @@ python scripts/tools/train.py \
 
 | Parameter | Default | Description |
 |-----------|---------|-------------|
-| `--nprocs` | 1 | Number of GPUs / processes |
+| `--nprocs` | 1 | Local process count for AstrAI's launcher; under `torchrun`, set it to global `WORLD_SIZE` for total-step calculation |
 | `--parallel_mode` | `fsdp` | `none`, `ddp`, or `fsdp` |
 | `--start_method` | `spawn` | Multiprocessing start method (`spawn`, `fork`, `forkserver`) |
 | `--backend` | `nccl` | Distributed backend (`nccl`, `gloo`) |
@@ -248,8 +256,8 @@ python scripts/tools/train.py \
 | `--master_port` | `29500` | Master node port |
 | `--device_type` | `cuda` | Device type |
 
-> `--tp_size` is parsed but **not yet wired** — tensor parallelism is future work. `ColumnParallelLinear` / `RowParallelLinear` exist in `astrai/parallel/module.py` but are not used by the model.
+> `--tp_size` is accepted by the CLI but discarded before configuration. Tensor parallelism is not implemented, and there is no tensor-parallel module or model integration.
 
 Full parameter reference: [CLI Reference](params.md). Training loop and strategies: [Training Guide](training.md).
 
-> Document Update Time: 2026-07-30
+> Document Update Time: 2026-08-02

docs/guides/evaluation.md

@@ -2,6 +2,29 @@
 
 AstrAI provides 7 evaluation scripts in `scripts/eval/` covering code generation, knowledge QA, perplexity, summarization, data quality, instruction following, and weight analysis.
 
+## Contents
+
+- [Prerequisites](#prerequisites)
+- [Overview](#overview)
+- [HumanEval](#humaneval-code-generation)
+- [MMLU](#mmlu-knowledge-qa)
+- [Perplexity](#perplexity-ppl)
+- [ROUGE](#rouge)
+- [IFD](#ifd-instruction-following-difficulty)
+- [IFEval](#ifeval-instruction-following)
+- [Weight Analysis](#weight-analysis)
+- [Tips](#tips)
+
+## Prerequisites
+
+HumanEval, MMLU, and IFEval import HuggingFace `datasets` to download their benchmark data. This package is not installed by AstrAI's base dependencies, so install it before running those scripts:
+
+```bash
+pip install datasets
+```
+
+The generation-based scripts require CUDA because they load the model on `cuda` with `bfloat16`. Direct-scoring and metric scripts support the devices shown below.
+
 ## Overview
 
 | Script | Metric | Model Invocation | External Dataset |
@@ -18,7 +41,15 @@ Two invocation patterns exist:
 - **Generation benchmarks** (HumanEval, IFEval): use `InferenceEngine` to generate responses, then score them.
 - **Scoring benchmarks** (MMLU, PPL, IFD): call `model()` directly under `torch.inference_mode()` for log-likelihood computation.
 
-Common defaults: `--param_path` defaults to `./params`; dtype defaults to `bfloat16` on CUDA, `float32` on CPU.
+| Script | Device support |
+|--------|----------------|
+| HumanEval | CUDA for generation; `--test_only` can score existing completions without loading a model |
+| IFEval | CUDA only |
+| MMLU | CUDA or CPU via `--device`; auto-selects CUDA when available |
+| PPL | CUDA or CPU via `--device`; auto-selects CUDA when available |
+| IFD | CUDA or CPU via `--device`; auto-selects CUDA when available |
+| ROUGE | CPU-only metric computation; no model is loaded |
+| Weight analysis | CUDA by default; CPU supported via `--device cpu` |
 
 ---
 
@@ -30,7 +61,7 @@ Generates completions for 164 programming problems, executes them against hidden
 python scripts/eval/evaluate_humaneval.py \
     --param_path ./params \
     --num_samples 20 \
-    --batch_size 32 \
+    --batch_size 64 \
     --max_tokens 512 \
     --output results/humaneval.json
 ```
@@ -47,7 +78,8 @@ python scripts/eval/evaluate_humaneval.py \
 | `--temperature` | 0.8 | Sampling temperature |
 | `--top_p` | 0.95 | Nucleus sampling threshold |
 | `--top_k` | 50 | Top-k sampling |
-| `--batch_size` | 32 | Generation batch size |
+| `--batch_size` | 64 | Generation batch size |
+| `--max_seq_len` | 4096 | KV cache sequence length |
 | `--test_workers` | 8 | ProcessPoolExecutor workers for test execution |
 | `--test_timeout` | 3.0 | Per-subprocess timeout (seconds) |
 | `--problems` | None | Restrict to specific problem indices |
@@ -66,7 +98,7 @@ python scripts/eval/evaluate_humaneval.py \
 python scripts/eval/evaluate_mmlu.py \
     --param_path ./params \
     --n_shot 5 \
-    --subjects math_algebra history_us \
+    --subjects abstract_algebra high_school_us_history \
     --output results/mmlu.json
 ```
 
@@ -82,12 +114,13 @@ python scripts/eval/evaluate_mmlu.py \
 | `--device` | auto | Device (`cuda` / `cpu`) |
 | `--dtype` | auto | `bfloat16` on CUDA, `float32` on CPU |
 | `--seed` | 0 | Seed for option permutation (0 = enabled, -1 = disabled) |
+| `--batch_size` | 4 | Questions per batch; each question produces four choice rows |
 
 **How it works**: For each question, builds a prompt with n-shot examples, then scores each choice (A/B/C/D) by computing the summed log-likelihood of the choice token given the context. The choice with the highest log-prob is the prediction.
 
 **Output**: stdout prints per-subject accuracy and overall. With `--output`, writes per-subject `{accuracy, correct, total}` + `_overall` aggregate.
 
-**Data**: Auto-downloads `cais/mmlu` from HuggingFace. Stored as per-subject CSVs in `<data_dir>/<split>/` and `<data_dir>/dev/` (for few-shot).
+**Data**: Auto-downloads `cais/mmlu` from HuggingFace. Stored as per-subject CSVs in `<data_dir>/<split>/` and `<data_dir>/dev/` (for few-shot). `--subjects` accepts canonical MMLU names such as `abstract_algebra`, `college_computer_science`, `high_school_us_history`, and `world_religions`.
 
 ---
 
@@ -100,7 +133,7 @@ python scripts/eval/evaluate_ppl.py \
     --param_path ./params \
     --input_path data.jsonl \
     --output_dir ppl_results/ \
-    --batch_size 4 \
+    --batch_size 64 \
     --max_length 2048
 ```
 
@@ -110,7 +143,7 @@ python scripts/eval/evaluate_ppl.py \
 | `--input_path` | required | Input file, glob, or directory |
 | `--output_dir` | required | Output directory for `summary.json` + token JSONL |
 | `--text_key` | `text` | Key for the text field in input data |
-| `--batch_size` | 4 | Batch size |
+| `--batch_size` | 64 | Batch size |
 | `--max_length` | 2048 | Max sequence length (tokens) |
 | `--token_level` | False | Store per-token log_probs + token-type analysis |
 | `--max_samples` | None | Random subsample per file |
@@ -119,7 +152,7 @@ python scripts/eval/evaluate_ppl.py \
 
 **Input**: JSONL or JSON files. Each item must have a field named by `--text_key` (default `text`). If `--input_path` is a directory, recursively collects `*.jsonl` and `*.json`.
 
-**Output**: `summary.json` with per-file stats (tokens, mean/median loss, perplexity, p50/p90/p95/p99). With `--token_level`, also writes per-token JSONL with token IDs and log-probs.
+**Output**: `summary.json` with per-file token count, mean loss, perplexity, and p50/p90/p95/p99 loss. Median loss is included only with `--token_level`; that mode also writes per-token JSONL with token IDs and log-probs.
 
 ---
 
@@ -210,7 +243,8 @@ python scripts/eval/evaluate_ifeval.py \
 | `--top_p` | 0.95 | Top-p sampling |
 | `--top_k` | 50 | Top-k sampling |
 | `--num_samples` | 1 | Samples per problem (best-of-n scoring) |
-| `--batch_size` | 1 | Inference batch size |
+| `--batch_size` | 64 | Inference batch size |
+| `--max_seq_len` | 4096 | KV cache sequence length |
 | `--limit` | None | Limit to first N problems (quick testing) |
 | `--dump_responses` | None | Path to dump raw responses as JSONL |
 
@@ -232,7 +266,7 @@ python scripts/eval/analyze_weights.py \
 
 | Parameter | Default | Description |
 |-----------|---------|-------------|
-| `--ckpt_dir` | required | Checkpoint dir with `model.safetensors` + `config.json` |
+| `--ckpt_dir` | required | Checkpoint directory containing `model.safetensors` |
 | `--compare` | None | Additional checkpoint dirs to compare |
 | `--no_svd` | False | Skip SVD; show only weight stats (faster) |
 | `--output` | None | Save results as JSON |
@@ -245,8 +279,8 @@ python scripts/eval/analyze_weights.py \
 ## Tips
 
 - **Quick test**: Use `--limit` (IFEval) or `--problems` (HumanEval) to run on a small subset first.
-- **Auto-download**: HumanEval, MMLU, and IFEval auto-download their datasets on first run. The other scripts expect user-provided data.
+- **Auto-download**: After installing `datasets`, HumanEval, MMLU, and IFEval auto-download their datasets on first run. The other scripts expect user-provided data.
 - **Output formats**: `--output` writes a single JSON for most scripts. PPL and IFD write an `--output_dir` containing `summary.json` plus per-file artifacts.
-- **CPU mode**: All scripts auto-detect CUDA. To force CPU, use `--device cpu --dtype float32`.
+- **CPU mode**: MMLU, PPL, and IFD support `--device cpu --dtype float32`; weight analysis supports `--device cpu`. HumanEval generation and IFEval are CUDA-only.
 
 > Document Update Time: 2026-07-30

docs/guides/inference.md

@@ -49,8 +49,7 @@ KVCache
   ├── seq_lens               [batch_size]
   ├── out_cache_loc          [batch, seq_len] — write indices for this forward
   ├── max_len                int — max(seq_lens), avoids GPU sync in decode
-  ├── page_table             [batch, max_len] — precomputed gather indices for decode (None for prefill)
-  └── decode_mask            [batch, max_len] bool — precomputed position validity mask (None for single-batch decode)
+  └── kv_indptr              [batch + 1] int32 — prefix sum of seq_lens, precomputed once per step
 ```
 
 Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k` to write, `k_buffer[layer_id, indices]` to gather.
@@ -87,7 +86,9 @@ Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_b
 - **CUDA kernel** (`rotary_emb.cu`): fused cos/sin lookup + rotation in a single kernel, used when the kernel is available, input is on CUDA, and `torch.is_grad_enabled()` is `False` (inference mode)
 - **Torch fallback**: complex multiply path (`torch.view_as_complex` → `torch.complex` multiply → `torch.view_as_real`), used during training (supports autograd backward) or when the CUDA kernel is not available
 
-`RotaryEmbedding` stores `cos_table`/`sin_table` as f32 buffers and returns a `(cos, sin)` tuple from `forward()`. Both attention backends share the same rotary dispatch — it is backend-agnostic.
+`RotaryEmbedding` stores a complex `freqs_cis` buffer and returns a tensor
+from `forward()`. Both attention backends share the same rotary dispatch — it
+is backend-agnostic.
 
 ## Continuous Batching
 
@@ -183,22 +184,58 @@ curl -X POST http://localhost:8000/v1/messages \
   -d '{"model":"astrai","system":"You are helpful.","messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
 ```
 
-Supports `stop_sequences` and streaming via `event: content_block_delta`.
+Supports `stop_sequences` and streaming via `event: content_block_delta`. Anthropic streams also end with the shared `data: [DONE]` sentinel after `event: message_stop`.
 
-### GenerationRequest Parameters
+### Request Parameters
+
+The HTTP protocols and direct engine API have distinct request models and defaults.
+
+**OpenAI** (`ChatCompletionRequest`):
 
 | Param | Type | Default | Description |
 |-------|------|---------|-------------|
+| `model` | str | `"astrai"` | Model name returned in responses |
 | `messages` | List[dict] | required | Chat messages (role, content) |
-| `top_k` | int | 50 | Top-k count |
-| `top_p` | float | 1.0 | Nucleus threshold |
-| `temperature` | float | 1.0 | Sampling temperature (> 0.0) |
-| `max_tokens` | Optional[int] | None | Max generation length |
-| `stream` | bool | False | Stream output |
+| `temperature` | Optional[float] | 1.0 | Sampling temperature (0.0-2.0) |
+| `top_p` | Optional[float] | 1.0 | Nucleus threshold (0.0-1.0) |
+| `top_k` | Optional[int] | 50 | Top-k count |
+| `max_tokens` | Optional[int] | 2048 | Max generation length |
+| `stream` | Optional[bool] | False | Stream output |
 | `stop` | Optional[Union[str, List[str]]] | None | Stop sequences |
-| `frequency_penalty` | float | 0.0 | Frequency penalty |
-| `tools` | Optional[List[dict]] | None | Tool definitions for function calling |
-| `tool_choice` | Optional[str] | None | Tool selection mode |
+| `n` | Optional[int] | 1 | Number of choices requested |
+| `presence_penalty` | Optional[float] | 0.0 | Presence penalty (-2.0 to 2.0) |
+| `frequency_penalty` | Optional[float] | 0.0 | Frequency penalty (-2.0 to 2.0) |
+| `logit_bias` | Optional[Dict[int, float]] | None | Per-token logit bias |
+| `user` | Optional[str] | None | End-user identifier |
+| `tools` | Optional[List[ToolDef]] | None | Tool definitions for function calling |
+| `tool_choice` | Optional[Union[str, Dict[str, Any]]] | `"auto"` | Tool selection mode or explicit tool choice |
+
+**Anthropic** (`MessagesRequest`):
+
+| Param | Type | Default | Description |
+|-------|------|---------|-------------|
+| `model` | str | `"astrai"` | Model name returned in responses |
+| `messages` | List[AnthropicMessage] | required | User/assistant messages |
+| `system` | Optional[str] | None | System prompt |
+| `max_tokens` | int | 1024 | Max generation length |
+| `temperature` | Optional[float] | 1.0 | Sampling temperature (0.0-2.0) |
+| `top_p` | Optional[float] | 1.0 | Nucleus threshold (0.0-1.0) |
+| `top_k` | Optional[int] | 50 | Top-k count |
+| `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
 
@@ -240,6 +277,8 @@ data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":
 
 event: message_stop
 data: {"type":"message_stop"}
+
+data: [DONE]
 ```
 
 ### Error Responses

docs/guides/params.md

@@ -13,9 +13,11 @@
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
+| `--config`, `-c` | YAML config file; explicit CLI options override YAML values | None |
 | `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`, `online_grpo`, `online_dpo`) | required |
 | `--data_root_path` | Dataset root directory | required |
 | `--param_path` | Model parameters or checkpoint path | required |
+| `--resume` | Resume training from `--param_path` | False |
 | `--n_epoch` | Total training epochs | 1 |
 | `--batch_per_device` | Batch size per device | 1 |
 | `--grad_accum_steps` | Gradient accumulation steps between optimizer steps | 1 |
@@ -26,7 +28,7 @@
 |-----------|-------------|---------|
 | `--warmup_ratio` | Fraction of total steps used for LR warmup | 0.05 |
 | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
-| `--max_grad_norm` | Maximum gradient norm for clipping (None disables) | 1.0 |
+| `--max_grad_norm` | Maximum gradient norm for clipping; the current CLI requires a positive number | 1.0 |
 
 ### Optimizer
 
@@ -36,9 +38,9 @@ non-matrix parameters through **AdamW** (`fused=True`).
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--optimizer` | Built-in optimizer (`muon_adamw`, `nora_nadamw`, `mano_adamw`) | `muon_adamw` |
-| `--weight_decay` | Weight decay (applied to Muon matrix params; non-matrix use 0) | 0.1 |
+| `--weight_decay` | Weight decay for optimizer parameter groups that are eligible for decay | 0.1 |
 | `--muon_momentum` | Muon momentum factor | 0.95 |
-| `--muon_nesterov` | Enable Nesterov momentum for Muon | True |
+| `--muon_nesterov`, `--no-muon_nesterov` | Enable or disable Nesterov momentum for Muon | enabled |
 | `--muon_ns_steps` | Newton-Schulz iteration steps for Muon | 5 |
 | `--muon_adjust_lr` | Muon LR adjustment strategy (`original`, `match_rms_adamw`) | `match_rms_adamw` |
 
@@ -56,15 +58,18 @@ under DTensor sharding and rejects layouts sharded along the last dimension.
 | `--nora_weight_decay` | Nora matrix weight decay | 0.0 |
 
 `mano_adamw` routes internal `Linear.weight` matrices to **Mano** (manifold
-normalized optimizer) and the remaining parameters to **NAdamW**. Mano projects
+normalized optimizer) and the remaining parameters to **AdamW**. Mano projects
 the momentum onto the tangent space of the Oblique manifold and normalizes it,
 alternating the projection axis (row/column) each step — replacing Muon's
 Newton-Schulz iteration with a cheaper normalization.
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--mano_momentum` | Mano momentum factor | 0.95 |
-| `--mano_nesterov` | Enable Nesterov momentum for Mano | True |
+| `--mano_momentum` | Accepted by the CLI but currently ignored by optimizer construction | 0.95 |
+| `--mano_nesterov`, `--no-mano_nesterov` | Accepted by the CLI but currently ignored by optimizer construction | enabled |
+
+The two Mano-specific flags are reserved for future wiring; do not rely on them
+to change optimizer behavior in the current release.
 
 Optimizer identity and hyperparameters are saved in checkpoint metadata. Optimizer
 states are intentionally not interchangeable: resume older MuonAdamW checkpoints
@@ -78,7 +83,7 @@ with `--optimizer=muon_adamw`.
 | `--stride` | Stride for sliding window over sequences | None |
 | `--random_seed` | Random seed for reproducibility | 3407 |
 | `--num_workers` | DataLoader worker processes | 4 |
-| `--no_pin_memory` | Disable pin_memory (enabled by default) | (flag) |
+| `--pin_memory`, `--no-pin_memory` | Enable or disable DataLoader pinned memory | enabled |
 
 ### Checkpoint & Resume
 
@@ -100,14 +105,20 @@ with `--optimizer=muon_adamw`.
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--log_dir` | Directory for metric logs | checkpoint/logs |
-| `--metrics` | Metrics to log (e.g. --metrics loss lr val_loss) | ["loss", "lr", "grad_norm"] |
+| `--metrics` | Repeatable metric option (for example, `--metrics loss --metrics lr --metrics val_loss`) | `loss`, `lr`, `grad_norm`, `grad_snr` |
 
 ### Gradient Checkpointing
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--gradient_checkpointing` | Enable activation checkpointing for DecoderBlock modules | False |
+| `--gradient_checkpointing`, `--no-gradient_checkpointing` | Enable or disable activation checkpointing for DecoderBlock modules | disabled |
+
+### Miscellaneous
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--compile` | Enable `torch.compile` with mode `default`, `reduce-overhead`, or `max-autotune`; omit to disable | None |
+| `--dry-run` | Validate the merged configuration and print the training plan without training | False |
 
 ### Distributed Training
 
@@ -120,21 +131,25 @@ with `--optimizer=muon_adamw`.
 | `--backend` | Distributed training backend | nccl |
 | `--master_addr` | Master node address | localhost |
 | `--master_port` | Master node port | 29500 |
+| `--tp_size` | Reserved tensor-parallel size; accepted but currently ignored | None |
 
 ### Strategy-specific
 
 | Parameter | Description | Default | Used by |
 |-----------|-------------|---------|---------|
-| `--dpo_beta` | DPO beta value | 0.1 | `dpo` |
+| `--dpo_beta` | DPO beta value | 0.1 | `dpo`, `online_dpo` |
 | `--label_smoothing` | Label smoothing for cross-entropy loss | 0.0 | `seq`, `sft` |
-| `--group_size` | GRPO group size | 4 | `grpo` |
-| `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo` |
-| `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo` |
+| `--group_size` | GRPO/rollout group size | 4 | `grpo`, `online_grpo`, `online_dpo` |
+| `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo`, `online_grpo` |
+| `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo`, `online_grpo` |
 | `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` |
 
 ### Online Rollout
 
-These options apply to `online_grpo` and `online_dpo`. Online strategies require
+`online_grpo` and `online_dpo` are factory aliases for the existing `grpo` and
+`dpo` strategy classes; online behavior is enabled by rollout components rather
+than separate strategy subclasses. These options apply to the online aliases.
+Online strategies require
 a `BaseRewardModel` factory in `TrainConfig`; `train.py` does not currently
 provide a command-line option for configuring one.
 
@@ -151,7 +166,7 @@ provide a command-line option for configuring one.
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--schedule_type` | LR scheduler type (`cosine`, `sgdr`, `wsd`) | cosine |
-| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default: 0.05 for cosine/SGDR, 0.0 for WSD) |
+| `--min_rate` | Minimum LR as fraction of base LR | None (all current schedulers use their effective default of 0.01) |
 | `--cycle_length` | SGDR first cycle length in steps | None (total_steps - warmup_steps) |
 | `--t_mult` | SGDR cycle length multiplier per restart | 2 |
 | `--stable_steps` | WSD stable plateau steps | None (80% of post-warmup steps) |
@@ -204,14 +219,6 @@ python scripts/tools/server.py --param_path ./params --device cuda --dtype bfloa
 
 See [Inference Guide](inference.md) for HTTP API documentation.
 
-# Preprocess
-
-```bash
-python scripts/tools/preprocess.py data/*.jsonl -o output/ -c config.json
-```
-
-See [Preprocessing Guide](preprocessing.md) for config file format and examples.
-
 ## Generate (`generate.py`)
 
 | Parameter | Type | Default | Description |
@@ -221,13 +228,12 @@ See [Preprocessing Guide](preprocessing.md) for config file format and examples.
 | `--output_json_file` | str | required | Path to the output JSONL file |
 | `--question_key` | str | `question` | Key for the question in input JSON |
 | `--response_key` | str | `response` | Key for the response in output JSON |
-| `--temperature` | float | `0.60` | Sampling temperature |
-| `--top_k` | int | `30` | Top-k filtering |
+| `--temperature` | float | `0.8` | Sampling temperature |
+| `--top_k` | int | `50` | Top-k filtering |
 | `--top_p` | float | `0.95` | Nucleus sampling threshold |
 | `--batch_size` | int | `1` | Batch size for generation |
 | `--num_samples` | int | `1` | Responses per prompt |
-| `--max_tokens` | int | model config `max_position_embeddings` | Maximum tokens to generate |
-| `--cache_len` | int | `2048` | KV cache length |
+| `--max_seq_len` | int | `2048` | KV cache sequence length |
 | `--frequency_penalty` | float | `0.0` | Frequency penalty |
 | `--rep_window` | int | `64` | Window size for frequency penalty |
 
@@ -243,14 +249,15 @@ python scripts/tools/generate.py \
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `input_files` | path(s) | required | Input JSONL file(s), supports glob (`data/*.jsonl`) |
+| `input_files` | path(s) | required | One or more existing `.jsonl` or `.json` paths. Wildcards work only when expanded by the invoking shell; the CLI does not expand globs itself. |
 | `--output_dir`, `-o` | path | required | Output directory for processed data |
 | `--config`, `-c` | path | required | Preprocessing pipeline config (JSON) |
 | `--tokenizer_path` | str | `params` | Path to tokenizer directory |
+| `--batch_size` | int | config value (`256` by default) | Override records processed per batch; must be at least 1 |
 
 Usage:
 ```bash
-python scripts/tools/preprocess.py data/*.jsonl -o output/ -c sft.json
+python scripts/tools/preprocess.py data/part-000.jsonl data/part-001.jsonl -o output/ -c sft.json
 ```
 
 See [Preprocessing Guide](preprocessing.md) for config file format and examples.

docs/guides/preprocessing.md

@@ -10,6 +10,7 @@ Declarative JSON-driven data preprocessing. `MaskBuilderFactory` supports three
 - [Configuration Reference](#configuration-reference) — all fields
 - [Mask Algorithm](#mask-algorithm)
 - [Output Layout](#output-layout)
+- [Training Compatibility](#training-compatibility)
 - [CLI](#cli)
 - [Python API](#python-api)
 
@@ -40,7 +41,7 @@ A single config file captures the entire pipeline, reusable and version-controll
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `field` | str | -- | JSONL key to read |
-| `action` | str | -- | `"train"` / `"mask"` / `"$role"` |
+| `action` | str | -- | `"train"` / `"mask"` / `"$role"` / `"value"`; `"value"` copies raw values without tokenization |
 | `template` | bool | `false` | Apply `chat_template` per message |
 | `add_special_tokens` | bool | `true` for first non-template section | Add special tokens during encode |
 
@@ -89,7 +90,7 @@ Config:
 }
 ```
 
-Output keys: `sequence` (int32), `loss_mask` (bool)
+Output keys: `sequence` (int32), `loss_mask` (bool), `position_ids` (int32)
 
 ### SFT Instruction
 
@@ -116,7 +117,7 @@ Config:
 }
 ```
 
-Output keys: `sequence`, `loss_mask`
+Output keys: `sequence`, `loss_mask`, `position_ids`
 
 ### Pretrain
 
@@ -142,7 +143,7 @@ Config:
 }
 ```
 
-Output keys: `sequence` (no `loss_mask` — all tokens trained)
+Output keys: `sequence`, `position_ids` (no `loss_mask` — all tokens trained)
 
 ### DPO
 
@@ -180,6 +181,11 @@ Config:
 
 Output keys: `chosen`, `chosen_mask`, `rejected`, `rejected_mask`
 
+The offline `Pipeline` can construct these keys, but its `.bin` output is not
+currently loadable for DPO training because the writer does not preserve
+per-record offsets. Train DPO directly from raw JSONL instead; see
+[Training Compatibility](#training-compatibility).
+
 ### GRPO
 
 Input JSONL:
@@ -228,6 +234,11 @@ Output keys: `prompts`, `prompts_mask`, `responses`, `masks`, `rewards` (float32
 - `mask_key: "masks"` — rename the auto-generated mask key (default: `responses_mask`)
 - `prompts_mask` is auto-generated (all masked) and unused by GRPOStrategy
 
+The offline `Pipeline` flattens GRPO response groups for `.bin` output without
+preserving their boundaries, and there is no automatic raw-JSONL GRPO processor
+in `DatasetFactory`. See
+[Training Compatibility](#training-compatibility) for the supported routes.
+
 ---
 
 ## Configuration Reference
@@ -257,7 +268,7 @@ When `sources` is set, `sections` is ignored.
 | `max_chars` | int | `2000000` | Skip text-mode items longer than this |
 | `max_items` | int or null | `null` | Stop after N documents |
 | `batch_size` | int | `256` | Records per tokenization batch |
-| `packing_strategy` | str | `"simple"` | Packing strategy: `"simple"`, `"bfd"`, `"bfd_split"` |
+| `packing_strategy` | str | `"simple"` | Packing is supported for single-output data with a `sequence` key: `"simple"`, `"bfd"`, or `"bfd_split"`. Multi-output DPO/GRPO data is not record-preserving packed output. |
 | `max_packed_len` | int | `8192` | Maximum length of a packed bin |
 | `truncation_mode` | str | `"keep_start"` | How to truncate sequences: `"keep_start"` or `"keep_end"` |
 
@@ -266,8 +277,8 @@ When `sources` is set, `sections` is ignored.
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `domain_key` | str or null | `null` | JSONL key for domain grouping |
-| `storage_format` | str | `"bin"` | `"bin"` (mmap). Reading also supports `"jsonl"` for on-the-fly tokenization |
-| `max_tokens_per_shard` | int | `100000000` | Flush threshold in cumulative tokens |
+| `storage_format` | str | `"bin"` | Pipeline output format. Only `"bin"` has a registered writer; `"jsonl"` is accepted by config validation but cannot be emitted by `Pipeline`. |
+| `max_tokens_per_shard` | int | `100000000` | Flush threshold counted from each record's primary flat sequence: `sequence` for single-output data, otherwise the first flat source output |
 | `dtype` | dict[str, str] | `{}` | Per-key tensor dtype override (e.g. `{"loss_mask": "bool"}`) |
 | `position_ids_mode` | str | `"doc_reset"` | How to compute position_ids: `"none"`, `"doc_reset"`, `"continuous"` |
 
@@ -304,11 +315,13 @@ output/
       meta.json
       sequence.bin
       loss_mask.bin
+      position_ids.bin
   wiki/
     shard_0000/
       meta.json
       sequence.bin
       loss_mask.bin
+      position_ids.bin
 ```
 
 ### Multi-Shard (`bin`)
@@ -322,13 +335,44 @@ output/
       meta.json
       sequence.bin
       loss_mask.bin
+      position_ids.bin
     shard_0001/
       meta.json
       sequence.bin
       loss_mask.bin
+      position_ids.bin
 ```
 
-For `bin` format, `MmapStore` discovers all shards under the domain directory via `rglob("meta.json")`. For `h5` format, `H5Store` discovers `.h5`/`.hdf5` files via recursive glob.
+`MmapStore` discovers binary shards recursively through their `meta.json` files.
+Each shard's metadata is a top-level object keyed by tensor name:
+
+```json
+{
+  "sequence": {"shape": [123456], "dtype": "int32"},
+  "loss_mask": {"shape": [123456], "dtype": "bool"},
+  "position_ids": {"shape": [123456], "dtype": "int32"}
+}
+```
+
+An optional `offsets` array may appear for record-oriented binary data written
+through `save_bin(..., record_keys=...)`; the preprocessing `BinWriter` does not
+currently request those offsets.
+
+---
+
+## Training Compatibility
+
+| Training type | Supported input route |
+|---------------|-----------------------|
+| `seq` | Offline preprocessed `.bin`, or raw `.jsonl` eagerly transformed by `JsonlStore` using `dataset_config.json` or the built-in `messages` config |
+| `sft` | Offline preprocessed `.bin`, or raw `.jsonl` through the same eager transform routes |
+| `dpo` | Raw `.jsonl` through the automatic lazy DPO processor selected by `DatasetFactory` when `tokenizer_path` is supplied, or a caller-provided record store |
+| `grpo` | A caller-provided, already-loaded `Store` with record-shaped `prompts`, `responses`, `masks`, and `rewards`; no automatic raw-JSONL processor is currently wired |
+
+Offline DPO and GRPO preprocessing configs describe the intended token fields,
+but their `.bin` output is not currently loadable for training. DPO binary
+shards lack per-record offsets. GRPO response groups are flattened before the
+binary writer and their record/group boundaries are not preserved.
 
 ---
 
@@ -336,15 +380,20 @@ For `bin` format, `MmapStore` discovers all shards under the domain directory vi
 
 ```bash
 # SFT
-python scripts/tools/preprocess.py data/sft/*.jsonl -o output/sft/ -c configs/sft_chat.json
+python scripts/tools/preprocess.py data/sft/part-000.jsonl -o output/sft/ -c configs/sft_chat.json --batch_size 128
 
 # DPO
-python scripts/tools/preprocess.py data/dpo/*.jsonl -o output/dpo/ -c configs/dpo.json --tokenizer_path params
+python scripts/tools/preprocess.py data/dpo/part-000.jsonl -o output/dpo/ -c configs/dpo.json --tokenizer_path params
 
 # GRPO
-python scripts/tools/preprocess.py data/grpo/*.jsonl -o output/grpo/ -c configs/grpo.json
+python scripts/tools/preprocess.py data/grpo/part-000.jsonl -o output/grpo/ -c configs/grpo.json
 ```
 
+Inputs may be `.jsonl` files or `.json` files containing one object or a list of
+objects. Each positional path must exist. A wildcard such as `data/*.jsonl`
+works only when the invoking shell expands it before Click receives the
+arguments; otherwise pass the files explicitly.
+
 ---
 
 ## Python API

docs/guides/training.md

@@ -41,7 +41,10 @@ RoPE embeds position into Q/K vectors via complex rotation:
 
 $$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
 
-`RotaryEmbedding` pre-computes `cos_table` and `sin_table` (f32, `[max_len, dim/2]`). `forward()` returns a `(cos, sin)` tuple indexed by `position_ids`. `apply_rotary_emb` applies the rotation: during training it uses torch complex multiply (autograd-compatible); during inference it auto-dispatches to a fused CUDA kernel when available.
+`RotaryEmbedding` pre-computes a complex `freqs_cis` buffer. `forward()` returns
+a tensor indexed by `position_ids`. `apply_rotary_emb` applies the rotation:
+during training it uses torch complex multiply (autograd-compatible); during
+inference it auto-dispatches to a fused CUDA kernel when available.
 
 ## Training Loop
 
@@ -52,11 +55,12 @@ on_train_begin
   model.train()
   on_epoch_begin
     for batch in dataloader:
-      on_batch_begin
       with executor.accumulate(model):
-        loss = strategy.compute_loss(batch)
-        context.loss = loss.item()
-        stand_loss = loss / executor.grad_accum_steps
+        on_batch_begin
+        loss_output = strategy(batch)
+        context.loss = loss_output["loss"].item()
+        context.metrics = loss_output["metrics"]
+        stand_loss = loss_output["loss"] / executor.grad_accum_steps
         executor.backward(stand_loss)
         context.consumed_samples += (
             context.config.batch_per_device * context.world_size
@@ -66,6 +70,7 @@ on_train_begin
         if executor.sync_gradients:
           on_optimizer_step
           optimizer.step()
+          strategy.on_optimizer_step()
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()
@@ -77,16 +82,18 @@ on_train_end
 
 | Hook | Fires | Default callback |
 |------|-------|-----------------|
-| `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
+| `on_train_begin` | Before training starts | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
 | `on_batch_begin` | Every batch | — |
-| `on_optimizer_step` | Every accumulation window | `GradientClippingCallback`, `MetricCallback`, `ProgressBarCallback` |
+| `on_optimizer_step` | Every accumulation window | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
 | `on_batch_end` | Every batch | `CheckpointCallback` |
 | `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
 | `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
-| `on_train_end` | Training ends (always via finally) | `CheckpointCallback`, `MetricCallback`, `GradientCheckpointingCallback` |
+| `on_train_end` | Training exits after `on_train_begin` completes (via `finally`) | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 
-Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm), `gradient_clipping` (always registered; computes grad norm, clips only when `max_grad_norm` is not `None`).
+Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm, rank-0), `gradient_clipping`. The gradient-clipping callback is always registered and always calls `executor.clip_grad_norm()` with the numeric `max_grad_norm` value.
+
+Strategies return `{"loss": Tensor, "metrics": Dict[str, float]}` when called by the trainer. Built-in metrics include the task-specific loss and, for MoE models, `moe_aux_loss` plus `moe_aux_loss_weighted`. Direct `compute_loss(batch)` calls continue to return a single loss tensor.
 
 ## Strategies
 
@@ -95,7 +102,7 @@ Default callbacks (in order): `gradient_checkpointing` (activation checkpointing
 Next-token cross-entropy with optional label smoothing:
 
 $$
-L_{\text{PT}} = -\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
+L_{\text{PT}} = -\frac{1}{T}\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
 $$
 
 Keys: `input_ids`, `target_ids`. Optional: `label_smoothing`.
@@ -105,7 +112,7 @@ Keys: `input_ids`, `target_ids`. Optional: `label_smoothing`.
 Masked cross-entropy (`ignore_index=-100`) over response tokens:
 
 $$
-L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
+L_{\text{SFT}} = -\frac{1}{L}\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
 $$
 
 Keys: `input_ids`, `target_ids`, `loss_mask`, `position_ids`. Optional: `label_smoothing`.
@@ -165,9 +172,9 @@ model factory.
 |------|-------|-------------|
 | Cosine | `CosineScheduler` | Linear warmup → cosine decay to `min_rate` |
 | SGDR | `SGDRScheduler` | Cosine annealing with warm restarts (`t_mult=2`) |
-| WSD | `WSDScheduler` | Warmup-Stable-Decay with sqrt cooldown |
+| WSD | `WSDScheduler` | Warmup-Stable-Decay with quadratic decay |
 
-Created by `SchedulerFactory.create(schedule_type, optimizer, **kwargs)`. Valid types: `"cosine"`, `"sgdr"`, `"wsd"`. Omit to use no scheduler.
+Created by `SchedulerFactory.create(schedule_type, optimizer, **kwargs)`. Valid types: `"cosine"`, `"sgdr"`, `"wsd"`. The training CLI always creates a scheduler and defaults `--schedule_type` to `"cosine"`.
 
 ## Gradient Checkpointing
 
@@ -185,10 +192,12 @@ Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoi
 
 ```
 Checkpoint(state_dict, epoch, consumed_samples, extra, meta, config)
-  ├── save(save_dir)    rank-0 only: meta.json (epoch/consumed_samples/timestamp) + config.json (model config) + model.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt)
+  ├── save(save_dir)    meta.json (epoch/consumed_samples/timestamp) + config.json (model config) + model.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt)
   └── load(save_dir, broadcast=False)    loads from local disk; set broadcast=True to broadcast metadata from rank-0
 ```
 
+`Checkpoint.save()` writes whenever it is called. During training, `CheckpointCallback` uses the executor checkpoint context so only rank 0 receives a state dict and calls `save()`.
+
 Optimizer/scheduler state persisted by default via `Checkpoint.extra`.  
 Model config (`context.model_config`) saved into `config.json` during training via `CheckpointCallback`.
 
@@ -232,4 +241,4 @@ nohup python scripts/tools/train.py \
 
 Full parameter reference at [params.md](params.md).
 
-> Document Update Time: 2026-07-31
+> Document Update Time: 2026-08-02

scripts/tools/benchmark.py

@@ -13,7 +13,6 @@ from astrai.model import AutoModel
 _DTYPES = ["bfloat16", "float16", "float32"]
 _CACHES = ["contiguous", "paged"]
 _BACKENDS = ["cuda", "torch_native"]
-CACHE_MAX_SEQ = 2048
 
 _BACKEND_MAP = {
     "cuda": ATTN_BACKEND.CUDA,
@@ -56,13 +55,13 @@ class GenerationBenchmark:
         self.config = config
         self.backend = backend
 
-    def _make_pool(self, batch_size: int) -> PagePool:
+    def _make_pool(self, batch_size: int, max_seq_len: int) -> PagePool:
         return PagePool(
             n_layers=self.config.num_hidden_layers,
             n_kv_heads=self.config.num_key_value_heads,
             head_dim=self.config.hidden_size // self.config.num_attention_heads,
             max_batch_size=batch_size,
-            max_seq_len=CACHE_MAX_SEQ,
+            max_seq_len=max_seq_len,
             device=self.device,
             dtype=self.dtype,
             page_size=1,
@@ -128,7 +127,7 @@ class GenerationBenchmark:
     ) -> BenchmarkResult:
         import time
 
-        pool = self._make_pool(batch_size)
+        pool = self._make_pool(batch_size, prompt_length)
         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)))
@@ -189,7 +188,10 @@ class GenerationBenchmark:
     ) -> BenchmarkResult:
         import time
 
-        pool = self._make_pool(batch_size)
+        # Decode grows seq_len monotonically up to prompt + 5 + gen*num_trials
+        # (warmup 5 steps, then one step per trial), so size the pool to cover it.
+        max_seq_len = prompt_length + 5 + gen_length * num_trials
+        pool = self._make_pool(batch_size, max_seq_len)
         task_ids = self._run_prefill(pool, batch_size, prompt_length)
 
         for i in range(5):

tests/module/test_forward_configs.py

@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+from astrai.model.components.mlp import MLP, DeepSeekMoE
 from astrai.model.transformer import AutoRegressiveLM
 from tests.helpers import TINY_CONFIG
 
@@ -32,6 +33,59 @@ CONFIGS = [
         },
         id="gqa_moe",
     ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+            "mlp_only_layers": [0],
+        },
+        id="gqa_moe_dense_first",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+            "decoder_sparse_step": 2,
+        },
+        id="gqa_moe_sparse_step",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+            "norm_topk_prob": True,
+        },
+        id="gqa_moe_norm_topk",
+    ),
+    pytest.param(
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+            "moe_intermediate_size": 24,
+            "shared_expert_intermediate_size": 20,
+        },
+        id="gqa_moe_custom_intermediate",
+    ),
     pytest.param(
         {
             **TINY_CONFIG,
@@ -105,3 +159,168 @@ def test_model_forward_with_padding(config_kwargs, device):
 
     assert output["logits"].shape == (batch_size, seq_len, config.vocab_size)
     assert not torch.isnan(output["logits"]).any()
+
+
+def test_moe_per_layer_ffn_resolution():
+    """Verify that mlp_only_layers and decoder_sparse_step resolve FFN types correctly."""
+    from astrai.config.model_config import AutoRegressiveLMConfig
+
+    # mlp_only_layers: first layer dense, rest MoE
+    config = AutoRegressiveLMConfig(
+        **{
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "mlp_only_layers": [0],
+        }
+    )
+    model = AutoRegressiveLM(config)
+    assert isinstance(model.layers[0].mlp, MLP)
+    assert not isinstance(model.layers[0].mlp, DeepSeekMoE)
+    assert isinstance(model.layers[1].mlp, DeepSeekMoE)
+
+    # decoder_sparse_step=2: every other layer is MoE
+    config2 = AutoRegressiveLMConfig(
+        **{
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "decoder_sparse_step": 2,
+        }
+    )
+    model2 = AutoRegressiveLM(config2)
+    # layer 0 (id=0): (0+1)%2=1 != 0 -> MLP
+    assert isinstance(model2.layers[0].mlp, MLP)
+    assert not isinstance(model2.layers[0].mlp, DeepSeekMoE)
+    # layer 1 (id=1): (1+1)%2=0 -> MoE
+    assert isinstance(model2.layers[1].mlp, DeepSeekMoE)
+
+    # decoder_sparse_step=1 (default): all layers MoE
+    config3 = AutoRegressiveLMConfig(
+        **{
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+        }
+    )
+    model3 = AutoRegressiveLM(config3)
+    for layer in model3.layers:
+        assert isinstance(layer.mlp, DeepSeekMoE)
+
+
+def test_moe_custom_intermediate_shape():
+    """Verify MoE uses custom intermediate sizes when specified."""
+    from astrai.config.model_config import AutoRegressiveLMConfig
+
+    config = AutoRegressiveLMConfig(
+        **{
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "moe_intermediate_size": 24,
+            "shared_expert_intermediate_size": 20,
+        }
+    )
+    model = AutoRegressiveLM(config)
+    moe_layer = model.layers[0].mlp
+    assert isinstance(moe_layer, DeepSeekMoE)
+    # routed experts use moe_intermediate_size
+    for expert in moe_layer.routed_experts:
+        assert expert.up.weight.shape[0] == 24
+        assert expert.gate.weight.shape[0] == 24
+        assert expert.down.weight.shape[1] == 24
+    # shared experts use shared_expert_intermediate_size
+    for expert in moe_layer.shared_experts:
+        assert expert.up.weight.shape[0] == 20
+        assert expert.gate.weight.shape[0] == 20
+        assert expert.down.weight.shape[1] == 20
+
+
+def test_moe_defaults_preserve_normalized_routing():
+    from astrai.config.model_config import AutoRegressiveLMConfig
+
+    config = AutoRegressiveLMConfig(
+        **TINY_CONFIG,
+        ffn_type="moe",
+        n_routed_experts=4,
+        n_shared_experts=1,
+        n_activated_experts=2,
+        topk_method="greedy",
+    )
+    model = AutoRegressiveLM(config)
+
+    assert config.norm_topk_prob is True
+    assert model.layers[0].mlp.norm_topk_prob is True
+
+
+def test_moe_aux_loss_only_emitted_during_training():
+    from astrai.config.model_config import AutoRegressiveLMConfig
+
+    config = AutoRegressiveLMConfig(
+        **TINY_CONFIG,
+        ffn_type="moe",
+        n_routed_experts=4,
+        n_shared_experts=1,
+        n_activated_experts=2,
+        topk_method="greedy",
+    )
+    model = AutoRegressiveLM(config)
+    input_ids = torch.randint(0, config.vocab_size, (2, 8))
+
+    outputs = model(input_ids)
+    assert outputs["aux_loss"].ndim == 0
+    assert outputs["aux_loss"].requires_grad
+    assert torch.isfinite(outputs["aux_loss"])
+
+    with torch.no_grad():
+        outputs = model(input_ids)
+    assert "aux_loss" not in outputs
+
+    model.eval()
+    outputs = model(input_ids)
+    assert "aux_loss" not in outputs
+
+
+def test_moe_component_forward_returns_ffn_output():
+    from astrai.model.components.mlp import DeepSeekMoE
+
+    moe = DeepSeekMoE(
+        dim=8,
+        dim_ffn=16,
+        n_routed_experts=4,
+        n_shared_experts=1,
+        n_activated_experts=2,
+    )
+
+    output = moe(torch.randn(2, 8, 8))
+
+    assert output["hidden_states"].shape == (2, 8, 8)
+    assert output["aux_loss"] is not None
+
+
+@pytest.mark.parametrize("decoder_sparse_step", [0, -1])
+def test_moe_rejects_invalid_decoder_sparse_step(decoder_sparse_step):
+    from pydantic import ValidationError
+
+    from astrai.config.model_config import AutoRegressiveLMConfig
+
+    with pytest.raises(ValidationError, match="decoder_sparse_step must be at least 1"):
+        AutoRegressiveLMConfig(
+            **TINY_CONFIG,
+            ffn_type="moe",
+            n_routed_experts=4,
+            n_activated_experts=2,
+            decoder_sparse_step=decoder_sparse_step,
+        )

tests/trainer/test_loss_output.py

@@ -0,0 +1,105 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from astrai.model.transformer import AutoRegressiveLM
+from astrai.trainer.strategy import BaseStrategy, SEQStrategy
+from astrai.trainer.train_callback import MetricCallback
+from tests.helpers import make_tiny_config
+
+
+def test_seq_strategy_combines_and_reports_moe_aux_loss(device):
+    config = make_tiny_config(
+        ffn_type="moe",
+        n_routed_experts=4,
+        n_shared_experts=1,
+        n_activated_experts=2,
+        topk_method="greedy",
+    )
+    model = AutoRegressiveLM(config).to(device=device)
+    strategy = SEQStrategy(model, device, moe_aux_loss_coef=0.25)
+    batch = {
+        "input_ids": torch.randint(0, config.vocab_size, (2, 8), device=device),
+        "target_ids": torch.randint(0, config.vocab_size, (2, 8), device=device),
+    }
+
+    output = strategy(batch)
+    legacy_loss = strategy.compute_loss(batch)
+
+    assert isinstance(legacy_loss, torch.Tensor)
+    assert set(output["metrics"]) == {
+        "loss",
+        "task_loss",
+        "moe_aux_loss",
+        "moe_aux_loss_weighted",
+    }
+    assert output["loss"].item() == pytest.approx(
+        output["metrics"]["task_loss"] + output["metrics"]["moe_aux_loss_weighted"],
+    )
+    assert output["metrics"]["moe_aux_loss_weighted"] == pytest.approx(
+        0.25 * output["metrics"]["moe_aux_loss"],
+    )
+    assert output["loss"].requires_grad
+    assert all(isinstance(metric, float) for metric in output["metrics"].values())
+
+
+def test_metric_callback_includes_dynamic_strategy_metrics(tmp_path):
+    callback = MetricCallback(
+        ckpt_dir=tmp_path,
+        save_interval=1,
+        metrics=["loss", "lr"],
+    )
+    context = SimpleNamespace(
+        metrics={"task_loss": 2.0, "moe_aux_loss": 1.0},
+        loss=2.01,
+        optimizer=SimpleNamespace(param_groups=[{"lr": 1e-3}]),
+        val_loss=None,
+        grad_norm=None,
+        grad_snr_tracker=None,
+        world_size=1,
+    )
+
+    metrics = callback._metrics(context, callback.metrics)
+
+    assert metrics == {
+        "loss": 2.01,
+        "lr": 1e-3,
+        "task_loss": 2.0,
+        "moe_aux_loss": 1.0,
+    }
+
+
+def test_metric_callback_only_computes_requested_metrics(tmp_path):
+    def fail_metric(context):
+        _ = context
+        raise AssertionError("unrequested metric was computed")
+
+    callback = MetricCallback(
+        ckpt_dir=tmp_path,
+        save_interval=1,
+        metrics=["loss"],
+    )
+    callback._metric_funcs["grad_snr"] = fail_metric
+    context = SimpleNamespace(
+        metrics={},
+        loss=2.0,
+        world_size=1,
+    )
+
+    metrics = callback._metrics(context, callback.metrics)
+
+    assert metrics == {"loss": 2.0}
+
+
+def test_legacy_strategy_tensor_loss_is_normalized():
+    class LegacyStrategy(BaseStrategy):
+        def compute_loss(self, batch):
+            return torch.tensor(2.0, requires_grad=True)
+
+    strategy = LegacyStrategy(torch.nn.Linear(1, 1), "cpu")
+
+    output = strategy({})
+
+    assert output["loss"].item() == 2.0
+    assert output["metrics"]["loss"] == 2.0

tests/trainer/test_online_strategy.py

@@ -159,21 +159,21 @@ def test_call_without_runner_falls_back_to_compute_loss_grpo(device):
         "masks": torch.ones(2, 4, 6, device=device),
         "rewards": torch.randn(2, 4, device=device),
     }
-    loss = strat(batch)
+    loss = strat(batch)["loss"]
     assert torch.isfinite(loss).item()
 
 
 def test_call_with_runner_returns_finite_loss_grpo(device):
     strat = _make_grpo(device)
     strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
-    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
+    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
     assert torch.isfinite(loss).item()
 
 
 def test_call_with_runner_returns_finite_loss_dpo(device):
     strat = _make_dpo(device)
     strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
-    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
+    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
     assert torch.isfinite(loss).item()
 
 
@@ -268,7 +268,7 @@ def test_step_called_when_sync_gradients_true(device):
 def test_loss_is_differentiable_dpo(device):
     strat = _make_dpo(device)
     strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
-    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
+    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
     loss.backward()
     has_grad = any(
         p.grad is not None and p.grad.abs().sum() > 0 for p in strat.model.parameters()
@@ -279,7 +279,7 @@ def test_loss_is_differentiable_dpo(device):
 def test_ref_model_not_updated_by_backward_dpo(device):
     strat = _make_dpo(device)
     strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
-    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
+    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
     loss.backward()
     for p in strat.ref_model.parameters():
         assert p.grad is None