[Bugfix] Sparse MLA: enable fp8_ds_mla dense prefill

benchmarks/kernels/bench_cp_gather_fp8.py

@@ -69,12 +69,11 @@ def make_inputs(total_tokens, num_reqs, block_size):
     # Output workspace
     dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
 
-    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
     workspace_starts_t = torch.tensor(
         workspace_starts, dtype=torch.int32, device="cuda"
     )
 
-    return cache, dst, block_table, seq_lens_t, workspace_starts_t
+    return cache, dst, block_table, workspace_starts_t
 
 
 def bench_scenario(label, num_reqs, total_tokens_list, save_path):
@@ -94,15 +93,15 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
         )
     )
     def bench_fn(total_tokens, provider, num_reqs):
-        cache, dst, block_table, seq_lens_t, ws_starts = make_inputs(
+        cache, dst, block_table, ws_starts = make_inputs(
             total_tokens, num_reqs, BLOCK_SIZE
         )
 
         quantiles = [0.5, 0.2, 0.8]
 
         ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
             lambda: ops.cp_gather_and_upconvert_fp8_kv_cache(
-                cache, dst, block_table, seq_lens_t, ws_starts, num_reqs
+                cache, dst, block_table, ws_starts, num_reqs
             ),
             quantiles=quantiles,
             rep=500,

csrc/cache.h

@@ -67,9 +67,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
     torch::Tensor const& src_cache,         // [NUM_BLOCKS, BLOCK_SIZE, 656]
     torch::Tensor const& dst,               // [TOT_TOKENS, 576]
     torch::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
-    torch::Tensor const& seq_lens,          // [BATCH]
     torch::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size);
+    int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
 
 // Indexer K quantization and cache function
 void indexer_k_quant_and_cache(

csrc/libtorch_stable/cache_kernels.cu

@@ -1174,7 +1174,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache(
     const int32_t num_reqs, const int32_t block_size,
     const int32_t total_tokens, const int64_t block_table_stride,
     const int64_t cache_block_stride, const int64_t cache_entry_stride,
-    const int64_t dst_entry_stride) {
+    const int64_t dst_entry_stride,
+    const int32_t* __restrict__ seq_starts) {  // Optional source offsets
   const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
   if (flat_warp_id >= total_tokens) return;
   const int lane_id = threadIdx.x & 31;
@@ -1192,7 +1193,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache(
 
   // Compute physical token address via block table
   const int out_token_id = flat_warp_id;
-  const int token_offset = out_token_id - workspace_starts[req_id];
+  int token_offset = out_token_id - workspace_starts[req_id];
+  if (seq_starts != nullptr) token_offset += seq_starts[req_id];
   const int cache_block_idx = token_offset / block_size;
   const int offset_in_block = token_offset % block_size;
   const int physical_block =
@@ -1383,9 +1385,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
     torch::stable::Tensor const& src_cache,    // [NUM_BLOCKS, BLOCK_SIZE, 656]
     torch::stable::Tensor const& dst,          // [TOT_TOKENS, 576]
     torch::stable::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
-    torch::stable::Tensor const& seq_lens,     // [BATCH]
     torch::stable::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size) {
+    int64_t batch_size,
+    std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
   torch::stable::accelerator::DeviceGuard device_guard(
       src_cache.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream();
@@ -1396,20 +1398,25 @@ void cp_gather_and_upconvert_fp8_kv_cache(
   STD_TORCH_CHECK(
       block_table.scalar_type() == torch::headeronly::ScalarType::Int,
       "block_table must be int32");
-  STD_TORCH_CHECK(seq_lens.scalar_type() == torch::headeronly::ScalarType::Int,
-                  "seq_lens must be int32");
   STD_TORCH_CHECK(
       workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int,
       "workspace_starts must be int32");
+  if (seq_starts.has_value()) {
+    STD_TORCH_CHECK(
+        seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,
+        "seq_starts must be int32");
+  }
 
   STD_TORCH_CHECK(src_cache.device() == dst.device(),
                   "src_cache and dst must be on the same device");
   STD_TORCH_CHECK(src_cache.device() == block_table.device(),
                   "src_cache and block_table must be on the same device");
-  STD_TORCH_CHECK(src_cache.device() == seq_lens.device(),
-                  "src_cache and seq_lens must be on the same device");
   STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(),
                   "src_cache and workspace_starts must be on the same device");
+  if (seq_starts.has_value()) {
+    STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
+                    "src_cache and seq_starts must be on the same device");
+  }
   auto dtype = src_cache.scalar_type();
   STD_TORCH_CHECK(
       dtype == torch::headeronly::ScalarType::Byte ||               // uint8
@@ -1438,6 +1445,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
   constexpr int warps_per_block = 8;
   const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block;
   const int block_size_threads = warps_per_block * 32;  // 256 threads
+  const int32_t* seq_starts_ptr =
+      seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
+                             : nullptr;
 
   vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid_size, block_size_threads, 0,
                                                stream>>>(
@@ -1446,7 +1456,7 @@ void cp_gather_and_upconvert_fp8_kv_cache(
       workspace_starts.const_data_ptr<int32_t>(),
       static_cast<int32_t>(batch_size), block_size, total_tokens,
       block_table_stride, cache_block_stride, cache_entry_stride,
-      dst_entry_stride);
+      dst_entry_stride, seq_starts_ptr);
 }
 
 // Macro to dispatch the kernel based on the data type.

csrc/libtorch_stable/ops.h

@@ -527,9 +527,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
                                                     // 656]
     torch::stable::Tensor const& dst,               // [TOT_TOKENS, 576]
     torch::stable::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
-    torch::stable::Tensor const& seq_lens,          // [BATCH]
     torch::stable::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size);
+    int64_t batch_size,
+    std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
 
 // Indexer K quantization and cache function
 void indexer_k_quant_and_cache(

csrc/libtorch_stable/torch_bindings.cpp

@@ -847,8 +847,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
 
   ops.def(
       "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
-      "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
-      "batch_size) -> ()");
+      "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? "
+      "seq_starts) -> ()");
 
   ops.def(
       "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "

tests/kernels/test_cp_gather_fp8.py

@@ -27,8 +27,8 @@ def _build_test_case(seq_lens, block_size, seed=42):
         seed: Random seed for reproducibility.
 
     Returns:
-        Tuple of (cache, block_table, seq_lens_t, workspace_starts_t,
-                  num_reqs, total_tokens, expected_output).
+        Tuple of (cache, block_table, workspace_starts_t, num_reqs,
+                  total_tokens, expected_output).
     """
     torch.manual_seed(seed)
 
@@ -112,15 +112,13 @@ def _build_test_case(seq_lens, block_size, seed=42):
             # Expected output: exact copy
             expected[out_idx, NOPE_DIM:] = rope
 
-    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
     workspace_starts_t = torch.tensor(
         workspace_starts, dtype=torch.int32, device="cuda"
     )
 
     return (
         cache,
         block_table,
-        seq_lens_t,
         workspace_starts_t,
         num_reqs,
         total_tokens,
@@ -197,15 +195,13 @@ def _build_test_case_fast(seq_lens, block_size, seed=42):
         flat_cache[:sl] = token_data[ws : ws + sl]
         block_start += nb
 
-    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
     workspace_starts_t = torch.tensor(
         workspace_starts, dtype=torch.int32, device="cuda"
     )
 
     return (
         cache,
         block_table,
-        seq_lens_t,
         workspace_starts_t,
         num_reqs,
         total_tokens,
@@ -232,7 +228,6 @@ def test_cp_gather_and_upconvert_fp8_kv_cache(seq_lens, block_size):
     (
         cache,
         block_table,
-        seq_lens_t,
         workspace_starts_t,
         num_reqs,
         total_tokens,
@@ -244,7 +239,7 @@ def test_cp_gather_and_upconvert_fp8_kv_cache(seq_lens, block_size):
     )
 
     ops.cp_gather_and_upconvert_fp8_kv_cache(
-        cache, dst, block_table, seq_lens_t, workspace_starts_t, num_reqs
+        cache, dst, block_table, workspace_starts_t, num_reqs
     )
 
     # NoPE: fp8 dequant has rounding error, so we allow small tolerance.
@@ -277,8 +272,6 @@ def test_cp_gather_fp8_shuffled_blocks():
     )
     block_table = torch.tensor([[3, 1]], dtype=torch.int32, device="cuda")
     workspace_starts = torch.tensor([0], dtype=torch.int32, device="cuda")
-    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
-
     expected = torch.zeros(
         total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda"
     )
@@ -314,7 +307,7 @@ def test_cp_gather_fp8_shuffled_blocks():
     )
 
     ops.cp_gather_and_upconvert_fp8_kv_cache(
-        cache, dst, block_table, seq_lens_t, workspace_starts, len(seq_lens)
+        cache, dst, block_table, workspace_starts, len(seq_lens)
     )
 
     torch.testing.assert_close(
@@ -323,6 +316,64 @@ def test_cp_gather_fp8_shuffled_blocks():
     assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
 
 
+@pytest.mark.parametrize(
+    "gather_seq_lens,seq_starts",
+    [
+        ([6, 5, 3], [3, 4, 5]),
+        ([0, 5, 3], [12, 4, 5]),
+    ],
+)
+def test_cp_gather_fp8_with_sequence_starts(gather_seq_lens, seq_starts):
+    """Gather request slices beginning at arbitrary cache positions."""
+    full_seq_lens = [12, 11, 9]
+    (
+        cache,
+        block_table,
+        _workspace_starts_t,
+        num_reqs,
+        _total_tokens,
+        full_expected,
+    ) = _build_test_case(full_seq_lens, block_size=4)
+
+    workspace_starts = torch.tensor(
+        [0, gather_seq_lens[0], sum(gather_seq_lens[:2])],
+        dtype=torch.int32,
+        device="cuda",
+    )
+    seq_starts_t = torch.tensor(seq_starts, dtype=torch.int32, device="cuda")
+    dst = torch.empty(
+        sum(gather_seq_lens),
+        NOPE_DIM + ROPE_DIM,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+
+    ops.cp_gather_and_upconvert_fp8_kv_cache(
+        cache,
+        dst,
+        block_table,
+        workspace_starts,
+        num_reqs,
+        seq_starts_t,
+    )
+
+    full_workspace_starts = [0, full_seq_lens[0], sum(full_seq_lens[:2])]
+    expected = torch.cat(
+        [
+            full_expected[
+                full_workspace_starts[i] + seq_starts[i] : full_workspace_starts[i]
+                + seq_starts[i]
+                + gather_seq_lens[i]
+            ]
+            for i in range(num_reqs)
+        ]
+    )
+    torch.testing.assert_close(
+        dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2
+    )
+    assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:])
+
+
 @pytest.mark.parametrize(
     "seq_lens,block_size",
     [
@@ -342,7 +393,6 @@ def test_cp_gather_fp8_large_seqlens(seq_lens, block_size):
     (
         cache,
         block_table,
-        seq_lens_t,
         workspace_starts_t,
         num_reqs,
         total_tokens,
@@ -354,7 +404,7 @@ def test_cp_gather_fp8_large_seqlens(seq_lens, block_size):
     )
 
     ops.cp_gather_and_upconvert_fp8_kv_cache(
-        cache, dst, block_table, seq_lens_t, workspace_starts_t, num_reqs
+        cache, dst, block_table, workspace_starts_t, num_reqs
     )
 
     torch.testing.assert_close(

tests/v1/attention/test_mla_backends.py

@@ -42,7 +42,10 @@
 )
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
-from vllm.v1.kv_cache_interface import MLAAttentionSpec
+from vllm.v1.kv_cache_interface import (
+    KVQuantMode,
+    MLAAttentionSpec,
+)
 
 BACKENDS_TO_TEST = [
     AttentionBackendEnum.CUTLASS_MLA,
@@ -55,6 +58,31 @@
 
 DEVICE_TYPE = current_platform.device_type
 
+
+@pytest.mark.parametrize(
+    ("cache_dtype", "expected_quant_mode"),
+    [
+        ("auto", KVQuantMode.NONE),
+        ("fp8_ds_mla", KVQuantMode.FP8_PER_TENSOR),
+    ],
+)
+def test_mla_kv_cache_spec_uses_layer_cache_dtype(
+    cache_dtype: str, expected_quant_mode: KVQuantMode
+):
+    layer = SimpleNamespace(kv_cache_dtype=cache_dtype, head_size=576)
+    vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=64), model_config=None
+    )
+
+    spec = MLAAttention.get_kv_cache_spec(layer, vllm_config)
+
+    assert isinstance(spec, MLAAttentionSpec)
+    assert spec.cache_dtype_str == cache_dtype
+    assert spec.kv_quant_mode == expected_quant_mode
+    if cache_dtype == "fp8_ds_mla":
+        assert spec.page_size_bytes == 64 * 656
+
+
 # Remove sm100 backends from the list if not using sm100
 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
     BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)