Measured with examples/sweep_arch.rs, one configuration per process, on an
otherwise idle RTX 3060 12 GB (desktop applications holding 1258 MiB).
Every earlier measurement in this repository's history was taken while a second training job held 5754 MiB and 100% of the GPU. Those numbers are superseded by this table and should not be compared against it.
scripts/bench_check.sh compares the best tokens/s of five runs against these
rows and fails below 97%. Measured 2026-10-07 at cd8b514 on the idle RTX 3060
(train step) and the Ryzen 5 7600X (decode, CPU). One run spreads about 3.5% on
both. bench_decode is sensitive to anything else on the CPU: with one
parallel rustc it measured 1464-2088 tokens/s, so run it on a quiet machine.
Re-record with scripts/bench_check.sh --record after a change that is meant
to move these numbers. The train-step row was re-recorded alone on
2026-10-07 at 6f48659 (idle GPU, best of 46927, 52707, 52890, 52783,
52815) after A7's BF16 MoE work; the CPU was busy, so the decode row was kept.
| Benchmark | Best tokens/s |
|---|---|
| bench_train_step 4 1024 cuda | 52890 |
| bench_decode 512 128 | 2434 |
FLOPs per token counts the parameter term and the attention term:
FLOPs/token = 6 * active_params + 12 * n_layers * seq_len * d_model
MFU = tok/s * FLOPs/token / 25e12 (RTX 3060 bf16 peak, fp32 accumulate)
| Configuration | Total | Active | tok/s | TFLOPS | MFU | Days for 50B |
|---|---|---|---|---|---|---|
| v32000 d512 L8 16x512 moe | 55.2M | 35.7M | 36302 | 8.69 | 34.8% | 15.9 |
| v16384 d512 L8 16x512 moe | 47.2M | 27.7M | 40115 | 7.68 | 30.7% | 14.4 |
| v16384 d512 L8 16x512 dense | 30.9M | 30.9M | 43704 | 9.21 | 36.8% | 13.2 |
| v16384 d512 L8 32x512 dense | 30.9M | 30.9M | 43952 | 9.26 | 37.0% | 13.2 |
| v16384 d512 L8 16x1024 dense | 30.9M | 30.9M | 34089 | 8.03 | 32.1% | 17.0 |
| v16384 d512 L12 16x512 dense | 42.2M | 42.2M | 29189 | 8.49 | 34.0% | 19.8 |
| v16384 d768 L12 16x512 dense | 88.7M | 88.7M | 15620 | 9.20 | 36.8% | 37.0 |
| v32000 d512 L8 128x128 moe | 55.2M | 35.7M | 39988 | 7.87 | 31.5% | 14.5 |
Contended figures for the same configurations ran at 16989, 18651, 21295 and 14453 tok/s: the clean card is 2.14x faster, and nothing else in this repository changed between the two sets.
| Configuration | Peak (including 1258 MiB desktop) |
|---|---|
| v16384 d512 L8 32x512 dense | 6779 MiB |
| v16384 d512 L8 16x1024 dense | 7355 MiB |
No configuration in the sweep ran out of memory, including the three that did under contention (batch 24, 12 layers, seq 1024 at batch 16). The 12 GB card has roughly 4.5 GB of headroom at the largest shape measured.
The plan in optimizePlan.md gated two tasks on this measurement:
- Task 6, bf16 activation storage — justified only below about 25% MFU, where the step is starved of bandwidth rather than arithmetic.
- Task 7, pointer-array batched attention GEMM — justified only if the per-launch overhead is a visible share of step time.
MFU sits between 30.7% and 37.0% across every shape measured, including the 128-sequence batch that issues the fewest launches per token and the d768 model that issues the largest GEMMs. A device this well fed does not have 1.3x sitting in its activation traffic, and the launch count is plainly not what bounds it. Both tasks are closed as not worth their complexity.
Re-open them only if a later change (much longer sequences, a much smaller
d_model) moves MFU back below 25%.
nsys profile --stats=true on an idle card, over
./target/release/bench --gpu --mixed-precision --d-model 768 --n-layers 16 \
--n-heads 12 --n-kv-heads 4 --head-dim 64 --d-ff 1408 --seq-len 1024 \
--batch-size 4 --steps 20 --gpu-memory-budget-mib 9000
Total GPU kernel time across the run is 7.00 s. The cuda_gpu_kern_sum
table reproduces the numbers Task 9 was written against:
| Kernel | Instances | Total | Avg | Share of GPU time |
|---|---|---|---|---|
causal_softmax_bwd |
336 | 627.3 ms | 1.867 ms | 9.0% |
causal_softmax_lse |
336 | 528.5 ms | 1.573 ms | 7.6% |
causal_probs_from_lse |
336 | 336.1 ms | 1.000 ms | 4.8% |
| Softmax total | 1491.9 ms | 21.3% |
336 instances is 16 layers x 21 steps, once per layer per step, so none of it
is setup cost. The nine cutlass::Kernel2<...> GEMM variants together account
for roughly 60% of GPU time; the two 64x64 tile variants among them, 1186 ms
and 17%, are the per-head attention GEMMs.
At this shape the score matrix is heads * sequences * seq_len = 49152 rows
of 1024 floats, of which the causal mask leaves about half visible: 25.2M
elements, 100.8 MB in FP32. Counting the passes each kernel makes over that
data against the card's 360 GB/s:
| Kernel | Traffic | Time | Achieved |
|---|---|---|---|
causal_softmax_lse |
3 reads + 2 writes of the visible half, plus one write zeroing the masked half — 605 MB | 1.573 ms | 384 GB/s |
causal_softmax_bwd |
reads of both g and p twice, one write of each half — 604 MB |
1.867 ms | 324 GB/s |
causal_probs_from_lse |
one read of the visible half, one write of the whole row — 302 MB | 1.000 ms | 302 GB/s |
All three run at or above 84% of the card's nominal bandwidth, and
causal_softmax_lse exceeds it on L2 hits. There is no meaningful headroom
inside these kernels. Tightening the pass count would move the total from
21.3% to perhaps 15%, and no rewrite that still materializes a
[seq_len, seq_len] FP32 matrix in global memory can do better than that,
because the traffic, not the arithmetic, is what costs.
That leaves two ways forward, and they are not the same size:
- Store the score matrix in BF16. Every byte of the traffic above halves,
and so does the FP32 traffic in the two attention GEMMs that read and write
the same buffer. The softmax kernels keep their structure and their FP32
arithmetic; only the loads and stores narrow. The existing
gemm_strided_batched_dispatchalready computes inCUBLAS_COMPUTE_32F_FAST_16BFand would needCUDA_R_16BFfor the score operand rather than a new code path. Expected: most of a 10% reduction in total GPU time, and the score buffer drops from 201 MB to 100 MB. - Never materialize the matrix — flash attention. This is what Task 9's
Step 2 asks for and what PyTorch's SDPA and TF's XLA fusion do. It removes
causal_probs_from_lseoutright, folds the softmax into the QK^T and AV passes, and is the only path that recovers the full 21%. It also replaces two cuBLAS tensor-core GEMMs with hand-written NVRTC kernels that have to beat them, which is the largest piece of work anywhere in this plan.
Option 1 is a precondition for neither and a fair fraction of the gain, so it is worth measuring before committing to option 2.
Five kernels became three, and the [seq_len, seq_len] matrix stopped being
written at all.
| Pass | Before | After |
|---|---|---|
| Forward | causal_softmax_lse plus two batched GEMMs per head |
flash_attention_fwd |
| Backward | causal_probs_from_lse, causal_softmax_bwd and four batched GEMMs per head |
flash_attention_delta, flash_attention_dq, flash_attention_dkv |
The forward kernel is blocked over query tiles. The backward pass needs two kernels because the query gradient is complete inside a query tile and the key and value gradients are complete inside a key tile; blocking the second one on the key/value head rather than the query head is what makes grouped-query attention exact without atomics. Both rebuild the scores from the stored log-sum-exp, which costs one extra matmul per tile pair and saves the whole round trip.
Same command as the profile above. The card was not idle: a game was running
on it throughout, so treat these as a paired comparison, not as absolute
numbers. RUSTING_BRAIN_NO_FLASH=1 is what selects the old path.
| Path | Run 1 | Run 2 |
|---|---|---|
| Fused | 12077 tok/s | 12064 tok/s |
| Three-kernel | 7087 tok/s | 7187 tok/s |
1.70x, and the final loss agrees to 0.006 across both paths (6.711-6.716 fused, 6.717-6.718 split) after 20 steps. For reference the three-kernel path measured 11905 tok/s on an idle card, so the fused path on a contended one is already past the old idle figure; the idle figure for the fused path is still to be taken.
Of that 1.70x, the forward fusion alone was 1.15x, measured the same way before the backward kernels existed. The backward half is the larger share, as the profile said it would be.
| Kernel | Share of GPU time |
|---|---|
cutlass ... 128x256_16x3_nt |
17.3% |
cutlass ... 256x128_16x3_nn |
17.3% |
cutlass ... 256x128_16x3_tn |
16.5% |
flash_attention_dkv |
8.0% |
flash_attention_dq |
5.8% |
flash_attention_fwd |
4.0% |
causal_softmax_lse, causal_softmax_bwd and causal_probs_from_lse are
gone, and so are the 8064-instance per-head attention GEMMs that used to run
beside them. Attention is now 17.8% of GPU time in three kernels, against
21.3% in the softmax kernels alone plus 17% in the per-head GEMMs before.
The three remaining cutlass entries are the model's own projections, and at
51% of GPU time between them they are what stands between this and PyTorch.
flash_attention_dkv is the slowest of the three fused kernels at a 1.04 ms
median against flash_attention_dq's 0.68 ms, which is what four matmuls per
tile pair and a tile count that falls from 16 to 1 across the grid look like.
Both device paths round their matmul operands to BF16, so neither is a
reference for the other; the FP32 host path is. On a two-layer, four-head
model at 100 tokens the fused forward lands 3.7e-3 from the host where the
three-kernel path lands 2.5e-3, on logits of magnitude 0.475. After one full
gradient-descent step the two land 1.0 to 1.4 times apart, the difference
being that the fused backward takes the softmax row sum from the output
gradient and the output rather than from probabilities it no longer keeps.
fused_attention_matches_the_three_kernel_path_or_skips_without_device and
its backward twin hold that ratio.
The fused-attention profile left the block's projection GEMMs at roughly half
of GPU time, all of them on cutlass_80_tensorop_s1688bf16gemm_*_align4. That
is the m16n8k8 tensor-core shape, which is what cuBLAS picks for FP32 operands
with a 32F_FAST_16BF compute type. The language-model head, whose operands
are already BF16, lands on s16816 ... align8 instead: twice the k per
instruction.
examples/gemm_probe.rs measured the three ways to close that gap on the four
projection shapes at batch 4, sequence 1024:
| Shape | units/inner | FP32 operands | BF16 operands | BF16 plus a per-call cast |
|---|---|---|---|---|
| qkv projection | 1280/768 | 12.82 TF | 13.07 TF | 9.22 TF |
| attention output | 768/768 | 9.83 TF | 11.67 TF | 7.36 TF |
| feed forward in | 2816/768 | 12.55 TF | 14.35 TF | 11.48 TF |
| feed forward out | 768/1408 | 12.49 TF | 12.85 TF | 8.09 TF |
Casting the operands inside the GEMM wrapper is a net loss - 0.65x to 0.91x
- so the kernel that produces an activation has to write it narrow. That is
what this change does:
rmsnorm_fwd,swiglu_fwd,swiglu_bwdandflash_attention_fwdtake anarrowflag and store throughstore_act, andActingpu_modelholds an activation as untyped bytes plus that flag. One byte buffer serves both precisions because a kernel launch only ever passes a device pointer and cuBLAS takes the element type as a runtime argument.
Weights and gradients, Adam moments and every accumulator stay FP32. The routed feed-forward stays FP32 end to end, because its gathers, scatters and row-wise reductions are FP32 kernels and it is not what these shapes exercise.
Idle card, same command as the profile above.
| Path | Run 1 | Run 2 | Run 3 |
|---|---|---|---|
| BF16 activations | 21516 tok/s | 21529 tok/s | 21484 tok/s |
FP32 activations (commit 1f3ad13) |
20152 tok/s | 20180 tok/s |
1.067x. Final loss 6.707 against 6.713, on logits the two paths compute with the same FP32 accumulators and different operand rounding.
185.6 ms of kernel time per step against 190.4 ms of wall clock, so launch overhead is 2.5% and not worth attacking.
| Group | Per step | Share |
|---|---|---|
| Projection and head GEMMs | 108.0 ms | 58.2% |
| Fused attention (three kernels) | 36.2 ms | 19.5% |
| Elementwise and the optimizer | 41.5 ms | 22.3% |
The GEMMs move roughly 2498 GFLOP per step - 1894 in the blocks, 604 in the
head - which at 108.0 ms is 23.1 TFLOPS against the card's 25.5 TFLOPS BF16
peak, 91%. There is nothing left in them. Half of them still land on
ampere_s1688gemm_bf16_128x128, the k8 kernel, but at that fraction of peak
the heuristic is not what bounds the step.
The three flash kernels move about 360 GFLOP per step at 36.2 ms, which is 10.0 TFLOPS, 39% of peak. That is where the remaining headroom is.
For reference: PyTorch 23121 tok/s, TensorFlow 20140 tok/s, this repository's pre-optimization idle baseline 11905 tok/s. At 21516 tok/s the step is 1.81x the baseline, past TensorFlow, and 7% short of PyTorch.
The previous step narrowed the activations that feed cuBLAS. qkv was left
FP32, even though its only readers were the three fused attention kernels and
they rounded every tile to BF16 as they loaded it. This step makes qkv and
grad_qkv narrow too, so the rounding happens once at the projection's store
instead of on every tile load.
eligible(mixed_precision, head_dim) = mixed_precision && head_dim == TILE, so
whenever the fused kernels run at all, every buffer they touch is BF16. The
runtime narrow flag is therefore gone from all four kernels in
cuda_flash.rs - the type is known at compile time. cuBLAS writes the
projection straight to BF16 (Ctype = CUDA_R_16BF, via linear_packed_act),
which removes a cast_act launch as well as the FP32 round trip.
20960 tok/s against 21524, all of it in flash_attention_dkv. It was not
occupancy: nvcc -arch=sm_86 -cubin -Xptxas -v reports 198 registers for dkv
both before and after, and deleting the narrow branch changed nothing.
The cause is that these kernels are latency bound, not bandwidth bound. Halving
the bytes buys nothing; issuing the same number of loads at 16 bits each,
plus a shift and an __uint_as_float per element, costs.
The fix is that two adjacent BF16 elements are exactly the 32-bit register an
m16n8k16 mma operand takes, so one 32-bit load fills a fragment with no
conversion at all:
__device__ __forceinline__ unsigned bf16_pair(const unsigned short* p, size_t i){
return *(const unsigned*)(p + i); // i must be even
}dkv now loads K and V straight into kf/vf, and tiles Q and dO two elements
a thread. dq tiles K and V the same way, and every grad_qkv and out store
is a packed 32-bit write. The forward kernel's K/V tile was tried this way and
measured worse - 8.0 ms to 10.9 ms - because its value tile is stored
transposed and the wider load buys bank conflicts; it stays scalar, with a
comment saying so.
One incidental exactness note: the attention scale moved off the queries and
onto the scores and the dk store. For head dim 64 the scale is 1/8, a power
of two, so this is bit-exact and it saves 64 multiplies per query row.
Idle card, 20 steps of the 102M model, same command as above.
| Path | Tokens/sec | Final loss |
|---|---|---|
BF16 qkv (579b040) |
21863 | 6.715 |
FP32 qkv (8c318ae) |
21524 | 6.710 |
1.016x. 144 library tests pass.
| Kernel | Per step |
|---|---|
ampere_s1688gemm_bf16_128x128_ldg8_stages_32x1_nt |
20.38 ms |
cutlass_80_tensorop_s16816gemm_bf16_256x128 |
19.37 ms |
flash_attention_dkv |
17.16 ms |
ampere_s1688gemm_bf16_128x128_ldg8_tn |
13.09 ms |
cutlass_80_tensorop_s16816gemm_bf16_128x256 |
12.79 ms |
flash_attention_dq |
10.96 ms |
adam |
9.04 ms |
flash_attention_fwd |
8.04 ms |
add_inplace |
5.76 ms |
swiglu_bwd |
4.68 ms |
rmsnorm_bwd |
4.28 ms |
cast_act |
4.04 ms |
swiglu_fwd |
3.41 ms |
rmsnorm_fwd |
2.80 ms |
rope_rotate |
2.66 ms |
rope_rotate is 32% cheaper than it was and cast_act 31%, both because they
now move half the bytes. The GEMMs are 59% of the step, the three flash kernels
19%, everything else 21%.
examples/flash_probe.rs runs the three kernels standalone on the benchmark
shape - forward 0.449 ms, grad query 0.656 ms, grad key/value 1.014 ms per
launch - which is the fast loop for working on them without a 20-step
benchmark in between.
PyTorch is 23121 tok/s. At 21863 the gap is 5.7%.
Four changes, each of them a pass over a whole activation or a whole parameter set that some other kernel was already positioned to absorb. None of them makes a kernel faster; they delete work.
A residual branch's backward pass produced two gradients for the same tensor and
added them in a separate kernel. rmsnorm_bwd already writes that tensor, so it
takes the other addend as an argument and folds it in:
float g_in=up[c]*w[c]*t-src[c]*shared;
if(add_res)g_in+=res[base+c];
dst[c]=g_in;21863 -> 22240 tok/s, 1.017x.
The gradient a norm hands down is a GEMM operand for the layer below it, and a mixed-precision GEMM wants it in BF16. That narrowing was its own pass. The same kernel writes it, because the value is in a register at that point:
if(copy_mode)store_act(copy,base+c,g_in,copy_mode-1);Which precision the copy wants is the receiving feed-forward's, not the
context flag: a routed layer's GEMMs are FP32 and ask for no copy at all, and
the fused attention path's merged is narrow only when flash attention runs.
Deriving it from the consumer (cache.merged.is_narrow(), and
cache.blocks[i].feed_forward_normed.is_narrow() for the norm above a block)
is what makes the FP32 and flash-disabled tests pass.
22240 -> 22383 tok/s, 1.006x.
The three flash backward kernels took grad_out in FP32 and converted on every
tile load. They now take it in BF16, which is both half the traffic and, because
two adjacent BF16 elements are the 32-bit register an m16n8k16 operand takes,
zero conversion:
df[kk][0]=bf16_pair(grad_out,e+c);22383 -> 22692 tok/s, 1.014x.
A residual branch's projection accumulates over the block input with beta = 1,
so that input had to be duplicated into the destination first: two
clone_dtod calls a block, 32 copies of 12.58 MB a step, 2.5 ms.
The first attempt removed them the other way round -- projection writes its own
buffer with beta = 0, and the next norm sums the two inputs -- and measured
neutral: 22875, 22747, 22681 tok/s against a 22692 baseline. The profile said
exactly why. The copies went from 2.52 ms a step to 0.08, and rmsnorm_fwd went
from 2.17 ms to 4.63. The two are the same number because the scheme still
materializes two tensors per branch; the beta = 1 read it also removed was
never costing anything, because it is hidden under a GEMM running at 22 TFLOPS.
What does pay is leaving the accumulate alone and moving only the copy. The norm above a branch reads the row it would have copied, so it writes the duplicate itself, out of registers:
for(int c=tid;c<cols;c+=ROW_THREADS){float v=src[c];if(do_copy)dup[c]=v;s+=v*v;}That is a store where there was a read-modify-write pass: 12.58 MB less per site.
22758 -> 22944 tok/s, 1.008x (mean of two 60-step runs each).
Every optimizer step ended by memsetting each parameter's gradient, and every
weight-gradient GEMM then accumulated into those zeros with beta = 1. That is
467 MB of zeroing a step.
A gradient buffer holding the last step's value is as good as a zeroed one, so
long as the first writer of the next step overwrites rather than accumulates.
Each DeviceParam carries a grad_dirty flag; grad_beta() hands out 0.0
once and 1.0 after, and the optimizer step clears the flag instead of the
memory. The callers that can only accumulate -- the embedding scatter's atomic
adds, the tied head's BF16 accumulate, and the optimizer itself for a parameter
that took no gradient -- call clear_grad(), which memsets only if nothing has
claimed the buffer yet.
The win is larger than the memsets alone, because beta = 0 also spares each
weight-gradient GEMM the read of its destination.
22944 -> 23411 tok/s, 1.020x (23399 and 23423 over two 60-step runs).
Idle card, 60 steps of the 102M model, same command as above.
| Path | Tokens/sec | Final loss |
|---|---|---|
This branch (0c5d951) |
23411 | 1.188 / 1.201 |
Before these four (579b040) |
21863 | 6.715 (20 steps) |
| PyTorch | 23121 | |
| TensorFlow | 20140 | |
| Idle-GPU baseline | 11905 |
1.97x the baseline, and 1.013x PyTorch. 144 library tests pass.
174.9 ms of kernel time a step.
| Kernel | Per step | Launches |
|---|---|---|
ampere_s1688gemm_bf16_128x128_ldg8_stages_32x1_nt |
19.57 ms | 48 |
cutlass_80_tensorop_s16816gemm_bf16_256x128_32x3_nn |
18.51 ms | 32 |
flash_attention_dkv |
15.45 ms | 16 |
ampere_s1688gemm_bf16_128x128_ldg8_tn |
11.58 ms | 16 |
cutlass_80_tensorop_s16816gemm_bf16_128x256_32x3_tn |
11.40 ms | 32 |
flash_attention_dq |
10.45 ms | 16 |
adam |
8.96 ms | 113 |
| LM head GEMMs (three kernels) | 22.72 ms | 3 |
flash_attention_fwd |
7.24 ms | 16 |
rmsnorm_bwd |
6.66 ms | 33 |
swiglu_bwd |
4.52 ms | 16 |
| memsets | 0.75 ms | 113 |
GEMMs are 60% of the step, the three flash kernels 19%, everything else 20%.
The GEMMs have no headroom. Mapping each cuBLAS and cutlass kernel to its shape
by grid dimensions gives 24.9 TFLOPS for the gate_up weight gradient, 24.0 for
qkv, 23.5 for the gate_up forward, 22.1 for the qkv and down forwards,
21.8 for the attention-output weight gradient and 19.9 for down's, against a
25.5 TFLOPS dense BF16 peak. 78% to 98% of the card.
The flash kernels run at roughly half of peak, which makes flash_attention_dkv
the largest single target left. Four ablations on it, each measured with
examples/flash_probe.rs against a 0.995 ms baseline:
| What was removed | Grad key/value |
|---|---|
| The two global tile loads | 1.023 ms |
The __expf epilogue |
1.027 ms |
The dk/dv mma pair |
0.556 ms |
The score/dpt mma pair |
0.825 ms |
Global traffic and the softmax epilogue cost nothing measurable: only removing
mma work moves the number, so the kernel is bound by the tensor-core and
shared-operand pipeline, not by bandwidth. With DKV_SHARED_BYTES at 37376 and
198 registers per thread the kernel is capped at two blocks per SM either way,
so the fix is not occupancy. The asymmetry between the two mma loops -- 0.44 ms
against 0.17 ms for identical mma and LDS.32 counts -- points at the second
loop's fragment packing, which is where to look next.
Bpe::encode on a sentencepiece-style tokenizer.json (Metaspace), release
build, one paragraph of English repeated. CPU, not a throughput row: it
only shows the shape of the cost.
| Characters | Ids | Time |
|---|---|---|
| 84 | 22 | 0.146 ms |
| 336 | 85 | 1.733 ms |
| 1,344 | 337 | 27.1 ms |
| 5,376 | 1,345 | 432 ms |
Four times the text costs about sixteen times the time: the merge loop is quadratic in the length of the piece, and Metaspace hands it the whole text. Roadmap item E8.
After E8 (heap merge), same file and machine, load average about 3.8. The 50 KB text is this README repeated. Ids hash the same before and after.
| Characters | Before | After | Per character after |
|---|---|---|---|
| 1,008 | 17.2 ms | 0.258 ms | 0.26 µs |
| 5,376 | 451 ms | 1.06 ms | 0.20 µs |
| 50,000 | 43.1 s | 12.6 ms | 0.25 µs |
Not measured: another process held the GPU. The term adds one FP32 multiply to the head kernel's softmax scale and, when the weight is nonzero, one atomic per row. At the default weight of 0 the host path skips its extra pass.
CPU, release build, one run each. The contraction check built a String of
the whole remaining text at every character. The text is
"The cat's hat, isn't it? 42 cats sat. " repeated 2,000 times (78,000
characters, 28,000 words). The words are the same before and after; the
tokenizer tests cover the split.
| Characters | Before | After |
|---|---|---|
| 78,000 | 617 ms | 1.42 ms |
Training itself is not timed: there is no fixed corpus to compare against yet.
Default model (55.2M total, 35.7M active, 6 routed layers of 8 experts, top
2), bench_train_step 4 1024 cuda, RTX 3060 with only gnome-clocks (41 MiB)
on it. Kernel times are nsys totals over the four steps the benchmark runs,
divided by four.
| Before | After | |
|---|---|---|
| Kernel time per step | 82.0 ms | 80.1 ms |
FP32-operand GEMMs (s1688bf16gemm, the expert GEMMs) |
18.4 ms | 6.1 ms |
| All GEMMs | 45.7 ms | 44.6 ms |
swiglu_fwd + swiglu_bwd |
3.5 ms | 3.2 ms |
cast_act |
1.6 ms | 1.4 ms |
| The two gathers | 1.5 ms | 1.3 ms |
| Best of five, tokens/s | 49,802 | 51,631 |
Most of the expert GEMMs' time moved into the BF16 GEMM rows rather than disappearing: at this shape (512 wide, about 1,000 rows per expert) the operand conversion was the smaller part of their cost.
Loss check: the default model at vocabulary 4096 (a byte-level BPE trained
with Bpe::train_json on the repository's markdown, 180k tokens), 16 x 256
tokens a step, Adam 3e-4, 500 steps, same seed, mixed precision on, before
and after this change.
| Step | FP32 routed | BF16 routed |
|---|---|---|
| 0 | 8.3616 | 8.3616 |
| 100 | 6.2913 | 6.2807 |
| 200 | 5.0341 | 5.0238 |
| 300 | 4.3433 | 4.3658 |
| 400 | 3.8325 | 3.7833 |
| 499 | 3.1830 | 3.1844 |
| Held-out | 4.9438 | 4.9239 |
bench_train_step 4 1024 cuda, default model, mixed precision. The MoE
block takes the BF16 normed input, so the router GEMM, the gather and the
shared expert read BF16 operands. The router logits stay FP32.
nsys kernel sum over 4 steps: 80.1 to 76.7 ms a step. The FP32-operand
s1688bf16 GEMM category is gone; the added cast_act costs 1.30 ms.
Throughput (GPU quiet at the time): 49,193, 52,872, 47,815, 52,823 and
52,914 tok/s. Best of five 52,914 against 51,631 after W1 (+2.5%). The
regression floor was not re-recorded: --record also re-records CPU decode,
which needs a quiet machine. strata_head took the GPU right after the
profile, so there are no further runs.
Loss, 500 steps of 16 x 256 on the repository's markdown with a 4096-entry BPE, Adam 3e-4, seed 7:
| Step | This change | W1 run |
|---|---|---|
| 0 | 8.3501 | 8.3616 |
| 100 | 6.2603 | 6.2807 |
| 200 | 4.8686 | 5.0238 |
| 300 | 4.1520 | 4.3658 |
| 400 | 3.6789 | 3.7833 |
| 499 | 2.9449 | 3.1844 |
| Held-out | 4.6321 | 4.9239 |
The W1 column was taken on an earlier copy of the corpus (the roadmap and
baseline files have changed since), so it is a sanity bound, not a paired
comparison. A same-corpus rerun of the previous commit ran out of memory
because strata held 10.5 GB of the GPU.
Release build on the Ryzen 5 7600X, one MoeLayer with d_model 512, 8
experts, top 2, expert width 640, and 2048 random token rows. Timed once in
a temporary test, which was not committed. The forward and backward times
spread by about 1.5x between runs; the two small pieces do not depend on
that spread.
| Piece | Time |
|---|---|
forward_train |
26-35 ms |
backward |
41-67 ms |
softmax_rows (mean of 100) |
91-95 µs |
auxiliary_losses (mean of 100) |
11 µs |
Both D1 and D2 are rejected. Each touches only [tokens, experts], which
is about 1,000 times smaller than the expert GEMMs.
No measurement: SpikeGuard is a scalar check per step on the host, off the
GPU path, and adds no work to the training step itself.
No measurement: no Stable Diffusion VAE checkpoint run was made, and
strata held 10.5 GB of the GPU, so peak memory against a whole decode
could not be compared. The tiled path calls decode per tile, so its peak
is that of one tile's decode plus the FP32 output and weight buffers.
nsys over bench_train_step 4 1024 cuda at 6f48659, idle GPU, grouped
per step:
| Kernel | Calls | ms/step |
|---|---|---|
gather_scale_rows (forward input, backward gradient) |
12 | 1.29 |
scatter_add_scaled (forward) |
6 | 0.93 |
scatter_add_rows (backward) |
6 | 0.92 |
row_dot_scatter (backward) |
6 | 0.63 |
moe_stats |
6 | 0.36 |
| Total kernel time | 78.1 |
The five kernels are 4.2 ms (5.4%) against 17.8 ms when W2 was written,
since W1 halved their bytes and the router work moved. The only redundant
read left is the gathered output-gradient rows, read by both
row_dot_scatter and the backward gather: about 0.5 ms (0.6%), below the
run-to-run spread, so W2 parts 1 and 2 are rejected.
Same nsys report as "MoE permutation passes re-measured" (6f48659, idle
GPU, 4 steps). adam: 828 launches, 20.43 ms in all, so 207 launches and
5.11 ms a step. Median launch 17.3 µs, largest 1.42 ms. The default model has
55,181,824 parameters (parameter_counts().total). Adam reads the value,
gradient and both moments and writes the value and both moments: 28 bytes a
parameter, 32 with the weight-decay re-read. That is 1.55 to 1.77 GB a step,
302 to 345 GB/s against the RTX 3060's 360 GB/s. A fused multi-tensor kernel
can only close the gap to the roofline, at most about 0.8 ms (1%), so A9 is
rejected.
Same idle-GPU nsys report at 6f48659 (bench_train_step 4 1024 cuda,
4 steps), exported with nsys stats -r cuda_gpu_trace. Each MoE layer's
kernels run in a fixed order: the forward gather, the per-expert gate/up and
down GEMMs, then scatter_add_scaled; in the backward pass
row_dot_scatter, the gradient gather, the per-expert GEMMs, then
scatter_add_rows. The GEMMs inside those brackets are the routed experts.
| ms/step | Launches/step | FLOPs/step | TFLOPS | |
|---|---|---|---|---|
| Routed forward | 3.92 | 96 | 53.2 G | 13.6 |
| Routed backward | 6.75 | 192 | 106.3 G | 15.8 |
| All other GEMMs | 32.50 | |||
| Total kernel time | 78.1 |
FLOPs: 6 MoE layers, 8192 routed rows each (4096 tokens, top 2), d_model
512, expert width 352, 2 * rows * 3 * 512 * 352 forward and twice that
backward. Each per-expert GEMM is about 1024 rows, and the grids are 24 to 48
blocks on the RTX 3060's 28 SMs: one partial wave, so the split costs
occupancy. At the dense GEMMs' 22 TFLOPS the routed set would take about
7.3 ms, so about 3.4 ms (4.4%) is recoverable, under the 10 ms bar the plan
set for writing a kernel. bench routes uniformly random tokens almost
evenly, so real text (step 3) may differ.
Default model on the GPU (default mixed precision), Adam 3e-4,
seed 7, 16 x 256 windows of the repository's markdown with a 4096-entry BPE.
Range of the per-expert share across the 8 experts, for each of the 6 MoE
layers, read with TransformerCache::expert_loads on the training batch.
An even split is 0.125.
| Step | Loss | Min-max share per layer |
|---|---|---|
| 0 | 8.355 | 0.058-0.242, 0.011-0.331, 0.035-0.304, 0.016-0.221, 0.019-0.406, 0.005-0.304 |
| 100 | 6.294 | 0.027-0.243, 0.050-0.207, 0.049-0.219, 0.026-0.194, 0.060-0.206, 0.010-0.234 |
| 200 | 5.004 | 0.082-0.148, 0.067-0.157, 0.092-0.153, 0.101-0.158, 0.095-0.194, 0.093-0.169 |
| 300 | 4.290 | 0.090-0.168, 0.093-0.160, 0.082-0.152, 0.055-0.178, 0.095-0.165, 0.103-0.146 |
| 499 | 3.151 | 0.111-0.145, 0.097-0.155, 0.089-0.152, 0.094-0.151, 0.085-0.148, 0.100-0.169 |
Held-out batch at step 499: every share is between 0.089 and 0.158. The load-balancing loss evens the routing out within about 200 steps; after that the busiest expert carries at most 1.35 times an even share, so the uniform benchmark is a fair stand-in for the expert GEMM estimate. The GPU had a 21 MiB test process on it; these are counts, not timings.
Same idle-GPU trace at 6f48659. The cast_act launches between each MoE
layer's forward gather and its scatter are the per-step weight packs: 144
launches (6 layers x 8 experts x gate, up and down) and 0.69 ms a step, out
of 1.31 ms for all 218 cast_act launches and 78.1 ms of kernel time. The
backward pass reuses the packed copies from the cache. Deferred: the fix
changes SwiGlu's public fields.
No measurement. The guard is one pass over the downloaded image to look for non-finite values, which is a host-side scan of the output only; it adds no kernel and does not change the decode itself.
Idle GPU (gnome-clocks only), bench_train_step 4 1024 cuda <head_dim>, the
default model with its 512-wide attention cut into heads of the given width
(2 KV heads), best of three:
| Head width | Fused | Three-kernel path (RUSTING_BRAIN_NO_FLASH) |
|---|---|---|
| 64 | 53,315 tok/s | 33,066 tok/s |
| 128 | 48,456 tok/s | 39,374 tok/s |
Kernel time per step at 128 (nsys, four steps): dkv 8.66 ms, dq 7.21 ms,
fwd 3.45 ms, delta 0.24 ms. At 64 the same build gives dkv 6.72, dq
3.60 and fwd 2.59 ms, unchanged from 6f48659. A first version that
divided a signed index by the width raised fwd_64 from 126 to 146
registers and doubled its time; the unsigned divide restores the old counts.
ptxas at 128: fwd 165 registers, dq 254, dkv 255 with 128 bytes spilled.
Idle GPU at 96ffcc5, bench_train_step 4 1024 cuda, nsys kernel summary over
four steps (77.5 ms of kernel time a step):
| Kernel | ms/step | Rate |
|---|---|---|
ampere_bf16_s16816gemm_bf16_256x128_..._tn (logits) |
5.83 | 23.0 TFLOP/s |
cutlass_80_tensorop_bf16_s16816gemm_relu_bf16_128x128_..._nn (grad of hidden) |
5.21 | 25.8 TFLOP/s |
cutlass_80_tensorop_bf16_s16816gemm_relu_bf16_128x256_..._nt (grad of weight) |
5.06 | 26.5 TFLOP/s |
ce_loss_grad_bf16 |
2.73 | ~288 GB/s |
Each GEMM is 4096 x 512 x 32000 (134 GFLOP). Rejected for speed: see roadmap A6.
Idle GPU at 53ce3b1, bench_decode 512 128 cuda (the default 55.2M model,
FP32 weights on the device, host KvCache): 5.08, 5.16 and 5.53 ms/token,
best 197 tokens/s. The host row on the small model measured 2346 tokens/s in
the same session.
nsys over the whole run (prefill plus 128 tokens), per decoded token:
| Item | Per token |
|---|---|
| Kernel time | ~1.1 ms (LM-head GEMV 0.20 ms) |
| Copy time | ~0.8 ms |
| Host-to-device copies | ~194 |
| Device-to-host copies | ~174 |
| Kernel launches | ~245 |
Kernels and copies explain under 2 ms of the 5.1 ms; the rest is the host
waiting on each device-to-host copy (cuMemcpyDtoHAsync totals 270 ms) and
per-op event, allocation and launch calls. Each layer runs attention on the
host cache, so every linear layer round-trips its activations. This is what
C2's device-resident cache removes.
Idle GPU, bench_decode 512 128 cuda, which now drives Decoder: 2.80, 2.60
and 2.61 ms/token, best 384 tokens/s, up from 197 (1.95x). The host decode
row was not re-measured: four rustc processes from another project held the
CPU (1651-1813 tokens/s, load average 5-6), and the host path is the same
forward_cached call as before.
nsys over the whole run, per decoded token: about 2.0 ms of kernels, of which
cast_act is 0.97 ms. That is Gpu::pack re-narrowing every weight to BF16
on every call, the one ponytail left in decode_block, forward_swiglu and
forward_moe. Device-to-host copies fell from about 174 a token to about 19.
Idle GPU (CPU load average about 5), bench_decode 512 128 cuda: 1.82, 1.60,
1.63, 1.62 and 1.61 ms/token, best 624 tokens/s, up from 384. nsys:
cast_act no longer appears; kernel time is 1.01 ms/token, of which the
LM-head GEMV is 0.20 ms. The remaining 0.6 ms of the 1.6 ms step is launch
and host overhead, which C1 (CUDA graphs) targets.
Idle GPU checked before and after every run (CPU load average about 2),
bench_decode 512 128 cuda <sequences>, default model. Three runs each,
ms per step and tokens/s over all sequences:
| sequences | ms/step | tokens/s |
|---|---|---|
| 1 | 1.62, 1.62 (first run 2.42, warm-up) | 618 |
| 4 | 2.01, 2.01, 2.02 | 1994 |
| 16 | 3.08, 3.09, 3.09 | 5194 |
| 64 | 7.04, 7.04, 7.05 | 9093 |
One sequence is unchanged from step 3 (624). A step costs 0.39 ms more at 4 sequences than at 1 and 5.4 ms more at 64, so the single-sequence step is almost all fixed launch overhead and weight reads.
No measurement: sample_consistency costs one model call per step like the
other solvers, and its own arithmetic is elementwise. The GPU was held by
another process (RustingStrata, 10.5 GB) this iteration.
No measurement: one model call per step like the other solvers, plus elementwise arithmetic. The GPU was held by another process (RustingStrata) this iteration.
No measurement: the GPU was held by another process (RustingStrata, 10.5 GB)
this iteration, and the speedup depends on how often a real draft agrees
with its target, which the random test models do not show. A round costs one
target pass over lookahead + 1 rows plus lookahead draft passes.
Profile first, idle GPU at 0f28b3e, nsys over bench_decode 512 128 cuda:
the LM-head GEMV (gemv2T_kernel_val on FP32 weights) took 30.8 ms over 129
calls, median 200 us, a fifth of the decode's kernel time; every other
projection already read BF16 packs. After narrowing the head once per
decoder, the same kernel on BF16 weights has a median of 120 us.
Same-session A/B, idle GPU (only gnome-clocks) before and after, CPU load
average about 6, bench_decode 512 128 cuda, warm runs:
| Build | ms/token | tokens/s |
|---|---|---|
0f28b3e (FP32 head) |
1.60, 1.60 | 626 |
| BF16 head | 1.50, 1.50 | 666 |
A background run during a cargo build on the same machine gave 155-488
tokens/s and was discarded.
No measurement: export is a one-off file write, and correctness was checked
against the gguf Python reader instead. The GPU was held by another process
(RustingEngine tests) this iteration.
nsys over 129 decode calls: 228,466 cuStreamWaitEvent (65 ms), 124,805
cuEventRecord (38.6 ms), 24,274 cudaLaunchKernel (50.5 ms). The event
calls come from cudarc's per-CudaSlice tracking in multi-stream mode.
Same-session A/B, bench_decode 512 128 cuda, idle GPU (gnome-clocks only)
before and after, second (warm) round:
| Build | 1 sequence | 16 sequences |
|---|---|---|
74fd609, tracking on |
1.53 ms, 655 tokens/s | 3.03 ms, 5280 tokens/s |
| tracking off | 1.52 ms, 656 tokens/s | 2.93 ms, 5455 tokens/s |
Not adopted: no gain at batch 1 and +3% at batch 16 do not pay for a global unsafe switch.
examples/checkpoint_memory.rs <batch> <seq_len>: default model, every block
checkpointed or none, peak of the device's default memory pool over one
training step (CU_MEMPOOL_ATTR_USED_MEM_HIGH, weights and optimizer state
included), then the mean of three timed steps. GPU idle (gnome-clocks only)
before and after; host load average about 11 from other work.
| Shape | Peak kept | Peak checkpointed | ms/step kept | ms/step checkpointed |
|---|---|---|---|---|
| 4 x 512 | 1508 MiB | 1078 MiB | 47.9, 43.6 | 53.3, 53.8 |
| 8 x 512 | 2037 MiB | 1251 MiB | - | - |
| 4 x 1024 | 2037 MiB | 1251 MiB | - | - |
| 16 x 512 | 2843 MiB | 1353 MiB | 131.7, 131.8 | 162.1, 160.9 |
Each further 8 x 512 tokens costs about 102 MiB checkpointed against 806 MiB
kept. The recompute adds 23% at 16 x 512, under the 30% the item estimated.
An earlier timing run overlapped a cargo clippy build and was discarded.
nsys of bench_decode 512 128 cuda 1 at 1c3dd1f: per layer and token,
eight FP32 score GEMVs (about 3.8 us each), attention_softmax_lse (4.9 us)
and eight FP32 value GEMVs (about 4.3 us), about 70 us, a third of kernel
time. A first fused kernel with one block per (head, sequence) took 105 us a
call at batch 1 (8 blocks, each walking 640 positions serially) and was slower
end to end; splitting the history into 64-position chunks fixed that:
decode_attention 16.7 us and decode_attention_merge 2.7 us a layer. Total
kernel time over the 129 decode calls is 107 ms, against about 195 ms before.
Same-session A/B, bench_decode 512 128 cuda <sequences>, idle GPU
(gnome-clocks only) before and after, second (warm) round, CPU load average
about 5:
| Build | 1 sequence | 16 sequences |
|---|---|---|
74fd609 |
1.51 ms, 663 tokens/s | 3.00 ms, 5335 tokens/s |
| fused | 1.16 ms, 864 tokens/s | 2.35 ms, 6819 tokens/s |
The first round's baseline batch-1 run (2.10 ms) was a cold outlier.
Profile at b014141 (nsys, bench_decode 512 128 cuda 1): the attention
output projection was the last FP32 GEMV in a mixed-precision decode step,
1024 calls at a median 6.0 us (5.7% of kernel time), against 5.5 us for the
BF16 weight GEMVs of four times the size. The merge kernel now writes BF16.
Same-session A/B, bench_decode 512 128 cuda <sequences>, three rounds, run
only once the GPU was idle (gnome-clocks only) and the CPU load average was
under 2. A first attempt ran beside another session's rustc build and a
forever_bear process at load average 9; batch-1 numbers then varied from 349
to 663 tokens/s and were discarded. Warm rounds (the first b014141 batch-1
run, 653 tokens/s, was cold):
| Build | 1 sequence | 16 sequences |
|---|---|---|
b014141 |
1.11-1.12 ms, 893-898 tokens/s | 2.25-2.27 ms, 7035-7100 tokens/s |
| BF16 output | 1.07-1.09 ms, 919-931 tokens/s | 2.25-2.26 ms, 7089-7123 tokens/s |
Profile at 0f20131 before the change: BF16 weight GEMVs 31.9% and the BF16
head 12.2% of a batch-1 decode step's kernel time, the bandwidth-bound part
int8 halves. bench_decode 512 128 cuda <sequences> [int8], three rounds
interleaved, started once the GPU was idle (gnome-clocks only) and the CPU
load average under 2; forever_bear reached the GPU near the end of the
third round, whose numbers match the first two.
| Build | 1 sequence | 16 sequences |
|---|---|---|
0f20131 |
891-909 tokens/s | 6985-7034 tokens/s |
| new, int8 off | 904-910 tokens/s | 7049-7071 tokens/s |
| new, int8 on | 998-1006 tokens/s | 6081-6105 tokens/s |
(One 0f20131 batch-1 run read 717 tokens/s, a cold outlier.)
nsys kernel sums over the 129 decode calls, idle GPU:
| Run | gemv_q8 |
Other |
|---|---|---|
| int8, 1 sequence | 4225 calls, 5.2 us avg, 28.5% | BF16 expert GEMVs 3.8 us avg, 19.2%; decode_attention 18.3% |
| int8, 16 sequences | 4225 calls, 21.8 us avg, 92 ms | decode_attention 99 ms |
| BF16, 16 sequences | - | cuBLAS BF16 GEMV kernels about 78 ms; decode_attention 107 ms |
At 16 rows the int8 kernel re-reads each weight four times (four rows a pass) and spends most of its time on per-row warp reductions and input loads, so it loses to cuBLAS.
A gemv_q8_rows kernel for more than four rows: each block of eight units
staged a 512-wide slice of every input row in shared memory as FP32, and each
lane read four consecutive weights in four 128-wide groups, so all loads were
contiguous across lanes. It passed the int8 parity test, now with a six-token
turn, and dropping one weight in four failed it. bench_decode 512 128 cuda <sequences> int8, three rounds interleaved, started on an idle GPU with the
load average under 2 (RustingStrata reached the GPU near the end; rounds 2
and 3 match round 1):
| Int8 decode | 9561323 |
Shared-memory kernel |
|---|---|---|
| 1 sequence | 953-1000 tokens/s | 984-1002 tokens/s (first run cold, 654) |
| 8 sequences | 4325-4360 tokens/s | 3517-3658 tokens/s |
| 16 sequences | 5978-6093 tokens/s | 4874-4900 tokens/s |
The 16-row input is 16 KB in BF16, so the old kernel's per-warp re-reads
most likely hit L1, and the estimate of L2 traffic behind the change was
wrong. Staging adds two block barriers and 8192 scalar loads a slice for 256
multiply-adds per lane, and 32 KB of shared memory caps a block at three per
SM. Nsight Compute cannot read the counters on this machine
(ERR_NVGPUCTRPERM, needs admin), so the next attempt must start from an
nsys sweep of kernel variants, not from a counter profile.
Profile at 9561323 (int8, one sequence): the routed experts' BF16 GEMVs
were 19.2% of kernel time, the largest share left at full precision. With
this change the experts and the router read int8 as well.
bench_decode 512 128 cuda <sequences> [int8], three rounds interleaved,
idle GPU, load average under 2, CLEAN. The machine ran slower than in the
slice-1 session (int8 off 820 against 905 tokens/s at one sequence), so
compare within this table only:
| Build | 1 sequence | 16 sequences |
|---|---|---|
9561323, int8 on |
893-897 tokens/s | 5568-5617 tokens/s |
| new, int8 off | 818-824 tokens/s | 6389-6518 tokens/s |
| new, int8 on | 971-979 tokens/s | 5626-5726 tokens/s |
(The first new int8 run at one sequence read 831, cold.) Int8 is now +19% over BF16 at one sequence, and still -12% at 16, the multi-row GEMV.
TransformerLm::quantize_lora_base stores every adapted projection's frozen
base as int8 with a scale per row; the device uploads those bytes and
dequantizes them into each weight pack where it used to cast FP32.
checkpoint_memory 4 512 lora, default model, rank-16 LoRA, no
checkpointing, two runs, load average 16 (another loop's rustc; the A/B
ran in the same process):
| Base | Peak device memory | Step |
|---|---|---|
| FP32 | 889 MiB | 46.1, 46.5 ms |
| int8 | 778 MiB | 46.6, 46.3 ms |
The 111 MiB is the adapted weights going from 4 bytes to 1 (about 38M of
the 101.7M; the embedding, the norms and the routers stay FP32). The step
does not slow: the dequantize kernel replaces a cast of the same length.
The BF16 packs a step keeps for its backward pass are unchanged, so at
this shape activations and packs, not the base, are most of the peak; the
saving grows with the model and is largest with checkpoint, which drops
the packs between blocks.
TransformerLm::quantize_lora_base_4bit(Codebook::Nf4): the adapted bases
at 4 bits with a scale per block of 64, uploaded as packed bytes and widened
into each weight pack by dequantize_q4. checkpoint_memory 4 512 lora,
default model, rank 16, RTX 3060, two runs, load average 4-7:
| Base | Peak device memory | Step |
|---|---|---|
| FP32 | 889 MiB | 43.2, 42.9 ms |
| int8 | 778 MiB | 43.2, 42.5 ms |
| NF4 | 762 MiB | 43.3, 42.6 ms |
The 16 MiB from int8 to NF4 matches the arithmetic: about 38M adapted weights go from 1 byte to 0.5625. As with int8, activations and the BF16 packs are most of the peak at this shape.
nsys on bench_decode 512 128 cuda 1 int8, RTX 3060 (360 GB/s), load 4:
gemv_q8 is 43% of kernel time (8069 calls, median 3.0 us). By grid:
512-unit projections (64 blocks) take 3.0 us for 256 KiB, about 87 GB/s;
2816-unit ones 7.1 us; the 32000-unit LM head 60 us for 16 MiB, about
270 GB/s. Only the head is bandwidth-bound, and it is about 6% of a token,
so halving the weight bytes would gain a few percent at most on the default
model. Not built for speed; it remains the way to fit a larger model.
Profile at 68f7c70, bench_decode 512 128 cuda 16 int8 under nsys:
gemv_q8 32.8% of kernel time (12,909 calls, 11.3 us average), the largest
kernel, ahead of decode_attention at 22.9%. At 16 rows it walks the weights
four rows at a time.
gemv_q8_mma (in src/cuda_flash.rs, which is built for compute_80): a
block of eight warps owns 16 weight rows as the mma's A operand and takes up
to 16 BF16 input rows as B, so each weight byte is read once. The warps split
the width in steps of 32 and add their partial sums in shared memory. The
first version gave each warp 16 rows and the whole width: a 512-unit layer
then had 32 warps in flight, and one sequence fell from 1040 to 735
tokens/s, 16 sequences were no faster than the old kernel.
bench_decode 512 128 cuda <sequences> int8, the old kernel against the new
one in the same build, rounds interleaved, gnome-clocks the only other GPU
process, load average 4 to 12 (other loops' builds):
| Sequences | Old int8 GEMV | Tensor-core GEMV | BF16 (int8 off) |
|---|---|---|---|
| 1 | 984-998 | 999-1003 | 837-855 |
| 2 | 1946-1950 | 1619-1914 | - |
| 4 | 3061-3089 | 3235-3246 | - |
| 8 | 4498-4518 | 5132-5148 | - |
| 16 | 5900-5930 | 7535-7700 | 6764-6788 |
(A third round with the load average rising to 12 read 5817 against 6801 at
16 sequences, and 952 against 870 at one.) The tensor-core kernel runs from
three rows on (Q8_MMA_ROWS); one and two rows keep the warp-per-unit
kernel, which ties it there.
quantized::Codebook::{Int4, Nf4} with one f32 absmax scale per block of
64 along a row, two indices a byte (byte j of a block holds weight j low and
weight j+32 high). Storage is 0.5625 bytes a weight at a width blocks divide,
against 1 byte plus a row scale for int8; narrow rows pad to a whole block.
Round-trip RMS error on standard-normal weights: symmetric int4 0.1074, NF4
0.0940 (nf4_loses_less_than_the_even_grid_on_normal_weights).
Host decode, bench_decode 64 64 cpu 1 with a temporary 100M-block-parameter
shape (d_model 1024, 8 layers, heads 16/16/64, d_ff 2816), rounds
interleaved:
| Version | int8 tok/s | nf4 tok/s | Load |
|---|---|---|---|
| Scalar codebook walk | 260-270 | 132-136 | 3-4 |
AVX2 vpermps lookup |
255-271 | 292-331 | 15-17 |
The scalar walk widened one weight at a time and was compute-bound at twice
int8's cost on half the bytes. Holding the 16 levels in two registers and
looking up eight indices per vpermps, blended on bit 3, makes the 4-bit
row cheaper than the int8 one. One outlier int8 round at 135 is left out.
examples/flash_probe.rs (fixed for the _64 kernel names and the causal
argument), batch 4 x 1024, 12 query heads over 4 key/value heads, head dim
64, RTX 3060, two interleaved rounds, load 9:
flash_attention_dkv_64 |
Registers | Shared | Probe |
|---|---|---|---|
| previous kernel | 198 | 37,376 B | 0.990, 0.992 ms |
B operands by ldmatrix.x4 |
168 | 37,376 B | 0.973, 0.974 ms |
... and ldmatrix.x4.trans from the query-major tiles |
168 | 18,944 B | 0.800, 0.800 ms |
The ldmatrix loads alone change little, so the second loop's fragment
packing was not the cost. Dropping the transposed tiles does two things:
the load phase stops issuing scattered 16-bit stores, and the block fits in
19 KB, so 168 registers and 19 KB allow three blocks an SM where 37 KB
allowed two.
In the model, nsys on bench_train_step 4 1024 cuda: flash_attention_dkv
847 to 637 us a launch; the forward (~484/444 us), dq (~335/325 us) and
delta (31 us) kernels are unchanged code and within noise. Throughput,
interleaved rounds, load 9 to 24 (other projects on the CPU; the first
before-round, 41661, was cold):
| Build | tok/s |
|---|---|
| before | 52149, 52208 |
| after | 52536 to 53452 |
The whole-step gain (about 1%) is smaller than the kernel time suggests;
the load was high, so re-measure on an idle machine before quoting it.
dkv_128 still sits at 255 registers (stack 128 to 184 bytes), so it gains
the shared memory but not the occupancy.
examples/checkpoint_memory.rs 4 512 lora, default model, rank 16, RTX 3060,
two interleaved rounds, with fold_lora returning at once in a throwaway
build (the adapter then has no effect; timing only). Load 6 to 9.
| Base | Fold | No fold |
|---|---|---|
| FP32 | 43.3, 42.6 ms | 40.5, 41.0 ms |
| int8 | 42.8, 42.9 ms | 40.5, 41.0 ms |
| NF4 | 43.4, 43.4 ms | 40.8, 40.8 ms |
The fold costs about 2.2 ms a step (5%). Moving the adapter to its readers
would replace one [units, inner] write per projection with a [rows, units] write in the forward pass and a [rows, inner] write in the input
gradient, at 2048 rows here, so it is not taken.
nsys on bench_decode 512 256 cuda at 79f9aac: 1.13 ms a token, of which
about 0.65 ms kernels. API side, 19 cuMemcpyDtoHAsync and 39
cuMemcpyHtoDAsync a token. Each MoE layer downloaded expert_of,
gate_of and the auxiliary-loss sums, so the stream drained six times a
token.
| Build | 1 seq | 4 seq | 16 seq |
|---|---|---|---|
| before | 784, 924 | 2725, 2730 | 6929, 6933 |
| every expert dense, zero gates | 877, 876 | 2623, 2634 | 7054, 7079 |
gemv_routed |
1357, 1363 | 3793, 3795 | 7918, 7930 |
Tokens/s, default model, 512-token prompt, 256 new tokens, RTX 3060. The
dense variant was measured in a separate session (its before-row read
711-897, 2620-2628, 6153-6731 there) and cut the downloads to one a token,
but its expert GEMVs rose from 63 to 138 ms of kernel time over the run.
Load 4 to 10 for the gemv_routed rounds, 9 to 20 for the dense ones.
time_a_t5_base_sized_encoder (release, ignored): a random T5-base-sized
gated encoder (d_model 768, 12 heads of 64, 12 layers, mlp 2048) over 512
tokens. Host T5Encoder::forward 528 ms; DeviceT5::forward, five runs
after a warm-up, 33.8, 29.0, 30.1, 28.6, 29.9 ms including the upload of the
embeddings and the bias and the download of the states. Worst difference
0.0073 against a largest state of 0.39. Load 32 (other jobs on the host),
RTX 3060.
time_a_dit_over_a_512_image (release, ignored): a random FLUX-shaped
denoiser (d_model 768, 6 heads of 128, 2 double and 4 single blocks, mlp
ratio 4) over 32x32 patches and 256 text tokens. Host Dit::forward_with
1.27 s; DeviceDit::forward_with, five runs after a warm-up, 15.7, 16.0,
15.9, 15.9, 17.1 ms, including the uploads and the download. Relative L2
error 0.0076. Load 7 to 5, RTX 3060.
On the tiny parity model the whole-model relative error is 0.020 at F32 and 0.008 at Q8: Q8 weights are read exactly on both sides, so the gap is BF16 weight rounding, not the kernels, which match the host to 1e-4.
BF16 against int8 weights, same test, after C6 part 3 (2026-10-07, load 8 to 15): F32/BF16 device 130 (first, cold), 23.7, 16.2, 15.7, 15.9 ms; Q8 device 15.8, 15.8, 16.4, 15.8, 16.8 ms. Hosts 1.16 s and 1.47 s. At 1280 tokens the GEMMs are compute-bound, so reading one byte a weight instead of two buys no time; it halves the weights' memory.
nsys profile --trace=cuda on cuda_image_bench /home/vasylt/Code/ai/ml/models/pony-v6-xl
with SIZES=1024 CPU=0 QUANTIZE=1: one UNet pass 0.251 s (first 0.40 s),
VAE decode 0.53 s, peak 7.45 GB. Kernel time 1847 ms: FP16 GEMMs 25.2% +
22.2% + 5.9% + 2.3%, im2col 18.7% (2804), image_attention_fwd 7.2%,
group_apply 3.6%, add_plane_bias 3.5%, transpose 0.96% (70 launches,
17.7 ms). Load 7.5 to 5.7.
Int8 decode still sent every MoE layer through forward_moe, whose routing
download drains the stream six times a token (the profile at 79f9aac, 19
device-to-host copies a token). gemv_routed_q8 is gemv_routed over int8
expert rows with a scale per row, so int8 decode now skips the round trip as
BF16 decode already did. bench_decode 512 128 cuda <n> int8, default
model, old and new binaries interleaved, three rounds, load average 15 to
40 (other loops building), GPU lock held:
| Build | 1 sequence | 16 sequences |
|---|---|---|
f97c1fc, int8 on |
1014-1049 tokens/s | 7732-7895 tokens/s |
| new, int8 on | 1535-1613 tokens/s | 7675-7804 tokens/s |
| new, int8 off (one run) | 1074 tokens/s | 8012 tokens/s |
Int8 is now +46% over BF16 at one sequence in this session; 16 sequences are unchanged, where the dense GEMMs dominate.
A note on method first. Under nsys the decode loop runs at 1.87 ms a
token against 0.62 ms unprofiled, because the trace slows the host's ~1000
event calls a token; the GPU then idles most of the time, its clocks drop,
and every small kernel reads about twice as slow. decode_attention read
11.5 us at a 32-160 position history and 13.6 us at 640 in the model trace,
and 12.9 us at 64 in a standalone loop with a short warmup, but 5-8 us in
the same loop after 30000 warmup launches. Compare kernel shares from such
a trace, not absolute microseconds, and time single kernels warm.
Standalone loop (decode_attention alone, 8 heads, 2 KV heads, width 64,
one sequence, L2-resident cache, 1000 launches after 30000 warmup):
| History | Before | After |
|---|---|---|
| 64 | 7.64 us | 5.05 us |
| 640 | 8.46 us | 6.59 us |
| 1664 | 20.46 us | 14.00 us |
Bisected at 64 positions: dropping the value pass saves the most (5.05 to 3.00), dropping the score pass 0.95 us.
bench_decode <prompt> 128 cuda 16, BF16, 45a8ef1 against the change,
three interleaved rounds, load average 18 to 33:
| Prompt | Before | After |
|---|---|---|
| 512 | 7780-7980 tokens/s | 8731-8835 tokens/s |
| 1536 | 4335-4380 tokens/s | 4548-5294 tokens/s |
One sequence (897-1133 against 742-1223) is inside the noise at this load.
nsys profile -t cuda bench_decode 512 128 cuda 1 int8, API call counts over
the whole run (129 decode calls plus the prompt):
| Call | b7aa1b5 |
Norms in the memo |
|---|---|---|
cuMemcpyHtoDAsync |
3092 | 916 |
cuMemAllocAsync |
19980 | 17804 |
cuLaunchKernel |
16891 | 16891 |
bench_decode, three interleaved rounds, load average 29 to 33 (the first
b7aa1b5 round ran cold and is left out): BF16 one sequence 1237-1242
against 1208-1358 tokens/s, int8 one sequence 1620-1677 against 1428-1681,
16 sequences 8700-8911 against 8904-8967. No change beyond the noise; the
uploads overlapped the kernels. Kept as a step toward graph capture.
nsys profile -t cuda bench_decode 512 256 cuda 1 int8, kernels in a
100 ms window near the end of decode, load 6-7:
| tokens/s under nsys | GPU busy | API calls a token | |
|---|---|---|---|
| own stream (multi-stream mode) | 637 | 25% | ~1000 cuStreamWaitEvent, ~500 cuEventRecord, ~260 event create/destroy |
| legacy default stream | 971 | 62% | none of those |
Unprofiled, five interleaved rounds, load 7 to 10 (tokens/s):
| run | before | after |
|---|---|---|
| 1 sequence BF16 | 1200-1406 | 1406-1428 |
| 1 sequence int8 | 1676-1712 | 1711-1747 |
| 16 sequences BF16 | 8746-8840 | 8772-8899 |
bench_train_step 4 1024 cuda |
53392-53428 | 53279-53630 |
The host event work overlapped the kernels, so the unprofiled gain is small.
Profiles taken under nsys before this change overstate how host-bound decode
is: nsys slows the API calls the tracking made. An earlier test that called
disable_event_tracking (C1 note) measured the same effect and also found it
within noise.
flash_attention_dq_128 kept a transposed copy of the key tile in shared
memory beside the key and value tiles: 52 KiB a block, one block per SM on
the RTX 3060 (100 KiB). The copy is gone; the dS.K product reads the key
tile transposed with ldmatrix.trans. ptxas registers are unchanged
(dq_128 254, dq_64 163 to 164), so the gain is occupancy.
nsys, bench_train_step 4 1024 cuda 128, four steps, load 5-8:
| kernel | before (median) | after (median) |
|---|---|---|
flash_attention_dq_128 |
1075 us | 620 us |
Four interleaved rounds, load 5 (tokens/s):
| head width | before | after |
|---|---|---|
| 128 | 48517-48770 | 49954-50304 |
| 64 | 53257-53636 | 52856-53643 |
Then flash_attention_dkv_128: the query tile in two halves of 32 so half
the scores are live, ptxas spills 180 to 48 bytes (255 registers either
way; dkv_64 keeps one pass, 168 registers, no spill). Clean GPU, load 3-4,
four interleaved rounds:
| before | after | |
|---|---|---|
dkv_128 median |
891 us | 783 us |
| step, heads of 128 | 49921-50451 tok/s | 50471-50589 tok/s |
| step, heads of 64 | 53134-53751 tok/s | 53558-53912 tok/s |
Splitting at 64 too made dkv_64 670 to 699 us, so the split is width 128
only.
Quiet GPU, load 4-6, SIZES=512,1024 CPU=0 QUANTIZE=1 cuda_image_bench /home/vasylt/Code/ai/ml/models/pony-v6-xl: one UNet pass 0.071 s at 512x512 and
0.254 s at 1024x1024 (VAE decode 0.13 and 0.54 s, peak 7.79 GB). Four times
the latent pixels costs 3.6 times the time, so a batch-2 pass at 512 can be
at most about 10% quicker than two passes; at 1024 the pass is already as
large. Not built.
Pony Diffusion V6 XL (SDXL), 512x512, prompt "score_9, a lighthouse on a
rocky coast in fog". Reference: ComfyUI at 1a510f04, run on the CPU in FP32
from the same diffusers files (UNet through comfy.sd.load_diffusion_model
and BaseModel.apply_model, CLIP through comfy.sd.load_clip, VAE through
convert_vae_state_dict). Ours: device Q8 unless marked. Relative L2 error.
| Part | Before C11 | After C11 |
|---|---|---|
| UNet eps, sigma 14.6 / 1 / 0.1 (same inputs) | 0.0010 / 0.0013 / 0.0060 | unchanged |
| VAE encoder, host F32 | 0.012 | unchanged |
| CLIP-L hidden states, host F32 | 7.5e-5 | unchanged |
| CLIP-G hidden states, device Q8 | 1.07 | 0.025 |
| CLIP-G hidden states, host F32 | 0.021 | |
| Pooled vector, device Q8 | 0.0066 | 0.0066 |
| Image-to-image, same prompt, strength 0.3, rel to the input image | 0.37 | 0.147 |
CLIP-G positions after the end token were off by rel 3 to 6 before the fix
(wrong padding token); after it, 0.07-0.09 in FP32, which is the tanh GELU
against the checkpoint's exact gelu. The VAE check against ComfyUI's
default CPU dtype (BF16) reads rel 0.108; that is ComfyUI's rounding. Load
1.3-1.7.
Tutorial program (tutorials/code/18_text_to_image.rs), 1024x1024, DPM++ 2M
25 steps, Q8, RTX 3060 shared with another job, load 1.4-1.7: loaded in 6.3 s,
text-to-image 13.2 s, image-to-image at strength 0.5 117.2 s, nearly all of
it the CPU VAE encoder (C12).
pony-v6-xl VAE, RTX 3060 shared with another job, load 2.3-2.7, three runs each after one warm-up inside the same process:
| Encode | CPU (before) | Device FP16 |
|---|---|---|
| 512x512 | 10-18 s | 0.060-0.063 s |
| 1024x1024 | 117-140 s | 0.249-0.253 s |
Against ComfyUI's FP32 CPU encoder on the same 512x512 image the device latent reads rel 0.066 (per-channel std and mean match to 3 digits; best scale 0.998). Host F32 reads 0.012 and ComfyUI's default BF16 0.108. Tutorial program, load 2.3-3.0: text-to-image 13.2 s, image-to-image at strength 0.5 6.9 s (117.2 s before).
examples/flash_probe.rs (batch 4 x 1024, 12 heads over 4, width 64, best of
five), idle GPU, load 2.2-2.8. The query-gradient kernel already reads its
key tile with ldmatrix.trans (A3), so only the forward kernel had A1's
transposed copy left.
| Forward kernel variant | ms |
|---|---|
| As is: scalar loads, values stored transposed | 0.446 |
Values key-major, read with ldmatrix.x4.trans |
0.440-0.442 |
| The same plus paired (32-bit) tile loads | 0.565-0.570 |
| Paired loads only, values still transposed | 0.492-0.493 |
Values key-major plus ldmatrix.x4 for the score keys |
0.507-0.509 |
The query gradient reads 0.633 ms and the key/value gradient 0.806 ms in the same runs. The best variant saves 5 us a call, about 0.04 ms of a 77 ms step, under the run-to-run spread, so nothing was kept.
examples/profile_shape.rs 4 2048 4096 30 (mixed, packed) and
bench_train_step 4 1024 cuda, old and new binaries alternated twice under
the GPU lock, load 0.8-1.0.
| Run | Before | After |
|---|---|---|
| Shape VAE step | 0.130-0.131 s | 0.085 s |
| LM step, 4 x 1024 | 53461-53562 tok/s | 53451-53487 tok/s |
Before, nsys showed a 44-47 ms GPU-idle cudaStreamSynchronize each shape
step, after a burst of cuMemFreeAsync: the default pool released its memory
at every sync (threshold 0). keep_freed_memory in device_context sets
CU_MEMPOOL_ATTR_RELEASE_THRESHOLD to u64::MAX. The LM step allocates
little per step, so it does not move.
checkpoint_memory <batch> <seq> adam8, default model, Adam, mixed
precision, under the GPU lock, load 1.5-2.1. Peak is the device pool's high
mark over one step, resident weights included.
| Shape | FP32 state | 8-bit state |
|---|---|---|
| 4 x 1024, peak | 2037 MiB | 1723 MiB |
| 4 x 1024, step (two runs) | 75.6, 75.6 ms | 75.2, 75.6 ms |
| 16 x 256, peak | 2037 MiB | 1723 MiB |
| 16 x 256, step | 68.4 ms | 68.1 ms |
adam8_loss_check 2000: the repository's markdown (192,562 training
tokens, 21,395 held out), BPE 4096, default model, 16 x 256, Adam 3e-4,
seed 7.
| Step | FP32 state | 8-bit state |
|---|---|---|
| 0 | 8.3532 | 8.3532 |
| 250 | 4.4934 | 4.4767 |
| 500 | 2.9560 | 2.9508 |
| 1000 | 0.6764 | 0.6204 |
| 1500 | 0.2060 | 0.2075 |
| 1999 | 0.1723 | 0.1824 |
| Held-out | 6.5538 | 6.6180 |
Both runs memorize the small corpus (42 passes), so the held-out loss rises for both; the 8-bit run stays within 1% of it.
examples/profile_shape.rs 4 2048 4096 30 (mixed, packed) and
bench_train_step 4 1024 cuda, old and new binaries alternated under the GPU
lock, load 2.3-4.0 (8 on the 15-minute average).
| Run | Before | After |
|---|---|---|
| Shape VAE step | 0.084-0.087 s | 0.062-0.065 s |
| LM step, 4 x 1024 | 54308-54654 tok/s | 54126-54511 tok/s |
The first build of the separate-buffer flash kernels cost the LM step 1-2%
(old 53126-54155, new 52480-53103 tok/s over five pairs). nsys put it all in
flash_attention_dkv_64: 651 to 811 us a call. Its registers went 168 to
173, which at 128 threads is two blocks per SM instead of three.
__launch_bounds__(128, 3) at width 64 holds it at 168 with an 8-byte spill:
690 us, while flash_attention_dq_64 went 485 to 460 us, so the backward pair
is even and the step is back within run-to-run spread.
SIZES=1024 CPU=0 QUANTIZE=1 cuda_image_bench pony-v6-xl, old and new
binaries alternated three rounds under the GPU lock, load 3.1-4.5.
| UNet pass | VAE decode | |
|---|---|---|
im2col + one GEMM |
0.244, 0.244, 0.247 s | 0.52, 0.52, 0.52 s |
im2row + k GEMMs |
0.241, 0.247, 0.242 s | 0.49, 0.49, 0.49 s |
nsys kernel sums (same command): before, im2col 346 ms over 2804
launches (18.5%); after, im2row 176 ms over 1038 launches (9.4%), with the
128x128 nn GEMM rising from 466 ms to 614 ms over 2185 launches. The UNet's
3x3 convolutions at 128x128 are too small for the gather to matter; the
decoder's at 512 and 1024 pixels gain.
BatchDecoder now pads uneven prompts and reads a padding flag per cached
position in decode_attention. Cost on equal-length decode, where every
flag is 1: bench_decode 512 256 cuda 1 int8 and bench_decode 512 128 cuda 16 (BF16), bfc294e against the change, three interleaved rounds, load 2.0
to 2.5, tokens/s:
| run | before | after |
|---|---|---|
| 1 sequence int8 | 1742-1751 (first round 1098, warm-up) | 1735-1747 |
| 16 sequences BF16 | 9276-9319 | 9181-9237 |
One sequence is unchanged. Sixteen sequences are about 1% slower in every
round: the flag load sits in the score loop of decode_attention. Left as
is; skipping the load when no sequence is padded would win it back if it
matters.
The device decode cache now starts empty and doubles as the session grows,
up to max_seq_len. Memory, from the allocation sizes, on the default
model (8 layers, rows of (8 + 2 x 2) x 64 FP32 values, 24,576 bytes a
position a sequence):
| session | before | after |
|---|---|---|
any, at max_seq_len 2048 |
50.3 MB a sequence | - |
| 20-token prompt | 50.3 MB | 0.5 MB |
| 512 + 256 tokens | 50.3 MB | 25.2 MB |
| 16 sequences, 512 + 128 | 805 MB | 403 MB (capacity 1024) |
Decode throughput, 54f7581 against the change, interleaved rounds, load 1.3 to 1.5,
tokens/s, first round of each binary dropped as warm-up:
| run | full allocation | grown |
|---|---|---|
bench_decode 512 256 cuda 1 int8 |
1744 | 1711-1742 |
bench_decode 512 128 cuda 16 |
9154-9223 | 9094-9144 |
bench_decode 16 1024 cuda 1 |
1468-1470 | 1465-1473 |
Sixteen sequences are 0.3-0.8% lower, about one rounding step of the ms/step figure; the two growths (512 then 641 positions here) copy the prefix once each.
The fused decode attention block now serves every query head of a key-value
head (four on the default model), so each cached key and value is loaded
once a block instead of four times. bench_decode, 2f02760 against the
change, two interleaved rounds, load 1.2 to 1.5, tokens/s:
| run | per query head | per key-value head |
|---|---|---|
64 256 cuda 1 int8 |
1774 (first round 1038, warm-up) | 1783-1810 |
512 256 cuda 1 |
1420-1449 | 1467-1477 |
2000 128 cuda 1 |
1259-1291 | 1389-1390 |
4000 256 cuda 1 int8 |
1197-1212 | 1425-1437 |
64 128 cuda 16 |
11964-12058 | 12523-12629 |
512 128 cuda 16 |
9163-9182 | 10348-10452 |
2000 128 cuda 16 |
5016-5024 | 6906-6915 |
Before, nsys put decode_attention at 23% of kernel time for the 4000-token
int8 run and 31% for 2000 x 16. The grid shrinks fourfold, but at 64
positions it still runs no slower.
bench_decode now times a second prefill after a warm-up one. 65355e8
against the change, two interleaved rounds plus a recheck, load 2.8 to 3.4,
prefill ms:
| run | score matrix | flash |
|---|---|---|
64 16 cuda 1 |
3.5-4.1 | 3.0 |
512 16 cuda 1 |
6.0-7.0 | 4.9 |
2000 16 cuda 1 |
27.2-30.0 | 12.5-12.6 |
4000 16 cuda 1 int8 |
96.5-98.4 | 32.1-37.6 |
512 16 cuda 16 |
43.8-44.9 | 35.0-42.5 |
1000 16 cuda 16 |
109.9-111.1 | 67.1-69.5 |
2000 16 cuda 16 |
363.5-365.2 | 148.5-148.6 |
Decode steps match (512 128 cuda 16 1.53-1.60 ms/step both); one round
showed 1.78 ms/step for 512 16 cuda 16 with the change, which the recheck
did not repeat.
Padded prompts, the same day: bench_decode with every other prompt half
as long (a throwaway copy of the example), 6566c19 against the change,
two interleaved rounds, load 2.3 to 2.5, prefill ms:
| run | score matrix | flash |
|---|---|---|
512 32 cuda 2 |
9.0-16.8 | 6.8 |
2000 32 cuda 2 |
59.1-91.5 | 19.0-19.1 |
512 32 cuda 16 |
47.2-48.4 | 31.7-32.3 |
2000 32 cuda 16 |
438.3-439.7 | 125.1-125.9 |
Decode steps match within noise (2000 x 16: 2.44-2.50 against 2.45-2.48 ms/step).
Later turns, the same day: bench_decode 2000 8 cuda <n> "" <turn> times a
second feed of turn ids a sequence after the decode (warmed by the same
turn on the warm-up decoder). Flash on every multi-row turn against the
score matrix, three interleaved rounds, load 19 to 38 (another loop's
rustc), ms:
| sequences, turn | score matrix | flash |
|---|---|---|
| 1, 8 | 1.9-3.6 | 2.2-4.1 |
| 1, 32 | 3.3-4.1 | 4.4-4.8 |
| 1, 128 | 5.1-6.1 | 4.1-6.7 |
| 1, 512 | 18.4-19.9 | 8.7-16.5 |
| 16, 8 | 10.0-11.1 | 24.8-26.0 |
| 16, 32 | 13.3-13.4 | 26.0-26.2 |
| 16, 128 | 34.9-35.4 | 32.5-34.1 |
| 16, 512 | 140.5-160.5 | 60.1-70.2 |
The flash path converts the whole cached history to BF16 each layer, which
a short turn does not earn back, so later turns take it from 128 rows. The
training kernel's extra mask test costs nothing: bench_train_step 4 1024 cuda 54321-54404 tokens/s, load 17.
Under mixed precision the device decode cache keeps only keys and values,
in BF16: 512 bytes a position a layer on the default model, against 3072
for FP32 [query, key, value] rows. bench_decode, 6e47a70 against the
change, three interleaved rounds, first (cold) round of the old build
dropped, load 5.4 to 6.5:
| run | FP32 rows | BF16 keys and values |
|---|---|---|
2000 128 cuda 1, ms/step |
0.70-0.71 | 0.71 (one round 0.96) |
2000 128 cuda 16, tokens/s |
7029-7041 | 7315-7365 |
2000 8 cuda 16, ms/step |
2.88-2.89 | 2.30-2.31 |
| then an 8-id turn, ms | 9.9-10.0 | 9.1-9.2 |
Multi-row turns read the cache in place through the flash kernel instead of a BF16 copy of the whole history, which is where the small-turn and short decode gains come from.
One adapter training step on CUDA (forward with adapters, device tape,
backward to every adapter gradient), cuda_mmdit::tests::time_a_dit_adapter_step
in a release build: d_model 768, 6 heads of 128, 2 double and 4 single
blocks, 1024 image tokens and 256 text tokens, adapters on every block
projection. Load 4.2.
| step | ms |
|---|---|
| first | 130 |
| next five | 56, 56, 56, 48, 45 |
1612 MiB of device memory in use. FLUX.1 itself is not timed: its BF16 weights (24 GB) do not fit the 12 GB card.
examples/chop_block_step.rs times one forward and backward pass through
RustingChop's seven convolution blocks (conv 3x3 without bias, GroupNorm with
min(8, c) groups, SiLU; strides 2, 1, 2, 1, 2, 1, 2; widths 32 to 128) plus
spatial softmax and global average pooling, on a 128 x 96 x 192 x 15 batch.
Median of five steps after one warm-up, RTX 3060, at c8851ce plus the example.
Load average 1.5 at the end of the run (another project's GPU job had just
released the lock).
| Path | Step (ms) |
|---|---|
| per-layer: each conv uploads its input and downloads its output, GroupNorm and SiLU on the host | 1002.3 |
| device-resident: only the input goes up | 98.9 |
The device-resident path is 10.1x faster. The head after the pooled features is not in either number.
One adapted forward and backward pass of Stable Diffusion 1.5's UNet at
512x512 (a 64x64 latent, 77 prompt tokens, batch 1) on the RTX 3060, FP16:
DeviceUnet::forward_adapted keeps the tape, backward_adapted sums every
adapter gradient on the device. Rank-2 adapters on to_q, to_k, to_v and
to_out.0 of every attention. The weights are random (the 1.5 checkpoint is
not on this machine); the cost of the step is its shapes. Test
cuda_image::lora::tests::an_sd15_step_at_512_is_timed (ignored), release
build, one warm-up and five timed steps. Load average 10.4 at the start and
8.1 at the end.
| Part | ms |
|---|---|
| forward with tape | 65.1 |
| backward to adapter gradients | 125.9 |
| step | 191.0 (all five 191.0) |
The optimizer update and the noise sampling are not in the number.