code: def cross_entropy_fwd_kernel(
      loss_ptr,  # data ptrs
      lse_ptr,
      logits_ptr,
      labels_ptr,
      smoothing,
      logit_scale,
      lse_square_scale,
      ignored_index,
      total_classes,
      class_start_idx,  # Useful for tensor parallel when each rank only has a subset of classes
      n_cols,  # shapes
      n_rows,
      logits_row_stride,  # strides
      BLOCK_SIZE: tl.constexpr,
      HAS_SMOOTHING: tl.constexpr,
      # if SPLIT (e.g. tensor parallel), don't include the LSE in the loss since it's not the final LSE
      SPLIT: tl.constexpr,
  ):
      row_idx = tl.program_id(0)
      col_block_idx = tl.program_id(1)
      logits_ptr = logits_ptr + row_idx * logits_row_stride.to(tl.int64)
      col_offsets = col_block_idx * BLOCK_SIZE + tl.arange(
          0, BLOCK_SIZE
      )
      label_idx = tl.load(labels_ptr + row_idx)
      logits = (
          tl.load(
              logits_ptr + col_offsets,
              mask=col_offsets < n_cols,
              other=-float("inf"),
          ).to(tl.float32)
          * logit_scale
      )
      max_logits = tl.max(logits, 0)
      if HAS_SMOOTHING:
          sum_logits = tl.sum(
              tl.where(col_offsets < n_cols, logits, 0.0), 0
          )
      lse = tl.log(tl.sum(tl.exp(logits - max_logits), 0)) + max_logits
      tl.store(lse_ptr + col_block_idx * n_rows + row_idx, lse)
      if label_idx == ignored_index:
          loss = 0.0
      else:
          label_idx -= class_start_idx
          if (
              label_idx >= col_block_idx * BLOCK_SIZE
              and label_idx
              < min(n_cols, (col_block_idx + 1) * BLOCK_SIZE)
          ):
              logits_label = (
                  tl.load(logits_ptr + label_idx) * logit_scale
              )
              if HAS_SMOOTHING:
                  loss = (
                      (lse if not SPLIT else 0.0)
                      - smoothing * sum_logits / total_classes
                      - (1 - smoothing) * logits_label
                  )
              else:
                  loss = (lse if not SPLIT else 0.0) - logits_label
          else:
              # If label is out of bounds, we set the CE loss to 0.0. But we still want the smoothing loss
              if HAS_SMOOTHING:
                  loss = smoothing * (
                      (lse if not SPLIT else 0.0)
                      - sum_logits / total_classes
                  )
              else:
                  loss = 0.0
          if not SPLIT:
              loss += lse_square_scale * lse * lse
      tl.store(loss_ptr + col_block_idx * n_rows + row_idx, loss)
  
file_path: flash-attention/flash_attn/ops/triton/cross_entropy.py
imports: from typing import Tuple, Optional, Union
  import torch
  from einops import rearrange
  import triton
  import triton.language
label_desc: Use image node_id 1 for calling a global function with example usage: cross_entropy_fwd_kernel(loss_ptr, lse_ptr, logits_ptr, labels_ptr, smoothing, logit_scale, lse_square_scale, ignored_index, total_classes, class_start_idx, n_cols, n_rows, logits_row_stride, BLOCK_SIZE, HAS_SMOOTHING, SPLIT) without return types