Avoid DSv4 compressed mixed multi-prefill FlashMLA crash

[foraxe]

Summary

This PR adds a narrow scheduler guard for DeepSeek V4 compressed attention.

When DSv4 compressed attention already has running decode requests, the scheduler now caps the effective prefill request count to 1. This avoids a mixed decode plus multi-new-prefill shape that currently produces invalid FlashMLA decode metadata on GB200.

The problematic live shape was:

Prefill batch, #new-seq: 1, #new-token: 4096, #running-req: 0
Prefill batch, #new-seq: 3, #new-token: 12288, #running-req: 1
CUDA error (.../get_decoding_sched_meta.cu:111): invalid argument

With this guard, the trailing prefills are admitted as separate prefill steps while decode is running:

DSv4 compressed prefill guard active: running_bs=1, effective_prefill_max_requests=1
Prefill batch, #new-seq: 1, #new-token: 4096, #running-req: 1
Prefill batch, #new-seq: 1, #new-token: 4096, #running-req: 2
Prefill batch, #new-seq: 1, #new-token: 4096, #running-req: 3

This is intentionally scoped to:

• DSv4 compressed attention
• compressed/compressed attention backend
• only while running_bs > 0
• only when the user did not already set prefill_max_requests <= 1

Validation

Validated on one GB200 node with full DeepSeek-V4-Flash TP=4, moe_runner_backend=flashinfer_mxfp4, disable_cuda_graph, disable_flashinfer_autotune, and SGLANG_JIT_DEEPGEMM_PRECOMPILE=0.

First, a single raw 4096+2 request completed and boundary instrumentation showed all layers 0..42 reached backend_post_flashmla, clearing the first FP8/DeepGEMM and first FlashMLA prefill paths as the blocker.

Then the original staggered failure shape was rerun:

• request 0: 4096+1500
• wait 5s
• requests 1-3: 4096+2

Result:

ok_count=4
error_count=0
elapsed_sec=251.926597

No #new-seq: 3 mixed-prefill batch occurred, and no get_decoding_sched_meta.cu:111 error occurred.

Notes

This PR does not claim to fix FlashMLA's internal handling of the mixed metadata shape. It prevents SGLang from constructing that currently unsupported DSv4 compressed-attention scheduling state. A follow-up investigation should compare the legal guarded metadata with the failing #running-req: 1 + #new-seq: 3 metadata to decide whether the deeper fix belongs in metadata construction or FlashMLA.

[gemini-code-assist]

Code Review

This pull request introduces a prefill guard for DeepSeek V4 models using compressed attention backends to limit prefill requests to 1 when the system is under load. I have suggested caching the configuration checks during initialization to avoid redundant evaluations in the hot scheduling loop, which will improve performance.

[foraxe]

Addressed the Gemini feedback in c2f72f48ec586b3a1a4e1f3217e6750e174dde0e by caching the DSv4 compressed/backend predicate during scheduler initialization and using the cached flag in the prefill hot path. Re-ran python3 -m py_compile python/sglang/srt/managers/scheduler.py and git diff --check.

[foraxe]

Hi, @Fridge003 , could you please take a look?

[foraxe]

Manually closing this PR as it is no longer being pursued.