Commit Graph
777 Commits
Author SHA1 Message Date
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 da6d94492d fix: parse yaml floats with yaml 1.2 schema
- register yaml 1.2 float resolver so scientific notation (2e-5) becomes float, not str
- replaces the decimal-point workaround in train configs
- add containerized training doc under docs/developer
2026-08-13 23:28:06 +08:00
ViperEkura 71b6e3aaaf feat: rework docker workflow for gpu-first training
- rewrite docker.sh with gpu default and --no-gpu override
- inject host uid/gid via ASTRAI_UID/GID in train.sh compose()
- filter readonly UID/GID lines when sourcing .env.train
- build image user via USER_UID/USER_GID args matching host uid/gid
- pass all GPUs (count: all) and filter by CUDA_VISIBLE_DEVICES inside the container
- forward NCCL vars through compose environment
2026-08-13 22:51:13 +08:00
ViperEkura f95722a277 feat: add containerized training workflow
- add a GPU trainer Compose profile with mounted data, models, and checkpoints
- add host commands for preflight, lifecycle, logs, status, and checkpoint cleanup
- resume from the latest complete checkpoint with external config or CLI arguments
2026-08-12 20:21:33 +08:00
ViperEkura 9f48cb8928 refactor: streamline Q block mapping
- bypass shared mapping for contiguous attention
- centralize paged Q tile broadcast in KV policy helpers
2026-08-10 08:40:18 +08:00
ViperEkura 9b58fef222 refactor: extract QTileMapper for prefill tile dispatch
- wrap one-thread map + shared broadcast + early exit
- both scalar and MMA prefill kernels use the shared helper
2026-08-09 23:18:03 +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 a5a3cc1fc2 refactor: unify attention param field names
- rename q_stride_* to q_*_stride to match mask stride convention
- rename mask_q_stride to mask_l_stride for consistent l-dim naming
- merge k/v and k_cache/v_cache into k_ptr/v_ptr; rename q to q_ptr
- KVSource policy selects contiguous vs paged mode at compile time
2026-08-09 20:23:58 +08:00
ViperEkura d565d44c43 fix: harden attention kernel boundaries
- fix scalar prefill head_dim=32 out-of-bounds via G=4 dispatch
- fix MMA decode 4D mask head indexing and invalid-row mask access
- add q_head/kv_head divisibility and head-dim contiguity checks
- validate split-KV scratch and decode out_buf layout in bindings
- set max dynamic shared memory for scalar decode D=256
- cover scalar prefill D=32 in pure C test
2026-08-09 14:53:24 +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 3416f98c58 fix: wire benchmark cache selection 2026-08-09 10:56:05 +08:00
ViperEkura d28552f878 refactor: use C++17 struct dispatch in csrc tests, tighten paged tolerances to 0.01
- Replace C++20 explicit lambda template parameters with file-scope structs (DecodeDispatch/PrefillDispatch etc.)
- Remove unused gs variable in run_decode_test
- Tighten paged test atol/rtol from 0.02 to 0.01 to match contiguous tests
2026-08-09 10:23:12 +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 a33ca04f60 fix: synchronize final decode async copy
- wait for the final split-KV tile before reading shared memory
- cover long decode with production context capacity
2026-08-09 00:31:47 +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 c52a2487ae fix: rename CUDA wheels with tag suffix to avoid upload clash
- all three release builds (pure, cu128, cu130) produce the same .whl filename, causing uploads to overwrite each other
- append the CUDA tag as a local version label (e.g. +cu128, +cu130)
2026-08-08 00:49:17 +08:00
ViperEkura 49aaa9a714 version: bump to 1.3.13 2026-08-07 23:56:01 +08:00
ViperEkura 056c1382ff docs: sync all documentation with current codebase
- remove GenerationRequest and generate_with_request references (class deleted)
- document cuda>flash>torch default priority and FlashAttnBackend
- add ASTR_BACKEND env var to backend docs, TorchNativeBackend (default) → (fallback)
- fix JsonlStore transform routing → DatasetFactory ownership
- fix CudaBackend fallback chain description (FlashAttn → TorchNative)
- add FlashAttnBackend to architecture strategy table
- add router_stats to DecoderOutput/FFNOutput TypedDict diagrams
- add decode_o_part/ml_part/decode_out to KVCache diagram
- add --append_eos/--no-append_eos to IFD evaluation parameter table
- update get-started CUDA kernel note (no longer requires explicit attn_backend activation)
- fix python -m scripts.tools.server (no __init__.py) → direct script call
2026-08-07 23:52:27 +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 6572be4f98 fix: prevent signal handler test from racing with training completion
- Set n_epoch=99999 so training runs until parent delivers signal instead of finishing too fast on CPU
- Drop ready-file deadline from 30s to 10s
2026-08-07 18:32:08 +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 0e7fe57d96 fix: use max_context_len for stable num_splits in paged decode
- PagedKV::host_kv_len now returns max_context_len instead of max_seq_len
- Eliminates grid-z instability for CUDA graph capture/replay
- Restore skip_no_kernel re-export accidentally removed by ruff --fix
2026-08-07 14:42:53 +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 a59ae8f32e fix: use c10::optional for o_part_buf/ml_part_buf decode kernel params 2026-08-06 20:50:48 +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