forked from ccf-ai-infra/GPUCodeForces
82 lines
3.4 KiB
Plaintext
82 lines
3.4 KiB
Plaintext
LayerNorm CUDA Implementation - Enhanced Optimized Version
|
||
|
||
Key optimization techniques used in this implementation:
|
||
|
||
1.Warp-Level Parallel Reduction: Implements efficient warp-level reduction for mean and variance calculations using warp shuffle operations
|
||
2.Vectorized Memory Access: Utilizes float4 vector loads/stores for coalesced memory access when feature dimension is divisible by 4
|
||
3.Shared Memory Hierarchy: Employs multi-level shared memory for intermediate results between warp and block levels
|
||
4.Dynamic Thread Configuration: Automatically adjusts thread block size based on feature dimension for optimal occupancy
|
||
5.Numerical Stability: Maintains numerical precision with robust variance calculation and epsilon handling
|
||
|
||
Bank Conflict Avoidance: Carefully structures shared memory access patterns to minimize bank conflicts
|
||
|
||
The custom kernel eliminates multiple memory passes by computing mean, variance, and normalization in a single fused operation with optimized memory hierarchy usage across warp, shared, and global memory levels.
|
||
|
||
Technical Features:
|
||
|
||
1.Warp Reduction: Efficient 32-thread warp reduction using __shfl_down_sync
|
||
2.Vectorization: Automatic fallback between vectorized (float4) and scalar operations
|
||
3.Smart Block Sizing: Adaptive thread configuration (64-512 threads) based on feature dimension
|
||
4.Memory Coalescing: Organized memory access patterns for maximum bandwidth utilization
|
||
5.Fused Operations: Combines statistics computation and normalization in one kernel
|
||
|
||
Performance Benefits:
|
||
|
||
1.Reduces global memory traffic by processing entire LayerNorm operation in-place
|
||
2.Eliminates intermediate tensor allocations between mean/variance calculations
|
||
3.Optimizes for various feature dimensions through adaptive thread configuration
|
||
4.Leverages CUDA memory hierarchy for maximum data reuse
|
||
|
||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
Simple model that performs LayerNorm normalization using PyTorch's built-in nn.LayerNorm.
|
||
"""
|
||
|
||
def __init__(self, normalized_shape=None, eps=1e-5, elementwise_affine=True):
|
||
super(Model, self).__init__()
|
||
# 如果未指定normalized_shape，将在forward中动态设置
|
||
self.normalized_shape = normalized_shape
|
||
self.eps = eps
|
||
self.elementwise_affine = elementwise_affine
|
||
self.layernorm = None
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Applies LayerNorm to the input tensor.
|
||
|
||
Args:
|
||
x (torch.Tensor): Input tensor of any shape.
|
||
|
||
Returns:
|
||
torch.Tensor: Output tensor with LayerNorm applied, same shape as input.
|
||
"""
|
||
# 如果layernorm未初始化，根据输入形状动态创建
|
||
if self.layernorm is None:
|
||
if self.normalized_shape is None:
|
||
# 默认对最后一个维度进行归一化
|
||
self.normalized_shape = x.shape[1:]
|
||
self.layernorm = nn.LayerNorm(
|
||
normalized_shape=self.normalized_shape,
|
||
eps=self.eps,
|
||
elementwise_affine=self.elementwise_affine
|
||
).to(x.device)
|
||
|
||
return self.layernorm(x)
|
||
|
||
|
||
batch_size = 16
|
||
dim = 16384
|
||
|
||
|
||
def get_inputs():
|
||
x = torch.randn(batch_size, dim)
|
||
return [x]
|
||
|
||
|
||
def get_init_inputs():
|
||
# 可以传入归一化形状、eps等参数，保持向后兼容
|
||
return [] # 使用默认参数 |