# 融合交叉熵实验项目

> **当前状态：** `celoss.py` 是用于推导和实验的未完成实现，当前不能作为正确、
> 稳定或高性能的交叉熵实现使用。代码保留了探索过程，适合用于分析和继续修复。

## 实现目标

代码尝试融合语言模型输出投影 `(B, N, D) @ (D, V)`、在线 softmax 和负对数似然，
避免在 DRAM 中保存完整的 `(B, N, V)` logits。最终目标是只输出每个 token 的损失，
再在包装函数中取平均值。

## 核心数据流

1. `naive_ce_loss()` 显式计算 logits，使用减最大值的 log-softmax 获得稳定的参考
   损失。
2. `fused_ce_loss_kernel()` 为一组 token 遍历词表 V 的分块；每个词表块内部再沿 D
   维分块执行矩阵乘法。
3. 内核在线更新 softmax 的全局最大值和分母，并在目标类别落入当前词表块时提取
   对应分子。
4. 遍历完成后计算目标类别概率及 NLL，将每个 token 的损失写入输出。
5. `fused_ce_loss()` 启动内核，并对 token 损失执行 PyTorch 平均归约。

## 已知限制

- 当前 Triton 内核尚未正确处理所有批次、步幅、尾块掩码和在线缩放细节。
- autotune 配置目前只启用一组很小的块参数，不能代表最终性能。
- 实现仅包含前向损失，没有连接 PyTorch Autograd，也没有反向内核。
- 文件中的测试和基准入口用于继续调试，不应视为通过验证的支持范围。

## 运行入口

```bash
python3 celoss.py
python3 celoss.py --benchmark
```

这些命令需要 CUDA 和 Triton；鉴于当前实现未完成，数值测试可能失败。

## 项目背景与视频

As I've said multiple times throughout these lessons, watching my videos and copying my code line by line does not count as learning, or as having completed this course. Before you can claim that you know how to write GPU kernels, you need to go actually design and write one (or many, preferably many) from first principles. 

For that reason, this is not a "lesson" in the way that the prior 9 were. Rather, it's a demonstration of an example project that you should do to test your knowledge. One alternative way to think about them while you watch is that you, the viewer, are a prospective employer who just asked me a very hard job interview question and I'm now attempting to answer it to demonstrate my skills.

In the first video I start with a goal derived from a vague memory of when I skimmed through Apple's cut cross-entropy loss [paper](https://arxiv.org/abs/2411.09009) months before having ever even written my first GPU kernel. Using that loose starting intuition I work out a plan for what will hopefully be a relatively efficient fused CE Loss kernel from first principles. Then for the second video I attempt to take that plan and put it into action. 

Will it even run? Was there something I wasn't accounting for? Will it be as fast (or hopefully faster than) PyTorch? Does it even resemble Apple's algorithm or does it take a different route entirely? Idk, this was all live off-the-cuff so we'll see.

*the answer is that I could not in the length of the two videos get it working. if i feel like putting more effort into making it actually run (and even more to making it actually fast) then i'll update this readme accordingly and maybe even make a third explanatory video. but for now it doesn't work*


[bilibili](https://www.bilibili.com/video/BV1fMyWBgERM?spm_id_from=333.788.videopod.episodes&vd_source=e98b669ccbafff4b5aa59dd6303b722f&p=9)

BTW, if you're looking for an idea for a project to do, maybe try to not only fix this kernel but fuse the forward & backward pass into one single kernel. What I mean by that is instead of outputting the final loss value, you can have your kernel just skip that step and go straight to outputting $\frac{\partial L}{\partial x}$ and $\frac{\partial L}{\partial E}$. Then, if someone were to use your new kernel, instead of doing the traditional .backward() on the final loss value, they'd actually do it on x and manually accumulate to E using the gradients you gave them. I've not actually done this myself, but I'm vaguely under the impression that this is what [Liger Kernel](https://github.com/linkedin/Liger-Kernel) does from one time when I skimmed a few lines of the wrapper function around their kernel.
