[Feat] add support sm90 fp8 mega moe - #422
lengrongfu wants to merge 3 commits into
Conversation
| DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous()); | ||
| DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); | ||
| DG_HOST_ASSERT(intermediate_hidden / 64 <= 64); |
There was a problem hiding this comment.
🔴 critical: 拒绝内核实际不支持的 hidden 对齐: 这里仅要求 hidden % 128 == 0，但例如 hidden=384 且平均每专家至少 64 个 token 时会选择 block_n=256，随后 MegaMoESM90Scheduler 的 L2_SHAPE_N % BLOCK_N == 0 静态断言使 JIT 编译失败；combine 分块也要求更严格的对齐。应在进入 JIT 前校验所有实际模板约束，或调整分块策略以支持这些当前被接受的形状。
🤖 v6
| constexpr int kGranMN = 128, kGranK = 128; | ||
| check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK, | ||
| num_experts_per_rank, false, true, torch::kFloat); | ||
| check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK, | ||
| num_experts_per_rank, false, true, torch::kFloat); |
There was a problem hiding this comment.
🔴 critical: 不要接受内核无法解释的转置权重缩放因子: 这里启用的 sm90_sfb_check 同时接受连续布局和最后两维转置后的布局，但新内核通过 base + n_block * k_blocks + k_block 直接读取缩放因子，完全忽略 tensor stride。调用者传入被该检查接受的转置 SF 时会静默使用错误的缩放值并产生错误结果；应只接受连续的 [E, N/128, K/128] 布局，或将 stride 传入内核并据此寻址。
🤖 v6
| int num_max_padded_sf_pool_tokens = 0; | ||
| for (int block_m: layout::kCandidateBlockM) { | ||
| num_max_padded_sf_pool_tokens = std::max( | ||
| num_max_padded_sf_pool_tokens, | ||
| layout::get_num_padded_sf_pool_tokens(num_max_pool_tokens, block_m) | ||
| ); |
There was a problem hiding this comment.
🟡 warning: 仅按 SM90 实际候选块大小分配 SF 池: 该循环遍历共享候选集合中的 block_m=8，因此把 SF 池扩展到 16 * num_max_pool_tokens；然而本次新增的 SM90 heuristic 只会选择 64 或 128，实际最大需求是 2 * num_max_pool_tokens。两个 SF 池因此被放大约 8 倍，在多卡、大 token/hidden 配置下会额外占用数 GB 对称显存并可能直接 OOM。
🤖 v6
| uint32_t kNumL2BlockNs = L2_SHAPE_N / BLOCK_N, | ||
| uint32_t kNumL1BlockKs = L1_SHAPE_K / BLOCK_K, | ||
| uint32_t kNumL2BlockKs = L2_SHAPE_K / BLOCK_K> | ||
| struct MegaMoESM90Scheduler { |
There was a problem hiding this comment.
🔵 suggestion: 与对端调度器相比移除了 3 条 2-CTA 相关的 DG_STATIC_ASSERT（kNumSMs / kNumL1BlockNs / kNumL2BlockNs 为偶数）。SM90 路径固定 cluster=1，放宽合理，但建议在注释中说明该放宽依赖 cluster_size==1 的前提，避免未来引入 2-CTA 配置时缺少编译期校验。
🤖 v5
| @@ -18,15 +18,21 @@ CUTLASS_DEVICE void cluster_sync_with_relaxed_arrive() { | |||
| cute::cluster_wait(); | |||
There was a problem hiding this comment.
🔵 suggestion: NVLink barrier 超时在对端 PR #36 中为 180s，本仓复用了共享的 kNumTimeoutCycles=60s（基线已有，非本 MR 引入）。若期望与对端调试行为对齐，可考虑为 NVLink 路径单独保留 180s 或加注释说明该差异。
🤖 v5
| DG_HOST_ASSERT(intermediate_hidden / 64 <= 64); | ||
|
|
||
| constexpr int kGranMN = 128, kGranK = 128; | ||
| check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK, |
There was a problem hiding this comment.
🟡 warning: check_sf_layout(..., false, true, torch::kFloat) 的 sm90_sfb_check 分支会同时接受 K-major 连续（stride(-1)==1, stride(-2)==size(-1)）和 MN-major 连续（stride(-1)==size(-2), stride(-2)==1）两种权重 SF 布局；但 sm90_fp8_mega_moe.cuh 中 l1/l2 权重 SF 是按 K-major 连续指针索引的（l1_weights_sf + expert * kL1SFPerExpert + k_block_idx + gate_n * kL1SFKBlocks）。传入被接受的 MN-major SF 时会静默读取错误的 scale，导致结果错误。建议这里收紧为只接受 K-major 连续布局，或让 kernel 同时处理两种 stride。
🤖 v4p
🤖 ds-review-bot Code Reviewv6新增路径会对部分已通过接口校验的形状在 JIT 阶段失败，并会接受导致静默错误结果的 SF 布局。此外，SF 工作区存在显著过度分配风险。 v5本 MR 将 sgl-project/DeepGEMM PR #36 的 SM90 (Hopper) FP8 MegaMoE 移植到本仓库。经与对端 MR 逐文件核对，所有跨仓契约均一致：(1) 对称缓冲区布局（input token / SF(hidden/32B) / topk_idx / topk_weights / L1 pool / L1 SF / l1_topk_weights / L2 pool / L2 SF(intermediate/16B, per-64 K float) / combine buffer 的排布及 SF 视图转置 stride {1, num_max_padded_sf_pool_tokens}）与对端 csrc/apis/sm90_mega.hpp 完全相同；(2) host API 校验（recipe (128,128,128)、intermediate_hidden/64<=64、(128,128) float 权重 SF 布局）及传给 kernel 的参数顺序一致，fp8_mega_moe_sm90 与对端 fp8_mega_moe 仅函数名不同，属各自绑定层内部命名，非跨仓契约问题；(3) layout::MegaMoESM90Workspace 与对端 layout::Workspace 字节布局逐字段一致，kCandidateBlockM {8,16,32,64,96,128,192}、kLCM=384、get_num_max_pool_tokens、get_num_padded_sf_pool_tokens 均相同；(4) kernel 与调度器相对对端仅为 MegaMoEScheduler→MegaMoESM90Scheduler、BlockPhase→MegaMoESM90BlockPhase 的机械重命名及注释差异，无逻辑分歧；(5) 权重预处理（仅 L1 权重 gran=8 gate/up 交织，L1/L2 SF 透传，kernel 内以 gate_n=sf_n_block_idx/2 处理非交织 SF）与对端等价；(6) token 对齐均为 kLCMCandidateBlockM=384；(7) barrier 模板化 WorkspaceT + SM90 DG_TRAP_ONLY_DEVICE_ASSERT 分支与对端 barrier 改动语义一致，顶层 commit 已修复依赖类型的 .template get_grid_sync_count_ptr 调用。另有两处可接受的非契约差异：Sm90SymmBuffer 额外支持 group.size()==1 走 torch.empty + SimpleNamespace（功能增强）；对端将共享的 get_num_experts_per_wave_for_mega_moe 私有化移植进 heuristics/sm90_mega_moe.hpp 且实现逐行一致。未发现接口/协议/数据格式/调用约定不一致，移植保真度高，建议合入。 v4p本 MR 将 sgl-project/DeepGEMM PR #36 的 SM90 (Hopper) FP8 MegaMoE 路径移植到当前仓库：新增 SM90 专用 workspace/scheduler、JIT 运行时与 kernel 实现，以及配套的 Python API（Sm90SymmBuffer、fp8_mega_moe_sm90 等）和 Hopper 测试。整体实现与原 PR 基本一致，结构上通过独立的 MegaMoESM90Workspace/Scheduler 避免了与现有 SM100 ring-buffer 布局冲突；主要问题是部分边界配置（>64 rank、MN-major 权重 SF）缺少校验或支持，可能导致编译失败或静默错误。 Files reviewed: 9 📍 未定位到 diff 的评论🟡 warning |
Ports upstream PR sgl-project#36 ("Sm90 mega moe on sgl dev"), which added a Hopper FP8xFP8 fused MoE GEMM kernel on the old `dev` branch layout (tvm_ffi_api.cpp / sgl_deep_gemm), onto nv_dev's current layout (csrc/apis/*.hpp, csrc/python_api.cpp, deep_gemm/mega). nv_dev already had MegaMoE, but only for SM100 (Blackwell), built on a ring-buffer / cluster-pair-interleaved persistent-grid scheduler and a UE8M0/FP4-oriented workspace layout. The SM90 kernel in PR deepseek-ai#36 instead uses a simpler pool-based workspace and per-expert-wave scheduler with float (non-UE8M0) scale factors, FP8-only weights (no FP4), and no 2-CTA clusters or shared-expert support -- these are not compatible data layouts, so the SM90 path is added as a fully parallel path alongside (not replacing) the SM100 one. New files (ported from the PR's `.cuh`/`.hpp`, adapted to nv_dev's current naming/dispatch conventions and symbol-renamed to avoid colliding with SM100's `layout::Workspace` / `sched::MegaMoEScheduler` / `sched::BlockPhase`): - deep_gemm/include/deep_gemm/layout/sm90_mega_moe.cuh (`layout::MegaMoESM90Workspace`, pool-based; reuses the existing shared `layout::Data`/`layout::Buffer`/`layout::TokenSrcMetadata` and `layout::get_num_max_pool_tokens` from layout/mega_moe.cuh) - deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh (`sched::MegaMoESM90Scheduler` / `sched::MegaMoESM90BlockPhase`) - deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh (the ported kernel, ~1935 lines, WGMMA/TMA-based Hopper impl) - csrc/jit_kernels/heuristics/sm90_mega_moe.hpp (`MegaMoESM90Config` and block/pipeline/wave heuristics, mirroring csrc/jit_kernels/heuristics/mega_moe.hpp's structure) - csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp (JIT host runtime, mirroring csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp) - tests/test_mega_moe_hopper.py (correctness test with an SM90 capability guard; skips cleanly on non-Hopper GPUs) Modified files: - csrc/apis/mega.hpp: adds `get_symm_buffer_size_for_sm90_mega_moe` and `fp8_mega_moe_sm90`, registered via the existing `deep_gemm::mega::register_apis` (already wired into csrc/python_api.cpp, so no python_api.cpp changes were needed). - deep_gemm/mega/__init__.py + deep_gemm/__init__.py: expose `Sm90SymmBuffer`, `get_symm_buffer_for_sm90_mega_moe`, `transform_weights_for_mega_moe_sm90`, `fp8_mega_moe_sm90`, following the existing SM100 exposure pattern. - deep_gemm/include/deep_gemm/comm/barrier.cuh: generalizes `grid_sync`/`nvlink_barrier` to a templated workspace type (instead of hard-coding `layout::Workspace`) so the SM90 kernel can reuse them with `layout::MegaMoESM90Workspace`; existing SM100 call sites are unaffected (the type is still deduced from the argument). Also ports the PR's ARCH 900-1000 guarded trap-instead -of-printf change to the NVLink barrier timeout path. Not ported: shared-expert support and the `situ` activation (the PR does not implement either for SM90). No GPU/CUDA toolchain is available in this environment, so this is a best-effort, close-reading port verified via `python3 -m py_compile`, brace/paren balance checks on all new/modified C++/CUDA files, and confirming `import deep_gemm` fails identically (missing compiled `_C` extension) before and after these changes -- i.e. no regression introduced. It has not been compiled or run on real Hopper hardware. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
workspace.get_grid_sync_count_ptr<kGridSyncIndex>() failed NVCC compilation because workspace is a dependent-type parameter (WorkspaceT&); without the 'template' keyword the '<' is parsed as less-than.
Signed-off-by: rongfu.leng <lenronfu@gmail.com>
4012402 to
11885a0
Compare
2 participants
Ref: sgl-project#36
vllm support deepseek-v4-flash-fp8 model PR: vllm-project/vllm#53527