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 []