Stanford CS336

项目结构

stanford-cs336/
├── csrc/                     # CUDA 源代码
│   ├── layernorm.cpp         # PyTorch 绑定
│   ├── layernorm_kernels.cu  # CUDA 内核实现
│   └── reduction_utils.cuh   # CUDA 辅助工具
├── cs336/                    # Python 包
│   └── ops/                  # 操作模块
├── tests/                    # 测试文件
├── setup.py                  # 编译脚本
└── README.md                 # 本文件

编译流程

编译步骤

1. 在项目根目录下运行以下命令编译 RMS Norm 模块：

   python setup.py build_ext --inplace

2. 编译成功后，会在当前目录生成 cs336/ops/layernorm.cpython-*.so 文件，这是编译好的 Python 扩展模块。

测试流程

我们提供了三个测试程序来验证 RMS Norm 实现的正确性和性能：

1. 基本功能测试

这个测试程序验证了 RMS Norm 在简单输入上的输出是否正确。 运行测试：

python tests/test_ops.py