# Kernel Optimization Guide: Triton & Inline CUDA This document outlines how to improve speedup by writing fused and optimized kernels using **Triton** and/or **inline PyTorch CUDA** in a single-file implementation. Solutions may use Triton kernels, inline CUDA code, or both. ## GPU Architecture NVIDIA A100 80GB ## Requirements - **Single-File Implementation:** Develop fused kernels within one file. - **Model Class:** The solution must include a class `Model` (an instance of `nn.Module`), with the main computation in its `forward` method. Do not write code for instantiating Model, getting inputs, or other usage boilerplate. - **Preserve Initialization:** Do not change the initialization of the `Model` class. - **Simplicity & Readability:** Write simple, easy-to-understand code and include clear comments. - **Focus on Efficiency:** Concentrate solely on efficient kernel coding. - **No Fallback Implementation:** Do not include any alternative or fallback code. - **Optimize For Single Forward Pass:** Do not cache values in initialization for use across different forward passes. You may optimize initialization to make each forward pass more efficient. - **No Logs:** Do not worry about writing logs. - **Assume Packages Available:** Do not check for packages; assume all required packages are available. - **Error Handling:** Any terminal output or errors will be reviewed by an LLM for feedback. Remember that just because you observe zero speedup or speedup less than 1, this is still a valid solution and the metric should be extracted. - **Don't Game the System:** Do not include placeholder kernels just to satisfy requirements; kernels must perform real work. Do not cache the final result; return it. - **Multiple Kernels Allowed:** You may define more than one kernel in the file. - **Do Not Use:** `torch.compile`, `torch.jit.script`, `torch.backends.cuda.matmul.allow_tf32`, `torch.backends.cudnn.allow_tf32`, `torch.backends.cudnn.benchmark`, `cudaGraph`, `cudaStream` or half/mixed precision as they obscure true results. - **Avoid Templates:** Use plain fused kernel functions without C++ templates. Do not use PyBind. ## Kernel Optimization Tips 1. **Programming model** - Each kernel launch creates a grid of independent program instances, each processing a tile (or block) of data. 2. **Planning** - Choose tile/block sizes to balance parallelism and resource limits. - Align tiles with memory layout for coalesced access. 3. **Memory access** - Use vectorized, aligned loads and stores. - Minimize divergent accesses and bank conflicts. 4. **Triton** - Use vector operations, minimal control flow, and atomics where appropriate. --- ## Example: Baseline (Element-Wise Add) Simple element-wise addition with no custom kernel: ```python import torch import torch.nn as nn class Model(nn.Module): def __init__(self) -> None: super().__init__() def forward(self, a, b): return a + b ``` --- ## Example: Optimized with Inline CUDA Fused element-wise addition using a custom CUDA kernel and PyTorch’s `load_inline`: ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline elementwise_add_source = ''' #include #include __global__ void elementwise_add_kernel(const float* a, const float* b, float* out, int size) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < size) { out[idx] = a[idx] + b[idx]; } } torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b) { auto size = a.numel(); auto out = torch::zeros_like(a); const int block_size = 256; const int num_blocks = (size + block_size - 1) / block_size; elementwise_add_kernel<<>>(a.data_ptr(), b.data_ptr(), out.data_ptr(), size); return out; } ''' elementwise_add_cpp_source = "torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b);" elementwise_add = load_inline( name="elementwise_add", cpp_sources=elementwise_add_cpp_source, cuda_sources=elementwise_add_source, functions=["elementwise_add_cuda"], verbose=True, extra_cflags=[""], extra_ldflags=[""], ) class Model(nn.Module): def __init__(self) -> None: super().__init__() self.elementwise_add = elementwise_add def forward(self, a, b): return self.elementwise_add.elementwise_add_cuda(a, b) ``` --- ## Example: Optimized with Triton Fused linear + reductions using a custom Triton kernel: ```python import torch import torch.nn as nn import triton import triton.language as tl @triton.jit def fused_kernel(x_ptr, weight_ptr, bias_ptr, out_ptr, stride_x: tl.constexpr, stride_out: tl.constexpr, in_features: tl.constexpr, out_features: tl.constexpr, pad_in: tl.constexpr): pid = tl.program_id(0) x_row = x_ptr + pid * stride_x out_addr = out_ptr + pid * stride_out offs = tl.arange(0, pad_in) x_vals = tl.load(x_row + offs, mask=(offs < in_features), other=0.0) aggr_w = tl.zeros([pad_in], dtype=tl.float32) for i in range(out_features): w_vals = tl.load(weight_ptr + i * in_features + offs, mask=(offs < in_features), other=0.0) aggr_w += w_vals dot_val = tl.sum(x_vals * aggr_w) aggr_bias = 0.0 for i in range(out_features): aggr_bias += tl.load(bias_ptr + i) tl.store(out_addr, dot_val + aggr_bias) class Model(nn.Module): def __init__(self, in_features, out_features): super(Model, self).__init__() self.linear = nn.Linear(in_features, out_features) self.in_features = in_features self.out_features = out_features def forward(self, x): batch_size = x.shape[0] output = torch.empty((batch_size, 1), device=x.device, dtype=x.dtype) if not x.is_contiguous(): x = x.contiguous() if not output.is_contiguous(): output = output.contiguous() stride_x = x.stride(0) stride_out = output.stride(0) pad_in = 1 << ((self.in_features - 1).bit_length()) grid = (batch_size,) fused_kernel[grid]( x, self.linear.weight, self.linear.bias, output, stride_x, stride_out, self.in_features, self.out_features, pad_in ) return output ``` --- The examples above are illustrative. Do not copy them verbatim; adapt the patterns to your task and use Triton, inline CUDA, or both as appropriate.

Mamba-SSD Speedup

Step Code View