Function kernel_fn

kernels/cross_entropy.py:70–105

Entry point called by bench.py. Must match reference.cross_entropy_ref signature. Args: logits: [batch_size, vocab_size] raw logits (float16 or float32) targets: [batch_size] integer class indices (long) Returns: Scalar mean cross-entropy loss

(logits: torch.Tensor, targets: torch.Tensor)

def kernel_fn(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """
    Entry point called by bench.py. Must match reference.cross_entropy_ref signature.

    Args:
        logits: [batch_size, vocab_size] raw logits (float16 or float32)
        targets: [batch_size] integer class indices (long)

    Returns:
        Scalar mean cross-entropy loss
    """
    assert logits.is_cuda and targets.is_cuda

    # Handle multi-dim: flatten to 2D
    if logits.ndim > 2:
        logits = logits.view(-1, logits.shape[-1])
        targets = targets.view(-1)

    n_rows, n_cols = logits.shape
    assert targets.shape[0] == n_rows

    losses = torch.empty(n_rows, device=logits.device, dtype=torch.float32)

    BLOCK_SIZE = triton.next_power_of_2(n_cols)

    grid = (n_rows,)
    cross_entropy_kernel[grid](
        logits,
        targets,
        losses,
        n_cols,
        logits.stride(0),
        BLOCK_SIZE=BLOCK_SIZE,
    )

    return losses.mean()