feat: wire up paged decode CUDA kernel to Python extension

- Add attn_paged_decode wrapper in ops.py with gather fallback - Register kernel in loader.py and export from __init__.py - Extract test_utils.cuh shared by all attention unit tests - Rename attn_paged_vs_contiguous.cu to attn_paged_decode_test.cu - Refactor decode/prefill tests to use common bf16 helpers and cpu ref - Fix k_cache dim check in attn_paged_decode.cu
feat: paged decode attention with split-KV (scalar + MMA)
2026-07-11 18:40:49 +08:00 · 2026-07-11 18:14:51 +08:00 · 2026-07-11 11:45:13 +08:00 · 2026-07-11 11:26:42 +08:00 · 2026-07-11 11:03:14 +08:00 · 2026-07-11 00:24:09 +08:00
39 changed files with 3327 additions and 270 deletions
--- a/.gitignore
+++ b/.gitignore
@ -7,8 +7,14 @@
 # Allow specific file types and root files
 !astrai/**/*.py
 !scripts/**/*.py
-!scripts/**/*.sh
 !tests/**/*.py
+!csrc/**/*.py
+
+!csrc/**/*.cu
+!csrc/**/*.h
+!csrc/**/*.cuh
+
+!scripts/**/*.sh

 # Allow GitHub files
 !/.github/**
@ -22,4 +28,9 @@
 !/CONTRIBUTING.md
 !/LICENSE
 !/pyproject.toml
-!/README.md
+!/README.md
+# Allow extension modules (only source .py)
+!/astrai/extension/**/*.py
+
+# Allow build files
+!/setup.py
--- a/README.md
+++ b/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 \
--- a/assets/docs/README-zh-CN.md
+++ b/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 \
--- a/assets/docs/architecture.md
+++ b/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
--- a/assets/docs/dataflow.md
+++ b/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
--- a/assets/docs/inference.md
+++ b/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
--- a/assets/docs/params.md
+++ b/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
--- a/assets/docs/preprocessing.md
+++ b/assets/docs/preprocessing.md
@ -361,4 +361,4 @@ Pipeline(
 ).run()
 ```

-> Document Update Time: 2026-06-03
+> Document Update Time: 2026-07-09
--- a/assets/docs/training.md
+++ b/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
--- a/astrai/extension/init.py
+++ b/astrai/extension/init.py
@ -0,0 +1,21 @@
+"""CUDA attention kernel wrappers with torch fallback.
+
+Public API:
+    - ``attn_decode`` — single-query decode attention
+    - ``attn_prefill`` — multi-query prefill attention
+    - ``attn_paged_decode`` — paged decode attention (direct page-table access)
+
+Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.attn_*``)
+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 attn_decode, attn_paged_decode, attn_prefill
+
+__all__ = [
+    "attn_decode",
+    "attn_paged_decode",
+    "attn_prefill",
+    "is_available",
+    "KERNEL_NAMES",
+]
--- a/astrai/extension/loader.py
+++ b/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 = ["attn_decode", "attn_prefill", "attn_paged_decode"]
+
+_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)
--- a/astrai/extension/ops.py
+++ b/astrai/extension/ops.py
@ -0,0 +1,149 @@
+"""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 _gather_kv_from_pages(
+    page_table: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_size: int,
+    kv_len: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Gather contiguous K/V from paged cache for torch SDPA fallback.
+
+    Shapes:
+        page_table : [batch, max_pages] (int64)
+        k_cache    : [n_pages, page_size, n_kv_heads, head_dim]
+        v_cache    : same as k_cache
+    Returns:
+        k, v : [batch, n_kv_heads, kv_len, head_dim]
+    """
+    batch, max_pages = page_table.shape
+    n_pages, ps, n_kv_heads, head_dim = k_cache.shape
+    if ps != page_size:
+        raise ValueError(f"k_cache page_size mismatch: {ps} vs {page_size}")
+
+    k = k_cache.new_empty(batch, n_kv_heads, kv_len, head_dim)
+    v = v_cache.new_empty(batch, n_kv_heads, kv_len, head_dim)
+
+    for b in range(batch):
+        for pos in range(kv_len):
+            log_pg = pos // page_size
+            pg_off = pos % page_size
+            phys = int(page_table[b, log_pg].item())
+            k[b, :, pos, :] = k_cache[phys, pg_off, :, :]
+            v[b, :, pos, :] = v_cache[phys, pg_off, :, :]
+    return k, v
+
+
+def attn_decode(
+    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["attn_decode"]:
+        return _modules["attn_decode"].attn_decode(
+            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 attn_prefill(
+    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["attn_prefill"]:
+        return _modules["attn_prefill"].attn_prefill(
+            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 attn_paged_decode(
+    q: torch.Tensor,
+    page_table: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_size: int,
+    kv_len: int,
+    mask: torch.Tensor | None = None,
+    is_causal: bool = False,
+    causal_offset: int = 0,
+    scale: float | None = None,
+) -> torch.Tensor:
+    if _available["attn_paged_decode"]:
+        return _modules["attn_paged_decode"].attn_paged_decode(
+            q,
+            page_table,
+            k_cache,
+            v_cache,
+            page_size,
+            kv_len,
+            mask=mask,
+            is_causal=is_causal,
+            causal_offset=causal_offset,
+            scale=scale,
+        )
+    k, v = _gather_kv_from_pages(page_table, k_cache, v_cache, page_size, kv_len)
+    return _torch_fallback(q, k, v, mask, is_causal, scale)
--- a/astrai/preprocessing/init.py
+++ b/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",
--- a/astrai/preprocessing/builder.py
+++ b/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)
--- a/astrai/serialization/dataset.py
+++ b/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:
--- a/astrai/trainer/train_callback.py
+++ b/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):
--- a/csrc/init.py
+++ b/csrc/init.py
@ -0,0 +1,2 @@
+# Source directory for CUDA kernels — build-time only.
+# Compiled .so files live in astrAI/_ext/.
--- a/csrc/build.py
+++ b/csrc/build.py
@ -0,0 +1,48 @@
+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",
+    "--threads=8",
+]
+
+
+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("attn_decode")
+register("attn_prefill")
+register("attn_paged_decode")
--- a/csrc/kernels/attn_common.h
+++ b/csrc/kernels/attn_common.h
@ -0,0 +1,54 @@
+#pragma once
+
+
+template<typename T, typename AT = float>
+struct AttentionParams {
+    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;
+    int num_splits;
+    float scale;
+
+    const T* __restrict__ q;
+    const T* __restrict__ k;
+    const T* __restrict__ v;
+    const bool* __restrict__ mask;
+
+    T* __restrict__ o;
+    AT* __restrict__ o_part;
+    AT* __restrict__ ml_part;
+};
+
+template<typename T, typename AT = float>
+struct PagedAttentionParams {
+    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;
+
+    int num_splits;
+    int page_size;
+    int max_pages;
+
+    const T* __restrict__ q;
+    const T* __restrict__ k_cache;
+    const T* __restrict__ v_cache;
+    const bool* __restrict__ mask;
+    const int64_t* __restrict__ page_table;
+    
+    T* __restrict__ o;
+    AT* __restrict__ o_part;
+    AT* __restrict__ ml_part;
+};
--- a/csrc/kernels/attn_decode.cu
+++ b/csrc/kernels/attn_decode.cu
@ -0,0 +1,111 @@
+#include "attn_decode_split_kv.cuh"
+#include "attn_entry_utils.cuh"
+
+#ifndef ASTRAI_NO_MMA
+#include "attn_decode_split_kv_mma.cuh"
+#endif
+
+static int decode_num_splits(int base_blocks, int tiles_total) {
+    int sm_count = 0;
+    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
+    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
+    return std::max(1, std::min(n, std::min(tiles_total, 32)));
+}
+
+// Scalar fallback: one warp per query head, split-KV across grid.z.
+static void launch_scalar_decode(AttentionParams<bf16>& p) {
+    int group_size = p.q_head / p.kv_head;
+    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
+    p.num_splits = decode_num_splits(p.batch * p.kv_head, chunks_total);
+
+    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
+    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
+    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
+    p.o_part = o_part.data_ptr<float>();
+    p.ml_part = ml_part.data_ptr<float>();
+
+    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
+    attn_decode_split_kv_kernel<<<dim3(p.batch * p.kv_head, 1, p.num_splits), dim3(32, group_size), smem>>>(p);
+    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+}
+
+#ifndef ASTRAI_NO_MMA
+// MMA head-packing requires G <= 16 (sQ has BR=16 rows). sm_80+ tensor-core
+// + cp.async wins even at G=1 (decode is memory-bound, not compute-bound).
+template <int HEAD_DIM, int BC>
+static void launch_mma_decode(AttentionParams<bf16>& p) {
+    int tiles_total = (p.kv_len + BC - 1) / BC;
+    p.num_splits = decode_num_splits(p.batch * p.kv_head, tiles_total);
+
+    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
+    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
+    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
+    p.o_part = o_part.data_ptr<float>();
+    p.ml_part = ml_part.data_ptr<float>();
+
+    attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
+        <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
+    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+}
+#endif
+
+template <int HEAD_DIM>
+static void dispatch_decode(AttentionParams<bf16>& p) {
+#ifndef ASTRAI_NO_MMA
+    int G = p.q_head / p.kv_head;
+    if (!p.use_mask && G >= 1 && G <= 16) {
+        launch_mma_decode<HEAD_DIM, 32>(p);
+        return;
+    }
+#endif
+    launch_scalar_decode(p);
+}
+
+torch::Tensor attn_decode(
+    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
+) {
+    AttentionParams<bf16> p;
+    attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
+    TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1");
+    TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
+
+    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("attn_decode", &attn_decode,
+        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)");
+}
--- a/csrc/kernels/attn_decode_split_kv.cuh
+++ b/csrc/kernels/attn_decode_split_kv.cuh
@ -0,0 +1,116 @@
+#pragma once
+#include <cuda_bf16.h>
+#include <float.h>
+#include "attn_common.h"
+
+using bf16 = __nv_bfloat16;
+constexpr int DC_CHUNK = 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;
+}
+
+__global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
+    int batch = blockIdx.x / p.kv_head;
+    int kv_head = blockIdx.x % p.kv_head;
+    int split = blockIdx.z;
+    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[];
+
+    // Split-KV: each split processes a contiguous subset of chunks
+    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
+    int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits;
+    int ch_begin = split * chunks_per_split;
+    int ch_end = min(chunks_total, ch_begin + chunks_per_split);
+
+    for (int ci = ch_begin; ci < ch_end; ci++) {
+        int chunk_start = ci * 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();
+    }
+
+    // ---- write UN-normalised partials for this split ----
+    size_t bh = (size_t)batch * p.q_head + q_head;
+    size_t slot = bh * p.num_splits + split;
+    int d0 = lane * hd_per_thread;
+    for (int i = 0; i < hd_per_thread; i++) {
+        int dd = d0 + i;
+        p.o_part[slot * p.head_dim + dd] = acc_reg[i];
+    }
+    if (lane == 0) {
+        p.ml_part[slot * 2] = m;
+        p.ml_part[slot * 2 + 1] = d;
+    }
+}
+
+// Reduce split-K partials into the final bf16 output. One block per (batch,
+// q_head); each thread folds across all splits with a single-pass
+// online-rescale reduction (expf + FMA counts halved vs 3-pass original).
+__global__ void attn_decode_combine_kernel(AttentionParams<bf16> p) {
+    int bh = blockIdx.x;
+    int d = threadIdx.x;
+    if (d >= p.head_dim) return;
+
+    size_t split_base = (size_t)bh * p.num_splits;
+    const float* mlp = p.ml_part + split_base * 2;
+    const float* op = p.o_part + split_base * p.head_dim;
+
+    float m = -FLT_MAX, l = 0.0f, acc = 0.0f;
+    for (int s = 0; s < p.num_splits; s++) {
+        float mi = mlp[s * 2];
+        if (mi <= -FLT_MAX) continue;
+        float li = mlp[s * 2 + 1];
+        float nm = fmaxf(m, mi);
+        float corr = __expf(m - nm);
+        float e = __expf(mi - nm);
+        acc = acc * corr + op[s * p.head_dim + d] * e;
+        l = l * corr + li * e;
+        m = nm;
+    }
+
+    float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
+    p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv);
+}
--- a/csrc/kernels/attn_decode_split_kv_mma.cuh
+++ b/csrc/kernels/attn_decode_split_kv_mma.cuh
@ -0,0 +1,172 @@
+#pragma once
+#include <cfloat>
+#include <cuda_bf16.h>
+#include "attn_common.h"
+#include "attn_mma_utils.cuh"
+
+using bf16 = __nv_bfloat16;
+
+// Split-K (FlashDecoding) tensor-core decode via GQA head-packing.
+//
+// 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). The KV sequence is partitioned across gridDim.z blocks so that
+// a decode with only batch*kv_head independent tasks can fill all SMs. Each
+// (batch, kv_head, split) block computes an UN-normalised partial (Oacc, m, l)
+// over its KV slice; the combine kernel below reduces across splits. Fixes the
+// "grid too small" bottleneck (0.04 waves/SM → many blocks) for long-context,
+// small-batch decode.
+//
+// Partial layout (float, contiguous):
+//   o_part : [batch, q_head, num_splits, HEAD_DIM]
+//   ml_part: [batch, q_head, num_splits, 2]  (m, l)
+//
+// 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 attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
+    constexpr int BR = 16;
+    constexpr int KD = HEAD_DIM / 16;
+    constexpr int NC8 = BC / 8;
+    constexpr int KT2 = BC / 16;
+    constexpr int DN8 = HEAD_DIM / 8;
+    constexpr int LD = HEAD_DIM;
+    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
+
+    const int lane = threadIdx.x;
+    const int gid = lane >> 2;
+    const int tid4 = lane & 3;
+
+    const int kv_head = blockIdx.x;
+    const int batch = blockIdx.y;
+    const int split = blockIdx.z;
+    const int G = p.q_head / p.kv_head;
+    const int q_head0 = kv_head * G;
+
+    __shared__ __align__(16) bf16 sK[BC * HEAD_DIM];
+    __shared__ __align__(16) bf16 sV[BC * HEAD_DIM];
+    __shared__ __align__(16) bf16 sQ[BR * HEAD_DIM];
+
+    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];
+        }
+        sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = val;
+    }
+    __syncwarp();
+
+    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_total = (p.kv_len + BC - 1) / BC;
+    const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
+    const int ti_begin = split * tiles_per_split;
+    const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
+    const int has_mask = p.use_mask && p.mask;
+
+    for (int ti = ti_begin; ti < ti_end; ti++) {
+        int kv0 = ti * BC;
+
+        bool full_tile = (kv0 + BC <= p.kv_len);
+        if (full_tile) {
+            constexpr int VEC = 8;
+            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();
+
+        float Sacc[NC8][4];
+        mma_compute_scores<KD, NC8>(Qa, sK, LD, SWIZ_MASK, lane, Sacc);
+
+        #pragma unroll
+        for (int n8 = 0; n8 < NC8; n8++)
+            Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
+            Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
+
+        int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
+        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
+                                    mask_base, p.mask, has_mask,
+                                    Sacc, Oacc, m0, m1, l0, l1, lane);
+
+        mma_pv_accumulate<DN8, KT2>(Sacc, sV, LD, SWIZ_MASK, lane, Oacc);
+        __syncwarp();
+    }
+
+    // ---- write UN-normalised partials for this split ----
+    auto split_slot = [&](int h) -> size_t {
+        size_t bh = (size_t)batch * p.q_head + h;
+        return bh * p.num_splits + split;
+    };
+#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 h = q_head0 + r0;
+            float* op = p.o_part + split_slot(h) * HEAD_DIM;
+            op[d] = Oacc[dn8][0];
+            op[d + 1] = Oacc[dn8][1];
+        }
+        if (r1 < G) {
+            int h = q_head0 + r1;
+            float* op = p.o_part + split_slot(h) * HEAD_DIM;
+            op[d] = Oacc[dn8][2];
+            op[d + 1] = Oacc[dn8][3];
+        }
+    }
+    if (tid4 == 0) {
+        int r0 = gid, r1 = gid + 8;
+        if (r0 < G) {
+            int h = q_head0 + r0;
+            float* mp = p.ml_part + split_slot(h) * 2;
+            mp[0] = m0; mp[1] = l0;
+        }
+        if (r1 < G) {
+            int h = q_head0 + r1;
+            float* mp = p.ml_part + split_slot(h) * 2;
+            mp[0] = m1; mp[1] = l1;
+        }
+    }
+}
+
--- a/csrc/kernels/attn_entry_utils.cuh
+++ b/csrc/kernels/attn_entry_utils.cuh
@ -0,0 +1,43 @@
+#pragma once
+#include <torch/extension.h>
+#include "attn_common.h"
+
+template<typename T>
+inline void attn_pack_params(
+    torch::Tensor q,
+    torch::Tensor k,
+    torch::Tensor v,
+    c10::optional<torch::Tensor> mask,
+    bool is_causal,
+    int64_t causal_offset,
+    c10::optional<double> scale,
+    AttentionParams<T>& p
+) {
+    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);
+
+    p.batch = (int)q.size(0);
+    p.q_head = (int)q.size(1);
+    p.kv_head = (int)k.size(1);
+    p.q_len = (int)q.size(2);
+    p.kv_len = (int)k.size(2);
+    p.head_dim = (int)q.size(3);
+    p.use_mask = mask.has_value() ? 1 : 0;
+    p.is_causal = is_causal ? 1 : 0;
+    p.causal_offset = (int)causal_offset;
+    p.scale = scale.has_value() ? (float)scale.value() : 1.0f / sqrtf((float)p.head_dim);
+    p.q = (const T*)q.data_ptr();
+    p.k = (const T*)k.data_ptr();
+    p.v = (const T*)v.data_ptr();
+    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;
+    }
+}
--- a/csrc/kernels/attn_mma_utils.cuh
+++ b/csrc/kernels/attn_mma_utils.cuh
@ -0,0 +1,237 @@
+#pragma once
+#include <cfloat>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+// 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));
+}
+
+// Predicated cp.async: copy 16 bytes when `pred`, otherwise zero-fill the
+// destination (src-size operand = 0 → no bytes read from src, so an
+// out-of-bounds src address is never dereferenced). Lets full and partial
+// tiles share one uniform async load path — no scalar fallback branch.
+__device__ __forceinline__ void cp_async_16_pred(bf16* smem_ptr,
+                                                  const void* gmem_ptr,
+                                                  bool pred) {
+    unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
+    int src_size = pred ? 16 : 0;
+    asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"
+                 :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
+}
+
+__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));
+}
+
+// ---------------------------------------------------------------------------
+// Shared MMA compute functions — used by both decode and prefill MMA kernels.
+// Extracted because S=Q@K^T, online softmax, and P@V are structurally identical
+// between the two kernels; only the per-row causal/mask bounds differ.
+// ---------------------------------------------------------------------------
+
+// S = Q @ K^T  (Qa pre-loaded and pre-scaled by the caller).
+// LD and SWIZ_MASK are constexpr in the calling kernel — passing them as
+// runtime ints lets the compiler fold them while keeping the signature clean.
+template <int KD, int NC8>
+__device__ inline void mma_compute_scores(
+    const unsigned Qa[KD][4],
+    const bf16* __restrict__ sK,
+    int LD, int SWIZ_MASK, int lane,
+    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 + Oacc rescale for one K/V tile. maxc0/maxc1 are the per-row
+// KV column bounds — prefill passes per-query-row causal limits while decode
+// passes the same value for both rows (q_len==1). Sacc is consumed in place
+// (replaced by P = exp(S - nm) for the subsequent P@V step).
+template <int NC8, int DN8>
+__device__ inline void mma_softmax_tile(
+    int kv0,
+    int maxc0, int maxc1,
+    int mask_base,
+    const bool* __restrict__ mask,
+    bool has_mask,
+    float Sacc[NC8][4],
+    float Oacc[DN8][4],
+    float& m0, float& m1,
+    float& l0, float& l1,
+    int lane)
+{
+    int tid4 = lane & 3;
+
+    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 && !mask[mask_base + cc]);
+        bool b1 = (c1 >= maxc0) || (has_mask && !mask[mask_base + c1]);
+        bool b2 = (cc >= maxc1) || (has_mask && !mask[mask_base + cc]);
+        bool b3 = (c1 >= maxc1) || (has_mask && !mask[mask_base + 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;
+
+    #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  (Sacc must contain P = attention weights after softmax).
+template <int DN8, int KT2>
+__device__ inline void mma_pv_accumulate(
+    float Sacc[][4],
+    const bf16* __restrict__ sV,
+    int LD, int SWIZ_MASK, int lane,
+    float Oacc[DN8][4])
+{
+    #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]);
+        }
+    }
+}
--- a/csrc/kernels/attn_paged_decode.cu
+++ b/csrc/kernels/attn_paged_decode.cu
@ -0,0 +1,152 @@
+#include "attn_paged_decode_split_kv.cuh"
+#ifndef ASTRAI_NO_MMA
+#include "attn_paged_decode_split_kv_mma.cuh"
+#endif
+
+#include <torch/extension.h>
+#include <c10/cuda/CUDAGuard.h>
+
+static int paged_decode_num_splits(int base_blocks, int tiles_total) {
+    int sm_count = 0;
+    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
+    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
+    return std::max(1, std::min(n, std::min(tiles_total, 32)));
+}
+
+static void launch_paged_scalar_decode(PagedAttentionParams<bf16>& p) {
+    int group_size = p.q_head / p.kv_head;
+    int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
+    p.num_splits = paged_decode_num_splits(p.batch * p.kv_head, chunks_total);
+
+    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
+    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
+    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
+    p.o_part = o_part.data_ptr<float>();
+    p.ml_part = ml_part.data_ptr<float>();
+
+    size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
+    paged_attn_decode_split_kv_kernel<<<
+        dim3(p.batch * p.kv_head, 1, p.num_splits),
+        dim3(32, group_size),
+        smem>>>(p);
+    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+}
+
+#ifndef ASTRAI_NO_MMA
+template <int HEAD_DIM, int BC>
+static void launch_paged_mma_decode(PagedAttentionParams<bf16>& p) {
+    int tiles_total = (p.kv_len + BC - 1) / BC;
+    p.num_splits = paged_decode_num_splits(p.batch * p.kv_head, tiles_total);
+
+    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
+    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
+    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
+    p.o_part = o_part.data_ptr<float>();
+    p.ml_part = ml_part.data_ptr<float>();
+
+    paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
+        <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
+    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+}
+#endif
+
+template <int HEAD_DIM>
+static void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
+#ifndef ASTRAI_NO_MMA
+    int G = p.q_head / p.kv_head;
+    if (!p.use_mask && G >= 1 && G <= 16 && p.page_size >= 32) {
+        launch_paged_mma_decode<HEAD_DIM, 32>(p);
+        return;
+    }
+#endif
+    launch_paged_scalar_decode(p);
+}
+
+torch::Tensor attn_paged_decode(
+    torch::Tensor q,
+    torch::Tensor page_table,
+    torch::Tensor k_cache,
+    torch::Tensor v_cache,
+    int64_t page_size,
+    int64_t kv_len,
+    c10::optional<torch::Tensor> mask,
+    bool is_causal = false,
+    int64_t causal_offset = 0,
+    c10::optional<double> scale = c10::nullopt
+) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+
+    int batch = q.size(0);
+    int q_head = q.size(1);
+    int head_dim = q.size(3);
+    int kv_head = k_cache.size(2);
+    int max_pages = page_table.size(1);
+
+    TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
+    TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bf16");
+    TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be bf16");
+    TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be bf16");
+    TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64");
+    TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)");
+    TORCH_CHECK(head_dim % 32 == 0, "head_dim must be multiple of 32");
+    TORCH_CHECK(k_cache.size(1) == page_size,
+                "k_cache dim 1 must equal page_size, got ",
+                k_cache.size(1), " vs ", page_size);
+    TORCH_CHECK(k_cache.size(0) >= 0, "k_cache must have at least 0 pages");
+
+    float scale_val = scale.has_value()
+        ? static_cast<float>(scale.value())
+        : 1.0f / std::sqrt(static_cast<float>(head_dim));
+
+    auto O = torch::empty_like(q);
+
+    PagedAttentionParams<bf16, float> p;
+    p.batch = batch;
+    p.q_head = q_head;
+    p.kv_head = kv_head;
+    p.q_len = 1;
+    p.kv_len = static_cast<int>(kv_len);
+    p.head_dim = head_dim;
+    p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
+    p.is_causal = is_causal ? 1 : 0;
+    p.causal_offset = static_cast<int>(causal_offset);
+    p.num_splits = 1;
+    p.scale = scale_val;
+    p.page_size = static_cast<int>(page_size);
+    p.max_pages = max_pages;
+    p.page_table = page_table.data_ptr<int64_t>();
+    p.k_cache = reinterpret_cast<const bf16*>(k_cache.data_ptr());
+    p.v_cache = reinterpret_cast<const bf16*>(v_cache.data_ptr());
+    p.q = reinterpret_cast<const bf16*>(q.data_ptr());
+    p.mask = p.use_mask ? mask.value().data_ptr<bool>() : nullptr;
+    p.o = reinterpret_cast<bf16*>(O.data_ptr());
+    p.o_part = nullptr;
+    p.ml_part = nullptr;
+
+    switch (p.head_dim) {
+        case 32:  dispatch_paged_decode<32>(p);  break;
+        case 64:  dispatch_paged_decode<64>(p);  break;
+        case 128: dispatch_paged_decode<128>(p); break;
+        case 256: dispatch_paged_decode<256>(p); break;
+        default:
+            TORCH_CHECK(false, "paged_decode: unsupported head_dim ", p.head_dim,
+                        " (supported: 32, 64, 128, 256)");
+    }
+
+    return O;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("attn_paged_decode", &attn_paged_decode,
+        py::arg("q"),
+        py::arg("page_table"),
+        py::arg("k_cache"),
+        py::arg("v_cache"),
+        py::arg("page_size"),
+        py::arg("kv_len"),
+        py::arg("mask") = py::none(),
+        py::arg("is_causal") = false,
+        py::arg("causal_offset") = 0,
+        py::arg("scale") = py::none(),
+        "Paged GQA decode — split-KV with direct page-table access.");
+}
--- a/csrc/kernels/attn_paged_decode_split_kv.cuh
+++ b/csrc/kernels/attn_paged_decode_split_kv.cuh
@ -0,0 +1,140 @@
+#pragma once
+#include <cuda_bf16.h>
+#include <float.h>
+#include "attn_common.h"
+
+using bf16 = __nv_bfloat16;
+constexpr int PDC_CHUNK = 64;
+
+__device__ inline float paged_warp_reduce_sum(float val) {
+    for (int offset = 16; offset > 0; offset >>= 1)
+        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
+    return val;
+}
+
+// Split-KV scalar decode: one warp per query head, grid.z partitions KV.
+__global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p) {
+    int batch = blockIdx.x / p.kv_head;
+    int kv_head = blockIdx.x % p.kv_head;
+    int split = blockIdx.z;
+    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;
+    #pragma unroll
+    for (int i = 0; i < hd_per_thread; i++)
+        q_reg[i] = __bfloat162float(p.q[q_off + i]);
+
+    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
+
+    extern __shared__ __align__(16) bf16 k_smem[];
+
+    int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
+    int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits;
+    int ch_begin = split * chunks_per_split;
+    int ch_end = min(chunks_total, ch_begin + chunks_per_split);
+
+    const int mask_base = batch * p.kv_len;
+
+    for (int ci = ch_begin; ci < ch_end; ci++) {
+        int chunk_start = ci * PDC_CHUNK;
+        int this_chunk = min(PDC_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) {
+            int s = i / p.head_dim;
+            int d_dim = i % p.head_dim;
+            int pos = chunk_start + s;
+            int logical_page = pos / p.page_size;
+            int page_offset = pos % p.page_size;
+            int phys_page = p.page_table[batch * p.max_pages + logical_page];
+            if (phys_page >= 0) {
+                int64_t off = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
+                            + (int64_t)page_offset * p.kv_head * p.head_dim
+                            + (int64_t)kv_head * p.head_dim
+                            + d_dim;
+                k_smem[i] = p.k_cache[off];
+            } else {
+                k_smem[i] = __float2bfloat16(0.0f);
+            }
+        }
+        __syncthreads();
+
+        for (int s = 0; s < this_chunk; s++) {
+            float partial = 0.0f;
+            #pragma unroll
+            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 = paged_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 pos = chunk_start + s;
+            int logical_page = pos / p.page_size;
+            int page_offset = pos % p.page_size;
+            int phys_page = p.page_table[batch * p.max_pages + logical_page];
+            if (phys_page >= 0) {
+                int64_t v_base = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
+                               + (int64_t)page_offset * p.kv_head * p.head_dim
+                               + (int64_t)kv_head * p.head_dim;
+                #pragma unroll
+                for (int i = 0; i < hd_per_thread; i++)
+                    acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v_cache[v_base + lane * hd_per_thread + i]) * beta;
+            } else {
+                #pragma unroll
+                for (int i = 0; i < hd_per_thread; i++)
+                    acc_reg[i] = acc_reg[i] * alpha + 0.0f * beta;
+            }
+            m = new_m;
+        }
+        __syncthreads();
+    }
+
+    size_t bh = (size_t)batch * p.q_head + q_head;
+    size_t slot = bh * p.num_splits + split;
+    int d0 = lane * hd_per_thread;
+    #pragma unroll
+    for (int i = 0; i < hd_per_thread; i++)
+        p.o_part[slot * p.head_dim + (d0 + i)] = acc_reg[i];
+    if (lane == 0) {
+        p.ml_part[slot * 2] = m;
+        p.ml_part[slot * 2 + 1] = d;
+    }
+}
+
+__global__ void paged_attn_decode_combine_kernel(PagedAttentionParams<bf16> p) {
+    int bh = blockIdx.x;
+    int d = threadIdx.x;
+    if (d >= p.head_dim) return;
+
+    size_t split_base = (size_t)bh * p.num_splits;
+    const float* mlp = p.ml_part + split_base * 2;
+    const float* op = p.o_part + split_base * p.head_dim;
+
+    float m = -FLT_MAX, l = 0.0f, acc = 0.0f;
+    for (int s = 0; s < p.num_splits; s++) {
+        float mi = mlp[s * 2];
+        if (mi <= -FLT_MAX) continue;
+        float li = mlp[s * 2 + 1];
+        float nm = fmaxf(m, mi);
+        float corr = __expf(m - nm);
+        float e = __expf(mi - nm);
+        acc = acc * corr + op[s * p.head_dim + d] * e;
+        l = l * corr + li * e;
+        m = nm;
+    }
+
+    float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
+    p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv);
+}
--- a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh
+++ b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh
@ -0,0 +1,174 @@
+#pragma once
+#include <cfloat>
+#include <cuda_bf16.h>
+#include "attn_common.h"
+#include "attn_mma_utils.cuh"
+
+using bf16 = __nv_bfloat16;
+
+// Paged split-KV tensor-core decode via GQA head-packing.
+// Identical algorithm to attn_decode_split_kv_mma_kernel but reads K/V
+// directly from the page pool through a page table, eliminating the gather
+// copy.  Each tile (BC=32) fits within a single page (page_size >= 32), so
+// the page-table lookup happens once per tile for cp.async.
+template <int HEAD_DIM, int BC>
+__global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16> p) {
+    constexpr int BR = 16;
+    constexpr int KD = HEAD_DIM / 16;
+    constexpr int NC8 = BC / 8;
+    constexpr int KT2 = BC / 16;
+    constexpr int DN8 = HEAD_DIM / 8;
+    constexpr int LD = HEAD_DIM;
+    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
+
+    const int lane = threadIdx.x;
+    const int gid = lane >> 2;
+    const int tid4 = lane & 3;
+
+    const int kv_head_idx = blockIdx.x;
+    const int batch = blockIdx.y;
+    const int split = blockIdx.z;
+    const int G = p.q_head / p.kv_head;
+    const int q_head0 = kv_head_idx * G;
+
+    __shared__ __align__(16) bf16 sK[BC * HEAD_DIM];
+    __shared__ __align__(16) bf16 sV[BC * HEAD_DIM];
+    __shared__ __align__(16) bf16 sQ[BR * HEAD_DIM];
+
+    // ---- load Q into registers via ldmatrix ----
+    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];
+        }
+        sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = val;
+    }
+    __syncwarp();
+
+    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 mask_base = batch * p.kv_len;
+    const int tiles_total = (p.kv_len + BC - 1) / BC;
+    const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
+    const int ti_begin = split * tiles_per_split;
+    const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
+    const int has_mask = p.use_mask && p.mask;
+
+    // Paged strides (constant for the block)
+    const int64_t page_stride = (int64_t)p.page_size * p.kv_head * HEAD_DIM;
+    const int64_t pos_stride  = (int64_t)p.kv_head * HEAD_DIM;
+    const int64_t head_off    = (int64_t)kv_head_idx * HEAD_DIM;
+
+    for (int ti = ti_begin; ti < ti_end; ti++) {
+        int kv0 = ti * BC;
+
+        // phys_page is constant for the whole tile (BC <= page_size).
+        int logical_page = kv0 / p.page_size;
+        int phys_page = p.page_table[batch * p.max_pages + logical_page];
+        bool page_valid = (phys_page >= 0);
+
+        bool full_tile = page_valid && (kv0 + BC <= p.kv_len);
+        if (full_tile) {
+            constexpr int VEC = 8;
+            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;
+                int page_off = kc % p.page_size;
+                int64_t gmem_base = (int64_t)phys_page * page_stride
+                                  + (int64_t)page_off * pos_stride
+                                  + head_off;
+                cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)],
+                            &p.k_cache[gmem_base + d]);
+                cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)],
+                            &p.v_cache[gmem_base + 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);
+                if (kc < p.kv_len && page_valid) {
+                    int page_off = kc % p.page_size;
+                    int64_t gmem_base = (int64_t)phys_page * page_stride
+                                      + (int64_t)page_off * pos_stride
+                                      + head_off;
+                    sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = p.k_cache[gmem_base + d];
+                    sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = p.v_cache[gmem_base + d];
+                } else {
+                    sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = z;
+                    sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = z;
+                }
+            }
+        }
+        __syncwarp();
+
+        float Sacc[NC8][4];
+        mma_compute_scores<KD, NC8>(Qa, sK, LD, SWIZ_MASK, lane, Sacc);
+
+        #pragma unroll
+        for (int n8 = 0; n8 < NC8; n8++)
+            Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
+            Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
+
+        int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
+        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
+                                    mask_base, p.mask, has_mask,
+                                    Sacc, Oacc, m0, m1, l0, l1, lane);
+
+        mma_pv_accumulate<DN8, KT2>(Sacc, sV, LD, SWIZ_MASK, lane, Oacc);
+        __syncwarp();
+    }
+
+    // ---- write UN-normalised partials for this split ----
+    auto split_slot = [&](int h) -> size_t {
+        size_t bh = (size_t)batch * p.q_head + h;
+        return bh * p.num_splits + split;
+    };
+    #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 h = q_head0 + r0;
+            float* op = p.o_part + split_slot(h) * HEAD_DIM;
+            op[d] = Oacc[dn8][0];
+            op[d + 1] = Oacc[dn8][1];
+        }
+        if (r1 < G) {
+            int h = q_head0 + r1;
+            float* op = p.o_part + split_slot(h) * HEAD_DIM;
+            op[d] = Oacc[dn8][2];
+            op[d + 1] = Oacc[dn8][3];
+        }
+    }
+    if (tid4 == 0) {
+        int r0 = gid, r1 = gid + 8;
+        if (r0 < G) {
+            int h = q_head0 + r0;
+            float* mp = p.ml_part + split_slot(h) * 2;
+            mp[0] = m0; mp[1] = l0;
+        }
+        if (r1 < G) {
+            int h = q_head0 + r1;
+            float* mp = p.ml_part + split_slot(h) * 2;
+            mp[0] = m1; mp[1] = l1;
+        }
+    }
+}
--- a/csrc/kernels/attn_prefill.cu
+++ b/csrc/kernels/attn_prefill.cu
@ -0,0 +1,81 @@
+#include "attn_prefill_split_q.cuh"
+#include "attn_entry_utils.cuh"
+
+#ifndef ASTRAI_NO_MMA
+#include "attn_prefill_split_q_mma.cuh"
+#endif
+
+template <int HEAD_DIM>
+static void dispatch_prefill(AttentionParams<bf16>& p) {
+#ifndef ASTRAI_NO_MMA
+    constexpr int WARPS = 4, BR = 16;
+    // KV tile: bigger tiles amortize the per-tile cp.async wait + barrier +
+    // loop overhead over more tensor-core work (this kernel is latency-bound,
+    // not compute/bandwidth-bound), so BC=32 wins ~6-8% over BC=16 for
+    // D<=128. D=256 stays at 16: BC=32 double-buffered would need 64KB smem,
+    // over the 48KB static cap. Both keep 3 blocks/SM (2 for D=256).
+    constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
+    // Register-hint MIN_BLOCKS tuned per HEAD_DIM's (BC=32) smem+register
+    // footprint: the largest blocks/SM that avoids register spills.
+    constexpr int MIN_BLOCKS = (HEAD_DIM <= 32) ? 6 : (HEAD_DIM <= 64) ? 4
+                             : (HEAD_DIM <= 128) ? 3 : 2;
+    dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
+    dim3 block(WARPS * 32, 1, 1);
+    // Static shared memory — no dynamic smem or cudaFuncSetAttribute needed.
+    // sK[BC*LD] + sV[BC*LD] + sQ[BR*LD], all sized by template params.
+    attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(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);
+    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
+#endif
+}
+
+torch::Tensor attn_prefill(
+    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
+) {
+    AttentionParams<bf16> p;
+    attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
+    TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
+
+    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("attn_prefill", &attn_prefill,
+        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)");
+}
--- a/csrc/kernels/attn_prefill_split_q.cuh
+++ b/csrc/kernels/attn_prefill_split_q.cuh
@ -0,0 +1,140 @@
+#pragma once
+#include <cfloat>
+#include <cuda_bf16.h>
+#include "attn_common.h"
+
+using bf16 = __nv_bfloat16;
+
+// 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 attn_prefill_split_q_kernel_t(AttentionParams<bf16> 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);
+
+    __shared__ __align__(16) bf16 sK[P_BC * HEAD_DIM];
+    __shared__ __align__(16) bf16 sV[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);
+    }
+}
--- a/csrc/kernels/attn_prefill_split_q_mma.cuh
+++ b/csrc/kernels/attn_prefill_split_q_mma.cuh
@ -0,0 +1,209 @@
+#pragma once
+#include <cfloat>
+#include <cuda_bf16.h>
+#include "attn_common.h"
+#include "attn_mma_utils.cuh"
+
+using bf16 = __nv_bfloat16;
+
+// Tensor-core prefill 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 fragments are loaded once
+// straight from global into the mma A-operand layout (no smem staging) and
+// kept resident in registers across the tile loop. S, O, and the online-softmax
+// stats (m, l) also live in registers.
+// Shared memory is statically sized via template parameters — no dynamic
+// allocation. 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, MIN_BLOCKS> with BC a
+// multiple of 16.
+//
+// Occupancy: __launch_bounds__ forces the compiler to fit MIN_BLOCKS blocks/SM,
+// spilling to local memory as needed. MIN_BLOCKS is tuned per HEAD_DIM to the
+// double-buffered smem footprint (2*BC*LD for each of K/V).
+//
+// Software pipeline: K/V are double-buffered and loaded via cp.async one tile
+// ahead, so the next tile streams from global memory while the current tile's
+// tensor-core math runs — hiding load latency (long_scoreboard). A single
+// __syncthreads per tile both publishes the freshly loaded tile cross-warp and
+// (because it runs before the next prefetch) guards the buffer being refilled,
+// so no second barrier is needed. Predicated cp.async (cp_async_16_pred)
+// zero-fills rows past kv_len, unifying full and partial tiles on one path.
+// BC=32 (D<=128) amortizes the per-tile wait+barrier+loop overhead over more
+// tensor-core work — this kernel is latency-bound (low occupancy from high
+// register pressure), so fewer, larger tiles beat many tiny ones.
+//
+// Optimizations: load Q fragments directly from global in mma A-operand layout
+// (no sQ staging, no prologue barriers); pre-scale Q by attention scale during Q load; packed bf16x2 output stores;
+// causal tile skipping (block-level prefetch bound + warp-level compute skip);
+// XOR swizzle (swiz_col) → eliminates ldmatrix bank conflicts without LD
+// padding (LD=HEAD_DIM).
+
+template <int HEAD_DIM, int WARPS, int BC, int MIN_BLOCKS>
+__global__ __launch_bounds__(WARPS * 32, MIN_BLOCKS)
+void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> 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;
+
+    // Static shared memory — sized by template parameters at compile time.
+    // K/V are double-buffered (STAGES=2): the next tile's cp.async load runs
+    // while the current tile's tensor-core math executes, hiding global-load
+    // latency (FA2-style software pipeline). No dynamic smem / carveout opt-in.
+    constexpr int STAGES = 2;
+    __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
+    __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
+
+    // Load the Q fragments straight from global into the mma A-operand layout
+    // (m16n8k16, row-major): no sQ staging area and no serialized per-warp
+    // prologue barriers. Each lane reads exactly the 8 Q elements ldmatrix
+    // would have produced, pre-scaled by the attention scale. Kept resident in
+    // registers across the tile loop.
+    //   frag[0]/[2]: row = qrow0 + gid ;  frag[1]/[3]: row = qrow0 + gid + 8
+    //   frag[0]/[1]: cols kt*16 + tid4*2 + {0,1} ; frag[2]/[3]: + 8
+    const int q_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
+    const int qra = qrow0 + gid;
+    const int qrb = qrow0 + gid + 8;
+    const bool va = qra < p.q_len, vb = qrb < p.q_len;
+    unsigned Qa[KD][4];
+#pragma unroll
+    for (int kt = 0; kt < KD; kt++) {
+        int c = kt * 16 + tid4 * 2;
+        const unsigned* pau = reinterpret_cast<const unsigned*>(
+            &p.q[q_base + qra * HEAD_DIM + c]);
+        const unsigned* pbu = reinterpret_cast<const unsigned*>(
+            &p.q[q_base + qrb * HEAD_DIM + c]);
+        Qa[kt][0] = va ? pau[0] : 0u;
+        Qa[kt][1] = vb ? pbu[0] : 0u;
+        Qa[kt][2] = va ? pau[4] : 0u;
+        Qa[kt][3] = vb ? pbu[4] : 0u;
+    }
+
+    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;
+
+    // Last active tile: block-level causal bound (all warps in the block share
+    // the K/V load, so the prefetch range is the block max, not per-warp).
+    int t_end = tiles - 1;
+    if (use_skip) {
+        int bt = block_max_kv / BC;
+        if (bt < t_end) t_end = bt;
+    }
+
+    constexpr int VEC = 8;  // bf16 per cp.async unit (16 bytes)
+    constexpr int TOTAL = BC * HEAD_DIM;
+
+    // Issue cp.async loads for tile `ti` into shared buffer `buf`. Predicated
+    // loads zero-fill rows past kv_len, so partial tiles need no scalar path.
+    auto load_tile = [&](int ti, int buf) {
+        int kv0 = ti * BC;
+        bf16* dK = sK + buf * BC * LD;
+        bf16* dV = sV + buf * BC * LD;
+#pragma unroll
+        for (int i = threadIdx.x * VEC; i < TOTAL; i += nthreads * VEC) {
+            int r = i / HEAD_DIM, d = i % HEAD_DIM;
+            int kc = kv0 + r;
+            bool valid = kc < p.kv_len;
+            int off = r * LD + swiz_col(d, r, SWIZ_MASK);
+            cp_async_16_pred(&dK[off], &p.k[kv_base + kc * HEAD_DIM + d], valid);
+            cp_async_16_pred(&dV[off], &p.v[kv_base + kc * HEAD_DIM + d], valid);
+        }
+        cp_async_commit();
+    };
+
+    // Prologue: kick off the first tile's load.
+    load_tile(0, 0);
+
+    for (int ti = 0; ti <= t_end; ti++) {
+        int buf = ti & 1;
+
+        // Wait for the current tile's async copies, then a single barrier: it
+        // both publishes this tile's data cross-warp AND guarantees the prior
+        // compute on the buffer we are about to refill has finished. Issuing
+        // the next tile's load *after* this barrier lets one barrier cover both
+        // hazards (vs two), while the load still overlaps this tile's math.
+        cp_async_wait_group<0>();
+        __syncthreads();
+        if (ti < t_end) load_tile(ti + 1, (ti + 1) & 1);
+
+        const bf16* bK = sK + buf * BC * LD;
+        const bf16* bV = sV + buf * BC * LD;
+        int kv0 = ti * BC;
+
+        // Warp-level causal skip
+        if (!use_skip || kv0 <= max_kv) {
+
+        // S = Q @ K^T + scale + online softmax + O += P @ V
+        float Sacc[NC8][4];
+        mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc);
+
+        // post-multiply scale in float (no bf16 precision loss from pre-scaling Q)
+        #pragma unroll
+        for (int n8 = 0; n8 < NC8; n8++)
+            Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
+            Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.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;
+        mma_softmax_tile<NC8, DN8>(kv0, maxc0, maxc1,
+                                    mb, p.mask, has_mask,
+                                    Sacc, Oacc, m0, m1, l0, l1, lane);
+
+        mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
+        }  // if active (warp-level causal skip)
+    }
+
+    // ---- write output ---- (packed bf16x2 stores: one 32-bit STG per pair,
+    // halves store count and removes the uncoalesced scalar-store penalty)
+    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) {
+            __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0,
+                                                     Oacc[dn8][1] * rl0);
+            *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr0 * HEAD_DIM + d]) = v;
+        }
+        if (qr1 < p.q_len) {
+            __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
+                                                     Oacc[dn8][3] * rl1);
+            *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr1 * HEAD_DIM + d]) = v;
+        }
+    }
+}
--- a/csrc/tests/attn_decode_test.cu
+++ b/csrc/tests/attn_decode_test.cu
@ -0,0 +1,225 @@
+/*
+Pure-C test:
+nvcc -I csrc -arch=sm_89 -O3 \
+    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
+    csrc/tests/attn_decode_test.cu -o test && ./test
+*/
+
+#include "test_utils.cuh"
+#include "../kernels/attn_decode_split_kv.cuh"
+#ifndef ASTRAI_NO_MMA
+#include "../kernels/attn_decode_split_kv_mma.cuh"
+#endif
+
+// Split-K scratch (torch-free): the production launcher allocates these from
+// torch; here we pass pre-allocated device buffers so the bench loop doesn't
+// pay a cudaMalloc per iteration. Size for the maximum split count (32).
+struct DecodeScratch {
+    float* o_part = nullptr;
+    float* ml_part = nullptr;
+};
+
+// Launch the production decode path (tensor-core head-packing MMA on sm_80+,
+// scalar fallback otherwise), mirroring dispatch_decode() in attn_decode.cu.
+#ifndef ASTRAI_NO_MMA
+static bool decode_use_mma(const AttentionParams<bf16>& p) {
+    int G = p.q_head / p.kv_head;
+    return !p.use_mask && G > 1 && G <= 16;
+}
+
+template <int HEAD_DIM, int BC>
+static void launch_mma_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
+    int tiles_total = (p.kv_len + BC - 1) / BC;
+    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
+    p.o_part = sc.o_part;
+    p.ml_part = sc.ml_part;
+
+    attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
+        <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
+    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+}
+#endif
+
+static void launch_scalar_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
+    int gs = p.q_head / p.kv_head;
+    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
+    p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
+    p.o_part = sc.o_part;
+    p.ml_part = sc.ml_part;
+
+    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
+    attn_decode_split_kv_kernel<<<dim3(p.batch * p.kv_head, 1, p.num_splits), dim3(32, gs), smem>>>(p);
+    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+}
+
+template <int HEAD_DIM>
+static void dispatch_decode_t(AttentionParams<bf16>& p, DecodeScratch& sc) {
+#ifndef ASTRAI_NO_MMA
+    if (decode_use_mma(p)) { launch_mma_decode<HEAD_DIM, 32>(p, sc); return; }
+#endif
+    launch_scalar_decode(p, sc);
+}
+
+static void dispatch_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
+    switch (p.head_dim) {
+        case 32:  dispatch_decode_t<32>(p, sc);  break;
+        case 64:  dispatch_decode_t<64>(p, sc);  break;
+        case 128: dispatch_decode_t<128>(p, sc); break;
+        case 256: dispatch_decode_t<256>(p, sc); break;
+        default:  printf("bench: unsupported D=%d\n", p.head_dim);
+    }
+}
+
+// Warmed-up, CUDA-event timed sweep over the production decode MMA path.
+// Decode (q_len==1) is memory-bound: the two matmuls are GEMV-shaped, so we
+// report both effective K/V read bandwidth and the (small) attention FLOP/s.
+// FLOP/s = 2 matmuls (q@K^T, P@V), each 2*B*Hq*kv*D flops.
+// Bytes    = K + V read = 2 * B*Hk*kv*D * sizeof(bf16).
+static void bench() {
+    const int cfgs[][5] = {
+        {1, 32, 4, 512, 128},    // B,Hq,Hk,seq,D
+        {1, 32, 4, 1024, 128},
+        {1, 32, 4, 2048, 128},
+        {1, 32, 4, 4096, 128},
+        {16, 32, 4, 2048, 128},
+        {32, 32, 4, 1024, 128},
+    };
+    int n = sizeof(cfgs)/sizeof(cfgs[0]);
+    const int WARMUP = 10, ITERS = 100;
+    printf("\n===== DECODE BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
+    printf("%-46s | %10s | %10s | %10s\n",
+           "config", "latency", "bandwidth", "throughput");
+    printf("---------------------------------------------------------------"
+           "----------------------------\n");
+
+    for (int ci = 0; ci < n; ci++) {
+        int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
+        int sl=cfgs[ci][3], D=cfgs[ci][4];
+        size_t nQ=(size_t)B*Hq*D, nKV=(size_t)B*Hk*sl*D;
+
+        bf16 *dQ,*dK,*dV,*dO,*tmp;
+        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+        size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
+        for (size_t i=0;i<nQ;i++)  tmp[i]=f2bf(randf());
+        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+
+        AttentionParams<bf16> p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
+        p.use_mask=0; 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=nullptr; p.o=dO;
+
+        DecodeScratch sc;
+        cudaMalloc(&sc.o_part, (size_t)B*Hq*32*D*sizeof(float));
+        cudaMalloc(&sc.ml_part, (size_t)B*Hq*32*2*sizeof(float));
+
+        for (int i=0;i<WARMUP;i++) dispatch_decode(p, sc);
+        cudaDeviceSynchronize();
+        cudaError_t err=cudaGetLastError();
+        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
+
+        cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
+        cudaEventRecord(s);
+        for (int i=0;i<ITERS;i++) dispatch_decode(p, sc);
+        cudaEventRecord(e); cudaEventSynchronize(e);
+        float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
+
+        double flops = 4.0*B*Hq*(double)sl*D;
+        double tflops = flops/(ms*1e-3)/1e12;
+        // HBM traffic: K + V read (B*Hk*sl*D each), bf16; Q/O negligible.
+        double bytes = 2.0 * (2.0*nKV);
+        double gbps = bytes/(ms*1e-3)/1e9;
+
+        char cfg[64];
+        snprintf(cfg, sizeof(cfg),
+                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
+                 B,Hq,Hk,1,sl,D,0);
+        printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
+               cfg, ms, gbps, tflops);
+
+        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
+        cudaFree(sc.o_part);cudaFree(sc.ml_part);
+        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
+    }
+}
+
+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);
+
+        AttentionParams<bf16> p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
+        p.use_mask=0; p.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=nullptr; p.o=dO;
+
+        // Split-K scratch (max 32 splits), sized for the production MMA path.
+        DecodeScratch sc;
+        cudaMalloc(&sc.o_part, (size_t)B*Hq*32*D*sizeof(float));
+        cudaMalloc(&sc.ml_part, (size_t)B*Hq*32*2*sizeof(float));
+
+        double t0=now_ms();
+        dispatch_decode(p, sc);
+        cudaDeviceSynchronize();
+        double kms=now_ms()-t0;
+        cudaError_t err=cudaGetLastError();
+        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
+
+        bf16* hOut=new bf16[nQ];
+        cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
+
+        float* ref=new float[nQ];
+        cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, 0, 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);cudaFree(dMask);
+        cudaFree(sc.o_part);cudaFree(sc.ml_part);
+        delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
+    }
+    printf("All tests passed!\n");
+    bench();
+    return 0;
+}
--- a/csrc/tests/attn_paged_decode_test.cu
+++ b/csrc/tests/attn_paged_decode_test.cu
@ -0,0 +1,346 @@
+// Compile:
+//   nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
+//        --extra-device-vectorization csrc/tests/attn_paged_decode_test.cu \
+//        -o /tmp/test_paged && /tmp/test_paged
+
+#include <cstring>
+#include "test_utils.cuh"
+#include "../kernels/attn_paged_decode_split_kv.cuh"
+#ifndef ASTRAI_NO_MMA
+#include "../kernels/attn_paged_decode_split_kv_mma.cuh"
+#endif
+
+// Copy contiguous K/V from page pool (reference gather)
+static void gather_kv_cpu(
+    const bf16* h_k_pool, const bf16* h_v_pool,
+    const int64_t* h_pt, int B, int Hkv, int kv_len,
+    int page_size, int head_dim,
+    bf16* h_k, bf16* h_v)
+{
+    int max_pages = (kv_len + page_size - 1) / page_size;
+    size_t page_stride = (size_t)page_size * Hkv * head_dim;
+    for (int b = 0; b < B; b++) {
+        for (int pos = 0; pos < kv_len; pos++) {
+            int log_pg = pos / page_size;
+            int pg_off = pos % page_size;
+            int phys = (int)h_pt[b * max_pages + log_pg];
+            for (int h = 0; h < Hkv; h++) {
+                size_t src_base = (size_t)phys * page_stride
+                                + (size_t)pg_off * Hkv * head_dim
+                                + h * head_dim;
+                size_t dst_base = ((size_t)b * Hkv + h) * kv_len * head_dim + (size_t)pos * head_dim;
+                memcpy(h_k + dst_base, h_k_pool + src_base, head_dim * sizeof(bf16));
+                memcpy(h_v + dst_base, h_v_pool + src_base, head_dim * sizeof(bf16));
+            }
+        }
+    }
+}
+
+template <int HEAD_DIM>
+static void launch_paged_decode(PagedAttentionParams<bf16, float>& p) {
+#ifndef ASTRAI_NO_MMA
+    int G_check = p.q_head / p.kv_head;
+    bool use_mma = !p.use_mask && G_check >= 1 && G_check <= 16 && p.page_size >= 32;
+    if (use_mma) {
+        int tiles_total = (p.kv_len + 32 - 1) / 32;
+        p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
+        paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, 32>
+            <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
+    } else
+#endif
+    {
+        int group_sz = p.q_head / p.kv_head;
+        int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
+        p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
+        size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
+        paged_attn_decode_split_kv_kernel<<<
+            dim3(p.batch * p.kv_head, 1, p.num_splits),
+            dim3(32, group_sz), smem>>>(p);
+    }
+    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+}
+
+template <int HEAD_DIM>
+static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int seed) {
+    printf("B=%d Hq=%d Hkv=%d kv_len=%d page_sz=%d head_dim=%d ... ", B, Hq, Hkv, kv_len, page_size, HEAD_DIM);
+    fflush(stdout);
+
+    int max_pages = (kv_len + page_size - 1) / page_size;
+    int n_phys_pages = B * max_pages;
+
+    size_t sz_q  = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
+    size_t sz_o  = sz_q;
+    size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t);
+    int max_splits = 32;
+    size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
+    size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
+
+    bf16 *d_q, *d_o_paged, *d_o_ref;
+    bf16 *d_k_pool, *d_v_pool;
+    int64_t* d_pt;
+    float *d_op, *d_ml;
+
+    cudaMalloc(&d_q, sz_q);
+    cudaMalloc(&d_o_paged, sz_o);
+    cudaMalloc(&d_o_ref, sz_o);
+    cudaMalloc(&d_k_pool, sz_kv);
+    cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_pt, sz_pt);
+    cudaMalloc(&d_op, sz_op);
+    cudaMalloc(&d_ml, sz_ml);
+
+    srand(seed);
+    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
+
+    bf16* h_q = (bf16*)malloc(sz_q);
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++)
+        h_q[i] = __float2bfloat16(rnd());
+    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
+
+    bf16* h_k_pool = (bf16*)malloc(sz_kv);
+    bf16* h_v_pool = (bf16*)malloc(sz_kv);
+    size_t ps = (size_t)page_size * Hkv * HEAD_DIM;
+    for (int pg = 0; pg < n_phys_pages; pg++) {
+        for (int off = 0; off < page_size; off++) {
+            for (int h = 0; h < Hkv; h++) {
+                for (int d = 0; d < HEAD_DIM; d++) {
+                    float v = sinf((float)(pg * 7919 + off * 1049 + h * 331 + d));
+                    size_t idx = (size_t)pg * ps + (size_t)off * Hkv * HEAD_DIM + h * HEAD_DIM + d;
+                    h_k_pool[idx] = __float2bfloat16(v);
+                    h_v_pool[idx] = __float2bfloat16(v * 0.3f);
+                }
+            }
+        }
+    }
+    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
+
+    int64_t* h_pt = (int64_t*)malloc(sz_pt);
+    int next_pg = 0;
+    for (int b = 0; b < B; b++)
+        for (int p = 0; p < max_pages; p++)
+            h_pt[b * max_pages + p] = next_pg++;
+    cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice);
+
+    bf16* h_k_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16));
+    bf16* h_v_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16));
+    gather_kv_cpu(h_k_pool, h_v_pool, h_pt, B, Hkv, kv_len, page_size, HEAD_DIM, h_k_cont, h_v_cont);
+
+    float* h_q_f = (float*)malloc((size_t)B * Hq * HEAD_DIM * sizeof(float));
+    float* h_k_f = (float*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(float));
+    float* h_v_f = (float*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
+    for (int i = 0; i < B * kv_len * Hkv * HEAD_DIM; i++) {
+        h_k_f[i] = bf2f(h_k_cont[i]);
+        h_v_f[i] = bf2f(h_v_cont[i]);
+    }
+
+    float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
+    cpu_attention_ref(h_q_f, h_k_f, h_v_f, nullptr, h_o_ref, B, Hq, Hkv, 1, kv_len, HEAD_DIM, 0, 0);
+
+    float scale_val = 1.0f / sqrtf((float)HEAD_DIM);
+    PagedAttentionParams<bf16, float> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv; p.q_len = 1;
+    p.kv_len = kv_len; p.head_dim = HEAD_DIM;
+    p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0;
+    p.num_splits = 1; p.scale = scale_val;
+    p.page_size = page_size; p.max_pages = max_pages;
+    p.page_table = d_pt;
+    p.k_cache = d_k_pool; p.v_cache = d_v_pool;
+    p.q = d_q; p.mask = nullptr; p.o = d_o_paged;
+    p.o_part = d_op; p.ml_part = d_ml;
+
+    launch_paged_decode<HEAD_DIM>(p);
+    cudaDeviceSynchronize();
+
+    bf16* h_o_bf16 = (bf16*)malloc(sz_o);
+    cudaMemcpy(h_o_bf16, d_o_paged, sz_o, cudaMemcpyDeviceToHost);
+    float* h_o_paged = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++)
+        h_o_paged[i] = __bfloat162float(h_o_bf16[i]);
+
+    float max_err = 0.0f;
+    int bad_idx = -1;
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
+        float e = fabsf(h_o_paged[i] - h_o_ref[i]);
+        if (e > max_err) { max_err = e; bad_idx = i; }
+    }
+
+    bool pass = max_err < 0.02f;
+
+    if (pass) {
+        printf("PASS (max_abs_err=%.4e)\n", max_err);
+    } else {
+        int b = bad_idx / (Hq * HEAD_DIM);
+        int h = (bad_idx / HEAD_DIM) % Hq;
+        int d = bad_idx % HEAD_DIM;
+        printf("FAIL (max_abs_err=%.4e at [%d,%d,%d]: ref=%.4f got=%.4f)\n",
+               max_err, b, h, d, h_o_ref[bad_idx], h_o_paged[bad_idx]);
+        printf("  ref[0..7]:");
+        for (int i = 0; i < 8 && i < HEAD_DIM; i++)
+            printf(" %.4f", h_o_ref[i]);
+        printf("\n  got[0..7]:");
+        for (int i = 0; i < 8 && i < HEAD_DIM; i++)
+            printf(" %.4f", h_o_paged[i]);
+        printf("\n");
+    }
+
+    free(h_q); free(h_k_pool); free(h_v_pool); free(h_pt);
+    free(h_k_cont); free(h_v_cont);
+    free(h_q_f); free(h_k_f); free(h_v_f);
+    free(h_o_ref); free(h_o_bf16); free(h_o_paged);
+    cudaFree(d_q); cudaFree(d_o_paged); cudaFree(d_o_ref);
+    cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
+    cudaFree(d_op); cudaFree(d_ml);
+
+    return pass ? 0 : 1;
+}
+
+struct TestCase {
+    int head_dim;
+    int B, Hq, Hkv, kv_len, page_size, seed;
+};
+
+static const TestCase TESTS[] = {
+    {128, 1, 1, 1, 8, 128, 1},
+    {128, 1, 4, 4, 128, 128, 2},
+    {128, 2, 4, 4, 256, 128, 3},
+    {128, 1, 4, 1, 64, 64, 4},
+    {128, 1, 8, 2, 64, 128, 5},
+    {128, 2, 16, 4, 128, 128, 6},
+    {64, 1, 4, 2, 32, 128, 7},
+    {256, 1, 2, 1, 16, 128, 8},
+    {32, 1, 4, 2, 32, 64, 9},
+    {128, 3, 8, 2, 256, 128, 10},
+    {128, 2, 32, 8, 512, 128, 11},
+#ifndef ASTRAI_NO_MMA
+    {128, 1, 16, 2, 256, 128, 12},
+    {128, 2, 32, 4, 512, 128, 13},
+#endif
+};
+
+static int dispatch_test(const TestCase& tc) {
+    switch (tc.head_dim) {
+        case 32:  return run_test<32>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
+        case 64:  return run_test<64>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
+        case 128: return run_test<128>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
+        case 256: return run_test<256>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
+        default:  return 1;
+    }
+}
+
+// Warmed-up, CUDA-event timed sweep over paged decode configs.
+// Reports per-call latency and effective K/V read bandwidth.
+// Bytes = K + V read through page table (B*Hk*kv*D each), bf16.
+template <int HEAD_DIM>
+static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) {
+    int max_pages = (kv_len + page_size - 1) / page_size;
+    int n_phys_pages = B * max_pages;
+
+    size_t sz_q  = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t);
+    int max_splits = 32;
+    size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
+    size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int64_t* d_pt;
+    float *d_op, *d_ml;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_pt, sz_pt);
+    cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
+
+    bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
+    cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
+    cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
+
+    int64_t* h_pt = (int64_t*)malloc(sz_pt);
+    int next_pg = 0;
+    for (int b = 0; b < B; b++)
+        for (int p = 0; p < max_pages; p++)
+            h_pt[b * max_pages + p] = next_pg++;
+    cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice);
+    free(h_pt);
+
+    float scale_val = 1.0f / sqrtf((float)HEAD_DIM);
+    PagedAttentionParams<bf16, float> pa;
+    pa.batch = B; pa.q_head = Hq; pa.kv_head = Hkv; pa.q_len = 1;
+    pa.kv_len = kv_len; pa.head_dim = HEAD_DIM;
+    pa.use_mask = 0; pa.is_causal = 0; pa.causal_offset = 0;
+    pa.num_splits = 1; pa.scale = scale_val;
+    pa.page_size = page_size; pa.max_pages = max_pages;
+    pa.page_table = d_pt;
+    pa.k_cache = d_k_pool; pa.v_cache = d_v_pool;
+    pa.q = d_q; pa.mask = nullptr; pa.o = d_o;
+    pa.o_part = d_op; pa.ml_part = d_ml;
+
+    const int WARMUP = 10, ITERS = 100;
+    for (int i = 0; i < WARMUP; i++) launch_paged_decode<HEAD_DIM>(pa);
+    cudaDeviceSynchronize();
+    CUDA_CHECK(cudaGetLastError());
+
+    cudaEvent_t s, e;
+    cudaEventCreate(&s); cudaEventCreate(&e);
+    cudaEventRecord(s);
+    for (int i = 0; i < ITERS; i++) launch_paged_decode<HEAD_DIM>(pa);
+    cudaEventRecord(e); cudaEventSynchronize(e);
+    float ms = 0; cudaEventElapsedTime(&ms, s, e); ms /= ITERS;
+
+    double flops = 4.0 * B * Hq * (double)kv_len * HEAD_DIM;
+    double tflops = flops / (ms * 1e-3) / 1e12;
+    size_t nKV = (size_t)B * Hkv * kv_len * HEAD_DIM;
+    double bytes = 2.0 * (2.0 * nKV);
+    double gbps = bytes / (ms * 1e-3) / 1e9;
+
+    char cfg[64];
+    snprintf(cfg, sizeof(cfg),
+             "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d page=%3d",
+             B, Hq, Hkv, 1, kv_len, HEAD_DIM, page_size);
+    printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
+           cfg, ms, gbps, tflops);
+
+    free(tmp);
+    cudaFree(d_q); cudaFree(d_o);
+    cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
+    cudaFree(d_op); cudaFree(d_ml);
+    cudaEventDestroy(s); cudaEventDestroy(e);
+}
+
+static void bench() {
+    printf("\n===== PAGED DECODE BENCH =====\n");
+    printf("%-46s | %10s | %10s | %10s\n",
+           "config", "latency", "bandwidth", "throughput");
+    printf("---------------------------------------------------------------"
+           "----------------------------\n");
+    bench_config<128>(1, 32, 4, 512, 128);
+    bench_config<128>(1, 32, 4, 1024, 128);
+    bench_config<128>(1, 32, 4, 2048, 128);
+    bench_config<128>(1, 32, 4, 4096, 128);
+    bench_config<128>(16, 32, 4, 2048, 128);
+    bench_config<128>(32, 32, 4, 1024, 128);
+}
+
+int main() {
+    int n = sizeof(TESTS) / sizeof(TESTS[0]);
+    int fail = 0;
+    printf("=== Paged Decode vs CPU reference (%d cases) ===\n\n", n);
+
+    for (int i = 0; i < n; i++) {
+        fail += dispatch_test(TESTS[i]);
+        if (fail) break;
+    }
+
+    if (fail) {
+        printf("\nFAILED (%d/%d tests failed)\n", fail, n);
+        return fail;
+    }
+    printf("\nAll %d tests passed!\n", n);
+    bench();
+    return 0;
+}
--- a/csrc/tests/attn_prefill_test.cu
+++ b/csrc/tests/attn_prefill_test.cu
@ -0,0 +1,178 @@
+/*
+Pure-C test:
+nvcc -I csrc -arch=sm_89 -O3 \
+    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
+    csrc/tests/attn_prefill_test.cu -o test && ./test
+*/
+
+#include "test_utils.cuh"
+#include "../kernels/attn_prefill_split_q.cuh"
+#ifndef ASTRAI_NO_MMA
+#include "../kernels/attn_prefill_split_q_mma.cuh"
+#endif
+
+// Launch the production prefill path (tensor-core MMA on sm_80+, else the
+// scalar fallback), mirroring dispatch_prefill() in attn_prefill.cu.
+template <int HEAD_DIM>
+static void launch_prefill(AttentionParams<bf16>& p) {
+#ifndef ASTRAI_NO_MMA
+    constexpr int WARPS = 4, BR = 16;
+    constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
+    constexpr int MIN_BLOCKS = (HEAD_DIM <= 32) ? 6 : (HEAD_DIM <= 64) ? 4
+                             : (HEAD_DIM <= 128) ? 3 : 2;
+    dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
+    dim3 block(WARPS * 32, 1, 1);
+    attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(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);
+    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
+#endif
+}
+
+static void dispatch_prefill(AttentionParams<bf16>& p) {
+    switch (p.head_dim) {
+        case 64:  launch_prefill<64>(p);  break;
+        case 128: launch_prefill<128>(p); break;
+        default:  printf("bench: unsupported D=%d\n", p.head_dim);
+    }
+}
+
+// Warmed-up, CUDA-event timed throughput sweep over the production MMA path.
+// Reports per-call latency and effective tensor-core TFLOP/s (2 matmuls:
+// QK^T and P@V, each 2*B*Hq*ql*kl*D flops; halved for causal).
+static void bench() {
+    const int cfgs[][7] = {
+        {1,32,4,512,512,128,0},
+        {1,32,4,1024,1024,128,0},
+        {1,32,4,2048,2048,128,0},
+        {1,32,4,2048,2048,128,1},
+        {4,32,4,2048,2048,128,1},
+        {1,32,4,4096,4096,128,1},
+    };
+    int n = sizeof(cfgs)/sizeof(cfgs[0]);
+    const int WARMUP = 10, ITERS = 50;
+    printf("\n===== PREFILL BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
+    printf("%-46s | %10s | %10s | %10s\n",
+           "config", "latency", "bandwidth", "throughput");
+    printf("---------------------------------------------------------------"
+           "----------------------------\n");
+
+    for (int ci = 0; ci < n; ci++) {
+        int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
+        int ql=cfgs[ci][3], kl=cfgs[ci][4], D=cfgs[ci][5], causal=cfgs[ci][6];
+        size_t nQ=(size_t)B*Hq*ql*D, nKV=(size_t)B*Hk*kl*D;
+
+        bf16 *dQ,*dK,*dV,*dO,*tmp;
+        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+        size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
+        for (size_t i=0;i<nQ;i++)  tmp[i]=f2bf(randf());
+        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+
+        AttentionParams<bf16> p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
+        p.use_mask=0; p.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;
+
+        for (int i=0;i<WARMUP;i++) dispatch_prefill(p);
+        cudaDeviceSynchronize();
+        cudaError_t err=cudaGetLastError();
+        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
+
+        cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
+        cudaEventRecord(s);
+        for (int i=0;i<ITERS;i++) dispatch_prefill(p);
+        cudaEventRecord(e); cudaEventSynchronize(e);
+        float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
+
+        double flops = 4.0*B*Hq*(double)ql*kl*D;
+        if (causal) flops *= 0.5;
+        double tflops = flops/(ms*1e-3)/1e12;
+        // HBM traffic: Q + O (B*Hq*ql*D each) + K + V (B*Hk*kl*D each), bf16.
+        double bytes = 2.0 * (2.0*nQ + 2.0*nKV);
+        double gbps = bytes/(ms*1e-3)/1e9;
+
+        char cfg[64];
+        snprintf(cfg, sizeof(cfg),
+                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
+                 B,Hq,Hk,ql,kl,D,causal);
+        printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
+               cfg, ms, gbps, tflops);
+
+        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
+        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
+    }
+}
+
+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);
+
+        AttentionParams<bf16> p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
+        p.use_mask=0; p.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;
+
+        double t0=now_ms();
+        dispatch_prefill(p);
+        cudaDeviceSynchronize();
+        double kms=now_ms()-t0;
+        cudaError_t err=cudaGetLastError();
+        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
+
+        bf16* hOut=new bf16[nQ];
+        cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
+
+        float* ref=new float[nQ];
+        cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal, 0);
+
+        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");
+    bench();
+    return 0;
+}
--- a/csrc/tests/test_utils.cuh
+++ b/csrc/tests/test_utils.cuh
@ -0,0 +1,90 @@
+#pragma once
+
+#include <cstdio>
+#include <cstdlib>
+#include <cmath>
+#include <chrono>
+#include <cuda_bf16.h>
+
+using bf16 = __nv_bfloat16;
+
+inline bf16 f2bf(float x) { return __float2bfloat16(x); }
+inline float bf2f(bf16 x) { return __bfloat162float(x); }
+
+inline float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
+
+inline double now_ms() {
+    using namespace std::chrono;
+    return duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count();
+}
+
+inline int compute_num_splits(int base_blocks, int tiles_total) {
+    int sm_count = 0;
+    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
+    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
+    if (n > tiles_total) n = tiles_total;
+    if (n > 32) n = 32;
+    if (n < 1) n = 1;
+    return n;
+}
+
+#define CUDA_CHECK(call) \
+    do { \
+        cudaError_t _e = (call); \
+        if (_e != cudaSuccess) { \
+            printf("CUDA error %s at %s:%d\n", cudaGetErrorString(_e), __FILE__, __LINE__); \
+            exit(1); \
+        } \
+    } while (0)
+
+// Generic CPU reference for multi-query / grouped-query attention.
+// Tensor shapes (all float*):
+//   Q : [B, Hq, q_len, D]
+//   K : [B, Hk, kv_len, D]
+//   V : [B, Hk, kv_len, D]
+//   O : [B, Hq, q_len, D]
+// mask: if q_len == 1, shape is [B, kv_len]; otherwise mask is not supported.
+static void cpu_attention_ref(
+    const float* Q, const float* K, const float* V, const bool* mask,
+    float* O, int B, int Hq, int Hk, int q_len, int kv_len, int D,
+    int is_causal, int causal_offset
+) {
+    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;
+            for (int qi = 0; qi < q_len; qi++) {
+                float mv = -INFINITY, sv = 0.0f;
+                float accum[256] = {0.0f};
+                int lim = kv_len;
+                if (is_causal) {
+                    int c = qi + causal_offset + 1;
+                    lim = (c < kv_len) ? c : kv_len;
+                }
+                for (int kj = 0; kj < lim; kj++) {
+                    if (mask != nullptr && q_len == 1) {
+                        if (!mask[b * kv_len + kj]) continue;
+                    }
+                    float dot = 0.0f;
+                    size_t q_idx = ((size_t)b * Hq + h) * q_len + qi;
+                    size_t kv_idx = ((size_t)b * Hk + kv_h) * kv_len + kj;
+                    for (int d = 0; d < D; d++)
+                        dot += Q[q_idx * D + d] * K[kv_idx * D + d];
+                    dot *= scale;
+                    float nm = fmaxf(mv, dot);
+                    float a = expf(mv - nm);
+                    float b_exp = expf(dot - nm);
+                    sv = sv * a + b_exp;
+                    for (int d = 0; d < D; d++)
+                        accum[d] = accum[d] * a + V[kv_idx * D + d] * b_exp;
+                    mv = nm;
+                }
+                float inv = 1.0f / sv;
+                size_t o_idx = ((size_t)b * Hq + h) * q_len + qi;
+                for (int d = 0; d < D; d++)
+                    O[o_idx * D + d] = accum[d] * inv;
+            }
+        }
+    }
+}
--- a/scripts/tools/benchmark.py
+++ b/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)
--- a/scripts/tools/train.py
+++ b/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(
--- a/setup.py
+++ b/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)
--- a/tests/data/conftest.py
+++ b/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")
--- a/tests/data/test_preprocess_builder.py
+++ b/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
Author	SHA1	Message	Date
ViperEkura	2c3cef1c87	feat: wire up paged decode CUDA kernel to Python extension - Add attn_paged_decode wrapper in ops.py with gather fallback - Register kernel in loader.py and export from __init__.py - Extract test_utils.cuh shared by all attention unit tests - Rename attn_paged_vs_contiguous.cu to attn_paged_decode_test.cu - Refactor decode/prefill tests to use common bf16 helpers and cpu ref - Fix k_cache dim check in attn_paged_decode.cu	2026-07-11 18:40:49 +08:00
ViperEkura	89ece26c25	feat: paged decode attention with split-KV (scalar + MMA) - PagedAttentionParams merged into attn_common.h - Scalar variant: warp-per-query-head split-KV, resolves page table per-position for the K/V shared-memory tile load - MMA variant (sm_80+): tensor-core head-packing with cp.async, single page-table lookup per tile (BC=32 fits within page_size>=32) - Standalone test: 14 cases across head_dim 32/64/128/256, GQA, multi-batch, both paths verified against CPU reference	2026-07-11 18:14:51 +08:00
ViperEkura	2c0b5d0b5e	perf: enable MMA decode path for G=1 full attention - inline decode_use_mma() into dispatch_decode() - drop G>1 guard, MMA works correctly for G>=1 - decode is memory-bound, tensor cores + cp.async still win at G=1	2026-07-11 11:45:13 +08:00
ViperEkura	a4ae7d17fb	perf: increase decode split-K parallelism for short sequences - Remove tiles_total/8 min-work cap that limited splits for small workloads - Simplify decode_num_splits to only use base_blocks and tiles_total - Short sequences now generate more blocks, improving SM utilization	2026-07-11 11:26:42 +08:00
ViperEkura	8a8550184f	refactor: template AttentionParams, rename .cuh to .h - Convert AttentionParams to a template struct supporting arbitrary types - Rename attn_common.cuh -> attn_common.h (no CUDA-specific code remains) - Include standard headers explicitly in each .cuh instead of via attn_common.cuh - Allow .h files in csrc/ via .gitignore	2026-07-11 11:03:14 +08:00
ViperEkura	b8b439b713	perf: fuse decode combine kernel to single-pass online-rescale reduction - Replace 3-scan loops (mstar, lstar, acc) with 1-pass online rescale - Halves __expf calls (num_splits vs 2*num_splits) and ml_part re-reads - Mathematically equivalent, no change to o_part traffic or output	2026-07-11 00:24:09 +08:00
ViperEkura	41cd40363a	perf: post-multiply attention scale in float instead of pre-scaling Q in bf16 - Replace bf16 pre-scale Q loading with direct 32-bit aligned bf16x2 reads - Apply scale in float32 after Q@K^T, before online softmax - Reduces causal max error from 2^-6 to 2^-8 with zero perf cost	2026-07-11 00:13:32 +08:00
ViperEkura	d923ebe38d	refactor: rename gqa_* to attn_, split-KV for all decode paths - Rename all csrc/kernels/gqa_.cuh/cu to attn_*, with _split_q / _split_kv strategy suffix and optional _mma compute suffix - Remove non-split MMA decode kernel, keep only split-KV path - Convert scalar decode fallback to split-KV (o_part/ml_part + combine) - Move combine kernel to attn_decode_split_kv.cuh (shared by both paths) - Rename GQAParams to AttentionParams - Update all C++ #include, PYBIND11, and Python extension references	2026-07-10 23:35:14 +08:00
ViperEkura	29b0423c4e	refactor: deduplicate kernel code with shared MMA and entry-point helpers	2026-07-10 20:52:27 +08:00
ViperEkura	88f8dca2c2	perf: enlarge prefill KV tile to BC=32 for D<=128 - Kernel is latency-bound (25% occupancy), not compute/bandwidth-bound - BC=16 wasted a cp.async wait + barrier + loop overhead per tiny tile - Double KV tile to BC=32 for D<=128; D=256 stays 16 (64KB > 48KB smem cap) - Retune MIN_BLOCKS per head_dim (32->6, 64->4, 128->3, 256->2) - Result: ~6-8% faster on L20, 0.93-1.20x vs torch SDPA, correctness unchanged	2026-07-10 17:43:23 +08:00
ViperEkura	9027fdc546	test: bench production MMA attention path with FLOP/s and bandwidth	2026-07-10 16:58:34 +08:00
ViperEkura	cbd140340d	perf: load prefill Q fragments directly from global (drop sQ staging) - read the 8 Q elements each lane needs straight from global into the mma A-operand layout, pre-scaled, instead of staging through shared sQ - removes the sQ smem area (20KB->16KB) and the serialized per-warp prologue with its WARPS __syncthreads barriers - result vs torch SDPA: prefill 0.70-0.82x -> 0.85-1.00x (matches torch at seq=128; 2048 1.251->1.159ms), correctness unchanged across head dims	2026-07-10 12:27:45 +08:00
ViperEkura	988e01314d	perf: pipeline prefill MMA kernel (double-buffered K/V + packed stores) - double-buffer K/V one tile ahead via cp.async to overlap load with tensor-core math (ncu long_scoreboard 2.12->0.53) - reorder wait->barrier->prefetch so one __syncthreads/tile covers both cross-warp publish and buffer-reuse (was two) - add predicated cp_async_16_pred (src-size=0 zero-fills OOB) to unify full/partial tiles, dropping the scalar fallback - halve BC to 16 to keep 3 blocks/SM despite the doubled smem - pack adjacent bf16 output into one 32-bit STG, removing the uncoalesced scalar-store penalty (14%->5% sectors) - result vs torch SDPA: prefill 0.61-0.78x -> 0.70-0.82x, spills eliminated	2026-07-10 12:19:14 +08:00
ViperEkura	7ba43a7c6f	perf: add split-K (FlashDecoding) to decode MMA kernel Decode has only batchkv_head independent tasks, so the grid was tiny (e.g. 16 blocks) leaving most SMs idle (ncu: 0.04 waves/SM, 11% DRAM). - Partition KV across gridDim.z blocks emitting unnormalised (O, m, l) partials, reduced by a new combine kernel - Choose split count to fill the device (~2 blocks/SM), capped by tile count and 32; fall back to single-pass direct-write when batchkv_head already saturates the SMs - Refactor decode dispatch into named helpers, de-duplicate scalar fallback Result: now DRAM-bound at 63% (99->543 GB/s), 2.1-2.5x over torch SDPA in the low-parallelism regime, on par at high parallelism	2026-07-10 11:43:18 +08:00
ViperEkura	dea59f7e1d	fix: restore resident Qa to fix sQ overwrite bug Moving Qa ldmatrix into the tile loop caused warps 0-2 to read warp 3's Q data from sQ (only the last warp's data survives the serialized load loop). Reverted to loading Qa during the init phase and keeping it resident; __launch_bounds__ still forces 128 regs (33% occupancy) with spill to local memory.	2026-07-10 00:49:58 +08:00
ViperEkura	85dc771460	perf: reduce MMA kernel registers, switch to static smem - Move Qa[KD][4] into tile loop (reload from sQ per tile) cutting ~32 resident registers for HEAD_DIM=128 - Replace extern __shared__ with static template-sized smem (no cudaFuncSetAttribute or dynamic allocation needed) - Add __launch_bounds__ with MIN_BLOCKS param, dispatch by HEAD_DIM (hd=128→4, hd=64→6, hd=32→6) - Remove dynamic smem from scalar kernel and C test - Result: hd=128 168→128 regs, 25%→33% occupancy	2026-07-10 00:39:47 +08:00
ViperEkura	2c5629b81d	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	2026-07-09 10:09:53 +08:00
ViperEkura	841a582b28	refactor: split mask builder by single/multi output - Extract SingleOutputMaskBuilder for SFT and pretrain configs - Extract MultiOutputMaskBuilder for DPO and GRPO configs - Keep SectionedMaskBuilder as backward-compatible facade - Register "single" and "multi" names in MaskBuilderFactory - Add parity and rejection tests for concrete builders	2026-07-08 21:18:34 +08:00
ViperEkura	c8567a6f65	fix: exclude embedding, lm_head, bias, and norm params from Muon optimizer, use AdamW	2026-07-08 19:42:20 +08:00
ViperEkura	8035be9b1f	fix: make MuonMix inherit from torch.optim.Optimizer	2026-07-08 17:03:05 +08:00
ViperEkura	e9b03f4fca	perf: apply cp.async, XOR swizzle, pre-scaled Q to decode MMA kernel Decode MMA kernel previously used scalar global→shared loads with LD=HEAD_DIM+8 padding and per-tile scale multiply. This commit brings it in line with the prefill MMA kernel (which already had these optimizations): - cp.async K/V loads (bypasses registers, halves load instructions) - XOR swizzle: LD=HEAD_DIM instead of HEAD_DIM+8 (zero waste smem) - Pre-scale Q during load (removes per-tile scale multiply in softmax) - Clean up prefill MMA kernel comments (no code change) ~2x speedup on decode (0.47ms→0.24ms at seq_len=512)	2026-07-08 16:15:14 +08:00
ViperEkura	fd65b9bc23	feat: support HEAD_DIM=32 and split extension into loader/ops - add case 32 to decode/prefill dispatch switch - fix swiz_col out-of-bounds for HEAD_DIM=32: XOR mask now limited to chunk count (3 for 32, 7 for >=64) instead of always 7, which produced column offsets >= LD=32 and corrupted shared memory - restructure decode dispatch to #ifndef/#else/#endif matching prefill - split astrai/extension/__init__.py into loader.py (kernel .so discovery) and ops.py (wrapper functions + torch SDPA fallback); __init__.py now re-exports the public API	2026-07-08 14:14:11 +08:00
ViperEkura	9ebaea840f	perf: cp.async K/V loads, shared sQ staging, causal skip, XOR swizzle - cp.async global→shared for K/V full-tile loads, eliminates 99.6% of shared-store bank conflicts (612K→2.7K per ncu) - add cp_async_16/commit/wait_all/wait_group<N> helpers in mma utils - shared sQ staging (single area, serialized per-warp load), cuts smem from (2BC + WARPSBR)LD to (2BC + BR)*LD bf16 - pre-scale Q by attention scale during Q load, removes per-tile scale multiply in softmax loop - causal tile skipping: block-level early break + warp-level skip - scalar fallback only for last partial tile - XOR swizzle (swiz_col) at 8-bf16 chunk granularity, eliminates ldmatrix bank conflicts without LD padding, LD=HEAD_DIM (zero smem waste), saves 1280 bytes/block vs HEAD_DIM+8 padding	2026-07-08 12:30:32 +08:00
ViperEkura	6adc221c10	refactor: extract shared MMA utils into gqa_mma_utils.cuh - Move mma16816, ld2, pk2, pkb, ldmatrix_x4/x2/x2_trans to shared header - gqa_prefill_attn_mma.cuh and gqa_decode_attn_mma.cuh both include it	2026-07-07 23:01:15 +08:00
ViperEkura	9e63cb9ed0	feat: MMA head-packing decode kernel with scalar fallback dispatch - Add gqa_decode_attn_mma.cuh for tensor-core decode path - Add dispatch_decode<> selecting MMA vs scalar based on G and mask - Add TORCH_CHECK for unsupported head_dim instead of silent scalar launch	2026-07-07 22:56:02 +08:00
ViperEkura	4225518cf3	perf: add fast-math and vectorization nvcc/cxx build flags - centralize CXX_FLAGS/NVCC_FLAGS in csrc/build.py as single source - add --use_fast_math, --ptxas-options=-O3,-v, --extra-device-vectorization - add -march=native -funroll-loops host flags - setup.py reads shared cxx_flags/nvcc_flags from registry - sync pure-C test build commands with new flags	2026-07-07 22:28:32 +08:00
ViperEkura	c50adbaac0	feat : replace AdamW with MuonMix (Muon + AdamW) optimizer - Muon for 2D matrix params, AdamW for 1D (norm/bias/embed) - MuonMix wrapper handles combined step/zero_grad/state_dict - New CLI args: weight_decay, muon_momentum, muon_nesterov, muon_ns_steps, muon_adjust_lr - Removed adamw_beta1/adamw_beta2/adamw_weight_decay - Moved optimizer/strategy params from signature to **kwargs	2026-07-07 14:10:36 +08:00
ViperEkura	536dbc0c9a	fix: set tqdm postfix before update so first step shows metrics	2026-07-07 00:14:13 +08:00
ViperEkura	4af7acd449	fix: support single .h5 file loading in load_h5	2026-07-07 00:11:16 +08:00
ViperEkura	53ed52b4b8	refactor: extension dispatch layer with CUDA/torch fallback - Add gqa_decode_attn/gqa_prefill_attn dispatch functions - Internal _available/__modules with underscore prefix - CUDA kernel path with F.scaled_dot_product_attention fallback - GQA head expansion in fallback path	2026-07-06 21:07:16 +08:00
ViperEkura	f1cc7cedce	feat: ldmatrix + smem padding for mma prefill kernel - replace scalar fragment loads with ldmatrix.sync.x4/x2 - add smem row-stride padding (LD = HEAD_DIM + 8) to eliminate 8-way bank conflicts from HEAD_DIM being a 32-bank multiple - switch build flag from positive to negative: -DASTRAI_NO_MMA for pre-sm_80 only; mma is the default path - vectorize scalar path smem loads with float4 ld8 - fix pure-C test configs for ld8 alignment	2026-07-06 20:55:22 +08:00
ViperEkura	ddc4bd1cf6	feat: tensor-core mma prefill with build-time dispatch - add register-resident flash-attention kernel using mma.sync.m16n8k16 - dispatch mma vs scalar at build time: pre-sm_80 defines -DASTRAI_NO_MMA, else defaults to mma - scalar path vectorized with float4 smem loads (ld8)	2026-07-06 20:33:24 +08:00
ViperEkura	cc36530c73	perf: group-split register-blocking gqa_prefill kernel - one query row per group of G=8 lanes, each owning HEAD_DIM/G dims of qreg[]/acc[] in registers - removes full 32-lane warp_reduce_sum; S dot reduces over only G lanes - templated on <HEAD_DIM,G,ROWS,P_BC>, block=(G,ROWS)=(8,32) - per-group shuffle mask so causal loop-bound divergence doesn't deadlock the shuffle - update pure-C test to the templated launch	2026-07-06 18:33:08 +08:00
ViperEkura	11fa807cfc	fix: correct prefill mask index, unify GQA kernel interface - Fix mask indexing: batchq_lenkv_len -> batch*kv_len - Add csrc/kernels/gqa_common.cuh with shared GQAParams struct - Unify decode/prefill Python API: both accept (q,k,v,mask=None,...) - Decode now supports optional mask, is_causal, causal_offset, scale - Rename struct fields: B->batch, Hq->q_head, Hk->kv_head, D->head_dim - Use py::arg() for correct None/defaults handling in pybind11 - Update pure C tests and build instructions (-arch=sm_89)	2026-07-06 17:21:23 +08:00
ViperEkura	bcdd93e0eb	feat: split kernel defs from bindings, add prefill tiled kernel and pure C tests - Split .cuh/.cu for gqa_decode_attn and gqa_prefill_attn - gqa_prefill_attn: tiled shared-memory K/V, fused load, compute-opt, mask support - Add pure C tests under csrc/tests/ for fast nvcc-only iteration - Update .gitignore for build artifacts	2026-07-06 16:14:55 +08:00
ViperEkura	579b8c3129	fix: correct gqa_decode_attn reduction + add gqa_prefill_attn - gqa_decode_attn: rewrite to per-KV-head, K in smem - gqa_prefill_attn: new kernel for Q_len > 1 with GQA	2026-07-06 13:45:18 +08:00
ViperEkura	d7da51569f	docs: update install instructions in EN/CN README	2026-07-06 12:25:36 +08:00
ViperEkura	e8e228d035	feat: add optional CUDA kernel system (csrc/) + fused GQA decode attention Structure: csrc/ -- .cu sources + build.py registry astrai/extension/ -- compiled .so + __init__.py (import dispatcher) setup.py -- CUDAExtension from csrc/build.py REGISTRY Control: CSRC_KERNELS=true\|false env var at install time. Fallback: astrai.extension.available dict for runtime detection.	2026-07-06 12:09:58 +08:00
ViperEkura	2579658e15	chore : shields release badge from /release to /tag	2026-07-05 20:34:40 +08:00
ViperEkura	f0cd0134c6	fix : update benchmark for v1.3.8 cache API, add argparse and cache type switch - Replaced old KVCage API with PageCache/ContiguousCache - Added --cache contiguous\|paged switch for decoding comparison - Added argparse for all params (batch/prompt/gen/device/dtype) - Fixed PageCache decode crash by extending pages for full sequence	2026-07-05 20:30:26 +08:00
ViperEkura	abb96996f8	docs : sync 6 doc files to actual code - architecture.md: removed TrainConfig.log_interval, split KVCache into PageCache/ContiguousCache with CacheView/PageCacheView/ContiguousCacheView, added JsonlStore, fixed GradientCheckpointingCallback type, CheckpointCallback typo, ProgressBarCallback hooks - training.md: added position_ids to SFT keys, fixed callback hook table, removed merged ValidationCallback - inference.md: documented ContiguousCache default vs PageCache paged - dataflow.md: added JsonlStore to storage backends and format detection - params.md: removed nonexistent --log_interval - preprocessing.md: updated timestamp	2026-07-05 19:35:18 +08:00