God Does Not Play Dice: Enforcing Batch Invariance in Agentic Kernel Generation
Organizations: STAR Lab, George Mason University · National University of Singapore
Abstract
Large language model (LLM) serving systems dynamically batch user requests for GPU execution, so the same request may be processed alongside different requests across runs. If the GPU kernels lack batch invariance (BI), this changing batch context can alter the numerical output for an unchanged input, undermining reproducibility. Existing kernel generation agents focus on correctness and efficiency but overlook BI. We find 19 confirmed BI violations among 83 batch-applicable released artifacts. To address this limitation, we propose \tool, which combines contract-guided static dependency analysis with dynamic validation of bitwise consistency across batch contexts. \tool is designed as a plug-in component for general kernel agents, allowing their original optimization capabilities to be reused while adding BI enforcement. We integrate it as successive acceptance gates in the optimization loop, rejecting structurally risky or dynamically non-invariant candidates before performance benchmarking and returning feedback to the agent. Experiments across three agent frameworks under baseline and BI-aware generation on real world LLM serving tasks show that \tool improves the BI pass rate from 67.6% to 91.9% while retaining 91.8% of the baseline geometric mean speedup. These results demonstrate that \tool can be incorporated into kernel optimization by agents as an additional reliability objective alongside correctness and speed.
Figures & tables
| Artifact source / benchmark | Batch- applicable / released | Confirmed BI violations | BI pass (%) |
|---|---|---|---|
| KernelBench.com / KernelBench.com | 79 / 126 | 18 | 77.2 |
| K-Search / FlashInfer-Bench | 4 / 5 | 1 | 75.0 |
| Total | 83 / 131 | 19 | 77.1 |
| "kernel" | ||
| "{" "}" | ||
| "extends" | ||
| "invariant_axis" | ||
| "batch" "symbols" | ||
| "[" "]" ";" |
| Framework | Generation condition | BI pass (%) | BI intervention (%) | GM speedup ( ) | Speedup retained (%) | Correctness pass (%) |
|---|---|---|---|---|---|---|
| K-Search | Baseline | 75.7 | - | 40.94 | 86.4 | 94.9 |
| BI-aware | 97.4 | 23.1 | 35.39 | 97.4 | ||
| Iterative Agent | Baseline | 48.6 | - | 25.84 | 87.7 | 94.9 |
| BI-aware | 89.2 | 43.6 | 22.66 | 94.9 | ||
| Evolve Agent | Baseline | 78.4 | - | 22.85 | 102.0 | 94.9 |
| BI-aware | 89.2 | 25.6 | 23.31 | 94.9 |
| Component | BI violation detection (%) | False rejection (%) | Median evaluation time (ms/candidate) | ||
| Overall | Unique | Overlap | |||
| Static module | 83 | 37 | 46 | 2 | 7.69 |
| Dynamic module | 55 | 9 | 46 | 0 | 2.24 |
| BIEvolve | 92 | – | – | 2 | 10.21 |
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
| Kernel family | Evidence category | Work partitioning | Reduction order | Algorithm selection | Expert routing |
|---|---|---|---|---|---|
| Attention | High | Confirmed | Confirmed | Confirmed | None |
| Linear attention | Medium | Confirmed | None | Confirmed | None |
| Mixture of experts | Low | None | Risk | Risk | Risk |
| Top- | Low | Risk | None | Risk | None |
| Recurrent | Low | None | None | Risk | None |
| Triangle multiplication | Low | None | Risk | Risk | None |
| BI outcome of the baseline final kernel | Tasks, | Late-stage speedup gain (%) | |
|---|---|---|---|
| Baseline | BI-aware | ||
| Confirmed violation | 9 | 6.11 | 26.30 |
| No confirmed violation | 28 | 13.87 | 13.51 |
GPU Kernel Generation
GPU (graphics processing unit) kernel generation automates writing and optimizing device code for faster computation across accelerators.
See all 67 papers