Skip to content

Avoid DSv4 compressed mixed multi-prefill FlashMLA crash #23741

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
foraxe wants to merge 2 commits into from
Closed

Avoid DSv4 compressed mixed multi-prefill FlashMLA crash #23741

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9271fc4
Avoid DSv4 compressed mixed multi-prefill crash
foraxe Apr 26, 2026
27467df
Cache DSv4 compressed prefill guard predicate
foraxe Apr 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 26 additions & 2 deletions python/sglang/srt/managers/scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
from torch.cuda import StreamContext as CudaStreamContext
from torch.distributed import barrier

from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.configs.model_config import ModelConfig, is_deepseek_compressed
from sglang.srt.constrained.grammar_manager import GrammarManager
from sglang.srt.disaggregation.decode import (
DecodePreallocQueue,
Expand Down Expand Up @@ -329,6 +329,11 @@ def __init__(

# Init model configs
self.init_model_config()
self._dsv4_compressed_prefill_guard_warned = False
Comment thread
foraxe marked this conversation as resolved.
self._is_dsv4_compressed_with_compressed_backend = (
is_deepseek_compressed(self.model_config.hf_config)
and self.server_args.get_attention_backends() == ("compressed", "compressed")
)

# Init metrics stats
self.init_metrics(tp_rank, pp_rank, dp_rank)
Expand Down Expand Up @@ -1950,6 +1955,23 @@ def get_num_allocatable_reqs(self, running_bs):
res = min(res, self.req_to_token_pool.available_size())
return res

def _get_effective_prefill_max_requests(self, running_bs: int) -> Optional[int]:
prefill_max_requests = self.server_args.prefill_max_requests
if (
running_bs > 0
and self._is_dsv4_compressed_with_compressed_backend
and (prefill_max_requests is None or prefill_max_requests > 1)
):
if not self._dsv4_compressed_prefill_guard_warned:
logger.warning(
"DSv4 compressed prefill guard active: running_bs=%s, "
"effective_prefill_max_requests=1",
running_bs,
)
self._dsv4_compressed_prefill_guard_warned = True
return 1
return prefill_max_requests
Comment thread
foraxe marked this conversation as resolved.

def get_new_batch_prefill(self) -> Optional[ScheduleBatch]:
prefill_delayer_single_pass = None
if self.prefill_delayer:
Expand Down Expand Up @@ -2019,6 +2041,8 @@ def _get_new_batch_prefill_raw(
if dynamic_size is not None:
chunked_prefill_size = dynamic_size

prefill_max_requests = self._get_effective_prefill_max_requests(running_bs)

# Prefill policy
adder = PrefillAdder(
self.page_size,
Expand All @@ -2030,7 +2054,7 @@ def _get_new_batch_prefill_raw(
chunked_prefill_size,
running_bs if self.is_mixed_chunk else 0,
self.priority_scheduling_preemption_threshold,
prefill_max_requests=self.server_args.prefill_max_requests,
prefill_max_requests=prefill_max_requests,
prefill_delayer_single_pass=prefill_delayer_single_pass,
dllm_config=self.dllm_config,
)
Expand Down
Loading