docs: fix documentation errors across README and assets/docs

- Correct training CLI args: remove non-existent --adamw_beta1/2, fix --weight_decay
- Fix optimizer section: document MuonMix instead of plain AdamW
- Fix inference.md decode phase description (all groups, not largest)
- Fix dataflow.md H5Store description (no share_memory_ in code)
- Fix architecture.md class diagram: remove Task.stream_callback,
  EncoderConfig.use_gated_attention; update stop_ids docs
- Fix --min_rate default value description to match code (0.01)
- Update all document timestamps to 2026-07-09

.gitignore

@@ -7,8 +7,13 @@
 # Allow specific file types and root files
 !astrai/**/*.py
 !scripts/**/*.py
-!scripts/**/*.sh
 !tests/**/*.py
+!csrc/**/*.py
+
+!csrc/**/*.cu
+!csrc/**/*.cuh
+
+!scripts/**/*.sh
 
 # Allow GitHub files
 !/.github/**
@@ -22,4 +27,9 @@
 !/CONTRIBUTING.md
 !/LICENSE
 !/pyproject.toml
-!/README.md
+!/README.md
+# Allow extension modules (only source .py)
+!/astrai/extension/**/*.py
+
+# Allow build files
+!/setup.py

README.md

@@ -9,7 +9,7 @@
 <div align="center">
   <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
   <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/v/tag/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
   <img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
   <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
 </div>
@@ -59,8 +59,9 @@ End-to-end walkthrough in 5 steps:
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
-pip install -e .
-# pip install -e ".[dev]"    # optional: dev dependencies (pytest, ruff)
+pip install -e .                                          # pure PyTorch (no CUDA kernels)
+# CSRC_KERNELS=true pip install -e . --no-build-isolation  # optional: fused CUDA kernels
+# pip install -e ".[dev]"                                  # dev dependencies (pytest, ruff)
 ```
 
 **2. Download model**
@@ -102,9 +103,7 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --adamw_beta1=0.9 \
-    --adamw_beta2=0.95 \
-    --adamw_weight_decay=0.01 \
+    --weight_decay=0.1 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \

assets/docs/README-zh-CN.md

@@ -15,7 +15,7 @@
 <div align="center">
   <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
   <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
+  <img src="https://img.shields.io/github/v/tag/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
   <img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
   <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
 </div>
@@ -65,8 +65,9 @@
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
-pip install -e .
-# pip install -e ".[dev]"    # 可选：开发依赖（pytest, ruff）
+pip install -e .                                          # 纯 PyTorch（不含 CUDA 内核）
+# CSRC_KERNELS=true pip install -e . --no-build-isolation  # 可选：融合 CUDA 内核加速
+# pip install -e ".[dev]"                                  # 可选：开发依赖（pytest, ruff）
 ```
 
 **2. 下载模型**
@@ -108,9 +109,7 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --adamw_beta1=0.9 \
-    --adamw_beta2=0.95 \
-    --adamw_weight_decay=0.01 \
+    --weight_decay=0.1 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \

assets/docs/architecture.md

@@ -63,7 +63,6 @@ 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
@@ -125,7 +124,6 @@ classDiagram
             +str ckpt_dir
             +int ckpt_interval
             +str log_dir
-            +int log_interval
             +List[str] metrics
             +Optional[LoRAConfig] lora
             +int random_seed
@@ -559,7 +557,7 @@ classDiagram
         }
 
         class GradientCheckpointingCallback {
-            +tuple modules
+            +Optional[List[type]] modules
             +on_train_begin(context)
             +on_train_end(context)
         }
@@ -573,31 +571,29 @@ classDiagram
             +on_batch_end(context)
             +on_train_end(context)
             +on_error(context)
-            +save_extra(context) dict$
+            +save_extra(context) dict
         }
 
         class ProgressBarCallback {
             +int num_epoch
             +int log_interval
             +IO file
+            +tqdm progress_bar
             +on_epoch_begin(context)
-            +on_batch_end(context)
+            +on_optimizer_step(context)
             +on_epoch_end(context)
         }
 
-        class MetricLoggerCallback {
+        class MetricCallback {
             +Path log_dir
             +int save_interval
-            +int log_interval
             +List[str] metrics
-            +on_batch_end(context)
+            +int val_step
+            +on_optimizer_step(context)
+            +on_epoch_end(context)
             +on_train_end(context)
             +on_error(context)
-        }
-
-        class ValidationCallback {
             -_run_validation(context)
-            +on_optimizer_step(context)
         }
 
         class CallbackFactory {
@@ -684,20 +680,44 @@ classDiagram
         }
 
         class KVCache {
-            -PagePool _pool
-            -Storage _storage
-            -TaskTable _table
-            +int page_size
+            <<abstract>>
             +task_alloc(task_id, prompt_ids) bool
             +task_free(task_id)
             +task_extend(task_id, pos) bool
             +task_cached(task_id) int
             +task_record_hashes(task_id, prompt_ids, start_logical_page)
-            +make_table_tensor(task_ids, device) Tensor
-            +bind(page_table, total_len) KvcacheView
+            +bind_tasks(task_ids, total_len, device) CacheView
         }
 
-        class KvcacheView {
+        class PageCache {
+            +int page_size
+            -PagePool _pool
+            -Storage _storage
+            -TaskTable _table
+            +task_alloc(task_id, prompt_ids) bool
+            +task_free(task_id)
+            +task_extend(task_id, pos) bool
+            +task_cached(task_id) int
+            +task_record_hashes(task_id, prompt_ids, start_logical_page)
+            +bind_tasks(task_ids, total_len, device) PageCacheView
+        }
+
+        class ContiguousCache {
+            +int max_seq_len
+            +Tensor k, v
+            +task_alloc(task_id, prompt_ids) bool
+            +task_free(task_id)
+            +task_extend(task_id, pos) bool
+            +bind_tasks(task_ids, total_len, device) ContiguousCacheView
+        }
+
+        class CacheView {
+            <<abstract>>
+            +write(layer_id, k, v)
+            +gather(layer_id) Tuple[Tensor, Tensor]
+        }
+
+        class PageCacheView {
             -Storage _storage
             +Tensor _page_table
             +int _total_len
@@ -705,6 +725,14 @@ classDiagram
             +gather(layer_id) Tuple[Tensor, Tensor]
         }
 
+        class ContiguousCacheView {
+            -ContiguousCache _cache
+            +Tensor _batch_indices
+            +int _total_len
+            +write(layer_id, k, v)
+            +gather(layer_id) Tuple[Tensor, Tensor]
+        }
+
         class TaskTable {
             +set(task_id, page_table, cached)
             +get(task_id) List[int]
@@ -714,23 +742,22 @@ classDiagram
             +table_tensor(task_ids, device) Tensor
         }
 
-        class Task {
-            +str task_id
-            +List prompt_ids
-            +Optional[int] max_tokens
-            +float temperature
-            +float top_p
-            +int top_k
-            +TaskStatus status
-            +List output_ids
-            +int input_tokens
-            +int output_tokens
-            +float arrival_time
-            +Optional[float] finish_time
-            +Optional[Callable] stream_callback
-            +int next_pos
-            +is_finished(stop_ids) bool
-        }
+    class Task {
+        +str task_id
+        +List prompt_ids
+        +Optional[int] max_tokens
+        +float temperature
+        +float top_p
+        +int top_k
+        +TaskStatus status
+        +List output_ids
+        +int input_tokens
+        +int output_tokens
+        +float arrival_time
+        +Optional[float] finish_time
+        +int next_pos
+        +is_finished(stop_ids) bool
+    }
 
         class TaskStatus {
             <<enumeration>>
@@ -1035,14 +1062,14 @@ classDiagram
     TrainCallback <|-- GradientCheckpointingCallback
     TrainCallback <|-- CheckpointCallback
     TrainCallback <|-- ProgressBarCallback
-    TrainCallback <|-- MetricLoggerCallback
-    TrainCallback <|-- ValidationCallback
+    TrainCallback <|-- MetricCallback
     BaseDataset <|-- SEQDataset
     BaseDataset <|-- SFTDataset
     BaseDataset <|-- DPODataset
     BaseDataset <|-- GRPODataset
     Store <|-- H5Store
     Store <|-- MmapStore
+    Store <|-- JsonlStore
     BaseSamplingStrategy <|-- TemperatureStrategy
     BaseSamplingStrategy <|-- TopKStrategy
     BaseSamplingStrategy <|-- TopPStrategy
@@ -1075,11 +1102,15 @@ classDiagram
     ResponseBuilder <|-- OpenAIResponseBuilder
     ResponseBuilder <|-- AnthropicResponseBuilder
     BaseMaskBuilder <|-- SectionedMaskBuilder
+    KVCache <|-- PageCache
+    KVCache <|-- ContiguousCache
+    CacheView <|-- PageCacheView
+    CacheView <|-- ContiguousCacheView
 
     %% --- Composition (strong ownership, part destroyed with whole) ---
-    KVCache *-- PagePool
-    KVCache *-- Storage
-    KVCache *-- TaskTable
+    PageCache *-- PagePool
+    PageCache *-- Storage
+    PageCache *-- TaskTable
     InferenceEngine *-- InferenceScheduler
     InferenceScheduler *-- KVCache
     InferenceScheduler *-- Executor
@@ -1107,7 +1138,8 @@ classDiagram
     TrainContext o-- BaseScheduler
     TrainContext o-- Checkpoint
     TrainContext o-- BaseExecutor
-    KvcacheView o-- Storage
+    PageCacheView o-- Storage
+    ContiguousCacheView o-- ContiguousCache
     SamplingPipeline o-- BaseSamplingStrategy
     BaseDataset o-- Store
     Pipeline o-- PipelineConfig
@@ -1129,6 +1161,7 @@ classDiagram
     DecoderBlock ..> FFNFactory : uses
     StoreFactory ..> H5Store : creates
     StoreFactory ..> MmapStore : creates
+    StoreFactory ..> JsonlStore : creates
     ConfigFactory ..> AutoRegressiveLMConfig : creates
     ConfigFactory ..> EncoderConfig : creates
     ExecutorFactory ..> NoneExecutor : creates
@@ -1142,7 +1175,8 @@ classDiagram
     TrainContextBuilder ..> ResumableDistributedSampler : creates
     Checkpoint ..> Checkpoint : serializes
     CheckpointCallback ..> Checkpoint : creates
-    KVCache ..> KvcacheView : binds
+    PageCache ..> PageCacheView : binds
+    ContiguousCache ..> ContiguousCacheView : binds
     InferenceEngine ..> GenerationRequest : uses
     InferenceEngine ..> GenerateResult : creates
     OpenAIResponseBuilder ..> ChatCompletionRequest : receives
@@ -1171,12 +1205,12 @@ classDiagram
 |--------|------------|-------------|
 | **astrai.config** | BaseConfig, BaseModelConfig, AutoRegressiveLMConfig, EncoderConfig, ConfigFactory, TrainConfig, PipelineConfig, InputConfig, ProcessingConfig, OutputConfig | Configuration management (to_dict/from_dict, to_file/from_file) |
 | **astrai.preprocessing** | BaseMaskBuilder, MaskBuilderFactory, SectionedMaskBuilder, Pipeline, filter_by_length, PackingStrategy, PackingStrategyFactory, PositionIdStrategy, PositionIdStrategyFactory, StoreWriter, StoreWriterFactory | Declarative JSON-driven data preprocessing |
-| **astrai.dataset** | BaseDataset–GRPODataset, Store–MmapStore, StoreFactory, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
+| **astrai.dataset** | BaseDataset–GRPODataset, Store–JsonlStore/MmapStore/H5Store, StoreFactory, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
 | **astrai.serialization** | Checkpoint | Model serialization |
 | **astrai.model** | AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, 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)–ValidationCallback, CallbackFactory | Training workflow |
-| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, KVCache–KvcacheView, Allocator–Storage, Task, TaskManager, TaskStatus, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, ChatMessage–MessagesRequest, app | Inference service |
+| **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory | Training workflow |
+| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, KVCache–ContiguousCache/PageCache, CacheView–ContiguousCacheView/PageCacheView, Allocator–Storage, Task, TaskManager, TaskStatus, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, ChatMessage–MessagesRequest, app | Inference service |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel & gradient accumulation |
 | **astrai.factory** | BaseFactory | Component registration |
 | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers |
@@ -1195,7 +1229,7 @@ classDiagram
 | **Context** | `TrainContext` | Unified training state bag |
 | **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
 | **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor`, `FSDPExecutor` | Gradient accumulation & model distribution |
-| **Storage** | `Store`, `H5Store`, `MmapStore` | Format-agnostic data access with multi-segment support |
+| **Storage** | `Store`, `H5Store`, `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 |
 
@@ -1207,10 +1241,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 `KVCache` + `SamplingPipeline`
 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
-7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (H5Store/MmapStore) loads data with explicit `_length` and multi-segment `_data`
+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`
 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-05-30
+> Document Update Time: 2026-07-09

assets/docs/dataflow.md

@@ -48,23 +48,26 @@ The output `meta.json` records the storage format, key names, dtype, total token
 
 `detect_format(load_path)` inspects the path:
 
-- If `load_path` is a file: checks suffix — `.h5`/`.hdf5` → `"h5"`, unknown suffix raises `ValueError`
-- If `load_path` is a directory: recursively globs for `*.h5`/`*.hdf5` files → `"h5"`, or `*.bin` + `**/meta.json` → `"bin"`
+- 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"`
 
 ### 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("h5")    → H5Store
+StoreFactory.create("bin")   → MmapStore
+StoreFactory.create("jsonl") → JsonlStore
 ```
 
-**H5Store**: Reads HDF5 files, supports `share_memory_()` for multi-process DataLoader workers (copies tensors to shared memory).
+**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage.
 
 **MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`.
 
-Both backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based indexing).
+**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.
+
+All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based indexing).
 
 ## Data Keys by Training Type
 
@@ -106,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-06-19
+> Document Update Time: 2026-07-09

assets/docs/inference.md

@@ -23,29 +23,40 @@ RoPE is applied **before** KV cache write, not after — otherwise position enco
 
 ## KVCache System
 
-Seven classes working together:
+Seven classes working together, with two concrete cache implementations:
+
+### ContiguousCache (default)
 
 ```
-KVCache (facade)
-  ├── PagePool         orchestrates page allocation + prefix matching
-  │     ├── Allocator   bitmask-based page allocator + ref-count + LRU eviction (inside PagePool)
-  │     └── PrefixCache hash-based prefix matching (page_hash via polynomial hash) (inside PagePool)
-  ├── TaskTable        maps task_id → page_table + cached token count
-  ├── Storage          k_cache / v_cache tensors (n_layers × n_pages × page_size × n_kv_heads × head_dim)
-  └── KvcacheView      bundles Storage + page_table + total_len for attention layers (returned by bind())
+ContiguousCache (simple contiguous per-slot cache)
+  ├── ContiguousCacheView  bundles k/v tensors + slot indices for attention layers
 ```
 
-`KVCache.bind(page_table, total_len)` returns a `KvcacheView` used by attention layers via `write()` / `gather()`.
+Created by default when no cache is passed to `InferenceScheduler`. Each task occupies a fixed slot of `[max_seq_len, n_kv_heads, head_dim]`. Simple and efficient for small-to-medium batch sizes.
+
+### PageCache (paged with prefix sharing)
+
+```
+PageCache (paged KV cache with prefix sharing, alternative)
+  ├── PagePool               orchestrates page allocation + prefix matching
+  │     ├── Allocator         bitmask-based page allocator + ref-count + LRU
+  │     └── PrefixCache       hash-based prefix matching (page_hash via polynomial hash)
+  ├── TaskTable              maps task_id → page_table + cached token count
+  ├── Storage                k_cache / v_cache tensors (n_layers × n_pages × page_size × n_kv_heads × head_dim)
+  └── PageCacheView          bundles Storage + page_table + total_len for attention layers
+```
+
+`isinstance(cache, KVCache)` checks dispatch to the correct view. Both implement the abstract `KVCache` interface used by `Executor` and `InferenceScheduler`.
 
 ## Continuous Batching
 
 `InferenceScheduler` runs a daemon thread with a 4-phase loop:
 
 ```
-1. Cleanup → Remove finished tasks, free KV pages
-2. Refill  → Pop from waiting_queue, task_alloc pages, activate
+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  → Pick largest same-position group, single-token forward
+4. Decode  → Run single-token forward for each same-position group
 ```
 
 ## Sampling (Strategy Pattern)
@@ -238,4 +249,4 @@ async for token in engine.generate_async("Hello", ...):    # -> AsyncGenerator[s
     print(token)
 ```
 
-> Document Update Time: 2026-06-19
+> Document Update Time: 2026-07-09

assets/docs/params.md

@@ -28,13 +28,17 @@
 | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
 | `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
 
-### Optimizer (AdamW)
+### Optimizer (MuonMix)
+
+Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`fused=True`).
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--adamw_beta1` | AdamW beta1 | 0.9 |
-| `--adamw_beta2` | AdamW beta2 | 0.95 |
-| `--adamw_weight_decay` | AdamW weight decay | 0.01 |
+| `--weight_decay` | Weight decay (applied to Muon matrix params; non-matrix use 0) | 0.1 |
+| `--muon_momentum` | Muon momentum factor | 0.95 |
+| `--muon_nesterov` | Enable Nesterov momentum for Muon | True |
+| `--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
 
@@ -67,7 +71,6 @@
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--log_dir` | Directory for metric logs | checkpoint/logs |
-| `--log_interval` | Number of optimizer steps between metric logs | 1 |
 | `--metrics` | Metrics to log (e.g. --metrics loss lr val_loss) | ["loss", "lr", "grad_norm"] |
 
 ### Gradient Checkpointing
@@ -105,7 +108,7 @@
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--schedule_type` | LR scheduler type (`cosine`, `sgdr`, `wsd`) | cosine |
-| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default) |
+| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default: 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 (required for wsd) |
@@ -127,9 +130,7 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --adamw_beta1=0.9 \
-    --adamw_beta2=0.95 \
-    --adamw_weight_decay=0.01 \
+    --weight_decay=0.1 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \
@@ -200,4 +201,4 @@ See [Preprocessing Guide](preprocessing.md) for config file format and examples.
 
 ---
 
-> Document Update Time: 2026-06-19
+> Document Update Time: 2026-07-09

assets/docs/preprocessing.md

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

assets/docs/training.md

@@ -80,13 +80,13 @@ on_train_end
 | `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
 | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
 | `on_batch_begin` | Every batch | — |
-| `on_optimizer_step` | Every accumulation window | `GradientClippingCallback`, `MetricLoggerCallback`, `ValidationCallback` |
-| `on_batch_end` | Every batch | `CheckpointCallback`, `MetricLoggerCallback`, `ProgressBarCallback` |
-| `on_epoch_end` | End of each epoch | `ProgressBarCallback` |
-| `on_error` | On exception during training | `CheckpointCallback`, `MetricLoggerCallback` |
-| `on_train_end` | Training ends (always via finally) | `CheckpointCallback`, `MetricLoggerCallback`, `GradientCheckpointingCallback` |
+| `on_optimizer_step` | Every accumulation window | `GradientClippingCallback`, `MetricCallback`, `ProgressBarCallback` |
+| `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` |
 
-Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `validation` (periodic validation on val_dataset), `metric_logger` (JSONL, rank-0), `progress_bar` (tqdm), `gradient_clipping`.
+Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm), `gradient_clipping`.
 
 ## Strategies
 
@@ -108,7 +108,7 @@ $$
 L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
 $$
 
-Keys: `input_ids`, `target_ids`, `loss_mask`. Optional: `label_smoothing`.
+Keys: `input_ids`, `target_ids`, `loss_mask`, `position_ids`. Optional: `label_smoothing`.
 
 ### DPO (Direct Preference Optimization)
 
@@ -201,9 +201,7 @@ nohup python scripts/tools/train.py \
     --warmup_ratio=0.05 \
     --max_lr=1e-4 \
     --max_grad_norm=1.0 \
-    --adamw_beta1=0.9 \
-    --adamw_beta2=0.95 \
-    --adamw_weight_decay=0.01 \
+    --weight_decay=0.1 \
     --window_size=2048 \
     --ckpt_interval=10000 \
     --ckpt_dir=./checkpoint \
@@ -214,4 +212,4 @@ nohup python scripts/tools/train.py \
 
 Full parameter reference at [params.md](params.md).
 
-> Document Update Time: 2026-05-30
+> Document Update Time: 2026-07-09

astrai/extension/__init__.py

@@ -0,0 +1,19 @@
+"""CUDA attention kernel wrappers with torch fallback.
+
+Public API:
+    - ``gqa_decode_attn`` — single-query decode attention
+    - ``gqa_prefill_attn`` — multi-query prefill attention
+
+Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.gqa_*``)
+when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``.
+"""
+
+from astrai.extension.loader import KERNEL_NAMES, is_available
+from astrai.extension.ops import gqa_decode_attn, gqa_prefill_attn
+
+__all__ = [
+    "gqa_decode_attn",
+    "gqa_prefill_attn",
+    "is_available",
+    "KERNEL_NAMES",
+]

astrai/extension/loader.py

@@ -0,0 +1,36 @@
+"""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

@@ -0,0 +1,86 @@
+"""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,7 +1,9 @@
 from astrai.preprocessing.builder import (
     BaseMaskBuilder,
     MaskBuilderFactory,
+    MultiOutputMaskBuilder,
     SectionedMaskBuilder,
+    SingleOutputMaskBuilder,
 )
 from astrai.preprocessing.packing import (
     PackingStrategy,
@@ -20,12 +22,14 @@ 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,8 +1,10 @@
 """Mask building for preprocessing pipeline.
 
 :class:`SectionRenderer` converts section specs into token ids and loss
-masks (template / text / value extraction).  :class:`SectionedMaskBuilder`
-orchestrates single-output / multi-output (DPO / GRPO) assembly.
+masks (template / text / value extraction).  :class:`SingleOutputMaskBuilder`
+handles single-output (SFT / pretrain), :class:`MultiOutputMaskBuilder`
+handles multi-output (DPO / GRPO), and :class:`SectionedMaskBuilder`
+orchestrates both modes as a façade.
 """
 
 from abc import ABC, abstractmethod
@@ -212,42 +214,17 @@ class MaskBuilderFactory(BaseFactory["BaseMaskBuilder"]):
     pass
 
 
-@MaskBuilderFactory.register("sectioned")
-class SectionedMaskBuilder(BaseMaskBuilder):
-    """Config-driven builder supporting single and multi-output modes.
+@MaskBuilderFactory.register("single")
+class SingleOutputMaskBuilder(BaseMaskBuilder):
+    """Build a single output sequence with optional loss mask.
 
-    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"``)
+    Expects ``config.input.sections`` (list of section specs).
     """
 
-    def __init__(self):
-        self.renderer = SectionRenderer()
+    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 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
@@ -266,9 +243,22 @@ class SectionedMaskBuilder(BaseMaskBuilder):
             result["loss_mask"] = mask
         return result
 
-    def _build_multi(
-        self, item: dict, sources_spec: dict, config, tokenizer
-    ) -> Optional[dict]:
+
+@MaskBuilderFactory.register("multi")
+class MultiOutputMaskBuilder(BaseMaskBuilder):
+    """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
 
@@ -313,3 +303,22 @@ class SectionedMaskBuilder(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,7 +26,10 @@ def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
     tensor_group: Dict[str, List[Tensor]] = {}
 
     root_path = Path(file_path)
-    h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
+    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:
         with h5py.File(h5_file, "r") as f:

astrai/trainer/train_callback.py

@@ -209,7 +209,6 @@ 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}",
@@ -220,6 +219,7 @@ 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/__init__.py

@@ -0,0 +1,2 @@
+# Source directory for CUDA kernels — build-time only.
+# Compiled .so files live in astrAI/_ext/.

csrc/build.py

@@ -0,0 +1,46 @@
+from pathlib import Path
+
+
+def _arch_flags() -> list[str]:
+    import torch
+
+    if torch.cuda.is_available():
+        cap = torch.cuda.get_device_capability()
+    else:
+        cap = (8, 0)
+    ver = f"{cap[0]}{cap[1]}"
+    flags = [f"-gencode=arch=compute_{ver},code=sm_{ver}"]
+    # tensor-core mma path (mma.sync.m16n8k16.bf16) requires sm_80+; decide the
+    # kernel dispatch at build time via this define rather than at runtime.
+    if cap[0] < 8:
+        flags.append("-DASTRAI_NO_MMA")
+    return flags
+
+
+_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": [*NVCC_FLAGS, *_arch_flags()],
+        "extra_link_args": kwargs.pop("extra_link_args", []),
+        **kwargs,
+    }
+
+
+register("gqa_decode_attn")
+register("gqa_prefill_attn")

csrc/kernels/gqa_common.cuh

@@ -0,0 +1,35 @@
+#pragma once
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cfloat>
+#include <algorithm>
+
+using bf16 = __nv_bfloat16;
+using std::min;
+
+constexpr int DC_CHUNK = 64;
+constexpr int Br = 32, Bc = 64;
+
+__device__ inline float warp_reduce_sum(float val) {
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
+    return val;
+}
+
+struct GQAParams {
+    int batch;
+    int q_head;
+    int kv_head;
+    int q_len;
+    int kv_len;
+    int head_dim;
+    int use_mask;
+    int is_causal;
+    int causal_offset;
+    float scale;
+    const bf16* __restrict__ q;
+    const bf16* __restrict__ k;
+    const bf16* __restrict__ v;
+    const bool* __restrict__ mask;
+    bf16* __restrict__ o;
+};

csrc/kernels/gqa_decode_attn.cu

@@ -0,0 +1,114 @@
+#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, 
+    torch::Tensor v,
+    c10::optional<torch::Tensor> mask,
+    bool is_causal = false, 
+    int64_t causal_offset = 0,
+    c10::optional<double> scale = c10::nullopt
+) {
+    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
+    TORCH_CHECK(q.dtype() == torch::kBFloat16);
+    TORCH_CHECK(k.dtype() == torch::kBFloat16);
+    TORCH_CHECK(v.dtype() == torch::kBFloat16);
+    TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1");
+
+    GQAParams p;
+    p.batch = q.size(0);
+    p.q_head = q.size(1);
+    p.kv_head = k.size(1);
+    p.q_len = 1;
+    p.kv_len = k.size(2);
+    p.head_dim = q.size(3);
+    TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
+    p.use_mask = mask.has_value();
+    p.is_causal = (int)is_causal;
+    p.causal_offset = (int)causal_offset;
+    p.scale = scale.has_value() ? (float)scale.value() : 1.0f / sqrtf((float)p.head_dim);
+    p.q = (const bf16*)q.data_ptr();
+    p.k = (const bf16*)k.data_ptr();
+    p.v = (const bf16*)v.data_ptr();
+    if (p.use_mask) {
+        TORCH_CHECK(mask.value().dtype() == torch::kBool);
+        TORCH_CHECK(mask.value().dim() == 2);
+        TORCH_CHECK(mask.value().size(0) == p.batch);
+        TORCH_CHECK(mask.value().size(1) == p.kv_len);
+        p.mask = mask.value().data_ptr<bool>();
+    } else {
+        p.mask = nullptr;
+    }
+
+    auto O = torch::empty_like(q);
+    p.o = (bf16*)O.data_ptr();
+
+    switch (p.head_dim) {
+        case 32:
+            dispatch_decode<32>(p);
+            break;
+        case 64:
+            dispatch_decode<64>(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;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("gqa_decode_attn", &gqa_decode_attn,
+        py::arg("q"),
+        py::arg("k"),
+        py::arg("v"),
+        py::arg("mask") = py::none(),
+        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)");
+}

csrc/kernels/gqa_decode_attn.cuh

@@ -0,0 +1,59 @@
+#pragma once
+#include "gqa_common.cuh"
+
+__global__ void gqa_decode_attn_kernel(GQAParams p) {
+    int batch = blockIdx.x / p.kv_head;
+    int kv_head = blockIdx.x % p.kv_head;
+    int group_size = blockDim.y;
+    int q_head = kv_head * group_size + threadIdx.y;
+    int lane = threadIdx.x;
+    int hd_per_thread = p.head_dim / 32;
+
+    float q_reg[8];
+    int q_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
+    for (int i = 0; i < hd_per_thread; i++)
+        q_reg[i] = __bfloat162float(p.q[q_off + i]);
+
+    int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * p.head_dim;
+    int mask_base = batch * p.kv_len;
+
+    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
+
+    extern __shared__ __align__(16) bf16 k_smem[];
+
+    for (int chunk_start = 0; chunk_start < p.kv_len; chunk_start += DC_CHUNK) {
+        int this_chunk = min(DC_CHUNK, p.kv_len - chunk_start);
+
+        int total = this_chunk * p.head_dim;
+        for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y)
+            k_smem[i] = p.k[kv_base + chunk_start * p.head_dim + i];
+        __syncthreads();
+
+        for (int s = 0; s < this_chunk; s++) {
+            float partial = 0.0f;
+            for (int i = 0; i < hd_per_thread; i++)
+                partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]);
+            partial = warp_reduce_sum(partial) * p.scale;
+
+            if (p.use_mask && p.mask && !p.mask[mask_base + chunk_start + s])
+                partial = -FLT_MAX;
+            if (p.is_causal && (chunk_start + s) > p.causal_offset)
+                partial = -FLT_MAX;
+
+            float new_m = fmaxf(m, partial);
+            float alpha = expf(m - new_m);
+            float beta  = expf(partial - new_m);
+            d = d * alpha + beta;
+
+            int v_off = kv_base + (chunk_start + s) * p.head_dim + lane * hd_per_thread;
+            for (int i = 0; i < hd_per_thread; i++)
+                acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i]) * beta;
+            m = new_m;
+        }
+        __syncthreads();
+    }
+
+    int out_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread;
+    for (int i = 0; i < hd_per_thread; i++)
+        p.o[out_off + i] = __float2bfloat16(acc_reg[i] / d);
+}

csrc/kernels/gqa_decode_attn_mma.cuh

@@ -0,0 +1,219 @@
+#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

@@ -0,0 +1,93 @@
+#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

@@ -0,0 +1,100 @@
+#include "gqa_prefill_attn.cuh"
+#include <torch/extension.h>
+
+#ifndef ASTRAI_NO_MMA
+#include "gqa_prefill_attn_mma.cuh"
+#endif
+
+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;
+    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 + 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);
+#else
+    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, 1);
+    size_t smem = 2 * P_BC * HEAD_DIM * sizeof(bf16);
+    gqa_prefill_attn_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block, smem>>>(p);
+#endif
+}
+
+torch::Tensor gqa_prefill_attn(
+    torch::Tensor q,
+    torch::Tensor k,
+    torch::Tensor v,
+    c10::optional<torch::Tensor> mask,
+    bool is_causal = false,
+    int64_t causal_offset = 0,
+    c10::optional<double> scale = c10::nullopt
+) {
+    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
+    TORCH_CHECK(q.dtype() == torch::kBFloat16);
+    TORCH_CHECK(k.dtype() == torch::kBFloat16);
+    TORCH_CHECK(v.dtype() == torch::kBFloat16);
+
+    GQAParams p;
+    p.batch = q.size(0);
+    p.q_head = q.size(1);
+    p.kv_head = k.size(1);
+    p.q_len = q.size(2);
+    p.kv_len = k.size(2);
+    p.head_dim = q.size(3);
+    TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
+    p.use_mask = mask.has_value();
+    p.is_causal = (int)is_causal;
+    p.causal_offset = (int)causal_offset;
+    p.scale = scale.has_value() ? (float)scale.value() : 1.0f / sqrtf((float)p.head_dim);
+    p.q = (const bf16*)q.data_ptr();
+    p.k = (const bf16*)k.data_ptr();
+    p.v = (const bf16*)v.data_ptr();
+    if (p.use_mask) {
+        TORCH_CHECK(mask.value().dtype() == torch::kBool);
+        TORCH_CHECK(mask.value().dim() == 2);
+        TORCH_CHECK(mask.value().size(0) == p.batch);
+        TORCH_CHECK(mask.value().size(1) == p.kv_len);
+        p.mask = mask.value().data_ptr<bool>();
+    } else {
+        p.mask = nullptr;
+    }
+
+    auto O = torch::empty_like(q);
+    p.o = (bf16*)O.data_ptr();
+
+    switch (p.head_dim) {
+        case 32:
+            dispatch_prefill<32>(p);
+            break;
+        case 64:
+            dispatch_prefill<64>(p);
+            break;
+        case 128:
+            dispatch_prefill<128>(p);
+            break;
+        case 256:
+            dispatch_prefill<256>(p);
+            break;
+        default:
+            TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim,
+                        " (supported: 32,64,128,256)");
+    }
+    return O;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("gqa_prefill_attn", &gqa_prefill_attn,
+        py::arg("q"),
+        py::arg("k"),
+        py::arg("v"),
+        py::arg("mask") = py::none(),
+        py::arg("is_causal") = false,
+        py::arg("causal_offset") = 0,
+        py::arg("scale") = py::none(),
+        "GQA prefill (tensor-core mma on sm_80+, scalar fallback)");
+}

csrc/kernels/gqa_prefill_attn.cuh

@@ -0,0 +1,137 @@
+#pragma once
+#include "gqa_common.cuh"
+
+// v9: group-split register blocking. G threads cooperate on one query row,
+// each owning HEAD_DIM/G dims of qreg[]/acc[]. Small per-thread footprint keeps
+// occupancy high; the S dot product is reduced across the G-lane group with a
+// short shuffle chain (log2(G) shuffles) instead of a full 32-lane warp reduce.
+// Online (per-kv) softmax — cheap because acc[] is only HEAD_DIM/G long.
+// Templated on <HEAD_DIM, G, ROWS, P_BC>. Block = (G, ROWS). G power-of-two,
+// G*ROWS a multiple of 32 with groups warp-aligned.
+
+template <int G>
+__device__ __forceinline__ float group_reduce_sum(float v, unsigned mask) {
+#pragma unroll
+    for (int o = G / 2; o > 0; o >>= 1)
+        v += __shfl_xor_sync(mask, v, o);
+    return v;
+}
+
+// load 8 contiguous bf16 from (16-byte aligned) smem as one float4, unpack to
+// 8 floats — cuts shared-load instructions 8x vs scalar bf16 loads.
+__device__ __forceinline__ void ld8(const bf16* p, float* o) {
+    float4 raw = *reinterpret_cast<const float4*>(p);
+    const __nv_bfloat162* h = reinterpret_cast<const __nv_bfloat162*>(&raw);
+#pragma unroll
+    for (int j = 0; j < 4; j++) {
+        float2 f = __bfloat1622float2(h[j]);
+        o[2 * j] = f.x;
+        o[2 * j + 1] = f.y;
+    }
+}
+
+template <int HEAD_DIM, int G, int ROWS, int P_BC>
+__global__ void gqa_prefill_attn_kernel_t(GQAParams p) {
+    constexpr int DPT = HEAD_DIM / G;
+
+    int q_tile = blockIdx.x;
+    int q_head = blockIdx.y;
+    int batch  = blockIdx.z;
+    int gpos   = threadIdx.x;  // 0..G-1  (which d-chunk)
+    int row    = threadIdx.y;  // 0..ROWS-1
+    int q_row  = q_tile * ROWS + row;
+
+    int kv_head = q_head / (p.q_head / p.kv_head);
+
+    extern __shared__ __align__(16) bf16 smem[];
+    bf16* sK = smem;
+    bf16* sV = sK + P_BC * HEAD_DIM;
+
+    float qreg[DPT];
+    if (q_row < p.q_len) {
+        int q_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT;
+#pragma unroll
+        for (int i = 0; i < DPT; i++)
+            qreg[i] = __bfloat162float(p.q[q_off + i]) * p.scale;
+    }
+
+    float m = -FLT_MAX, l = 0.0f;
+    float acc[DPT];
+#pragma unroll
+    for (int i = 0; i < DPT; i++)
+        acc[i] = 0.0f;
+
+    int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * HEAD_DIM;
+    int tiles   = (p.kv_len + P_BC - 1) / P_BC;
+    int tt      = G * ROWS;
+    int lid     = row * G + gpos;
+
+    // per-group shuffle mask: only the G lanes of this row's group participate,
+    // so causal masking (differing loop bounds across rows in a warp) is safe.
+    int lane_in_warp = lid & 31;
+    unsigned gmask = (G == 32) ? 0xFFFFFFFFu
+                               : (((1u << G) - 1u) << (lane_in_warp & ~(G - 1)));
+
+    for (int ti = 0; ti < tiles; ti++) {
+        int kv0  = ti * P_BC;
+        int tlen = min(P_BC, p.kv_len - kv0);
+
+        for (int i = lid; i < tlen * HEAD_DIM; i += tt) {
+            int gidx = kv_base + (kv0 + i / HEAD_DIM) * HEAD_DIM + (i % HEAD_DIM);
+            sK[i] = p.k[gidx];
+            sV[i] = p.v[gidx];
+        }
+        __syncthreads();
+
+        int lim = tlen;
+        if (p.is_causal && q_row < p.q_len) {
+            int ep = q_row + p.causal_offset + 1;
+            if (kv0 >= ep)
+                lim = 0;
+            else if (kv0 + tlen > ep)
+                lim = ep - kv0;
+        }
+
+        for (int s = 0; s < lim; s++) {
+            const bf16* kr = sK + s * HEAD_DIM + gpos * DPT;
+            float part = 0.0f;
+#pragma unroll
+            for (int i = 0; i < DPT; i += 8) {
+                float k8[8];
+                ld8(kr + i, k8);
+#pragma unroll
+                for (int j = 0; j < 8; j++)
+                    part = fmaf(qreg[i + j], k8[j], part);
+            }
+            float dot = group_reduce_sum<G>(part, gmask);
+
+            if (p.use_mask && p.mask && !p.mask[batch * p.kv_len + kv0 + s])
+                dot = -FLT_MAX;
+
+            float nm = fmaxf(m, dot);
+            float al = __expf(m - nm);
+            float be = __expf(dot - nm);
+            l = l * al + be;
+
+            const bf16* vr = sV + s * HEAD_DIM + gpos * DPT;
+#pragma unroll
+            for (int i = 0; i < DPT; i += 8) {
+                float v8[8];
+                ld8(vr + i, v8);
+#pragma unroll
+                for (int j = 0; j < 8; j++)
+                    acc[i + j] = fmaf(v8[j], be, acc[i + j] * al);
+            }
+            m = nm;
+        }
+        __syncthreads();
+    }
+
+    if (q_row < p.q_len) {
+        int o_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT;
+        float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f;
+#pragma unroll
+        for (int i = 0; i < DPT; i++)
+            p.o[o_off + i] = __float2bfloat16(acc[i] * rl);
+    }
+}

csrc/kernels/gqa_prefill_attn_mma.cuh

@@ -0,0 +1,246 @@
+#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
+// cores via mma.sync.m16n8k16 (f32 accumulate). Q stays resident in registers;
+// S, O, and the online-softmax stats (m, l) live in registers too — nothing is
+// staged through shared memory except the cooperatively-loaded K/V tiles. The
+// mma fragment layout is used directly: the S accumulator (f32) maps element-
+// 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)
+// → cuts smem; pre-scale Q by attention scale during Q load; cp.async global→
+// shared for K/V; scalar fallback only for the last partial tile; causal tile
+// skipping (block-level early break + warp-level skip); XOR swizzle (swiz_col)
+// → eliminates ldmatrix bank conflicts without LD padding (LD=HEAD_DIM).
+
+template <int HEAD_DIM, int WARPS, int BC>
+__global__ void gqa_prefill_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 (swiz_col) handles bank conflicts
+    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);  // chunk bits, stay within LD
+
+    const int warp = threadIdx.x / 32;
+    const int lane = threadIdx.x % 32;
+    const int gid = lane >> 2;   // 0..7  → rows gid, gid+8
+    const int tid4 = lane & 3;   // 0..3
+    const int nthreads = WARPS * 32;
+
+    const int q_head = blockIdx.y;
+    const int batch = blockIdx.z;
+    const int kv_head = q_head / (p.q_head / p.kv_head);
+    const int qrow0 = (blockIdx.x * WARPS + warp) * BR;
+
+    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]
+
+    // Q resident A-fragments (loaded once per warp via shared staging).
+    // 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]
+                                         : __float2bfloat16(0.0f);
+                sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(qv, scale_bf16);
+            }
+            __syncwarp();
+        #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)]);
+        }
+        __syncthreads();  // prevent next warp from overwriting sQ prematurely
+    }
+
+    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 tiles = (p.kv_len + BC - 1) / BC;
+    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)
+                                     ? 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;
+            }
+        }
+        __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 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 (in registers) ----
+        // Q is pre-scaled, so Sacc already includes the attention scale.
+        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 c1 = cc + 1;
+            bool b0 = (cc >= maxc0) || (has_mask && !p.mask[mb + cc]);
+            bool b1 = (c1 >= maxc0) || (has_mask && !p.mask[mb + c1]);
+            bool b2 = (cc >= maxc1) || (has_mask && !p.mask[mb + cc]);
+            bool b3 = (c1 >= maxc1) || (has_mask && !p.mask[mb + c1]);
+            float s0 = b0 ? -FLT_MAX : Sacc[n8][0];
+            float s1 = b1 ? -FLT_MAX : Sacc[n8][1];
+            float s2 = b2 ? -FLT_MAX : Sacc[n8][2];
+            float s3 = b3 ? -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;
+
+        // rescale O accumulator by per-row correction
+#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]);
+            }
+        }
+        }  // if active (warp-level causal skip)
+        __syncthreads();
+    }
+
+    // ---- write output ----
+    float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
+    float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
+    const int o_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
+#pragma unroll
+    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] =
+                __float2bfloat16(Oacc[dn8][0] * 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] =
+                __float2bfloat16(Oacc[dn8][2] * rl1);
+            p.o[o_base + qr1 * HEAD_DIM + d + 1] =
+                __float2bfloat16(Oacc[dn8][3] * rl1);
+        }
+    }
+}

csrc/tests/gqa_decode_test.cu

@@ -0,0 +1,130 @@
+/*
+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>
+#include <sys/time.h>
+#include "../kernels/gqa_decode_attn.cuh"
+
+static double now_ms() {
+    struct timeval tv;
+    gettimeofday(&tv, NULL);
+    return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
+}
+
+static void cpu_decode(const float* Q, const float* K, const float* V,
+                        const bool* mask, float* O,
+                        int B, int Hq, int Hk, int seq_len, int D) {
+    float scale = 1.0f / sqrtf((float)D);
+    int n_rep = Hq / Hk;
+    for (int b = 0; b < B; b++) {
+        for (int h = 0; h < Hq; h++) {
+            int kv_h = h / n_rep;
+            float mv = -INFINITY, sv = 0.0f;
+            float accum[256] = {0};
+            for (int s = 0; s < seq_len; s++) {
+                if (!mask[b * seq_len + s]) continue;
+                float dot = 0.0f;
+                for (int d = 0; d < D; d++)
+                    dot += Q[((b * Hq + h) * 1 + 0) * D + d]
+                         * K[((b * Hk + kv_h) * seq_len + s) * D + d];
+                dot *= scale;
+                float nm = fmaxf(mv, dot);
+                float al = expf(mv - nm);
+                float be = expf(dot - nm);
+                sv = sv * al + be;
+                for (int d = 0; d < D; d++)
+                    accum[d] = accum[d] * al
+                             + V[((b * Hk + kv_h) * seq_len + s) * D + d] * be;
+                mv = nm;
+            }
+            float inv = 1.0f / sv;
+            for (int d = 0; d < D; d++)
+                O[((b * Hq + h) * 1 + 0) * D + d] = accum[d] * inv;
+        }
+    }
+}
+
+static bf16 f2bf(float x) { return __float2bfloat16(x); }
+static float bf2f(bf16 x) { return __bfloat162float(x); }
+static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
+
+int main() {
+    const int configs[][5] = {
+        {1, 2, 1, 64, 32},    // B,Hq,Hk,seq_len,D
+        {1, 32, 4, 512, 128},
+        {1, 32, 4, 1024, 128},
+    };
+    int n_cfgs = sizeof(configs) / sizeof(configs[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], gs = Hq / Hk;
+        printf("=== B=%d Hq=%d Hk=%d seq=%d D=%d gs=%d ===\n", B,Hq,Hk,sl,D,gs);
+
+        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);
+
+        GQAParams 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=1; p.is_causal=0; p.causal_offset=0;
+        p.scale=1.0f/sqrtf((float)D);
+        p.q=dQ; p.k=dK; p.v=dV; p.mask=dMask; p.o=dO;
+
+        size_t smem=DC_CHUNK*D*sizeof(bf16);
+        dim3 block(32, gs);
+        dim3 grid(B*Hk);
+        printf("grid=(%d,1,1) block=(%d,%d,1) smem=%zu\n",
+               grid.x, block.x, block.y, smem);
+
+        double t0=now_ms();
+        gqa_decode_attn_kernel<<<grid,block,smem>>>(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_decode(hQ,hK,hV,hMask,ref,B,Hq,Hk,sl,D);
+
+        float max_err=0;
+        for (size_t i=0;i<nQ;i++){
+            float d=fabsf(bf2f(hOut[i])-ref[i]);
+            if(d>max_err) max_err=d;
+        }
+        printf("kernel: %.3f ms  max_err: %.6e\n\n",kms,max_err);
+
+        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
+        delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
+    }
+    printf("All tests passed!\n");
+    return 0;
+}

csrc/tests/gqa_prefill_test.cu

@@ -0,0 +1,133 @@
+/*
+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>
+#include <sys/time.h>
+#include "../kernels/gqa_prefill_attn.cuh"
+
+static double now_ms() {
+    struct timeval tv;
+    gettimeofday(&tv, NULL);
+    return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
+}
+
+static void cpu_attention(const float* Q, const float* K, const float* V, float* O,
+                          int B, int Hq, int Hk, int q_len, int kv_len, int D,
+                          int is_causal, int causal_off) {
+    float scale = 1.0f / sqrtf((float)D);
+    int n_rep = Hq / Hk;
+    for (int b = 0; b < B; b++) {
+        for (int h = 0; h < Hq; h++) {
+            for (int qi = 0; qi < q_len; qi++) {
+                int kv_h = h / n_rep;
+                float mv = -INFINITY, sv = 0.0f;
+                float accum[256] = {0};
+                int lim = is_causal ? min(kv_len, qi + causal_off + 1) : kv_len;
+                for (int kj = 0; kj < lim; kj++) {
+                    float dot = 0.0f;
+                    for (int d = 0; d < D; d++)
+                        dot += Q[((b*Hq + h)*q_len + qi)*D + d]
+                             * K[((b*Hk + kv_h)*kv_len + kj)*D + d];
+                    dot *= scale;
+                    float nm = fmaxf(mv, dot);
+                    float al = expf(mv - nm);
+                    float be = expf(dot - nm);
+                    sv = sv * al + be;
+                    for (int d = 0; d < D; d++)
+                        accum[d] = accum[d] * al
+                                 + V[((b*Hk + kv_h)*kv_len + kj)*D + d] * be;
+                    mv = nm;
+                }
+                float inv = 1.0f / sv;
+                for (int d = 0; d < D; d++)
+                    O[((b*Hq + h)*q_len + qi)*D + d] = accum[d] * inv;
+            }
+        }
+    }
+}
+
+static __nv_bfloat16 f2bf(float x) { return __float2bfloat16(x); }
+static float bf2f(__nv_bfloat16 x) { return __bfloat162float(x); }
+static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
+
+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]);
+
+    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];
+        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);
+
+        GQAParams 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.is_causal=causal; p.causal_offset=0;
+        p.scale=1.0f/sqrtf((float)D);
+        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
+
+        constexpr int G=8, ROWS=32, P_BC=32;
+        dim3 grid((ql+ROWS-1)/ROWS, Hq, B);
+        dim3 block(G, ROWS, 1);
+        size_t smem=2*P_BC*D*sizeof(bf16);
+        printf("grid=(%d,%d,%d) block=(%d,%d,%d) smem=%zu\n",
+               grid.x,grid.y,grid.z, block.x,block.y,block.z, smem);
+
+        double t0=now_ms();
+        switch (D) {
+            case 64:  gqa_prefill_attn_kernel_t<64, G,ROWS,P_BC><<<grid,block,smem>>>(p); break;
+            case 128: gqa_prefill_attn_kernel_t<128,G,ROWS,P_BC><<<grid,block,smem>>>(p); break;
+            default: printf("unsupported D=%d\n",D); return 1;
+        }
+        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(hQ,hK,hV,ref,B,Hq,Hk,ql,kl,D,causal,0);
+
+        float max_err=0;
+        for (size_t i=0;i<nQ;i++) {
+            float d=fabsf(bf2f(hOut[i])-ref[i]);
+            if(d>max_err) max_err=d;
+        }
+        printf("kernel: %.3f ms  max_err: %.6e\n\n",kms,max_err);
+
+        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
+        delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
+    }
+    printf("All tests passed!\n");
+    return 0;
+}

scripts/tools/benchmark.py

@@ -1,12 +1,13 @@
 """Benchmark AutoRegressiveLM with KVCache"""
 
+import argparse
 from dataclasses import dataclass
 from typing import Any, Dict
 
 import torch
 
 from astrai.config import AutoRegressiveLMConfig
-from astrai.inference import KVCache
+from astrai.inference import ContiguousCache, PageCache
 from astrai.model.transformer import AutoRegressiveLM
 
 
@@ -24,41 +25,14 @@ class GenerationBenchmark:
         config: AutoRegressiveLMConfig,
         device: str = "cuda",
         dtype: torch.dtype = torch.bfloat16,
-        page_size: int = 128,
+        cache_type: str = "contiguous",
     ):
         self.config = config
         self.device = device
         self.dtype = dtype
+        self.cache_type = cache_type
         self.model = AutoRegressiveLM(config).to(device=device, dtype=dtype)
         self.model.eval()
-        head_dim = config.dim // config.n_heads
-        n_pages = (config.max_len * 4 + page_size - 1) // page_size
-        self._page_cache = KVCache(
-            config.n_layers,
-            n_pages,
-            page_size,
-            config.n_kv_heads,
-            head_dim,
-            device,
-            dtype,
-        )
-
-    def _prepare_inputs(self, batch_size: int, prompt_length: int, total_length: int):
-        prompt_ids = torch.randint(
-            low=0,
-            high=self.config.vocab_size,
-            size=(batch_size, prompt_length),
-            device=self.device,
-            dtype=torch.long,
-        )
-        gen_ids = torch.randint(
-            low=0,
-            high=self.config.vocab_size,
-            size=(batch_size, total_length - prompt_length),
-            device=self.device,
-            dtype=torch.long,
-        )
-        return prompt_ids, gen_ids
 
     @torch.inference_mode()
     def run_prefill_benchmark(
@@ -68,8 +42,12 @@ class GenerationBenchmark:
         num_trials: int = 10,
     ) -> BenchmarkResult:
         for _ in range(3):
-            prompt_ids, _ = self._prepare_inputs(
-                batch_size, prompt_length, prompt_length
+            prompt_ids = torch.randint(
+                0,
+                self.config.vocab_size,
+                (batch_size, prompt_length),
+                device=self.device,
+                dtype=torch.long,
             )
             _ = self.model(prompt_ids)
         torch.cuda.synchronize()
@@ -78,12 +56,15 @@ class GenerationBenchmark:
         total_tokens = batch_size * prompt_length * num_trials
 
         for trial in range(num_trials):
-            prompt_ids, _ = self._prepare_inputs(
-                batch_size, prompt_length, prompt_length
+            prompt_ids = torch.randint(
+                0,
+                self.config.vocab_size,
+                (batch_size, prompt_length),
+                device=self.device,
+                dtype=torch.long,
             )
             start = torch.cuda.Event(enable_timing=True)
             end = torch.cuda.Event(enable_timing=True)
-
             start.record()
             _ = self.model(prompt_ids)
             end.record()
@@ -107,6 +88,7 @@ class GenerationBenchmark:
                 "prompt_length": prompt_length,
                 "dtype": str(self.dtype),
                 "device": self.device,
+                "cache": "none",
             },
         )
 
@@ -120,29 +102,56 @@ class GenerationBenchmark:
     ) -> BenchmarkResult:
         total_time = 0.0
         total_tokens = batch_size * gen_length * num_trials
-        page_size = self._page_cache.page_size
 
         for trial in range(num_trials):
-            prompt_ids, gen_ids = self._prepare_inputs(
-                batch_size,
-                prompt_length,
-                prompt_length + gen_length,
-            )
-
-            n_pages = (prompt_length + gen_length + page_size - 1) // page_size
-            total = n_pages * batch_size
-            pages = []
-            for _ in range(total):
-                p = self._page_cache._pool.alloc()
-                assert p >= 0, "OOM"
-                pages.append(p)
-            page_table = torch.tensor(
-                [pages[i * n_pages : (i + 1) * n_pages] for i in range(batch_size)],
-                dtype=torch.long,
+            prompt_ids = torch.randint(
+                0,
+                self.config.vocab_size,
+                (batch_size, prompt_length),
                 device=self.device,
+                dtype=torch.long,
+            )
+            gen_ids = torch.randint(
+                0,
+                self.config.vocab_size,
+                (batch_size, gen_length),
+                device=self.device,
+                dtype=torch.long,
             )
 
-            cv = self._page_cache.bind(page_table, total_len=prompt_length)
+            head_dim = self.config.dim // self.config.n_heads
+            max_seq = prompt_length + gen_length
+
+            if self.cache_type == "contiguous":
+                cache = ContiguousCache(
+                    self.config.n_layers,
+                    batch_size,
+                    max_seq,
+                    self.config.n_kv_heads,
+                    head_dim,
+                    self.device,
+                    self.dtype,
+                )
+            else:
+                page_size = 128
+                n_pages = (max_seq + page_size - 1) // page_size * batch_size
+                cache = PageCache(
+                    self.config.n_layers,
+                    n_pages,
+                    page_size,
+                    self.config.n_kv_heads,
+                    head_dim,
+                    self.device,
+                    self.dtype,
+                )
+
+            task_ids = [f"b{i}" for i in range(batch_size)]
+            for tid in task_ids:
+                cache.task_alloc(tid, [0] * max_seq)
+                for p in range(max_seq):
+                    cache.task_extend(tid, p)
+
+            cv = cache.bind_tasks(task_ids, prompt_length, self.device)
             _ = self.model(
                 prompt_ids,
                 paged_cache=cv,
@@ -152,37 +161,35 @@ class GenerationBenchmark:
                 .unsqueeze(0)
                 .expand(batch_size, -1),
             )
-
             torch.cuda.synchronize()
 
             start = torch.cuda.Event(enable_timing=True)
             end = torch.cuda.Event(enable_timing=True)
-
             start.record()
-            current_pos = prompt_length
+
             for i in range(gen_length):
-                input_token = gen_ids[:, i : i + 1]
-                cv = self._page_cache.bind(page_table, total_len=current_pos + 1)
+                pos = prompt_length + i
+                cv = cache.bind_tasks(task_ids, pos + 1, self.device)
                 _ = self.model(
-                    input_token,
+                    gen_ids[:, i : i + 1],
                     paged_cache=cv,
                     position_ids=torch.full(
                         (batch_size, 1),
-                        current_pos,
+                        pos,
                         dtype=torch.long,
                         device=self.device,
                     ),
                 )
-                current_pos += 1
+
             end.record()
             torch.cuda.synchronize()
 
+            for tid in task_ids:
+                cache.task_free(tid)
+
             trial_time = start.elapsed_time(end) / 1000
             total_time += trial_time
 
-            for idx in pages:
-                self._page_cache._pool.free(idx)
-
             print(
                 f"  Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s "
                 f"({gen_length / trial_time:.1f} tok/s)"
@@ -199,6 +206,7 @@ class GenerationBenchmark:
                 "gen_length": gen_length,
                 "dtype": str(self.dtype),
                 "device": self.device,
+                "cache": self.cache_type,
             },
         )
 
@@ -216,6 +224,42 @@ def print_benchmark_result(result: BenchmarkResult):
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="AutoRegressiveLM benchmark")
+    parser.add_argument(
+        "--device", type=str, default="cuda", help="Device (default: cuda)"
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bfloat16",
+        choices=["bfloat16", "float16", "float32"],
+        help="Dtype",
+    )
+    parser.add_argument(
+        "--cache",
+        type=str,
+        default="contiguous",
+        choices=["contiguous", "paged"],
+        help="KV cache type",
+    )
+    parser.add_argument("--batch_size", type=int, default=4, help="Batch size")
+    parser.add_argument("--prompt_length", type=int, default=512, help="Prompt length")
+    parser.add_argument("--gen_length", type=int, default=128, help="Generation length")
+    parser.add_argument("--num_trials", type=int, default=5, help="Number of trials")
+    parser.add_argument(
+        "--prefill_only", action="store_true", help="Run prefill benchmark only"
+    )
+    parser.add_argument(
+        "--decode_only", action="store_true", help="Run decoding benchmark only"
+    )
+    args = parser.parse_args()
+
+    dtype_map = {
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+        "float32": torch.float32,
+    }
+
     config = AutoRegressiveLMConfig(
         vocab_size=10000,
         dim=1536,
@@ -227,23 +271,29 @@ if __name__ == "__main__":
         norm_eps=1e-5,
     )
 
-    benchmark = GenerationBenchmark(config)
+    benchmark = GenerationBenchmark(
+        config, device=args.device, dtype=dtype_map[args.dtype], cache_type=args.cache
+    )
 
     print("=" * 80)
-    print("Running AutoRegressiveLM Generation Benchmark (KVCache)")
+    print(
+        f"Running AutoRegressiveLM Benchmark (device={args.device}, dtype={args.dtype})"
+    )
     print("=" * 80)
 
-    prefill_result = benchmark.run_prefill_benchmark(
-        batch_size=4,
-        prompt_length=512,
-        num_trials=5,
-    )
-    print_benchmark_result(prefill_result)
+    if not args.decode_only:
+        prefill_result = benchmark.run_prefill_benchmark(
+            batch_size=args.batch_size,
+            prompt_length=args.prompt_length,
+            num_trials=args.num_trials,
+        )
+        print_benchmark_result(prefill_result)
 
-    gen_result = benchmark.run_decoding_benchmark(
-        batch_size=4,
-        prompt_length=512,
-        gen_length=128,
-        num_trials=5,
-    )
-    print_benchmark_result(gen_result)
+    if not args.prefill_only:
+        gen_result = benchmark.run_decoding_benchmark(
+            batch_size=args.batch_size,
+            prompt_length=args.prompt_length,
+            gen_length=args.gen_length,
+            num_trials=args.num_trials,
+        )
+        print_benchmark_result(gen_result)

scripts/tools/train.py

@@ -1,9 +1,11 @@
 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
@@ -12,6 +14,84 @@ 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.")
@@ -64,22 +144,35 @@ def parse_args() -> argparse.Namespace:
         help="Max gradient norm for clipping.",
     )
     parser.add_argument(
-        "--adamw_beta1",
+        "--weight_decay",
         type=float,
-        default=0.9,
-        help="Beta1 for AdamW optimizer.",
+        default=0.1,
+        help="Weight decay (applied to Muon matrix params; non-matrix use 0).",
     )
     parser.add_argument(
-        "--adamw_beta2",
+        "--muon_momentum",
         type=float,
         default=0.95,
-        help="Beta2 for AdamW optimizer.",
+        help="Momentum factor for Muon optimizer.",
     )
     parser.add_argument(
-        "--adamw_weight_decay",
-        type=float,
-        default=0.01,
-        help="Weight decay for AdamW optimizer.",
+        "--muon_nesterov",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Enable Nesterov momentum for Muon.",
+    )
+    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."
@@ -265,21 +358,8 @@ def create_model(config):
     return AutoRegressiveLM(config).to(dtype=torch.bfloat16)
 
 
-def create_optimizer(model, **kwargs) -> optim.Optimizer:
-    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_optimizer(model, **kwargs) -> MuonMix:
+    return MuonMix(model, **kwargs)
 
 
 def create_scheduler(
@@ -310,7 +390,6 @@ 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,
@@ -323,16 +402,7 @@ 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,
@@ -353,6 +423,7 @@ 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)
@@ -368,12 +439,12 @@ def train(
         window_size = config.max_len
 
     strategy_kwargs = {
-        "beta": dpo_beta,
-        "label_smoothing": label_smoothing,
-        "clip_eps": grpo_clip_eps,
-        "kl_coef": grpo_kl_coef,
-        "group_size": group_size,
-        "sync_interval": grpo_sync_interval,
+        "beta": kwargs.pop("dpo_beta"),
+        "label_smoothing": kwargs.pop("label_smoothing"),
+        "clip_eps": kwargs.pop("grpo_clip_eps"),
+        "kl_coef": kwargs.pop("grpo_kl_coef"),
+        "group_size": kwargs.pop("group_size"),
+        "sync_interval": kwargs.pop("grpo_sync_interval"),
     }
 
     executor_kwargs = {
@@ -391,11 +462,12 @@ def train(
 
     optimizer_fn = partial(
         create_optimizer,
-        **{
-            "lr": max_lr,
-            "betas": (adamw_beta1, adamw_beta2),
-            "weight_decay": adamw_weight_decay,
-        },
+        lr=kwargs.pop("max_lr"),
+        weight_decay=kwargs.pop("weight_decay"),
+        momentum=kwargs.pop("muon_momentum"),
+        nesterov=kwargs.pop("muon_nesterov"),
+        ns_steps=kwargs.pop("muon_ns_steps"),
+        adjust_lr_fn=kwargs.pop("muon_adjust_lr"),
     )
 
     total_steps = compute_total_steps(

setup.py

@@ -0,0 +1,61 @@
+import os
+import sys
+from pathlib import Path
+
+from setuptools import setup
+from setuptools.command.build_ext import build_ext as _build_ext
+
+sys.path.insert(0, str(Path(__file__).parent))
+os.makedirs("astrai/extension", exist_ok=True)
+
+
+def _should_build():
+    force = os.environ.get("CSRC_KERNELS", "").strip().lower()
+    if force == "true":
+        return True
+    if force == "false":
+        return False
+    try:
+        import shutil
+
+        import torch
+
+        return shutil.which("nvcc") is not None and torch.cuda.is_available()
+    except Exception:
+        return False
+
+
+ext_modules = []
+cmdclass = {}
+
+if _should_build():
+    import torch
+    from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+    from csrc.build import REGISTRY
+
+    _torch_lib = torch.utils.cpp_extension.library_paths()[0]
+
+    for name, info in REGISTRY.items():
+        ext_modules.append(
+            CUDAExtension(
+                f"astrai.extension.{name}",
+                info["sources"],
+                extra_compile_args={
+                    "cxx": info["cxx_flags"],
+                    "nvcc": info["nvcc_flags"],
+                },
+                extra_link_args=[f"-Wl,-rpath,{_torch_lib}"],
+            )
+        )
+    cmdclass["build_ext"] = BuildExtension
+
+if not cmdclass:
+
+    class _NullBuildExt(_build_ext):
+        def build_extensions(self):
+            pass
+
+    cmdclass["build_ext"] = _NullBuildExt
+
+setup(ext_modules=ext_modules, cmdclass=cmdclass)

tests/data/conftest.py

@@ -10,7 +10,11 @@ from astrai.config.preprocess_config import (
     PipelineConfig,
     ProcessingConfig,
 )
-from astrai.preprocessing.builder import SectionedMaskBuilder
+from astrai.preprocessing.builder import (
+    MultiOutputMaskBuilder,
+    SectionedMaskBuilder,
+    SingleOutputMaskBuilder,
+)
 from astrai.tokenize import AutoTokenizer
 
 _SPECIAL_TOKENS_CONFIG = {
@@ -210,6 +214,16 @@ 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,7 +8,9 @@ from astrai.config.preprocess_config import (
 )
 from astrai.preprocessing.builder import (
     MaskBuilderFactory,
+    MultiOutputMaskBuilder,
     SectionedMaskBuilder,
+    SingleOutputMaskBuilder,
 )
 from tests.data.conftest import (
     _CHAT_SECTIONS,
@@ -272,12 +274,18 @@ 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():
-    builder_obj = MaskBuilderFactory.create("sectioned")
-    assert isinstance(builder_obj, SectionedMaskBuilder)
+    single = MaskBuilderFactory.create("single")
+    assert isinstance(single, SingleOutputMaskBuilder)
+    multi = MaskBuilderFactory.create("multi")
+    assert isinstance(multi, MultiOutputMaskBuilder)
+    sectioned = MaskBuilderFactory.create("sectioned")
+    assert isinstance(sectioned, SectionedMaskBuilder)
 
 
 def test_dpo_chat_basic(chat_tokenizer, builder):
@@ -367,3 +375,59 @@ 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