[Bug]: DeepSeek-V4-Flash-DSpark fails on H200/SM90 with FlashMLA KV cache shape mismatch

[yy-fighting]

Your current environment

• vLLM version: 0.23.1rc1.dev788+gfa4321de3
• torch: 2.11.0+cu129
• GPU: NVIDIA H200 (SM90), 8× GPU
• Model: deepseek-ai/DeepSeek-V4-Flash-DSpark (local path /data/models/DeepSeek-V4-Flash-DSpark)
• Command used:

vllm serve /data/models/DeepSeek-V4-Flash-DSpark \
  --trust-remote-code \
  --tensor-parallel-size 8 \
  --kv-cache-dtype fp8 \
  --block-size 256 \
  --max-model-len auto \
  --max-num-seqs 256 \
  --tokenizer-mode deepseek_v4 \
  --reasoning-parser deepseek_v4 \
  --spec-method dspark \
  --spec-model /data/models/DeepSeek-V4-Flash-DSpark \
  --spec-tokens 5 \
  --port 30004

🐛 Describe the bug

Serving DeepSeek-V4-Flash-DSpark without speculative decoding works. Enabling DSpark speculative decoding (--spec-method dspark) causes engine initialization to fail during KV-cache warmup of the draft model.

The model weights load successfully, including the DSpark draft weights:

(Worker_TP0 pid=...) INFO ... [dspark.py:455] DSpark draft model loaded: 96 params

Then the worker crashes in DeepseekV4FlashMLAAttention._forward_decode -> flash_mla_with_kvcache with:

RuntimeError: kv must have shape (num_blocks, page_block_size, h_kv, bytes_per_token)

Full traceback:

(Worker_TP0 pid=4109550) ERROR ... [multiproc_executor.py:1004]     self.forward_mqa(q, kv, positions, out)
(Worker_TP0 pid=4109550) ERROR ... [multiproc_executor.py:1004]   File ".../vllm/models/deepseek_v4/nvidia/flashmla.py", line 219, in _forward_decode
(Worker_TP0 pid=4109550) ERROR ... [multiproc_executor.py:1004]     out, _ = flash_mla_with_kvcache(
...
RuntimeError: kv must have shape (num_blocks, page_block_size, h_kv, bytes_per_token)

What I tried

1. --kv-cache-dtype fp8 (default/expected): fails with the shape error above.
2. --kv-cache-dtype bfloat16: fails earlier with AssertionError: DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, got bfloat16.
3. --attention-backend FLASHINFER_MLA_SPARSE_DSV4 for both target and draft: fails with Error in function 'TllmGenFmhaRunner' ... Unsupported architecture on H200/SM90.

So on H200/SM90 there appears to be no working attention backend for DeepSeek-V4-Flash-DSpark with DSpark.

Related issues/PRs

• [Bug]: The spe function of the DeepSeek-v4-flash-DSpark model fails to run #47418 reports a different failure (draft_id_to_target_id AttributeError) on H20.
• [Bugfix][DeepSeek V4] Avoid SM100 FlashMLA cache shape mismatch #47610 fixes the same kv must have shape error but only for SM10/B200 by switching the default to FlashInfer, which does not help SM90/H200 because FlashInfer MLA sparse DSV4 reports Unsupported architecture here.

Expected behavior

DeepSeek-V4-Flash-DSpark should start successfully with --spec-method dspark on H200/SM90, or at least have a documented working attention-backend combination.

Additional context

The draft config contains:

{
  "dspark_block_size": 5,
  "dspark_target_layer_ids": [40, 41, 42],
  "compress_ratios": [..., 4, 128, 4, 0, 0, 0]
}

The DSpark draft layers are created at indices num_hidden_layers + i (43, 44, 45) in DSparkDeepseekV4Model, while _remap_dspark_name maps checkpoint mtp.{i}.* to model.layers.{i}.* (0, 1, 2). This looks like a potential weight-loading mismatch in addition to the KV-cache shape issue.

[yy-fighting]

Update: root causes identified and fixed locally

I spent some time debugging the DeepSeek-V4-Flash-DSpark startup failure on H200/SM90 and got the model serving successfully after two small fixes. Below is a summary in case it helps narrow down the PR.

Environment

• vLLM: 0.23.1rc1.dev788+gfa4321de3 (vllm-dspark wheel)
• PyTorch: 2.11.0+cu129
• GPU: 8× NVIDIA H200 (SM90), TP=8
• Model: /data/models/DeepSeek-V4-Flash-DSpark
• Command: same as in the issue body, with --spec-method dspark --spec-model <same> --spec-tokens 5

Root cause 1: DSpark draft weight path mismatch

In vllm/models/deepseek_v4/nvidia/dspark.py, _remap_dspark_name was mapping checkpoint names mtp.{i}.* to model.layers.{num_hidden_layers + i}.*. The DeepseekV4DecoderLayer instances are created with that prefixed name (for compress_ratio handling), but the actual nn.ModuleList index is 0/1/2, so the parameter names registered in PyTorch are model.layers.{stage}.*. The loader therefore looked up weights under the wrong key and the draft layers ended up randomly initialized / not loaded.

Fix: map mtp.{i}.* block weights to model.layers.{i}.*:

def _remap_dspark_name(self, name: str) -> str | None:
    ...
    # Draft layers live in a ModuleList, so their actual parameter names are
    # model.layers.{stage}.* even though the prefix passed to the decoder
    # layer is layers.{num_hidden_layers + stage} (for compress_ratio).
    return f"model.layers.{stage}.{rest}"

File: vllm/models/deepseek_v4/nvidia/dspark.py:462-491

Root cause 2: KV-cache shape mismatch for fp8_ds_mla

DeepseekV4Attention.get_kv_cache_spec returned an MLAAttentionSpec without setting kv_quant_mode. Downstream in gpu_model_runner.py, the spec's cache_dtype_str was then passed as "auto" to DeepseekV4FlashMLABackend.get_kv_cache_shape, which returned (num_blocks, block_size, 512). For --kv-cache-dtype fp8 with the fp8_ds_mla layout, the correct per-token slot size is 584 bytes, so KV allocation was too small and triggered the shape mismatch / assert.

Fix: pass kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype) in the MLAAttentionSpec:

return MLAAttentionSpec(
    block_size=vllm_config.cache_config.block_size,
    num_kv_heads=1,
    head_size=self.head_dim,
    dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
    compress_ratio=self.compress_ratio,
    cache_dtype_str=self.kv_cache_dtype,
    alignment=576 if uses_fp8_ds_mla_layout else None,
    model_version="deepseek_v4",
    kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
)

File: vllm/models/deepseek_v4/attention.py:605-624

(The same two changes were also applied to the corresponding files under vllm-main.)

Verification

1. Qwen3 DSpark sanity check (/data/models/Qwen3-4B + deepseek-ai/dspark_qwen3_4b_block7) served and completed requests correctly.
2. DeepSeek-V4-Flash-DSpark now starts successfully with the command from the issue and responds to /v1/completions.

Serving benchmark (after fixes)

Run with the built-in vllm bench serve on ShareGPT, 50 prompts, max-concurrency=16, sharegpt-output-len=256:

Metric	Value
Request throughput	2.38 req/s
Output token throughput	610.47 tok/s
Acceptance rate	28.20%
Mean acceptance length	2.41
Mean TTFT	2162 ms
Mean TPOT	16.85 ms

Speculative decoding metrics from the server log during the run settled in the 25–36% acceptance-rate range, which is in line with the DSpark paper expectations for this model.

Next step

I can open a PR with the two fixes if the maintainers agree with the diagnosis. Let me know if you want me to prepare it.

[yy-fighting]

Closing this issue — both root causes have already been fixed in the latest vllm-main:

1. _remap_dspark_name in vllm/models/deepseek_v4/nvidia/dspark.py now maps mtp.{i}.* block weights to model.layers.{i}.*.
2. DeepseekV4Attention.get_kv_cache_spec in vllm/models/deepseek_v4/attention.py now passes kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype) to MLAAttentionSpec.

I verified these fixes locally on H200/SM90 with DeepSeek-V4-Flash-DSpark and the model serves successfully.

Thanks to the maintainers who already landed the fixes.

[yy-fighting]

Re-opening — I was too quick to close this.

After checking the actual upstream main branch:

• vllm/models/deepseek_v4/nvidia/dspark.py does have the _remap_dspark_name fix (model.layers.{stage}).
• vllm/models/deepseek_v4/attention.py does not yet have the kv_quant_mode fix in DeepseekV4Attention.get_kv_cache_spec.

The kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype) line is still missing upstream, so the KV cache shape mismatch for fp8_ds_mla remains unfixed in main. The nightly wheel I was testing from (0.23.1rc1.dev788+gfa4321de3) apparently includes that patch, which is why my local service started successfully.

I'll prepare a proper PR for the attention.py fix.

[sgsdxzy]

Could you please re-open the issue? I am still encountering the RuntimeError: kv must have shape (num_blocks, page_block_size, h_kv, bytes_per_token) on latest main (8974ed8) even after adding kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype). I think a recent commit broke it completely.

[flyerming]

RuntimeError: Worker failed with error 'kv must have shape (num_blocks, page_block_size, h_kv, bytes_per_token)', please check the stack trace above for the root cause
/usr/lib/python3.12/multiprocessing/resource_tracker.py:279: UserWarning: resource_tracker: There appear to be 8 leaked shared_memo

I tried adding kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype) with 0.23.1rc1.dev803+gd9c1767cd.

[yy-fighting]

I can confirm the root cause on the current nightly wheel (0.23.1rc1.dev799+g69715823d) with 8×H200.

The RuntimeError: kv must have shape (num_blocks, page_block_size, h_kv, bytes_per_token) is not just about the main MLA cache; it is because the SWA cache is also reshaped with the wrong shape.

DeepseekV4Attention.get_kv_cache_spec() already resolves --kv-cache-dtype fp8 to fp8_ds_mla and can return the correct 584-byte layout, but DeepseekV4SWACache.get_kv_cache_spec() in vllm/v1/attention/backends/mla/sparse_swa.py does not set kv_quant_mode. As a result, gpu_model_runner / attn_utils passes cache_dtype_str=\"auto\" to DeepseekSparseSWABackend.get_kv_cache_shape(), which falls back to the plain (num_blocks, block_size, 512) shape instead of the required (num_blocks, block_size, 584). FlashMLA then rejects the SWA cache tensor.

The following two minimal changes let the model pass the KV-cache warmup on H200/SM90:

1. vllm/models/deepseek_v4/attention.py — DeepseekV4Attention.get_kv_cache_spec():
   Add kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype) to the MLAAttentionSpec.

2. vllm/v1/attention/backends/mla/sparse_swa.py — DeepseekV4SWACache.get_kv_cache_spec():
   Import get_kv_quant_mode and add kv_quant_mode=get_kv_quant_mode(cache_dtype) to the SlidingWindowMLASpec.

After applying both changes I was able to start DeepSeek-V4-Flash-DSpark with --spec-method dspark on H200. The subsequent DeepGEMM warmup still fails with NVCC compilation failed for me, but that appears to be a separate issue; using VLLM_DEEP_GEMM_WARMUP=skip works around it and the server responds to /v1/completions.

I can open a PR with these two changes if the maintainers agree with the diagnosis.

[ChuanhongLi]

RuntimeError: Worker failed with error 'kv must have shape (num_blocks, page_block_size, h_kv, bytes_per_token)', please check the stack trace above for the root cause /usr/lib/python3.12/multiprocessing/resource_tracker.py:279: UserWarning: resource_tracker: There appear to be 8 leaked shared_memo

I tried adding kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype) with 0.23.1rc1.dev803+gd9c1767cd.

Same error!

[xfalcox]

Same issue on 2 x H200 without DSpark

2026-07-06T17:49:57.296161458Z (EngineCore pid=43257) RuntimeError: Worker failed with error 'kv must have shape (num_blocks, page_block_size, h_kv, bytes_per_token)', please check the stack trace above for the root cause