From af2456b11328a15ea7f2f8e48aeb7ed30a61e32f Mon Sep 17 00:00:00 2001 From: pavelzak Date: Sat, 15 Aug 2026 23:59:31 -0700 Subject: [PATCH 1/2] [Bugfix][Hardware][NVIDIA] Fix DSV4 sparse MLA spec-decode shapes on SM120 FlashInfer path FlashInfer's dsv4 sparse decode API disambiguates decode-vs-prefill by query rank. On SM120 (GB10) with speculative decoding (DSpark), the flattened 3-D [tokens, heads, 512] query for next_n > 1 batches is misrouted to the varlen prefill kernel, whose SM120 build asserts num_tokens > 64 ("Decode ... must go through sparse_mla_sm120_decode_dsv4"). Four fixes to DeepseekV4FlashInferSM120Attention: - Pass a 4-D [batch, next_n, heads, 512] query (and matching output view) for spec decode batches so FlashInfer routes them to the decode kernels. - Reshape the companion sparse indices/lens tensors to [batch, next_n, ...] alongside the 4-D query (FlashInfer validates their layout against the query). - Route <= 64-token prefill segments (the DSpark draft's k-token pass, short chunked-prefill tails) through per-request decode-form [1, q_len, ...] calls, since the SM120 sparse prefill kernel asserts num_tokens > 64. - Skip empty prefill chunks (zero-token spans in query_start_loc): the FlashInfer sparse kernel crashes reshaping 0 elements. Validated on 2x DGX Spark (GB10, SM121) TP=2 serving DeepSeek-V4-Flash-0731 with DSpark speculative decoding. Signed-off-by: pavelzak Co-Authored-By: Claude Fable 5 --- .../deepseek_v4/nvidia/flashinfer_sparse.py | 74 ++++++++++++++++++- 1 file changed, 73 insertions(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index e73156c1c121..1e5e00421acc 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -784,13 +784,38 @@ def _forward_decode( raise RuntimeError( "Compressed sparse MLA decode requires compressed sparse indices." ) + # flashinfer's dsv4 decode API expects query as + # [batch, q_len_per_request, heads, 512]. With speculative decoding + # (next_n = k+1 tokens per request) a flattened [tokens, heads, 512] + # is ambiguous: flashinfer's normalizer misroutes it to the varlen + # prefill kernel, whose SM120 build asserts num_tokens > 64 + # ("Decode ... must go through sparse_mla_sm120_decode_dsv4"). + out_arg = output + if num_decodes > 0 and num_decode_tokens > num_decodes: + assert num_decode_tokens % num_decodes == 0, ( + f"ragged spec decode batch: {num_decode_tokens} tokens over " + f"{num_decodes} requests" + ) + next_n = num_decode_tokens // num_decodes + q = q.view(num_decodes, next_n, *q.shape[1:]) + out_arg = output.view(num_decodes, next_n, *output.shape[1:]) + # Companion tensors must match the [batch, next_n, ...] layout the + # 4-D query implies (flashinfer validates indices against it). + swa_indices = swa_indices.reshape(num_decodes, next_n, -1) + swa_lens = swa_lens.reshape(num_decodes, next_n) + if extra_sparse_indices is not None: + extra_sparse_indices = extra_sparse_indices.reshape( + num_decodes, next_n, -1 + ) + if extra_sparse_lengths is not None: + extra_sparse_lengths = extra_sparse_lengths.reshape(num_decodes, next_n) flashinfer_trtllm_batch_decode_sparse_mla_dsv4( query=q, swa_kv_cache=swa_cache, workspace_buffer=self._get_workspace(q.device), sparse_indices=swa_indices, compressed_kv_cache=extra_cache, - out=output, + out=out_arg, bmm1_scale=self.scale, sinks=self.attn_sink, kv_layout="NHD", @@ -898,12 +923,59 @@ def _forward_prefill( ) q_chunk = q[query_start:query_end] + if q_chunk.shape[0] == 0: + # Empty chunk (zero-token span in query_start_loc): the + # flashinfer sparse kernel crashes reshaping 0 elements. + continue swa_indices_chunk = swa_metadata.prefill_swa_indices[query_start:query_end] swa_lens_chunk = swa_metadata.prefill_swa_lens[query_start:query_end] if extra_kv_paged is not None and extra_sparse_indices_chunk is None: raise RuntimeError( "Compressed sparse MLA prefill requires compressed sparse indices." ) + if q_chunk.shape[0] <= 64: + # SM120's sparse prefill kernel asserts num_tokens > 64. + # Small segments (the DSpark draft's k-token pass, short + # chunked-prefill tails) must use the uniform decode-form + # [1, q_len, ...] call, which flashinfer routes to its + # sparse_mla_sm120_decode_dsv4 kernels (q_len <= 64 legal). + for ri in range(chunk_start, chunk_end): + rs = int(query_start_loc_cpu[num_decodes + ri] - prefill_token_base) + re_ = int( + query_start_loc_cpu[num_decodes + ri + 1] - prefill_token_base + ) + if re_ <= rs: + continue + ql = re_ - rs + esi = ( + extra_sparse_indices[rs:re_].reshape(1, ql, -1) + if extra_sparse_indices is not None + else None + ) + esl = ( + extra_sparse_lengths[rs:re_].reshape(1, ql) + if extra_sparse_lengths is not None + else None + ) + flashinfer_trtllm_batch_decode_sparse_mla_dsv4( + query=q[rs:re_].reshape(1, ql, *q.shape[1:]), + swa_kv_cache=swa_kv_paged, + workspace_buffer=self._get_workspace(q.device), + sparse_indices=swa_metadata.prefill_swa_indices[rs:re_].reshape( + 1, ql, -1 + ), + compressed_kv_cache=extra_kv_paged, + out=output[rs:re_].reshape(1, ql, *output.shape[1:]), + bmm1_scale=self.scale, + sinks=self.attn_sink, + kv_layout="NHD", + swa_topk_lens=swa_metadata.prefill_swa_lens[rs:re_].reshape( + 1, ql + ), + extra_sparse_indices=esi, + extra_sparse_topk_lens=esl, + ) + continue flashinfer_trtllm_batch_decode_sparse_mla_dsv4( query=q_chunk, swa_kv_cache=swa_kv_paged, From 701d830fa121053aef3a5281b73c7397e14eafbf Mon Sep 17 00:00:00 2001 From: pavelzak Date: Wed, 2 Sep 2026 21:10:38 -0700 Subject: [PATCH 2/2] [Bugfix] force contiguity on reshaped spec-decode index/length tensors Reported by @kitch2400 on #52499: under CUDA graph capture with C128A layers, c128a_global_decode_topk_indices can be a slice of a larger, alignment-padded workspace buffer kept fixed-size for graph safety. reshape() onto that slice can produce a tensor whose shape matches but whose layout doesn't satisfy flashinfer's CHECK_INPUT_AND_TYPE contiguity check (eidx.IsContiguous() at sparse_mla_sm120.cu), because reshape() only guarantees a valid *view* when the source strides permit one - it doesn't force copy-to-contiguous in every non-trivial case. .contiguous() makes the copy explicit and unconditional. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01C7YveGPhnT8v6qrAEucmVW Signed-off-by: pavelzak --- .../deepseek_v4/nvidia/flashinfer_sparse.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 1e5e00421acc..11e53d45f4ad 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -801,14 +801,22 @@ def _forward_decode( out_arg = output.view(num_decodes, next_n, *output.shape[1:]) # Companion tensors must match the [batch, next_n, ...] layout the # 4-D query implies (flashinfer validates indices against it). - swa_indices = swa_indices.reshape(num_decodes, next_n, -1) - swa_lens = swa_lens.reshape(num_decodes, next_n) + # .contiguous() is required, not cosmetic: c128a_global_decode_ + # topk_indices can be a slice of a larger, alignment-padded + # workspace buffer (kept fixed-size for CUDA graph capture), so + # its reshape()-compatible view may not satisfy flashinfer's + # CHECK_INPUT_AND_TYPE contiguity check on the C++ side even + # though the shape matches. + swa_indices = swa_indices.reshape(num_decodes, next_n, -1).contiguous() + swa_lens = swa_lens.reshape(num_decodes, next_n).contiguous() if extra_sparse_indices is not None: extra_sparse_indices = extra_sparse_indices.reshape( num_decodes, next_n, -1 - ) + ).contiguous() if extra_sparse_lengths is not None: - extra_sparse_lengths = extra_sparse_lengths.reshape(num_decodes, next_n) + extra_sparse_lengths = extra_sparse_lengths.reshape( + num_decodes, next_n + ).contiguous() flashinfer_trtllm_batch_decode_sparse_mla_dsv4( query=q, swa_kv_cache=swa_cache,