ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: separate extension ops and backends

Browse Source
This commit is contained in:
ViperEkura
2026-08-16 21:15:52 +08:00
parent 3406157431
commit 6ac3b51496
18 changed files with 82 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/guides/inference.md
+1 -1
View File
@@ -90,7 +90,7 @@ Fallback: when `CudaBackend` cannot handle an input (wrong dtype or head_dim), `
### Rotary Embedding Backend
Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.py`, which auto-dispatches:
Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/backend/rotary.py`, which auto-dispatches:
- **CUDA kernel** (`rotary_emb.cu`): fused cos/sin lookup + rotation in a single kernel, used when the kernel is available, the input is bf16 on CUDA, and `torch.is_grad_enabled()` is `False` (inference mode)
- **Torch fallback**: complex multiply path (`torch.view_as_complex` → `torch.complex` multiply → `torch.view_as_real`), used during training (supports autograd backward) or when the CUDA kernel is not available