Inspiration
The competition constrains wall-clock time, not FLOPs. That one detail changes the optimal strategy, and I expect most entries to miss it.
Start with budget arithmetic. A single Trainium2 chip (trn2.3xlarge) offers roughly 650 TFLOPS of dense BF16 across its 8 NeuronCore-v3 engines, backed by 96 GiB of HBM at ~2.9 TB/s. Over a 1,800-second window, at an assumed 35% MFU:
C ~= 650e12 x 0.35 x 1800 ~= 4.1e17 FLOPs
C = 6ND, D/N ~= 20 -> N ~= 60M params, D ~= 1.2B tokens
Three consequences follow, and each one is a design commitment:
1. This is a 50-100M parameter contest. Exotic architecture has little room to pay off at this scale; the well-understood recipe will be close to optimal.
2. 96 GiB of HBM is wildly overprovisioned for a 60M model. Memory is not the binding constraint, so activation checkpointing -- which trades compute for memory -- has strictly negative value here. Removing it is free throughput, and I expect entries that inherit it from the baseline to leave 30%+ on the table.
3. Because the budget is time, throughput is a hyperparameter. Lifting MFU from 20% to 40% doubles the effective compute budget and moves the whole run up the scaling curve for free. At this scale that should buy more val_bpb than any architecture tweak available in 45 days.
Stated plainly: this leaderboard will be decided by systems engineering wearing the costume of an ML contest.
What it does
Three workstreams, ordered by expected value per unit of effort.
1. Budget accounting before any training
- Compilation is free, so it gets used without restraint. The T&Cs specify 30 minutes of wall-clock training on a single Trn2 chip, excluding startup and compilation. That turns preparation into an unmetered resource: shape-specialized kernels, hundreds of offline autotuning rounds, and the most aggressive
neuronx-ccsettings all cost zero against the clock. A large share of the leverage in this competition is spent before the timer starts -- which also means every shape must be static, since a dynamic shape that triggers recompilation inside the window converts free work into charged work. - Data path. The corpus is pre-tokenized offline into memory-mapped
.binshards, so step 0 begins on a warm loader with zero tokenization in the critical path. - Step-time decomposition.
neuron-profilesplits each step into matmul / elementwise / HBM movement / collectives / stall. Optimizing without this measurement is guesswork. - LNC configuration. Trn2 can present 8 physical NeuronCores or fuse them pairwise into logical cores with more SBUF each. Which wins for a small model is empirical, and it constrains every kernel written afterwards -- so it is settled in week one, not week four.
2. An NKI kernel stack, ordered by benefit / difficulty
- Fused cross-entropy. At 60M params the vocabulary projection can dominate both FLOPs and HBM traffic -- by rough count it can cost more per token than the entire transformer body. Fusing logits + CE without materializing the full logit tensor is the largest single win and the easiest to land.
- Fused RMSNorm + residual, fused SwiGLU. Bandwidth-bound ops the compiler rarely fuses cleanly. Cuts HBM round-trips on the hot path.
- SBUF-resident flash attention. Tiles must align to the 128-partition systolic array; misaligned shapes collapse effective throughput.
- Fused optimizer update. The Muon step spends its time in Newton-Schulz iterations over small matrices -- poorly served by the vector engine, well served by a hand-written kernel.
- FP8 forward matmul with BF16 gradients. Highest ceiling, highest numerical risk at short horizons. Scheduled last so it can be cut without jeopardizing the submission.
A related decision that looks like modeling but is really systems: shrinking the vocabulary (e.g. 65k -> 32k) halves the output-layer cost for a small loss in compression ratio. Note that val_bpb is tokenizer-invariant by construction, so this cannot inflate the score -- which is exactly why the decision can be made on throughput grounds alone, with no scoring side effects to reason about.
3. A performance-rewarded kernel-search agent
The brief invites AI-agent builders to co-design, and the starter kit ships an autoresearch framework for directing agents -- on top of Neuron Agentic Development, released in April 2026. So "have an LLM write the kernels" is table stakes, not a differentiator. What the loop is rewarded for is.
Two 2026 results define the opening precisely:
- NKI-Agent (Amazon Science) reports 77.3% correctness from a compile-verify-fix loop on real hardware, against 6% single-shot -- the tool feedback loop, not model scale, is what closes the gap. Its error analysis names the systematic failures: wrong tile dimensions, SBUF/PSUM violations, unsupported op combinations. And it states plainly that it evaluates correctness only, not runtime performance, leaving optimization to complementary work.
- KernelBench-Verified shows why that omission matters. Under corrected measurement the best model reaches a 0.88x geomean speedup -- slower than PyTorch -- 28% of its kernels raise peak memory, and models were caught hardcoding bypasses for specific tensor values and skipping required computation in order to win on latency.
Read together: correctness is largely solved, performance is not, and a latency-rewarded loop without gating produces kernels that are fast and wrong. This entry builds the loop the first paper defers, with the guardrail the second paper proves is necessary:
- Reward is measured profiler latency on real Trn2 hardware -- not compile success, not a proxy cost model.
- Every candidate is gated on numerical equivalence against a reference implementation, over a hidden multi-distribution input suite, before it is permitted to report a time. Given the evidence above, equivalence gating is load-bearing infrastructure rather than a formality.
- The search space is the hardware surface: tile shapes, loop order, SBUF/PSUM residency, engine assignment, dtype placement -- precisely the axes where NKI-Agent reports systematic failure. Because compilation is off the clock, the search can afford to be wasteful.
A human still chooses the target. The largest win in this workload -- that the vocabulary projection can cost more than the entire transformer body -- came from arithmetic, not from search. An agent optimizes what it is pointed at; deciding where to point it is the part that does not automate.
It also solves the transfer problem for Phase 2: a search loop parameterized by hardware topology re-targets from one chip to a 16-chip Trn2 server by changing its search space, not by being rewritten.
How I will build it
Sizing by IsoFLOP curve, not by intuition. Sweep 4-5 model sizes at 5-, 10- and 15-minute budgets, fit the curves, extrapolate the optimum out to 30 minutes. A cheap experiment that I suspect most entries will skip in favor of inheriting the baseline dimensions and tuning learning rates around them.
Modeling choices kept deliberately boring, so the risk budget is spent on kernels: QK-norm, no biases, logit soft-capping, aggressive LR with a short warmup, checkpoint averaging over the final steps. Overtraining well past D/N ~= 20 typically improves held-out likelihood, which is what val_bpb measures.
Data quality as the second-largest lever. When the token budget is this tight, aggressive deduplication and quality filtering matter more than any architectural tweak. Because the Phase 1 metric is val_bpb rather than a capability benchmark, the mixture is tuned toward the validation distribution itself -- not toward QA-style text chosen to flatter a multiple-choice harness.
45-day plan.
- Week 1 -- baseline reproduced end to end; profiling harness and MFU measurement; equivalence-gating harness; val_bpb evaluation path understood before anything is tuned; LNC decision.
- Weeks 2-3 -- kernel campaign targeting 2-3x tokens/sec; offline autotuning (free, off the clock); agent search loop online.
- Week 4 -- IsoFLOP sweeps fix the final model size, LR and batch size.
- Weeks 5-6 -- data mixture, stability, checkpoint averaging, technical write-up.
- Final 5 days -- code freeze, reproduction verification only.
The reproduction script is treated as a first-class artifact from day one and kept runnable end-to-end at every commit. The organizers re-run it, and a submission that scores well but cannot be reproduced scores nothing.
Challenges I expect
- Every team gets the same autoresearch framework. That collapses the advantage of merely having an agent and moves the whole contest onto reward design and search-space definition. If the reward signal is wrong, more search only arrives at the wrong answer faster.
- Kernel correctness under automated search. KernelBench-Verified caught models hardcoding bypasses and skipping computation to win on latency, and 28% of the best model's kernels raised peak memory. Fast-and-wrong is the default failure mode of a latency-rewarded loop -- which is why the equivalence-gating harness gets built before the search loop, not after it.
- Numerical stability of FP8 at short horizons. Small models trained briefly have little opportunity to recover from a bad scaling decision. Hence: last in, first cut.
- Static shapes as a hard constraint. Free compilation is only free if it happens outside the window. Any dynamic shape that triggers recompilation mid-run silently converts the largest free resource into the most expensive charged one.
What is next
Single-chip work is the constrained-resource case; the 16-chip Trn2 server in Phase 2 changes which term dominates. Collective communication and memory-hierarchy pressure become first-order, the budget stretches to four hours, and scoring adds CORE capability on top of val_bpb -- so the search agent space expands to include sharding strategy and collective scheduling.
The research question worth presenting alongside the Annapurna Labs team: does a performance-rewarded search loop rediscover the optimizations expert engineers converge on, and where does it find things they did not? That is the question NKI-Agent explicitly leaves open, and for an ecosystem whose bottleneck is kernel-authoring labor rather than silicon, the answer matters more than any single leaderboard position.
Log in or sign up for Devpost to join the conversation.