Compare commits
No commits in common. "2c5629b81ded391fc36b1dd9dfc865d000cbf656" and "53ed52b4b8281aa4bb3483b5d4ec17c9cdbdb39e" have entirely different histories.
2c5629b81d
...
53ed52b4b8
|
|
@ -103,7 +103,9 @@ nohup python scripts/tools/train.py \
|
|||
--warmup_ratio=0.05 \
|
||||
--max_lr=1e-4 \
|
||||
--max_grad_norm=1.0 \
|
||||
--weight_decay=0.1 \
|
||||
--adamw_beta1=0.9 \
|
||||
--adamw_beta2=0.95 \
|
||||
--adamw_weight_decay=0.01 \
|
||||
--window_size=2048 \
|
||||
--ckpt_interval=10000 \
|
||||
--ckpt_dir=./checkpoint \
|
||||
|
|
|
|||
|
|
@ -109,7 +109,9 @@ nohup python scripts/tools/train.py \
|
|||
--warmup_ratio=0.05 \
|
||||
--max_lr=1e-4 \
|
||||
--max_grad_norm=1.0 \
|
||||
--weight_decay=0.1 \
|
||||
--adamw_beta1=0.9 \
|
||||
--adamw_beta2=0.95 \
|
||||
--adamw_weight_decay=0.01 \
|
||||
--window_size=2048 \
|
||||
--ckpt_interval=10000 \
|
||||
--ckpt_dir=./checkpoint \
|
||||
|
|
|
|||
|
|
@ -63,6 +63,7 @@ classDiagram
|
|||
+Optional[int] n_heads
|
||||
+Optional[int] n_kv_heads
|
||||
+Optional[bool] use_qk_norm
|
||||
+Optional[bool] use_gated_attention
|
||||
+str ffn_type
|
||||
+Optional[dict] rope_scaling
|
||||
+Optional[str] pooling_type
|
||||
|
|
@ -755,6 +756,7 @@ classDiagram
|
|||
+int output_tokens
|
||||
+float arrival_time
|
||||
+Optional[float] finish_time
|
||||
+Optional[Callable] stream_callback
|
||||
+int next_pos
|
||||
+is_finished(stop_ids) bool
|
||||
}
|
||||
|
|
@ -1247,4 +1249,4 @@ classDiagram
|
|||
10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
|
||||
11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers
|
||||
|
||||
> Document Update Time: 2026-07-09
|
||||
> Document Update Time: 2026-07-05
|
||||
|
|
|
|||
|
|
@ -61,7 +61,7 @@ StoreFactory.create("bin") → MmapStore
|
|||
StoreFactory.create("jsonl") → JsonlStore
|
||||
```
|
||||
|
||||
**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage.
|
||||
**H5Store**: Reads HDF5 files, supports `share_memory_()` for multi-process DataLoader workers (copies tensors to shared memory).
|
||||
|
||||
**MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`.
|
||||
|
||||
|
|
@ -109,4 +109,4 @@ DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_typ
|
|||
|
||||
Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.
|
||||
|
||||
> Document Update Time: 2026-07-09
|
||||
> Document Update Time: 2026-07-05
|
||||
|
|
|
|||
|
|
@ -56,7 +56,7 @@ PageCache (paged KV cache with prefix sharing, alternative)
|
|||
1. Cleanup → Remove finished tasks, free KV cache slots/pages
|
||||
2. Refill → Pop from waiting_queue, task_alloc resources, activate
|
||||
3. Prefill → Group by (prompt_len, start_pos), run full forward
|
||||
4. Decode → Run single-token forward for each same-position group
|
||||
4. Decode → Pick largest same-position group, single-token forward
|
||||
```
|
||||
|
||||
## Sampling (Strategy Pattern)
|
||||
|
|
@ -249,4 +249,4 @@ async for token in engine.generate_async("Hello", ...): # -> AsyncGenerator[s
|
|||
print(token)
|
||||
```
|
||||
|
||||
> Document Update Time: 2026-07-09
|
||||
> Document Update Time: 2026-07-05
|
||||
|
|
|
|||
|
|
@ -28,17 +28,13 @@
|
|||
| `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
|
||||
| `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
|
||||
|
||||
### Optimizer (MuonMix)
|
||||
|
||||
Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`fused=True`).
|
||||
### Optimizer (AdamW)
|
||||
|
||||
| Parameter | Description | Default |
|
||||
|-----------|-------------|---------|
|
||||
| `--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` |
|
||||
| `--adamw_beta1` | AdamW beta1 | 0.9 |
|
||||
| `--adamw_beta2` | AdamW beta2 | 0.95 |
|
||||
| `--adamw_weight_decay` | AdamW weight decay | 0.01 |
|
||||
|
||||
### Data Loading
|
||||
|
||||
|
|
@ -108,7 +104,7 @@ Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`f
|
|||
| Parameter | Description | Default |
|
||||
|-----------|-------------|---------|
|
||||
| `--schedule_type` | LR scheduler type (`cosine`, `sgdr`, `wsd`) | cosine |
|
||||
| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default: 0.01) |
|
||||
| `--min_rate` | Minimum LR as fraction of base LR | None (scheduler default) |
|
||||
| `--cycle_length` | SGDR first cycle length in steps | None (total_steps - warmup_steps) |
|
||||
| `--t_mult` | SGDR cycle length multiplier per restart | 2 |
|
||||
| `--stable_steps` | WSD stable plateau steps | None (required for wsd) |
|
||||
|
|
@ -130,7 +126,9 @@ nohup python scripts/tools/train.py \
|
|||
--warmup_ratio=0.05 \
|
||||
--max_lr=1e-4 \
|
||||
--max_grad_norm=1.0 \
|
||||
--weight_decay=0.1 \
|
||||
--adamw_beta1=0.9 \
|
||||
--adamw_beta2=0.95 \
|
||||
--adamw_weight_decay=0.01 \
|
||||
--window_size=2048 \
|
||||
--ckpt_interval=10000 \
|
||||
--ckpt_dir=./checkpoint \
|
||||
|
|
@ -201,4 +199,4 @@ See [Preprocessing Guide](preprocessing.md) for config file format and examples.
|
|||
|
||||
---
|
||||
|
||||
> Document Update Time: 2026-07-09
|
||||
> Document Update Time: 2026-07-05
|
||||
|
|
@ -361,4 +361,4 @@ Pipeline(
|
|||
).run()
|
||||
```
|
||||
|
||||
> Document Update Time: 2026-07-09
|
||||
> Document Update Time: 2026-07-05
|
||||
|
|
|
|||
|
|
@ -201,7 +201,9 @@ nohup python scripts/tools/train.py \
|
|||
--warmup_ratio=0.05 \
|
||||
--max_lr=1e-4 \
|
||||
--max_grad_norm=1.0 \
|
||||
--weight_decay=0.1 \
|
||||
--adamw_beta1=0.9 \
|
||||
--adamw_beta2=0.95 \
|
||||
--adamw_weight_decay=0.01 \
|
||||
--window_size=2048 \
|
||||
--ckpt_interval=10000 \
|
||||
--ckpt_dir=./checkpoint \
|
||||
|
|
@ -212,4 +214,4 @@ nohup python scripts/tools/train.py \
|
|||
|
||||
Full parameter reference at [params.md](params.md).
|
||||
|
||||
> Document Update Time: 2026-07-09
|
||||
> Document Update Time: 2026-07-05
|
||||
|
|
|
|||
|
|
@ -1,19 +1,91 @@
|
|||
"""CUDA attention kernel wrappers with torch fallback.
|
||||
import importlib
|
||||
import logging
|
||||
|
||||
Public API:
|
||||
- ``gqa_decode_attn`` — single-query decode attention
|
||||
- ``gqa_prefill_attn`` — multi-query prefill attention
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.gqa_*``)
|
||||
when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``.
|
||||
"""
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from astrai.extension.loader import KERNEL_NAMES, is_available
|
||||
from astrai.extension.ops import gqa_decode_attn, gqa_prefill_attn
|
||||
_available: dict[str, bool] = {}
|
||||
_modules: dict[str, object] = {}
|
||||
|
||||
__all__ = [
|
||||
"gqa_decode_attn",
|
||||
"gqa_prefill_attn",
|
||||
"is_available",
|
||||
"KERNEL_NAMES",
|
||||
]
|
||||
for _name in ["gqa_decode_attn", "gqa_prefill_attn"]:
|
||||
try:
|
||||
_mod = importlib.import_module(f".{_name}", package=__package__)
|
||||
_available[_name] = True
|
||||
_modules[_name] = _mod
|
||||
except ImportError:
|
||||
_available[_name] = False
|
||||
_modules[_name] = None
|
||||
|
||||
|
||||
def _expand_kv_heads(
|
||||
k: torch.Tensor, v: torch.Tensor, q_head: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Expand K/V heads to match Q heads for GQA fallback."""
|
||||
kv_head = k.size(1)
|
||||
if kv_head == q_head:
|
||||
return k, v
|
||||
group = q_head // kv_head
|
||||
k = k.repeat_interleave(group, dim=1)
|
||||
v = v.repeat_interleave(group, dim=1)
|
||||
return k, v
|
||||
|
||||
|
||||
def _torch_fallback(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
mask: torch.Tensor | None,
|
||||
is_causal: bool,
|
||||
scale: float | None,
|
||||
) -> torch.Tensor:
|
||||
k, v = _expand_kv_heads(k, v, q.size(1))
|
||||
attn_mask = mask[:, None, None, :] if mask is not None else None
|
||||
return F.scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=attn_mask, is_causal=is_causal and mask is None, scale=scale
|
||||
)
|
||||
|
||||
|
||||
def gqa_decode_attn(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
mask: torch.Tensor | None = None,
|
||||
is_causal: bool = False,
|
||||
causal_offset: int = 0,
|
||||
scale: float | None = None,
|
||||
) -> torch.Tensor:
|
||||
if _available["gqa_decode_attn"]:
|
||||
return _modules["gqa_decode_attn"].gqa_decode_attn(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
mask=mask,
|
||||
is_causal=is_causal,
|
||||
causal_offset=causal_offset,
|
||||
scale=scale,
|
||||
)
|
||||
return _torch_fallback(q, k, v, mask, is_causal, scale)
|
||||
|
||||
|
||||
def gqa_prefill_attn(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
mask: torch.Tensor | None = None,
|
||||
is_causal: bool = False,
|
||||
causal_offset: int = 0,
|
||||
scale: float | None = None,
|
||||
) -> torch.Tensor:
|
||||
if _available["gqa_prefill_attn"]:
|
||||
return _modules["gqa_prefill_attn"].gqa_prefill_attn(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
mask=mask,
|
||||
is_causal=is_causal,
|
||||
causal_offset=causal_offset,
|
||||
scale=scale,
|
||||
)
|
||||
return _torch_fallback(q, k, v, mask, is_causal, scale)
|
||||
|
|
|
|||
|
|
@ -1,36 +0,0 @@
|
|||
"""Dynamic discovery and loading of compiled CUDA kernel modules.
|
||||
|
||||
Each kernel is registered in ``csrc/build.py`` and built into a ``.so`` placed
|
||||
in this package directory. On import we try to load each one; kernels that
|
||||
failed to build (or are running on a CPU-only machine) are marked unavailable
|
||||
so the wrapper functions can fall back to ``torch`` SDPA.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
KERNEL_NAMES = ["gqa_decode_attn", "gqa_prefill_attn"]
|
||||
|
||||
_available: dict[str, bool] = {}
|
||||
_modules: dict[str, object] = {}
|
||||
|
||||
for _name in KERNEL_NAMES:
|
||||
try:
|
||||
_mod = importlib.import_module(f".{_name}", package=__package__)
|
||||
_available[_name] = True
|
||||
_modules[_name] = _mod
|
||||
except ImportError:
|
||||
_available[_name] = False
|
||||
_modules[_name] = None
|
||||
|
||||
|
||||
def is_available(name: str) -> bool:
|
||||
"""Return ``True`` if the compiled kernel ``name`` was loaded."""
|
||||
return _available.get(name, False)
|
||||
|
||||
|
||||
def get_module(name: str) -> object:
|
||||
"""Return the loaded kernel module for ``name``, or ``None`` if unavailable."""
|
||||
return _modules.get(name)
|
||||
|
|
@ -1,86 +0,0 @@
|
|||
"""GQA attention wrapper functions — one entry point per compiled kernel.
|
||||
|
||||
Each wrapper dispatches to its CUDA kernel (loaded in ``loader.py``) when
|
||||
available, otherwise falls back to ``torch`` SDPA.
|
||||
|
||||
Add new kernel wrappers here; split into per-variant files only if this file
|
||||
grows large.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from astrai.extension.loader import _available, _modules
|
||||
|
||||
|
||||
def _expand_kv_heads(
|
||||
k: torch.Tensor, v: torch.Tensor, q_head: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Expand K/V heads to match Q heads for GQA fallback."""
|
||||
kv_head = k.size(1)
|
||||
if kv_head == q_head:
|
||||
return k, v
|
||||
group = q_head // kv_head
|
||||
k = k.repeat_interleave(group, dim=1)
|
||||
v = v.repeat_interleave(group, dim=1)
|
||||
return k, v
|
||||
|
||||
|
||||
def _torch_fallback(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
mask: torch.Tensor | None,
|
||||
is_causal: bool,
|
||||
scale: float | None,
|
||||
) -> torch.Tensor:
|
||||
"""Reference attention via ``scaled_dot_product_attention``."""
|
||||
k, v = _expand_kv_heads(k, v, q.size(1))
|
||||
attn_mask = mask[:, None, None, :] if mask is not None else None
|
||||
return F.scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=attn_mask, is_causal=is_causal and mask is None, scale=scale
|
||||
)
|
||||
|
||||
|
||||
def gqa_decode_attn(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
mask: torch.Tensor | None = None,
|
||||
is_causal: bool = False,
|
||||
causal_offset: int = 0,
|
||||
scale: float | None = None,
|
||||
) -> torch.Tensor:
|
||||
if _available["gqa_decode_attn"]:
|
||||
return _modules["gqa_decode_attn"].gqa_decode_attn(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
mask=mask,
|
||||
is_causal=is_causal,
|
||||
causal_offset=causal_offset,
|
||||
scale=scale,
|
||||
)
|
||||
return _torch_fallback(q, k, v, mask, is_causal, scale)
|
||||
|
||||
|
||||
def gqa_prefill_attn(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
mask: torch.Tensor | None = None,
|
||||
is_causal: bool = False,
|
||||
causal_offset: int = 0,
|
||||
scale: float | None = None,
|
||||
) -> torch.Tensor:
|
||||
if _available["gqa_prefill_attn"]:
|
||||
return _modules["gqa_prefill_attn"].gqa_prefill_attn(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
mask=mask,
|
||||
is_causal=is_causal,
|
||||
causal_offset=causal_offset,
|
||||
scale=scale,
|
||||
)
|
||||
return _torch_fallback(q, k, v, mask, is_causal, scale)
|
||||
|
|
@ -1,9 +1,7 @@
|
|||
from astrai.preprocessing.builder import (
|
||||
BaseMaskBuilder,
|
||||
MaskBuilderFactory,
|
||||
MultiOutputMaskBuilder,
|
||||
SectionedMaskBuilder,
|
||||
SingleOutputMaskBuilder,
|
||||
)
|
||||
from astrai.preprocessing.packing import (
|
||||
PackingStrategy,
|
||||
|
|
@ -22,14 +20,12 @@ from astrai.preprocessing.writer import (
|
|||
__all__ = [
|
||||
"BaseMaskBuilder",
|
||||
"MaskBuilderFactory",
|
||||
"MultiOutputMaskBuilder",
|
||||
"PackingStrategy",
|
||||
"PackingStrategyFactory",
|
||||
"Pipeline",
|
||||
"PositionIdStrategy",
|
||||
"PositionIdStrategyFactory",
|
||||
"SectionedMaskBuilder",
|
||||
"SingleOutputMaskBuilder",
|
||||
"StoreWriter",
|
||||
"StoreWriterFactory",
|
||||
"filter_by_length",
|
||||
|
|
|
|||
|
|
@ -1,10 +1,8 @@
|
|||
"""Mask building for preprocessing pipeline.
|
||||
|
||||
:class:`SectionRenderer` converts section specs into token ids and loss
|
||||
masks (template / text / value extraction). :class:`SingleOutputMaskBuilder`
|
||||
handles single-output (SFT / pretrain), :class:`MultiOutputMaskBuilder`
|
||||
handles multi-output (DPO / GRPO), and :class:`SectionedMaskBuilder`
|
||||
orchestrates both modes as a façade.
|
||||
masks (template / text / value extraction). :class:`SectionedMaskBuilder`
|
||||
orchestrates single-output / multi-output (DPO / GRPO) assembly.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
|
@ -214,17 +212,42 @@ class MaskBuilderFactory(BaseFactory["BaseMaskBuilder"]):
|
|||
pass
|
||||
|
||||
|
||||
@MaskBuilderFactory.register("single")
|
||||
class SingleOutputMaskBuilder(BaseMaskBuilder):
|
||||
"""Build a single output sequence with optional loss mask.
|
||||
@MaskBuilderFactory.register("sectioned")
|
||||
class SectionedMaskBuilder(BaseMaskBuilder):
|
||||
"""Config-driven builder supporting single and multi-output modes.
|
||||
|
||||
Expects ``config.input.sections`` (list of section specs).
|
||||
Single-output::
|
||||
|
||||
{"input": {"sections": [
|
||||
{"field": "messages", "action": "$role", "template": true}
|
||||
]}}
|
||||
→ {"sequence": [...], "loss_mask": [...], "domain": "..."}
|
||||
|
||||
Multi-output (DPO / GRPO)::
|
||||
|
||||
{"input": {"sources": {
|
||||
"chosen": {"sections": [{"field": "chosen", "action": "$role", "template": true}]},
|
||||
"rejected": {"sections": [{"field": "rejected", "action": "$role", "template": true}]},
|
||||
}}}
|
||||
→ {"chosen": [...], "chosen_mask": [...], "rejected": [...], "rejected_mask": [...], "domain": "..."}
|
||||
|
||||
Output spec fields::
|
||||
|
||||
sections – list of section specs (same format as single-output)
|
||||
list_field – True when JSONL field holds a list (GRPO responses)
|
||||
mask_key – explicit loss-mask output key (default: ``"{output_key}_mask"``)
|
||||
"""
|
||||
|
||||
def __init__(self, renderer: Optional[SectionRenderer] = None):
|
||||
self.renderer = renderer or SectionRenderer()
|
||||
def __init__(self):
|
||||
self.renderer = SectionRenderer()
|
||||
|
||||
def build(self, item: dict, config, tokenizer) -> Optional[dict]:
|
||||
sources_spec = getattr(config.input, "sources", None)
|
||||
if sources_spec:
|
||||
return self._build_multi(item, sources_spec, config, tokenizer)
|
||||
return self._build_single(item, config, tokenizer)
|
||||
|
||||
def _build_single(self, item: dict, config, tokenizer) -> Optional[dict]:
|
||||
sections = config.input.sections
|
||||
if not sections:
|
||||
return None
|
||||
|
|
@ -243,22 +266,9 @@ class SingleOutputMaskBuilder(BaseMaskBuilder):
|
|||
result["loss_mask"] = mask
|
||||
return result
|
||||
|
||||
|
||||
@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
|
||||
|
||||
def _build_multi(
|
||||
self, item: dict, sources_spec: dict, config, tokenizer
|
||||
) -> Optional[dict]:
|
||||
result: dict = {}
|
||||
any_output = False
|
||||
|
||||
|
|
@ -303,22 +313,3 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
|
|||
|
||||
result["domain"] = _extract_domain(item, config.output.domain_key)
|
||||
return result
|
||||
|
||||
|
||||
@MaskBuilderFactory.register("sectioned")
|
||||
class SectionedMaskBuilder(BaseMaskBuilder):
|
||||
"""Façade that dispatches to SingleOutputMaskBuilder or MultiOutputMaskBuilder.
|
||||
|
||||
Preserves backward compatibility for existing configs and code that rely
|
||||
on the ``"sectioned"`` factory name.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._single = SingleOutputMaskBuilder()
|
||||
self._multi = MultiOutputMaskBuilder()
|
||||
|
||||
def build(self, item: dict, config, tokenizer) -> Optional[dict]:
|
||||
sources_spec = getattr(config.input, "sources", None)
|
||||
if sources_spec:
|
||||
return self._multi.build(item, config, tokenizer)
|
||||
return self._single.build(item, config, tokenizer)
|
||||
|
|
|
|||
|
|
@ -26,9 +26,6 @@ def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
|
|||
tensor_group: Dict[str, List[Tensor]] = {}
|
||||
|
||||
root_path = Path(file_path)
|
||||
if root_path.is_file() and root_path.suffix in (".h5", ".hdf5"):
|
||||
h5_files = [root_path]
|
||||
else:
|
||||
h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
|
||||
|
||||
for h5_file in h5_files:
|
||||
|
|
|
|||
|
|
@ -209,6 +209,7 @@ class ProgressBarCallback(TrainCallback):
|
|||
|
||||
@only_on_rank(0)
|
||||
def on_optimizer_step(self, context: TrainContext):
|
||||
self.progress_bar.update(1)
|
||||
postfix = {
|
||||
"step": context.optimizer_step,
|
||||
"loss": f"{context.loss:.4f}",
|
||||
|
|
@ -219,7 +220,6 @@ class ProgressBarCallback(TrainCallback):
|
|||
if context.val_loss is not None:
|
||||
postfix["val_loss"] = f"{context.val_loss:.4f}"
|
||||
self.progress_bar.set_postfix(postfix)
|
||||
self.progress_bar.update(1)
|
||||
|
||||
@only_on_rank(0)
|
||||
def on_epoch_end(self, context: TrainContext):
|
||||
|
|
|
|||
|
|
@ -20,23 +20,13 @@ def _arch_flags() -> list[str]:
|
|||
_kernels_dir = Path("csrc/kernels")
|
||||
REGISTRY: dict[str, dict] = {}
|
||||
|
||||
CXX_FLAGS = ["-O3", "-march=native", "-funroll-loops"]
|
||||
NVCC_FLAGS = [
|
||||
"-O3",
|
||||
"--expt-relaxed-constexpr",
|
||||
"--use_fast_math",
|
||||
"--ptxas-options=-O3,-v",
|
||||
"--extra-device-vectorization",
|
||||
]
|
||||
|
||||
|
||||
def register(name: str, sources: list[str] | None = None, **kwargs):
|
||||
if sources is None:
|
||||
sources = [str(_kernels_dir / f"{name}.cu")]
|
||||
REGISTRY[name] = {
|
||||
"sources": sources,
|
||||
"cxx_flags": [*CXX_FLAGS],
|
||||
"nvcc_flags": [*NVCC_FLAGS, *_arch_flags()],
|
||||
"nvcc_flags": ["-O3", "--expt-relaxed-constexpr", *_arch_flags()],
|
||||
"extra_link_args": kwargs.pop("extra_link_args", []),
|
||||
**kwargs,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,43 +1,6 @@
|
|||
#include "gqa_decode_attn.cuh"
|
||||
#include <torch/extension.h>
|
||||
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
#include "gqa_decode_attn_mma.cuh"
|
||||
#endif
|
||||
|
||||
template <int HEAD_DIM>
|
||||
static void dispatch_decode(GQAParams& p) {
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
constexpr int BC = 32, BR = 16, LD = HEAD_DIM; // XOR swizzle → no padding
|
||||
int G = p.q_head / p.kv_head;
|
||||
// head-packing tensor-core path needs 1 < G <= 16 (MMA M dim) and no mask;
|
||||
// everything else uses the scalar kernel
|
||||
if (!p.use_mask && G > 1 && G <= 16) {
|
||||
dim3 grid(p.kv_head, p.batch, 1);
|
||||
dim3 block(32, 1, 1);
|
||||
// sK + sV + sQ, each BC/BR * LD (single buffer for high occupancy)
|
||||
int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
|
||||
cudaFuncSetAttribute(gqa_decode_attn_mma_kernel<HEAD_DIM, BC>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
gqa_decode_attn_mma_kernel<HEAD_DIM, BC><<<grid, block, smem>>>(p);
|
||||
return;
|
||||
}
|
||||
// scalar fallback (per-KV-head, one warp per query head)
|
||||
int group_size = p.q_head / p.kv_head;
|
||||
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
|
||||
dim3 block(32, group_size);
|
||||
dim3 grid(p.batch * p.kv_head);
|
||||
gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
|
||||
#else
|
||||
// scalar fallback (per-KV-head, one warp per query head)
|
||||
int group_size = p.q_head / p.kv_head;
|
||||
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
|
||||
dim3 block(32, group_size);
|
||||
dim3 grid(p.batch * p.kv_head);
|
||||
gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
|
||||
#endif
|
||||
}
|
||||
|
||||
torch::Tensor gqa_decode_attn(
|
||||
torch::Tensor q,
|
||||
torch::Tensor k,
|
||||
|
|
@ -81,23 +44,12 @@ torch::Tensor gqa_decode_attn(
|
|||
auto O = torch::empty_like(q);
|
||||
p.o = (bf16*)O.data_ptr();
|
||||
|
||||
switch (p.head_dim) {
|
||||
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)");
|
||||
}
|
||||
int group_size = p.q_head / p.kv_head;
|
||||
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
|
||||
dim3 block(32, group_size);
|
||||
dim3 grid(p.batch * p.kv_head);
|
||||
|
||||
gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
|
||||
return O;
|
||||
}
|
||||
|
||||
|
|
@ -110,5 +62,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|||
py::arg("is_causal") = false,
|
||||
py::arg("causal_offset") = 0,
|
||||
py::arg("scale") = py::none(),
|
||||
"GQA decode (tensor-core head-packing on sm_80+, scalar fallback)");
|
||||
"GQA decode (per-KV-head, shared K)");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,219 +0,0 @@
|
|||
#pragma once
|
||||
#include "gqa_common.cuh"
|
||||
#include "gqa_mma_utils.cuh"
|
||||
|
||||
// Tensor-core decode via GQA head-packing with cp.async loads.
|
||||
//
|
||||
// Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core work
|
||||
// on its own. But GQA gives us G = q_head / kv_head query heads that all share
|
||||
// one kv_head. We pack those G heads into the M=16 rows of mma.sync.m16n8k16,
|
||||
// turning G independent GEMVs into a single GEMM that reuses each loaded K/V tile
|
||||
// across all G heads (K/V load is the decode bottleneck, so the reuse is the win,
|
||||
// not the flops). Fragment layout is identical to the prefill mma kernel; the
|
||||
// only differences are (1) the M rows come from different heads at position 0
|
||||
// instead of different sequence positions of one head, and (2) causal masking is
|
||||
// a single scalar bound shared by every row. One warp owns one (batch, kv_head);
|
||||
// requires G <= 16.
|
||||
//
|
||||
// Optimizations:
|
||||
// - cp.async global→shared for K/V (bypasses registers, cuts instruction count)
|
||||
// - XOR swizzle (swiz_col): LD=HEAD_DIM, zero waste, no bank conflicts
|
||||
// - pre-scaled Q: Q scaled during load, softmax skips per-tile multiply
|
||||
// - single-buffer: keeps smem small for high occupancy
|
||||
|
||||
template <int HEAD_DIM, int BC>
|
||||
__global__ void gqa_decode_attn_mma_kernel(GQAParams p) {
|
||||
constexpr int BR = 16;
|
||||
constexpr int KD = HEAD_DIM / 16; // Q/K k-tiles
|
||||
constexpr int NC8 = BC / 8; // S n-tiles (N=8 each)
|
||||
constexpr int KT2 = BC / 16; // P k-tiles (K=16 each)
|
||||
constexpr int DN8 = HEAD_DIM / 8; // O n-tiles (N=8 each)
|
||||
constexpr int LD = HEAD_DIM; // XOR swizzle handles bank conflicts, zero waste
|
||||
constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
|
||||
|
||||
const int lane = threadIdx.x; // single warp
|
||||
const int gid = lane >> 2; // 0..7 → rows gid, gid+8
|
||||
const int tid4 = lane & 3;
|
||||
|
||||
const int kv_head = blockIdx.x;
|
||||
const int batch = blockIdx.y;
|
||||
const int G = p.q_head / p.kv_head;
|
||||
const int q_head0 = kv_head * G;
|
||||
|
||||
extern __shared__ __align__(16) bf16 smem[];
|
||||
bf16* sK = smem; // [BC][LD]
|
||||
bf16* sV = sK + BC * LD; // [BC][LD]
|
||||
bf16* sQ = sV + BC * LD; // [BR][LD]
|
||||
|
||||
// ---- stage Q into shared (pre-scaled, swizzled) ----
|
||||
bf16 scale_bf16 = __float2bfloat16(p.scale);
|
||||
for (int i = lane; i < BR * HEAD_DIM; i += 32) {
|
||||
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||
bf16 val = __float2bfloat16(0.0f);
|
||||
if (r < G) {
|
||||
int qh = q_head0 + r;
|
||||
val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d]; // q_len == 1
|
||||
}
|
||||
sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16);
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
// Q resident A-fragments
|
||||
unsigned Qa[KD][4];
|
||||
int qrow_l = (lane & 7) + (lane & 8);
|
||||
int qcol_l = (lane & 16) ? 8 : 0;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KD; kt++)
|
||||
ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
|
||||
|
||||
float Oacc[DN8][4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < DN8; j++)
|
||||
Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
|
||||
float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
|
||||
|
||||
const int kv_base = (batch * p.kv_head + kv_head) * p.kv_len * HEAD_DIM;
|
||||
const int mask_base = batch * p.kv_len;
|
||||
const int tiles = (p.kv_len + BC - 1) / BC;
|
||||
const int has_mask = p.use_mask && p.mask;
|
||||
|
||||
for (int ti = 0; ti < tiles; ti++) {
|
||||
int kv0 = ti * BC;
|
||||
|
||||
// ---- load K/V tile to shared (cp.async on full tiles) ----
|
||||
bool full_tile = (kv0 + BC <= p.kv_len);
|
||||
if (full_tile) {
|
||||
constexpr int VEC = 8; // 8 bf16 = 16 bytes per cp.async
|
||||
int total = BC * HEAD_DIM;
|
||||
#pragma unroll
|
||||
for (int i = lane * VEC; i < total; i += 32 * VEC) {
|
||||
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||
int kc = kv0 + r;
|
||||
cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)],
|
||||
&p.k[kv_base + kc * HEAD_DIM + d]);
|
||||
cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)],
|
||||
&p.v[kv_base + kc * HEAD_DIM + d]);
|
||||
}
|
||||
cp_async_commit();
|
||||
cp_async_wait_all();
|
||||
} else {
|
||||
for (int i = lane; i < BC * HEAD_DIM; i += 32) {
|
||||
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||
int kc = kv0 + r;
|
||||
bf16 z = __float2bfloat16(0.0f);
|
||||
sK[r * LD + swiz_col(d, r, SWIZ_MASK)] =
|
||||
(kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z;
|
||||
sV[r * LD + swiz_col(d, r, SWIZ_MASK)] =
|
||||
(kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z;
|
||||
}
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
// S = Q @ K^T (Q already pre-scaled, so Sacc includes scale)
|
||||
float Sacc[NC8][4];
|
||||
#pragma unroll
|
||||
for (int n8 = 0; n8 < NC8; n8++) {
|
||||
Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f;
|
||||
int krow_l = n8 * 8 + (lane & 7);
|
||||
int kcol_h = (lane & 8) ? 8 : 0;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KD; kt++) {
|
||||
unsigned b[2];
|
||||
ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]);
|
||||
mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]);
|
||||
}
|
||||
}
|
||||
|
||||
// ---- online softmax (Q pre-scaled → no per-tile scale multiply) ----
|
||||
float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
|
||||
#pragma unroll
|
||||
for (int n8 = 0; n8 < NC8; n8++) {
|
||||
int cc = kv0 + n8 * 8 + 2 * tid4;
|
||||
bool bc0 = (cc >= p.kv_len) ||
|
||||
(has_mask && !p.mask[mask_base + cc]);
|
||||
bool bc1 = (cc + 1 >= p.kv_len) ||
|
||||
(has_mask && !p.mask[mask_base + cc + 1]);
|
||||
bool cz = p.is_causal;
|
||||
int off = p.causal_offset;
|
||||
bool bad0 = bc0 || (cz && cc > off);
|
||||
bool bad1 = bc1 || (cz && (cc + 1) > off);
|
||||
float s0 = bad0 ? -FLT_MAX : Sacc[n8][0];
|
||||
float s1 = bad1 ? -FLT_MAX : Sacc[n8][1];
|
||||
float s2 = bad0 ? -FLT_MAX : Sacc[n8][2];
|
||||
float s3 = bad1 ? -FLT_MAX : Sacc[n8][3];
|
||||
Sacc[n8][0] = s0; Sacc[n8][1] = s1; Sacc[n8][2] = s2; Sacc[n8][3] = s3;
|
||||
rmax0 = fmaxf(rmax0, fmaxf(s0, s1));
|
||||
rmax1 = fmaxf(rmax1, fmaxf(s2, s3));
|
||||
}
|
||||
rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1));
|
||||
rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2));
|
||||
rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1));
|
||||
rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2));
|
||||
|
||||
float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1);
|
||||
float corr0 = (nm0 == -FLT_MAX) ? 1.0f : __expf(m0 - nm0);
|
||||
float corr1 = (nm1 == -FLT_MAX) ? 1.0f : __expf(m1 - nm1);
|
||||
|
||||
float rsum0 = 0.0f, rsum1 = 0.0f;
|
||||
#pragma unroll
|
||||
for (int n8 = 0; n8 < NC8; n8++) {
|
||||
float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][0] - nm0);
|
||||
float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][1] - nm0);
|
||||
float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][2] - nm1);
|
||||
float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][3] - nm1);
|
||||
Sacc[n8][0] = p0; Sacc[n8][1] = p1; Sacc[n8][2] = p2; Sacc[n8][3] = p3;
|
||||
rsum0 += p0 + p1;
|
||||
rsum1 += p2 + p3;
|
||||
}
|
||||
rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1);
|
||||
rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2);
|
||||
rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1);
|
||||
rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2);
|
||||
l0 = l0 * corr0 + rsum0;
|
||||
l1 = l1 * corr1 + rsum1;
|
||||
m0 = nm0; m1 = nm1;
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < DN8; j++) {
|
||||
Oacc[j][0] *= corr0; Oacc[j][1] *= corr0;
|
||||
Oacc[j][2] *= corr1; Oacc[j][3] *= corr1;
|
||||
}
|
||||
|
||||
// O += P @ V
|
||||
#pragma unroll
|
||||
for (int kt2 = 0; kt2 < KT2; kt2++) {
|
||||
unsigned Pa[4];
|
||||
Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]);
|
||||
Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]);
|
||||
Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]);
|
||||
Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]);
|
||||
int vrow_l = kt2 * 16 + (lane & 15);
|
||||
#pragma unroll
|
||||
for (int dn8 = 0; dn8 < DN8; dn8++) {
|
||||
unsigned b[2];
|
||||
ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]);
|
||||
mma16816(Oacc[dn8], Pa, b, Oacc[dn8]);
|
||||
}
|
||||
}
|
||||
__syncwarp(); // sK/sV reused next tile
|
||||
}
|
||||
|
||||
// ---- write output ----
|
||||
float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
|
||||
float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
|
||||
#pragma unroll
|
||||
for (int dn8 = 0; dn8 < DN8; dn8++) {
|
||||
int d = dn8 * 8 + 2 * tid4;
|
||||
int r0 = gid, r1 = gid + 8;
|
||||
if (r0 < G) {
|
||||
int o_off = (batch * p.q_head + q_head0 + r0) * HEAD_DIM + d;
|
||||
p.o[o_off] = __float2bfloat16(Oacc[dn8][0] * rl0);
|
||||
p.o[o_off + 1] = __float2bfloat16(Oacc[dn8][1] * rl0);
|
||||
}
|
||||
if (r1 < G) {
|
||||
int o_off = (batch * p.q_head + q_head0 + r1) * HEAD_DIM + d;
|
||||
p.o[o_off] = __float2bfloat16(Oacc[dn8][2] * rl1);
|
||||
p.o[o_off + 1] = __float2bfloat16(Oacc[dn8][3] * rl1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,93 +0,0 @@
|
|||
#pragma once
|
||||
|
||||
// Shared MMA utilities for tensor-core GQA kernels.
|
||||
// mma.sync.m16n8k16 PTX wrappers, ldmatrix helpers, and bf16 packing.
|
||||
|
||||
// mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
|
||||
__device__ __forceinline__ void mma16816(float* d, const unsigned* a,
|
||||
const unsigned* b, const float* c) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
|
||||
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};"
|
||||
: "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3])
|
||||
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
|
||||
"f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
|
||||
}
|
||||
|
||||
// read two adjacent bf16 from smem as one packed .b32 (elem0 low, elem1 high)
|
||||
__device__ __forceinline__ unsigned ld2(const bf16* p) {
|
||||
return *reinterpret_cast<const unsigned*>(p);
|
||||
}
|
||||
|
||||
// pack two floats into one bf16x2 as .b32
|
||||
__device__ __forceinline__ unsigned pk2(float a, float b) {
|
||||
__nv_bfloat162 v = __floats2bfloat162_rn(a, b);
|
||||
return *reinterpret_cast<unsigned*>(&v);
|
||||
}
|
||||
|
||||
// pack two (non-contiguous) bf16 into one .b32
|
||||
__device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) {
|
||||
__nv_bfloat162 v;
|
||||
v.x = a;
|
||||
v.y = b;
|
||||
return *reinterpret_cast<unsigned*>(&v);
|
||||
}
|
||||
|
||||
// ldmatrix: cooperatively load mma fragments from smem (one instruction per
|
||||
// 16x16 / 16x8 tile) with the exact register layout mma expects — replaces the
|
||||
// scalar per-thread fragment packing, cutting shared-load instructions and bank
|
||||
// conflicts. Each lane supplies the shared address of one 8-wide row.
|
||||
__device__ __forceinline__ void ldmatrix_x4(unsigned* r, const bf16* p) {
|
||||
unsigned a = __cvta_generic_to_shared(p);
|
||||
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
|
||||
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
|
||||
: "r"(a));
|
||||
}
|
||||
__device__ __forceinline__ void ldmatrix_x2(unsigned* r, const bf16* p) {
|
||||
unsigned a = __cvta_generic_to_shared(p);
|
||||
asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
|
||||
: "=r"(r[0]), "=r"(r[1])
|
||||
: "r"(a));
|
||||
}
|
||||
__device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) {
|
||||
unsigned a = __cvta_generic_to_shared(p);
|
||||
asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];"
|
||||
: "=r"(r[0]), "=r"(r[1])
|
||||
: "r"(a));
|
||||
}
|
||||
|
||||
// XOR swizzle for shared-memory column at 8-bf16 chunk granularity.
|
||||
// Eliminates ldmatrix bank conflicts without LD padding: consecutive rows
|
||||
// land in distinct bank groups. swiz_col(d, r, mask) = ((d>>3)^(r&mask))<<3 | (d&7).
|
||||
// mask must cover log2(HEAD_DIM/8) chunk bits but stay within LD: use 7 for
|
||||
// HEAD_DIM>=64 (8+ chunks), 3 for HEAD_DIM=32 (4 chunks). Default 7 keeps
|
||||
// existing HEAD_DIM>=64 call sites working unchanged.
|
||||
__device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) {
|
||||
return ((d >> 3) ^ (r & mask)) << 3 | (d & 7);
|
||||
}
|
||||
|
||||
// cp.async: copy 16 bytes (8 bf16) from global to shared memory directly,
|
||||
// bypassing registers. Eliminates shared-store bank conflicts and cuts
|
||||
// load-loop instruction count in half (1 cp.async vs 1 LDG + 1 STS).
|
||||
// Requires sm_80+.
|
||||
__device__ __forceinline__ void cp_async_16(bf16* smem_ptr, const void* gmem_ptr) {
|
||||
unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
|
||||
asm volatile("cp.async.ca.shared.global [%0], [%1], 16;"
|
||||
:: "r"(smem_addr), "l"(gmem_ptr));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void cp_async_commit() {
|
||||
asm volatile("cp.async.commit_group;");
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void cp_async_wait_all() {
|
||||
asm volatile("cp.async.wait_all;");
|
||||
}
|
||||
|
||||
// Wait until at most N commit groups are still in flight. Used for
|
||||
// double-buffered pipelining: wait_group<1> lets the next tile's cp.async
|
||||
// continue while ensuring the current tile's data is ready.
|
||||
template <int N>
|
||||
__device__ __forceinline__ void cp_async_wait_group() {
|
||||
asm volatile("cp.async.wait_group %0;" :: "n"(N));
|
||||
}
|
||||
|
|
@ -8,11 +8,10 @@
|
|||
template <int HEAD_DIM>
|
||||
static void dispatch_prefill(GQAParams& p) {
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
constexpr int WARPS = 4, BC = 32, BR = 16, LD = HEAD_DIM;
|
||||
constexpr int WARPS = 4, BC = 32, BR = 16, LD = HEAD_DIM + 8;
|
||||
dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
|
||||
dim3 block(WARPS * 32, 1, 1);
|
||||
// sK + sV (each BC*LD) + shared sQ staging (BR*LD)
|
||||
int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
|
||||
int smem = (2 * BC * LD + WARPS * BR * LD) * (int)sizeof(bf16);
|
||||
cudaFuncSetAttribute(gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC><<<grid, block, smem>>>(p);
|
||||
|
|
@ -68,9 +67,6 @@ torch::Tensor gqa_prefill_attn(
|
|||
p.o = (bf16*)O.data_ptr();
|
||||
|
||||
switch (p.head_dim) {
|
||||
case 32:
|
||||
dispatch_prefill<32>(p);
|
||||
break;
|
||||
case 64:
|
||||
dispatch_prefill<64>(p);
|
||||
break;
|
||||
|
|
@ -82,7 +78,7 @@ torch::Tensor gqa_prefill_attn(
|
|||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim,
|
||||
" (supported: 32,64,128,256)");
|
||||
" (supported: 64,128,256)");
|
||||
}
|
||||
return O;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
#pragma once
|
||||
#include "gqa_common.cuh"
|
||||
#include "gqa_mma_utils.cuh"
|
||||
|
||||
// Tensor-core prefill, register-resident flash attention (raw mma.sync PTX).
|
||||
// One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
|
||||
|
|
@ -11,12 +10,57 @@
|
|||
// for-element onto the P matrix_a (bf16) operand, so softmax needs no shuffle
|
||||
// repack; row reductions fold across the 4-lane thread group. Templated on
|
||||
// <HEAD_DIM, WARPS, BC> with BC a multiple of 16.
|
||||
//
|
||||
// Optimizations: shared sQ staging (single area, serialized per-warp load)
|
||||
// → cuts smem; pre-scale Q by attention scale during Q load; cp.async global→
|
||||
// shared for K/V; scalar fallback only for the last partial tile; causal tile
|
||||
// skipping (block-level early break + warp-level skip); XOR swizzle (swiz_col)
|
||||
// → eliminates ldmatrix bank conflicts without LD padding (LD=HEAD_DIM).
|
||||
|
||||
// mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
|
||||
// (only compiled when ASTRAI_HAS_MMA is set, i.e. built for sm_80+)
|
||||
__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);
|
||||
}
|
||||
__device__ __forceinline__ unsigned pk2(float a, float b) {
|
||||
__nv_bfloat162 v = __floats2bfloat162_rn(a, b);
|
||||
return *reinterpret_cast<unsigned*>(&v);
|
||||
}
|
||||
// pack two (non-contiguous) bf16 into one .b32
|
||||
__device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) {
|
||||
__nv_bfloat162 v;
|
||||
v.x = a;
|
||||
v.y = b;
|
||||
return *reinterpret_cast<unsigned*>(&v);
|
||||
}
|
||||
|
||||
// ldmatrix: cooperatively load mma fragments from smem (one instruction per
|
||||
// 16x16 / 16x8 tile) with the exact register layout mma expects — replaces the
|
||||
// scalar per-thread fragment packing, cutting shared-load instructions and bank
|
||||
// conflicts. Each lane supplies the shared address of one 8-wide row.
|
||||
__device__ __forceinline__ void ldmatrix_x4(unsigned* r, const bf16* p) {
|
||||
unsigned a = __cvta_generic_to_shared(p);
|
||||
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
|
||||
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
|
||||
: "r"(a));
|
||||
}
|
||||
__device__ __forceinline__ void ldmatrix_x2(unsigned* r, const bf16* p) {
|
||||
unsigned a = __cvta_generic_to_shared(p);
|
||||
asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
|
||||
: "=r"(r[0]), "=r"(r[1])
|
||||
: "r"(a));
|
||||
}
|
||||
__device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) {
|
||||
unsigned a = __cvta_generic_to_shared(p);
|
||||
asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];"
|
||||
: "=r"(r[0]), "=r"(r[1])
|
||||
: "r"(a));
|
||||
}
|
||||
|
||||
template <int HEAD_DIM, int WARPS, int BC>
|
||||
__global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
|
||||
|
|
@ -25,8 +69,9 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
|
|||
constexpr int NC8 = BC / 8; // S n-tiles (N=8 each)
|
||||
constexpr int KT2 = BC / 16; // P k-tiles (K=16 each)
|
||||
constexpr int DN8 = HEAD_DIM / 8; // O n-tiles (N=8 each)
|
||||
constexpr int LD = HEAD_DIM; // XOR swizzle (swiz_col) handles bank conflicts
|
||||
constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1); // chunk bits, stay within LD
|
||||
constexpr int LD = HEAD_DIM + 8; // padded smem row stride (kills ldmatrix
|
||||
// bank conflicts: consecutive rows land
|
||||
// in distinct banks instead of colliding)
|
||||
|
||||
const int warp = threadIdx.x / 32;
|
||||
const int lane = threadIdx.x % 32;
|
||||
|
|
@ -42,31 +87,24 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
|
|||
extern __shared__ __align__(16) bf16 smem[];
|
||||
bf16* sK = smem; // [BC][LD]
|
||||
bf16* sV = sK + BC * LD; // [BC][LD]
|
||||
bf16* sQ = sV + BC * LD; // shared staging [BR][LD]
|
||||
bf16* sQ = sV + BC * LD + warp * (BR * LD); // per-warp [BR][LD]
|
||||
|
||||
// Q resident A-fragments (loaded once per warp via shared staging).
|
||||
// Pre-scale by attention scale so softmax doesn't need to multiply later.
|
||||
// stage Q into smem (zero-padded past q_len)
|
||||
const int q_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM;
|
||||
unsigned Qa[KD][4];
|
||||
bf16 scale_bf16 = __float2bfloat16(p.scale);
|
||||
int qrow_l = (lane & 7) + (lane & 8); // 0..15
|
||||
int qcol_l = (lane & 16) ? 8 : 0;
|
||||
for (int w = 0; w < WARPS; w++) {
|
||||
if (warp == w) {
|
||||
for (int i = lane; i < BR * HEAD_DIM; i += 32) {
|
||||
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||
int qr = qrow0 + r;
|
||||
bf16 qv = (qr < p.q_len) ? p.q[q_base + qr * HEAD_DIM + d]
|
||||
: __float2bfloat16(0.0f);
|
||||
sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(qv, scale_bf16);
|
||||
sQ[r * LD + d] = (qr < p.q_len) ? p.q[q_base + qr * HEAD_DIM + d] : __float2bfloat16(0.0f);
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
// Q resident A-fragments: Qa[kt][0..3] (loaded once via ldmatrix.x4)
|
||||
unsigned Qa[KD][4];
|
||||
int qrow_l = (lane & 7) + (lane & 8); // 0..15
|
||||
int qcol_l = (lane & 16) ? 8 : 0;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KD; kt++)
|
||||
ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
|
||||
}
|
||||
__syncthreads(); // prevent next warp from overwriting sQ prematurely
|
||||
}
|
||||
ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + kt * 16 + qcol_l]);
|
||||
|
||||
float Oacc[DN8][4];
|
||||
#pragma unroll
|
||||
|
|
@ -79,90 +117,60 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
|
|||
const int qr0 = qrow0 + gid; // row for c0/c1
|
||||
const int qr1 = qrow0 + gid + 8; // row for c2/c3
|
||||
|
||||
// Causal tile-skip bounds (no-op when is_causal == 0)
|
||||
const int use_skip = p.is_causal;
|
||||
const int max_kv = qrow0 + BR - 1 + p.causal_offset;
|
||||
const int block_max_kv =
|
||||
blockIdx.x * WARPS * BR + WARPS * BR - 1 + p.causal_offset;
|
||||
const int has_mask = p.use_mask && p.mask;
|
||||
const int mb = batch * p.kv_len;
|
||||
|
||||
for (int ti = 0; ti < tiles; ti++) {
|
||||
int kv0 = ti * BC;
|
||||
|
||||
// Block-level causal early break
|
||||
if (use_skip && kv0 > block_max_kv) break;
|
||||
|
||||
// ---- load K/V tile to shared memory (cp.async on full tiles) ----
|
||||
bool full_tile = (kv0 + BC <= p.kv_len);
|
||||
if (full_tile) {
|
||||
constexpr int VEC = 8; // bf16 per cp.async unit (16 bytes)
|
||||
int total = BC * HEAD_DIM;
|
||||
#pragma unroll
|
||||
for (int i = threadIdx.x * VEC; i < total; i += nthreads * VEC) {
|
||||
int r = i / HEAD_DIM;
|
||||
int d = i % HEAD_DIM;
|
||||
int kc = kv0 + r;
|
||||
cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)], &p.k[kv_base + kc * HEAD_DIM + d]);
|
||||
cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)], &p.v[kv_base + kc * HEAD_DIM + d]);
|
||||
}
|
||||
cp_async_commit();
|
||||
cp_async_wait_all();
|
||||
} else {
|
||||
for (int i = threadIdx.x; i < BC * HEAD_DIM; i += nthreads) {
|
||||
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||
int kc = kv0 + r;
|
||||
bf16 z = __float2bfloat16(0.0f);
|
||||
sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = (kc < p.kv_len)
|
||||
? 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;
|
||||
}
|
||||
sK[r * LD + d] = (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z;
|
||||
sV[r * LD + d] = (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Warp-level causal skip
|
||||
if (!use_skip || kv0 <= max_kv) {
|
||||
|
||||
// S = Q @ K^T → Sacc[n8][0..3] (n8: 8 kv cols each)
|
||||
float Sacc[NC8][4];
|
||||
#pragma unroll
|
||||
for (int n8 = 0; n8 < NC8; n8++) {
|
||||
Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f;
|
||||
int krow_l = n8 * 8 + (lane & 7);
|
||||
int kcol_h = (lane & 8) ? 8 : 0;
|
||||
int kv = kv0 + n8 * 8 + gid;
|
||||
int krow_l = n8 * 8 + (lane & 7); // kv within tile
|
||||
int kcol_h = (lane & 8) ? 8 : 0; // which k-half
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KD; kt++) {
|
||||
unsigned b[2];
|
||||
ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]);
|
||||
ldmatrix_x2(b, &sK[krow_l * LD + kt * 16 + kcol_h]);
|
||||
mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]);
|
||||
}
|
||||
(void)kv;
|
||||
}
|
||||
|
||||
// ---- online softmax (in registers) ----
|
||||
// Q is pre-scaled, so Sacc already includes the attention scale.
|
||||
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;
|
||||
// scale + mask, then per-row (gid, gid+8) max over held cols
|
||||
float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
|
||||
#pragma unroll
|
||||
for (int n8 = 0; n8 < NC8; n8++) {
|
||||
int cc = kv0 + n8 * 8 + 2 * tid4;
|
||||
int c1 = cc + 1;
|
||||
bool b0 = (cc >= maxc0) || (has_mask && !p.mask[mb + cc]);
|
||||
bool b1 = (c1 >= maxc0) || (has_mask && !p.mask[mb + c1]);
|
||||
bool b2 = (cc >= maxc1) || (has_mask && !p.mask[mb + cc]);
|
||||
bool b3 = (c1 >= maxc1) || (has_mask && !p.mask[mb + c1]);
|
||||
float s0 = b0 ? -FLT_MAX : Sacc[n8][0];
|
||||
float s1 = b1 ? -FLT_MAX : Sacc[n8][1];
|
||||
float s2 = b2 ? -FLT_MAX : Sacc[n8][2];
|
||||
float s3 = b3 ? -FLT_MAX : Sacc[n8][3];
|
||||
Sacc[n8][0] = s0; Sacc[n8][1] = s1;
|
||||
Sacc[n8][2] = s2; Sacc[n8][3] = s3;
|
||||
int cc = kv0 + n8 * 8 + 2 * tid4; // col for c0/c2
|
||||
bool bc0 = (cc >= p.kv_len) ||
|
||||
(p.use_mask && p.mask && !p.mask[batch * p.kv_len + cc]);
|
||||
bool bc1 = (cc + 1 >= p.kv_len) ||
|
||||
(p.use_mask && p.mask && !p.mask[batch * p.kv_len + cc + 1]);
|
||||
bool cz = p.is_causal;
|
||||
int off = p.causal_offset;
|
||||
bool bad0 = bc0 || (cz && cc > qr0 + off);
|
||||
bool bad1 = bc1 || (cz && (cc + 1) > qr0 + off);
|
||||
bool bad2 = bc0 || (cz && cc > qr1 + off);
|
||||
bool bad3 = bc1 || (cz && (cc + 1) > qr1 + off);
|
||||
float s0 = bad0 ? -FLT_MAX : Sacc[n8][0] * p.scale;
|
||||
float s1 = bad1 ? -FLT_MAX : Sacc[n8][1] * p.scale;
|
||||
float s2 = bad2 ? -FLT_MAX : Sacc[n8][2] * p.scale;
|
||||
float s3 = bad3 ? -FLT_MAX : Sacc[n8][3] * p.scale;
|
||||
Sacc[n8][0] = s0; Sacc[n8][1] = s1; Sacc[n8][2] = s2; Sacc[n8][3] = s3;
|
||||
rmax0 = fmaxf(rmax0, fmaxf(s0, s1));
|
||||
rmax1 = fmaxf(rmax1, fmaxf(s2, s3));
|
||||
}
|
||||
// reduce max across the 4-lane group (tid4)
|
||||
rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1));
|
||||
rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2));
|
||||
rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1));
|
||||
|
|
@ -175,16 +183,11 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
|
|||
float rsum0 = 0.0f, rsum1 = 0.0f;
|
||||
#pragma unroll
|
||||
for (int n8 = 0; n8 < NC8; n8++) {
|
||||
float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f
|
||||
: __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;
|
||||
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;
|
||||
}
|
||||
|
|
@ -211,16 +214,15 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
|
|||
Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]);
|
||||
Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]);
|
||||
Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]);
|
||||
int vrow_l = kt2 * 16 + (lane & 15);
|
||||
int vrow_l = kt2 * 16 + (lane & 15); // kv within tile (0..15)
|
||||
#pragma unroll
|
||||
for (int dn8 = 0; dn8 < DN8; dn8++) {
|
||||
unsigned b[2];
|
||||
ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]);
|
||||
ldmatrix_x2_trans(b, &sV[vrow_l * LD + dn8 * 8]);
|
||||
mma16816(Oacc[dn8], Pa, b, Oacc[dn8]);
|
||||
}
|
||||
}
|
||||
} // if active (warp-level causal skip)
|
||||
__syncthreads();
|
||||
__syncthreads(); // sK/sV reused next tile
|
||||
}
|
||||
|
||||
// ---- write output ----
|
||||
|
|
@ -231,16 +233,12 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) {
|
|||
for (int dn8 = 0; dn8 < DN8; dn8++) {
|
||||
int d = dn8 * 8 + 2 * tid4;
|
||||
if (qr0 < p.q_len) {
|
||||
p.o[o_base + qr0 * HEAD_DIM + d] =
|
||||
__float2bfloat16(Oacc[dn8][0] * rl0);
|
||||
p.o[o_base + qr0 * HEAD_DIM + d + 1] =
|
||||
__float2bfloat16(Oacc[dn8][1] * rl0);
|
||||
p.o[o_base + qr0 * HEAD_DIM + d] = __float2bfloat16(Oacc[dn8][0] * rl0);
|
||||
p.o[o_base + qr0 * HEAD_DIM + d + 1] = __float2bfloat16(Oacc[dn8][1] * rl0);
|
||||
}
|
||||
if (qr1 < p.q_len) {
|
||||
p.o[o_base + qr1 * HEAD_DIM + d] =
|
||||
__float2bfloat16(Oacc[dn8][2] * rl1);
|
||||
p.o[o_base + qr1 * HEAD_DIM + d + 1] =
|
||||
__float2bfloat16(Oacc[dn8][3] * rl1);
|
||||
p.o[o_base + qr1 * HEAD_DIM + d] = __float2bfloat16(Oacc[dn8][2] * rl1);
|
||||
p.o[o_base + qr1 * HEAD_DIM + d + 1] = __float2bfloat16(Oacc[dn8][3] * rl1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,10 +1,4 @@
|
|||
/*
|
||||
Pure-C test:
|
||||
nvcc -I csrc -arch=sm_89 -O3 \
|
||||
--use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
|
||||
csrc/tests/gqa_decode_test.cu -o test && ./test
|
||||
*/
|
||||
|
||||
// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_decode_test.cu -o test && ./test
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cmath>
|
||||
|
|
|
|||
|
|
@ -1,10 +1,4 @@
|
|||
/*
|
||||
Pure-C test:
|
||||
nvcc -I csrc -arch=sm_89 -O3 \
|
||||
--use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
|
||||
csrc/tests/gqa_prefill_test.cu -o test && ./test
|
||||
*/
|
||||
|
||||
// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_prefill_test.cu -o test && ./test
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cmath>
|
||||
|
|
|
|||
|
|
@ -1,11 +1,9 @@
|
|||
import argparse
|
||||
import os
|
||||
from functools import partial
|
||||
from typing import Any, Dict
|
||||
|
||||
import torch
|
||||
import torch.optim as optim
|
||||
from torch import Tensor, nn
|
||||
|
||||
from astrai.config import AutoRegressiveLMConfig, TrainConfig
|
||||
from astrai.dataset import DatasetFactory
|
||||
|
|
@ -14,84 +12,6 @@ from astrai.model.components.decoder_block import DecoderBlock
|
|||
from astrai.trainer import SchedulerFactory, Trainer
|
||||
|
||||
|
||||
class MuonMix(optim.Optimizer):
|
||||
"""Combined Muon (matrix) + AdamW (non-matrix) optimizer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
lr: float = 3e-4,
|
||||
weight_decay: float = 0.1,
|
||||
momentum: float = 0.95,
|
||||
nesterov: bool = True,
|
||||
ns_steps: int = 5,
|
||||
adjust_lr_fn: str = "match_rms_adamw",
|
||||
):
|
||||
defaults = dict(
|
||||
lr=lr,
|
||||
weight_decay=weight_decay,
|
||||
momentum=momentum,
|
||||
nesterov=nesterov,
|
||||
ns_steps=ns_steps,
|
||||
adjust_lr_fn=adjust_lr_fn,
|
||||
)
|
||||
params = [p for p in model.parameters() if p.requires_grad]
|
||||
super().__init__(params, defaults)
|
||||
|
||||
matrix_params: list[Tensor] = []
|
||||
other_params: list[Tensor] = []
|
||||
for name, param in model.named_parameters():
|
||||
if not param.requires_grad:
|
||||
continue
|
||||
if (
|
||||
param.dim() >= 2
|
||||
and "norm" not in name
|
||||
and "bias" not in name
|
||||
and "embed" not in name
|
||||
and "lm_head" not in name
|
||||
):
|
||||
matrix_params.append(param)
|
||||
else:
|
||||
other_params.append(param)
|
||||
|
||||
self.muon = optim.Muon(
|
||||
matrix_params,
|
||||
lr=lr,
|
||||
weight_decay=weight_decay,
|
||||
momentum=momentum,
|
||||
nesterov=nesterov,
|
||||
ns_steps=ns_steps,
|
||||
adjust_lr_fn=adjust_lr_fn,
|
||||
)
|
||||
self.adamw = optim.AdamW(
|
||||
[{"params": other_params, "weight_decay": 0.0}],
|
||||
lr=lr,
|
||||
betas=(0.9, 0.95),
|
||||
fused=True,
|
||||
)
|
||||
|
||||
self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]
|
||||
|
||||
@torch.no_grad()
|
||||
def step(self, closure=None):
|
||||
self.muon.step(closure)
|
||||
self.adamw.step(closure)
|
||||
|
||||
def zero_grad(self, set_to_none: bool = True):
|
||||
self.muon.zero_grad(set_to_none)
|
||||
self.adamw.zero_grad(set_to_none)
|
||||
|
||||
def state_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"muon": self.muon.state_dict(),
|
||||
"adamw": self.adamw.state_dict(),
|
||||
}
|
||||
|
||||
def load_state_dict(self, state_dict: Dict[str, Any]):
|
||||
self.muon.load_state_dict(state_dict["muon"])
|
||||
self.adamw.load_state_dict(state_dict["adamw"])
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
|
||||
parser = argparse.ArgumentParser(description="Train the AutoRegressiveLM model.")
|
||||
|
|
@ -144,35 +64,22 @@ def parse_args() -> argparse.Namespace:
|
|||
help="Max gradient norm for clipping.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--weight_decay",
|
||||
"--adamw_beta1",
|
||||
type=float,
|
||||
default=0.1,
|
||||
help="Weight decay (applied to Muon matrix params; non-matrix use 0).",
|
||||
default=0.9,
|
||||
help="Beta1 for AdamW optimizer.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--muon_momentum",
|
||||
"--adamw_beta2",
|
||||
type=float,
|
||||
default=0.95,
|
||||
help="Momentum factor for Muon optimizer.",
|
||||
help="Beta2 for AdamW optimizer.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--muon_nesterov",
|
||||
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.",
|
||||
"--adamw_weight_decay",
|
||||
type=float,
|
||||
default=0.01,
|
||||
help="Weight decay for AdamW optimizer.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--random_seed", type=int, default=3407, help="Random seed for reproducibility."
|
||||
|
|
@ -358,8 +265,21 @@ def create_model(config):
|
|||
return AutoRegressiveLM(config).to(dtype=torch.bfloat16)
|
||||
|
||||
|
||||
def create_optimizer(model, **kwargs) -> MuonMix:
|
||||
return MuonMix(model, **kwargs)
|
||||
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_scheduler(
|
||||
|
|
@ -390,6 +310,7 @@ def train(
|
|||
train_type: str,
|
||||
param_path: str,
|
||||
data_root_path: str,
|
||||
max_lr: float,
|
||||
n_epoch: int,
|
||||
batch_per_device: int,
|
||||
start_epoch: int,
|
||||
|
|
@ -402,7 +323,16 @@ def train(
|
|||
val_step: int,
|
||||
metrics: list[str],
|
||||
log_dir: str,
|
||||
dpo_beta: float,
|
||||
grpo_clip_eps: float,
|
||||
grpo_kl_coef: float,
|
||||
group_size: int,
|
||||
grpo_sync_interval: int,
|
||||
adamw_beta1: float,
|
||||
adamw_beta2: float,
|
||||
adamw_weight_decay: float,
|
||||
max_grad_norm: float,
|
||||
label_smoothing: float,
|
||||
random_seed: int,
|
||||
num_workers: int,
|
||||
pin_memory: bool,
|
||||
|
|
@ -423,7 +353,6 @@ def train(
|
|||
t_mult: int,
|
||||
stable_steps: int,
|
||||
decay_steps: int,
|
||||
**kwargs,
|
||||
):
|
||||
assert train_type in ["seq", "sft", "dpo", "grpo"]
|
||||
assert os.path.exists(param_path)
|
||||
|
|
@ -439,12 +368,12 @@ def train(
|
|||
window_size = config.max_len
|
||||
|
||||
strategy_kwargs = {
|
||||
"beta": kwargs.pop("dpo_beta"),
|
||||
"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"),
|
||||
"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,
|
||||
}
|
||||
|
||||
executor_kwargs = {
|
||||
|
|
@ -462,12 +391,11 @@ def train(
|
|||
|
||||
optimizer_fn = partial(
|
||||
create_optimizer,
|
||||
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"),
|
||||
**{
|
||||
"lr": max_lr,
|
||||
"betas": (adamw_beta1, adamw_beta2),
|
||||
"weight_decay": adamw_weight_decay,
|
||||
},
|
||||
)
|
||||
|
||||
total_steps = compute_total_steps(
|
||||
|
|
|
|||
5
setup.py
5
setup.py
|
|
@ -41,10 +41,7 @@ if _should_build():
|
|||
CUDAExtension(
|
||||
f"astrai.extension.{name}",
|
||||
info["sources"],
|
||||
extra_compile_args={
|
||||
"cxx": info["cxx_flags"],
|
||||
"nvcc": info["nvcc_flags"],
|
||||
},
|
||||
extra_compile_args={"cxx": ["-O3"], "nvcc": info["nvcc_flags"]},
|
||||
extra_link_args=[f"-Wl,-rpath,{_torch_lib}"],
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -10,11 +10,7 @@ from astrai.config.preprocess_config import (
|
|||
PipelineConfig,
|
||||
ProcessingConfig,
|
||||
)
|
||||
from astrai.preprocessing.builder import (
|
||||
MultiOutputMaskBuilder,
|
||||
SectionedMaskBuilder,
|
||||
SingleOutputMaskBuilder,
|
||||
)
|
||||
from astrai.preprocessing.builder import SectionedMaskBuilder
|
||||
from astrai.tokenize import AutoTokenizer
|
||||
|
||||
_SPECIAL_TOKENS_CONFIG = {
|
||||
|
|
@ -214,16 +210,6 @@ def builder():
|
|||
return SectionedMaskBuilder()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def single_builder():
|
||||
return SingleOutputMaskBuilder()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def multi_builder():
|
||||
return MultiOutputMaskBuilder()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tokenizer_dir(temp_dir, test_tokenizer):
|
||||
d = os.path.join(temp_dir, "tok")
|
||||
|
|
|
|||
|
|
@ -8,9 +8,7 @@ from astrai.config.preprocess_config import (
|
|||
)
|
||||
from astrai.preprocessing.builder import (
|
||||
MaskBuilderFactory,
|
||||
MultiOutputMaskBuilder,
|
||||
SectionedMaskBuilder,
|
||||
SingleOutputMaskBuilder,
|
||||
)
|
||||
from tests.data.conftest import (
|
||||
_CHAT_SECTIONS,
|
||||
|
|
@ -274,18 +272,12 @@ def test_sectioned_text_too_short(test_tokenizer, builder):
|
|||
|
||||
def test_factory_registered():
|
||||
names = MaskBuilderFactory.list_registered()
|
||||
assert "single" in names
|
||||
assert "multi" in names
|
||||
assert "sectioned" in names
|
||||
|
||||
|
||||
def test_factory_create():
|
||||
single = MaskBuilderFactory.create("single")
|
||||
assert isinstance(single, SingleOutputMaskBuilder)
|
||||
multi = MaskBuilderFactory.create("multi")
|
||||
assert isinstance(multi, MultiOutputMaskBuilder)
|
||||
sectioned = MaskBuilderFactory.create("sectioned")
|
||||
assert isinstance(sectioned, SectionedMaskBuilder)
|
||||
builder_obj = MaskBuilderFactory.create("sectioned")
|
||||
assert isinstance(builder_obj, SectionedMaskBuilder)
|
||||
|
||||
|
||||
def test_dpo_chat_basic(chat_tokenizer, builder):
|
||||
|
|
@ -375,59 +367,3 @@ def test_grpo_single_reward(chat_tokenizer, builder):
|
|||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert result["rewards"] == [0.9]
|
||||
|
||||
|
||||
def test_single_builder_matches_facade(chat_tokenizer, builder, single_builder):
|
||||
config = make_chat_config()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
]
|
||||
}
|
||||
facade_result = builder.build(item, config, chat_tokenizer)
|
||||
single_result = single_builder.build(item, config, chat_tokenizer)
|
||||
assert single_result == facade_result
|
||||
|
||||
|
||||
def test_single_builder_rejects_multi_config(chat_tokenizer, single_builder):
|
||||
config = make_dpo_chat_config()
|
||||
item = {
|
||||
"chosen": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
],
|
||||
"rejected": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "5"},
|
||||
],
|
||||
}
|
||||
assert single_builder.build(item, config, chat_tokenizer) is None
|
||||
|
||||
|
||||
def test_multi_builder_matches_facade(chat_tokenizer, builder, multi_builder):
|
||||
config = make_dpo_chat_config()
|
||||
item = {
|
||||
"chosen": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
],
|
||||
"rejected": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "5"},
|
||||
],
|
||||
}
|
||||
facade_result = builder.build(item, config, chat_tokenizer)
|
||||
multi_result = multi_builder.build(item, config, chat_tokenizer)
|
||||
assert multi_result == facade_result
|
||||
|
||||
|
||||
def test_multi_builder_rejects_single_config(chat_tokenizer, multi_builder):
|
||||
config = make_chat_config()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
]
|
||||
}
|
||||
assert multi_builder.build(item, config, chat_tokenizer) is None
|
||||
|
|
|
|||
Loading…
Reference in New Issue