From 741584f8ca89c98cc6f18d2b4c6ec1e6d7e021b3 Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Tue, 27 Oct 2026 11:35:01 -0700 Subject: [PATCH] [Bugfix][MLA] Gather nvfp4_ds_mla chunked-context KV with the NVFP4 kernel (#58744) Signed-off-by: Stu Cao Signed-off-by: Lucas Wilkinson Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson Co-authored-by: vllm-agent --- csrc/cache_kernels.cu | 11 +- csrc/nvfp4_ds_mla_cache.h | 8 +- csrc/nvfp4_ds_mla_cache_kernels.cu | 19 +- csrc/ops.h | 3 +- csrc/torch_bindings.cpp | 3 +- tests/kernels/attention/test_cache.py | 43 ++++ tests/test_config.py | 21 ++ .../test_mla_prefill_context_nvfp4.py | 241 ++++++++++++++++++ vllm/_custom_ops.py | 4 + vllm/config/vllm.py | 10 + .../layers/attention/mla_attention.py | 9 + 11 files changed, 358 insertions(+), 14 deletions(-) create mode 100644 tests/v1/attention/test_mla_prefill_context_nvfp4.py diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index dbcb6ebaef9..9567d87f443 100644 --- a/csrc/cache_kernels.cu +++ b/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 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(), dst.data_ptr(), block_table.const_data_ptr(), workspace_starts.const_data_ptr(), + seq_starts.has_value() ? seq_starts.value().const_data_ptr() + : nullptr, static_cast(batch_size), block_size, total_tokens, block_table_stride, cache_block_stride, cache_entry_stride, dst_entry_stride, stream); diff --git a/csrc/nvfp4_ds_mla_cache.h b/csrc/nvfp4_ds_mla_cache.h index 2b5f757846a..b9aa9464813 100644 --- a/csrc/nvfp4_ds_mla_cache.h +++ b/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 diff --git a/csrc/nvfp4_ds_mla_cache_kernels.cu b/csrc/nvfp4_ds_mla_cache_kernels.cu index b1f369e5ab0..3b80f73de6c 100644 --- a/csrc/nvfp4_ds_mla_cache_kernels.cu +++ b/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<<>>( 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 diff --git a/csrc/ops.h b/csrc/ops.h index 1c5ddc0beb4..401f6b4996a 100644 --- a/csrc/ops.h +++ b/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 seq_starts = std::nullopt); // Indexer K quantization and cache function void indexer_k_quant_and_cache( diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 1c84b53ece2..fef3076056e 100644 --- a/csrc/torch_bindings.cpp +++ b/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 " diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 4bc0ce1a67b..149aedd30fc 100644 --- a/tests/kernels/attention/test_cache.py +++ b/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) diff --git a/tests/test_config.py b/tests/test_config.py index 0d21e52dfd0..1f71e1c966d 100644 --- a/tests/test_config.py +++ b/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 diff --git a/tests/v1/attention/test_mla_prefill_context_nvfp4.py b/tests/v1/attention/test_mla_prefill_context_nvfp4.py new file mode 100644 index 00000000000..4c5a484d56b --- /dev/null +++ b/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) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 09a45982698..c8574a53bcc 100644 --- a/vllm/_custom_ops.py +++ b/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, ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index bbeb12c20cb..078da543555 100644 --- a/vllm/config/vllm.py +++ b/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 diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index c8b9498af16..173c40dc89d 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/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 "