3 Commits

Author	SHA1	Message	Date
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 (2*BC + WARPS*BR)*LD to (2*BC + 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