You write custom CUDA kernels to replace the pytorch operators in the given 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 matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.  
  
// Key optimization techniques used in this implementation:
        // 1. Operator Fusion: Fused log_softmax + negative log likelihood into a single kernel
        // 2. Shared Memory Optimization: Utilizes shared memory for logits and exp computations
        // 3. Fast Math Functions: Employs optimized mathematical operations including expf, logf
        // 4. Hierarchical Reduction: Implements warp-level and block-level reductions for statistics
        // 5. Memory Access Coalescing: Organized thread-block mapping for optimal global memory access patterns
        // 6. Numerical Stability: Proper handling of max subtraction for stable softmax computation

        // The custom CUDA implementation provides significant performance improvements over the native PyTorch version
        // by eliminating intermediate tensor allocations and leveraging GPU-specific optimizations.

        // Specific Technical Optimizations:
        // Memory Hierarchy Optimization:
        // 1. Shared Memory: Stores logits and exp values for fast intra-block access
        // 2. Global Memory: Coalesced access patterns for input and target tensors
        // 3. Register Utilization: Extensive use of registers for temporary computations

        // Computational Optimizations:
        // 1. Fast Exponential: expf() with numerical stability considerations
        // 2. Hierarchical Reduction: Warp-level and block-level reductions for max and sum
        // 3. Parallel Statistics: Concurrent computation of max, sum, and final loss

        // Parallelism Strategy:
        // 1. Grid Structure: One block per sample in batch (B blocks)
        // 2. Block Configuration: 256 threads per block for optimal occupancy
        // 3. Workload Distribution: Dynamic workload balancing across classes

        // Numerical Precision:
        // 1. Maintains mathematical equivalence with reference implementation
        // 2. Proper max subtraction for numerical stability in softmax
        // 3. Exact loss computation preserved despite performance optimizations

        // The implementation demonstrates how custom CUDA kernels can dramatically accelerate cross-entropy loss

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 = 4096
NUM_CLASSES = 1000  # 假设分类类别数
FEATURE_DIM = NUM_CLASSES  # CrossEntropyLoss输入最后一维为类别数

class Model(nn.Module):

    def __init__(self):
        super().__init__()
        # CrossEntropyLoss 会自动包含 LogSoftmax + NLLLoss
        self.criterion = nn.CrossEntropyLoss(reduction='mean')

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return self.criterion(input, target)

def get_inputs():
    # CrossEntropyLoss 输入：logits [N, C]
    input_scores = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)

    # 目标标签：每个样本一个类别索引（0 ~ C-1）
    target = torch.randint(0, FEATURE_DIM, (BATCH_SIZE,), dtype=torch.long)

    return [input_scores, target]

def get_init_inputs():
    return []