Akahsizrr/kernelstrain2

kernelstrain2 — kernelinfer

Synthetic long-horizon GPU kernel-optimization trajectories, inference-only.

Each row is a complete optimization session (~30–200 steps, simulating 1–5 hours of work): the model receives an inference-kernel task, emits a full CUDA or Triton candidate, and gets simulated compile / correctness / benchmark feedback — iterating, self-correcting, and tuning launch parameters until convergence. No reasoning traces; assistant turns are kernel code only.

Scope

Inference workloads only — no training kernels. Operations cover:

• —decode/prefill attention: flashdecode, flashprefill, swadecode, mladecode, treeattn, chunkedprefill, crossattn, fusedqkv_rope
• —quantized GEMM/GEMV: fp16, int8, int4, fp4, fp8 (incl. blockwise DeepGEMM style), split-K, wgmma/TMA persistent variants
• —KV-cache ops: pagedkvcopy, kvappend, kvint8_quant
• —samplers: topk, nucleus (top-p), min-p, logits_softmax
• —MoE: scatter, combine, grouped_gemm
• —SSM: mambascan, mambastep
• —misc: rmsnorm, rope, swigluact, embeddinggather, conv1dcausal, conv2dnhwc, gemv_fp16

41 ops, CUDA + Triton implementations.

Hardware targets

H100 (sm90a), H200 (sm90a), A100 (sm80), A10 (sm86), RTX A6000 (sm86), RTX 5090 (sm120). Arch-gating is modeled correctly (wgmma/TMA only on sm90a; fp8 on sm89+/sm90/sm120; fp4 on sm_120; per-block smem caps 99/164/227KB).

Format

JSONL, one session per line:

• —messages — training data: system + user (task, then [t+H:MM:SS] RESULT feedback turns) + assistant (full kernel source, no prose, no CoT)
• —trace — compact per-step metadata (tag/kind/params/status/time/kept)
• —top-level: op, impl (cuda|triton), gpu, arch, dtype, shape, signature, desc, baseline_ms, best{step,time_ms,speedup}, attempts, elapsed_s, elapsed_hms, num_steps

batches/ holds 100-session shards; sessions.jsonl is the merged view; tasks_holdout.jsonl = 195 held-out task specs (no solutions).

Stats

• —1,899 sessions, all unique (op, gpu, shape, dtype) combos
• —~180K total steps; avg 94.6 steps/session
• —avg 2.18h simulated (min 1.03h, max 4.93h); avg 13.4x speedup
• —step statuses: 151.8K ok / 22.4K compileerror / 5.3K wrongresult
• —~1.5GB

Provenance

Fully synthetic: benchmark/correctness numbers come from an analytic roofline-style model — no kernels were ever compiled or executed. Intended for training small models on long-horizon iterative optimization behavior. Per-batch QA reports in the generating repo (kernelinfer/qa_reports/); all 19 batches scored 10/10 on mechanical checks and ≥8.5/10 on qualitative review after fixes.