ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
718 Commits 1 Branch 14 Tags
a5a3cc1fc2f068efa998872f45ed343d0f4deeb5
Commit Graph
14 Commits
Author SHA1 Message Date
ViperEkura 0c1b7664c1 refactor: split infer core into subpackages by concern 
- Eliminate core/ directory into cache/, runtime/, network/ subpackages plus flat modules
- Split cache.py (647 lines) into cache/{buffer,strategy,pool}.py by layer
- Add explicit ContiguousStrategy, make AllocationStrategy a real ABC
- Move TaskCacheState to cache/strategy.py, drop string forward references
- Rename api/ to network/, server.py to app.py
- Move sample.py into runtime/ alongside executor and graph
- Simplify TaskCacheManager.__init__ to single pool param
- Expose pool.strategy and pool.req_pool as public properties
- Fix KVCache import in attention_backend.py (TYPE_CHECKING guard)
- Fix steady-state decode reading uninitialized position_ids on first step
 2026-08-08 23:43:05 +08:00
ViperEkura 3e67b4f88d perf: add fused CUDA rotary embedding kernel 
- Single-kernel rotary embedding (cos/sin lookup + rotation) replaces PyTorch complex-multiply path (3 kernel launches + f32 upcast per call)
- RotaryEmbedding now stores cos_table/sin_table and returns (cos, sin) f32 tuple instead of a complex tensor
- apply_rotary_emb in rotary_backend.py auto-dispatches: CUDA kernel if available, else torch complex-multiply fallback; backend-agnostic (both attention backends benefit)
- Kernel: 256-thread blocks, grid-stride loop, vectorized __nv_bfloat162 load/store, f32 compute, bf16 out
- Standalone kernel 6-9x faster than torch across decode/prefill shapes, max diff 0 (decode) to 3e-2 (large prefill, bf16)
- Benchmark (L20, bf16, CUDA backend): B=1 9.48->7.25ms (+31%), B=4 10.73->7.67ms (+40%), B=8 10.77->7.81ms (+38%), B=16 10.79->7.83ms (+38%)
 2026-07-31 15:27:31 +08:00
ViperEkura 3067a8e1a6 feat: unify attention backend with multi-dim mask support 
- Add attention() functional entry delegating to active backend
- GQA/MLA forward calls attention() instead of inline cache/SDPA
- CUDA kernels support 2D/3D/4D mask via mask_h_stride field
- CudaBackend.fwd_decode builds 2D padding mask for mixed seq_lens
- KVCache.max_len precomputed in bind_tasks to avoid GPU sync
- batch==1 decode short-circuits mask=None
- Split tests into conftest, test_backend, test_backend_equivalence, test_kernel_mask
- 440 tests pass, L20 decode 1.44-1.60x speedup vs torch native
 2026-07-30 20:38:34 +08:00
ViperEkura df979b4469 refactor: use single-index access and update docs for cache architecture 
- Replace all buffer[layer_id][loc] double indexing with buffer[layer_id, loc] single advanced indexing in cache.py and attention.py
- Revert KVStorage buffers back to 4D [n_layers, size, n_kv_heads, head_dim], remove leftover 3D reshape/view in MLA path
- Update docs/guides/inference.md, docs/developer/internals.md, docs/developer/architecture.md to reflect new PagePool/KVStorage/ReqToTokenPool/KVCache classes
 2026-07-30 17:47:04 +08:00
ViperEkura deb2d7e127 refactor: rebuild KV cache with three-layer separation architecture 
- Replace CacheView/ContiguousCache/PageCache with SGLang-inspired design: KVStorage (flat token-level NHD buffers [n_layers, size, H, D]), ReqToTokenPool (index table [req_idx, pos] -> token_slot), Allocator + PrefixCache (slot allocation with LRU and prefix sharing)
- Add KVCache as pure dataclass passed to model: k_buffer, v_buffer, req_to_token, req_pool_indices, seq_lens, out_cache_loc
- PagePool orchestrates all three layers, supports contiguous mode (pre-allocated per-request blocks, default) and paged mode (page_size=1 or >1 with dynamic allocation and prefix caching)
- Attention layers now do raw buffer indexing instead of opaque write/gather method calls on CacheView objects
- Update executor.bind_tasks signature: seq_lens list + start_pos
- Rename paged_cache -> kv_cache throughout model/ and inference/
 2026-07-30 17:19:06 +08:00
ViperEkura d7ac66fb73 refactor: simplify attention mask handling 2026-07-20 20:36:16 +08:00
ViperEkura 5416c2e8fb perf: replace paged KV cache with contiguous ContiguousCache, decode all groups 
- Add KVCache/CacheView abstract base classes in cache.py
- Add ContiguousCache (contiguous per-slot buffer, default) alongside PageCache (paged, renamed from old KVCache)
- Merge make_table_tensor + bind into bind_tasks on KVCache interface
- Remove task_cached/task_record_hashes from base class (PageCache-only)
- Scheduler: decode all position groups instead of just the largest (eliminates 63% group skip rate)
- Scheduler: accept optional cache param for swapping implementations
- Model layer type hints use CacheView base class
- Batch 1-32: 1-7% speedup from eliminating Storage.gather overhead
- All 183 inference tests pass
 2026-07-05 11:34:36 +08:00
ViperEkura 88ec63121d feat : GPT-2 residual scaling weight init 
- Linear: normal(0, init_std) replaces kaiming_uniform_(a=sqrt(5))
- o_proj / mlp.down: init_std = 0.02 / sqrt(2 * n_layers)
- MoE: expert down scaled by 1/sqrt(1/n_shared + 1/K)
- Embedding: normal(0, 0.02), unchanged
 2026-06-25 15:08:31 +08:00
ViperEkura 6ae1828449 refactor : 清理工厂和配置系统中的死代码与冗余抽象 
- 删除 Registry 中未使用的 category/priority 字段，_entries 简化为直接存储类引用
- 修正 __init_subclass__ 避免叶子类（AutoRegressiveLM 等）创建空注册表
- 删除 5 个工厂的薄 create() 覆写，统一使用 BaseFactory.create(name, *args, **kwargs)
- 删除 3 处零调用的 available_types/available_strategies 别名死代码
- 删除零调用的 BaseModelConfig.to_file 死代码
- 将 BaseConfig.from_json/to_json 重命名为 from_file/to_file，消除与子类重复
- 移除两个 inference builder 中总是被覆写的 prompt_tokens=0
 2026-06-07 11:39:50 +08:00
ViperEkura a44fd22a99 fix: 修复训练与模型参数传递问题 
- state_dict_fn 传入 CheckpointCallback，修复多卡 DDP 下 key 前缀丢失
- MLA 增加 use_qk_norm 支持，消除参数静默丢失
- moe_topk_method 统一命名为 topk_method
- checkpoint 回调移至最前
 2026-05-17 11:20:13 +08:00
ViperEkura 48a53121ba refactor: 工厂 kwargs 过滤及组件参数清理 
- BaseFactory.create() 按 __init__ 签名过滤多余 kwargs
- 移除 GQA/MLA/MLP/DeepSeekMoE 中多余的 **kwargs
- MLP/DeepSeekMoE 参数名统一为 dim_ffn
- scheduler max_seq_len 增加 None 显式判断
- 默认 max_prompt_len 提升至 2048
 2026-05-16 16:47:41 +08:00
ViperEkura 0ba8c70ce1 fix: 修复 MLA 多个 bug 并缩小测试模型参数 
- MLA kv_b_proj 输出维度和 q_rope 切分偏移修复
- 打通 MLA 配置从 ModelConfig 到 DecoderBlock 的传递路径
- rope_theta 配置不再被忽略，MLA 使用 qk_rope_head_dim
- tie_weight 使用 is True 避免 None 隐式生效
- norm_eps/rope base 类型标注修正
- 测试模型参数缩小 (dim=8, head_dim=4)
- 新增 6 种架构配置 × 2 场景的前向传播测试
 2026-05-16 14:57:43 +08:00
ViperEkura e12f1a7ee5 feat: BaseModelConfig + DeepSeekMoE + 工厂模式替代 if/else 
- BaseModelConfig: fields() 精确字段匹配 + 类型矫正 + 未知key警告
- DeepSeekMoE: 共享专家 + 路由专家 + top-K 门控
- AttnFactory/FFNFactory: 装饰器注册，DecoderBlock 零分支
- config 用 attn_type/ffn_type 驱动组件选择
 2026-05-15 20:34:52 +08:00
ViperEkura ef25efffa2 refactor: 拆分 module.py 为 components 子包 
- rope/linear/norm/embedding/mlp/attention/decoder_block 各自独立文件
- 依赖单向无循环
- 公开接口不变，外部无需修改
 2026-05-15 20:08:36 +08:00