Commit Graph
24 Commits
Author SHA1 Message Date
ViperEkura cb60713a72 feat: add per-task throughput and latency metrics
- extract TaskTiming + MetricsCollector out of Task/TaskManager
- unify prefill/decode timing into single record() context manager
- expose avg_ttft_ms, avg_decode_tps, avg_e2e_latency_ms via /stats
2026-08-08 12:10:06 +08:00
ViperEkura 5c180cfa90 fix : handle zero-token batch generation
- return empty results without running inference for non-positive limits
- keep scheduler batch outputs aligned with requested max_tokens
- add engine and scheduler regression coverage
2026-08-06 12:31:09 +08:00
ViperEkura b2230fefd8 feat : add radix prefix cache
- replace hash-only lookup with page-granular radix matching
- keep partial pages private and cache only materialized KV prefixes
- integrate completed-request caching and add radix behavior tests
2026-08-06 11:45:52 +08:00
ViperEkura 654e6eb0d1 fix : correct prefill sampling and record alignment
- sample the first token from prefill logits without duplicating the prompt tail
- reject incomplete multi-output records before preprocessing alignment
- cover cached generation and partial DPO records with regression tests
2026-08-05 22:20:29 +08:00
ViperEkura 88751d0b08 refactor: share prefill+decode step between scheduler paths
- Extract _step() as the single prefill-group + task_extend + decode primitive
- _run_generation_loop and run_batch now both call it, so the two cannot drift
- run_batch now records prefix hashes (paged mode) and uses input order for
  decode, matching the loop thread
2026-08-03 13:45:27 +08:00
ViperEkura 50cfd0d555 perf: reduce decode overhead in scheduler and executor
- Precompute page_table and decode_mask on KVCache once per step in PagePool.bind_tasks, instead of per-layer in CudaBackend/TorchNativeBackend
- Skip frequency penalty history tensor construction when all penalties are 0 in Executor.execute_decode
- Omit FrequencyPenaltyStrategy from sampling pipeline when penalty is 0
- Deduplicate get_active_tasks calls in scheduler loop (3 to 1), remove redundant sorted() on decode tasks
- Benchmark (L20, bf16, CUDA backend): B=1 9.48->9.40ms (+1%), B=4 10.73->9.89ms (+8.6%), B=8 10.77->10.13ms (+6.4%)
2026-07-31 14:50:16 +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 53c804e233 refactor : merge max_prompt_len into max_seq_len, replace assert with raise
- Engine/Scheduler/TaskManager: merge max_prompt_len into max_seq_len
- train.py: replace bare assert with ValueError/FileNotFoundError
- server.py: add --max_seq_len CLI option
- engine.py: remove dead page_size param
2026-07-27 08:05:11 +08:00
ViperEkura 0c86c89af4 refactor : align config field names with Hugging Face
- dim -> hidden_size, n_layers -> num_hidden_layers
- dim_ffn -> intermediate_size, n_heads -> num_attention_heads
- n_kv_heads -> num_key_value_heads, max_len -> max_position_embeddings
- norm_eps -> rms_norm_eps, tie_weight -> tie_word_embeddings
- update model, inference, training, scripts, tests, docs
2026-07-20 22:05:31 +08:00
ViperEkura 06eeeead79 refactor: map instruction/input/output to chat roles
- RolloutGenerator._instruction_to_messages builds system/user/assistant list (instruction->system, input->user, output->assistant), replacing single-user-turn concatenation
- Remove _iter_samples helper; _prepare_prompts zips parallel list-of-strings fields directly per the collate_fn contract
- Tests adopt a system-aware chat template and pin the three-field role mapping
- Drop unused imports caught by ruff F401 (torch.Tensor in scheduler.py, iter_raw_records in pipeline.py, Tuple in evaluate_rouge.py)
2026-07-20 13:55:25 +08:00
ViperEkura 95c43368ae refactor: unify rollout onto inference engine KV-cache path
- RolloutGenerator now delegates prefill/decode to InferenceScheduler.run_batch (sync API, no background thread), sharing one KV-cache code path with the inference server and eliminating O(n^2) recompute in rollout
- Add sample(return_logprobs=) and Executor.execute_decode(return_logprobs=) to expose behaviour-policy log-probs through the engine; Task gains output_logprobs
- RolloutResult now subclasses RawRollout (adds rewards only), removing duplicated fields
- RolloutRunner.__call__ returns (result, is_fresh) instead of relying on object identity, removing the fragile refresh-detection contract
- Remove O(n^2) generate_responses helper and dead code (_tokenize_prompts, unused old_model arg)
- train_context.py wires InferenceScheduler directly instead of hand-rolling SamplingPipeline
- Tests: +11 covering return_logprobs, run_batch, and KV-cache-backed rollout semantics; 404 pass
2026-07-20 12:52:20 +08:00
ViperEkura 9d3ccfdffc fix: incremental decode to avoid U+FFFD in streaming
- StreamDecoder buffers incomplete multi-byte sequences
- Task.decode_new_token replaces per-token decode in scheduler
- flush_remaining emits final buffered text on task finish
2026-07-18 13:05:36 +08:00
ViperEkura a24a7b4da5 perf: merge decode batch for 10x throughput
- merge all active decode tasks into single forward pass (was grouped by next_pos)
- add per-task write_positions to ContiguousCacheView for correct KV writes
- override ContiguousCache.task_cached (base returned 0, caused prefill loops)
- add --cache_len/--frequency_penalty/--rep_window to generate.py
- chunked batch processing with tqdm progress

bench (1.2B model, 128 prompts, 64 tok, batch=128):
  before: 77.2s, ~111 tok/s
  after:   7.1s, ~1210 tok/s (10.9x)
2026-07-18 08:50:46 +08:00
ViperEkura 849e1e00a3 refactor: clean up inference design patterns
1. KVCache base: add default task_cached/task_record_hashes, remove getattr from scheduler
2. Remove page_size param from scheduler constructor (ContiguousCache-only)
3. InferenceEngine expose cache param for KVCache injection
4. Rename page_cache -> kv_cache in Executor
5. Move stream_callback from Task to TaskManager._callbacks dict
6. TaskManager.clear_queues clears callbacks
2026-07-05 11:41:54 +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 3e234c46f6 fix: 使用 threading.Event 替代裸 bool,补全公共 API
- scheduler 停止信号改用 threading.Event,跨解释器安全
- 移除 _fatal_error 和 check_health,异常仅用 logger.error 记录
- 补全 astrai/__init__.py,暴露所有主要模块
2026-06-18 15:38:35 +08:00
ViperEkura f521a30b22 fix : FSDP 优化器顺序、温度除零、调度器静默死亡、ref模型设备
- executor: use_orig_params 硬编码 True,FSDP 不替换 Parameter 对象
- strategy: DPO/GRPO ref 模型创建后移到 device
- sample: TemperatureStrategy clamp 1e-8,engine 验证改为 >0
- scheduler: 异常不 re-raise 避免 daemon 静默死亡,stop() 发回调给 waiting 任务
2026-05-29 21:57:44 +08:00
ViperEkura cb8dcb97ea refactor : 移除 -> None 返回值标注,拆分 FSDP 参数,新增 mmap 数据集存储
- 删除所有 def 函数 -> None 返回值类型标注
- FSDPExecutor 参数从 **kwargs 拆为显式声明,None 值自动过滤
- 新增 MmapStorage (bin) 存储后端,基于 numpy.memmap 零拷贝加载
- 新增 save_bin/load_bin/json_to_bin 工具函数
- detect_format 支持 bin 格式自动检测
2026-05-28 13:57:06 +08:00
ViperEkura 47c37e4876 refactor: 推理协议层重构为策略/建造者模式
- ProtocolHandler 改为具体类,格式化委托给 ResponseBuilder
- 新增 api/protocols/ 目录,含 OpenAIResponseBuilder、AnthropicResponseBuilder
- GenContext、StopInfo 参数对象替代 StreamContext
- 消除 Builder 的实例可变状态(accumulated、_yielded)
- SSE 工具和停止检测收归 ProtocolHandler 统一管理
- prepare() 方法合并原来的 build_prompt、create_response_id
- 参数校验去重:仅 GenerationRequest.init 负责校验
- Prefill 阶段提前短路完全命中的缓存任务
2026-05-26 00:12:57 +08:00
ViperEkura ff509ff39f fix: decode后task_extend失败时提前中止,scheduler崩溃时通知waiting任务 2026-05-20 19:23: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 e3382f6bb5 fix: 修复推理引擎 batch decode 中多项正确性与并发问题
- scheduler: decode 分组由幂次分桶改为精确 next_pos,消除 KV cache 位置错乱
- task: activate() 加锁操作 active_tasks,消除数据竞争
- engine: wait_completion 加超时,防止分配失败时永久死锁
- sample: TopKStrategy 向量化为 per-sample threshold,尊重各 task 的 top_k
- cache: Storage.write/gather 中 -1 页改用 mask 处理,防数据污染
- executor: prefill 逐 task 循环改为单次 tensor 调用
2026-05-14 21:31:39 +08:00
ViperEkura 205b40bd28 refactor: 重构 cache 和 inference 参数体系,分离存储与分配
- 合并 GenerationRequest/GenerationParams,统一 max_tokens 参数名
- PagePool/PrefixCache 分离为 Allocator + PrefixCache + PagePool
- 拆分 KV 存储为独立 Storage 类,PagedCache → KVCache,CacheView → KvcacheView
- Allocator.inc_ref 移除 LRU 防止竞争,Storage.write 增加负页防御
- Allocator/PrefixCache/TaskTable 加 threading.Lock 保证线程安全
- server.py uvicorn.run 改为传 app 对象修复导入错误
- benchmark.py 适配 KVCache 新 API
2026-05-14 20:05:08 +08:00
ViperEkura 2196c34c52 refactor: 重构 inference 模块架构,引入设计模式并分组文件
- 新增 protocol.py 协议层,Template Method 模式消除流/非流分支 45% 重复
- SSEBuilder 统一 SSE 构造,StopChecker 独立 stop_sequence 检测
- AnthropicHandler 追踪已产出文本,修复 stop 时重复 delta
- server.py 路由从约 100 行缩减至 3 行
- 拆分为 core/(cache/executor/scheduler/task)和 api/(protocol/server)
- 外部保持二级导入路径(from astrai.inference import Name)
- 删除所有分隔线注释,代码按语义自然分组
2026-05-14 17:42:37 +08:00