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

docs: sync all documentation with current codebase

Browse Source 
- remove GenerationRequest and generate_with_request references (class deleted)
- document cuda>flash>torch default priority and FlashAttnBackend
- add ASTR_BACKEND env var to backend docs, TorchNativeBackend (default) → (fallback)
- fix JsonlStore transform routing → DatasetFactory ownership
- fix CudaBackend fallback chain description (FlashAttn → TorchNative)
- add FlashAttnBackend to architecture strategy table
- add router_stats to DecoderOutput/FFNOutput TypedDict diagrams
- add decode_o_part/ml_part/decode_out to KVCache diagram
- add --append_eos/--no-append_eos to IFD evaluation parameter table
- update get-started CUDA kernel note (no longer requires explicit attn_backend activation)
- fix python -m scripts.tools.server (no __init__.py) → direct script call
This commit is contained in:
ViperEkura
2026-08-07 23:52:27 +08:00
parent f163520fff
commit 056c1382ff
8 changed files with 39 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/developer/internals.md
+4 -2
View File
@@ -178,8 +178,10 @@ Three-layer separation (SGLang-inspired):
Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/attention_backend.py`):
- **`TorchNativeBackend`** (default): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
- **`CudaBackend`**: decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool). Falls back to `TorchNativeBackend` when kernel unavailable.
- **`CudaBackend`** (default): decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool). Falls back to `FlashAttnBackend` when dtype unsupported.
- **`FlashAttnBackend`**: optional flash-attn dispatch with `flash_attn_with_kvcache` fast path for contiguous cache; falls back to KV gather + `flash_attn_func`.
- **`TorchNativeBackend`** (always-available fallback): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
- Default priority: cuda > flash > torch. Set `ASTR_BACKEND=cuda|torch_native|flash` to override.
Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.py`, which auto-dispatches to the fused CUDA kernel (`rotary_emb.cu`) during inference or torch complex multiply during training (for autograd compatibility). Both attention backends share the same rotary dispatch.