Write a custom CUDA kernel to optimize SphereFace Loss (A-Softmax).

Formula: Loss = -log( exp(s * psi(theta_yi)) / Sum(exp(s * cos(theta_j_modified))) )
Where:
- input is the cosine matrix.
- For non-target classes: logit = s * cos(theta).
- For target class yi: logit = s * psi(theta_yi).
- psi(theta) is the A-Softmax function: (-1)^k * cos(m*theta) - 2k, where theta is in [k*pi/m, (k+1)*pi/m].

Problem Analysis:
1. Complex Logic overhead: Computing psi(theta) involves acos, floor, cos, and conditional logic. Doing this in Python via vectorized masking creates huge intermediate tensors and overhead.
2. Memory Bandwidth: Similar to other margin losses, the standard implementation involves scatter/gather and global memory round-trips before the final CrossEntropy.

Optimization Strategy: Fused Piecewise Function and Reduction

1. Fused Logic: Implement the piecewise psi(theta) function as a __device__ helper.
   - k = floor(angle * m / PI)
   - val = pow(-1, k) * cos(m * angle) - 2 * k
   This calculation is performed only when the column index matches the target label.

2. One-Block-per-Row: Launch one block per sample to handle the row-wise softmax reduction.

3. Online Reduction: Iterate through the row (vectorized float4 loads) to compute Max and SumExp. Apply the psi(theta) transformation on-the-fly for the target class index.

4. Result: Compute the NLL loss based on the modified target logit and the reduction results, writing only the final scalar loss per sample.
  
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
import math

BATCH_SIZE = 512
NUM_CLASSES = 10000 
SHAPE = (BATCH_SIZE, NUM_CLASSES)

# SphereFace 超参数
SCALE_S = 64.0
MARGIN_M = 4 

class SphereFaceLoss(nn.Module):
    """
    Standard PyTorch implementation of SphereFace (A-Softmax).
    psi(theta) = (-1)^k * cos(m*theta) - 2k
    """
    def __init__(self, s=64.0, m=4, reduction='mean'):
        super(SphereFaceLoss, self).__init__()
        self.s = s
        self.m = m
        self.reduction = reduction
        self.pi = math.pi

    def forward(self, cosine: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        # cosine: (N, C) should be clamped to (-1, 1)
        # label: (N)
        
        # 1. Gather target cosines: cos(theta_yi)
        target_cosine = cosine.gather(1, label.view(-1, 1)).squeeze(1)
        
        # 2. Compute theta: acos(x)
        # clamp to avoid nan
        input_theta = torch.acos(torch.clamp(target_cosine, -1.0 + 1e-7, 1.0 - 1e-7))
        
        # 3. Compute k
        # theta in [k*pi/m, (k+1)*pi/m]
        # k = floor(theta * m / pi)
        k = (input_theta * self.m / self.pi).floor()
        
        # 4. Compute psi(theta)
        # (-1)^k * cos(m*theta) - 2k
        minus_one_pow_k = torch.pow(-1, k)
        cos_m_theta = torch.cos(input_theta * self.m)
        psi = minus_one_pow_k * cos_m_theta - 2 * k
        
        # 5. Scatter back to logits
        # In SphereFace, logits = s * cosine for j!=yi, and s * psi for j==yi
        
        # Create a copy or use scatter
        one_hot = torch.zeros_like(cosine)
        one_hot.scatter_(1, label.view(-1, 1), 1.0)
        
        # logits = (1 - one_hot) * cosine + one_hot * psi
        logits = cosine * (1.0 - one_hot) + psi.unsqueeze(1) * one_hot
        
        # 6. Scale
        logits = logits * self.s
        
        # 7. CrossEntropy
        loss = F.cross_entropy(logits, label, reduction='none')
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, s=64.0, m=4, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = SphereFaceLoss(s=s, m=m, reduction=reduction)
    
    def forward(self, cosine, label):
        return self.loss_fn(cosine, label)

def get_inputs():
    cosine = torch.randn(SHAPE, dtype=torch.float32)
    cosine = torch.clamp(cosine, -0.99, 0.99)
    label = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [cosine.contiguous(), label.contiguous()]

def get_init_inputs():
    return [SCALE_S, MARGIN_M, 'none']