@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 82af928c41

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Attention][Spec Decode] FlashMLA spec decode support (#26541)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Oct 14, 2025 at 15:38 UTC 82af928c41886a9e7e56467a08c9b99982f7c980

5 files changed +214 -91

tests/v1/attention/test_mla_backends.py

+156 -71

@@ -1,6 +1,12 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -"""Tests for v1 MLA backends without GPUModelRunner dependency."""

3 +"""Tests for v1 MLA backends without GPUModelRunner dependency.

4 +

5 +Known Issues:

6 +- FLASH_ATTN_MLA backend occasionally produces NaN values in

7 + test_backend_correctness[mixed_small] when run after

8 + test_backend_correctness[small_prefill], but passes when run alone.

9 +"""

11 import pytest

12 import torch

@@ -14,6 +20,8 @@ from tests.v1.attention.utils import (

20 )

21 from vllm import _custom_ops as ops

22 from vllm.attention.backends.registry import _Backend

23 +from vllm.attention.ops.flashmla import is_flashmla_dense_supported

24 +from vllm.config.vllm import set_current_vllm_config

25 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv

26 from vllm.v1.attention.backends.utils import CommonAttentionMetadata

27 from vllm.v1.kv_cache_interface import FullAttentionSpec

@@ -29,6 +37,10 @@ BACKENDS_TO_TEST = [

37 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:

38 BACKENDS_TO_TEST.remove(_Backend.CUTLASS_MLA)

40 +# Remove FLASHMLA from the list if not supported

41 +if not is_flashmla_dense_supported()[0]:

42 + BACKENDS_TO_TEST.remove(_Backend.FLASHMLA)

43 +

44 torch.manual_seed(42)

@@ -66,6 +78,12 @@ BATCH_SPECS = {

78 "large_prefill": BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),

79 "single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),

80 "single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),

81 + "spec_decode_small": BatchSpec(

82 + seq_lens=[128, 256, 512, 1024], query_lens=[4, 4, 4, 4]

83 + ),

84 + "spec_decode_medium": BatchSpec(

85 + seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[8, 8, 8, 8, 8, 8]

86 + ),

87 }

@@ -239,61 +257,64 @@ def run_attention_backend(

257

258 builder_cls, impl_cls = try_get_attention_backend(backend)

259

242 - # Build metadata

243 - builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)

244 - attn_metadata = builder.build(

245 - common_prefix_len=0,

246 - common_attn_metadata=common_attn_metadata,

247 - )

260 + # Set the current vllm config so that get_current_vllm_config() works

261 + # in the backend implementations

262 + with set_current_vllm_config(vllm_config):

263 + # Build metadata

264 + builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)

265 + attn_metadata = builder.build(

266 + common_prefix_len=0,

267 + common_attn_metadata=common_attn_metadata,

268 + )

269

249 - # Instantiate MLA implementation

250 - num_heads = vllm_config.model_config.get_num_attention_heads(

251 - vllm_config.parallel_config

252 - )

253 - num_kv_heads = vllm_config.model_config.get_num_kv_heads(

254 - vllm_config.parallel_config

255 - )

256 - head_size = vllm_config.model_config.get_head_size()

257 - scale = 1.0 / (head_size**0.5)

258 - impl = impl_cls(

259 - num_heads=num_heads,

260 - head_size=head_size,

261 - scale=scale,

262 - num_kv_heads=num_kv_heads,

263 - alibi_slopes=None,

264 - sliding_window=None,

265 - kv_cache_dtype="auto",

266 - logits_soft_cap=None,

267 - attn_type="decoder",

268 - kv_sharing_target_layer_name=None,

269 - q_lora_rank=None,

270 - kv_lora_rank=kv_lora_rank,

271 - qk_nope_head_dim=qk_nope_head_dim,

272 - qk_rope_head_dim=qk_rope_head_dim,

273 - qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,

274 - v_head_dim=v_head_dim,

275 - kv_b_proj=mock_kv_b_proj,

276 - )

270 + # Instantiate MLA implementation

271 + num_heads = vllm_config.model_config.get_num_attention_heads(

272 + vllm_config.parallel_config

273 + )

274 + num_kv_heads = vllm_config.model_config.get_num_kv_heads(

275 + vllm_config.parallel_config

276 + )

277 + head_size = vllm_config.model_config.get_head_size()

278 + scale = 1.0 / (head_size**0.5)

279 + impl = impl_cls(

280 + num_heads=num_heads,

281 + head_size=head_size,

282 + scale=scale,

283 + num_kv_heads=num_kv_heads,

284 + alibi_slopes=None,

285 + sliding_window=None,

286 + kv_cache_dtype="auto",

287 + logits_soft_cap=None,

288 + attn_type="decoder",

289 + kv_sharing_target_layer_name=None,

290 + q_lora_rank=None,

291 + kv_lora_rank=kv_lora_rank,

292 + qk_nope_head_dim=qk_nope_head_dim,

293 + qk_rope_head_dim=qk_rope_head_dim,

294 + qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,

295 + v_head_dim=v_head_dim,

296 + kv_b_proj=mock_kv_b_proj,

297 + )

298

278 - # Process weights to create W_UK_T and W_UV attributes needed by MLA

279 - act_dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)

280 - impl.process_weights_after_loading(act_dtype)

299 + # Process weights to create W_UK_T and W_UV attributes needed by MLA

300 + act_dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)

301 + impl.process_weights_after_loading(act_dtype)

302

282 - # Create mock layer and output buffer

283 - mock_layer = MockAttentionLayer(device)

284 - num_tokens = query.shape[0]

285 - output = torch.empty(

286 - num_tokens, num_heads * v_head_dim, dtype=query.dtype, device=query.device

287 - )

303 + # Create mock layer and output buffer

304 + mock_layer = MockAttentionLayer(device)

305 + num_tokens = query.shape[0]

306 + output = torch.empty(

307 + num_tokens, num_heads * v_head_dim, dtype=query.dtype, device=query.device

308 + )

309

289 - # Run forward pass

290 - # NOTE: The query, key, and value are already shaped correctly

291 - # in the calling test function.

292 - output = impl.forward(

293 - mock_layer, query, kv_c, k_pe, kv_cache, attn_metadata, output=output

294 - )

310 + # Run forward pass

311 + # NOTE: The query, key, and value are already shaped correctly

312 + # in the calling test function.

313 + output = impl.forward(

314 + mock_layer, query, kv_c, k_pe, kv_cache, attn_metadata, output=output

315 + )

316

296 - return output

317 + return output

318

319

320 @pytest.mark.parametrize(

@@ -309,6 +330,8 @@ def run_attention_backend(

330 "large_prefill",

331 "single_decode",

332 "single_prefill",

333 + "spec_decode_small",

334 + "spec_decode_medium",

335 ],

336 )

337 @pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-V2-Lite-Chat"])

@@ -328,10 +351,39 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):

351 simulated paged KV cache.

352 5. Comparing the vLLM backend's output to the ground-truth SDPA output.

353 """

354 + from vllm.v1.attention.backends.mla.common import QueryLenSupport

355 +

356 batch_spec = BATCH_SPECS[batch_spec_name]

357 + is_spec_decode_test = batch_spec_name.startswith("spec_decode")

358 + spec_decode_backends = {_Backend.FLASH_ATTN_MLA, _Backend.FLASHMLA}

359 +

360 + block_size = 16

361 + required_blocks = sum(

362 + (seq_len + block_size - 1) // block_size for seq_len in batch_spec.seq_lens

363 + )

364 + # Add 1 for null block at index 0, and some buffer

365 + num_gpu_blocks = required_blocks + 1 + 100

366 +

367 vllm_config = create_vllm_config(

333 - model_name=model, max_model_len=max(batch_spec.seq_lens), num_gpu_blocks=2048

368 + model_name=model,

369 + max_model_len=max(batch_spec.seq_lens),

370 + num_gpu_blocks=num_gpu_blocks,

371 + block_size=block_size,

372 )

373 +

374 + # For spec decode tests, add a speculative_config to set the reorder_batch_threshold

375 + if is_spec_decode_test:

376 + from vllm.config import SpeculativeConfig

377 +

378 + # Get the query length from the batch spec (they should all be uniform)

379 + query_len = batch_spec.query_lens[0]

380 + # Set num_speculative_tokens to query_len - 1

381 + # (since threshold is 1 + num_spec_tokens)

382 + # Use ngram method which doesn't require a draft model

383 + vllm_config.speculative_config = SpeculativeConfig(

384 + method="ngram", num_speculative_tokens=query_len - 1

385 + )

386 +

387 device = torch.device("cuda:0")

388

389 kv_cache_spec = create_standard_kv_cache_spec(vllm_config)

@@ -395,11 +447,37 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):

447 # K_PE (rope component): [s_len, 1, qk_rope_head_dim]

448 k_pe_full = torch.randn(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)

449

398 - # Determine if this is decode or prefill

450 + # Determine if this sequence uses the decode pipeline or prefill

451 + # pipeline for each backend

452 + # NOTE: For spec decode tests with uniform query_len > 1, backends that

453 + # support spec decode (FLASH_ATTN_MLA with varlen support, FLASHMLA with

454 + # uniform support) will use the decode pipeline (MQA-style), while

455 + # backends that only support single-token queries will use the prefill

456 + # pipeline (MHA-style). This ensures the reference implementation

457 + # matches each backend's actual decode/prefill pipeline path.

458 is_decode = []

400 - for i, backend in enumerate(BACKENDS_TO_TEST):

459 + for backend_idx, backend in enumerate(BACKENDS_TO_TEST):

460 builder_cls, _ = try_get_attention_backend(backend)

402 - is_decode.append(q_len <= builder_cls.reorder_batch_threshold)

461 + if is_spec_decode_test:

462 + query_len_support = getattr(

463 + builder_cls, "query_len_support", QueryLenSupport.SINGLE_ONLY

464 + )

465 + supports_spec = query_len_support != QueryLenSupport.SINGLE_ONLY

466 + is_decode.append(supports_spec)

467 + else:

468 + threshold = getattr(builder_cls, "reorder_batch_threshold", None)

469 + query_len_support = getattr(

470 + builder_cls, "query_len_support", QueryLenSupport.SINGLE_ONLY

471 + )

472 + within_threshold = q_len <= threshold if threshold else False

473 + if (

474 + within_threshold

475 + and query_len_support == QueryLenSupport.UNIFORM

476 + and i > 0

477 + ):

478 + first_q_len = query_lens[0]

479 + within_threshold = q_len == first_q_len

480 + is_decode.append(within_threshold)

481

482 # Split q into nope and rope components

483 q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)

@@ -478,11 +556,11 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):

556 sdpa_out_i_prefill = sdpa_out_i_prefill.transpose(1, 2).squeeze(0)

557 sdpa_out_i_prefill = sdpa_out_i_prefill.flatten(start_dim=-2)

558

481 - for i, backend in enumerate(BACKENDS_TO_TEST):

482 - if is_decode[i]:

483 - all_sdpa_outputs[i].append(sdpa_out_i_decode)

559 + for backend_idx, backend in enumerate(BACKENDS_TO_TEST):

560 + if is_decode[backend_idx]:

561 + all_sdpa_outputs[backend_idx].append(sdpa_out_i_decode)

562 else:

485 - all_sdpa_outputs[i].append(sdpa_out_i_prefill)

563 + all_sdpa_outputs[backend_idx].append(sdpa_out_i_prefill)

564

565 # Inputs for vLLM MLA backends are just the new tokens

566 all_q_vllm.append(q_c)

@@ -497,9 +575,9 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):

575 query_vllm = torch.cat(all_q_vllm, dim=0)

576 kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)

577 k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)

500 - sdpa_outputs = []

501 - for i, backend in enumerate(BACKENDS_TO_TEST):

502 - sdpa_outputs.append(torch.cat(all_sdpa_outputs[i], dim=0))

578 + sdpa_outputs = {}

579 + for backend_idx, backend in enumerate(BACKENDS_TO_TEST):

580 + sdpa_outputs[backend] = torch.cat(all_sdpa_outputs[backend_idx], dim=0)

581

582 # Create mock kv_b_proj using the same weights as reference implementation

583 from vllm.model_executor.layers.linear import ColumnParallelLinear

@@ -516,7 +594,7 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):

594 kv_b_proj_weight = kv_b_proj_weight.view(

595 kv_lora_rank, num_q_heads * (qk_nope_head_dim + v_head_dim)

596 )

519 - mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T)

597 + mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T, requires_grad=False)

598

599 # Create metadata using original batch spec

600 common_attn_metadata = create_common_attn_metadata(

@@ -537,7 +615,11 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):

615 )

616

617 # 4. Run vLLM backends and compare

540 - for i, backend_name in enumerate(BACKENDS_TO_TEST):

618 + for backend_idx, backend_name in enumerate(BACKENDS_TO_TEST):

619 + # Skip backends that don't support spec decode for spec decode tests

620 + if is_spec_decode_test and backend_name not in spec_decode_backends:

621 + continue

622 +

623 backend_output = run_attention_backend(

624 backend_name,

625 kv_cache_spec,

@@ -556,14 +638,17 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):

638 mock_kv_b_proj,

639 )

640

641 + # Use backend_idx to get the correct SDPA output for this backend

642 + expected_output = sdpa_outputs[backend_name]

643 +

644 # Check shape and dtype consistency

560 - assert backend_output.shape == sdpa_outputs[i].shape, (

645 + assert backend_output.shape == expected_output.shape, (

646 f"[{backend_name}] shape {backend_output.shape} != "

562 - f"SDPA shape {sdpa_outputs[i].shape}"

647 + f"SDPA shape {expected_output.shape}"

648 )

564 - assert backend_output.dtype == sdpa_outputs[i].dtype, (

649 + assert backend_output.dtype == expected_output.dtype, (

650 f"[{backend_name}] dtype {backend_output.dtype} != "

566 - f"SDPA dtype {sdpa_outputs[i].dtype}"

651 + f"SDPA dtype {expected_output.dtype}"

652 )

653

654 assert torch.isfinite(backend_output).all(), (

@@ -574,12 +659,12 @@ def test_backend_correctness(dist_init, batch_spec_name: str, model: str):

659 rtol = 1e-2

660 atol = 5e-1

661

577 - max_diff = torch.max(torch.abs(backend_output - sdpa_outputs[i])).item()

662 + max_diff = torch.max(torch.abs(backend_output - expected_output)).item()

663 max_rel_diff = torch.max(

579 - torch.abs(backend_output - sdpa_outputs[i]) / torch.abs(sdpa_outputs[i])

664 + torch.abs(backend_output - expected_output) / torch.abs(expected_output)

665 ).item()

666 all_close = torch.allclose(

582 - backend_output, sdpa_outputs[i], rtol=rtol, atol=atol

667 + backend_output, expected_output, rtol=rtol, atol=atol

668 )

669

670 assert all_close, (

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

+37 -12

@@ -190,6 +190,7 @@ return curr_o @ W_O

190 import functools

191 from abc import abstractmethod

192 from dataclasses import dataclass, field

193 +from enum import Enum

194 from typing import ClassVar, Generic, TypeVar

195

196 import torch

@@ -227,6 +228,24 @@ from vllm.v1.attention.backends.utils import (

228 )

229 from vllm.v1.kv_cache_interface import AttentionSpec

230

231 +

232 +class QueryLenSupport(Enum):

233 + """Defines the level of query length support for an attention backend's

234 + decode pipeline.

235 +

236 + - SINGLE_ONLY: Decode pipeline only supports single-token queries

237 + (query_len=1)

238 + - UNIFORM: Decode pipeline supports uniform multi-token queries

239 + (all requests must have same query_len > 1)

240 + - VARLEN: Decode pipeline supports variable-length queries

241 + (mixed query lengths in same batch)

242 + """

243 +

244 + SINGLE_ONLY = "single_only"

245 + UNIFORM = "uniform"

246 + VARLEN = "varlen"

247 +

248 +

249 try:

250 from vllm.vllm_flash_attn import flash_attn_varlen_func

251

@@ -460,19 +479,18 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

479 understand this class

480 """

481

463 - # Whether the backend supports reordering the batch such that

464 - # short sequences (i.e. verification for speculative decoding) are

465 - # classified as decode requests.

466 - # If True, this will increase `reorder_batch_threshold` (below) when

467 - # speculative decoding is enabled, and set `require_uniform=True` when

468 - # when reordering the batch. Non-uniform decode requests will

469 - # fall back to prefill in this case.

470 - supports_uniform_spec_as_decode: ClassVar[bool] = False

482 + # Defines the level of query length support for this backend.

483 + # - SINGLE_ONLY: Only single-token queries (no spec decode support)

484 + # - UNIFORM: Supports uniform multi-token queries (spec decode with uniform lengths)

485 + # - VARLEN: Supports variable-length queries (spec decode with mixed lengths)

486 + # If set to UNIFORM or VARLEN, this will increase `reorder_batch_threshold` when

487 + # speculative decoding is enabled.

488 + query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.SINGLE_ONLY

489

490 # The threshold for reordering the batch into decode and prefill requests.

491 # If > 1, the batch will be reordered such that requests with

492 # query length <= threshold are classified as decode requests.

475 - # Use `supports_uniform_spec_as_decode` (above) to set this automatically

493 + # Use `query_len_support` (above) to set this automatically

494 # when speculative decoding is enabled.

495 reorder_batch_threshold: int = 1

496

@@ -599,11 +617,18 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

617 device=device,

618 )

619

602 - supports_spec_as_decode = self.supports_uniform_spec_as_decode

620 + supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY

621 self._init_reorder_batch_threshold(

604 - self.reorder_batch_threshold, supports_spec_as_decode

622 + self.reorder_batch_threshold, supports_spec_decode

623 )

624

625 + # Validate consistency between query_len_support and reorder_batch_threshold

626 + if self.query_len_support == QueryLenSupport.SINGLE_ONLY:

627 + assert self.reorder_batch_threshold == 1, (

628 + f"reorder_batch_threshold must be 1 when query_len_support is "

629 + f"SINGLE_ONLY, got {self.reorder_batch_threshold}"

630 + )

631 +

632 def _build_fi_prefill_wrappers(self, prefill: FlashInferPrefillMetadata):

633 qo_indptr = prefill.query_start_loc

634

@@ -745,7 +770,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

770 split_decodes_and_prefills(

771 common_attn_metadata,

772 decode_threshold=self.reorder_batch_threshold,

748 - require_uniform=self.supports_uniform_spec_as_decode,

773 + require_uniform=(self.query_len_support != QueryLenSupport.VARLEN),

774 )

775 )

776

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

+3 -2

@@ -24,6 +24,7 @@ from vllm.v1.attention.backends.mla.common import (

24 MLACommonImpl,

25 MLACommonMetadata,

26 MLACommonMetadataBuilder,

27 + QueryLenSupport,

28 )

29 from vllm.v1.attention.backends.utils import AttentionCGSupport

30 from vllm.v1.kv_cache_interface import AttentionSpec

@@ -66,8 +67,8 @@ class FlashAttnMLAMetadata(MLACommonMetadata[FlashAttnMLADecodeMetadata]):

68 class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]):

69 cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

69 -

70 - reorder_batch_threshold: int = 512

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

71 + reorder_batch_threshold: int = 512 # process small prefills with decode pathway

73 def __init__(

74 self,

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

+2 -4

@@ -13,6 +13,7 @@ from vllm.v1.attention.backends.mla.common import (

13 MLACommonImpl,

14 MLACommonMetadata,

15 MLACommonMetadataBuilder,

16 + QueryLenSupport,

17 )

18 from vllm.v1.attention.backends.utils import AttentionCGSupport

@@ -22,11 +23,8 @@ FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024

25 class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):

25 - # enable spec-as-decode optimization

26 - supports_uniform_spec_as_decode: ClassVar[bool] = True

27 -

28 - # enable full CUDA Graph support for decode-only capture

26 cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

27 + query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM

30 class FlashInferMLABackend(MLACommonBackend):

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

+16 -2

@@ -20,8 +20,13 @@ from vllm.v1.attention.backends.mla.common import (

20 MLACommonImpl,

21 MLACommonMetadata,

22 MLACommonMetadataBuilder,

23 + QueryLenSupport,

24 +)

25 +from vllm.v1.attention.backends.utils import (

26 + AttentionCGSupport,

27 + reshape_attn_output_for_spec_decode,

28 + reshape_query_for_spec_decode,

29 )

24 -from vllm.v1.attention.backends.utils import AttentionCGSupport

30 from vllm.v1.kv_cache_interface import AttentionSpec

32 logger = init_logger(__name__)

@@ -62,6 +67,9 @@ class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):

68 class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

69 cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

70 + query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM

71 + reorder_batch_threshold: int = 512 # process small prefills with decode pathway

72 + # ^ TODO(matt): tune this

74 def __init__(

75 self,

@@ -216,8 +224,12 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

224 q = torch.cat(q, dim=-1)

225

226 assert isinstance(q, torch.Tensor)

227 +

228 + num_decodes = attn_metadata.num_decodes

229 + q = reshape_query_for_spec_decode(q, num_decodes)

230 +

231 o, lse = flash_mla_with_kvcache(

220 - q=q.unsqueeze(1), # Add seqlen dim of 1 (decode)

232 + q=q,

233 k_cache=kv_c_and_k_pe_cache.unsqueeze(-2), # Add head dim of 1

234 block_table=attn_metadata.decode.block_table,

235 cache_seqlens=attn_metadata.decode.seq_lens,

@@ -230,4 +242,6 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

242 descale_k=layer._k_scale.reshape(1),

243 )

244

245 + o = reshape_attn_output_for_spec_decode(o)

246 +

247 return o, lse