README.md

@@ -103,7 +103,9 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --weight_decay=0.1 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \

assets/docs/README-zh-CN.md

@@ -109,7 +109,9 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --weight_decay=0.1 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \

assets/docs/architecture.md

@@ -63,6 +63,7 @@ classDiagram
             +Optional[int] n_heads
             +Optional[int] n_kv_heads
             +Optional[bool] use_qk_norm
+            +Optional[bool] use_gated_attention
             +str ffn_type
             +Optional[dict] rope_scaling
             +Optional[str] pooling_type
@@ -755,6 +756,7 @@ classDiagram
             +int output_tokens
             +float arrival_time
             +Optional[float] finish_time
+            +Optional[Callable] stream_callback
             +int next_pos
             +is_finished(stop_ids) bool
         }
@@ -1247,4 +1249,4 @@ classDiagram
 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-09
+> Document Update Time: 2026-07-05

assets/docs/dataflow.md

@@ -61,7 +61,7 @@ StoreFactory.create("bin")   → MmapStore
 StoreFactory.create("jsonl") → JsonlStore
 ```
 
-**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage.
+**H5Store**: Reads HDF5 files, supports `share_memory_()` for multi-process DataLoader workers (copies tensors to shared memory).
 
 **MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`.
 
@@ -109,4 +109,4 @@ DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_typ
 
 Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.
 
-> Document Update Time: 2026-07-09
+> Document Update Time: 2026-07-05

assets/docs/inference.md

@@ -56,7 +56,7 @@ PageCache (paged KV cache with prefix sharing, alternative)
 1. Cleanup → Remove finished tasks, free KV cache slots/pages
 2. Refill  → Pop from waiting_queue, task_alloc resources, activate
 3. Prefill → Group by (prompt_len, start_pos), run full forward
-4. Decode  → Run single-token forward for each same-position group
+4. Decode  → Pick largest same-position group, single-token forward
 ```
 
 ## Sampling (Strategy Pattern)
@@ -249,4 +249,4 @@ async for token in engine.generate_async("Hello", ...):    # -> AsyncGenerator[s
     print(token)
 ```
 
-> Document Update Time: 2026-07-09
+> Document Update Time: 2026-07-05

assets/docs/params.md

@@ -28,17 +28,13 @@
 | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
 | `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
 
-### Optimizer (MuonMix)
+### Optimizer (AdamW)
-
-Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`fused=True`).
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--weight_decay` | Weight decay (applied to Muon matrix params; non-matrix use 0) | 0.1 |
+| `--adamw_beta1` | AdamW beta1 | 0.9 |
-| `--muon_momentum` | Muon momentum factor | 0.95 |
+| `--adamw_beta2` | AdamW beta2 | 0.95 |
-| `--muon_nesterov` | Enable Nesterov momentum for Muon | True |
+| `--adamw_weight_decay` | AdamW weight decay | 0.01 |
-| `--muon_ns_steps` | Newton-Schulz iteration steps for Muon | 5 |
-| `--muon_adjust_lr` | Muon LR adjustment strategy (`original`, `match_rms_adamw`) | `match_rms_adamw` |
 
 ### Data Loading
 
@@ -108,7 +104,7 @@ Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`f
 | 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.01) |
+| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default) |
 | `--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 (required for wsd) |
@@ -130,7 +126,9 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --weight_decay=0.1 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \
@@ -201,4 +199,4 @@ See [Preprocessing Guide](preprocessing.md) for config file format and examples.
 
 ---
 
-> Document Update Time: 2026-07-09
+> Document Update Time: 2026-07-05

assets/docs/preprocessing.md

@@ -361,4 +361,4 @@ Pipeline(
 ).run()
 ```
 
-> Document Update Time: 2026-07-09
+> Document Update Time: 2026-07-05

assets/docs/training.md

@@ -201,7 +201,9 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --weight_decay=0.1 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
+    --adamw_weight_decay=0.01 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \
@@ -212,4 +214,4 @@ nohup python scripts/tools/train.py \
 
 Full parameter reference at [params.md](params.md).
 
-> Document Update Time: 2026-07-09
+> Document Update Time: 2026-07-05

astrai/extension/__init__.py

@@ -1,19 +1,91 @@
-"""CUDA attention kernel wrappers with torch fallback.
+import importlib
+import logging
 
-Public API:
+import torch
-    - ``gqa_decode_attn`` — single-query decode attention
+import torch.nn.functional as F
-    - ``gqa_prefill_attn`` — multi-query prefill attention
 
-Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.gqa_*``)
+logger = logging.getLogger(__name__)
-when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``.
-"""
 
-from astrai.extension.loader import KERNEL_NAMES, is_available
+_available: dict[str, bool] = {}
-from astrai.extension.ops import gqa_decode_attn, gqa_prefill_attn
+_modules: dict[str, object] = {}
 
-__all__ = [
+for _name in ["gqa_decode_attn", "gqa_prefill_attn"]:
-    "gqa_decode_attn",
+    try:
-    "gqa_prefill_attn",
+        _mod = importlib.import_module(f".{_name}", package=__package__)
-    "is_available",
+        _available[_name] = True
-    "KERNEL_NAMES",
+        _modules[_name] = _mod
-]
+    except ImportError:
+        _available[_name] = False
+        _modules[_name] = None
+
+
+def _expand_kv_heads(
+    k: torch.Tensor, v: torch.Tensor, q_head: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Expand K/V heads to match Q heads for GQA fallback."""
+    kv_head = k.size(1)
+    if kv_head == q_head:
+        return k, v
+    group = q_head // kv_head
+    k = k.repeat_interleave(group, dim=1)
+    v = v.repeat_interleave(group, dim=1)
+    return k, v
+
+
+def _torch_fallback(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    mask: torch.Tensor | None,
+    is_causal: bool,
+    scale: float | None,
+) -> torch.Tensor:
+    k, v = _expand_kv_heads(k, v, q.size(1))
+    attn_mask = mask[:, None, None, :] if mask is not None else None
+    return F.scaled_dot_product_attention(
+        q, k, v, attn_mask=attn_mask, is_causal=is_causal and mask is None, scale=scale
+    )
+
+
+def gqa_decode_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    mask: torch.Tensor | None = None,
+    is_causal: bool = False,
+    causal_offset: int = 0,
+    scale: float | None = None,
+) -> torch.Tensor:
+    if _available["gqa_decode_attn"]:
+        return _modules["gqa_decode_attn"].gqa_decode_attn(
+            q,
+            k,
+            v,
+            mask=mask,
+            is_causal=is_causal,
+            causal_offset=causal_offset,
+            scale=scale,
+        )
+    return _torch_fallback(q, k, v, mask, is_causal, scale)
+
+
+def gqa_prefill_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    mask: torch.Tensor | None = None,
+    is_causal: bool = False,
+    causal_offset: int = 0,
+    scale: float | None = None,
+) -> torch.Tensor:
+    if _available["gqa_prefill_attn"]:
+        return _modules["gqa_prefill_attn"].gqa_prefill_attn(
+            q,
+            k,
+            v,
+            mask=mask,
+            is_causal=is_causal,
+            causal_offset=causal_offset,
+            scale=scale,
+        )
+    return _torch_fallback(q, k, v, mask, is_causal, scale)

astrai/extension/loader.py

@@ -1,36 +0,0 @@
-"""Dynamic discovery and loading of compiled CUDA kernel modules.
-
-Each kernel is registered in ``csrc/build.py`` and built into a ``.so`` placed
-in this package directory. On import we try to load each one; kernels that
-failed to build (or are running on a CPU-only machine) are marked unavailable
-so the wrapper functions can fall back to ``torch`` SDPA.
-"""
-
-import importlib
-import logging
-
-logger = logging.getLogger(__name__)
-
-KERNEL_NAMES = ["gqa_decode_attn", "gqa_prefill_attn"]
-
-_available: dict[str, bool] = {}
-_modules: dict[str, object] = {}
-
-for _name in KERNEL_NAMES:
-    try:
-        _mod = importlib.import_module(f".{_name}", package=__package__)
-        _available[_name] = True
-        _modules[_name] = _mod
-    except ImportError:
-        _available[_name] = False
-        _modules[_name] = None
-
-
-def is_available(name: str) -> bool:
-    """Return ``True`` if the compiled kernel ``name`` was loaded."""
-    return _available.get(name, False)
-
-
-def get_module(name: str) -> object:
-    """Return the loaded kernel module for ``name``, or ``None`` if unavailable."""
-    return _modules.get(name)

astrai/extension/ops.py

@@ -1,86 +0,0 @@
-"""GQA attention wrapper functions — one entry point per compiled kernel.
-
-Each wrapper dispatches to its CUDA kernel (loaded in ``loader.py``) when
-available, otherwise falls back to ``torch`` SDPA.
-
-Add new kernel wrappers here; split into per-variant files only if this file
-grows large.
-"""
-
-import torch
-import torch.nn.functional as F
-
-from astrai.extension.loader import _available, _modules
-
-
-def _expand_kv_heads(
-    k: torch.Tensor, v: torch.Tensor, q_head: int
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Expand K/V heads to match Q heads for GQA fallback."""
-    kv_head = k.size(1)
-    if kv_head == q_head:
-        return k, v
-    group = q_head // kv_head
-    k = k.repeat_interleave(group, dim=1)
-    v = v.repeat_interleave(group, dim=1)
-    return k, v
-
-
-def _torch_fallback(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    mask: torch.Tensor | None,
-    is_causal: bool,
-    scale: float | None,
-) -> torch.Tensor:
-    """Reference attention via ``scaled_dot_product_attention``."""
-    k, v = _expand_kv_heads(k, v, q.size(1))
-    attn_mask = mask[:, None, None, :] if mask is not None else None
-    return F.scaled_dot_product_attention(
-        q, k, v, attn_mask=attn_mask, is_causal=is_causal and mask is None, scale=scale
-    )
-
-
-def gqa_decode_attn(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    mask: torch.Tensor | None = None,
-    is_causal: bool = False,
-    causal_offset: int = 0,
-    scale: float | None = None,
-) -> torch.Tensor:
-    if _available["gqa_decode_attn"]:
-        return _modules["gqa_decode_attn"].gqa_decode_attn(
-            q,
-            k,
-            v,
-            mask=mask,
-            is_causal=is_causal,
-            causal_offset=causal_offset,
-            scale=scale,
-        )
-    return _torch_fallback(q, k, v, mask, is_causal, scale)
-
-
-def gqa_prefill_attn(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    mask: torch.Tensor | None = None,
-    is_causal: bool = False,
-    causal_offset: int = 0,
-    scale: float | None = None,
-) -> torch.Tensor:
-    if _available["gqa_prefill_attn"]:
-        return _modules["gqa_prefill_attn"].gqa_prefill_attn(
-            q,
-            k,
-            v,
-            mask=mask,
-            is_causal=is_causal,
-            causal_offset=causal_offset,
-            scale=scale,
-        )
-    return _torch_fallback(q, k, v, mask, is_causal, scale)

astrai/preprocessing/__init__.py

@@ -1,9 +1,7 @@
 from astrai.preprocessing.builder import (
     BaseMaskBuilder,
     MaskBuilderFactory,
-    MultiOutputMaskBuilder,
     SectionedMaskBuilder,
-    SingleOutputMaskBuilder,
 )
 from astrai.preprocessing.packing import (
     PackingStrategy,
@@ -22,14 +20,12 @@ from astrai.preprocessing.writer import (
 __all__ = [
     "BaseMaskBuilder",
     "MaskBuilderFactory",
-    "MultiOutputMaskBuilder",
     "PackingStrategy",
     "PackingStrategyFactory",
     "Pipeline",
     "PositionIdStrategy",
     "PositionIdStrategyFactory",
     "SectionedMaskBuilder",
-    "SingleOutputMaskBuilder",
     "StoreWriter",
     "StoreWriterFactory",
     "filter_by_length",

astrai/preprocessing/builder.py

@@ -1,10 +1,8 @@
 """Mask building for preprocessing pipeline.
 
 :class:`SectionRenderer` converts section specs into token ids and loss
-masks (template / text / value extraction).  :class:`SingleOutputMaskBuilder`
+masks (template / text / value extraction).  :class:`SectionedMaskBuilder`
-handles single-output (SFT / pretrain), :class:`MultiOutputMaskBuilder`
+orchestrates single-output / multi-output (DPO / GRPO) assembly.
-handles multi-output (DPO / GRPO), and :class:`SectionedMaskBuilder`
-orchestrates both modes as a façade.
 """
 
 from abc import ABC, abstractmethod
@@ -214,17 +212,42 @@ class MaskBuilderFactory(BaseFactory["BaseMaskBuilder"]):
     pass
 
 
-@MaskBuilderFactory.register("single")
+@MaskBuilderFactory.register("sectioned")
-class SingleOutputMaskBuilder(BaseMaskBuilder):
+class SectionedMaskBuilder(BaseMaskBuilder):
-    """Build a single output sequence with optional loss mask.
+    """Config-driven builder supporting single and multi-output modes.
 
-    Expects ``config.input.sections`` (list of section specs).
+    Single-output::
+
+        {"input": {"sections": [
+            {"field": "messages", "action": "$role", "template": true}
+        ]}}
+        → {"sequence": [...], "loss_mask": [...], "domain": "..."}
+
+    Multi-output (DPO / GRPO)::
+
+        {"input": {"sources": {
+            "chosen": {"sections": [{"field": "chosen", "action": "$role", "template": true}]},
+            "rejected": {"sections": [{"field": "rejected", "action": "$role", "template": true}]},
+        }}}
+        → {"chosen": [...], "chosen_mask": [...], "rejected": [...], "rejected_mask": [...], "domain": "..."}
+
+    Output spec fields::
+
+        sections      – list of section specs (same format as single-output)
+        list_field    – True when JSONL field holds a list (GRPO responses)
+        mask_key      – explicit loss-mask output key (default: ``"{output_key}_mask"``)
     """
 
-    def __init__(self, renderer: Optional[SectionRenderer] = None):
+    def __init__(self):
-        self.renderer = renderer or SectionRenderer()
+        self.renderer = SectionRenderer()
 
     def build(self, item: dict, config, tokenizer) -> Optional[dict]:
+        sources_spec = getattr(config.input, "sources", None)
+        if sources_spec:
+            return self._build_multi(item, sources_spec, config, tokenizer)
+        return self._build_single(item, config, tokenizer)
+
+    def _build_single(self, item: dict, config, tokenizer) -> Optional[dict]:
         sections = config.input.sections
         if not sections:
             return None
@@ -243,22 +266,9 @@ class SingleOutputMaskBuilder(BaseMaskBuilder):
             result["loss_mask"] = mask
         return result
 
-
+    def _build_multi(
-@MaskBuilderFactory.register("multi")
+        self, item: dict, sources_spec: dict, config, tokenizer
-class MultiOutputMaskBuilder(BaseMaskBuilder):
+    ) -> Optional[dict]:
-    """Build multiple output sequences (DPO / GRPO).
-
-    Expects ``config.input.sources`` (dict of output_key → spec).
-    """
-
-    def __init__(self, renderer: Optional[SectionRenderer] = None):
-        self.renderer = renderer or SectionRenderer()
-
-    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
-        sources_spec = getattr(config.input, "sources", None)
-        if not sources_spec:
-            return None
-
         result: dict = {}
         any_output = False
 
@@ -303,22 +313,3 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
 
         result["domain"] = _extract_domain(item, config.output.domain_key)
         return result
-
-
-@MaskBuilderFactory.register("sectioned")
-class SectionedMaskBuilder(BaseMaskBuilder):
-    """Façade that dispatches to SingleOutputMaskBuilder or MultiOutputMaskBuilder.
-
-    Preserves backward compatibility for existing configs and code that rely
-    on the ``"sectioned"`` factory name.
-    """
-
-    def __init__(self):
-        self._single = SingleOutputMaskBuilder()
-        self._multi = MultiOutputMaskBuilder()
-
-    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
-        sources_spec = getattr(config.input, "sources", None)
-        if sources_spec:
-            return self._multi.build(item, config, tokenizer)
-        return self._single.build(item, config, tokenizer)

astrai/serialization/dataset.py

@@ -26,9 +26,6 @@ def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
     tensor_group: Dict[str, List[Tensor]] = {}
 
     root_path = Path(file_path)
-    if root_path.is_file() and root_path.suffix in (".h5", ".hdf5"):
-        h5_files = [root_path]
-    else:
     h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
 
     for h5_file in h5_files:

astrai/trainer/train_callback.py

@@ -209,6 +209,7 @@ class ProgressBarCallback(TrainCallback):
 
     @only_on_rank(0)
     def on_optimizer_step(self, context: TrainContext):
+        self.progress_bar.update(1)
         postfix = {
             "step": context.optimizer_step,
             "loss": f"{context.loss:.4f}",
@@ -219,7 +220,6 @@ class ProgressBarCallback(TrainCallback):
         if context.val_loss is not None:
             postfix["val_loss"] = f"{context.val_loss:.4f}"
         self.progress_bar.set_postfix(postfix)
-        self.progress_bar.update(1)
 
     @only_on_rank(0)
     def on_epoch_end(self, context: TrainContext):

csrc/build.py

@@ -20,23 +20,13 @@ def _arch_flags() -> list[str]:
 _kernels_dir = Path("csrc/kernels")
 REGISTRY: dict[str, dict] = {}
 
-CXX_FLAGS = ["-O3", "-march=native", "-funroll-loops"]
-NVCC_FLAGS = [
-    "-O3",
-    "--expt-relaxed-constexpr",
-    "--use_fast_math",
-    "--ptxas-options=-O3,-v",
-    "--extra-device-vectorization",
-]
-
 
 def register(name: str, sources: list[str] | None = None, **kwargs):
     if sources is None:
         sources = [str(_kernels_dir / f"{name}.cu")]
     REGISTRY[name] = {
         "sources": sources,
-        "cxx_flags": [*CXX_FLAGS],
+        "nvcc_flags": ["-O3", "--expt-relaxed-constexpr", *_arch_flags()],
-        "nvcc_flags": [*NVCC_FLAGS, *_arch_flags()],
         "extra_link_args": kwargs.pop("extra_link_args", []),
         **kwargs,
     }

csrc/kernels/gqa_decode_attn.cu

@@ -1,43 +1,6 @@
 #include "gqa_decode_attn.cuh"
 #include <torch/extension.h>
 
-#ifndef ASTRAI_NO_MMA
-#include "gqa_decode_attn_mma.cuh"
-#endif
-
-template <int HEAD_DIM>
-static void dispatch_decode(GQAParams& p) {
-#ifndef ASTRAI_NO_MMA
-    constexpr int BC = 32, BR = 16, LD = HEAD_DIM;  // XOR swizzle → no padding
-    int G = p.q_head / p.kv_head;
-    // head-packing tensor-core path needs 1 < G <= 16 (MMA M dim) and no mask;
-    // everything else uses the scalar kernel
-    if (!p.use_mask && G > 1 && G <= 16) {
-        dim3 grid(p.kv_head, p.batch, 1);
-        dim3 block(32, 1, 1);
-        // sK + sV + sQ, each BC/BR * LD (single buffer for high occupancy)
-        int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
-        cudaFuncSetAttribute(gqa_decode_attn_mma_kernel<HEAD_DIM, BC>,
-                             cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
-        gqa_decode_attn_mma_kernel<HEAD_DIM, BC><<<grid, block, smem>>>(p);
-        return;
-    }
-    // scalar fallback (per-KV-head, one warp per query head)
-    int group_size = p.q_head / p.kv_head;
-    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
-    dim3 block(32, group_size);
-    dim3 grid(p.batch * p.kv_head);
-    gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
-#else
-    // scalar fallback (per-KV-head, one warp per query head)
-    int group_size = p.q_head / p.kv_head;
-    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
-    dim3 block(32, group_size);
-    dim3 grid(p.batch * p.kv_head);
-    gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
-#endif
-}
-
 torch::Tensor gqa_decode_attn(
     torch::Tensor q,
     torch::Tensor k, 
@@ -81,23 +44,12 @@ torch::Tensor gqa_decode_attn(
     auto O = torch::empty_like(q);
     p.o = (bf16*)O.data_ptr();
 
-    switch (p.head_dim) {
+    int group_size = p.q_head / p.kv_head;
-        case 32:
+    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
-            dispatch_decode<32>(p);
+    dim3 block(32, group_size);
-            break;
+    dim3 grid(p.batch * p.kv_head);
-        case 64:
+
-            dispatch_decode<64>(p);
+    gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
-            break;
-        case 128:
-            dispatch_decode<128>(p);
-            break;
-        case 256:
-            dispatch_decode<256>(p);
-            break;
-        default:
-            TORCH_CHECK(false, "decode: unsupported head_dim ", p.head_dim,
-                        " (supported: 32, 64, 128, 256)");
-    }
     return O;
 }
 
@@ -110,5 +62,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("is_causal") = false,
         py::arg("causal_offset") = 0,
         py::arg("scale") = py::none(),
-        "GQA decode (tensor-core head-packing on sm_80+, scalar fallback)");
+        "GQA decode (per-KV-head, shared K)");
 }

csrc/kernels/gqa_decode_attn_mma.cuh

@@ -1,219 +0,0 @@
-#pragma once
-#include "gqa_common.cuh"
-#include "gqa_mma_utils.cuh"
-
-// Tensor-core decode via GQA head-packing with cp.async loads.
-//
-// Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core work
-// on its own. But GQA gives us G = q_head / kv_head query heads that all share
-// one kv_head. We pack those G heads into the M=16 rows of mma.sync.m16n8k16,
-// turning G independent GEMVs into a single GEMM that reuses each loaded K/V tile
-// across all G heads (K/V load is the decode bottleneck, so the reuse is the win,
-// not the flops). Fragment layout is identical to the prefill mma kernel; the
-// only differences are (1) the M rows come from different heads at position 0
-// instead of different sequence positions of one head, and (2) causal masking is
-// a single scalar bound shared by every row. One warp owns one (batch, kv_head);
-// requires G <= 16.
-//
-// Optimizations:
-//   - cp.async global→shared for K/V (bypasses registers, cuts instruction count)
-//   - XOR swizzle (swiz_col): LD=HEAD_DIM, zero waste, no bank conflicts
-//   - pre-scaled Q: Q scaled during load, softmax skips per-tile multiply
-//   - single-buffer: keeps smem small for high occupancy
-
-template <int HEAD_DIM, int BC>
-__global__ void gqa_decode_attn_mma_kernel(GQAParams p) {
-    constexpr int BR = 16;
-    constexpr int KD = HEAD_DIM / 16;  // Q/K k-tiles
-    constexpr int NC8 = BC / 8;        // S n-tiles (N=8 each)
-    constexpr int KT2 = BC / 16;       // P k-tiles (K=16 each)
-    constexpr int DN8 = HEAD_DIM / 8;  // O n-tiles (N=8 each)
-    constexpr int LD = HEAD_DIM;       // XOR swizzle handles bank conflicts, zero waste
-    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
-
-    const int lane = threadIdx.x;  // single warp
-    const int gid = lane >> 2;     // 0..7 → rows gid, gid+8
-    const int tid4 = lane & 3;
-
-    const int kv_head = blockIdx.x;
-    const int batch = blockIdx.y;
-    const int G = p.q_head / p.kv_head;
-    const int q_head0 = kv_head * G;
-
-    extern __shared__ __align__(16) bf16 smem[];
-    bf16* sK = smem;                  // [BC][LD]
-    bf16* sV = sK + BC * LD;          // [BC][LD]
-    bf16* sQ = sV + BC * LD;          // [BR][LD]
-
-    // ---- stage Q into shared (pre-scaled, swizzled) ----
-    bf16 scale_bf16 = __float2bfloat16(p.scale);
-    for (int i = lane; i < BR * HEAD_DIM; i += 32) {
-        int r = i / HEAD_DIM, d = i % HEAD_DIM;
-        bf16 val = __float2bfloat16(0.0f);
-        if (r < G) {
-            int qh = q_head0 + r;
-            val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d];  // q_len == 1
-        }
-        sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16);
-    }
-    __syncwarp();
-
-    // Q resident A-fragments
-    unsigned Qa[KD][4];
-    int qrow_l = (lane & 7) + (lane & 8);
-    int qcol_l = (lane & 16) ? 8 : 0;
-#pragma unroll
-    for (int kt = 0; kt < KD; kt++)
-        ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
-
-    float Oacc[DN8][4];
-#pragma unroll
-    for (int j = 0; j < DN8; j++)
-        Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
-    float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
-
-    const int kv_base = (batch * p.kv_head + kv_head) * p.kv_len * HEAD_DIM;
-    const int mask_base = batch * p.kv_len;
-    const int tiles = (p.kv_len + BC - 1) / BC;
-    const int has_mask = p.use_mask && p.mask;
-
-    for (int ti = 0; ti < tiles; ti++) {
-        int kv0 = ti * BC;
-
-        // ---- load K/V tile to shared (cp.async on full tiles) ----
-        bool full_tile = (kv0 + BC <= p.kv_len);
-        if (full_tile) {
-            constexpr int VEC = 8;  // 8 bf16 = 16 bytes per cp.async
-            int total = BC * HEAD_DIM;
-#pragma unroll
-            for (int i = lane * VEC; i < total; i += 32 * VEC) {
-                int r = i / HEAD_DIM, d = i % HEAD_DIM;
-                int kc = kv0 + r;
-                cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)],
-                            &p.k[kv_base + kc * HEAD_DIM + d]);
-                cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)],
-                            &p.v[kv_base + kc * HEAD_DIM + d]);
-            }
-            cp_async_commit();
-            cp_async_wait_all();
-        } else {
-            for (int i = lane; i < BC * HEAD_DIM; i += 32) {
-                int r = i / HEAD_DIM, d = i % HEAD_DIM;
-                int kc = kv0 + r;
-                bf16 z = __float2bfloat16(0.0f);
-                sK[r * LD + swiz_col(d, r, SWIZ_MASK)] =
-                    (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z;
-                sV[r * LD + swiz_col(d, r, SWIZ_MASK)] =
-                    (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z;
-            }
-        }
-        __syncwarp();
-
-        // S = Q @ K^T  (Q already pre-scaled, so Sacc includes scale)
-        float Sacc[NC8][4];
-#pragma unroll
-        for (int n8 = 0; n8 < NC8; n8++) {
-            Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f;
-            int krow_l = n8 * 8 + (lane & 7);
-            int kcol_h = (lane & 8) ? 8 : 0;
-#pragma unroll
-            for (int kt = 0; kt < KD; kt++) {
-                unsigned b[2];
-                ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]);
-                mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]);
-            }
-        }
-
-        // ---- online softmax (Q pre-scaled → no per-tile scale multiply) ----
-        float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
-#pragma unroll
-        for (int n8 = 0; n8 < NC8; n8++) {
-            int cc = kv0 + n8 * 8 + 2 * tid4;
-            bool bc0 = (cc >= p.kv_len) ||
-                       (has_mask && !p.mask[mask_base + cc]);
-            bool bc1 = (cc + 1 >= p.kv_len) ||
-                       (has_mask && !p.mask[mask_base + cc + 1]);
-            bool cz = p.is_causal;
-            int off = p.causal_offset;
-            bool bad0 = bc0 || (cz && cc > off);
-            bool bad1 = bc1 || (cz && (cc + 1) > off);
-            float s0 = bad0 ? -FLT_MAX : Sacc[n8][0];
-            float s1 = bad1 ? -FLT_MAX : Sacc[n8][1];
-            float s2 = bad0 ? -FLT_MAX : Sacc[n8][2];
-            float s3 = bad1 ? -FLT_MAX : Sacc[n8][3];
-            Sacc[n8][0] = s0; Sacc[n8][1] = s1; Sacc[n8][2] = s2; Sacc[n8][3] = s3;
-            rmax0 = fmaxf(rmax0, fmaxf(s0, s1));
-            rmax1 = fmaxf(rmax1, fmaxf(s2, s3));
-        }
-        rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1));
-        rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2));
-        rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1));
-        rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2));
-
-        float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1);
-        float corr0 = (nm0 == -FLT_MAX) ? 1.0f : __expf(m0 - nm0);
-        float corr1 = (nm1 == -FLT_MAX) ? 1.0f : __expf(m1 - nm1);
-
-        float rsum0 = 0.0f, rsum1 = 0.0f;
-#pragma unroll
-        for (int n8 = 0; n8 < NC8; n8++) {
-            float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][0] - nm0);
-            float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][1] - nm0);
-            float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][2] - nm1);
-            float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][3] - nm1);
-            Sacc[n8][0] = p0; Sacc[n8][1] = p1; Sacc[n8][2] = p2; Sacc[n8][3] = p3;
-            rsum0 += p0 + p1;
-            rsum1 += p2 + p3;
-        }
-        rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1);
-        rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2);
-        rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1);
-        rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2);
-        l0 = l0 * corr0 + rsum0;
-        l1 = l1 * corr1 + rsum1;
-        m0 = nm0; m1 = nm1;
-
-#pragma unroll
-        for (int j = 0; j < DN8; j++) {
-            Oacc[j][0] *= corr0; Oacc[j][1] *= corr0;
-            Oacc[j][2] *= corr1; Oacc[j][3] *= corr1;
-        }
-
-        // O += P @ V
-#pragma unroll
-        for (int kt2 = 0; kt2 < KT2; kt2++) {
-            unsigned Pa[4];
-            Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]);
-            Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]);
-            Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]);
-            Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]);
-            int vrow_l = kt2 * 16 + (lane & 15);
-#pragma unroll
-            for (int dn8 = 0; dn8 < DN8; dn8++) {
-                unsigned b[2];
-                ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]);
-                mma16816(Oacc[dn8], Pa, b, Oacc[dn8]);
-            }
-        }
-        __syncwarp();  // sK/sV reused next tile
-    }
-
-    // ---- write output ----
-    float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
-    float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
-#pragma unroll
-    for (int dn8 = 0; dn8 < DN8; dn8++) {
-        int d = dn8 * 8 + 2 * tid4;
-        int r0 = gid, r1 = gid + 8;
-        if (r0 < G) {
-            int o_off = (batch * p.q_head + q_head0 + r0) * HEAD_DIM + d;
-            p.o[o_off] = __float2bfloat16(Oacc[dn8][0] * rl0);
-            p.o[o_off + 1] = __float2bfloat16(Oacc[dn8][1] * rl0);
-        }
-        if (r1 < G) {
-            int o_off = (batch * p.q_head + q_head0 + r1) * HEAD_DIM + d;
-            p.o[o_off] = __float2bfloat16(Oacc[dn8][2] * rl1);
-            p.o[o_off + 1] = __float2bfloat16(Oacc[dn8][3] * rl1);
-        }
-    }
-}

csrc/kernels/gqa_mma_utils.cuh

@@ -1,93 +0,0 @@
-#pragma once
-
-// Shared MMA utilities for tensor-core GQA kernels.
-// mma.sync.m16n8k16 PTX wrappers, ldmatrix helpers, and bf16 packing.
-
-// mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
-__device__ __forceinline__ void mma16816(float* d, const unsigned* a,
-                                          const unsigned* b, const float* c) {
-    asm volatile(
-        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
-        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};"
-        : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3])
-        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
-          "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
-}
-
-// read two adjacent bf16 from smem as one packed .b32 (elem0 low, elem1 high)
-__device__ __forceinline__ unsigned ld2(const bf16* p) {
-    return *reinterpret_cast<const unsigned*>(p);
-}
-
-// pack two floats into one bf16x2 as .b32
-__device__ __forceinline__ unsigned pk2(float a, float b) {
-    __nv_bfloat162 v = __floats2bfloat162_rn(a, b);
-    return *reinterpret_cast<unsigned*>(&v);
-}
-
-// pack two (non-contiguous) bf16 into one .b32
-__device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) {
-    __nv_bfloat162 v;
-    v.x = a;
-    v.y = b;
-    return *reinterpret_cast<unsigned*>(&v);
-}
-
-// ldmatrix: cooperatively load mma fragments from smem (one instruction per
-// 16x16 / 16x8 tile) with the exact register layout mma expects — replaces the
-// scalar per-thread fragment packing, cutting shared-load instructions and bank
-// conflicts. Each lane supplies the shared address of one 8-wide row.
-__device__ __forceinline__ void ldmatrix_x4(unsigned* r, const bf16* p) {
-    unsigned a = __cvta_generic_to_shared(p);
-    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
-                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
-                 : "r"(a));
-}
-__device__ __forceinline__ void ldmatrix_x2(unsigned* r, const bf16* p) {
-    unsigned a = __cvta_generic_to_shared(p);
-    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
-                 : "=r"(r[0]), "=r"(r[1])
-                 : "r"(a));
-}
-__device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) {
-    unsigned a = __cvta_generic_to_shared(p);
-    asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];"
-                 : "=r"(r[0]), "=r"(r[1])
-                 : "r"(a));
-}
-
-// XOR swizzle for shared-memory column at 8-bf16 chunk granularity.
-// Eliminates ldmatrix bank conflicts without LD padding: consecutive rows
-// land in distinct bank groups. swiz_col(d, r, mask) = ((d>>3)^(r&mask))<<3 | (d&7).
-// mask must cover log2(HEAD_DIM/8) chunk bits but stay within LD: use 7 for
-// HEAD_DIM>=64 (8+ chunks), 3 for HEAD_DIM=32 (4 chunks). Default 7 keeps
-// existing HEAD_DIM>=64 call sites working unchanged.
-__device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) {
-    return ((d >> 3) ^ (r & mask)) << 3 | (d & 7);
-}
-
-// cp.async: copy 16 bytes (8 bf16) from global to shared memory directly,
-// bypassing registers. Eliminates shared-store bank conflicts and cuts
-// load-loop instruction count in half (1 cp.async vs 1 LDG + 1 STS).
-// Requires sm_80+.
-__device__ __forceinline__ void cp_async_16(bf16* smem_ptr, const void* gmem_ptr) {
-    unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
-    asm volatile("cp.async.ca.shared.global [%0], [%1], 16;"
-                 :: "r"(smem_addr), "l"(gmem_ptr));
-}
-
-__device__ __forceinline__ void cp_async_commit() {
-    asm volatile("cp.async.commit_group;");
-}
-
-__device__ __forceinline__ void cp_async_wait_all() {
-    asm volatile("cp.async.wait_all;");
-}
-
-// Wait until at most N commit groups are still in flight. Used for
-// double-buffered pipelining: wait_group<1> lets the next tile's cp.async
-// continue while ensuring the current tile's data is ready.
-template <int N>
-__device__ __forceinline__ void cp_async_wait_group() {
-    asm volatile("cp.async.wait_group %0;" :: "n"(N));
-}

csrc/kernels/gqa_prefill_attn.cu

@@ -8,11 +8,10 @@
 template <int HEAD_DIM>
 static void dispatch_prefill(GQAParams& p) {
 #ifndef ASTRAI_NO_MMA
-    constexpr int WARPS = 4, BC = 32, BR = 16, LD = HEAD_DIM;
+    constexpr int WARPS = 4, BC = 32, BR = 16, LD = HEAD_DIM + 8;
     dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
     dim3 block(WARPS * 32, 1, 1);
-    // sK + sV (each BC*LD) + shared sQ staging (BR*LD)
+    int smem = (2 * BC * LD + WARPS * BR * LD) * (int)sizeof(bf16);
-    int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
     cudaFuncSetAttribute(gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC>,
                          cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
     gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC><<<grid, block, smem>>>(p);
@@ -68,9 +67,6 @@ torch::Tensor gqa_prefill_attn(
     p.o = (bf16*)O.data_ptr();
 
     switch (p.head_dim) {
-        case 32:
-            dispatch_prefill<32>(p);
-            break;
         case 64:
             dispatch_prefill<64>(p);
             break;
@@ -82,7 +78,7 @@ torch::Tensor gqa_prefill_attn(
             break;
         default:
             TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim,
-                        " (supported: 32,64,128,256)");
+                        " (supported: 64,128,256)");
     }
     return O;
 }

csrc/kernels/gqa_prefill_attn_mma.cuh

@@ -1,6 +1,5 @@
 #pragma once
 #include "gqa_common.cuh"
-#include "gqa_mma_utils.cuh"
 
 // Tensor-core prefill, register-resident flash attention (raw mma.sync PTX).
 // One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
@@ -11,12 +10,57 @@
 // for-element onto the P matrix_a (bf16) operand, so softmax needs no shuffle
 // repack; row reductions fold across the 4-lane thread group. Templated on
 // <HEAD_DIM, WARPS, BC> with BC a multiple of 16.
-//
+
-// Optimizations: shared sQ staging (single area, serialized per-warp load)
+// mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
-// → cuts smem; pre-scale Q by attention scale during Q load; cp.async global→
+// (only compiled when ASTRAI_HAS_MMA is set, i.e. built for sm_80+)
-// shared for K/V; scalar fallback only for the last partial tile; causal tile
+__device__ __forceinline__ void mma16816(float* d, const unsigned* a,
-// skipping (block-level early break + warp-level skip); XOR swizzle (swiz_col)
+                                          const unsigned* b, const float* c) {
-// → eliminates ldmatrix bank conflicts without LD padding (LD=HEAD_DIM).
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};"
+        : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3])
+        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+          "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+}
+
+// read two adjacent bf16 from smem as one packed .b32 (elem0 low, elem1 high)
+__device__ __forceinline__ unsigned ld2(const bf16* p) {
+    return *reinterpret_cast<const unsigned*>(p);
+}
+__device__ __forceinline__ unsigned pk2(float a, float b) {
+    __nv_bfloat162 v = __floats2bfloat162_rn(a, b);
+    return *reinterpret_cast<unsigned*>(&v);
+}
+// pack two (non-contiguous) bf16 into one .b32
+__device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) {
+    __nv_bfloat162 v;
+    v.x = a;
+    v.y = b;
+    return *reinterpret_cast<unsigned*>(&v);
+}
+
+// ldmatrix: cooperatively load mma fragments from smem (one instruction per
+// 16x16 / 16x8 tile) with the exact register layout mma expects — replaces the
+// scalar per-thread fragment packing, cutting shared-load instructions and bank
+// conflicts. Each lane supplies the shared address of one 8-wide row.
+__device__ __forceinline__ void ldmatrix_x4(unsigned* r, const bf16* p) {
+    unsigned a = __cvta_generic_to_shared(p);
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
+                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
+                 : "r"(a));
+}
+__device__ __forceinline__ void ldmatrix_x2(unsigned* r, const bf16* p) {
+    unsigned a = __cvta_generic_to_shared(p);
+    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
+                 : "=r"(r[0]), "=r"(r[1])
+                 : "r"(a));
+}
+__device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) {
+    unsigned a = __cvta_generic_to_shared(p);
+    asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];"
+                 : "=r"(r[0]), "=r"(r[1])
+                 : "r"(a));
+}
 
 template <int HEAD_DIM, int WARPS, int BC>
 __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
@@ -25,8 +69,9 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
     constexpr int NC8 = BC / 8;        // S n-tiles (N=8 each)
     constexpr int KT2 = BC / 16;       // P k-tiles (K=16 each)
     constexpr int DN8 = HEAD_DIM / 8;  // O n-tiles (N=8 each)
-    constexpr int LD = HEAD_DIM;   // XOR swizzle (swiz_col) handles bank conflicts
+    constexpr int LD = HEAD_DIM + 8;   // padded smem row stride (kills ldmatrix
-    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);  // chunk bits, stay within LD
+                                       // bank conflicts: consecutive rows land
+                                       // in distinct banks instead of colliding)
 
     const int warp = threadIdx.x / 32;
     const int lane = threadIdx.x % 32;
@@ -42,31 +87,24 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
     extern __shared__ __align__(16) bf16 smem[];
     bf16* sK = smem;                                          // [BC][LD]
     bf16* sV = sK + BC * LD;                                  // [BC][LD]
-    bf16* sQ = sV + BC * LD;               // shared staging [BR][LD]
+    bf16* sQ = sV + BC * LD + warp * (BR * LD);               // per-warp [BR][LD]
 
-    // Q resident A-fragments (loaded once per warp via shared staging).
+    // stage Q into smem (zero-padded past q_len)
-    // Pre-scale by attention scale so softmax doesn't need to multiply later.
     const int q_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
-    unsigned Qa[KD][4];
-    bf16 scale_bf16 = __float2bfloat16(p.scale);
-    int qrow_l = (lane & 7) + (lane & 8);       // 0..15
-    int qcol_l = (lane & 16) ? 8 : 0;
-    for (int w = 0; w < WARPS; w++) {
-        if (warp == w) {
     for (int i = lane; i < BR * HEAD_DIM; i += 32) {
         int r = i / HEAD_DIM, d = i % HEAD_DIM;
         int qr = qrow0 + r;
-                bf16 qv = (qr < p.q_len) ? p.q[q_base + qr * HEAD_DIM + d]
+        sQ[r * LD + d] = (qr < p.q_len) ? p.q[q_base + qr * HEAD_DIM + d] : __float2bfloat16(0.0f);
-                                         : __float2bfloat16(0.0f);
-                sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(qv, scale_bf16);
     }
     __syncwarp();
+
+    // Q resident A-fragments: Qa[kt][0..3] (loaded once via ldmatrix.x4)
+    unsigned Qa[KD][4];
+    int qrow_l = (lane & 7) + (lane & 8);       // 0..15
+    int qcol_l = (lane & 16) ? 8 : 0;
 #pragma unroll
     for (int kt = 0; kt < KD; kt++)
-                ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
+        ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + kt * 16 + qcol_l]);
-        }
-        __syncthreads();  // prevent next warp from overwriting sQ prematurely
-    }
 
     float Oacc[DN8][4];
 #pragma unroll
@@ -79,90 +117,60 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
     const int qr0 = qrow0 + gid;      // row for c0/c1
     const int qr1 = qrow0 + gid + 8;  // row for c2/c3
 
-    // Causal tile-skip bounds (no-op when is_causal == 0)
-    const int use_skip = p.is_causal;
-    const int max_kv = qrow0 + BR - 1 + p.causal_offset;
-    const int block_max_kv =
-        blockIdx.x * WARPS * BR + WARPS * BR - 1 + p.causal_offset;
-    const int has_mask = p.use_mask && p.mask;
-    const int mb = batch * p.kv_len;
-
     for (int ti = 0; ti < tiles; ti++) {
         int kv0 = ti * BC;
 
-        // Block-level causal early break
-        if (use_skip && kv0 > block_max_kv) break;
-
-        // ---- load K/V tile to shared memory (cp.async on full tiles) ----
-        bool full_tile = (kv0 + BC <= p.kv_len);
-        if (full_tile) {
-            constexpr int VEC = 8;  // bf16 per cp.async unit (16 bytes)
-            int total = BC * HEAD_DIM;
-#pragma unroll
-            for (int i = threadIdx.x * VEC; i < total; i += nthreads * VEC) {
-                int r = i / HEAD_DIM;
-                int d = i % HEAD_DIM;
-                int kc = kv0 + r;
-                cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)], &p.k[kv_base + kc * HEAD_DIM + d]);
-                cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)], &p.v[kv_base + kc * HEAD_DIM + d]);
-            }
-            cp_async_commit();
-            cp_async_wait_all();
-        } else {
         for (int i = threadIdx.x; i < BC * HEAD_DIM; i += nthreads) {
             int r = i / HEAD_DIM, d = i % HEAD_DIM;
             int kc = kv0 + r;
             bf16 z = __float2bfloat16(0.0f);
-                sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = (kc < p.kv_len)
+            sK[r * LD + d] = (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z;
-                                     ? p.k[kv_base + kc * HEAD_DIM + d] : z;
+            sV[r * LD + d] = (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z;
-                sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = (kc < p.kv_len)
-                                     ? p.v[kv_base + kc * HEAD_DIM + d] : z;
-            }
         }
         __syncthreads();
 
-        // Warp-level causal skip
-        if (!use_skip || kv0 <= max_kv) {
-
         // S = Q @ K^T  → Sacc[n8][0..3]   (n8: 8 kv cols each)
         float Sacc[NC8][4];
 #pragma unroll
         for (int n8 = 0; n8 < NC8; n8++) {
             Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f;
-            int krow_l = n8 * 8 + (lane & 7);
+            int kv = kv0 + n8 * 8 + gid;
-            int kcol_h = (lane & 8) ? 8 : 0;
+            int krow_l = n8 * 8 + (lane & 7);        // kv within tile
+            int kcol_h = (lane & 8) ? 8 : 0;         // which k-half
 #pragma unroll
             for (int kt = 0; kt < KD; kt++) {
                 unsigned b[2];
-                ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]);
+                ldmatrix_x2(b, &sK[krow_l * LD + kt * 16 + kcol_h]);
                 mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]);
             }
+            (void)kv;
         }
 
         // ---- online softmax (in registers) ----
-        // Q is pre-scaled, so Sacc already includes the attention scale.
+        // scale + mask, then per-row (gid, gid+8) max over held cols
-        int maxc0 = p.is_causal ? min(p.kv_len, qr0 + p.causal_offset + 1)
-                                : p.kv_len;
-        int maxc1 = p.is_causal ? min(p.kv_len, qr1 + p.causal_offset + 1)
-                                : p.kv_len;
         float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
 #pragma unroll
         for (int n8 = 0; n8 < NC8; n8++) {
-            int cc = kv0 + n8 * 8 + 2 * tid4;
+            int cc = kv0 + n8 * 8 + 2 * tid4;  // col for c0/c2
-            int c1 = cc + 1;
+            bool bc0 = (cc >= p.kv_len) ||
-            bool b0 = (cc >= maxc0) || (has_mask && !p.mask[mb + cc]);
+                       (p.use_mask && p.mask && !p.mask[batch * p.kv_len + cc]);
-            bool b1 = (c1 >= maxc0) || (has_mask && !p.mask[mb + c1]);
+            bool bc1 = (cc + 1 >= p.kv_len) ||
-            bool b2 = (cc >= maxc1) || (has_mask && !p.mask[mb + cc]);
+                       (p.use_mask && p.mask && !p.mask[batch * p.kv_len + cc + 1]);
-            bool b3 = (c1 >= maxc1) || (has_mask && !p.mask[mb + c1]);
+            bool cz = p.is_causal;
-            float s0 = b0 ? -FLT_MAX : Sacc[n8][0];
+            int off = p.causal_offset;
-            float s1 = b1 ? -FLT_MAX : Sacc[n8][1];
+            bool bad0 = bc0 || (cz && cc > qr0 + off);
-            float s2 = b2 ? -FLT_MAX : Sacc[n8][2];
+            bool bad1 = bc1 || (cz && (cc + 1) > qr0 + off);
-            float s3 = b3 ? -FLT_MAX : Sacc[n8][3];
+            bool bad2 = bc0 || (cz && cc > qr1 + off);
-            Sacc[n8][0] = s0; Sacc[n8][1] = s1;
+            bool bad3 = bc1 || (cz && (cc + 1) > qr1 + off);
-            Sacc[n8][2] = s2; Sacc[n8][3] = s3;
+            float s0 = bad0 ? -FLT_MAX : Sacc[n8][0] * p.scale;
+            float s1 = bad1 ? -FLT_MAX : Sacc[n8][1] * p.scale;
+            float s2 = bad2 ? -FLT_MAX : Sacc[n8][2] * p.scale;
+            float s3 = bad3 ? -FLT_MAX : Sacc[n8][3] * p.scale;
+            Sacc[n8][0] = s0; Sacc[n8][1] = s1; Sacc[n8][2] = s2; Sacc[n8][3] = s3;
             rmax0 = fmaxf(rmax0, fmaxf(s0, s1));
             rmax1 = fmaxf(rmax1, fmaxf(s2, s3));
         }
+        // reduce max across the 4-lane group (tid4)
         rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1));
         rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2));
         rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1));
@@ -175,16 +183,11 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
         float rsum0 = 0.0f, rsum1 = 0.0f;
 #pragma unroll
         for (int n8 = 0; n8 < NC8; n8++) {
-            float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f
+            float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][0] - nm0);
-                                                 : __expf(Sacc[n8][0] - nm0);
+            float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][1] - nm0);
-            float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f
+            float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][2] - nm1);
-                                                 : __expf(Sacc[n8][1] - nm0);
+            float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][3] - nm1);
-            float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f
+            Sacc[n8][0] = p0; Sacc[n8][1] = p1; Sacc[n8][2] = p2; Sacc[n8][3] = p3;
-                                                 : __expf(Sacc[n8][2] - nm1);
-            float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f
-                                                 : __expf(Sacc[n8][3] - nm1);
-            Sacc[n8][0] = p0; Sacc[n8][1] = p1;
-            Sacc[n8][2] = p2; Sacc[n8][3] = p3;
             rsum0 += p0 + p1;
             rsum1 += p2 + p3;
         }
@@ -211,16 +214,15 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
             Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]);
             Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]);
             Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]);
-            int vrow_l = kt2 * 16 + (lane & 15);
+            int vrow_l = kt2 * 16 + (lane & 15);  // kv within tile (0..15)
 #pragma unroll
             for (int dn8 = 0; dn8 < DN8; dn8++) {
                 unsigned b[2];
-                ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]);
+                ldmatrix_x2_trans(b, &sV[vrow_l * LD + dn8 * 8]);
                 mma16816(Oacc[dn8], Pa, b, Oacc[dn8]);
             }
         }
-        }  // if active (warp-level causal skip)
+        __syncthreads();  // sK/sV reused next tile
-        __syncthreads();
     }
 
     // ---- write output ----
@@ -231,16 +233,12 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
     for (int dn8 = 0; dn8 < DN8; dn8++) {
         int d = dn8 * 8 + 2 * tid4;
         if (qr0 < p.q_len) {
-            p.o[o_base + qr0 * HEAD_DIM + d] =
+            p.o[o_base + qr0 * HEAD_DIM + d] = __float2bfloat16(Oacc[dn8][0] * rl0);
-                __float2bfloat16(Oacc[dn8][0] * rl0);
+            p.o[o_base + qr0 * HEAD_DIM + d + 1] = __float2bfloat16(Oacc[dn8][1] * rl0);
-            p.o[o_base + qr0 * HEAD_DIM + d + 1] =
-                __float2bfloat16(Oacc[dn8][1] * rl0);
         }
         if (qr1 < p.q_len) {
-            p.o[o_base + qr1 * HEAD_DIM + d] =
+            p.o[o_base + qr1 * HEAD_DIM + d] = __float2bfloat16(Oacc[dn8][2] * rl1);
-                __float2bfloat16(Oacc[dn8][2] * rl1);
+            p.o[o_base + qr1 * HEAD_DIM + d + 1] = __float2bfloat16(Oacc[dn8][3] * rl1);
-            p.o[o_base + qr1 * HEAD_DIM + d + 1] =
-                __float2bfloat16(Oacc[dn8][3] * rl1);
         }
     }
 }

csrc/tests/gqa_decode_test.cu

@@ -1,10 +1,4 @@
-/*
+// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_decode_test.cu -o test && ./test
-Pure-C test:
-nvcc -I csrc -arch=sm_89 -O3 \
-    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
-    csrc/tests/gqa_decode_test.cu -o test && ./test
-*/
-
 #include <cstdio>
 #include <cstdlib>
 #include <cmath>

csrc/tests/gqa_prefill_test.cu

@@ -1,10 +1,4 @@
-/*
+// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_prefill_test.cu -o test && ./test
-Pure-C test:
-nvcc -I csrc -arch=sm_89 -O3 \
-    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
-    csrc/tests/gqa_prefill_test.cu -o test && ./test
-*/
-
 #include <cstdio>
 #include <cstdlib>
 #include <cmath>

scripts/tools/train.py

@@ -1,11 +1,9 @@
 import argparse
 import os
 from functools import partial
-from typing import Any, Dict
 
 import torch
 import torch.optim as optim
-from torch import Tensor, nn
 
 from astrai.config import AutoRegressiveLMConfig, TrainConfig
 from astrai.dataset import DatasetFactory
@@ -14,84 +12,6 @@ from astrai.model.components.decoder_block import DecoderBlock
 from astrai.trainer import SchedulerFactory, Trainer
 
 
-class MuonMix(optim.Optimizer):
-    """Combined Muon (matrix) + AdamW (non-matrix) optimizer."""
-
-    def __init__(
-        self,
-        model: nn.Module,
-        lr: float = 3e-4,
-        weight_decay: float = 0.1,
-        momentum: float = 0.95,
-        nesterov: bool = True,
-        ns_steps: int = 5,
-        adjust_lr_fn: str = "match_rms_adamw",
-    ):
-        defaults = dict(
-            lr=lr,
-            weight_decay=weight_decay,
-            momentum=momentum,
-            nesterov=nesterov,
-            ns_steps=ns_steps,
-            adjust_lr_fn=adjust_lr_fn,
-        )
-        params = [p for p in model.parameters() if p.requires_grad]
-        super().__init__(params, defaults)
-
-        matrix_params: list[Tensor] = []
-        other_params: list[Tensor] = []
-        for name, param in model.named_parameters():
-            if not param.requires_grad:
-                continue
-            if (
-                param.dim() >= 2
-                and "norm" not in name
-                and "bias" not in name
-                and "embed" not in name
-                and "lm_head" not in name
-            ):
-                matrix_params.append(param)
-            else:
-                other_params.append(param)
-
-        self.muon = optim.Muon(
-            matrix_params,
-            lr=lr,
-            weight_decay=weight_decay,
-            momentum=momentum,
-            nesterov=nesterov,
-            ns_steps=ns_steps,
-            adjust_lr_fn=adjust_lr_fn,
-        )
-        self.adamw = optim.AdamW(
-            [{"params": other_params, "weight_decay": 0.0}],
-            lr=lr,
-            betas=(0.9, 0.95),
-            fused=True,
-        )
-
-        self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]
-
-    @torch.no_grad()
-    def step(self, closure=None):
-        self.muon.step(closure)
-        self.adamw.step(closure)
-
-    def zero_grad(self, set_to_none: bool = True):
-        self.muon.zero_grad(set_to_none)
-        self.adamw.zero_grad(set_to_none)
-
-    def state_dict(self) -> Dict[str, Any]:
-        return {
-            "muon": self.muon.state_dict(),
-            "adamw": self.adamw.state_dict(),
-        }
-
-    def load_state_dict(self, state_dict: Dict[str, Any]):
-        self.muon.load_state_dict(state_dict["muon"])
-        self.adamw.load_state_dict(state_dict["adamw"])
-
-
 def parse_args() -> argparse.Namespace:
 
     parser = argparse.ArgumentParser(description="Train the AutoRegressiveLM model.")
@@ -144,35 +64,22 @@ def parse_args() -> argparse.Namespace:
         help="Max gradient norm for clipping.",
     )
     parser.add_argument(
-        "--weight_decay",
+        "--adamw_beta1",
         type=float,
-        default=0.1,
+        default=0.9,
-        help="Weight decay (applied to Muon matrix params; non-matrix use 0).",
+        help="Beta1 for AdamW optimizer.",
     )
     parser.add_argument(
-        "--muon_momentum",
+        "--adamw_beta2",
         type=float,
         default=0.95,
-        help="Momentum factor for Muon optimizer.",
+        help="Beta2 for AdamW optimizer.",
     )
     parser.add_argument(
-        "--muon_nesterov",
+        "--adamw_weight_decay",
-        action=argparse.BooleanOptionalAction,
+        type=float,
-        default=True,
+        default=0.01,
-        help="Enable Nesterov momentum for Muon.",
+        help="Weight decay for AdamW optimizer.",
-    )
-    parser.add_argument(
-        "--muon_ns_steps",
-        type=int,
-        default=5,
-        help="Newton-Schulz iteration steps for Muon.",
-    )
-    parser.add_argument(
-        "--muon_adjust_lr",
-        type=str,
-        default="match_rms_adamw",
-        choices=["original", "match_rms_adamw"],
-        help="Muon learning rate adjustment strategy.",
     )
     parser.add_argument(
         "--random_seed", type=int, default=3407, help="Random seed for reproducibility."
@@ -358,8 +265,21 @@ def create_model(config):
     return AutoRegressiveLM(config).to(dtype=torch.bfloat16)
 
 
-def create_optimizer(model, **kwargs) -> MuonMix:
+def create_optimizer(model, **kwargs) -> optim.Optimizer:
-    return MuonMix(model, **kwargs)
+    decay_params = []
+    no_decay_params = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if param.dim() < 2 or "norm" in name or "bias" in name:
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+    param_groups = [
+        {"params": decay_params, "weight_decay": kwargs.pop("weight_decay", 0.01)},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+    return optim.AdamW(param_groups, fused=True, **kwargs)
 
 
 def create_scheduler(
@@ -390,6 +310,7 @@ def train(
     train_type: str,
     param_path: str,
     data_root_path: str,
+    max_lr: float,
     n_epoch: int,
     batch_per_device: int,
     start_epoch: int,
@@ -402,7 +323,16 @@ def train(
     val_step: int,
     metrics: list[str],
     log_dir: str,
+    dpo_beta: float,
+    grpo_clip_eps: float,
+    grpo_kl_coef: float,
+    group_size: int,
+    grpo_sync_interval: int,
+    adamw_beta1: float,
+    adamw_beta2: float,
+    adamw_weight_decay: float,
     max_grad_norm: float,
+    label_smoothing: float,
     random_seed: int,
     num_workers: int,
     pin_memory: bool,
@@ -423,7 +353,6 @@ def train(
     t_mult: int,
     stable_steps: int,
     decay_steps: int,
-    **kwargs,
 ):
     assert train_type in ["seq", "sft", "dpo", "grpo"]
     assert os.path.exists(param_path)
@@ -439,12 +368,12 @@ def train(
         window_size = config.max_len
 
     strategy_kwargs = {
-        "beta": kwargs.pop("dpo_beta"),
+        "beta": dpo_beta,
-        "label_smoothing": kwargs.pop("label_smoothing"),
+        "label_smoothing": label_smoothing,
-        "clip_eps": kwargs.pop("grpo_clip_eps"),
+        "clip_eps": grpo_clip_eps,
-        "kl_coef": kwargs.pop("grpo_kl_coef"),
+        "kl_coef": grpo_kl_coef,
-        "group_size": kwargs.pop("group_size"),
+        "group_size": group_size,
-        "sync_interval": kwargs.pop("grpo_sync_interval"),
+        "sync_interval": grpo_sync_interval,
     }
 
     executor_kwargs = {
@@ -462,12 +391,11 @@ def train(
 
     optimizer_fn = partial(
         create_optimizer,
-        lr=kwargs.pop("max_lr"),
+        **{
-        weight_decay=kwargs.pop("weight_decay"),
+            "lr": max_lr,
-        momentum=kwargs.pop("muon_momentum"),
+            "betas": (adamw_beta1, adamw_beta2),
-        nesterov=kwargs.pop("muon_nesterov"),
+            "weight_decay": adamw_weight_decay,
-        ns_steps=kwargs.pop("muon_ns_steps"),
+        },
-        adjust_lr_fn=kwargs.pop("muon_adjust_lr"),
     )
 
     total_steps = compute_total_steps(

setup.py

@@ -41,10 +41,7 @@ if _should_build():
             CUDAExtension(
                 f"astrai.extension.{name}",
                 info["sources"],
-                extra_compile_args={
+                extra_compile_args={"cxx": ["-O3"], "nvcc": info["nvcc_flags"]},
-                    "cxx": info["cxx_flags"],
-                    "nvcc": info["nvcc_flags"],
-                },
                 extra_link_args=[f"-Wl,-rpath,{_torch_lib}"],
             )
         )

tests/data/conftest.py

@@ -10,11 +10,7 @@ from astrai.config.preprocess_config import (
     PipelineConfig,
     ProcessingConfig,
 )
-from astrai.preprocessing.builder import (
+from astrai.preprocessing.builder import SectionedMaskBuilder
-    MultiOutputMaskBuilder,
-    SectionedMaskBuilder,
-    SingleOutputMaskBuilder,
-)
 from astrai.tokenize import AutoTokenizer
 
 _SPECIAL_TOKENS_CONFIG = {
@@ -214,16 +210,6 @@ def builder():
     return SectionedMaskBuilder()
 
 
-@pytest.fixture
-def single_builder():
-    return SingleOutputMaskBuilder()
-
-
-@pytest.fixture
-def multi_builder():
-    return MultiOutputMaskBuilder()
-
-
 @pytest.fixture
 def tokenizer_dir(temp_dir, test_tokenizer):
     d = os.path.join(temp_dir, "tok")

tests/data/test_preprocess_builder.py

@@ -8,9 +8,7 @@ from astrai.config.preprocess_config import (
 )
 from astrai.preprocessing.builder import (
     MaskBuilderFactory,
-    MultiOutputMaskBuilder,
     SectionedMaskBuilder,
-    SingleOutputMaskBuilder,
 )
 from tests.data.conftest import (
     _CHAT_SECTIONS,
@@ -274,18 +272,12 @@ def test_sectioned_text_too_short(test_tokenizer, builder):
 
 def test_factory_registered():
     names = MaskBuilderFactory.list_registered()
-    assert "single" in names
-    assert "multi" in names
     assert "sectioned" in names
 
 
 def test_factory_create():
-    single = MaskBuilderFactory.create("single")
+    builder_obj = MaskBuilderFactory.create("sectioned")
-    assert isinstance(single, SingleOutputMaskBuilder)
+    assert isinstance(builder_obj, SectionedMaskBuilder)
-    multi = MaskBuilderFactory.create("multi")
-    assert isinstance(multi, MultiOutputMaskBuilder)
-    sectioned = MaskBuilderFactory.create("sectioned")
-    assert isinstance(sectioned, SectionedMaskBuilder)
 
 
 def test_dpo_chat_basic(chat_tokenizer, builder):
@@ -375,59 +367,3 @@ def test_grpo_single_reward(chat_tokenizer, builder):
     }
     result = builder.build(item, config, chat_tokenizer)
     assert result["rewards"] == [0.9]
-
-
-def test_single_builder_matches_facade(chat_tokenizer, builder, single_builder):
-    config = make_chat_config()
-    item = {
-        "messages": [
-            {"role": "user", "content": "What is 2+2?"},
-            {"role": "assistant", "content": "4"},
-        ]
-    }
-    facade_result = builder.build(item, config, chat_tokenizer)
-    single_result = single_builder.build(item, config, chat_tokenizer)
-    assert single_result == facade_result
-
-
-def test_single_builder_rejects_multi_config(chat_tokenizer, single_builder):
-    config = make_dpo_chat_config()
-    item = {
-        "chosen": [
-            {"role": "user", "content": "What is 2+2?"},
-            {"role": "assistant", "content": "4"},
-        ],
-        "rejected": [
-            {"role": "user", "content": "What is 2+2?"},
-            {"role": "assistant", "content": "5"},
-        ],
-    }
-    assert single_builder.build(item, config, chat_tokenizer) is None
-
-
-def test_multi_builder_matches_facade(chat_tokenizer, builder, multi_builder):
-    config = make_dpo_chat_config()
-    item = {
-        "chosen": [
-            {"role": "user", "content": "What is 2+2?"},
-            {"role": "assistant", "content": "4"},
-        ],
-        "rejected": [
-            {"role": "user", "content": "What is 2+2?"},
-            {"role": "assistant", "content": "5"},
-        ],
-    }
-    facade_result = builder.build(item, config, chat_tokenizer)
-    multi_result = multi_builder.build(item, config, chat_tokenizer)
-    assert multi_result == facade_result
-
-
-def test_multi_builder_rejects_single_config(chat_tokenizer, multi_builder):
-    config = make_chat_config()
-    item = {
-        "messages": [
-            {"role": "user", "content": "What is 2+2?"},
-            {"role": "assistant", "content": "4"},
-        ]
-    }
-    assert multi_builder.build(item, config, chat_tokenizer) is None