Write a custom CUDA kernel to optimize PolyLoss (Poly-1 formulation).

Formula: Loss = CrossEntropy + epsilon * (1 - Pt)
Where Pt is the softmax probability of the target class: Pt = exp(logit_t) / sum(exp(logits)).
And CrossEntropy = -log(Pt).
So, Loss = -log(Pt) + epsilon * (1 - Pt).

Problem Analysis:
1. Memory Explosion: A standard PyTorch implementation calculates F.softmax(logits) first. For a batch size N and classes C, this materializes a full (N, C) tensor of probabilities, even though the loss only depends on the specific probability of the target class Pt.
2. Bandwidth Waste: Reading and writing the full probability matrix consumes massive global memory bandwidth.
3. Redundant Computation: Calculating CrossEntropy separately from the probability term (1 - Pt) involves redundant Log-Sum-Exp operations.

Optimization Strategy: Fused Softmax-Reduction Kernel

The strategy is to fuse the Softmax normalization logic and the Loss calculation into a single kernel that operates row-wise.

1. One-Block-per-Row: Launch one CUDA block for each sample in the batch. The block keeps the row data in Shared Memory (or registers) to minimize Global Memory access.

2. Shared Memory Caching:
Load the entire row of logits (size C) into Shared Memory.
This reduces global memory reads to exactly 1 pass per element.

3. Fused Reduction (Log-Sum-Exp):
Pass 1 (Max): Find the maximum logit M in the row for numerical stability.
Pass 2 (Sum): Compute the sum of exponentials S = sum(exp(x_i - M)).
Simultaneously, identify the target logit x_t corresponding to the target label.

4. Direct Loss Computation:
Compute log_pt = (x_t - M) - log(S).
Compute pt = exp(log_pt).
Compute loss = -log_pt + epsilon * (1 - pt).
Write only the final scalar loss to global memory.

This approach completely eliminates the (N, C) intermediate probability tensor, reducing memory writes by a factor of C.
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn
import torch.nn.functional as F


BATCH_SIZE = 4096
NUM_CLASSES = 1000 # ImageNet-1K
SHAPE = (BATCH_SIZE, NUM_CLASSES)

EPSILON = 2.0 
REDUCTION = 'none'

class Poly1CrossEntropyLoss(nn.Module):
    """
    PolyLoss: A Polynomial Expansion Perspective of Classification Loss Functions (ICLR 2022)
    Equation: L = -log(Pt) + epsilon * (1 - Pt)
    """
    def __init__(self, epsilon=2.0, reduction='mean'):
        super(Poly1CrossEntropyLoss, self).__init__()
        self.epsilon = epsilon
        self.reduction = reduction

    def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        # logits: (N, C)
        # labels: (N) long
        
        # 1. 计算 Pt (Target Probabilities)
        probs = F.softmax(logits, dim=-1)
        
        pt = probs.gather(1, labels.unsqueeze(1)).squeeze(1)
        
        # 2. 计算 Cross Entropy: -log(Pt)
        ce_loss = F.cross_entropy(logits, labels, reduction='none')
        
        # 3. Poly-1 Term: eps * (1 - Pt)
        poly1 = self.epsilon * (1 - pt)
        
        # 4. Final Loss
        loss = ce_loss + poly1
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, epsilon=2.0, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = Poly1CrossEntropyLoss(epsilon=epsilon, reduction=reduction)
    
    def forward(self, logits, labels):
        return self.loss_fn(logits, labels)

def get_inputs():
    logits = torch.randn(SHAPE, dtype=torch.float32)
    labels = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [logits.contiguous(), labels.contiguous()]

def get_init_inputs():
    return [EPSILON, REDUCTION]