Commit Graph
340 Commits
Author SHA1 Message Date
ViperEkura 69f35c46e0 fix: quantize amax from raw values, not scaled fp8 values
- amax for delayed scale was the quantized max (always ~448), so scale collapsed to 1
- this made fp8 gradients diverge (cosine 0.05) and training stall
- stop w/x transpose-quantize amax from polluting the grad scale
2026-08-14 14:26:16 +08:00
ViperEkura 0378e62e17 refactor: split fp8 into fp8_ops adapter and fp8 policy module
- fp8_ops is the only module touching the pybind (kernel interface)
- fp8.py keeps scaling state, delayed amax and aten::linear dispatch
- remove circular imports between old fp8_ops/fp8_state/fp8_dispatch
2026-08-14 12:25:29 +08:00
ViperEkura 5244f1a8fc feat: add te-style scaled fp8 training via fp8_autocast
- per-tensor scales applied inside cublasLt via A_SCALE/B_SCALE
- delayed scaling: weight amax history ring, refresh every 16 steps
- quantize kernels emit atomic amax, device-side scale updates
- fp8_autocast context toggles aten::linear dispatch like torch.autocast
- fallback to bf16 when M/N not 16-aligned (fp8 gemm constraint)
- x/g scales delayed one step, reuse free atomic amax (no abs/max reduce)
2026-08-14 12:14:04 +08:00
ViperEkura 15862d4b56 perf: fuse fp8 linear fwd and bwd into single kernel calls
- fp8_linear_forward: cast + cublasLt GEMM + transpose + bias in one call
- fp8_linear_backward: scale-free, dtype derived from input tensor
- drops per-op Python dispatch (was ~6-8 launches per linear) and amax syncs
- 1024x1024 linear: 6.8x slow -> 0.67x (36.7us vs 24.8us bf16)
- small-model e2e still 1.71x slow; 15bt estimate ~0.78x (linear-heavy)
2026-08-14 01:24:37 +08:00
ViperEkura c6a82a5029 refactor: align linear backward dtype with weight
- cast gradients and inputs to weight.dtype instead of hardcoded bf16
- single code path covers bf16 and fp32 models, no branch needed
- gradient dtype now matches the leaf parameter dtype exactly
2026-08-14 01:01:58 +08:00
ViperEkura a5b238dd86 feat: add fp8 training via cublasLt dispatch
- fp8_mm kernel (csrc): cublasLt fp8 e4m3 gemm, TN layout mapped zero-copy
- custom::fp8_mm custom op: meta/cuda/cpu kernels + scale-corrected bf16 autograd
- aten::linear and linear_backward dispatch on CUDA key, zero model changes
- per-tensor scale or raw cast; single-GPU smoke loss matches bf16
2026-08-14 00:39:49 +08:00
ViperEkura c5fba9c238 perf: flatten paged prefill tile dispatch
- remove the host-provided max_q_len argument
- dispatch only the ragged prefill tile upper bound
- validate the rebuilt CUDA backend end to end
2026-08-09 23:12:53 +08:00
ViperEkura cd31f1f62f refactor: tidy attention params and launcher interfaces
- rename output pointer field o to o_ptr for consistency with q_ptr/k_ptr/v_ptr
- regroup AttentionParams fields by responsibility and fix misleading comments
- drop unused max_seq_len/total_q fields and paged decode max_seq_len arg
- drop redundant group_size param from decode launchers (computed from p)
2026-08-09 20:52:06 +08:00
ViperEkura 596c35fd71 fix: report gradient snr in db 2026-08-09 13:40:27 +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 c1d05ae11d perf: benchmark decode via real inference engine
- route decode benchmark through InferenceEngine generate path
- add enable_cuda_graph toggle to engine, scheduler, and executor
- make benchmark --cuda-graph/--no-cuda-graph control the toggle
- hoist local time imports to module top
2026-08-09 11:47:14 +08:00
ViperEkura cf4f5ab9f6 feat: add persistent DataLoader workers
- Keep training workers alive between epochs when enabled.
- Avoid invalid prefetch settings for single-process loading.
2026-08-09 11:38:50 +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 7f0e8bb8c2 fix: let flash backend handle 4D causal prefill mask
- Treat 4D masks as causal (flash handles it natively), keep rejecting custom non-causal masks
- Enables flash backend in benchmark --compare and real prefill path
2026-08-08 23:51:27 +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 ca50fe4721 refactor: remove inference redundancy and fix cache leaks
- drop Executor unused tokenizer field, _head_dim, stale metrics docstring
- unify greedy sampling via SamplingPipeline.sample, drop top-level duplicate
- drop Task.flush_remaining no-op and unreachable prompt-length branch
- drop ProtocolHandler redundant chunks list (reuse body)
- fix page_size=1 token-slot leak on task_free
- clear _task_pages/_task_slots on alloc-failure paths
- reset _bind_state on task_free to avoid stale steady-state reuse
- remove unreachable contiguous branches in paged-only helpers
2026-08-08 21:45:01 +08:00
ViperEkura d9240ab149 refactor: split train context build steps
- separate checkpoint, model, data, and strategy setup\n- keep build orchestration concise and readable
2026-08-08 18:15:14 +08:00
ViperEkura d7cd69fef5 feat: add streaming IterableDataset for pretraining
- StreamingSeqDataset yields windows sequentially through each shard
- Shard-level shuffle, distributed and multi-worker shard partitioning
- __len__ returns total window count for scheduler total_steps
- Better OS page-cache locality than random-access map-style datasets
2026-08-08 16:15:10 +08:00
ViperEkura 9bff61fb91 perf: use cudaEvent for precise GPU timing in debug logs
- cudaEvent.elapsed_time gives microsecond precision vs perf_counter
- cudaEvent measures actual GPU execution, not just kernel launch
- falls back to time.perf_counter on CPU-only devices
2026-08-08 13:18:11 +08:00
ViperEkura 0b661bae85 fix: remove blocking cleanup from streaming generator
- stream finally froze main thread on cache.task_free
- scheduler handles cleanup in next loop iteration instead
2026-08-08 13:12:40 +08:00
ViperEkura ae9fd546ef perf: merge prefill warmup into _warmup_cuda_graphs
- 64-token prefill forward triggers cuBLAS auto-tuning at init
- reduces first-chat prefill from ~520ms to ~27ms
- warmup decode also drops from ~215ms to ~71ms
2026-08-08 13:10:58 +08:00
ViperEkura e3ea850dc9 fix: default backend race, raise on explicit fallback
- _default_backend lazy init protected with threading.Lock
- attention() raises when explicit backend cannot handle call
- FlashAttnBackend rejects prefill with non-None attn_mask
- training test uses TORCH_NATIVE backend directly
2026-08-08 13:00:58 +08:00
ViperEkura 6e5088cc7d refactor: remove prefill from CUDA graph warmup
- decode capture works without pre-filled KV values
- reduces init time and eliminates unused prefill forward
2026-08-08 12:49:07 +08:00
ViperEkura cbc584470d refactor: centralize logging in astrai.logging, replace ASTRAI_TIMED with log level
- 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
2026-08-08 12:39:27 +08:00
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 49aaa9a714 version: bump to 1.3.13 2026-08-07 23:56:01 +08:00
ViperEkura f163520fff refactor: break JsonlStore→preprocessing circular dependency
- 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
2026-08-07 23:22:32 +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 184fbbce5c refactor: extract shared steady-state increment detection
- 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
2026-08-07 23:00:25 +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 05739629fc feat: add timed() context manager and backend supports()
- Each backend exposes static supports(**kwargs) for capability query
- CudaBackend.supports checks head_dim + kernel availability
- FlashAttnBackend/TorchNativeBackend always return True
- timed() context manager gated by ASTRAI_TIMED=1 env var, logs via logger.info
- Wraps warmup prefill/decode, execute_prefill, and execute_decode
2026-08-07 20:51:30 +08:00
ViperEkura e0f7fa8e13 feat: enable CUDA graph by default with init-time warmup
- 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
2026-08-07 20:04:10 +08:00
ViperEkura af25833fab fix: add out_buf to attn_paged_decode for CUDA graph capture compatibility
- 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
2026-08-07 19:45:59 +08:00
ViperEkura 81788faef4 perf: use flash_attn_with_kvcache for contiguous cache decode
- 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
2026-08-07 18:21:02 +08:00
ViperEkura 55ee258e95 style: fix ruff lint warnings
- 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
2026-08-07 14:17:48 +08:00
ViperEkura ef1bb6f401 refactor: unify greedy check with _is_greedy helper
- Replace batch-scattered temperature==0 checks with (temperature == 0).all()

- Reuse _is_greedy in standalone sample() function
2026-08-07 14:14:15 +08:00
ViperEkura 6f49738991 feat: auto-select best available attention backend
- Default backend resolves to highest-priority available: flash -> cuda -> torch
- attention() falls back per-call for training/fp32/unsupported head_dim
- Re-apply index_copy_ for CUDA KV cache writes (index_put_ race mitigation)
2026-08-07 13:48:59 +08:00
ViperEkura 6054b8dbd4 feat: add CUDA-graph capture for decode forward
- 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
2026-08-06 19:57:12 +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 d0c5debbab perf: preload V in decode split-kv shared mem and cache partial tensors
- 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)
2026-08-06 18:27:00 +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 a317a4756b refactor: stateless MoE routing with grouped dispatch
- replace per-expert mask scan with sort+bincount grouped dispatch
- carry router stats in forward output instead of module state
- keep MoE diagnostics working under DDP/FSDP wrappers
- remove unused _load_balancing_loss helper
2026-08-05 18:42:12 +08:00
ViperEkura 9b7e6c205f feat: add moe auxloss and metrics 2026-08-05 18:12:28 +08:00
ViperEkura 8152760b5f refactor : use factory for attention backends
- register built-in backends through BaseFactory
- derive benchmark choices from registered backends
- cover string selection and invalid backend names
2026-08-05 15:37:22 +08:00
ViperEkura 8c052c99ee feat: add optional FlashAttention (FA2/FA3) backend
- 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
2026-08-05 15:27:26 +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 d0e5d910de perf: reduce remaining per-step allocations
- 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
2026-08-03 01:10:06 +08:00