Commit Graph
63 Commits
Author SHA1 Message Date
ViperEkura 88c06db096 fix: resolve audited training and inference bugs
- reject prompts that encode to zero tokens in add_task instead of admitting a task whose prefill can never run, and surface empty-id run_batch calls as prompt_empty errors
- deliver the STOP stream callback when cancelling a live task so clients observe termination instead of hanging until socket timeout
- strip the torch.compile _orig_mod. prefix at every unwrap_model site and when loading checkpoints so FSDP state dicts and saved weights no longer leak the wrapper name into downstream keys
- reject online_* train strategies with nprocs > 1 at config validation time, explaining the NCCL all-gather deadlock they would otherwise hit mid-run
- apply the frequency penalty before temperature scaling (OpenAI semantics) so the penalty survives temperature=0 instead of being annihilated by the 1e8 logit blowup, and exclude penalty pipelines from the greedy fast path
- return logprobs from the raw pre-strategy distribution so they match training-side policy logprobs for PPO/GRPO importance ratios
2026-09-02 21:25:01 +08:00
0z5a e58a728b80 feat: version rollout weight updates
Track a monotonic policy version across optimizer steps, scheduler updates, and rollout results. Serialize synchronous generation with weight acknowledgements and invalidate reusable prefix KV entries so cached samples remain attributable to the behavior policy that generated them.
2026-09-02 19:01:41 +08:00
0z5a 01bcd0d105 perf: batch ragged prefill requests
- Pack prompts with a shared prefix start and attention backend into one forward.

- Select per-request final logits from cumulative query lengths.

- Cover ragged tokens, logprobs, scheduling, and documentation.
2026-09-02 15:00:16 +08:00
0z5a 9c3ef0c2a1 fix: cancel abandoned generation tasks
- Propagate stream closure and stop-sequence termination into scheduler cancellation
- Defer active KV release to the scheduler owner and close metrics safely
- Expose lifecycle counters and cover waiting, active, and allocation-race cleanup
2026-09-02 12:55:05 +08:00
0z5a 90de5bc1bd fix: report failed rollout requests
- Return structured finish and error reasons for synchronous generation
- Reject failed online rollout batches instead of training on empty responses
- Verify allocation and extension failures release metrics and KV state
2026-09-02 12:52:50 +08:00
ViperEkura 1cf7d6c76b perf: fill steady-state decode input ids via d2d copy
- add InferenceWorkspace.fill_input_ids_from_device copying device tokens straight into the fixed-address input_ids buffer
- cache each decode step's sampled tokens on-device in DecodeSteadyState.last_tokens; when the task signature is unchanged the next step reuses them, replacing the tolist -> python list -> elementwise host fill -> pageable h2d round-trip
- _sample_logits returns (host payload, device tokens); prefill discards the device tensor
- signature change (task join/leave/first decode) still takes the host path; both dispatch paths covered by tests

Benchmark: NVIDIA L20, BF16, 1B model + 0.11B test model (4 layers, hidden 512), contiguous KV cache, CUDA Graph, greedy, prompt 512, generation 256, engine decode via scripts/tools/benchmark.py (alternating A/B, 2-4 paired runs)
- 0.11B batch 32: 21429 -> 24415 tok/s mean (1.14x, +13.9%), 4/4 paired runs faster
- 1B batch 32: 4242 -> 4388 tok/s (1.034x, +3.4%), 7.54 -> 7.29 ms/step
- batch 1: no measurable change (<0.5%)
2026-08-31 14:24:51 +08:00
ViperEkura dcc96de12a test: refactor tests and fix inference edge cases
- Convert protocol and MoE test classes to plain functions
- Add real server/engine integration and generate_async tests
- Isolate test model per test and use pytest tmp_path
- Reset FastAPI engine state after inference tests
- Fix generate_async StopIteration handling on Python 3.12
- Fix HF adapter MoE dense/shared and Gemma qk_norm mapping
- Correct dev dependency httpx2 to httpx
2026-08-21 22:59:51 +08:00
ViperEkura 84753d3e08 refactor: deduplicate and restructure test suite
- extract preprocessing config factories into tests/data/factories.py
- keep conftest.py fixtures-only; stop importing builders from it
- promote temp_dir fixture to root conftest for cross-directory reuse
- unify duplicate BPE tokenizer builders into build_test_tokenizer
- merge grpo/dpo online e2e tests into one parametrized integration test
- extract engine mock factory and shared model batch builders
- drop local tempfile usage in favor of shared fixtures

No behavior change: 519 tests pass.
2026-08-20 01:53:28 +08:00
ViperEkura c79d34eee1 refactor: simplify training and inference interfaces
- avoid constructing model_fn more than once when reading config
- keep inference package exports focused on public entry points
- rename extra strategy arguments to strategy_kwargs
2026-08-19 20:55:13 +08:00
ViperEkura f252af495c refactor: remove dead code and deduplicate scheduler setup 2026-08-19 14:58:37 +08:00
ViperEkura 1bcd8f53ab perf: precompute ragged Q tile scheduling 2026-08-16 23:32:46 +08:00
ViperEkura 3406157431 refactor: standardize packed 3d inference
- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch
2026-08-16 13:24:02 +08:00
ViperEkura 3fb4b8ab13 perf: use int32 paged KV indices
- store page-table, request-row, and cache-location indices as int32
- preserve CUDA graph replay with bit-exact logits and KV cache coverage
- improve B=1 decode latency by 1-6% across 1K-32K contexts on L20
2026-08-15 13:17:06 +08:00
ViperEkura 47b3ed4e44 feat: propagate attention backend across scheduler threads
- InferenceEngine/Scheduler accept an explicit backend
- capture request-level attn_backend context onto Task
- split prefill/decode batches by backend instance
- ASTR_BACKEND env overrides ContextVar as process-wide policy
- report resolved backend and CUDA-graph state in benchmark
2026-08-09 13:32:40 +08:00
ViperEkura be90dfe2bd fix: isolate continuous batch decode state
- Match steady-state metadata to the active task IDs
- Rebuild request mappings for cached prefix pages
- Add regressions for batch refill and prefix reuse
2026-08-09 01:01:41 +08:00
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 3fa7e66676 refactor: decouple task cache from PagePool and unify steady-state detection
- TaskCacheRegistry -> TaskCacheManager (independent, held by scheduler)
- TaskCacheState co-locates 5 parallel dicts into one dataclass
- AllocationStrategy base class + PagedStrategy subclass (page_size is a parameter)
- _rollback() helper for unified cleanup (no duplicate free paths)
- Task._kv_len + prefill_done property (explicit, no output_tokens proxy)
- Steady-state detection single-sourced in TaskCacheManager.bind()
- PagePool is now pure physical layer (no task knowledge)
- Removed dead _page_to_hash dict in RadixCache
2026-08-08 22:43:06 +08:00
ViperEkura 1b1f1a0707 fix: add dtype guard to FlashAttnBackend capability check
- _backend_supports now rejects fp32 for FlashAttnBackend (flash-attn only supports fp16/bf16), preventing runtime crash on fallback chain
- rename test_default_backend_is_torch_native to reflect multi-backend reality
- scheduler test fixture uses bf16 model (matches production, avoids unnecessary 3-step fallback chain)
2026-08-07 23:08:22 +08:00
ViperEkura 02469887f5 refactor: simplify inference engine and backend dispatch
- merge _generate_streaming/_generate_non_streaming into single _generate() with stream flag
- delete dead GenerationRequest class and generate_with_request method
- inline _next_token helper into generate_async
- replace flash-attn double-checked locking with functools.lru_cache
- extract _write_and_gather_kv helper shared by TorchNative/FlashAttn backends
- inline _kv_cache_is_contiguous into its sole call site in FlashAttnBackend
- change default backend priority from flash>cuda>torch to cuda>flash>torch
- add ASTR_BACKEND env var to override default backend at resolve time
- add supports_graph() static method to AttentionBackend ABC, override in CudaBackend
- replace isinstance(get_backend(), CudaBackend) with get_backend().supports_graph() in executor
- add torch.cuda.is_available() guard to CudaBackend.supports()
2026-08-07 22:28:48 +08:00
ViperEkura 6f67ba8942 perf: move decode split partials to InferenceWorkspace
- Replace per-.cu-file static cached tensors with workspace-managed pre-allocated buffers

- InferenceWorkspace now owns decode_o_part / decode_ml_part (mirrors FlashInfer's workspace pattern)

- KVCache carries the buffers through the backend -> C++ kernel chain

- C++ kernels accept optional pre-allocated buffers; fallback to alloc_split_partials for backward compat

- Pre-allocates once at Executor init, zero allocation in the decode hot loop

- Prerequisite for CUDA-graph capture (all kernel addresses are stable)
2026-08-06 19:12:09 +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 a03504a280 perf: preallocate inference decode buffers
- add InferenceWorkspace with fixed-shape per-step buffers (input_ids, decode mask, KV bind metadata) for CUDA-graph capture
- bind_tasks derives seq_lens from the pool's own _task_len tracking, dropping the seq_lens parameter
- update decode metadata in-place (position_ids, seq_lens, kv_indptr) instead of re-allocating per step
- task_extend advances _task_len in contiguous mode so the pool tracks current length
- skip log_softmax when logprobs are not requested
2026-08-03 00:55:26 +08:00
ViperEkura a27c8a819d test: prune low-value and duplicate tests
- Remove tautological test_trainer assertions that never trained
- Drop grpo isfinite-only smokes and merge frozen-model checks via parametrize
- Merge duplicate tool_parser cases (find/streaming/factory) with parametrize
- Collapse duplicate dataset store/detect_format tests
- Remove misleading scheduler/task tests that asserted the opposite of their names
- Merge signal-handler SIGTERM/SIGINT into one parametrized case
- Drop cross-file grpo strategy duplication kept in online_strategy
2026-08-01 16:01:20 +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 115192c67c refactor: remove H5 storage backend in favor of mmap bin
- Remove H5Store, H5Writer, save_h5/load_h5 and h5py dependency
- MmapStore (bin) is the sole pre-tokenized storage backend
- Move setup_logging after imports to fix E402 in __init__.py
- Clean up unused imports across test files
- Move inline test imports to file top
2026-07-29 12:50:27 +08:00
ViperEkura 5ba21f4eb3 refactor: eliminate test duplication via shared helpers
- Add tests/helpers.py with shared config, dataset, tokenizer, executor, and assertion helpers
- Replace 15 copies of device one-liner with session-scoped fixture
- Collapse 5 near-identical Dataset subclasses into RandomTokenDataset
- Remove duplicate _make_config/_make_model/_make_frozen and FakeTokenizer/FakeExecutor definitions
- Make test_callbacks and test_early_stopping use existing train_config_factory
- Replace 6 duplicate meta.json read blocks with load_shard_meta
- Fix mkdtemp leaks in test_lora.py with TemporaryDirectory
2026-07-27 22:34:53 +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 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 d08a92c7bd feat: add frequency penalty to inference sampling pipeline
- Add FrequencyPenaltyStrategy (logit -= penalty * count)
- Per-task rep_window for penalty history lookup
- Wire through engine, task, executor, API layer
- Add --frequency_penalty and --rep_window to stream_chat.py
- 9 unit tests for frequency penalty strategy
2026-07-17 21:28:31 +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 25d4ea3f91 refactor : 压缩测试代码,消除重复
- fixture 替代重复实例化和 tokenizer 落盘
- parametrize 合并同构测试
- helper 消除 save_h5 + DatasetFactory.load 样板
- 净减 272 行
2026-06-19 14:54:39 +08:00
ViperEkura 9e31d4ef2b feat : BaseToolParser.feed 增加可选 token_ids 参数
- format_chunk ABC 改为 (token, **kwargs),body/token_ids 通过 kw 传入
- ProtocolHandler._handle_stream 逐 token encode 并透传
- Anthropic builder 用 **kwargs 吸收不使用的参数,零变更
- 新增 3 个 token_ids 参数测试
2026-06-06 11:19:30 +08:00
ViperEkura 52aa4d01d5 feat : 推理层增加 vLLM 风格工具调用解析
- 新增 BaseToolParser 抽象基类,定义 feed/parse_complete 流式接口
- 新增 SimpleJsonToolParser,解析 {"name":"...","arguments":{...}} 格式
- 新增 ToolParserFactory,基于 BaseFactory 实现可插拔注册
- 集成 parser 到 OpenAIResponseBuilder,支持流式/非流式工具调用
- 扩展 ChatMessage 和 ChatCompletionRequest,增加 tools/tool_choice 字段
- 重构 format_chunk 接口,传入累积文本支持全量重新解析
- 新增 74 个单元测试,覆盖扫描/查找/流式解析/完整解析/工厂
2026-06-06 08:54:10 +08:00
ViperEkura dc7d2cfbca refactor : FastAPI 懒加载单例,消除模块级副作用
- import astrai.inference 不再在模块加载时创建 FastAPI 实例
- 路由移至 APIRouter;get_app() 首次调用时懒构造单例
- _create_engine 和 run_server 的 param_path 改为必填
- 更新测试改用 get_app() 替代模块级 app
2026-06-04 15:52:27 +08:00
ViperEkura dd1b39f435 fix: ProgressBar默认输出到stdout
- file参数默认值改为None, 内部用 or sys.stdout 兜底
- 清理inference API中未使用的import (Optional, time, field)
- 删除test_protocol中未使用的ctx变量
2026-05-26 13:27:05 +08:00
ViperEkura 94d6e713e9 test: 补充推理协议层单测覆盖
- StopChecker、GenContext、StopInfo 单测
- OpenAIResponseBuilder / AnthropicResponseBuilder 全部方法
- Anthropic 停止序列裁剪逻辑(含 unyielded 边界)
- GenerationRequest 参数校验含负值边界
- Scheduler prefill 短路验证
2026-05-26 00:21:52 +08:00
ViperEkura 785d65436c fix: 修复 to_dict list 类型丢失与 OpenAI stop 参数失效
- to_dict() 增加 list 类型序列化支持,metrics 等字段不再丢失
- OpenAIHandler 补充 get_stop_sequences/on_token,读取 request.stop 并检测停止序列
- 文档类图补充缺失字段、修正关系分类、ChatCompletionRequest 字段增加 Optional
2026-05-19 21:07:07 +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
ViperEkura 6d6ef99e66 perf: 消除 PagedCache.write 中的 position_ids GPU 同步,解码提速 15%
- CacheView.write 用 total_len - k.size(1) 推导 start_pos,替代 position_ids[0,0].item()

- 移除 GQA/MLA/DecoderBlock 中不再使用的 position_ids 参数

- PagedCache.write 参数 position_ids:Tensor → start_pos:int
2026-05-14 15:37:48 +08:00
ViperEkura c0effc9f5b refactor: 位置编码改用 position_ids [B,S],简化 attention mask 构建
- RotaryEmbedding/CacheView 接受 position_ids 替代 start_pos

- process_attention_mask 用 position_ids >= arange 做逐位置 causal

- 训练/无 KV cache 时 position_ids=None 内部自动处理

- 移除 executor/benchmark 中冗余的 input_mask 构造
2026-05-14 13:26:31 +08:00
ViperEkura df0845e916 chore: 解耦 Executor/Scheduler/TaskManager,修复 stop 页泄漏,移除 ServerState 全局单例 2026-05-12 13:47:55 +08:00
ViperEkura 7440e9c809 style: 重命名 test_scheduler_concurrency 为 test_scheduler 2026-05-12 12:24:36 +08:00
ViperEkura 7d4029c2a4 test: inference 模块补全单元测试,cache/sample/engine/task
- test_cache: page_hash, PagePool, PrefixCache, TaskTable, PagedCache write/gather
- test_sample: TemperatureStrategy, TopKStrategy, TopPStrategy, SamplingPipeline, sample()
- test_engine: _Result 线程安全, generate stream/non-stream batch/single
- test_task: Task 生命周期, TaskManager 队列操作
- 4 新文件, +771 行, 116 total tests
2026-05-12 12:17:57 +08:00
ViperEkura 5203b7f53e perf: 测试优化,model 改为 session 共享,scheduler 用 Event 替代 sleep
- 拆出 session-scoped test_tokenizer + test_model,14 次创建 → 1 次
- 删除无用 test_env fixture
- 固定模型维度,消除随机性
- 添加 pytest markers 配置
2026-05-12 11:35:18 +08:00