anki-rwkv · factory overview · model architectures · R01 campaign · 8k/2k run · data quality

Making RWKV training several times faster, bit for bit

Every iteration tests one optimisation against the last accepted configuration. A change only counts if the trained weights stay bitwise identical and throughput really improves.

Current accepted throughput
–
reviews/s
Speedup vs original baseline
–
Iterations
–
Accepted improvements
–

Throughput over iterations

Step line: accepted reference after each iteration. Dots: each candidate’s full-run reviews/s, coloured by decision. Dashed markers are milestones; hover or tap anything for details.

Cumulative speedup

Accepted reference divided by the original no-prefetch baseline.

Replicated-block gains

Bars: mean per-update gain from interleaved reference/candidate blocks. Markers: one-sided lower confidence bound. Promotion needs a positive lower bound.

Milestones

What changed and why it mattered. Curated milestones by default; pick decisions to list every matching iteration. Colours match the chart above.

    How one update flows through CPU and GPU

    The machine: AMD Ryzen 7 9850X3D (8 physical cores) and an RTX 5080. Five preparation workers each own a physical core; the main process uses the other three. Numbers below come from the accepted configuration of iteration 0090.

    CPU and GPU layout of one training update Review histories go from disk to five preparation workers, into a queue of ready batches, then to the main process, which copies each batch to the GPU and launches forward, backward and the AdamW update. Loss values come back two updates later. CPU · 8 physical cores GPU · RTX 5080 Disk Revlog parquet Worker 1 Worker 2 Worker 3 Worker 4 Worker 5 Batch queue Up to 10 ready Main process 3 cores, 1 thread Copy in Forward + loss Backward Clip + AdamW Teal: one worker per core Dashed: loss read back The main process never waits for the GPU: it reads the loss of update n−2 while it launches update n.

    Preparation workers

    1. Read a user’s review history from parquet.
    2. Cut it into contexts of 128 reviews.
    3. Encode card, deck, note, preset and user IDs and the time features.
    4. Pack 512 contexts (up to 65,536 reviews) into one batch, crossing user boundaries so batches stay full (F29).
    5. Write the batch into one shared-memory buffer (F20/B01).

    Main process

    1. Take the next ready batch, pin it, and start one asynchronous copy to the GPU (F20).
    2. Launch the forward pass, loss and backward pass through native C++ dispatch instead of Python autograd (F65).
    3. Launch grouped gradient copies from BF16 to the FP32 master weights, clipping and AdamW (F23).
    4. Read the loss and gradient norm of update n−2, not n, so it never stalls (F33/F33b).

    GPU

    1. Expand the encoded features in device memory (F14).
    2. Run the RWKV recurrence forward and backward with exact NVRTC kernels (F31, F90). The WKV backward is the largest kernel, about 76 ms per large update.
    3. Apply the weight update; the FP32 master weights then refresh the BF16 copy used for the next update.

    Why everything overlaps

    Switch between the overlapped pipeline used today and a naive loop that does one step at a time.

    Updates per second
    –
    GPU busy
    –
    Reviews per second
    –
    Prepare a batch Copy and launch kernels GPU computing Read loss (sync)

    The GPU is still idle about 40% of the time. One update launches thousands of small kernels, and on Windows (WDDM) every launch costs CPU time, so the main thread’s dispatch is now the bottleneck. That is why most recent gains came from launching less (native dispatch, F65), synchronizing less (F33, F37) and grouping copies (F23), rather than from faster maths.

    How to read this

    All iterations

    IterTrackReviews/s× baselineDecisionBlock meanLower bound