GPU Compilers and Kernel Languages

Your vector-add kernel passed every correctness check yesterday. Suppose the same source now takes 300 milliseconds on its first call, or fails on a newer GPU with "no kernel image is available." Reading the source again won't explain either symptom. The GPU never executes CUDA C++ or Python directly.

The accelerator field guide mapped execution and memory owners across hardware. The parallel primitives lab then made synchronization scope, reduction trees, and prefix dependencies measurable. Keep those habits. Here we follow compiler artifacts between source and hardware, because each boundary owns a different class of bugs.

One addition, several representations

For 1,000 pairs of numbers, add each pair into a separate output slot. At index zero, $0 + (-3) = -3$; at index one, $0.25 + (-2) = -1.75$. In general, write $out_i = x_i + y_i$ for each valid index $i$.

CUDA C++ expresses the work from one thread's point of view. With 256 threads per block, 1,000 elements require four blocks: 1,024 threads, of which 24 must skip memory access. The same fixture below also checks lengths immediately around a block boundary.

Save this as vector_add.cu. It requires a CUDA toolkit to build and compatible NVIDIA hardware to run. The checks use small, exactly representable FP32 values; exact equality here isn't a general tolerance policy for reductions or matrix multiplication.

Successful execution prints PASS n=1, PASS n=255, PASS n=256, PASS n=257, and PASS n=1000. These are expected results, not measurements from this review's CPU-only environment. Distinct input values help catch misplaced writes; checking finiteness prevents a NaN from disappearing inside a maximum-error reduction. This fixture doesn't test large-index overflow, arbitrary strides, or aliasing.

At source level, you can inspect the bounds check and the index formula. You can't yet prove which machine instructions were emitted, how many physical registers they require, or whether a deployed binary contains code for the target GPU.

NVIDIA's offline compiler driver, nvcc, coordinates separate host and device compilation. For this SIMT device kernel, the compilation path reaches Parallel Thread Execution (PTX), a versioned virtual instruction set architecture (ISA). PTX acts as a stable, hardware-agnostic intermediate representation. Crucially, PTX operates over an unbounded virtual register space (such as %r0, %r1, %f0, %f1). It doesn't know how many physical registers exist on the target GPU, nor does it concern itself with warp scheduler stalls or instruction issue slots.^{[1]Reference 1NVIDIA CUDA Compiler Driver NVCChttps://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/}

The optimizing assembler, ptxas, transforms that virtual PTX into a CUDA binary, called a cubin, for a specific streaming multiprocessor target such as sm_90. A cubin is a standard ELF container holding encoded GPU machine instructions plus symbols, relocations, and resource metadata.^{[2]Reference 2CUDA Binary Utilitieshttps://docs.nvidia.com/cuda/cuda-binary-utilities/} In this translation, ptxas isn't a passive 1-to-1 translator:

• It performs physical register allocation, mapping unbounded virtual registers into the SM's finite register file through graph coloring and live-range analysis.
• It performs instruction scheduling, reordering instructions to interleave independent math between memory load issues and their arrival, hiding memory and arithmetic latency.
• It injects hardware control codes, setting per-instruction stall counts, yield flags, and scoreboard dependency barriers that direct hardware warp schedulers.
• It manages register spilling, inserting spill and fill instructions to thread-private .local memory if register demand exceeds physical constraints.

Engineers commonly call the human-readable native instruction listing SASS. SASS isn't an executable artifact or a portable source format. It's a disassembly of target-specific machine instructions stored in a cubin.

To see this transformation concretely, compare the virtual PTX emitted for the core vector add against the native SASS disassembled from the resulting sm_90 cubin:

Notice the register reuse in SASS: ptxas reused physical register R2 for both the first loaded float and the final addition result, freeing up registers for other concurrent warps.

There can also be an earlier compiler layer. NVVM intermediate representation (IR) is NVIDIA's GPU-oriented subset and extension of LLVM IR. Language frontends can generate NVVM IR, and libNVVM compiles compatible NVVM IR into PTX.^{[3]Reference 3NVVM IR Specificationhttps://docs.nvidia.com/cuda/nvvm-ir-spec/[4]Reference 4libNVVM APIhttps://docs.nvidia.com/cuda/libnvvm-api/} That makes this conceptual path useful:

$$\text{source} \rightarrow \text{LLVM/NVVM IR} \rightarrow \text{PTX} \rightarrow \text{ptxas} \rightarrow \text{cubin (SASS)}$$

Don't turn a conceptual compiler layer into a deployment promise. nvcc publicly documents source-to-PTX and PTX-to-cubin stages; it doesn't promise a stable, human-readable NVVM file for every build. Another frontend may use different IRs before it reaches PTX, and a non-NVIDIA backend won't target PTX at all. CUDA Tile, introduced later, has a separate Tile IR path; the SIMT path here isn't a universal pipeline for every CUDA language.

You find an FFMA instruction in nvdisasm output. Which layer has supplied the strongest evidence that this native operation exists?

Answer

The target cubin. nvdisasm reads encoded machine instructions from a cubin and renders them as SASS. A source expression or PTX instruction can suggest an FFMA, but a later compiler stage may still fuse, split, remove, or replace it.

nvcc builds a package, not one universal binary

A CUDA translation unit may contain CPU host code and GPU device code. nvcc separates those paths, sends the host path to a supported C++ compiler, compiles device code for requested virtual and real architectures, and embeds device images into the host object. The container of one or more device images is a fatbinary, usually shortened to fatbin.^{[1]Reference 1NVIDIA CUDA Compiler Driver NVCChttps://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/}

The distinction between compute_90 and sm_90 records intent:

Ordinary PTX targets such as compute_90 provide a forward path to compatible later NVIDIA targets, not to older GPUs missing required features. The driver must also understand the emitted PTX ISA version. A toolkit upgrade can therefore raise the required driver version even when compute_90 stays unchanged.

Target suffixes matter: architecture-specific PTX such as compute_90a doesn't have the same forward-compatibility guarantee, and family-specific targets such as compute_100f restrict compatibility to their documented family. Cubins are narrower still, with compatibility rules within supported major compute-capability families, not across major versions. Check the exact target, not just the word “PTX.”^{[1]Reference 1NVIDIA CUDA Compiler Driver NVCChttps://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/}

The runtime flow branches after packaging:

This branch explains two otherwise confusing observations. A cold process may pay compilation or lazy-loading cost before the first kernel. A package may run on the build machine yet fail on another GPU because neither a compatible cubin nor usable PTX was embedded.

Produce artifacts you can inspect

On a supported CUDA development system, compile the CUDA source for one native Hopper target and keep PTX for future targets. --keep asks nvcc to retain supported intermediate files, while the two --generate-code entries package both sm_90 cubin and compute_90 PTX. Offline compilation and disassembly don't require a GPU, but they do require these NVIDIA tools and a supported host toolchain. Apple Clang alone can't run this lab. Use your deployment targets instead of copying sm_90 onto unrelated hardware.

Create standalone PTX and cubin files as well. Keeping both makes the boundary explicit and gives nvdisasm a cubin it can read directly.

Now inspect the package before judging performance. cuobjdump accepts a standalone cubin or a host executable, object, library, or external fatbin. nvdisasm accepts standalone cubins and adds richer control-flow and register-liveness views.^{[2]Reference 2CUDA Binary Utilitieshttps://docs.nvidia.com/cuda/cuda-binary-utilities/}

Read these outputs as a chain of evidence:

1. PTX header: Check .version, .target, address size, and kernel entry name. A mismatched .target can explain a deployment failure before launch.
2. PTX body: Look for global loads, bounds predication, arithmetic, and the store. This confirms virtual operations, not final scheduling.
3. Resource report: Record registers, static shared memory, local memory, and stack per kernel. Local memory can indicate spills, although not every local-memory use is a spill.
4. SASS listing: Confirm native loads, arithmetic, stores, predicates, and architecture-specific instructions. Compare against compiler flags and source line info.
5. Profiler trace: Measure executed instructions, memory behavior, and time on the target GPU. Static disassembly can't tell you which path dominates wall time.

⚠️ Common mistake: A shorter SASS listing isn't automatically faster. Instruction latency, issue rate, dependency chains, memory transactions, occupancy, and input-dependent control flow still decide runtime behavior.

Register pressure, occupancy cliffs, and spilling

When you inspect cuobjdump --dump-resource-usage or parse compiler logs, the most critical number is the register count per thread. That single metric dictates whether your kernel saturates GPU execution units or falls off a severe occupancy cliff.

Every modern NVIDIA Streaming Multiprocessor (SM), across architectures from Ampere (sm_80) and Ada (sm_89) to Hopper (sm_90) and Blackwell (sm_100), houses a physical register file of 64K (65,536) 32-bit registers. At the same time, an SM can host up to 2,048 active concurrent threads (organized as 64 warps of 32 threads).

To achieve 100% theoretical occupancy, all 2,048 threads must fit within the SM's register budget simultaneously:

$$\frac{65{,}536 \text{ registers}}{2{,}048 \text{ threads}} = 32 \text{ registers/thread}$$

When each thread uses 32 registers or fewer, all 64 warps reside on the SM at once. The hardware warp schedulers have maximum freedom to swap out a stalled warp (waiting on global memory or math pipelines) and issue independent instructions from ready warps.

What happens when register demand creeps upward? The SM can't dynamically allocate fractions of registers. Instead, hardware allocates registers in fixed warp granularities, producing steep occupancy cliffs:

These transitions aren't gentle slopes. Adding a single local variable or unrolling an inner loop that pushes register count from 32 to 33 immediately drops thread occupancy by 25%.

The spilling disaster: local memory is DRAM

When a complex kernel needs more registers than the hardware allows, or when you force a low register limit via compiler flags, ptxas doesn't abort. It spills excess registers into Local Memory (.local).

Don't let the name mislead you: local memory is not fast on-chip SRAM. Local memory is thread-private memory backed by off-chip DRAM, cached in L1 and L2.

When 1,000 or more active threads on an SM spill variables to .local memory simultaneously, the SM's 128KB to 256KB L1 cache lines thrash relentlessly. Spilled values get evicted down to off-chip DRAM. Every register read and write that turned into a spill now incurs memory controller roundtrips, triggering a catastrophic 10x to 50x latency penalty that ruins kernel throughput.

Parsing ptxas -v reports

To catch register pressure and spills before running a kernel, pass -Xptxas -v (or --ptxas-options=-v) to nvcc:

The compiler outputs an exact resource receipt for every compiled kernel:

Every field in this report delivers an operational truth:

• Used 32 registers: Physical 32-bit registers allocated per thread. At exactly 32 registers, the kernel hits the 100% occupancy threshold for 2,048 threads per SM.
• 0 bytes smem: Statically allocated shared memory per block.
• 0 bytes cmem[0]: Constant memory bank 0 usage (kernel parameters and pointer addresses passed from the CPU host).
• Spilled 0 bytes to local memory, 0 bytes read back: Zero is the target. If you see Spilled 16 bytes to local memory, 16 bytes read back, each thread spilled 4 32-bit registers (16 bytes). Across a full grid of 100,000 threads, those 4 words translate into hundreds of thousands of unnecessary memory bus accesses.

You have two primary compiler knobs to manage this tradeoff:

1. --maxrregcount=N: Forces ptxas to cap registers per thread at $N$. If the kernel requires more than $N$ registers, the compiler spills to .local memory to preserve target occupancy.
2. __launch_bounds__(maxThreadsPerBlock, minBlocksPerMultiprocessor): A function qualifier placed in CUDA C++ source above __global__. It promises the compiler your launch bounds, allowing ptxas to calculate the exact register ceiling for that block size without accidentally triggering an occupancy drop.

Keep an artifact manifest

A benchmark result becomes easier to reproduce when its compiler state travels beside it. This is a template, not a captured build. Replace every placeholder and record at least these fields:

Hash the source, compile options, target, compiler version, and any compile-time constants into a JIT cache key. Shape-specialized kernel systems also need shape, stride, data type, alignment, and algorithmic mode when those values change generated code. Caching only by function name can silently load a valid cubin for the wrong contract.

For example, suppose two variants have 1,000 elements but strides 1 and 2. They address different memory locations, so a specialization that assumes stride 1 mustn't be reused for stride 2. Similarly, changing the target architecture suffix from sm_90 to sm_90a unlocks architecture-specific Hopper features like Tensor Memory Accelerator (TMA) asynchronous copies and Warpgroup Matrix Multiply-Accumulate (WGMMA) instructions. A cached binary compiled for sm_90 won't use those instructions, while an sm_90a cubin will crash if loaded on a GPU without the architecture-accelerated feature set.

This CPU exercise builds a canonical key from an explicit compilation recipe, including launch geometry and pipeline stages. It doesn't compile a kernel or implement a production artifact store.

The marker string stands in for source bytes only in this exercise. A real cryptographic cache key must cover transitive includes, linked device libraries, compiler components, backend options, and any assumed pointer alignment (since 16-byte alignment enables 128-bit vector loads). Preserve flag order because repeated flags can override earlier ones. Runtime tensor values needn't enter the key unless compilation specializes on them. Store the artifact's hash and compatibility metadata separately; a cache hit isn't evidence that an untrusted binary is safe to load.

JIT is part of request latency

Just-in-time (JIT) compilation moves work from build time to runtime. NVIDIA's runtime compiler, NVRTC, accepts CUDA C++ strings and can emit PTX, cubin, CUDA Tile IR, or link-time optimization IR depending on target and options. The CUDA Driver API can load the result, while nvJitLink can combine PTX or link-time IR and produce a linked cubin.^{[5]Reference 5NVRTChttps://docs.nvidia.com/cuda/nvrtc/}

PTX loaded at runtime is compiled by the device driver's JIT compiler. Generated cubins enter a compute cache, and a driver upgrade invalidates that cache so the new JIT compiler can regenerate code.^{[6]Reference 6CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

That runtime behavior creates an operational contract:

Multi-process compilation races in distributed clusters

In production distributed training and serving (such as 8 GPUs on an HGX node running PyTorch DistributedDataParallel, DeepSpeed, Megatron-LM, or vLLM), all 8 worker processes execute in parallel.

When all 8 ranks encounter an uncompiled kernel or a new sequence length at the exact same millisecond, they all detect a cache miss simultaneously. If all 8 processes attempt to compile and write to the same disk location (such as ~/.triton/cache/<hash>.so or ~/.nv/ComputeCache) without coordination, severe race conditions occur:

• Corrupted binary loads: Process A is halfway through writing the compiled ELF shared object when Process B checks the directory, finds the file name, and immediately calls dlopen() or cuModuleLoad(). Process B crashes with an invalid ELF header, truncated file error, or segmentation fault.
• Overlapping clobbers: Multiple processes open the same file descriptor for writing simultaneously, corrupting disk blocks and producing a broken cache entry that breaks all future runs until manually purged.

Production GPU engines resolve this race through three complementary strategies:

1. Atomic temporary write and replace: Each worker writes its compiled binary to a unique temporary file on the exact same filesystem (for example, <hash>.tmp.<pid>.<uuid>), flushes it to disk with os.fsync(), and performs an atomic rename via os.replace(). Under POSIX semantics, rename is guaranteed to be atomic within the same filesystem. Any competing reader sees either the complete existing file or the complete newly replaced file, never a half-written fragment.
2. Advisory file locking: Before initiating compilation, the worker acquires an exclusive advisory lock using fcntl.flock or a cross-platform lock file (<hash>.lock). The first process to acquire the lock performs the compile. The remaining 7 processes block. Once the lock releases, the waiting processes wake up, see the populated cache file, and take an immediate cache hit without recompiling.
3. Rank-0 compilation barrier: In managed distributed setups, rank 0 is assigned compilation responsibility during an explicit initialization phase. All workers wait at a torch.distributed.barrier(). Once rank 0 finishes populating the on-disk cache, workers 1 through 7 safely load the compiled artifacts in read-only mode.

Use the driver's switches to test both deployment branches. The first command forces embedded PTX through JIT and disables its disk cache for the run. The second disables PTX JIT, so success requires a compatible embedded cubin.^{[7]Reference 7CUDA Environment Variableshttps://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/environment-variables.html}

Run each mode in a fresh process. These shell assignments affect only that invocation; unset any conflicting JIT variables you previously exported. These are packaging tests, not kernel benchmarks.

CUDA module loading is lazy by default in current documentation. First use can therefore include module or kernel loading even when a compatible cubin exists. Warm the exact kernel or use eager loading when a latency-sensitive service needs that work outside the request path.^{[7]Reference 7CUDA Environment Variableshttps://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/environment-variables.html}

A service packages sm_90 cubin plus compute_90 PTX. Its first request on a future compatible GPU is slow, but later requests are fast. What should you separate in the trace?

Answer

Separate process initialization, module loading, PTX-to-cubin JIT, cache write, first kernel launch, and steady-state execution. A fast warm kernel doesn't disprove expensive compilation or loading on the first request.

Same kernel, different unit of thought

CUDA C++ asks you to describe one thread, then organize threads into blocks. Triton asks you to describe one blocked program instance operating on vectors of offsets. Both versions below compute the same masked vector add.

The Triton kernel forms a block of 256 indices, loads two blocks under a mask, adds them, and stores the result. The compiler maps that blocked program to GPU threads and memory operations. The kernel follows Triton's vector-add tutorial; the wrapper below adds explicit input checks.^{[8]Reference 8Triton Programming Guide: Introductionhttps://triton-lang.org/main/programming-guide/chapter-1/introduction.html[9]Reference 9Vector Additionhttps://triton-lang.org/main/getting-started/tutorials/01-vector-add.html}

Run this on a supported Linux/PyTorch/Triton CUDA environment. It should print six shape passes and one strided-view rejection. The GPU code wasn't executed in this review; don't treat the CPU checks below as Triton compilation or memory-safety evidence. Inputs requiring autograd aren't supported by this wrapper: a custom backward rule would be separate work.

The arithmetic stayed fixed. Ownership changed:

Triton removes some thread-level bookkeeping. It doesn't remove the need to reason about coalescing, block size, occupancy, specialization, or numerical behavior. A blocked expression is an optimization hypothesis until generated artifacts and measurements support it.

How Triton compiles: the MLIR lowering pipeline

Triton isn't an interpreter or a simple string-template emitter. It's an optimizing compiler built on the Multi-Level Intermediate Representation (MLIR) framework. It lowers abstract block mathematics into native GPU machine code through a sequence of explicit intermediate dialects:

1. Triton IR (tt dialect): Parsed directly from the Python AST and @triton.jit decorator. This level captures high-level tile operations (tt.load, tt.store, tt.add, tt.dot). Operations at this stage are hardware-agnostic: they express math over multidimensional tiles without committing to thread counts, warp allocations, or physical memory layouts.
2. TritonGPU IR (ttg dialect): This stage contains Triton's core GPU optimization intelligence. The compiler takes the abstract tiles and maps them to physical GPU resources:
   • Distributed layout assignment: The compiler assigns formal layout encodings (such as #triton_gpu.blocked, #triton_gpu.shared, and #triton_gpu.dot_op) that dictate how each element in a tile is distributed across warps and individual thread lanes.
   • Shared memory allocation and XOR swizzling: Matrix tiles loaded into shared memory often suffer from shared memory bank conflicts when threads access columns across regular power-of-two strides. Triton automatically allocates shared memory buffers and applies an XOR address-swizzling pattern to the indices. This scatters memory accesses across all 32 independent banks, completely eliminating bank serialization without manual programmer padding.
   • Asynchronous software pipelining: For reduction or GEMM loops, Triton automatically constructs circular multi-stage buffers in shared memory (num_stages = 3 or 4). It emits asynchronous global-to-shared copy instructions (cp.async on Ampere/Hopper), so that while the SM computes iteration $k$ on Tensor Cores, it prefetches iteration $k+1$ from global memory in the background.
   • Vectorized memory coalescing: When the compiler inspects offsets = pid * BLOCK + tl.arange(0, BLOCK), it proves that the memory slice is contiguous and aligned. It automatically synthesizes wide 128-bit vector load instructions (LDG.E.128 in SASS) that fetch 16 bytes per thread in a single transaction. This guarantees coalesced memory bus saturation without requiring manual float4 pointer casting in user code.
3. LLVM IR / NVVM IR: The optimized TritonGPU dialect lowers into LLVM IR annotated with target GPU intrinsics (such as barrier synchronization, special registers, and warp shuffles).
4. PTX virtual ISA: The LLVM NVPTX backend compiles the LLVM IR into PTX virtual assembly.
5. Native SASS via ptxas: Finally, the NVIDIA driver invokes ptxas to allocate physical registers from the SM's 64K register file, schedule instructions, and assemble the final .cubin binary.

Check ownership without a GPU

A CPU enumeration can verify the index arithmetic independently of either compiler. It can't detect generated-code bugs, GPU races, or invalid device memory access. Predict the counts before running it: length 257 requires two 256-lane blocks, with 255 lanes masked out.

Now change the block size to 128 and repeat. Length 257 needs three blocks, with 127 masked lanes. A passing multiple-of-block test alone misses the off-by-one error; that's why both GPU fixtures include 255 and 257.

Six authoring surfaces, six control contracts

A kernel language isn't a ranking from easy to powerful. Each surface chooses which decisions belong to author, compiler, library, and runtime.

CUDA C++

CUDA C++ uses single instruction, multiple threads (SIMT) semantics. You control thread and block geometry, address calculations, synchronization, memory spaces, and low-level intrinsics. That directness fits irregular algorithms, architecture experiments, and cases where a higher layer hides a critical decision. Its portability surface is NVIDIA GPUs, with source, PTX, and fatbins handling generations inside that ecosystem.^{[6]Reference 6CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

Evidence should reach the cubin. Keep ptxas resource output, disassemble native code, profile target hardware, and compare results against a trusted implementation.

Triton block programs

Triton is a Python-embedded domain-specific language (DSL) whose program instances operate on blocks. Its compiler performs block-level data-flow analysis and can automate coalescing, vectorization, shared-memory management, synchronization, and instruction selection.^{[8]Reference 8Triton Programming Guide: Introductionhttps://triton-lang.org/main/programming-guide/chapter-1/introduction.html} Current project documentation lists NVIDIA and AMD GPU support, but backend support and generated code still vary by hardware and Triton release.^{[10]Reference 10Triton Language and Compilerhttps://github.com/triton-lang/triton}

Triton fits custom fusion and dense or structured kernels where block ownership is natural. Inspect specialized IR and target assembly when available, because one Python function can compile into many shape, type, and meta-parameter variants.

CUTLASS and CuTe

CUTLASS supplies CUDA C++ templates for high-performance linear algebra. CuTe, used inside modern CUTLASS, models hierarchical tensor layouts, tensors, copy atoms, and matrix multiply-accumulate atoms. Its C++ surface keeps detailed control over thread-data mapping and architecture features. CuTe DSL brings the same low-level concepts into a Python JIT stack while preserving explicit memory, thread, and data hierarchy.^{[11]Reference 11CUTLASS and CuTe Documentationhttps://docs.nvidia.com/cutlass/latest/overview.html}

Consider this family when matrix multiplication, convolution, attention building blocks, or unusual layouts need explicit layout and hardware-atom control. Whether it exposes a needed capability depends on the specific APIs and releases being compared. Expect a steeper layout-algebra learning curve and a larger specialization space. Use CUTLASS profiler or an equivalent harness, inspect PTX and SASS, and keep compile-time policies in the benchmark manifest.

CUDA Tile

CUDA Tile changes CUDA's unit of thought from one SIMT thread to one block operating collectively on immutable multidimensional tiles. The compiler chooses the number of threads per block and maps tile operations onto registers, shared memory, tensor cores, and other hardware resources. SIMT and tile kernels can coexist in one application.^{[12]Reference 12Writing Tile Kernelshttps://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/writing-tile-kernels.html}

Tooling snapshot, verified September 2, 2026: NVIDIA's CUDA 13.3 documentation lists CUDA Tile in Python through cuda.tile and CUDA Tile C++ in the toolkit from 13.3 onward. Tile shapes must be compile-time powers of two in the documented model. Treat earlier toolkits and different language releases as separate environments, not compatible assumptions.^{[12]Reference 12Writing Tile Kernelshttps://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/writing-tile-kernels.html}

CUDA Tile offers source portability across NVIDIA generations by hiding thread mapping, not cross-vendor portability. Inspect Tile IR or emitted device images where supported, record compiler version, then profile on each architecture you claim to support.

Pallas

Pallas extends JAX with custom kernels for GPU and tensor processing unit (TPU) backends. Kernels use references (Refs) to memory, launch grids, block specifications (BlockSpecs), and backend-specific pipelining or hardware APIs. A kernel is embedded in a larger JAX program through an API such as pl.kernel, pl.pallas_call, or a backend-specific wrapper; their signatures and supported transformations differ.^{[13]Reference 13Pallas: a JAX Kernel Languagehttps://docs.jax.dev/en/latest/pallas/}

Pallas is useful when JAX composition or GPU and TPU integration matters. A custom kernel doesn't automatically inherit every JAX transformation; check the chosen API's support and supply derivative rules when needed. Shared concepts don't guarantee one source or tuning configuration performs equally across backends. JAX describes Pallas as under active development, and its quickstart still warns that the API is experimental. Hardware coverage is operation-specific: JAX 0.11 added Ampere matrix instructions and asynchronous copies to Mosaic GPU. The older Pallas Triton backend is deprecated in JAX 0.11 and scheduled for removal.^{[14]Reference 14Pallas Changeloghttps://docs.jax.dev/en/latest/pallas/CHANGELOG.html} Pin JAX and backend versions, test interpret mode where supported, inspect lowering, and profile each target.

NKI

Neuron Kernel Interface (NKI) targets AWS Trainium, Inferentia2, Trainium2, and Trainium3 NeuronCores. Its high-level nki.language API handles tensor indexing and placement, while lower-level nki.isa exposes hardware operations for computation, data movement, dynamic control flow, and cross-core communication. The documented execution pattern moves inputs from high-bandwidth memory (HBM) into the on-chip state buffer (SBUF), computes on NeuronCore engines, then stores outputs to HBM.^{[15]Reference 15About Neuron Kernel Interfacehttps://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/get-started/about/index.html}

NKI fits workloads committed to AWS Neuron hardware that need custom operations or tighter control than framework compilation supplies. It isn't a CUDA portability layer. Keep Neuron compiler and instance type in the artifact record, use framework-level correctness comparisons, and collect Neuron profiler evidence.

The comparison is easier to use after each model has a concrete meaning:

A Triton kernel and a CuTe rewrite pass the same correctness fixture. Triton is faster on H100, while CuTe is faster on B200. What must the review artifact contain before either implementation becomes the default?

Answer

Keep both source revisions, compiler and backend versions, specialization keys, compile flags, native targets, correctness tolerances, cold and warm timings, emitted-code evidence, and per-target profiler results. Choose against the deployment GPU mix and latency objective. One architecture's timing or a PTX instruction count can't establish a portable winner.

Diagnose the first failing boundary

Compiler failures become manageable when you ask which representation last satisfied its contract.

Don't skip directly from source to profiler. A profiler identifies an executed bottleneck, but it can't explain a missing image that never loaded. Disassembly proves emitted code, but it can't prove correct answers. Correctness tests prove values for their fixtures, but they can't establish speed. The evidence layers complement one another.

Choose by the decision you must own

Start with a framework or vendor library if it meets the need. For a custom kernel, choose by the decision you need to express:

1. Use a framework or vendor library when an existing operation already meets correctness and performance needs.
2. Use Triton when custom block structure, fusion, or masking is central and compiler-managed thread mapping is acceptable.
3. Use Pallas when the kernel must compose naturally with JAX or span its GPU and TPU ecosystem.
4. Use NKI when NeuronCore memory placement or engines are the target, not an afterthought.
5. Use CUDA Tile when block-level tile semantics fit and NVIDIA's compiler should own intra-block threads.
6. Use CUTLASS/CuTe when tensor layouts, asynchronous copies, and matrix hardware atoms need explicit control.
7. Use CUDA C++ when per-thread behavior, synchronization, special instructions, or unsupported irregularity requires it.

These are alternatives, not a ranked ladder. A prototype may begin in Triton, expose a compiler limitation, and move one hot kernel to CuTe or CUDA C++. Another kernel may move upward after a compiler learns the missing transformation.

If you have a compatible NVIDIA environment, practice the full artifact path before moving on. Build the vector-add executable with an sm_90 cubin and compute_90 PTX, save its manifest and cuobjdump image inventory, then run forced-PTX and cubin-only checks in fresh processes. Capture the resource report, SASS, correctness output, and separate cold and warm timings in one review artifact. Run the Triton version on the same five nonempty shapes. Both implementations should match the same reference exactly, and the strided-view test must reject before launch. For a memory-safety check, run the CUDA executable under compute-sanitizer --tool memcheck ./build/vector_add. Capture failures as well as passes. On CPU-only machines, complete the indexing and cache-identity exercises without claiming a GPU build or benchmark.

Mastery check