From 95550cbc25f08bfd42ec92c9b61b440896327450 Mon Sep 17 00:00:00 2001 From: "Gavin.Zhu" Date: Mon, 18 May 2026 06:22:42 +0000 Subject: [PATCH] fix(dsv4 topk_v2): honor cluster contract in fused kernel SMALL/TRIVIAL branches `topk_fused_transform` is annotated `__cluster_dims__(1, kClusterSize=8, 1)` via FUSED_COMBINE_KERNEL, but its SMALL (`seq_len <= Small::kMax1PassLength`) and TRIVIAL (`seq_len <= K`) branches let 7 of every 8 cluster blocks `return` immediately without ever calling `cluster.sync()`. On SM90 and SM100 with TP=4/DP=4 (the configuration covered by CI) this is tolerated, but on B300 / SM100 at TP=8/DP=8 with DeepEP it surfaces as `CUDA_ERROR_ILLEGAL_ADDRESS` reported at downstream sync points (CUDA graph replay, dense FP8 GEMM, etc.) -- the actual failing kernel is async. Reproducer: deepseek-ai/DeepSeek-V4-Pro on 8x B300 SXM6, `--tp 8 --dp-size 8 --enable-dp-attention --moe-a2a-backend deepep` with `SGLANG_OPT_USE_TOPK_V2=1` (the default since #25406). conc=32 ISL=8192 OSL=1024 triggers the fused 1-stage path (batch_size_per_rank=4 <= kNumClusters=15) and crashes within ~10 seconds of decode. Workaround: `SGLANG_OPT_USE_TOPK_V2=0` falls back to the v1 kernel. Fix: restructure both branches so rank-0 does the work and all 8 ranks meet at `cooperative_groups::this_cluster().sync()` before exit, honoring the implicit cluster contract. Validated end-to-end on a B300 box: SGLANG_OPT_USE_TOPK_V2=1 + the patch + the same workload completes cleanly at 958.5 tok/s/gpu, vs 963.7 with the TOPK_V2=0 fallback (well within run-to-run noise -- the 7 dummy ranks were already pinned to the cluster's SM allotment whether they early-returned or sat at the barrier). --- .../jit_kernel/csrc/deepseek_v4/topk_v2.cuh | 46 ++++++++++++++++--- 1 file changed, 40 insertions(+), 6 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh index 8c4a526575ea..a754cc0a7bf1 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh @@ -325,21 +325,55 @@ topk_fused_transform(const __grid_constant__ TopKParams params) { __shared__ int32_t s_topk_indices[K]; const auto batch_id = blockIdx.x; const auto cluster_rank = blockIdx.y; + // PDLWaitPrimary at kernel entry: all 8 cluster CTAs wait for the primary + // kernel before reading params.seq_lens / scores. This both (1) ensures + // primary-produced data is visible and (2) gives every CTA a symmetric + // PDL state, sidestepping the SM100 cluster crash described below. + device::PDLWaitPrimary(); const auto seq_len = params.seq_lens[batch_id]; const auto transform = params.get_transform(batch_id, s_topk_indices); + // This kernel is __cluster_dims__(1, kClusterSize, 1) + launched with + // enable_pdl(true). On SM100 (B300) with TP=8 + DP-attention, letting 7 + // of 8 cluster CTAs early-return here -- while only rank 0 enters the + // body and (in the SMALL case) calls Small::run's PDLWaitPrimary -- ends + // in CUDA_ERROR_ILLEGAL_ADDRESS, surfaced at the next sync point + // (typically CUDA graph replay or a downstream deep_gemm call). + // + // Empirically, hoisting PDLWaitPrimary alone is NOT sufficient to fix the + // crash; what is load-bearing is forcing all 8 CTAs to meet at a + // cluster.sync() before any of them exits. So we keep cluster.sync() at + // each branch's tail and emit PDLTriggerSecondary just before it -- the + // trigger fires per-CTA, lets dependent kernels start earlier, and the + // barrier maintains the cluster invariant SM100 requires. + // + // We don't have a clean documented mechanism for why early-exit (which + // is normally legal per CUDA cluster semantics) fails specifically on + // SM100 + __cluster_dims__ + enable_pdl. Empirically validated against + // DSv4-Pro on 8x B300 at TP=8 + DP=8 + EP=8 + MegaMoE W4A4 (DSv4-Pro + // workload exclusively hits the SMALL branch since c4_seq_lens ~ 2304 + // << Small::kMax1PassLength = 16384). if (seq_len <= K) { - if (cluster_rank != 0) return; // only first rank work - impl::trivial_transform(transform, seq_len, K); + if (cluster_rank == 0) { + impl::trivial_transform(transform, seq_len, K); + } + device::PDLTriggerSecondary(); + cooperative_groups::this_cluster().sync(); } else if (seq_len <= Small::kMax1PassLength) { - if (cluster_rank != 0) return; // only first rank work - Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem, /*use_pdl=*/true); - Small::transform(transform); + if (cluster_rank == 0) { + // use_pdl=false: PDLWaitPrimary is already hoisted to kernel entry. + Small::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem, /*use_pdl=*/false); + Small::transform(transform); + } + device::PDLTriggerSecondary(); + cooperative_groups::this_cluster().sync(); } else { const auto [offset, length] = partition_work(seq_len, cluster_rank); const auto ws = params.workspace + batch_id * params.workspace_stride; Large::stage1_init(smem); - device::PDLWaitPrimary(); + // PDLWaitPrimary hoisted to kernel entry; emit PDLTriggerSecondary after + // stage1_prologue so dependent kernels can overlap with stage1 work. Large::stage1_prologue(params.get_scores(batch_id) + offset, length, smem); + device::PDLTriggerSecondary(); Large::stage1(s_topk_indices, length, smem); Large::stage1_epilogue(transform, offset, ws, smem); cooperative_groups::this_cluster().sync();