From 8077fb1df75ac0ad1ebe2c11cb9824bb00f2558d Mon Sep 17 00:00:00 2001 From: ybyang <10629930+whybeyoung@users.noreply.github.com> Date: Thu, 11 Jun 2026 20:52:45 +0800 Subject: [PATCH] fix(deepgemm): align PP-parallel warmup bs to CP padding (#27922) --- .../sglang/srt/layers/deep_gemm_wrapper/compile_utils.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index e5c39f850..5e8139f6b 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -463,10 +463,15 @@ def pp_parallel_deep_gemm_warmup(model_runner) -> None: # n_splits ~= n_sms / ceil(bs/block_m) with block_m=64; sweep 5 bs to # cover the brackets real /generate hits (smallest decode shape, # mid-low, two mid, and n_splits=1 for ~5K+ token prefill). Ceil-align - # to attn_cp_size for DSA prefill CP's seq_len % cp_size == 0 assert. + # bs to the CP padding alignment (cp_size, or 2*cp_size for DSA + # in-seq-split). _dummy_run does not pad q/hidden like the real flow, so + # an unaligned bs makes DSA's padded num_splits longer than the q tokens + # and trips FlashMLA's "num_splits must have shape (b+1)" check. + from sglang.srt.layers.utils.cp_utils import get_cp_padding_align_size + n_sms = torch.cuda.get_device_properties(model_runner.device).multi_processor_count block_m = 64 - cp = max(model_runner.attn_cp_size, 1) + cp = max(get_cp_padding_align_size(), 1) batch_sizes = sorted( { ceil_align(bs, cp)