From 2b8593ff245d533e1712f19d5de18e6da6bdfe56 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 21 Jun 2026 17:33:49 +0000 Subject: [PATCH] B300 NSA: default NSA prefill to flashmla under context parallelism (fix k_rope CP crash) _set_default_nsa_backends defaulted CP prefill to trtllm on Blackwell+fp8-KV, but the trtllm NSA prefill path mis-shards k_rope under CP: the rope input carries all CP ranks' tokens (nnz*cp) while the per-rank rope kernel asserts k_rope_in.size(0)==nnz -> 'Check failed: 256 vs 64' at cp=4. The function's own comment already notes flashmla_sparse pads heads 64->128 and works under CP; this just auto-selects it under CP instead of requiring an explicit --nsa-prefill-backend override. Matches the gov journal (2026-06-08 entry: trtllm-for-CP-prefill was the wrong backend). Decode/non-CP unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/server_args.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 561fcb88d..9f1e54fc8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1664,12 +1664,18 @@ class ServerArgs: (user_set_prefill and self.nsa_prefill_backend != "trtllm") or (user_set_decode and self.nsa_decode_backend != "trtllm") ) - if self.dp_size == 1 and major >= 10 and not user_wants_flashmla: - # Default to trtllm on Blackwell without DP attention. The "no DP - # attention" limitation does NOT apply to context-parallel prefill - # (flashmla_sparse pads heads 64->128 and works under CP), which is - # why an explicit flashmla choice above takes this out of the trtllm - # branch entirely. + if ( + self.dp_size == 1 + and major >= 10 + and not user_wants_flashmla + and not self.enable_nsa_prefill_context_parallel + ): + # Default to trtllm on Blackwell without DP attention. NSA-prefill + # context parallelism is excluded here on purpose: the trtllm NSA + # prefill path mis-shards k_rope under CP (the rope input carries all + # CP ranks' tokens, nnz*cp, but the per-rank rope kernel asserts + # k_rope_in.size(0)==nnz). flashmla_sparse pads heads 64->128 and works + # under CP, so CP falls through to the flashmla branch below. if not user_set_prefill: self.nsa_prefill_backend = "trtllm" if not user_set_decode: