Repository navigation
Public Release 26/09 - #432
Conversation
| } | ||
|
|
||
| void set_num_sms(const int& new_num_sms) { | ||
| DG_HOST_ASSERT(0 < new_num_sms and new_num_sms <= jit->device.get_num_sms()); |
There was a problem hiding this comment.
🟡 warning: 保留 set_num_sms(0) 的重置语义: 旧实现允许传入 0，并由 get_num_sms() 将其解释为恢复设备的全部 SM；这里改成严格大于零后，已有调用 set_num_sms(0) 会直接抛异常，而且一旦限制过 SM 数，调用方也无法再通过原有方式恢复默认值。建议继续接受零作为重置值。
🤖 v6
| DG_HOST_ASSERT(cu_seq_len_k_start.dim() == 1 and cu_seq_len_k_start.size(0) == num_q_tokens); | ||
| DG_HOST_ASSERT(cu_seq_len_k_end.dim() == 1 and cu_seq_len_k_end.size(0) == num_q_tokens); | ||
| DG_HOST_ASSERT(sparse_kv_block_indices.scalar_type() == torch::kInt and sparse_kv_block_indices.is_contiguous()); | ||
| DG_HOST_ASSERT(cu_seq_len_k_start.scalar_type() == torch::kInt and cu_seq_len_k_start.is_contiguous()); | ||
| DG_HOST_ASSERT(cu_seq_len_k_end.scalar_type() == torch::kInt and cu_seq_len_k_end.is_contiguous()); |
There was a problem hiding this comment.
🟡 warning: 校验稀疏元数据输入位于同一 CUDA 设备: 当 sparse_kv_block_indices 位于 CUDA、但 cu_seq_len_k_start 或 cu_seq_len_k_end 是 CPU 张量或位于另一块 GPU 时，这些检查仍会通过，随后其 data_ptr<int>() 被直接传给 CUDA kernel，可能导致非法地址并破坏 CUDA 上下文；应像 get_mqa_logits_metadata 一样检查 CUDA 设备及设备一致性。分页版本的 context_lens、block_table 和 indices 也有同样问题。
🤖 v6
|
|
||
| # Backward compat: derive `mma_type` from `use_fp8_dispatch` if provided | ||
| if use_fp8_dispatch is not None: | ||
| assert use_fp8_dispatch == (mma_type.split('x')[0] == 'fp8') |
There was a problem hiding this comment.
🔴 critical: 注释声称“从 use_fp8_dispatch 推导 mma_type”，但实现只是 assert use_fp8_dispatch == (mma_type.split('x')[0] == 'fp8')，并没有推导。旧调用方若传 use_fp8_dispatch=False（而 mma_type 保持默认 'fp8xfp4'）会直接断言失败，等于破坏了本想保留的向后兼容。建议真正推导（False 时映射到对应旧语义的 mma_type），或在不兼容时抛出带明确迁移指引的错误而非裸 assert。
🤖 v5
| DG_HOST_UNREACHABLE("Unsupported MoE scoring function"); | ||
| } | ||
|
|
||
| static const torch::Tensor& get_score_barriers(const torch::TensorOptions& options) { |
There was a problem hiding this comment.
🟡 warning: get_score_barriers 使用函数内 static std::unordered_map<CUDAStream, Tensor> 持有 per-stream 的 CUDA tensor：(1) 无锁保护，多线程并发调用存在数据竞争；(2) 这些 tensor 在进程退出的静态析构阶段销毁，此时 CUDA driver 可能已卸载，会触发与本 MR 中为 cuBLASLt workspace 专门规避的 cudaErrorCudartUnloading 同类问题（见 runtime.hpp 的 DG_USE_TEMP_CUBLASLT_WORKSPACE 注释）。建议纳入 Runtime 生命周期管理，或提供同样的规避手段（temp 分配/有意泄漏）。
🤖 v5
| DG_CUBLASLT_CHECK(cublasLtCreate(&cublaslt_handle)); | ||
| } | ||
|
|
||
| ~Runtime() noexcept(false) { |
There was a problem hiding this comment.
🟡 warning: ~Runtime() noexcept(false) 中调用 DG_CUBLASLT_CHECK(cublasLtDestroy(...)) 可能抛异常。Runtime 是 LazyInit 单例，析构发生在静态销毁阶段，此时抛异常会直接 std::terminate；且若 CUDA 已先于它卸载，cublasLtDestroy 很可能返回错误。建议析构中只记录错误、不抛出。另外 cublaslt_workspaces 这个按 stream 增长的 map 同样没有并发保护。
🤖 v5
| // FP8 GEMM alias names | ||
| // FP4 and FP8 GEMM alias names | ||
| m.attr("fp4_gemm_nt") = m.attr("fp8_fp4_gemm_nt"); | ||
| m.attr("fp8_gemm_nt") = m.attr("fp8_fp4_gemm_nt"); |
There was a problem hiding this comment.
🔵 suggestion: 这里把历史名称 fp8_gemm_nt（以及 m_grouped_fp8_gemm_* 等）重新绑定到通用的 fp8_fp4_gemm_nt 实现，意味着既有只传 FP8 张量的调用现在会走统一的 FP8/FP4 派发，可能选择到不同的 heuristic/kernel。方向本身没问题，但属于对既有用户的静默行为变化，建议在 release notes / changelog 中显式说明，避免性能或数值预期出现无感知的变化。
🤖 v4f
| const bool& use_unaligned_ks) { | ||
| DG_HOST_ASSERT(jit->device.get_arch_major() == 10); | ||
| const auto [num_q_tokens, num_max_sparse_blocks] = get_shape<2>(sparse_kv_block_indices); | ||
| DG_HOST_ASSERT(num_q_tokens > 0 and num_kv_tokens >= 0); |
There was a problem hiding this comment.
🟡 warning: get_sparse_mqa_logits_metadata 允许 num_kv_tokens >= 0，而实际消费该 metadata 的 fp8_fp4_sparse_mqa_logits 要求 num_kv_tokens > 0（第 316 行）。这会导致可以为「空 KV」构造出一份 metadata，随后在 kernel 端才失败（且该路径注释说明会跳过 header 校验以避免同步），错误发生点与原因不够直观。建议在此处收紧为 > 0，或明确支持并测试空 KV 的情况。
🤖 v4f
| cublaslt_gemm_tn, cublaslt_gemm_tt, | ||
| cublaslt_nvfp4_gemm_nt, | ||
| batched_syrk, batched_symm, | ||
| ) |
There was a problem hiding this comment.
🔵 suggestion: 本次去掉了包裹 _C 内核导入的 try/except ImportError，因此当 _C 未构建或某个符号缺失（例如旧工具链）时，import deep_gemm 会直接硬失败，而不再像以前那样优雅降级。结合新的 CUDA≥12.9 / C++20 要求，这可能是刻意为之，但确实移除了旧版本的兼容逃生通道。建议保留一个带明确错误信息的保护，或至少在文档中说明导入失败的行为。
🤖 v4f
| @@ -28,14 +25,10 @@ Despite its lightweight design, DeepGEMM's performance matches or exceeds expert | |||
|
|
|||
There was a problem hiding this comment.
🔵 suggestion: Requirements 写的是 PyTorch 2.3 or higher，但 Mega MoE 相关代码/注释明确标注 symmetric memory 需要 PyTorch >= 2.9（deep_gemm/mega/init.py 中 # NOTES: requires PyTorch >= 2.9）。建议按功能分别列出版本下限（基础 GEMM vs Mega MoE 对称内存），避免用户按 README 的最低 2.3 安装后在 Mega MoE 路径上遇到运行时失败。
🤖 v4f
| } | ||
|
|
||
| static const torch::Tensor& get_score_barriers(const torch::TensorOptions& options) { | ||
| const auto stream = at::cuda::getCurrentCUDAStream(); |
There was a problem hiding this comment.
🔵 suggestion: get_score_barriers 以 stream 为 key 缓存一份一次清零的 score barrier，并断言创建时不在 CUDA graph capture 中，但后续每次 launch 之间并不会重新清零/重置。该实现的正确性完全依赖 kernel 在每轮结束时把 barrier 相位恢复干净；请确认（并最好注释/断言）kernel 确实在异常路径、提前返回下也会正确复位，否则连续调用可能产生跨调用残留。
🤖 v4f
🤖 ds-review-bot Code Reviewv6该补丁引入了一个公开配置 API 的兼容性回退，并且新的稀疏注意力元数据接口可能将主机或跨设备指针传入 CUDA kernel。 v5Public Release 26/09（v2.6.1 → v2.8.0，115+ 文件，约 +12k/−5k 行）。主要内容：1) 新增 Mega MHC 融合大核（HC prenorm GEMM + Sinkhorn + RMSNorm + FP8 量化输出，csrc/apis/mega_mhc.hpp 及对应 impls/epilogue/layout/scheduler）；2) 新增 Mega Gate 融合 MoE 门控（BF16 gate GEMM + top-k，支持 sigmoid/sqrtsoftplus/identity）；3) 新增 Sparse MQA Logits（含 paged 版本）及其 metadata kernel；4) Mega MoE 新增 FP8×FP8 路由专家权重支持（mma_type 参数）；5) 内置 JIT 模块整体迁移到新的 DeepJIT 子模块（third-party/deep_jit），移除 fmt 依赖并改用 C++20 <format>，环境变量语义相应调整（DG_JIT_* 回退 DJ_JIT_*、多路径缓存、拆分 PTXAS 检查开关）；6) 构建门槛提升（CUDA ≥ 12.9、PyTorch ≥ 2.3、C++20、CMake ≥ 3.18）；7) 新增 469 行 SF 格式规范文档（docs/scaling-factor-format.md，内容详尽并列出 source of truth 文件）与大量新测试（test_mega_gate/test_mega_mhc 等），可验证性明显提升。已核对：README 中环境变量说明与 DeepJIT 实际行为一致（默认缓存目录 $HOME/.dj、多路径查找、DG_JIT_CHECK_NO_SPILLS/NO_LOCAL_MEMORY 拆分）；config.hpp 的 get_pdl 解引用 optional 当前安全（DeepJIT 默认初始化为 false）。需要作者跟进的问题见 comments：use_fp8_dispatch 向后兼容处理与注释不符且会使旧调用断言失败；两处以 static/成员 map 持有 CUDA tensor 的全局状态存在线程安全与进程退出时 CUDA 已卸载的析构风险（与本次为 cuBLASLt workspace 加的 DG_USE_TEMP_CUBLASLT_WORKSPACE 规避的问题同源）；~Runtime() noexcept(false) 可能在静态析构期抛出导致 terminate；SymmBuffer.destroy() 清理不完整；wheel 打包整个 csrc/docs 的体积问题等。均不阻塞发布。 v4f本次 Public Release 26/09（约 115 个变更文件，+12k/-5k 行）是一次较大的架构与功能更新。核心改动是把运行时 JIT 完整迁移到新的 DeepJIT 子模块：删除内置的 csrc/jit/*（cache/compiler/device_runtime/handle/include_parser/kernel_runtime）与 csrc/utils/{hash,lazy_init,system} 等，移除 NVRTC 与 third-party/fmt 依赖，改用 std::format，C++ 标准提升到 C++20，最低要求提升到 CUDA 12.9 / PyTorch 2.3，并把 csrc/docs 打进 wheel 以便报错定位。功能层面新增了 Sparse MQA logits（含 metadata 构造）、MegaGate（BF16 gate GEMM 融合 MoE top-k 路由）、MegaMHC（mHC + Sinkhorn + RMSNorm + 可选 FP8 量化）、FP8×FP8 Mega MoE、k_grouped_fp4_gemm_nt_contiguous、cuBLASLt NVFP4 GEMM 与 batched_syrk/symm，新增 use_deterministic_algorithms / psum-layout 支持，并统一了 fp8/fp4 GEMM 的命名别名。配套新增了 docs/scaling-factor-format.md 与 test_mega_gate.py / test_mega_mhc.py。整体架构方向清晰、重构合理，多数新内核只在 SM100 上可用并有较密的 host 端校验；本次审查未在已阅读的范围内发现明确的数值/逻辑错误，但缓冲区生命周期与复用、JIT/运行时资源管理、以及若干断言与文档一致性上存在可改进点，详见各条 comment。 Files reviewed: 115 |
Merge upstream 66081d4 (Public Release 26/09, deepseek-ai#432) into dev while retaining SGL packaging, API compatibility, and downstream kernels. Synchronize upstream tests and include the validated integration fixes. Keep the new sparse MQA implementation aligned with upstream. Squash the merge and its three follow-up commits into one two-parent merge commit. Preserve the exact validated tree; no source changes.
* Fix a race condition where in-flight tensormaps are updated in-place (deepseek-ai#343) * Public release 26/07 (deepseek-ai#377) * Public Release 26/09 (deepseek-ai#432) * Public release 26/09 * Update News --------- Co-authored-by: Ivan Komarov <Ivan.Komarov@dfyz.info> Co-authored-by: Pengle Zhang <53843527+guyan364@users.noreply.github.com> Co-authored-by: Zhean Xu <94977922+zheanxu@users.noreply.github.com>
…ring (#16) * Fix a race condition where in-flight tensormaps are updated in-place (deepseek-ai#343) * Public release 26/07 (deepseek-ai#377) * Public Release 26/09 (deepseek-ai#432) * Public release 26/09 * Update News * Fix Mega MoE task info slot release ordering (deepseek-ai#441) --------- Co-authored-by: Ivan Komarov <Ivan.Komarov@dfyz.info> Co-authored-by: Pengle Zhang <53843527+guyan364@users.noreply.github.com> Co-authored-by: Zhean Xu <94977922+zheanxu@users.noreply.github.com>
#447) * Public Release 26/09 (#432) * Public release 26/09 * Update News * feat(sm120): add sparse MQA support on upstream main APIs Port contiguous and paged sparse kernels to DeepJIT while preserving SM100 dispatch, stream-local metadata, entry-balanced scheduling and register guards. Restore histogram synchronization and capture-safe UE8M0 validation, with independent SM120 numerical and CUDA graph regression fixtures. * feat(sm120): restore native GEMM paths on main APIs Restore the original BF16/FP8/FP4 kernels, scheduler, heuristics, and split-K paths from 139f504 through DeepJIT adapters. Support alpha, independent C/D strides, GPU-only KPSUM, and bounded padding cleanup. Separately repair inherited mixed-tail, producer handoff, grouped row ownership, and TMA output-width defects. Preserve the original TMA/MMA pipelines and dispatch choices except for the documented safety eligibility guard. This is a bounded restoration checkpoint, not full migration or universal performance parity. BF16 odd-N, FP4 K-grouped, and remaining public integrations are unfinished. * feat(sm120): preserve nv_dev features alongside main API SM120 kernels * fix(kernels): address migration review and restore grouped TMA pipeline Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * fix(sm120): harden GEMM execution and complete HeadSplits output mapping - Wait for in-flight TMA descriptor reads before overwriting the per-CTA K-grouped BF16 descriptor, and add full-SM randomized alias/separate-C regressions. - Require packed INT32 SFB on the SM120 skip-head path instead of silently reading FP32 scales as UE8M0. - Apply HeadSplits index mapping with logical-column bounds in scalar stores and split-K reduction, replacing the former capability guards while keeping alpha/C/stride contracts. - Store BF16 batched einsum and HC prenorm directly to supported strided outputs, keeping the existing fallback for unsupported layouts. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * test(sm120): collect SM120 tests into dedicated files Restore tests/test_{attention,bf16,fp8_fp4,hyperconnection}.py to their upstream byte state and consolidate all SM120-specific coverage into tests/test_sm120_*.py plus the shared tests/sm120_exercise.py helper. This eliminates per-release merge conflicts with upstream's wholesale rewrites of the shared test files. Switch gating from collection-time string skipif (which initialized CUDA during pytest collection on non-SM120 machines and broke the sanitizer runner's direct-call path) to upstream's call-time test_filter convention, and give each standalone file its own __main__ runner matching upstream script style. Moved test bodies are AST-identical; test counts and results are unchanged (741 passed both before and after on SM120, memcheck/racecheck subsets clean). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * feat(sm120): vendor device layer from DeepGEMM-sm120 v0.1.3 Adopt the standalone DeepGEMM-sm120 repository as the single source of truth for the SM120 device headers: add the CUDA>=13 compile-time guard (silent block_scale drop on pre-13 ptxas), relocate tensor_map_replace_global_dim_in_smem into common/sm120_utils.cuh so the kernels no longer depend on ptx/tma.cuh providing it, and inline the FP4 smem pack factor in layout/sparse_mqa_logits.cuh. Marker-free vendoring: the files carry no canonical-source comments or other foreign markers; provenance and drift control live on the canonical repo's side (per-tag manifest, check_vendor.py --fork, scheduled downstream watcher). Verified: 741 SM120 tests pass on sm_120a from a fresh JIT cache. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * fix(sm120): gate ptxas -O2 for sparse MQA logits on CUDA 13.3+ CUDA 13.3's ptxas crashes (SIGSEGV; DeepJIT compile subprocess exit 139) when compiling the non-warp-specialized MXFP4 sparse MQA logits variants at -O3 with any --register-usage-level (lucifer1004/DeepGEMM-sm120#2); CUDA 13.2, the validated toolchain, is unaffected. Add a toolchain-gated per-kernel override via DeepJIT CompilerOptions::extra_nvcc_flags: when nvcc reports >= 13.3, the non-warp-specialized instantiations compile with ptxas -O2, which avoids the crash and reproduces the 13.2 register allocation (REG:101 for the reference instantiation). The warp-specialized variant is not affected by the crash and keeps the default flags (its compiled reg==64 contract also verified under -O2). Validated on sm_120a: sparse MQA suite 9/9 under CUDA 13.3.73 with the gate active (previously 3 failing tests), and 9/9 under CUDA 13.2.86 with the gate inactive (unchanged default path). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * feat(sm120): FP8 paged MQA logits PAGE_KV=32 support Vendor DeepGEMM-sm120 v0.1.4, which derives BLOCK_KV = min(PAGE_KV, 64) in the FP8 paged kernel (mirroring the FP4 sibling): a 32-row page gets a 32-row compute tile and 4 KV groups at SPLIT_KV=128, instead of a 64-row tile straddling two non-contiguous physical pages. DSv4.1 indexer caches mix 64- and 32-state pages (vllm-project#14). Host glue: relax the FP8 launcher and fused-cache API gates to admit block_kv=32 for arch 12, derive tile_kv = min(block_kv, 64) like the FP4 launcher, and drop to two KV stages for page32 + 64 heads + paired queries, where three stages exceed the 99 KiB SMEM budget by 4 bytes. Tests: the paged-MQA contract matrix gains (fp8, page32), and a focused case covers the two-stage fallback (page32, 64 heads, paired/varlen). Validated on sm_120a from a fresh JIT cache: paged MQA suite 20/20 passed, incl. graph replay and legacy-API cross-checks. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * perf(sm90): clamp paged MQA scheduler range once at construction The stale/OOB-metadata protection added during the migration lived in the per-task hot path (compound has_next_task/exist_q_atom_idx checks per fetch), regressing paged MQA logits on SM90: ~1% on upstream shapes and 13-16% at next_n=4. Move the protection to the constructor: clamp end_q_atom_idx to the in-batch sentinel (zeroing end_kv_idx at the sentinel) and collapse out-of-range starts, then make the per-task end test a single ordered 64-bit compare on (q_atom, kv) packed keys. The ordered compare also keeps the hang fix: an advance overshooting the end is caught by >=, never by exact equality. context_lens/indices are only dereferenced when current_q_atom_idx is provably in batch (end is clamped at or below the sentinel, and the sentinel range always has end_kv == 0). Patch from local review. Validated on H200: test_paged_mqa_logits 24/24 cases pass; next_n=4 (H=64, D=128, page 64) 140.6us vs 154.6us pre-patch at L=8192 and 1093.3us vs 1233.1us at L=65536 (-9%/-11%), upstream-shape sweep at parity or better. Cross-compiled for sm_90a. Reviewer's full 105-case SM90 suite passes on their side. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * perf(layout): pow2 shift fast path in k-grouped SF pack row accounting pack_fp32_into_ue8m0's per-group loop ran three runtime integer divisions per group per thread on the psum path (align of previous end, align of group K, ceil_div to SF rows). gran_k and k_alignment are powers of two at every call site (host-asserted), so detect that once per block and use shift forms instead; the generic division path is kept for any future non-pow2 alignment. The earlier suggestion to drop the align entirely on SM90/100 was unsafe: psum tail K sizes are not k_alignment-multiples, and without the align the row count comes up short (B300 fp8 psum accumulate regression, diff 0.024 vs 0.008 tolerance). Validated: pack output bitwise identical to the division form on a 128-group psum shape (B300); kernel time 38.7 -> 22.1 us/call (-43%) on the same shape. SM100 k-grouped fp8/bf16 contiguous suites pass on B300 (fresh JIT cache), SM120 k-grouped regressions pass on sm_120a, and the TU cross-compiles for sm_103a/sm_120a. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * docs(layout): note why the pow2 shift forms are exact on SM90/100 psum end offsets are k_alignment-multiples and gran_k divides k_alignment there, so aligned_group_k % gran_k == 0 and the rounding shift is exact division; the ceil semantics matter only on SM120. Comment-only change. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * perf(layout): move gran_k/k_alignment into the JIT template for SF pack Follow-up to 3aca951: the runtime pow2 branch regressed non-pow2 k_alignment (reported by RayWang96: k_alignment=384, a third of their k-grouped cases, 91 -> 124 us vs main on B300). Instead of branching at runtime, pass gran_k/k_alignment as JIT template parameters: the per-group align/ceil_div fold to shifts (pow2) or multiply-high sequences (non-pow2) at compile time, covering both alignment classes with no runtime branch. Runtime kernel args drop accordingly. Validated on B300: pack output bitwise identical to the pre-change form at k_alignment 128 and 384; pack kernel 38.8 -> 14.0 us/call at 128 (-64%) and 65.2 -> 32.3 us/call at 384 (-50%); k-grouped fp8/bf16 contiguous suites pass from a fresh JIT cache. SM120 k-grouped regressions pass on sm_120a; the TU cross-compiles for sm_103a/sm_120a. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * fix(sm120): vendor v0.1.5 generic-proxy WAR fences from DeepGEMM-sm120 Vendor lucifer1004/DeepGEMM-sm120 v0.1.5 (b31a688): 14 fence_view_async_shared() sites across 8 SM120 kernel headers, ordering generic-proxy smem reads (ldmatrix/ld_shared) before the empty-barrier arrive that hands each stage back to the TMA producer. Ports #453 (bf16 GEMM) and fixes the same hazard found by audit in bmk_bnk_mn, tf32_hc_prenorm, fp8_fp4_gemm_1d1d, and the fp8/fp4 dense+paged MQA logits kernels. Sparse MQA is exempt (legacy cp.async producer). Perf A/B on RTX PRO 6000 (bf16 GEMM, fp8 dense/paged MQA, 3 runs): all deltas within run-to-run noise. Fork suite: 235 passed (GPU4, fresh JIT cache). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --------- Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Zhean Xu <94977922+zheanxu@users.noreply.github.com>
2 participants
This update introduces new kernels, features, and performance improvements:
New kernels
New features
Performance improvements
Breaking changes
use_deterministic_algorithms()API, which defaults toFalseso some kernels can run faster.Trueto restore deterministic, batch-invariant behavior for RL and similar workloads.Contributors