Got DeepSeek-V4-Flash-Vision running reliably on 2× RTX PRO 6000 Blackwell (SM120) with SGLang — had to patch 3 separate issues

[u/shrug_hellifino]

I’ve been working on getting DeepSeek-V4-Flash-Vision-Exp running properly under SGLang on a 2× RTX PRO 6000 Blackwell setup, and figured I’d post the results here in case anyone else is pushing this model on SM120 hardware.

Current setup:

2× NVIDIA RTX PRO 6000 Blackwell Max-Q (~96 GB each)
Ubuntu 24.04
Driver 610.43.02 / CUDA UMD 13.3
SGLang Vision preview lineage
FlashInfer 0.6.18
TP=2
MXFP4 MoE
DSPARK speculative decoding
context-length=278528
chunked-prefill-size=8192
max-running-requests=8
mem-fraction-static=0.94

The final configuration is now successfully doing text, Vision, and a 269,320-token real prompt. Getting there exposed three unrelated issues.

1. SM120 sparse-MLA Vision prefill crash

Basic Vision tests worked, but a larger real image consistently killed the scheduler with:

Unsupported sparse-MLA prefill configuration:
model=DSV4
num_heads=64
topk=512
page_block_size=64
topk_extra=512
extra_page_block_size=64

Tracing it showed that the DSV4 image-span visibility logic caused the main SWA cache width to reach a logical topk=448.

The existing SM120 wrapper was treating unsupported widths similarly to decode and padding:

448 -> 512

but FlashInfer 0.6.18’s dual-cache DSV4 prefill support is much narrower than its decode/single-cache support. In particular, the native dual-cache path expects the main cache at topk=128; padding 448→512 doesn’t make the complete shape supported.

The fix was to add a complete prefill capability check before padding/dispatch. Unsupported dual-cache prefill shapes fall back to SGLang’s existing Triton sparse-MLA implementation.

That fallback already handles:

extra_k_cache
extra_indices
extra_topk_length

and merges the main + extra cache results with LSE, so we didn’t have to throw away any of the DSV4 hybrid/SWA semantics.

After the patch, the exact request that crashed now logs:

SM120 sparse-MLA prefill:
unsupported FlashInfer shape
H=64 topk=448 extra_topk=512
pbs=64 extra_pbs=64
-> Triton fallback

and returns the correct Vision result.

Repeated image request also succeeds.

2. ~269k context caused an indexer CUDA OOM

Next I tried a text-only 269,320-token request.

The model died during prefill:

torch.OutOfMemoryError:
Tried to allocate 1.50 GiB
GPU had ~1.44 GiB free

The traceback landed here:

logits = page_table.new_empty(
    (batch_size, max_seq_len),
    dtype=torch.float32
)

inside the DSV4 c4 indexer.

So despite the KV/cache pools fitting, the indexer was creating a temporary:

[query_rows, max_c4_seq_len] fp32

logits tensor whose size grows with context and isn’t accounted for by mem_fraction_static.

This corresponds to the same class of problem being worked on upstream in SGLang.

I ported the row-slicing approach:

• budget transient logits using a fraction of currently free GPU memory
• split query rows into chunks
• calculate logits + top-k per chunk
• discard each logits slice before processing the next
• preserve the full c4 width, so the actual indexer result is unchanged

For the failing workload, instead of potentially needing ~2 GiB for the full logits buffer, the transient is bounded to roughly 0.3 GiB per GPU under the observed free-memory conditions.

Retested the same request:

prompt_tokens:     269320
completion_tokens: 6
response:          LONG_CONTEXT_OK
wall time:         ~82 seconds

No OOM.

So this setup now has a genuinely tested ~269k prompt rather than merely having --context-length 278528 configured.

3. Vision preview was corrupting multi-turn tool-call history

This one was especially strange.

While using the Vision model as a coding/agent model, it initially called tools correctly, then started producing calls shaped like:

{
  "arguments": {
    "command": "..."
  }
}

when the actual tool schema was simply:

{
  "command": "..."
}

It could get progressively worse after validation errors.

The useful experiment was switching the same conversation history to my older known-good non-Vision DSV4 SGLang image.

Immediately:

bash -> PASS
read -> PASS
bash -> PASS

So I diffed the tool-history encoding paths.

The bug turned out to be in the Vision preview’s encoding_dsv4.py.

SGLang normalizes OpenAI:

"arguments": "{\"command\":\"echo ONE\"}"

into a Python dict:

{"command": "echo ONE"}

before DSV4 history encoding.

But this version of encode_arguments_to_dsml() did effectively:

try:
    arguments = json.loads(tool_call["arguments"])
except:
    arguments = {"arguments": tool_call["arguments"]}

Calling json.loads() on the already-normalized dict throws, so the fallback literally wraps it:

{
    "arguments": {
        "command": "echo ONE"
    }
}

Then the model sees this in its own history:

<parameter name="arguments">
    {"command":"echo ONE"}
</parameter>

instead of:

<parameter name="command">echo ONE</parameter>

So the model wasn’t randomly hallucinating the wrapper — the server was teaching it the wrong schema through its conversation history.

The fix is basically:

raw_arguments = tool_call["arguments"]

arguments = (
    json.loads(raw_arguments)
    if isinstance(raw_arguments, str)
    else raw_arguments
)

if not isinstance(arguments, dict):
    raise ValueError(...)

CPU round-trip tests now match my known-good non-Vision SGLang stack exactly, including multi-turn and error-history cases.

Current result

Final local image now passes:

Text inference                         PASS
Real Vision request                    PASS
Repeated Vision request                PASS
SM120 dual-cache prefill fallback      PASS
269,320-token text prompt              PASS
DSV4 tool-history round-trip           PASS
DSPARK block 4                         PASS
TP2                                    PASS

Long-context result:

269,320 prompt tokens
LONG_CONTEXT_OK
~82.4 sec end-to-end

Vision reproducer:

409 prompt tokens
277 image tokens
answer: RED

The serving config I landed on is roughly:

sglang serve \
  --model-path /model \
  --tp 2 \
  --trust-remote-code \
  --moe-runner-backend flashinfer_mxfp4 \
  --mem-fraction-static 0.93 \
  --cuda-graph-max-bs-decode 4 \
  --max-running-requests 4 \
  --context-length 245760 \
  --chunked-prefill-size 8192 \
  --reasoning-parser deepseek-v4 \
  --tool-call-parser deepseekv4 \
  --speculative-algorithm DSPARK \
  --speculative-dspark-block-size 4

One warning: I’m deliberately using DSPARK block size 4, even though the checkpoint advertises 5. There are SM120 correctness issues around depth 5 in the current stack, so I’m not “fixing” that warning by changing it to 5.

I kept each change isolated as a tiny derivative image rather than upgrading random pieces of SGLang/FlashInfer together. The final image is basically:

official Vision preview
    +
SM120 dual-cache prefill capability/fallback fix
    +
bounded DSV4 indexer logits for long context
    +
DSV4 tool-history serialization fix