Custom LLM Kernel Optimization in Production: Triton vs. CUDA C++ vs. torch.compile vs. CUTLASS

Serving large language models at scale requires extracting maximum performance from modern GPU architectures like NVIDIA Ampere, Hopper, and Blackwell. While early production deployments relied on standard PyTorch eager execution and standard cuBLAS calls, high-throughput serving systems such as vLLM, SGLang, and TensorRT-LLM depend on specialized fused GPU kernels to eliminate memory bandwidth bottlenecks and saturate Tensor Cores.

Engineering teams face four primary paradigms for kernel development: PyTorch's native torch.compile (Inductor), OpenAI's Triton, NVIDIA's CUTLASS 3.x with CuTe, and raw CUDA C++ with inline PTX assembly. Each operates at a distinct level of hardware abstraction, trading off engineering velocity against low-level micro-architectural control.

Understanding where each tool fits across the compute-bound and memory-bound stages of LLM inference is essential for optimizing time-to-first-token (TTFT), inter-token latency (ITL), and total serving cost.

The Roofline Reality: Compute vs. Memory Bandwidth

To understand why custom kernels are necessary, consider the operational phases of transformer inference on modern accelerator hardware such as the NVIDIA H100 SXM5 GPU:

Hardware Specs: The NVIDIA H100 SXM5 provides 3.35 TB/s of High Bandwidth Memory (HBM3) bandwidth, 989 TFLOPS of dense FP16 Tensor Core compute, and 1,979 TFLOPS of dense FP8 Tensor Core compute.
Ridge Point Threshold: The roofline ridge point represents the arithmetic intensity required to transition from memory bandwidth-bound to compute-bound execution. For FP16, this threshold is approximately 295 FLOPs per byte (989 TFLOPS / 3.35 TB/s). For FP8, the threshold doubles to approximately 590 FLOPs per byte (1,979 TFLOPS / 3.35 TB/s).

LLM inference consists of two distinct workloads with opposite performance profiles:

Prefill Phase (Prompt Ingestion): High batch and sequence token counts process simultaneously. Matrix multiplications ( $Q \times K^T$ , $S \times V$ , MLP feed-forward projections) exhibit high arithmetic intensity ( $O(N^2)$ or large $M \times K \times N$ GEMMs). This phase is compute-bound and requires maximizing Tensor Core utilization.
Decode Phase (Token Generation): Tokens are generated autoregressively one step at a time ( $M=1$ or small effective batch sizes). The operations degenerate into matrix-vector multiplications (GEMV), dynamic KV cache indexing, and elementwise operations (RoPE, RMSNorm, SwiGLU activations, logit sampling). Arithmetic intensity drops well below 50 FLOPs per byte, making decode strictly memory bandwidth-bound.

Under standard PyTorch eager execution, each elementwise operation launches a separate GPU kernel. A sequence of RMSNorm -> Residual Add -> Rotary Embedding incurs three distinct kernel launches (each with 3 to 5 microseconds of launch overhead) and round-trips data between GPU registers and High Bandwidth Memory (HBM) three times. Fusing these operations into a single kernel keeps intermediate activations in fast on-chip SRAM or registers, cutting memory traffic by 60% to 80% and eliminating launch latency.

The Four Kernel Development Paradigms

Production inference runtimes combine multiple kernel authoring approaches to balance developer productivity with execution efficiency.

1. OpenAI Triton: Block-Level Pythonic Programming

OpenAI Triton introduced a programming model based on blocked arrays (tiles) rather than individual threads. Instead of managing thread-level SIMT indexing, shared memory allocation, and warp synchronization manually, engineers write Python functions using the @triton.jit decorator.

Triton's compiler handles:

Automatic memory coalescing when loading 2D blocks from global memory.
Allocation and layout transformation of shared memory (SRAM).
Software pipelining across execution iterations to overlap global memory loads with compute.

import triton
import triton.language as tl

@triton.jit
def _fused_rmsnorm_kernel(
    X_ptr, Y_ptr, W_ptr,
    stride_x, stride_y,
    N_COLS: tl.constexpr,
    EPS: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    row_idx = tl.program_id(0)
    row_start_x = X_ptr + row_idx * stride_x
    row_start_y = Y_ptr + row_idx * stride_y
    
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < N_COLS
    
    # Load input row into SRAM/registers
    x = tl.load(row_start_x + cols, mask=mask, other=0.0).to(tl.float32)
    
    # Compute variance across row
    var = tl.sum(x * x, axis=0) / N_COLS
    rrms = 1.0 / tl.sqrt(var + EPS)
    
    # Load weights and scale
    w = tl.load(W_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    y = (x * rrms) * w
    
    tl.store(row_start_y + cols, y.to(tl.bfloat16), mask=mask)

Where Triton Dominates: Fused point-wise operators (RMSNorm, LayerNorm), Rotary Position Embedding (RoPE), dynamic quantization and dequantization (FP8/INT8/INT4), MoE top-k routing and token dispatch, and custom attention variants (such as FlashLinearAttention and SageAttention).

Limitations: Triton abstracts away hardware-specific micro-architectural primitives. On NVIDIA Hopper architectures, Triton has historically required extra overhead to pass Tensor Memory Accelerator (TMA) descriptors compared to native C++ implementations, according to PyTorch's Deep Dive on the Hopper TMA Unit for FP8 GEMMs.

2. CUTLASS 3.x & CuTe: C++ Template Meta-Programming

NVIDIA CUTLASS is an open-source template library for high-performance matrix multiplication and deep learning primitives on NVIDIA GPUs. Version 3.x introduced CuTe, a domain-specific C++ library that formalizes multi-dimensional layouts, coordinates, and hierarchical tensor algebra.

CUTLASS operates at the block, warp, and warp-group levels:

Warp Specialization: On Hopper architectures, threads within a threadblock are divided into dedicated roles. Producer warps execute asynchronous memory loads via the Tensor Memory Accelerator (TMA), while consumer warps execute matrix multiply-accumulate operations (WGMMA) directly against shared memory.
Asynchronous Pipelining: Hardware-native barriers (mbarrier) allow consumer warps to compute on Tile $N$ while producer warps stream Tile $N+1$ directly from HBM to shared memory without consuming register file space.

Where CUTLASS Dominates: High-throughput dense and sparse GEMMs, mixed-precision FP8/FP4 matrix multiplications with block scaling, and frontier attention mechanisms like FlashAttention-3, which achieved up to 740 TFLOPS in FP16 and over 1.2 PFLOPS in FP8 by leveraging CuTe and Hopper warp-specialization.

Trade-Offs: CUTLASS 3.x has a steep learning curve due to complex C++ template metaprogramming, slow compilation times, and intricate type systems that require deep understanding of GPU memory hierarchies.

3. torch.compile & Inductor: Automated Graph Fusion

Introduced in PyTorch 2.0 and expanded in PyTorch 2.6, torch.compile uses TorchDynamo to intercept Python frame execution, build an FX graph, and lower it through the TorchInductor compiler.

TorchInductor automatically analyzes memory access patterns and generates optimized C++ or Triton kernels:

Vertical Fusion: Combines linear sequences of point-wise operations into single fused loops.
Horizontal Fusion: Merges parallel independent operations that share input tensors to reduce kernel dispatch overhead.
CUDA Graphs Integration: Captures static execution subgraphs to eliminate CPU-side kernel launch overhead during repeated decode steps, as documented in the vLLM torch.compile integration architecture.
FlexAttention: PyTorch's FlexAttention compiles user-defined score modifications (e.g. document masking, relative position biases, prefix-LM masks) directly into fused Triton attention kernels without writing custom C++ or CUDA code.

Where torch.compile Dominates: Whole-model compilation, standard activation chains (e.g. MLP blocks, attention projection post-processing), and zero-boilerplate model modernization.

Failure Modes: Dynamic batch sizes and varying prompt sequence lengths can trigger repeated graph recompilations (compilation storms) and graph breaks when hitting unsupported Python constructs, falling back to eager execution.

4. Hand-Crafted CUDA C++ & PTX: Direct Hardware Control

Writing raw CUDA C++ using nvcc and inline PTX assembly gives engineers complete, deterministic control over every thread, register, and hardware instruction.

Key low-level mechanisms available in CUDA C++ include:

Warp Shuffle Primitives: Intrinsic functions (__shfl_down_sync, __shfl_xor_sync) allow direct register-to-register data exchange between threads within a warp without touching shared memory.
Direct PTX Instruction Access: Low-level instructions such as cp.async (Ampere async global-to-shared copy) and wgmma.mma_async (Hopper warp-group matrix multiply) can be invoked directly.
Shared Memory Swizzling: Explicit mathematical mapping of tensor indices to the 32 shared memory banks, preventing multi-way bank conflicts during strided memory reads.

Where Raw CUDA Dominates: Low-latency decode kernels (such as FlashInfer's specialized PagedAttention decode routines), custom KV cache layouts, speculative decoding draft-verification engines, and proprietary hardware extensions in TensorRT-LLM.

Trade-Offs: High implementation effort, manual register pressure management, complex debugging, and architectures that must be rewritten or re-tuned for each new GPU generation.

Architectural Comparison

The following breakdown summarizes how these four paradigms compare across key technical dimensions:

Abstraction Level:
torch.compile: Graph-level (automated IR lowering and operator fusion).
Triton: Block-level (coalesced tile operations, automated thread/SRAM mapping).
CUTLASS 3.x / CuTe: Tensor layout and warp-group level (explicit layout algebra and hardware pipeline scheduling).
Raw CUDA C++ / PTX: Thread-level SIMT (explicit registers, warps, and shared memory bank management).
Primary Strength:
torch.compile: Zero boilerplate, rapid full-model optimization, native PyTorch ecosystem integration.
Triton: High developer velocity, concise Pythonic kernels for custom elementwise and reduction operations.
CUTLASS 3.x / CuTe: Speed-of-light GEMM performance, full access to Hopper/Blackwell hardware units (TMA, WGMMA).
Raw CUDA C++ / PTX: Deterministic latency, fine-grained register and warp control for memory-bound decode loops.
Primary Bottleneck or Risk:
torch.compile: Graph breaks and recompilation storms on dynamic shapes.
Triton: Lower control over micro-architectural scheduling and multi-stage hardware pipelines.
CUTLASS 3.x / CuTe: Extreme C++ template complexity and slow build times.
Raw CUDA C++ / PTX: Fragile maintainability, high development costs, and architectural lock-in.
Typical LLM Inference Use Cases:
torch.compile: General transformer layers, FlexAttention masks, model glue code.
Triton: Fused RMSNorm, RoPE, FP8 dynamic quantization, MoE router dispatch.
CUTLASS 3.x / CuTe: FP8/FP16 prefill GEMM, FlashAttention-3, grouped GEMM for MoE experts.
Raw CUDA C++ / PTX: PagedAttention decode (GEMV), speculative verification, TensorRT-LLM core plugins.

The Production Kernel Strategy

Leading inference frameworks do not choose a single paradigm in isolation. Instead, they implement a multi-tiered kernel architecture:

Tier 1 (Automated Baseline): Use torch.compile with CUDA Graphs for model wiring, standard linear layer chaining, and dynamic masking via FlexAttention.
Tier 2 (Custom Pointwise & Reductions): Write Triton kernels for fused layer norms, rotary embeddings, MoE routing gates, and custom quantization/dequantization passes.
Tier 3 (Frontier Matrix & Attention Compute): Employ CUTLASS 3.x / CuTe for dense and grouped GEMMs (especially W8A8 and FP4 on Hopper and Blackwell) and prefill attention backends.
Tier 4 (Latency-Critical Decode Loops): Maintain optimized CUDA C++ / PTX kernels for single-token decode operations, PagedAttention cache lookups, and specialized hardware synchronization primitives.

This layered approach allows engineering teams to maximize developer productivity on common model components while concentrating low-level optimization efforts on the specific kernels that govern throughput and tail latency.

Custom LLM Kernel Optimization in Production: Triton vs. CUDA C++ vs. torch.compile vs. CUTLASS

The Roofline Reality: Compute vs. Memory Bandwidth

The Four Kernel Development Paradigms

1. OpenAI Triton: Block-Level Pythonic Programming

2. CUTLASS 3.x & CuTe: C++ Template Meta-Programming

3. torch.compile & Inductor: Automated Graph Fusion

4. Hand-Crafted CUDA C++ & PTX: Direct Hardware Control

Architectural Comparison

The Production Kernel Strategy

Sources

Agent Egress Security in Production: Network Sandboxing, Secretless Token Rewriting, and DNS Exfiltration Defenses

Skanska Signs .2B Contract to Build Four AI Data Centers in Southeast USA

UK Heterogeneous Compute Startup Callosum Raises 00M Seed Backed by Atomico and Sovereign AI Fund