Skip to content

roadmap.md

Latest commit

 

History

History
1463 lines (1266 loc) · 153 KB

roadmap.md

File metadata and controls

1463 lines (1266 loc) · 153 KB
Raw

RustingBrain Roadmap

Where the library goes after 2.0.2. This file is the index: it names every piece of work worth doing, sizes it, orders it, and points at the document that holds the detail when one already exists. It does not repeat those documents.

Detail lives in For
featuresPlan.md Feature candidates with costs, and the "not on this list" rationale
docs/improvement-plan.md MoE training throughput, work items W0-W9
docs/optimize-plan.md Dense training throughput, Tasks 0-9
docs/baseline.md Measured numbers and the current per-kernel profile
2Dto3DPlan.md Image-to-3D stages A-H and the deferred texture work
docs/interoceptive_spec.md The interoceptive trading agent and its acceptance criteria

How to read this

Effort is for one person: S is days, M is one to two weeks, L is a month or more.

Priority:

  • P0: blocks a release, or fixes something wrong today.
  • P1: the next things to build. Most users will notice them.
  • P2: worth doing once P1 is done, or when someone asks for it.
  • P3: parked. Written down so nobody has to argue it again.

Ground rules that apply to every item:

  1. Profile first. No performance item starts without a fresh profile (nsys or examples/profile_step.rs) that shows the target is where the time goes. The estimates below come from earlier profiles. Re-measure before trusting them.
  2. Training runs on CUDA. Every training-throughput number and every end-to-end training check runs on the GPU path. The CPU path is the reference for parity tests, not a training target.
  3. Parity before speed. A new kernel lands only with a test that compares it against the CPU reference or the kernel it replaces.
  4. Read the reject lists first. docs/improvement-plan.md has two lists, "Already landed — do not redo" and "Already measured and rejected — do not redo". Reopen an item from them only with a new measurement.

Where things stand

  • Dense transformer training on one RTX 3060: 21.8k tokens/s at 4 x 1024 for a 101.7M-parameter model. That ties or beats PyTorch with torch.compile at every measured shape. GEMMs run at 78-98% of the card's BF16 peak.
  • Kernel time per step: GEMMs 60%, the three flash-attention kernels 19%, everything else 20%. The flash kernels run at about half of peak, so they are the largest target left.
  • MoE costs 1.37x a dense step for 1.14x the arithmetic. The routed path runs in BF16 under mixed precision (A7 W1).
  • The library has MoE, LoRA (CPU and CUDA), quantization-aware training, int8 inference, ViT, a masked-LM (bidirectional) mode, Lion and Adam, streaming datasets, the text-to-image pipeline (MMDiT, UNet, VAE, CLIP, T5, LLaMA-family prompt encoders), image-to-3D, ONNX import and export, and the interoceptive RL agent.
  • Single GPU only. No multi-GPU, no distributed training, no CPU offload.
  • Structural tax: no autograd and no N-D tensor type. CPU and CUDA paths are written twice. Every new layer is two forward passes, two hand-written backward passes and a finite-difference test.

Milestone 2.1 — finish what is open (P0)

Release the work that is already in the tree, and close what is half done.

# Item Effort Notes
2.1.1 Land the uncommitted diff done S Landed as the topic commits from 1e0de06 to 92c732a; the tree is clean. Original note:
2.1.2 Interoceptive acceptance criteria done S Clippy is clean with and without cuda. The stress-reflex test already existed (agent.rs:246). tests/interoceptive_no_alloc.rs counts allocations over 1000 steps, resets and replay samples, and a mutation check confirmed that it fails when one allocation is added. Original scope: docs/interoceptive_spec.md §5: zero warnings, the stress-reflex logit test (margin_stress 0.05 against 0.95, EmergencyFlatten logit rises above the threshold), and no heap allocation per step in the replay loop. Prove the last one with a counting allocator in a test, not by reading the code.
2.1.3 Fix Matrix::random done, smaller than written S Network already seeds its own He init, and no production code calls Matrix::random. It is now documented as unseeded test filler. Original note: src/matrix.rs:98 draws from thread_rng over 0..1, ignores the seed, and makes every weight positive. Make it take an RNG and use a zero-centred range, or deprecate it in favour of Param::he_uniform. The legacy Network path uses it.
2.1.4 Stream the checkpoint writer done S write_f32s in src/transformer.rs, covered by the existing round-trip tests. Original note: save_optimizer_state and save_bin (src/transformer.rs) build a full Vec<u8> copy of each tensor before they write it. Write little-endian chunks into the BufWriter instead. This matters on long runs that save every 200 steps.
2.1.5 Update the README Scope section done for the facts S The convolution claim and the model list are fixed. "Not planned" for encoder-decoder and multi-GPU is still decision 1. Original note: The section says convolutions have no backward pass and encoder-decoder is "not planned". Check both claims against the tree and against the decisions in this file.
2.1.6 Changelog and release S Prepared as 3.0.0, not 2.1.0: TransformerConfig, RopeSpec and ShapeVaeConfig gained public fields. The old "Unreleased" section had already shipped in 2.0.1 and 2.0.2 (including the ImagePipeline::new change), so it now has that heading. Publishing is left to the owner.

Consumer requests

Added 2026-10-09 from the consumers' missingFeatures.md files (/home/vasylt/Code/rusting/<project>/missingFeatures.md) and the owner's detection track for the CS2 player detector. These come before the other tracks. P0 blocks a consumer (no stopgap possible, or a plan item waits on it); P1 has a working stopgap in the consumer. A stopgap with a gradient check is the reference implementation. When an item lands, the consumer entry's - Status: line becomes landed in RustingBrain <commit> with the API under it.

RustingEngine's list asks for game-engine tooling (CLI, scene editing, capture), none of it machine learning, so it adds nothing here.

Order: items wanted by the most projects first, smallest first within that, except K8 after the CPU layers it is tested against: K1, K5, K7, K6, K3, K4, K2, K8, K10, K11, K13c, K13b, K13a, K13d, K12, K9, K13e, K13f.

# Item Priority Effort Asked by Acceptance
K1 Public SiLU derivative: ffn::silu_derivative made pub done 2026-10-09: test ffn::tests::silu_derivative_matches_finite_differences (ten points from -8 to 9, within 1e-3) and a doc test calling it from outside the crate P1 S RustingChop "Public SiLU derivative", RustingCS2 (copied stopgap), detection track (SiLU everywhere in YOLO11) A doc test calls it from outside the crate; finite-difference test at ten points.
K2 conv::TrainableGroupNorm::new(groups, channels, eps) on ImageBatch, gamma/beta: Param, forward -> (ImageBatch, GroupNormCache), backward, params_mut() (CPU; device kernels come with K8) done 2026-10-09: new and forward return a Result. Test conv::tests::trainable_group_norm_matches_the_inference_one_and_finite_differences: forward bit-equal to conv::GroupNorm, input/gamma/beta gradients within 1e-3 (relative to the largest) of central differences P1 S RustingChop "Trainable GroupNorm", RustingCS2 Input, gamma and beta gradients within 1e-3 of finite differences; matches conv::GroupNorm forward exactly.
K3 conv::spatial_softmax / spatial_softmax_backward (coordinates linspace(-1, 1), x then y) and conv::global_avg_pool / global_avg_pool_backward done 2026-10-09 (host; device kernels are K8c): tests conv::tests::spatial_softmax_matches_finite_differences_and_finds_a_peak and conv::tests::global_avg_pool_matches_finite_differences (every input gradient within 1% of the largest) P1 S RustingChop "Spatial softmax", "Global average pool", RustingCS2 Finite-difference checks of both backwards; output layout as the stopgap src/stopgap/spatial.rs.
K4 param::Dropout::new(rate), forward(&Matrix, &mut StdRng, training) -> (Matrix, DropoutMask), backward(&mask, &Matrix), inverted scaling done 2026-10-09: new returns a Result (rates outside [0, 1) refused). Test param::tests::dropout_keeps_the_mean_and_routes_the_gradient_through_its_mask (1e5 ones at 0.15: mean within 1%, dropped share within 0.5%) P1 S RustingChop "Dropout", RustingCS2 Identity when not training; mean kept within 1% over 1e5 entries; backward masks the same entries.
K5 losses::bce_with_logits(logits, targets, smoothing, weights), losses::soft_cross_entropy(logits, targets, weights), losses::huber(pred, target, delta, weights), each the batch-mean loss and dL/dlogits done 2026-10-09: tests losses::logit_tests::* (each gradient within 1% of the largest against central differences, with and without weights; BCE at logits +-100 gives loss 100 and gradients +-0.5, soft cross-entropy 200 and finite) P1 S RustingChop "Logit losses with per-sample weights", RustingCS2, detection track (BCE-with-logits for the class head) Finite-difference checks with and without weights and smoothing; stable at logits of +-100.
K6 TrainableConv2d without a bias (without_bias() or bias: Option<Param>) done 2026-10-09: TrainableConv2d::without_bias() and has_bias(); the pub bias: Param field stays (zero and unused) so existing code still compiles. Test conv::tests::a_convolution_without_bias_adds_and_trains_none (output exactly the biased one minus its bias, same weight gradients, no bias in params_mut) P1 S RustingChop "Conv without bias", RustingCS2 A bias-free conv's gradients match finite differences; params_mut has no bias.
K7 Guard the 1-based optimizer step: step == 0 is an error or a debug_assert!, and the docs of apply_step, step_clipped and Param::step say steps count from 1 done 2026-10-09: Param::step asserts step > 0 in every build (every optimizer path goes through it), docs say steps count from 1; test param::tests::step_zero_panics_instead_of_poisoning_the_parameters P1 S RustingChop "Document (or guard) the 1-based optimizer step", RustingCS2 A test that step 0 panics in debug (or returns an error) instead of turning parameters NaN.
K8 Device-resident image training: DeviceImageBatch, TrainableConv2d::forward_device/backward with bias on the device, device GroupNorm, SiLU, spatial softmax and global average pool, public upload and download; BF16 activations optional. Split: a (M) the batch type and conv forward/backward with bias a done 2026-10-09: conv::DeviceImageBatch, CudaDevice::upload/download, TrainableConv2d::forward_device/backward_device, bias and its gradient on the device; test conv::tests::a_device_chain_stays_on_the_card_and_matches_the_host (two biased convs, one with rebuilt columns: zero transfers inside the chain, outputs and all gradients within 1e-4 abs + 1e-4 rel of the host); b (S) GroupNorm and SiLU kernels b done 2026-10-09: TrainableGroupNorm::to_cuda_on/forward_device/backward_device with DeviceGroupNormCache, conv::silu_device/silu_backward_device; test conv::tests::device_group_norm_and_silu_match_the_host (2x16x17x12, 4 groups: zero transfers inside the chain, output, input, gamma and beta gradients within 1e-4 abs + 1e-4 rel of the host); c (S) spatial softmax and pooling kernels c done 2026-10-09: conv::spatial_softmax_device/spatial_softmax_backward_device with DeviceSpatialSoftmaxCache, conv::global_avg_pool_device/global_avg_pool_backward_device; test conv::tests::device_spatial_softmax_and_pooling_match_the_host (3x17x19x5: zero transfers, outputs and input gradients within 1e-5 abs + 1e-4 rel of the host); d (S) step time of RustingChop's 7-block model at batch 128 against the per-layer path, in docs/baseline.md d done 2026-10-09: examples/chop_block_step.rs, per-layer 1002.3 ms, device-resident 98.9 ms per forward and backward step at batch 128 (10.1x, load average 1.5) P0 L RustingChop "Device-resident image training", RustingCS2 (same root cause), detection track (graph with device-resident activations) Every device op within FP tolerance of its CPU function; no host round trip inside the chain (checked by a transfer counter); measured step time.
K9 recurrent::Gru::new(input, hidden, rng), forward_seq, backward_seq (through time, with dh0), step for batch 1, params_mut, CUDA gate GEMMs done 2026-10-09: as wanted, plus to_cpu and hidden; PyTorch's nn.GRU equations and initialization, rows batch-major. Tests recurrent::tests::every_gradient_matches_finite_differences_over_sixteen_steps (x, h0 and all four parameters, worst gap 1.7e-4), stepping_matches_the_sequence (1e-6), device_gate_gemms_match_the_host_or_skip_without_a_device (256-wide cell, batch 4, 16 steps: outputs and every gradient within 2.9e-6; load average 3.8) P1 M RustingCS2 "Recurrent cell with backward through time" Every gradient within 1e-3 of finite differences over 16 steps; CUDA within FP tolerance of CPU.
K10 Glb::base_color_image(&self, primitive) -> Option<(&[u8], &str)> (bytes and MIME type), plus Glb::json() and Glb::binary() done 2026-10-09: the primitive is named by (mesh, primitive) indices. Test gltf::tests::a_textured_primitive_hands_back_its_embedded_image (embedded PNG bytes and image/png back; none for a flat-colour or material-less primitive) P0 S RustingMedia 2 "Base-colour texture bytes out of a GLB" (blocks its PLAN 4) A GLB with an embedded PNG returns the same bytes and image/png; none for an untextured primitive.
K11 Consistent outward winding from marching_tetrahedra and _sparse done 2026-10-09: each face is turned so its normal points from the tetrahedron's inside corners to its outside ones; test mesh::tests::marching_tetrahedra_winds_every_face_outward (sphere radius 0.6, resolution 48, dense and sparse: every face normal agrees with the radial direction, winding number at an inside point within 0.05 of 1) P1 S RustingMedia 1 "Consistent winding from marching_tetrahedra" On a sphere field (radius 0.6, resolution 48) every face normal agrees with the field gradient at its centroid; winding number about 1 inside.
K12 DINOv2 ViT-S/14 and ViT-B/14 (register-free) safetensors in vit_encoder, patch tokens out done 2026-10-09: VitEncoder::load_dinov2(path, image_size, precision) reads a facebook/dinov2-* checkpoint (no pre-norm, exact GELU, LayerScale folded into the out-projection and second FFN layer, position table resampled bicubically to the image size); IMAGENET_MEAN/IMAGENET_STD; row 0 of encode is the class token, rows 1.. the patch tokens. Test vit_encoder::tests::dinov2_matches_transformers (ignored, RB_DINOV2_REF from scripts/dinov2_reference.py): small 4.4e-4 at scale 27.9, base 1.0e-3 at scale 24.3, worst 4.2e-5 relative; device tower cuda_image::tests::a_whole_vit_tower_matches_the_host_or_skips_without_a_device now also runs the DINOv2 layout (load average 4.0) P0 M RustingMedia 3 "DINOv2 weights in vit_encoder" (blocks its PLAN 5) Patch tokens within 1e-3 relative of the HF transformers reference on two images.
K13a Done 2026-10-09 as K13a-i to K13a-iii. Detection layers with CUDA backward: trainable BatchNorm2d (train and eval, running stats) and conv+BN fusion for inference, SiLU, nearest 2x upsample, channel concat, 5x5 stride-1 max pool (SPPF), the C2PSA attention, composed into one model with backward through the whole graph (shares K8's device graph) P0 M Detection track Each layer's gradients within 1e-3 of finite differences on the CPU and CUDA within FP tolerance of CPU.
K13a-i Trainable BatchNorm2d (train and eval, running stats), conv+BN fusion, nearest upsample, channel concat and split, 5x5 stride-1 max pool, residual add, each on the host and the device done 2026-10-09: conv::TrainableBatchNorm2d (new, yolo = eps 1e-3 momentum 0.03, forward(&mut, input, training), backward, fuse_into(&conv), forward_device/backward_device, running statistics as a [2, C] Param kept on the device), upsample_nearest_batch/_backward, concat_channels/split_channels, max_pool_same/_backward, and *_device versions plus add_device. Tests conv::tests::batch_norm_matches_finite_differences_in_both_modes_and_tracks_running_statistics (input gradient within 1e-3 of central differences training and evaluating; running mean and unbiased variance within 1e-6 of PyTorch's update), a_fused_convolution_equals_the_convolution_then_the_evaluating_norm (1e-5), upsample_concat_split_and_max_pool_backwards_are_their_adjoints (max pool within 1e-3 of central differences), and device_detection_layers_match_the_host (every output, gradient and running statistic within 1e-4 + 1e-4 relative; load average 4.9) P0 S Detection track As K13a.
K13a-ii Depthwise 3x3 convolution with backward (C2PSA's positional encoding, the Detect head's class branch) and the pixel attention inside C2PSA, host and device P0 S Detection track Done. DepthwiseConv2d, pixel_attention/_backward and their _device forms. conv::tests::depthwise_convolution_matches_finite_differences (strides 1 and 2, input and weight gaps < 1e-3); conv::tests::pixel_attention_matches_ultralytics (attended map, values and qkv gradient within 1e-5 of torch autograd, fixture /home/vasylt/Code/games/cs2/fixtures/k13a_attention.json); conv::tests::device_depthwise_convolution_matches_the_host and device_pixel_attention_matches_the_host (within 1e-4 relative; load average 1.7).
K13a-iii C2PSA attention block and the full YOLO11 graph with backward through it done 2026-10-09: yolo::Yolo11 (new(Yolo11Scale, classes, rng), load_ultralytics, forward/backward on the host, to_cuda_on, forward_device/outputs_device/backward_device, state_mut, params_mut), yolo::decode (DFL expectation, sigmoid scores, ready for NMS) and yolo::detection_loss (ultralytics v8DetectionLoss: assigner, BCE, CIoU and DFL, gradient with respect to the three maps). Tests yolo::tests::yolo11_matches_ultralytics (a 1/16-width YOLO11 against ultralytics 8.4.174: eval and train outputs, input gradient and running statistics within 1e-4 relative, parameter gradients within 2e-3, the worst a BatchNorm gamma summing thousands of terms at 1.2e-3), detection_loss_matches_ultralytics (loss terms within 1e-5, map gradients within 1e-4 relative, fixture /home/vasylt/Code/games/cs2/fixtures/export_k13e_loss.py), and device_yolo11_matches_the_host (train-mode outputs and input gradient within 1e-3 relative, measured 2e-4 and 3e-4 from batch-statistic summation order; running statistics 1e-4; parameter gradients 2e-3; load average 3.9) P0 M Detection track As K13a.
K13b Detection ops: xywh/xyxy, IoU, GIoU and CIoU with backward, distribution focal loss, task-aligned assigner, NMS (CPU, then CUDA), mAP50 and mAP50-95 as ultralytics computes them done 2026-10-09 (host): detection::box_iou(pred, target, Overlap) with its gradient, distribution_focal_loss, make_anchors, box_to_ltrb/ltrb_to_box, TaskAlignedAssigner::yolo(strides).assign, non_max_suppression, mean_average_precision. Tests detection::tests::overlap_gradients_match_finite_differences (IoU and GIoU within 2e-4) and detection::tests::detection_ops_match_ultralytics against ultralytics 8.4.174 (/home/vasylt/Code/games/cs2/fixtures/export_k13b_ops.py): 64 box pairs' IoU/GIoU/CIoU values and gradients within 1e-5, DFL loss within 1e-5 and gradient within 1e-6, assignment of 84 anchors to 3 boxes exact (23 foreground, scores within 1e-5), NMS of 300 boxes keeps the same 70 rows bit-equal, per-threshold matches exact on 6 images and mAP50 0.4752 / mAP50-95 0.2885 within 1e-6. CIoU's aspect weight is held constant in the gradient as ultralytics does. CUDA NMS moves to K13f, where its cost is measured P0 M Detection track Loss gradients against finite differences; assigner, NMS and mAP against ultralytics outputs exported as fixtures to /home/vasylt/Code/games/cs2/fixtures/.
K13c Data: YOLO-format dataset reader, letterbox resize, mosaic and horizontal flip, JPEG decode through the images feature done 2026-10-09: detection::read_yolo_split, parse_yolo_labels, rect_shape, DetectionImage::open/resized/letterbox/flip_horizontal, detection::mosaic; test detection::tests::the_pipeline_matches_ultralytics_on_the_player_dataset (8 labelled val images through ultralytics 8.4.174's rect validation pipeline: shape 736x1312 and padding equal, boxes within 1e-5 normalised, 1536 sampled pixels worst 1 level and mean 0.098 apart; one mosaic of four train images at a fixed centre: 5 boxes within 0.01 px), fixtures from /home/vasylt/Code/games/cs2/fixtures/export_k13c_data.py. The random affine after mosaic and HSV jitter are not in it P0 S Detection track Reads /home/vasylt/Code/games/cs2/dataset/ (1107 train, 278 val); letterboxed boxes match ultralytics' on fixtures.
K13d YOLO11s weight import and Gate 1: best.pt and yolo11s.pt exported to safetensors once by a script in /home/vasylt/Code/games/cs2/fixtures/ done 2026-10-09, Gate 1 passed: /home/vasylt/Code/games/cs2/fixtures/export_k13d_yolo11s.py writes k13d_player_v1.safetensors (the trained EMA weights), k13d_yolo11s.safetensors and the first val batch with its raw maps. yolo::tests::player_model_matches_ultralytics_on_sixteen_val_images: worst raw-map gap 1.8e-6 of the map's largest value over 16 images. examples/yolo_val.rs on all 278 val images: mAP50 0.7394 (ultralytics 0.739), mAP50-95 0.5803 (0.576), 40.0 ms/image of which the model is 36.7 ms, unfused and with backward caches kept (load average 3.8) P0 M Detection track On 16 val images raw head outputs within 1e-3 relative of PyTorch; val mAP50 0.739 within 0.005.
K13e Gate 2: examples/yolo_detect.rs fine-tunes from yolo11s.pt on the dataset done 2026-10-09: ultralytics' trainer at optimizer=auto (mosaic/affine/HSV/flip, close_mosaic 10, AdamW 0.002 with decay on conv weights only, 3 warmup epochs, linear decay, accumulation to 64, clip 10, EMA validated). 25 epochs at batch 4 (batch 8 at 1280 runs out of memory in fp32 even with Yolo11::set_recompute_columns): EMA val mAP50 0.7336, mAP50-95 0.5467 in 1.572 h; mAP50 0.5975/0.6070/0.6744/0.7085 at epochs 5/10/15/20 (ultralytics 0.583/0.652 at 10/20). 4.9 samples/s against ultralytics' 27.5 (AMP, 40 s/epoch): 100 epochs would take ~6.5 h, so the full schedule needs the owner. Load average 1-4 during the run. Augmentation test detection::tests::an_identity_affine_and_unit_hsv_keep_the_image P0 L Detection track Val mAP50 >= 0.72; samples/s and wall-clock logged against ultralytics' 1.118 h for 100 epochs. Runs over 2 hours need the owner.
K13f Gate 3: batch-1 inference at 1280 on the 3060 done 2026-10-09, Gate 3 passed: Yolo11::fuse, TrainableConv2d::infer_device, bf16 im2col and GEMMs for the fused 3x3 convolutions under CudaDevice::with_mixed_precision, Yolo11::save_ultralytics, Validation::stages_ms. examples/yolo_val.rs ... 1 tf32 on the 278 val images at batch 1, rect letterboxed (e.g. 1280x736): upload 0.56, forward 13.56, download 0.47 ms = 14.6 ms for the model (unfused fp32 was 37.7 ms); host decode 3.34 ms and NMS 0.02 ms more, 17.9 ms in all; mAP50 0.7390, mAP50-95 0.5800 (fp32 0.7394/0.5803). Load average 2.6 at the start and 15 at the end (other projects). Tests yolo::tests::device_yolo11_matches_the_host (fused host and device within 1e-4 of unfused eval, mixed within 5e-2) and yolo::tests::a_saved_model_loads_back_exactly (round trip exact, fused too; the player checkpoint saved back equals ultralytics' file tensor for tensor). API in CHANGELOG.md. Next levers if it matters: decode on the device (3.3 ms), copy_channels for Concat (1.5 ms), an implicit GEMM instead of im2col P0 M Detection track Under 15 ms with time per stage; then the load, train and run API goes in CHANGELOG.md.
K13g YOLO training throughput, so the full 100-epoch Gate 2 schedule fits the 2-hour pre-approval (about 15 samples/s at batch 4; ultralytics does 27.5 with AMP) P1 M Detection track Added 2026-10-09. yolo_detect --steps 40 --batch 4 at 1280 before: 5.3 samples/s (per batch data 67 ms, forward 221, loss 66, backward and step 401; load 5-6). nsys over 30 batches: dw_bwd_weight 11.7 ms a call (15% of kernel time), im2col_train 11%, col2im_train 9%, memsets 62 ms a batch (column-sized buffers zeroed before a beta-0 GEMM), FP32 SIMT GEMMs about 30%. Step 1 done: the depthwise weight gradient reads channels coalesced (one thread per 64 pixels, tap and channel, atomics into the zeroed sum); matmul_device and col2im allocate without a memset; 1x1 stride-1 convolutions skip im2col and col2im. 6.5 samples/s (forward 197, backward and step 275 ms), losses unchanged to 4 digits. Tests: conv::tests::device_depthwise_convolution_matches_the_host, a_device_chain_stays_on_the_card_and_matches_the_host (its second convolution is now 1x1), the yolo:: and detection:: device tests. Step 2 done: yolo_detect trains on a mixed-precision device (every FP32 GEMM runs as BF16 on the tensor cores with an FP32 accumulator, which is what ultralytics' default AMP does); detection_loss decodes, assigns and differentiates each image on its own rayon task; the input batch is built on the prefetch thread; the batch norm and depthwise backward buffers that a kernel overwrites whole are allocated without a memset. 9.0 samples/s at load 4-5 (forward 156 ms, loss 25, backward and step 225), losses within 0.5% of FP32 (box 0.967 vs 0.966, cls 10.61 vs 10.62, dfl 0.929 vs 0.931). Tests: yolo::tests::detection_loss_matches_ultralytics, the yolo::, detection:: and conv::tests::device device tests. nsys after step 2: kernels 379 ms a batch; im2col_train 43 ms, col2im_train 32, the batch norm and SiLU elementwise kernels about 100, bias and norm column sums as GEMV about 33. Step 3 done: SiLU runs inside the device batch norm kernels (TrainableBatchNorm2d::forward_device_with, a new bn_bwd_prep kernel), 9.3 samples/s. Step 4 done: on a mixed-precision context a training convolution that needs im2col builds bf16 columns, runs all three GEMMs on bf16 operands, writes its column gradients as bf16 and sums them with a new col2im_bf16. 10.1 samples/s at load 6-7 (forward 136 ms, loss 22, backward and step 203), losses within 0.3% of FP32. Test: conv::tests::a_mixed_precision_convolution_tracks_the_host (columns kept and rebuilt, channels in fours and not; output, input, weight and bias gradients within 1% in norm). Step 5 done: col2im_bf16 sums four channels a thread when the channel count allows, 10.6 samples/s. Step 6 done: the device batch norm takes its means, variance and gamma and beta gradients from a col_sums kernel that reads channels coalesced, instead of GEMVs against ones over squares and products buffers; the backward pass takes both of its sums from one read. 11.3 samples/s at load 4-6 (data 30 ms, forward 124, loss 27, backward and step 171), losses within 0.2% of step 4. Tests: the yolo::, detection:: and conv::tests device tests; a sign mutation in the fused sum fails yolo::tests::device_yolo11_matches_the_host. nsys after step 6: kernels 283 ms a batch; im2col_bf16 36.5, col_sums 34, bn_apply 21, copy_channels 19, bn_bwd_prep 18, bn_bwd_apply 18, col2im_bf16 13, add_f32 11; device-to-device copies 19.5 ms. The 40-batch timing runs inside the warmup, where the optimizer steps almost every batch; with steps every 16 batches (the rest of a run) step 6 does 12.1 samples/s. Step 7 done: the batch norm backward no longer writes the gradient before its SiLU: the column sums and bn_bwd_apply take it through the SiLU as they read it (bn_gy), and bn_bwd_prep is gone. 11.5 samples/s at load 8-10 (backward and step 155 ms). Tests: as step 6; dropping the SiLU in bn_bwd_apply fails yolo::tests::device_yolo11_matches_the_host. Step 8 done: a convolution inside Yolo11::forward_device no longer copies its input into its cache (TrainableConv2d::forward_device_keeping(input, false), backward_device_from); the tape already holds the input, so the backward pass reads it there. 12.8 samples/s at load 6-7 (forward 104 ms, loss 29, backward and step 151). Tests: the yolo::, detection:: and conv::tests device tests. Step 9 done: the device batch norm sums its mean and variance in one pass, shifted by each channel's first value so the variance does not cancel away (the running mean as the shift was 4.2e-3 off on a fresh norm over values near 30, no shift 2.4e-3, the first value 6e-5). col_sums 27.2 to 20.5 ms a batch in nsys; 12.6-13.0 samples/s at load 4-8, inside the run-to-run spread. Test: conv::tests::a_device_batch_norm_over_offset_values_matches_the_host (within 5e-4 of the host). An attempt at im2col_bf16 with one thread per pixel and channel over all taps measured 34.7 ms against 35.5 and was dropped. Step 10 done: nsys showed the device idle 12% of a batch, a 20 ms gap of it the host loss; timed inside, the loss spent 9.5 ms gathering every anchor's row into a fresh 35 MB buffer and 2.5 ms scattering the gradients back one row at a time. detection_loss now reads each anchor's row from its map in place and scatters an image's gradients in one parallel copy a map: loss 30 to 19 ms a batch at load 6. Test: yolo::tests::detection_loss_matches_ultralytics. Step 11 done: yolo_detect --warmup N times the steady state, where the optimizer steps once every 16 batches of four: 14.3 samples/s at load 3.5 before this step. A CPU timing of detection_loss on random maps (four images, 20 boxes each) put decode at 9.7 ms, the assigner at 9.6 and the gradient copy-out at 7. Each image now decodes its anchors in parallel without a heap allocation per anchor, the assigner picks each box's top candidates with an exact select (metric down, anchor up) instead of a full sort and runs the boxes in parallel, and the gradients are written into the three map buffers directly: 28 to 9 ms on the CPU, and 20 to 7 ms a batch on the bench. Steady state 15.1 and 15.4 samples/s at load 6.5, which puts 100 epochs of 1107 images at about 2.0 h plus validation. Tests: yolo::tests::detection_loss_matches_ultralytics, the detection:: assigner tests. Step 12 done: the steady-state profile put every elementwise kernel near the card's bandwidth (bn_apply ~300 GB/s), so the cut is in bytes: the device batch norm keeps its input instead of writing the normalized values, and its backward kernels recompute them from the input, mean and inverse deviation (bit for bit the forward's formula). Forward 96 to 90-93 ms a batch, 15.15 to 15.3 samples/s averaged over two A/B pairs at load 6. Tests: the 98 device tests, among them conv::tests::a_device_batch_norm_over_offset_values_matches_the_host and yolo::tests::device_yolo11_matches_the_host. Step 13 done: YOLO's device backward pass took a slice's gradient back by concatenating it between two zeroed blocks and then adding whole tensors, and copied an Add's gradient for its second input. copy_channels gained an add flag, and gradients now accumulate in place: a slice's into its channel range of the input's gradient, an Add's into whichever input already has one (a copy only when neither does), and every other sum in place instead of into a new buffer. 15.3 to 15.75 samples/s averaged over two A/B pairs at load 5-6. Tests: the 98 device tests, yolo::tests::device_yolo11_matches_the_host among them. Tried after it: uploading the input as bytes and dividing on the device (20 MB a batch instead of 78) measured 15.75 against 15.7 samples/s, inside the noise, and was dropped. A real epoch with validation at batch 4 (yolo_detect --epochs 1 --val-every 1 --close-mosaic 0, load 5-7) took 71 s of training at 15.7 samples/s plus about 5 s of validation, so 100 epochs project to 1.99 h: inside the limit with no margin, and batch 8 runs out of memory. Also tried: not copying out a slice that only feeds concatenations (the concatenation and its backward read the source's channels in place) cut copy_channels from 24.2 to 21.4 ms a batch and kernel time from 233.6 to 232.2 ms, but measured 15.3-15.4 against 15.6-15.8 samples/s in two A/B pairs, and was dropped. Open, needs the owner: the rest of the kernel time is bandwidth-bound elementwise work (batch norm, channel copies) and im2col_bf16 plus col2im_bf16 (47.6 ms a batch, 18%); 12 ms more is fp32 pointwise GEMMs for which cuBLAS picks SIMT or align-1 kernels. What remains is an implicit-GEMM convolution, which cuDNN 9 (installed in /usr/lib) provides through cudarc's cudnn feature, a new system dependency that needs approval; writing one by hand is a large job of uncertain return. Otherwise the 100-epoch run (projected 1.99 h) needs the owner's go-ahead past the 2-hour limit. Rejected 2026-10-10, the fp32 pointwise GEMMs: a fresh nsys over 60 steady-state batches (kernels 233.4 ms a batch, GEMMs 77.4) put 7.3 ms a batch in simt_sgemm ... align1 kernels. A shape log showed them on the backbone's 1x1 convolutions (e.g. m64 n409600 k64), whose pointers and leading dimensions are all 16-byte aligned: cuBLAS simply picks SIMT for them under COMPUTE_32F_FAST_16BF. COMPUTE_32F_FAST_TF32 for the fp32-operand GEMMs measured 15.15 against 15.75 samples/s in two A/B pairs at load 1-4, and was dropped. A bf16 cast of the input before each pointwise GEMM does not pay either: with k of 64 to 96 these GEMMs are bandwidth-bound (773 us against about 600 us for the bytes alone at m64 n409600 k64), and the cast pass reads and writes more than the kernel change saves. The stem convolution (k 27, wmma align1, 0.97 ms forward) writes 210 MB of output, so padding its k to 32 saves under 0.4 ms.

Track A — Training performance (CUDA)

Every item here is gated on ground rule 1.

A1. flash_attention_dkv fragment packing — P1, M — DONE 2026-10-07

Done. The fragment-packing guess was wrong: loading the second loop's (and the first loop's) B operands with ldmatrix.x4 instead of four LDS.32 cut registers from 198 to 168 but left the probe at 0.973 against 0.990 ms. What paid was dropping the dimension-major copies of the query and output-gradient tiles: the second loop now reads the query-major tiles with ldmatrix.x4.trans, the load phase loses its four 16-bit transposed stores per pair, and shared memory falls from 37 KB to 19 KB, so with 168 registers three blocks fit on an SM instead of two. Probe 0.990-0.992 to 0.800 ms (-18%); in the model (nsys, bench_train_step 4 1024 cuda) 847 to 637 us a launch (-25%), whole step 52149-52208 to 52536-53452 tok/s under load 9-24. Tests: both fused_attention_*_three_kernel_path_* (64 and 128, causal and bidirectional); reading every transposed matrix from the first column block fails the backward one. The 5 ms target was past the card: at 0.995 ms the kernel already ran at about 54% of the 25.5 TFLOPS BF16 peak. Evidence in docs/baseline.md ("dkv without transposed tiles").

This is the largest single kernel left, at 15.45 ms a step. The ablations in docs/baseline.md ("What is left") show that the kernel is bound by mma work and not by bandwidth. The second mma loop costs 0.44 ms against 0.17 ms for the first, with identical mma and LDS.32 counts. Start with the fragment packing in that second loop. Use examples/flash_probe.rs against its 0.995 ms baseline. Target: the same mma throughput as the first loop, about 5 ms a step.

A2. flash_attention_dq and flash_attention_fwd — P2, M — REJECTED 2026-10-08

These cost 10.45 ms and 7.24 ms. Apply whatever A1 learns about the second loop. Do not start before A1 lands.

Rejected after measuring, 2026-10-08. A1's lesson was to drop the transposed shared-memory copies and read with ldmatrix.trans. The query-gradient kernel already does (A3). In the forward kernel, values key-major with ldmatrix.x4.trans moved the probe from 0.446 to 0.441 ms, about 0.04 ms of a 77 ms step. Paired tile loads made it slower (0.49-0.57 ms), and so did ldmatrix for the score keys (0.51 ms). Nothing was kept. Evidence in docs/baseline.md ("Forward and query-gradient kernels after A1").

A3. Flash attention for other head widths — P1, M — DONE 2026-10-08

cuda_flash::eligible requires head_dim == 64 and mixed precision. Every other width falls back to the three-kernel path that used to be 21% of a step. Add head_dim 128, which most published LLaMA-family checkpoints use and which B5 needs. Add 32 only if a model needs it.

Done for 128. The four training kernels are templates over the head width and are built for 64 and 128 (flash_attention_*_64, flash_attention_*_128); the image kernel stays at 64. The 128-wide backward kernels opt in to more than 48 KiB of shared memory (dq 52 KiB, dkv 70.5 KiB). Parity: both fused_attention_*_three_kernel_path_* tests run at 64 and 128, causal and bidirectional, against the FP32 host. Idle GPU, batch 4 x 1024, the default model cut into 4 heads of 128: 48,456 tok/s fused against 39,374 tok/s on the three-kernel path (+23%). The 64-wide kernels compile to the same register counts as before (126/163/198) and the default model runs 53,315 tok/s. Left: flash_attention_dq_128 uses 254 registers and dkv_128 spills 128 bytes; at 128 they take 7.2 and 8.7 ms a step against 3.6 and 6.7 ms at 64. dq at 128 done 2026-10-08. Its limit was shared memory, not registers: a second, transposed copy of the key tile made 52 KiB a block, so one block fit per SM. It now reads the key tile transposed with ldmatrix.trans, as dkv already does: 34 KiB, two blocks per SM (254 registers still fit two). nsys, bench_train_step 4 1024 cuda 128, load 5: dq_128 median 1075 to 620 us a call; four interleaved rounds 48517-48770 to 49954-50304 tokens/s (+2.6%), head 64 flat (53257-53636 against 52856-53643). Test: fused_attention_backward_matches_the_three_kernel_path_or_skips_without_device; feeding the wrong ldmatrix half to the first output tile fails it. dkv at 128, 2026-10-08: its key, value and two gradient fragments are 192 registers, and the whole query tile's scores 64 more, so it spilled 180 bytes. At width 128 it now takes the query tile in two halves of 32: 48 bytes spilled, dkv_128 median 891 to 783 us, step 49921-50451 to 50471-50589 tokens/s. At 64 one pass stays (the split cost 4% there). Test: the same backward test; dropping the half from the gradient products' row offset fails it. A3 done.

A4. Non-causal flash in BF16 for training — P1, M — DONE (2026-09-23, found 2026-10-07)

Done in 52bbf36 ("train a bidirectional model on the device"): every fused kernel takes a causal flag, and a bidirectional() TransformerLm (masked LM, ViT-shaped encoder) trains on the fused BF16 path. Tests: gpu_transformer::tests::fused_attention_matches_the_three_kernel_path_or_skips_without_device and ..._backward_..., both run causal and bidirectional at head widths 64 and 128. What remains is the flow transformer and the shape decoder, which run FP32 throughout and so cannot reach a BF16 kernel; that is A5.

The FP16 non-causal variant exists for inference in the image path. Training a bidirectional model (masked LM, ViT, the flow transformer) has no fused path. gpu_cross.rs and gpu_flow.rs both say "FP32 throughout". This item unblocks A5.

A5. BF16 for gpu_cross and gpu_flow — P2, M — done 2026-10-08

The 3D training path runs in FP32, so it gets half the tensor-core rate or less. A4 is done; profile examples/train_shape.rs first. Its self-attention can then call the bidirectional flash kernels directly.

Profile, 2026-10-08 (examples/profile_shape.rs 4 2048 4096, mixed, packed, nsys): 0.129 s a step against 0.206 in FP32, about 80 ms of it kernels. Every step had a 44-47 ms gap with the GPU idle: one cudaStreamSynchronize after the gradient norm's first readback, preceded by a burst of cuMemFreeAsync. The default memory pool's release threshold is zero, so each sync handed the step's freed buffers back to the driver and the next step mapped them in again. device_context now sets the threshold to u64::MAX (keep_freed_memory): the shape step went 0.130-0.131 to 0.085 s (-35%), bench_train_step 4 1024 cuda stayed 53451-53562 tokens/s. Load 0.8-1.0. No test: the call fails the context if the driver refuses it.

Done 2026-10-08. A second nsys profile (same command, load 2.6) put the three-kernel softmax at about 17 ms a step and the per-head FP32-operand attention GEMMs at about 20 ms more: about 40% of the step. The flash kernels now take separate query and key/value buffers, widths and lengths (forward and query gradient over query tiles, key/value gradient over key tiles), and self-attention passes one buffer as both. gpu_cross runs every activation and weight in BF16 and attention through those kernels when the context is mixed precision and the head is 64 or 128 wide; otherwise it keeps the FP32 three-kernel path. Shape VAE step 0.084-0.087 to 0.062-0.065 s (-25%); LM step unchanged once the width-64 key/value gradient kernel was held to three blocks per SM (see docs/baseline.md). Test: gpu_cross::tests::fused_cross_attention_matches_the_host_or_skips_without_device, two shapes with more key tiles than query tiles and the reverse, worst error 0.47% of the largest value; a forward kernel that stops one key tile short fails it (0.93%), as does a key/value kernel that walks the wrong number of query tiles.

A6. Fused LM head and cross-entropy — P1, M — rejected for speed 2026-10-07

Measured on the idle GPU at 96ffcc5 (bench_train_step 4 1024 cuda, nsys, four steps): the LM head is three cuBLAS GEMMs of 4096 x 512 x 32000, 134 GFLOP each, taking 5.83, 5.21 and 5.06 ms a step: 23.0, 25.8 and 26.5 TFLOP/s, at the 3060's BF16 tensor-core rate with an FP32 accumulator. ce_loss_grad_bf16 is 2.73 ms: it reads the 262 MB BF16 logits twice and writes them once, about 288 GB/s of 360. A fused kernel can only remove that 2.73 ms (3.5% of 77.5 ms), and only if its own GEMM keeps up with cuBLAS, which no hand GEMM in this crate does (the routed experts reach 13.6 to 15.8 TFLOP/s). The logits buffer is 262 MB in BF16 at this shape, not the 512 MB below, which is 2% of the card. Re-open for memory, not speed, if a run is ever bound by that buffer.

Original note:

The LM head GEMMs are 22.72 ms a step. src/transformer.rs:1367 materializes the full [rows, vocab] logits on the device. A chunked fused GEMM-plus-softmax-plus-loss that never writes the logits saves both time and the largest activation buffer. At vocab 32k and 4 x 1024 tokens, that buffer is 512 MB in FP32. This also helps if W8 (halve the vocabulary) is rejected.

A7. MoE work items W1-W4 — P1, S to M each

These are already scoped in docs/improvement-plan.md:

  • W1: run the routed path in BF16 done 2026-10-07. Under mixed precision the gathered rows, packed expert weights, SwiGLU hidden and the backward upstreams are BF16; the expert outputs, the scatters and the router stay FP32. gather_scale_rows writes either precision. Measured on the default model at 4 x 1024: kernel time 82.0 to 80.1 ms a step, the FP32-operand expert GEMMs 18.4 to 6.1 ms, best of five 49,802 to 51,631 tok/s (+3.7%, under the estimate, which was for a 768-wide model). Loss over 500 steps on the repository's markdown tracks the FP32-routed run (3.184 against 3.183 at step 499, held-out 4.924 against 4.944). Evidence in docs/baseline.md ("MoE routed path in BF16").
  • W2: fuse the permutation passes parts 1 and 2 rejected 2026-10-07 after re-measuring. On the idle GPU at 6f48659, the five permutation kernels take 4.2 ms of a 78.1 ms step (gather 1.29, scatter_add_scaled 0.93, scatter_add_rows 0.92, row_dot_scatter 0.63, moe_stats 0.36), down from 17.8 ms when the item was written. Both planned fusions are stale: scatter_add_scaled runs in the forward pass and row_dot_scatter in the backward, so they cannot share a read, and the backward gather reads the output gradient, not the input the forward cache already keeps. The one redundant read left is row_dot_scatter and the backward gather both reading the gathered output-gradient rows: about 0.5 ms (0.6%), under the run-to-run spread. Part 3 (no forward gather) stays behind A8. Evidence in docs/baseline.md ("MoE permutation passes re-measured").
  • W3: fold pack_gate_up into the parameter layout deferred to the next major version, 2026-10-07. Re-measured in the idle-GPU trace at 6f48659: the per-step packs (pack_gate_up and pack_down, the cast_act launches inside each MoE forward bracket) are 144 launches and 0.69 ms of a 78.1 ms step (0.9%), not 2.5 ms, since W1 halved the bytes. The transient buffer is 5.8 MB of BF16 per routed layer at the default 8 experts x 352 x 512. The fix replaces SwiGlu's public gate and up fields (moe::Expert is SwiGlu), which 3.x cannot do. Take it in 4.0, or earlier if a many-expert model makes the transient buffer matter.
  • W4: measure the expert GEMM split, and only then optimize it. Step 1 done 2026-10-07: read off the nsys kernel trace of the idle-GPU run at 6f48659, taking the GEMMs between each gather and its scatter as the routed experts (no NVTX needed; the order is fixed). Routed expert GEMMs are 10.67 ms of a 78.1 ms step: 3.92 ms forward in 96 launches (about 13.6 TFLOPS) and 6.75 ms backward in 192 (about 15.8 TFLOPS), against roughly 22 TFLOPS for the dense GEMMs. Each expert GEMM is about 1024 rows by 352 or 704 columns, and cuBLAS launches 24 to 48 blocks for it on a 28-SM card, so the split costs occupancy, not tile choice. At dense efficiency the whole routed set would take about 7.3 ms: about 3.4 ms (4.4%) recoverable, under the plan's 10 ms bar for step 2. A grouped GEMM over all eight experts is the fix if it is ever taken, and it needs A8. Evidence in docs/baseline.md ("Routed expert GEMMs attributed"). Step 3 done 2026-10-07: TransformerCache::expert_loads() returns each MoE layer's per-expert share of routed slots, from the host or the device cache. Tests: transformer::tests::expert_loads_report_each_moe_layer_and_sum_to_one, gpu_transformer::tests::expert_loads_on_the_device_match_the_host_or_skip_without_device (fails with the device shares mis-normalised). On the repository's markdown (4096-entry BPE, 500 steps) the busiest expert takes up to 0.41 at initialisation, but 0.145-0.169 per layer from step 200 on against an even 0.125, so step 1's estimate holds for trained routing. Evidence in docs/baseline.md ("Expert loads on real text").

Also move the router GEMM onto the tensor cores. Done 2026-10-07. Under mixed precision the MoE block now takes the BF16 normed input like the dense feed-forward: the router multiplies BF16 operands into FP32 logits, the gather reads the BF16 rows directly and the shared expert runs in BF16. Kernel time 80.1 to 76.7 ms a step; the FP32-operand s1688bf16 GEMMs are gone. Best of five 51,631 to 52,914 tok/s (+2.5%). Loss over 500 steps ends at 2.945 (held-out 4.632) against 3.184 (4.924) for the W1 run. Evidence in docs/baseline.md ("MoE router and shared expert in BF16").

A8. MoE capacity factor — P2, M

Blocked until 4.0 (checked 2026-10-08): MoeConfig and MoeLayer are public structs with all fields public, so the capacity factor cannot become a field in 3.x. A setting kept outside the config would not survive a save. Its payoff is also small now: the grouped expert GEMM it enables was measured at about 3.4 ms of 78 ms (W4 step 1).

This is the TODO at src/moe.rs:200. A fixed-capacity variant gives statically shaped expert buffers. It is the precondition for the third part of W2, for a batched-pointer expert GEMM, and for C1 (CUDA graphs). Keep dropless as the default.

A9. Fused multi-tensor optimizer step — P2, S — rejected 2026-10-07

Measured on the idle GPU at 6f48659 (bench_train_step 4 1024 cuda): adam is 207 launches and 5.11 ms a step over 55.18M parameters. At 28 to 32 bytes of traffic per parameter that is 302 to 345 GB/s, 84% to 96% of the 3060's 360 GB/s. The median launch takes 17 µs, so launch cost is not what bounds it, and the GPU queue never drains during a step. One kernel over a pointer table could save at most the gap to the roofline, about 0.2 to 0.8 ms (0.3% to 1%), under the run-to-run spread. Lion moves fewer bytes and is bound the same way. Evidence in docs/baseline.md ("Adam re-measured").

Original note:

Adam is 113 launches and 8.96 ms a step. One kernel over a pointer table of all parameters removes 112 launches. Measure first: Adam is bandwidth-bound, so the gain may only be the launch overhead. Do the same for Lion.

A10. Activation checkpointing — P1, M — DONE 2026-10-07

Recompute each block's forward pass during the backward pass instead of storing it. This trades about 30% more compute for most of the activation memory. It is the cheapest way to train longer sequences or a larger model on a 12 GB card. Make it per-layer and opt-in on the builder.

Done 2026-10-07. TransformerBuilder::checkpoint(blocks) (and TransformerLm::set_checkpoint) marks blocks whose device cache keeps only the block input; backward_stack reruns forward_block from it before that block's backward. Expert loads are taken in the forward pass so a checkpointed routed layer still reports them. Device only: the host path and the vision stack keep everything. Test gpu_transformer::tests::checkpointed_blocks_train_like_kept_ones_or_skip_without_device (FP32 and BF16, a dense and a routed block; fails when the recompute reads zeros). Default model, every block checkpointed: peak 2843 to 1353 MiB at 16 x 512, step 131.8 to 161.5 ms (+23%). Evidence in docs/baseline.md ("Activation checkpointing").

A11. 8-bit optimizer state — P2, M — DONE 2026-10-08

Block-wise quantized Adam moments cut optimizer memory by 75%. Gate it on the loss-curve check from docs/optimize-plan.md Task 4 Step 7. Note W6 in docs/improvement-plan.md: optimizer-state traffic "can hurt".

Done 2026-10-08. TransformerLm::set_8bit_optimizer_state(true), after to_cuda, holds every trainable weight's Adam moments as one byte each: blocks of 256 values with an FP32 absmax and bitsandbytes' dynamic map (signed for the first moment, unsigned for the second). The adam8 kernel dequantizes, steps, reduces the block's new absmax and requantizes in one pass. Opt-in and device only, so Optimizer stays as it is; a Lion step on 8-bit state fails. Saving and loading optimizer state dequantize and requantize. Tests: gpu_transformer::tests::an_8bit_adam_step_matches_the_host_or_skips_without_device (five steps against the same scheme on the host, gradients over six decades, a partial block, and the round trip back to FP32) and a_model_with_8bit_optimizer_state_trains_like_fp32_or_skips_without_device; a 1% wrong rescale in the kernel fails both. Default model: peak 2037 to 1723 MiB at 4 x 1024 and 16 x 256 (314 MiB, the 6 bytes a weight predicts for 55.2M weights), step time unchanged (75.2-75.6 ms). Loss over 2000 steps (examples/adam8_loss_check.rs, repository markdown, BPE 4096, 16 x 256, Adam 3e-4) tracks FP32 state: 0.1824 against 0.1723 at the end, held-out 6.618 against 6.554. The corpus is small enough that both runs overfit it; a larger-corpus check is still worth doing before a long run. Evidence in docs/baseline.md ("8-bit Adam state"). Seen once while verifying, in one of six full device runs: checkpointed_blocks_train_like_kept_ones_or_skip_without_device failed its 1e-5 loss match (it does not use 8-bit state); it passed seven reruns. The MoE scatters add with float atomics, whose order varies; the tolerance may be too tight for that.

A12. Overlap host round trips — P3

W5 measured the ceiling at 5%. Leave it parked until everything above lands.

Track B — Training features

B1. Data-parallel training on two GPUs — P1, L

This is the item the README apologizes for. Plan:

  1. One process with one context per device, and one model replica each.
  2. All-reduce the gradients with NCCL through cudarc. Bucket the gradients and overlap the all-reduce with the backward pass.
  3. Each rank draws from TokenFile with its rank folded into the window offsets, so the "resume at step N, same data" property still holds.
  4. Acceptance: on two cards, the loss curve matches a single card at the same effective batch, and throughput is 1.8x or better.

Do not build tensor parallelism, pipeline parallelism or multi-node until two cards on one host works. See featuresPlan.md "Multi-GPU".

B2. ZeRO-1 optimizer-state sharding — P2, M

Shard the optimizer state across the data-parallel ranks. Do this after B1.

B3. More optimizers — P2, S each

  • AdamW with decoupled weight decay already decoupled: Param subtracts learning_rate * weight_decay * value apart from the Adam update (src/param.rs), the CUDA kernel matches it (cuda_adam_weight_decay_one_batch_update_matches_cpu_or_skips_without_device), and Optimizer::Adam documents it as AdamW-style. Checked 2026-10-07.
  • Muon for the 2D weight matrices, with AdamW for the rest. Measure it against Adam on the 101.7M benchmark model.
  • Schedule-free AdamW. It removes the schedule argument from long runs.
  • Blocked until 4.0 (checked 2026-10-08): Optimizer is an exhaustive public enum, so a Muon or schedule-free variant is a breaking change. The semver CI job (J3) would flag it.
  • Adafactor. It has factored second moments, and is cheap in memory.

Ship each one only if it beats Adam or Lion on loss per token in a real run.

B4. LoRA completion and QLoRA — P1, M — DONE 2026-10-08

Closed 2026-10-08: every bullet below is done, measured out or already existed.

LoRA works on CPU and CUDA. What is left:

  • src/param.rs:810: a device-resident adapter accumulates its gradient on the host. Keep it on the device.
  • src/gpu_model.rs:1390: the adapter is folded into the weight once per forward pass. Use two thin GEMMs instead, which also lets the base weight stay frozen in int8. Measured and rejected 2026-10-07. The frozen int8 and NF4 bases already landed with the fold (QLoRA slices below). Skipping fold_lora entirely, default model, rank 16, 4 x 512, takes the step from 42.6-43.4 to 40.5-41.0 ms, so the fold is at most 2.2 ms (5%). Two thin GEMMs per reader would write a [rows, units] term in the forward pass and a [rows, inner] term in the input gradient, where the fold writes one [units, inner] weight; with 2048 rows against 512 to 1536 wide weights that is more traffic, not less. What the fold pays is mostly launches (two narrowing casts and a rank-16 GEMM a projection); a cast kernel that adds scale * up . down while it writes the slot would save at most that 5%, for LoRA runs only. Evidence in docs/baseline.md ("LoRA fold cost").
  • src/param.rs:810 nothing to do (checked 2026-10-07): the comment itself says no device path reaches that host method; device training runs in gpu_model.
  • QLoRA: an int8 or NF4 frozen base with BF16 adapters. Split:
    • Slice 1, int8 base done 2026-10-07. TransformerLm::quantize_lora_base stores every adapted projection's frozen base as int8 (the Precision::Q8 scheme) on the host and the device; pack_fresh dequantizes it into the weight pack and folds the adapter in as before. The model trains, generates and save_loras, and refuses merge_lora, save_bin and ONNX export. Test: gpu_transformer::tests::a_gpu_lora_step_on_an_int8_base_matches_the_host_or_skips_without_device (dense and MoE, FP32 and BF16, against the host on the same bytes; dequantizing half the base fails it). Default model, rank 16, 4 x 512: peak 889 to 778 MiB, step unchanged (46.1-46.5 to 46.3-46.6 ms). Evidence in docs/baseline.md ("LoRA on an int8 base").
    • Slice 2, NF4 base done 2026-10-07 (D3(c)). TransformerLm::quantize_lora_base_4bit(Codebook) stores the adapted bases at 4 bits with a scale per block of 64; the device uploads the packed bytes and a dequantize_q4 kernel widens them into the weight pack through the codebook. has_int8_base became has_quantized_base. Test: gpu_transformer::tests::a_gpu_lora_step_on_an_nf4_base_matches_the_host_or_skips_without_device (FP32 and BF16 against the host on the same nibbles; swapping the kernel's nibble halves fails it). Default model, rank 16, 4 x 512: peak 889 (FP32) / 778 (int8) / 762 MiB (NF4), step unchanged at 42.5-43.3 ms. Evidence in docs/baseline.md ("LoRA on an NF4 base").
  • Save and load an adapter on its own, without the base model. Already existed: TransformerLm::save_lora and load_lora.

B5. Load published LLaMA-family weights into TransformerLm — P1, M — DONE 2026-10-07

Closed 2026-10-07. All three prerequisites and the loader are done (below), and B4 trains adapters on the loaded model, including on an int8 or NF4 base. One gap is recorded, not open work: no test compares logits against the reference implementation on real published weights. The only LLaMA-family checkpoint on this machine is a 3B model (maya1, LLaMA 3 layout, llama3 RoPE scaling) whose FP32 copy does not fit beside the 12 GB of free RAM, and there is no Python transformers to produce the reference. The loader's tests check it against the crate's own TextEncoder reader instead (hf::tests::the_reference_llama_reader_agrees_with_what_was_saved).

TextEncoder already reads Qwen2 and LLaMA safetensors for inference. Map the same tensors into a trainable TransformerLm, so a user can fine-tune a published model with B4. This needs A3 (head_dim 128 is the common case), RoPE scaling (B6) and a tokenizer round trip through Bpe.

Status: weights done. TransformerLm::load_hf (src/hf.rs) reads LLaMA and Mistral checkpoints, single-file or sharded. It refuses biases (so Qwen2), per-head q/k norms, YaRN or dynamic rope_scaling and a sliding window shorter than max_seq_len, instead of loading them wrong. Linear and llama3 scaling load. head_dim 128 runs on the fused attention path (A3, 2026-10-07).

Tokenizer: sentencepiece BPE done. Bpe reads the LLaMA 2 / Mistral / TinyLlama flavour (▁ for spaces, <0xNN> byte fallback, the Metaspace pre-tokenizer or the older Prepend normalizer). Against the Hugging Face tokenizers library on a real 32k vocabulary, 310 lines of the README encode to the same ids and decode to the same text in both file formats. Test: tokenizer::tests::sentencepiece_bpe_matches_the_reference_tokenizer. LLaMA 3 tokenizer done. Bpe recognises the LLaMA 3 split regex (digits in runs of up to three, case-insensitive contractions) and its Qwen2 variant (single digits), and honours ignore_merges. Against the Hugging Face tokenizers library on the real LLaMA 3 (128k) and Qwen-Image (151k) files, 474 lines (this README, src/tokenizer.rs, mixed-script text) encode to the same ids and decode to the same text. This also fixed Qwen files, whose empty end_of_word_suffix was read as a CLIP suffix. Test: tokenizer::tests::the_llama3_split_matches_the_reference_tokenizer. head_dim 128 on the fast path: done with A3.

B6. Long context — P2, M

  • RoPE scaling: linear, NTK-aware and YaRN. B5 needs these to read published configurations correctly. Linear and llama3 done (RopeScaling, TransformerConfig::rope_scaling, read and written by src/hf.rs). YaRN and dynamic NTK are still open: YaRN also scales attention logits, and dynamic NTK rebuilds the tables per sequence length. Blocked until 4.0: RopeScaling and TransformerConfig are exhaustive, so a YaRN variant or a sliding-window field is a breaking change (checked 2026-10-07).
  • Sliding-window attention in the flash kernel (a mask on the tile loop).
  • Sequence packing with document masks in TokenBatch, so packed documents do not attend to each other.
  • Both blocked until 4.0 (checked 2026-10-08): a window is model configuration, so it belongs in the exhaustive TransformerConfig, and attention reads its row map from Layout, a public struct with all fields public, so a document map is a new field there. TokenBatch's fields are private, but it reaches attention only through Layout.

B7. Training stability tools — P2, S each

  • Track the gradient norm and the per-layer update-to-weight ratio done: grad_norm and step_clipped already return the norm. TransformerLm::step_with_update_ratios steps and returns |dw| / |w| per parameter, in params_mut order, on host and device. Tests: transformer::tests::update_ratios_measure_the_step_taken (exact under SGD), gpu_transformer::tests::update_ratios_on_the_device_match_the_host_or_skip_without_device (fails without the post-step sync). TransformerLm has no fit_with; the caller logs the vector. 2026-10-07.
  • Skip a step when the loss is NaN done: optimizers::step_clipped (and so TransformerLm::step_clipped, FlowTransformer::step_clipped and the interoceptive agent) skips the update when the gradient norm is NaN or infinite, clears the gradients, and returns the norm so the caller sees it. Before, a NaN norm compared false against max_norm and the NaN step went through, poisoning the weights and Adam moments. The model step counter does not advance on a skip. Test: transformer::tests::a_non_finite_gradient_skips_the_step (fails without the guard). Skipping a loss spike and rolling back to the last good state done: SpikeGuard::new(ratio, beta, warmup) keeps a bias-corrected moving mean of accepted losses and flags a non-finite loss or one above ratio times the mean; a spike leaves the mean alone and consecutive() counts runs. The caller skips with accumulate_step plus zero_grad, and rolls back with the existing load_bin and load_optimizer_state (pattern in the doc example). Tests: optimizers::tests::a_spike_is_skipped_without_moving_the_mean, doctest. 2026-10-07.
  • z-loss on the LM head done: TransformerBuilder::z_loss(weight) and TransformerLm::set_z_loss add weight * logsumexp(logits)^2 per predicted position to train_step and accumulate_step, reported in TotalLoss::auxiliary_loss. Default 0 leaves the loss unchanged. On CUDA the term folds into the fused ce_loss_grad kernels (FP32 and BF16). Tests: causal_lm_loss::tests::z_loss_gradient_matches_finite_differences, gpu_transformer::tests::a_z_loss_step_on_the_device_matches_the_host_or_skips_without_device (fails with the device term zeroed). No throughput row: the GPU was shared (2026-10-07).

B8. Preference and RL fine-tuning — P3, revisit

featuresPlan.md excludes RLHF and DPO. The interoceptive module now ships a DQN loop inside the crate, which weakens the "different library" argument. DPO needs no reward model and no sampling loop: it is a loss over pairs of sequences. Reopen this if B5 lands and someone wants to align a fine-tuned model.

Track C — Inference performance

C1. CUDA graphs for decode — P1, M — DONE 2026-10-08 (graphs rejected)

A decode step is many small kernels, so launch overhead dominates. Capture one decode step per shape and replay it. This needs a device-resident KV cache (C2) and static buffers. cuda_image.rs says "no graph capture" too. The UNet loop is also a candidate.

Profile 2026-10-07 (nsys, 129 decode calls): host time is dominated by cudarc's per-buffer event tracking (about 1770 cuStreamWaitEvent and 970 cuEventRecord calls a token, near 0.9 ms of API time), more than the 190 kernel launches. Turning tracking off (CudaContext::disable_event_tracking) gave 655 vs 656 tokens/s at batch 1 and 5280 vs 5455 at batch 16 on a clean GPU, because that host work overlaps the kernels. Rejected: within noise at batch 1, and it is unsafe for every buffer shared across streams. Graph capture still needs static buffers to remove the overlap ceiling.

Part 1 done 2026-10-07: routed experts without the host round trip. The profile at 79f9aac showed 19 device-to-host copies a token at batch 1: every MoE layer downloaded its routing (and its auxiliary-loss sums) so cuBLAS could shape one GEMM per expert, which drained the stream six times a token and would also block graph capture. A decode step of up to 16 rows now runs decode_moe: a gemv_routed kernel reads each slot's expert on the device for the gate/up product, and again for the down product summed over a token's slots with its gates; the auxiliary losses are skipped. Running every expert densely with zero gates instead (no new GEMV) removed the copies too but read four times the expert weights and left batch 1 flat. Int8 decode and prompts keep the old path. Default model, 512-token prompt, 256 new tokens, two interleaved rounds, load 4 to 10: 1 sequence 784-924 to 1357-1363 tokens/s, 4 sequences 2725-2730 to 3793-3795, 16 sequences 6929-6933 to 7918-7930. Tests: the decoder parity tests (a_device_decoder_matches_the_host_decoder_or_skips_without_device, a_device_batch_decoder_matches_host_decoders_or_skips_without_device, whose tiny model has two MoE layers and runs both the vector and scalar paths); a mutant that ignores the routing gates fails both tests. Evidence in docs/baseline.md ("Routed experts on the device in decode").

Part 2 (next): capture a decode step as a CUDA graph. Still needed: the remaining per-step host-to-device uploads (RMSNorm weights, token ids, positions) and per-call allocations must become static buffers, and the sequence length must come from device memory so one graph serves every step. Split 2026-10-08:

  • 2a, resident norm weights. Done 2026-10-08. Gpu::decode_norm keeps each RMSNorm weight in the decoder's PackMemo, so a token no longer uploads 17 of them. nsys, int8, one sequence, 129 decode calls: cuMemcpyHtoDAsync 3092 to 916 calls (24 to 7 a token). Throughput is unchanged within the load noise (one sequence 1208-1358 against 1237-1242 tokens/s BF16, 16 sequences 8904-8967 against 8700-8911): the host work overlapped the kernels. Tests: the device decoder parity tests.
  • 2a', one stream. Done 2026-10-08. A fresh nsys profile (int8, one sequence) had the GPU busy 25% of a decode window: cudarc waited on and recorded an event for every kernel argument, about 1000 cuStreamWaitEvent and 500 cuEventRecord calls a token, because GpuContext made its own stream and that put the context in multi-stream mode. GpuContext now runs on the legacy default stream, so there is one ordered stream and no tracking. Under nsys the GPU is busy 62% and decode runs 637 to 971 tokens/s. Unprofiled, five interleaved rounds at load 7: one sequence BF16 1200-1406 to 1406-1428, int8 1676-1712 to 1711-1747, 16 sequences 8746-8840 to 8772-8899; bench_train_step 4 1024 flat (53392-53428 against 53279-53630 tokens/s). Small unprofiled because the host work overlapped the kernels; with it gone, 2c has less to win at batch 1. Tests: the full suite with RUSTING_BRAIN_REQUIRE_CUDA=1.
  • 2b, static buffers and 2c, capture and replay rejected for speed 2026-10-08. After 2a', nvidia-smi sampled every 50 ms during an unprofiled 3000-token decode (one sequence, 64-token prompt) reads the GPU busy 95-97% (median 95 int8 at 1574 tokens/s, 97 BF16 at 1351): decode is GPU-bound, so removing launches can win about 5% at most. Redo only if the kernels get much faster or a profile shows the GPU idle again.

C2. Device-resident KV cache and batched decode — P1, M — DONE

This is W9 in docs/improvement-plan.md, still not measured. Keep the KV cache on the device across Decoder calls, and decode several sequences in one batch. Measure tokens/s with examples/bench_decode.rs before and after.

Step 1 done: bench_decode <prompt> <new> cuda decodes the default model on the device. Before the cache work it runs 197 tokens/s (5.1 ms/token) with about 1.1 ms of kernels and 0.8 ms of copies a token; the rest is host round-trips, about 174 device-to-host copies a token (docs/baseline.md). Step 2 done: a Decoder on a device model keeps every layer's projected rows on the device (gpu_model::DeviceKv) and runs the whole block there; generate goes through Decoder. 384 tokens/s, 1.95x. Parity test a_device_decoder_matches_the_host_decoder_or_skips_without_device covers a prompt, single tokens and a second multi-token turn in FP32 and BF16. Step 3 done: DeviceKv keeps every weight pack for the decoder's life (PackMemo, lent to the context per call). 624 tokens/s, 3.2x the step-1 number; kernels are 1.01 ms of a 1.6 ms token, the rest is launch overhead (C1). Step 4 done: TransformerLm::batch_decoder(n) returns a BatchDecoder that decodes n equal-length sequences in step. On a device the cache rows interleave position-major, so one step's rows stay contiguous and the projection, RoPE (new divisor argument) and logits run once for all of them; attention runs as strided batched GEMMs. Default model, 512-token prompt: 618 tokens/s at 1 sequence, 1994 at 4, 5194 at 16, 9093 at 64. Parity tests a_batch_decoder_matches_one_decoder_per_sequence (host) and a_device_batch_decoder_matches_host_decoders_or_skips_without_device (FP32 and BF16); breaking the RoPE divisor fails the device test. Ragged batches need C3. Step 5 done (2026-10-07): under mixed precision the device decode keeps a BF16 copy of the LM head (DeviceKv::head) and narrows the final norm, so the head GEMV reads half the bytes. nsys put the FP32 head at a fifth of a step's kernel time (median 200 us); BF16 runs 120 us. Same-session A/B on an idle GPU: 626 to 666 tokens/s at one sequence. Breaking the last-rows offset fails the three device decoder parity tests.

C3. Paged KV cache — P2, M — part 1 done 2026-10-08

Allocate the KV cache in fixed blocks so sequences of different lengths share a pool. Do this after C2, and only if batched decode with uneven lengths is a real use case.

2026-10-08: C14 made uneven prompts a supported use; they are padded, so this item would now win back the padding's cache capacity.

Split 2026-10-08. The device cache was allocated at max_seq_len for every sequence when a decoder was made, whatever the session's length, and C16's fork and C17's beam search copied all of it.

  • Part 1, done 2026-10-08: grow the device cache on demand. DeviceKv starts empty and reserve grows every layer's buffer and the padding slots to the positions a call needs, at least doubling and at most max_seq_len, copying the cached prefix. The layout is unchanged, so no kernel changed. Default model (8 layers, 8 + 2 x 2 heads of 64, FP32 rows): 24,576 bytes a position a sequence, so 50.3 MB a sequence before any token; now a 20-token prompt holds 0.5 MB and a 512 + 256 session 25.2 MB. Tests: the existing device decoder, batch decoder, long decode and fork tests cross several growths; skipping the prefix copy fails three of them. Decode throughput within 0.3% (docs/baseline.md, "KV cache grown on demand (C3 part 1)").
  • Part 2: blocks shared across sequences. A block table per sequence, so uneven sequences in one batch take only their own length. Every attention path (decode_attention, the score GEMMs over a strided history) reads a contiguous position-major history today, so this changes the kernels and is L, not M. Only worth it for continuous batching, which the "Not doing" list keeps out (serving). Parked.

C4. Speculative decoding — P2, M — greedy done 2026-10-07, sampled done 2026-10-08

Use a small draft model and verify with the large one. This needs C2. It is pure library code once batched verification exists.

Greedy done: TransformerLm::generate_speculative(draft, prompt, n, lookahead) drafts up to lookahead ids, scores them in one target pass (the device decode now returns the logits of the last positions rows), keeps the matching prefix plus the target's own id, and rewinds both decoders (KvCache::truncate, DeviceKv::truncate). Tests speculative_decoding_matches_greedy_decoding (host) and device_speculative_decoding_matches_greedy_or_skips_without_device: equal to generate with a greedy sampler for lookahead 0, 1, 3 and 8, with the target as its own draft and with an unrelated one, run to max_seq_len. Accepting every proposal, or rewinding one position too few, fails them. Sampled done 2026-10-08: TransformerLm::generate_speculative_with(draft, prompt, n, lookahead, sampler) keeps a proposal x with probability min(1, p(x)/q(x)) and replaces the first refusal with a draw from max(0, p - q), so the output is distributed as sampling the target alone. generate_speculative is now this with a greedy sampler. Sampler gained crate-private probabilities, draw and uniform; pick keeps its behaviour. Test sampled_speculative_decoding_draws_from_the_target: 4000 runs of 3 ids at lookahead 2 with a draft at total-variation distance 0.82 from the target, empirical marginals of the first two ids within 0.05 of the exact ones. Keeping every proposal fails it (distance 0.82), and resampling a refusal from p instead of p - q fails it (0.13). The greedy tests pass unchanged on host and device. Left: a tokens/s measurement with a real draft/target pair on a quiet GPU; there is no trained pair in this repository to measure with.

C5. Weight-only int8 and int4 GEMM on CUDA — P1, M — DONE 2026-10-08

Closed 2026-10-08: int8 decode covers every block linear and the head (slices 1-4). Int4 was measured first (D3(b), 2026-10-07) and is not worth it for speed on the default model; take it as D3(b) when a model too big for int8 needs the memory.

Precision::Q8 is CPU only for the language model. A dequantize-in-register GEMV kernel makes decode memory-bound on 1 byte (or half a byte) per weight instead of 2. Decode on a 3060 is bandwidth-bound, so this is close to a 2x or 4x gain for batch 1.

Slice 1 done (2026-10-07): TransformerLm::set_int8_decode makes device decoders pack the attention projections, the dense and shared-expert feed-forwards and the LM head as int8 with a scale per row, the Precision::Q8 scheme, quantized on the device from the FP32 weights. gemv_q8 (one warp per output row, 16 weights a lane, up to four input rows per pass) serves up to 16 rows; a prompt dequantizes into a scratch copy and takes cuBLAS. Test: an_int8_device_decoder_matches_a_quantized_host_or_skips_without_device against a host model with the same weights quantized (dense, untied, FP32 within 1e-3, BF16 within 3e-2); reading the BF16 input out of order or scaling the dequantized weights by 1.01 fails it. Default model: one sequence 905 to 1000 tokens/s (+10%); 16 sequences 7050 to 6090 (-13%), the int8 GEMV at 21.8 us a call against cuBLAS's BF16 kernels. Opt-in, so off by default. Evidence in docs/baseline.md ("Int8 device decode"). Next slices: the routed experts (6 of the default model's 8 layers are MoE, and their BF16 GEMVs are 19% of the int8 step's kernel time); a multi-row kernel that beats cuBLAS at 16 rows (x in shared memory, or int8 tensor cores); int4. Rejected 2026-10-07: staging the input rows in shared memory for more than four rows (6050 to 4880 tokens/s at 16 sequences; the re-reads it removed hit L1). The parity test now has a six-token turn, which covers the GEMV's loop over row groups. Evidence in docs/baseline.md ("Rejected: shared-memory input rows"). Slice 2 done (2026-10-07): the routed experts and the router go int8 too, so int8 decode now follows the host Precision::Q8 scheme for every block linear. Gpu::project_at reads one expert out of the packed set. The parity test runs on the MoE tiny() model; reading expert 0's scales or weights for every expert fails it. Default model, same-session A/B: one sequence 895 to 975 tokens/s (+9%; int8 off reads 820 in that session, so int8 is now +19%), 16 sequences 5600 to 5680. Evidence in docs/baseline.md ("Int8 routed experts"). Slice 3 done (2026-10-07): the multi-row kernel. gemv_q8_mma runs a decode step of three or more sequences on the tensor cores: a block of eight warps owns 16 weight rows as the mma's A operand, takes up to 16 BF16 input rows as B and splits the width over its warps. Test: gpu_model::tests::an_int8_projection_matches_the_host_for_every_batch_or_skips_without_device (1 to 16 rows, a partial unit tile, against the host product on the device's int8 bytes); feeding the second sequence tile the wrong inputs or dropping one warp's partial sum fails it. Default model: 16 sequences 5900 to 7600 tokens/s (BF16 reads 6780 in that session), 8 sequences 4510 to 5140, one sequence unchanged. Evidence in docs/baseline.md ("Int8 decode GEMV on the tensor cores"). Slice 4 done (2026-10-08): int8 decode takes the routed GEMV too. Under int8 the MoE layers still ran forward_moe, which downloads the routing every layer; gemv_routed_q8 now reads each slot's int8 expert row and its scale on the device, as gemv_routed does at BF16, when both widths are a multiple of 16. Test: the int8 decoder parity test now also runs an expert width of 16 (12 keeps the old path); reading expert row 0's scale for every row fails it. Default model, bench_decode 512 128 cuda <n> int8, three interleaved rounds, load 15-40: one sequence 1014-1049 to 1535-1613 tokens/s (+53%), 16 sequences 7732-7895 to 7675-7804 (flat). Evidence in docs/baseline.md ("Int8 routed GEMV in decode"). Left: int4, closed above.

C6. Image models on CUDA, the rest — P1, M — DONE 2026-10-07

Split 2026-10-07 into four parts.

  • Part 1, the T5 encoder on the device: cuda_t5::DeviceT5. Done 2026-10-07. It runs on the BF16 language-model kernels, not the FP16 image path, because T5 1.1 activations pass 65504; swiglu_fwd gained a tanh-GELU gate for it. Test: cuda_t5::tests::a_device_t5_matches_the_host_or_skips_without_device (tiny gated config, F32 and Q8, within 2% of the largest state); a mutant that drops the position bias fails it (F32 worst 0.037 of 0.42). The ignored time_a_t5_base_sized_encoder (d_model 768, 12 layers, 512 tokens): host 528 ms, device 28.6-33.8 ms, worst 0.0073 of 0.39. Evidence in docs/baseline.md ("T5 encoder on the device").
  • Part 2: the transformer denoiser (Dit) on the device. Done 2026-10-07. cuda_mmdit::DeviceDit, BF16 weights, FP32 residual streams, fused BF16 attention (heads of 64 or 128), and three small FP32 kernels: the modulated layer norm, the head norm with the rotation, and a bias/GELU column slice. Every block's modulation is one stacked GEMV. Tests: cuda_mmdit::tests::the_mmdit_kernels_match_the_host (each kernel against the host to 1e-4; a mutant that drops the rotation fails it, worst 0.43) and cuda_mmdit::tests::a_device_dit_matches_the_host_or_skips_without_device (whole model at F32 and Q8, relative L2 0.020 and 0.008 against a 0.03 bound). The ignored time_a_dit_over_a_512_image (d_model 768, 2 double and 4 single blocks, 1024 patches, 256 text tokens): host 1.27 s, device 15.7-17.1 ms. Evidence in docs/baseline.md ("MMDiT on the device").
  • Part 3: wire both into the pipeline and measure BF16 weights against the byte path, per featuresPlan.md Stage 6. Done 2026-10-07. ImagePipeline::to_cuda now attaches a transformer denoiser and a gated T5 encoder on one BF16 GpuContext; the original ungated T5 stays on the host. Test: pipeline::tests::a_transformer_pipeline_on_the_device_matches_the_host_or_skips_without_device (a tiny Stable Diffusion 3 with heads of 64: prompt states and a two-step image against the host, relative L2 within 2% and 5%, and not bit-equal, so the host path did not run). BF16 against int8 weights on the device, time_a_dit_over_a_512_image: 15.7-16.2 ms against 15.8-16.8 ms a pass, so the byte path costs no time and halves the weights; both stay.
  • Part 4: src/cuda_image.rs:438: replace the naive transpose with a tiled shared-memory transpose, if the profile shows it. Rejected 2026-10-07. nsys on cuda_image_bench (SDXL, 1024x1024, one UNet pass plus the VAE decode): transpose is 70 launches, 17.7 ms, 0.96% of 1847 ms of kernel time, so a perfect transpose saves under 1%. The same profile puts im2col at 18.7%; that is C10.
  • FP16 overflow guard on the SDXL decoder (cuda_image.rs:25): make sure that it fails loudly and does not produce a black image. Done 2026-10-07. DeviceVae::decode returns an error when any output value is not finite, and names the host path as the fallback. Test: cuda_image::tests::a_vae_decode_that_overflows_fp16_fails_instead_of_returning_nan (a latent past 65504 decodes finite on the host and errors on the device).

C7. Diffusion sampling — P2, S each

  • More solvers: DPM++ 2M Karras sigmas done (DPM++ 2M was already there; Scheduler::karras_sigmas adds the spacing and sample_with_sigmas runs any caller-chosen levels. Test: diffusion::tests::karras_sigmas_span_the_stable_diffusion_range checks diffusers' 14.6146 / 0.0292 range for SD betas and a 10-step DPM++ run. 2026-10-07), UniPC done 2026-10-07: sample_unipc runs diffusers' default UniPC (bh2, second order, predicting the image, first order on the first and last step) for both schedules. Test: diffusion::tests::unipc_matches_the_diffusers_scheduler matches UniPCMultistepScheduler from diffusers 0.40 within 1e-5 at 2, 8 and 20 steps on a nonlinear toy model; dropping the first-order corrector or fixing the second-order weights at 0.5 fails it. LCM done 2026-10-07 (timesteps confirmed against diffusers 0.40 LCMScheduler): Scheduler::lcm_sigmas(steps) picks the timesteps diffusers' LCMScheduler picks from its 50-step teacher grid (4 steps: 999, 759, 499, 259), and sample_consistency runs multistep consistency sampling (take the answer, renoise it to the next level). Solver is an exhaustive public enum, so this is a function, not a variant. The c_skip boundary blend is left out: below 1e-5 at every grid timestep. Tests: diffusion::tests::lcm_sigmas_sit_on_the_timesteps_diffusers_picks and consistency_sampling_renoises_each_answer_to_the_next_level (fails when the renoise uses the current level). Not run on a real LCM checkpoint.
  • CFG batching: run the guided and unguided passes as one forward pass of batch 2. Rejected 2026-10-08. Quiet GPU, cuda_image_bench pony-v6-xl (SDXL), QUANTIZE=1: one UNet pass 0.071 s at 512x512 and 0.254 s at 1024x1024, so four times the pixels costs 3.6 times the time. A batch of two at 512 is at best a 1024-sized pass, about 10% under two passes, and at 1024 the kernels are already that big, so nearly nothing. The device UNet is single-image throughout (convolutions, group norm, attention), so batching it is an L item for at most 10% at the small size. Redo if a small-image model becomes the main use.
  • Tiled VAE decode, so large images fit in 12 GB done: VaeDecoder::decode_tiled(latent, tile, overlap) decodes overlapping latent squares through decode (so on the device too) and blends the overlaps with linear ramps. Tests: vae::tests::tiles_blend_back_into_the_whole_image (a local decode must stitch back exactly; fails with unweighted normalisation) and a_tiled_decode_matches_the_whole_decode_when_one_tile_covers_it. Not measured on a real SD VAE: no checkpoint run, the GPU was shared. 2026-10-07. Seen once while verifying, unrelated: mesh::simplify_tests::an_error_bound_removes_the_flat_interior_and_keeps_the_shape failed in one full CPU run and passed in six reruns; it may be flaky.

C8. Fused decode attention — P1, S — DONE 2026-10-07

Found by the C1 profile: a single-row decode step ran attention as one cuBLAS GEMV per head for the scores, a softmax, and one per head for the values, FP32 throughout. On the default model that is 17 launches a layer and about 70 us, a third of the token's kernel time. decode_attention now splits the history into 64-position chunks, one block per (head, sequence, chunk), reading keys and values in place from the device cache, and decode_attention_merge rescales and adds the chunks. Several new rows (a prompt, a speculative check) keep the three-kernel path. Tests: the existing device decoder, batch decoder and speculative tests, plus a_long_device_decode_matches_the_host_or_skips_without_device (130-token prompt, three chunks, two sequences); a wrong score scale fails the first two, a merge without the rescale fails the last. Attention 70 to 19 us a layer; 663 to 864 tokens/s for one sequence (+30%) and about 5350 to 6870 for 16 (+28%). Evidence in docs/baseline.md ("Fused decode attention"). Follow-up 2026-10-08: the kernel scored one position per warp with a shuffle reduction and summed values one position at a time, so its loads waited on each other. Four threads per position for the scores and an unrolled value loop put every load in flight at once. Warm-clock microbench (8 heads, 2 KV heads, width 64): 7.64 to 5.05 us at 64 positions, 8.46 to 6.59 at 640, 20.46 to 14.00 at 1664. bench_decode, BF16, three interleaved rounds, load 18-33: 16 sequences 7780-7980 to 8731-8835 tokens/s (+11%) after a 512-token prompt, 4335-4380 to 4548-5294 after 1536; one sequence within the load noise. Tests: the device decoder parity tests, including a_long_device_decode_matches_the_host_or_skips_without_device; halving the score reduction fails three of them. Evidence in docs/baseline.md ("Decode attention loads in flight").

C9. BF16 decode attention output — P1, S — DONE 2026-10-07

Found by the profile after C8: under mixed precision the attention output projection was the only decode GEMV still in FP32 (1024 calls, 6.0 us each). decode_attention_merge now writes its output BF16 when the context is in mixed precision, so the projection reads the BF16 weight copy. The cuBLAS path for several new rows still writes FP32. Test: the mixed-precision case of a_device_decoder_matches_the_host_decoder_or_skips_without_device, whose tolerance went from 5e-2 to 5e-3 (measured error at most 1.3e-3); a merge that scales its BF16 output by 1.5 now fails it, and passed the old bound. 893-898 to 919-931 tokens/s for one sequence (+3.5%); 16 sequences unchanged within noise. Evidence in docs/baseline.md ("BF16 decode attention output").

C10. Implicit-GEMM convolution for the image path — P2, L

Added 2026-10-07 from the C6 part 4 profile. im2col is 18.7% of the SDXL 1024x1024 bench's kernel time (2804 launches), behind only the GEMMs it feeds. The kernel is already coalesced and writes nine taps a thread, so what is left is not writing the column matrix at all: a convolution kernel that gathers its B tiles from the input map straight into shared memory. cuDNN would do it but is a new dependency, which needs the owner. Acceptance: a parity test against Conv2d::forward for 3x3 stride 1 and 2 with padding, at FP16 and BF16; cuda_image_bench 1024x1024 UNet pass and VAE decode faster than 0.251 s and 0.53 s under the same load.

Split 2026-10-08. The image path has no BF16 module (its activations are FP16 throughout), so parity is checked at FP16 only.

  • Part 1 — DONE 2026-10-08: row-wise lowering for stride-1 convolutions. im2row expands only the horizontal taps, so its buffer holds k copies of the input instead of k * k; each vertical tap is then a whole-row offset into that buffer, and the convolution is k accumulating cuBLAS GEMMs over weights repacked at upload into [ky][out][in * k + kx]. Stride-2 convolutions keep im2col. Test: a_device_convolution_matches_the_host_or_skips_without_a_device now also covers 3x3 unpadded and 5x5 padding 2, beside the chunked test; making every tap overwrite instead of accumulate fails it. cuda_image_bench 1024x1024 (load 3.1-4.5): VAE decode 0.52 to 0.49 s in all three rounds; UNet pass 0.244-0.247 to 0.241-0.247 s, unchanged within noise. The gather drops from 346 ms (im2col, 2804 launches) to 176 ms (im2row, 1038 launches); about half of that comes back as GEMM time. Evidence in docs/baseline.md ("Row-wise convolution lowering").
  • Part 2: a true implicit-GEMM kernel that gathers B tiles from the map into shared memory, for the UNet's 0.251 s target. Open; the risks are that the CHW layout leaves shifted taps unaligned for 16-byte loads and that no hand GEMM in this crate matches cuBLAS yet (13.6-15.8 against about 22 TFLOP/s). Parked 2026-10-08 on the part 1 profile. Split by launch shape, the UNet's im2row launches (320 to 1280 channels) total about 22 ms over the whole bench run, 1-2% of a 0.245 s pass, so the UNet target cannot come from the gather. The VAE's (128 to 512 channels) total about 150 ms, at about 250 GB/s, against roughly 300 ms of GEMM in the same decode. A hand implicit GEMM at this crate's 60-70% of cuBLAS would lose about as much GEMM time as it saves in gather. Worth taking only through cuDNN, which is a new dependency. Owner: may the image path depend on cuDNN (through cudarc's cudnn feature) for its 3x3 convolutions?

C11. Pad SDXL's second prompt with its own pad token — P0, S — DONE 2026-10-08

Found while writing the J7 image tutorial: every SDXL text-to-image result carried a cross-hatched texture, and image-to-image at strength 0.3 with the same prompt moved the picture by rel 0.37 and painted streaks over it. A reference check against ComfyUI (already installed, no new dependency) on Pony Diffusion V6 XL cleared the UNet (eps rel 0.0010 to 0.0060 at sigma 14.6, 1 and 0.1, device Q8 against CPU FP32), the VAE encoder (rel 0.012 against FP32) and CLIP-L (rel 7.5e-5), and found CLIP-G off by rel 1.07 on every position after the prompt. ClipTextEncoder::tokenize padded with the end token; SDXL's tokenizer_2 pads with !, and the UNet cross-attends to the padded positions. ImagePipeline::load now reads pad_token from each CLIP tokenizer's tokenizer_config.json and calls the new ClipTextEncoder::set_pad_token. After: CLIP-G rel 0.021 in FP32, the texture is gone, and the same image-to-image run moves by rel 0.147 with no streaks. Tests: a_named_pad_token_fills_after_the_end_token, and two_clip_towers_are_read_side_by_side_the_way_an_xl_model_wants now checks both towers' padding (it fails without the fix). Stable Diffusion 2.x and SD3 name ! too and are fixed by the same change. The CLIP-G residual of 0.021 comes from the tanh GELU approximation where the checkpoint asks for exact gelu; that is P3 until a measurement shows it in an image. Evidence in docs/baseline.md ("Image pipeline reference check").

C12. VAE encoder on the device — P2, M — DONE 2026-10-08

Added 2026-10-08 from the J7 tutorial run. VaeEncoder has no attach_device, so image-to-image encodes on the CPU: 1024x1024 takes about 117-140 s against 13 s for the whole 25-step text-to-image run on the RTX 3060, and 512x512 takes 10-18 s. The decoder's device blocks (resnet, mid attention, group norm, convolution) cover everything but the stride-2 downsample. Acceptance: a parity test against VaeEncoder::encode at FP16; 1024x1024 encode under 2 s under the same load.

Done: VaeEncoder::attach_device uploads a DeviceVaeEncoder that reuses the decoder's resnet and attention blocks, and ImagePipeline::to_cuda attaches it. The downsample's right-and-bottom-only padding is a pad_end on the device convolution: im2col already reads zeros past the edge, so only the output size changes. Test a_whole_vae_encode_matches_the_host_or_skips_without_a_device (fails with pad_end 0). On pony-v6-xl, RTX 3060 shared with another job, load 2.3-2.7: 1024x1024 encode 117-140 s to 0.25 s, 512x512 10-18 s to 0.06 s. Against ComfyUI's FP32 encoder the device latent reads rel 0.066 (host F32 0.012, ComfyUI's own BF16 0.108). The tutorial's repaint at strength 0.5 went from 117.2 s to 6.9 s.

C13. Min-p sampling and presence/frequency penalties — P2, S — DONE 2026-10-08

Added 2026-10-08. Sampler has temperature, top-k, top-p and a multiplicative repetition penalty. Min-p (keep tokens whose probability is at least p times the top one) is the truncation most local-inference users now reach for first, and the additive presence and frequency penalties are the OpenAI-style knobs that chat front ends pass through. All are a few lines in Sampler::kept, so pick, probabilities and sampled speculative decoding get them at once. Acceptance: unit tests that min-p keeps exactly the tokens at or above the threshold and always the argmax, and that the two penalties subtract presence once and frequency per occurrence.

Done: Sampler::min_p, Sampler::presence_penalty, Sampler::frequency_penalty. Min-p applies after the temperature and top-k, before top-p, as in Hugging Face's warper order. Tests min_p_keeps_the_tokens_at_or_above_its_share_of_the_top and presence_counts_once_and_frequency_per_occurrence (counting the presence penalty per occurrence fails the second).

C14. Uneven prompt lengths in BatchDecoder — P2, M — DONE 2026-10-08

Added 2026-10-08. BatchDecoder made every call feed every sequence the same number of tokens, so batched serving of real prompts, which never have equal lengths, needed the caller to pad and mask by hand, and the library had no mask to do it with. Acceptance: a batch decoder fed prompts of different lengths gives each sequence the logits of a single-sequence host decoder fed that prompt alone, on the host and on a device, including a prompt padded by more than one 64-position attention chunk.

Done: on a device, shorter prompts are padded on the left with id 0, and a per-position slot flag in DeviceKv masks padding keys out of attention: in decode_attention for the fused path, and in a new mask_padding kernel before the softmax of the score path. Rotary positions are relative, so a left-padded sequence scores as it would alone. Only the prompts may differ: padding later in the history would change the distance to earlier tokens, and feeding uneven lengths after the prompt is an error. Padding still takes max_seq_len capacity, which C3 would remove. Tests a_device_batch_decoder_matches_host_decoders_or_skips_without_device (prompts of 4, 2 and 6 tokens) and a_long_device_decode_matches_the_host_or_skips_without_device (130 and 40 tokens, 90 padding positions). Dropping either mask fails both. Decode throughput: see docs/baseline.md, "Padded batch prompts (C14)".

C15. Next-token logits on Decoder and BatchDecoder — P2, S — DONE 2026-10-08

Added 2026-10-08. The decoders kept their logits private, so every choice Sampler does not make (constrained or grammar-guided output, a logit bias, log-probabilities for ranking or classification, beam search) had to rerun the whole forward pass through TransformerLm::forward. Acceptance: a public accessor for the next-token logits and a way to append a caller-chosen id with the forward pass deferred as next defers it, with a test that pushing the argmax reproduces greedy decoding.

Done: Decoder::logits, Decoder::push, BatchDecoder::logits(sequence), BatchDecoder::push(ids). Test pushing_the_argmax_of_the_logits_decodes_greedily (fails if push does not queue the id for the next pass).

C16. Forking a Decoder — P2, S — DONE 2026-10-08

Added 2026-10-08. Every request that starts with the same system prompt or few-shot preamble prefilled it again, and beam search, which branches a session, could not be built on the decoders at all. Acceptance: a fork that continues independently of its original, each matching a fresh session fed the same history, on the host and on a device.

Done: Decoder::fork. On a device the fork copies the cached rows and shares the packed weights and the narrowed LM head through Arc. Tests a_forked_decoder_continues_independently and a_forked_device_decoder_continues_independently_or_skips_without_device (a fork that drops the device caches fails it). BatchDecoder has no fork yet; add one when batched beam search wants it.

C17. Beam search — P2, S — DONE 2026-10-08

Added 2026-10-08. The crate had sampling and speculative decoding but no beam search, the usual decoder for translation, summarisation and other tasks where the single most probable output matters more than variety. Acceptance: one beam is greedy decoding, and with a beam per vocabulary id two steps return the most probable pair found by exhaustive search.

Done: TransformerLm::generate_beam(prompt, max_new_tokens, beams, end) on Decoder::fork, logits and push (C15, C16). A continuation that samples end finishes and gives up its beam; the result is the one with the best mean log-probability per token. Test beam_search_finds_the_most_probable_continuation (keeping one candidate a step, or scoring each step without its prefix, fails it). Each step forks every kept beam, a full cache copy; reusing a parent for its last child is the first saving if beam search turns out hot.

C18. Per-token log-probabilities — P2, S — DONE 2026-10-08

Added 2026-10-08. Scoring a given text (ranking candidate answers, classifying by the most probable label, a per-document perplexity) needed the caller to run forward_train and take the log-softmax by hand. Acceptance: log_probabilities(ids) whose negated mean equals evaluate's loss on that sequence and whose entries match a decoder's logits, on the host and on a device.

Done: TransformerLm::log_probabilities. Tests log_probabilities_match_the_loss_and_the_decoder and device_log_probabilities_match_the_loss_or_skip_without_device (scoring each row's own input id instead of the next fails them). It runs the training forward and downloads [len, vocab] logits, as evaluate does.

C19. Grouped-query decode attention — P2, M — DONE 2026-10-08

Added 2026-10-08. Profiles of bench_decode showed the fused decode attention at 23% of the kernel time for one int8 sequence after 4000 tokens and 31% for 16 sequences after 2000. Each block served one query head, so with 8 heads over 2 key-value heads every cached key and value was read by four blocks. Acceptance: one block per (key-value head, sequence, chunk) scoring every query head of its group, passing the existing device decoder tests, a mutation that breaks the grouping failing them, and tokens/s no lower at 64 positions and higher at 2000 and 4000.

Done: decode_attention in src/cuda_training.rs serves up to DECODE_GROUP (8) query heads a block; larger groups take the cuBLAS path. The device decoder tests (a_long_device_decode_matches_the_host_or_skips_without_device and the other 33 matching decod, beam, fork, log_prob) pass; scoring every query head of a group with the group's first query fails 6 of them. bench_decode tokens/s, 2f02760 against the change: 4000 + 256 int8 one sequence 1197-1212 to 1425-1437, 2000 + 128 x 16 5016-5024 to 6906-6915, 512 + 128 x 16 9163-9182 to 10348-10452, 64 + 256 int8 one sequence 1774 to 1783-1810. Numbers in docs/baseline.md.

C20. Flash attention for decoder prompts — P2, S — DONE 2026-10-08

Added 2026-10-08. After C19, nsys put the prompt path's attention_softmax_lse at 21% of kernel time for 16 sequences after a 2000-token prompt, plus the two GEMMs around a [history, history] score matrix a head (2 GB at that size). The time to the first token is that prefill. Acceptance: an even first prompt under mixed precision at a head width the training flash kernel takes runs that kernel, matching the host decoder; bench_decode reports the prefill time and it drops at 2000 positions.

Done: decode_block launches the training flash forward once per sequence on a BF16 copy of the prompt's cache rows, taking the interleaved layout as rows sequences times as wide. Later feeds, padded prompts, FP32 and other head widths keep the cuBLAS path. Test a_device_prompt_at_head_width_64_matches_the_host_or_skips_without_device (dropping the causal flag, the per-sequence offset or the widened output stride each fails it). Prefill, 65355e8 against the change: 2000 x 16 363-365 to 149 ms, 1000 x 16 110-111 to 67-70, 512 x 16 44-45 to 35-38, 2000 x 1 27-30 to 12.5, 4000 x 1 int8 97-98 to 32-38. Decode steps are unchanged.

Follow-up 2026-10-08: left-padded prompts of uneven length take the flash kernel too, each sequence launched over its real positions only, with the padding rows of the output zeroed. The test now also pads one sequence by 90 positions; launching over the padding or skipping the zeroing fails it. Prefill with every other prompt half as long, 6566c19 against the change: 2000 x 16 438-440 to 125-126 ms, 512 x 16 47-48 to 32, 2000 x 2 59-92 to 19.

Follow-up 2026-10-08: a later turn of 128 rows or more (a chat turn, a long speculative check) takes the flash kernel too. The kernel's causal flag now carries the key position of the first query plus one, so query row r sees keys up to r + shift; training passes 1 and is unchanged. bench_decode 2000 8 cuda 16 "" 512 turn 140-160 to 60-70 ms; below 128 rows the BF16 copy of the history costs more than it saves (16 x 8: 10 ms against 25), so those keep cuBLAS. Test: the same prompt test now feeds a 130-row turn after the cache (error under 4e-3, bound 1e-2); masking only each query tile's last key tile fails it (0.038 against 0.026).

C21. BF16 key-value cache for device decode — P2, M — DONE 2026-10-08

Added 2026-10-08 from an nsys profile of bench_decode 2000 128 cuda 16 after C20: decode_attention is 23% of kernel time, 128 us a layer a step, reading about 34 MB of FP32 cache rows a layer (about 264 GB/s of 360), so it is bound by the bytes. The cache also keeps every row's query, which no later step reads: 3072 bytes a position a layer on the default model where keys and values in BF16 take 512. Under mixed precision with the fused kernels available, keep only keys and values, in BF16; queries go to a scratch buffer for the call. Multi-row turns then read the cache in place through the flash kernel, with no BF16 copy of the history. Acceptance: the device decoder parity tests pass under mixed precision with the new cache (and a mutation in the new store fails them); bench_decode 2000 128 cuda 16 tokens/s higher and one-sequence decode no lower; KV bytes a position a layer 3072 to 512 on the default model; no small-turn regression (bench_decode 2000 8 cuda 16 "" 8). Done: the mixed-precision cache holds BF16 keys and values only, 512 bytes a position a layer on the default model against 3072. Tests: transformer::tests::a_device_decoder_matches_the_host_decoder_or_skips_without_device, a_device_batch_decoder_matches_host_decoders_or_skips_without_device, a_device_prompt_at_head_width_64_matches_the_host_or_skips_without_device, an_int8_device_decoder_matches_a_quantized_host_or_skips_without_device, device_log_probabilities_match_the_loss_or_skip_without_device and device_speculative_decoding_matches_greedy_or_skips_without_device pass; storing each new row one position late fails the head-width-64 prompt test. bench_decode 2000 128 cuda 16: 7040 to 7360 tokens/s; 8-step turns 2.89 to 2.30 ms a step; one sequence unchanged at 0.71 ms a step (load 5.4-6.5, docs/baseline.md).

Track D — CPU path

The CPU path is the parity reference and the no-GPU fallback. Keep it correct and reasonably fast, and do not chase peak throughput on it.

# Item Effort Notes
D1 MoE loop order rejected 2026-10-07 S Measured: at 2048 tokens, d_model 512, 8 experts, auxiliary_losses takes 11 µs against a 35 ms forward and 67 ms backward pass (0.01%). The column walk covers only [tokens, 8]. See docs/baseline.md ("CPU MoE pieces").
D2 One exp per element in softmax_rows rejected 2026-10-07 S Measured at the same shape: softmax_rows takes 91 µs of a 102 ms forward-plus-backward (0.09%), so halving its exp calls saves under 0.05%. Below D4, which was rejected at 0.4%.
D3 int4 / NF4 weight storage M Shared by B4 (QLoRA) and C5. Block-wise scales, same layout on CPU and CUDA. Split 2026-10-07: (a) host storage: a 4-bit index per weight, two to a byte, into a 16-entry codebook (symmetric int4 or NF4), with one f32 absmax per block of 64 along a row; matmul, dequantize, and a test that the product matches the dequantized float product and that NF4 has the lower error on normal weights. (b) quantize a model's linears to it and decode on the device with a 4-bit GEMV (C5's int4). (c) NF4 frozen base for LoRA (B4 slice 2). (a) done 2026-10-07: quantized::{Codebook, Quantized::from_matrix_4bit} and TransformerLm::quantize_4bit (block linears at 4 bits, embedding and untied head at int8; training, BIN Q8 save and device decode refuse). Tests a_4bit_product_matches_the_float_product_of_its_own_weights, nf4_loses_less_than_the_even_grid_on_normal_weights (RMS int4 0.1074, NF4 0.0940), a_4bit_block_of_zeros_stays_zero_and_its_neighbour_keeps_its_extremes, a_4bit_weight_takes_under_six_tenths_of_a_byte, a_4bit_model_matches_a_float_model_on_the_same_grid_and_refuses_to_train. Host decode at a 100M shape: NF4 292-331 tok/s against int8 255-271 (docs/baseline.md). (c) done 2026-10-07: see B4 slice 2. (b) measured 2026-10-07, not started: at one sequence gemv_q8 is 43% of device decode kernel time on an RTX 3060, but the 512-unit projections run at about 87 GB/s of 360, latency-bound, so halving their bytes buys little; only the LM head (60 us, about 270 GB/s, 6% of a token) is bandwidth-bound. Worth it for memory on a model too big for int8, not for speed on the default model.
D4 Drop the cloned inputs in caches S Measured at 0.4% and rejected. It only becomes worth it with many experts. P3.
D5 Conv2d backward on CUDA — Already exists: TrainableConv2d has a backward pass and a device path. The uncommitted diff fixes its bias, which never trained on the device.

Track E — Data

BatchSource, DatasetStream, JsonlStream, TokenStream and TokenFile exist already. What is missing:

# Item Priority Effort Notes
E1 Weighted dataset mixing P1 S Done: TokenMix in src/token_file.rs. Original plan: draw from N TokenFiles by weight, determined by the step number so that a resume is exact. Every real pretraining mix needs this.
E2 Parquet reader P2 M Most public corpora on the Hub ship as Parquet. Put it behind a feature flag so the default build stays small. Read one column of text, a row group at a time.
E3 WebDataset (tar shards) done P2 S rusting_brain::webdataset::WebDataset streams Samples (key plus (extension, bytes) members) from a tar shard, by hand with no new dependency: ustar, old GNU, GNU long names and pax path. Tests: members_group_into_samples_by_key, a_truncated_or_corrupt_shard_is_an_error, fuzz::webdataset_survives_corruption; also read correctly from shards written by GNU tar and by Python tarfile in pax and GNU format, with 150-character names. Decoding the images is left to vit_encoder::read_image or the caller. Original note: Image and text pairs for C6 training and for ViT. tar is a small crate, or about 100 lines by hand.
E4 Shuffled chunk windows P3 S src/token_file.rs:288 returns windows in file order. Evaluation wants that, so add it only if something else needs a shuffled full pass.
E5 Multi-line CSV fields P3 S src/dataset.rs:924 reads single-line rows only.
E6 Memory-mapped shards P3 S src/shards.rs:32 seeks instead of mapping. Change it only if a profile shows that the reads cost time.
E7 Tokenizer training P2 M Done 2026-10-07 (byte-level BPE). Bpe::train_json(texts, vocab_size, specials) trains a GPT-2 style byte-level BPE (heap of pair counts, lazy stale entries, deterministic ties) and writes a tokenizer.json that Bpe::from_json loads. Evidence: tokenizer::tests::a_trained_tokenizer_loads_and_round_trips (round trip on unseen Unicode, frequent word becomes one token, specials keep ids 0..n and no merge crosses them, same corpus gives the same file) and training_refuses_a_vocabulary_too_small_or_a_repeated_special. Also fixed an O(n²) in pre_tokenize (617 ms to 1.4 ms on 78k chars). Open: WordPiece/Unigram training; loading the output in HF tokenizers not checked (package not installed here).
E8 Linear-time sentencepiece BPE merge done P1 S Bpe::merge keeps the pairs in a heap over a linked list, n log n. the_heap_merge_matches_the_pair_by_pair_loop checks it against the old loop on random merges and text up to 1,000 symbols; on a real sentencepiece file the ids of 1 KB, 5 KB and 50 KB texts hash the same before and after, and 50 KB went from 43.1 s to 12.6 ms (docs/baseline.md). Original note: Bpe::merge (src/tokenizer.rs) rescans every adjacent pair after each merge, and a Metaspace or Prepend file (LLaMA 2, Mistral, TinyLlama) hands it the whole text as one piece. Measured in docs/baseline.md ("Tokenizer merge"): 5,376 characters take 432 ms, four times the text takes sixteen times as long, so one 100 KB document in TokenStream would take minutes. Replace the loop with a linked list and a rank-ordered heap. Acceptance: ids identical to the current loop on the existing tokenizer tests and on a 50 KB text, and time per character flat from 1 KB to 50 KB.

Track F — Interop and formats

# Item Priority Effort Notes
F1 Write safetensors P1 S Already existed: safetensors::write and write_as (F32, BF16). Writing lets a trained TransformerLm or LoRA adapter load in PyTorch, llama.cpp converters and the Hub. The format is a JSON header plus raw bytes.
F2 Hugging Face config export P1 S Done: TransformerLm::save_hf writes config.json and model.safetensors; a test reads them back through the separate TextEncoder LLaMA reader. Original plan: write a LLaMA-compatible config.json next to F1's output, so that transformers loads the result directly. This is the reverse of B5.
F3 GGUF export done 2026-10-07 P2 M Run trained models in llama.cpp and Ollama. TransformerLm::save_gguf(path, q8_0) writes a llama-architecture GGUF v3 file, F32 or Q8_0 (norms stay F32), with Q and K rows permuted for llama.cpp's adjacent-pair RoPE, and linear or llama3 rope scaling (rope_freqs). Checked against the gguf 0.19 Python reader: every tensor present with the right shape, Q8_0 bytes identical to gguf.quants.quantize, F32 tensors equal to save_hf after the converter's permute, rope_freqs equal to the converter's formula. Tests in gguf::tests; the RoPE-convention test fails without the permutation. Left: tokenizer metadata (written as no_vocab), and a load in llama.cpp itself, which is not installed here.
F4 ONNX export for MoE P3 M save_onnx covers dense networks and the dense LM. MoE needs dynamic routing in the graph, which ONNX handles poorly. Park it until someone asks.
F5 Checkpoint format version P1 S Already existed: every binary file starts with a versioned magic (RBWTS001, RBOPT001, RBLOR001, RBEMA001), save_bin stores the config as a JSON header, and load_bin checks the parameter count and every shape. Original plan: put a version field and a model-shape header in binary checkpoints, so that an old checkpoint fails with a clear error and not with garbage weights. Check first whether one already exists.

Track G — Model families and modalities

# Item Priority Effort Notes
G1 CLIP image tower done 2026-10-08, smaller than written P2 M The tower already existed as vit_encoder::VitEncoder. Added VitEncoder::embed (class token through visual_projection), clip::logit_scale and clip::zero_shot (softmax over scaled cosines; retrieval ranks by the same cosines). Against transformers 5.16.1 CLIPModel on a small random checkpoint (eos 3): image and text embeddings within 2.2e-7, zero-shot probabilities within 1.2e-7. That check found a bug: older configurations (OpenAI CLIP, Stable Diffusion 1.x, FLUX.1's CLIP-L) say bos 0, eos 2, which transformers treats as legacy and pools at the highest id; ClipTextConfig::from_file now maps that case to the vocabulary's last two ids (49406, 49407). Tests: vit_encoder::tests::the_embedding_is_the_projected_class_token, clip::tests::zero_shot_is_a_softmax_over_scaled_cosines, clip::tests::a_published_config_states_the_shape (legacy ids). Not run on a published checkpoint: none with both towers is on disk. Original note: src/clip.rs:13 has no image tower. With it, the library can do zero-shot classification and image-text retrieval, and it reuses VisionTransformer.
G2 ControlNet and inpainting — done 2026-10-08 P2 L Left out in featuresPlan.md Stage 7. Both condition on a second image. Part 1 (S): diffusion::Inpaint wraps any Denoiser with a picture latent and one keep weight a latent row; where it keeps, the prediction is the one whose clean image is the picture, so kept rows walk their own noise path to it under every solver and schedule, and the model paints the rest to match. ImagePipeline::sample_latent_inpaint(.., keep, strength, ..). Tests: diffusion::tests::inpainting_keeps_the_picture_and_repaints_the_rest (flow Euler, shifted flow DPM++ 2M, DDIM ancestral: kept rows land on the picture within 1e-3, repainted rows on the model's image, a 0.5 row halfway; a wrong-length mask errors) and the inpainting step in pipeline::tests::a_picture_can_start_the_run_instead_of_noise. Part 2 (M): inpainting UNets (in_channels 9 against out_channels 4) load and read the mask and the blanked picture's latent after the noisy latent through Unet::set_inpaint, with only the noisy latent scaled, as diffusers' scheduler does. ImagePipeline::generate_inpaint(prompt, image, keep, strength, ..) takes a pixel mask: a trained UNet gets the binarized mask sampled at each latent cell's corner and the latent of the picture with the repainted part zeroed, the way diffusers builds them; any other model goes through Inpaint with the mask averaged over each cell. Tests: unet::tests::an_inpainting_model_reads_the_mask_after_the_latent, pipeline::tests::inpainting_runs_through_a_trained_unet_and_any_other_model (both paths, wrong mask size errors, the mask does not leak into the next run), and the ignored unet::tests::an_inpainting_model_matches_diffusers against scripts/unet_inpaint_reference.py (diffusers 0.40, worst error 5.5e-7 on a mean output of 0.075; scaling the mask channels too fails it at 0.072). Part 3 (M): ControlNet loads a diffusers ControlNet (a trained copy of the UNet's way down, the controlnet_cond_embedding stack that brings an RGB picture in [0, 1] down to the latent, and one zero convolution a rung) and Unet::set_controls(vec![Control { net, image, scale }]) adds its scaled corrections to every skip and to the middle; several ControlNets add up. The UNet's way down moved into a shared Encoder. On a CUDA device the ControlNets run there too (DeviceControlNet, uploaded by set_controls; ponytail: note on Unet::upload_controls). Tests: unet::tests::a_controlnet_steers_the_unet_only_as_hard_as_its_scale (scale 0 gives the plain output bit for bit, scale 1 changes it, a picture of the wrong size errors), cuda_image::tests::a_controlnet_on_the_device_matches_the_host_or_skips_without_a_device (two ControlNets, within 0.05 of the host's largest value, and well closer to the host than the uncontrolled output is) and the ignored unet::tests::a_controlnet_matches_diffusers against scripts/controlnet_reference.py (diffusers 0.40, scale 0.7: worst error 8.2e-7 on a mean output of 0.108; leaving the scale out fails it at 0.040, leaving the middle correction out at 0.033).
G3 Diffusion fine-tuning (LoRA on a UNet or DiT) — done: part 1 (loading LoRA files) and part 2 for the DiT (training) 2026-10-08; UNet training is G10, done 2026-10-09 P2 L Training for the image path. It needs D5, A4 and B4 on the image blocks. Split 2026-10-08. Part 1 (M): load published LoRA files into an image model. Most people meet image LoRAs as files to apply, not to train. ShardedSafeTensors::add_lora(path, prefixes, scale) reads PEFT/diffusers (lora_A/lora_B) and kohya (lora_down/lora_up/alpha) pairs, matches them to the checkpoint's weights by name (dots or underscores), and adds scale * alpha / rank * up @ down to each weight as it is read; ImagePipeline::load_with_loras applies them to the denoiser. Acceptance: a test against PEFT's merge on a tiny diffusers UNet within 1e-5; keys for another component are skipped, unknown adapter kinds and unmatched names are errors. Done: tests safetensors::tests::a_lora_adds_up_times_down_to_the_weight_it_names (kohya and PEFT spellings give the exact merged weight; a stray name, a LoHa key, or asking for the text encoder's adapter is an error; another part's file changes nothing), pipeline::tests::a_lora_file_changes_the_models_it_is_loaded_into (strength 0 gives the plain image bit for bit; a file's text encoder half changes the picture further, through lora_te_ names on the CLIP towers; a text-encoder-only file is an error), and the ignored unet::tests::a_lora_file_merges_as_peft_does against scripts/unet_lora_reference.py (PEFT 0.19 fuse_lora at strength 0.6 over 72 adapters on attention, proj_in, the feed-forward and 3x3 resnet convolutions: worst error 1.1e-6 on a mean output of 0.19, the same for the kohya copy; ignoring alpha fails it at 0.008). kohya's SDXL files name the original UNet's blocks (input_blocks_4_1, middle_block_0_in_layers_2); add_lora aliases them to the diffusers names from the blocks the checkpoint holds, and a_lora_file_merges_as_peft_does loads such a file (kohya_sdxl_lora.safetensors, whose names diffusers' own converter maps back to the same modules) to the same output; starting the up blocks at index 1 fails it. kohya's FLUX files name the reference DiT's modules, whose fused img_attn_qkv and linear1 diffusers splits into to_q/to_k/to_v (and proj_mlp); one adapter may now feed several weights, which take its output rows in turn, and the pipeline takes lora_unet_ for a DiT too. Test mmdit::tests::a_kohya_lora_for_the_reference_names_adapts_both_spellings_alike (one kohya file on the reference and diffusers checkpoints gives the same prediction within 1e-5, and a different one from strength 0; swapping to_q and to_k in the split fails it). The final modulation (final_layer.adaLN_modulation, written in the other half order) is not aliased and stays an error. Part 2 (L): training, split 2026-10-08. The image models are inference-only (no backward anywhere in mmdit.rs or unet.rs); the LM's param::Lora already has the adapter math. The DiT goes first: it has fewer layer kinds than the UNet (linear, affine-free layer norm, adaLN modulation, per-head RMS norm, RoPE, attention, tanh GELU). 2a (M): a host reference for one LoRA step on Dit: a forward that keeps its activations and a backward to adapters on every block projection (qkv, projection, mlp_in, mlp_out, fused_in, fused_out), base frozen, conditioning frozen. It is the oracle for 2b, not a training path (training runs on the GPU). Acceptance: on a tiny diffusers FluxTransformer2DModel with PEFT adapters on the same modules, the adapter gradients of a fixed loss match PyTorch autograd within 1e-4 relative. 2a done 2026-10-08: dit_lora::DitLora reads a diffusers-FLUX PEFT file onto the fused projections (to_q/to_k/to_v are row ranges of qkv or fused_in, proj_mlp the rest of fused_in), forward keeps a Tape, and backward adds adapter gradients through every block, the RoPE (by its inverse rotation), both norms and attention. Test: the ignored dit_lora::tests::dit_lora_gradients_match_pytorch against scripts/dit_lora_reference.py (diffusers 0.40, PEFT, two double and two single blocks, all 68 adapter tensors; worst prediction error 3.4e-6, worst gradient error 5.4e-5 relative; dropping the layer norm's projection term fails it at 0.89, skipping the inverse RoPE on single-block queries at 1.4e-3). 2b (L): the same step on CUDA, matched to 2a, with step time measured. Split 2026-10-08: the device attention is the fused kernel, which takes heads of 64 or 128, so 2b is matched to 2a's host code on a tiny config with 64-wide heads rather than to the PyTorch file directly. 2b-i (M): DeviceDit runs the adapters; acceptance: the prediction matches DitLora::forward within the BF16 tolerance DeviceDit already meets. 2b-ii (L), split 2026-10-08. 2b-ii-a (S): the backward kernels the forward's own kernels need: the modulated layer norm, the head norm with RoPE, and the tanh GELU; acceptance: each matches the host reference's backward (dit_lora::layer_norm_backward, head_norm_backward after the inverse rotation, gelu_derivative) within 1e-4. 2b-ii-b (L): a device tape of what the backward reads and the device backward to adapter gradients, with those kernels, the existing flash backward, gate_residual_backward and linear_backward_input; acceptance: every adapter gradient matches DitLora::backward within 2e-2 relative, and one step at 1024 image tokens on the default FLUX shape scaled to fit 12 GB is timed in docs/baseline.md. 2b done 2026-10-08: DeviceDit::forward_with runs the adapters, backward keeps a device tape and returns every adapter gradient, with new kernels for the modulated layer norm, head norm with RoPE and tanh GELU backward. Tests: cuda_mmdit::tests::device_adapters_match_the_host_or_skip_without_device and device_adapter_gradients_match_the_host_or_skip_without_device (relative gradient error 0.011 overall, worst tensor 0.047 at double.0.text.projection; the bound is 2e-2 overall and 1e-1 a tensor, BF16 noise on single small tensors runs past 2e-2). Dropping the key norm from the double blocks' head norm backward fails it at 0.031 overall, 1.06 on single.1.fused_in. One step (forward, backward, adapter gradients) at 1024 image and 256 text tokens on d_model 768, 6 heads, 2 double and 4 single blocks: 45-56 ms after a 130 ms first step, 1612 MiB in use (time_a_dit_adapter_step, load 4.2). The full FLUX shape was not timed: its BF16 weights do not fit 12 GB. 2c (M), split 2026-10-08: FLUX.1's BF16 weights (24 GB) do not fit 12 GB and the device backward refuses int8 bases, so the training path is checked on small DiTs here. 2c-i (S): fresh adapters on every module a diffusers FLUX LoRA covers (DitLora::new), one rectified-flow step on the device (DeviceDit::flow_step: noise to sigma, regress noise - latent, clipped Adam on the adapters), and DitLora::write_peft; acceptance: a written file reads back to the same prediction within 1e-5, and 40 steps on one latent halve its loss. 2c-ii (M): a train_dit_lora binary over cached latents and prompt states that writes a PEFT file part 1 loads and diffusers loads too, run on the tiny diffusers DiT. The UNet follows as its own row once 2a-2c land. 2c done 2026-10-08: DitLora::new puts fresh adapters on the 34 modules diffusers' FLUX LoRAs cover (A random, B zero), DeviceDit::flow_step runs one rectified-flow step with clipped Adam, and DitLora::write_peft writes the diffusers PEFT names. examples/train_dit_lora.rs encodes captioned images and captions once, drops the encoders and trains on the device. Tests: dit_lora::tests::written_adapters_read_back_to_the_same_prediction (writing lora_A for lora_B fails it), cuda_mmdit::tests::flow_steps_fit_one_latent_or_skip_without_device (40 steps at 1e-2: loss 1.52 to 0.0020), and the ignored dit_lora::tests::a_trained_lora_loads_as_diffusers_loads_it with scripts/dit_lora_trained_check.py: 200 steps at 64x64 on three captioned pictures over the tiny diffusers DiT (loss 1.14 over steps 1-50 to 0.33 over 151-200), diffusers 0.40 load_lora_adapter loads all 34 adapters, the LoRA moves its prediction by 1.7 relative, and both read_peft and load_with_loras match it within 1.1e-6 relative.
G10 UNet LoRA training (Stable Diffusion 1.5 and SDXL) — a (host reference), b-i (device backward kernels), b-ii (adapted device forward) and b-iii-a (device attention backward) and b-iii-b (device backward to adapter gradients) done 2026-10-08; b-iii-c and c (the train_unet_lora example) done 2026-10-09 — G10 done P2 L Added 2026-10-08 after G3 part 2 trained the DiT. Most published image LoRAs are trained on SD 1.5 and SDXL UNets, and unet.rs, cuda_image.rs and DeviceUnet are inference-only. diffusers' own train_text_to_image_lora targets the attention projections (to_q, to_k, to_v, to_out.0), so the backward needs weight gradients only there, and input gradients through the resnet blocks (convolution, group norm, SiLU), the cross-attention and the up and down samplers. a (M): a host reference: a UNet forward that keeps its activations and a backward to adapters on the attention projections, base and conditioning frozen. Acceptance: on a tiny diffusers UNet2DConditionModel with PEFT adapters on those modules, every adapter gradient of a fixed loss matches PyTorch autograd within 1e-4 relative. b (L): the same on DeviceUnet in FP16; acceptance: adapter gradients within 2e-2 relative of a, and one step on SD 1.5 at 512x512 timed in docs/baseline.md. Split: b-i (M) FP16 device backward kernels (the reverse of im2col, group norm, affine layer norm, GEGLU, SiLU and nearest upsampling), each matched to its host function in unet_lora within an FP16 tolerance; b-ii (M) DeviceUnet forward with adapters, keeping a device tape, matched to UnetLora::forward; b-iii (L) the device backward to adapter gradients and the timed SD 1.5 step, split 2026-10-08: SD 1.5's heads are 40, 80 and 160 wide and the fused attention takes 64 or 128, so b-iii-a (M) a device attention backward for any head width, a chunk of query rows at a time, with key and value gradients summed in FP32; acceptance: within 2% of dit_lora::attention_backward for self- and cross-attention over several chunks. b-iii-b (L) the device backward to adapter gradients over the tape; acceptance: within 2e-2 relative of UnetLora::backward. b-iii-c (S) one SD 1.5 step at 512x512 timed in docs/baseline.md done 2026-10-09: 191.0 ms (forward 65.1, backward 125.9; random weights of the 1.5 shapes, load 8 to 10), cuda_image::lora::tests::an_sd15_step_at_512_is_timed (ignored). c (M): a train_unet_lora example (noise-prediction loss on the scheduler's noise levels, captioned images encoded once, as train_dit_lora); acceptance: on the tiny UNet the loss falls and diffusers' load_lora_weights loads the written file to the crate's own prediction within 1e-5. done 2026-10-09: UnetLora::new (he-uniform down, zero up, scale alpha/rank) and write_peft (diffusers PEFT keys, the scale folded into lora_B), DeviceUnet::noise_step (noisy latent, adapted pass, eps MSE, device backward, Adam with the gradient norm clipped to 1), and examples/train_unet_lora.rs. Tests: unet_lora::tests::written_adapters_read_back_to_the_same_prediction; cuda_image::lora::tests::noise_steps_lower_the_loss_or_skip_without_a_device (fixed noise at sigma 1.5: 1.227 to 1.139 in 30 steps, 1.097 in 200; rank-4 adapters on the tiny UNet's one attention level fit slowly); the ignored pipeline::tests::publish_a_tiny_sd and unet_lora::tests::a_trained_unet_lora_loads_as_diffusers_loads_it against scripts/unet_lora_trained_check.py (300 steps of the example on a diffusers random UNet, 48 adapters loaded by load_lora_adapter; read_peft within 8.5e-7 relative of diffusers, the LoRA moving the prediction 0.74). a done 2026-10-08: unet_lora::UnetLora reads a diffusers PEFT file onto those projections, forward keeps a Tape, and backward adds adapter gradients back through every transformer (attention, cross-attention with its frozen prompt, affine layer norms, GEGLU, the group norm and projections in and out), every resnet (convolutions by the reverse of im2col, group norms, SiLU), the down and up samplers and the skips. Tests: unet_lora::tests::adapter_gradients_match_finite_differences (all 88 adapters of a random-weight tiny UNet, each factor's slope along its own gradient within 1% plus 1e-6 of f32 noise), and the ignored unet_lora_gradients_match_pytorch against scripts/unet_lora_gradients.py (diffusers 0.40 and PEFT on the two-level UNet of unet_lora_reference.py, 96 adapter tensors: prediction within 3.0e-6, worst gradient 4.9e-5 relative). Dropping the group norm backward's projection term fails both (1.31 relative on the first to_k); skipping the skip gradients on the way down fails the first. The script runs diffusers' GEGLU through the tanh GELU, as unet.rs does; diffusers itself uses the exact one, which is why the inference test against diffusers allows 1e-3. b-i done 2026-10-08: ImageGpu gained conv_backward (weight^T . grad lowered to FP16 columns a chunk of output rows at a time, then the reverse of im2col or of one im2row tap as FP32 atomics, with a 1x1 convolution one GEMM), group_norm_backward, layer_norm_backward, gelu_gate_backward, silu_backward and upsample_backward; act_lhs_transposed takes out_narrow, because an FP16-accumulating context writes only FP16. Test: cuda_image::tests::backward_kernels_match_the_host_or_skip_without_a_device (every kernel against its host function in unet_lora, within 1% of the largest reference value; the five convolution shapes of the forward test, each in one chunk and one row per chunk). Dropping the group norm backward's projection term fails it at 0.996 relative. b-ii done 2026-10-08: cuda_image::DeviceAdapters::upload puts a UnetLora on the device, and DeviceUnet::forward_adapted adds each adapter after its projection (`[q
G4 3D texture P2 L Deferred in 2Dto3DPlan.md §6. The rasterizer route is the open decision.
G5 Eikonal regularization for the shape VAE P3 M src/shape_vae.rs:34 needs dL/dquery_point.
G6 T5 attention mask P3 S src/t5.rs:13 has no mask. Padded prompts in a batch need one.
G7 Encoder-decoder (seq2seq) training P3 L The README says "not planned". T5 inference exists, and cross-attention exists in gpu_cross. Revisit if translation or summarization becomes a goal.
G8 Time-series models for the interoceptive agent done 2026-10-09 P2 M The MarketTrunk is an MLP. A small causal transformer over a window of ticks reuses TransformerLm blocks. Done: interoceptive::CausalTrunk embeds each candle of the window as a token, runs TransformerBlocks (RoPE, SwiGLU, causal attention, one sequence per row through Layout), reads the last candle's state after an RMS norm, joins the position and trend features there and ends in one ReLU layer to hidden. AgentConfig::causal: Option<CausalConfig> picks it (None, the MLP, by default); QNetwork::trunk is now the Trunk enum. Test: interoceptive::agent::tests::a_causal_trunk_backward_matches_finite_differences (every parameter of a two-layer, three-candle trunk under the Q network; skipping the RMS norm's backward fails it at 0.166 against 0.196). Not done: a training comparison against the MLP, which needs the agent trained (owner decision 6).
G9 Textual inversion embeddings — done 2026-10-08 P2 M Stable Diffusion users load learned prompt tokens (<cat-toy>) as often as LoRA files, and the CLIP path has no way to add a token. Read diffusers ({"<token>": [d]}), A1111 (emb_params [n, d]) and SDXL (clip_l/clip_g) safetensors files, add the token to each CLIP tokenizer and its rows to the tower's embedding table, and expand a multi-vector token as diffusers does. Acceptance: an ignored test matches diffusers load_textual_inversion + maybe_convert_prompt on a tiny CLIP tower with the SDXL tokenizer, hidden states and pooled vector within 1e-5. Done: Bpe::add_token, ClipTextEncoder::add_token, clip::read_embedding, ImagePipeline::add_embedding. A legacy config (eos_token_id = 2) pools at the highest id, as Transformers does, which differs from the first end token once an added id passes it. Tests: clip::tests::textual_inversion_matches_diffusers (ignored, RB_TI_REF from scripts/textual_inversion_reference.py; ids exact, worst error under 1e-5 on both towers; it fails when the argmax pool or the expansion is removed) and tokenizer::tests::an_added_token_stands_for_all_its_ids_and_cannot_be_added_twice.

Track H — Interoceptive agent

After 2.1.2:

# Item Priority Effort Notes
H1 Real-data backtest harness — part 1 done P1 M Part 1 (evaluation): interoceptive::backtest runs the greedy policy once over a span of candles with fees and slippage and no shocks, and reports EquityStats (total return, per-step Sharpe, maximum drawdown) and the liquidation flag beside a buy-and-hold baseline; walk_forward makes the train/test windows. Tests: interoceptive::backtest::tests::*. Part 2 (open): train on each window and report across windows. The agent trains on the CPU only (H6 is P3), so this part needs an owner decision on a CPU training run, or H6. Original note: ticks_from_csv exists. Add a walk-forward split (train on one window, evaluate on the next), costs and slippage, and a report (Sharpe, maximum drawdown, liquidation count) against a buy-and-hold baseline. Without this, a training curve says nothing about the agent.
H2 Double DQN and a target network done P1 S The target network already existed; the target was plain max_a Q_target. InteroceptiveAgent::train_step now uses double_bellman_targets: the online network picks the next action, the target network scores it. Test: interoceptive::loss::tests::td_loss_touches_only_the_chosen_action_and_done_cuts_the_bootstrap (a row where online and target disagree gets 1.5, not the max-based 2.5). Cost: one more online forward per update.
H3 Prioritized replay done P2 S ReplayBuffer::prioritized(capacity, obs_dim, alpha, beta) keeps a sum tree; sample_prioritized_into draws stratified and reports importance weights, update_priorities takes InteroceptiveAgent::last_td_errors, and Trainer uses it when the buffer is prioritized. Tests: prioritized_sampling_follows_the_td_errors (frequencies within 0.01 of p/Σp over 100k draws, weights (N P)^-β normalized to 1, new pushes at the max priority), weights_scale_each_row_of_the_td_loss, and the prioritized pass of a_short_run_collects_trains_and_finishes_episodes. Beta annealing is left to the caller (set_beta). Original note: A change to ReplayBuffer.
H4 Actor-critic (PPO) P2 M Continuous position sizing instead of discrete actions.
H5 Ablation: interoception on against off P1 S The spec's whole claim is that steering the trunk from outside beats concatenating the state onto the input. Train both with the same seeds and the same shocks, and report the difference.
H6 CUDA trunk P3 M Only if training time becomes the bottleneck. The networks are small, so the CPU and rayon may win.

Track I — Metal

metal_training covers dense networks only.

# Item Priority Effort
I1 Transformer forward and inference on Metal P3 L
I2 Transformer training on Metal P3 L

Park both unless a Mac user asks. Every CUDA kernel would need a Metal twin, and that is the structural tax at its worst.

Track J — Engineering and quality

# Item Priority Effort Notes
J1 Throughput regression check done P1 S scripts/bench_check.sh builds both benches with --features cuda, takes the best tokens/s of five runs, and fails below 97% of the rows in docs/baseline.md ("Regression check"); --record prints new rows. It warns when another process holds the GPU or the load average is 2 or more. Baseline at cd8b514: train step 50329 tok/s (CUDA, single runs 48353-50001), decode 2434 tok/s (CPU, single runs 2348-2434). Checked: train passed at 50472; a 60000 baseline row failed it; decode failed at 2261 with a parallel rustc at load 5.7, which is the warning's case. Original note: Write a script that runs bench_train_step and bench_decode and compares the result against the numbers in docs/baseline.md. Fail if a result is more than 3% slower. Run it before every release. CI has no GPU, so run it locally.
J2 Parity tests in CI done P1 S ci.yml already ran cargo test, clippy -D warnings and the MSRV check. Added cargo doc --no-deps --features cuda,onnx,images with RUSTDOCFLAGS=-D warnings, and fixed the 12 broken or private intra-doc links it found. Every CUDA test skip now goes through cuda_training::skip_without_cuda; RUSTING_BRAIN_REQUIRE_CUDA=1 turns a skip into a failure. Checked with CUDA_VISIBLE_DEVICES=: param:: tests pass 8/8 without the variable and fail 2 with it. Original note: Check that ci.yml runs cargo test, clippy -D warnings, the MSRV build (1.85) and cargo doc. The CUDA tests skip without a device. Make sure they skip and do not pass silently.
J3 cargo semver-checks in CI done 2026-10-08 P2 S A semver job in ci.yml runs obi1kenobi/cargo-semver-checks-action@v2 against the latest tag (git describe --tags --abbrev=0) with default features. Locally, cargo semver-checks --baseline-rev v3.0.0 --default-features (0.51.0): 229 checks pass, 31 skip. With --features cuda,onnx,images it fails to build here: the tool resolves dependencies without Cargo.lock, and kstring 2.0.5 wants rustc 1.96 against the local 1.95; add those features when CI's stable is past that. Original note: 2.0 broke the API once. Make the next break a choice.
J4 Fuzz the file readers done P2 S src/fuzz.rs feeds 2,000 seeded corruptions (bit flips, overwritten runs, cuts, splices, huge numbers in text headers) of one valid file to each reader: safetensors, GLB, BPE and unigram tokenizer.json, .npy, IDX, load_bin and the EMA file. Runs in cargo test (fuzz::*, 0.2 s). It found: a 58 GB allocation from a corrupt BPE token id (ids now refused at 2^24), an index panic on a truncated v2 .npy header, shape overflows in .npy and IDX, load_bin, load_lora and Ema::load allocating whatever length or shape the file claimed (now checked against the file size, and load_bin checks the configuration's weight count, RoPE table and LoRA rank before building; the rank check came from a later review and has its own test). Seeded mutation, not cargo fuzz: that is a new tool and dependency; ask if coverage-guided runs are wanted. Original note: safetensors, gltf, npy, idx, tokenizer.json and the binary checkpoints all read untrusted bytes. Use cargo fuzz on each header parser. A malformed file must return an error and must not panic or allocate without a limit.
J5 Stale-binary guard library side done P1 S build.rs records git rev-parse --short HEAD and reruns when HEAD or a branch ref moves; rusting_brain::BUILD is "<version> (<hash>)". Test: build_names_the_checkout (BUILD equals the current HEAD after a commit). Open, outside this repo: RustingLLM train and bench print rusting_brain::BUILD next to the parameter counts (src/bin/train.rs:234), and its README names cargo build --release --features cuda as the one build command. Original note: W0 in improvement-plan.md. A trainer that was nine hours stale cost 1.63x. Check whether the guard landed.
J6 Kernel source of truth P3 L featuresPlan.md asks whether a shared kernel description pays for itself. Decide when the fourth new layer type needs two hand-written backward passes. Until then, accept the tax.
J7 Tutorials (TokenFile, LoRA and image slices done) P2 S each One chapter each for LoRA fine-tuning, TokenFile pretraining end to end, the image pipeline, and the interoceptive agent (blocked: the agent trains on the CPU only; owner decision 6). TokenFile slice, 2026-10-07: chapter 16's script now writes train.bin/validation.bin and trains from TokenFile::batch; chapters 14 and 16 and the README load tokenizers with Bpe instead of the tokenizers crate. Run on CUDA: the resumed step-200 checkpoint gives held-out 3.5091, the live value. The run found the save_bin stale-weights bug (fixed in 014106e). LoRA slice, 2026-10-07: chapter 17 (tutorials/17_lora_fine_tuning.md) trains a base, adapts it, saves, reloads and merges the adapter. On CUDA: domain loss 4.86 to 0.15, reloaded 0.1528, merged 0.1528, general loss 0.11 to 1.73 (the forgetting section). The run found that LoRA failed on CUDA under mixed precision (fixed in cc85271). Image slice, 2026-10-08: chapter 18 (tutorials/18_text_to_image.md) loads a diffusers directory at Q8, moves it to the GPU, samples with DPM++ 2M and a negative prompt, and repaints the result. On an RTX 3060 with Pony Diffusion V6 XL: 1024x1024 in 13.2 s, repaint at strength 0.5 in 117.2 s (CPU VAE encoder; 6.9 s after C12). The run found the SDXL padding bug (C11, fixed in fdf9f0d).
J8 ponytail: debt review done for 3.0 P2 S 44 comments reviewed on 2026-10-07. Already tracked: CLIP image tower (G1), UNet conditioning (G2), YaRN, Conv2d backward (TrainableConv2d exists), TokenStream order (E4), CSV newlines (E5), shard mmap (E6). New: E8, the quadratic sentencepiece merge, measured before filing. The rest name limits with a working way around them (evaluate in training-size batches, skip_batches reads the prefix, materialized LoRA weight, FP16 decoder range, naive transpose kernel) or are inference-only scope notes, and stay as comments. Original note: The tree has about 40 ponytail: comments. Each one names a known limit. Review them once per release, and move the ones that now matter into this file.
J9 Flaky gpu_transformer::tests::checkpointed_blocks_train_like_kept_ones_or_skip_without_device P2 S Seen 2026-10-09: in about one full cargo test --features cuda --lib run in three, the mixed-precision case differs at step 2 by 1.7e-4 and 3.9e-5 (bound 1e-5); alone it passes. Find which backward kernel is order-dependent (FP32 atomics feeding BF16 weights are the first suspect) and either make it deterministic or bound the BF16 case by its noise. 2026-10-09 retry: not reproduced in 30 runs alone, 15 runs of gpu_transformer:: and 12 full CUDA suite runs (load average 3 to 21). The tiny model's heads are 4 wide, so J11's flash switch does not explain it. Left open until it shows again.
J10 Flaky mesh::simplify_tests::an_error_bound_removes_the_flat_interior_and_keeps_the_shape done 2026-10-09 P2 S The flat-grid assertion (faces == 2) failed in 3 of 200 runs of the test alone. Cause: HashMap order decided which of the zero-cost edges came out of the heap first, and an edge refused for folding was dropped for good, so some orders got stuck above two faces. collapse_edges now keeps refused edges and weighs them again once the queue has nothing cheap left, for as long as a round still collapses something. 0 failures in 400 runs; mesh:: 29 pass.
J11 Flaky fused-attention tests in the full CUDA suite done 2026-10-09 P2 S Found while retrying J9: transformer::tests::a_device_prompt_at_head_width_64_matches_the_host_or_skips_without_device and pipeline::tests::a_transformer_pipeline_on_the_device_matches_the_host_or_skips_without_device failed in 5 of 6 full runs with "not a BF16 key-value cache without flash attention" and "needs the fused attention kernel". Cause: the two gpu_transformer flash parity tests flipped the process-wide cuda_flash::DISABLED. It is now cuda_flash::OFF_ON_THIS_THREAD, a test-only thread-local; the environment variable stays process-wide. 0 failures in 6 full runs after.

Track R — Research 2026-10-10: what the consumers will ask for next

Found by reading the consumers' plans for work that will need this crate, not from their missingFeatures.md files. Order: P1 first, smallest first.

# Item Priority Effort Notes
R1 Warmup-stable-decay schedule done 2026-10-10 P1 S Schedule::warmup_stable_decay. Test warmup_stable_decay_holds_the_peak_then_falls_linearly_to_the_floor: steps 9-80 are exactly the peak, step 90 is 0.5, step 100 is the floor, and a 200-step schedule gives the same rates through step 80. Original note: RustingLLM PLAN.md:210 wants WSD: hold the peak, then decay over the last fraction of the run. Schedule only has warmup plus cosine. Add Schedule::warmup_stable_decay(peak, warmup, steps, decay_fraction), linear decay to the floor, and keep floor. Acceptance: a unit test that the rate is flat at the peak through the stable phase, falls linearly to the floor over the last fraction, and that rates before the decay start do not change when steps grows (the property that lets a run be extended).
R2 Weight decay on matrices only done 2026-10-10 P1 S Param::step drops the decay for one-row parameters before the CPU or device update. Tests weight_decay_shrinks_matrices_and_leaves_one_row_parameters_alone (Adam and Lion: a 2x2 matrix becomes 0.95x with lr 0.1, wd 0.5; a 1x3 gain is unchanged) and weight_decay_on_the_device_matches_the_host_or_skips_without_device (within 1e-6). Full CUDA suite 656 passed. Original note: RustingLLM PLAN.md:214 asks for decay 0.1 "on matrices only" and to check what Adam decays. Network already skips biases, but Param::step (CPU, param.rs:632) and the device Adam and Lion kernels (gpu_transformer.rs:1528) decay every parameter, norm gains included. Skip decay for one-row parameters (norm gains, biases), the nanoGPT rule. Acceptance: with a zero gradient, one Adam step leaves a norm gain unchanged and shrinks a matrix by lr * wd, on CPU and CUDA.
R3 Digit splitting in the BPE trainer done 2026-10-10 P1 S Bpe::train_json_split_digits trains on Qwen2's split (split_llama3(text, 1), already parity-tested against the Qwen2 rule) and writes it as a Split plus ByteLevel sequence, the shape of Qwen2's own file. Test a_digit_split_tokenizer_never_learns_a_token_spanning_two_digits: the GPT-2 trainer learns multi-digit tokens on the same corpus and this one learns none, " 2024" is 5 tokens, the file loads with the split and round-trips text. Not checked against Hugging Face tokenizers itself: it is not installed here. Original note: RustingLLM PLAN.md:184 retrains a 32k BPE with digits split one by one. Bpe::train_json uses the GPT-2 pre-tokenizer only. Add the Hugging Face Digits { individual_digits: true } pre-tokenizer as a training option, written into the JSON so Bpe::from_json encodes the same way. Acceptance: a trained tokenizer never merges two digits, and its JSON round-trips to the same ids.
R4 Dataset::shuffle without cloning rows done 2026-10-10 P2 S Rows are moved with mem::take. Measured on 200,000 rows x 64 features, 20 seeded shuffles, release build, load average 13-17 (uptime 1 day 10:53): 30.9-35.9 ms per shuffle before, 4.9-7.3 ms after. Test a_seeded_shuffle_permutes_the_rows_in_pairs_and_repeats (there was no shuffle test). fit_with still clones the dataset once per fit, not per epoch. Original note: RustingTrade improvementPlan.md P08. shuffle clones every row (dataset.rs:451); DatasetStream::reorder already moves them with mem::take. Do the same. Acceptance: same order as before for the same seed (existing tests), and a measured shuffle time before and after on a large dataset.

Checked and not filed:

  • RustingMedia PLAN.md:104, logit-normal timesteps: flow_transformer::sample_timesteps is already logit-normal with a shift.
  • RustingTrade P09, device-side validation: CudaTrainingSession::prepare_validation already keeps the validation set on the device. Checkpoint cadence lives in RustingTrade's own trainer.
  • RustingCS2 PLAN.md:263, audio input: no capture exists yet. Add STFT and log-mel features when it does.
  • RustingLLM DPO (PLAN.md:255) is owner decision 2.

Suggested order

consumers        K1-K13 (Consumer requests), before everything below
2.1  (now)        2.1.1-2.1.6, H5, J1, J5
2.2  (fine-tune)  F1, F2, F5, A3, B5, B4, B6 (RoPE scaling), E1
2.3  (speed)      A1, A6, A7, A10, C2, C5, C1
2.4  (scale)      B1, B2, A8, A11, E2, E7
2.5  (images)     A4, A5, C6, G1, G3, G10, D5
later             everything at P2/P3, re-ranked by what users ask for

Why this order:

  • 2.1 ships what already exists and fixes correctness problems.
  • 2.2 is the largest gain for a one-card user: take a published model, fine-tune it with LoRA, and export it to the Hub. This uses code that already exists (TextEncoder, LoRA, SafeTensors).
  • 2.3 goes after the measured remainder of the profile: the flash kernels are 19% of a step and the LM head is 13%. It also makes inference worth using.
  • 2.4 is the multi-GPU work. It is large, and it pays off only for users with two cards.
  • 2.5 gives the image path the same treatment the text path got.

Decisions that need an owner

  1. Encoder-decoder (G7): keep "not planned" in the README, or reopen it?
  2. DPO and RL fine-tuning (B8): the interoceptive module changes the argument in featuresPlan.md. Decide whether alignment belongs in the crate.
  3. Vocabulary size (W8): a 15% gain, but it is a model decision and not a kernel decision.
  4. N-D tensor type: a 3.0 breaking change, or never? It is the precondition for J6 and for a cleaner convolution path.
  5. Parquet (E2): add the arrow/parquet dependency tree behind a feature flag, or write a minimal reader?
  6. CPU training for the interoceptive agent: H1 part 2, H5 and the J7 agent chapter all need the agent trained, and it trains on the CPU only. Allow a CPU run for these, or promote H6 (CUDA trunk, now P3) so they can run on the device? Asked 2026-10-08.
  7. cuDNN for image convolutions (C10 part 2): may the image path depend on cuDNN, through cudarc's cudnn feature, for its 3x3 convolutions? A hand implicit GEMM would only about break even. Asked 2026-10-08.
  8. YOLO Gate 2 at the full schedule (K13g): the 100-epoch fine-tune projects to 1.99 h at 15.7 samples/s, inside the 2-hour pre-approval with no margin. Approve a run that may go a few minutes past 2 h, or allow cuDNN (decision 7) for the convolutions first. Asked 2026-10-10.

Not doing

These are the same as featuresPlan.md "Not on this list", and they still hold:

  • Serving, batching servers and HTTP APIs. ONNX, GGUF and safetensors export are the seam.
  • A Python binding.
  • An autograd engine. This is a rewrite and not a feature (see J6).
  • Tensor parallelism, pipeline parallelism and multi-node, until B1 works on two cards.