- move setup_logging to astrai/logging.py
- timed() now uses logger.isEnabledFor(DEBUG) instead of separate env var
- enable ASTR_LOG_LEVEL=DEBUG to see per-step timing logs
- call setup_logging() in stream_chat.py
- 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
- move JSONL transform auto-creation from JsonlStore.load to DatasetFactory.load via _build_jsonl_transform helper
- remove TokenizeTransform and PipelineConfig imports from storage module
- JsonlStore.load now requires explicit transform= for eager mode
- DatasetFactory.load remains the public API with identical convenience behavior
- add _BindState dataclass and _is_steady_increment() to cache.py
- replace _bind_sig/_bind_seq_lens dual fields with single _bind_state
- replace DecodeSteadyState bare tuple with named dataclass
- use _is_steady_increment() in both PagePool.bind_tasks and Executor.execute_decode
- 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()
- Pre-allocate decode_out in InferenceWorkspace so attn_paged_decode does not call torch::empty inside graph capture
- Run live forward before graph capture to get valid output (graph pool memory is zeroed after capture block exits)
- Greedy generation with graph replay is bit-exact across all batch sizes
- _warmup_cuda_graphs pre-captures graphs at init for power-of-two batch sizes
- Graph enabled only when CudaBackend + supported head_dim + warmup succeeds
- Decode speedups vs no-graph: B=1 2.09x, B=4 1.80x, B=8 1.94x, B=16 1.76x
- Pre-allocate decode_out in InferenceWorkspace so attn_paged_decode does not call torch::empty inside graph capture
- Wire decode_out through KVCache, PagePool.bind_tasks, and CudaBackend.fwd_decode
- Run live forward before graph capture to get valid output (graph pool memory is zeroed after capture block exits)
- Greedy generation with graph replay is bit-exact across all batch sizes
- Decode speedups vs no-graph: B=1 2.09x, B=4 1.80x, B=8 1.94x, B=16 1.76x
- Decode with contiguous cache uses flash_attn_with_kvcache instead of materializing full KV via gather + flash_attn_func
- _backend_supports allows FlashAttnBackend for decode (q_len==1) even with explicit mask
- Decode speedups vs TorchNative (B=1,4,8,16 mean): cuda 1.55x, flash 1.40x, torch_native 1.00x
- Read K/V directly from flat pool via cache_batch_idx + cache_seqlens, zero-copy view reshape
- Remove unused local variable b in attention_backend.py
- Remove unused variable rank0_sd in test_broadcast_state_dict.py
- Remove unused imports across test files
- New CudaGraphContext class: warmup -> capture -> replay lifecycle
- One graph per batch_size key, all inputs at fixed workspace addresses
- Added position_ids buffer to InferenceWorkspace (required for graph capture)
- Graph only activates when CUDA backend is the current backend
- Default off (opt-in) due to slight numerical divergence in graph replay
- Sampling stays outside the graph (torch.multinomial uses mutable RNG)
- Resolved circular import: KVCache -> TYPE_CHECKING in attention_backend.py
- 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)
- Preload V into shared memory alongside K to eliminate per-element KV address lookups in the inner softmax/accum loop (doubles smem)
- Cache split-KV partial tensors (o_part, ml_part) with static tensors instead of per-call allocation in both decode and paged-decode paths
- Force is_causal=True in CUDA decode backend (decode is always causal)
- 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
- add FlashAttnBackend (ATTN_BACKEND.FLASH) using flash_attn_func with KV-cache gather + GQA, mirroring TorchNativeBackend
- add flash_attn_available() probe gated on compute capability plus a real-kernel smoke test, cached at first use
- lazy-import flash-attn via importlib so it stays an optional dependency, raising clear errors when unusable
- add 'flash' optional extra (flash-attn>=2.6) and export the new backend
- 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
- hoist prefill qo_indptr into the workspace so CudaBackend.fwd_prefill does not rebuild it per layer
- cache has_freq in SamplingBatchInfo to drop the per-step GPU any() sync
- drop pin_memory host staging for input_ids; sync copy suffices for a small batch
- 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
- SamplingBatchInfo: sample params built once per task set (top_k int32, pinned async H2D)
- position_ids advances by +1 on steady-state decode instead of re-building
- DecodeBindCache: bind_tasks increments seq_lens/kv_indptr, reuses req_pool_indices
- saves ~240us of python/launch overhead per decode step
- Converts detached strategy metrics before returning loss output
- Removes redundant item conversion from the trainer loop
- Updates the documented contract and regression tests
- Propagates MoE load-balancing loss through model outputs
- Logs task, auxiliary, and weighted losses across strategies
- Computes only explicitly requested callback metrics
- Preserves tensor compute_loss API and adds regression tests
- bind_tasks builds kv_indptr (prefix sum of seq_lens) a single time
- fwd_decode/fwd_prefill reuse it instead of rebuilding per layer
- Removes 24 cumsum launches per decode step (was ~1ms/step at B=4)
- Decode B=4: 9.60 -> 7.82 ms/step (-18.5%), +22.8% tok/s
- CudaBackend.fwd_decode passes attn_mask directly instead of kv_cache.decode_mask
- TorchNativeBackend derives pos_mask from attn_mask[:,0,0] on decode
- Drop decode_mask and page_table fields from KVCache and bind_tasks
- ManoAdamW now uses torch.optim.AdamW(fused=True, betas=(0.9, 0.95)) matching MuonAdamW, eliminating a confounding variable in optimizer comparison experiments
- only NoraNAdamW retains NAdamW, which is correct per the Nora paper design
- add GradSNRTracker to metric_util.py computing SNR = E[g]^2 / Var(g) via per-parameter EMA moments
- add grad_snr_tracker field to TrainContext (instantiated by default)
- register grad_snr in MetricCallback, update tracker on each optimizer step before metrics are recorded
- add grad_snr to default --metrics in train.py CLI
- implement Mano (v2) with axis-rotating tangent projection and manifold normalization, replacing Newton-Schulz iteration
- composite ManoAdamW reuses partition_optimizer_parameters and composite helpers
- register mano_adamw in OptimizerFactory, export Mano and ManoAdamW
- add --mano_momentum and --mano_nesterov CLI options in Optimizer group
- add mano_adamw hyperparameters branch in train.py
- document mano_adamw in params.md
- add tests for single-step projection, axis alternation, factory registration, closure, and resume
- add astrai/optim/composite.py with shared step/zero_grad/state_dict/param_groups helpers and OptimizerFactory
- rename MuonMix to MuonAdamW (matches registered name muon_adamw) and file to muon_adamw.py
- use @OptimizerFactory.register decorator in each optimizer module instead of post-import registration in __init__
- fix closure being invoked once per sub-optimizer in MuonAdamW.step (now exactly once via composite_step)
- NoraNAdamW.step now forwards closure correctly
- alloc_split_partials now uses torch::empty: the split kernel writes every slot it owns, so the per-call zeros/full memset was pure overhead (2 kernels per layer per step)
- decode split-KV MMA kernels now run a true multi-stage cp.async pipeline (wait_group<STAGES-1> instead of wait_group<0>), keeping STAGES-1 tile loads in flight; the old wait_group<0> serialized load and compute so deeper STAGES made no difference
- add a fallback path when ntiles < STAGES to avoid a race on the last tile
- apply_rotary_emb now checks torch.is_grad_enabled() before dispatching to CUDA kernel
- Training (grad enabled) uses torch complex multiply path which supports autograd backward
- Inference (inference_mode/no_grad) uses CUDA kernel as before
- Without this fix, training backward would crash — the CUDA kernel has no autograd backward()
- CUDAExtension module names changed from astrai.extension.<name> to astrai.extension.lib.<name>
- Compiled .so files now land in astrai/extension/lib/ instead of alongside Python source
- loader.py imports from .lib.<name> subpackage
- Add astrai/extension/lib/__init__.py to make lib a proper package
- Separates compiled artifacts from Python source for cleaner directory structure
ViperEkura
Gaolingx
QueenAmish