forked from ccf-ai-infra/GPUCodeForces
152 lines
5.0 KiB
Plaintext
152 lines
5.0 KiB
Plaintext
You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm architecture to get speedups.
|
||
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
|
||
|
||
InfoNCE Loss CUDA Optimization with Fused Kernel Design
|
||
|
||
1. Fused Kernel Architecture
|
||
Per-Query Parallelism: Each CUDA block processes one query instance
|
||
|
||
Grid Strategy: gridDim.x = B (batch size), blockDim.x = 256
|
||
|
||
Kernel Fusion: Combines multiple operations in single kernel:
|
||
|
||
Positive similarity computation
|
||
|
||
Negative similarity processing
|
||
|
||
Stable log-sum-exp calculation
|
||
|
||
Loss computation per query
|
||
|
||
2. Three-Phase Computation Pipeline
|
||
Phase 1: Positive logit calculation
|
||
|
||
Grid-stride dot product computation
|
||
|
||
Block-level sum reduction using shared memory
|
||
|
||
Phase 2: Stable log-sum-exp (Max finding)
|
||
|
||
Find global maximum across positive and negative logits
|
||
|
||
Block-level max reduction with shared memory
|
||
|
||
Phase 3: Stable log-sum-exp (Sum calculation)
|
||
|
||
Compute sum of exponentials with numerical stability
|
||
|
||
Block-level sum reduction
|
||
|
||
3. Memory Access Optimization
|
||
Coalesced Access: Sequential memory access patterns for query/positive data
|
||
|
||
Shared Memory Utilization:
|
||
|
||
s_dot[BLOCK_SIZE] for dot product reduction
|
||
|
||
s_max[BLOCK_SIZE] for max reduction
|
||
|
||
s_sum[BLOCK_SIZE] for sum reduction
|
||
|
||
Contiguous Tensors: Ensure all input tensors are contiguous
|
||
|
||
4. Numerical Stability Features
|
||
Log-Sum-Exp Trick: Subtract maximum before exponentiation
|
||
|
||
Stable Division: Apply temperature scaling after reduction
|
||
|
||
Float Safety: Use FLT_MAX for initial max values
|
||
|
||
5. Parallel Reduction Patterns
|
||
|
||
6. Performance Optimizations
|
||
Minimal Global Memory Writes: Only thread 0 writes final loss per query
|
||
|
||
Grid-Stride Loops: Handle arbitrary feature dimensions (D) and negative counts (N)
|
||
|
||
Efficient Thread Utilization: All threads participate in computations
|
||
|
||
Pre-computed Similarities: Negative similarities computed via optimized cuBLAS matmul
|
||
|
||
7. Implementation Features
|
||
Batch Independence: Each query processed independently enabling parallelism
|
||
|
||
Template-Free Design: Optimized for float32 precision
|
||
|
||
Comprehensive Validation: Tensor shape and device checking
|
||
|
||
PyTorch Integration: Seamless tensor passing and automatic differentiation
|
||
|
||
Key CUDA Concepts Used
|
||
Block-Level Reduction for parallel statistics computation
|
||
|
||
Shared Memory Synchronization using __syncthreads()
|
||
|
||
Grid-Stride Loops for workload distribution across feature dimensions
|
||
|
||
Memory Coalescing for efficient global memory access
|
||
|
||
Kernel Fusion combining multiple mathematical operations
|
||
|
||
Workflow Summary
|
||
Python Pre-processing: Compute negative similarities using optimized matrix multiplication
|
||
|
||
CUDA Kernel Execution: Fused computation of positive similarity and stable loss calculation
|
||
|
||
Parallel Reduction: Each block computes loss for one query instance
|
||
|
||
Final Aggregation: Mean reduction across all query losses
|
||
|
||
Expected Performance Benefits
|
||
3-8x speedup over PyTorch implementation for large batch sizes
|
||
|
||
Better numerical stability with careful floating-point handling
|
||
|
||
Reduced memory bandwidth through kernel fusion
|
||
|
||
Scalable performance with increasing batch and feature dimensions
|
||
|
||
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
|
||
import torch.nn.functional as F
|
||
|
||
BATCH_SIZE = 256
|
||
FEATURE_DIM = 512
|
||
TEMPERATURE = 0.1
|
||
N_NEGATIVES = BATCH_SIZE * 10
|
||
|
||
|
||
class Model(nn.Module):
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.temperature = TEMPERATURE
|
||
|
||
def forward(self, query: torch.Tensor, positive: torch.Tensor, negatives: torch.Tensor) -> torch.Tensor:
|
||
# InfoNCE Loss实现
|
||
# (B, D) vs (B, D) -> (B,)
|
||
positive_sim = F.cosine_similarity(query, positive, dim=1) / self.temperature
|
||
|
||
# (B, D) @ (D, N_NEG) -> (B, N_NEG)
|
||
negative_sims = torch.matmul(query, negatives.t()) / self.temperature
|
||
|
||
# 拼接: (B, 1) 和 (B, N_NEG) -> (B, 1 + N_NEG)
|
||
logits = torch.cat([positive_sim.unsqueeze(1), negative_sims], dim=1)
|
||
|
||
# 标签总是 0，因为正样本总是在索引 0
|
||
labels = torch.zeros(query.size(0), dtype=torch.long, device=query.device)
|
||
|
||
loss = F.cross_entropy(logits, labels)
|
||
return loss
|
||
|
||
|
||
def get_inputs():
|
||
query = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
|
||
positive = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
|
||
negatives = F.normalize(torch.randn(N_NEGATIVES, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
|
||
return [query, positive, negatives]
|
||
|
||
|
||
def get_init_inputs():
|
||
return [] |