Track 1 (1.167 min on 8xH100, -5.15% same-node): Reduced-width QK and packed FP8 attention

[Glitchfix]

Summary

This PR reduces the Q/K head width from 128 to 96 while keeping V and the residual stream at 128 dimensions. It packs the Q, K, and V projections into a shared FP8 projection and adds Triton kernels for the associated layout, normalization, RoPE, and backward work.

The routing vectors used by attention do not necessarily need the same width as the value vectors carrying content. Normalized Q/K vectors can use a smaller geometry, reducing projection work without narrowing V or the model residual.

The branch has been rebased onto current main. The DC attention layer conflicts are resolved and the candidate has been re-measured against a same-node baseline, which the earlier revision of this PR could not provide. Resolving the conflicts also exposed some redundant memory traffic that the earlier revision still had.

Changes

• Reduce Q/K head width from 128 to 96; keep V at 128.
• Store Q/K separately from the rectangular V/O bank and project packed QKV in one FP8 operation.
• Cache row-major and transposed FP8 weight layouts once per optimizer update, emitting both from a single weight read.
• Fuse QK normalization, RoPE, key offset, FA3 padding, and layout conversion in Triton.
• Extend the FP8 MLP work to the full ReLU-squared forward and backward path across all four MLP GEMMs and both projections.
• Reconstruct relu(pre) from the stored squared activation instead of saving and reloading the pre-activation.
• Quantize the logical QK/V packs without materializing concatenations or full-size abs temporaries, and keep the QK-norm and RoPE lane swaps register-resident.
• Let the DC correction consume reduced-width Q/K views directly while retaining 128-wide V, which removes the padding copies the earlier revision needed.
• Use 45 final-lr extension iterations, for 1315 timed training steps.
• Continue using the stock kernels-community/flash-attn3. This PR does not include a custom FA3 implementation.

Results

Four candidate seeds and a same-node f411b3d baseline, run back to back on one Prime Intellect node with 8x NVIDIA H100 80GB HBM3 (datacrunch, FIN-02). Stock requirements.txt: Python 3.12, PyTorch 2.10.0+cu128 for CUDA 12.8, Triton 3.6.0. The training and validation token streams and the validation calculation are unchanged, requirements.txt is unmodified, and the torch.compile call sites and _inductor.config lines are identical to f411b3d.

Candidate

Seed	Steps	Validation loss	Training time
2	1315	3.2732	70.626 s
4	1315	3.2753	70.019 s
42	1315	3.2760	69.737 s
1337	1315	3.2744	69.993 s
Mean	1315	3.274725	70.094 s (1.1682 min)

Loss sample standard deviation 0.001209. The one-sample, one-sided test against 3.28 gives t=-8.723776 and p=0.001585, which satisfies the Track 1 loss requirement.

Same-node baseline (f411b3d)

Run	Steps	Validation loss	Training time
1	1285	3.2765	73.728 s
2	1285	3.2767	74.065 s
Mean	1285	3.276600	73.897 s (1.2316 min)

Speed

70.094 / 73.897 = 0.94854, a 5.15% reduction in training time on identical hardware, which projects 1.1667 minutes against the 1.23-minute record that f411b3d holds. The baseline measured 1.2316 min on this node against 1.23 published, a 0.13% difference, so the node is not unusually fast or slow.

Per step the candidate is 7.31% faster, 53.30 ms against 57.51 ms. The 45 extension iterations mean 1315 steps against main's 1285, and that 2.33% step-count increase gives back roughly a third of the per-step gain.

The earlier revision of this PR reported 77.812 s over 1425 steps and could only be compared against a published time from a different node, which gave 1.75%. Part of the absolute difference since then is the shorter schedule this branch picked up in the rebase. The 5.15% above is the portion measured on one node with both arms run back to back, and the FP8 backward coverage and reduced memory traffic listed under Changes are what moved it.

Alternate configuration

records/track_1_short/2026-08-03_FP8MLPBackwardPackedQKV/alternate_ext65/ holds a 65-extension configuration that also passes, at mean loss 3.270500 and 71.892 s (ratio 0.97433, projecting 1.1984 min), with its own three-run baseline in alternate_ext65_baseline/. The 45-extension result above is the one being submitted. Since the leaderboard time is a ratio, the shorter schedule gives the better record provided its loss sample is tight enough to reach significance.

Reproduction and validation

Candidate runs were launched as:

TRAIN_SEED=<2|4|42|1337> ./run.sh

NUM_EXTENSION_ITERATIONS=45 and QK_HEAD_DIM=96 are the checked-in defaults, so a plain ./run.sh reproduces the submitted 1315-step configuration.

Recalculate the submitted statistics with:

python records/track_1_short/2026-08-03_FP8MLPBackwardPackedQKV/statistics.py

Additional checks completed:

• Confirmed the validation block is byte-identical to current main.
• Direct forward/backward equivalence checks for reduced-width QK, packed FP8 QKV, and the fused FP8 MLP paths.
• python -m py_compile train_gpt.py triton_kernels.py

Raw logs, per-seed results, environment output, and the statistics script are in records/track_1_short/2026-08-03_FP8MLPBackwardPackedQKV.

[kroggen]

Interesting how the difference was way bigger (~3.68 s) on 8xH200

[Glitchfix]

Interesting how the difference was way bigger (~3.68 s) on 8xH200

actually the gain shown was against in the H100 run was comparison with the leaderboard run
I could only have the node for a short time to run the baseline and the candidate run for a fair comparison

the bandwidth happens to be the bottleneck for packing and transpose operation benefit, I think the gain would be much better in newer GPUs
interesting for sure 😄

[ClassicLarry]

This PR is the next one up to get merged into main. However there are some conflicts with the DC attention layer.
If you retest on main and figure out how you want to handle the conflicts, I can merge faster.

[Glitchfix]

@ClassicLarry let me resolve the conflicts
most of it comes from the unequal qk width that I have implemented
I found there can be certain things from the rebased master that I can reuse
I will retest as soon as I find 8xh100 free node to run it

[Glitchfix]

@ClassicLarry I have resolved the conflicts
the conflict was indeed the unequal qk width as expected
I aslo reused the rebased DC kernel and its FP8 down projection
I extended FP8 to the MLP backward pass as well

I also found some more optimizations
I did suspected since the 8xh200 run gained a lot more than h100 despite the same architecture and no real memory advantage for a run this size

I found that dropping the saved pre activation since relu(pre) is recoverable from the stored squared activation, and the both FP8 weight layouts that I have been improved they now come out of one weight read instead of two this was something I had missed earlier

The retest is substantially faster than the current master's baseline of 1.23min
It stands at 3.274725 mean loss, p=0.001585, 70.094s against 73.897s for master on the same node
so ~1.167min, 5.15% faster than the current baseline 😄

I have updated the description to reflect the same

[Glitchfix]

@kroggen I ran this on 8xH200 again
it is suprising how fast it was able to go for the same 45 ext, the delta is huge while the

Seed	Validation loss	Training time
2	3.2796	66.406 s
4	3.2739	66.383 s
42	3.2739	66.603 s
1337	3.2810	66.586 s
Mean	3.277100	66.495 s (1.1082 min)

while the master baseline ran in 72.028s

I am really curious to optimize these on bigger GPUs, need more GPU credits 😄

[kroggen]

Yeah, I was expecting the opposite: smaller difference on H200 because the master baseline would also be faster. They are the same architecture with just difference in HBM size and speed (4.8 TB/s vs 3.35 TB/s)

Maybe the baseline is more memory-bound on H100 and more compute-bound on H200