@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 67f9046e4a

[Bugfix] Sparse MLA: enable fp8_ds_mla dense prefill (#48642)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Matthew Bonanni committed Jul 16, 2026 at 18:44 UTC 67f9046e4a12d939ffad0abc917585a00ef5b42d

13 files changed +516 -85

benchmarks/kernels/bench_cp_gather_fp8.py

+3 -4

@@ -69,12 +69,11 @@ def make_inputs(total_tokens, num_reqs, block_size):

69 # Output workspace

70 dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda")

72 - seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")

72 workspace_starts_t = torch.tensor(

73 workspace_starts, dtype=torch.int32, device="cuda"

74 )

77 - return cache, dst, block_table, seq_lens_t, workspace_starts_t

76 + return cache, dst, block_table, workspace_starts_t

79 def bench_scenario(label, num_reqs, total_tokens_list, save_path):

@@ -94,7 +93,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):

93 )

94 )

95 def bench_fn(total_tokens, provider, num_reqs):

97 - cache, dst, block_table, seq_lens_t, ws_starts = make_inputs(

96 + cache, dst, block_table, ws_starts = make_inputs(

97 total_tokens, num_reqs, BLOCK_SIZE

98 )

@@ -102,7 +101,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):

101

102 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(

103 lambda: ops.cp_gather_and_upconvert_fp8_kv_cache(

105 - cache, dst, block_table, seq_lens_t, ws_starts, num_reqs

104 + cache, dst, block_table, ws_starts, num_reqs

105 ),

106 quantiles=quantiles,

107 rep=500,

csrc/cache.h

+1 -2

@@ -67,9 +67,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(

67 torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]

68 torch::Tensor const& dst, // [TOT_TOKENS, 576]

69 torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]

70 - torch::Tensor const& seq_lens, // [BATCH]

70 torch::Tensor const& workspace_starts, // [BATCH]

72 - int64_t batch_size);

71 + int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);

73 // Indexer K quantization and cache function

74 void indexer_k_quant_and_cache(

csrc/libtorch_stable/cache_kernels.cu

+19 -9

@@ -1174,7 +1174,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache(

1174 const int32_t num_reqs, const int32_t block_size,

1175 const int32_t total_tokens, const int64_t block_table_stride,

1176 const int64_t cache_block_stride, const int64_t cache_entry_stride,

1177 - const int64_t dst_entry_stride) {

1177 + const int64_t dst_entry_stride,

1178 + const int32_t* __restrict__ seq_starts) { // Optional source offsets

1179 const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;

1180 if (flat_warp_id >= total_tokens) return;

1181 const int lane_id = threadIdx.x & 31;

@@ -1192,7 +1193,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache(

1193

1194 // Compute physical token address via block table

1195 const int out_token_id = flat_warp_id;

1195 - const int token_offset = out_token_id - workspace_starts[req_id];

1196 + int token_offset = out_token_id - workspace_starts[req_id];

1197 + if (seq_starts != nullptr) token_offset += seq_starts[req_id];

1198 const int cache_block_idx = token_offset / block_size;

1199 const int offset_in_block = token_offset % block_size;

1200 const int physical_block =

@@ -1383,9 +1385,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(

1385 torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]

1386 torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]

1387 torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]

1386 - torch::stable::Tensor const& seq_lens, // [BATCH]

1388 torch::stable::Tensor const& workspace_starts, // [BATCH]

1388 - int64_t batch_size) {

1389 + int64_t batch_size,

1390 + std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {

1391 torch::stable::accelerator::DeviceGuard device_guard(

1392 src_cache.get_device_index());

1393 const cudaStream_t stream = get_current_cuda_stream();

@@ -1396,20 +1398,25 @@ void cp_gather_and_upconvert_fp8_kv_cache(

1398 STD_TORCH_CHECK(

1399 block_table.scalar_type() == torch::headeronly::ScalarType::Int,

1400 "block_table must be int32");

1399 - STD_TORCH_CHECK(seq_lens.scalar_type() == torch::headeronly::ScalarType::Int,

1400 - "seq_lens must be int32");

1401 STD_TORCH_CHECK(

1402 workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int,

1403 "workspace_starts must be int32");

1404 + if (seq_starts.has_value()) {

1405 + STD_TORCH_CHECK(

1406 + seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,

1407 + "seq_starts must be int32");

1408 + }

1409

1410 STD_TORCH_CHECK(src_cache.device() == dst.device(),

1411 "src_cache and dst must be on the same device");

1412 STD_TORCH_CHECK(src_cache.device() == block_table.device(),

1413 "src_cache and block_table must be on the same device");

1409 - STD_TORCH_CHECK(src_cache.device() == seq_lens.device(),

1410 - "src_cache and seq_lens must be on the same device");

1414 STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(),

1415 "src_cache and workspace_starts must be on the same device");

1416 + if (seq_starts.has_value()) {

1417 + STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),

1418 + "src_cache and seq_starts must be on the same device");

1419 + }

1420 auto dtype = src_cache.scalar_type();

1421 STD_TORCH_CHECK(

1422 dtype == torch::headeronly::ScalarType::Byte || // uint8

@@ -1438,6 +1445,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(

1445 constexpr int warps_per_block = 8;

1446 const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block;

1447 const int block_size_threads = warps_per_block * 32; // 256 threads

1448 + const int32_t* seq_starts_ptr =

1449 + seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()

1450 + : nullptr;

1451

1452 vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid_size, block_size_threads, 0,

1453 stream>>>(

@@ -1446,7 +1456,7 @@ void cp_gather_and_upconvert_fp8_kv_cache(

1456 workspace_starts.const_data_ptr<int32_t>(),

1457 static_cast<int32_t>(batch_size), block_size, total_tokens,

1458 block_table_stride, cache_block_stride, cache_entry_stride,

1449 - dst_entry_stride);

1459 + dst_entry_stride, seq_starts_ptr);

1460 }

1461

1462 // Macro to dispatch the kernel based on the data type.

csrc/libtorch_stable/ops.h

+2 -2

@@ -527,9 +527,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(

527 // 656]

528 torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]

529 torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]

530 - torch::stable::Tensor const& seq_lens, // [BATCH]

530 torch::stable::Tensor const& workspace_starts, // [BATCH]

532 - int64_t batch_size);

531 + int64_t batch_size,

532 + std::optional<torch::stable::Tensor> seq_starts = std::nullopt);

533

534 // Indexer K quantization and cache function

535 void indexer_k_quant_and_cache(

csrc/libtorch_stable/torch_bindings.cpp

+2 -2

@@ -847,8 +847,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {

847

848 ops.def(

849 "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "

850 - "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "

851 - "batch_size) -> ()");

850 + "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? "

851 + "seq_starts) -> ()");

852

853 ops.def(

854 "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "

tests/kernels/test_cp_gather_fp8.py

+63 -13

@@ -27,8 +27,8 @@ def _build_test_case(seq_lens, block_size, seed=42):

27 seed: Random seed for reproducibility.

29 Returns:

30 - Tuple of (cache, block_table, seq_lens_t, workspace_starts_t,

31 - num_reqs, total_tokens, expected_output).

30 + Tuple of (cache, block_table, workspace_starts_t, num_reqs,

31 + total_tokens, expected_output).

32 """

33 torch.manual_seed(seed)

@@ -112,7 +112,6 @@ def _build_test_case(seq_lens, block_size, seed=42):

112 # Expected output: exact copy

113 expected[out_idx, NOPE_DIM:] = rope

114

115 - seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")

115 workspace_starts_t = torch.tensor(

116 workspace_starts, dtype=torch.int32, device="cuda"

117 )

@@ -120,7 +119,6 @@ def _build_test_case(seq_lens, block_size, seed=42):

119 return (

120 cache,

121 block_table,

123 - seq_lens_t,

122 workspace_starts_t,

123 num_reqs,

124 total_tokens,

@@ -197,7 +195,6 @@ def _build_test_case_fast(seq_lens, block_size, seed=42):

195 flat_cache[:sl] = token_data[ws : ws + sl]

196 block_start += nb

197

200 - seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")

198 workspace_starts_t = torch.tensor(

199 workspace_starts, dtype=torch.int32, device="cuda"

200 )

@@ -205,7 +202,6 @@ def _build_test_case_fast(seq_lens, block_size, seed=42):

202 return (

203 cache,

204 block_table,

208 - seq_lens_t,

205 workspace_starts_t,

206 num_reqs,

207 total_tokens,

@@ -232,7 +228,6 @@ def test_cp_gather_and_upconvert_fp8_kv_cache(seq_lens, block_size):

228 (

229 cache,

230 block_table,

235 - seq_lens_t,

231 workspace_starts_t,

232 num_reqs,

233 total_tokens,

@@ -244,7 +239,7 @@ def test_cp_gather_and_upconvert_fp8_kv_cache(seq_lens, block_size):

239 )

240

241 ops.cp_gather_and_upconvert_fp8_kv_cache(

247 - cache, dst, block_table, seq_lens_t, workspace_starts_t, num_reqs

242 + cache, dst, block_table, workspace_starts_t, num_reqs

243 )

244

245 # NoPE: fp8 dequant has rounding error, so we allow small tolerance.

@@ -277,8 +272,6 @@ def test_cp_gather_fp8_shuffled_blocks():

272 )

273 block_table = torch.tensor([[3, 1]], dtype=torch.int32, device="cuda")

274 workspace_starts = torch.tensor([0], dtype=torch.int32, device="cuda")

280 - seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")

281 -

275 expected = torch.zeros(

276 total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"

277 )

@@ -314,7 +307,7 @@ def test_cp_gather_fp8_shuffled_blocks():

307 )

308

309 ops.cp_gather_and_upconvert_fp8_kv_cache(

317 - cache, dst, block_table, seq_lens_t, workspace_starts, len(seq_lens)

310 + cache, dst, block_table, workspace_starts, len(seq_lens)

311 )

312

313 torch.testing.assert_close(

@@ -323,6 +316,64 @@ def test_cp_gather_fp8_shuffled_blocks():

316 assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])

317

318

319 +@pytest.mark.parametrize(

320 + "gather_seq_lens,seq_starts",

321 + [

322 + ([6, 5, 3], [3, 4, 5]),

323 + ([0, 5, 3], [12, 4, 5]),

324 + ],

325 +)

326 +def test_cp_gather_fp8_with_sequence_starts(gather_seq_lens, seq_starts):

327 + """Gather request slices beginning at arbitrary cache positions."""

328 + full_seq_lens = [12, 11, 9]

329 + (

330 + cache,

331 + block_table,

332 + _workspace_starts_t,

333 + num_reqs,

334 + _total_tokens,

335 + full_expected,

336 + ) = _build_test_case(full_seq_lens, block_size=4)

337 +

338 + workspace_starts = torch.tensor(

339 + [0, gather_seq_lens[0], sum(gather_seq_lens[:2])],

340 + dtype=torch.int32,

341 + device="cuda",

342 + )

343 + seq_starts_t = torch.tensor(seq_starts, dtype=torch.int32, device="cuda")

344 + dst = torch.empty(

345 + sum(gather_seq_lens),

346 + NOPE_DIM + ROPE_DIM,

347 + dtype=torch.bfloat16,

348 + device="cuda",

349 + )

350 +

351 + ops.cp_gather_and_upconvert_fp8_kv_cache(

352 + cache,

353 + dst,

354 + block_table,

355 + workspace_starts,

356 + num_reqs,

357 + seq_starts_t,

358 + )

359 +

360 + full_workspace_starts = [0, full_seq_lens[0], sum(full_seq_lens[:2])]

361 + expected = torch.cat(

362 + [

363 + full_expected[

364 + full_workspace_starts[i] + seq_starts[i] : full_workspace_starts[i]

365 + + seq_starts[i]

366 + + gather_seq_lens[i]

367 + ]

368 + for i in range(num_reqs)

369 + ]

370 + )

371 + torch.testing.assert_close(

372 + dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2

373 + )

374 + assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])

375 +

376 +

377 @pytest.mark.parametrize(

378 "seq_lens,block_size",

379 [

@@ -342,7 +393,6 @@ def test_cp_gather_fp8_large_seqlens(seq_lens, block_size):

393 (

394 cache,

395 block_table,

345 - seq_lens_t,

396 workspace_starts_t,

397 num_reqs,

398 total_tokens,

@@ -354,7 +404,7 @@ def test_cp_gather_fp8_large_seqlens(seq_lens, block_size):

404 )

405

406 ops.cp_gather_and_upconvert_fp8_kv_cache(

357 - cache, dst, block_table, seq_lens_t, workspace_starts_t, num_reqs

407 + cache, dst, block_table, workspace_starts_t, num_reqs

408 )

409

410 torch.testing.assert_close(

tests/v1/attention/test_mla_backends.py

+29 -1

@@ -42,7 +42,10 @@ from vllm.v1.attention.backends.mla.prefill import (

42 )

43 from vllm.v1.attention.backends.registry import AttentionBackendEnum

44 from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported

45 -from vllm.v1.kv_cache_interface import MLAAttentionSpec

45 +from vllm.v1.kv_cache_interface import (

46 + KVQuantMode,

47 + MLAAttentionSpec,

48 +)

50 BACKENDS_TO_TEST = [

51 AttentionBackendEnum.CUTLASS_MLA,

@@ -55,6 +58,31 @@ BACKENDS_TO_TEST = [

59 DEVICE_TYPE = current_platform.device_type

61 +

62 +@pytest.mark.parametrize(

63 + ("cache_dtype", "expected_quant_mode"),

64 + [

65 + ("auto", KVQuantMode.NONE),

66 + ("fp8_ds_mla", KVQuantMode.FP8_PER_TENSOR),

67 + ],

68 +)

69 +def test_mla_kv_cache_spec_uses_layer_cache_dtype(

70 + cache_dtype: str, expected_quant_mode: KVQuantMode

71 +):

72 + layer = SimpleNamespace(kv_cache_dtype=cache_dtype, head_size=576)

73 + vllm_config = SimpleNamespace(

74 + cache_config=SimpleNamespace(block_size=64), model_config=None

75 + )

76 +

77 + spec = MLAAttention.get_kv_cache_spec(layer, vllm_config)

78 +

79 + assert isinstance(spec, MLAAttentionSpec)

80 + assert spec.cache_dtype_str == cache_dtype

81 + assert spec.kv_quant_mode == expected_quant_mode

82 + if cache_dtype == "fp8_ds_mla":

83 + assert spec.page_size_bytes == 64 * 656

84 +

85 +

86 # Remove sm100 backends from the list if not using sm100

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

88 BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)

tests/v1/attention/test_sparse_mla_backends.py

+321 -1

@@ -40,10 +40,16 @@ from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (

40 )

41 from vllm.v1.attention.backends.mla.flashmla_sparse import (

42 FlashMLASparseBackend,

43 + FlashMLASparseImpl,

44 + FlashMLASparseMetadata,

45 + FlashMLASparseMetadataBuilder,

46 triton_convert_req_index_to_global_index,

47 )

48 from vllm.v1.attention.backends.mla.indexer import split_indexer_prefill_chunks

46 -from vllm.v1.attention.backends.utils import split_prefill_chunks

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

50 + split_decodes_and_prefills,

51 + split_prefill_chunks,

52 +)

53 from vllm.v1.attention.ops import flashmla

55 SPARSE_BACKEND_BATCH_SPECS = {

@@ -721,6 +727,120 @@ def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_s

727 torch.testing.assert_close(result, reference_result, rtol=0, atol=0)

728

729

730 +@pytest.mark.skipif(

731 + torch.cuda.get_device_capability() < (9, 0),

732 + reason="FlashMLASparseBackend requires CUDA 9.0 or higher",

733 +)

734 +def test_triton_convert_rejects_req_id_longer_than_token_indices():

735 + """Guard against the #47327 regression: the kernel grid is sized by

736 + req_id but the output is allocated like token_indices, so a full-batch

737 + req_id combined with an MQA-subset token_indices wrote past the end of

738 + the output buffer. The wrapper must reject the length mismatch instead

739 + of corrupting memory."""

740 + device = torch.device(DEVICE_TYPE)

741 + num_topk_tokens = 128

742 + block_size = 64

743 + block_table = torch.arange(40, dtype=torch.int32, device=device).view(4, 10)

744 +

745 + # Full batch: 2 decode tokens + 10 prefill tokens

746 + req_id_full = torch.tensor(

747 + [0, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3], dtype=torch.int32, device=device

748 + )

749 + num_mqa_tokens = 2

750 + token_indices = torch.randint(

751 + 0,

752 + block_size * 10,

753 + (num_mqa_tokens, num_topk_tokens),

754 + dtype=torch.int32,

755 + device=device,

756 + )

757 +

758 + with pytest.raises(AssertionError, match="must cover the same tokens"):

759 + triton_convert_req_index_to_global_index(

760 + req_id_full,

761 + block_table,

762 + token_indices,

763 + BLOCK_SIZE=block_size,

764 + NUM_TOPK_TOKENS=num_topk_tokens,

765 + )

766 +

767 + # The sliced call is the intended usage and must match the reference.

768 + result = triton_convert_req_index_to_global_index(

769 + req_id_full[:num_mqa_tokens],

770 + block_table,

771 + token_indices,

772 + BLOCK_SIZE=block_size,

773 + NUM_TOPK_TOKENS=num_topk_tokens,

774 + )

775 + reference = _triton_convert_reference_impl(

776 + req_id_full[:num_mqa_tokens],

777 + block_table,

778 + token_indices,

779 + block_size,

780 + num_topk_tokens,

781 + )

782 + torch.testing.assert_close(result, reference, rtol=0, atol=0)

783 +

784 +

785 +@pytest.mark.skipif(

786 + torch.cuda.get_device_capability() < (9, 0),

787 + reason="FlashMLASparseBackend requires CUDA 9.0 or higher",

788 +)

789 +def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():

790 + """Guard against the #47327 regression: when the dense-MHA prefill split

791 + is active, forward_mqa only receives the leading decode tokens, but

792 + _forward_bf16_kv passed the full-batch req_id_per_token to the index

793 + conversion, making it write past the end of its output buffer. The call

794 + site must slice req_id_per_token to the MQA tokens."""

795 + device = torch.device(DEVICE_TYPE)

796 + num_topk_tokens = 128

797 + block_size = 64

798 + num_batch_tokens = 12

799 + num_mqa_tokens = 2

800 +

801 + attn_metadata = SimpleNamespace(

802 + req_id_per_token=torch.tensor(

803 + [0, 1] + [2] * 5 + [3] * 5, dtype=torch.int32, device=device

804 + ),

805 + block_table=torch.arange(40, dtype=torch.int32, device=device).view(4, 10),

806 + block_size=block_size,

807 + )

808 + assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens

809 +

810 + q = torch.zeros(num_mqa_tokens, 4, 576, dtype=torch.bfloat16, device=device)

811 + kv_cache = torch.zeros(40 * block_size, 576, dtype=torch.bfloat16, device=device)

812 + topk_indices = torch.randint(

813 + 0,

814 + block_size * 10,

815 + (num_mqa_tokens, num_topk_tokens),

816 + dtype=torch.int32,

817 + device=device,

818 + )

819 +

820 + captured = {}

821 +

822 + def _stub_kernel(q, kv, indices, lengths):

823 + captured["indices"] = indices

824 + return torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device)

825 +

826 + stub_impl = SimpleNamespace(_bf16_flash_mla_kernel=_stub_kernel)

827 +

828 + out = FlashMLASparseImpl._forward_bf16_kv(

829 + stub_impl, q, kv_cache, topk_indices, attn_metadata

830 + )

831 +

832 + assert out.shape[0] == num_mqa_tokens

833 + assert captured["indices"].shape[0] == num_mqa_tokens

834 + reference = _triton_convert_reference_impl(

835 + attn_metadata.req_id_per_token[:num_mqa_tokens],

836 + attn_metadata.block_table,

837 + topk_indices,

838 + block_size,

839 + num_topk_tokens,

840 + )

841 + torch.testing.assert_close(captured["indices"], reference, rtol=0, atol=0)

842 +

843 +

844 @pytest.mark.parametrize(

845 "seq_lens,max_buf,expected",

846 [

@@ -1124,3 +1244,203 @@ def test_triton_convert_returns_valid_counts():

1244 )

1245 assert isinstance(result_only, torch.Tensor)

1246 torch.testing.assert_close(result_only, result, rtol=0, atol=0)

1247 +

1248 +

1249 +def test_flashmla_cache_dtype_aliases_use_ds_layout():

1250 + from vllm.model_executor.layers.attention.mla_attention import (

1251 + _canonicalize_sparse_mla_kv_cache_dtype,

1252 + )

1253 +

1254 + # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer

1255 + # stores kv_cache_dtype, so they cannot bypass the gate.

1256 + for alias in ("fp8", "fp8_e4m3"):

1257 + assert (

1258 + _canonicalize_sparse_mla_kv_cache_dtype(FlashMLASparseBackend, alias)

1259 + == "fp8_ds_mla"

1260 + )

1261 +

1262 +

1263 +def test_flashmla_fp8_metadata_reuses_common_batch_split():

1264 + builder = SimpleNamespace(

1265 + device=torch.device(DEVICE_TYPE),

1266 + vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=8)),

1267 + )

1268 + common_metadata = SimpleNamespace(

1269 + num_actual_tokens=1,

1270 + seq_lens_cpu_upper_bound=torch.tensor([1]),

1271 + query_start_loc_cpu=torch.tensor([0, 1]),

1272 + block_table_tensor=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE),

1273 + )

1274 + metadata = FlashMLASparseMetadata(

1275 + num_reqs=1,

1276 + max_query_len=1,

1277 + max_seq_len=1,

1278 + num_actual_tokens=1,

1279 + query_start_loc=torch.tensor([0, 1], device=DEVICE_TYPE),

1280 + slot_mapping=torch.tensor([0], device=DEVICE_TYPE),

1281 + block_table=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE),

1282 + req_id_per_token=torch.zeros(1, dtype=torch.int32, device=DEVICE_TYPE),

1283 + num_decodes=0,

1284 + num_prefills=1,

1285 + num_decode_tokens=0,

1286 + )

1287 +

1288 + fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode(

1289 + builder, common_metadata, metadata

1290 + )

1291 +

1292 + assert fp8_metadata.num_decodes == 0

1293 + assert fp8_metadata.num_prefills == 1

1294 + assert fp8_metadata.num_decode_tokens == 0

1295 + assert fp8_metadata.num_prefill_tokens == 1

1296 +

1297 +

1298 +def test_flashmla_common_metadata_requires_uniform_decodes():

1299 + common_metadata = SimpleNamespace(

1300 + max_query_len=3,

1301 + num_reqs=3,

1302 + num_actual_tokens=6,

1303 + query_start_loc_cpu=torch.tensor([0, 1, 3, 6]),

1304 + is_prefilling=None,

1305 + )

1306 +

1307 + split = split_decodes_and_prefills(

1308 + common_metadata,

1309 + decode_threshold=128,

1310 + require_uniform=FlashMLASparseMetadataBuilder.require_uniform_decodes,

1311 + )

1312 +

1313 + assert split == (1, 2, 1, 5)

1314 +

1315 +

1316 +def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch):

1317 + monkeypatch.setattr(

1318 + "vllm.v1.attention.backends.mla.flashmla_sparse.get_mla_metadata",

1319 + lambda: (object(), None),

1320 + )

1321 + builder = SimpleNamespace(

1322 + device=torch.device(DEVICE_TYPE),

1323 + dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE),

1324 + max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE),

1325 + )

1326 + query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660])

1327 + common_metadata = SimpleNamespace(

1328 + num_actual_tokens=660,

1329 + query_start_loc_cpu=query_start_loc_cpu,

1330 + seq_lens=torch.arange(7, device=DEVICE_TYPE),

1331 + )

1332 + metadata = FlashMLASparseMetadata(

1333 + num_reqs=7,

1334 + max_query_len=110,

1335 + max_seq_len=110,

1336 + num_actual_tokens=660,

1337 + query_start_loc=query_start_loc_cpu.to(DEVICE_TYPE),

1338 + slot_mapping=torch.arange(660, device=DEVICE_TYPE),

1339 + block_table=torch.zeros(7, 1, dtype=torch.int32, device=DEVICE_TYPE),

1340 + req_id_per_token=torch.zeros(660, dtype=torch.int32, device=DEVICE_TYPE),

1341 + num_decodes=7,

1342 + num_prefills=0,

1343 + num_decode_tokens=660,

1344 + )

1345 +

1346 + fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode(

1347 + builder, common_metadata, metadata

1348 + )

1349 +

1350 + assert fp8_metadata.num_decodes == 6

1351 + assert fp8_metadata.num_decode_tokens == 660

1352 + assert fp8_metadata.decode is not None

1353 + assert fp8_metadata.decode.decode_query_len == 110

1354 + torch.testing.assert_close(

1355 + fp8_metadata.decode.seq_lens, torch.arange(6, device=DEVICE_TYPE)

1356 + )

1357 +

1358 +

1359 +@pytest.mark.parametrize("use_mixed_batch", [False, True])

1360 +def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: bool):

1361 + num_decode_tokens = 2

1362 + num_batch_tokens = 5

1363 + q = torch.empty(num_decode_tokens, 2, 3, device=DEVICE_TYPE)

1364 + topk_indices = torch.empty(num_decode_tokens, 4, device=DEVICE_TYPE)

1365 + kernel_q_shapes = []

1366 +

1367 + def convert_indices(*args, **kwargs): # noqa: ARG001

1368 + assert not kwargs.get("HAS_PREFILL_WORKSPACE", False)

1369 + if not kwargs.get("return_valid_counts", False):

1370 + return topk_indices

1371 + valid_counts = torch.full(

1372 + (num_decode_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE

1373 + )

1374 + return topk_indices, valid_counts

1375 +

1376 + monkeypatch.setattr(

1377 + "vllm.v1.attention.backends.mla.flashmla_sparse."

1378 + "triton_convert_req_index_to_global_index",

1379 + convert_indices,

1380 + )

1381 +

1382 + def run_kernel(**kwargs):

1383 + kernel_q_shapes.append(kwargs["q"].shape)

1384 + return kwargs["q"][..., :1], None

1385 +

1386 + if use_mixed_batch:

1387 + fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(

1388 + scheduler_metadata=object(), # type: ignore[arg-type]

1389 + dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),

1390 + cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),

1391 + )

1392 + else:

1393 + FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode

1394 + fp8_metadata = FP8Meta(

1395 + num_decodes=1,

1396 + num_prefills=1,

1397 + num_decode_tokens=num_decode_tokens,

1398 + num_prefill_tokens=num_batch_tokens - num_decode_tokens,

1399 + decode=FP8Meta.Decode(

1400 + seq_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),

1401 + kernel_metadata=object(), # type: ignore[arg-type]

1402 + decode_query_len=num_decode_tokens,

1403 + ),

1404 + prefill=FP8Meta.Prefill(

1405 + request_ids=torch.empty(

1406 + num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE

1407 + ),

1408 + workspace_starts=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE),

1409 + chunks=[],

1410 + ),

1411 + )

1412 + metadata = SimpleNamespace(

1413 + fp8_extra_metadata=fp8_metadata,

1414 + fp8_use_mixed_batch=use_mixed_batch,

1415 + num_actual_tokens=num_batch_tokens,

1416 + req_id_per_token=torch.empty(

1417 + num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE

1418 + ),

1419 + block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE),

1420 + block_size=64,

1421 + )

1422 + impl = SimpleNamespace(

1423 + kv_cache_dtype="fp8_ds_mla",

1424 + topk_indices_buffer=topk_indices,

1425 + num_heads=2,

1426 + kv_lora_rank=1,

1427 + _fp8_flash_mla_kernel=run_kernel,

1428 + )

1429 + impl._forward_fp8_kv_mixed_batch = MethodType(

1430 + FlashMLASparseImpl._forward_fp8_kv_mixed_batch, impl

1431 + )

1432 + impl._forward_fp8_kv_separate_prefill_decode = MethodType(

1433 + FlashMLASparseImpl._forward_fp8_kv_separate_prefill_decode, impl

1434 + )

1435 +

1436 + output, lse = FlashMLASparseImpl.forward_mqa(

1437 + impl,

1438 + q,

1439 + torch.empty(0, device=DEVICE_TYPE),

1440 + metadata,

1441 + None,

1442 + )

1443 +

1444 + assert kernel_q_shapes == [(1, num_decode_tokens, 2, 3)]

1445 + assert output.shape == (num_decode_tokens, 2, 1)

1446 + assert lse is None

vllm/_custom_ops.py

+3 -3

@@ -2759,9 +2759,9 @@ def cp_gather_and_upconvert_fp8_kv_cache(

2759 src_cache: torch.Tensor,

2760 dst: torch.Tensor,

2761 block_table: torch.Tensor,

2762 - seq_lens: torch.Tensor,

2762 workspace_starts: torch.Tensor,

2763 batch_size: int,

2764 + seq_starts: torch.Tensor | None = None,

2765 ) -> None:

2766 """Gather and upconvert FP8 KV cache to BF16 workspace.

2767

@@ -2769,12 +2769,12 @@ def cp_gather_and_upconvert_fp8_kv_cache(

2769 src_cache: FP8 KV cache [num_blocks, block_size, 656]

2770 dst: BF16 output workspace [total_tokens, 576]

2771 block_table: Block indices [num_reqs, max_blocks]

2772 - seq_lens: Sequence lengths [num_reqs]

2772 workspace_starts: Workspace start offsets [num_reqs]

2773 batch_size: Number of requests

2774 + seq_starts: Optional source sequence offsets [num_reqs]

2775 """

2776 torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache(

2777 - src_cache, dst, block_table, seq_lens, workspace_starts, batch_size

2777 + src_cache, dst, block_table, workspace_starts, batch_size, seq_starts

2778 )

2779

2780

vllm/model_executor/layers/attention/mla_attention.py

+25 -6

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

1749 self.determine_chunked_prefill_workspace_size(vllm_config)

1750 )

1751

1752 + use_packed_fp8_cache = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"

1753 if self.dcp_world_size > 1:

1754 # Note(hc): The local kvcache is incomplete when DCP is triggered,

1755 # an additional kvcache allgather across the DCP group is therefore

@@ -1761,7 +1762,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1762 + self.chunked_prefill_workspace_size // self.dcp_world_size,

1763 self.model_config.get_head_size(),

1764 ),

1764 - dtype=self.model_config.dtype,

1765 + dtype=torch.bfloat16

1766 + if use_packed_fp8_cache

1767 + else self.model_config.dtype,

1768 device=device,

1769 )

1770 else:

@@ -1770,7 +1773,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1773 self.chunked_prefill_workspace_size,

1774 self.model_config.get_head_size(),

1775 ),

1773 - dtype=self.q_data_type,

1776 + dtype=torch.bfloat16 if use_packed_fp8_cache else self.q_data_type,

1777 device=device,

1778 )

1779

@@ -2115,7 +2118,16 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):

2118

2119 for i in range(iters):

2120 toks = prefill_metadata.chunked_context.seq_tot[i]

2118 - if not use_fp8_prefill:

2121 + if self.kv_cache_dtype == "fp8_ds_mla":

2122 + ops.cp_gather_and_upconvert_fp8_kv_cache(

2123 + src_cache=kv_c_and_k_pe_cache,

2124 + dst=workspace[:toks],

2125 + block_table=prefill_metadata.block_table,

2126 + workspace_starts=prefill_metadata.chunked_context.cu_seq_lens[i],

2127 + batch_size=attn_metadata.num_prefills,

2128 + seq_starts=prefill_metadata.chunked_context.starts[i],

2129 + )

2130 + elif not use_fp8_prefill:

2131 ops.gather_and_maybe_dequant_cache(

2132 src_cache=kv_c_and_k_pe_cache,

2133 dst=workspace,

@@ -2230,9 +2242,16 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):

2242 padded_local_cu_seq_lens = (

2243 prefill_metadata.chunked_context.padded_local_cu_seq_lens[i]

2244 )

2233 - if is_quantized_kv_cache(self.kv_cache_dtype) and (

2234 - self.kv_cache_dtype != "fp8_ds_mla"

2235 - ):

2245 + if self.kv_cache_dtype == "fp8_ds_mla":

2246 + ops.cp_gather_and_upconvert_fp8_kv_cache(

2247 + src_cache=kv_c_and_k_pe_cache,

2248 + dst=workspace[:toks],

2249 + block_table=prefill_metadata.block_table,

2250 + workspace_starts=padded_local_cu_seq_lens,

2251 + batch_size=attn_metadata.num_prefills,

2252 + seq_starts=prefill_metadata.chunked_context.starts[i],

2253 + )

2254 + elif is_quantized_kv_cache(self.kv_cache_dtype):

2255 assert k_scale is not None

2256 ops.gather_and_maybe_dequant_cache(

2257 src_cache=kv_c_and_k_pe_cache,

vllm/model_executor/layers/attention/sparse_mla_attention.py

+3 -1

@@ -3,7 +3,7 @@

3 """Shared forward_mha implementation and metadata builder for sparse MLA backends."""

5 from shutil import which

6 -from typing import TYPE_CHECKING, Generic, TypeVar

6 +from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar

8 import numpy as np

9 import torch

@@ -38,6 +38,7 @@ T = TypeVar("T", bound=AttentionMetadata)

39 class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):

40 metadata_cls: type[T]

41 + require_uniform_decodes: ClassVar[bool] = False

43 def __init__(

44 self,

@@ -173,6 +174,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):

174 num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(

175 common_attn_metadata,

176 decode_threshold=self.reorder_batch_threshold or 1,

177 + require_uniform=self.require_uniform_decodes,

178 )

179 (

180 prefill_query_start_loc,

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

+40 -41

@@ -32,7 +32,6 @@ from vllm.v1.attention.backends.mla.sparse_utils import (

32 from vllm.v1.attention.backends.utils import (

33 reshape_attn_output_for_spec_decode,

34 reshape_query_for_spec_decode,

35 - split_decodes_and_prefills,

35 split_prefill_chunks,

36 )

37 from vllm.v1.attention.ops.flashmla import (

@@ -57,8 +56,8 @@ logger = init_logger(__name__)

56 # the FP8 decode kernel for decode.

57 # Currently we use #1 when the number of heads per rank is low (i.e. TP) since the BF16

58 # prefill kernel requires padding the number of heads to 128 while the decode does not

60 -# so when the per ranke head count is below MIN_HEADS_FOR_BF16_PREFILL we use the mixed

61 -# batch mode (#2).

59 +# so when the per-rank head count is below MIN_HEADS_FOR_BF16_PREFILL we use the mixed

60 +# batch mode (#1).

61 MIN_HEADS_FOR_BF16_PREFILL = 32

63 """

@@ -182,10 +181,6 @@ class FlashMLASparseMetadata(AttentionMetadata):

181

182 @dataclass

183 class Prefill:

185 - # Sequence lengths (context + query) for prefill requests

186 - # Shape: [num_prefill_reqs]

187 - seq_lens: torch.Tensor

188 -

184 # Request ID for each token: -1 for decode tokens, request index

185 # (0, 1, 2, ...) for prefill tokens.

186 # Shape: [num_actual_tokens]

@@ -204,7 +199,6 @@ class FlashMLASparseMetadata(AttentionMetadata):

199 Prefill requests may be chunked to fit within the fixed workspace size.

200 """

201

207 - seq_lens: torch.Tensor

202 tokens_slice: slice

203 block_table: torch.Tensor

204 req_start_idx: int

@@ -239,6 +233,7 @@ class FlashMLASparseMetadataBuilder(

233 SparseMLACommonMetadataBuilder[FlashMLASparseMetadata]

234 ):

235 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

236 + require_uniform_decodes: ClassVar[bool] = True

237 metadata_cls = FlashMLASparseMetadata

238

239 def __init__(

@@ -322,10 +317,11 @@ class FlashMLASparseMetadataBuilder(

317 self,

318 common_attn_metadata: CommonAttentionMetadata,

319 ) -> "FlashMLASparseMetadata.FP8KernelMetadata":

325 - """Build FP8 metadata treating all tokens as one mixed batch.

320 + """Build FP8 metadata treating MQA tokens as one batch.

321

327 - This matches main branch's approach and avoids the BF16 prefill kernel

328 - which has head padding overhead when num_heads is small (high TP case).

322 + The scheduler initializes lazily from the runtime query shape, which may

323 + be the full batch or only decodes when prefills use dense MHA. This avoids

324 + the BF16 prefill kernel's head-padding overhead at high TP.

325 """

326 num_tokens = common_attn_metadata.num_actual_tokens

327

@@ -353,20 +349,29 @@ class FlashMLASparseMetadataBuilder(

349 def _build_fp8_separate_prefill_decode(

350 self,

351 common_attn_metadata: CommonAttentionMetadata,

352 + metadata: FlashMLASparseMetadata,

353 ) -> "FlashMLASparseMetadata.FP8SeparatePrefillDecode":

354 num_tokens = common_attn_metadata.num_actual_tokens

355

356 (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = (

360 - split_decodes_and_prefills(

361 - common_attn_metadata,

362 - decode_threshold=self.reorder_batch_threshold or 1,

363 - require_uniform=True,

364 - )

357 + metadata.num_decodes,

358 + metadata.num_prefills,

359 + metadata.num_decode_tokens,

360 + num_tokens - metadata.num_decode_tokens,

361 )

362

363 + decode_query_len = 0

364 + active_num_decodes = num_decodes

365 + if num_decodes > 0:

366 + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

367 + decode_query_len = (query_start_loc_cpu[1] - query_start_loc_cpu[0]).item()

368 + assert decode_query_len > 0

369 + active_num_decodes = num_decode_tokens // decode_query_len

370 + assert active_num_decodes * decode_query_len == num_decode_tokens

371 +

372 FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode

373 fp8_metadata = FP8Meta(

369 - num_decodes=num_decodes,

374 + num_decodes=active_num_decodes,

375 num_prefills=num_prefills,

376 num_decode_tokens=num_decode_tokens,

377 num_prefill_tokens=num_prefill_tokens,

@@ -374,7 +379,6 @@ class FlashMLASparseMetadataBuilder(

379

380 # Extract prefill sequence lengths (context + query, not just query)

381 # Decode requests come first in the batch, prefill requests follow

377 - prefill_seq_lens = None

382 prefill_request_id = None

383 prefill_workspace_starts = None

384 prefill_chunks = None

@@ -386,11 +390,9 @@ class FlashMLASparseMetadataBuilder(

390 # slice below), so no D2H sync is needed.

391 seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound

392 assert seq_lens_cpu is not None

389 - seq_lens = common_attn_metadata.seq_lens

393 query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

394

395 prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:]

393 - prefill_seq_lens = seq_lens[num_decodes:]

396

397 # Build prefill_request_id: -1 for decode, request index for

398 # prefill. This enables a single

@@ -438,7 +440,6 @@ class FlashMLASparseMetadataBuilder(

440 offset = prefill_workspace_starts_cpu[chunk_start].item()

441 prefill_workspace_starts_cpu[chunk_start:chunk_end] -= offset

442

441 - chunk_seq_lens = prefill_seq_lens[chunk_start:chunk_end]

443 chunk_tot_seqlen = prefill_seq_lens_cpu[chunk_start:chunk_end].sum()

444 token_start = query_start_loc_cpu[num_decodes + chunk_start].item()

445 token_end = query_start_loc_cpu[num_decodes + chunk_end].item()

@@ -452,7 +453,6 @@ class FlashMLASparseMetadataBuilder(

453

454 prefill_chunks.append(

455 FP8Meta.Prefill.Chunk(

455 - seq_lens=chunk_seq_lens,

456 tokens_slice=tokens_slice,

457 block_table=chunk_block_table,

458 req_start_idx=chunk_start,

@@ -466,27 +466,22 @@ class FlashMLASparseMetadataBuilder(

466 )

467

468 fp8_metadata.prefill = FP8Meta.Prefill(

469 - seq_lens=prefill_seq_lens,

469 request_ids=prefill_request_id,

470 workspace_starts=prefill_workspace_starts,

471 chunks=prefill_chunks,

472 )

473

474 if num_decodes > 0:

476 - # Compute decode_query_len for spec decode (uniform due to require_uniform)

477 - query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

478 - decode_query_len = (query_start_loc_cpu[1] - query_start_loc_cpu[0]).item()

479 -

475 # Use padded head count since that's what the kernel will see

476 scheduler_metadata, _ = get_mla_metadata()

477

478 kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata(

479 scheduler_metadata=scheduler_metadata,

485 - dummy_block_table=self.dummy_block_table[:num_decodes],

486 - cache_lens=self.max_model_len_tensor[:num_decodes],

480 + dummy_block_table=self.dummy_block_table[:active_num_decodes],

481 + cache_lens=self.max_model_len_tensor[:active_num_decodes],

482 )

483 fp8_metadata.decode = FP8Meta.Decode(

489 - seq_lens=common_attn_metadata.seq_lens[:num_decodes],

484 + seq_lens=common_attn_metadata.seq_lens[:active_num_decodes],

485 kernel_metadata=kernel_meta,

486 decode_query_len=decode_query_len,

487 )

@@ -510,7 +505,7 @@ class FlashMLASparseMetadataBuilder(

505 )

506 else:

507 metadata.fp8_extra_metadata = self._build_fp8_separate_prefill_decode(

513 - common_attn_metadata

508 + common_attn_metadata, metadata

509 )

510

511 return metadata

@@ -597,9 +592,10 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

592 attn_metadata: FlashMLASparseMetadata,

593 ) -> torch.Tensor:

594 # Convert per-request indices to global slots (decode) or workspace

600 - # offsets (prefill).

595 + # offsets (prefill). req_id_per_token covers the whole batch; slice it

596 + # to the MQA tokens (q may exclude prefill tokens routed to dense MHA).

597 topk_indices, topk_length = triton_convert_req_index_to_global_index(

602 - attn_metadata.req_id_per_token,

598 + attn_metadata.req_id_per_token[: topk_indices.shape[0]],

599 attn_metadata.block_table,

600 topk_indices,

601 BLOCK_SIZE=attn_metadata.block_size,

@@ -624,11 +620,18 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

620 fp8_metadata = attn_metadata.fp8_extra_metadata

621 assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode)

622 num_decodes = fp8_metadata.num_decodes

623 + num_mqa_tokens = q.shape[0]

624 + num_decode_tokens = fp8_metadata.num_decode_tokens

625 + num_prefill_tokens = num_mqa_tokens - num_decode_tokens

626 + assert num_prefill_tokens in (0, fp8_metadata.num_prefill_tokens), (

627 + "FP8 sparse MLA expects either the decode subset or the full batch"

628 + )

629

630 prefill_request_ids = None

631 prefill_workspace_starts = None

632 has_prefill_workspace = False

631 - if fp8_metadata.prefill is not None:

633 + if num_prefill_tokens > 0:

634 + assert fp8_metadata.prefill is not None

635 prefill_request_ids = fp8_metadata.prefill.request_ids

636 prefill_workspace_starts = fp8_metadata.prefill.workspace_starts

637 has_prefill_workspace = True

@@ -640,7 +643,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

643 # prefill_workspace_starts has been adjusted in-place per chunk so

644 # prefill indices automatically come out chunk-local

645 topk_indices, topk_length = triton_convert_req_index_to_global_index(

643 - attn_metadata.req_id_per_token,

646 + attn_metadata.req_id_per_token[: topk_indices.shape[0]],

647 attn_metadata.block_table,

648 topk_indices,

649 BLOCK_SIZE=attn_metadata.block_size,

@@ -676,9 +679,6 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

679 # -> (num_decode_tokens, num_heads, head_dim_v)

680 return reshape_attn_output_for_spec_decode(attn_out)

681

679 - num_decode_tokens = fp8_metadata.num_decode_tokens

680 - num_prefill_tokens = fp8_metadata.num_prefill_tokens

681 -

682 # Pure decode: direct call without allocation

683 if num_decode_tokens > 0 and num_prefill_tokens == 0:

684 assert fp8_metadata.decode is not None

@@ -686,7 +686,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

686 else:

687 # Mixed or pure prefill: allocate output tensor

688 attn_out = q.new_empty(

689 - (attn_metadata.num_actual_tokens, self.num_heads, self.kv_lora_rank),

689 + (num_mqa_tokens, self.num_heads, self.kv_lora_rank),

690 dtype=q.dtype,

691 device=q.device,

692 )

@@ -704,7 +704,6 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

704 kv_c_and_k_pe_cache,

705 chunk_workspace,

706 chunk.block_table,

707 - chunk.seq_lens,

707 chunk.workspace_starts,

708 len(chunk.block_table),

709 )

@@ -738,7 +737,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

737 # Convert per-request indices to global slots (decode) or workspace

738 # offsets (prefill).

739 topk_indices = triton_convert_req_index_to_global_index(

741 - attn_metadata.req_id_per_token,

740 + attn_metadata.req_id_per_token[: topk_indices.shape[0]],

741 attn_metadata.block_table,

742 topk_indices,

743 BLOCK_SIZE=attn_metadata.block_size,

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

@@ -154,6 +154,11 @@ def triton_convert_req_index_to_global_index(

154 assert req_id.dtype == torch.int32

155 assert block_table.dtype == torch.int32

156 assert token_indices.dtype == torch.int32

157 + assert req_id.shape[0] == token_indices.shape[0], (

158 + f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) "

159 + "must cover the same tokens; the grid is sized by req_id but the output "

160 + "is allocated like token_indices, so a longer req_id writes out of bounds"

161 + )

162 assert token_indices.shape[1] == NUM_TOPK_TOKENS

163 assert NUM_TOPK_TOKENS % BLOCK_N == 0, (

164 f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})"