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: