Change
93218ca
93218cab6ce9d1ce0a0801f347309e7f576bc644 · commit on GitHub
pytorch-blog-feed: changed (253509 bytes, HTTP 200)
raw/pytorch-blog-feed/response.xml modified
- Source
- pytorch-blog-feed
- Lines added
- +568
- Lines removed
- -336
- Stored bytes at this commit
- 253,509
- Timestamp
- origin
- Raw artifact at this commit
- raw/pytorch-blog-feed/response.xml
Recorded headers
| observed_at | 2026-10-20T04:51:23.408Z |
|---|---|
| origin_date | 2026-10-20T04:39:34.000Z |
| status | 200 |
| final URL | https://pytorch.org/blog/feed/ |
| etag | "3688e36a9e1502dd79af47474a39ecdf" |
| last-modified | Tue, 20 Oct 2026 04:20:36 GMT |
| date | Tue, 20 Oct 2026 04:51:23 GMT |
| age | 709 |
| cache-control | public, max-age=60, s-maxage=43200, stale-while-revalidate=86400, stale-if-error=604800 |
| cf-cache-status | null |
| content-encoding | null |
| content-length | 253509 |
@
@@ -12,7 +12,7 @@ <atom:link href="https://pytorch.org/blog/feed/" rel="self" type="application/rss+xml" /> <link>https://pytorch.org/blog/</link> <description></description>-
<lastBuildDate>Fri, 16 Oct 2026 16:30:46 +0000</lastBuildDate>+
<lastBuildDate>Mon, 19 Oct 2026 22:53:51 +0000</lastBuildDate> <language>en-US</language> <sy:updatePeriod> hourly </sy:updatePeriod>@
@@ -28,6 +28,570 @@ <height>32</height></image> <item>+
<title>One Search, Two Layers: Compiler Controls Inside PyTorch Helion using NVIDIA CompileIQ</title>+
<link>https://pytorch.org/blog/one-search-two-layers-compiler-controls-inside-pytorch-helion-using-nvidia-compileiq/</link>+
+
<dc:creator><![CDATA[Aditya Srikanth (NVIDIA), Rupak Roy (NVIDIA), Abhinav Singh (NVIDIA), Jason Ansel (Meta), Clement Anthonioz-Blanc (Meta) and Jongsok Choi (Meta)]]></dc:creator>+
<pubDate>Mon, 19 Oct 2026 15:45:05 +0000</pubDate>+
<category><![CDATA[Blog]]></category>+
<guid isPermaLink="false">https://pytorch.org/?p=188657</guid>+
+
<description><![CDATA[<p>With Helion’s release last year, it aimed at enabling developers to write GPU kernels in PyTorch-native Python and let its powerful autotuner do the work that used to require hand-written...</p>+
<p>The post <a href="https://pytorch.org/blog/one-search-two-layers-compiler-controls-inside-pytorch-helion-using-nvidia-compileiq/">One Search, Two Layers: Compiler Controls Inside PyTorch Helion using NVIDIA CompileIQ</a> appeared first on <a href="https://pytorch.org">PyTorch</a>.</p>+
]]></description>+
<content:encoded><![CDATA[<p>With <a href="https://pytorch.org/blog/helion/">Helion’s</a> release last year, it aimed at enabling developers to write GPU kernels in PyTorch-native Python and let its powerful <a href="https://pytorch.org/blog/accelerating-autotuning-in-helion/">autotuner</a…+
<p>The question that now needs to be answered is this: If a kernel has already been autotuned and optimized to the best of a developer’s ability, is there anything left?</p>+
<p>There is. It lives below the kernel, in the compiler.</p>+
<p>Every kernel autotuner in the ecosystem whether it is Helion’s, Triton’s, the one inside torch.compile, they all search the same thing: tile shapes, warp counts, pipeline stages, memory access strategy. All of those are decisions about what work each thread block does. Underneath them…+
<p>Helion now enables this, and this post is about what that second compiler-specific space contains, why searching it jointly with the kernel configuration space shows further improvements on kernels that were already deemed to be optimal.</p>+
<p>The searching is done by NVIDIA’s CompileIQ, which optimizes ptxas controls the way an autotuner optimizes tile sizes and packages the result as a portable Advanced Control File (ACF). Most users won’t run that search themselves: CompileIQ publishes pre-tuned ACF catalogs that Helion&…+
<h2>Starting off with highly-optimized kernels & optimizing the last-mile</h2>+
<p>It matters a great deal what baseline one compares against. A “we made kernels faster” claim is easy if the starting point is a naive kernel, and close to meaningless for real-world workloads.</p>+
<p>The kernels we use here come from the <a href="https://github.com/pytorch/helion/tree/main/pretuned_kernels">pretuned_kernels/</a> in the Helion repository, roughly 27 kernels covering the shapes that show up in real model code: FP8 and NVFP4 quantization, RMSNorm variants, RoPE, grouped GEMM, ca…+
<p>The important property to note here is that these are not autotuned on the fly. Each one ships with ahead-of-time heuristic files per architecture: _helion_aot_rms_norm_cuda_sm90.py for H100, _helion_aot_rms_norm_cuda_sm100.py for B200 etc. that Helion selects at runtime. The autotuning search al…+
<p>That is our baseline: a kernel whose configuration space has already been swept, on the hardware it’s running on, by the kernel authors and the Helion compiler developers. If there’s headroom below that, it’s <em>real</em> headroom and that’s the last-mile focus of this post.</p…+
<h2>What Helion’s autotuner searches, and where it stops</h2>+
<p>A Helion kernel is ordinary PyTorch with one new idea – tiles:</p>+
<pre><code class="language-python">import torch, helion, helion.language as hl
+
+
@helion.kernel()
+
def softmax(x: torch.Tensor) -> torch.Tensor:
+
n, _m = x.size()
+
out = torch.empty_like(x)
+
for tile_n in hl.tile(n):
+
out[tile_n, :] = torch.nn.functional.softmax(x[tile_n, :], dim=1)
+
return out
+
</code></pre>+
<p>The user never says how big a tile is. That’s the point: hl.tile declares that the iteration space is tiled, and the autotuner decides the rest. A single kernel definition implicitly generates a search space of thousands of Triton configurations that span block_sizes, loop_orders, l2_groupi…+
<p>Then the winning Triton kernel goes to ptxas, and everything from that point down happens at PTXAS’ default settings. The kernel autotuner’s search space is defined by what the kernel language can express, and no kernel language expresses “allocate registers differently”. This <…+
<p>It is imperative to mention that this exploration is genuinely hard. There are complex interactions between various compiler controls: a change to register allocation shifts occupancy, which changes whether a scheduling decision helps or hurts. The right setting is workload-specific, and there…+
<p><!-- IMAGE PLACEHOLDER: image1 - re-upload from the Google Doc via the WordPress media library and insert here --></p>+
<h2><a href="https://pytorch.org/wp-content/uploads/2026/10/lowering-stack-2.svg"><img fetchpriority="high" decoding="async" class="alignnone wp-image-188664 size-full" src="https://pytorch.org/wp-content/uploads/2026/10/lowering-stack-2.svg" alt="" width="1000" height="640" /></a></h2>+
<h2>The joint search: Helion Autotuner and PTXAS’ Advanced Controls</h2>+
<p>Helion exposes the underlying compiler control space through two settings. Controls are packaged in an <a href="https://nvidia.github.io/CompileIQ/v1.0/booster_packs.html">Advanced Control File</a> (ACF, typically ending with .acf) which is a binary artifact applied at compile time, portable and …+
<p>The first setting pins a known-good ACF as part of a config the user has already chosen:</p>+
<pre><code class="language-python">softmax_tuned = helion.kernel(
+
config=helion.Config(
+
advanced_controls_file="h100_example1.acf",
+
block_sizes=[4],
+
num_warps=8,
+
num_stages=2,
+
reduction_loops=[512],
+
indexing=["pointer", "pointer", "tensor_descriptor",
+
"tensor_descriptor", "pointer", "pointer"],
+
),
+
static_shapes=True,
+
)(_softmax)
+
</code></pre>+
<p>The second is a more interesting case. It hands a list of ACFs to Helion’s own autotuner and lets them be searched as just another dimension, alongside tile sizes and warp counts:</p>+
<pre><code class="language-python">softmax_search = helion.kernel(
+
autotune_search_acf=[
+
"acfs/helion-booster-pack/candidate_03.acf",
+
"acfs/helion-booster-pack/candidate_11.acf",
+
"acfs/helion-booster-pack/candidate_17.acf",
+
],
+
autotune_compile_timeout=120,
+
static_shapes=True,
+
)(_softmax)
+
</code></pre>+
<p>Two things are worth pointing out in this mode:</p>+
<ul>+
<li>The empty string is appended for you. Helion automatically adds “”, meaning “no ACF, compiler defaults”, to the list that is supplied. The default configuration is therefore always a candidate in the search, and the autotuner will only select an ACF if it actually won on …+
<li>Sequential tuning finds a local optimum; this joint method doesn’t. This is the substantive argument for doing it in one loop rather than two. Suppose you autotune the kernel first, get block_sizes=[128, 128] with num_stages=4, and then go searching for compiler controls. You are now searc…+
</ul>+
<p>The cost is search time. Each ACF multiplies the number of compile-and-measure trials, so this is not something to leave on by default for every kernel in a codebase. It is something to run once, offline, on a kernel that matters, and then pin the result with advanced_controls_file, exactly the w…+
<h2>Where do these Advanced Control Files come from</h2>+
<p>Now that we’ve understood the benefits of joint autotuning of Helion parameters with PTXAS’ advanced control files (ACF), the natural question one may have is how can these ACFs be generated. That is where NVIDIA’s <a href="https://github.com/NVIDIA/CompileIQ">CompileIQ</a> comes into the picture…+
<p>CompileIQ treats compiler-control selection as a hyperparameter optimization problem. It pulls a curated search space for the target compiler and toolkit version, then runs an evolutionary search: propose a population of candidate control sets, compile each one, check correctness against a refere…+
<p>For most users however, the additional search budget is unnecessary, because the searching has been done. CompileIQ publishes <a href="https://nvidia.github.io/CompileIQ/v1.0/booster_packs.html">Booster Packs</a> – pre-tuned ACF catalogs – on its <a href="https://github.com/NVIDIA/Com…+
<p>If you have a workload the pack doesn’t cover, pip install compileiq and run your own targeted search. That path is documented in NVIDIA’s <a href="https://developer.nvidia.com/blog/extract-more-kernel-performance-with-nvidia-compileiq-auto-tuning">CompileIQ blog post</a> and in vario…+
<h2>Helion and CompileIQ: Experimental Results</h2>+
<p>This section highlights many possible ways in which one can run a CompileIQ optimization campaign on Helion kernels. As a sample case study we applied CompileIQ to the pretuned nvfp4_gemv_fp4in_kernel, a Helion W4A4 (NVFP4 weight x NVFP4 activation) decode GEMV kernel that already ships a checked…+
<h3>Environment and Methodology</h3>+
<p>Benchmarking was conducted on a GB200 system with sm_100 (Blackwell, compute capability 10.0), utilizing ptxas 13.3.73 (CUDA 13.3), CompileIQ 1.0.3, Triton 3.8.0, and an editable source installation of Helion. All timing evaluations were executed via triton.testing.do_bench with the GPU clock loc…+
<h3>Kernel Description</h3>+
<p>nvfp4_gemv_fp4in_kernel in helion/pretuned_kernels/nvfp4_gemv/nvfp4_gemv.py, computing out = (weight @ x) * alpha for a decode (one output row) GEMV. Both the packed NVFP4 weight and activation bytes are unpacked per K-group via hl.load_float4_e2m1fn_x16_to_float16, accumulated in fp16 over a 16-…+
<p>This kernel includes a checked-in AOT configuration specifically for the N=8192, K=28672 “down_proj” decode-GEMV shape in _helion_aot_nvfp4_gemv_cuda_sm100.py:</p>+
<pre><code class="language-python">_CONFIG_FP4IN_DEFAULT = {
+
"block_sizes": [4, 256],
+
"num_warps": 2,
+
"num_stages": 3,
+
"range_multi_buffers": [None, True],
+
}
+
</code></pre>+
<h3>Using Runtime as CompileIQ Objective Function</h3>+
<p>Using runtime as the CompileIQ objective function lets us optimize execution speed directly. The objective function takes a candidate ACF configuration, applies it to our fixed Helion kernel, and evaluates it: the kernel runs once, its output is checked against the PyTorch reference implementatio…+
<pre><code class="language-python">def objective(config) -> float:
+
acf_path = ... # write `config` to a .acf, or None for baseline
+
kernel = _make_kernel(acf_path)
+
got = run()
+
_check(got, ref, "fp4in", N, K) # correctness gate
+
return do_bench(run, warmup=25, rep=100, return_mode="mean")
+
</code></pre>+
<p>We initialize a Search instance that couples our scoring routine to an evolutionary autotuner over PTXAS compiler controls, then invoke .start(). It sweeps candidate ACFs across generations generations of a pool_size-sized population, breeding the top-performing control configurations forward int…+
<pre><code class="language-python">config = SearchConfiguration(problem_type="min", generations=8, pool_size=15)
+
tuner = Search(
+
objective_function=objective,
+
search_space=PtxasSearchSpace(version=cuda_version, variant="default"),
+
search_config=config,
+
dump_results="nvfp4_gemv_results.csv",
+
worker_type=WorkerTypes.ISOLATED, # crashing ACF candidates shouldn't kill the search
+
)
+
+
with gpu_benchmark_mode(clock_mhz=target_clock_mhz, raise_on_failure=False):
+
results = tuner.start(task_timeout=60)
+
</code></pre>+
<p>Executing an evolutionary search with CompileIQ across PTXAS ACFs, spanning 8 generations with a pool size of 15 candidates over our fixed Helion configuration on GB200 yielded an optimized ACF control set. This candidate reduced execution time from a baseline of 0.0971 ms down to 0.0928 ms, deli…+
<pre><code class="language-plaintext">sample baseline run: 0.09713 ms
+
Best runtime (CompileIQ Search): 0.0928 ms
+
Speedup: 1.047x
+
</code></pre>+
<h3>Validation of CompileIQ Search Results</h3>+
<p>To confirm our findings, every speedup is re-measured independently using increased warmup and repetition iterations under fixed GPU clock frequencies. Each configuration undergoes strict correctness verification against the reference model alongside a Welch’s t-test and Cohen’s d fil…+
<pre><code class="language-plaintext"># run 1
+
baseline: n=1744 mean=0.09752ms std=0.00218
+
optimized: n=2018 mean=0.09301ms std=0.00133
+
speedup_mean: 1.0485
+
p_value: 0.0
+
cohens_d: 2.501
+
significant: True
+
+
# run 2 (independent re-run)
+
baseline: n=1757 mean=0.09744ms std=0.00216
+
optimized: n=2028 mean=0.09299ms std=0.00130
+
speedup_mean: 1.0478
+
p_value: 0.0
+
cohens_d: 2.489
+
significant: True
+
</code></pre>+
<p>The two executions remain virtually identical within experimental variance: with d ≈ 2.5 demonstrating a massive effect size that easily surpasses our statistical threshold of 0.2 , verifying that this ACF performance gain is genuine.</p>+
<p>The two runs are virtually identical within experimental variance, and d ≈ 2.5 in both is a massive effect size nearly 12.5x past our d > 0.2 gate, confirming this ACF performance gain is genuine, not noise.</p>+
<h3>Tuning against a hardware metric instead of runtime</h3>+
<p>While execution speed is a common optimization goal, CompileIQ is not limited to runtime alone. Integrating directly with Nsight Compute (ncu) enables the objective function to target any hardware metric exposed by the tool. To determine the key architectural bottleneck limiting kernel performanc…+
<pre><code class="language-plaintext">Section: GPU Speed of Light Throughput
+
Memory Throughput % 78.99
+
L1/TEX Cache Throughput % 89.35
+
Compute (SM) Throughput % 37.29
+
+
Section: Occupancy
+
Theoretical Occupancy % 25.0
+
Achieved Occupancy % 21.63
+
+
Est. Speedup: 21.01%
+
... theoretical occupancy (25.0%) is limited by the number of required
+
registers.
+
</code></pre>+
<p>Profiling shows the kernel is register-limited. NCU’s own cost model attributes a 21% ‘Est. Speedup’ opportunity to occupancy, and this ceiling is set by register count, something PTXAS’s register-allocation controls can directly influence.</p>+
<p>Switching the objective from runtime to occupancy only touches two things: what the objective function scores, and the search’s problem_type:</p>+
<pre><code class="language-python"># objective(): do_bench(...) becomes an ncu metric read
+
return _measure_occupancy(acf_path) # ncu: sm__warps_active.avg.pct_of_peak_sustained_active
+
+
# search config: minimize runtime -> maximize occupancy
+
config = SearchConfiguration(problem_type="max", generations=4, pool_size=6)
+
</code></pre>+
<p>Pointing the search at achieved occupancy itself across 4 generations × pool size 6, rather than runtime, found a candidate that moved the proxy metric by +0.16pp (21.60% → 21.76% achieved occupancy, p = 9.2e-23, d = 13.17) and, independently re-measured, translated into a modest but statisticall…+
<h3>Multi-Objective Tuning</h3>+
<p>CompileIQ can also perform multi-objective tuning, where the scoring routine evaluates candidates along multiple dimensions simultaneously by returning a tuple rather than a scalar float. Instead of isolating a single top performer, the autotuning process maps out candidates along the Pareto fron…+
<pre><code class="language-python">def objective(config) -> tuple[float, float]:
+
runtime_ms = do_bench(run, warmup=25, rep=100, return_mode="mean")
+
l1_throughput = _measure_l1throughput(acf_path)
+
return runtime_ms, -l1_throughput
+
</code></pre>+
<p>8 generations × pool size 15 converged to a single dominant candidate on the Pareto front. Unlike a tradeoff curve, this single point beats every other evaluated candidate on both axes simultaneously in the search’s own scoring. Validating it on a fresh cache at shipping-grade trial counts:…+
<pre><code class="language-plaintext">=== RUNTIME (ms, lower is better) ===
+
p_value=3.516e-295 cohens_d=87.024
+
speedup_mean=1.0689x
+
=== L1TEX THROUGHPUT (%, higher is better; internally negated) ===
+
p_value=7.123e-19 cohens_d=3.686
+
baseline L1TEX THROUGHPUT mean=81.5087%
+
candidate L1TEX THROUGHPUT mean=85.9620%
+
</code></pre>+
<p>Multi-objective tuning found an ACF delivering a validated 6.9% runtime speedup, accompanied by a simultaneous 4.5pp L1TEX throughput gain (81.51% → 85.96%). Both results are statistically significant (p ≪ 0.05, Cohen’s d ≫ 0.2), with neither metric being traded off against the other.</p>+
<h3>Joint ACF Optimization Across Multiple Input Shapes</h3>+
<p>While earlier evaluations focused on a single kernel shape, N=8192, K=28672, real-world model deployment demands coverage across multiple configurations. The nvfp4_gemv implementation ships with several pre-tuned Helion AOT configurations, including a dedicated setup for N=15360, K=5120. Instead …+
<pre><code class="language-python">def objective(config) -> tuple[float, ...]:
+
"""One shared ACF, one runtime score per shape (both minimize)."""
+
acf_path = ... # write `config` to a temp .acf, or None for baseline
+
return tuple(_bench_shape(acf_path, shape) for shape in SHAPES)
+
</code></pre>+
<p>Executing an 8-generation search with a pool size of 15 across both shapes yielded a Pareto front containing 4 optimal candidates:</p>+
<table>+
<thead>+
<tr>+
<th>Candidate</th>+
<th>down_proj</th>+
<th>15360_5120</th>+
</tr>+
</thead>+
<tbody>+
<tr>+
<td><strong>pareto_0</strong></td>+
<td>0.0927 ms</td>+
<td>0.0309 ms</td>+
</tr>+
<tr>+
<td><strong>pareto_1</strong></td>+
<td>0.0927 ms</td>+
<td>0.0309 ms</td>+
</tr>+
<tr>+
<td><strong>pareto_2</strong></td>+
<td>0.0944 ms</td>+
<td>0.0308 ms</td>+
</tr>+
<tr>+
<td><strong>pareto_3</strong></td>+
<td>0.0917 ms</td>+
<td>0.0321 ms</td>+
</tr>+
</tbody>+
</table>+
<p>Compared to the default benchmark times of 0.0972 ms for down_proj and 0.0315 ms for 15360_5120, pareto_0 emerges as the standout configuration that consistently outpaces the baseline across both target tensor dimensions.</p>+
<p>To confirm these performance gains, we execute an independent validation of pareto_0 on each individual shape. Both configurations comfortably pass our statistical validation (1.05x speedup with a p-value of 1.1e-224 and Cohen’s d of 67.71 for down_proj, and 1.01x speedup with a p-value of …+
<h2>Applicability of this workflow for end-to-end model tuning</h2>+
<p>While the focus of this entire post is about optimizing the underlying compiler for extracting more performance from individual kernels, it is important to zoom out and highlight how a workflow like this could be utilized for end-to-end tuning of models on PyTorch.</p>+
<p>The first step is to profile and characterize the most important kernels that are part of the model which you wish to optimize. Then, each of these kernels can then be a target of CompileIQ optimization, with an associated Advanced Control File to extract maximum performance from that instantiati…+
<p>We want to be upfront about the expected gains: On kernels that are already highly optimized, compiler-control tuning has shown runtime improvements on the order of 2-7% which may still mean a lot for critical kernels in top models. Less optimized kernels have achieved double-digit performance ga…+
<p>The important points to note are that it is found automatically, and requires little to no compiler expertise, it is verified against your own correctness check, and the result is a file you can commit next to the kernel and reapply on every build. It is the difference between a few percent that …+
<p>It also compounds exactly where you’d want it to. There are high-value kernels inside inference loops that run continuously across large fleets. Even a few percent at the bottom of the stack, applied to the operations that dominate a decode step, is valuable at that scale.</p>+
<h3>Caveats</h3>+
<ul>+
<li><strong>ACFs are hardware- and toolkit-specific:</strong> An ACF tuned for sm_90 is not valid on sm_100. <a href="https://nvidia.github.io/CompileIQ/stable/compilers_overview.html#safety-correctness-read-this-first">Applying a mismatched ACF</a> can produce incorrect results or runtime errors, n…+
<li><strong>Search-time is a nontrivial cost:</strong> Every ACF in autotune_search_acf multiplies your autotuning trials. Use it offline on kernels that matter, then pin the winner.</li>+
<li><strong>Gains are workload-specific:</strong> The result depends on the workload, the shapes, the architecture, and how much headroom the kernel has left. The only number that matters for your kernel is the one you measure on your kernel.</li>+
</ul>+
<h2>Your turn to try this out!</h2>+
<p>If you’ve never written a GPU kernel: start at <a href="https://github.com/pytorch/helion/tree/main/pretuned_kernels">pretuned_kernels/</a>. Pick something close to what you need, read it and then run it. The kernels are short, and they run at tuned speed immediately, with no autotuning wai…+
<p>If you already use Helion: download the <a href="https://github.com/NVIDIA/CompileIQ/releases?q=booster-packs&expanded=true">helion-booster-pack</a>, add autotune_search_acf=[…] to one kernel you care about, and autotune it. The no-ACF baseline stays in the search, so the downside is bo…+
<p>If you work on kernels for a living: the takeaway here is not the tool itself. It is that the compiler’s control surface is a legitimate optimization space, it is searchable, and most people are not yet searching it. That is true of any compiler that exposes tunable controls, on any backend…+
<h3>Licensing</h3>+
<p>Helion is open source and hosted by the PyTorch Foundation. CompileIQ is free to use and it is distributed under the NVIDIA Software License Agreement, a proprietary EULA. The ACF files it produces are freely distributable, which is why booster packs can be published and why an ACF can be committ…+
<h3>Links</h3>+
<ul>+
<li>Helion+
<ul>+
<li><a href="https://github.com/pytorch/helion">https://github.com/pytorch/helion</a></li>+
<li><a href="https://helionlang.com">https://helionlang.com</a></li>+
<li><a href="https://github.com/pytorch/helion/tree/main/examples/acfs">https://github.com/pytorch/helion/tree/main/examples/acfs</a></li>+
<li><a href="https://github.com/pytorch/helion/tree/main/pretuned_kernels">https://github.com/pytorch/helion/tree/main/pretuned_kernels</a></li>+
</ul>+
</li>+
<li>CompileIQ+
<ul>+
<li><a href="https://github.com/NVIDIA/CompileIQ">https://github.com/NVIDIA/CompileIQ</a></li>+
<li><a href="https://nvidia.github.io/CompileIQ/stable/">https://nvidia.github.io/CompileIQ/stable/</a></li>+
<li><a href="https://github.com/NVIDIA/CompileIQ/releases?q=booster-packs&expanded=true">https://github.com/NVIDIA/CompileIQ/releases?q=booster-packs&expanded=true</a></li>+
<li><a href="https://developer.nvidia.com/blog/extract-more-kernel-performance-with-nvidia-compileiq-auto-tuning">https://developer.nvidia.com/blog/extract-more-kernel-performance-with-nvidia-compileiq-auto-tuning</a></li>+
</ul>+
</li>+
</ul>+
<p>The post <a href="https://pytorch.org/blog/one-search-two-layers-compiler-controls-inside-pytorch-helion-using-nvidia-compileiq/">One Search, Two Layers: Compiler Controls Inside PyTorch Helion using NVIDIA CompileIQ</a> appeared first on <a href="https://pytorch.org">PyTorch</a>.</p>+
]]></content:encoded>+
+
+
+
</item>+
<item>+
<title>Introducing TorchNeuron: Native PyTorch backend for AWS Trainium</title>+
<link>https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/</link>+
+
<dc:creator><![CDATA[Nafea Bshara, Robert Hundt, Ron Diamant, Colin Brace, Rahul Solanki, Maen Suleiman, Khaled ElGalaind, Temesghen Kahsai, Cyril Bortolato, Indu Thangakrishnan, Shauheen Zahirazami, Sarang Ozarde, Adarsh Singh, Hanbo Wang, Esha Lakhotia, Jianying Lang]]></dc:creator>+
<pubDate>Mon, 19 Oct 2026 12:45:23 +0000</pubDate>+
<category><![CDATA[Blog]]></category>+
<guid isPermaLink="false">https://pytorch.org/?p=188670</guid>+
+
<description><![CDATA[<p>Frontier models span language, video and audio, on architectures from dense and mixture-of-experts transformers to diffusion, state-space, world and structure-prediction models. PyTorch is where AI research happens, where the models...</p>+
<p>The post <a href="https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/">Introducing TorchNeuron: Native PyTorch backend for AWS Trainium</a> appeared first on <a href="https://pytorch.org">PyTorch</a>.</p>+
]]></description>+
<content:encoded><![CDATA[<p>Frontier models span language, video and audio, on architectures from dense and mixture-of-experts transformers to diffusion, state-space, world and structure-prediction models. PyTorch is where AI research happens, where the models are defined, and increasingl…+
<p>Today, to address the growing demand from the developers and research community on Trainium, we are excited to share an even simpler approach to program the chip as a standard PyTorch device, ahead of its open source launch at PyTorch Conference North America on Oct. 20. AWS Trainium is a portfol…+
<p>TorchNeuron supports eager execution for research and debugging including new adaptive eager mode with improved performance, torch.compile for production performance, and the <a href="https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/index.html">Neuron Kernel Interface (NKI)</a> for cus…+
<blockquote><p><em>“PyTorch’s vision is straightforward: developers should have a consistent experience across hardware platforms. Native support for AWS Trainium expands that choice, enabling researchers and engineers to experiment and iterate using the eager execution and distributed A…+
<h2>Getting started: Easy device switching</h2>+
<p>Getting started is straightforward. Run your code on Trainium using the pointer: <code>.to('neuron')</code>.</p>+
<pre><code class="language-python"># Run a training loop on Trainium
+
model = MyModel().to('neuron') # the only change
+
optimizer = torch.optim.AdamW(model.parameters())
+
for batch in dataloader:
+
optimizer.zero_grad()
+
loss = criterion(model(batch), targets)
+
loss.backward()
+
optimizer.step()
+
</code></pre>+
<p>Modern code written against torch.accelerator, PyTorch’s recommended device-agnostic API, runs unmodified. A compiled model needs the Neuron backend named, once with <code>torch.compiler.set_default_backend("neuron")</code>, after which <code>torch.compile(model)</code> takes it, or per cal…+
<pre><code class="language-python"># Run a compiled training loop on Trainium
+
torch.compiler.set_default_backend("neuron") # call once, before any compile
+
model = MyModel().to('neuron')
+
model.compile() # no backend argument
+
optimizer = torch.optim.AdamW(model.parameters())
+
+
for batch in dataloader:
+
optimizer.zero_grad()
+
loss = criterion(model(batch), targets)
+
loss.backward()
+
optimizer.step()
+
</code></pre>+
<h2>Capabilities</h2>+
<p>TorchNeuron registers Trainium through PrivateUse1, PyTorch’s official extension point for custom hardware, so the standard APIs work as written: <code>torch.nn</code>, <code>torch.tensor</code>, <code>torch.autocast</code>, <code>torch.distributed</code>, and <code>torch.profiler</code> al…+
<p><!-- IMAGE PLACEHOLDER: torchneuron-architecture.png - upload via Media Library. Suggested alt text: "High-level architecture diagram of TorchNeuron" --></p>+
<p><em><img decoding="async" class="alignnone wp-image-188675 size-full" src="https://pytorch.org/wp-content/uploads/2026/10/AWS-TorchNeuron-image-1.png" alt="" width="1496" height="1005" srcset="https://pytorch.org/wp-content/uploads/2026/10/AWS-TorchNeuron-image-1.png 1496w, https://pytorch.org/wp…+
<p><em>High-level architecture diagram of TorchNeuron (</em><a href="https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/pytorch-native-overview.html"><em>details</em></a><em>).</em></p>+
<h3>Eager</h3>+
<p>Eager runs operations immediately, so Python control flow, prints, and breakpoints all work with no compile or trace step. This is the mode for research and debugging. Eager mode dispatches operations individually. Executed directly, each operation compiles its own Neuron Executable File Format (…+
<h3>torch.compile</h3>+
<p><code>torch.compile(model, backend="neuron")</code> captures the model into a graph the Neuron compiler optimizes, and is the mode for stable, high-performance models moving to production. torch.compile can be used to optimize the whole or part of the model and works seamlessly with eager mode. F…+
<p>The Neuron inference stacks are built on the same <code>torch.compile</code> support in TorchNeuron. vLLM Neuron is a plugin to vLLM and serves LLMs; vLLM Omni Neuron is a plugin to vLLM Omni and serves diffusion models for image and video generation.</p>+
<h3>Ops and dtypes</h3>+
<p>Broad aten operator coverage for ops seen across text, multimodal and vision language, image and video generation and understanding, state space models, and world models. Models can use BF16, FP32, and FP16, with integer and boolean types, and <code>torch.autocast</code> and mixed precision code …+
<h3>Distributed APIs</h3>+
<p>Distributed inference and training can use the standard PyTorch APIs and the high-level parallelism modules compose with your model as written: DDP, FSDP2, SimpleFSDP, tensor parallelism, and pipeline parallelism, along with common compositions like FSDP2 with tensor parallelism, ZeRO that are es…+
<h3>Streams and Events</h3>+
<p>Torch streams and events are supported natively. <code>torch_neuronx.Stream</code> and <code>torch_neuronx.Event</code> map PyTorch’s stream semantics onto Trainium’s hardware queues. Within a stream, data movement and computation stay in order automatically. Separate streams carry no…+
<h3>Neuron Kernel Interface (NKI)</h3>+
<p>The <a href="https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/index.html">Neuron Kernel Interface (NKI)</a> lets researchers and developers develop, optimize, and execute custom operators for both eager and torch.compile. Researchers can prototype ideas that require low level hardware …+
<p>TorchNeuron also provides Trainium optimized, high-performance operators implemented with NKI and released open source as <a href="https://github.com/aws-neuron/nki-library/blob/main/README.md">NKI-lib</a>, covering mainstream LLM and multimodal scenarios for both dense and MoE models, such as Fl…+
<pre><code class="language-python">import torch
+
import nki
+
import nki.language as nl
+
from torch_neuronx import nki_op
+
+
# Declaring and implementing NKI kernels
+
# Note in NKI kernels, the output tensor should be defined as parameter to the kernel
+
@nki.jit
+
def sin_kernel(in_ptr0, out_ptr):
+
input_tile = nl.load(in_ptr0[0:128])
+
output_tile = nl.sin(input_tile)
+
nl.store(out_ptr[0:128], value=output_tile)
+
return out_ptr
+
+
# defining cos since it will be used in backward path of sin
+
@nki.jit
+
def cos_kernel(in_ptr0, out_ptr):
+
input_tile = nl.load(in_ptr0[0:128])
+
output_tile = nl.cos(input_tile)
+
nl.store(out_ptr[0:128], value=output_tile)
+
return out_ptr
+
+
# after this line, there is no NKI code, just native PyTorch, registering custom ops
+
@nki_op("mylib::pt_nki_sin", mutates_args={})
+
def pt_nki_sin(x: torch.Tensor) -> torch.Tensor:
+
out = torch.empty_like(x)
+
sin_kernel(x, out)
+
return out
+
+
@nki_op("mylib::pt_nki_cos", mutates_args={})
+
def pt_nki_cos(x: torch.Tensor) -> torch.Tensor:
+
out = torch.empty_like(x)
+
cos_kernel(x, out)
+
return out
+
+
def pt_nki_sin_backward(ctx, grad):
+
x, = ctx.saved_tensors
+
return grad * pt_nki_cos(x) # cos is derivative of sin
+
+
def pt_sin_setup_context(ctx, inputs, output):
+
x, = inputs
+
ctx.save_for_backward(x)
+
+
# Register backward pass and context function used during autograd for nki_sin
+
pt_nki_sin.register_autograd(pt_nki_sin_backward, setup_context=pt_sin_setup_context)
+
+
if __name__ == "__main__":
+
x = torch.randn(128, 1, device="neuron", requires_grad=True)
+
# Forward pass uses sin_kernel
+
y = pt_nki_sin(x)
+
# Backward pass, autograd calls uses nki_sin_backward() -> cos_kernel()
+
loss = y.sum()
+
loss.backward()
+
</code></pre>+
<h3>Profiling with TorchNeuron</h3>+
<p>TorchNeuron integrates with the standard <code>torch.profiler.profile</code> API extended with a NeuronConfig object that controls Neuron-specific data collection. This allows users to capture and correlate traces from framework-level CPU operations (PyTorch dispatch, Python stacks) as well as Ne…+
<p>Profiles captured can be visualized in the Neuron Explorer UI, where users can trace execution across framework, runtime, and device, jump from any event to the exact source line (source code viewer), navigate the workload through its call hierarchy (hierarchy viewer), and follow dependency chain…+
<p>Neuron Explorer also emits a <a href="https://awsdocs-neuron.readthedocs-hosted.com/en/latest/tools/neuron-explorer/profile-schema-reference.html">Parquet output</a> with per-op and per-engine metrics, so the same data behind the UI can be queried directly or consumed by agents for automated bott…+
<p>Following code snippet shows an example (<a href="https://awsdocs-neuron.readthedocs-hosted.com/en/latest/tools/neuron-explorer/how-to-profile-workload.html#pytorch-profiling">details</a>):</p>+
<pre><code class="language-python">import torch
+
from torch.profiler import profile, ProfilerActivity
+
from torch_neuronx.profiling import NeuronConfig, ProfileMode, NeuronProfiler
+
+
# Configure Neuron profiling and pass appropriate path as OUTPUT_DIR
+
nc = NeuronConfig(
Diff display stops at 400 lines. The line counts above are from the whole diff. 55 lines shown here cut at 300 characters. The raw artifact at this commit is linked above.