Skip to content

[perf](sft): accelerate FP8 training with shared repack and decoded tiles - #2193

Draft
yyj6666667 wants to merge 4 commits into
kvcache-ai:mainfrom
yyj6666667:perf/sft-fp8-shared-repack-pr
Draft

[perf](sft): accelerate FP8 training with shared repack and decoded tiles#2193
yyj6666667 wants to merge 4 commits into
kvcache-ai:mainfrom
yyj6666667:perf/sft-fp8-shared-repack-pr

Conversation

@yyj6666667

@yyj6666667 yyj6666667 commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Summary

Follow-up to #2141 and #2188, limited to five production headers:

  • Shared asynchronous backward-repack ownership, completion and weight-version tracking.
  • SFT integration that allows FP8 repack to overlap GPU work while protecting shared CPU pools.
  • Direct packed-FP8 transpose, including the corresponding block-scale grid.
  • A shared inference/SFT GEMM policy with bounded 128x128 decoded-tile reuse (32 KiB per worker).

Frozen weights and the shared backward BufferB remain FP8 plus scales.
No benchmark harness, added test files, documentation, CMake changes, NVTX tracing
or extra profiler statistics are included in the aggregate diff. The existing
upstream tests are unchanged.

Performance and verification

DeepSeek V3.1 FP8 LoRA SFT on two EPYC 9355 CPUs and eight RTX 5090 GPUs:
sequence 1024, batch 1/GPU, GAS 1, 64 CPU workers/TP2, LoRA rank 8/alpha 16,
CPU retain / GPU recompute. A/B/B/A means were 62.1535 -> 136.4012 global
tokens/s (2.1946x)
, with 2 warmup + 16 measured updates per run.
All four short fixture runs had finite loss and gradients, with loss
12.671875 -> 11.8515625; this is not a long-term convergence claim.

Performance was measured before PR packaging and this scope reduction, not
retimed on the core-only binary. The core-only revision 7db13789 was separately
rebuilt on base 9b0da97c with GCC 11.4 / CUDA 12.8 / AVX512-BF16/VBMI and checked
using external validation fixtures: 274 bitwise FP8 GEMM comparisons, parallel
packed-byte/scale transpose, repack lifecycle, and 30 native pipeline cases.
Validation fixtures are not part of this PR's file diff.

Draft integration boundary

Upstream #2186 subsequently added a separate AMX_K2_SFT_MOE_TP that still
exposes prepare_backward_bb_for_async(); this wrapper expects
prepare_shared_backward_weights() and invalidate_backward_repack().
That interface integration and a rebuilt/tested merge head remain required
before merge. The newer aa2cd307 configuration and upstream CI are not claimed
as validated. This PR does not add RAWINT4/INT8 performance work.

Keep only shared backward-repack ownership, the SFT integration, and FP8 packed transpose/GEMM changes. Remove the benchmark harness, added test files, documentation, and optional profiling/tracing additions from the upstream diff. The original commits and private experiment archive preserve the removed material; validation runs outside the PR.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant