cache-dit-custom-kernel

# Cache-DiT 自定义算子开发与优化 ## 使用范围与路由 本 skill 处理 `csrc/kernels/` 下**所有**算子族（现有 `comm`、`attn`、`svdq`、`models`，目前全部是 CUDA/CuTe C++ kernel）的接入、移植、构建、profile 与优化。读完本节再动手： | 任务 | 判据 | 读哪里 | |---|---|---| | 接入新算子 / 新算子族 | 要新建 `csrc/kernels/<族>/` 下的文件，或要动 `setup.py` 的源清单/扩展/schema 注册 | `references/kernel-integration-chain.md` | | 接入**模型专属**算子族（`models`） | 算子只服务单个 vendored 模型，要动 `csrc/kernels/models/` 与 `kernels/models/`，或要在模型侧加 `fused_ops.py` 调用点 | 本文 §9 + `references/kernel-integration-chain.md` §2 | | 从上游仓库移植算子 | 有 nunchaku / deepcompressor / flashinfer 参照实现 | `references/operator-migration.md` | | 改现有算子实现 | 只改 `.cu/.cuh/.h`，schema、`setup.py`、Python 包装都不动 | 本文 §5 + 该族 Python 模块 + 对应测试 | | 性能优化 | 数值契约与链路不变 | 本文 §6；SVDQ 另读下列三份 ref | | SVDQuant 原理 | 推导、误差上界、scale 粒度、低秩残差 | `references/svdquant-mathematics.md` | SVDQ 专项参考，修改 NVFP4 CuTe 路径前先读： - `references/svdq-nvfp4-cute.md`：当前实现、有效/无效尝试和实测性能。 - `references/rfc-cuda-optimizations.md`：优化门槛、待办、状态更新规则。 - `references/svdq-nvfp4-cute-tma-sync-analysis.md`：TMA/mbarrier、warp/block 同步、K-block 软件流水及 CUTLASS/FlashInfer SM120 对照；做 SVDQ CuTe GEMM 调度与性能优化前先读。 原理路由：先读 `references/svdquant-mathematics.md` 的 §1–4 建立记号并复核推导与成立条件，需要落到当前实现时读 §5，评估 TurboQuant / DVD-Quant 是否可借鉴时读 §6–8；性能结论与数值契约仍以 `references/svdq-nvfp4-cute.md` 为准。 配合 `cuda-auto-tune`（NCU 驱动迭代优化）、`cuda-cpp-kernel`（CUDA/PTX/架构细节）、`cutlass-cpp-kernel`（CUTLASS/CuTe C++ 模板）、`cute-dsl-kernel`（CuTe DSL Python kernel 编写）使用。本仓目前没有 CuTe DSL 算子；将来引入时，kernel 编写走 `cute-dsl-kernel`，接入、注册与构建仍走本 skill。NCU 优化必须先 profile 后修改，再复采验证。 **算子形态现状与后端边界**：仓库当前**只有纯 CUDA/CuTe C++ kernel**（`csrc/kernels/`），没有 Triton，也没有 CuTe DSL 算子。后续**可能**引入 CuTe DSL（`@cute.kernel`）算子，但必须先补齐 Python 侧注册、扩展构建与打包这条链再落地；**Triton 明确不支持**，不要新增，也不要恢复 `kernels/backend.py` 之类后端选择层。上游只有 Triton 实现时，读算法后用 CUDA 重写。细则与反模式见 `references/operator-migration.md` §1、§9。 ## 1. 分层架构（所有算子族共用） ```text 内部调用点 -> torch.ops.cache_dit_cuda_ops.<op> Python 侧 torch.library：dispatch / fake / compile -> kernels/<族>.py @torch.library.impl(..., "CUDA") 元数据校验 + 建输出 -> _call_<op>() out-param 桥 + 签名归一化 -> _C_ops.ops.<op> pybind 子模块（单扩展，四族共用） -> csrc/kernels/<族>/<族>.cpp -> <kernel>.cu TORCH_CHECK + launch ``` 两条最容易找错地方的机制： - `cache_dit_cuda_ops` 命名空间**只存在于 Python**（`torch.library.define` 在模块 import 时创建）；`csrc/` 里没有任何 `TORCH_LIBRARY`，pybind 只暴露一个 `ops` 子模块，由唯一的 `csrc/kernels/ops.cpp` 建立，各族只提供 `register_<族>_ops(py::module_&)`。两套注册表靠 `_call_*` 桥接。 - 四族现状、extension 名、namespace、arch 门槛、内部调用点与文件名契约（`<族>/<族>.h` + `<族>/<族>.cpp`），只在 `references/kernel-integration-chain.md` §2 维护，不要在这里重抄。 C++ 入口一律 `void` + out-param；schema 默认函数式，所以 out-param 对 torch 不可见、不需要 `mutates_args`。校验以 C++ 的 `TORCH_CHECK` 为权威（`svdq` 是历史例外，两侧都几乎不校验，见 ref §3.3）。 ## 2. 新算子接入 checklist 细节、模板代码与实测陷阱见 `references/kernel-integration-chain.md`；这里只给可逐项打勾的顺序： 1. `csrc/kernels/<族>/<kernel>.cu`：kernel + 入口 `TORCH_CHECK`。 2. 声明头：`void <op>(at::Tensor ...)`，`namespace cache_dit::<族>`。 3. registrar TU `<族>/<族>.cpp`：`void cache_dit::<族>::register_<族>_ops(pybind11::module_& ops)` 里 `ops.def("<op>", ...)`，并在 `csrc/kernels/ops.cpp` 的 `PYBIND11_MODULE` 里加一行调用。**全仓只有一个 `PYBIND11_MODULE`**。 4. `setup.py`：加 `_required_<族>_sources()`、并进 `_required_cuda_sources()`、族目录加进 `_get_cuda_extension()` 的 `include_dirs`。没有 per-family extension 了；arch 谓词只有 `_is_supported_cuda_sm_target`（ref §5.1/§5.3）。 5. `src/cache_dit/kernels/<族>.py` 五段：`from .cuda_ext import _get_required_ops_module` → `_call_<op>` → `torch.library.define` + `impl(CUDA)` + `register_fake` → 公开包装 → `__all__`。**不要再写 per-family loader / availability getter**。 6. `src/cache_dit/kernels/__init__.py`：re-export 公开包装；可用性只有 `cuda_ext_is_available()` / `cuda_ext_get_load_error()`（定义在 `kernels/cuda_ext.py`）；`__all__` 等于全部非下划线**算子**成员（子模块 `attn` / `comm` / `cuda_ext` / `models` / `svdquant` 是包属性但不进 `__all__`）。 7. 内部调用点：可用性门禁用 `cuda_ext_is_available()` 或捕获 `RuntimeError`；错误信息要带构建命令与 arch 门槛（sm89+、排除 sm90）。 8. 测试：`test_ops_dispatch.py` 参数表条目、fake 用例、参考实现对比、`fullgraph=True` compile smoke（模板 `tests/kernels/test_cuda_fp8_comm_ops.py`）。 三条硬规则：schema 形参顺序即调用顺序（写反不报错、静默算错）；fake 的 shape/dtype 必须与 impl 一致（否则 eager 绿、compile 炸）；`torch.library.define`/`impl` 不幂等，不要 reload 这些模块。 ## 3. SVDQ 专项：架构与命名 SVDQ 调用链： ```text SVDQW4A4Linear.quantize / forward_quant -> src/cache_dit/kernels/svdquant.py -> torch.library schema + fake/CUDA impl -> svdq.cpp -> svdq.h（registrar + dispatcher 实现） -> svdq::native_kernels 或 svdq::cute_kernels ``` 目录与命名契约： - `csrc/kernels/svdq/native/`：唯一 native 实现，使用 `svdq::native_kernels`。**已扁平化，没有 `core/`、`int4/`、`fp4/` 子目录**： - `gemm_w4a4.{cu,cuh}`：W4A4 主 GEMM。INT4 与 NVFP4 共用同一条 host 分派与 epilogue 链，device 侧按 `USE_FP4` 在 `gemm_w4a4_kernel`（INT4）与 `gemm_w4a4_fp4_kernel`（NVFP4）之间选择。 - `gemm_w4a4_impl.cuh` / `gemm_w4a4_launch.cuh`：launch 接线（`GEMM_W4A4_Launch<Config, USE_FP4>`）与显式实例化助手。 - `gemm_w4a4_int4.cu` / `gemm_w4a4_fp4.cu`：唯一两个负责把 kernel 实例化进 `.so` 的 TU。 - `gemm_base.cuh`、`mma.cuh`、`lora.cuh`、`epilogues.cuh`、`gemm_utils.cuh`、`utils.cuh`、`dispatch.h`：共享构件；`epilogues.cuh` 现在只剩 `EpilogueGelu`（其余 `EpilogueQKVProj` / `RMSNormRope` / `PackQKV` / `LiteLA` 已随无调用者的 native 分支删除）。 - `zgemm.h`：对外只声明 `gemm_w4a4` 与 `quantize_w4a4`。 - `csrc/kernels/svdq/cute/fp4/gemm/`：`gemm_w4a4.{cu,cuh,h}`，NVFP4 主 GEMM、BF16 LoRA-up 融合和 TMA pipeline。 - `csrc/kernels/svdq/cute/fp4/quantize/`：`quantize_uni.cu` 是通用 producer；`quantize_impl.cuh` 是 rank 特化模板实现；`quantize_r32/64/128.cu` 只实例化各自 rank，便于并行编译。公共入口声明在 `quantize.h`。 - cute 侧文件与 C++ kernel/host/launch 名不重复携带 `_cute`。`svdq.h`/`svdq.cpp` 只保留 `gemm_w4a4`、`quantize_w4a4` 两个 dispatcher（**声明在头、定义在 `.cpp`**，头里的非 inline 定义会让 `ops.cpp` 与 `svdq.cpp` 链接期 `multiple definition`），由 `runtime.h` 的 `RuntimeParams`（runtime/fp4/alpha/act_unsigned）与 `LoraLayout` 分派到 native/cute namespace；dispatcher 命名空间是 `cache_dit::svdq`，而 kernel 命名空间仍是全局 `svdq::native_kernels` / `svdq::cute_kernels`（所以 `svdq.cpp` 里必须写 `::svdq::cute_kernels::...`，否则会解析到 `cache_dit::svdq` 自身）。namespace 叶子名也统一为 `gemm_w4a4` / `quantize_w4a4`，native 侧只有这一份 ABI（无 ext 变体）。Python 侧同样只有 `svdq_gemm_w4a4` 与 `svdq_quantize_w4a4` 两个公共入口，没有私有宽 ABI 旁路。 - `csrc/kernels/svdq/runtime.h`：`RuntimeKernel` / `RuntimeParams` 与 `LoraLayout` 的**唯一**定义处。原 `runtime_params.h` 已删除；`LoraLayout` 原先定义在 cute 的 `gemm_w4a4.h`，现由 `runtime.h` 提供、cute 头 include 它。只有 host 侧依赖这个头，device kernel 体不依赖。 - 结构体名区分实现：`W4A4QuantizeUniversalKernel` 只在 `quantize_uni.cu`，`W4A4QuantizeRankKernel` 只在 `quantize_impl.cuh`。 - `common.h`、`Tensor.h`、`torch.h` 的共享类型仍在全局命名空间；host 侧用 `::Tensor` 避免与 `cute::Tensor` 歧义。 公开 `runtime_kernel` 只接受 `None/native/cute`。INT4 与 NVFP4 都走同一个 native `gemm_w4a4`，差别只在 packed 权重/scale 格式（`fp4` 开关）；显式 `cute` 时 Python 默认 `lora_layout="row_major"`；显式 native 与 `fused_mlp` 保留 FP32 fragment-D 契约。GEMM 尚未达到相对 native 的 5% 验收目标，见 §8.1。 SVDQ 在通用模板之外的偏离（C++ 侧几乎无 `TORCH_CHECK`、形状/通道校验在 Python；`RuntimeParams`/`LoraLayout` 仍属全局 `svdq`）见 `references/kernel-integration-chain.md` §2 与 §4.2，照抄前先看。 ## 4. SVDQ 专项：独立 NVFP4 CuTe 工作流 保持以下数据流，不得重新引入旧工作流适配层： ```text quantize_r32/64/128（或 rank32/K<=3072 的 quantize_uni） -> native packed qact/ascales + row-major BF16 D[M,R] -> gemm_w4a4_kernel -> A/SFA G2R + B/SFB/D/L1 TMA + NVFP4 MMA + BF16 LoRA-up MMA ``` CuTe producer 的两条合法 D 路径： - fused：CTA 内完成 FP32 D 归约，无 global scratch/finalize；当前 rank32 M>8192 使用。 - split-K：CTA 内 rank streaming，D 写自己的 fragment-major FP32 scratch，kernel 后用向量化 finalize 转 BF16；rank64/128 和较小 M 的 rank32 使用。这里的 scratch 是 CuTe producer 内部实现，不是允许读取 native producer 的 FP32 fragment D。 禁止在 CuTe forward 中调用 native activation quantizer、读取 native FP32 LoRA scratch、执行旧布局反解或增加 `lora_act_to_bf16`。旧 checkpoint 只能在加载时一次性转换 `qweight/wscales/proj_up/proj_down/smooth_factor`；state dict 保持原 native pack 格式。 当前 producer 已迁移到 CuTe `TiledMMA/TiledCopy` + swizzled shared staging，并按 rank/K/M 选择 fused、K128 async split-K 或 K256 split-K。修改时必须保留 native qact/ascales 布局、D 的 FP32 累加/单次 BF16 rounding、K64 tail 和 rank 特化入口。 ## 5. 可重复构建 先初始化依赖并激活目标 Python 环境： ```bash git submodule update --init --recursive CACHE_DIT_BUILD_CUDA=1 \ CACHE_DIT_CUDA_ARCH_LIST=120a \ bash tools/build_fast.sh install -e . --no-build-isolation --no-deps ``` `csrc/third_party/cutlass` 是 NVIDIA 官方 v4.6.1 submodule，默认构建直接使用它。`CACHE_DIT_CUTLASS_DIR` 只用于显式开发覆盖；不要假设机器上存在同级 CUTLASS checkout。 `tools/build_fast.sh` 默认 `MAX_JOBS=32`，通过 `tools/nvcc` 将 nvcc 交给 ccache。当前环境已有 ccache；用 `bash tools/build_fast.sh` 调用，因为脚本没有可执行位。需要强制缓存能力时设 `CACHE_DIT_REQUIRE_CCACHE=1`。仅在确有必要时设 `CACHE_DIT_CLEAN=1`，它会删除构建目录。 每次 kernel 迭代使用同一入口。不要因为源码改动触发全量重编就放弃缓存；记录构建是否 cold、是否换过 compiler/shim，避免把首次 ccache 填充时间误报为加速收益。 **头文件依赖**：SVDQ 扩展现在通过 `CUDAExtension(depends=[csrc/kernels/svdq/**/*.{h,cuh}])` 跟踪 SVDQ 头文件；修改 `quantize_impl.cuh` / `gemm_w4a4.cuh` 应触发扩展重建。改该目录之外的共享头，或用旧 wheel/缓存构建时，仍需显式 `CACHE_DIT_CLEAN=1`。新增算子族要同步的 `setup.py` 位置见 `references/kernel-integration-chain.md` §5.1。 ## 6. 优化工作流 1. 明确验收对象：producer、GEMM、完整 layer 或 e2e，不能混用。 2. 固定 shape、dtype、GPU、warmup、迭代数和输入；profile 前检查 GPU 占用。 3. 先取 nsys kernel-only 和 NCU full 基线。nsys 用于性能判定，NCU 用于归因。 4. 从 roofline、occupancy、寄存器/local memory、global sectors、shared bank conflict、PC sampling 和 warp stall 中确定一个主瓶颈。 5. 一次只做一个有假设支撑的改动；保持数值契约不变。 6. 重建后先跑数值回归，再用同口径 nsys 复测；NCU 复采只解释指标变化。 7. 将有效和无效尝试都写入性能参考；无效方案给出证据，不要反复尝试。 8. 更新 RFC 状态，只有完成独立 kernel 与 e2e 验收才可标完成。 ### 测量口径 - **nsys kernel-only median/mean**：kernel 性能主指标；报告实例数、min/max 和稳态 warmup。 - **CUDA Event op-level**：包含分配、清零和 launch 的 Python 算子调用成本，不能替代 kernel-only。 - **NCU duration**：受 replay、时钟和 profiler 影响，只用于优化前后归因，不作为交付性能。 - **FLUX e2e**：固定模型、prompt、seed、分辨率、steps 和 `--compile`，至少多次重复；同时做图像检查和 PSNR。 不要把冷启动首轮离群值当成稳定收益，也不要用单 GEMM 收益推断 e2e 收益。 ## 7. 验证与回归 通用测试矩阵（含每个用例的真实强度与局限：dispatch 参数表是手工维护的、source manifest 只覆盖 svdq、`tests/kernels/` 不进 CI）见 `references/kernel-integration-chain.md` §7。 SVDQ 回归命令： ```bash python -m pytest -q \ tests/kernels/test_ops_dispatch.py \ tests/kernels/test_svdquant_cute.py \ tests/kernels/test_svdquant_runtime.py \ tests/kernels/test_svdquant_source_manifest.py \ tests/quantization/test_quantize_config.py ``` `test_svdquant_source_manifest.py` 是 CPU-only 护栏（不 import torch、扩展或 setup.py）：它用 `ast` 抽取并执行 `setup.py` 的源清单函数，双向比对 `csrc/kernels/svdq/` 下的 `.cu/.cpp`。改动该目录的文件名或源清单后必须跑；CI 不收集它（workflows 不许改），只能本地执行。 CUDA 数值改动还需运行 producer、独立 GEMM、高层 Linear 和 INT4 隔离回归。改名时同时检查 Python 私有 pybind 调用、公开 torch.library 名称、fake 实现、CLI、文档和临时 harness。运行 `pre-commit run --files ...`、`git diff --check` 和 `read_lints`。 FLUX 验收使用同一 seed 对比默认 native、显式 cute 和 BF16 baseline。精度结论需说明量化差异；不能仅凭“正常出图”宣称无损。 ## 8. SVDQ 当前性能结论与边界 ### 8.1 最新状态（2026-11-11，RTX PRO 5000） - `main` 已合入 PR #1284（mainloop 索引简化）；`dev` 在其上包含 P2 dispatch commit `3ad684c6`。P0 8M×1N、P1 S4/独立 LoRA ring 均已实测失败并回退，证据见 `references/svdq-nvfp4-cute.md` §30。 - 官方 `bench/kernels/bench_svdquant.py`，5 shape × rank32/64/128：本次会话复测 producer geomean **1.1344×**（15/15，1.0784×–1.2129×）；这些格均不进入小 M fused 新分支，属于无关回归检查，不能把相对历史会话的增量归因于本项改动。既有三轮中位的 row-major kernel-only GEMM **1.0248×**（12/15）、row-major e2e **1.0553×**（14/15）、cute+native-D e2e **1.0462×**（14/15）仍是独立历史口径，不与本次 producer 数字拼接收益。 - 当前 row-major dispatch：`M>=16384 || (M>=8192 && N>=6144)` 走 main-first，其余 LoRA-first；受控 7 轮交错 A/B 中 main-first 只在 6/15 格胜出。native fragment D 始终 main-first。 - GEMM NCU（15 格当前 dispatch）：Tensor **84.2%**、SM 80.2%、DRAM 12.9%、L2 hit 99.9%、occupancy 17.75%、regs 168、smem 85 KiB、local load/store=0。stall：math_pipe_throttle 36.2% + wait 30.8%，long_scoreboard 18.4%，sleeping 7.5%。 - GEMM 已接近当前架构局部最优。A/SFA 全 TMA+SMEM 是最初实现，已被 hybrid G2R 击败（同 shape nsys 约 137.4us vs 117.2us）；不要再把 A/SFA TMA 作为优化方向。8M×1N 慢 20.6%，S4/独立 LoRA ring 慢约 1%，均已证伪。 - producer 仅对 **rank64、M≤512、K≤3072** 使用 fused，对 **rank128、M≤256、K≤3072** 使用 K128 async split-K，其余保持原 dispatch。rank64 NCU 20.58→18.18us，真实 FLUX 形状加权 producer 52.92→51.01ms/step；rank128/M256 NCU 24.80→12.77us，官方 Event 36.74→29.98us，M512 保持 K256。M=1 仍慢于 native，不能宣称小 M 已全面追平。 - split-K bank-conflict 路线已证伪：B CACHEGLOBAL 将 LDGSTS 冲突 282.9K→0，但 NCU 57.25→57.66us；64-bit copy 将冲突升至 603.8K 且更慢。不要再以冲突计数下降作为保留改动的理由。`quantize_uni.cu` 保持现状；rank32/16K fused 已到 DRAM 90.2% 带宽墙。 - 结论适用 PRO 5000；未在 PRO 6000/H800 复测前不要外推绝对阈值。split-K D 使用 relaxed FP32 atomics，不保证逐位可复现，D 用误差门槛。 - 后续优先级、精确 NCU 数值和禁止重复清单见 `references/svdq-nvfp4-cute.md` §31 与 `references/rfc-cuda-optimizations.md` §6。 ### 8.2 历史基线（截至 2026-11-05，M=4096/N=K=3072/rank=32） 以下内容保留用于追溯早期 hybrid/TMA 迭代，不代表 2026-11-10 的当前实现： - 当前 CuTe GEMM 已恢复为 A/SFA G2R + B/SFB/D/L1 TMA，当前 CuTe producer 直接输出 native A/SFA 与 BF16 D，无 finalize。 - 当前统一 CUDA Event 基线见 `references/svdq-nvfp4-cute.md` §17（GPU4，Event50×3轮）。M4096 rank32 kernel-only 120.22us vs native120.24us；完整层 M1024/M4096 分别 62.99/164.62us vs native63.65/168.34us。 - **rank32 大 shape 仍落后 native**：`8192x3072`/`4096x6144`/`8192x6144`/`16Kx8Kx8K` 的 cute/native 为1.026/1.020/1.047/1.064。rank128 在所有大 shape 略优于 native（0.32%–1.43%）；rank64 在 `8192x6144`/16K 退化。下一条主矛盾是 rank32/64 大 shape。 - hybrid 的机制是拆掉 A/SFA smem 中转：相对全 TMA，shared 指令/wavefront/bank conflict 降22%/18%/29%，`math_pipe_throttle` 6.00→4.61，静态 FENCE 9→2；代价是 A 重复 G2R load，global 请求升4倍、L1 hit 94.5→55.1%，但 L2 命中87%，总量仍远低于 native。 - 当前 hybrid 重测逐 k-block `A load -> B/SFB copy -> MMA`：168 regs、零 spill；12格净变化约0，2K rank32 +0.61%，8K/16K 小优0.2%–0.5%，在1%门槛内，已保留。旧“0.12%不保留”是旧 hybrid/nsys 口径。 - stages 仍只有 S3：S2 被 output overlay 64KiB > pipeline 55.3KiB 拒绝，S4 110.6KiB > 99KiB；G2R 只释放未使用的 SFA 3KiB，A/B 区仍被 D/L1 overlay 占用。要加 stages 必须先解耦 overlay。 - 剩余 0.737M bank conflict 已定位为 **LoRA BF16 D/L1 的标量 LDS**，不是 SFB。canonical `tile_to_shape(SW64 atom,(BM,rankTile,stages))` + LDSM 已将它降到 **483**，零 spill；M4096 约快0.45%、16K 约快0.42%，已保留。BK64 fallback 必须用 SW32 atom；不要重做 conflict 消除。 - FlashInfer CuTeDSL 四 shape 比 native 慢9%–15%，不移植。A-only TMA 与 native pack 不兼容（rel_l2=1.30）；producer coalesced 写回无收益（46.81→47.01us）。 - 16K NCU 排除 wave 尾部和 L2（两边37.24 waves、cute L2 95.3%）；剩余差距在 shared 装载（16.4×指令）和 TMA/mbarrier sleeping/barrier。删除3条 load-side `fence_view_async_shared` 仅±0.2%，已保留为等价简化，但没有缩小大 shape 差距。 - 三条新主循环方案均失败并回退：A/SFA前移±0.4%噪声；A/SFA双缓冲SASS过闸（R221/224、零spill）但全矩阵慢2.4%–6.1%；BK64 S5/S6/S7在16K慢26%–27%。不要再重复这三条。