AstrAI/astrai/extension
ViperEkura 57729fd92d refactor: stride-based attn interface with layout and causal mask
- Replace is_causal + causal_offset with unified causal_offset (-1 = off, >=0 = first Q pos)
- Causal and mask can now coexist (was mutually exclusive)
- Add stride-based addressing for Q/KV/O (layout-agnostic, zero-copy)
- Add layout param ("bhld"/"blhd") parsed in Python, passed as int to C++
- Support 2D [batch, kv_len] and 3D [batch, q_len, kv_len] mask
- Vectorize paged KV gather in Python fallback (was per-token Python loop)
- Extract shared helpers: compute_num_splits, alloc_split_partials, dispatch_head_dim
- Unify paged_decode entry via attn_pack_paged_params
- Update mma_softmax_tile for 3D mask with per-row qrow indexing
2026-07-14 21:34:42 +08:00
..
__init__.py refactor: stride-based attn interface with layout and causal mask 2026-07-14 21:34:42 +08:00
loader.py feat: wire up paged decode CUDA kernel to Python extension 2026-07-11 18:40:49 +08:00
ops.py refactor: stride-based attn interface with layout and causal mask 2026-07-14 21:34:42 +08:00