TensorTonic
Problems
Cross-Entropy Loss as a CUDA Kernel
Loss FunctionsGPU Programming
Medium

Compute the mean cross-entropy loss over a batch of BBB examples with CCC classes. Given a logits matrix of shape (B,C)(B, C)(B,C) and integer targets of length BBB with values in [0,C)[0, C)[0,C), the per-row log-softmax is:

log⁡softmax(z)j=zj−max⁡kzk−log⁡ ⁣∑kexp⁡(zk−max⁡kzk)\log\text{softmax}(z)_j = z_j - \max_k z_k - \log\!\sum_k \exp(z_k - \max_k z_k)logsoftmax(z)j​=zj​−kmax​zk​−logk∑​exp(zk​−kmax​zk​)

and the mean loss across the batch is:

L=−1B∑i=0B−1log⁡softmax(logitsi)targeti\mathcal{L} = -\frac{1}{B} \sum_{i=0}^{B-1} \log\text{softmax}(\text{logits}_i)_{\text{target}_i}L=−B1​i=0∑B−1​logsoftmax(logitsi​)targeti​​

logits\texttt{logits}logits is a device pointer to a row-major B×CB \times CB×C float32 buffer, target\texttt{target}target is a device pointer to a length-BBB int32 buffer, and loss\texttt{loss}loss is a single float32 scalar. Use the max-subtract trick for numerical stability.

Loading visualization...

Examples

Input: B = 2, C = 3, logits = [[1, 2, 3], [3, 2, 1]], target = [2, 0]

Output: 0.40760.40760.4076

Both rows have log-sum-exp ≈0.4076\approx 0.4076≈0.4076, the target logit equals the row max in each case, so −log⁡softmaxtarget≈0.4076-\log\text{softmax}_{\text{target}} \approx 0.4076−logsoftmaxtarget​≈0.4076 for each row, averaging to the same value.

Input: B = 1, C = 4, logits = [[10, 0, 0, 0]], target = [0]

Output: 0.0001360.0001360.000136

The target logit dominates the row, so softmaxtarget≈1\text{softmax}_{\text{target}} \approx 1softmaxtarget​≈1 and the loss is nearly zero.

Hint 1

Subtract the row max from each logit before calling expf, otherwise large logits overflow float32.

Hint 2

Use one block per row to compute the max and log-sum-exp with shared-memory reductions, then have thread 0 contribute that row's loss with atomicAdd into a single scalar.

Hint 3

A small finalize kernel that divides the accumulated sum by BBB keeps the launcher clean and avoids host-side math on the scalar.

Requirements

  • Subtract the per-row max from logits before exponentiating for numerical stability.
  • Compute the per-row log-sum-exp log⁡ ⁣∑kexp⁡(zk−max⁡kzk)\log\!\sum_k \exp(z_k - \max_k z_k)log∑k​exp(zk​−maxk​zk​) via a block-wide reduction.
  • Aggregate per-row contributions −(zti−max⁡−lse)-(z_{t_i} - \max - \text{lse})−(zti​​−max−lse) across all rows and divide by BBB to get the mean.
  • Do not change the solve\texttt{solve}solve signature; pointers are device pointers owned by the harness.
  • Write the final scalar mean loss to loss[0]\texttt{loss[0]}loss[0].

Constraints

  • 1≤B≤40961 \le B \le 40961≤B≤4096
  • 1≤C≤40961 \le C \le 40961≤C≤4096
  • logits\texttt{logits}logits is a device pointer to B⋅CB \cdot CB⋅C float32 values in row-major order.
  • target\texttt{target}target is a device pointer to BBB int32 values in the range [0,C)[0, C)[0,C).
  • loss\texttt{loss}loss is a device pointer to a single float32 scalar.
Try Similar Problems
Cosine Similarity as a CUDA KernelMediumCross Entropy Loss (Mean Reduction) in TritonMediumDropout as a CUDA KernelEasyScaled Dot-Product Attention in CUDAHardArgmax as a CUDA KernelMediumRMS Normalization as a CUDA KernelMediumMean and Variance as a CUDA KernelMediumLayer Normalization as a CUDA KernelMediumSoftmax as a CUDA KernelMediumMatrix Multiplication as a CUDA KernelMediumL1 Normalization as a CUDA KernelMediumCUDA Dot Product: Multi-Block ReductionMedium

Accepts: array

Accepts: array

You must run your code first.

Kernel performance

Submit a correct solution to see your kernel performance.

Cross-Entropy Loss as a CUDA Kernel

Cross-Entropy Loss as a CUDA Kernel

Loss FunctionsGPU Programming
Medium

Compute the mean cross-entropy loss over a batch of BBB examples with CCC classes. Given a logits matrix of shape (B,C)(B, C)(B,C) and integer targets of length BBB with values in [0,C)[0, C)[0,C), the per-row log-softmax is:

log⁡softmax(z)j=zj−max⁡kzk−log⁡ ⁣∑kexp⁡(zk−max⁡kzk)\log\text{softmax}(z)_j = z_j - \max_k z_k - \log\!\sum_k \exp(z_k - \max_k z_k)logsoftmax(z)j​=zj​−kmax​zk​−logk∑​exp(zk​−kmax​zk​)

and the mean loss across the batch is:

L=−1B∑i=0B−1log⁡softmax(logitsi)targeti\mathcal{L} = -\frac{1}{B} \sum_{i=0}^{B-1} \log\text{softmax}(\text{logits}_i)_{\text{target}_i}L=−B1​i=0∑B−1​logsoftmax(logitsi​)targeti​​

logits\texttt{logits}logits is a device pointer to a row-major B×CB \times CB×C float32 buffer, target\texttt{target}target is a device pointer to a length-BBB int32 buffer, and loss\texttt{loss}loss is a single float32 scalar. Use the max-subtract trick for numerical stability.

Loading visualization...

Examples

Input: B = 2, C = 3, logits = [[1, 2, 3], [3, 2, 1]], target = [2, 0]

Output: 0.40760.40760.4076

Both rows have log-sum-exp ≈0.4076\approx 0.4076≈0.4076, the target logit equals the row max in each case, so −log⁡softmaxtarget≈0.4076-\log\text{softmax}_{\text{target}} \approx 0.4076−logsoftmaxtarget​≈0.4076 for each row, averaging to the same value.

Input: B = 1, C = 4, logits = [[10, 0, 0, 0]], target = [0]

Output: 0.0001360.0001360.000136

The target logit dominates the row, so softmaxtarget≈1\text{softmax}_{\text{target}} \approx 1softmaxtarget​≈1 and the loss is nearly zero.

Hint 1

Subtract the row max from each logit before calling expf, otherwise large logits overflow float32.

Hint 2

Use one block per row to compute the max and log-sum-exp with shared-memory reductions, then have thread 0 contribute that row's loss with atomicAdd into a single scalar.

Hint 3

A small finalize kernel that divides the accumulated sum by BBB keeps the launcher clean and avoids host-side math on the scalar.

Requirements

  • Subtract the per-row max from logits before exponentiating for numerical stability.
  • Compute the per-row log-sum-exp log⁡ ⁣∑kexp⁡(zk−max⁡kzk)\log\!\sum_k \exp(z_k - \max_k z_k)log∑k​exp(zk​−maxk​zk​) via a block-wide reduction.
  • Aggregate per-row contributions −(zti−max⁡−lse)-(z_{t_i} - \max - \text{lse})−(zti​​−max−lse) across all rows and divide by BBB to get the mean.
  • Do not change the solve\texttt{solve}solve signature; pointers are device pointers owned by the harness.
  • Write the final scalar mean loss to loss[0]\texttt{loss[0]}loss[0].

Constraints

  • 1≤B≤40961 \le B \le 40961≤B≤4096
  • 1≤C≤40961 \le C \le 40961≤C≤4096
  • logits\texttt{logits}logits is a device pointer to B⋅CB \cdot CB⋅C float32 values in row-major order.
  • target\texttt{target}target is a device pointer to BBB int32 values in the range [0,C)[0, C)[0,C).
  • loss\texttt{loss}loss is a device pointer to a single float32 scalar.
Try Similar Problems
Cosine Similarity as a CUDA KernelMediumCross Entropy Loss (Mean Reduction) in TritonMediumDropout as a CUDA KernelEasyScaled Dot-Product Attention in CUDAHardArgmax as a CUDA KernelMediumRMS Normalization as a CUDA KernelMediumMean and Variance as a CUDA KernelMediumLayer Normalization as a CUDA KernelMediumSoftmax as a CUDA KernelMediumMatrix Multiplication as a CUDA KernelMediumL1 Normalization as a CUDA KernelMediumCUDA Dot Product: Multi-Block ReductionMedium

Accepts: array

Accepts: array

You must run your code first.

Kernel performance

Submit a correct solution to see your kernel performance.

Problems
Loading...
1 / 1