Commit Graph
202 Commits
Author SHA1 Message Date
ViperEkura 5e76fbd1bf perf: fp8 rings, lean autocast, gemm staging
- Finalize scale rings inside the quantize kernels: a last-block epilogue (threadfence + counter elect) folds amax into hist, reduces the window and publishes the next scale on device, zero extra launches; _ScaleRing packs [hist | scale | counter] into one CUDA buffer.
- Split FP8QuantizeParams out of FP8Params so each operator owns its fields; linear_forward/backward_fp8 take optional ring arguments.
- Drop the inference weight-quantization cache; the optimizer bumps the weight version every step, so a cache would miss anyway.
- Zero amax scratch via empty + cudaMemsetAsync instead of torch::zeros, cutting a ~50us fill_ dispatch per quantize.
- Stage crosswise-B operands K-major with cp.async (contract >= 8192) and PRMT-transpose per k_seg region in smem, interleaved with the MMAs; the sync LDG + byte-scatter path it replaces was long-scoreboard bound (ncu 4.6 vs 0.4 stalls/issue).
- Load crosswise-A direct with an in-register PRMT transpose; its operands are typically L2-resident and the staging round trip measured as a net loss.
- Enable grouped rasterization for the congruous NT forward (shared B stripe keeps the weight operand hot in L2) and make the smem budget layout-aware (Fp8GemmSmem) while holding two CTAs per SM.
- Annotate ops/fp8.py return types; drop weight-cache and decorator tests, hoist their imports to module level.

e2e 12L/dim1024/B4xT512 fused AdamW: fp8 137.8ms/step vs bf16 210.3ms, 1.53x. Kernel vs cuBLASLt _scaled_mm: fwd 1.03-1.09x, dX 1.33-1.47x, dW 1.30-1.39x (from 1.10/1.42-1.49/1.52-1.56x), before the pre-transposed copies cuBLASLt needs for dX/dW. fp8 train step vs bf16: 1.34x at 2048 tokens (was 1.25x), 1.08x at 512.
2026-08-25 14:24:11 +08:00
ViperEkura 4dc5e923e0 perf: finalize fp8 scale rings inside quantize kernels
- last-block epilogue (threadfence + counter elect) folds amax into hist[idx], reduces the window and publishes the next scale on device — zero extra launches per linear layer
- _ScaleRing packs [hist | scale | counter] into one CUDA buffer; the eager hist-write / max / scale-copy chain and update() are gone
- split FP8QuantizeParams out of FP8Params so each operator owns its fields; linear_forward/backward_fp8 take optional ring arguments
- e2e 12L/dim1024/B4xT512 (fused AdamW): fp8 137.8ms/step vs bf16 210.3ms, 1.53x; fwd 1.82x, bwd 1.50x
2026-08-25 11:12:14 +08:00
ViperEkura 998b443aa3 refactor: dedupe fp8 meta state into per-operand rings
- collapse FP8TensorMeta's 12 slots + 6 copy-paste methods into three _ScaleRing objects (hist/idx/scale/initialized + update/seed)
- skip meta allocation entirely on the DynamicScaling path (zero rings, scales measured inline)
- drop write-only FP8State._last_device and unused E4M3_MAX alias
2026-08-24 21:19:45 +08:00
ViperEkura 7da1439c9e feat: static fp8 weights and bias with fused epilogue
- linear_forward_fp8 accepts pre-quantized w8 (matching fmt) and skips the weight quantize; amax_w returns 0 on that path since no bf16 values are seen
- bias is now fused into the GEMM epilogue for both dtypes, replacing the separate torch-level add (one elementwise kernel per linear removed)
- FP8Params.bias becomes void* with a new bias_scale slot: null scale = raw bf16 bias, non-null = fp8 storage dequantized in the epilogue after the operand scaling and before any output quantization
- ops/fp8.py relaxes the w dtype check to bf16-or-fp8 and passes bias_scale through
- regression test covers w8/b8, w8/bf16-bias and the amax_w = 0 contract vs an explicit quantization reference
2026-08-24 19:25:23 +08:00
ViperEkura 29e5f571af fix: own fp8 linear backward via autograd Function
- backward used to read the global fp8 flag at loss.backward() time, so calling it outside fp8_autocast silently fell back to bf16 mm (953 ms cublas per step, 49.9% of the model step)
- _LinearFp8(torch.autograd.Function) now owns the fwd/bwd pair: forward captures fmt/recipe/meta on ctx inside the autocast region, backward reads only ctx (scales from the meta rings, masks from ctx.needs_input_grad), so backward is fp8 wherever it runs
- register the aten::linear impl on AutogradCUDA (replaces torch's generated linear formula that calls aten::linear_backward into the bf16 fallback) and keep the CUDA key for inference_mode
- drop the aten::linear_backward override and fp8_linear_backward (dead paths)
- regression test asserts the fp8 backward fires outside the autocast region and grads match the bf16 reference by direction/norm (E5M2 noise)
- model step (0.67B, CE loss, batch 4x1024): backward GEMMs 953 -> 618 ms (1.54x), full step ~1.2x
2026-08-24 19:03:01 +08:00
ViperEkura 4b10d3ca37 perf: vectorize fp8 quantize and swizzle gemm smem 2026-08-23 20:31:44 +08:00
ViperEkura 4244df2785 perf: pure FP8 fwd/bwd and lean non-transposed GEMM
- drop the fused kernel; forward/backward are quantize + a pre-quantized GEMM
- rename module fp8_mm -> fp8_ops (mm.cu -> ops.cu)
- kernels/launchers fp8_gemm_kernel / launch_fp8_gemm; drop PqTraits/gather_trans/pack_fp8x4_vector
- remove the in-kernel transposed-operand branches (TransA/TransB)
- backward: quantize g once (amax_g here), explicit fp8 transposes, fast non-transposed GEMMs (dX = g@w^T, dW = g^T@x^T)
- each pass uses a single FP8 format (E4M3 fwd / E5M2 bwd)
2026-08-23 15:38:30 +08:00
ViperEkura a29bdfae46 refactor: rework attention backend resolution
- explicit attn_backend() context wins over ASTR_BACKEND env
- polymorphic available()/supports_call() replace isinstance dispatch
- cache singleton backend instances to avoid hot-path allocation
- training (fwd=None) resolves cuda > flash > torch by capability
- flash dense supports mask-free calls only; masked training falls back to torch
2026-08-23 14:47:02 +08:00
ViperEkura 75304d084d refactor: unify CUDA skip guards in tests
- add skip_no_fp8 (CUDA + fp8_mm kernel + cc 8.9+) to tests/conftest.py
- use skip_no_cuda / skip_no_kernel / skip_no_fp8 directly in test modules
- drop _GPU alias and tests.extension.conftest re-exports
- remove unused imports (Union in hf_adapter, make_grpo_config in data conftest)
2026-08-22 21:08:06 +08:00
ViperEkura 16a55bb474 refactor: reorganize CUDA kernels into per-family directories
- move attention kernels to csrc/kernels/attention/ and rotary to rotary/
- add shared common/mma.cuh (mma_sync, ldmatrix) and device.cuh (sm checks)
- split fp8_mm into three-layer fp8/common.h, gemm.cuh, mm.cu
- fix fused FP8 GEMM ldmatrix lane indexing to fix OOB shared reads
- update extension ops, loader, and kernel tests
2026-08-22 20:40:31 +08:00
ViperEkura cb21af38ba feat: unify Docker serving configuration in YAML
- Add server.py --config serve.yaml; explicit CLI flags override YAML
- Add scripts/serve.sh and serve_runtime.py for the Compose lifecycle
- Template server/cpu ports and param mounts in docker-compose.yml
- Document schema in docs/developer/docker-serving.md and params guide
- Add tests for runtime parsing and server CLI merge logic
2026-08-21 23:16:45 +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 7d27f3e078 feat: load HuggingFace checkpoints via key/config conversion
- Add astrai.serialization.hf_adapter mapping LLaMA-style HF keys to AstrAI names (input_layernorm, gate_proj, MoE experts/shared_experts) with config aliases for dense and MoE (Mixtral/DeepSeek-V3) layouts; reject biased projections, mismatched head_dim and MLA
- Give AutoModel.from_pretrained weights_format=auto|astrai|hf with auto-detection; read sharded safetensors via model.safetensors.index.json
- Adapt preloaded weights/config in train_context and benchmark CLI
2026-08-20 11:34:59 +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 3d3ea47d37 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

Benchmark: NVIDIA L20, BF16, 1B model, paged KV cache, CUDA Graph, prompt 512, generation 128 (median of 3 alternating runs)
- batch 1: 234.5 -> 242.6 tok/s (1.034x, +3.4%)
- batch 8: 1243.1 -> 1286.6 tok/s (1.035x, +3.5%)
2026-08-19 00:36:53 +08:00
ViperEkura f7f14d0e5f test: add online GRPO end-to-end training test 2026-08-19 00:07:16 +08:00
ViperEkura 7580d80d45 perf: accelerate FP8 backward with fused fast kernel
- route dX/dW through the fused 128x64 fast kernel via contiguous transposes
- drop the legacy 64x64 kernel, cutting dX 1.55->0.38 ms and dW 1.28->0.26 ms
- sync all threads after cp.async.wait_group to fix sporadic NaN in large GEMMs
- add fp8_mm_prequant_fp8 custom op for FP8-in/FP8-out GEMM
2026-08-18 23:46:43 +08:00
ViperEkura cb51a3587b perf: optimize fused FP8 GEMM kernel 2026-08-18 19:56:15 +08:00
ViperEkura 1bcd8f53ab perf: precompute ragged Q tile scheduling 2026-08-16 23:32:46 +08:00
ViperEkura 6ac3b51496 refactor: separate extension ops and backends 2026-08-16 21:15:52 +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 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 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 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 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 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 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 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 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 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 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 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 020e2eff4e refactor: emit strategy metrics as floats
- Converts detached strategy metrics before returning loss output
- Removes redundant item conversion from the trainer loop
- Updates the documented contract and regression tests
2026-08-02 06:38:28 +08:00
ViperEkura 1c7369f293 feat: add MoE auxiliary loss metrics
- 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
2026-08-02 06:30:43 +08:00
ViperEkura d7db37a70f fix: preserve MoE routing defaults 2026-08-02 05:30:26 +08:00
Gaolingx 6d98bb4f9f 20260801-moe model impl
need to add aux loss for load balancing
2026-08-01 22:48:58 +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