IIntroduction to AI Systems Performance EngineeringTechnical textbook
Chapter 11: Triton, CUTLASS, and PyTorch Operators
11.1

Kernel implementation options

Select an implementation level from the required schedule control and complete software lifecycle

Before you begin

Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.

  • Kernel resource model
  • Compiler basics

A profile of one model shows a 90-microsecond chain of normalization, scaling, and residual-add kernels. The engineer first writes the expression with ordinary PyTorch operators and runs torch.compile. If the generated timeline still contains unwanted intermediate transfers, the engineer can implement the block schedule in Triton, compose a matrix path from CUTLASS, or write CUDA C++ for an instruction or synchronization feature unavailable at higher levels. Each option can be measured against the same reference function and input shapes.

The implementation level is the software interface through which the engineer describes GPU work. PyTorch composition describes tensor operations, Triton describes blocked program instances, CUTLASS and CuTe describe reusable layouts and matrix-building blocks, and CUDA exposes threads plus the platform instruction surface. Compiler dumps, generated code, profiler counters, correctness tests, and maintenance requirements show what control each level actually provides for the operation under study.

Performance engineers work through layers of leverage. Framework composition offers portability and automatic differentiation but limited control over fusion and scheduling. Triton exposes block-level programs while a compiler maps them to GPU threads. CUTLASS offers reusable C++ templates for deeply structured matrix kernels. CUDA provides the broadest direct platform control and the greatest responsibility for every detail.

Choosing a lower implementation layer can expose additional scheduling control, but it also expands the correctness surface, build matrix, compilation time, portability burden, and maintenance cost. For the normalization example, remain with compiled PyTorch when it produces the required fusion and reaches the measured bandwidth limit. Use Triton when the missing control is row ownership, reduction structure, masks, block size, or fusion. Use CUTLASS or CuTe when a GEMM-like operation needs a custom layout, datatype, epilogue, or matrix pipeline. Use CUDA C++ when the required instruction, synchronization scope, cluster operation, or memory primitive cannot be expressed reliably through the other interfaces.

The implementation decision is easiest when framed around a specific missing capability. Examples include a required shared-memory layout, a warp-group matrix instruction, a cross-block synchronization method, a custom backward, or graph-level fusion across several operations. A statement such as 'CUDA is faster' is not enough because the answer depends on the operation, shapes, compiler, and target device.

Begin with the framework or vendor implementation

A composition of existing framework operations is the semantic baseline. Framework composition provides autograd, datatype behavior, device dispatch, shape handling, and integration with compilation. Before writing a custom kernel, compile the operation in its real graph and inspect whether the framework already fuses the expression or selects a mature library call.

Vendor libraries are the appropriate baseline for common dense linear algebra and communication. They contain architecture-specific kernels and dispatch logic tested across a broad input range. A custom path is justified when the workload has a repeatable gap: an unsupported fusion, an unusual shape family, a special layout, or an operation that the library does not implement.

Use Triton for block-structured custom kernels

Triton programs describe blocks of indices and values. The blocked programming model fits elementwise fusion, row reductions, normalization, softmax, and many tiled matrix operations. The compiler handles much of the lane mapping and instruction selection, which makes schedule experiments shorter than equivalent CUDA implementations.

Triton still exposes consequential choices: program partition, pointer order, masks, block sizes, warp count, pipeline stages, and compile-time specialization. Use it when these controls are sufficient. Inspect generated code when the block-level explanation predicts efficient access but the measured kernel spills, scatters loads, or misses a target instruction.

Use CUTLASS or CuTe for structured matrix kernels

CUTLASS provides composable C++ building blocks for layouts, copies, matrix operations, collective main loops, epilogues, and device-level dispatch. CUTLASS is appropriate when the operation resembles GEMM but needs a custom datatype, layout, epilogue, grouped schedule, or architecture-specific pipeline.

The template surface is substantial because it represents a large design space. Reusing tested components can be safer than writing every fragment mapping and transfer primitive directly. The cost is compile complexity and a requirement to understand the layout and collective contracts. Prefer an existing operation or example that closely matches the target before creating a new component hierarchy.

Use CUDA when direct platform control is required

CUDA C++ exposes the broad platform surface: thread and cluster cooperation, memory scopes, asynchronous operations, inline PTX when necessary, and direct integration with native libraries and tooling. CUDA C++ is appropriate when a higher-level system cannot express the required operation or generates persistently inferior code for a material workload.

The additional control creates additional obligations. The implementation owns bounds, synchronization, architecture guards, compilation, ABI, error handling, and often separate backward code. Porting to a new architecture can require a new schedule even when the source still compiles. Document the exact capability that required CUDA so future maintainers can reevaluate the decision.

Preserve semantics, eligibility, and fallback

Keep a readable reference implementation that defines outputs, accepted inputs, numerical behavior, mutation, and errors. The optimized path needs an eligibility predicate for device, datatype, shape, strides, alignment, and framework mode. Test conditions immediately inside and outside each predicate boundary.

Evaluate lifecycle as well as peak speed. Include backward, autocast, fake tensors, dynamic shapes, compilation, packaging, and supported architectures. A small gain on one rare shape may not justify a private extension. A frequently executed workload with a large fleet cost can justify a narrower and more expensive implementation.

Worked example: Choosing a path for fused normalization

A PyTorch chain of residual addition, RMS normalization, and gating writes several intermediates and is significant in end-to-end time.

  1. First test compiler fusion. If the graph captures cleanly and generated code reaches the bandwidth roof across target shapes, custom code may add little value.
  2. Write down the semantic contract: residual output behavior, epsilon, input and accumulation datatypes, accepted strides, and gradient requirements. Use the reference to test every later version.
  3. If scheduling or reduction is poor, prototype in Triton. Its block program can express row ownership, reduction, and fused elementwise work with a compact source and autotuned block choices.
  4. Use CUDA when higher-level tools cannot express a required warp primitive, asynchronous pipeline, or architecture-specific operation and the measured performance effect is material.
  5. For every level, compare correctness, shape coverage, compilation and startup, integration, and end-to-end effect—not only the best steady kernel timing.
  6. Retain the previous level as a fallback. A CUDA specialization for one architecture does not remove the need for correct execution elsewhere.

Conclusion: The selected implementation is justified by a named capability and a measured result. In this example, compiled PyTorch remains the implementation when it produces the required fused schedule; Triton is appropriate when explicit row and reduction scheduling fixes the measured gap; CUDA is appropriate only when a required platform operation remains unavailable. Every custom path retains the reference implementation and a tested fallback.

Common errors

  • Selecting a tool by prestige. The narrowest implementation is not the most valuable if a higher layer already reaches the relevant roof.
  • Ignoring framework integration until the kernel is finished. Autograd, compilation, layout, and dispatch are part of the operator design.
  • Removing the high-level reference. A readable semantic baseline is essential for tests, fallbacks, and future maintainers.
  • Moving to a lower layer before identifying the missing control. The rewrite can reproduce the same schedule with more code and the same limit.

Reference summary

  • CUDA C++: direct control, complete platform surface, highest debugging burden, and architecture-specific tuning options.
  • Triton: program instances operate on blocks of values; concise experimentation, autotuning, and compiler-managed lane mapping.
  • CUTLASS/CuTe: composable layouts, copies, matrix operations, and pipelines for high-performance tensor kernels.
  • Vendor libraries: mature dispatch, broad testing, and often the baseline to beat.
  • Framework/compiler generated kernels: broad graph context and automatic fusion, with less predictable code generation.

For the normalization, scaling, and residual chain used in this chapter, first compile the ordinary PyTorch expression and inspect its launches and high-bandwidth memory (HBM) traffic. Use Triton when the missing requirement is control over row ownership, blocked loads, a reduction, or fusion. Use CUTLASS or CuTe when the operation is a matrix multiplication that needs a custom layout, datatype, epilogue, or architecture-specific pipeline. Use CUDA C++ when the required instruction, synchronization scope, or memory operation cannot be expressed reliably through those interfaces.

Knowledge check

What evidence would justify rewriting a correct Triton kernel in CUDA rather than continuing to tune the Triton version?

Show the answer guide
  • A rewrite is justified only after the Triton implementation is correct, tuned across representative shapes, and shown by repeated measurement to miss a stated performance objective.
  • Profiler and generated-code evidence must identify a limitation that further Triton changes cannot express or remove, such as required asynchronous-copy control, unsupported instruction or layout mapping, excessive register allocation, a synchronization pattern, or a compiler code-generation defect.
  • A CUDA prototype must demonstrate a material end-to-end gain on the workload distribution, not only on one favorable microbenchmark, while preserving numerical accuracy and required edge cases.
  • The expected benefit must justify added implementation, integration, testing, portability, and maintenance cost. If a Triton or compiler fix can remove the limitation with similar results, the rewrite is not yet established.
Practice

Write an evidence-based abstraction decision

Choose a real Triton kernel such as fused layer normalization for shapes [1,4096], [32,4096], [1024,4096], and [8192,4096]. Tune at least six valid block, warp, or stage configurations. Inspect generated code and profile the best version. Compare it with an optimized CUDA or library implementation on the same inputs and integration path.

Evidence to keep: Submit the Triton code and tuning space, correctness tests, a benchmark matrix for all declared shapes, generated-code or profiler evidence for the limiting mechanism, and a one-page decision record that either justifies a CUDA rewrite or explicitly rejects it. The decision must quantify expected end-to-end benefit and maintenance cost.

Triton illustrates a useful middle layer: the programmer describes blocks of data and computation while the compiler chooses much of the lane-level mapping.

References

Sources and further reading

The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.

  1. Triton Tutorials
  2. NVIDIA CUTLASS
  3. Custom C++ and CUDA Operators
  4. torch.compile Programming Model
  5. Nsight Compute Profiling Guide