Contribution Tasks

JeffreyXiang/FlexGEMM

候选人的 GPU kernel / 训练性能专长与 FlexGEMM 高度匹配，但当前无 GPU 算力。因此任务聚焦于：(1) Triton autotune cache 的纯逻辑 bug（#26，可在 Mac 上通过导入和 mock 验证）；(2) Roadmap #23 中可拆解的 CPU 可验证子任务（cache schema 校验、all_triton 合并影响分析、接口文档化）；(3) 构建/文档/测试基建类贡献（#33 Windows WHL、#19 构建错误文档化）。避免所有需要真实 GPU 路径的任务（#32、#13、#14、#30、#8）。

当前方向：维护者正从 v1.0 向 v2.0 过渡：dev/all_triton 分支旨在用纯 Triton 替代 CUDA extension 以实现跨平台；同时 main 分支仍在快速迭代（kernels/ 49 次提交）。近期热点是 ROCm 适配（PR #18）、autotune cache 合并（commit 2026-04-22）、CUDA 13 兼容（PR #15）。Issue 响应较慢（多数 0 评论），外部贡献者少。

★ 154Fork 562 个候选任务Gemini：LongCat-2.0

更新于 2026-10-08T07:37:23+00:00 · 打开仓库 ↗

一、项目定位

FlexGEMM 是一个基于 Triton 的 3D 稀疏卷积 GEMM 后端，核心目标是在 Submanifold Convolution 与体素稀疏卷积的训练/推理路径上，用 Triton 重写高性能稀疏 GEMM kernel，从而在 A100 等 GPU 上稳定超越 spconv、torchsparse、fvdb 等现有方案。项目名字里的 'Flex' 来自它对 Explicit / Implicit / Masked Implicit 三类稀疏卷积算法的统一抽象，以及可选的 Split-K 并行策略；README 明确把 'Triton-First Architecture' 作为第一卖点，强调跨平台（CUDA + ROCm）兼容性。当前版本为 v1.1.0（2026-04-02），距离 v1.0.0（2026-01-13）仅三个月即发版，说明作者仍在快速迭代但已进入对外可交付阶段。代码结构上，Python 包位于 flex_gemm/，其中 ops/ 暴露 sparse_submanifold_conv3d 等高层算子，kernels/ 下同时存在 Triton kernel（flex_gemm/kernels/triton/）与手写 CUDA 扩展（flex_gemm/kernels/cuda/，通过 torch.utils.cpp_extension.CUDAExtension 编译），后者承载 hashmap、serialize、grid_sample、neighbor_map 等底层原语。从 commit 热点看，flex_gemm/kernels 以 49 次提交遥遥领先，说明性能核心仍在频繁调整；tests/submconv（7 次）和 setup.py（6 次）次之，反映 ROCm 适配与构建系统是新近重点。整体判断：这是一个处于 '可用但仍在打磨' 阶段的研究型性能库，作者同时维护算法、kernel、构建、跨平台兼容与 autotune cache，单人主导但已有外部 PR 合入（gustavomassa、cuzelac、ZJLi2013）。

解决什么问题、给谁用

解决 3D 稀疏卷积（尤其是 Submanifold Convolution 与常规稀疏卷积）在训练阶段的吞吐瓶颈。现有方案（spconv 2.x/3.x、torchsparse、fvdb）在 FP16/TF32 下仍存在 kernel launch 开销大、Split-K 并行不充分、对不规则稀疏模式适配差等问题。FlexGEMM 用 Triton 重写 GEMM 核心，配合 Explicit / Implicit / Masked Implicit 三种算法变体与可选 Split-K，在 A100 80GB 上对 C=64~1024、RES=8~1024 的体素数据实现约 2× 训练加速，且不增加显存占用。目标用户是做 3D 视觉 / 点云 / 辐射场 / 3D 生成模型（如 TRELLIS.2）的工程师与研究员，典型场景包括：高分辨率体素化点云训练、Submanifold Sparse ConvNet（如 SparseUNet、MinkowskiNet 风格）的前向与反向、以及需要自定义稀疏卷积算子的扩散式 3D 生成管线。

1) 3D 深度学习研究员：需要在 MinkowskiNet / SparseUNet / VoxelNet 等 backbone 上获得更高训练吞吐，或在 TRELLIS.2 这类 3D 生成框架中替换默认稀疏卷积后端。2) 点云感知工程师：在自动驾驶、机器人场景下处理大规模体素化点云，需要 FP16/TF32 下稳定训练。3) 性能库开发者 / 编译器工程师：关注 Triton 在稀疏 GEMM 上的工程实践，想借鉴 Explicit / Implicit / Masked Implicit 算法抽象与 Split-K 实现。4) 跨平台用户：需要在 NVIDIA CUDA 与 AMD ROCm 双栈上获得统一接口（项目已合入 ROCm PR 并处理 __syncwarp / tf32 兼容）。典型使用方式如 README 示例：调用 flex_gemm.ops.spconv.sparse_submanifold_conv3d，通过 set_algorithm() 选择 Algorithm.MASKED_IMPLICIT_GEMM_SPLITK，传入体素坐标、特征与卷积核完成前向与反向。

核心能力

能力	在哪	成熟度

阶段：可用但仍在快速打磨：v1.0.0 于 2026-01-13 发布，v1.1.0 于 2026-04-02 发布，三个月内完成 ROCm 适配、autotune cache 合并、CUDA 13.0 兼容、__syncwarp 修复等；commit 仍在密集落在 kernels 与构建系统上，说明性能与兼容性仍在演进。README 未标 Experimental，但项目规模小、单人主导、算子覆盖面窄，属于 '研究型性能库进入可交付' 阶段。

技术栈：Python ≥ 3.8；PyTorch ≥ 2.4.0；Triton ≥ 3.2.0（Linux/macOS）/ triton-windows ≥ 3.2.0（Windows）；CUDA Toolkit（nvcc，用于编译 CUDAExtension）；ROCm/HIP（可选，通过 BUILD_TARGET=rocm 启用）；C++17（cxx 与 nvcc 均指定 -std=c++17）；OpenMP（-fopenmp / /openmp）；setuptools ≥ 64 + torch.utils.cpp_extension.BuildExtension；filelock（依赖之一）

规模：Release：v1.0.0（2026-01-13）、v1.1.0（2026-04-02），两个大版本间隔约 3 个月。目录与文件：根目录 9 个顶层条目；flex_gemm/ 下 53 个文件（含 kernels、ops、utils 子树）；tests/ 下 15 个文件；examples/ 下 3 个文件；assets/ 下 3 张 benchmark 图。Commit 热点：flex_gemm/kernels 49 次、tests/submconv 7 次、setup.py 6 次、flex_gemm/ops 6 次、tests/spconv 4 次，显示性能核心与跨平台构建是近期主要工作量。最近推送：2026-06-25（pushed_at）。Star 数：证据中未提供，需验证。

二、架构与代码地图

FlexGEMM 的代码组织可划分为五层。最上层是接入层（ops/），对外暴露 sparse_submanifold_conv3d、sparse_conv3d、grid_sample、serialize 等高层 PyTorch 友好算子，内部维护算法选择（Algorithm 枚举）、neighbor_cache 复用与 autotune 配置分发。往下是调度层（ops/spconv/ 与 ops/grid_sample/ 下的 wrapper 文件），负责把输入的稀疏坐标、特征、权重张量规整为 Triton kernel 所需的 indice、pointer、neighbor_map 等索引结构，并决定走 Explicit / Implicit / Masked Implicit 中哪条路径、是否启用 Split-K。第三层是执行层，即 Triton kernel 集合（flex_gemm/kernels/triton/spconv/ 与 grid_sample/），全部用 @triton.jit 写成，承担实际的 GEMM 计算与前后向累加，是性能核心；提交热点 49 次集中在此。第四层是硬件原语层（flex_gemm/kernels/cuda/），通过 torch.utils.cpp_extension.CUDAExtension 编译，提供 hashmap、z_order/hilbert 编码、grid_sample、neighbor_map 等底层 CUDA 例程，被 ops 层与 Triton kernel 共同调用。最底层是工具层（flex_gemm/utils/ 与 autotune_cache.json），承载 Autotuner、构建期 cache 合并、跨平台编译标志等支撑逻辑。整体数据流从上到下依次流经 ops → triton/cuda kernels → GPU，返回时沿原路回传梯度。

FlexGEMM 五层架构：上层 ops 经调度层分派到 Triton 执行层，辅以 CUDA 硬件原语层，工具层贯穿编译与 autotune 配置

模块 / 路径	职责 · 入口 · 依赖
ops.spconv `flex_gemm/ops/spconv/` 约 3 文件 + 多算法 wrapper	稀疏卷积高层算子入口，暴露 sparse_submanifold_conv3d / sparse_conv3d，管理 Algorithm 枚举与 neighbor_cache 复用入口：sparse_submanifold_conv3d, sparse_conv3d, set_algorithm, get_algorithm, Algorithm 依赖：kernels.triton.spconv, kernels.cuda.spconv, ops.utils README 示例直接调用此模块；Algorithm 枚举决定走 Implicit / Masked Implicit / Split-K 哪条路径
ops.grid_sample `flex_gemm/ops/grid_sample/` 约 3 文件	稀疏 grid_sample 高层算子，封装 Triton 与 PyTorch 参考实现入口：grid_sample, grid_sample_torch 依赖：kernels.triton.grid_sample 提供 fwd/bwd 两条路径，含 config.py 做 kernel 配置
ops.serialize `flex_gemm/ops/serialize.py` 1 文件	体素坐标序列化 / 反序列化，提供 z_order 与 hilbert 编码入口入口：encode_seq, decode_seq 依赖：kernels.cuda.serialize 根据空间尺寸自动选 int32/int64 编码宽度
ops.utils `flex_gemm/ops/utils.py` 1 文件	共享工具，如 hashmap 初始化入口：init_hashmap 依赖：torch 按体积极举选 uint32/uint64 hashmap 类型
kernels.triton.spconv `flex_gemm/kernels/triton/spconv/` 约 12 个 py 文件 + config + __init__	Triton 稀疏卷积 GEMM 核心，承载 Implicit / Masked Implicit / Split-K 多算法变体入口：sparse_conv_implicit_gemm, sparse_conv_implicit_gemm_splitk, sparse_conv_masked_implicit_gemm, sparse_conv_masked_implicit_gemm_splitk, sparse_submanifold_conv_fwd_implicit_gemm, sparse_submanifold_conv_fwd_implicit_gemm_splitk, sparse_submanifold_conv_fwd_masked_implicit_gemm, sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk, sparse_submanifold_conv_bwd_implicit_gemm, sparse_submanifold_conv_bwd_implicit_gemm_splitk, sparse_submanifold_conv_bwd_masked_implicit_gemm, sparse_submanifold_conv_bwd_masked_implicit_gemm_splitk 依赖：kernels.triton.utils, ops.spconv config 提交热点 49 次集中在此目录，是性能调优主战场
kernels.triton.grid_sample `flex_gemm/kernels/triton/grid_sample/` 约 3 文件 + config	Triton grid_sample 核心，fwd/bwd 加权求和入口：indice_weighed_sum_fwd, indice_weighed_sum_bwd 依赖：kernels.triton.utils 与 ops.grid_sample 配对
kernels.triton.utils `flex_gemm/kernels/triton/utils.py` 1 文件	Triton kernel 共享工具（指针运算、加载宏等）入口：需验证依赖：triton 被 triton/spconv 与 triton/grid_sample 共用
kernels.cuda.hash `flex_gemm/kernels/cuda/hash/` 约 3 文件（api.h, hash.cu, hash.cuh）	CUDA hashmap 插入 / 查找原语，支持 1D key 与 3D 坐标入口：hashmap_insert, hashmap_lookup, hashmap_insert_3d, hashmap_lookup_3d, hashmap_insert_3d_idx_as_val 依赖：torch::extension 被 ops.utils.init_hashmap 与 spconv 索引构建调用
kernels.cuda.serialize `flex_gemm/kernels/cuda/serialize/` 约 5 文件（api.cu, api.h, z_order.h, hilbert.h, utils.h）	CUDA z_order / hilbert 编解码，含 CPU 回退路径入口：z_order_encode, z_order_decode, hilbert_encode, hilbert_decode 依赖：torch::extension api.cu 内同时分发 cuda:: 与 cpu:: 实现
kernels.cuda.grid_sample `flex_gemm/kernels/cuda/grid_sample/` 约 3 文件（api.h, grid_sample.cu, grid_sample.h）	CUDA grid_sample 参考实现，与 Triton 版并存入口：需验证依赖：torch::extension setup.py 将其编译进 flex_gemm.kernels.cuda
kernels.cuda.spconv `flex_gemm/kernels/cuda/spconv/` 约 6 文件（api.h, subm_neighbor_map.cu/h, sparse_neighbor_map.cu/h, migemm_neighmap_pp.cu/h, utils.h）	CUDA 邻居映射与 miGEMM 预处理原语入口：subm_neighbor_map, sparse_neighbor_map, migemm_neighmap_pp 依赖：torch::extension 为 Triton kernel 构建 indice/pointer 等索引结构
utils.autotuner `flex_gemm/utils/autotuner.py` 1 文件	Triton autotune 封装与 cache 管理入口：需验证依赖：triton, autotune_cache.json setup.py 安装时合并 autotune_cache.json 到 ~/.flex_gemm/
build `setup.py` 1 文件	构建入口，编译 CUDAExtension、合并 autotune cache、处理 CUDA/ROCm/Windows 跨平台标志入口：setup() 依赖：torch.utils.cpp_extension, pyproject.toml 提交 6 次，含 ROCm 适配、CUDA 13 兼容、cache 合并等近期热点

目录树（按文件数）

flex_gemm/ 53 个文件
__init__.py, kernels, ops, utils
tests/ 15 个文件
grid_sample.py, hashmap.py, serialize.py, spconv, submconv, utils.py
assets/ 3 个文件
benchmark_train_fp16.png, benchmark_train_fp32.png, benchmark_train_tf32.png
examples/ 3 个文件
spconv.py, submconv.py, utils.py
.gitignore/ 1 个文件
LICENSE/ 1 个文件
README.md/ 1 个文件
autotune_cache.json/ 1 个文件
pyproject.toml/ 1 个文件
setup.py/ 1 个文件

一次调用怎么流过这些模块

以训练阶段一次 Submanifold Convolution 前向为例。用户调用 flex_gemm.ops.spconv.sparse_submanifold_conv3d(feats, coords, shape, weight, bias)，进入 ops 层 wrapper；wrapper 先通过 set_algorithm 选定的 Algorithm 枚举（如 MASKED_IMPLICIT_GEMM_SPLITK）决定调度路径，同时借助 ops.utils.init_hashmap 或 kernels.cuda.hash 查询 / 构建 neighbor_cache，避免重复生成邻居索引。接着 ops 层把 feats、weight 与坐标索引打包，转成 Triton kernel 所需的 indice、pointer、neighbor_map 等稀疏结构——这一步会调用 kernels.cuda.spconv 的 subm_neighbor_map / sparse_neighbor_map 原语。随后调度层根据算法变体把请求分派到 kernels.triton.spconv 中对应的 @triton.jit 函数（如 sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk），Triton kernel 在 GPU 上执行 GEMM 与累加，并利用 kernels/triton/utils.py 中的共享宏完成指针加载与 Split-K 规约。计算结果 out_feats 沿原路返回 ops 层并透传给调用者，neighbor_cache 被缓存以供下次复用。反向传播时，梯度沿相同链路回传，ops 层调用对应 bwd 版本的 Triton kernel（如 sparse_submanifold_conv_bwd_masked_implicit_gemm_splitk），最终更新 weight.grad 与 feats.grad。整个流程的关键数据结构包括：Algorithm 枚举（调度选择）、neighbor_cache（索引复用）、autotune_cache.json（kernel 配置缓存）、以及 Triton kernel 内部的 indice/pointer 隐式张量。

关键类型与函数

名称	路径	用途
Algorithm	flex_gemm/ops/spconv/	枚举 Explicit / Implicit / Masked Implicit 及 Split-K 变体，控制调度路径
sparse_submanifold_conv3d	flex_gemm/ops/spconv/submanifold_conv3d.py	Submanifold Convolution 高层入口，README 示例直接调用
sparse_conv3d	flex_gemm/ops/spconv/sparse_conv3d.py	常规 3D 稀疏卷积高层入口
set_algorithm / get_algorithm	flex_gemm/ops/spconv/	全局算法选择开关，决定后续走哪条 Triton kernel 路径
sparse_conv_implicit_gemm	flex_gemm/kernels/triton/spconv/sparse_conv_implicit_gemm.py	Triton Implicit GEMM 前向 kernel（常规稀疏卷积）
sparse_conv_implicit_gemm_splitk	flex_gemm/kernels/triton/spconv/sparse_conv_implicit_gemm_splitk.py	Triton Implicit GEMM Split-K 变体，降低大 K 下的累加瓶颈
sparse_conv_masked_implicit_gemm	flex_gemm/kernels/triton/spconv/sparse_conv_masked_implicit_gemm.py	Triton Masked Implicit GEMM 前向 kernel
sparse_conv_masked_implicit_gemm_splitk	flex_gemm/kernels/triton/spconv/sparse_conv_masked_implicit_gemm_splitk.py	Masked Implicit + Split-K 组合变体
sparse_submanifold_conv_fwd_implicit_gemm	flex_gemm/kernels/triton/spconv/sparse_submanifold_conv_fwd_implicit_gemm.py	Submanifold 前向 Implicit GEMM kernel
sparse_submanifold_conv_fwd_implicit_gemm_splitk	flex_gemm/kernels/triton/spconv/sparse_submanifold_conv_fwd_implicit_gemm_splitk.py	Submanifold 前向 Implicit GEMM Split-K 变体
sparse_submanifold_conv_fwd_masked_implicit_gemm	flex_gemm/kernels/triton/spconv/sparse_submanifold_conv_fwd_masked_implicit_gemm.py	Submanifold 前向 Masked Implicit GEMM kernel
sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk	flex_gemm/kernels/triton/spconv/sparse_submanifold_conv_fwd_masked_implicit_gemm_splitk.py	Submanifold 前向 Masked Implicit + Split-K 变体
sparse_submanifold_conv_bwd_implicit_gemm	flex_gemm/kernels/triton/spconv/sparse_submanifold_conv_bwd_implicit_gemm.py	Submanifold 反向 Implicit GEMM kernel
sparse_submanifold_conv_bwd_implicit_gemm_splitk	flex_gemm/kernels/triton/spconv/sparse_submanifold_conv_bwd_implicit_gemm_splitk.py	Submanifold 反向 Implicit GEMM Split-K 变体
sparse_submanifold_conv_bwd_masked_implicit_gemm	flex_gemm/kernels/triton/spconv/sparse_submanifold_conv_bwd_masked_implicit_gemm.py	Submanifold 反向 Masked Implicit GEMM kernel
sparse_submanifold_conv_bwd_masked_implicit_gemm_splitk	flex_gemm/kernels/triton/spconv/sparse_submanifold_conv_bwd_masked_implicit_gemm_splitk.py	Submanifold 反向 Masked Implicit + Split-K 变体
indice_weighed_sum_fwd / bwd	flex_gemm/kernels/triton/grid_sample/	Triton grid_sample 前向 / 反向加权求和 kernel
hashmap_insert_3d / lookup_3d	flex_gemm/kernels/cuda/hash/	CUDA 3D 坐标 hashmap 插入与查找，支撑 neighbor_cache
z_order_encode / decode	flex_gemm/kernels/cuda/serialize/	体素坐标 Z-order 编解码，用于空间排序与序列化
hilbert_encode / decode	flex_gemm/kernels/cuda/serialize/	体素坐标 Hilbert 编解码，替代 Z-order 的局部性优化
subm_neighbor_map / sparse_neighbor_map	flex_gemm/kernels/cuda/spconv/	CUDA 邻居映射原语，为 Triton kernel 构建稀疏索引
migemm_neighmap_pp	flex_gemm/kernels/cuda/spconv/migemm_neighmap_pp.cu	miGEMM 邻居映射预处理，配合 Implicit GEMM 路径
init_hashmap	flex_gemm/ops/utils.py	按体积极举初始化 uint32/uint64 hashmap 容器
encode_seq / decode_seq	flex_gemm/ops/serialize.py	高层序列化入口，自动选择 int32/int64 与 z_order/hilbert
autotuner	flex_gemm/utils/autotuner.py	Triton autotune 封装与 cache 管理

扩展点

ops/spconv/ 与 kernels/triton/spconv/：新增算法变体（如 Explicit GEMM 或新 Split-K 策略）：先在 ops/spconv 的 Algorithm 枚举中加值，再在 kernels/triton/spconv/ 添加对应 fwd/bwd @triton.jit kernel，并在 ops wrapper 中补充分派逻辑
kernels/triton/spconv/ 与 kernels/cuda/spconv/：新增后端（如 ROCm/HIP 专用路径或新 GPU 架构优化）：在 kernels/cuda/spconv 添加对应 .cu 实现，setup.py 已通过 IS_HIP_EXTENSION 与 GPU_ARCHS 暴露跨平台编译开关；Triton 路径天然跨平台，只需在 kernels/triton/spconv 内调整配置
kernels/triton/grid_sample/ 与 ops/grid_sample/：新增 grid_sample 变体（如三线性以外的新插值模式）：在 kernels/triton/grid_sample/ 添加 fwd/bwd kernel，ops/grid_sample/ 暴露新 wrapper
ops/serialize.py 与 kernels/cuda/serialize/：新增序列化模式（如 Morton 以外的空间填充曲线）：在 kernels/cuda/serialize/ 添加新 .h 实现，ops/serialize.py 的 encode_seq/decode_seq 扩展 mode 参数
utils/autotuner.py 与 autotune_cache.json：新增 kernel 配置或调优策略：在 autotuner.py 中扩展搜索空间，新配置通过 autotune_cache.json 持久化与分发

建议阅读顺序

README.md
flex_gemm/ops/spconv/__init__.py 与 submanifold_conv3d.py
flex_gemm/ops/spconv/sparse_conv3d.py
flex_gemm/ops/grid_sample/grid_sample.py
flex_gemm/ops/serialize.py
flex_gemm/ops/utils.py
flex_gemm/kernels/triton/spconv/config.py 与 sparse_submanifold_conv_fwd_implicit_gemm.py
flex_gemm/kernels/triton/spconv/ 下各 fwd/bwd kernel 文件
flex_gemm/kernels/triton/grid_sample/ 与 utils.py
flex_gemm/kernels/cuda/ 下 hash、serialize、spconv、grid_sample 子目录
flex_gemm/utils/autotuner.py 与 setup.py
tests/submconv/ 与 tests/spconv/

三、本地跑起来（没有 GPU 的 Mac）

安装

# 1. 克隆仓库
git clone https://github.com/JeffreyXiang/FlexGEMM.git
cd FlexGEMM
# 2. 创建 Python 3.10+ 虚拟环境（Apple Silicon 推荐 3.10~3.12）
python -m venv .venv
source .venv/bin/activate
# 3. 安装 PyTorch（CPU / MPS 版，Apple Silicon 自带 MPS 后端）
pip install torch>=2.4.0
# 4. 关键：跳过 CUDA 扩展编译
# setup.py 默认会编译 CUDAExtension，在 macOS 上会失败
# 需要临时绕过：注释掉 ext_modules 或设置环境变量（需验证）
# 方案 A：直接修改 setup.py，将 ext_modules=[...] 改为 ext_modules=[]
# 方案 B：尝试通过环境变量跳过（需验证是否支持）
# 5. 安装 flex_gemm（跳过构建隔离以使用本地修改）
pip install . --no-build-isolation
# 6. 验证安装
python -c "import flex_gemm; print(flex_gemm.__version__)"
# 注意：triton>=3.2.0 在 macOS 上无官方支持，pip 会安装但无法执行 kernel
# 如果安装失败，可尝试：pip install triton --no-deps（需验证）

哪些路径能真跑

['# 可在 CPU / MPS 上执行的路径：', '# 1. ops/serialize.py - 坐标编码/解码逻辑（但底层调用 kernels.cuda，需验证是否有 CPU 回退）', '# 从 serialize/api.cu 看，z_order_encode/decode 有 cpu:: 实现，理论上支持 CPU', '# 2. ops/utils.py - init_hashmap 等纯 PyTorch 工具函数', '# 3. flex_gemm/kernels/triton/ - Triton kernel 定义可导入但无法执行（macOS 无 Triton 后端）', '# 4. tests/ 目录 - 测试框架可导入，但所有 GPU 测试会跳过或失败', '', '# 只能读代码或需 Colab T4 验证的路径：', '# 1. flex_gemm/kernels/cuda/ - 所有 CUDA 扩展（hashmap、grid_sample、neighbor_map）', '# 2. flex_gemm/kernels/triton/spconv/ - 所有稀疏卷积 Triton kernel', '# 3. flex_gemm/kernels/triton/grid_sample/ - grid_sample Triton kernel', '# 4. examples/spconv.py, examples/submconv.py - 示例代码需要 CUDA', '# 5. tests/spconv/, tests/submconv/ - 所有前向/反向测试']

最小可运行

# 1. 验证基础导入（CPU 环境）
python -c "from flex_gemm.ops.utils import init_hashmap; print('utils OK')"
# 2. 验证序列化模块（如果 CPU 回退可用）
python -c "
import torch
from flex_gemm.ops.serialize import encode_seq, decode_seq
coords = torch.tensor([[0, 10, 20, 30], [0, 15, 25, 35]], dtype=torch.int32)
shape = torch.Size([1, 1, 64, 64, 64])
codes = encode_seq(coords, shape, mode='z_order')
decoded = decode_seq(codes, shape, mode='z_order')
assert torch.equal(coords, decoded), 'Round-trip failed'
print('serialize OK')
"
# 3. 验证 hashmap 基础操作（如果 CPU 回退可用）
python -c "
import torch
from flex_gemm.ops.utils import init_hashmap
from flex_gemm.kernels.cuda import hashmap_insert_3d, hashmap_lookup_3d
shape = [1, 1, 32, 32, 32]
hashmap_keys, hashmap_vals = init_hashmap(shape, 1024, device='cpu')
coords = torch.tensor([[0, 1, 2, 3]], dtype=torch.int32)
vals = torch.tensor([42], dtype=torch.uint32)
hashmap_insert_3d(hashmap_keys, hashmap_vals, coords, vals, 32, 32, 32)
result = hashmap_lookup_3d(hashmap_keys, hashmap_vals, coords, 32, 32, 32)
print('hashmap OK')
"
# 4. 验证 ops 层导入（不执行 kernel）
python -c "
from flex_gemm.ops.spconv import sparse_submanifold_conv3d, Algorithm
print('spconv ops OK')
"
# 5. 验证 Triton kernel 可导入（不执行）
python -c "
from flex_gemm.kernels.triton.spconv import sparse_submanifold_conv_fwd_implicit_gemm
print('triton kernel import OK')
"
# 注意：以上命令假设 serialize 和 hashmap 有 CPU 回退实现
# 如果失败，说明这些模块也依赖 CUDA，需在 Colab T4 上验证

测试

['# 测试框架：pytest（从 tests/ 目录结构推断，需验证是否有 pytest.ini/setup.cfg）', '# 测试目录：tests/，包含 grid_sample.py, hashmap.py, serialize.py, spconv/, submconv/', '', '# 运行 CPU 子集（跳过 GPU 测试）：', '# 1. 先尝试运行所有测试，观察哪些自动跳过', 'pytest tests/ -v --tb=short 2>&1 | head -100', '', '# 2. 只运行非 GPU 测试（如果测试有标记）', 'pytest tests/ -v -m "not gpu and not cuda" # 需验证是否有 pytest markers', '', '# 3. 单独运行特定模块', 'pytest tests/serialize.py -v # 如果有 CPU 回退', 'pytest tests/hashmap.py -v # 如果有 CPU 回退', '', '# 4. 跳过 spconv/submconv（这些肯定需要 GPU）', 'pytest tests/ --ignore=tests/spconv --ignore=tests/submconv -v', '', '# 预计耗时：CPU 测试 < 1 分钟，GPU 测试需 Colab T4（5-10 分钟）', '# 注意：tests/ 目录下没有 conftest.py（从 tree 推断），可能缺少 fixtures']

调试

# 1. 断点位置建议：
# - flex_gemm/ops/spconv/submanifold_conv3d.py：高层算子入口
# - flex_gemm/kernels/triton/spconv/*.py：Triton kernel 内部（需 GPU）
# - flex_gemm/kernels/cuda/ext.cpp：CUDA 扩展绑定
# 2. 日志开关：
# - Triton 调试：export TRITON_DEBUG=1（需验证版本支持）
# - PyTorch CUDA：export CUDA_LAUNCH_BLOCKING=1（GPU 环境）
# - 详细编译日志：pip install -e . --no-build-isolation -v
# 3. 环境变量：
# - BUILD_TARGET=cuda/rocm/auto：控制 setup.py 编译目标
# - GPU_ARCHS：ROCm 架构指定（如 gfx90a）
# - TRITON_CACHE_DIR：Triton 编译缓存位置
# 4. Profiling 入口：
# - examples/spconv.py, examples/submconv.py：基准测试脚本
# - tests/submconv/train_benchmark.py, inf_benchmark.py：性能基准
# - 使用 torch.profiler 包裹 ops 调用
# 5. Autotune cache：
# - 安装后位于 ~/.flex_gemm/autotune_cache.json
# - 删除可强制重新调优：rm ~/.flex_gemm/autotune_cache.json

CI

['# CI 系统：从仓库结构推断（需验证是否有 .github/workflows/）', '# 证据：tree 中未列出 .github 目录，可能使用外部 CI 或未公开', '', '# 从提交历史看：', '# - PR #18 (rocm)：ROCm 适配，说明有跨平台测试需求', '# - PR #9 (dev/sparse_conv)：稀疏卷积功能开发', '# - PR #15 (fix/syncwarp)：bug 修复，说明有代码审查', '', '# 可能的 CI 检查（从 tests/ 目录推断）：', '# 1. 单元测试：pytest tests/', '# 2. 导入测试：验证 flex_gemm 可导入', '# 3. 示例运行：examples/spconv.py, examples/submconv.py', '', '# PR 可能被卡住的情况：', '# 1. 缺少 GPU 测试环境（如果 CI 无 GPU runner）', '# 2. 代码风格检查（未找到 .flake8/.pylintrc，可能无）', '# 3. 文档更新（README 中的示例代码需同步）', '', '# 建议：提交 PR 前在 Colab T4 上运行完整测试套件']

坑

# 1. macOS 上 CUDA 扩展编译必然失败
# 解决：修改 setup.py 跳过 ext_modules，或等待作者提供 CPU 回退
# 2. Triton 在 macOS 上无官方支持
# pyproject.toml 要求 triton>=3.2.0，但 Apple Silicon 无法执行 Triton kernel
# 解决：仅用于代码阅读，实际执行需 Colab T4
# 3. autotune_cache.json 合并逻辑
# setup.py 安装时会合并 ~/.flex_gemm/autotune_cache.json
# 如果 cache 损坏可能导致导入失败，删除即可重建
# 4. PyTorch 版本兼容性
# pyproject.toml 要求 torch>=2.4.0，但 Apple Silicon 的 PyTorch 版本可能滞后
# 解决：使用 pytorch-nightly 或验证当前版本
# 5. 序列化模块的 CPU/GPU 分支
# serialize/api.cu 有 cpu:: 和 cuda:: 两套实现
# 但 ops/serialize.py 直接调用 kernels.cuda.z_order_encode
# 需验证是否自动回退到 CPU，或需要手动指定 device='cpu'
# 6. 测试依赖外部数据
# tests/spconv_fwd.py 中的 sphere_coords 等函数生成测试数据
# 这些函数可能依赖 CUDA，需检查 tests/utils.py
# 7. 缺少 CONTRIBUTING.md
# 仓库无详细贡献指南，PR 前建议先开 Issue 沟通
# 8. 提交热点集中在 kernels/（49 次）
# 说明性能核心仍在频繁调整，API 可能不稳定

四、维护者与社区

v1.0.0（2026-01-13）→ v1.1.0（2026-04-02），两个 Release 间隔约 79 天，约 2.5 个月一次 minor 发版。最后一次 commit 在 2026-04-22（Merge autotune cache on install），之后 2026-06-25 仓库仍有 push 但无新 Release，说明 v1.1.0 后进入维护/修 bug 阶段，发版节奏放缓。commit 热点集中在 flex_gemm/kernels（49 次），远高于其他模块，说明性能核心仍在频繁调整；tests/submconv（7 次）、setup.py（6 次）、flex_gemm/ops（6 次）次之，反映 ROCm 适配与构建系统是新近重点。

谁	角色	依据
JeffreyXiang（Jianfeng Xiang）	创始人 / 唯一核心维护者，仓库 owner，所有 PR 最终由他合并（Merge pull request #9/#11/#15/#18 均为其账号），README 作者邮箱 belljig@outlook.com 与 pyproject.toml 一致。	仓库 URL JeffreyXiang/FlexGEMM；pyproject.toml authors = [{ name = "Jianfeng Xiang", email = "belljig@outlook.com" }]；所有 Merge commit 的 committer 均为 JeffreyXiang；README 技术博客链接指向 jeffreyxiang.github.io。
ZJLi2013	ROCm 平台贡献者，提交了 ROCm 适配 PR #18 并被合并（2026-04-20），涉及 AMD/HIP 兼容性修复。	Merge pull request #18 from ZJLi2013/rocm；commit 包含 'Disable tf32 input_precision on AMD/ROCm (only ieee supported)' 与 'Fix __syncwarp for ROCm/HIP: AMD wavefront lockstep via wave_barrier'。
cuzelac	kernel bug fix 贡献者，提交了 PR #15 修复 warp-level reduction 中缺失的 __syncwarp()。	Merge pull request #15 from cuzelac/fix/syncwarp-reduce-kernel；commit 'Add missing __syncwarp() in warp-level reduction'（2026-03-08）。
gustavomassa	早期贡献者，提交了 PR #11（2026-02-02 合并），具体改动内容需查看 PR diff。	Merge pull request #11 from gustavomassa/main（2026-02-02）。

流程与 Review 风格

README 的 Contributing 章节给出了最简流程：Fork → 从 main 建分支 → 遵循项目风格 → 跑 tests/ → 开 PR 并附详细描述。无 CONTRIBUTING.md 文件，无 CLA/DCO 要求，无 pre-commit 配置（根目录无 .pre-commit-config.yaml）。pyproject.toml 与 setup.py 共存，说明仍使用 setuptools 而非纯 PEP 621。Issue 模板与 PR 模板均未在根目录或 .github/ 中提供（需验证 .github 是否存在）。无分支保护规则的直接证据，但所有合并均通过 PR 完成（无直接 push main 的 commit）。

从 PR 合并记录看，JeffreyXiang 作为唯一合并者，review 风格偏向「合并后修复」——例如 PR #15 合并后，同一作者又提交了 'Add -allow-unsupported-compiler for CUDA 13.0+ compatibility'（2026-03-08），说明合并时可能未完全覆盖边界情况。Issue 响应速度中等：Issue #8（CUDA 13.0 support）有 2 条评论，后续确实有 commit 修复；Issue #19（Build error at sparse_neighbor_map.cu）有 1 条评论但无后续 commit 直接关联。Issue #33（Windows WHL）、#32（RTX 50-series sm_120）、#30（CPU fallback）均 0 评论，说明维护者对非核心平台问题的响应较慢。整体判断：对性能相关 bug（syncwarp、CUDA 13 兼容）响应快，对平台适配类需求响应慢或依赖外部 PR。

渠道

README 中仅提供一个外部链接：技术博客 https://jeffreyxiang.github.io/en/blogs/flexgem
README 提及的 Real-world Demo 项目：https://github.com/microsoft/TRELLIS.2
Issue 与 PR 是唯一的直接沟通渠道（GitHub 原生）

这里的规矩

得体做法：1. 提 bug 时附上复现脚本与 GPU 型号/驱动版本（从 Issue #32 标题 'device not ready on RTX 50-series (Blackwell, sm_120)' 可推断维护者期望看到硬件细节）；2. PR 从 main 建分支，描述中说明动机与测试结果；3. 性能相关改动附上 benchmark 数据（README 强调 A100 80GB 上的 FP16/TF32/FP32 对比）；4. 跨平台改动（ROCm/Windows）明确标注 BUILD_TARGET 环境变量。
不得体做法：1. 无 GPU 验证就提 kernel 性能 PR（维护者明显以 A100 实测为基准）；2. 在 Issue 中只说 'it doesn't work' 而无复现步骤（参考 Issue #13 'Silent failure in RTX 5080' 0 评论，可能因信息不足被忽略）；3. 直接向 main 分支 push（所有合并均通过 PR）；4. 提与核心性能无关的 refactor 而无 benchmark 证明无回归。

维护者现在最想要的帮助

Issue #23 'FlexGEMM 2.0：改进计划 (Roadmap)' —— 维护者主动开了 roadmap issue，说明在规划下一代，但内容需查看 issue 正文（当前证据仅显示标题）。
Issue #30 'Feature Proposal: CPU fallbacks for CUDA kernels (low VRAM support)' —— 有 PR #29 'WIP: CPU fallbacks for CUDA kernels' 但状态为 WIP，说明 CPU fallback 是维护者认可但尚未完成的方向，对无 GPU 的候选人来说是低风险切入机会。
Issue #32 'CUDA driver error: device not ready on RTX 50-series (Blackwell, sm_120)' —— 新硬件支持需求，维护者尚未响应，可能是未来工作。
Issue #26 '[dev/all_triton branch] Inconsistent autotune cache keys for kernels wrapped with @triton.heuristics' —— 内部 bug，说明 autotune cache 机制仍需打磨。
Issue #33 'Please provide a WHL for Windows' —— 分发需求，PR #31 已修复 Windows/MSVC build，但 WHL 包尚未提供。
从 commit 热点看，flex_gemm/kernels（49 次）仍是焦点，说明性能优化（尤其是 Triton kernel 与 Split-K 策略）仍是核心诉求。
ROCm 适配（PR #18 刚合并 2026-04-20）可能仍有后续工作，如 Issue #32 的 sm_120 支持可类比。

五、切入方案

建议长期负责：ops/serialize + ops/utils + 全项目 CPU fallback 路径 + 测试/CI 基础设施
serialize（z_order/hilbert 编码）和 utils（hashmap 初始化）是唯一有 CPU 实现分支的模块（serialize/api.cu 中 cpu:: 命名空间），且是 ops/spconv 与 ops/grid_sample 的共同依赖。补全 CPU fallback 路径后，你能在 macOS 上跑通完整数据流（encode → conv → decode），无需 GPU 即可验证算法正确性。同时，测试/CI 是项目最大短板（无 CI、无单元测试），贡献可见度高、合并风险低，是通向维护者身份的最快路径。长期可向上延伸到 Triton kernel 的 reference 实现与 autotune cache 管理。

它现在缺什么（你无 GPU 也能补）

缺口	依据	为什么是你
无 CPU/MPS 可运行路径，macOS 用户无法本地验证任何 kernel	setup.py 默认编译 CUDAExtension，在 macOS 上必然失败；runbook 已记录需注释掉 ext_modules 才能安装。kernels/__init__.py 直接 from . import cuda，import 即触发 CUDA 扩展加载。serialize/api.cu 虽有 cpu:: 命名空间分支（z_order_encode/decode），但 ops/serialize.py 直接调用 kernels.cuda.z_order_encode，无设备判断回退。	你擅长端侧部署与跨平台适配，可在无 GPU 条件下通过补 CPU fallback 路径、加设备判断分支、写 macOS 安装文档，直接扩大用户群。
测试覆盖严重不足，无单元测试、无 CI、无回归保护	tests/ 下 15 个文件全部依赖 CUDA（device='cuda'），无 pytest 配置，无 .github/workflows/，无 CONTRIBUTING.md。commit 热点显示 tests/submconv 仅 7 次、tests/spconv 仅 4 次，且全部是 benchmark/ablation 脚本而非单元测试。	可在 CPU 路径上先建立纯 PyTorch 参考实现（如 Triton kernel 的 torch 回写版本），为 serialize、hashmap、grid_sample 补单元测试，搭 GitHub Actions CI（CPU-only），这是无 GPU 条件下最高杠杆工作。
autotune_cache.json 无版本化/无验证机制，安装时静默覆盖风险	setup.py 末尾 _deep_merge 逻辑在 install 时合并 autotune_cache.json 到 ~/.flex_gemm/，commit 'Merge autotune cache on install instead of overwriting'（2026-04-22）说明作者意识到覆盖问题但方案仍粗糙。cache 文件无 schema 校验、无版本号、无 per-GPU 标识。	可设计 cache schema + 校验脚本 + per-device 命名，防止错误 cache 污染，属于构建/工具链层面，无需 GPU。
ROCm 适配刚起步，缺乏系统性测试与文档	PR #18（ZJLi2013）合入于 2026-04-20，仅两次 commit（Disable tf32 on ROCm、Fix __syncwarp），README 未提 ROCm 安装步骤，setup.py 虽支持 BUILD_TARGET=rocm 但无文档。	可补 ROCm 安装文档、在 CI 中加 ROCm 构建检查（若 GitHub Actions 支持）、或至少写已知限制清单。
缺少与 spconv/torchsparse/fvdb 的自动化基准对比脚本	README 展示 benchmark_train_fp16/tf32/fp32.png 三张图，但仓库内无自动化 benchmark 脚本；examples/ 仅有 spconv.py、submconv.py 两个极简示例，无对比框架。	可写一套可选依赖的 benchmark harness（import spconv 才跑对比），在 CPU 模式下跳过，为未来性能回归提供基础设施。
Triton kernel 缺少数值正确性参考实现与误差边界文档	flex_gemm/kernels/triton/spconv/ 下 12 个 kernel 文件（fwd/bwd × implicit/masked_implicit × splitk），每个 kernel 直接 @triton.jit，无 torch 参考实现、无 atol/rtol 说明、无数值精度测试。	可用纯 PyTorch 重写每个 Triton kernel 的数学等价版本作为 reference，配合 tests/ 中的 ablation 脚本做数值对比，无 GPU 也能验证算法正确性。
Windows 支持声明存在但无实际验证路径	pyproject.toml 区分 triton 与 triton-windows，setup.py 有 platform.system() == "Windows" 分支，但无 Windows CI、无 Windows 安装文档、无已知限制说明。	可补 Windows 安装指南、在 CI 加 Windows 构建 job（即使不跑 kernel），降低贡献门槛。
Issue/PR 模板缺失，贡献流程不清晰	仓库根目录无 .github/ISSUE_TEMPLATE/、无 PULL_REQUEST_TEMPLATE.md、无 CONTRIBUTING.md（contributing 字段为空）。外部贡献者（gustavomassa、cuzelac、ZJLi2013）的 PR 合入全靠口头沟通。	可起草 CONTRIBUTING.md、Issue/PR 模板，标准化贡献流程，这是维护者角色的基础工作。

第一批 PR

题目	范围	为什么安全
docs: add macOS (CPU/MPS) installation guide	README.md 新增 'Installing on macOS' 章节 + setup.py 注释说明	纯文档改动，不触及任何代码逻辑；runbook 已验证安装步骤可行，只需文档化
feat(serialize): add CPU fallback for z_order_encode/decode	flex_gemm/ops/serialize.py 加设备判断 + flex_gemm/ops/serialize_cpu.py 纯 PyTorch 实现	serialize/api.cu 已有 cpu:: 命名空间证明作者设计时考虑 CPU 路径；新文件不修改现有 CUDA 代码，失败不影响主路径
test: add CPU-only unit tests for serialize roundtrip	tests/test_serialize_cpu.py 新文件，覆盖 z_order/hilbert encode-decode	新增测试文件，不修改现有代码；可在 CI 中先跑 CPU 测试，GPU 测试保持 skip
ci: add GitHub Actions CPU-only build and test	.github/workflows/ci.yml 新文件 + setup.py 支持 FLEX_GEMM_SKIP_CUDA_EXT=1 环境变量	CI 配置独立于业务代码；环境变量跳过 CUDA 扩展是常见模式（参考 PyTorch 自身 setup.py）
docs: add CONTRIBUTING.md with development setup and PR guidelines	CONTRIBUTING.md 新文件，含 macOS/Linux/Windows 三平台 setup、测试命令、代码风格	纯文档，无代码风险；可参考 PyTorch/spconv 的 CONTRIBUTING 模板

怎么知道自己站住了

PR #1-#5 在 30 天内被合并，且无 revert
JeffreyXiang 在 PR 中主动请求 review 或 assign issue 给你
CPU CI 绿灯通过，且后续 PR 自动触发 CI
你被加入 CODEOWNERS 或 collaborator 列表
其他贡献者开始引用你写的文档或测试模板
autotune cache versioning proposal 被接受并进入 roadmap
你能在 Colab T4 上独立跑通完整 benchmark 并复现 README 性能数据

风险与对策

维护者响应慢或合并瓶颈（单人维护项目）：前 30 天主动在 Issue 中 ping 并展示完整 PR 链；若 2 周无响应，转向 fork + 公开 roadmap 吸引社区
CPU fallback 路径与 Triton kernel 数值精度不一致：每个 CPU 实现必须配 roundtrip 测试 + 与 Triton kernel 的数值对比（在 Colab T4 上跑），atol 文档化
autotune cache 格式变更破坏现有用户：PR 中保持向后兼容（旧 cache 仍可读，新字段可选），并在 CHANGELOG 中明确迁移指南
CI 中 CUDA 扩展编译失败阻塞所有 PR：CI 分两个 job：cpu-only（必过）+ gpu-optional（允许失败），用 GitHub Actions 的 continue-on-error
macOS 上 import flex_gemm 即崩溃（kernels.cuda 加载失败）：在 kernels/__init__.py 中加 try/except 延迟加载 cuda 子模块，import 失败时降级为纯 CPU 模式并警告
Triton kernel 的 PyTorch reference 实现性能差导致测试超时：reference 实现仅用于正确性验证，测试中限制输入规模（N≤1000），并标注 'not for benchmark'
ROCm 文档与实际行为不一致（无 AMD GPU 验证）：文档明确标注 'Community-contributed, not actively maintained by core team'，并链接 PR #18 作为来源
项目进入维护停滞（最后一次 push 2026-06-25 后无活动）：90 天内评估 commit 频率；若停滞，将贡献迁移至 fork 或转向兼容项目（如 fvdb、torchsparse）

六、怎么介入这个项目

无 CONTRIBUTING.md，参考 README 末尾的 'Contributing' 章节：fork → branch from main → 遵循代码风格 → 跑 tests/ → 开 PR
无 CLA/DCO 要求（MIT 许可证）
Issue 评论前可先确认理解是否正确，维护者响应可能较慢（数天到数周）
PR 描述应包含：问题复现、修复方案、测试结果
无 CI，需本地验证后说明测试环境

建议顺序

fix(autotune): 修复 @triton.heuristics 包裹 kernel 的 cache key 命名空间错误（#26）
feat(autotune): 为 autotune_cache.json 添加 schema 校验与版本化（#23 子任务）
docs: 为 #19 Windows 构建错误添加已知问题与修复说明
docs: 提供 Windows WHL 或明确预编译分发策略（#33）
docs: 分析 dev/all_triton 与 main 的接口差异并输出合并检查清单（#23 子任务）
test: 为 autotune cache 和 serialize 添加 CPU 可运行的单元测试

成为长期维护者的路径

维护者 JeffreyXiang 单人主导，commit 频率高但 Issue 响应慢（多数 0 评论）
外部贡献者极少（gustavomassa、cuzelac、ZJLi2013 各有 1 个 PR 合入）
无 CONTRIBUTING.md、无 Issue/PR 模板、无 CI
建议首次贡献选小范围、高确定性的 bug fix 或文档，建立信任后再碰 Roadmap 级任务

七、任务卡

任务 1优先 · medium · Mac · 2-3 个晚上

fix(autotune): 修复 @triton.heuristics 包裹 kernel 的 cache key 命名空间错误

Issue #26 · [dev/all_triton branch] Inconsistent autotune cache keys for kernels wrapped with @triton.heuristics ↗

无人认领0 条评论更新 2026-05-09

用到的专长：Triton 运行时与代码生成、autotune 机制、cache 管理

目标：修复 flex_gemm/utils/autotuner.py 中 _get_function_cache_key 函数，使其正确处理 @triton.heuristics 包裹的 kernel，确保 cache key 始终为 flex_gemm.kernels.triton.spconv.<file>.<kernel_name> 格式。

为什么值得长期做：autotune cache 是 FlexGEMM 性能核心（安装时合并到 ~/.flex_gemm/），cache key 不一致会导致 heuristic-wrapped kernel 的缓存失效或污染，直接影响推理性能。此 bug 涉及代码生成与 Triton 运行时交互，是候选人 kernel 专长的直接应用。

怎么介入：Issue 无 assignee、无 open PR、0 条评论，可安全认领。维护者 2026-05-09 创建后无活动，建议评论确认修复目标分支。
第一个 PR 的边界：仅修改 flex_gemm/utils/autotuner.py 的 _get_function_cache_key 函数 + 新增 tests/test_autotune_cache_key.py。不改动现有 cache 格式或 kernel 代码。

第一步：在 Mac 上克隆仓库，导入 flex_gemm.utils.autotuner，检查 _get_function_cache_key 实现（第 359-362 行），理解当前如何从 fn 提取 module。

本机怎么复现 / 验证：Mac 上：git clone + pip install -e . --no-build-isolation（需注释 setup.py 中 ext_modules 或设置 FLEX_GEMM_SKIP_CUDA_EXT=1），然后 python -c 'from flex_gemm.utils.autotuner import _get_function_cache_key; from triton.runtime.autotuner import Heuristics; ...' 检查 key 生成逻辑。

认领留言（英文，可直接贴到 Issue）

Hi @JeffreyXiang, I'd like to work on fixing the inconsistent autotune cache key issue (#26). The root cause is clear: when a kernel is wrapped with @triton.heuristics, the Heuristics object's __module__ is 'triton.runtime.autotuner' instead of the actual kernel module. I'll fix _get_function_cache_key to unwrap Heuristics and get the original JITFunction's module. Before I start: should this fix target main or dev/all_triton branch? I'll have a PR ready within a few days.

大致实施方案

阅读 flex_gemm/utils/autotuner.py 第 359-362 行及周边代码，理解 _get_function_cache_key 逻辑
检查 triton.runtime.autotuner.Heuristics 类的 __module__ 属性，确认其值为 triton.runtime.autotuner
修改 _get_function_cache_key：当 fn 是 Heuristics 实例时，尝试从 fn.fn 或 fn.inner_fn 获取原始 JITFunction 的 __module__
添加单元测试：mock 一个 @triton.heuristics 包裹的函数，验证 cache key 格式正确
在 dev/all_triton 分支上测试（或 cherry-pick 到 main 验证兼容性）

可能涉及的目录或文件

flex_gemm/utils/autotuner.py（核心修改）
tests/ 下新增 test_autotune_cache_key.py（新增）
可能需要检查 flex_gemm/kernels/triton/spconv/ 下哪些 kernel 使用了 @triton.heuristics

验收方式

Mac 上运行：python -c 'from flex_gemm.utils.autotuner import _get_function_cache_key; ...' 验证 key 格式
新增 tests/test_autotune_cache_key.py（新增），mock Heuristics 对象测试
检查现有 autotune_cache.json 中是否有 triton.runtime.autotuner.* 前缀的条目

开工前问题与风险

向维护者确认

此修复应基于 main 还是 dev/all_triton 分支？Issue 标题标注了 dev/all_triton，但 main 可能也有同样问题
是否有其他 kernel 装饰器组合（如 @triton.heuristics + @triton.autotune）需要同时处理？

风险

Heuristics 对象的内部结构可能随 Triton 版本变化，需做版本兼容
修改 cache key 格式可能导致现有用户的 autotune_cache.json 失效（需文档说明）

任务 2可选 · medium · Mac · 3–4 个晚上

Roadmap #23 子任务：为 Triton autotune cache 设计版本化 schema 与校验

Issue #23 · FlexGEMM 2.0：改进计划 (Roadmap) ↗

无人认领0 条评论更新 2026-05-11

用到的专长：构建系统与工具链设计，熟悉 Triton autotune cache 机制。

目标：为 autotune_cache.json 设计一个带 schema_version、kernel_signature、device_fingerprint 字段的 schema，并写一个校验脚本（Python，纯 PyTorch 即可），在 setup.py 合并前拒绝不兼容条目。

为什么值得长期做：Roadmap 明确把 autotune cache 改进列为 2.0 子任务，当前 cache 无 schema、无版本号、无 per-GPU 标识，安装时静默合并（setup.py _deep_merge）。这是构建/工具链层面工作，无需 GPU，但直接影响所有用户的安装可靠性与维护者的 cache 演进自由度。

怎么介入：Issue #23 是 Roadmap 类型，无 assignee、无 open PR，正文明确列出 autotune cache 改进子任务。先评论确认 schema 方向再动手。
第一个 PR 的边界：第一个 PR 只包含 schema 定义（JSON 示例 + JSON Schema 文件）、校验脚本、setup.py 调用点、3 个单元测试。不改动 autotuner.py 的 key 生成逻辑。

第一步：1. 在 Mac 上 clone 仓库，阅读 flex_gemm/utils/autotuner.py 与 setup.py 末尾的 _deep_merge 逻辑；2. 用 TRITON_INTERPRET=1 python -c 'import flex_gemm' 在 CPU 上跑通 import，观察当前 cache 结构；3. 在 Issue #23 下评论确认维护者是否接受 schema 方案。

本机怎么复现 / 验证：git clone ... && cd FlexGEMM && pip install -e . --no-build-isolation; TRITON_INTERPRET=1 python -c 'import flex_gemm; print(flex_gemm.__file__)'; 查看 ~/.flex_gemm/autotune_cache.json 结构; python scripts/validate_autotune_cache.py（新增后）

认领留言（英文，可直接贴到 Issue）

Hi, I'd like to pick up the autotune cache schema subtask from #23. Plan: add schema_version + per-entry kernel_signature/device_fingerprint fields, write scripts/validate_autotune_cache.py, and call it from setup.py before _deep_merge. I'll keep backward compat (warn + skip on mismatch) unless you prefer hard fail. First PR will be schema + validator only (~200 LOC). Does this direction look OK? Target: ~1 week.

大致实施方案

阅读 autotuner.py 中 cache key 生成逻辑（_get_function_cache_key）与 setup.py 的 _deep_merge
设计 JSON schema 草案：顶层加 schema_version、generated_at；每条 cache entry 加 kernel_qualname、device_fingerprint（可选）、config_hash
写 scripts/validate_autotune_cache.py：读取 ~/.flex_gemm/autotune_cache.json，校验 schema_version 与必要字段，失败时 exit(1) 并打印原因
在 setup.py 的 _deep_merge 前调用校验脚本，不兼容则跳过合并并 warn
新增 tests/test_autotune_cache_schema.py（新增），用临时 JSON 文件跑 pytest，覆盖合法 / 缺失字段 / 版本不匹配三种情况
更新 README 或新增 docs/AUTOTUNE_CACHE.md 说明 schema 含义

可能涉及的目录或文件

flex_gemm/utils/autotuner.py
setup.py
scripts/validate_autotune_cache.py（新增）
tests/test_autotune_cache_schema.py（新增）

验收方式

pytest tests/test_autotune_cache_schema.py 全部通过
手动构造一个 schema_version 不匹配的 cache 文件，运行 setup.py install 应打印 warning 而非崩溃
TRITON_INTERPRET=1 python -c 'import flex_gemm' 在 Mac 上不报 import 错误

开工前问题与风险

向维护者确认

维护者是否接受在 cache 顶层加 schema_version 字段？还是希望保持向后兼容、用 optional 字段？
per-device fingerprint 是否必要（当前 cache 是全局一份）？
校验失败时是 hard fail 还是 warn + skip？

风险

schema 设计过严会导致现有用户 cache 全部失效，需与维护者对齐
setup.py 中调用外部脚本可能影响 pip install 体验，需测试干净环境

JeffreyXiang/FlexGEMM

一、项目定位

解决什么问题、给谁用

同类项目与差别

核心能力

二、架构与代码地图

一次调用怎么流过这些模块

关键类型与函数

扩展点

最近在动的地方

建议阅读顺序

三、本地跑起来（没有 GPU 的 Mac）

安装

哪些路径能真跑

最小可运行

测试

调试

CI

坑

四、维护者与社区

流程与 Review 风格

渠道

这里的规矩

维护者现在最想要的帮助

五、切入方案

它现在缺什么（你无 GPU 也能补）

第 1–30 天：看懂并露面

第 31–60 天：稳定产出

第 61–90 天：接管一块

第一批 PR

怎么知道自己站住了

六、怎么介入这个项目

建议顺序

七、任务卡

fix(autotune): 修复 @triton.heuristics 包裹 kernel 的 cache key 命名空间错误

向维护者确认

风险

Roadmap #23 子任务：为 Triton autotune cache 设计版本化 schema 与校验

向维护者确认

风险