Jasper-KLearning to optimize GPU kernels.

Overview

GPU kernel optimization is essential to turning AI hardware capability into real training and inference performance. Faster matrix multiplications, normalization, and other core operations can increase throughput, shorten execution time, and reduce the cost of running large models. Achieving these gains is a demanding engineering problem: each implementation must coordinate parallelism, memory access, data reuse, and scheduling for a specific workload and architecture while preserving numerical correctness. An optimization that works for one input shape or device can regress on another.

Current AI coding models still fall short of making this process reliably automatic. Generating correct code is only the first step; finding a faster implementation requires repeated experiments on real hardware, careful interpretation of performance feedback, and informed revisions. The key challenge is to learn optimization strategies that transfer to new tasks. This makes kernel optimization a compelling testbed for reinforcement learning (RL), where correctness and measured execution performance provide direct, verifiable training signals.

We introduce Jasper-K, a family of large language models (LLMs) post-trained with RL for automated GPU kernel optimization. Our framework trains models through multi-turn optimization trajectories, combining exploration across candidate implementations with iterative refinement guided by compilation, correctness checks, execution-time measurements, and profiling feedback. This feedback loop teaches models to pursue both correctness and speed, turning execution experience into improved optimization capability. The results below show how this approach translates into gains across base models, competitive performance against frontier models, and higher end-to-end inference throughput.

Our main contributions are as follows:

1. RL post-training improves kernel optimization across three base models. We develop an RL post-training recipe that improves kernel optimization performance across three base models on both KernelBench (Ouyang et al., 2025) and the LeetCuda evaluation suite, under a unified evaluation protocol and matched inference-time optimization budgets. On KernelBench, Jasper-K-27B, Jasper-K-32B, and Jasper-K-35B (MoE) achieve relative improvements of 51.64%, 2.85%, and 2.83%, respectively, in METR Best GM Speedup (Lin, 2025) over their corresponding base models, using a budget of 16 four-turn trajectories (16 × 4) per task. On LeetCuda, the same three models improve METR Best GM Speedup @16 × 4 by 19.39%, 5.74%, and 17.26%, respectively, over their corresponding base models. We further compare Jasper-K-32B against Kevin (Baronio et al., 2025), which shares the QwQ-32B base model, to assess the relative effectiveness of their respective post-training recipes.

2. RL-fine-tuned Jasper-K-27B achieves competitive performance against frontier models. On LeetCuda, Jasper-K-27B, obtained by RL post-training Qwen3.8-27B with the base model's reasoning effort set to medium, achieves 2.580× METR Best GM Speedup @16 × 4, exceeding GPT-6.1 Sol at medium reasoning effort by 9.23%. It also delivers performance on par with GPT-6.1 Sol, DeepSeek V4.1 Flash-1017, and Claude Opus 5.5 at high reasoning effort. These results highlight how task-specific RL post-training can make a 27B-parameter model competitive with the evaluated frontier models on specialized kernel optimization workloads under the reported settings.

3. Jasper-K delivers performance gains on model-derived kernels and end-to-end inference. We further evaluate Jasper-K-27B on operator workloads extracted from two target models, Qwen3.8-27B and DeepSeek V4.1 Flash-1017, using Qwen3.8-27B as the baseline kernel generator. In the Qwen3.8-27B experiment, the same model plays two distinct roles: its operators form the target workload, and it generates the baseline kernel implementations. Compared with kernels generated by this baseline, Jasper-K-27B achieves 14.49% and 55.6% higher geometric mean best speedups on the two workloads, respectively. When integrated into Qwen3.8-27B inference, Jasper-K-generated kernels improve end-to-end throughput by 3.98% over the configuration using base-model-generated kernels and by 27.94% over the original, unoptimized configuration.

The inference harnesses for all three Jasper-K models are open source on GitHub, with model weights publicly available on Hugging Face.

Jasper-K marks an early milestone toward our vision of recursive self-improvement through full-stack RL. We aim to build a feedback loop in which models optimize the AI infrastructure that supports their own training and inference, and the resulting efficiency gains fuel further training and model improvement. In the near future, we will extend this RL framework to larger and more capable base models. We expect this next stage to deliver stronger kernel optimization performance and greater gains on real workloads, advancing this cycle of improvement.

Our Results

RL gains over corresponding base models

We first evaluate whether RL post-training improves kernel optimization within each model family. We pair Jasper-K-27B, Jasper-K-32B, and Jasper-K-35B (MoE) with their respective base models, Qwen3.8-27B, QwQ-32B, and Qwen3.6-35B-A3B (MoE), on both KernelBench and LeetCuda. Jasper-K improves over its corresponding base models on both benchmarks. This consistent pattern across model families and task sets provides evidence for the effectiveness of our RL post-training recipe.

We have trained and evaluated our kernel optimization models on several of the most advanced AI chips developed in China. This blog presents results from one of these hardware platforms. All compared models are evaluated on this platform, using the same execution environment, task definitions, evaluation harness, correctness criteria, timing protocol, and feedback interface, with a fixed budget of 16 rollouts and up to four turns per rollout. For the training setup aligned with Kevin, we train on 180 KernelBench Level 1 and 2 tasks; the remaining 20 tasks are reserved for training monitoring and model selection. KernelBench evaluation uses the original, unmodified benchmark, which includes training tasks and is therefore not an exclusively held-out generalization test. LeetCuda extends the evaluation to our in-house operator tasks, originating from the reinforcement learning team at the University of Alberta, Canada. See the complete dataset split.

On KernelBench, Jasper-K-27B increases METR Best GM Speedup from 2.289× to 3.471× (+51.64%), Jasper-K-32B from 1.265× to 1.301× (+2.85%), and Jasper-K-35B (MoE) from 1.308× to 1.345× (+2.83%), relative to their respective base models. On LeetCuda, METR Best GM Speedup increases from 2.161× to 2.580× (+19.39%) for Jasper-K-27B, from 1.323× to 1.399× (+5.74%) for Jasper-K-32B, and from 1.373× to 1.610× (+17.26%) for Jasper-K-35B (MoE). These improvements on LeetCuda provide additional evidence that the benefits of RL post-training extend beyond KernelBench.

Each chart reports Best GM Speedup for 16 trajectories × up to 4 turns per task. Solid bars show RAW; the translucent extension reaches METR after applying a 1× reference-fallback floor per task. Relative gains in the text use METR. Read the score definitions and worked example.

KernelBench L1–L3 250 tasks · Best GM Speedup @16 × 4

Jasper-KKevinOther baselines

RAW · solidMETR · translucent extension

Swipe to view all models左右滑动查看全部模型↔

3.5×2.625×1.75×0.875×0×

Qwen3.827B family

RAW 3.466×METR 3.471×

Jasper-K-27B

RAW 2.174×METR 2.289×

Qwen3.827B · medium

QwQ32B family

RAW 0.984×METR 1.301×

Jasper-K-32B

RAW 0.965×METR 1.285×

Kevin32B

RAW 0.931×METR 1.265×

QwQ32B

Qwen3.635B A3B family

RAW 1.223×METR 1.345×

Jasper-K-35B(MoE)

RAW 0.932×METR 1.308×

Qwen3.635B A3B

LeetCuda 50 test tasks · Best GM Speedup @16 × 4

Jasper-KFrontier models & KevinOther baselines

RAW · solidMETR · translucent extension

Swipe to view all models左右滑动查看全部模型↔

3×2.25×1.5×0.75×0×

Qwen3.827B + frontier references

RAW 2.557×METR 2.58×

Jasper-K-27B

RAW 2.562×METR 2.565×

Claude Opus 5.5high

RAW 2.554×METR 2.554×

GPT-6.1 Solhigh

RAW 2.546×METR 2.546×

DeepSeekV4.1 Flash-1017high

RAW 2.362×METR 2.362×

GPT-6.1 Solmedium

RAW 2.16×METR 2.161×

Qwen3.827B · medium

QwQ32B family

RAW 1.155×METR 1.399×

Jasper-K-32B

RAW 0.881×METR 1.361×

Kevin32B

RAW 0.963×METR 1.323×

QwQ32B

Qwen3.635B A3B family

RAW 1.565×METR 1.61×

Jasper-K-35B(MoE)

RAW 1.104×METR 1.373×

Qwen3.635B A3B

Comparison with Kevin-32B and cross-benchmark generalization.

We compare Jasper-K-32B with Kevin-32B (Baronio et al., 2025) and their shared QwQ-32B base model on KernelBench and LeetCuda. On KernelBench, Jasper-K-32B achieves RAW and METR Best GM Speedups of 0.984× and 1.301×, respectively, outperforming Kevin-32B (0.965× / 1.285×) and QwQ-32B (0.931× / 1.265×). On LeetCuda, Jasper-K-32B achieves 1.155× RAW and 1.399× METR, compared with 0.881× / 1.361× for Kevin-32B and 0.963× / 1.323× for QwQ-32B.

Notably, Kevin-32B exhibits a different pattern across the two benchmarks. While its RL post-training improves both RAW and METR scores over the base model on KernelBench, its LeetCuda RAW score falls 8.52% below that of QwQ-32B, despite achieving a higher METR score. Since RAW retains slowdowns relative to the reference implementation whereas METR floors per-task Speedups at 1×, this discrepancy indicates that Kevin's gains under the fallback-based metric do not translate into an improvement when slowdowns are preserved.

This observation raises the question of whether RL post-training can induce benchmark-specific specialization at the expense of performance robustness across different workload distributions. In contrast, Jasper-K-32B consistently improves upon its base model under both metrics on both benchmarks, suggesting that the benefits of its post-training recipe extend beyond KernelBench. Nevertheless, because Kevin was trained under different hardware conditions, these results do not establish whether the observed differences arise from benchmark specialization, hardware adaptation, or other aspects of the training process. A controlled comparison would be needed to isolate these factors.

Comparison with frontier models at different reasoning efforts

On LeetCuda, Jasper-K-27B achieves 2.580× METR Best GM Speedup, outperforming GPT-6.1 Sol at medium reasoning effort by 9.23% and performing on par with GPT-6.1 Sol, Claude Opus 5.5, and DeepSeek V4.1 Flash-1017 at high reasoning effort. All comparisons use the same execution environment, evaluation protocol, feedback interface, and 16 × 4 search budget. These results show that task-specific RL post-training can make a 27B-parameter model highly competitive with the evaluated frontier models on specialized kernel optimization workloads.

From Faster Kernels to Faster Model Inference

01 / QWEN3.8 · END-TO-END

Faster Qwen3.8 Inference

• A FFN down
• B FFN gate / up
• C GDN input
• D Attention output
• E Attention QKV
• F Residual + RMSNorm
• G SiLU × gate
• H GDN norm + gate
• I Q/K norm + RoPE / packing

02 / QWEN3.8 · NINE KERNEL FAMILIES

14.49% Geomean advantage

03 / DEEPSEEK V4.1 FLASH · FIVE FAMILIES

55.60% Geomean advantage

Evaluation Protocol

Hardware and reference baseline

We evaluate kernel implementations on the selected AI accelerator platform using PyTorch 2.10.0. Each candidate is timed on a single device; multiple devices are used to evaluate independent candidates in parallel. Speedup is defined as the execution time of the task’s PyTorch eager reference implementation divided by that of the candidate.

$$\mathrm{S}\mathrm{p}\mathrm{e}\mathrm{e}\mathrm{d}\mathrm{u}\mathrm{p}=\frac{T_{\mathrm{P}\mathrm{y}\mathrm{T}\mathrm{o}\mathrm{r}\mathrm{c}\mathrm{h}\mathrm{e}\mathrm{a}\mathrm{g}\mathrm{e}\mathrm{r}}}{T_{\mathrm{c}\mathrm{a}\mathrm{n}\mathrm{d}\mathrm{i}\mathrm{d}\mathrm{a}\mathrm{t}\mathrm{e}}}.\backslash mathrm\{Speedup\}\; =\; \backslash frac\{T\_\{\backslash mathrm\{PyTorch\backslash \; eager\}\}\}\{T\_\{\backslash mathrm\{candidate\}\}\}.$$

Benchmark score aggregation

For each task $pp$, let $B_{p}B\_p$ be the largest measured speedup among its valid, correct candidates across all 16 trajectories and up to four turns per trajectory. If none has a valid correctness and timing result, the reported benchmark aggregation assigns $B_p=1B\_p=1$ and counts the task as unsolved. All $NN$ benchmark tasks remain in the denominator. We report $$\mathrm{R}\mathrm{A}\mathrm{W}=\exp \left(\frac{1}{N}\sum _{p=1}^N\ln B_p\right),\backslash mathrm\{RAW\}=\backslash exp\backslash !\backslash left(\backslash frac\{1\}\{N\}\backslash sum\_\{p=1\}^\{N\}\backslash ln\; B\_p\backslash right),$$ $$\mathrm{M}\mathrm{E}\mathrm{T}\mathrm{R}=\exp (\frac{1}{N}\sum _{p=1}^N\ln \max (1,B_p)).\backslash mathrm\{METR\}=\backslash exp\backslash !\backslash left(\backslash frac\{1\}\{N\}\backslash sum\_\{p=1\}^\{N\}\backslash ln\backslash max(1,B\_p)\backslash right).$$RAW retains measured slowdowns on solved tasks. METR credits improvements above the reference while treating slower or unsolved tasks as reference fallback. A fallback contribution of $1\times 1\backslash times$ denotes no measured gain and does not count as a correct generated solution. Both aggregates use the same candidate eligibility and numerical correctness checks.

The best candidate is selected separately for every task before the geometric mean is taken. Thus, Best GM Speedup is neither an average over all generated candidates nor an end-to-end model throughput measurement. The $16\times 416\backslash times4$ budget counts candidate-generation opportunities; the four turns within a trajectory are sequential revisions that share execution feedback. Let $M$ denote the benchmark-level Best GM Speedup under the chosen RAW or METR convention. The relative gain over a comparison model is $100\times (M_{\text{Jasper-K}}/M_{\mathrm{c}\mathrm{o}\mathrm{m}\mathrm{p}\mathrm{a}\mathrm{r}\mathrm{i}\mathrm{s}\mathrm{o}\mathrm{n}}-1)\%100\backslash times(M\_\{\backslash mathrm\{Jasper-K\}\}/M\_\{\backslash mathrm\{comparison\}\}-1)\backslash \%$, where $M_{\text{Jasper-K}}$ and $M_{\text{comparison}}$ are the corresponding scores for Jasper-K and the comparison model, respectively. Both scores use the same metric convention, task set, and search budget.

Library admissibility

We apply a static screening stage to reject library-only implementations that merely delegate the computation to existing library operators without implementing a custom kernel. For example, a candidate that simply calls torch.matmul followed by torch.relu is rejected. Hybrid implementations that combine custom kernels with library operations are permitted, as are implementations written entirely with custom kernels. This policy requires candidates to contribute custom kernel code while retaining flexibility to reuse library operations where appropriate.

Precision and error tolerances

Our experiments primarily use FP32 inputs and outputs, while permitting lower-precision intermediate computation when the resulting outputs satisfy our correctness criteria. We retain PyTorch’s default TF32 policy in our evaluation environment: TF32 is disabled for matrix multiplication and permitted for cuDNN operations.

We adopt this policy because overly tight tolerances can exclude useful optimizations involving reduced intermediate precision, whereas overly loose tolerances can admit incorrect kernels. Each candidate must pass five randomized correctness trials. Under our numerical correctness checks, at least 99% of the elements in each floating-point output must satisfy the applicable tolerances, and global relative L1 and L2 errors must each be no greater than 0.01 for FP32 outputs. We use workload-specific calibrated tolerances when available, with absolute and relative tolerances of 0.01 as the fallback for unmatched workloads. Together, these checks allow numerical variation from alternative implementations while constraining both elementwise deviations and aggregate error.

Isolation of environment settings

We observed that candidate-generated code could modify process-wide numerical settings without restoring them, including changes to torch.set_float32_matmul_precision and torch.backends.cuda.matmul.allow_tf32. When evaluation workers were reused, these changes persisted and affected subsequent evaluations, substantially altering numerical outputs and execution times. This contamination distorted both training rewards and inference-time benchmark results, making numerical-state isolation crucial for reliable evaluation. We therefore reset backend precision settings to a prescribed configuration before each evaluation to prevent settings inherited from previous candidates from influencing subsequent results.

Datasets

We use KernelBenchFixed as our sole source of RL training data. Following the same 180/20 training–validation split structure as Kevin, we use 180 Level 1 and Level 2 tasks for policy training and reserve 20 tasks for validation and checkpoint selection. Final testing uses two separate 50-task sets: KernelBench Level 3 and the LeetCuda benchmark. The original KernelBench dataset is available on Hugging Face.

Training data and task corrections

The training and validation pools come from the L1 and L2 portions of KernelBenchFixed. Across this 200-task pool, KernelBenchFixed revises 5 L1 tasks and 16 L2 tasks to address constant or highly concentrated outputs and selected algebraic shortcuts that bypass much of the intended computation. For example, it changes the reduction axis in a task where subtracting the mean over a singleton dimension always produces zero, and replaces uniform inputs with Gaussian inputs in a max-reduction task whose outputs otherwise concentrate near one. Selected matrix-multiplication tasks include a nonlinearity before a reduction to prevent the computation from collapsing into a matrix–vector product. These revisions encourage learning kernel implementation strategies rather than exploiting degenerate workloads.

Validation and checkpoint selection

The remaining 10 L1 and 10 L2 tasks form the heldout validation set. They provide feedback on training progress and help us choose which checkpoint to evaluate on the test sets. Validation tasks do not contribute policy gradients. Checkpoint selection uses validation performance; the two test sets are reserved for evaluating the selected checkpoint.

Test sets and generalization

Our KernelBench test set consists of all 50 L3 tasks, using their original task definitions and reference implementations. L3 tasks are excluded from RL training and validation. This evaluates transfer from the L1/L2 training tasks to the L3 workloads.

LeetCuda is our in-house benchmark of 50 kernel optimization tasks, originating from the reinforcement learning team at the University of Alberta, Canada and is used only for testing. Its tasks are included in neither the RL training set nor the validation set used for checkpoint selection. By extending evaluation beyond the public KernelBench tasks to independently authored workloads, LeetCuda provides complementary evidence of generalization beyond the tasks used in post-training. All compared models are evaluated on the same LeetCuda tasks under the same execution, correctness, timing, feedback, and search-budget settings.

Jasper-K: Multi-turn RL to Explore

We believe the key to unlocking the full potential of RL fine-tuning is to sustain exploration throughout training. For kernel optimization, this means encouraging the model to keep searching for faster implementations even after it has found one that works. To this end, we design Jasper-K’s post-training as an integrated recipe for sustained exploration. A multi-turn RL framework enables feedback-driven search, while improvement-based rewards and tail-aware exploration encourage continued progress and give rare breakthroughs greater influence on training. We complement these components with optional profiling-augmented feedback to guide optimization decisions, sampling-distribution replay to better align policy updates with the rollout sampling distribution, and correctness-gated length control together with explicit thinking budgets to manage response length and preserve capacity for kernel generation.

Multi-turn RL framework

Inspired by Kevin (Baronio et al., 2025), we post-train Jasper-K using multi-rollout, multi-turn GRPO. For each kernel task, the sampling policy generates $GG$ trajectories, each consisting of $TT$ turns of code generation and execution. At each turn, the model proposes a kernel implementation, which is compiled, checked for numerical correctness, and benchmarked for execution time if valid. The resulting feedback, optionally including profiling information, guides subsequent revisions. Across rollouts, the model can explore different optimization strategies. Within a rollout, it can recover from compilation failures, identify performance bottlenecks, and adapt its strategy as it learns more about the kernel. This expands exploration beyond individual kernel implementations to the strategies the model uses to find better ones. Policy parameters remain fixed during trajectory collection, so execution feedback guides subsequent revisions through the model’s context. For the policy update, we associate each turn $tt$ of rollout $ii$ with a reward $r_{i,t}r\_\{i,t\}$ and compute its return as $$R_{i,t}=\sum _{k=t}^T{\gamma }^{k-t}{r}_{i,k},R\_\{i,t\}=\backslash sum\_\{k=t\}^\{T\}\backslash gamma^\{k-t\}r\_\{i,k\},$$ where $\gamma \backslash gamma$ is the discount factor. This allows a revision to receive credit for gains it enables in later turns.

Improvement-based rewards

For kernel optimization, what matters at the end of a trajectory is the best correct implementation discovered along the way. Let $s_{i,t}s\_\{i,t\}$ denote the nonnegative score of the kernel proposed at turn $tt$ of rollout $ii$, with speed-based credit awarded only to implementations that pass correctness checks. Writing $b_{i,t}={\max }_{1\le k\le t}{s}_{i,k}b\_\{i,t\}=\backslash max\_\{1\backslash leq\; k\backslash leq\; t\}s\_\{i,k\}$ for the best score achieved so far and setting $b_{i,0}=0b\_\{i,0\}=0$, we define the best-so-far incremental reward as $$r_{i,t}=b_{i,t}-b_{i,t-1}.r\_\{i,t\}=b\_\{i,t\}-b\_\{i,t-1\}.$$With $\gamma =1\backslash gamma=1$, these increments telescope to $R_{i,t}=b_{i,T}-b_{i,t-1}R\_\{i,t\}=b\_\{i,T\}-b\_\{i,t-1\}$. In particular, the full trajectory return $R_{i,1}=b_{i,T}R\_\{i,1\}=b\_\{i,T\}$ is exactly the best score discovered across all turns.

Once a strong kernel has been found, reproducing it earns no additional reward, while an unsuccessful revision leaves the previous best intact. This gives the model an incentive to explore alternatives with greater upside. For example, given a previous best score of $11$, a guaranteed score of $1.21.2$ earns an incremental reward of $0.20.2$. An alternative that scores either $22$ or $00$ with equal probability has a lower expected standalone score of $11$, yet earns an expected incremental reward of $0.50.5$.

Tail-aware exploration

Improvement-based rewards credit each new best within a trajectory. Inspired by Tail-Likelihood Reinforcement Learning (TailRL) (Ramasubramanian et al., 2026), we apply tail-aware advantage weighting at the turn level to give greater weight to unusually large gains across rollouts. For a fixed kernel task and turn $tt$, we compare the turn-specific returns $R_{i,t}R\_\{i,t\}$ across retained rollouts. Since $R_{i,t}R\_\{i,t\}$ includes improvements achieved in later turns, a revision can receive a high advantage even if it yields no immediate speedup. For a fixed kernel task and turn $tt$, let $nn$ be the number of retained rollouts. Index them in ascending order of return, so that $R_{(1),t}\le \cdots \le R_{(n),t}R\_\{(1),t\}\; \backslash leq\; \backslash cdots\; \backslash leq\; R\_\{(n),t\}$. For each rank $j=2,\dots ,nj=2,\backslash ldots,n$, the gap between adjacent TailRL advantages is $$A_{(j),t}^{\mathrm{t}\mathrm{a}\mathrm{i}\mathrm{l}}-A_{(j-1),t}^{\mathrm{t}\mathrm{a}\mathrm{i}\mathrm{l}}=\frac{1}{n-j+1}(R_{(j),t}-R_{(j-1),t}).A^\{\backslash mathrm\{tail\}\}\_\{(j),t\}\; -\; A^\{\backslash mathrm\{tail\}\}\_\{(j-1),t\}\; =\; \backslash frac\{1\}\{n-j+1\}\backslash bigl(R\_\{(j),t\}\; -\; R\_\{(j-1),t\}\backslash bigr).$$For example, in a group of four rollouts, a return gap of $\mathrm{\Delta }\backslash Delta$ between the top two translates into an advantage gap of $\mathrm{\Delta }\backslash Delta$, while the same return gap between the middle two translates into an advantage gap of $\frac{\mathrm{\Delta }}{2}\backslash frac\{\backslash Delta\}\{2\}$. This gives rare breakthroughs greater influence on training, consistent with Best@K evaluation, where a single exceptional kernel can determine the result.

We use $A_{i,t}^{\mathrm{t}\mathrm{a}\mathrm{i}\mathrm{l}}A^\{\backslash mathrm\{tail\}\}\_\{i,t\}$ in the clipped GRPO objective for the response tokens generated at turn $tt$ of rollout $ii$. Task prompts and execution feedback are excluded from the loss. After each update, the updated weights are synchronized to the sampling policy for the next batch of trajectories.

Profiling-augmented feedback

We believe that tool-generated feedback can support more effective exploration by providing additional evidence about the performance of generated code. Beyond correctness and overall speedup, we optionally augment execution feedback with compact profiling summaries, including the number of GPU kernel launches and the fraction of GPU kernel time spent in the dominant kernels. These summaries are incorporated into the context of subsequent turns, helping the model identify bottlenecks and decide where to focus its optimization efforts. For example, a recorded training trajectory received the following profiling feedback, shown here with abbreviated kernel names:

Profile per forward pass:
12 GPU kernel launches
GPU busy: 100% of wall time
GPU kernel time breakdown:
  GEMM (aten::addmm): 97%, 1 launch
  fused_scale_clamp: 1%, 1 launch
  Other kernels: 2%
For reference only.

In a subsequent revision summary, the model explicitly referred to the 97% GEMM share and reasoned that further elementwise fusion alone would offer limited gains. This example illustrates how profiling can inform the model’s choice of optimization strategy beyond what an aggregate speedup reveals.

The benefit of this information, however, depends on the base model. For Qwen3.8-27B, we observe better performance when profiling feedback is included during training than when it is omitted. For less capable models, additional profiling information can instead be counterproductive. We hypothesize that this difference reflects the model’s ability to interpret performance evidence and translate it into useful code revisions: stronger models can exploit richer feedback, whereas weaker models may struggle to distinguish actionable bottlenecks from incidental details. This motivates matching the richness of tool feedback to the capabilities of the base model.

Sampling-distribution replay

Sustained exploration also requires stable policy updates under the distribution used to collect rollouts. We adopt sampling-distribution replay to account for the temperature scaling and top-$kk$/top-$pp$ truncation applied during generation. The sampling policy records log-probabilities under this processed distribution. During training, we approximately reconstruct the retained token set using the same sampling parameters and normalize the policy’s token probabilities within that set when computing the clipped GRPO objective. Gradients with respect to excluded logits are zero. Because the trainer reconstructs the set from its own logits, its membership can differ from the sampler’s; we explicitly include the sampled token whenever it falls outside the reconstructed set. This confines the direct probability redistribution induced by negative advantages to the retained candidates, limiting its spread into low-probability tokens excluded by the sampling procedure. The mechanism is intended to improve training stability while preserving exploration among plausible continuations.

Correctness-gated length control

To discourage unnecessarily long responses while preserving useful exploration, we add a relative length reward conditioned on correctness (Kimi Team et al., 2025). For each task and turn $tt$, let $L_{i,t}L\_\{i,t\}$ be the number of generated response tokens in rollout $ii$, including reasoning, code, and summary tokens. Among eligible rollouts, we compute $${\lambda }_{i,t}=\frac{1}{2}-\frac{L_{i,t}-L_{\min ,t}}{L_{\max ,t}-L_{\min ,t}},\backslash lambda\_\{i,t\}\; =\; \backslash frac\{1\}\{2\}\; -\; \backslash frac\{L\_\{i,t\}-L\_\{\backslash min,t\}\}\; \{L\_\{\backslash max,t\}-L\_\{\backslash min,t\}\},$$ setting ${\lambda }_{i,t}=0\backslash lambda\_\{i,t\}=0$ when all responses have the same length. The length reward is This rewards relatively short correct responses and penalizes relatively long responses, while preventing short incorrect answers from earning a length bonus. We set $\beta =0\backslash beta=0$ for the first 10 training batches and $\beta =0.1\backslash beta=0.1$ thereafter, bounding the per-turn length contribution to $[-0.05,0.05][-0.05,0.05]$. The resulting training reward is ${\tilde{r}}_{i,t}=r_{i,t}+r_{i,t}^{\mathrm{l}\mathrm{e}\mathrm{n}}\backslash tilde\; r\_\{i,t\}=r\_\{i,t\}+r^\{\backslash mathrm\{len\}\}\_\{i,t\}$, where $r_{i,t}r\_\{i,t\}$ is the improvement-based reward defined above. We use ${\tilde{r}}_{i,t}\backslash tilde\; r\_\{i,t\}$ to compute discounted returns and advantages. Because lengths are compared within the same task and turn, this mechanism provides relative pressure toward concise solutions without prescribing a fixed target response length.

Explicit thinking budgets

As RL training progresses, the model’s reasoning traces can become increasingly long. Although the length reward encourages concise responses, it provides a soft incentive rather than an absolute bound on reasoning length. Consequently, an increasing fraction of rollouts may exhaust the response budget before producing a kernel, reducing the amount of usable training data and execution feedback. We therefore impose a hard thinking-token budget within the total response budget. If the model has not emitted the end-of-thinking marker when this limit is reached, we append a short transition followed by </think> and resume generation with the remaining token budget. The inserted control tokens are excluded from the policy loss, while model-generated reasoning and answer tokens remain eligible for training. In our configuration, the thinking budget is 32,768 tokens within a total response budget of 40,960 tokens, leaving approximately 8,192 tokens, minus the inserted transition, for code and summary generation. This mechanism complements the length reward by reserving generation capacity for an executable implementation, reducing truncation before kernel generation. In our experiments, introducing this constraint allowed training to continue for several times as many batches before excessive truncation became a limiting factor.

What RL Teaches a Kernel Optimizer

What does RL post-training change in an actual optimization attempt? On this held-out task, the Qwen3.8-27B base model before RL reaches a best speedup of only 1.017× across 16 rollouts and four turns per rollout. After RL post-training, Jasper-K-27B discovers a correct kernel at 1.855× under the same search budget—an 82.4% increase in best speedup. The four-turn trajectory below shows how it gets there.

The workload is a 1024 × 8192 input multiplied by an 8192 × 8192 weight matrix, followed by BatchNorm, a scalar bias, division, and Swish. Speedups are measured against the task’s reference implementation. The displayed code extracts the decisive lines from the generated implementations; comments are shortened or added for clarity.

After RL: one trajectory, four revisions

1. Turn 1 · fusion1.012×
2. Turn 2 · direct FP16Precision failure
3. Turn 3 · three terms1.238×
4. Turn 4 · two terms1.855×

# FP32 matrix multiplication remains unchanged.
y = F.linear(x, self.matmul.weight, self.matmul.bias)
# scale and shift fold BN, scalar bias, and division.
out = _fused_ext.fused_affine_swish(y, scale, shift)

float val = input[idx];
float t = val * scale[col] + shift[col];
out[idx] = t / (1.0f + __expf(-t));

# Both the inputs and GEMM output are FP16.
x_half = x.half()
y_half = torch.mm(x_half, self.weight_fp16.t())

// Casting back cannot restore lost precision.
float val = __half2float(input[idx]);
float t = val * scale[col] + shift[col];
output[idx] = t / (1.0f + expf(-t));

# Split transposed weights at initialization.
Wt = self.matmul.weight.t().contiguous()
Wt_hi = Wt.half()
Wt_lo = (Wt - Wt_hi.float()).half()

# Split the input into FP16 high and residual parts.
_ext.split_fp32_2xfp16(x, Xh, Xl)
# C is FP32; False overwrites, True accumulates.
_ext.fp16_gemm_fp32(Xh, self.Wt_hi, C, False)
_ext.fp16_gemm_fp32(Xh, self.Wt_lo, C, True)
_ext.fp16_gemm_fp32(Xl, self.Wt_hi, C, True)

# Keep weight correction; omit input correction.
Xh = x.half()
# C retains FP32 accumulation and output.
_ext.fp16_gemm_fp32(Xh, self.Wt_hi, C, False)
_ext.fp16_gemm_fp32(Xh, self.Wt_lo, C, True)

# Fold the remaining operations into one kernel.
out = _ext.fused_scale_shift_swish(C, scale, shift)

Future Work

Our results suggest that RL post-training can improve a model's ability to explore kernel implementations and learn from execution feedback. They also leave open how broadly these capabilities transfer, how effectively models can direct their own experiments, and how training can continue to expose them to new optimization challenges. We plan to investigate these questions along three directions.

Transfer across hardware and workloads. We plan to extend training and evaluation to NVIDIA GPUs and include SOL-Bench in our evaluation. Beyond broadening benchmark coverage, this will help distinguish optimization strategies that transfer across platforms from adaptations to specific hardware. Comparisons under matched correctness criteria and optimization budgets will also provide a clearer assessment of Jasper-K relative to other kernel optimization approaches.

Scaling Jasper-K to larger models. Scaling has been a fundamental driver of progress in large language models. We plan to extend our RL methods to larger open-weight models, such as DeepSeek V4.1 Flash 1017, GLM-5.3, and Kimi K3.1 Pro, to explore whether RL can further improve kernel optimization and enable additional specialized capabilities without compromising general-purpose capabilities.

Full-stack RL-driven self-improving optimization systems. Beyond the extensions discussed above, we aim to explore full-stack RL for recursive self-improvement (RSI), extending optimization beyond the kernel-generation policy to task generation, agent harness design, and exploration strategies. Our long-term goal is to leverage RL to continuously improve both kernel implementations and the optimization process itself.

Additional Experimental Results

Best performance is one part of the picture. We also report how often a model finds a correct or faster kernel, its average performance across trajectories, and its performance on the held-out KernelBench L3 tasks. The tables cover all models in the corresponding charts and use the same 16 rollouts × up to 4 turns budget.

What each metric tells us