Skip to content

[Feat][Kernel] Fused CE LogProb without materializing logits + fused backward #155

Closed

Description

@KJLdefeated

Why

Avoid landing large [B, S, V] logits for large-vocabulary models, so the saved memory can be used for larger batches or longer CoT responses. The forward path streams vocab blocks with online softmax; the backward path recomputes tiles to trade compute for memory instead of storing full logits/probabilities.

Implementation

  • CUDA: use SM90 TMA/WGMMA-oriented streaming where available, with SM80 cp.async/mma.sync fallback; keep online-softmax state in registers and shared memory.
  • Triton: provide a portable semantic baseline and tolerance test target for CUDA/ROCm native implementations.

#122

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions