KITE/SST: Small public-model matched-FLOP diagnostic¶
This is a scaled-down KITE/SST mechanism check, not a reproduction of the MoE scale or original benchmarks. Script: scripts/benchmark_kite_equal_flops.py.
Use raw public WikiText-2 with independent train, validation, and test files, UTF-8 bytes, and a 32-token next-token task. Seeds are 42, 43, and 44. A width-32 two-layer Prefiller grows into a 2+2-layer SST in two stages. The control is four-layer dense. Both share initial embeddings, the output head, and first two layer weights; parameter counts are 64,704 and 68,800.
SST trains for 30+30 steps; the control takes 17+28 steps to match forward/backward FLOPs detected by the PyTorch profiler. Both restart AdamW at stage boundaries with learning rate 0.001. Control FLOPs / SST FLOPs = 0.9993. Counting excludes optimizer work, LayerNorm, and memory movement, so hardware work is not strictly equal. Test is not used for method or hyperparameter selection.
| Seed | SST val NLL | Dense val NLL | SST test NLL | Dense test NLL | SST last token ms | Dense last token ms |
|---|---|---|---|---|---|---|
| 42 | 3.7512 | 4.0244 | 3.7661 | 4.0508 | 0.279 | 0.346 |
| 43 | 3.5698 | 3.7940 | 3.6004 | 3.8177 | 0.279 | 0.343 |
| 44 | 3.5245 | 3.7784 | 3.5059 | 3.7568 | 0.281 | 0.356 |
PYTHONPATH=src python scripts/benchmark_kite_equal_flops.py \
--data-root data --seeds 42,43,44 \
--source-steps 30 --continuation-steps 30
On this fixed small task and approximate budget, SST has lower NLL and local last-token latency than that dense control. Local latency is not GPU latency and cannot be extrapolated to 67B scaling or GPU throughput.
A100 CUDA inference microbenchmark¶
scripts/benchmark_kite_gpu_latency.py separately measures randomly initialized width-32 2+2 SST and four-layer dense on seeds 42, 43, and 44 with 20 repetitions after warmup per configuration. Inputs are generated byte tokens, not WikiText-2-trained weights, and there are no language-quality metrics. decode-one times one token after a full prefix. SST prefill-last-only computes only Prefiller outputs relevant to the last token and is separate from full-sequence prefill. Maximum absolute last-token logit error versus full prefill is 2.46e-7.
| Context tokens | SST full prefill | SST last token prefill | SST decode-one | Dense prefill | Dense decode-one |
|---|---|---|---|---|---|
| 32 | 1.421 | 1.417 | 1.435 | 1.374 | 1.447 |
| 128 | 1.423 | 1.422 | 1.438 | 1.378 | 1.454 |
| 256 | 1.441 | 1.431 | 1.455 | 1.406 | 1.457 |
Times are ms, the median of each seed's repeat-measurement median. The sanitized GPU receipt records versions and commands. In this very small, short-context, launch-bound benchmark, no clear CUDA latency advantage was observed. This verifies only that real small-model CUDA prefill and decode run. Paper scale, long context, throughput, and trained weights still require validation; this is not reproduction of the paper's effect.