@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 88aa0d287d

[Spec decode] Support variable-length decode for Kimi-K3 adaptive ver (#52988)

Signed-off-by: Albert Cheng <albecheng@nvidia.com> Signed-off-by: Albert Cheng (Engrg-Hardware 1) <albecheng@login-bia01.bia.clusters.nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>

Albert Cheng committed Sep 22, 2026 at 20:43 UTC 88aa0d287dd3abac89741bbea349350a4d49194e

13 files changed +203 -22

tests/v1/attention/test_flashinfer_mla_dcp.py

+10

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +from types import MethodType

5 from unittest.mock import MagicMock

7 import pytest

@@ -61,6 +62,7 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):

62 attn_metadata.max_seq_len = 1

63 attn_metadata.decode.block_table = torch.zeros(2, 1, dtype=torch.int32)

64 attn_metadata.decode.seq_lens = torch.ones(2, dtype=torch.int32)

65 + attn_metadata.decode.max_query_len = 1

67 query = torch.ones(2, 24, 576, dtype=torch.bfloat16)

68 kv_cache = torch.ones(1, 128, 576, dtype=torch.bfloat16)

@@ -100,6 +102,9 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):

102 impl.dcp_world_size = 8

103 impl.dcp_rank = 3

104 impl.cp_kv_cache_interleave_size = 1

105 + impl._flattened_decode_metadata = MethodType(

106 + flashinfer_mla.FlashInferMLAImpl._flattened_decode_metadata, impl

107 + )

108

109 num_reqs, query_len = 2, 3

110 num_tokens = num_reqs * query_len

@@ -114,6 +119,11 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal):

119 attn_metadata.decode.block_table = block_table

120 attn_metadata.decode.seq_lens = seq_lens

121 attn_metadata.decode.dcp_tot_seq_lens = global_causal_seq_lens

122 + attn_metadata.decode.max_query_len = query_len

123 + attn_metadata.decode.query_start_loc = torch.arange(

124 + 0, num_tokens + 1, query_len, dtype=torch.int32

125 + )

126 + attn_metadata.decode.query_len = 0

127

128 query = torch.ones(num_tokens, 24, 576, dtype=torch.bfloat16)

129 kv_cache = torch.ones(2, 128, 576, dtype=torch.bfloat16)

tests/v1/attention/test_mla_backends.py

@@ -1267,6 +1267,7 @@ def test_flashmla_dcp_decode_metadata_uses_gathered_query_heads(

1267 query_start_loc_cpu=query_start_loc,

1268 query_start_loc_device=query_start_loc,

1269 num_decode_tokens=2,

1270 + max_query_len=1,

1271 dcp_tot_seq_lens_device=None,

1272 )

1273

tests/v1/attention/test_rocm_aiter_mla_mtp_split.py

@@ -278,6 +278,7 @@ def test_dcp_fp8_verify_build_uses_segmented(monkeypatch):

278 query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),

279 query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),

280 num_decode_tokens=2 * qlen,

281 + max_query_len=qlen,

282 dcp_tot_seq_lens_device=torch.tensor([10, 12], dtype=torch.int32),

283 )

284

@@ -633,6 +634,7 @@ def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch):

634 query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32),

635 query_start_loc_device=torch.tensor([0, 4, 8], dtype=torch.int32),

636 num_decode_tokens=8,

637 + max_query_len=4,

638 dcp_tot_seq_lens_device=None,

639 )

640

@@ -683,6 +685,7 @@ def test_min_kv_seq_len_ignores_cudagraph_padding_rows(monkeypatch):

685 query_start_loc_cpu=query_start_loc,

686 query_start_loc_device=query_start_loc,

687 num_decode_tokens=num_reqs * mtp_qlen,

688 + max_query_len=mtp_qlen,

689 dcp_tot_seq_lens_device=None,

690 )

691

@@ -733,6 +736,7 @@ def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr(

736 query_start_loc_cpu=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32),

737 query_start_loc_device=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32),

738 num_decode_tokens=seq_lens.numel() * mtp_qlen,

739 + max_query_len=mtp_qlen,

740 dcp_tot_seq_lens_device=None,

741 )

742

@@ -801,6 +805,7 @@ def test_decode_expands_kernel_block_page_indices(monkeypatch):

805 query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.int32),

806 query_start_loc_device=torch.tensor([0, 1, 2], dtype=torch.int32),

807 num_decode_tokens=seq_lens.numel(),

808 + max_query_len=1,

809 dcp_tot_seq_lens_device=None,

810 )

811

@@ -892,6 +897,7 @@ def test_persistent_metadata_gate(

897 query_start_loc_cpu=query_start_loc,

898 query_start_loc_device=query_start_loc,

899 num_decode_tokens=num_reqs * qo_len,

900 + max_query_len=qo_len,

901 dcp_tot_seq_lens_device=None,

902 )

903

@@ -961,6 +967,7 @@ def test_persistent_metadata_gate_without_gluon_build(

967 query_start_loc_cpu=query_start_loc,

968 query_start_loc_device=query_start_loc,

969 num_decode_tokens=num_reqs * qo_len,

970 + max_query_len=qo_len,

971 dcp_tot_seq_lens_device=None,

972 )

973

@@ -993,6 +1000,7 @@ def _build_non_causal(monkeypatch, *, num_heads, kv_cache_dtype, qlen, mtp_qlen)

1000 query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),

1001 query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),

1002 num_decode_tokens=2 * qlen,

1003 + max_query_len=qlen,

1004 dcp_tot_seq_lens_device=None,

1005 )

1006 return metadata, get_mla_metadata_v1

tests/v1/worker/test_mamba_hybrid_model_state.py

@@ -34,6 +34,7 @@ def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> Non

34 query_start_loc_np=torch.tensor([0, 1], dtype=torch.int32).numpy(),

35 query_start_loc=torch.tensor([0, 1], dtype=torch.int32),

36 num_scheduled_tokens=torch.tensor([1], dtype=torch.int32),

37 + max_query_len=None,

38 seq_lens_cpu_upper_bound=torch.tensor([1537], dtype=torch.int32),

39 seq_lens=torch.tensor([1537], dtype=torch.int32),

40 is_prefilling_np=torch.tensor([False]).numpy(),

vllm/model_executor/layers/attention/mla_attention.py

@@ -2471,6 +2471,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

2471 query_start_loc_cpu: torch.Tensor,

2472 query_start_loc_device: torch.Tensor,

2473 num_decode_tokens: int,

2474 + max_query_len: int,

2475 dcp_tot_seq_lens_device: torch.Tensor | None,

2476 ) -> MLACommonDecodeMetadata:

2477 return MLACommonDecodeMetadata(

@@ -2635,6 +2636,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

2636 query_start_loc_cpu=query_start_loc_cpu[: num_decodes + 1],

2637 query_start_loc_device=query_start_loc[: num_decodes + 1],

2638 num_decode_tokens=num_decode_tokens,

2639 + max_query_len=min(

2640 + common_attn_metadata.max_query_len, self.reorder_batch_threshold

2641 + ),

2642 dcp_tot_seq_lens_device=dcp_tot_seq_lens_device,

2643 )

2644

vllm/models/dots3_note/nvidia/attention.py

+2 -1

@@ -353,9 +353,10 @@ class Dots3NoteMLAMetadataBuilder(TritonMLAMetadataBuilder):

353 query_start_loc_cpu: torch.Tensor,

354 query_start_loc_device: torch.Tensor,

355 num_decode_tokens: int,

356 + max_query_len: int,

357 dcp_tot_seq_lens_device: torch.Tensor | None,

358 ) -> Dots3NoteDecodeMetadata:

358 - del max_seq_len, query_start_loc_device, num_decode_tokens

359 + del max_seq_len, query_start_loc_device, num_decode_tokens, max_query_len

360 query_len = int(query_start_loc_cpu[1] - query_start_loc_cpu[0])

361 return Dots3NoteDecodeMetadata(

362 block_table=block_table_tensor,

vllm/models/kimi_k3/nvidia/dspark_mla.py

+29 -1

@@ -15,7 +15,10 @@ from vllm.model_executor.layers.linear import (

15 ReplicatedLinear,

16 )

17 from vllm.model_executor.layers.logits_processor import LogitsProcessor

18 -from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead

18 +from vllm.model_executor.models.qwen3_dspark import (

19 + DSparkConfidenceHead,

20 + DSparkMarkovHead,

21 +)

22 from vllm.model_executor.models.utils import (

23 AutoWeightsLoader,

24 WeightsMapper,

@@ -188,6 +191,18 @@ class K3DSparkModel(nn.Module):

191 self.config.markov_rank,

192 prefix=maybe_prefix(prefix, "markov_head"),

193 )

194 + self.confidence_head: DSparkConfidenceHead | None = None

195 + if getattr(self.config, "enable_confidence_head", False):

196 + with_markov = getattr(self.config, "confidence_head_with_markov", False)

197 + input_dim = self.config.hidden_size + (

198 + self.config.markov_rank if with_markov else 0

199 + )

200 + self.confidence_head = DSparkConfidenceHead(

201 + input_dim,

202 + prefix=maybe_prefix(prefix, "confidence_head"),

203 + bias=True,

204 + with_markov=with_markov,

205 + )

206 self._max_num_context_tokens = (

207 vllm_config.scheduler_config.max_num_batched_tokens

208 )

@@ -436,6 +451,12 @@ class K3DSparkForCausalLM(nn.Module):

451 assert vllm_config.speculative_config is not None

452 self.draft_model_config = vllm_config.speculative_config.draft_model_config

453 self.config = self.draft_model_config.hf_config

454 + if getattr(self.config, "enable_confidence_head", False):

455 + # The class mapper drops the head as training-only; keep it for drafts

456 + # that actually built one, since adaptive verification reads it.

457 + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper(

458 + orig_to_new_substr={"confidence_head": "confidence_head"}

459 + )

460 target_layer_num = vllm_config.model_config.get_total_num_hidden_layers()

461 self.model = K3DSparkModel(

462 vllm_config=vllm_config,

@@ -494,6 +515,13 @@ class K3DSparkForCausalLM(nn.Module):

515 def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:

516 return self.model.markov_head.bias(markov_embed, self.logits_processor)

517

518 + def compute_confidence(

519 + self, head_hidden: torch.Tensor, markov_embed: torch.Tensor

520 + ) -> torch.Tensor:

521 + """Per-position acceptance probability for each drafted token."""

522 + assert self.model.confidence_head is not None

523 + return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed))

524 +

525 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

526 loader = AutoWeightsLoader(self)

527 # read: 1. all weights. 2. context kv weights

vllm/models/kimi_k3/nvidia/kda_metadata.py

+17 -1

@@ -25,7 +25,7 @@ from vllm.model_executor.layers.mamba.checkpoint import (

25 from vllm.platforms import current_platform

26 from vllm.triton_utils import tl, triton

27 from vllm.utils.torch_utils import async_tensor_h2d

28 -from vllm.v1.attention.backend import CommonAttentionMetadata

28 +from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata

29 from vllm.v1.attention.backends.gdn_attn import (

30 GDNAttentionBackend,

31 GDNAttentionMetadata,

@@ -305,6 +305,10 @@ class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):

305

306

307 class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

308 + # Overrides GDN's UNIFORM_BATCH: adaptive verification requires ALWAYS from every

309 + # builder, and KDA reads per-request offsets off device within a fixed k+1 window,

310 + # so one k+1 graph replays any 1..k+1 mix.

311 + _cudagraph_support = AttentionCGSupport.ALWAYS

312 mamba_aligned_state_indices: torch.Tensor | None = None

313

314 def __init__(

@@ -407,6 +411,12 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

411 query_start_loc_cpu.diff() > 0

412 )

413 spec_sequence_masks_cpu |= active_decode_mask_cpu

414 + # Adaptive's cost-table profiler probes eager dummy rows longer than

415 + # RecoverSSM's num_spec+1 activation capacity; route them out instead

416 + # of tripping the capacity check below.

417 + spec_sequence_masks_cpu &= (

418 + query_start_loc_cpu.diff() <= self.num_spec + 1

419 + )

420 # Native KDA can use its regular decode path when no draft token

421 # was scheduled. RecoverSSM must preserve its extended conv window.

422 if (

@@ -750,3 +760,9 @@ class KimiK3KDAAttentionBackend(GDNAttentionBackend):

760 @staticmethod

761 def get_builder_cls() -> type[KimiK3KDAMetadataBuilder]:

762 return KimiK3KDAMetadataBuilder

763 +

764 + @classmethod

765 + def supports_device_cpu_query_lens_mismatch(cls) -> bool:

766 + # Pure spec-decode reads its plan from DEVICE offsets, so CPU/device

767 + # mismatch is fine; the mixed branch still assumes CPU totals match.

768 + return True

vllm/v1/attention/backends/mla/flashattn_mla.py

@@ -194,8 +194,11 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]

194 query_start_loc_cpu: torch.Tensor,

195 query_start_loc_device: torch.Tensor,

196 num_decode_tokens: int,

197 + max_query_len: int,

198 dcp_tot_seq_lens_device: torch.Tensor | None,

199 ) -> FlashAttnMLADecodeMetadata:

200 + # Measured, not the promised bound: _schedule_decode asserts metadata size

201 + # against a persistent buffer.

202 query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

203 max_query_len = query_lens_cpu.max().item()

204

vllm/v1/attention/backends/mla/flashinfer_mla.py

+116 -14

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +from dataclasses import dataclass

5 from typing import TYPE_CHECKING, Any, ClassVar

7 import torch

@@ -15,6 +16,7 @@ from vllm.config.cache import CacheDType

16 from vllm.logger import init_logger

17 from vllm.model_executor.layers.attention.mla_attention import (

18 MLACommonBackend,

19 + MLACommonDecodeMetadata,

20 MLACommonImpl,

21 MLACommonMetadata,

22 MLACommonMetadataBuilder,

@@ -115,9 +117,28 @@ def _get_multi_ctas_kv_counter_buffer(

117 return _fi_multi_ctas_kv_counter

118

119

118 -class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):

119 - _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

120 - query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM

120 +@dataclass

121 +class FlashInferMLADecodeMetadata(MLACommonDecodeMetadata):

122 + flattened_block_table: torch.Tensor | None = None

123 + flattened_seq_lens: torch.Tensor | None = None

124 + flattened_row_req: torch.Tensor | None = None

125 + # Row count the flattened tensors were built for (0 = not built), not a query len.

126 + query_len: int = 0

127 + query_start_loc: torch.Tensor | None = None

128 + max_query_len: int = 1

129 +

130 +

131 +@dataclass

132 +class FlashInferMLAMetadata(MLACommonMetadata[FlashInferMLADecodeMetadata]):

133 + pass

134 +

135 +

136 +class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[FlashInferMLAMetadata]):

137 + # Adaptive verification requires ALWAYS from every builder, matching upstream's

138 + # DeepseekV4FlashMLAMetadataBuilder. The kernels tile ragged queries from the device

139 + # query offsets (flashinfer #3238), so one k+1 graph replays any 1..k+1 mix.

140 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS

141 + query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.VARLEN

142 # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa.

143 supports_non_causal_multi_token_decode: ClassVar[bool] = True

144 supports_non_causal_multi_token_dcp: ClassVar[bool] = True

@@ -143,10 +164,31 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):

164 layer_names,

165 vllm_config,

166 device,

146 - MLACommonMetadata,

167 + FlashInferMLAMetadata,

168 supports_dcp_with_varlen=True,

169 )

170

171 + def _build_decode(

172 + self,

173 + block_table_tensor: torch.Tensor,

174 + seq_lens_device: torch.Tensor,

175 + max_seq_len: int,

176 + query_start_loc_cpu: torch.Tensor,

177 + query_start_loc_device: torch.Tensor,

178 + num_decode_tokens: int,

179 + max_query_len: int,

180 + dcp_tot_seq_lens_device: torch.Tensor | None,

181 + ) -> FlashInferMLADecodeMetadata:

182 + # Promised bound, not measured: a capture dummy puts one token/req, so

183 + # measuring would bake an undersized graph. Ragged only when max_query_len>1.

184 + return FlashInferMLADecodeMetadata(

185 + block_table=block_table_tensor,

186 + seq_lens=seq_lens_device,

187 + dcp_tot_seq_lens=dcp_tot_seq_lens_device,

188 + query_start_loc=query_start_loc_device,

189 + max_query_len=max_query_len,

190 + )

191 +

192

193 class FlashInferMLABackend(MLACommonBackend):

194 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

@@ -208,7 +250,7 @@ class FlashInferMLABackend(MLACommonBackend):

250 return None

251

252

211 -class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

253 +class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]):

254 can_return_lse_for_decode: bool = True

255 supports_dcp: bool = True

256 # DCP is the only path that consumes LSE. It uses monolithic CuTeDSL,

@@ -276,7 +318,7 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

318 self,

319 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

320 kv_c_and_k_pe_cache: torch.Tensor,

279 - attn_metadata: MLACommonMetadata,

321 + attn_metadata: FlashInferMLAMetadata,

322 layer: AttentionLayer,

323 ) -> tuple[torch.Tensor, torch.Tensor | None]:

324 assert kv_c_and_k_pe_cache.numel() > 0

@@ -289,14 +331,32 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

331 block_table = attn_metadata.decode.block_table

332 seq_lens = attn_metadata.decode.seq_lens

333

334 + # Led by the promised bound (baked into the graph), not

335 + # num_decode_tokens // num_decodes -- that average is a per-request length

336 + # only for uniform batches.

337 + multi_token_decode = (

338 + attn_metadata.decode.max_query_len > 1

339 + or attn_metadata.num_decode_tokens > attn_metadata.num_decodes

340 + )

341 + cum_seq_lens_q: torch.Tensor | None = None

342 + max_q_len: int | None = None

343 + row_req: torch.Tensor | None = None

344 +

345 if not attn_metadata.causal:

346 # FlashInfer decode has no causal flag. Flatten each non-causal

347 # query block into independent single-token rows.

295 - query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes

348 q = q.unsqueeze(1)

297 - if query_len > 1:

298 - block_table = block_table.repeat_interleave(query_len, dim=0)

299 - seq_lens = seq_lens.repeat_interleave(query_len)

349 + if multi_token_decode:

350 + block_table, seq_lens, row_req = self._flattened_decode_metadata(

351 + attn_metadata, q.shape[0]

352 + )

353 + elif attn_metadata.decode.max_query_len > 1 and self.dcp_world_size == 1:

354 + # Causal spec decode: keep q compact and let the kernel tile each

355 + # request's length (uniform 1+k and adaptive ragged). Mirrors vllm #52157.

356 + # DCP keeps the uniform reshape below: it needs LSE, which the kernels

357 + # do not return on the ragged path (flashinfer #3238).

358 + cum_seq_lens_q = attn_metadata.decode.query_start_loc

359 + max_q_len = attn_metadata.decode.max_query_len

360 # trtllm API requires extra dimension q_len_per_request for MTP

361 elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:

362 logger.warning_once(

@@ -327,10 +387,8 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

387 if self.dcp_world_size > 1:

388 causal_seqlens_kv_global = attn_metadata.decode.dcp_tot_seq_lens

389 assert causal_seqlens_kv_global is not None

330 - if not attn_metadata.causal and query_len > 1:

331 - causal_seqlens_kv_global = causal_seqlens_kv_global.repeat_interleave(

332 - query_len

333 - )

390 + if row_req is not None:

391 + causal_seqlens_kv_global = causal_seqlens_kv_global[row_req]

392 extra_kwargs.update(

393 enable_dcp=True,

394 cp_world=self.dcp_world_size,

@@ -342,6 +400,16 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

400 # trtllm-gen rejects MLA head counts it can't tile (e.g. 96);

401 # fall back to cute-dsl for those.

402 decode_backend = _select_mla_decode_backend(runtime_num_heads)

403 + if cum_seq_lens_q is not None:

404 + # Neither decode backend returns LSE on the ragged path

405 + # (flashinfer #3238); DCP, the only LSE consumer, took the uniform

406 + # branch above, so this only guards a future caller wiring the two.

407 + assert not return_lse, (

408 + "FlashInferMLA ragged decode cannot return LSE; DCP and adaptive "

409 + "variable-length decode are mutually exclusive."

410 + )

411 + extra_kwargs["cum_seq_lens_q"] = cum_seq_lens_q

412 + extra_kwargs["max_q_len"] = max_q_len

413 if decode_backend:

414 extra_kwargs["backend"] = decode_backend

415 elif kv_c_and_k_pe_cache.shape[-2] in (32, 64):

@@ -386,3 +454,37 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

454 o = o.view(-1, o.shape[-2], o.shape[-1])

455

456 return o, lse

457 +

458 + def _flattened_decode_metadata(

459 + self,

460 + attn_metadata: FlashInferMLAMetadata,

461 + num_rows: int,

462 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

463 + """Expand per-request decode tensors to one row per query token.

464 +

465 + Cached across the layers of a group, which all see the same batch.

466 + """

467 + decode = attn_metadata.decode

468 + assert decode is not None

469 + if decode.query_len != num_rows:

470 + cu = decode.query_start_loc

471 + assert cu is not None

472 + # searchsorted on the device offsets, not repeat_interleave(uniform_len):

473 + # the latter only lines up for uniform batches and hands ragged rows

474 + # another request's KV (silent garbage).

475 + rows = torch.arange(num_rows, device=cu.device, dtype=cu.dtype)

476 + row_req = torch.searchsorted(cu[1:], rows, right=True).clamp_(

477 + max=decode.block_table.shape[0] - 1

478 + )

479 + decode.flattened_row_req = row_req

480 + decode.flattened_block_table = decode.block_table[row_req]

481 + decode.flattened_seq_lens = decode.seq_lens[row_req]

482 + decode.query_len = num_rows

483 + assert decode.flattened_block_table is not None

484 + assert decode.flattened_seq_lens is not None

485 + assert decode.flattened_row_req is not None

486 + return (

487 + decode.flattened_block_table,

488 + decode.flattened_seq_lens,

489 + decode.flattened_row_req,

490 + )

vllm/v1/attention/backends/mla/flashmla.py

@@ -158,6 +158,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

158 query_start_loc_cpu: torch.Tensor,

159 query_start_loc_device: torch.Tensor,

160 num_decode_tokens: int,

161 + max_query_len: int,

162 dcp_tot_seq_lens_device: torch.Tensor | None,

163 ) -> FlashMLADecodeMetadata:

164 query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

@@ -1007,6 +1007,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

1007 query_start_loc_cpu: torch.Tensor,

1008 query_start_loc_device: torch.Tensor,

1009 num_decode_tokens: int,

1010 + max_query_len: int,

1011 dcp_tot_seq_lens_device: torch.Tensor | None,

1012 ) -> AiterMLADecodeMetadata:

1013 causal = self._decode_causal

vllm/v1/worker/gpu/model_states/mamba_hybrid.py

+10 -5

@@ -250,7 +250,11 @@ class MambaHybridModelState(DefaultModelState):

250 num_reqs = input_batch.num_reqs

251 num_tokens = input_batch.num_tokens

252 query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)

253 - max_query_len = input_batch.num_scheduled_tokens.max().item()

253 + # Prefer the promised bound: a capture dummy's measured max is its even

254 + # split, not the length the graph must replay.

255 + max_query_len = input_batch.max_query_len

256 + if max_query_len is None:

257 + max_query_len = input_batch.num_scheduled_tokens.max().item()

258 seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound

259 if for_capture:

260 # Capture with worst-case max_seq_len so the graph is valid at any replay.

@@ -278,10 +282,11 @@ class MambaHybridModelState(DefaultModelState):

282 num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32)

283 num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req

284 if num_draft_tokens_per_req is not None:

281 - # A row is a spec-decode row only when its whole prompt is already

282 - # computed, i.e. exactly one non-draft (decode) token is scheduled.

283 - is_decode = (

284 - input_batch.num_scheduled_tokens == num_draft_tokens_per_req + 1

285 + # Test request state, not num_scheduled_tokens == draft_count+1:

286 + # adaptive rewrites num_scheduled_tokens to an even split, so that

287 + # equality rarely holds and would demote every verify row to decode.

288 + is_decode = (~input_batch.is_prefilling_np) & (

289 + input_batch.num_scheduled_tokens > 0

290 )

291 spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode

292 num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where(