执行摘要:此PR清理了7个在AOT编译中不可达的CUDA头文件,净删2852行,降低代码库维护负担。通过交叉检查编译单元和manifest确保安全,并验证了kernel inventory测试和H200 AOT编译无回归。
功能与动机:PR body明确指出要删除在AOT编译单元中无法到达的头文件。这些头文件已成孤立代码,维护它们不仅浪费资源,还可能误导开发者认为这些代码仍在活跃使用。删除后可以缩小内核代码搜索空间,加快编译。
实现拆解:作者通过以下步骤完成删除:
- 交叉检查:针对所有AOT构建根目录(如common_ops_sm90_build、common_ops_sm100_build)和manifest,逐一验证每个头文件的包含关系,确认7个头文件没有被任何编译单元引用。
- 删除头文件:从
python/sglang/kernels/aot/csrc/ 中删除以下文件:
gemm/marlin/marlin_template.h(主干Marlin kernel模板,1629行)
gemm/marlin/dequant.h(反量化逻辑,504行)
elementwise/pos_enc.cuh(位置编码kernel,467行)
gemm/marlin/marlin.cuh(Marlin辅助,96行)
gemm/marlin/marlin_dtypes.cuh(Marlin数据类型,82行)
cutlass_extensions/gemm/dispatch_policy.hpp(FP8调度策略,38行)
gemm/marlin/kernel.h(Marlin kernel声明,36行)
- 保留JIT副本:确保
python/sglang/kernels/jit/ 下的对应文件不受影响,JIT编译功能继续工作。
- 验证:运行
test/registered/kernels/test_kernel_inventory.py 全部通过(5 passed)。在H200 GPU上构建SM90和SM100的AOT共享库并加载测试,通过rotary和GPTQ-Marlin的JIT smake测试。
关键源码片段:作为示例,下面展示被删除的 dispatch_policy.hpp 的完整内容。该文件定义了FP8 Block Scaling专属的主循环调度策略,但由于AOT构建路径未包含此头,其功能已通过JIT路径实现。
// Adapted from https://github.com/vllm-project/vllm/blob/main/csrc/cutlass_extensions/gemm/dispatch_policy.hpp
#include "cutlass/gemm/dispatch_policy.hpp"
namespace cutlass::gemm {
// FP8 related policies (including Blocked Scaled Accumulation)
// `ScaleGranularityM` specifies scaling granularity along M
template <int ScaleGranularityM = 0>
struct KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum
: KernelTmaWarpSpecializedCooperative {};
// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, Warp
// specialized dynamic schedule for FP8 kernels with Block Scaling
template <int Stages_, class ClusterShape_ = Shape<_1, _1, _1>,
class KernelSchedule = KernelTmaWarpSpecialized,
int ScaleGranularityM = 0>
struct MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8
: MainloopSm90TmaGmmaWarpSpecialized<Stages_, ClusterShape_,
KernelSchedule> {
static_assert(
cute::is_same_v<KernelSchedule,
KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<
ScaleGranularityM>>,
"KernelSchedule must be one of the warp specialized policies");
};
} // namespace cutlass::gemm
该文件被删除是因为它在AOT manifest中不可达。其中定义的 ClusterShape_ 和 KernelSchedule 符号在当前代码库中不再有 AOT 引用。
评论区精华:无实质性的审查评论。
风险与影响:
- 风险:极低。删除的头文件均不可达,且已通过测试验证。
- 影响:对用户零影响。对开发者:减少了代码库体积,但需要确保不会重新引入这些头文件的依赖。当前JIT版本完整保留。
关联脉络:此PR是内核代码库持续清理的一部分。虽然没有直接关联的PR,但与近期内核重构(如PR#32813、#32886、#32887)共同提升了内核模块的维护性。
参与讨论