Optimizing CUDA like a Human: Micro-Profiling Tools as Expert Surrogates for LLM-Based GPU Kernel Optimization

# Optimizing CUDA like a Human: Micro-Profiling Tools as Expert Surrogates for LLM-Based GPU Kernel Optimization
Source: [https://arxiv.org/html/2606.26453](https://arxiv.org/html/2606.26453)
Jiading Gai\*jiadingg@amazon\.com AmazonShuai Zhang\*shuaizs@amazon\.com AmazonKaj Bostrombostromk@amazon\.com AmazonJin Huangjinhun@amazon\.com AmazonVihang Patilpvihang@amazon\.com AmazonHaoyang Fanghaoyfang@amazon\.com AmazonBernie Wangyuyawang@amazon\.com AmazonHuzefa Rangwala†rhuzefa@gmail\.com SiemensGeorge Karypis†karypis@umn\.edu University of Minnesota

###### Abstract

We presentKernelPro, a closed\-loop multi\-agent system that automatically generates, profiles, and iteratively optimizes GPU kernel code by integrating large language model \(LLM\) code generation with hardware profiler feedback and pluggable bottleneck detection tools\.KernelProintroduces four contributions: \(1\) a*semantic feedback operator*that encodes expert heuristics as pluggable micro\-profiling tools, transforming raw hardware metrics into actionable natural language guidance, \(2\) a*two\-stage tool invocation architecture*where roofline\-based bottleneck classification filters which specialized analysis tools execute, combining kernel\-level \(ncu\), instruction\-level \(SASS\), and system\-level \(nsys\) profiling, \(3\) a*domain\-adapted MCTS*with progressive widening, asymmetric branching, log\-reward calibration, dead\-end pruning, and search memory for cross\-iteration learning, and \(4\)*direct CuTe source\-level code generation*via autonomous code search over the CUTLASS/CuTe codebase, mimicking how expert engineers write high\-performance GPU kernels in raw CUDA\+CuTe\. On KernelBench,KernelProachieves geometric mean speedups of 2\.42×\\times/4\.69×\\times/5\.30×\\timeson Levels 1/2/3, establishing state\-of\-the\-art \(SOTA\) performance across all difficulty levels\. On VeOmni’s expert\-optimized MoE training kernels,KernelProachieves 1\.23×\\timesover hand\-tuned Triton by generating a from\-scratch raw\-CUDA\+CuTe Hopper WGMMA kernel\. Ablation studies demonstrate that each design component independently and significantly improves optimization quality: micro\-profiling tools \(p<0\.0001p<0\.0001vs raw metrics\), MCTS search \(26% higher geometric mean vs greedy,p=0\.004p=0\.004\), and proactive tool orchestration \(23% improvement,p=0\.035p=0\.035\)\. Finally,KernelProis the first CUDA kernel coding agent to optimize*energy efficiency*beyond the speed\-only focus of prior systems, demonstrating an 11\.6% measured energy reduction at matched speed\.

11footnotetext:\*Equal contribution\.Code will be released upon publication\.

## 1Introduction

Optimizing GPU kernels demands expertise that spans hardware architecture, memory hierarchies, instruction scheduling, and vendor\-specific programming models\. Experienced engineers follow a disciplined cycle: profile the kernel, recognize bottleneck patterns in the metrics, diagnose root causes, apply targeted transformations, and re\-profile to verify\. The bottleneck shifts after each successful optimization, and the process repeats until hardware limits are reached\. What makes experts effective is not the ability to read numbers from a profiler, but the*pattern recognition*and*diagnostic reasoning*they apply to those numbers–heuristics refined over years and often codified into personal scripts and tools\(NVIDIA Corporation,[2024](https://arxiv.org/html/2606.26453#bib.bib24); Yang et al\.,[2020](https://arxiv.org/html/2606.26453#bib.bib33)\)\.

Recent LLM\-based CUDA optimization systems\(Zhang et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib37); Dong et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib8); Li et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib20); Chen et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib7)\)have shown that language models can generate competitive kernel code when given profiling data\. However, these systems present raw or lightly summarized metrics directly to the LLM and rely on the model to implicitly replicate expert reasoning\. This conflates two distinct capabilities:*interpreting hardware telemetry*\(a structured, rule\-governed task\) and*generating optimized code*\(a creative, context\-dependent task\)\. By asking the LLM to do both at once, existing agents forgo the systematic intermediate analysis that makes human experts reliable–they miss optimizations, produce inconsistent diagnoses, and cannot easily incorporate new analysis patterns without retraining\.

We introduceKernelPro, a closed\-loop multi\-agent system that separates these concerns by encoding expert heuristics as*pluggable micro\-profiling tools*–executable analysis functions that transform raw hardware metrics into actionable natural language guidance before the code\-generation LLM sees them\. Each tool implements a specific diagnostic pattern: a trigger condition \(e\.g\., tensor core utilization below 10%\), analysis logic \(e\.g\., check whether the kernel performs reducible matrix operations\), and a prescriptive recommendation \(e\.g\., “rewrite using CUTLASS GEMM for a major expected speedup”\)\. A two\-stage invocation pipeline first classifies the kernel’s bottleneck type via roofline analysis, then dispatches only the relevant tools–reducing prompt noise and focusing the LLM on pertinent optimizations\.

KernelProfurther contributes a domain\-adapted Monte Carlo Tree Search \(MCTS\) strategy that balances exploitation of promising optimization paths with exploration of alternative approaches\. Unlike the flat iterative refinement used by prior LLM\-based CUDA optimization systems, MCTS with UCT selection avoids premature convergence to local optima on complex multi\-step optimization problems\. Combined with a search memory mechanism that distills findings across iterations,KernelProachieves monotonically improving performance curves under fixed compute budgets\.

On the KernelBench benchmark\(Ouyang et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib26)\),KernelProachieves geometric mean speedups of2\.42×2\.42\\times/4\.69×4\.69\\times/5\.30×5\.30\\timeson Levels 1/2/3, surpassing KernelBlaster\(Dong et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib8)\), the prior SOTA at time of writing, on all levels\. Following the principle that design choices should be empirically established rather than assumed, we validate each ofKernelPro’s components through controlled ablations—each assessed for significance with a paired Wilcoxon signed\-rank test\(Hollander et al\.,[2014](https://arxiv.org/html/2606.26453#bib.bib10)\)—on a fixed 42\-task KernelBench subset \(Levels 1–3\), selected for transformer relevance and used throughout this paper \(Section[5\.3](https://arxiv.org/html/2606.26453#S5.SS3)\)\. For instance, micro\-profiling tools provide 125% higher speedup than raw metrics \(p<0\.0001p<0\.0001\), and MCTS significantly outperforms greedy search \(Wilcoxonp=0\.004p=0\.004, 26% higher geometric mean\)\.

Overall,KernelProcontributes and synthesizes five ideas to make LLM\-based optimization more grounded, interpretable, and effective on modern GPU architectures:

1. 1\.Asemantic feedback operatorthat formalizes expert heuristics as pluggable, LLM\-invocable micro\-profiling tools, transforming hardware metrics into natural language optimization guidance\.
2. 2\.Atwo\-stage tool invocation architecturewhere roofline\-based bottleneck classification \(Stage 1\) filters which specialized analysis tools execute \(Stage 2\), combining kernel\-level, instruction\-level, and system\-level profiling\.
3. 3\.Adomain\-adapted MCTSfor LLM\-based CUDA optimization with progressive widening, asymmetric branching, log\-reward calibration, and dead\-end pruning\. Search memory enables cross\-iteration learning\.
4. 4\.Direct CuTe source\-level code generationvia autonomous code search over the CUTLASS/CuTe codebase\(NVIDIA,[2023](https://arxiv.org/html/2606.26453#bib.bib23)\)\(CUTLASS is a CUDA template library of high\-performance kernels; CuTe is its underlying tensor/layout\-algebra layer for composing kernels from scratch\), mimicking how expert engineers write high\-performance GPU kernels in raw CUDA\+CuTe \(full generated source in Appendix[H](https://arxiv.org/html/2606.26453#A8)\)\.
5. 5\.Thefirst energy\-aware CUDA kernel coding agent: whereas prior LLM\-based kernel agents optimize for speed alone,KernelProextends its reward and profiling tools to*energy efficiency*as a secondary objective via a lexicographic energy\-aware reward\. As a preliminary study, a matched\-speed A/B finds a measured11\.6%11\.6\\%energy reduction at identical speed via lower\-energy instruction selection \(Section[5\.5](https://arxiv.org/html/2606.26453#S5.SS5)\); a full evaluation is left to future work\.

## 2Related Work

Several recent systems apply LLMs to CUDA kernel optimization, differing primarily in how hardware profiling data reaches the LLM\. Some rely solely on end\-to\-end timing without hardware profiling\(Andrews & Witteveen,[2025](https://arxiv.org/html/2606.26453#bib.bib1); Zhang et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib38); Wiedemann et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib31); Chen et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib6)\)\. Others pass raw or statistically filteredncumetrics directly into the LLM’s prompt and rely on the LLM itself to interpret them\(Zhang et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib37); Dong et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib8); Li et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib20); Team & Meta,[2025](https://arxiv.org/html/2606.26453#bib.bib29); Han et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib9)\)\. cuPilot\(Chen et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib7)\)adds roofline\-based bottleneck classification but stops at a label \(memory\-bound vs\. compute\-bound\) without generating actionable directives from the underlying counters\. In all cases, the LLM remains the sole interpreter of profiling data–no prior LLM\-based system provides programmatic translation of hardware counters into optimization directives\. The challenge predates LLMs: classical performance advisors such as GPA\(Zhou et al\.,[2021](https://arxiv.org/html/2606.26453#bib.bib39)\)use instruction sampling and data\-flow analysis to attribute stalls to their root causes and emit human\-readable optimization suggestions, and independently observe that raw profiler output “provide\[s\] little insight into how to improve the code\.”KernelProadopts the same diagnose\-then\-prescribe philosophy but*closes the loop*: its micro\-profiling tools translate counters into natural\-language directives*before*any LLM involvement, and an LLM then applies the fix directly rather than leaving it to a human engineer\. Our ablation \(Appendix[C](https://arxiv.org/html/2606.26453#A3)\) confirms that this tool\-based interpretation significantly outperforms both no\-feedback and raw\-metrics\-only feedback\.

CUTLASS and CuTe code generation\.Several agents target NVIDIA’s CUTLASS library, but at different abstraction levels\. AVO\(Chen et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib6)\)and StitchCUDA\(Li et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib20)\)generate high\-level CUTLASS template instantiations–selecting tile sizes, pipeline stages, and epilogue fusions from a predefined configuration space\.KernelProoperates at a lower level: it generates raw CUDA\+CuTe source code by autonomously searching the CUTLASS/CuTe codebase for layout algebras, copy atoms, and MMA atoms, then composing them into novel kernels – the same workflow used by expert kernel engineers\.

Tree search and evolutionary methods for LLM code optimization\.A growing body of work integrates Monte Carlo Tree Search with LLMs\. MCTSr\(Zhang et al\.,[2024](https://arxiv.org/html/2606.26453#bib.bib36)\)applies UCT to mathematical olympiad problems, where each node represents an answer version refined via self\-evaluation\. DeepSearch\(Wu et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib32)\)embeds MCTS into the RL training loop with global frontier selection, targeting model weight updates rather than inference\-time optimization\. TreeRL\(Hou et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib13)\)proposes entropy\-guided token\-level branching that forks from high\-uncertainty tokens\. Tree\-GRPO\(Ji et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib15)\)uses complete agent steps as tree nodes and derives process supervision signals from intra\-tree relative advantages\. These systems treat tree search as a mechanism for*training*LLMs or generating*reasoning traces*\. AlphaEvolve\(Novikov et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib22)\)takes an evolutionary approach, using population\-based LLM sampling with scalar fitness to optimize algorithms including TPU kernels, but requires millions of samples over days of compute\.

KernelProintroduces a domain\-adapted MCTS for LLM\-based GPU kernel optimization, where each node represents a complete compiled\-and\-profiled CUDA kernel\. This coarse granularity demands different design choices: progressive widening\(Auger et al\.,[2013](https://arxiv.org/html/2606.26453#bib.bib2)\)gates expansion because each node requires compilation, execution, and profiling, ensuring exploration grows sublinearly with visit count; log\-reward calibration eliminates the need for learned value functions; asymmetric branching factors reflect the empirical observation that valid optimizations outnumber valid repairs; dead\-end pruning removes exhausted subtrees; and search memory enables cross\-iteration learning\.KernelProadopts MCTSr’s ROOT re\-expansion \(fresh seed injection when existing subtrees plateau\) but replaces self\-reward with measured speedup, eliminating reward hacking risk\. While writing this paper, we observed that OptiML\(Bhattacharjee et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib4)\)has also independently applied MCTS with UCT for CUDA kernel optimization, which we view as convergent evidence that tree search is a natural fit for this domain\.

## 3The Architecture ofKernelPro

### 3\.1High\-Level Pipeline

KernelProoperates as a multi\-agent system with three primary agents coordinated by a search orchestrator\. TheBenchmarking Agent\(Stage 1\) analyzes the reference implementation to establish baseline performance metrics and classify the kernel’s bottleneck type \(compute\-bound, memory\-bound, latency\-bound, or mixed\) via roofline analysis\. TheProgrammer Agentgenerates optimized CUDA code in three modes: initial generation, debugging compilation/correctness errors, and iterative optimization based on profiling feedback\. TheProfiling Agent\(Stage 2\) executes micro\-profiling tools filtered by the Stage 1 bottleneck classification, producing semantic feedback that guides the next optimization iteration\. This two\-stage architecture–bottleneck classification followed by filtered tool execution–reduces prompt noise and focuses the LLM on relevant optimizations\.

Stage 1 \(one\-time\)Iterative LoopStage 2Micro\-Profiling ToolsCUDA/Triton/PyTorchKernelBenchmarking Agent\(Roofline Classification\)Bottleneck TypeSearch OrchestratorMCTS / GreedyProgrammer Agentinit∣\\middebug∣\\midoptimizeNNCandidatesBest ValidatedCUDA KernelSearch Memorycross\-iteration learningCompile & Executevalidate correctnessProfiling Agent\(Semantic Feedback\)ncu\(8 tools\)nsys\(5 tools\)SASS\(2 tools\)input kernelperf analysisselect nodeconvergedCUDA kernel codebottleneck type filteringinjectmemoizememoizesemantic feedbackncu metrics \+ \.cubinControl / Main FlowSemantic FeedbackMemory FlowTool InvocationBottleneck Filtering

Figure 1:KernelProagentic optimization workflow\. Stage 1 \(Benchmarking Agent\) performs one\-time roofline\-based bottleneck classification\. The iterative loop comprises search, programming, compilation, correctness validation, profiling, and semantic feedback\. The Search Orchestrator \(MCTS or greedy\) maintains an expansion tree over candidate solutions, while the Stage 2 Profiling Agent invokes bottleneck\-filtered micro\-profiling tools \(8ncu\+ 5nsys\+ 2 SASS analyzers\)\. Search Memory persists cross\-iteration findings outside the loop\. When the search converges, the workflow emits the best validated CUDA kernel\.All agents communicate with LLMs through abackend\-agnostic interfacebuilt on the OpenAI Agents SDK\(OpenAI,[2025](https://arxiv.org/html/2606.26453#bib.bib25)\), so the same agent code runs unmodified on AWS Bedrock or vLLM\(Kwon et al\.,[2023](https://arxiv.org/html/2606.26453#bib.bib18)\)\. The Bedrock path requires the most adaptation: a custom adapter translates tool definitions, multi\-turn tool\-use/tool\-result exchanges, and structured responses between the SDK’s tool\-calling protocol and the Bedrock Claude Messages API, whereas vLLM’s OpenAI\-compatible endpoint is consumed directly with no protocol translation\.

Proactive Tool Orchestration for the Profiling Agent\.A deliberate design choice inKernelProis the use of*proactive*tool orchestration for the Profiling Agent, rather than the*reactive*pattern typical of LLM agent frameworks\(Yao et al\.,[2023](https://arxiv.org/html/2606.26453#bib.bib35)\)\. In reactive systems, the LLM decides which tools to invoke–a stochastic process that may skip critical analyses\. The Profiling Agent instead orchestrates its 15 micro\-profiling tools*deterministically*: all relevant tools are executed based on the bottleneck classification, injecting their guidance into the prompt before LLM invocation\. This mirrors how human CUDA experts systematically analyze all relevant metrics rather than selectively querying information\. The proactive approach guarantees comprehensive analysis and consistent optimization quality, which is essential for reliable CUDA optimization outcomes\. A concrete example: in Task 41 \(Appendix[G\.1](https://arxiv.org/html/2606.26453#A7.SS1)\), the tools turn a memory\-bottleneck diagnosis into a ranked list of specific, actionable directives \(epilogue fusion, BF16 tensor cores, register\-pressure and coalescing fixes\); the agent applies the top recommendations together for a2\.8×2\.8\\timessingle\-step improvement\. A reactive agent invoking only a few tools would surface a fraction of this guidance\.

### 3\.2Agent Decomposition

#### 3\.2\.1Benchmarking Agent \(Stage 1\)

The Benchmarking Agent runsonce per problemat the start of optimization\. It parses the reference code to identify the mathematical operations involved \(GEMM, Softmax, convolution, etc\.\), computes a roofline estimate of theoretical peak performance from compute intensity and hardware specs, and classifies the problem as compute\-bound, memory\-bound, or mixed based on the CI/ridge ratio\. It then profiles the reference implementation withncuto validate this theoretical classification against measured utilization\. The output is aReferenceAnalysisobject that guides downstream tool selection for the Programmer and Profiling Agents\.

#### 3\.2\.2Programmer Agent

The Programmer Agent generates candidate kernel implementations throughKernelPro’s backend\-agnostic model interface, operating in three modes:

- •Initial: Generate first\-attempt CUDA implementations from reference kernel
- •Debugging: Fix compilation errors or correctness failures in previous attempts
- •Optimize: Improve performance of working solutions based on profiler feedback

Prompts are constructed via amodular PromptBuilderthat composes task\-specific instructions from reusable components: base optimization strategies, dataset\-specific formats, mode\-specific instructions, and curated examples\. This modularity allows different baseline types \(PyTorch, Triton\(Tillet et al\.,[2019](https://arxiv.org/html/2606.26453#bib.bib30)\), CUDA\) to share common optimization knowledge while receiving tailored guidance\. A notable feature isGPU\-aware prompt injection: at agent creation time,KernelProdetects the target GPU architecture and injects architecture\-specific guidance—most importantly the correct tensor\-core instruction shapes \(§[3\.2\.4](https://arxiv.org/html/2606.26453#S3.SS2.SSS4)\)—so generated code targets the right instruction format without requiring the LLM to memorize hardware details\.

The agent outputs a structuredMultipleCompletionsobject containingNNcandidate solutions per iteration, each with its optimization strategy annotation\. This structured output enables downstream validation and trajectory recording\. A two\-phase truncation\-recovery mechanism handles cases where generated code exceeds the output token limit \(detected via amax\_tokensstop reason\): Phase A retries with doubledmax\_tokens; if still truncated, Phase B retries at that doubled budget with an appended directive to restart concisely\.

#### 3\.2\.3Profiling Agent \(Stage 2\)

The Profiling Agent transforms raw hardware metrics into semantic optimization guidance\. Rather than simply forwardingncuoutput to the LLM, it constructs a rich analysis prompt through a multi\-stage process\.

First, the agent collects metrics from three profiling sources:ncuprovides kernel\-level metrics \(throughput, occupancy, stalls\),nsysprovides system\-level timeline data \(launch overhead, synchronization\), andSASS analysisprovides instruction\-level information by disassembling the compiled\.sofile\. To obtain the binary path,KernelProintercepts PyTorch’sload\_inlinecompilation pipeline, capturing the ephemeral\.soartifact for downstreamncuand SASS analysis\.ncuprofiling targets only theKernelPro\-generated\_\_global\_\_kernels identified via binary symbol analysis \(cuobjdump \-symbols, filtering forSTB\_GLOBALentry points\), with fallback to duration\-based kernel selection for library\-call solutions \(e\.g\., cuBLAS\)\. This multi\-level fusion is essential–ncualone cannot distinguish whether a kernel uses Ampere\-style async copies or Hopper TMA, but SASS reveals this directly\.

Second, the agent invokes themicro\-profiling tool registry\. The key functionToolRegistry\.generate\_prompt\_guidance\(\)iterates through registered tools \(filtered by Stage 1 bottleneck type\), executes each tool’sanalyze\(\)method on the collected metrics, and aggregates triggered findings into markdown\-formatted guidance\. This guidance is injected into the LLM prompt under an “Automated Bottleneck Analysis” section\.

The resulting prompt structure follows a consistent format: SASS analysis results \(tensor core instruction counts, memory patterns\), followed by an nsys timeline summary \(launch overhead percentage, sync calls\)—where gap detection uses an adaptive threshold ofmax⁡\(1%​of end\-to\-end time,5​μ​s\)\\max\(1\\%\\ \\text\{of end\-to\-end time\},\\ 5\\mu\\text\{s\}\)so that only gaps meaningful relative to total runtime are flagged, avoiding false positives on long\-running kernels—then tool\-generated bottleneck analysis with severity ratings and recommendations, and finally ncu metrics organized by category \(compute, memory, occupancy, stalls\)\.

The LLM then synthesizes this information into a structuredProfileAnalysisoutput containing bottleneck classification, severity assessment, and prioritized optimization recommendations\. This structured output feeds back to the Programmer Agent for the next iteration\.

Tool invocation pattern\.The Profiling Agent uses*proactive*tool orchestration\. With 15 micro\-profiling tools, reactive invocation would introduce stochastic selection–the LLM might call 3 tools and skip 10, missing critical analyses\. Instead,KernelProdeterministically executes all relevant tools based on the Stage 1 bottleneck classification, injecting their combined guidance into the prompt*before*LLM invocation\. This guarantees comprehensive coverage and consistent optimization quality across runs\. Stage 1 classifies the bottleneck using two independent methods for robustness: a theoretical roofline bound \(arithmetic intensity vs\. the hardware ridge point\) and the measured compute\-vs\-memory throughput fromncu\. Agreement between the two gives high confidence; on disagreement,KernelProdefers to the theoretical bound, since it reflects the kernel’s inherent ceiling rather than the current \(possibly unoptimized\) implementation\. Table[1](https://arxiv.org/html/2606.26453#S3.T1)shows the resulting tool filtering\.

Table 1:Tool filtering by bottleneck type
#### 3\.2\.4Open\-Source Library Integration

KernelProtreats open\-source GPU libraries asfirst\-class citizensthrough an extensible integration architecture\. Prior CUDA AI agents either stitch pre\-compiled library calls \(cuBLAS, cuDNN\) or generate raw CUDA/Triton–none support source\-level code generation with libraries like CUTLASS\.KernelPro’s design enables LLM\-based code generation with any GPU programming library through two mechanisms: \(1\)GPU\-aware prompt injection–detecting target hardware and injecting library\-specific configurations, and \(2\) alibrary\-aware build and runtime environment–ensuring generated code compiles, executes, validates, and profiles seamlessly\.

We demonstrate this architecture withCUTLASS, an industry\-grade library for custom tensor\-core kernels \(adopted in systems such as PyTorch and TensorRT\); itsCuTelayer provides the tensor and layout abstractions \(copy atoms, MMA atoms, layout algebra\) used to compose such kernels from scratch\. Generating this code is challenging for LLMs: it must match hardware\-specific instruction shapes and satisfy architectural conventions that are easy to violate and sparsely documented, and it depends on a precise build configuration\.

For CUTLASS,KernelPro’sprompt injectiondetects the target GPU at runtime and injects architecture\-specific instruction shapes, include paths, and compiler flags\.KernelPro’sbuild and runtime environmenthandles compilation via PyTorch’sload\_inlinewith CUTLASS\-specific flags, tracks compilation events, validates correctness against reference implementations, and profiles compiled kernels withncu/nsys\. This closed\-loop integration ensures that LLM\-generated CUTLASS/CuTe code can compile, execute, profile, and evaluate seamlessly\. Appendix[G](https://arxiv.org/html/2606.26453#A7)presents two complete trajectories demonstrating this integration end\-to\-end: a custom BF16 epilogue functor achieving12\.1×12\.1\\times\(Task 41\) and a dual\-GEMM pipeline with hand\-written kernels achieving2\.79×2\.79\\times\(Task 46\)\.

LLM\-autonomous code search\.A key challenge in CUTLASS/CuTe code generation is that LLMs frequently produce code violating sparsely\-documented architectural constraints \(e\.g\., barrier lifecycle rules, warp specialization patterns, TMA pipeline depth\)\. Rather than relying on static prompt injection of examples,KernelProequips the Programmer Agent withLLM\-autonomous code search toolsthat allow it to query the CUTLASS/CuTe codebase directly during generation:

- •search\_cutlass: recursive grep over file contents with optional path/glob scoping \(search calls are budgeted per iteration to encourage convergence toward generation; reads and listings are unbudgeted\)
- •read\_cutlass\_file: Read specific files with line ranges
- •list\_cutlass\_directory: Directory listing for navigation
- •check\_smem\_layout: pre\-flight validation that compiles a minimaltile\_to\_shapeprobe withnvccagainst the CUTLASS headers, catching shape/stride divisibility errors—the dominant CuTe compile failure—before a full kernel build \(unbudgeted\)

This design follows the principle that the LLM should*discover*relevant patterns autonomously rather than receiving pre\-selected examples\. The LLM decides what to search based on the optimization context—e\.g\., retrievingwgmmapipeline patterns and MMA/copy atoms from the live CuTe source when generating Hopper tensor\-core code—and composes them into novel kernels\. This is the source\-level workflow that expert engineers use, and the mechanism behindKernelPro’s from\-scratch raw\-CuTe kernels \(Section[5\.4](https://arxiv.org/html/2606.26453#S5.SS4)\)\. Contemporaneous work \(GrepSeek;Salemi et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib27)\) independently finds that agent\-issued executable search over a raw corpus can outperform index\-based retrieval\.

Additionally, a constraint mining tool pre\-extracts hard architectural requirements from CUTLASS source \(barrier lifecycle, warp specialization rules, vectorized epilogue patterns, tile swizzle\) and injects them as non\-negotiable code generation constraints\. A companion tool validates generated code against these requirements before compilation, catching architectural errors early\.

#### 3\.2\.5Multi\-Language Input Support

KernelProsupportsPyTorch\-to\-CUDA,Triton\-to\-CUDA, andCUDA\-to\-CUDAtranslation\. This flexibility serves diverse optimization workflows–from AI researchers prototyping in PyTorch to kernel engineers fine\-tuning existing Triton or CUDA implementations\.

PyTorchModel \+ nn\.ModuleTriton@triton\.jit \+ run\(\)CUDAexisting kernelSource LanguagesDataset Adapters\(Normalization\)kernelbench\.py triton2cuda\.py cuda2cuda\.py \(future\)Unified Problem Formatclass Model\(nn\.Module\) \+ get\_inputs\(\) \+ get\_init\_inputs\(\)Modular Prompt Builder\(baseline\_type routing\)dataset\_formats/\{pytorch\|triton\}\.txt \+ shared optimization strategiesOutput: CUDA C\+\+torch\.utils\.cpp\_extension\.load\_inline

Figure 2:Multi\-language input pipeline: all source languages are normalized to a unified format before being routed through language\-specific prompts to produce optimized CUDA kernel\.While prior CUDA agents accept only PyTorch input \(via KernelBench\) or only existing CUDA code,KernelPronatively supports PyTorch, Triton, and CUDA inputs—enabling optimization of kernels at any stage of the development pipeline \(Figure[2](https://arxiv.org/html/2606.26453#S3.F2)\)\. All inputs are normalized to a unifiedclass Modelinterface with language\-specific prompt routing, producing CUDA C\+\+ output viaload\_inline\.

This multi\-language capability addresses a key production need: automating kernel optimization and hardware migration regardless of the source abstraction\. Triton kernels hitting DSL\-imposed performance ceilings \(e\.g\., Ampere\-era instructions on Hopper, Section[5\.4](https://arxiv.org/html/2606.26453#S5.SS4)\), PyTorch prototypes requiring production CUDA implementations, and legacy CUDA kernels needing re\-optimization for new GPU architectures can all be optimized through a single system—eliminating the need for scarce CUDA experts at each translation boundary\.

### 3\.3Search Strategy

KernelProuses a solution\-levelMonte Carlo Tree Search \(MCTS\)strategy to navigate the CUDA optimization landscape\. Each node is a complete compiled\-and\-profiled kernel rather than a single reasoning step\. Unlike greedy search, MCTS systematically explores the multi\-modal landscape \(tiling vs vectorization vs algorithmic restructuring\) while concentrating compute budget on the most promising subtrees\. Algorithm[1](https://arxiv.org/html/2606.26453#alg1)presents the full MCTS loop\.

Algorithm 1MCTS Search with Progressive Widening and Search Memory\.Green linesdenoteKernelPro\-specific adaptations for CUDA kernel optimization\.1:Input:Problem

PP, max iterations

TT, candidates per expansion

NN, UCT constant

CuctC\_\{\\text\{uct\}\}
2:

ROOT←VirtualNode​\(\)\\text\{ROOT\}\\leftarrow\\text\{VirtualNode\}\(\),

ℳ←∅\\mathcal\{M\}\\leftarrow\\emptyset⊳\\trianglerightVirtual root, search memory

3:

ref\_analysis←BenchmarkingAgent\.analyze​\(P\)\\text\{ref\\\_analysis\}\\leftarrow\\text\{BenchmarkingAgent\.analyze\}\(P\)⊳\\trianglerightStage 1: roofline classification

4:for

t=0t=0to

T−1T\-1do

5:

node←UCT\-Select​\(ROOT,Cuct\)\\text\{node\}\\leftarrow\\text\{UCT\-Select\}\(\\text\{ROOT\},C\_\{\\text\{uct\}\}\)⊳\\trianglerightTraverse tree via UCT \+ progressive widening

6:ifnode = ROOTthenmode←initial\\text\{mode\}\\leftarrow\\text\{initial\};parent←None\\text\{parent\}\\leftarrow\\text\{None\}⊳\\trianglerightFresh seed injection

7:else ifnode\.status = FAILEDthenmode←debug\\text\{mode\}\\leftarrow\\text\{debug\};parent←node\\text\{parent\}\\leftarrow\\text\{node\}⊳\\trianglerightRepair attempt

8:elsemode←optimize\\text\{mode\}\\leftarrow\\text\{optimize\};parent←node\\text\{parent\}\\leftarrow\\text\{node\}⊳\\trianglerightRefine working solution

9:

candidates←ProgrammerAgent\.generate​\(P,parent,parent\.feedback,ℳ,mode,N\)\\text\{candidates\}\\leftarrow\\text\{ProgrammerAgent\.generate\}\(P,\\text\{parent\},\\text\{parent\.feedback\},\\mathcal\{M\},\\text\{mode\},N\)
10:foreach

ccin candidatesdo

11:

result←Execute​\(c\)\\text\{result\}\\leftarrow\\text\{Execute\}\(c\)⊳\\trianglerightCompile, run, validate

12:ifresult\.validthen

13:

metrics←Profile​\(c\)\\text\{metrics\}\\leftarrow\\text\{Profile\}\(c\)⊳\\trianglerightncu\+nsyshardware profiling

14:

sass\_info←SASSAnalyze\(c\.cubin\)\\text\{sass\\\_info\}\\leftarrow\\text\{SASSAnalyze\}\(c\.\\text\{cubin\}\)⊳\\trianglerightDisassemble compiled binary

15:

guidance←RunTools​\(metrics,sass\_info,bottleneck\)\\text\{guidance\}\\leftarrow\\text\{RunTools\}\(\\text\{metrics\},\\text\{sass\\\_info\},\\text\{bottleneck\}\)⊳\\trianglerightExpert surrogates

16:

analysis←ProfilingAgent\.analyze​\(c,metrics,guidance\)\\text\{analysis\}\\leftarrow\\text\{ProfilingAgent\.analyze\}\(c,\\text\{metrics\},\\text\{guidance\}\)
17:

c\.feedback←analysisc\.\\text\{feedback\}\\leftarrow\\text\{analysis\}⊳\\trianglerightTool\-guided feedback replaces rollout

18:endif

19:

r​\(c\)←Reward​\(c\)r\(c\)\\leftarrow\\text\{Reward\}\(c\)⊳\\trianglerightln⁡\(speedup\)\\ln\(\\text\{speedup\}\),−2\.0\-2\.0, or−3\.0\-3\.0

20:

Backpropagate​\(c,r​\(c\)\)\\text\{Backpropagate\}\(c,r\(c\)\)⊳\\trianglerightUpdate visits and rewards to ROOT

21:

AttachChild​\(node,c\)\\text\{AttachChild\}\(\\text\{node\},c\)
22:endfor

23:

MarkTerminals​\(node\)\\text\{MarkTerminals\}\(\\text\{node\}\)⊳\\trianglerightPrune if≥\\geq3 all\-failed repair children

24:

ℳ←ExtractMemory​\(t,candidates,result,ℳ\)\\mathcal\{M\}\\leftarrow\\text\{ExtractMemory\}\(t,\\text\{candidates\},\\text\{result\},\\mathcal\{M\}\)⊳\\trianglerightBackground summarizer

25:endfor

26:return

GlobalBest​\(ROOT\)\\text\{GlobalBest\}\(\\text\{ROOT\}\)⊳\\trianglerightBest solution across entire tree

#### 3\.3\.1UCT Selection

KernelPro’s MCTS maintains a search tree where each node represents a concrete CUDA solution \(compiled, profiled, and evaluated\)\. Selection follows the UCT algorithm\(Kocsis & Szepesvári,[2006](https://arxiv.org/html/2606.26453#bib.bib17)\): starting from a virtual ROOT node, the tree is traversed by repeatedly selecting the child that maximizes the upper confidence boundUCT​\(c\)=r¯c\+Cuct​ln⁡Nparent/Nc\\text\{UCT\}\(c\)=\\bar\{r\}\_\{c\}\+C\_\{\\text\{uct\}\}\\sqrt\{\\ln N\_\{\\text\{parent\}\}/N\_\{c\}\}, wherer¯c\\bar\{r\}\_\{c\}is the node’s mean reward,NcN\_\{c\}its visit count,NparentN\_\{\\text\{parent\}\}the parent’s visit count, andCuct=2C\_\{\\text\{uct\}\}=\\sqrt\{2\}\. Unlike evolutionary approaches that maintain a candidate population, tree\-structured selection gives proper credit to underexplored subtrees–a node buried three levels deep in a promising branch competes fairly against shallow alternatives because the exploration bonus is computed relative to its parent’s visits, not the global budget\.

#### 3\.3\.2Progressive Widening with Domain\-Adapted Branching

Standard MCTS assumes a finite action set and expands a node until all its actions are tried\.KernelPro’s action–“generate another CUDA solution”–has an effectively unbounded branching factor, and each child is expensive to evaluate \(compilation, execution,ncuprofiling, and SASS analysis\)\.KernelProtherefore uses progressive widening\(Auger et al\.,[2013](https://arxiv.org/html/2606.26453#bib.bib2)\)to gate expansion: a node withNNvisits may have at mostCpw⋅NαC\_\{\\text\{pw\}\}\\cdot N^\{\\alpha\}children \(α=0\.5\\alpha=0\.5\), so a node earns additional candidates only as repeated visits confirm its promise\.

The key domain adaptation is asymmetric branching factors tied to node status\. For ROOT \(fresh seed generation\) and FAILED nodes \(repair attempts\), the widening constant isCpw=2C\_\{\\text\{pw\}\}=2–moderate diversity is appropriate because failed solutions typically need a specific fix rather than broad exploration\. For SUCCESSFUL nodes \(optimization refinements\),Cpw=3C\_\{\\text\{pw\}\}=3–the search allocates more children because correct CUDA kernels have richer optimization surfaces with many orthogonal improvement axes \(memory hierarchy tuning, instruction\-level parallelism, tensor core utilization etc\)\. This asymmetry encodes the empirical observation that the space of valid optimizations is broader than the space of valid repairs\.

At each internal node during traversal,KernelProcomputes a*virtual child UCT score*: the hypothetical UCT of a new child that would inherit its parent’s mean reward as a prior and receive a fair share of visits\. If this virtual score exceeds the UCT of the best existing child, the search expands the node \(generates a new solution as its child\) rather than descending further\. This mechanism provides a principled expansion criterion–new children are only created when the uncertainty about unexplored alternatives outweighs the expected value of exploiting known children\.

#### 3\.3\.3Reward Calibration

Backpropagation requires a scalar reward signal that meaningfully ranks solutions across the tree\.KernelProuses a three\-tier reward function calibrated on a diverse set of transformer\-centric GPU kernels:

r​\(n\)=\{ln⁡\(speedup​\(n\)\)correct−2\.0incorrect output−3\.0crash/compile errorr\(n\)=\\begin\{cases\}\\ln\(\\text\{speedup\}\(n\)\)&\\text\{correct\}\\\\ \-2\.0&\\text\{incorrect output\}\\\\ \-3\.0&\\text\{crash/compile error\}\\end\{cases\}\(1\)
The logarithmic transform for correct solutions prevents high\-speedup outliers \(e\.g\., 100×\\timeson trivially parallelizable kernels\) from dominating the tree–a 2×\\timesspeedup on a well\-optimized kernel is arguably harder than 50×\\timeson an unoptimized one, and log\-scaling reflects this\. The gap between tiers \(\-2\.0 for incorrect vsln⁡\(1\)=0\\ln\(1\)=0at the correctness boundary\) ensures that “almost correct” solutions never dominate the tree over genuinely correct ones, preventing the search from wasting budget refining broken code\. Rewards propagate from leaf to ROOT via backpropagation, updating visit counts and cumulative rewards at every ancestor\.

#### 3\.3\.4Dead\-End Pruning and ROOT Re\-expansion

Two mechanisms prevent the search from stalling\. First,terminal marking: when a failed node accumulates three or more repair children that all fail, it is marked terminal and excluded from future selection\. This prunes dead\-end subtrees where the underlying algorithmic approach is fundamentally broken, propagating upward through the tree as entire branches become exhausted\.

Second,ROOT re\-expansion: when UCT selects the virtual ROOT node itself \(rather than any of its children\), the system generates entirely fresh solutions from scratch rather than refining existing ones\. This mechanism–inspired by MCTSr\(Zhang et al\.,[2024](https://arxiv.org/html/2606.26453#bib.bib36)\)–injects exploration diversity as the tree deepens, preventing the search from over\-committing to early seeds that happen to look promising\. The progressive widening gate on ROOT ensures this happens at a controlled rate: fresh seeds are generated only when the virtual child’s UCT score at ROOT exceeds that of all existing top\-level solutions, indicating that unexplored regions of the solution space may outperform known optima\.

A controlled ablation comparing MCTS against greedy search under matched compute budgets \(Appendix[E](https://arxiv.org/html/2606.26453#A5)\) confirms that tree\-structured exploration yields a 26% higher geometric mean speedup across 42 KernelBench tasks \(p=0\.004p=0\.004\), with the largest gains on problems requiring multi\-step optimization chains where greedy converges prematurely\. The NetVLAD trajectory \(Appendix[G](https://arxiv.org/html/2606.26453#A7), Task 46\) provides concrete evidence: 43 consecutive failures precede the first working solution—a search depth that would exhaust any fixed\-budget single\-shot approach—yet once a correct kernel emerges, profiling\-guided refinement converges rapidly \(0\.75×→2\.79×0\.75\\times\\to 2\.79\\timesin two refinement steps\)\.

### 3\.4Search Memory: Cross\-Iteration Learning

A fundamental challenge in iterative LLM\-based optimization is that each iteration starts*tabula rasa*–the LLM has no memory of prior attempts beyond the immediate parent’s profiling feedback\. This leads to three failure modes we observe in practice: \(1\)repeated errors, where the LLM re\-attempts approaches that failed in earlier iterations; \(2\)lost discoveries, where useful search findings \(file paths, API patterns\) from tool calls are not carried forward; and \(3\)plateau behavior, where the LLM cycles between similar strategies without building toward a solution\.

KernelProaddresses this with asearch memorysystem \(Figure[1](https://arxiv.org/html/2606.26453#S3.F1)\)–a lightweight, append\-only knowledge store scoped to a single optimization run\. The system has three stages\. First, apost\-iteration extractor: after each iteration completes \(success or failure\), a separate background LLM call receives the iteration context–tool calls made, code approach attempted, and evaluation outcome–and produces structured entries for five memory sections\. This separation ensures the code\-generation LLM operates without meta\-cognitive overhead\. Second, astructured memory storeorganizes entries into semantically distinct sections:Search Findings\(file paths, API patterns, and architectural details discovered via tool calls\),Errors & Corrections\(what was tried and why it failed–compilation errors, correctness failures, performance regressions\),Successful Patterns\(approaches that achieved measurable speedup, with the technique and result\),Key Files\(source files containing relevant implementations or examples\), andLearnings\(higher\-level principles distilled from the iteration, e\.g\., “shared memory limits occupancy for this kernel size”\)\. Third,prompt injection: the accumulated search memory is formatted as a structured block and injected into the Programmer Agent’s system prompt at the start of each iteration\. The memory grows monotonically across iterations, providing an increasingly rich context for optimization decisions\.

A controlled ablation across 42 KernelBench tasks shows that search memory produces statistically equivalent final speedup while providing an early\-convergence trend\. Full ablation experiments are reported in Appendix[D](https://arxiv.org/html/2606.26453#A4)\.

## 4Micro\-Profiling Tools

KernelPro’s 15 micro\-profiling tools encode GPU optimization expertise as composable analyzers that transform raw profiler metrics into actionable guidance\. Each tool implements a trigger\-analyze\-recommend pattern: triggering on specific bottleneck signatures, analyzing severity and root cause, and generating targeted recommendations for the LLM\. The tools are organized by profiler source \(Table[2](https://arxiv.org/html/2606.26453#S4.T2)\)—kernel\-level \(ncu\), instruction\-level \(SASS binary analysis\), and system\-level \(nsys\)—because different bottlenecks manifest at different levels of abstraction and no single profiler provides complete visibility\.

Table 2:Micro\-profiling tool categories by profiler source### 4\.1Tools as Expert Surrogates

As motivated in Section[1](https://arxiv.org/html/2606.26453#S1),KernelProframes micro\-profiling tools asexecutable surrogates for human expert workflow–encoding the profile\-recognize\-diagnose\-prescribe cycle as callable tools\. While demonstrated for CUDA, this pattern generalizes to any domain with measurable metrics, known bottleneck patterns, and actionable fixes \(e\.g\., database query optimization, compiler tuning\)\.

#### 4\.1\.1The 5\-Tuple Formalization

EachKernelPromicro\-profiling tool encodes one expert heuristic as a formal 5\-tuple:

τ=\(name,ℳreq,θ,trigger,analyze\)\\tau=\(\\text\{name\},\\mathcal\{M\}\_\{\\text\{req\}\},\\theta,\\text\{trigger\},\\text\{analyze\}\)\(2\)
whereℳreq⊆ℳ\\mathcal\{M\}\_\{\\text\{req\}\}\\subseteq\\mathcal\{M\}is the required metric subset \(from the∼50\{\\sim\}50ncumetricsKernelProcollects\),θ\\thetaare threshold parameters,trigger:ℳ→\{0,1\}\\text\{trigger\}:\\mathcal\{M\}\\to\\\{0,1\\\}is the activation predicate, andanalyze:ℳ→𝒪\\text\{analyze\}:\\mathcal\{M\}\\to\\mathcal\{O\}produces structured output guidance\. The output space𝒪\\mathcal\{O\}is a fixed schema comprising: a*severity*level \(critical/high/medium/low\), a*root cause*explanation of why the bottleneck exists, ranked*recommendations*with expected improvement estimates, and optionally a*code example*demonstrating the fix\. This schema is what distinguishes semantic feedback from raw metric forwarding–the LLM receives not “occupancy = 6%” but “occupancy is critically low \(6%\) because shared memory usage limits concurrent blocks to 2; switch to warp\-level reduction for an expected 3–5×\\timesimprovement\.” A full worked example is provided in Appendix[B](https://arxiv.org/html/2606.26453#A2)\. The following table illustrates how the 5\-tuple maps to expert reasoning:

Table 3:Mapping human expert workflow to tool components
#### 4\.1\.2Proactive vs\. Reactive Tool Orchestration

KernelProusesproactivetool orchestration for the Profiling Agent: rather than letting the LLM decide which tools to call \(the reactive pattern in standard agent frameworks\),KernelProdeterministically executes all bottleneck\-relevant tools and injects their combined guidance into the prompt before LLM invocation\. Crucially, “bottleneck\-relevant” is not “all tools”–the bottleneck classification from Stage 1 feeds a tool\-affinity filter that excludes irrelevant analyses \(e\.g\., memory\-bound kernels skip tensor core utilization checks, compute\-bound kernels skip coalescing analysis\), preventing signal dilution while guaranteeing comprehensive coverage of the actual bottleneck\. This transforms expert knowledge from an optional resource into a guaranteed input\. The reactive pattern’s weakness–stochastic tool selection leading to missed analyses and inconsistent coverage–is particularly harmful for CUDA optimization, where the relevant bottleneck is often non\-obvious and experts systematically check all metrics rather than guessing which to examine\. Our controlled ablation \(Appendix[F](https://arxiv.org/html/2606.26453#A6)\) confirms this empirically: proactive orchestration achieves 23% higher geometric mean speedup than reactive function calling \(p=0\.035p=0\.035\)\.

#### 4\.1\.3Pluggability and User Extensibility

KernelPro’s tools arepluggable–they can be dynamically added, removed, or modified without changing the core system via a decorator\-based registration pattern \(@ToolRegistry\.register\)\. Users subclassProfilingTool, declare required metrics and thresholds, and implement ananalyze\(\)method; the tool is then automatically incorporated into the proactive orchestration pipeline\. This makesKernelProa platform rather than a fixed tool: teams can encode domain\-specific heuristics \(e\.g\., MoE expert load imbalance detection\), add support for new GPU architectures before official updates, or register proprietary optimizations as private tools\. The full registration API and a worked example are provided in Appendix[B](https://arxiv.org/html/2606.26453#A2)\(Listing[2](https://arxiv.org/html/2606.26453#LST2)\)\.

### 4\.2Tool Specifications

Each tool follows the 5\-tuple formalization described above\. The complete tool taxonomy \(15 tools organized by category and bottleneck type\) and detailed specifications are provided in Appendix[B](https://arxiv.org/html/2606.26453#A2), including full reference implementations \(Listings[2](https://arxiv.org/html/2606.26453#LST2)and[3](https://arxiv.org/html/2606.26453#LST3)\) that make the trigger\-analyze\-recommend pattern concrete\. We detail two representative tools below: one NCU\-based analyzer and the SASS binary analysis tools\.

#### 4\.2\.1Warp Stall Analyzers

KernelProincludes two warp stall analyzers built onncumetrics\. The primaryWarpStallToolanalyzes why warps are not issuing instructions, identifying the dominant stall reason\. It requires five warp\-stall metricsℳreq=\{mem\_dep,short\_scoreboard,long\_scoreboard,barrier,branch\_resolving\}\\mathcal\{M\}\_\{\\text\{req\}\}=\\\{\\texttt\{mem\\\_dep\},\\;\\texttt\{short\\\_scoreboard\},\\;\\texttt\{long\\\_scoreboard\},\\;\\texttt\{barrier\},\\;\\texttt\{branch\\\_resolving\}\\\}, wherem\.sm\.sdenotes the percentage of active cycles stalled on reasonss\. The tool finds the dominant stalls∗=arg⁡maxs∈ℳreq⁡m\.ss^\{\*\}=\\arg\\max\_\{s\\in\\mathcal\{M\}\_\{\\text\{req\}\}\}m\.sand triggers whenm\.s∗m\.s^\{\*\}exceedsθhigh=40%\\theta\_\{\\text\{high\}\}=40\\%\(severity escalates to critical at60%60\\%\)\. A companionBarrierStallTooltriggers at a lower threshold \(30%30\\%\) specifically for barrier synchronization issues, which are both a performance and correctness concern\.

ℳreq\\displaystyle\\mathcal\{M\}\_\{\\text\{req\}\}=\{mem\_dep,short\_scoreboard,long\_scoreboard,\\displaystyle=\\bigl\\\{\\,\\texttt\{mem\\\_dep\},\\;\\texttt\{short\\\_scoreboard\},\\;\\texttt\{long\\\_scoreboard\},barrier,branch\_resolving\}\\displaystyle\\qquad\\;\\;\\texttt\{barrier\},\\;\\texttt\{branch\\\_resolving\}\\,\\bigr\\\}s∗\\displaystyle s^\{\*\}=arg⁡maxs∈ℳreq⁡m\.s\\displaystyle=\\arg\\max\_\{s\\in\\mathcal\{M\}\_\{\\text\{req\}\}\}m\.sθ\\displaystyle\\theta=\{high:40%,critical:60%\}\\displaystyle=\\\{\\text\{high\}:40\\%,\\;\\text\{critical\}:60\\%\\\}trigger​\(m\)\\displaystyle\\text\{trigger\}\(m\)=𝟙\[m\.s∗\>θhigh\]\\displaystyle=\\mathds\{1\}\\bigl\[m\.s^\{\*\}\>\\theta\_\{\\text\{high\}\}\\bigr\]analyze​\(m\)\\displaystyle\\text\{analyze\}\(m\)→\(severity bym\.s∗;root cause and fix specific tos∗\)\\displaystyle\\rightarrow\\bigl\(\\text\{severity by \}m\.s^\{\*\};\\ \\text\{root cause and fix specific to \}s^\{\*\}\\bigr\)
When triggered, the tool translates the raw stall percentage into a structured diagnostic: a*root cause*explaining why the stall occurs \(e\.g\., “warps are waiting for global memory operations to complete—400\+ cycle latency with no prefetching”\), followed by ranked*actionable recommendations*tailored to the dominant stall type\. Memory dependency stalls produce guidance such as “use shared memory to cache frequently accessed data” and “implement vectorized loads \(float4\) for better bandwidth”; barrier stalls recommend “use warp\-level primitives \(\_\_shfl\_sync,\_\_ballot\_sync\) instead of\_\_syncthreads\(\)” with a code example showing the correct pattern; scoreboard stalls suggest increasing instruction\-level parallelism or using async copies \(cp\.async\) on Ampere\+; and branch resolving stalls recommend replacing conditionals with arithmetic predication \(e\.g\.,result = cond \* val1 \+ \(1\-cond\) \* val2\)\. This translation is what enables the LLM to act on profiler output without requiring GPU architecture expertise—the tool bridges the gap between “47% memory dependency stalls” \(a number\) and “add shared memory tiling with\_\_syncthreads\(\)between load and compute phases” \(an action\)\. The complete implementation, including the per\-stall\-type root causes and recommendations, is given in Listing[3](https://arxiv.org/html/2606.26453#LST3)\(Appendix[B](https://arxiv.org/html/2606.26453#A2)\)\.

#### 4\.2\.2SASS Binary Analysis Tools

Whilencutools analyze*runtime counters*, SASS tools analyze the*compiled binary itself*\(what hardware instructions were actually emitted bynvcc\)\.ncucan report “tensor core utilization = 0%” but cannot explain whether the kernel lacks tensor core instructions entirely, uses an older generation \(HMMA vs WGMMA\), or delegates to cuBLAS \(whose tensor\-core kernels run as separate library launches, absent from the user kernel’s compiled\.so\)\.KernelProextracts SASS viacuobjdump \-\-dump\-sasson the compiled\.soand counts instruction patterns across five tensor core generations \(WGMMA, HMMA, IMMA, DMMA, BMMA\), memory operations \(LDGSTS, TMA\), and register spilling \(STL/LDL\)\. Full 5\-tuple specifications for both SASS tools are given in Appendix[B](https://arxiv.org/html/2606.26453#A2)\.

WGMMAInstructionToolprovides definitive hardware feature detection\. When no tensor core instructions appear in SASS, it checks for cuBLAS dynamic symbols vianm \-D—distinguishing “no TC instructions because the kernel is element\-wise” from “no TC instructions because cuBLAS handles them at runtime\.” On Hopper targets, if HMMA instructions are found but no WGMMA, it recommends switching to Hopper warpgroup MMA \(expected 1\.3–1\.5×\\timesimprovement\)\. In the VeOmni case study \(Section[5\.4](https://arxiv.org/html/2606.26453#S5.SS4)\) this mechanism flagged 37 candidates that compiled but emitted zero tensor\-core instructions and anHMMA\-with\-register\-spill regression, steering the search toward the winning WGMMA kernel\.

RegisterSpillDetectorcombines SASS evidence \(STL/LDL counts prove spills exist\) withncumetrics to determine whether spills*actually*cause harm\. Not all spills are pathological, so the tool triggers only under three multi\-signal conditions—spill\-induced latency bubbles, bandwidth contention from local\-memory traffic, or register\-bound occupancy that starves the scheduler—avoiding false alarms on harmless spills while catching cases where register pressure is the root bottleneck\. \(Exact thresholds in Appendix[B](https://arxiv.org/html/2606.26453#A2)\.\)

## 5Results and Evaluation

We evaluateKernelProalong three axes—SOTA comparison on KernelBench, controlled ablations of each design component, and production validation on expert\-optimized kernels—and close with a preliminary energy\-aware extension\. Throughout all experiments, a kernel is*numerically correct*if its output satisfiestorch\.allclose\(out, ref, atol=1e\-2, rtol=1e\-3\)\.

### 5\.1Experimental Setup

Benchmark\.KernelBench\(Ouyang et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib26)\)comprises 250 GPU kernel tasks at three difficulty levels: Level 1 \(100 single operators\), Level 2 \(100 fused/composed operators\), and Level 3 \(50 full model architectures\)\. Speedup is measured as wall\-clock time of the PyTorch eager baseline divided by the optimized CUDA kernel\.

Configuration\.Unless otherwise noted, experiments use Claude Sonnet 4\.6 on NVIDIA A100 GPUs \(AWS p4 instances\), with 15 seeds per task \(5 temperatures×\\times3 rounds\), 30 iterations, 2 candidates per iteration, and MCTS search\. A full comparison with concurrent systems \(hardware, metric type, task subsets\) is provided in Appendix[A](https://arxiv.org/html/2606.26453#A1)\.

### 5\.2State\-of\-the\-Art Comparison on KernelBench

Table 4:KernelProvs\. KernelBlaster on KernelBench \(geometric mean speedup, A100\)\. KernelBlaster was the prior SOTA at time of submission\.KernelProoutperforms KernelBlaster on all three levels, with the largest margin on Level 3 \(\+253%\)—full model architectures where MCTS\-guided multi\-step optimization and micro\-profiling tools provide the most actionable guidance\.KernelProachieves complete task coverage: 100/100 on Level 1, 100/100 on Level 2, and 50/50 on Level 3\. Despite KernelBlaster’s use of RL\-style exploration and a persistent knowledge base,KernelProachieves higher speedups with a simpler approach grounded in micro\-profiling tools as expert surrogates\. A comprehensive comparison with five additional concurrent systems \(CudaForge, StitchCUDA, KernelFoundry, AVO, and others\) confirmingKernelPro’s advantage across hardware platforms and metric types is provided in Appendix[A](https://arxiv.org/html/2606.26453#A1)\.

Robustness to outliers\.Since some Level 2 and Level 3 kernels have trivially parallelizable reference implementations \(yielding speedups\>\>100×\\times\), we verify thatKernelPro’s advantage is not driven by a few extreme outliers\. Table[5](https://arxiv.org/html/2606.26453#S5.T5)reports geometric mean speedups under progressively aggressive capping thresholds\.

Table 5:Geometric mean speedup under capping thresholds\. Each task’s speedup is clamped tomin⁡\(speedup,cap\)\\min\(\\text\{speedup\},\\text\{cap\}\)before computing the geometric mean\. Results demonstrate broad\-based gains across all levels\.Level 1 speedups show minimal sensitivity to capping \(maximum task speedup is 21\.4×\\times; geometric mean is unchanged through Cap@20\), confirming uniformly distributed gains\. Level 3 retains 4\.31×\\timeseven at Cap@10, demonstrating that its strong performance is broad\-based rather than driven by outliers\. Level 2 shows the largest sensitivity to capping \(4\.69×\\times→\\to3\.25×\\timesat Cap@10\) due to several reference implementations with extreme inefficiency; even under Cap@5,KernelProstill achieves 2\.61×\\times—well above all competing systems\.

### 5\.3Component Ablation

We isolate the contribution of eachKernelProdesign component through controlled ablations on 42 KernelBench tasks \(L1: 19, L2: 16, L3: 7\) selected as a representative subset of transformer\-relevant workloads, enabling controlled ablation at scale across multiple conditions, seeds, and iterations that would be prohibitive on the full 250\-task benchmark\. Each ablation varies a single component while holding all others fixed\. All ablations use the same 42\-task set; the effective sample sizennreported in Wilcoxon tests may be smaller because the test excludes tied pairs and tasks where a condition fails to produce a valid kernel\. Table[6](https://arxiv.org/html/2606.26453#S5.T6)summarizes the results; full experimental details are in the referenced appendices\.

Table 6:Consolidated ablation: independent contribution of eachKernelProcomponent\. Effect size is the ratio of geometric mean speedup \(treatment/control\)\. All use Wilcoxon signed\-rank test \(one\-sided\)\.Three components independently and significantly improve optimization quality \(p<0\.05p<0\.05\)\. We discuss the most informative findings from each ablation below; MCTS provides the second\-largest effect \(\+26%\) by enabling escape from local optima that trap greedy search\. Proactive orchestration \(\+23%\) ensures comprehensive tool coverage versus stochastic reactive invocation\. Search memory shows a positive but non\-significant trend on final speedup \(\+6%,p=0\.181p=0\.181\), while improving early convergence\.

#### 5\.3\.1Raw Metrics Are Harmful

The tool ablation reveals a counterintuitive result: providing rawncumetrics without interpretation performs significantly*worse*than providing no profiling feedback at all \(geometric mean 1\.77×\\timesvs 3\.35×\\times, Wilcoxonp=0\.0007p=0\.0007\)\. This condition passes∼\{\\sim\}50 raw hardware counters as name\-value pairs—comparable to the approach used by CudaForge\(Zhang et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib37)\)\. Unstructured metric dumps distract the LLM: it attempts to optimize counters that do not translate to wall\-clock improvements, or misinterprets raw values without roofline context\.KernelPro’s full pipeline \(4\.00×\\times\) outperforms both controls, confirming that*interpretation*—not data access—drives optimization quality\. The Task 41 trajectory \(Appendix[G\.1](https://arxiv.org/html/2606.26453#A7.SS1)\) illustrates this concretely: the tools translate the raw counters \(15\.88 sectors/request, 8\.7 GB read traffic, 244 registers/thread\) into concrete directives—fuse the epilogue, switch to BF16 tensor cores—that produce a2\.8×2\.8\\timessingle\-step improvement; the same counters presented unstructured carry no such actionable directive\. This finding generalizes beyond KernelBench: on 6 representative Triton inference kernels \(Appendix[C\.3](https://arxiv.org/html/2606.26453#A3.SS3)\), tools achieve 36% higher speedup in 31% fewer iterations\.

Table 7:Tool ablation: correctness and speedup across 42 KernelBench tasks\. Unsolved tasks scored as 1\.0×\\times\.Table[8](https://arxiv.org/html/2606.26453#S5.T8)reports pairwise Wilcoxon signed\-rank tests confirming the statistical ordering Ours\>\>No Feedback\>\>Rawncu\. Three insights emerge: \(1\) the interpretation gap \(Raw→\\toOurs\) is larger than the data gap \(None→\\toOurs\), confirming that*how*profiling data is presented matters more than*whether*it is present; \(2\) Rawncuis significantly*worse*than No Feedback \(p=0\.0007p=0\.0007\), meaning unstructured metrics actively degrade optimization quality—a direct challenge to the design of systems like CudaForge that pass raw counters; \(3\)n=39n=39–4040out of 42 tasks show non\-zero differences, indicating the effect is broad\-based rather than driven by a few outlier tasks\.

Table 8:Pairwise Wilcoxon signed\-rank tests \(one\-sided,H1H\_\{1\}: first\>\>second\)\.
#### 5\.3\.2Which Tools Drive the Largest Gains?

Per\-tool impact analysis across 42 tasks reveals which tools most reliably translate into optimization gains\. Table[9](https://arxiv.org/html/2606.26453#S5.T9)summarizes tool coverage and effectiveness; the*hit rate*—fraction of fires producing\>\>1\.5×\\timesimprovement—measures how reliably the LLM acts on each tool’s guidance\.

Table 9:Per\-tool impact across 42 KernelBench tasks\. “Fires” counts optimization turns where the tool’s recommendation appears; “Hit Rate” is the fraction producing\>\>1\.5×\\timesimprovement\.The tensor core switch is the single highest\-impact action: for compute\-bound kernels where the LLM generates CUDA\-core\-only code, the TensorCoreUnderutilizationTool’s recommendation to use cuBLAS/CUTLASS produces an average6\.25×\\timesgain across 9 tasks \(max 9\.6×\\times\)—Task 41’s4\.11×→11\.49×4\.11\\times\\to 11\.49\\timessingle\-iteration leap to a fused BF16 CUTLASS GEMM \(Appendix[G\.1](https://arxiv.org/html/2606.26453#A7.SS1)\) exemplifies this pattern\. RegisterSpillDetector has the highest per\-fire reliability \(18\.2%\), because its diagnostic \(“168 registers/thread, spilling to local memory”\) is immediately actionable—the LLM can directly restructure loops to reduce live variables\. Memory optimization tools \(coalescing, vectorization\) fire most broadly \(37–41 of 42 tasks\) but convert at lower rates \(4–6%\), reflecting the difficulty of non\-trivial data layout changes\.

#### 5\.3\.3MCTS Escapes Local Optima

The 26% aggregate improvement from MCTS masks dramatic per\-task effects\. On tasks with multi\-modal optimization landscapes, MCTS achieves up to10\.0×\\timesthe speedup of greedy search on the same task \(L1 task 7: MCTS 25\.0×\\timesvs greedy 2\.5×\\times; L2 task 33: MCTS 8\.0×\\timesvs greedy 1\.0×\\times; L3 task 31: MCTS 25\.0×\\timesvs greedy 4\.2×\\times\)\. In these cases, greedy converges to a correct but suboptimal solution and cannot escape its local basin, while MCTS explores alternative subtrees that discover fundamentally superior optimization paths \(e\.g\., switching from shared\-memory tiling to warp\-shuffle reduction\)\. Greedy wins on 8 tasks—typically unimodal landscapes where depth\-first exploitation suffices and MCTS wastes budget on unnecessary exploration\.

### 5\.4Production Validation: VeOmni MoE

We further applyKernelProtoVeOmni, a production MoE training stack whose grouped\-GEMM kernels are hand\-optimized in Triton by expert engineers\. We target thedW1weight\-gradient kernel \(group\_gemm\_same\_mn\) from GPT\-OSS\-120B’sgate\_up\_projbackward pass: per experteeit computesdW1e=Xe⊤​Ge\\mathrm\{dW1\}\_\{e\}=X\_\{e\}^\{\\top\}G\_\{e\}\(M=2880M\{=\}2880,N=5760N\{=\}5760,E=128E\{=\}128\), with a*ragged*contractionKeK\_\{e\}that varies per expert under top\-4 routing \(totalR=∑eKe=57,768R=\\sum\_\{e\}K\_\{e\}=57\{,\}768rows; per\-expert load from a measured\-routing imbalance model, Appendix[H](https://arxiv.org/html/2606.26453#A8)\)\. This is harder than a forward GEMM—the imbalance falls on the*reduction*axis, so each expert accumulates over a variable\-lengthKeK\_\{e\}while writing a fixed\[M,N\]\[M,N\]tile\.

On a single H100,KernelProachieved a1\.23×1\.23\\timesspeedupover the expert\-tuned Triton baseline; beyond this best kernel, it found16 other distinct correct kernelsthat also beat the baseline\. Table[10](https://arxiv.org/html/2606.26453#S5.T10)traces the MCTS search climbing from a0\.07×0\.07\\timesfirst attempt to1\.23×1\.23\\timesover 18 iterations, each leap driven by the bottleneckKernelPro’s profiling stage surfaced at the prior node\. The gain comes not from tensor cores per se—the Triton baseline already issues Hopper warp\-group MMA—but from shape\-specific tuning of the memory path and occupancy for this transposed, ragged\-KKlayout\.

Table 10:KernelProsearch progression on the VeOmnidW1weight\-gradient kernel \(H100\)\. Bottlenecks are those flagged byKernelPro’s Stage\-2 profiling at each milestone\.##### Native CuTe \+ raw\-CUDA generation\.

This kernel is the clearest production demonstration ofKernelPro’s headline claim \(Contribution 4\): autonomous*source\-level*CuTe code generation, not high\-level template instantiation or library stitching\. The winning solution containszero library calls—no cuBLAS, no CUTLASS template—and is instead a from\-scratch Hopper kernel thatKernelProcomposed by searching the CUTLASS/CuTe codebase during generation\. At the source level it assembles: anSM90\_64x128x16WGMMAwarp\-group MMA atom driven bycute::gemmwith explicitwarpgroup\_arrive/commit\_batch/waitfences; afour\-stagecp\.asyncsoftware pipelineprefetching tiles into shared memory;128\-byte\-swizzled shared memory\(GMMA::Layout\_MN\_SW128\_Atom\) for bank\-conflict\-free WGMMA operand reads; andCuTe layout\-algebra predication\(make\_identity\_tensorwithcopy\_if\) to mask the ragged per\-expertKeK\_\{e\}boundaries\. Notably,KernelProreasoned explicitly that a single TMA descriptor would mis\-read neighboring experts’ rows under raggedKKand*chose*cp\.asyncover TMA—an architectural trade\-off normally reserved for human kernel engineers\. Listing[1](https://arxiv.org/html/2606.26453#LST1)shows the resulting WGMMA mainloop\.

1

2usingAtom=typenameMmaSel<TA\>::type;

3TiledMMAmma=make\_tiled\_mma\(Atom\{\}\);

4\.\.\.

5for\(intkt=0;kt<k\_tiles;\+\+kt\)\{

6cp\_async\_wait<NS\-2\>\(\);\_\_syncthreads\(\);

7intkload=kt\+\(NS\-1\);

8if\(kload<k\_tiles\)\{

9fillA\(kload\);copy\_if\(copyA,tApA,tAgA\(\_,\_,\_,kload\),tAsA\(\_,\_,\_,smem\_pipe\_write\)\);

10fillB\(kload\);copy\_if\(copyB,tBpB,tBgB\(\_,\_,\_,kload\),tBsB\(\_,\_,\_,smem\_pipe\_write\)\);

11\}

12cp\_async\_fence\(\);smem\_pipe\_write=\(smem\_pipe\_write\+1\)%NS;

13warpgroup\_arrive\(\);

14cute::gemm\(mma,tCrA\(\_,\_,\_,smem\_pipe\_read\),tCrB\(\_,\_,\_,smem\_pipe\_read\),tCrC\);

15warpgroup\_commit\_batch\(\);warpgroup\_wait<0\>\(\);

16smem\_pipe\_read=\(smem\_pipe\_read\+1\)%NS;

17\}

Listing 1:KernelPro\-generated WGMMA mainloop for the VeOmnidW1weight\-gradient kernel \(excerpt\)\. No library GEMM is invoked; the matrix multiply is a hand\-written SM90 warp\-group MMA\.
##### SASS analysis in the trajectory\.

This run also exercisesKernelPro’sinstruction\-levelSASS tools \(Section[4\.2\.2](https://arxiv.org/html/2606.26453#S4.SS2.SSS2)\) in a way kernel\-levelncucannot\. Because the LLM frequently emits code it*believes*uses warp\-group tensor cores, theWGMMAInstructionTooldisassembles each compiled\.soand reports the tensor\-core instruction mix directly\. Across the search it flagged37 candidatesthat compiled and ran but emitted zero tensor\-core instructions \(Found 0 TC ops \(WGMMA:0 HMMA:0\)\)—scalar fallbacks that no utilization metric would distinguish from a slow\-but\-real tensor\-core kernel—and surfaced generation regressions such as a candidate using the older AmpereHMMAinstructions with250 register\-spill instructions\(16 TC ops \(HMMA:16\), STL:125 LDL:125\)\. The decisive late iterations were in turn driven by an occupancy diagnosis—KernelPro’s profiling reported the kernel*latency\-bound at18\.8%18\.8\\%warp occupancy, limited by168168registers/thread and shared\-memory pressure*—which steered the search toward the smaller\-footprint, higher\-occupancy64×128×1664\{\\times\}128\{\\times\}16configuration that crossed the Triton baseline\. The optimization is thus guided by instruction\- and occupancy\-level evidence, not runtime counters alone\. The completeKernelPro\-generated kernel—mainloop, prologue/epilogue, and host launch—is provided in Appendix[H](https://arxiv.org/html/2606.26453#A8); Listing[1](https://arxiv.org/html/2606.26453#LST1)shows only the WGMMA mainloop excerpt\.

### 5\.5Energy\-Aware Kernel Optimization: A Preliminary Study

Power is increasingly the limiting factor in AI scaling\(International Energy Agency,[2024](https://arxiv.org/html/2606.26453#bib.bib14)\), and GPU energy is dominated by data movement rather than compute\(Horowitz,[2014](https://arxiv.org/html/2606.26453#bib.bib12)\)\. This makes energy a first\-class measure of kernel efficiency—yet existing LLM\-based kernel agents target latency alone\.KernelProis the first CUDA kernel coding agent to optimize energy efficiency; we explore whether its profiling\-driven search can reduce kernel energy, and where that reduction is separable from speedup\.

KernelPro’s design extends naturally beyond speedup\. We treat energy efficiency as a*secondary*objective under strict lexicographic priority: speedup always wins, and energy is optimized only to break ties among speed\-equivalent solutions, via a single energy term added to the reward:

r​\(n\)=ln⁡\(speedup​\(n\)\)\+ε⋅ln⁡\(energy\_reduction​\(n\)\),ε≪1,r\(n\)=\\ln\(\\text\{speedup\}\(n\)\)\+\\varepsilon\\cdot\\ln\(\\text\{energy\\\_reduction\}\(n\)\),\\quad\\varepsilon\\ll 1,\(3\)for correct solutions \(ε\\varepsilonsmall enough that any speedup gain dominates any energy gain\)\. The search, tools, and pipeline are otherwise unchanged\.

##### Kernel\-level energy model\.

GPU energy is dominated by data movement, not compute: a DRAM access costs∼\\sim100–200×\\timesa register access\(Horowitz,[2014](https://arxiv.org/html/2606.26453#bib.bib12)\), and the register file is the single largest dynamic\-power component on modern GPUs\(Kandiah et al\.,[2021](https://arxiv.org/html/2606.26453#bib.bib16)\)\. Device\-level power sensors \(nvidia\-smi\) are too coarse and noisy for per\-kernel attribution on A100/H100\(Yang et al\.,[2024](https://arxiv.org/html/2606.26453#bib.bib34)\), so we instead estimate energy deterministically fromncucountersKernelProalready collects, as a weighted cost over four documented bottlenecks, with weights set by the per\-operation energy hierarchy of Horowitz\(Horowitz,[2014](https://arxiv.org/html/2606.26453#bib.bib12)\)and AccelWattch\(Kandiah et al\.,[2021](https://arxiv.org/html/2606.26453#bib.bib16)\):

Eproxy=20​pJ/B⋅BDRAM⏟B1: DRAM bytes\+100​pJ⋅Sld⋅max⁡\(0,r4−1\)⏟B2: uncoalesced\+80​pJ⋅Slocal⏟B3: spill traffic\+3​pJ⋅I⏟B4: instructionsE\_\{\\text\{proxy\}\}=\\underbrace\{20\\,\\text\{pJ/B\}\\cdot B\_\{\\text\{DRAM\}\}\}\_\{\\text\{B1: DRAM bytes\}\}\+\\underbrace\{100\\,\\text\{pJ\}\\cdot S\_\{\\text\{ld\}\}\\\!\\cdot\\\!\\max\(0,\\tfrac\{r\}\{4\}\\\!\-\\\!1\)\}\_\{\\text\{B2: uncoalesced\}\}\+\\underbrace\{80\\,\\text\{pJ\}\\cdot S\_\{\\text\{local\}\}\}\_\{\\text\{B3: spill traffic\}\}\+\\underbrace\{3\\,\\text\{pJ\}\\cdot I\}\_\{\\text\{B4: instructions\}\}\(4\)whereBDRAMB\_\{\\text\{DRAM\}\}is DRAM bytes read\+written,SldS\_\{\\text\{ld\}\}global\-load sectors,rrsectors\-per\-request,SlocalS\_\{\\text\{local\}\}local\-memory \(spill\) sectors, andIIinstructions executed\. The resultingenergy\_reductionratio \(baseline/candidate cost\) is intrinsic to the kernel and independent of wall\-clock time, so it rewards reductions in hardware activity that timing alone cannot see\. To confirm that proxy\-driven choices translate to real device energy, we separately measure winning kernels in millijoules using a rigorous locked\-clock, idle\-subtracted NVML protocol\. We also add four*energy\-only*micro\-profiling tools that flag latency\-hidden energy waste \(bank conflicts, register spills, uncoalesced access, redundant barriers\) the speed\-only tools ignore\. Appendix[I](https://arxiv.org/html/2606.26453#A9)details the NVML protocol, the energy model, the four tools, and the literature mapping each to a documented bottleneck\.

To isolate the energy\-aware reward from speedup, we run paired arms \(i\.e\., speed\-only versus energy\-aware\), and measure the winning kernels in millijoules under the rigorous protocol of Appendix[I\.1](https://arxiv.org/html/2606.26453#A9.SS1)\. Most energy savings coincide with speed savings—reducing data movement cuts both—so the speed\-only search already captures them; the interesting case is where they decouple\. On theSwishactivation \(KernelBench L1, pure elementwise, memory\-bound\), both arms reach an identical2\.52×2\.52\\timesspeedup with identical memory traffic \(1 vectorized load \+ 1 store per thread\), yet the energy\-aware kernel draws1802 mJ versus 2038 mJ—11\.6% less dynamic energy\(Table[11](https://arxiv.org/html/2606.26453#S5.T11)\)\. The mechanism is visible in SASS: the energy\-aware kernel compiles to56 instructions versus 216for the same computation, having chosen the fast reciprocal intrinsic \(\_\_fdividef, oneMUFU\) over an IEEE\-accurate division \(32FFMA\+ branch\-heavy refinement\) and a read\-only cache load \(\_\_ldg\)\. Because the kernel is bandwidth\-bound, those∼\\sim160 extra instructions hide under memory stalls and cost*no*latency—but the SMs still execute and draw power for them\. The speed\-only objective is blind to this; the energy term is what selects the cheaper instruction stream\. This is a concrete instance of energy reducible at fixed speed; a systematic characterization is left to future work\.

Table 11:Energy\-aware matched\-speed A/B on theSwishkernel \(KernelBench L1\)\. Both arms reach identical speedup and memory traffic; dynamic energy is measured under the locked\-clock, idle\-subtracted NVML protocol \(Appendix[I\.1](https://arxiv.org/html/2606.26453#A9.SS1)\)\. SASS instruction counts are for the vectorized kernel\.

## 6Conclusion

KernelProdemonstrates that the key bottleneck in LLM\-based GPU kernel optimization is not code generation capability but*how profiling data reaches the LLM*\. By encoding expert heuristics as pluggable micro\-profiling tools—executable analysis functions that transform raw hardware counters into actionable natural language guidance—KernelProreplaces the implicit reasoning required by prior systems with explicit, systematic diagnostic workflows\. Controlled ablations establish a clear contribution hierarchy: micro\-profiling tools provide the largest gain \(125% over raw metrics\), with the counterintuitive finding that rawncumetrics*actively harm*optimization quality relative to no feedback at all—directly challenging systems that pass unstructured counters to LLMs; domain\-adapted MCTS and proactive tool orchestration each contribute significant further gains\. Together these yield geometric mean speedups of2\.42×2\.42\\times/4\.69×4\.69\\times/5\.30×5\.30\\timeson KernelBench Levels 1/2/3 with 100% task coverage, surpassing the prior SOTA on all levels\. Beyond academic benchmarks,KernelProachieves 1\.23×\\timesover an expert\-optimized Triton baseline on a production MoE weight\-gradient kernel \(H100\), generating a from\-scratch raw\-CUDA\+CuTe Hopper WGMMA kernel, not a library call\.KernelPro’s CUTLASS/CuTe integration—with LLM\-autonomous code search and architecture constraint mining—further enables direct source\-level tensor core code generation, grounding the LLM in authoritative library source rather than memorized patterns \(Appendix[G](https://arxiv.org/html/2606.26453#A7)presents two complete trajectories\)\. Finally, as the first CUDA kernel coding agent to optimize energy efficiency,KernelProreduces measured energy by11\.6%11\.6\\%at matched speed via lower\-energy instruction selection—a preliminary result whose systematic characterization we leave to future work\.KernelProhas submitted 6 PRs to FlashInfer \(1 merged, 5 under review\), and scaling this into a sustained PR contribution pipeline is a key future direction\.

## References

- Andrews & Witteveen \(2025\)Martin Andrews and Sam Witteveen\.GPU kernel scientist: An LLM\-driven framework for iterative kernel optimization\.In*ES\-FoMo III Workshop at ICML*, 2025\.URL[https://arxiv\.org/abs/2506\.20807](https://arxiv.org/abs/2506.20807)\.
- Auger et al\. \(2013\)David Auger, Adrien Couetoux, and Olivier Teytaud\.Continuous upper confidence trees with polynomial exploration – consistency\.In*Proceedings of the European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases \(ECML\-PKDD\)*, pp\. 194–209, 2013\.doi:10\.1007/978\-3\-642\-40988\-2\_13\.
- Bai et al\. \(2024\)Yushi Bai, Shangqing Tu, Jiajie Zhang, Hao Peng, Xiaozhi Wang, Xin Lv, Shulin Cao, Jiazheng Xu, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li\.LongBench v2: Towards deeper understanding and reasoning on realistic long\-context multitasks\.*arXiv preprint arXiv:2412\.15204*, 2024\.URL[https://arxiv\.org/abs/2412\.15204](https://arxiv.org/abs/2412.15204)\.
- Bhattacharjee et al\. \(2026\)Arijit Bhattacharjee, Heng Ping, Son Vu Le, Paul Bogdan, Nesreen K\. Ahmed, and Ali Jannesari\.OptiML: An end\-to\-end framework for program synthesis and CUDA kernel optimization\.*arXiv preprint arXiv:2602\.12305*, 2026\.URL[https://arxiv\.org/abs/2602\.12305](https://arxiv.org/abs/2602.12305)\.
- Chatterjee et al\. \(2017\)Niladrish Chatterjee, Mike O’Connor, Donghyuk Lee, Daniel R Johnson, Stephen W Keckler, Minsoo Rhu, and William J Dally\.Architecting an energy\-efficient DRAM system for GPUs\.In*2017 IEEE International Symposium on High Performance Computer Architecture \(HPCA\)*, pp\. 73–84\. IEEE, 2017\.
- Chen et al\. \(2026\)Terry Chen, Zhifan Ye, Bing Xu, Zihao Ye, Timmy Liu, Ali Hassani, Tianqi Chen, Andrew Kerr, Haicheng Wu, Yang Xu, Yu\-Jung Chen, Hanfeng Chen, Aditya Kane, Ronny Krashinsky, Ming\-Yu Liu, Vinod Grover, Luis Ceze, Roger Bringmann, John Tran, Wei Liu, Fung Xie, Michael Lightstone, and Humphrey Shi\.AVO: Agentic variation operators for autonomous evolutionary search\.In*arXiv preprint arXiv:2603\.24517*, 2026\.URL[https://arxiv\.org/abs/2603\.24517](https://arxiv.org/abs/2603.24517)\.
- Chen et al\. \(2025\)Yongchao Chen, Yueying Li, Yue Zhang, Shreyas Singh, Tian Lan, and Yongle Zhang\.cuPilot: A strategy\-coordinated multi\-agent framework for CUDA kernel evolution\.In*arXiv preprint arXiv:2512\.16465*, 2025\.URL[https://arxiv\.org/abs/2512\.16465](https://arxiv.org/abs/2512.16465)\.
- Dong et al\. \(2026\)Kris Shengjun Dong, Sahil Modi, Dima Nikiforov, Sana Damani, Edward Lin, Siva Kumar Sastry Hari, and Christos Kozyrakis\.KernelBlaster: Continual cross\-task CUDA optimization via memory\-augmented in\-context reinforcement learning\.In*arXiv preprint arXiv:2602\.14293*, 2026\.URL[https://arxiv\.org/abs/2602\.14293](https://arxiv.org/abs/2602.14293)\.
- Han et al\. \(2026\)Yuxuan Han, Meng\-Hao Guo, Zhengning Liu, Wenguang Chen, and Shi\-Min Hu\.Making LLMs optimize multi\-scenario CUDA kernels like experts\.In*arXiv preprint arXiv:2603\.07169*, 2026\.URL[https://arxiv\.org/abs/2603\.07169](https://arxiv.org/abs/2603.07169)\.
- Hollander et al\. \(2014\)Myles Hollander, Douglas A\. Wolfe, and Eric Chicken\.*Nonparametric Statistical Methods*\.John Wiley & Sons, 3rd edition, 2014\.
- Hong & Kim \(2010\)Sunpyo Hong and Hyesoon Kim\.An integrated GPU power and performance model\.In*Proceedings of the 37th Annual International Symposium on Computer Architecture \(ISCA\)*, pp\. 280–289, 2010\.
- Horowitz \(2014\)Mark Horowitz\.1\.1 computing’s energy problem \(and what we can do about it\)\.In*2014 IEEE International Solid\-State Circuits Conference Digest of Technical Papers \(ISSCC\)*, pp\. 10–14\. IEEE, 2014\.
- Hou et al\. \(2025\)Zhenyu Hou, Ziniu Hu, Yujiang Li, Rui Lu, Jie Tang, and Yuxiao Dong\.TreeRL: LLM reinforcement learning with on\-policy tree search\.*arXiv preprint arXiv:2506\.11902*, 2025\.URL[https://arxiv\.org/abs/2506\.11902](https://arxiv.org/abs/2506.11902)\.
- International Energy Agency \(2024\)International Energy Agency\.Electricity 2024: Analysis and forecast to 2026\.Technical report, International Energy Agency \(IEA\), Paris, 2024\.URL[https://www\.iea\.org/reports/electricity\-2024](https://www.iea.org/reports/electricity-2024)\.
- Ji et al\. \(2025\)Yuxiang Ji, Ziyu Ma, Yong Wang, Guanhua Chen, Xiangxiang Chu, and Liaoni Wu\.Tree search for LLM agent reinforcement learning\.In*arXiv preprint arXiv:2509\.21240*, 2025\.URL[https://arxiv\.org/abs/2509\.21240](https://arxiv.org/abs/2509.21240)\.
- Kandiah et al\. \(2021\)Vijay Kandiah, Scott Peverelle, Mahmoud Khairy, Junrui Pan, Amogh Manjunath, Timothy G Rogers, Tor M Aamodt, and Nikos Hardavellas\.Accelwattch: A power modeling framework for modern gpus\.In*MICRO\-54: 54th Annual IEEE/ACM International Symposium on Microarchitecture*, pp\. 738–753, 2021\.
- Kocsis & Szepesvári \(2006\)Levente Kocsis and Csaba Szepesvári\.Bandit based Monte\-Carlo planning\.In*Proceedings of the 17th European Conference on Machine Learning \(ECML\)*, pp\. 282–293, 2006\.doi:10\.1007/11871842\_29\.
- Kwon et al\. \(2023\)Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E\. Gonzalez, Hao Zhang, and Ion Stoica\.Efficient memory management for large language model serving with PagedAttention\.In*Proceedings of the 29th ACM Symposium on Operating Systems Principles \(SOSP\)*, 2023\.URL[https://arxiv\.org/abs/2309\.06180](https://arxiv.org/abs/2309.06180)\.
- Leng et al\. \(2013\)Jingwen Leng, Tayler Hetherington, Ahmed ElTantawy, Syed Gilani, Nam Sung Kim, Tor M Aamodt, and Vijay Janapa Reddi\.GPUWattch: Enabling energy optimizations in GPGPUs\.In*Proceedings of the 40th Annual International Symposium on Computer Architecture \(ISCA\)*, pp\. 487–498, 2013\.
- Li et al\. \(2026\)Shiyang Li, Zijian Zhang, Winson Chen, Yuebo Luo, Mingyi Hong, and Caiwen Ding\.StitchCUDA: An automated multi\-agents end\-to\-end GPU programing framework with rubric\-based agentic reinforcement learning\.In*Proceedings of the 43rd International Conference on Machine Learning \(ICML\)*, Proceedings of Machine Learning Research\. PMLR, 2026\.URL[https://arxiv\.org/abs/2603\.02637](https://arxiv.org/abs/2603.02637)\.
- Nayak & Basu \(2024\)Ajay Nayak and Arkaprava Basu\.Over\-synchronization in GPU programs\.In*57th IEEE/ACM International Symposium on Microarchitecture \(MICRO\)*\. IEEE, 2024\.
- Novikov et al\. \(2025\)Alexander Novikov, Ngân V u, Marvin Eisenberger, Emilien Dupont, Po\-Sen Huang, Adam Zsolt Wagner, Sergey Shirobokov, Borislav Kozlovskii, Francisco J\. R\. Ruiz, Abbas Mehrabian, M\. Pawan Kumar, Abigail See, Swarat Chaudhuri, George Holland, Alex Davies, Sebastian Nowozin, Pushmeet Kohli, and Matej Balog\.AlphaEvolve: A coding agent for scientific and algorithmic discovery\.*arXiv preprint arXiv:2506\.13131*, 2025\.URL[https://arxiv\.org/abs/2506\.13131](https://arxiv.org/abs/2506.13131)\.
- NVIDIA \(2023\)NVIDIA\.CUTLASS: CUDA templates for linear algebra subroutines\.[https://github\.com/NVIDIA/cutlass](https://github.com/NVIDIA/cutlass), 2023\.
- NVIDIA Corporation \(2024\)NVIDIA Corporation\.CUDA C\+\+ Best Practices Guide, 2024\.URL[https://docs\.nvidia\.com/cuda/cuda\-c\-best\-practices\-guide/](https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/)\.
- OpenAI \(2025\)OpenAI\.OpenAI Agents SDK\.[https://github\.com/openai/openai\-agents\-python](https://github.com/openai/openai-agents-python), 2025\.
- Ouyang et al\. \(2025\)Anne Ouyang, Simon Guo, Simran Arora, Alex L\. Zhang, William Hu, Christopher Ré, and Azalia Mirhoseini\.KernelBench: Can LLMs write efficient GPU kernels?In*arXiv preprint arXiv:2502\.10517*, 2025\.URL[https://arxiv\.org/abs/2502\.10517](https://arxiv.org/abs/2502.10517)\.
- Salemi et al\. \(2026\)Alireza Salemi, Chang Zeng, Atharva Nijasure, Jui\-Hui Chung, Razieh Rahimi, Fernando Diaz, and Hamed Zamani\.GrepSeek: Training search agents for direct corpus interaction\.*arXiv preprint arXiv:2605\.29307*, 2026\.URL[https://arxiv\.org/abs/2605\.29307](https://arxiv.org/abs/2605.29307)\.
- Spector et al\. \(2024\)Benjamin Spector, Aaryan Singhal, Simran Arora, and Chris Re\.GPUs Go Brrr, 2024\.URL[https://hazyresearch\.stanford\.edu/blog/2024\-05\-12\-tk](https://hazyresearch.stanford.edu/blog/2024-05-12-tk)\.Hazy Research Blog, introducing ThunderKittens\.
- Team & Meta \(2025\)Ansor Team and Meta\.KernelEvolve: Scaling agentic kernel coding for heterogeneous AI accelerators at meta\.In*arXiv preprint arXiv:2512\.23236*, 2025\.URL[https://arxiv\.org/abs/2512\.23236](https://arxiv.org/abs/2512.23236)\.
- Tillet et al\. \(2019\)Philippe Tillet, H\. T\. Kung, and David Cox\.Triton: An intermediate language and compiler for tiled neural network computations\.In*Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages \(MAPL\)*, 2019\.
- Wiedemann et al\. \(2026\)Nina Wiedemann, Quentin Leboutet, Michael Paulitsch, Diana Wofk, and Benjamin Ummenhofer\.KernelFoundry: Hardware\-aware evolutionary GPU kernel optimization\.In*arXiv preprint arXiv:2603\.12440*, 2026\.URL[https://arxiv\.org/abs/2603\.12440](https://arxiv.org/abs/2603.12440)\.
- Wu et al\. \(2025\)Fang Wu, Weihao Xuan, Heli Qi, Ximing Lu, Aaron Tu, Li Erran Li, and Yejin Choi\.DeepSearch: Overcome the bottleneck of reinforcement learning with verifiable rewards via Monte Carlo tree search\.In*arXiv preprint arXiv:2509\.25454*, 2025\.URL[https://arxiv\.org/abs/2509\.25454](https://arxiv.org/abs/2509.25454)\.
- Yang et al\. \(2020\)Charlene Yang, Thorsten Kurth, and Samuel Williams\.Hierarchical roofline analysis for GPUs: Accelerating performance optimization for the NERSC\-9 Perlmutter system\.*Concurrency and Computation: Practice and Experience*, 32\(20\):e5547, 2020\.doi:10\.1002/cpe\.5547\.
- Yang et al\. \(2024\)Zeyu Yang, Karel Adámek, and Wesley Armour\.Part\-time power measurements: nvidia\-smi’s lack of attention\.*arXiv preprint arXiv:2312\.02741*, 2024\.
- Yao et al\. \(2023\)Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik R\. Narasimhan, and Yuan Cao\.ReAct: Synergizing reasoning and acting in language models\.In*International Conference on Learning Representations \(ICLR\)*, 2023\.URL[https://arxiv\.org/abs/2210\.03629](https://arxiv.org/abs/2210.03629)\.
- Zhang et al\. \(2024\)Di Zhang, Jianbo Wu, Jingdi Lei, Tong Che, Jiatong Li, Tong Xie, Xiaoshui Huang, Shufei Zhang, Marco Pavone, Yuqiang Li, Wanli Ouyang, and Dongzhan Zhou\.Accessing GPT\-4 level mathematical olympiad solutions via Monte Carlo Tree Self\-refine with LLaMa\-3 8B\.*arXiv preprint arXiv:2406\.07394*, 2024\.URL[https://arxiv\.org/abs/2406\.07394](https://arxiv.org/abs/2406.07394)\.
- Zhang et al\. \(2025\)Zijian Zhang, Rong Wang, Shiyang Li, Yuebo Luo, Mingyi Hong, and Caiwen Ding\.CudaForge: An agent framework with hardware feedback for CUDA kernel optimization\.In*arXiv preprint arXiv:2511\.01884*, 2025\.URL[https://arxiv\.org/abs/2511\.01884](https://arxiv.org/abs/2511.01884)\.
- Zhang et al\. \(2026\)Zijian Zhang, Shiyang Li, Rong Wang, Yuebo Luo, Mingyi Hong, and Caiwen Ding\.CUDA Agent: Large\-scale agentic RL for high\-performance CUDA kernel generation\.In*arXiv preprint arXiv:2602\.24286*, 2026\.URL[https://arxiv\.org/abs/2602\.24286](https://arxiv.org/abs/2602.24286)\.
- Zhou et al\. \(2021\)Keren Zhou, Xiaozhu Meng, Ryuichi Sai, and John Mellor\-Crummey\.GPA: A GPU performance advisor based on instruction sampling\.In*Proceedings of the 2021 IEEE/ACM International Symposium on Code Generation and Optimization \(CGO\)*, pp\. 115–125, 2021\.URL[https://arxiv\.org/abs/2009\.04061](https://arxiv.org/abs/2009.04061)\.

## Appendix AComparison with Concurrent Work

Table[12](https://arxiv.org/html/2606.26453#A1.T12)provides an overview of concurrent LLM\-based GPU kernel optimization systems: KernelBlaster\(Dong et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib8)\), CudaForge\(Zhang et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib37)\), StitchCUDA\(Li et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib20)\), KernelFoundry\(Wiedemann et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib31)\), and AVO\(Chen et al\.,[2026](https://arxiv.org/html/2606.26453#bib.bib6)\)\. Cross\-system comparisons warrant several caveats\. First, although speedups are all measured relative to a PyTorch baseline on the same GPU, different architectures offer different optimization ceilings \(e\.g\., Hopper tensor cores vs\. Ampere\), so absolute ratios are not directly comparable across hardware\. Second, the systems differ in which KernelBench tasks they evaluate on; for instance, StitchCUDA trains on 80% of KernelBench and evaluates on the remaining 20%, so its numbers reflect in\-distribution performance rather than zero\-shot evaluation on the full benchmark\. Third, the reported metric differs:KernelProand KernelBlaster report geometric mean speedup, whereas CudaForge reports arithmetic mean and StitchCUDA reports an end\-to\-end average\. We plan to open\-sourceKernelPro’s generated solutions for all 250 KernelBench tasks, pending approval\.

Table 12:Landscape of LLM\-based GPU kernel optimization systems on KernelBench\. Speedup is over PyTorch eager baseline on the same GPU\. “Tasks” denotes the number of test tasks evaluated per level\. Metric type: GeoMean \(G\), ArithMean \(A\)\.‡StitchCUDA uses GRPO\-trained Qwen3\-32B as Coder \(trained on 80% of KernelBench\), GPT\-5\.2 as Planner/Verifier\.

## Appendix BMicro\-Profiling Tool Specifications

Table 13:KernelPromicro\-profiling tool taxonomy and trigger conditions### B\.1Detailed Tool Formalizations

##### WarpStallTool\.

Identifies*why*warps fail to issue by finding the dominant stall reason among fivencuwarp\-stall counters, then emitting a fix tailored to that reason—the clearest instance of diagnostic reasoning over raw numbers\. Theanalyzestep dispatches on the dominant stall: memory\-dependency→\\rightarrowshared\-memory caching and vectorized \(float4\) loads; barrier→\\rightarrowwarp\-level primitives \(\_\_shfl\_sync/\_\_ballot\_sync\) in place of\_\_syncthreads\(\); long\-scoreboard→\\rightarrowasync prefetch \(cp\.asyncon Ampere\+, TMA on Hopper\); branch\-resolving→\\rightarrowarithmetic predication\.

ℳreq\\displaystyle\\mathcal\{M\}\_\{\\text\{req\}\}=\{mem\_dep,short\_scoreboard,long\_scoreboard,\\displaystyle=\\\{\\texttt\{mem\\\_dep\},\\texttt\{short\\\_scoreboard\},\\texttt\{long\\\_scoreboard\},barrier,branch\_resolving\}\(stall % per active warp\)\\displaystyle\\qquad\\texttt\{barrier\},\\texttt\{branch\\\_resolving\}\\\}\\quad\(\\text\{stall \\% per active warp\}\)θ\\displaystyle\\theta=\{high:40%,critical:60%\}\\displaystyle=\\\{\\text\{high\}\{:\}\\,40\\%,\\ \\text\{critical\}\{:\}\\,60\\%\\\}s∗\\displaystyle s^\{\*\}=arg⁡maxs∈ℳreq⁡m\.s\\displaystyle=\\arg\\max\_\{s\\in\\mathcal\{M\}\_\{\\text\{req\}\}\}m\.strigger​\(m\)\\displaystyle\\text\{trigger\}\(m\)=𝟙\[m\.s∗\>θhigh\]\\displaystyle=\\mathds\{1\}\[\\,m\.s^\{\*\}\>\\theta\_\{\\text\{high\}\}\\,\]analyze→\(severity bym\.s∗;rec:fix specific tos∗\)\\displaystyle\\rightarrow\(\\text\{severity by \}m\.s^\{\*\};\\ \\text\{rec\}\{:\}\\ \\text\{fix specific to \}s^\{\*\}\)

##### WGMMAInstructionTool \(SASS\)\.

Tensor core utilization is the single most impactful optimization for matrix\-heavy GPU workloads\. Critically, on Hopper GPUs, onlyWGMMAinstructions can achieve peak throughput–olderHMMAinstructions cap at∼\\sim63% of peak, leaving 37% performance on the table\(Spector et al\.,[2024](https://arxiv.org/html/2606.26453#bib.bib28)\)\. This tool provides definitive tensor core detection by analyzing compiled SASS binary viacuobjdump \-\-dump\-sass, identifying not just*whether*tensor cores are used but*which generation*of instructions, enabling architecture\-specific optimization guidance\. Supports all tensor core instruction variants:

- •WGMMA: Hopper warpgroup MMA \(sm\_90\) – FP16/BF16/FP8/INT8/TF32
- •HMMA: Half\-precision MMA \(sm\_70\+\) – FP16/BF16
- •IMMA/DMMA/BMMA: Specialized MMA – INT8/INT4, FP64, INT1

Additionally detects cuBLAS usage vianm \-Dsymbol lookup \(cuBLAS auto\-selects tensor core kernels on Volta\+ GPUs\)\. The two\-stage approach: \(1\) fastnm \-Dcheck for library calls, \(2\) full SASS disassembly for user\-written kernels\.

ℳreq\\displaystyle\\mathcal\{M\}\_\{\\text\{req\}\}=\{sass\_analysis\}\\displaystyle=\\\{\\texttt\{sass\\\_analysis\}\\\}tc\_count=wgmma\+hmma\+imma\+dmma\+bmma\\displaystyle=\\texttt\{wgmma\}\+\\texttt\{hmma\}\+\\texttt\{imma\}\+\\texttt\{dmma\}\+\\texttt\{bmma\}trigger​\(m\)\\displaystyle\\text\{trigger\}\(m\)=𝟙​\[tc\_count\>0∨uses\_cublas\]\\displaystyle=\\mathds\{1\}\[\\text\{tc\\\_count\}\>0\\lor\\texttt\{uses\\\_cublas\}\]analyze→\(which TC generation is emitted; recommend the\\displaystyle\\rightarrow\(\\text\{which TC generation is emitted; recommend the\}architecture\-appropriate path, e\.g\. WGMMA on Hopper\)\\displaystyle\\qquad\\text\{architecture\-appropriate path, e\.g\.\\ WGMMA on Hopper\}\)

##### RegisterSpillDetector \(SASS\)\.

Detection is multi\-signal: from the SASS binary it countsSTL\(store local\) andLDL\(load local\) instructions, then flags a spill \(stl\_count\+ldl\_count\>0\\texttt\{stl\\\_count\}\+\\texttt\{ldl\\\_count\}\>0\) as harmful only when one of three conditions holds, with the exact cutoffs below:

- •Case A \(Latency\-bound\): long\-scoreboard stall\>15%\>15\\%*and*compute throughput<50%<50\\%*and*memory throughput<50%<50\\%\(spills add latency while neither roofline is saturated\)\.
- •Case B \(Memory\-bound by spills\): memory throughput\>60%\>60\\%*and*local\-memory traffic ratio\>30%\>30\\%\(spills consume bandwidth needed for useful data\)\.
- •Case C \(Occupancy\-limited\): registers are the binding occupancy limit*and*the scheduler is starved \(issue activity<30%<30\\%*or*eligible warps/cycle<0\.6<0\.6\)\.

ℳreq\\displaystyle\\mathcal\{M\}\_\{\\text\{req\}\}=\{stl\_count,ldl\_count,long\_scoreboard,compute\_tput,\\displaystyle=\\\{\\texttt\{stl\\\_count\},\\texttt\{ldl\\\_count\},\\texttt\{long\\\_scoreboard\},\\texttt\{compute\\\_tput\},memory\_tput,local\_traffic\_ratio,occupancy\_limit\_regs,\\displaystyle\\qquad\\texttt\{memory\\\_tput\},\\texttt\{local\\\_traffic\\\_ratio\},\\texttt\{occupancy\\\_limit\\\_regs\},issue\_active,eligible\_warps\}\\displaystyle\\qquad\\texttt\{issue\\\_active\},\\texttt\{eligible\\\_warps\}\\\}trigger​\(m\)\\displaystyle\\text\{trigger\}\(m\)=𝟙​\[stl\_count\+ldl\_count\>0\]∧\(A∨B∨C\)\\displaystyle=\\mathds\{1\}\[\\texttt\{stl\\\_count\}\+\\texttt\{ldl\\\_count\}\>0\]\\;\\land\\;\(\\text\{A\}\\lor\\text\{B\}\\lor\\text\{C\}\)analyze→\(severity by case; rec:reduce register pressure\\displaystyle\\rightarrow\(\\text\{severity by case; rec\}\{:\}\\ \\text\{reduce register pressure\}via\_\_launch\_bounds\_\_,\-maxrregcount, or kernel split\)\\displaystyle\\qquad\\text\{via \{\\\_\\\_launch\\\_bounds\\\_\\\_\}, \{\-maxrregcount\}, or kernel split\}\)

##### TensorCoreUnderutilizationTool\.

Flags compute\-bound kernels that run on CUDA cores instead of tensor cores\. As a 5\-tuple:

ℳreq\\displaystyle\\mathcal\{M\}\_\{\\text\{req\}\}=\{tensor\_core\_util\}\\displaystyle=\\\{\\texttt\{tensor\\\_core\\\_util\}\\\}θ\\displaystyle\\theta=\{low:10%,moderate:50%\}\\displaystyle=\\\{\\text\{low\}\{:\}\\,10\\%,\\ \\text\{moderate\}\{:\}\\,50\\%\\\}trigger​\(m\)\\displaystyle\\text\{trigger\}\(m\)=𝟙\[m\.tensor\_core\_util<θlow\]\\displaystyle=\\mathds\{1\}\[m\.\\texttt\{tensor\\\_core\\\_util\}<\\theta\_\{\\text\{low\}\}\]analyze→\(severity:critical,\\displaystyle\\rightarrow\(\\text\{severity\}\{:\}\\ \\textsc\{critical\},rec:“rewrite the matmul as a cuBLAS/CUTLASS GEMM”\)\\displaystyle\\qquad\\text\{rec\}\{:\}\\ \\text\{\`\`rewrite the matmul as a cuBLAS/CUTLASS GEMM''\}\)

##### LaunchOverheadDetector\.

Flags kernels whose end\-to\-end time is dominated by launch overhead rather than compute\. As a 5\-tuple:

ℳreq\\displaystyle\\mathcal\{M\}\_\{\\text\{req\}\}=\{launch\_overhead\_pct​\(pL\),avg\_launch\_ns,kernel\_launches\}\\displaystyle=\\\{\\texttt\{launch\\\_overhead\\\_pct\}\\ \(p\_\{L\}\),\\ \\texttt\{avg\\\_launch\\\_ns\},\\ \\texttt\{kernel\\\_launches\}\\\}θ\\displaystyle\\theta=\{high:20%,critical:40%\}\\displaystyle=\\\{\\text\{high\}\{:\}\\,20\\%,\\ \\text\{critical\}\{:\}\\,40\\%\\\}trigger​\(m\)\\displaystyle\\text\{trigger\}\(m\)=𝟙​\[pL\>θhigh\]\\displaystyle=\\mathds\{1\}\[\\,p\_\{L\}\>\\theta\_\{\\text\{high\}\}\\,\]analyze→\(severity bypL,rec:“CUDA Graph capture or kernel fusion”,\\displaystyle\\rightarrow\(\\text\{severity by \}p\_\{L\},\\ \\text\{rec\}\{:\}\\ \\text\{\`\`CUDA Graph capture or kernel fusion''\},est\. speedup=\(1−pL/100\)−1\)\\displaystyle\\qquad\\text\{est\.\\ speedup\}=\(1\-p\_\{L\}/100\)^\{\-1\}\)
The five tools above are representative; we refer readers to the open\-source release for the complete set and their full definitions\.

### B\.2Defining a New Tool

KernelPro’s tools are pluggable: a new heuristic is added by subclassingProfilingTool, declaring its requiredncumetrics and thresholds, implementinganalyze\(\), and registering it with the@ToolRegistry\.registerdecorator—after which the proactive orchestrator invokes it automatically \(no changes to the core system\)\. Listing[2](https://arxiv.org/html/2606.26453#LST2)shows a complete tool \(the register\-pressure detector\);analyze\(\)returns a structuredToolOutputcarrying the severity, root cause, and ranked recommendations that become the natural\-language feedback injected into the prompt\.

1@ToolRegistry\.register

2classHighRegisterUsageTool\(ProfilingTool\):

3name="high\_register\_usage"

4priority=70

5REG\_WARNING,REG\_CRITICAL=96,128

6

7@property

8defrequired\_metrics\(self\)\-\>List\[str\]:

9return\["launch\_\_registers\_per\_thread"\]

10

11defanalyze\(self,kernel\_name:str,

12metrics:Dict\[str,Any\]\)\-\>ToolOutput:

13regs=self\.get\_metric\(metrics,

14"launch\_\_registers\_per\_thread",default=32\)

15ifregs<self\.REG\_WARNING:

16returnToolOutput\(triggered=False\)

17critical=regs\>=self\.REG\_CRITICAL

18returnToolOutput\(

19triggered=True,

20severity=Severity\.HIGHifcriticalelseSeverity\.MEDIUM,

21title="HighRegisterUsage",

22summary=f"\{regs:\.0f\}registers/thread\-mayspilltolocalmemory",

23root\_cause="Registerdemandexceedsthephysicalfile;spillsto"

24"localmemoryare10\-100xslowerthanregisters\.",

25recommendations=\[

26"Use\_\_launch\_bounds\_\_toguidethecompiler",

27"Capregisterswith\-maxrregcount",

28"Splitthekernelormovedatatosharedmemory"\],

29expected\_improvement="2\-10xifeliminatingspilling"\)

Listing 2:Defining and registering a micro\-profiling tool\. Theanalyzemethod returns a structuredToolOutput\(severity, root cause, ranked recommendations\) that the orchestrator renders into prompt feedback\.Tools span a wide complexity range under the sameanalyze→\\rightarrowToolOutputinterface\. Listing[2](https://arxiv.org/html/2606.26453#LST2)is a minimal single\-threshold detector; Listing[3](https://arxiv.org/html/2606.26453#LST3)shows production\-grade diagnostic logic—WarpStallToolreads five warp\-stall counters, identifies the*dominant*stall reason viaargmax, and dispatches a root cause and recommendations tailored to that reason \(the same expert “which stall, and why” reasoning a human applies toncuoutput\)\. We reproduce its complete implementation\.

1@ToolRegistry\.register

2classWarpStallTool\(ProfilingTool\):

3name="warp\_stall\_analyzer"

4description="Analyzeswarpstallreasonstoidentifybottlenecks"

5version="1\.0\.0"

6priority=80

7

8

9STALL\_HIGH=40\.0

10STALL\_CRITICAL=60\.0

11

12@property

13defrequired\_metrics\(self\)\-\>List\[str\]:

14return\[

15"smsp\_\_warp\_issue\_stalled\_memory\_dependency\_per\_warp\_active\.pct",

16"smsp\_\_warp\_issue\_stalled\_short\_scoreboard\_per\_warp\_active\.pct",

17"smsp\_\_warp\_issue\_stalled\_long\_scoreboard\_per\_warp\_active\.pct",

18"smsp\_\_warp\_issue\_stalled\_barrier\_per\_warp\_active\.pct",

19"smsp\_\_warp\_issue\_stalled\_branch\_resolving\_per\_warp\_active\.pct"

20\]

21

22defanalyze\(self,kernel\_name:str,metrics:Dict\[str,Any\]\)\-\>ToolOutput:

23

24stalls=\{

25"memory\_dependency":self\.get\_metric\(metrics,

26"smsp\_\_warp\_issue\_stalled\_memory\_dependency\_per\_warp\_active\.pct",default=0\),

27"short\_scoreboard":self\.get\_metric\(metrics,

28"smsp\_\_warp\_issue\_stalled\_short\_scoreboard\_per\_warp\_active\.pct",default=0\),

29"long\_scoreboard":self\.get\_metric\(metrics,

30"smsp\_\_warp\_issue\_stalled\_long\_scoreboard\_per\_warp\_active\.pct",default=0\),

31"barrier":self\.get\_metric\(metrics,

32"smsp\_\_warp\_issue\_stalled\_barrier\_per\_warp\_active\.pct",default=0\),

33"branch\_resolving":self\.get\_metric\(metrics,

34"smsp\_\_warp\_issue\_stalled\_branch\_resolving\_per\_warp\_active\.pct",default=0\)

35\}

36

37

38dominant\_stall=max\(stalls,key=stalls\.get\)

39dominant\_value=stalls\[dominant\_stall\]

40

41

42ifdominant\_value<self\.STALL\_HIGH:

43returnToolOutput\(triggered=False\)

44

45is\_critical=dominant\_value\>=self\.STALL\_CRITICAL

46severity=Severity\.HIGHifis\_criticalelseSeverity\.MEDIUM

47

48

49root\_cause=self\.\_get\_stall\_root\_cause\(dominant\_stall\)

50recommendations=self\.\_get\_stall\_recommendations\(dominant\_stall,stalls\)

51

52returnToolOutput\(

53triggered=True,

54severity=severity,

55title=f"HighWarpStalls\(\{dominant\_stall\.replace\(’\_’,’’\)\.title\(\)\}\)",

56summary=f"\{dominant\_value:\.1f\}%ofwarpsstalledon\{dominant\_stall\}",

57metrics\_observed=\{

58"Memorydependencystalls":f"\{stalls\[’memory\_dependency’\]:\.1f\}%",

59"Shortscoreboardstalls":f"\{stalls\[’short\_scoreboard’\]:\.1f\}%",

60"Longscoreboardstalls":f"\{stalls\[’long\_scoreboard’\]:\.1f\}%",

61"Barrierstalls":f"\{stalls\[’barrier’\]:\.1f\}%",

62"Branchresolvingstalls":f"\{stalls\[’branch\_resolving’\]:\.1f\}%",

63"Dominantstall":dominant\_stall

64\},

65metrics\_thresholds=\{"Anystalltype":f"<\{self\.STALL\_HIGH\}%\(target\)"\},

66root\_cause=root\_cause,

67recommendations=recommendations,

68expected\_improvement="1\.5\-3xdependingonstalltype"

69\)

70

71def\_get\_stall\_root\_cause\(self,stall\_type:str\)\-\>str:

72causes=\{

73"memory\_dependency":"Warpswaitformemoryoperationstocomplete"

74"\(memory\-bound\):400\+cyclegloballatency,uncoalescedaccesses,"

75"orinsufficientparallelismtohidelatency\.",

76"short\_scoreboard":"Warpswaitforshared\-memory/texture\(MIO\)ops:"

77"shared\-memorybankconflicts,texturecachemisses,syncoverhead\.",

78"long\_scoreboard":"Warpswaitforlong\-latencyopstrackedbythelong"

79"scoreboard:outstandinggloballoads,texturefetches,atomics\.",

80"barrier":"Warpswaitat\_\_syncthreads\(\)fortheblocktocatchup:"

81"loadimbalance,toomanysyncpoints,ordivergentexecutiontimes\.",

82"branch\_resolving":"Warpswaitforbranchconditionstoresolve:"

83"data\-dependentorthread\-ID\-dependentdivergentpaths\."

84\}

85returncauses\.get\(stall\_type,"Unknownstalltype"\)

86

87def\_get\_stall\_recommendations\(self,stall\_type:str,all\_stalls:Dict\)\-\>List\[str\]:

88base\_recs=\{

89"memory\_dependency":\[

90"Usesharedmemorytocachefrequentlyaccesseddata",

91"Implementvectorizedloads\(float4\)forbetterbandwidth",

92"Ensurecoalescedmemoryaccesspatterns",

93"Increaseoccupancytohidememorylatency\(morewarps\)",

94"Use\_\_ldg\(\)forread\-onlydata\(usestexturecache\)",

95"Considersoftwarepipeliningtooverlapcomputeandmemory"

96\],

97"short\_scoreboard":\[

98"Checkforsharedmemorybankconflicts",

99"Reducesharedmemoryaccessesifpossible",

100"Usepaddingtoavoidbankconflicts:‘\_\_shared\_\_floatA\[32\]\[33\]‘",

101"Considerusingregistersinsteadofsharedmemory",

102"Profilewithncu\-\-setfulltoseebankconflictcounts"

103\],

104"long\_scoreboard":\[

105"Similartomemorydependency\-optimizememoryaccess",

106"UseasynccopyforprefetchingifonAmpere\+\(cp\.async\)",

107"ConsiderTMAonHopperforbulktransfers",

108"Increasewarp\-levelparallelism"

109\],

110"barrier":\[

111"Reducenumberof\_\_syncthreads\(\)calls",

112"Usewarp\-levelprimitives\(\_\_shfl,\_\_ballot\)whenpossible",

113"Balanceworkloadacrossthreadsinablock",

114"Considerwarp\-synchronousprogrammingforsmalldata",

115"NEVERput\_\_syncthreads\(\)insideconditionalblocks"

116\],

117"branch\_resolving":\[

118"Reorganizedatatominimizebranchdivergence",

119"Usearithmeticinsteadofbranches:result=cond\*v1\+\(1\-cond\)\*v2",

120"Sortdatatogroupsimilarcomputationpaths",

121"Usepredicationforshortconditionalblocks",

122"Considerseparatekernelsfordifferentcodepaths"

123\]

124\}

125returnbase\_recs\.get\(stall\_type,\["Profilefurthertoidentifyrootcause"\]\)

Listing 3:WarpStallTool: complete implementation\. It selects the dominant stall reason among fivencucounters viaargmaxand returns a reason\-specific root cause and recommendations\.

## Appendix CMicro\-Profiling Tool Effectiveness: KernelBench and Production Kernels

We isolate the contribution ofKernelPro’s micro\-profiling tools with two controlled with\-tools\-versus\-without\-tools ablations: one on 42 KernelBench tasks \(§[C\.1](https://arxiv.org/html/2606.26453#A3.SS1)\) and one on six production Triton inference kernels \(§[C\.3](https://arxiv.org/html/2606.26453#A3.SS3)\), confirming the effect holds on both benchmark and real serving workloads\.

### C\.1KernelBench Tasks

We conduct a controlled ablation on the 42\-task subset \(Section[5\.3](https://arxiv.org/html/2606.26453#S5.SS3)\) to isolate the contribution of micro\-profiling tools\. Three configurations compare progressive levels of profiling feedback:

- •No Feedback: The LLM receives only source code and compilation/correctness errors \(\-\-skip\-profiling\)\. Represents naive LLM retry\.
- •RawncuMetrics: Full NCU counter dump \(∼\\sim50 raw metric values\) appended to the prompt with no interpretation or bottleneck classification \(\-\-raw\-metrics\-only\)\.
- •Ours: Rawncumetrics plus allKernelPromicro\-profiling tools \(Stage 1 roofline classification \+ Stage 2 filtered tool analysis with structured optimization guidance\)\.

All conditions use Sonnet 4\.6 on A100 GPUs with 15 seeds per task \(5 temperatures×\\times3 independent rounds\), 30 iterations, 2 candidates per iteration, and greedy search\. We report the best speedup across all seeds for each task\-configuration pair, with outlier speedups capped to limit disproportionate influence\.

##### Task IDs\.

The 42 KernelBench tasks used in this ablation are:

- •Level 1\(19\): 1, 6, 7, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 23, 24, 26, 40, 88, 97
- •Level 2\(16\): 3, 9, 12, 14, 18, 19, 22, 29, 30, 33, 34, 36, 37, 39, 40, 41
- •Level 3\(7\): 28, 30, 31, 32, 43, 44, 50

Table[14](https://arxiv.org/html/2606.26453#A3.T14)shows the fraction of tasks where each configuration produces a correct kernel faster than the PyTorch eager baseline\.

Table 14:Correctness ablation: fraction of tasks solved \(speedup\>1×\>1\\times\)\.Ours achieves a 100% solve rate across all 42 tasks, solving 31% more tasks than No Feedback and 17% more than RawncuMetrics\. On several tasks, only Ours achieves substantial speedup \(e\.g\., L2 task 40: 7\.5×\\timesvs 1\.0×\\times/1\.0×\\times; L2 task 41: 5\.8×\\timesvs 1\.0×\\times/1\.0×\\times\), demonstrating that profiling\-guided insight unlocks optimizations inaccessible to unguided approaches\. The L2 task 41 trajectory \(Appendix[G\.1](https://arxiv.org/html/2606.26453#A7.SS1)\) provides a detailed walkthrough of how tool guidance produces its12\.1×12\.1\\timesfinal speedup—the2\.8×2\.8\\timessingle\-step improvement, where the agent acts on the tools’ top\-ranked directives \(epilogue fusion plus BF16 tensor cores\), exemplifies the mechanism by which structured feedback unlocks optimization paths invisible to unguided search\.

Table[15](https://arxiv.org/html/2606.26453#A3.T15)compares optimization quality across configurations\.

Table 15:Speedup ablation \(outliers capped, unsolved tasks scored as 1\.0×\\times\)\.Table[16](https://arxiv.org/html/2606.26453#A3.T16)shows the fraction of tasks exceeding progressively higher speedup thresholds\. Ours dominates across all thresholds, with the gap widening at higher values\.

Table 16:fast\_p: fraction of tasks with speedup\>p\>p\.We use the Wilcoxon signed\-rank test\(Hollander et al\.,[2014](https://arxiv.org/html/2606.26453#bib.bib10)\), a non\-parametric paired test that makes no distributional assumptions, to evaluate whether configurations differ in optimization quality\. For each task, we take the best speedup \(with outliers capped\) as that configuration’s score; unsolved tasks are scored as 1\.0×\\times\. The test is one\-sided \(H1H\_\{1\}: first configuration produces higher speedups\);nnreports the number of non\-zero differences \(Wilcoxon excludes tied pairs\)\. Results \(Table[8](https://arxiv.org/html/2606.26453#S5.T8)\) confirm the ordering Ours\>\>No Feedback\>\>Rawncu\.

We highlight three key observations from the ablation results\.

##### Full pipeline is necessary for reliable optimization\.

Ours solves 100% of tasks–the structured profiling pipeline consistently identifies actionable bottlenecks that the model can address\.

##### Raw metrics without interpretation are harmful\.

RawncuMetrics performs significantly*worse*than No Feedback \(p=0\.0007p=0\.0007\)\. This condition passes∼50\{\\sim\}50raw NCU counters without interpretation–comparable to CudaForge\(Zhang et al\.,[2025](https://arxiv.org/html/2606.26453#bib.bib37)\), which presents statistically pre\-filtered metrics as verbatim name\-value pairs\. Unstructured metric dumps appear to distract the model: it attempts to optimize counters that do not translate to wall\-clock improvements, or misinterprets raw values without roofline context\. This validatesKernelPro’s core design principle that*interpretation*, not mere data access, drives optimization quality\.

##### The interpretation layer is the key differentiator\.

The gap from Rawncuto Ours \(raw→\\toanalyzed\) is larger than from No Feedback to Ours \(none→\\toanalyzed\)\. The value lies not in*collecting*profiling data, but in*presenting it as actionable, structured optimization guidance*via the Stage 1/Stage 2 pipeline\. This reveals thatKernelPro’s tool design is not simply “give the LLM more data”–it is specifically about*structured interpretation*of profiling data\. Raw metrics actively degrade performance because the LLM lacks the domain context to translate hardware counters into optimization actions; only the micro\-profiling tools’ detect\-analyze\-recommend pattern bridges this gap reliably\.

##### Profiling prevents correct\-but\-slow plateaus\.

Trajectory analysis reveals thatKernelPronever produces a correct kernel that fails to exceed baseline \(0/62 trajectories stuck at 1\.0×\\times\), while No Feedback gets stuck in 10% of cases\. When both configurations recover from an initial compilation error, No Feedback often produces a kernel that merely matches baseline performance and cannot improve further without knowing what to optimize\.KernelProimmediately receives structured guidance \(e\.g\., “reduce register usage from 168 to<<64 per thread”\) and follows a clear optimization path\. This explains whyKernelProachieves higher fast\_1 despite similar raw correctness rates: profiling does not fix bugs, but it prevents the model from settling on correct\-but\-slow solutions by always providing an actionable optimization target\.

### C\.2Per\-Tool Impact Analysis

Per\-tool coverage and hit rates are reported in Table[9](https://arxiv.org/html/2606.26453#S5.T9)\(main body\)\. Below we provide additional per\-bottleneck breakdowns\.

Table[17](https://arxiv.org/html/2606.26453#A3.T17)shows the dominant optimization action triggered by each bottleneck type, measured as the single largest improvement per task\.

Table 17:Most impactful optimization action per bottleneck type across 42 tasks\.Three observations emerge:

##### Tensor Core switch is the highest\-impact single action\.

For compute\-bound kernels, the tools identify that the LLM\-generated CUDA kernel uses CUDA cores instead of Tensor Cores and recommend switching to cuBLAS/CUTLASS\. This single recommendation produces an average 6\.25×\\timesgain across 9 tasks \(max 9\.6×\\times\)\. Without the tool’s quantitative signal \(e\.g\., “TC utilization: 0\.7%, threshold\>\>10%”\), the LLM has no way to know its manual kernel is orders of magnitude slower than the hardware’s peak capability\. The Task 41 trajectory \(Appendix[G\.1](https://arxiv.org/html/2606.26453#A7.SS1)\) traces this pattern end\-to\-end: at iteration 13 the agent acts on the tools’ top directives—fusing the epilogue and moving to a BF16 tensor\-core GEMM—for a4\.11×→11\.49×4\.11\\times\\to 11\.49\\timesleap; the full generated solution is reproduced in the same appendix\.

##### Register spill detection has the highest per\-fire reliability\.

RegisterSpill fires in only 16/42 tasks but converts to\>\>1\.5×\\timesimprovement 18\.2% of the time—the highest among all tools\. The diagnostic “register spill detected: 168 registers/thread, spilling to local memory” is highly actionable: the LLM can directly reduce register pressure by restructuring loops or reducing live variables\. Example: L1\_1 goes from 1\.0×\\timesto 10\.6×\\timesafter the LLM eliminates register spills\.

##### Memory optimizations provide broad but incremental gains\.

MemoryCoalescing and Vectorization fire in 37–41 of 42 tasks but have low hit rates \(4–6%\)\. Memory access patterns are the most common bottleneck on A100 but are harder for the LLM to fix correctly—coalescing requires non\-trivial data layout changes, while vectorization requires careful alignment handling\.

### C\.3Production Inference Kernels

To confirm the effect holds beyond KernelBench, we repeat the with\-tools\-versus\-without\-tools ablation on six representative Triton inference kernels covering common transformer serving patterns \(NVIDIA A100\)\. The kernels are:layer\_norm\(fused forward layer normalization with affine transform\),vec\_matmul\(single\-vector matrix multiplication for decoding\),persistent\_vec\_matmul\(persistent\-loop variant that reuses thread blocks across output tiles\),gemm\_split\_k\(split\-K GEMM with atomic accumulation for high\-parallelism reduction\),fused\_moe\(fused Mixture\-of\-Experts token routing and batched expert GEMM from vLLM\), andkernel\_paged\_attention\(paged multi\-head attention with online softmax over non\-contiguous KV\-cache blocks from vLLM, supporting GQA\)\. Both configurations receive identicalncu/nsys/SASS raw metrics; the only difference is whether the tool\-generated guidance is included in the LLM prompt\.

Table 18:Triton inference kernel ablation: tools enable better speedup in fewer iterations\.Tools achieved 36% higher geometric mean speedup \(2\.02×\\timesvs 1\.49×\\times\) in 31% fewer iterations \(5\.0 vs 7\.2\)\.fused\_moefailed entirely without tools after 10 iterations but achieved 3\.26×\\timeswith tools in just 2 iterations\. Thevec\_matmulcase illustrates the “fast but wrong” trap: without tools, the LLM repeatedly generated incorrect solutions achieving 1\.69–1\.80×\\timesspeedups that failed correctness, unable to diagnose why faster attempts broke\. With tool guidance identifying the bottleneck \(“shared memory limiting occupancy to 2%”\), it switched to warp\-level reduction \(\_\_shfl\_down\_sync\) and found a correct 2\.03×\\timessolution in 4 iterations\.

## Appendix DSearch Memory Ablation

We conduct a controlled ablation to isolate the contribution of search memory–KernelPro’s cross\-iteration learning mechanism described in Section[3\.4](https://arxiv.org/html/2606.26453#S3.SS4)\. Two configurations compare the fullKernelPropipeline with and without memory:

- •Memory ON \(Treatment\): FullKernelPropipeline with session memory extraction\. Accumulated findings \(errors\_and\_corrections,successful\_patterns,search\_findings\) are injected into subsequent iterations\.
- •Memory OFF \(Control\): Same pipeline with\-\-disable\-session\-memory\. Each iteration sees only the problem spec and profiling feedback from the parent node\. Represents memoryless search\.

All conditions use Sonnet 4\.6 on H100 GPUs with 25 iterations, 2 candidates per iteration, greedy search, and temperature 0\.2\. Both conditions run fullncuprofiling and the profiling agent–the sole variable is whether cross\-iteration findings are extracted and injected\. Outlier speedups are capped at 25×\\times\(consistent with Appendix[C](https://arxiv.org/html/2606.26453#A3)\)\. This ablation uses the same 42\-task subset as the other component ablations \(Section[5\.3](https://arxiv.org/html/2606.26453#S5.SS3)\)\.

### D\.1Search Memory Does Not Change Final Speedup

Table[19](https://arxiv.org/html/2606.26453#A4.T19)reports the best speedup achieved by each condition \(capped at 25×\\times\)\.

Table 19:Search memory ablation: optimization quality\.The Wilcoxon signed\-rank test finds no significant difference in final optimization quality \(z=0\.911z=0\.911,p=0\.181p=0\.181,n=36n=36\)\. The geometric mean ratio \(ON/OFF\) is1\.06×1\.06\\times–a small positive effect that does not reach significance due to high per\-task variance\. The 90% CI for the ratio \(\[0\.936,1\.200\]\[0\.936,1\.200\]\) bounds any quality difference to at most 20%, so the flat ceiling is a bounded result rather than merely underpowered\.

### D\.2Search Memory Speeds Early Convergence

Table 20:Convergence: geometric mean best@NN\(capped at 25×\\times\)\. Wilcoxon onlog⁡\(best@​N\)\\log\(\\text\{best@\}N\)\.†Marginal \(p<0\.10p<0\.10\)\. best@NN\.

While final speedup is equivalent, memory shows a marginal advantage in*early*convergence\. Table[20](https://arxiv.org/html/2606.26453#A4.T20)reports best\-so\-far speedup at progressive search depths\.

Memory shows its strongest effect in the first 10 iterations \(p=0\.072p=0\.072\), with treatment winning on 15/36 tasks versus 10/36 for control atN=10N=10\(meanlog⁡\(T/C\)=\+0\.043\\log\(T/C\)=\+0\.043\)\. This advantage diminishes at deeper search depths, consistent with memory’s role in avoiding redundant early exploration–memory prevents re\-exploration of dead\-end strategies during early iterations, but sufficiently long search compensates for the lack of cross\-iteration learning\. Its primary contribution is thus reaching good solutions faster rather than finding better ones given unlimited time\. In practice this is the useful regime: under a fixed iteration budget–where users want the best kernel within a time constraint rather than the absolute best given unlimited search–the faster early convergence translates to reduced wall\-clock time at no quality cost\.

## Appendix ESearch Strategy Ablation: MCTS vs Greedy

We conduct a controlled ablation to isolate the contribution of MCTS tree search versus greedy search\.KernelProis the only LLM\-based CUDA optimization agent to employ structured tree search; all related systems use flat iterative refinement or evolutionary population\-based search\. This ablation tests whether principled exploration via UCT selection produces better optimizations than pure greedy exploitation under a fixed compute budget of 60 candidates \(30 iterations×\\times2 candidates per expansion\)\.

Both conditions use the fullKernelProprofiling pipeline \(Stage 1 roofline \+ Stage 2 filtered tools\), session memory, Sonnet 4\.6 on A100 GPUs, and identical hyperparameters: 15 seeds per task \(5 temperatures×\\times3 rounds\), 30 iterations, 2 candidates per iteration\. The sole variable is the search strategy:

- •Greedy: Each iteration expands only the current best node, with fresh restart after consecutive failures\. Repair targets are selected via a composite priority*\(execution stage, error magnitude, speedup\)*–kernels that execute but produce wrong output rank above compilation failures, with smaller numerical errors and higher partial speedups preferred as tiebreakers, concentrating repair budget on candidates closest to correctness\.
- •MCTS: Tree search with progressive widening and UCT selection \(c=2c=\\sqrt\{2\}\)\. Terminal marking prunes dead\-end branches\.

Table[21](https://arxiv.org/html/2606.26453#A5.T21)reports aggregate metrics on the 42 KernelBench tasks\. MCTS achieves a 26% higher geometric mean speedup \(4\.60×\\timesvs 3\.65×\\times\) and wins on 29 tasks while greedy wins on only 8\. The Wilcoxon signed\-rank test rejects the null hypothesis atp=0\.004p=0\.004, providing strong evidence that MCTS produces higher\-quality optimizations under matched compute budgets\.

Table 21:Search strategy ablation: MCTS vs Greedy \(42 KernelBench tasks\)\.Table[22](https://arxiv.org/html/2606.26453#A5.T22)shows the fraction of tasks exceeding progressively higher speedup thresholds\. MCTS dominates across all thresholds, with the gap largest at moderate speedups \(1\.5–5×\\times\) where MCTS solves 8–11% more tasks than greedy\. This suggests MCTS is most valuable for tasks requiring multi\-step optimization chains where greedy converges prematurely to a local optimum\.

Table 22:fast\_p: fraction of tasks with best speedup\>p\>p\.The largest MCTS advantages occur on tasks requiring exploration to escape local optima: L1 task 7 \(10\.0×\\timesratio: MCTS 25\.0×\\timesvs greedy 2\.5×\\times\), L2 task 33 \(8\.0×\\timesratio: MCTS 8\.0×\\timesvs greedy 1\.0×\\times\), L3 task 31 \(6\.0×\\timesratio: MCTS 25\.0×\\timesvs greedy 4\.2×\\times\)\. In these cases, greedy converges to a correct but suboptimal solution and cannot escape, while MCTS explores alternative subtrees that discover superior optimization paths\. Task 46 \(Appendix[G\.2](https://arxiv.org/html/2606.26453#A7.SS2)\) demonstrates the extreme case: 43 consecutive failures before the first valid solution—a search depth that tree search’s persistence enables while budget\-limited single\-shot approaches would have long abandoned\. Greedy wins on 8 tasks, with the largest margin on L2 task 22 \(greedy 9\.1×\\timesvs MCTS 1\.1×\\times\)\. These tend to be tasks where the optimization landscape is unimodal–depth\-first exploitation suffices and MCTS wastes budget on unnecessary exploration\.

The results confirm that principled tree search is a significant contributor toKernelPro’s performance\. The 26% improvement in geometric mean speedup represents the gap between exploiting a single promising path \(greedy\) and systematically exploring the optimization landscape \(MCTS\)\. Combined with the micro\-profiling tool ablation \(Appendix[C](https://arxiv.org/html/2606.26453#A3)\), this establishes that both components–tool\-guided feedback and MCTS search–independently and significantly improve optimization quality\.

## Appendix FTool Invocation Mode Ablation

We conduct a controlled ablation to isolate the contribution of proactive tool orchestration–KernelPro’s deterministic execution of all relevant micro\-profiling tools before LLM invocation \(Section[3\.2\.3](https://arxiv.org/html/2606.26453#S3.SS2.SSS3)\)\. This ablation tests whether guaranteed comprehensive analysis outperforms the standard reactive pattern where the LLM decides which tools to invoke via function calling\.

Both conditions use the fullKernelPromicro\-profiling tool suite \(15 tools\), greedy search \(30 iterations×\\times2 candidates\), session memory, and Sonnet 4\.6 on A100 GPUs at five temperature settings \(0\.2, 0\.33, 0\.45, 0\.57, 0\.7\) with 2 independent rounds per task \(10 seeds total\)\. The sole variable is the tool invocation pattern:

- •Proactive: All relevant micro\-profiling tools are executed deterministically based on bottleneck classification\. Results are injected into the prompt before LLM invocation\.
- •Reactive: The LLM decides which tools to invoke via standard function calling\. The same 15 tools are available but invocation is stochastic\.

Table[23](https://arxiv.org/html/2606.26453#A6.T23)reports aggregate metrics across 42 tasks \(best\-of\-seeds per task\)\. Proactive mode achieves a 23% higher geometric mean speedup and wins on 10 of 14 decisive tasks \(p=0\.035p=0\.035\)\. The effect is consistent across difficulty levels \(L1, L2, L3\)\.

Table 23:Tool invocation mode ablation: Proactive vs Reactive \(42 tasks, best\-of\-seeds\)\.Table[24](https://arxiv.org/html/2606.26453#A6.T24)shows the fraction of tasks exceeding progressively higher speedup thresholds\. Proactive mode dominates across all thresholds, with the largest gap at moderate speedups where proactive orchestration unlocks optimizations that reactive invocation misses\.

Table 24:fast\_p by tool invocation mode: fraction of tasks with best speedup\>p\>p\.The results validate the proactive design choice described in Section[3\.2\.3](https://arxiv.org/html/2606.26453#S3.SS2.SSS3)\. Under reactive invocation, the LLM calls an average of 3–4 tools per iteration, missing critical analyses that would reveal optimization opportunities\. Proactive orchestration guarantees that all relevant tools execute, providing the comprehensive profiling context that enables higher\-quality optimizations\. This is particularly important for CUDA optimization where the relevant bottleneck is often non\-obvious–an expert systematically checks all metrics rather than guessing which to examine\. The Task 41 trajectory \(Appendix[G\.1](https://arxiv.org/html/2606.26453#A7.SS1)\) provides a concrete illustration: the tools emit six ranked, actionable directives for the diagnosed bottleneck, and the agent compounds several of them in one step \(epilogue fusion through the lower\-ranked BF16 switch\) for a2\.8×2\.8\\timesimprovement—reactive invocation calling 3–4 tools would surface only the top of this list, forgoing the rest\.

Combined with the micro\-profiling tool ablation \(Appendix[C](https://arxiv.org/html/2606.26453#A3)\) and search strategy ablation \(Appendix[E](https://arxiv.org/html/2606.26453#A5)\), this establishes that three components–interpreted tool feedback, proactive orchestration, and MCTS search–each independently and significantly improve optimization quality \(p<0\.05p<0\.05for all three\)\.

## Appendix GDesign Component Case Studies: CUTLASS Optimization Trajectories

To illustrate howKernelPro’s design components work together in practice, we present two complementary KernelBench trajectories that produce CUTLASS\-based solutions\. Level 2 Task 41 demonstrates*deep refinement*—a single decisive architectural choice followed by iterative profiling\-driven polishing\. Level 3 Task 46 demonstrates*resilient exploration*—sustained search through 43 failed attempts before achieving a working solution, then rapid profiling\-guided convergence\. Together they showcase howKernelPro’s iterative search, micro\-profiling tools, and CUTLASS integration enable optimization trajectories that single\-shot generation cannot replicate\.

### G\.1Level 2 / Task 41: GEMM \+ BatchNorm \+ GELU \+ ReLU

##### Problem\.

The reference implements a linear layer \(16384×4096→409616384\\times 4096\\to 4096\) followed by batch normalization, GELU activation, and ReLU—four separate PyTorch operations producing four kernel launches with intermediate DRAM round\-trips\. The compute\-bound GEMM dominates runtime, but the memory\-bound elementwise tail contributes significant overhead at this batch size\.

##### Trajectory overview\.

KernelProexplores 41 iterations total, of which 7 produce valid solutions \(34 fail compilation or correctness checks\)\. The successful lineage progresses through iterations 9→\\to12→\\to13→\\to26→\\to28→\\to37→\\to40, achieving speedups of3\.91×→4\.11×→11\.49×→11\.73×→12\.00×→12\.09×→12\.10×3\.91\\times\\to 4\.11\\times\\to 11\.49\\times\\to 11\.73\\times\\to 12\.00\\times\\to 12\.09\\times\\to 12\.10\\times\. Final speedup:12\.1×\\times\.

##### The breakthrough \(iteration 13\)\.

The most significant optimization event occurs at the third successful iteration, where speedup jumps from4\.11×4\.11\\timesto11\.49×11\.49\\times—a2\.8×2\.8\\timesimprovement in a single step\. The micro\-profiling tools diagnose a memory bottleneck \(15\.88 sectors/request indicating strided loads, 8\.7 GB read traffic dominating runtime, 244 registers/thread limiting occupancy\) and rank a set of remedies, the top two of which the agent applies together\. It completes the fusion of the trailing batch\-normalization and activation into the GEMM epilogue—a customGeluReluActfunctor plus BatchNorm affine folded into the GEMM weights \(α⋅W\\alpha\\cdot W,α⋅bias\+β\\alpha\\cdot\\text\{bias\}\+\\beta\)—eliminating the last standalone elementwise kernel and its round trip through DRAM; and it switches the GEMM from a TF32 path \(floatinputs,GemmShape<128,128,32\>\) to BF16 tensor cores \(cutlass::bfloat16\_t,GemmShape<128,128,16\>\), which halves the dominant read traffic and roughly doubles tensor\-core throughput\. Both changes attack the same diagnosed bottleneck—off\-chip data movement—and together account for the leap; the precision switch supplies most of the magnitude, while the fusion removes the residual inter\-kernel traffic and is what makes the BF16 path a single end\-to\-end CUTLASS kernel rather than a library call\.

##### Profiling\-driven polish \(iterations 26–40\)\.

After the breakthrough, the profiling tools shift from identifying the primary bottleneck to precision tuning\. At iteration 26, ncu identifies the bottleneck as compute\-bound with 230 registers/thread; the agent explores a larger tile \(GemmShape<128,256,32\>\) to increase arithmetic intensity\. At iteration 28, profiling reveals this larger tile underperforms due to wave quantization effects at the problem’s dimensions, and the agent reverts toGemmShape<128,128,32\>with different warp shapes\. Subsequent iterations explore secondary optimizations—fusing the BF16 cast into the epilogue, adjusting pipeline stages, and tuning warp tile shapes \(64×\\times128×\\times64vs32×\\times64×\\times64\)—each guided by specific ncu metrics\. The final four iterations collectively contribute0\.61×0\.61\\timesadditional speedup through these incremental refinements\.

##### CUTLASS integration\.

The final solution showcases deep CUTLASS customization: a customGeluReluActepilogue functor template \(handling both scalarfloatand vectorizedcutlass::Array<float, N\>\), batch normalization folded into GEMM weights at initialization, and BF16 precision with FP32 accumulation\. The agent demonstrates understanding of CUTLASS’s type system—definingLinearCombinationGeneric<GeluReluAct, float, 4, float, float\>with 4\-wide vectorized epilogue access—which requires knowledge of CUTLASS’s template conventions that static few\-shot examples alone would not generalize\.

KernelPro– Level 2 Task 41—12\.10×12\.10\\timesSpeedup[⬇](data:text/plain;base64,aW1wb3J0IG9zCm9zLmVudmlyb25bJ1RPUkNIX0NVREFfQVJDSF9MSVNUJ10gPSAnOC4wJwppbXBvcnQgdG9yY2gKaW1wb3J0IHRvcmNoLm5uIGFzIG5uCmZyb20gdG9yY2gudXRpbHMuY3BwX2V4dGVuc2lvbiBpbXBvcnQgbG9hZF9pbmxpbmUKCmN1ZGFfc3JjID0gcicnJwojaW5jbHVkZSA8dG9yY2gvZXh0ZW5zaW9uLmg+CiNpbmNsdWRlIDxBVGVuL2N1ZGEvQ1VEQUNvbnRleHQuaD4KI2luY2x1ZGUgPGN1ZGFfcnVudGltZS5oPgojaW5jbHVkZSA8Y3VkYV9iZjE2Lmg+CiNpbmNsdWRlIDxjdXRsYXNzL2N1dGxhc3MuaD4KI2luY2x1ZGUgPGN1dGxhc3MvYXJyYXkuaD4KI2luY2x1ZGUgPGN1dGxhc3MvbnVtZXJpY190eXBlcy5oPgojaW5jbHVkZSA8Y3V0bGFzcy9nZW1tL2RldmljZS9nZW1tX3VuaXZlcnNhbC5oPgojaW5jbHVkZSA8Y3V0bGFzcy9lcGlsb2d1ZS90aHJlYWQvbGluZWFyX2NvbWJpbmF0aW9uX2dlbmVyaWMuaD4KI2luY2x1ZGUgPGN1dGxhc3MvZ2VtbS90aHJlYWRibG9jay90aHJlYWRibG9ja19zd2l6emxlLmg+CiNpbmNsdWRlIDxjbWF0aD4KCi8vIEZhc3QgdmVjdG9yaXplZCBGUDMyIC0+IEJGMTYgY2FzdDogMTZCIGxvYWRzLCA4QiBzdG9yZXMsIHByb2Nlc3NlcyA0IGVsZW1lbnRzL3RocmVhZApfX2dsb2JhbF9fIF9fbGF1bmNoX2JvdW5kc19fKDI1NiwgNCkKdm9pZCBjYXN0X2ZwMzJfdG9fYmYxNl92ZWM0KGNvbnN0IGZsb2F0NCogX19yZXN0cmljdF9fIGluLCBfX252X2JmbG9hdDE2MiogX19yZXN0cmljdF9fIG91dCwgaW50IG40KSB7CiAgICBpbnQgaWR4ID0gYmxvY2tJZHgueCAqIGJsb2NrRGltLnggKyB0aHJlYWRJZHgueDsKICAgIGludCBzdHJpZGUgPSBibG9ja0RpbS54ICogZ3JpZERpbS54OwogICAgZm9yIChpbnQgaSA9IGlkeDsgaSA8IG40OyBpICs9IHN0cmlkZSkgewogICAgICAgIGZsb2F0NCBmID0gaW5baV07CiAgICAgICAgX19udl9iZmxvYXQxNjIgYSA9IF9fZmxvYXRzMmJmbG9hdDE2Ml9ybihmLngsIGYueSk7CiAgICAgICAgX19udl9iZmxvYXQxNjIgYiA9IF9fZmxvYXRzMmJmbG9hdDE2Ml9ybihmLnosIGYudyk7CiAgICAgICAgKChfX252X2JmbG9hdDE2Miopb3V0KVtpKjIgKyAwXSA9IGE7CiAgICAgICAgKChfX252X2JmbG9hdDE2Miopb3V0KVtpKjIgKyAxXSA9IGI7CiAgICB9Cn0KCnRvcmNoOjpUZW5zb3IgY2FzdF90b19iZjE2X2Zhc3QodG9yY2g6OlRlbnNvciB4KSB7CiAgICBUT1JDSF9DSEVDSyh4LmlzX2N1ZGEoKSAmJiB4LmR0eXBlKCkgPT0gdG9yY2g6OmtGbG9hdDMyLCAibmVlZCBmcDMyIGN1ZGEiKTsKICAgIHggPSB4LmNvbnRpZ3VvdXMoKTsKICAgIGF1dG8gb3V0ID0gdG9yY2g6OmVtcHR5KHguc2l6ZXMoKSwgeC5vcHRpb25zKCkuZHR5cGUodG9yY2g6OmtCRmxvYXQxNikpOwogICAgaW50NjRfdCBuID0geC5udW1lbCgpOwogICAgaWYgKG4gJSA0ICE9IDAgfHwgKCh1aW50cHRyX3QpeC5kYXRhX3B0cigpICUgMTYpICE9IDApIHsKICAgICAgICByZXR1cm4geC50byh0b3JjaDo6a0JGbG9hdDE2KTsKICAgIH0KICAgIGludCBuNCA9IG4gLyA0OwogICAgaW50IGJsb2NrID0gMjU2OwogICAgaW50IGdyaWQgPSBzdGQ6Om1pbigyMDQ4LCAobjQgKyBibG9jayAtIDEpIC8gYmxvY2spOwogICAgYXV0byBzdHJlYW0gPSBhdDo6Y3VkYTo6Z2V0Q3VycmVudENVREFTdHJlYW0oKTsKICAgIGNhc3RfZnAzMl90b19iZjE2X3ZlYzQ8PDxncmlkLCBibG9jaywgMCwgc3RyZWFtPj4+KAogICAgICAgIHJlaW50ZXJwcmV0X2Nhc3Q8Y29uc3QgZmxvYXQ0Kj4oeC5kYXRhX3B0cjxmbG9hdD4oKSksCiAgICAgICAgcmVpbnRlcnByZXRfY2FzdDxfX252X2JmbG9hdDE2Mio+KG91dC5kYXRhX3B0cigpKSwgbjQpOwogICAgcmV0dXJuIG91dDsKfQoKdGVtcGxhdGUgPHR5cGVuYW1lIFQ+IHN0cnVjdCBHZWx1UmVsdUFjdDsKdGVtcGxhdGUgPD4gc3RydWN0IEdlbHVSZWx1QWN0PGZsb2F0PiB7CiAgICBzdGF0aWMgY29uc3QgYm9vbCBrSXNIZWF2eSA9IHRydWU7CiAgICBDVVRMQVNTX0hPU1RfREVWSUNFIGZsb2F0IG9wZXJhdG9yKCkoZmxvYXQgeCkgY29uc3QgewogICAgICAgIGlmICh4IDw9IDAuZikgcmV0dXJuIDAuZjsKICAgICAgICBmbG9hdCB4MyA9IHggKiB4ICogeDsKICAgICAgICBmbG9hdCBpbm5lciA9IDAuNzk3ODg0NTYwOGYgKiAoeCArIDAuMDQ0NzE1ZiAqIHgzKTsKICAgICAgICByZXR1cm4gMC41ZiAqIHggKiAoMS4wZiArIDo6dGFuaGYoaW5uZXIpKTsKICAgIH0KfTsKdGVtcGxhdGUgPGludCBOPiBzdHJ1Y3QgR2VsdVJlbHVBY3Q8Y3V0bGFzczo6QXJyYXk8ZmxvYXQsIE4+PiB7CiAgICBzdGF0aWMgY29uc3QgYm9vbCBrSXNIZWF2eSA9IHRydWU7CiAgICBDVVRMQVNTX0hPU1RfREVWSUNFIGN1dGxhc3M6OkFycmF5PGZsb2F0LCBOPiBvcGVyYXRvcigpKAogICAgICAgIGN1dGxhc3M6OkFycmF5PGZsb2F0LCBOPiBjb25zdCYgeCkgY29uc3QgewogICAgICAgIGN1dGxhc3M6OkFycmF5PGZsb2F0LCBOPiByOyBHZWx1UmVsdUFjdDxmbG9hdD4gb3A7CiAgICAgICAgQ1VUTEFTU19QUkFHTUFfVU5ST0xMIGZvciAoaW50IGkgPSAwOyBpIDwgTjsgKytpKSByW2ldID0gb3AoeFtpXSk7CiAgICAgICAgcmV0dXJuIHI7CiAgICB9Cn07Cgp1c2luZyBFcGlsb2d1ZU9wID0gY3V0bGFzczo6ZXBpbG9ndWU6OnRocmVhZDo6TGluZWFyQ29tYmluYXRpb25HZW5lcmljPAogICAgR2VsdVJlbHVBY3QsIGZsb2F0LCA0LCBmbG9hdCwgZmxvYXQ+OwoKdXNpbmcgR2VtbURlZmF1bHQgPSBjdXRsYXNzOjpnZW1tOjpkZXZpY2U6OkdlbW1Vbml2ZXJzYWw8CiAgICBjdXRsYXNzOjpiZmxvYXQxNl90LCBjdXRsYXNzOjpsYXlvdXQ6OlJvd01ham9yLAogICAgY3V0bGFzczo6YmZsb2F0MTZfdCwgY3V0bGFzczo6bGF5b3V0OjpDb2x1bW5NYWpvciwKICAgIGZsb2F0LCBjdXRsYXNzOjpsYXlvdXQ6OlJvd01ham9yLCBmbG9hdCwKICAgIGN1dGxhc3M6OmFyY2g6Ok9wQ2xhc3NUZW5zb3JPcCwgY3V0bGFzczo6YXJjaDo6U204MCwKICAgIGN1dGxhc3M6OmdlbW06OkdlbW1TaGFwZTwxMjgsMTI4LDMyPiwKICAgIGN1dGxhc3M6OmdlbW06OkdlbW1TaGFwZTw2NCw2NCwzMj4sCiAgICBjdXRsYXNzOjpnZW1tOjpHZW1tU2hhcGU8MTYsOCwxNj4sCiAgICBFcGlsb2d1ZU9wLAogICAgY3V0bGFzczo6Z2VtbTo6dGhyZWFkYmxvY2s6OkdlbW1JZGVudGl0eVRocmVhZGJsb2NrU3dpenpsZTw4PiwgNCwgOCwgOD47Cgp1c2luZyBHZW1tU21hbGwgPSBjdXRsYXNzOjpnZW1tOjpkZXZpY2U6OkdlbW1Vbml2ZXJzYWw8CiAgICBjdXRsYXNzOjpiZmxvYXQxNl90LCBjdXRsYXNzOjpsYXlvdXQ6OlJvd01ham9yLAogICAgY3V0bGFzczo6YmZsb2F0MTZfdCwgY3V0bGFzczo6bGF5b3V0OjpDb2x1bW5NYWpvciwKICAgIGZsb2F0LCBjdXRsYXNzOjpsYXlvdXQ6OlJvd01ham9yLCBmbG9hdCwKICAgIGN1dGxhc3M6OmFyY2g6Ok9wQ2xhc3NUZW5zb3JPcCwgY3V0bGFzczo6YXJjaDo6U204MCwKICAgIGN1dGxhc3M6OmdlbW06OkdlbW1TaGFwZTw2NCwxMjgsNjQ+LAogICAgY3V0bGFzczo6Z2VtbTo6R2VtbVNoYXBlPDMyLDY0LDY0PiwKICAgIGN1dGxhc3M6OmdlbW06OkdlbW1TaGFwZTwxNiw4LDE2PiwKICAgIEVwaWxvZ3VlT3AsCiAgICBjdXRsYXNzOjpnZW1tOjp0aHJlYWRibG9jazo6R2VtbUlkZW50aXR5VGhyZWFkYmxvY2tTd2l6emxlPDg+LCA0LCA4LCA4PjsKCnVzaW5nIEdlbW1CaWcgPSBjdXRsYXNzOjpnZW1tOjpkZXZpY2U6OkdlbW1Vbml2ZXJzYWw8CiAgICBjdXRsYXNzOjpiZmxvYXQxNl90LCBjdXRsYXNzOjpsYXlvdXQ6OlJvd01ham9yLAogICAgY3V0bGFzczo6YmZsb2F0MTZfdCwgY3V0bGFzczo6bGF5b3V0OjpDb2x1bW5NYWpvciwKICAgIGZsb2F0LCBjdXRsYXNzOjpsYXlvdXQ6OlJvd01ham9yLCBmbG9hdCwKICAgIGN1dGxhc3M6OmFyY2g6Ok9wQ2xhc3NUZW5zb3JPcCwgY3V0bGFzczo6YXJjaDo6U204MCwKICAgIGN1dGxhc3M6OmdlbW06OkdlbW1TaGFwZTwxMjgsMjU2LDMyPiwKICAgIGN1dGxhc3M6OmdlbW06OkdlbW1TaGFwZTw2NCw2NCwzMj4sCiAgICBjdXRsYXNzOjpnZW1tOjpHZW1tU2hhcGU8MTYsOCwxNj4sCiAgICBFcGlsb2d1ZU9wLAogICAgY3V0bGFzczo6Z2VtbTo6dGhyZWFkYmxvY2s6OkdlbW1JZGVudGl0eVRocmVhZGJsb2NrU3dpenpsZTw4PiwgMywgOCwgOD47Cgp0b3JjaDo6VGVuc29yIGZ1c2VkX2xpbmVhcl9ibl9nZWx1X3JlbHUodG9yY2g6OlRlbnNvciB4LCB0b3JjaDo6VGVuc29yIFcsIHRvcmNoOjpUZW5zb3IgYikgewogICAgVE9SQ0hfQ0hFQ0soeC5pc19jdWRhKCkgJiYgVy5pc19jdWRhKCkgJiYgYi5pc19jdWRhKCksICJpbnB1dHMgbXVzdCBiZSBDVURBIik7CiAgICB4ID0geC5jb250aWd1b3VzKCk7IFcgPSBXLmNvbnRpZ3VvdXMoKTsgYiA9IGIuY29udGlndW91cygpOwogICAgaW50IE0gPSB4LnNpemUoMCksIEsgPSB4LnNpemUoMSksIE4gPSBXLnNpemUoMCk7CiAgICBhdXRvIG91dCA9IHRvcmNoOjplbXB0eSh7TSwgTn0sIGIub3B0aW9ucygpKTsKCiAgICBhdXRvIGxhdW5jaCA9IFsmXShhdXRvIGdlbW1fb3BfdGFnKSAtPiBjdXRsYXNzOjpTdGF0dXMgewogICAgICAgIHVzaW5nIEdlbW1UID0gdHlwZW5hbWUgZGVjbHR5cGUoZ2VtbV9vcF90YWcpOjp0eXBlOwogICAgICAgIEdlbW1UIGdlbW1fb3A7CiAgICAgICAgdHlwZW5hbWUgR2VtbVQ6OkFyZ3VtZW50cyBhcmdzKAogICAgICAgICAgICBjdXRsYXNzOjpnZW1tOjpHZW1tVW5pdmVyc2FsTW9kZTo6a0dlbW0sIHtNLCBOLCBLfSwgMSwgezEuMGYsIDEuMGZ9LAogICAgICAgICAgICByZWludGVycHJldF9jYXN0PGN1dGxhc3M6OmJmbG9hdDE2X3QqPih4LmRhdGFfcHRyKCkpLAogICAgICAgICAgICByZWludGVycHJldF9jYXN0PGN1dGxhc3M6OmJmbG9hdDE2X3QqPihXLmRhdGFfcHRyKCkpLAogICAgICAgICAgICBiLmRhdGFfcHRyPGZsb2F0PigpLCBvdXQuZGF0YV9wdHI8ZmxvYXQ+KCksCiAgICAgICAgICAgIChpbnQ2NF90KU0qSywgKGludDY0X3QpTipLLCAwLCAoaW50NjRfdClNKk4sIEssIEssIDAsIE4pOwogICAgICAgIHJldHVybiBnZW1tX29wKGFyZ3MpOwogICAgfTsKCiAgICBjdXRsYXNzOjpTdGF0dXMgc3RhdHVzOwogICAgYm9vbCBzbWFsbF9tID0gKE0gPD0gNjQpLCBiaWcgPSAoTSA+PSA1MTIpICYmIChOID49IDUxMikgJiYgKEsgPj0gNTEyKTsKICAgIGlmIChzbWFsbF9tKSB7IHN0cnVjdCBUYWcgeyB1c2luZyB0eXBlID0gR2VtbVNtYWxsOyB9IHRhZzsgc3RhdHVzID0gbGF1bmNoKHRhZyk7IH0KICAgIGVsc2UgaWYgKGJpZykgeyBzdHJ1Y3QgVGFnIHsgdXNpbmcgdHlwZSA9IEdlbW1CaWc7IH0gdGFnOyBzdGF0dXMgPSBsYXVuY2godGFnKTsgfQogICAgZWxzZSB7IHN0cnVjdCBUYWcgeyB1c2luZyB0eXBlID0gR2VtbURlZmF1bHQ7IH0gdGFnOyBzdGF0dXMgPSBsYXVuY2godGFnKTsgfQogICAgVE9SQ0hfQ0hFQ0soc3RhdHVzID09IGN1dGxhc3M6OlN0YXR1czo6a1N1Y2Nlc3MsICJDVVRMQVNTIEdFTU0gZmFpbGVkIik7CiAgICByZXR1cm4gb3V0Owp9CicnJwoKY3BwX3NyYyA9ICgidG9yY2g6OlRlbnNvciBmdXNlZF9saW5lYXJfYm5fZ2VsdV9yZWx1KHRvcmNoOjpUZW5zb3IgeCwgdG9yY2g6OlRlbnNvciBXLCB0b3JjaDo6VGVuc29yIGIpO1xuIgogICAgICAgICAgICJ0b3JjaDo6VGVuc29yIGNhc3RfdG9fYmYxNl9mYXN0KHRvcmNoOjpUZW5zb3IgeCk7XG4iKQoKbW9kID0gbG9hZF9pbmxpbmUoCiAgICBuYW1lPSdmdXNlZF9saW5ibl9nZWx1cmVsdV9iZjE2X3Y1JywgY3BwX3NvdXJjZXM9Y3BwX3NyYywgY3VkYV9zb3VyY2VzPWN1ZGFfc3JjLAogICAgZnVuY3Rpb25zPVsnZnVzZWRfbGluZWFyX2JuX2dlbHVfcmVsdScsICdjYXN0X3RvX2JmMTZfZmFzdCddLAogICAgZXh0cmFfY3VkYV9jZmxhZ3M9WyctSS9wYXRoL3RvL2N1dGxhc3MvaW5jbHVkZScsICctSS9wYXRoL3RvL2N1dGxhc3MvdG9vbHMvdXRpbC9pbmNsdWRlJywKICAgICAgICAnLS1leHB0LXJlbGF4ZWQtY29uc3RleHByJywgJy1zdGQ9YysrMTcnLCAnLWdlbmNvZGU9YXJjaD1jb21wdXRlXzgwLGNvZGU9c21fODAnLAogICAgICAgICctTzMnLCAnLUROREVCVUcnLCAnLS11c2VfZmFzdF9tYXRoJ10sIHZlcmJvc2U9RmFsc2UpCgpjbGFzcyBNb2RlbE5ldyhubi5Nb2R1bGUpOgogICAgZGVmIF9faW5pdF9fKHNlbGYsIGluX2ZlYXR1cmVzLCBvdXRfZmVhdHVyZXMpOgogICAgICAgIHN1cGVyKCkuX19pbml0X18oKQogICAgICAgIHNlbGYuZ2VtbSA9IG5uLkxpbmVhcihpbl9mZWF0dXJlcywgb3V0X2ZlYXR1cmVzKQogICAgICAgIHNlbGYuYmF0Y2hfbm9ybSA9IG5uLkJhdGNoTm9ybTFkKG91dF9mZWF0dXJlcykKICAgICAgICBzZWxmLl9mdXNlZF9XX2JmMTYgPSBOb25lCiAgICAgICAgc2VsZi5fZnVzZWRfYl9mcDMyID0gTm9uZQoKICAgIGRlZiBfYnVpbGRfZnVzZWQoc2VsZiwgZGV2aWNlKToKICAgICAgICB3aXRoIHRvcmNoLm5vX2dyYWQoKToKICAgICAgICAgICAgZ2FtbWEgPSBzZWxmLmJhdGNoX25vcm0ud2VpZ2h0LnRvKGRldmljZSkuZmxvYXQoKQogICAgICAgICAgICBiZXRhX3AgPSBzZWxmLmJhdGNoX25vcm0uYmlhcy50byhkZXZpY2UpLmZsb2F0KCkKICAgICAgICAgICAgbWVhbiA9IHNlbGYuYmF0Y2hfbm9ybS5ydW5uaW5nX21lYW4udG8oZGV2aWNlKS5mbG9hdCgpCiAgICAgICAgICAgIHZhciA9IHNlbGYuYmF0Y2hfbm9ybS5ydW5uaW5nX3Zhci50byhkZXZpY2UpLmZsb2F0KCkKICAgICAgICAgICAgc2NhbGUgPSBnYW1tYSAvIHRvcmNoLnNxcnQodmFyICsgc2VsZi5iYXRjaF9ub3JtLmVwcykKICAgICAgICAgICAgVyA9IHNlbGYuZ2VtbS53ZWlnaHQudG8oZGV2aWNlKS5mbG9hdCgpCiAgICAgICAgICAgIGJpYXMgPSBzZWxmLmdlbW0uYmlhcy50byhkZXZpY2UpLmZsb2F0KCkKICAgICAgICAgICAgc2VsZi5fZnVzZWRfV19iZjE2ID0gKFcgKiBzY2FsZS51bnNxdWVlemUoMSkpLnRvKHRvcmNoLmJmbG9hdDE2KS5jb250aWd1b3VzKCkKICAgICAgICAgICAgc2VsZi5fZnVzZWRfYl9mcDMyID0gKChiaWFzIC0gbWVhbikgKiBzY2FsZSArIGJldGFfcCkuY29udGlndW91cygpCgogICAgZGVmIGZvcndhcmQoc2VsZiwgeCk6CiAgICAgICAgaWYgc2VsZi5fZnVzZWRfV19iZjE2IGlzIE5vbmUgb3Igc2VsZi5fZnVzZWRfV19iZjE2LmRldmljZSAhPSB4LmRldmljZToKICAgICAgICAgICAgc2VsZi5fYnVpbGRfZnVzZWQoeC5kZXZpY2UpCiAgICAgICAgeF9iZjE2ID0gbW9kLmNhc3RfdG9fYmYxNl9mYXN0KHguY29udGlndW91cygpKSBpZiB4LmR0eXBlID09IHRvcmNoLmZsb2F0MzIgZWxzZSB4LnRvKHRvcmNoLmJmbG9hdDE2KQogICAgICAgIHJldHVybiBtb2QuZnVzZWRfbGluZWFyX2JuX2dlbHVfcmVsdSh4X2JmMTYsIHNlbGYuX2Z1c2VkX1dfYmYxNiwgc2VsZi5fZnVzZWRfYl9mcDMyKQ==)importosos\.environ\[’TORCH\_CUDA\_ARCH\_LIST’\]=’8\.0’importtorchimporttorch\.nnasnnfromtorch\.utils\.cpp\_extensionimportload\_inlinecuda\_src=r’’’\#include<torch/extension\.h\>\#include<ATen/cuda/CUDAContext\.h\>\#include<cuda\_runtime\.h\>\#include<cuda\_bf16\.h\>\#include<cutlass/cutlass\.h\>\#include<cutlass/array\.h\>\#include<cutlass/numeric\_types\.h\>\#include<cutlass/gemm/device/gemm\_universal\.h\>\#include<cutlass/epilogue/thread/linear\_combination\_generic\.h\>\#include<cutlass/gemm/threadblock/threadblock\_swizzle\.h\>\#include<cmath\>//FastvectorizedFP32\-\>BF16cast:16Bloads,8Bstores,processes4elements/thread\_\_global\_\_\_\_launch\_bounds\_\_\(256,4\)voidcast\_fp32\_to\_bf16\_vec4\(constfloat4\*\_\_restrict\_\_in,\_\_nv\_bfloat162\*\_\_restrict\_\_out,intn4\)\{intidx=blockIdx\.x\*blockDim\.x\+threadIdx\.x;intstride=blockDim\.x\*gridDim\.x;for\(inti=idx;i<n4;i\+=stride\)\{float4f=in\[i\];\_\_nv\_bfloat162a=\_\_floats2bfloat162\_rn\(f\.x,f\.y\);\_\_nv\_bfloat162b=\_\_floats2bfloat162\_rn\(f\.z,f\.w\);\(\(\_\_nv\_bfloat162\*\)out\)\[i\*2\+0\]=a;\(\(\_\_nv\_bfloat162\*\)out\)\[i\*2\+1\]=b;\}\}torch::Tensorcast\_to\_bf16\_fast\(torch::Tensorx\)\{TORCH\_CHECK\(x\.is\_cuda\(\)&&x\.dtype\(\)==torch::kFloat32,"needfp32cuda"\);x=x\.contiguous\(\);autoout=torch::empty\(x\.sizes\(\),x\.options\(\)\.dtype\(torch::kBFloat16\)\);int64\_tn=x\.numel\(\);if\(n%4\!=0\|\|\(\(uintptr\_t\)x\.data\_ptr\(\)%16\)\!=0\)\{returnx\.to\(torch::kBFloat16\);\}intn4=n/4;intblock=256;intgrid=std::min\(2048,\(n4\+block\-1\)/block\);autostream=at::cuda::getCurrentCUDAStream\(\);cast\_fp32\_to\_bf16\_vec4<<<grid,block,0,stream\>\>\>\(reinterpret\_cast<constfloat4\*\>\(x\.data\_ptr<float\>\(\)\),reinterpret\_cast<\_\_nv\_bfloat162\*\>\(out\.data\_ptr\(\)\),n4\);returnout;\}template<typenameT\>structGeluReluAct;template<\>structGeluReluAct<float\>\{staticconstboolkIsHeavy=true;CUTLASS\_HOST\_DEVICEfloatoperator\(\)\(floatx\)const\{if\(x<=0\.f\)return0\.f;floatx3=x\*x\*x;floatinner=0\.7978845608f\*\(x\+0\.044715f\*x3\);return0\.5f\*x\*\(1\.0f\+::tanhf\(inner\)\);\}\};template<intN\>structGeluReluAct<cutlass::Array<float,N\>\>\{staticconstboolkIsHeavy=true;CUTLASS\_HOST\_DEVICEcutlass::Array<float,N\>operator\(\)\(cutlass::Array<float,N\>const&x\)const\{cutlass::Array<float,N\>r;GeluReluAct<float\>op;CUTLASS\_PRAGMA\_UNROLLfor\(inti=0;i<N;\+\+i\)r\[i\]=op\(x\[i\]\);returnr;\}\};usingEpilogueOp=cutlass::epilogue::thread::LinearCombinationGeneric<GeluReluAct,float,4,float,float\>;usingGemmDefault=cutlass::gemm::device::GemmUniversal<cutlass::bfloat16\_t,cutlass::layout::RowMajor,cutlass::bfloat16\_t,cutlass::layout::ColumnMajor,float,cutlass::layout::RowMajor,float,cutlass::arch::OpClassTensorOp,cutlass::arch::Sm80,cutlass::gemm::GemmShape<128,128,32\>,cutlass::gemm::GemmShape<64,64,32\>,cutlass::gemm::GemmShape<16,8,16\>,EpilogueOp,cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8\>,4,8,8\>;usingGemmSmall=cutlass::gemm::device::GemmUniversal<cutlass::bfloat16\_t,cutlass::layout::RowMajor,cutlass::bfloat16\_t,cutlass::layout::ColumnMajor,float,cutlass::layout::RowMajor,float,cutlass::arch::OpClassTensorOp,cutlass::arch::Sm80,cutlass::gemm::GemmShape<64,128,64\>,cutlass::gemm::GemmShape<32,64,64\>,cutlass::gemm::GemmShape<16,8,16\>,EpilogueOp,cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8\>,4,8,8\>;usingGemmBig=cutlass::gemm::device::GemmUniversal<cutlass::bfloat16\_t,cutlass::layout::RowMajor,cutlass::bfloat16\_t,cutlass::layout::ColumnMajor,float,cutlass::layout::RowMajor,float,cutlass::arch::OpClassTensorOp,cutlass::arch::Sm80,cutlass::gemm::GemmShape<128,256,32\>,cutlass::gemm::GemmShape<64,64,32\>,cutlass::gemm::GemmShape<16,8,16\>,EpilogueOp,cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8\>,3,8,8\>;torch::Tensorfused\_linear\_bn\_gelu\_relu\(torch::Tensorx,torch::TensorW,torch::Tensorb\)\{TORCH\_CHECK\(x\.is\_cuda\(\)&&W\.is\_cuda\(\)&&b\.is\_cuda\(\),"inputsmustbeCUDA"\);x=x\.contiguous\(\);W=W\.contiguous\(\);b=b\.contiguous\(\);intM=x\.size\(0\),K=x\.size\(1\),N=W\.size\(0\);autoout=torch::empty\(\{M,N\},b\.options\(\)\);autolaunch=\[&\]\(autogemm\_op\_tag\)\-\>cutlass::Status\{usingGemmT=typenamedecltype\(gemm\_op\_tag\)::type;GemmTgemm\_op;typenameGemmT::Argumentsargs\(cutlass::gemm::GemmUniversalMode::kGemm,\{M,N,K\},1,\{1\.0f,1\.0f\},reinterpret\_cast<cutlass::bfloat16\_t\*\>\(x\.data\_ptr\(\)\),reinterpret\_cast<cutlass::bfloat16\_t\*\>\(W\.data\_ptr\(\)\),b\.data\_ptr<float\>\(\),out\.data\_ptr<float\>\(\),\(int64\_t\)M\*K,\(int64\_t\)N\*K,0,\(int64\_t\)M\*N,K,K,0,N\);returngemm\_op\(args\);\};cutlass::Statusstatus;boolsmall\_m=\(M<=64\),big=\(M\>=512\)&&\(N\>=512\)&&\(K\>=512\);if\(small\_m\)\{structTag\{usingtype=GemmSmall;\}tag;status=launch\(tag\);\}elseif\(big\)\{structTag\{usingtype=GemmBig;\}tag;status=launch\(tag\);\}else\{structTag\{usingtype=GemmDefault;\}tag;status=launch\(tag\);\}TORCH\_CHECK\(status==cutlass::Status::kSuccess,"CUTLASSGEMMfailed"\);returnout;\}’’’cpp\_src=\("torch::Tensorfused\_linear\_bn\_gelu\_relu\(torch::Tensorx,torch::TensorW,torch::Tensorb\);\\n""torch::Tensorcast\_to\_bf16\_fast\(torch::Tensorx\);\\n"\)mod=load\_inline\(name=’fused\_linbn\_gelurelu\_bf16\_v5’,cpp\_sources=cpp\_src,cuda\_sources=cuda\_src,functions=\[’fused\_linear\_bn\_gelu\_relu’,’cast\_to\_bf16\_fast’\],extra\_cuda\_cflags=\[’\-I/path/to/cutlass/include’,’\-I/path/to/cutlass/tools/util/include’,’\-\-expt\-relaxed\-constexpr’,’\-std=c\+\+17’,’\-gencode=arch=compute\_80,code=sm\_80’,’\-O3’,’\-DNDEBUG’,’\-\-use\_fast\_math’\],verbose=False\)classModelNew\(nn\.Module\):def\_\_init\_\_\(self,in\_features,out\_features\):super\(\)\.\_\_init\_\_\(\)self\.gemm=nn\.Linear\(in\_features,out\_features\)self\.batch\_norm=nn\.BatchNorm1d\(out\_features\)self\.\_fused\_W\_bf16=Noneself\.\_fused\_b\_fp32=Nonedef\_build\_fused\(self,device\):withtorch\.no\_grad\(\):gamma=self\.batch\_norm\.weight\.to\(device\)\.float\(\)beta\_p=self\.batch\_norm\.bias\.to\(device\)\.float\(\)mean=self\.batch\_norm\.running\_mean\.to\(device\)\.float\(\)var=self\.batch\_norm\.running\_var\.to\(device\)\.float\(\)scale=gamma/torch\.sqrt\(var\+self\.batch\_norm\.eps\)W=self\.gemm\.weight\.to\(device\)\.float\(\)bias=self\.gemm\.bias\.to\(device\)\.float\(\)self\.\_fused\_W\_bf16=\(W\*scale\.unsqueeze\(1\)\)\.to\(torch\.bfloat16\)\.contiguous\(\)self\.\_fused\_b\_fp32=\(\(bias\-mean\)\*scale\+beta\_p\)\.contiguous\(\)defforward\(self,x\):ifself\.\_fused\_W\_bf16isNoneorself\.\_fused\_W\_bf16\.device\!=x\.device:self\.\_build\_fused\(x\.device\)x\_bf16=mod\.cast\_to\_bf16\_fast\(x\.contiguous\(\)\)ifx\.dtype==torch\.float32elsex\.to\(torch\.bfloat16\)returnmod\.fused\_linear\_bn\_gelu\_relu\(x\_bf16,self\.\_fused\_W\_bf16,self\.\_fused\_b\_fp32\)

### G\.2Level 3 / Task 46: NetVLAD with Ghost Clusters

##### Problem\.

The reference implements NetVLAD with ghost clusters: a multi\-step pipeline involving matrix multiplication \(2048×512→372048\\times 512\\to 37clusters\), batch normalization, softmax, cluster slicing, weighted sum, batched matrix multiplication \(residual computation\), L2 normalization, and flatten\. The operations span both compute\-bound \(two matrix multiplications\) and memory\-bound \(softmax, normalization, transpose\) regimes, making this a challenging heterogeneous optimization target\.

##### Trajectory overview\.

KernelProexplores 46 iterations, with only the final 3 producing valid solutions \(43 fail compilation or correctness checks\)\. The first valid solution at iteration 42 achieves only0\.75×0\.75\\times—a*regression*below baseline\. Profiling\-guided refinement then produces2\.38×2\.38\\timesat iteration 44 and2\.79×2\.79\\timesat iteration 45\. Final speedup:2\.79×\\times\.

##### Resilience through 43 failures\.

This trajectory represents the most extreme exploration ratio in all 250 KernelBench tasks: 46 iterations attempted for 3 successes \(15\.3×\\timesexploration ratio\)\. The 43 failed iterations include incorrect CUTLASS configurations \(wrong layouts for the batched GEMM\), numerical errors in fused softmax normalization, and compilation failures from incompatible template instantiations\. The search strategy’s persistence—continuing to generate candidates informed by prior failures—is critical here: a budget\-limited single\-shot approach would have abandoned this problem long before iteration 42\. When iteration 42 finally compiles and runs correctly, it incorporates layout choices \(RowMajor for the first GEMM, ColumnMajor for the batched GEMM\) that reflect lessons accumulated across the prior failed attempts\.

##### Profiling\-driven recovery from regression \(iterations 42–45\)\.

The initial valid solution at0\.75×0\.75\\timesdemonstrates a common failure mode: correct CUTLASS compilation that performs worse than PyTorch’s optimized defaults due to poor memory access patterns\. The micro\-profiling tools quantify the problem precisely: 17\.86 sectors/request \(ideal: 4, indicating∼4×\{\\sim\}4\\timeswasted bandwidth\) and L1 hit rate of only 41\.55%\. The recommendations are specific and actionable: “fix memory access pattern to enable coalescing,” “use vectorized loads \(float4\) where layout permits,” and “use shared memory staging to batch and reorder accesses\.” At iteration 44, the agent restructures memory access patterns based on this feedback, achieving2\.38×2\.38\\times—a3\.2×3\.2\\timesimprovement from profiling guidance alone\. At iteration 45, ncu identifies shared memory as the new occupancy limiter and recommends reducing per\-block shared memory usage\. The agent responds with\[32\]\[33\]bank\-conflict\-free padding on shared memory tiles and\_\_launch\_bounds\_\_annotations, reaching the final2\.79×2\.79\\times\.

##### CUTLASS integration\.

Unlike Task 41’s single\-GEMM deep customization, Task 46 deploys CUTLASS broadly across a heterogeneous pipeline: two distinct GEMM configurations \(128×\\times128×\\times16 for the assignment matrix multiplication and 64×\\times64×\\times16 for the batched residual computation\), each with layout\-aware configuration matched to the operation’s access patterns\. These CUTLASS GEMMs serve as building blocks alongside hand\-written kernels—a fused batch normalization \+ softmax \+ slice kernel using warp\-shuffle reductions, a tiled 32×\\times32 transpose with bank\-conflict padding, and float4 vectorized loads for the normalization passes\. This hybrid architecture \(CUTLASS for compute\-bound GEMMs, hand\-tuned CUDA for memory\-bound operations\) demonstratesKernelPro’s ability to match optimization strategy to each operation’s bottleneck characteristics\.

KernelPro– Level 3 Task 46—2\.79×2\.79\\timesSpeedup[⬇](data:text/plain;base64,aW1wb3J0IG9zCm9zLmVudmlyb25bJ1RPUkNIX0NVREFfQVJDSF9MSVNUJ10gPSAnOC4wJwppbXBvcnQgbWF0aAppbXBvcnQgdG9yY2gKaW1wb3J0IHRvcmNoLm5uIGFzIG5uCmZyb20gdG9yY2gudXRpbHMuY3BwX2V4dGVuc2lvbiBpbXBvcnQgbG9hZF9pbmxpbmUKCmN1ZGFfc3JjID0gcicnJwojaW5jbHVkZSA8dG9yY2gvZXh0ZW5zaW9uLmg+CiNpbmNsdWRlIDxBVGVuL2N1ZGEvQ1VEQUNvbnRleHQuaD4KI2luY2x1ZGUgPGN1ZGFfcnVudGltZS5oPgojaW5jbHVkZSA8Y3V0bGFzcy9nZW1tL2RldmljZS9nZW1tLmg+CiNpbmNsdWRlIDxjdXRsYXNzL2dlbW0vZGV2aWNlL2dlbW1fYmF0Y2hlZC5oPgoKdXNpbmcgR2VtbVRGMzIgPSBjdXRsYXNzOjpnZW1tOjpkZXZpY2U6OkdlbW08CiAgICBmbG9hdCwgY3V0bGFzczo6bGF5b3V0OjpSb3dNYWpvciwKICAgIGZsb2F0LCBjdXRsYXNzOjpsYXlvdXQ6OlJvd01ham9yLAogICAgZmxvYXQsIGN1dGxhc3M6OmxheW91dDo6Um93TWFqb3IsCiAgICBmbG9hdCwgY3V0bGFzczo6YXJjaDo6T3BDbGFzc1RlbnNvck9wLCBjdXRsYXNzOjphcmNoOjpTbTgwLAogICAgY3V0bGFzczo6Z2VtbTo6R2VtbVNoYXBlPDEyOCwgMTI4LCAxNj4sCiAgICBjdXRsYXNzOjpnZW1tOjpHZW1tU2hhcGU8NjQsIDY0LCAxNj4sCiAgICBjdXRsYXNzOjpnZW1tOjpHZW1tU2hhcGU8MTYsIDgsIDg+PjsKCnVzaW5nIEdlbW1CYXRjaGVkVEYzMiA9IGN1dGxhc3M6OmdlbW06OmRldmljZTo6R2VtbUJhdGNoZWQ8CiAgICBmbG9hdCwgY3V0bGFzczo6bGF5b3V0OjpDb2x1bW5NYWpvciwKICAgIGZsb2F0LCBjdXRsYXNzOjpsYXlvdXQ6OlJvd01ham9yLAogICAgZmxvYXQsIGN1dGxhc3M6OmxheW91dDo6Um93TWFqb3IsCiAgICBmbG9hdCwgY3V0bGFzczo6YXJjaDo6T3BDbGFzc1RlbnNvck9wLCBjdXRsYXNzOjphcmNoOjpTbTgwLAogICAgY3V0bGFzczo6Z2VtbTo6R2VtbVNoYXBlPDY0LCA2NCwgMTY+LAogICAgY3V0bGFzczo6Z2VtbTo6R2VtbVNoYXBlPDMyLCAzMiwgMTY+LAogICAgY3V0bGFzczo6Z2VtbTo6R2VtbVNoYXBlPDE2LCA4LCA4Pj47Cgp0b3JjaDo6VGVuc29yIGdlbW1fcnIodG9yY2g6OlRlbnNvciBBLCB0b3JjaDo6VGVuc29yIEIpIHsKICAgIGludCBNID0gQS5zaXplKDApLCBLID0gQS5zaXplKDEpLCBOID0gQi5zaXplKDEpOwogICAgYXV0byBDID0gdG9yY2g6OmVtcHR5KHtNLCBOfSwgQS5vcHRpb25zKCkpOwogICAgR2VtbVRGMzIgZ2VtbTsKICAgIGN1ZGFTdHJlYW1fdCBzdHJlYW0gPSBhdDo6Y3VkYTo6Z2V0Q3VycmVudENVREFTdHJlYW0oKTsKICAgIEdlbW1URjMyOjpBcmd1bWVudHMgYXJncyh7TSxOLEt9LAogICAgICAgIHtBLmRhdGFfcHRyPGZsb2F0PigpLCBLfSwge0IuZGF0YV9wdHI8ZmxvYXQ+KCksIE59LAogICAgICAgIHtDLmRhdGFfcHRyPGZsb2F0PigpLCBOfSwge0MuZGF0YV9wdHI8ZmxvYXQ+KCksIE59LCB7MS4wZiwgMC4wZn0pOwogICAgZ2VtbShhcmdzLCBudWxscHRyLCBzdHJlYW0pOwogICAgcmV0dXJuIEM7Cn0KCnRvcmNoOjpUZW5zb3IgZ2VtbV9iYXRjaGVkX2NyKHRvcmNoOjpUZW5zb3IgQSwgdG9yY2g6OlRlbnNvciBCLAogICAgICAgICAgICAgICAgICAgICAgICAgICAgICBpbnQgQnN6LCBpbnQgTSwgaW50IE5faW5uZXIsIGludCBOX291dCkgewogICAgYXV0byBDID0gdG9yY2g6OmVtcHR5KHtCc3osIE0sIE5fb3V0fSwgQi5vcHRpb25zKCkpOwogICAgR2VtbUJhdGNoZWRURjMyIGdlbW07CiAgICBjdWRhU3RyZWFtX3Qgc3RyZWFtID0gYXQ6OmN1ZGE6OmdldEN1cnJlbnRDVURBU3RyZWFtKCk7CiAgICBHZW1tQmF0Y2hlZFRGMzI6OkFyZ3VtZW50cyBhcmdzKHtNLCBOX291dCwgTl9pbm5lcn0sCiAgICAgICAge0EuZGF0YV9wdHI8ZmxvYXQ+KCksIE19LCAoaW50NjRfdClOX2lubmVyICogTSwKICAgICAgICB7Qi5kYXRhX3B0cjxmbG9hdD4oKSwgTl9vdXR9LCAoaW50NjRfdClOX2lubmVyICogTl9vdXQsCiAgICAgICAge0MuZGF0YV9wdHI8ZmxvYXQ+KCksIE5fb3V0fSwgKGludDY0X3QpTSAqIE5fb3V0LAogICAgICAgIHtDLmRhdGFfcHRyPGZsb2F0PigpLCBOX291dH0sIChpbnQ2NF90KU0gKiBOX291dCwKICAgICAgICB7MS4wZiwgMC4wZn0sIEJzeik7CiAgICBnZW1tKGFyZ3MsIG51bGxwdHIsIHN0cmVhbSk7CiAgICByZXR1cm4gQzsKfQoKX19nbG9iYWxfXyB2b2lkIGJuX3NvZnRtYXhfc2xpY2Vfa2VybmVsKAogICAgY29uc3QgZmxvYXQqIF9fcmVzdHJpY3RfXyBsb2dpdHMsIGNvbnN0IGZsb2F0KiBfX3Jlc3RyaWN0X18gc2NhbGUsCiAgICBjb25zdCBmbG9hdCogX19yZXN0cmljdF9fIGJpYXMsIGZsb2F0KiBfX3Jlc3RyaWN0X18gb3V0LAogICAgaW50IHJvd3MsIGludCBLX3RvdGFsLCBpbnQgSykgewogICAgaW50IHJvdyA9IGJsb2NrSWR4Lng7CiAgICBpbnQgdGlkID0gdGhyZWFkSWR4Lng7CiAgICBleHRlcm4gX19zaGFyZWRfXyBmbG9hdCBzZGF0YVtdOwogICAgZmxvYXQgbG9jYWxfbWF4ID0gLUlORklOSVRZOwogICAgZm9yIChpbnQgaSA9IHRpZDsgaSA8IEtfdG90YWw7IGkgKz0gYmxvY2tEaW0ueCkgewogICAgICAgIGZsb2F0IHYgPSBsb2dpdHNbcm93ICogS190b3RhbCArIGldICogc2NhbGVbaV0gKyBiaWFzW2ldOwogICAgICAgIHNkYXRhW2ldID0gdjsKICAgICAgICBsb2NhbF9tYXggPSBmbWF4Zihsb2NhbF9tYXgsIHYpOwogICAgfQogICAgX19zeW5jdGhyZWFkcygpOwogICAgaW50IGxhbmUgPSB0aWQgJiAzMTsgaW50IHdpZCA9IHRpZCA+PiA1OwogICAgaW50IG53YXJwcyA9IChibG9ja0RpbS54ICsgMzEpID4+IDU7CiAgICBmb3IgKGludCBvZmYgPSAxNjsgb2ZmID4gMDsgb2ZmID4+PSAxKQogICAgICAgIGxvY2FsX21heCA9IGZtYXhmKGxvY2FsX21heCwgX19zaGZsX3hvcl9zeW5jKDB4ZmZmZmZmZmYsIGxvY2FsX21heCwgb2ZmKSk7CiAgICBfX3NoYXJlZF9fIGZsb2F0IHdhcnBfbWF4WzMyXTsgX19zaGFyZWRfXyBmbG9hdCBzX21heDsKICAgIGlmIChsYW5lID09IDApIHdhcnBfbWF4W3dpZF0gPSBsb2NhbF9tYXg7CiAgICBfX3N5bmN0aHJlYWRzKCk7CiAgICBpZiAod2lkID09IDApIHsKICAgICAgICBmbG9hdCB2ID0gKGxhbmUgPCBud2FycHMpID8gd2FycF9tYXhbbGFuZV0gOiAtSU5GSU5JVFk7CiAgICAgICAgZm9yIChpbnQgb2ZmID0gMTY7IG9mZiA+IDA7IG9mZiA+Pj0gMSkgdiA9IGZtYXhmKHYsIF9fc2hmbF94b3Jfc3luYygweGZmZmZmZmZmLCB2LCBvZmYpKTsKICAgICAgICBpZiAobGFuZSA9PSAwKSBzX21heCA9IHY7CiAgICB9CiAgICBfX3N5bmN0aHJlYWRzKCk7CiAgICBmbG9hdCBsb2NhbF9zdW0gPSAwLmY7CiAgICBmb3IgKGludCBpID0gdGlkOyBpIDwgS190b3RhbDsgaSArPSBibG9ja0RpbS54KSB7CiAgICAgICAgZmxvYXQgZSA9IF9fZXhwZihzZGF0YVtpXSAtIHNfbWF4KTsKICAgICAgICBzZGF0YVtpXSA9IGU7IGxvY2FsX3N1bSArPSBlOwogICAgfQogICAgX19zeW5jdGhyZWFkcygpOwogICAgZm9yIChpbnQgb2ZmID0gMTY7IG9mZiA+IDA7IG9mZiA+Pj0gMSkKICAgICAgICBsb2NhbF9zdW0gKz0gX19zaGZsX3hvcl9zeW5jKDB4ZmZmZmZmZmYsIGxvY2FsX3N1bSwgb2ZmKTsKICAgIF9fc2hhcmVkX18gZmxvYXQgd2FycF9zdW1bMzJdOyBfX3NoYXJlZF9fIGZsb2F0IHNfc3VtOwogICAgaWYgKGxhbmUgPT0gMCkgd2FycF9zdW1bd2lkXSA9IGxvY2FsX3N1bTsKICAgIF9fc3luY3RocmVhZHMoKTsKICAgIGlmICh3aWQgPT0gMCkgewogICAgICAgIGZsb2F0IHYgPSAobGFuZSA8IG53YXJwcykgPyB3YXJwX3N1bVtsYW5lXSA6IDAuZjsKICAgICAgICBmb3IgKGludCBvZmYgPSAxNjsgb2ZmID4gMDsgb2ZmID4+PSAxKSB2ICs9IF9fc2hmbF94b3Jfc3luYygweGZmZmZmZmZmLCB2LCBvZmYpOwogICAgICAgIGlmIChsYW5lID09IDApIHNfc3VtID0gdjsKICAgIH0KICAgIF9fc3luY3RocmVhZHMoKTsKICAgIGZsb2F0IGludiA9IDEuZiAvIHNfc3VtOwogICAgZm9yIChpbnQgaSA9IHRpZDsgaSA8IEs7IGkgKz0gYmxvY2tEaW0ueCkgb3V0W3JvdyAqIEsgKyBpXSA9IHNkYXRhW2ldICogaW52Owp9Cgp0b3JjaDo6VGVuc29yIGJuX3NvZnRtYXhfc2xpY2UodG9yY2g6OlRlbnNvciBsb2dpdHMsIHRvcmNoOjpUZW5zb3Igc2NhbGUsIHRvcmNoOjpUZW5zb3IgYmlhcywgaW50IEspIHsKICAgIGludCByb3dzID0gbG9naXRzLnNpemUoMCksIEtfdG90YWwgPSBsb2dpdHMuc2l6ZSgxKTsKICAgIGF1dG8gb3V0ID0gdG9yY2g6OmVtcHR5KHtyb3dzLCBLfSwgbG9naXRzLm9wdGlvbnMoKSk7CiAgICBibl9zb2Z0bWF4X3NsaWNlX2tlcm5lbDw8PHJvd3MsIDEyOCwgS190b3RhbCAqIHNpemVvZihmbG9hdCksCiAgICAgICAgYXQ6OmN1ZGE6OmdldEN1cnJlbnRDVURBU3RyZWFtKCk+Pj4oCiAgICAgICAgbG9naXRzLmRhdGFfcHRyPGZsb2F0PigpLCBzY2FsZS5kYXRhX3B0cjxmbG9hdD4oKSwgYmlhcy5kYXRhX3B0cjxmbG9hdD4oKSwKICAgICAgICBvdXQuZGF0YV9wdHI8ZmxvYXQ+KCksIHJvd3MsIEtfdG90YWwsIEspOwogICAgcmV0dXJuIG91dDsKfQoKX19nbG9iYWxfXyB2b2lkIHN1bV9hc3NpZ25tZW50X2tlcm5lbF92MigKICAgIGNvbnN0IGZsb2F0KiBfX3Jlc3RyaWN0X18gYXNzaWdubWVudCwgZmxvYXQqIF9fcmVzdHJpY3RfXyBhX3N1bSwgaW50IE4sIGludCBLKSB7CiAgICBpbnQgYiA9IGJsb2NrSWR4Lng7IGludCB0eCA9IHRocmVhZElkeC54OyBpbnQgdHkgPSB0aHJlYWRJZHgueTsgaW50IGJ5ID0gYmxvY2tEaW0ueTsKICAgIGludCBiYXNlID0gYiAqIE4gKiBLOwogICAgZmxvYXQgcyA9IDAuZjsKICAgIGZvciAoaW50IG4gPSB0eTsgbiA8IE47IG4gKz0gYnkpIHMgKz0gYXNzaWdubWVudFtiYXNlICsgbiAqIEsgKyB0eF07CiAgICBleHRlcm4gX19zaGFyZWRfXyBmbG9hdCBzc3VtW107CiAgICBzc3VtW3R5ICogSyArIHR4XSA9IHM7CiAgICBfX3N5bmN0aHJlYWRzKCk7CiAgICBpZiAodHkgPT0gMCkgewogICAgICAgIGZsb2F0IHRvdGFsID0gMC5mOwogICAgICAgIGZvciAoaW50IGogPSAwOyBqIDwgYnk7IGorKykgdG90YWwgKz0gc3N1bVtqICogSyArIHR4XTsKICAgICAgICBhX3N1bVtiICogSyArIHR4XSA9IHRvdGFsOwogICAgfQp9Cgp0b3JjaDo6VGVuc29yIHN1bV9hc3NpZ25tZW50KHRvcmNoOjpUZW5zb3IgYXNzaWdubWVudCwgaW50IEJzeiwgaW50IE4sIGludCBLKSB7CiAgICBhdXRvIGFfc3VtID0gdG9yY2g6OmVtcHR5KHtCc3osIEt9LCBhc3NpZ25tZW50Lm9wdGlvbnMoKSk7CiAgICBpbnQgYnkgPSAyNTYgLyBLOyBpZiAoYnkgPCAxKSBieSA9IDE7IGlmIChieSA+IDMyKSBieSA9IDMyOwogICAgZGltMyBibG9jayhLLCBieSk7CiAgICBzdW1fYXNzaWdubWVudF9rZXJuZWxfdjI8PDxCc3osIGJsb2NrLCAoc2l6ZV90KUsqYnkqc2l6ZW9mKGZsb2F0KSwKICAgICAgICBhdDo6Y3VkYTo6Z2V0Q3VycmVudENVREFTdHJlYW0oKT4+PigKICAgICAgICBhc3NpZ25tZW50LmRhdGFfcHRyPGZsb2F0PigpLCBhX3N1bS5kYXRhX3B0cjxmbG9hdD4oKSwgTiwgSyk7CiAgICByZXR1cm4gYV9zdW07Cn0KCl9fZ2xvYmFsX18gX19sYXVuY2hfYm91bmRzX18oMTI4LCA4KSB2b2lkIGNvbXB1dGVfaW52X2tlcm5lbCgKICAgIGNvbnN0IGZsb2F0KiBfX3Jlc3RyaWN0X18gdmxhZCwgY29uc3QgZmxvYXQqIF9fcmVzdHJpY3RfXyBhX3N1bSwKICAgIGNvbnN0IGZsb2F0KiBfX3Jlc3RyaWN0X18gY2x1c3RlcnMyX3QsIGZsb2F0KiBfX3Jlc3RyaWN0X18gaW52X291dCwgaW50IEssIGludCBEKSB7CiAgICBpbnQgYiA9IGJsb2NrSWR4Lng7IGludCB0aWQgPSB0aHJlYWRJZHgueDsgaW50IGJzID0gYmxvY2tEaW0ueDsKICAgIGV4dGVybiBfX3NoYXJlZF9fIGZsb2F0IHNfYXN1bVtdOwogICAgZm9yIChpbnQgaSA9IHRpZDsgaSA8IEs7IGkgKz0gYnMpIHNfYXN1bVtpXSA9IGFfc3VtW2IgKiBLICsgaV07CiAgICBfX3N5bmN0aHJlYWRzKCk7CiAgICBpbnQgYl9vZmYgPSBiICogSyAqIEQ7IGludCBENCA9IEQgPj4gMjsKICAgIGNvbnN0IGZsb2F0NCogdmxhZDQgPSByZWludGVycHJldF9jYXN0PGNvbnN0IGZsb2F0NCo+KHZsYWQgKyBiX29mZik7CiAgICBjb25zdCBmbG9hdDQqIGM0ID0gcmVpbnRlcnByZXRfY2FzdDxjb25zdCBmbG9hdDQqPihjbHVzdGVyczJfdCk7CiAgICBmbG9hdCBzdW1zcSA9IDAuZjsKICAgIGZvciAoaW50IGsgPSAwOyBrIDwgSzsgaysrKSB7CiAgICAgICAgZmxvYXQgYXMgPSBzX2FzdW1ba107CiAgICAgICAgZm9yIChpbnQgZCA9IHRpZDsgZCA8IEQ0OyBkICs9IGJzKSB7CiAgICAgICAgICAgIGZsb2F0NCB2NCA9IF9fbGRnKCZ2bGFkNFtrKkQ0ICsgZF0pOwogICAgICAgICAgICBmbG9hdDQgY3YgPSBfX2xkZygmYzRbaypENCArIGRdKTsKICAgICAgICAgICAgZmxvYXQgYTA9djQueC1hcypjdi54LCBhMT12NC55LWFzKmN2LnksIGEyPXY0LnotYXMqY3YueiwgYTM9djQudy1hcypjdi53OwogICAgICAgICAgICBzdW1zcSArPSBhMCphMCArIGExKmExICsgYTIqYTIgKyBhMyphMzsKICAgICAgICB9CiAgICB9CiAgICBpbnQgbGFuZSA9IHRpZCAmIDMxOyBpbnQgd2lkID0gdGlkID4+IDU7IGludCBud2FycHMgPSAoYnMrMzEpPj41OwogICAgZm9yIChpbnQgb2ZmPTE2OyBvZmY+MDsgb2ZmPj49MSkgc3Vtc3EgKz0gX19zaGZsX3hvcl9zeW5jKDB4ZmZmZmZmZmYsIHN1bXNxLCBvZmYpOwogICAgX19zaGFyZWRfXyBmbG9hdCB3YXJwX3NbOF07CiAgICBpZiAobGFuZSA9PSAwKSB3YXJwX3Nbd2lkXSA9IHN1bXNxOwogICAgX19zeW5jdGhyZWFkcygpOwogICAgaWYgKHdpZCA9PSAwKSB7CiAgICAgICAgZmxvYXQgdiA9IChsYW5lIDwgbndhcnBzKSA/IHdhcnBfc1tsYW5lXSA6IDAuZjsKICAgICAgICBmb3IgKGludCBvZmY9MTY7IG9mZj4wOyBvZmY+Pj0xKSB2ICs9IF9fc2hmbF94b3Jfc3luYygweGZmZmZmZmZmLCB2LCBvZmYpOwogICAgICAgIGlmIChsYW5lID09IDApIGludl9vdXRbYl0gPSByc3FydGYoZm1heGYodiwgMWUtMjRmKSk7CiAgICB9Cn0KCl9fZ2xvYmFsX18gX19sYXVuY2hfYm91bmRzX18oMjU2LCA2KSB2b2lkIHRyYW5zcG9zZV9ub3JtYWxpemVfa2VybmVsKAogICAgY29uc3QgZmxvYXQqIF9fcmVzdHJpY3RfXyB2bGFkLCBjb25zdCBmbG9hdCogX19yZXN0cmljdF9fIGFfc3VtLAogICAgY29uc3QgZmxvYXQqIF9fcmVzdHJpY3RfXyBjbHVzdGVyczJfdCwgY29uc3QgZmxvYXQqIF9fcmVzdHJpY3RfXyBpbnZfYXJyLAogICAgZmxvYXQqIF9fcmVzdHJpY3RfXyBvdXQsIGludCBLLCBpbnQgRCkgewogICAgaW50IGIgPSBibG9ja0lkeC56OyBpbnQgdGlsZV9rID0gYmxvY2tJZHgueDsgaW50IHRpbGVfZCA9IGJsb2NrSWR4Lnk7CiAgICBpbnQgdHggPSB0aHJlYWRJZHgueDsgaW50IHR5ID0gdGhyZWFkSWR4Lnk7CiAgICBpbnQgazAgPSB0aWxlX2sgKiAzMjsgaW50IGQwID0gdGlsZV9kICogMzI7CiAgICBfX3NoYXJlZF9fIGZsb2F0IHRpbGVbMzJdWzMzXTsKICAgIF9fc2hhcmVkX18gZmxvYXQgc19hc3VtWzMyXTsKICAgIGlmICh0eSA9PSAwKSB7IGludCBrID0gazArdHg7IHNfYXN1bVt0eF0gPSAoazxLKSA/IGFfc3VtW2IqSytrXSA6IDAuZjsgfQogICAgX19zeW5jdGhyZWFkcygpOwogICAgZmxvYXQgaW52ID0gaW52X2FycltiXTsKICAgICNwcmFnbWEgdW5yb2xsCiAgICBmb3IgKGludCByID0gMDsgciA8IDQ7IHIrKykgewogICAgICAgIGludCBrX2xvY2FsID0gdHkqNCtyLCBkX2xvY2FsID0gdHgsIGsgPSBrMCtrX2xvY2FsLCBkID0gZDArZF9sb2NhbDsKICAgICAgICBmbG9hdCB2YWwgPSAwLmY7CiAgICAgICAgaWYgKGsgPCBLICYmIGQgPCBEKSB7CiAgICAgICAgICAgIHZhbCA9IF9fbGRnKCZ2bGFkW2IqSypEICsgaypEICsgZF0pIC0gc19hc3VtW2tfbG9jYWxdKl9fbGRnKCZjbHVzdGVyczJfdFtrKkQrZF0pOwogICAgICAgIH0KICAgICAgICB0aWxlW2tfbG9jYWxdW2RfbG9jYWxdID0gdmFsOwogICAgfQogICAgX19zeW5jdGhyZWFkcygpOwogICAgI3ByYWdtYSB1bnJvbGwKICAgIGZvciAoaW50IHIgPSAwOyByIDwgNDsgcisrKSB7CiAgICAgICAgaW50IGRfbG9jYWwgPSB0eSo0K3IsIGtfbG9jYWwgPSB0eCwgZCA9IGQwK2RfbG9jYWwsIGsgPSBrMCtrX2xvY2FsOwogICAgICAgIGlmIChkIDwgRCAmJiBrIDwgSykgb3V0W2IqRCpLICsgZCpLICsga10gPSB0aWxlW2tfbG9jYWxdW2RfbG9jYWxdICogaW52OwogICAgfQp9Cgp0b3JjaDo6VGVuc29yIGZpbmFsaXplKHRvcmNoOjpUZW5zb3IgdmxhZCwgdG9yY2g6OlRlbnNvciBhX3N1bSwKICAgICAgICAgICAgICAgICAgICAgICB0b3JjaDo6VGVuc29yIGNsdXN0ZXJzMl90LCBpbnQgQnN6LCBpbnQgSywgaW50IEQpIHsKICAgIGF1dG8gb3V0ID0gdG9yY2g6OmVtcHR5KHtCc3osIEQqS30sIHZsYWQub3B0aW9ucygpKTsKICAgIGF1dG8gaW52ID0gdG9yY2g6OmVtcHR5KHtCc3p9LCB2bGFkLm9wdGlvbnMoKSk7CiAgICBjdWRhU3RyZWFtX3Qgc3RyZWFtID0gYXQ6OmN1ZGE6OmdldEN1cnJlbnRDVURBU3RyZWFtKCk7CiAgICBjb21wdXRlX2ludl9rZXJuZWw8PDxCc3osIDEyOCwgSypzaXplb2YoZmxvYXQpLCBzdHJlYW0+Pj4oCiAgICAgICAgdmxhZC5kYXRhX3B0cjxmbG9hdD4oKSwgYV9zdW0uZGF0YV9wdHI8ZmxvYXQ+KCksCiAgICAgICAgY2x1c3RlcnMyX3QuZGF0YV9wdHI8ZmxvYXQ+KCksIGludi5kYXRhX3B0cjxmbG9hdD4oKSwgSywgRCk7CiAgICBkaW0zIGdyaWQoKEsrMzEpLzMyLCAoRCszMSkvMzIsIEJzeik7IGRpbTMgYmxvY2soMzIsIDgpOwogICAgdHJhbnNwb3NlX25vcm1hbGl6ZV9rZXJuZWw8PDxncmlkLCBibG9jaywgMCwgc3RyZWFtPj4+KAogICAgICAgIHZsYWQuZGF0YV9wdHI8ZmxvYXQ+KCksIGFfc3VtLmRhdGFfcHRyPGZsb2F0PigpLAogICAgICAgIGNsdXN0ZXJzMl90LmRhdGFfcHRyPGZsb2F0PigpLCBpbnYuZGF0YV9wdHI8ZmxvYXQ+KCksCiAgICAgICAgb3V0LmRhdGFfcHRyPGZsb2F0PigpLCBLLCBEKTsKICAgIHJldHVybiBvdXQ7Cn0KJycnCgpjcHBfc3JjID0gcicnJwp0b3JjaDo6VGVuc29yIGdlbW1fcnIodG9yY2g6OlRlbnNvciBBLCB0b3JjaDo6VGVuc29yIEIpOwp0b3JjaDo6VGVuc29yIGdlbW1fYmF0Y2hlZF9jcih0b3JjaDo6VGVuc29yIEEsIHRvcmNoOjpUZW5zb3IgQiwgaW50IEJzeiwgaW50IE0sIGludCBOX2lubmVyLCBpbnQgTl9vdXQpOwp0b3JjaDo6VGVuc29yIGJuX3NvZnRtYXhfc2xpY2UodG9yY2g6OlRlbnNvciBsb2dpdHMsIHRvcmNoOjpUZW5zb3Igc2NhbGUsIHRvcmNoOjpUZW5zb3IgYmlhcywgaW50IEspOwp0b3JjaDo6VGVuc29yIHN1bV9hc3NpZ25tZW50KHRvcmNoOjpUZW5zb3IgYXNzaWdubWVudCwgaW50IEJzeiwgaW50IE4sIGludCBLKTsKdG9yY2g6OlRlbnNvciBmaW5hbGl6ZSh0b3JjaDo6VGVuc29yIHZsYWQsIHRvcmNoOjpUZW5zb3IgYV9zdW0sIHRvcmNoOjpUZW5zb3IgY2x1c3RlcnMyX3QsIGludCBCc3osIGludCBLLCBpbnQgRCk7CicnJwoKbW9kID0gbG9hZF9pbmxpbmUoCiAgICBuYW1lPSduZXR2bGFkX29wdDMnLCBjcHBfc291cmNlcz1jcHBfc3JjLCBjdWRhX3NvdXJjZXM9Y3VkYV9zcmMsCiAgICBmdW5jdGlvbnM9WydnZW1tX3JyJywnZ2VtbV9iYXRjaGVkX2NyJywnYm5fc29mdG1heF9zbGljZScsJ3N1bV9hc3NpZ25tZW50JywnZmluYWxpemUnXSwKICAgIGV4dHJhX2N1ZGFfY2ZsYWdzPVsnLUkvcGF0aC90by9jdXRsYXNzL2luY2x1ZGUnLCAnLUkvcGF0aC90by9jdXRsYXNzL3Rvb2xzL3V0aWwvaW5jbHVkZScsCiAgICAgICAgJy0tZXhwdC1yZWxheGVkLWNvbnN0ZXhwcicsICctc3RkPWMrKzE3JywgJy1nZW5jb2RlPWFyY2g9Y29tcHV0ZV84MCxjb2RlPXNtXzgwJywKICAgICAgICAnLU8zJywgJy1ETkRFQlVHJywgJy0tdXNlX2Zhc3RfbWF0aCddLCB2ZXJib3NlPUZhbHNlKQoKY2xhc3MgTW9kZWxOZXcobm4uTW9kdWxlKToKICAgIGRlZiBfX2luaXRfXyhzZWxmLCBjbHVzdGVyX3NpemUsIGZlYXR1cmVfc2l6ZSwgZ2hvc3RfY2x1c3RlcnMpOgogICAgICAgIHN1cGVyKE1vZGVsTmV3LCBzZWxmKS5fX2luaXRfXygpCiAgICAgICAgc2VsZi5mZWF0dXJlX3NpemUgPSBmZWF0dXJlX3NpemUKICAgICAgICBzZWxmLmNsdXN0ZXJfc2l6ZSA9IGNsdXN0ZXJfc2l6ZQogICAgICAgIGluaXRfc2MgPSAoMSAvIG1hdGguc3FydChmZWF0dXJlX3NpemUpKQogICAgICAgIGNsdXN0ZXJzID0gY2x1c3Rlcl9zaXplICsgZ2hvc3RfY2x1c3RlcnMKICAgICAgICBzZWxmLmNsdXN0ZXJzID0gbm4uUGFyYW1ldGVyKGluaXRfc2MgKiB0b3JjaC5yYW5kbihmZWF0dXJlX3NpemUsIGNsdXN0ZXJzKSkKICAgICAgICBzZWxmLmJhdGNoX25vcm0gPSBubi5CYXRjaE5vcm0xZChjbHVzdGVycykKICAgICAgICBzZWxmLmNsdXN0ZXJzMiA9IG5uLlBhcmFtZXRlcihpbml0X3NjICogdG9yY2gucmFuZG4oMSwgZmVhdHVyZV9zaXplLCBjbHVzdGVyX3NpemUpKQogICAgICAgIHNlbGYub3V0X2RpbSA9IHNlbGYuY2x1c3Rlcl9zaXplICogZmVhdHVyZV9zaXplCiAgICAgICAgc2VsZi5fY2FjaGVkID0gRmFsc2UKCiAgICBkZWYgX2J1aWxkX2NhY2hlKHNlbGYsIGRldmljZSwgZHR5cGUpOgogICAgICAgIGJuID0gc2VsZi5iYXRjaF9ub3JtCiAgICAgICAgc2NhbGUgPSAoYm4ud2VpZ2h0IC8gdG9yY2guc3FydChibi5ydW5uaW5nX3ZhciArIGJuLmVwcykpLnRvKGRldmljZT1kZXZpY2UsIGR0eXBlPWR0eXBlKS5jb250aWd1b3VzKCkKICAgICAgICBiaWFzID0gKGJuLmJpYXMgLSBibi5ydW5uaW5nX21lYW4gKiBzY2FsZSkudG8oZGV2aWNlPWRldmljZSwgZHR5cGU9ZHR5cGUpLmNvbnRpZ3VvdXMoKQogICAgICAgIHNlbGYucmVnaXN0ZXJfYnVmZmVyKCdfYm5fc2NhbGUnLCBzY2FsZSwgcGVyc2lzdGVudD1GYWxzZSkKICAgICAgICBzZWxmLnJlZ2lzdGVyX2J1ZmZlcignX2JuX2JpYXMnLCBiaWFzLCBwZXJzaXN0ZW50PUZhbHNlKQogICAgICAgIHNlbGYucmVnaXN0ZXJfYnVmZmVyKCdfY2x1c3RlcnNfYycsIHNlbGYuY2x1c3RlcnMuZGV0YWNoKCkuY29udGlndW91cygpLCBwZXJzaXN0ZW50PUZhbHNlKQogICAgICAgIGMyXzJkID0gc2VsZi5jbHVzdGVyczIuZGV0YWNoKCkudmlldyhzZWxmLmZlYXR1cmVfc2l6ZSwgc2VsZi5jbHVzdGVyX3NpemUpCiAgICAgICAgc2VsZi5yZWdpc3Rlcl9idWZmZXIoJ19jbHVzdGVyczJfdCcsIGMyXzJkLnQoKS5jb250aWd1b3VzKCksIHBlcnNpc3RlbnQ9RmFsc2UpCiAgICAgICAgc2VsZi5fY2FjaGVkID0gVHJ1ZQoKICAgIGRlZiBmb3J3YXJkKHNlbGYsIHgsIG1hc2s9Tm9uZSk6CiAgICAgICAgaWYgbm90IHNlbGYuX2NhY2hlZCBvciBzZWxmLl9ibl9zY2FsZS5kZXZpY2UgIT0geC5kZXZpY2U6CiAgICAgICAgICAgIHNlbGYuX2J1aWxkX2NhY2hlKHguZGV2aWNlLCB4LmR0eXBlKQogICAgICAgIEIsIE4sIEQgPSB4LnNpemUoMCksIHguc2l6ZSgxKSwgc2VsZi5mZWF0dXJlX3NpemUKICAgICAgICBLID0gc2VsZi5jbHVzdGVyX3NpemUKICAgICAgICB4X2ZsYXQgPSB4LmNvbnRpZ3VvdXMoKS52aWV3KC0xLCBEKQogICAgICAgIGxvZ2l0cyA9IG1vZC5nZW1tX3JyKHhfZmxhdCwgc2VsZi5fY2x1c3RlcnNfYykKICAgICAgICBhc3NpZ25tZW50ID0gbW9kLmJuX3NvZnRtYXhfc2xpY2UobG9naXRzLCBzZWxmLl9ibl9zY2FsZSwgc2VsZi5fYm5fYmlhcywgSykKICAgICAgICBhX3N1bSA9IG1vZC5zdW1fYXNzaWdubWVudChhc3NpZ25tZW50LCBCLCBOLCBLKQogICAgICAgIHZsYWQgPSBtb2QuZ2VtbV9iYXRjaGVkX2NyKGFzc2lnbm1lbnQsIHguY29udGlndW91cygpLCBCLCBLLCBOLCBEKQogICAgICAgIHJldHVybiBtb2QuZmluYWxpemUodmxhZCwgYV9zdW0sIHNlbGYuX2NsdXN0ZXJzMl90LCBCLCBLLCBEKQ==)importosos\.environ\[’TORCH\_CUDA\_ARCH\_LIST’\]=’8\.0’importmathimporttorchimporttorch\.nnasnnfromtorch\.utils\.cpp\_extensionimportload\_inlinecuda\_src=r’’’\#include<torch/extension\.h\>\#include<ATen/cuda/CUDAContext\.h\>\#include<cuda\_runtime\.h\>\#include<cutlass/gemm/device/gemm\.h\>\#include<cutlass/gemm/device/gemm\_batched\.h\>usingGemmTF32=cutlass::gemm::device::Gemm<float,cutlass::layout::RowMajor,float,cutlass::layout::RowMajor,float,cutlass::layout::RowMajor,float,cutlass::arch::OpClassTensorOp,cutlass::arch::Sm80,cutlass::gemm::GemmShape<128,128,16\>,cutlass::gemm::GemmShape<64,64,16\>,cutlass::gemm::GemmShape<16,8,8\>\>;usingGemmBatchedTF32=cutlass::gemm::device::GemmBatched<float,cutlass::layout::ColumnMajor,float,cutlass::layout::RowMajor,float,cutlass::layout::RowMajor,float,cutlass::arch::OpClassTensorOp,cutlass::arch::Sm80,cutlass::gemm::GemmShape<64,64,16\>,cutlass::gemm::GemmShape<32,32,16\>,cutlass::gemm::GemmShape<16,8,8\>\>;torch::Tensorgemm\_rr\(torch::TensorA,torch::TensorB\)\{intM=A\.size\(0\),K=A\.size\(1\),N=B\.size\(1\);autoC=torch::empty\(\{M,N\},A\.options\(\)\);GemmTF32gemm;cudaStream\_tstream=at::cuda::getCurrentCUDAStream\(\);GemmTF32::Argumentsargs\(\{M,N,K\},\{A\.data\_ptr<float\>\(\),K\},\{B\.data\_ptr<float\>\(\),N\},\{C\.data\_ptr<float\>\(\),N\},\{C\.data\_ptr<float\>\(\),N\},\{1\.0f,0\.0f\}\);gemm\(args,nullptr,stream\);returnC;\}torch::Tensorgemm\_batched\_cr\(torch::TensorA,torch::TensorB,intBsz,intM,intN\_inner,intN\_out\)\{autoC=torch::empty\(\{Bsz,M,N\_out\},B\.options\(\)\);GemmBatchedTF32gemm;cudaStream\_tstream=at::cuda::getCurrentCUDAStream\(\);GemmBatchedTF32::Argumentsargs\(\{M,N\_out,N\_inner\},\{A\.data\_ptr<float\>\(\),M\},\(int64\_t\)N\_inner\*M,\{B\.data\_ptr<float\>\(\),N\_out\},\(int64\_t\)N\_inner\*N\_out,\{C\.data\_ptr<float\>\(\),N\_out\},\(int64\_t\)M\*N\_out,\{C\.data\_ptr<float\>\(\),N\_out\},\(int64\_t\)M\*N\_out,\{1\.0f,0\.0f\},Bsz\);gemm\(args,nullptr,stream\);returnC;\}\_\_global\_\_voidbn\_softmax\_slice\_kernel\(constfloat\*\_\_restrict\_\_logits,constfloat\*\_\_restrict\_\_scale,constfloat\*\_\_restrict\_\_bias,float\*\_\_restrict\_\_out,introws,intK\_total,intK\)\{introw=blockIdx\.x;inttid=threadIdx\.x;extern\_\_shared\_\_floatsdata\[\];floatlocal\_max=\-INFINITY;for\(inti=tid;i<K\_total;i\+=blockDim\.x\)\{floatv=logits\[row\*K\_total\+i\]\*scale\[i\]\+bias\[i\];sdata\[i\]=v;local\_max=fmaxf\(local\_max,v\);\}\_\_syncthreads\(\);intlane=tid&31;intwid=tid\>\>5;intnwarps=\(blockDim\.x\+31\)\>\>5;for\(intoff=16;off\>0;off\>\>=1\)local\_max=fmaxf\(local\_max,\_\_shfl\_xor\_sync\(0xffffffff,local\_max,off\)\);\_\_shared\_\_floatwarp\_max\[32\];\_\_shared\_\_floats\_max;if\(lane==0\)warp\_max\[wid\]=local\_max;\_\_syncthreads\(\);if\(wid==0\)\{floatv=\(lane<nwarps\)?warp\_max\[lane\]:\-INFINITY;for\(intoff=16;off\>0;off\>\>=1\)v=fmaxf\(v,\_\_shfl\_xor\_sync\(0xffffffff,v,off\)\);if\(lane==0\)s\_max=v;\}\_\_syncthreads\(\);floatlocal\_sum=0\.f;for\(inti=tid;i<K\_total;i\+=blockDim\.x\)\{floate=\_\_expf\(sdata\[i\]\-s\_max\);sdata\[i\]=e;local\_sum\+=e;\}\_\_syncthreads\(\);for\(intoff=16;off\>0;off\>\>=1\)local\_sum\+=\_\_shfl\_xor\_sync\(0xffffffff,local\_sum,off\);\_\_shared\_\_floatwarp\_sum\[32\];\_\_shared\_\_floats\_sum;if\(lane==0\)warp\_sum\[wid\]=local\_sum;\_\_syncthreads\(\);if\(wid==0\)\{floatv=\(lane<nwarps\)?warp\_sum\[lane\]:0\.f;for\(intoff=16;off\>0;off\>\>=1\)v\+=\_\_shfl\_xor\_sync\(0xffffffff,v,off\);if\(lane==0\)s\_sum=v;\}\_\_syncthreads\(\);floatinv=1\.f/s\_sum;for\(inti=tid;i<K;i\+=blockDim\.x\)out\[row\*K\+i\]=sdata\[i\]\*inv;\}torch::Tensorbn\_softmax\_slice\(torch::Tensorlogits,torch::Tensorscale,torch::Tensorbias,intK\)\{introws=logits\.size\(0\),K\_total=logits\.size\(1\);autoout=torch::empty\(\{rows,K\},logits\.options\(\)\);bn\_softmax\_slice\_kernel<<<rows,128,K\_total\*sizeof\(float\),at::cuda::getCurrentCUDAStream\(\)\>\>\>\(logits\.data\_ptr<float\>\(\),scale\.data\_ptr<float\>\(\),bias\.data\_ptr<float\>\(\),out\.data\_ptr<float\>\(\),rows,K\_total,K\);returnout;\}\_\_global\_\_voidsum\_assignment\_kernel\_v2\(constfloat\*\_\_restrict\_\_assignment,float\*\_\_restrict\_\_a\_sum,intN,intK\)\{intb=blockIdx\.x;inttx=threadIdx\.x;intty=threadIdx\.y;intby=blockDim\.y;intbase=b\*N\*K;floats=0\.f;for\(intn=ty;n<N;n\+=by\)s\+=assignment\[base\+n\*K\+tx\];extern\_\_shared\_\_floatssum\[\];ssum\[ty\*K\+tx\]=s;\_\_syncthreads\(\);if\(ty==0\)\{floattotal=0\.f;for\(intj=0;j<by;j\+\+\)total\+=ssum\[j\*K\+tx\];a\_sum\[b\*K\+tx\]=total;\}\}torch::Tensorsum\_assignment\(torch::Tensorassignment,intBsz,intN,intK\)\{autoa\_sum=torch::empty\(\{Bsz,K\},assignment\.options\(\)\);intby=256/K;if\(by<1\)by=1;if\(by\>32\)by=32;dim3block\(K,by\);sum\_assignment\_kernel\_v2<<<Bsz,block,\(size\_t\)K\*by\*sizeof\(float\),at::cuda::getCurrentCUDAStream\(\)\>\>\>\(assignment\.data\_ptr<float\>\(\),a\_sum\.data\_ptr<float\>\(\),N,K\);returna\_sum;\}\_\_global\_\_\_\_launch\_bounds\_\_\(128,8\)voidcompute\_inv\_kernel\(constfloat\*\_\_restrict\_\_vlad,constfloat\*\_\_restrict\_\_a\_sum,constfloat\*\_\_restrict\_\_clusters2\_t,float\*\_\_restrict\_\_inv\_out,intK,intD\)\{intb=blockIdx\.x;inttid=threadIdx\.x;intbs=blockDim\.x;extern\_\_shared\_\_floats\_asum\[\];for\(inti=tid;i<K;i\+=bs\)s\_asum\[i\]=a\_sum\[b\*K\+i\];\_\_syncthreads\(\);intb\_off=b\*K\*D;intD4=D\>\>2;constfloat4\*vlad4=reinterpret\_cast<constfloat4\*\>\(vlad\+b\_off\);constfloat4\*c4=reinterpret\_cast<constfloat4\*\>\(clusters2\_t\);floatsumsq=0\.f;for\(intk=0;k<K;k\+\+\)\{floatas=s\_asum\[k\];for\(intd=tid;d<D4;d\+=bs\)\{float4v4=\_\_ldg\(&vlad4\[k\*D4\+d\]\);float4cv=\_\_ldg\(&c4\[k\*D4\+d\]\);floata0=v4\.x\-as\*cv\.x,a1=v4\.y\-as\*cv\.y,a2=v4\.z\-as\*cv\.z,a3=v4\.w\-as\*cv\.w;sumsq\+=a0\*a0\+a1\*a1\+a2\*a2\+a3\*a3;\}\}intlane=tid&31;intwid=tid\>\>5;intnwarps=\(bs\+31\)\>\>5;for\(intoff=16;off\>0;off\>\>=1\)sumsq\+=\_\_shfl\_xor\_sync\(0xffffffff,sumsq,off\);\_\_shared\_\_floatwarp\_s\[8\];if\(lane==0\)warp\_s\[wid\]=sumsq;\_\_syncthreads\(\);if\(wid==0\)\{floatv=\(lane<nwarps\)?warp\_s\[lane\]:0\.f;for\(intoff=16;off\>0;off\>\>=1\)v\+=\_\_shfl\_xor\_sync\(0xffffffff,v,off\);if\(lane==0\)inv\_out\[b\]=rsqrtf\(fmaxf\(v,1e\-24f\)\);\}\}\_\_global\_\_\_\_launch\_bounds\_\_\(256,6\)voidtranspose\_normalize\_kernel\(constfloat\*\_\_restrict\_\_vlad,constfloat\*\_\_restrict\_\_a\_sum,constfloat\*\_\_restrict\_\_clusters2\_t,constfloat\*\_\_restrict\_\_inv\_arr,float\*\_\_restrict\_\_out,intK,intD\)\{intb=blockIdx\.z;inttile\_k=blockIdx\.x;inttile\_d=blockIdx\.y;inttx=threadIdx\.x;intty=threadIdx\.y;intk0=tile\_k\*32;intd0=tile\_d\*32;\_\_shared\_\_floattile\[32\]\[33\];\_\_shared\_\_floats\_asum\[32\];if\(ty==0\)\{intk=k0\+tx;s\_asum\[tx\]=\(k<K\)?a\_sum\[b\*K\+k\]:0\.f;\}\_\_syncthreads\(\);floatinv=inv\_arr\[b\];\#pragmaunrollfor\(intr=0;r<4;r\+\+\)\{intk\_local=ty\*4\+r,d\_local=tx,k=k0\+k\_local,d=d0\+d\_local;floatval=0\.f;if\(k<K&&d<D\)\{val=\_\_ldg\(&vlad\[b\*K\*D\+k\*D\+d\]\)\-s\_asum\[k\_local\]\*\_\_ldg\(&clusters2\_t\[k\*D\+d\]\);\}tile\[k\_local\]\[d\_local\]=val;\}\_\_syncthreads\(\);\#pragmaunrollfor\(intr=0;r<4;r\+\+\)\{intd\_local=ty\*4\+r,k\_local=tx,d=d0\+d\_local,k=k0\+k\_local;if\(d<D&&k<K\)out\[b\*D\*K\+d\*K\+k\]=tile\[k\_local\]\[d\_local\]\*inv;\}\}torch::Tensorfinalize\(torch::Tensorvlad,torch::Tensora\_sum,torch::Tensorclusters2\_t,intBsz,intK,intD\)\{autoout=torch::empty\(\{Bsz,D\*K\},vlad\.options\(\)\);autoinv=torch::empty\(\{Bsz\},vlad\.options\(\)\);cudaStream\_tstream=at::cuda::getCurrentCUDAStream\(\);compute\_inv\_kernel<<<Bsz,128,K\*sizeof\(float\),stream\>\>\>\(vlad\.data\_ptr<float\>\(\),a\_sum\.data\_ptr<float\>\(\),clusters2\_t\.data\_ptr<float\>\(\),inv\.data\_ptr<float\>\(\),K,D\);dim3grid\(\(K\+31\)/32,\(D\+31\)/32,Bsz\);dim3block\(32,8\);transpose\_normalize\_kernel<<<grid,block,0,stream\>\>\>\(vlad\.data\_ptr<float\>\(\),a\_sum\.data\_ptr<float\>\(\),clusters2\_t\.data\_ptr<float\>\(\),inv\.data\_ptr<float\>\(\),out\.data\_ptr<float\>\(\),K,D\);returnout;\}’’’cpp\_src=r’’’torch::Tensorgemm\_rr\(torch::TensorA,torch::TensorB\);torch::Tensorgemm\_batched\_cr\(torch::TensorA,torch::TensorB,intBsz,intM,intN\_inner,intN\_out\);torch::Tensorbn\_softmax\_slice\(torch::Tensorlogits,torch::Tensorscale,torch::Tensorbias,intK\);torch::Tensorsum\_assignment\(torch::Tensorassignment,intBsz,intN,intK\);torch::Tensorfinalize\(torch::Tensorvlad,torch::Tensora\_sum,torch::Tensorclusters2\_t,intBsz,intK,intD\);’’’mod=load\_inline\(name=’netvlad\_opt3’,cpp\_sources=cpp\_src,cuda\_sources=cuda\_src,functions=\[’gemm\_rr’,’gemm\_batched\_cr’,’bn\_softmax\_slice’,’sum\_assignment’,’finalize’\],extra\_cuda\_cflags=\[’\-I/path/to/cutlass/include’,’\-I/path/to/cutlass/tools/util/include’,’\-\-expt\-relaxed\-constexpr’,’\-std=c\+\+17’,’\-gencode=arch=compute\_80,code=sm\_80’,’\-O3’,’\-DNDEBUG’,’\-\-use\_fast\_math’\],verbose=False\)classModelNew\(nn\.Module\):def\_\_init\_\_\(self,cluster\_size,feature\_size,ghost\_clusters\):super\(ModelNew,self\)\.\_\_init\_\_\(\)self\.feature\_size=feature\_sizeself\.cluster\_size=cluster\_sizeinit\_sc=\(1/math\.sqrt\(feature\_size\)\)clusters=cluster\_size\+ghost\_clustersself\.clusters=nn\.Parameter\(init\_sc\*torch\.randn\(feature\_size,clusters\)\)self\.batch\_norm=nn\.BatchNorm1d\(clusters\)self\.clusters2=nn\.Parameter\(init\_sc\*torch\.randn\(1,feature\_size,cluster\_size\)\)self\.out\_dim=self\.cluster\_size\*feature\_sizeself\.\_cached=Falsedef\_build\_cache\(self,device,dtype\):bn=self\.batch\_normscale=\(bn\.weight/torch\.sqrt\(bn\.running\_var\+bn\.eps\)\)\.to\(device=device,dtype=dtype\)\.contiguous\(\)bias=\(bn\.bias\-bn\.running\_mean\*scale\)\.to\(device=device,dtype=dtype\)\.contiguous\(\)self\.register\_buffer\(’\_bn\_scale’,scale,persistent=False\)self\.register\_buffer\(’\_bn\_bias’,bias,persistent=False\)self\.register\_buffer\(’\_clusters\_c’,self\.clusters\.detach\(\)\.contiguous\(\),persistent=False\)c2\_2d=self\.clusters2\.detach\(\)\.view\(self\.feature\_size,self\.cluster\_size\)self\.register\_buffer\(’\_clusters2\_t’,c2\_2d\.t\(\)\.contiguous\(\),persistent=False\)self\.\_cached=Truedefforward\(self,x,mask=None\):ifnotself\.\_cachedorself\.\_bn\_scale\.device\!=x\.device:self\.\_build\_cache\(x\.device,x\.dtype\)B,N,D=x\.size\(0\),x\.size\(1\),self\.feature\_sizeK=self\.cluster\_sizex\_flat=x\.contiguous\(\)\.view\(\-1,D\)logits=mod\.gemm\_rr\(x\_flat,self\.\_clusters\_c\)assignment=mod\.bn\_softmax\_slice\(logits,self\.\_bn\_scale,self\.\_bn\_bias,K\)a\_sum=mod\.sum\_assignment\(assignment,B,N,K\)vlad=mod\.gemm\_batched\_cr\(assignment,x\.contiguous\(\),B,K,N,D\)returnmod\.finalize\(vlad,a\_sum,self\.\_clusters2\_t,B,K,D\)

## Appendix HVeOmnidW1Generated Kernel \(Full Source\)

##### What the kernel computes\.

In MoE training, each token is routed to a few ofEEexperts and each expert applies its own FFN\. VeOmni packs all routed tokens into a single contiguous buffer ordered by expert and runs a*grouped*GEMM, where the per\-expert row counts are encoded as a cumulative\-sum vectorcumsum\_K\. ThedW1kernel is the backward\-pass*weight gradient*of the first FFN projection \(gate\_up\_proj\): for each experteeit accumulatesdW1e=Xe⊤​Ge\\mathrm\{dW1\}\_\{e\}=X\_\{e\}^\{\\top\}G\_\{e\}over that expert’s token rows, whereXeX\_\{e\}are the expert’s inputs andGeG\_\{e\}the gradient of its pre\-activation\. The contraction lengthKeK\_\{e\}\(the number of tokens routed to expertee\) therefore varies per expert and is the only ragged dimension;M=2880M\{=\}2880andN=5760N\{=\}5760are fixed by the GPT\-OSS\-120B hidden/intermediate sizes\.

##### Estimating the token\-to\-expert distribution\.

Because the speedup depends on the per\-expert load \(theKeK\_\{e\}\), the benchmark must feed realistic routing imbalance rather than a uniform split\. We measured GPT\-OSS\-120B routing on3030LongBench\-v2\(Bai et al\.,[2024](https://arxiv.org/html/2606.26453#bib.bib3)\)samples×\\times3636MoE layers, recomputing each layer’s top\-KKassignment from the router logits and counting tokens per expert\. The per\-expert counts are heavily skewed within a single \(sample, layer\)—coefficient of variationCV≈2\.13\\mathrm\{CV\}\\approx 2\.13, versusCV≈\(E−1\)/R≈0\.011\\mathrm\{CV\}\\approx\\sqrt\{\(E\-1\)/R\}\\approx 0\.011for a uniform split—and balance out only when pooled across inputs and layers \(CV≈0\.277\\mathrm\{CV\}\\approx 0\.277\)\. Fitting generators to the empirical sorted\-load curve byL1L\_\{1\}distance, a Dirichlet–Multinomial fits best \(L1=0\.129L\_\{1\}\{=\}0\.129, vs\.0\.1470\.147for a logistic\-normal\)\. We therefore sample each call’s routing as

p∼Dirichlet​\(α​1E\),α=0\.24,E=128;\(K1,…,KE\)∼Multinomial​\(R,p\),p\\sim\\mathrm\{Dirichlet\}\(\\alpha\\,\\mathbf\{1\}\_\{E\}\),\\quad\\alpha\{=\}0\.24,\\ E\{=\}128;\\qquad\(K\_\{1\},\\dots,K\_\{E\}\)\\sim\\mathrm\{Multinomial\}\(R,\\,p\),withR=57,768R\{=\}57\{,\}768total routed rows\. The smallα\\alphareproduces the measured heavy tail—a few hot experts, a long warm shoulder, and several empty experts \(Ke=0K\_\{e\}\{=\}0\)—exercising the short\-tail and empty\-expert cases a uniform\-shape GEMM never encounters\.

##### The generated kernel\.

The box below lists the completeKernelPro\-generated CUDA source for the winningdW1kernel \(Section[5\.4](https://arxiv.org/html/2606.26453#S5.SS4)\), which attains1\.23×1\.23\\timesover the expert\-tuned Triton baseline on H100\.KernelPro’s Programmer Agent generated it at the CuTe source level—composing copy atoms, an MMA atom, and swizzled layouts into a novel kernel rather than instantiating a high\-level CUTLASS GEMM template or calling a prebuilt library \(no cuBLAS, nocutlass::gemm::device\)\. The only GEMM primitive iscute::gemmover a singleSM90\_64x128x16\_F32F16F16\_SSwarp\-group MMA atom, fed by a four\-stagecp\.asyncshared\-memory pipeline withGMMA::Layout\_MN\_SW128\_Atomswizzling; the ragged per\-expert contractionKeK\_\{e\}is handled by CuTe layout\-algebra predication \(make\_identity\_tensor\+copy\_if\), including the empty\-expert \(Ke=0K\_\{e\}\{=\}0\) early\-out\. The MMA atom is selected through aMmaSeltemplate specialized for bothhalf\_tandbfloat16\_t\(one reusable kernel, not a one\-off\), and a*vectorized SMEM\-staged epilogue*restages accumulators into swizzled shared memory \(GMMA::Layout\_K\_SW128\_Atom\) and stores them with predicated 128\-bitUniversalCopy<uint128\_t\>writes to coalesce the output\. This is direct evidence forKernelPro’s native raw\-CUDA \+ CuTe source\-level generation \(Contribution 4\): the kernel is composed from CuTe primitives, not produced by template instantiation or library calls\. The listing shows only the device kernel and its launcher; the \(boilerplate\) PyTorch binding is omitted\.

KernelPro– VeOmnidW1Weight\-Gradient Grouped GEMM—1\.23×1\.23\\timesSpeedup[⬇](data:text/plain;base64,I2luY2x1ZGUgPHRvcmNoL2V4dGVuc2lvbi5oPgojaW5jbHVkZSA8QVRlbi9jdWRhL0NVREFDb250ZXh0Lmg+CiNpbmNsdWRlIDxjdWRhX3J1bnRpbWUuaD4KCiNpbmNsdWRlIDxjdXRlL3RlbnNvci5ocHA+CiNpbmNsdWRlIDxjdXRlL2F0b20vbW1hX2F0b20uaHBwPgojaW5jbHVkZSA8Y3V0ZS9hdG9tL2NvcHlfYXRvbS5ocHA+CiNpbmNsdWRlIDxjdXRlL2FsZ29yaXRobS9nZW1tLmhwcD4KI2luY2x1ZGUgPGN1dGxhc3MvbnVtZXJpY190eXBlcy5oPgoKdXNpbmcgbmFtZXNwYWNlIGN1dGU7Cgp0ZW1wbGF0ZSA8Y2xhc3MgVD4gc3RydWN0IE1tYVNlbDsKdGVtcGxhdGUgPD4gc3RydWN0IE1tYVNlbDxjdXRsYXNzOjpoYWxmX3Q+IHsKICB1c2luZyB0eXBlID0gU005MF82NHgxMjh4MTZfRjMyRjE2RjE2X1NTPEdNTUE6Ok1ham9yOjpNTiwgR01NQTo6TWFqb3I6Ok1OPjsKfTsKdGVtcGxhdGUgPD4gc3RydWN0IE1tYVNlbDxjdXRsYXNzOjpiZmxvYXQxNl90PiB7CiAgdXNpbmcgdHlwZSA9IFNNOTBfNjR4MTI4eDE2X0YzMkJGMTZCRjE2X1NTPEdNTUE6Ok1ham9yOjpNTiwgR01NQTo6TWFqb3I6Ok1OPjsKfTsKCnRlbXBsYXRlIDxpbnQgQk0sIGludCBCTiwgaW50IEJLLCBpbnQgTlMsIGNsYXNzIFRBLCBjbGFzcyBUQiwgY2xhc3MgVEM+Cl9fZ2xvYmFsX18gX19sYXVuY2hfYm91bmRzX18oMTI4LCAzKQp2b2lkIGR3MV9rZXJuZWwoaW50IE0sIGludCBOLAogICAgICAgICAgICAgICAgY29uc3QgVEEqIF9fcmVzdHJpY3RfXyBhLAogICAgICAgICAgICAgICAgY29uc3QgVEIqIF9fcmVzdHJpY3RfXyBiLAogICAgICAgICAgICAgICAgVEMqIF9fcmVzdHJpY3RfXyBjLAogICAgICAgICAgICAgICAgY29uc3QgbG9uZyBsb25nKiBfX3Jlc3RyaWN0X18gY3Vtc3VtX0spCnsKICBpbnQgZSA9IGJsb2NrSWR4Lno7CiAgbG9uZyBsb25nIGtfZW5kICAgPSBjdW1zdW1fS1tlXTsKICBsb25nIGxvbmcga19zdGFydCA9IChlID09IDApID8gMExMIDogY3Vtc3VtX0tbZS0xXTsKICBpbnQgSyA9IChpbnQpKGtfZW5kIC0ga19zdGFydCk7CgogIGNvbnN0IFRBKiBBID0gYSArIGtfc3RhcnQgKiAobG9uZyBsb25nKU07CiAgY29uc3QgVEIqIEIgPSBiICsga19zdGFydCAqIChsb25nIGxvbmcpTjsKICBUQyogICAgICAgQyA9IGMgKyAobG9uZyBsb25nKWUgKiAobG9uZyBsb25nKU0gKiAobG9uZyBsb25nKU47CgogIGF1dG8gY3RhX3RpbGVyID0gbWFrZV9zaGFwZShJbnQ8Qk0+e30sIEludDxCTj57fSwgSW50PEJLPnt9KTsKCiAgVGVuc29yIG1BID0gbWFrZV90ZW5zb3IobWFrZV9nbWVtX3B0cihBKSwgbWFrZV9zaGFwZShNLCBLKSwgbWFrZV9zdHJpZGUoSW50PDE+e30sIE0pKTsKICBUZW5zb3IgbUIgPSBtYWtlX3RlbnNvcihtYWtlX2dtZW1fcHRyKEIpLCBtYWtlX3NoYXBlKE4sIEspLCBtYWtlX3N0cmlkZShJbnQ8MT57fSwgTikpOwogIFRlbnNvciBtQyA9IG1ha2VfdGVuc29yKG1ha2VfZ21lbV9wdHIoQyksIG1ha2Vfc2hhcGUoTSwgTiksIG1ha2Vfc3RyaWRlKE4sIEludDwxPnt9KSk7CgogIGF1dG8gY3RhX2Nvb3JkID0gbWFrZV9jb29yZChibG9ja0lkeC54LCBibG9ja0lkeC55LCBfKTsKICBUZW5zb3IgZ0EgPSBsb2NhbF90aWxlKG1BLCBjdGFfdGlsZXIsIGN0YV9jb29yZCwgU3RlcDxfMSwgWCwgXzE+e30pOwogIFRlbnNvciBnQiA9IGxvY2FsX3RpbGUobUIsIGN0YV90aWxlciwgY3RhX2Nvb3JkLCBTdGVwPCBYLF8xLCBfMT57fSk7CiAgVGVuc29yIGdDID0gbG9jYWxfdGlsZShtQywgY3RhX3RpbGVyLCBjdGFfY29vcmQsIFN0ZXA8XzEsXzEsICBYPnt9KTsKCiAgVGVuc29yIGNBaWQgPSBsb2NhbF90aWxlKG1ha2VfaWRlbnRpdHlfdGVuc29yKG1ha2Vfc2hhcGUoTSwgSykpLCBjdGFfdGlsZXIsIGN0YV9jb29yZCwgU3RlcDxfMSwgWCwgXzE+e30pOwogIFRlbnNvciBjQmlkID0gbG9jYWxfdGlsZShtYWtlX2lkZW50aXR5X3RlbnNvcihtYWtlX3NoYXBlKE4sIEspKSwgY3RhX3RpbGVyLCBjdGFfY29vcmQsIFN0ZXA8IFgsXzEsIF8xPnt9KTsKICBUZW5zb3IgY0NpZCA9IGxvY2FsX3RpbGUobWFrZV9pZGVudGl0eV90ZW5zb3IobWFrZV9zaGFwZShNLCBOKSksIGN0YV90aWxlciwgY3RhX2Nvb3JkLCBTdGVwPF8xLF8xLCAgWD57fSk7CgogIGF1dG8gc0FfbGF5b3V0ID0gdGlsZV90b19zaGFwZShHTU1BOjpMYXlvdXRfTU5fU1cxMjhfQXRvbTxUQT57fSwgbWFrZV9zaGFwZShJbnQ8Qk0+e30sIEludDxCSz57fSwgSW50PE5TPnt9KSk7CiAgYXV0byBzQl9sYXlvdXQgPSB0aWxlX3RvX3NoYXBlKEdNTUE6OkxheW91dF9NTl9TVzEyOF9BdG9tPFRCPnt9LCBtYWtlX3NoYXBlKEludDxCTj57fSwgSW50PEJLPnt9LCBJbnQ8TlM+e30pKTsKCiAgZXh0ZXJuIF9fc2hhcmVkX18gY2hhciBzbWVtX3Jhd1tdOwogIGNvbnN0ZXhwciBzaXplX3QgYUJ5dGVzID0gKHNpemVfdClCTSAqIEJLICogTlMgKiBzaXplb2YoVEEpOwogIGNvbnN0ZXhwciBzaXplX3Qgb2ZmQiAgID0gKChhQnl0ZXMgKyAxMjcpIC8gMTI4KSAqIDEyODsKICBUQSogc0FwdHIgPSByZWludGVycHJldF9jYXN0PFRBKj4oc21lbV9yYXcpOwogIFRCKiBzQnB0ciA9IHJlaW50ZXJwcmV0X2Nhc3Q8VEIqPihzbWVtX3JhdyArIG9mZkIpOwoKICBUZW5zb3Igc0EgPSBtYWtlX3RlbnNvcihtYWtlX3NtZW1fcHRyKHNBcHRyKSwgc0FfbGF5b3V0KTsKICBUZW5zb3Igc0IgPSBtYWtlX3RlbnNvcihtYWtlX3NtZW1fcHRyKHNCcHRyKSwgc0JfbGF5b3V0KTsKCiAgdXNpbmcgQ3BBID0gQ29weV9BdG9tPFNNODBfQ1BfQVNZTkNfQ0FDSEVHTE9CQUxfWkZJTEw8dWludDEyOF90PiwgVEE+OwogIHVzaW5nIENwQiA9IENvcHlfQXRvbTxTTTgwX0NQX0FTWU5DX0NBQ0hFR0xPQkFMX1pGSUxMPHVpbnQxMjhfdD4sIFRCPjsKICBUaWxlZENvcHkgY29weUEgPSBtYWtlX3RpbGVkX2NvcHkoQ3BBe30sIExheW91dDxTaGFwZTxfMTYsXzg+Pnt9LCBMYXlvdXQ8U2hhcGU8XzgsXzE+Pnt9KTsKICBUaWxlZENvcHkgY29weUIgPSBtYWtlX3RpbGVkX2NvcHkoQ3BCe30sIExheW91dDxTaGFwZTxfMTYsXzg+Pnt9LCBMYXlvdXQ8U2hhcGU8XzgsXzE+Pnt9KTsKCiAgVGhyQ29weSB0Y2EgPSBjb3B5QS5nZXRfc2xpY2UodGhyZWFkSWR4LngpOwogIFRockNvcHkgdGNiID0gY29weUIuZ2V0X3NsaWNlKHRocmVhZElkeC54KTsKICBUZW5zb3IgdEFnQSA9IHRjYS5wYXJ0aXRpb25fUyhnQSk7CiAgVGVuc29yIHRCZ0IgPSB0Y2IucGFydGl0aW9uX1MoZ0IpOwogIFRlbnNvciBzQV93ID0gYXNfcG9zaXRpb25faW5kZXBlbmRlbnRfc3dpenpsZV90ZW5zb3Ioc0EpOwogIFRlbnNvciBzQl93ID0gYXNfcG9zaXRpb25faW5kZXBlbmRlbnRfc3dpenpsZV90ZW5zb3Ioc0IpOwogIFRlbnNvciB0QXNBID0gdGNhLnBhcnRpdGlvbl9EKHNBX3cpOwogIFRlbnNvciB0QnNCID0gdGNiLnBhcnRpdGlvbl9EKHNCX3cpOwogIFRlbnNvciB0QWNBID0gdGNhLnBhcnRpdGlvbl9TKGNBaWQpOwogIFRlbnNvciB0QmNCID0gdGNiLnBhcnRpdGlvbl9TKGNCaWQpOwoKICBUZW5zb3IgdEFwQSA9IG1ha2VfdGVuc29yPGJvb2w+KHNoYXBlKHRBc0EoXyxfLF8sMCkpKTsKICBUZW5zb3IgdEJwQiA9IG1ha2VfdGVuc29yPGJvb2w+KHNoYXBlKHRCc0IoXyxfLF8sMCkpKTsKCiAgdXNpbmcgQXRvbSA9IHR5cGVuYW1lIE1tYVNlbDxUQT46OnR5cGU7CiAgVGlsZWRNTUEgbW1hID0gbWFrZV90aWxlZF9tbWEoQXRvbXt9KTsKICBUaHJNTUEgdGhyX21tYSA9IG1tYS5nZXRfc2xpY2UodGhyZWFkSWR4LngpOwogIFRlbnNvciB0Q3NBID0gdGhyX21tYS5wYXJ0aXRpb25fQShzQSk7CiAgVGVuc29yIHRDc0IgPSB0aHJfbW1hLnBhcnRpdGlvbl9CKHNCKTsKICBUZW5zb3IgdENnQyA9IHRocl9tbWEucGFydGl0aW9uX0MoZ0MpOwogIFRlbnNvciB0Q3JBID0gdGhyX21tYS5tYWtlX2ZyYWdtZW50X0EodENzQSk7CiAgVGVuc29yIHRDckIgPSB0aHJfbW1hLm1ha2VfZnJhZ21lbnRfQih0Q3NCKTsKICBUZW5zb3IgdENyQyA9IHRocl9tbWEubWFrZV9mcmFnbWVudF9DKHRDZ0MpOwogIGNsZWFyKHRDckMpOwoKICBpbnQga190aWxlcyA9IChLICsgQksgLSAxKSAvIEJLOwoKICBhdXRvIGZpbGxBID0gWyZdKGludCBrdCkgewogICAgVGVuc29yIGNjID0gdEFjQShfLF8sXyxrdCk7CiAgICBDVVRFX1VOUk9MTAogICAgZm9yIChpbnQgaSA9IDA7IGkgPCBzaXplKHRBcEEpOyArK2kpIHRBcEEoaSkgPSBlbGVtX2xlc3MoY2MoaSksIG1ha2VfY29vcmQoTSwgSykpOwogIH07CiAgYXV0byBmaWxsQiA9IFsmXShpbnQga3QpIHsKICAgIFRlbnNvciBjYyA9IHRCY0IoXyxfLF8sa3QpOwogICAgQ1VURV9VTlJPTEwKICAgIGZvciAoaW50IGkgPSAwOyBpIDwgc2l6ZSh0QnBCKTsgKytpKSB0QnBCKGkpID0gZWxlbV9sZXNzKGNjKGkpLCBtYWtlX2Nvb3JkKE4sIEspKTsKICB9OwoKICBpbnQgc21lbV9waXBlX3dyaXRlID0gMDsKICBpbnQgc21lbV9waXBlX3JlYWQgID0gMDsKCiAgQ1VURV9VTlJPTEwKICBmb3IgKGludCBzID0gMDsgcyA8IE5TIC0gMTsgKytzKSB7CiAgICBpZiAocyA8IGtfdGlsZXMpIHsKICAgICAgZmlsbEEocyk7IGNvcHlfaWYoY29weUEsIHRBcEEsIHRBZ0EoXyxfLF8scyksIHRBc0EoXyxfLF8sc21lbV9waXBlX3dyaXRlKSk7CiAgICAgIGZpbGxCKHMpOyBjb3B5X2lmKGNvcHlCLCB0QnBCLCB0QmdCKF8sXyxfLHMpLCB0QnNCKF8sXyxfLHNtZW1fcGlwZV93cml0ZSkpOwogICAgfQogICAgY3BfYXN5bmNfZmVuY2UoKTsKICAgIHNtZW1fcGlwZV93cml0ZSA9IChzbWVtX3BpcGVfd3JpdGUgKyAxKSAlIE5TOwogIH0KCiAgQ1VURV9OT19VTlJPTEwKICBmb3IgKGludCBrdCA9IDA7IGt0IDwga190aWxlczsgKytrdCkgewogICAgY3BfYXN5bmNfd2FpdDxOUyAtIDI+KCk7CiAgICBfX3N5bmN0aHJlYWRzKCk7CgogICAgaW50IGtsb2FkID0ga3QgKyAoTlMgLSAxKTsKICAgIGlmIChrbG9hZCA8IGtfdGlsZXMpIHsKICAgICAgZmlsbEEoa2xvYWQpOyBjb3B5X2lmKGNvcHlBLCB0QXBBLCB0QWdBKF8sXyxfLGtsb2FkKSwgdEFzQShfLF8sXyxzbWVtX3BpcGVfd3JpdGUpKTsKICAgICAgZmlsbEIoa2xvYWQpOyBjb3B5X2lmKGNvcHlCLCB0QnBCLCB0QmdCKF8sXyxfLGtsb2FkKSwgdEJzQihfLF8sXyxzbWVtX3BpcGVfd3JpdGUpKTsKICAgIH0KICAgIGNwX2FzeW5jX2ZlbmNlKCk7CiAgICBzbWVtX3BpcGVfd3JpdGUgPSAoc21lbV9waXBlX3dyaXRlICsgMSkgJSBOUzsKCiAgICB3YXJwZ3JvdXBfZmVuY2Vfb3BlcmFuZCh0Q3JDKTsKICAgIHdhcnBncm91cF9hcnJpdmUoKTsKICAgIGN1dGU6OmdlbW0obW1hLCB0Q3JBKF8sXyxfLHNtZW1fcGlwZV9yZWFkKSwgdENyQihfLF8sXyxzbWVtX3BpcGVfcmVhZCksIHRDckMpOwogICAgd2FycGdyb3VwX2NvbW1pdF9iYXRjaCgpOwogICAgd2FycGdyb3VwX3dhaXQ8MD4oKTsKICAgIHdhcnBncm91cF9mZW5jZV9vcGVyYW5kKHRDckMpOwoKICAgIHNtZW1fcGlwZV9yZWFkID0gKHNtZW1fcGlwZV9yZWFkICsgMSkgJSBOUzsKICB9CgogIC8vIC0tLS0tLS0tLS0gVmVjdG9yaXplZCBTTUVNLXN0YWdlZCBlcGlsb2d1ZSAoY29hbGVzY2VkIDEyOC1iaXQgc3RvcmVzKSAtLS0tLS0tLS0tCiAgYXV0byBzQ19sYXlvdXQgPSB0aWxlX3RvX3NoYXBlKEdNTUE6OkxheW91dF9LX1NXMTI4X0F0b208VEM+e30sIG1ha2Vfc2hhcGUoSW50PEJNPnt9LCBJbnQ8Qk4+e30pKTsKICBUZW5zb3Igc0MgPSBtYWtlX3RlbnNvcihtYWtlX3NtZW1fcHRyKHJlaW50ZXJwcmV0X2Nhc3Q8VEMqPihzbWVtX3JhdykpLCBzQ19sYXlvdXQpOwoKICBUZW5zb3IgdENyQ19vdXQgPSBtYWtlX3RlbnNvcjxUQz4oc2hhcGUodENyQykpOwogIENVVEVfVU5ST0xMCiAgZm9yIChpbnQgaSA9IDA7IGkgPCBzaXplKHRDckMpOyArK2kpIHRDckNfb3V0KGkpID0gc3RhdGljX2Nhc3Q8VEM+KGZsb2F0KHRDckMoaSkpKTsKCiAgVGVuc29yIHRDc0MgPSB0aHJfbW1hLnBhcnRpdGlvbl9DKHNDKTsKICBfX3N5bmN0aHJlYWRzKCk7CiAgY29weSh0Q3JDX291dCwgdENzQyk7CiAgX19zeW5jdGhyZWFkcygpOwoKICB1c2luZyBDcEMgPSBDb3B5X0F0b208VW5pdmVyc2FsQ29weTx1aW50MTI4X3Q+LCBUQz47CiAgVGlsZWRDb3B5IGNvcHlDID0gbWFrZV90aWxlZF9jb3B5KENwQ3t9LAogICAgICBMYXlvdXQ8U2hhcGU8XzgsXzE2PiwgU3RyaWRlPF8xNixfMT4+e30sCiAgICAgIExheW91dDxTaGFwZTxfMSxfOD4+e30pOwogIFRockNvcHkgdGNjID0gY29weUMuZ2V0X3NsaWNlKHRocmVhZElkeC54KTsKICBUZW5zb3IgdENzQzIgPSB0Y2MucGFydGl0aW9uX1Moc0MpOwogIFRlbnNvciB0Q2dDMiA9IHRjYy5wYXJ0aXRpb25fRChnQyk7CiAgVGVuc29yIHRDY0MyID0gdGNjLnBhcnRpdGlvbl9EKGNDaWQpOwogIFRlbnNvciB0Q3BDICA9IG1ha2VfdGVuc29yPGJvb2w+KHNoYXBlKHRDc0MyKF8sXyxfKSkpOwogIENVVEVfVU5ST0xMCiAgZm9yIChpbnQgaSA9IDA7IGkgPCBzaXplKHRDcEMpOyArK2kpIHRDcEMoaSkgPSBlbGVtX2xlc3ModENjQzIoaSksIG1ha2VfY29vcmQoTSwgTikpOwogIGNvcHlfaWYoY29weUMsIHRDcEMsIHRDc0MyLCB0Q2dDMik7Cn0KCnRlbXBsYXRlIDxjbGFzcyBUQSwgY2xhc3MgVEIsIGNsYXNzIFRDPgp2b2lkIGxhdW5jaCh0b3JjaDo6VGVuc29yIGEsIHRvcmNoOjpUZW5zb3IgYiwgdG9yY2g6OlRlbnNvciBjLCB0b3JjaDo6VGVuc29yIGN1bXN1bV9LLAogICAgICAgICAgICBpbnQgTSwgaW50IE4sIGludCBFKSB7CiAgY29uc3RleHByIGludCBCTSA9IDEyOCwgQk4gPSAxMjgsIEJLID0gMzIsIE5TID0gNDsKICBzaXplX3QgYUJ5dGVzID0gKHNpemVfdClCTSAqIEJLICogTlMgKiBzaXplb2YoVEEpOwogIHNpemVfdCBvZmZCICAgPSAoKGFCeXRlcyArIDEyNykgLyAxMjgpICogMTI4OwogIHNpemVfdCBiQnl0ZXMgPSAoc2l6ZV90KUJOICogQksgKiBOUyAqIHNpemVvZihUQik7CiAgc2l6ZV90IG1haW5sb29wX3NtZW0gPSBvZmZCICsgYkJ5dGVzOwogIHNpemVfdCBlcGlfc21lbSA9IChzaXplX3QpQk0gKiBCTiAqIHNpemVvZihUQyk7CiAgaW50IHNtZW0gPSAoaW50KShtYWlubG9vcF9zbWVtID4gZXBpX3NtZW0gPyBtYWlubG9vcF9zbWVtIDogZXBpX3NtZW0pOwoKICBkaW0zIGdyaWQoKE0gKyBCTSAtIDEpIC8gQk0sIChOICsgQk4gLSAxKSAvIEJOLCBFKTsKICBkaW0zIGJsb2NrKDEyOCk7CgogIGF1dG8ga3B0ciA9ICZkdzFfa2VybmVsPEJNLCBCTiwgQkssIE5TLCBUQSwgVEIsIFRDPjsKICBjdWRhRnVuY1NldEF0dHJpYnV0ZShrcHRyLCBjdWRhRnVuY0F0dHJpYnV0ZU1heER5bmFtaWNTaGFyZWRNZW1vcnlTaXplLCBzbWVtKTsKCiAgYXV0byBzdHJlYW0gPSBhdDo6Y3VkYTo6Z2V0Q3VycmVudENVREFTdHJlYW0oKTsKICBrcHRyPDw8Z3JpZCwgYmxvY2ssIHNtZW0sIHN0cmVhbT4+PigKICAgICAgTSwgTiwKICAgICAgcmVpbnRlcnByZXRfY2FzdDxjb25zdCBUQSo+KGEuZGF0YV9wdHIoKSksCiAgICAgIHJlaW50ZXJwcmV0X2Nhc3Q8Y29uc3QgVEIqPihiLmRhdGFfcHRyKCkpLAogICAgICByZWludGVycHJldF9jYXN0PFRDKj4oYy5kYXRhX3B0cigpKSwKICAgICAgcmVpbnRlcnByZXRfY2FzdDxjb25zdCBsb25nIGxvbmcqPihjdW1zdW1fSy5kYXRhX3B0cjxpbnQ2NF90PigpKSk7Cn0KCnRvcmNoOjpUZW5zb3IgZHcxX2ZvcndhcmQodG9yY2g6OlRlbnNvciBhLCB0b3JjaDo6VGVuc29yIGIsIHRvcmNoOjpUZW5zb3IgY3Vtc3VtX0ssCiAgICAgICAgICAgICAgICAgICAgICAgICAgaW50NjRfdCBudW1fZXhwZXJ0cywgaW50NjRfdCBNLCBpbnQ2NF90IE4pIHsKICBUT1JDSF9DSEVDSyhhLmlzX2N1ZGEoKSAmJiBiLmlzX2N1ZGEoKSwgImlucHV0cyBtdXN0IGJlIENVREEiKTsKICBhID0gYS5jb250aWd1b3VzKCk7CiAgYiA9IGIuY29udGlndW91cygpOwogIGN1bXN1bV9LID0gY3Vtc3VtX0sudG8odG9yY2g6OmtJbnQ2NCkuY29udGlndW91cygpOwogIGludCBFID0gKGludCludW1fZXhwZXJ0czsKICBhdXRvIG91dCA9IHRvcmNoOjplbXB0eSh7RSwgTSwgTn0sIGEub3B0aW9ucygpKTsKICBpZiAoYS5zY2FsYXJfdHlwZSgpID09IHRvcmNoOjprSGFsZikgewogICAgbGF1bmNoPGN1dGxhc3M6OmhhbGZfdCwgY3V0bGFzczo6aGFsZl90LCBjdXRsYXNzOjpoYWxmX3Q+KGEsIGIsIG91dCwgY3Vtc3VtX0ssIChpbnQpTSwgKGludClOLCBFKTsKICB9IGVsc2UgaWYgKGEuc2NhbGFyX3R5cGUoKSA9PSB0b3JjaDo6a0JGbG9hdDE2KSB7CiAgICBsYXVuY2g8Y3V0bGFzczo6YmZsb2F0MTZfdCwgY3V0bGFzczo6YmZsb2F0MTZfdCwgY3V0bGFzczo6YmZsb2F0MTZfdD4oYSwgYiwgb3V0LCBjdW1zdW1fSywgKGludClNLCAoaW50KU4sIEUpOwogIH0gZWxzZSB7CiAgICBUT1JDSF9DSEVDSyhmYWxzZSwgInVuc3VwcG9ydGVkIGR0eXBlIik7CiAgfQogIHJldHVybiBvdXQ7Cn0=)\#include<torch/extension\.h\>\#include<ATen/cuda/CUDAContext\.h\>\#include<cuda\_runtime\.h\>\#include<cute/tensor\.hpp\>\#include<cute/atom/mma\_atom\.hpp\>\#include<cute/atom/copy\_atom\.hpp\>\#include<cute/algorithm/gemm\.hpp\>\#include<cutlass/numeric\_types\.h\>usingnamespacecute;template<classT\>structMmaSel;template<\>structMmaSel<cutlass::half\_t\>\{usingtype=SM90\_64x128x16\_F32F16F16\_SS<GMMA::Major::MN,GMMA::Major::MN\>;\};template<\>structMmaSel<cutlass::bfloat16\_t\>\{usingtype=SM90\_64x128x16\_F32BF16BF16\_SS<GMMA::Major::MN,GMMA::Major::MN\>;\};template<intBM,intBN,intBK,intNS,classTA,classTB,classTC\>\_\_global\_\_\_\_launch\_bounds\_\_\(128,3\)voiddw1\_kernel\(intM,intN,constTA\*\_\_restrict\_\_a,constTB\*\_\_restrict\_\_b,TC\*\_\_restrict\_\_c,constlonglong\*\_\_restrict\_\_cumsum\_K\)\{inte=blockIdx\.z;longlongk\_end=cumsum\_K\[e\];longlongk\_start=\(e==0\)?0LL:cumsum\_K\[e\-1\];intK=\(int\)\(k\_end\-k\_start\);constTA\*A=a\+k\_start\*\(longlong\)M;constTB\*B=b\+k\_start\*\(longlong\)N;TC\*C=c\+\(longlong\)e\*\(longlong\)M\*\(longlong\)N;autocta\_tiler=make\_shape\(Int<BM\>\{\},Int<BN\>\{\},Int<BK\>\{\}\);TensormA=make\_tensor\(make\_gmem\_ptr\(A\),make\_shape\(M,K\),make\_stride\(Int<1\>\{\},M\)\);TensormB=make\_tensor\(make\_gmem\_ptr\(B\),make\_shape\(N,K\),make\_stride\(Int<1\>\{\},N\)\);TensormC=make\_tensor\(make\_gmem\_ptr\(C\),make\_shape\(M,N\),make\_stride\(N,Int<1\>\{\}\)\);autocta\_coord=make\_coord\(blockIdx\.x,blockIdx\.y,\_\);TensorgA=local\_tile\(mA,cta\_tiler,cta\_coord,Step<\_1,X,\_1\>\{\}\);TensorgB=local\_tile\(mB,cta\_tiler,cta\_coord,Step<X,\_1,\_1\>\{\}\);TensorgC=local\_tile\(mC,cta\_tiler,cta\_coord,Step<\_1,\_1,X\>\{\}\);TensorcAid=local\_tile\(make\_identity\_tensor\(make\_shape\(M,K\)\),cta\_tiler,cta\_coord,Step<\_1,X,\_1\>\{\}\);TensorcBid=local\_tile\(make\_identity\_tensor\(make\_shape\(N,K\)\),cta\_tiler,cta\_coord,Step<X,\_1,\_1\>\{\}\);TensorcCid=local\_tile\(make\_identity\_tensor\(make\_shape\(M,N\)\),cta\_tiler,cta\_coord,Step<\_1,\_1,X\>\{\}\);autosA\_layout=tile\_to\_shape\(GMMA::Layout\_MN\_SW128\_Atom<TA\>\{\},make\_shape\(Int<BM\>\{\},Int<BK\>\{\},Int<NS\>\{\}\)\);autosB\_layout=tile\_to\_shape\(GMMA::Layout\_MN\_SW128\_Atom<TB\>\{\},make\_shape\(Int<BN\>\{\},Int<BK\>\{\},Int<NS\>\{\}\)\);extern\_\_shared\_\_charsmem\_raw\[\];constexprsize\_taBytes=\(size\_t\)BM\*BK\*NS\*sizeof\(TA\);constexprsize\_toffB=\(\(aBytes\+127\)/128\)\*128;TA\*sAptr=reinterpret\_cast<TA\*\>\(smem\_raw\);TB\*sBptr=reinterpret\_cast<TB\*\>\(smem\_raw\+offB\);TensorsA=make\_tensor\(make\_smem\_ptr\(sAptr\),sA\_layout\);TensorsB=make\_tensor\(make\_smem\_ptr\(sBptr\),sB\_layout\);usingCpA=Copy\_Atom<SM80\_CP\_ASYNC\_CACHEGLOBAL\_ZFILL<uint128\_t\>,TA\>;usingCpB=Copy\_Atom<SM80\_CP\_ASYNC\_CACHEGLOBAL\_ZFILL<uint128\_t\>,TB\>;TiledCopycopyA=make\_tiled\_copy\(CpA\{\},Layout<Shape<\_16,\_8\>\>\{\},Layout<Shape<\_8,\_1\>\>\{\}\);TiledCopycopyB=make\_tiled\_copy\(CpB\{\},Layout<Shape<\_16,\_8\>\>\{\},Layout<Shape<\_8,\_1\>\>\{\}\);ThrCopytca=copyA\.get\_slice\(threadIdx\.x\);ThrCopytcb=copyB\.get\_slice\(threadIdx\.x\);TensortAgA=tca\.partition\_S\(gA\);TensortBgB=tcb\.partition\_S\(gB\);TensorsA\_w=as\_position\_independent\_swizzle\_tensor\(sA\);TensorsB\_w=as\_position\_independent\_swizzle\_tensor\(sB\);TensortAsA=tca\.partition\_D\(sA\_w\);TensortBsB=tcb\.partition\_D\(sB\_w\);TensortAcA=tca\.partition\_S\(cAid\);TensortBcB=tcb\.partition\_S\(cBid\);TensortApA=make\_tensor<bool\>\(shape\(tAsA\(\_,\_,\_,0\)\)\);TensortBpB=make\_tensor<bool\>\(shape\(tBsB\(\_,\_,\_,0\)\)\);usingAtom=typenameMmaSel<TA\>::type;TiledMMAmma=make\_tiled\_mma\(Atom\{\}\);ThrMMAthr\_mma=mma\.get\_slice\(threadIdx\.x\);TensortCsA=thr\_mma\.partition\_A\(sA\);TensortCsB=thr\_mma\.partition\_B\(sB\);TensortCgC=thr\_mma\.partition\_C\(gC\);TensortCrA=thr\_mma\.make\_fragment\_A\(tCsA\);TensortCrB=thr\_mma\.make\_fragment\_B\(tCsB\);TensortCrC=thr\_mma\.make\_fragment\_C\(tCgC\);clear\(tCrC\);intk\_tiles=\(K\+BK\-1\)/BK;autofillA=\[&\]\(intkt\)\{Tensorcc=tAcA\(\_,\_,\_,kt\);CUTE\_UNROLLfor\(inti=0;i<size\(tApA\);\+\+i\)tApA\(i\)=elem\_less\(cc\(i\),make\_coord\(M,K\)\);\};autofillB=\[&\]\(intkt\)\{Tensorcc=tBcB\(\_,\_,\_,kt\);CUTE\_UNROLLfor\(inti=0;i<size\(tBpB\);\+\+i\)tBpB\(i\)=elem\_less\(cc\(i\),make\_coord\(N,K\)\);\};intsmem\_pipe\_write=0;intsmem\_pipe\_read=0;CUTE\_UNROLLfor\(ints=0;s<NS\-1;\+\+s\)\{if\(s<k\_tiles\)\{fillA\(s\);copy\_if\(copyA,tApA,tAgA\(\_,\_,\_,s\),tAsA\(\_,\_,\_,smem\_pipe\_write\)\);fillB\(s\);copy\_if\(copyB,tBpB,tBgB\(\_,\_,\_,s\),tBsB\(\_,\_,\_,smem\_pipe\_write\)\);\}cp\_async\_fence\(\);smem\_pipe\_write=\(smem\_pipe\_write\+1\)%NS;\}CUTE\_NO\_UNROLLfor\(intkt=0;kt<k\_tiles;\+\+kt\)\{cp\_async\_wait<NS\-2\>\(\);\_\_syncthreads\(\);intkload=kt\+\(NS\-1\);if\(kload<k\_tiles\)\{fillA\(kload\);copy\_if\(copyA,tApA,tAgA\(\_,\_,\_,kload\),tAsA\(\_,\_,\_,smem\_pipe\_write\)\);fillB\(kload\);copy\_if\(copyB,tBpB,tBgB\(\_,\_,\_,kload\),tBsB\(\_,\_,\_,smem\_pipe\_write\)\);\}cp\_async\_fence\(\);smem\_pipe\_write=\(smem\_pipe\_write\+1\)%NS;warpgroup\_fence\_operand\(tCrC\);warpgroup\_arrive\(\);cute::gemm\(mma,tCrA\(\_,\_,\_,smem\_pipe\_read\),tCrB\(\_,\_,\_,smem\_pipe\_read\),tCrC\);warpgroup\_commit\_batch\(\);warpgroup\_wait<0\>\(\);warpgroup\_fence\_operand\(tCrC\);smem\_pipe\_read=\(smem\_pipe\_read\+1\)%NS;\}autosC\_layout=tile\_to\_shape\(GMMA::Layout\_K\_SW128\_Atom<TC\>\{\},make\_shape\(Int<BM\>\{\},Int<BN\>\{\}\)\);TensorsC=make\_tensor\(make\_smem\_ptr\(reinterpret\_cast<TC\*\>\(smem\_raw\)\),sC\_layout\);TensortCrC\_out=make\_tensor<TC\>\(shape\(tCrC\)\);CUTE\_UNROLLfor\(inti=0;i<size\(tCrC\);\+\+i\)tCrC\_out\(i\)=static\_cast<TC\>\(float\(tCrC\(i\)\)\);TensortCsC=thr\_mma\.partition\_C\(sC\);\_\_syncthreads\(\);copy\(tCrC\_out,tCsC\);\_\_syncthreads\(\);usingCpC=Copy\_Atom<UniversalCopy<uint128\_t\>,TC\>;TiledCopycopyC=make\_tiled\_copy\(CpC\{\},Layout<Shape<\_8,\_16\>,Stride<\_16,\_1\>\>\{\},Layout<Shape<\_1,\_8\>\>\{\}\);ThrCopytcc=copyC\.get\_slice\(threadIdx\.x\);TensortCsC2=tcc\.partition\_S\(sC\);TensortCgC2=tcc\.partition\_D\(gC\);TensortCcC2=tcc\.partition\_D\(cCid\);TensortCpC=make\_tensor<bool\>\(shape\(tCsC2\(\_,\_,\_\)\)\);CUTE\_UNROLLfor\(inti=0;i<size\(tCpC\);\+\+i\)tCpC\(i\)=elem\_less\(tCcC2\(i\),make\_coord\(M,N\)\);copy\_if\(copyC,tCpC,tCsC2,tCgC2\);\}template<classTA,classTB,classTC\>voidlaunch\(torch::Tensora,torch::Tensorb,torch::Tensorc,torch::Tensorcumsum\_K,intM,intN,intE\)\{constexprintBM=128,BN=128,BK=32,NS=4;size\_taBytes=\(size\_t\)BM\*BK\*NS\*sizeof\(TA\);size\_toffB=\(\(aBytes\+127\)/128\)\*128;size\_tbBytes=\(size\_t\)BN\*BK\*NS\*sizeof\(TB\);size\_tmainloop\_smem=offB\+bBytes;size\_tepi\_smem=\(size\_t\)BM\*BN\*sizeof\(TC\);intsmem=\(int\)\(mainloop\_smem\>epi\_smem?mainloop\_smem:epi\_smem\);dim3grid\(\(M\+BM\-1\)/BM,\(N\+BN\-1\)/BN,E\);dim3block\(128\);autokptr=&dw1\_kernel<BM,BN,BK,NS,TA,TB,TC\>;cudaFuncSetAttribute\(kptr,cudaFuncAttributeMaxDynamicSharedMemorySize,smem\);autostream=at::cuda::getCurrentCUDAStream\(\);kptr<<<grid,block,smem,stream\>\>\>\(M,N,reinterpret\_cast<constTA\*\>\(a\.data\_ptr\(\)\),reinterpret\_cast<constTB\*\>\(b\.data\_ptr\(\)\),reinterpret\_cast<TC\*\>\(c\.data\_ptr\(\)\),reinterpret\_cast<constlonglong\*\>\(cumsum\_K\.data\_ptr<int64\_t\>\(\)\)\);\}torch::Tensordw1\_forward\(torch::Tensora,torch::Tensorb,torch::Tensorcumsum\_K,int64\_tnum\_experts,int64\_tM,int64\_tN\)\{TORCH\_CHECK\(a\.is\_cuda\(\)&&b\.is\_cuda\(\),"inputsmustbeCUDA"\);a=a\.contiguous\(\);b=b\.contiguous\(\);cumsum\_K=cumsum\_K\.to\(torch::kInt64\)\.contiguous\(\);intE=\(int\)num\_experts;autoout=torch::empty\(\{E,M,N\},a\.options\(\)\);if\(a\.scalar\_type\(\)==torch::kHalf\)\{launch<cutlass::half\_t,cutlass::half\_t,cutlass::half\_t\>\(a,b,out,cumsum\_K,\(int\)M,\(int\)N,E\);\}elseif\(a\.scalar\_type\(\)==torch::kBFloat16\)\{launch<cutlass::bfloat16\_t,cutlass::bfloat16\_t,cutlass::bfloat16\_t\>\(a,b,out,cumsum\_K,\(int\)M,\(int\)N,E\);\}else\{TORCH\_CHECK\(false,"unsupporteddtype"\);\}returnout;\}

## Appendix IEnergy\-Aware Optimization: Methodology

This appendix expands the energy\-aware extension of Section[5\.5](https://arxiv.org/html/2606.26453#S5.SS5): howKernelPromeasures per\-kernel energy \(§[I\.1](https://arxiv.org/html/2606.26453#A9.SS1)\), how energy enters the search as a secondary objective \(§[I\.2](https://arxiv.org/html/2606.26453#A9.SS2)\), the energy\-only micro\-profiling tools added \(§[I\.3](https://arxiv.org/html/2606.26453#A9.SS3)\), and the literature on dominant energy\-heavy GPU bottlenecks with theKernelProtool targeting each \(§[I\.4](https://arxiv.org/html/2606.26453#A9.SS4)\)\.

### I\.1NVML Millijoule Measurement Protocol

KernelPromeasures real GPU energy via NVIDIA’s on\-board counternvmlDeviceGetTotalEnergyConsumption\(\), a monotonic millijoule accumulator and the only direct per\-device energy counter NVIDIA exposes\. Naive use is unreliable: on A100/H100 the on\-board power sensor samples only∼\\sim25% of the runtime, leaving the rest unmonitored, so accounting for this duty\-cycling can change measured energy by an average of35%35\\%\(up to65%65\\%\) relative to an external power meter\(Yang et al\.,[2024](https://arxiv.org/html/2606.26453#bib.bib34)\)\. We therefore measure under a controlled protocol whenever we report millijoules\. We lock GPU clocks \(e\.g\. 1410 MHz on A100\) to remove DVFS as a confound, warm up to thermal steady state, and use windows long enough \(∼\\sim2 s\) to take a stable median over 30 repetitions\. To recover the kernel’s own consumption we subtract idle draw, reporting*dynamic*energyEdyn=Eraw−Pidle​tE\_\{\\text\{dyn\}\}=E\_\{\\text\{raw\}\}\-P\_\{\\text\{idle\}\}\\,t\(Pidle≈76\.7P\_\{\\text\{idle\}\}\\approx 76\.7W on our A100,ttthe measured runtime\), removing the device\-wide constant draw NVML always includes\. Crucially, because NVML’s per\-process baseline drifts, we measure both arms’ kernels back\-to\-back in a single session\.

NVML yields robust*device\-total*dynamic energy but cannot attribute it to a single SM or to one kernel within a fused graph—so during search we rank candidates with the deterministic counter\-based proxy \(§[I\.2](https://arxiv.org/html/2606.26453#A9.SS2)\), reserving this protocol for post\-hoc validation of winning kernels\.

### I\.2Energy as a Lexicographic Secondary Objective

The reward extension \(Eq\.[3](https://arxiv.org/html/2606.26453#S5.E3)\) places energy under strict lexicographic priority:ε\\varepsilonis small enough that any speedup gain dominates any energy gain, soKernelPronever trades speed for energy—energy only differentiates speed\-equivalent candidates\. The extension is a single environment switch \(IFCO\_ENERGY\_AWARE\) that activates the reward term, an energy\-preferences prompt section, and the energy\-only tools; a passive mode \(IFCO\_ENERGY\_MEASURE\) records energy*without*acting on it, giving a control arm whose search remains byte\-identical to standardKernelPro\. Search strategy, tool filtering, profiling pipeline, and candidate generation are otherwise unchanged\.

Because NVML is device\-wide and noisy, theenergy\_reductionterm in the reward is the*deterministic*picojoule proxyEproxyE\_\{\\text\{proxy\}\}of Eq\.[4](https://arxiv.org/html/2606.26453#S5.E4)\(Section[5\.5](https://arxiv.org/html/2606.26453#S5.SS5)\), computed fromncucountersKernelProalready collects and weighted by the per\-operation energy hierarchy of Horowitz\(Horowitz,[2014](https://arxiv.org/html/2606.26453#bib.bib12)\)and AccelWattch\(Kandiah et al\.,[2021](https://arxiv.org/html/2606.26453#bib.bib16)\)\. The reduction ratioEproxy​\(baseline\)/Eproxy​\(candidate\)E\_\{\\text\{proxy\}\}\(\\text\{baseline\}\)/E\_\{\\text\{proxy\}\}\(\\text\{candidate\}\)is intrinsic to the kernel and independent of wall\-clock time, so it rewards reductions in hardware activity that timing alone cannot see\.

### I\.3Energy\-Only Micro\-Profiling Tools

KernelPro’s profiling runs a registry of micro\-profiling tools filtered by the Stage\-1 speed bottleneck\. The energy extension adds a separate registry partition that bypasses this filter and always fires \(energy waste is relevant regardless of the speed bottleneck\), routing findings to a subordinate “Energy Efficiency Notes” prompt section\. Each tool fires only when its pattern is present*and*latency\-hidden \(the corresponding stall metric is low\)—i\.e\. it targets energy waste that timing does*not*already penalize, since any pattern that hurt speed is already handled by the speed tools\. Severity is capped below that of speed\-critical findings, enforcing the lexicographic priority\. Table[25](https://arxiv.org/html/2606.26453#A9.T25)lists the four tools\.

Table 25:Energy\-only micro\-profiling tools\. Each fires only when latency\-hidden, so it targets energy waste invisible to timing\.TheSwishwin of Section[5\.5](https://arxiv.org/html/2606.26453#S5.SS5)came from the reward’s instruction\-count term \(B4\) rather than any of these four tools; surfacing that signal directly at the prompt level—e\.g\. a tool that reports SASS instruction count and flags IEEE operations replaceable by fast\-math intrinsics under a bandwidth\-bound roofline—is a natural extension we leave to future work\.

### I\.4Energy\-Heavy GPU Bottlenecks and Targeting Tools

Table[26](https://arxiv.org/html/2606.26453#A9.T26)summarizes the GPU energy\-efficiency literature, ranking the dominant sources of dynamic kernel energy and mapping each to theKernelPromechanism that targets it\. The dominant*kernel\-level*lever is reducing off\-chip data movement \(B1\): a DRAM access costs∼\\sim100–200×\\timesa register access, so data movement, not arithmetic, dominates energy\(Horowitz,[2014](https://arxiv.org/html/2606.26453#bib.bib12); Leng et al\.,[2013](https://arxiv.org/html/2606.26453#bib.bib19); Hong & Kim,[2010](https://arxiv.org/html/2606.26453#bib.bib11)\)\. This lever is largely*speed\-coupled*\(less data movement is also faster\), so the speed\-only search already captures most of it; the energy\-aware objective targets the*residual*—energy reducible at fixed speed \(bank conflicts, redundant barriers, coalescing/precision that leave wall\-clock unchanged\)\. Frequency/DVFS and occupancy, the largest*speed\-independent*levers reported in the literature, are runtime knobs outsideKernelPro’s kernel\-generation scope\.

Table 26:Dominant sources of dynamic energy in GPU kernels \(literature\) and theKernelPromechanism targeting each\. B1–B4 are the proxy terms of §[I\.2](https://arxiv.org/html/2606.26453#A9.SS2)\.