[Bugfix][MLA] Gather nvfp4_ds_mla chunked-context KV with the NVFP4 kernel (#58744)

Signed-off-by: Stu Cao <stu@primeintellect.ai>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

csrc/cache_kernels.cu

@@ -1879,7 +1879,7 @@ void cp_gather_and_upconvert_nvfp4_kv_cache(
     torch::stable::Tensor const& dst,               // [TOT_TOKENS, 576]
     torch::stable::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
     torch::stable::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size) {
+    int64_t batch_size, std::optional<torch::stable::Tensor> seq_starts) {
   // The kernel lives in nvfp4_ds_mla_cache_kernels.cu, which is only built when
   // the target list contains an SM100 arch.
 #if !defined(USE_ROCM) && defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
@@ -1902,6 +1902,13 @@ void cp_gather_and_upconvert_nvfp4_kv_cache(
   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() == seq_starts.value().device(),
+                    "src_cache and seq_starts must be on the same device");
+  }
   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(),
@@ -1929,6 +1936,8 @@ void cp_gather_and_upconvert_nvfp4_kv_cache(
       src_cache.const_data_ptr<uint8_t>(), dst.data_ptr(),
       block_table.const_data_ptr<int32_t>(),
       workspace_starts.const_data_ptr<int32_t>(),
+      seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
+                             : nullptr,
       static_cast<int32_t>(batch_size), block_size, total_tokens,
       block_table_stride, cache_block_stride, cache_entry_stride,
       dst_entry_stride, stream);

csrc/nvfp4_ds_mla_cache.h

@@ -29,9 +29,9 @@ void launch_concat_and_cache_nvfp4_ds_mla(const void* kv_c, const void* k_pe,
 // Gather an nvfp4_ds_mla cache into a bf16 [total_tokens, 576] workspace.
 void launch_cp_gather_and_upconvert_nvfp4_kv_cache(
     const uint8_t* src_cache, void* dst, const int32_t* block_table,
-    const int32_t* workspace_starts, int32_t num_reqs, int32_t block_size,
-    int32_t total_tokens, int64_t block_table_stride,
-    int64_t cache_block_stride, int64_t cache_entry_stride,
-    int64_t dst_entry_stride, cudaStream_t stream);
+    const int32_t* workspace_starts, const int32_t* seq_starts,
+    int32_t num_reqs, int32_t block_size, int32_t total_tokens,
+    int64_t block_table_stride, int64_t cache_block_stride,
+    int64_t cache_entry_stride, int64_t dst_entry_stride, cudaStream_t stream);
 
 }  // namespace vllm

csrc/nvfp4_ds_mla_cache_kernels.cu

@@ -183,6 +183,7 @@ __global__ void cp_gather_and_upconvert_nvfp4_kv_cache_kernel(
     __nv_bfloat16* __restrict__ dst,          // [total_tokens, 576]
     const int32_t* __restrict__ block_table,  // [num_reqs, BLOCK_INDICES]
     const int32_t* __restrict__ workspace_starts,  // [num_reqs]
+    const int32_t* __restrict__ seq_starts,        // [num_reqs] or nullptr
     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,
@@ -203,7 +204,10 @@ __global__ void cp_gather_and_upconvert_nvfp4_kv_cache_kernel(
   const int req_id = lo;
 
   const int out_token_id = flat_warp_id;
-  const int token_offset = out_token_id - workspace_starts[req_id];
+  // seq_starts offsets the read into the request's context, so a chunk that
+  // continues a request reads from its start instead of the first token.
+  const int seq_start = seq_starts == nullptr ? 0 : seq_starts[req_id];
+  const int token_offset = out_token_id - workspace_starts[req_id] + seq_start;
   const int cache_block_idx = token_offset / block_size;
   const int offset_in_block = token_offset % block_size;
   const int physical_block =
@@ -254,18 +258,19 @@ void launch_concat_and_cache_nvfp4_ds_mla(const void* kv_c, const void* k_pe,
 
 void launch_cp_gather_and_upconvert_nvfp4_kv_cache(
     const uint8_t* src_cache, void* dst, const int32_t* block_table,
-    const int32_t* workspace_starts, int32_t num_reqs, int32_t block_size,
-    int32_t total_tokens, int64_t block_table_stride,
-    int64_t cache_block_stride, int64_t cache_entry_stride,
-    int64_t dst_entry_stride, cudaStream_t stream) {
+    const int32_t* workspace_starts, const int32_t* seq_starts,
+    int32_t num_reqs, int32_t block_size, int32_t total_tokens,
+    int64_t block_table_stride, int64_t cache_block_stride,
+    int64_t cache_entry_stride, int64_t dst_entry_stride, cudaStream_t stream) {
   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
   cp_gather_and_upconvert_nvfp4_kv_cache_kernel<<<grid_size, block_size_threads,
                                                   0, stream>>>(
       src_cache, reinterpret_cast<__nv_bfloat16*>(dst), block_table,
-      workspace_starts, num_reqs, block_size, total_tokens, block_table_stride,
-      cache_block_stride, cache_entry_stride, dst_entry_stride);
+      workspace_starts, seq_starts, num_reqs, block_size, total_tokens,
+      block_table_stride, cache_block_stride, cache_entry_stride,
+      dst_entry_stride);
 }
 
 }  // namespace vllm

csrc/ops.h

@@ -826,7 +826,8 @@ void cp_gather_and_upconvert_nvfp4_kv_cache(
     torch::stable::Tensor const& dst,               // [TOT_TOKENS, 576]
     torch::stable::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
     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/torch_bindings.cpp

@@ -1132,7 +1132,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
 
   ops.def(
       "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, "
-      "Tensor block_table, Tensor workspace_starts, int batch_size) -> ()");
+      "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? "
+      "seq_starts=None) -> ()");
 
   ops.def(
       "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "

tests/kernels/attention/test_cache.py

@@ -1454,6 +1454,49 @@ def test_gather_and_maybe_dequant_cache_mla_with_seq_starts(
     torch.testing.assert_close(dst, expected)
 
 
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="nvfp4_ds_mla requires SM100",
+)
+@torch.inference_mode()
+def test_cp_gather_and_upconvert_nvfp4_kv_cache_with_seq_starts():
+    block_size, num_blocks = 16, 32
+    cache = torch.randint(
+        0, 256, (num_blocks, block_size, NVFP4_DS_MLA_ENTRY_SIZE), dtype=torch.uint8
+    ).cuda()
+    # Gathering every slot in order gives the per-slot reference values.
+    ref = torch.empty(num_blocks * block_size, 576, dtype=torch.bfloat16).cuda()
+    ops.cp_gather_and_upconvert_nvfp4_kv_cache(
+        cache,
+        ref,
+        torch.arange(num_blocks, dtype=torch.int32).cuda()[None],
+        torch.zeros(1, dtype=torch.int32).cuda(),
+        1,
+    )
+
+    seq_starts = torch.tensor([3, 32, 0], dtype=torch.int32)
+    seq_lens = torch.tensor([20, 10, 16], dtype=torch.int32)
+    block_table = torch.stack(
+        [torch.randperm(num_blocks, dtype=torch.int32) for _ in seq_lens]
+    )
+    slots = []
+    for bt, start, n in zip(block_table, seq_starts.tolist(), seq_lens.tolist()):
+        pos = torch.arange(start, start + n)
+        slots.append(bt[pos // block_size] * block_size + pos % block_size)
+    dst = torch.empty(int(seq_lens.sum()), 576, dtype=torch.bfloat16).cuda()
+    ops.cp_gather_and_upconvert_nvfp4_kv_cache(
+        cache,
+        dst,
+        block_table.cuda(),
+        (seq_lens.cumsum(0) - seq_lens).int().cuda(),
+        len(seq_lens),
+        seq_starts.cuda(),
+    )
+    torch.testing.assert_close(
+        dst, ref[torch.cat(slots).cuda()], atol=0, rtol=0, equal_nan=True
+    )
+
+
 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
 @pytest.mark.parametrize("use_seq_starts", [False, True])
 @pytest.mark.parametrize("device", CUDA_DEVICES)

tests/test_config.py

@@ -354,6 +354,27 @@ def test_mismatched_float_kv_cache_is_rejected(cache_dtype, model_dtype, should_
         config._reject_mismatched_float_kv_cache()
 
 
+@pytest.mark.parametrize("dcp_size", [1, 2])
+def test_nvfp4_ds_mla_rejects_decode_context_parallelism(dcp_size):
+    # The nvfp4_ds_mla chunked-context gather is not validated under DCP, so
+    # the combination fails at startup instead of on the first request that
+    # carries cached context.
+    kwargs = dict(
+        cache_config=CacheConfig(cache_dtype="nvfp4_ds_mla"),
+        device_config=DeviceConfig(device="cpu"),
+        parallel_config=ParallelConfig(
+            tensor_parallel_size=2,
+            decode_context_parallel_size=dcp_size,
+            distributed_executor_backend="mp",
+        ),
+    )
+    if dcp_size > 1:
+        with pytest.raises(ValueError, match="decode context parallelism"):
+            VllmConfig(**kwargs)
+    else:
+        VllmConfig(**kwargs)
+
+
 def test_per_request_spec_decode_metrics_requires_spec_decode():
     # The flag only makes sense with speculative decoding configured; enabling
     # it without --speculative-config should fail fast rather than silently

tests/v1/attention/test_mla_prefill_context_nvfp4.py

@@ -0,0 +1,241 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Chunked-context MLA prefill with an nvfp4_ds_mla KV cache.
+
+``MLACommonBaseImpl._compute_prefill_context`` gathers each chunk's context KV
+into a BF16 workspace. For nvfp4_ds_mla it must use the NVFP4 upconvert gather
+(the generic one rejects the dtype) and pass the chunk starts through, so a
+continuation chunk reads from its start rather than from the request's first
+token.
+
+The reference runs the same method on a BF16 cache that holds the dequantized
+values, through the generic gather, which applies ``seq_starts`` itself. Every
+chunk must then see bit-identical Q, K and V.
+"""
+
+from types import SimpleNamespace
+from typing import cast
+
+import pytest
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.attention.mla_attention import (
+    MLACommonBaseImpl,
+    MLACommonPrefillMetadata,
+    build_mla_chunked_context_metadata,
+)
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="nvfp4_ds_mla requires SM100 (Blackwell)",
+)
+
+_KV_LORA_RANK = 512
+_QK_NOPE = 128
+_QK_ROPE = 64
+_V_HEAD_DIM = 128
+_ENTRY = _KV_LORA_RANK + _QK_ROPE
+_NVFP4_ENTRY = 352  # 256 e2m1 NoPE + 64 e4m3 RoPE + 32 e4m3 scale factors
+_NUM_HEADS = 2
+_BLOCK_SIZE = 64
+_WORKSPACE_TOKENS = 512
+# The second chunk continues the first request and also includes a fresh one.
+_CONTEXT_LENS = [900, 100, 64, 0]
+_QUERY_LENS = [8, 4, 6, 5]
+
+
+class _RecordingPrefillBackend:
+    """Records the Q/K/V each context chunk attends over."""
+
+    def __init__(self) -> None:
+        self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
+
+    def run_prefill_context_chunk(self, *, chunk, q, k, v):
+        self.calls.append((q.float().clone(), k.float().clone(), v.float().clone()))
+        num_q = q.shape[0]
+        out = torch.full(
+            (num_q, _NUM_HEADS, _V_HEAD_DIM),
+            (k.float().mean() + v.float().mean()).item(),
+            device=q.device,
+            dtype=torch.bfloat16,
+        )
+        lse = torch.full(
+            (_NUM_HEADS, num_q), 1.0 + chunk.index, device=q.device, dtype=torch.float32
+        )
+        return out, lse
+
+
+class _KVBProj(torch.nn.Module):
+    """Stand-in for ``kv_b_proj`` with a BF16 weight (returns (out, bias))."""
+
+    def __init__(self, device: torch.device) -> None:
+        super().__init__()
+        weight = torch.randn(
+            _NUM_HEADS * (_QK_NOPE + _V_HEAD_DIM),
+            _KV_LORA_RANK,
+            device=device,
+            dtype=torch.bfloat16,
+        )
+        self.register_buffer("weight", weight * 0.05)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
+        return torch.nn.functional.linear(x, self.weight), None
+
+
+class _Impl:
+    """Only the attributes the context loop reads."""
+
+    _compute_prefill_context = MLACommonBaseImpl._compute_prefill_context
+    _concat_k_nope_k_pe = MLACommonBaseImpl._concat_k_nope_k_pe
+    _use_flashinfer_concat_mla_k = False
+
+    def __init__(self, kv_b_proj: _KVBProj, kv_cache_dtype: str) -> None:
+        self.kv_b_proj = kv_b_proj
+        self.kv_cache_dtype = kv_cache_dtype
+        self.kv_lora_rank = _KV_LORA_RANK
+        self.num_heads = _NUM_HEADS
+        self.qk_nope_head_dim = _QK_NOPE
+        self.qk_rope_head_dim = _QK_ROPE
+        self.v_head_dim = _V_HEAD_DIM
+
+
+def _build_prefill_metadata(
+    device: torch.device,
+    q_data_type: torch.dtype,
+    block_table: torch.Tensor,
+    backend: _RecordingPrefillBackend,
+) -> MLACommonPrefillMetadata:
+    query_start_loc_cpu = torch.zeros(len(_QUERY_LENS) + 1, dtype=torch.int32)
+    query_start_loc_cpu[1:] = torch.tensor(_QUERY_LENS, dtype=torch.int32).cumsum(0)
+    # Packed DS-MLA caches are always upconverted into a BF16 workspace.
+    workspace = torch.empty(
+        (_WORKSPACE_TOKENS, _ENTRY), dtype=torch.bfloat16, device=device
+    )
+    chunked_context = build_mla_chunked_context_metadata(
+        context_lens_cpu=torch.tensor(_CONTEXT_LENS, dtype=torch.int32),
+        prefill_query_start_loc_cpu=query_start_loc_cpu,
+        chunked_prefill_workspace=workspace,
+        chunked_prefill_workspace_size=_WORKSPACE_TOKENS,
+        block_size=_BLOCK_SIZE,
+        align_chunk_to_block=True,
+        device=device,
+        dcp_world_size=1,
+        dcp_local_block_size=1,
+        dcp_virtual_block_size=1,
+    )
+    assert chunked_context is not None
+    assert any(
+        c.is_continuation and c.num_requests > 1 for c in chunked_context.chunks
+    ), "the batch must exercise a mixed continuation chunk"
+    return MLACommonPrefillMetadata(
+        block_table=block_table,
+        query_start_loc=query_start_loc_cpu.to(device),
+        max_query_len=max(_QUERY_LENS),
+        chunked_context=chunked_context,
+        q_data_type=q_data_type,
+        output_dtype=torch.bfloat16,
+        prefill_backend=backend,
+    )
+
+
+@torch.inference_mode()
+def test_nvfp4_prefill_context_matches_dequantized_cache() -> None:
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    # The sparse MLA builders, the only ones that accept nvfp4_ds_mla, hand this
+    # path a BF16 query and a BF16 workspace.
+    q_data_type = torch.bfloat16
+
+    num_reqs = len(_CONTEXT_LENS)
+    max_blocks = (max(_CONTEXT_LENS) + max(_QUERY_LENS)) // _BLOCK_SIZE + 1
+    num_blocks = num_reqs * max_blocks
+    block_table = (
+        torch.randperm(num_blocks, device=device)
+        .to(torch.int32)
+        .view(num_reqs, max_blocks)
+    )
+
+    nvfp4_cache = torch.zeros(
+        (num_blocks, _BLOCK_SIZE, _NVFP4_ENTRY), dtype=torch.uint8, device=device
+    )
+    scale = torch.ones(1, dtype=torch.float32, device=device)
+    for req, context_len in enumerate(_CONTEXT_LENS):
+        if context_len == 0:
+            continue
+        pos = torch.arange(context_len, device=device)
+        slots = block_table[req, pos // _BLOCK_SIZE].long() * _BLOCK_SIZE
+        ops.concat_and_cache_mla(
+            torch.randn(
+                context_len, _KV_LORA_RANK, device=device, dtype=torch.bfloat16
+            ),
+            torch.randn(context_len, _QK_ROPE, device=device, dtype=torch.bfloat16),
+            nvfp4_cache,
+            slots + pos % _BLOCK_SIZE,
+            kv_cache_dtype="nvfp4_ds_mla",
+            scale=scale,
+        )
+
+    # Dequantize every slot once to build the BF16 reference cache.
+    bf16_cache = torch.empty(
+        (num_blocks * _BLOCK_SIZE, _ENTRY), dtype=torch.bfloat16, device=device
+    )
+    ops.cp_gather_and_upconvert_nvfp4_kv_cache(
+        src_cache=nvfp4_cache,
+        dst=bf16_cache,
+        block_table=torch.arange(num_blocks, dtype=torch.int32, device=device)[None],
+        workspace_starts=torch.zeros(1, dtype=torch.int32, device=device),
+        batch_size=1,
+    )
+    bf16_cache = bf16_cache.view(num_blocks, _BLOCK_SIZE, _ENTRY)
+
+    kv_b_proj = _KVBProj(device)
+    q = (
+        torch.randn(
+            (sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE),
+            device=device,
+            dtype=torch.bfloat16,
+        )
+        * 0.2
+    ).to(q_data_type)
+
+    backend = _RecordingPrefillBackend()
+    out, lse = cast(
+        MLACommonBaseImpl, _Impl(kv_b_proj, "nvfp4_ds_mla")
+    )._compute_prefill_context(
+        q,
+        nvfp4_cache,
+        SimpleNamespace(
+            prefill=_build_prefill_metadata(device, q_data_type, block_table, backend)
+        ),
+        scale,
+    )
+    ref_backend = _RecordingPrefillBackend()
+    ref_out, ref_lse = cast(
+        MLACommonBaseImpl, _Impl(kv_b_proj, "auto")
+    )._compute_prefill_context(
+        q,
+        bf16_cache,
+        SimpleNamespace(
+            prefill=_build_prefill_metadata(
+                device, q_data_type, block_table, ref_backend
+            )
+        ),
+        scale,
+    )
+
+    assert len(backend.calls) == len(ref_backend.calls) > 1
+    for chunk_idx, (call, ref_call) in enumerate(
+        zip(backend.calls, ref_backend.calls, strict=True)
+    ):
+        for name, actual, expected in zip(("q", "k", "v"), call, ref_call, strict=True):
+            torch.testing.assert_close(
+                actual,
+                expected,
+                atol=0,
+                rtol=0,
+                msg=lambda m, n=name, i=chunk_idx: f"chunk {i} {n} differs: {m}",
+            )
+    torch.testing.assert_close(out, ref_out, atol=0, rtol=0)
+    torch.testing.assert_close(lse, ref_lse, atol=0, rtol=0)

vllm/_custom_ops.py

@@ -3177,6 +3177,7 @@ def cp_gather_and_upconvert_nvfp4_kv_cache(
     block_table: torch.Tensor,
     workspace_starts: torch.Tensor,
     batch_size: int,
+    seq_starts: torch.Tensor | None = None,
 ) -> None:
     """Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace.
 
@@ -3186,6 +3187,8 @@ def cp_gather_and_upconvert_nvfp4_kv_cache(
         block_table: Block indices [num_reqs, max_blocks]
         workspace_starts: Workspace start offsets [num_reqs]
         batch_size: Number of requests
+        seq_starts: Per-request token offset into the context to start
+            reading from [num_reqs]; None reads from the first token.
 
     """
     torch.ops._C_cache_ops.cp_gather_and_upconvert_nvfp4_kv_cache(
@@ -3194,6 +3197,7 @@ def cp_gather_and_upconvert_nvfp4_kv_cache(
         block_table,
         workspace_starts,
         batch_size,
+        seq_starts,
     )
 
 

vllm/config/vllm.py

@@ -3421,6 +3421,16 @@ class VllmConfig:
 
     @model_validator(mode="after")
     def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig":
+        if (
+            self.cache_config.cache_dtype == "nvfp4_ds_mla"
+            and self.parallel_config.decode_context_parallel_size > 1
+        ):
+            raise ValueError(
+                "nvfp4_ds_mla KV cache is not supported with decode context "
+                "parallelism: the chunked-context prefill gather is not "
+                "validated under DCP. Use decode_context_parallel_size=1 or a "
+                "different --kv-cache-dtype."
+            )
         if self.model_config is None:
             return self
         # The ds_mla layouts are MLA-only by construction; the plain nvfp4

vllm/model_executor/layers/attention/mla_attention.py

@@ -2999,6 +2999,15 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
                     batch_size=chunk.num_requests,
                     seq_starts=chunk.starts,
                 )
+            elif self.kv_cache_dtype == "nvfp4_ds_mla":
+                ops.cp_gather_and_upconvert_nvfp4_kv_cache(
+                    src_cache=kv_c_and_k_pe_cache,
+                    dst=workspace[:toks],
+                    block_table=chunk_block_table,
+                    workspace_starts=chunk.cu_seq_lens,
+                    batch_size=chunk.num_requests,
+                    seq_starts=chunk.starts,
+                )
             elif current_platform.is_cpu():
                 assert not is_quantized_kv_cache(self.kv_cache_dtype), (
                     "CPU MLA context gather fallback only supports "