Prhub

#29286 [sgl-kernel/cpu]: exclude amx gemm source from arm build

原始 PR 作者 cyb70289 合并时间 2026-06-26 08:23 文件变更 3 提交数 1 评论 4 代码增减 +18 / -2

sgl-kernel infra arm

执行摘要

ARM 构建排除 AMX GEMM 源码

这是一个小型梳理提交，目的是将 Intel AMX GEMM 源码从 ARM 构建中排除。ARM 平台目前没有 AMX 指令集，这些源码在该平台无用且可能导致编译失败或依赖问题。

此 PR 是低风险、高价值的跨平台构建梳理。建议快速合入。感兴趣的工程师可以关注未来的跟踪 PR，该 PR 可能会真正为 ARM 实现 GEMM 内核以替换存根。

讨论亮点

Review 中 gemini-code-assist[bot] 建议将新存根函数的参数名注释掉以避免编译器警告，作者 cyb70289 已采纳该建议，并将参数名改为 /*A*/、/*mat1*/ 等形式。作者解释这些是“dummy functions, just to satisfy compiler. Will implement if to be used.”。整个讨论非常简短且意见一致。

实现拆解

  1. CMakeLists.txt 构建配置变更（sgl-kernel/csrc/cpu/CMakeLists.txt）：将 gemm_int8.cpp 追加到 SGLANG_CPU_X86_ONLY_SOURCES 变量中。该变量专门用于存放仅 x86 平台需要的源文件，在 ARM 构建时会被跳过。

  2. ARM 存根实现（sgl-kernel/csrc/cpu/aarch64/gemm_int8.cpp）：新增两个存根函数 per_token_quant_int8_cpu 和 int8_scaled_mm_cpu，它们直接抛出 TORCH_CHECK(false, "not implemented yet") 异常。这些存根用于满足 ARM 构建时对符号解析的需求，防止链接错误。

  3. 移除 ARM 条件编译守卫（sgl-kernel/csrc/cpu/gemm_int8.cpp）：删除了 #ifndef __aarch64__ 和 #endif 两行，使得 int8_scaled_mm_with_quant 函数在所有架构下都参与编译。在 ARM 上，该函数会调用存根，但由于编译时会报错提示未实现，实际上不会被正常使用。

文件 模块 状态 重要度
sgl-kernel/csrc/cpu/aarch64/gemm_int8.cpp CPU 内核 modified 5.86
sgl-kernel/csrc/cpu/gemm_int8.cpp CPU 内核 modified 4.54
sgl-kernel/csrc/cpu/CMakeLists.txt 构建配置 modified 1.3

关键符号

per_token_quant_int8_cpu int8_scaled_mm_cpu int8_scaled_mm_with_quant

关键源码片段

sgl-kernel/csrc/cpu/aarch64/gemm_int8.cpp core-logic

为 ARM 平台添加了两个存根函数，确保链接器可以解析符号。

// sgl-kernel/csrc/cpu/aarch64/gemm_int8.cpp (after PR)
​
// 在 anonymous namespace 的 int8_scaled_mm_impl 模板实例化之后
// 新增以下两个存根函数，用于在 ARM 构建中提供符号定义
​
// 存根函数，仅用于满足编译器链接需求，实际调用会抛出异常
std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& /*A*/) {
  TORCH_CHECK(false, "not implemented yet");
  return {at::Tensor(), at::Tensor()};
}
​
// 存根函数，所有参数均注释掉以避免 -Wunused-parameter 警告
at::Tensor int8_scaled_mm_cpu(
    at::Tensor& /*mat1*/,
    at::Tensor& /*mat2*/,
    at::Tensor& /*scales1*/,
    at::Tensor& /*scales2*/,
    const std::optional<at::Tensor>& /*bias*/,
    at::ScalarType /*out_dtype*/,
    bool /*is_vnni*/) {
  TORCH_CHECK(false, "not implemented yet");
  return at::Tensor();
}
​
// 原有的 int8_scaled_mm_with_quant 函数（未改动）
// 该函数在 ARM 上也会编译，但内部会调用上方的存根，运行时异常

评论区精华

存根函数参数命名与编译器警告 style

gemini-code-assist[bot] 建议将存根函数的参数名注释掉，以避免 `-Wunused-parameter` 或 `-Werror` 编译器的警告。

结论：作者已采纳，将参数名改为注释形式 `/*A*/`、`/*mat1*/` 等。 · 已解决

存根用途说明 documentation

作者在回复中说明这些函数是 "Dummy functions, just to satisfy compiler. Will implement if to be used."

结论：明确了存根的目的和后续计划。 · 已解决

风险与影响

风险较低。主要改动涉及构建配置和存根添加，不会影响 x86 平台的正常功能。ARM 平台上的 int8_scaled_mm_with_quant 函数现在编译时会包含完整的 x86 实现（通过条件编译移除），但在 ARM 上运行时，由于底层存根会抛出异常，实际调用会失败。这是一个预期行为——意图就是先保持存根以便未来实现。如果现有 ARM 用户依赖于 int8_scaled_mm_with_quant 的编译版本，可能会遇到运行时错误，但过去该函数在 ARM 上并未编译，所以回归可能性很小。

用户影响：ARM 平台用户不会看到行为变化——之前无法编译的代码现在依然不会运行（抛出异常）。x86 平台用户完全不受影响。
系统影响：构建系统更清晰，ARM 构建排除不必要的 x86 源文件，减少编译时间。
团队影响：简化了跨平台代码管理，使后续为 ARM 实现原生 GEMM 内核更容易——只需替换存根即可。

存根函数使用异常 ARM 运行时可能触发未实现错误

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。

完整报告

Markdown PDF

参与讨论