chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

B300 NSA: default NSA prefill to flashmla under context parallelism (fix k_rope CP crash)

Browse Source 
_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) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-21 17:33:49 +00:00
parent b3a4f4174f
commit 2b8593ff24
1 changed files with 12 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/sglang/srt/server_args.py
View File

@@ -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: