import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.cpp_extension import load_inline # CUDA source for optimized rearrangement kernels with vectorized memory access and cache hints attention_cuda_source = ''' #include #include __global__ void __launch_bounds__(1024) split_transpose_kernel( const float4* __restrict__ qkv, float4* __restrict__ q_out, float4* __restrict__ k_out, float4* __restrict__ v_out, const int T, const int nh, const int hs_v4) { const int t = blockIdx.x; const int b = blockIdx.y; const int nh_idx = threadIdx.y; const int hs_v4_idx = threadIdx.x; if (t < T && nh_idx < nh && hs_v4_idx < hs_v4) { const int C_v4 = nh * hs_v4; const int qkv_row_offset = (b * T + t) * (3 * C_v4); const int qkv_idx = qkv_row_offset + nh_idx * hs_v4 + hs_v4_idx; const int out_idx = b * (nh * T * hs_v4) + nh_idx * (T * hs_v4) + t * hs_v4 + hs_v4_idx; // Vectorized reads with __ldg hints for read-only cache optimization const float4 q_val = __ldg(&qkv[qkv_idx]); const float4 k_val = __ldg(&qkv[qkv_idx + C_v4]); const float4 v_val = __ldg(&qkv[qkv_idx + 2 * C_v4]); q_out[out_idx] = q_val; k_out[out_idx] = k_val; v_out[out_idx] = v_val; } } __global__ void __launch_bounds__(1024) merge_transpose_stride_kernel( const float4* __restrict__ y, float4* __restrict__ merged, const int T, const int nh, const int hs_v4, const int sB_v4, const int sNH_v4, const int sT_v4) { const int t = blockIdx.x; const int b = blockIdx.y; const int nh_idx = threadIdx.y; const int hs_v4_idx = threadIdx.x; if (t < T && nh_idx < nh && hs_v4_idx < hs_v4) { const int C_v4 = nh * hs_v4; const int y_idx = b * sB_v4 + nh_idx * sNH_v4 + t * sT_v4 + hs_v4_idx; const int merged_idx = (b * T + t) * C_v4 + nh_idx * hs_v4 + hs_v4_idx; merged[merged_idx] = __ldg(&y[y_idx]); } } void launch_split_transpose(torch::Tensor qkv, torch::Tensor q, torch::Tensor k, torch::Tensor v) { const int B = qkv.size(0); const int T = qkv.size(1); const int nh = q.size(1); const int hs_v4 = q.size(3) / 4; dim3 grid(T, B); dim3 block(hs_v4, nh); split_transpose_kernel<<>>( (const float4*)qkv.data_ptr(), (float4*)q.data_ptr(), (float4*)k.data_ptr(), (float4*)v.data_ptr(), T, nh, hs_v4 ); } void launch_merge_transpose(torch::Tensor y, torch::Tensor merged) { const int B = y.size(0); const int nh = y.size(1); const int T = y.size(2); const int hs_v4 = y.size(3) / 4; dim3 grid(T, B); dim3 block(hs_v4, nh); merge_transpose_stride_kernel<<>>( (const float4*)y.data_ptr(), (float4*)merged.data_ptr(), T, nh, hs_v4, (int)(y.stride(0) / 4), (int)(y.stride(1) / 4), (int)(y.stride(2) / 4) ); } ''' attention_cpp_source = ''' void launch_split_transpose(torch::Tensor qkv, torch::Tensor q, torch::Tensor k, torch::Tensor v); void launch_merge_transpose(torch::Tensor y, torch::Tensor merged); ''' # Compile optimized inline CUDA extension attention_ops = load_inline( name="attention_ops_vfinal", cpp_sources=attention_cpp_source, cuda_sources=attention_cuda_source, functions=["launch_split_transpose", "launch_merge_transpose"], verbose=False, ) class Model(nn.Module): def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen): super().__init__() assert n_embd % n_head == 0 self.c_attn = nn.Linear(n_embd, 3 * n_embd) self.c_proj = nn.Linear(n_embd, n_embd) self.attn_dropout = nn.Dropout(attn_pdrop) self.resid_dropout = nn.Dropout(resid_pdrop) self.register_buffer("bias", torch.tril(torch.ones(max_seqlen, max_seqlen)) .view(1, 1, max_seqlen, max_seqlen)) self.n_head = n_head self.n_embd = n_embd self.attn_pdrop_val = attn_pdrop def forward(self, x): B, T, C = x.size() nh = self.n_head hs = C // nh # 1. QKV projection using fused addmm to reduce kernel launch overhead qkv = torch.addmm(self.c_attn.bias, x.view(-1, C), self.c_attn.weight.t()).view(B, T, 3 * C) # 2. Optimized vectorized split-transpose using custom CUDA kernel q = torch.empty(B, nh, T, hs, device=x.device, dtype=x.dtype) k = torch.empty(B, nh, T, hs, device=x.device, dtype=x.dtype) v = torch.empty(B, nh, T, hs, device=x.device, dtype=x.dtype) attention_ops.launch_split_transpose(qkv, q, k, v) # 3. Core attention computation using FlashAttention-accelerated SDPA y = F.scaled_dot_product_attention( q, k, v, dropout_p=self.attn_pdrop_val if self.training else 0.0, is_causal=True ) # 4. Fused head-merging and transposition handle potentially non-contiguous attention output merged = torch.empty(B, T, C, device=x.device, dtype=x.dtype) attention_ops.launch_merge_transpose(y, merged) # 5. Final output projection and dropout proj = torch.addmm(self.c_proj.bias, merged.view(-1, C), self.c_proj.weight.t()).view(B, T, C) return self.resid_dropout(proj)

Attention-Kernel Speedup

Step Code View