kv: nvfp4_ds_mla — 4-bit NVFP4 KV cache for B12X sparse MLA (SM120)

csrc/libtorch_stable/cache_kernels.cu

@@ -7,6 +7,11 @@
 #include "quantization/vectorization_utils.cuh"
 #include "concat_mla_q.cuh"
 
+#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
+  #define NVFP4_ENABLE_ELTS16 1
+  #include "quantization/fp4/nvfp4_utils.cuh"
+#endif
+
 #ifdef USE_ROCM
   #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
 #else
@@ -546,6 +551,97 @@ __global__ void concat_and_cache_ds_mla_kernel(
       *reinterpret_cast<const uint64_t*>(result);
 }
 
+#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
+template <typename scalar_t>
+__global__ void concat_and_cache_nvfp4_mla_kernel(
+    const scalar_t* __restrict__ kv_c,  // [num_tokens, kv_lora_rank]
+    const scalar_t* __restrict__ k_pe,  // [num_tokens, pe_dim]
+    uint8_t* __restrict__ kv_cache,     // [num_blocks, block_size, 432]
+    const int64_t* __restrict__ slot_mapping,  // [num_tokens]
+    const int block_stride,                    //
+    const int entry_stride,                    //
+    const int kv_c_stride,                     //
+    const int k_pe_stride,                     //
+    const int kv_lora_rank,                    //
+    const int pe_dim,                          //
+    const int block_size                       //
+) {
+  using CudaType = typename CUDATypeConverter<scalar_t>::Type;
+  using PVec = PackedVec<CudaType, CVT_FP4_PACK16>;
+
+  static constexpr int kNopeBytes = 256;
+  static constexpr int kScaleBytes = 32;
+  static constexpr int kPadBytes = 16;
+  static constexpr int kRopeOffset = kNopeBytes + kScaleBytes + kPadBytes;
+  static constexpr int kFp4GroupSize = CVT_FP4_SF_VEC_SIZE;
+  static constexpr int kEltsPerThread = CVT_FP4_ELTS_PER_THREAD;
+  static constexpr int kThreadsPerScale = kFp4GroupSize / kEltsPerThread;
+
+  const int64_t token_idx = blockIdx.x;
+  const int64_t slot_idx = slot_mapping[token_idx];
+  if (slot_idx < 0) {
+    return;
+  }
+
+  const int64_t block_idx = slot_idx / block_size;
+  const int64_t block_offset = slot_idx % block_size;
+  uint8_t* __restrict__ token_dst =
+      kv_cache + block_idx * block_stride + block_offset * entry_stride;
+
+  const CudaType* __restrict__ token_src =
+      reinterpret_cast<const CudaType*>(kv_c) + token_idx * kv_c_stride;
+
+  const int group_count = kv_lora_rank / kFp4GroupSize;
+  const int thread_group_count = blockDim.x / kThreadsPerScale;
+  const int thread_group = threadIdx.x / kThreadsPerScale;
+  const int thread_group_lane = threadIdx.x % kThreadsPerScale;
+
+  for (int group = thread_group; group < group_count;
+       group += thread_group_count) {
+    PVec in_vec;
+    const CudaType* __restrict__ src =
+        token_src + group * kFp4GroupSize + thread_group_lane * kEltsPerThread;
+
+#pragma unroll
+    for (int i = 0; i < kEltsPerThread / 2; ++i) {
+      in_vec.elts[i] =
+          reinterpret_cast<const typename PackedTypeConverter<CudaType>::Type*>(
+              src)[i];
+    }
+
+    uint8_t scale_byte;
+    uint8_t* scale_out = (thread_group_lane == 0) ? &scale_byte : nullptr;
+    fp4_packed_t packed =
+        cvt_warp_fp16_to_fp4<CudaType, kThreadsPerScale>(in_vec, 1.0f,
+                                                         scale_out);
+
+#if CVT_FP4_PACK16
+    uint8_t* data_dst = token_dst + group * 8;
+    reinterpret_cast<uint64_t*>(data_dst)[0] =
+        (uint64_t(packed.hi) << 32) | uint64_t(packed.lo);
+#else
+    uint8_t* data_dst = token_dst + group * 8 + thread_group_lane * 4;
+    reinterpret_cast<uint32_t*>(data_dst)[0] = packed;
+#endif
+
+    if (scale_out != nullptr) {
+      token_dst[kNopeBytes + group] = scale_byte;
+    }
+  }
+
+  for (int i = threadIdx.x; i < kPadBytes; i += blockDim.x) {
+    token_dst[kNopeBytes + kScaleBytes + i] = 0;
+  }
+
+  scalar_t* __restrict__ rope_dst =
+      reinterpret_cast<scalar_t*>(token_dst + kRopeOffset);
+  const scalar_t* __restrict__ rope_src = k_pe + token_idx * k_pe_stride;
+  for (int i = threadIdx.x; i < pe_dim; i += blockDim.x) {
+    rope_dst[i] = rope_src[i];
+  }
+}
+#endif
+
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __global__ void indexer_k_quant_and_cache_kernel(
     const scalar_t* __restrict__ k,  // [num_tokens, head_dim]
@@ -839,13 +935,23 @@ void reshape_and_cache_flash(
           kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size,         \
           reinterpret_cast<const float*>(scale.data_ptr()));
 
+void concat_and_cache_nvfp4_mla(
+    torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe,
+    torch::stable::Tensor& kv_cache, torch::stable::Tensor& slot_mapping,
+    torch::stable::Tensor& scale);
+
 void concat_and_cache_mla(
     torch::stable::Tensor& kv_c,      // [num_tokens, kv_lora_rank]
     torch::stable::Tensor& k_pe,      // [num_tokens, pe_dim]
     torch::stable::Tensor& kv_cache,  // [num_blocks, block_size, (kv_lora_rank
                                       // + pe_dim)]
     torch::stable::Tensor& slot_mapping,  // [num_tokens] or [num_actual_tokens]
     const std::string& kv_cache_dtype, torch::stable::Tensor& scale) {
+  if (kv_cache_dtype == "nvfp4_ds_mla") {
+    concat_and_cache_nvfp4_mla(kv_c, k_pe, kv_cache, slot_mapping, scale);
+    return;
+  }
+
   // NOTE(woosuk): In vLLM V1, key.size(0) can be different from
   // slot_mapping.size(0) because of padding for CUDA graphs.
   // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
@@ -902,6 +1008,62 @@ void concat_and_cache_mla(
   }
 }
 
+void concat_and_cache_nvfp4_mla(
+    torch::stable::Tensor& kv_c,      // [num_tokens, kv_lora_rank]
+    torch::stable::Tensor& k_pe,      // [num_tokens, pe_dim]
+    torch::stable::Tensor& kv_cache,  // [num_blocks, block_size, 432]
+    torch::stable::Tensor& slot_mapping,  // [num_tokens] or [num_actual_tokens]
+    torch::stable::Tensor& scale) {
+  (void)scale;
+  int num_tokens = slot_mapping.size(0);
+  int kv_lora_rank = kv_c.size(1);
+  int pe_dim = k_pe.size(1);
+
+  STD_TORCH_CHECK(kv_lora_rank == 512,
+                  "kv_lora_rank must be 512 for nvfp4_ds_mla");
+  STD_TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for nvfp4_ds_mla");
+  STD_TORCH_CHECK(kv_cache.element_size() == 1,
+                  "kv_cache must be uint8 for nvfp4_ds_mla");
+  STD_TORCH_CHECK(kv_cache.size(2) == 432,
+                  "kv_cache.size(2) must be 432 bytes for nvfp4_ds_mla");
+  STD_TORCH_CHECK(kv_c.element_size() == 2,
+                  "kv_c.element_size() must be 2 for nvfp4_ds_mla");
+  STD_TORCH_CHECK(k_pe.element_size() == 2,
+                  "k_pe.element_size() must be 2 for nvfp4_ds_mla");
+
+#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
+  int block_size = kv_cache.size(1);
+  int kv_c_stride = kv_c.stride(0);
+  int k_pe_stride = k_pe.stride(0);
+  int block_stride = kv_cache.stride(0);
+  int entry_stride = kv_cache.stride(1);
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      kv_c.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+
+  dim3 grid(num_tokens);
+  dim3 block(128);
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
+      kv_c.scalar_type(), "concat_and_cache_nvfp4_mla", [&] {
+        vllm::concat_and_cache_nvfp4_mla_kernel<scalar_t>
+            <<<grid, block, 0, stream>>>(
+                reinterpret_cast<scalar_t*>(kv_c.data_ptr()),
+                reinterpret_cast<scalar_t*>(k_pe.data_ptr()),
+                reinterpret_cast<uint8_t*>(kv_cache.data_ptr()),
+                slot_mapping.const_data_ptr<int64_t>(), block_stride,
+                entry_stride, kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim,
+                block_size);
+      });
+#else
+  (void)num_tokens;
+  STD_TORCH_CHECK(
+      false,
+      "nvfp4_ds_mla KV cache requires SM100+ (Blackwell). "
+      "Please rebuild vllm with a Blackwell-compatible CUDA target.");
+#endif
+}
+
 namespace vllm {
 
 template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>

csrc/libtorch_stable/ops.h

@@ -515,6 +515,12 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c,
                           const std::string& kv_cache_dtype,
                           torch::stable::Tensor& scale);
 
+void concat_and_cache_nvfp4_mla(torch::stable::Tensor& kv_c,
+                                torch::stable::Tensor& k_pe,
+                                torch::stable::Tensor& kv_cache,
+                                torch::stable::Tensor& slot_mapping,
+                                torch::stable::Tensor& scale);
+
 // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
 void concat_and_cache_mla_rope_fused(
     torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,

csrc/libtorch_stable/torch_bindings.cpp

@@ -835,6 +835,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
       "                     str kv_cache_dtype,"
       "                     Tensor scale) -> ()");
 
+  ops.def(
+      "concat_and_cache_nvfp4_mla(Tensor kv_c, Tensor k_pe,"
+      "                           Tensor! kv_cache,"
+      "                           Tensor slot_mapping,"
+      "                           Tensor scale) -> ()");
+
   // Rotate Q and K, then write to kv cache for MLA
   ops.def(
       "concat_and_cache_mla_rope_fused("
@@ -933,6 +939,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
   ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache));
   ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash));
   ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
+  ops.impl("concat_and_cache_nvfp4_mla",
+           TORCH_BOX(&concat_and_cache_nvfp4_mla));
   ops.impl("concat_and_cache_mla_rope_fused",
            TORCH_BOX(&concat_and_cache_mla_rope_fused));
   ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));

tests/kernels/attention/test_cache.py

@@ -840,6 +840,124 @@ def test_concat_and_cache_ds_mla(
         torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1)
 
 
+@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
+@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
+@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEEDS)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+@torch.inference_mode()
+def test_concat_and_cache_nvfp4_mla(
+    kv_lora_rank: int,
+    qk_rope_head_dim: int,
+    num_tokens: int,
+    block_size: int,
+    num_blocks: int,
+    dtype: torch.dtype,
+    seed: int,
+    device: str,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("nvfp4_ds_mla requires CUDA")
+    if current_platform.is_rocm():
+        pytest.skip("nvfp4_ds_mla is not supported on ROCm")
+    if not current_platform.has_device_capability(100):
+        pytest.skip("nvfp4_ds_mla requires SM100+ (Blackwell)")
+    if dtype.itemsize != 2:
+        pytest.skip("nvfp4_ds_mla only supports 16-bit input")
+    if kv_lora_rank != 512:
+        pytest.skip("nvfp4_ds_mla requires kv_lora_rank == 512")
+    from tests.kernels.quantization.nvfp4_utils import break_fp4_bytes
+
+    kv_cache_dtype = "nvfp4_ds_mla"
+    set_random_seed(seed)
+    torch.set_default_device(device)
+    torch.accelerator.set_device_index(device)
+
+    # 432 B/token record: 256 B packed E2M1 NoPE + 32 B E4M3 group-16
+    # scales + 16 B alignment pad + 128 B 16-bit RoPE.
+    group_size = 16
+    nope_bytes = kv_lora_rank // 2
+    num_groups = kv_lora_rank // group_size
+    pad_bytes = 16
+    rope_offset = nope_bytes + num_groups + pad_bytes
+    entry_size = rope_offset + 2 * qk_rope_head_dim
+    assert entry_size == 432
+
+    total_slots = num_blocks * block_size
+    slot_mapping_lst = random.sample(range(total_slots), num_tokens)
+    slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
+
+    kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device)
+    k_pe = torch.randn(num_tokens, qk_rope_head_dim, dtype=dtype, device=device)
+
+    # The kernel quantizes with an implicit global scale of 1.0; the scale
+    # argument keeps the concat_and_cache_mla signature family but is unused.
+    scale = torch.tensor(1.0, dtype=torch.float32, device=device)
+    kv_cache = torch.zeros(
+        num_blocks, block_size, entry_size, dtype=torch.uint8, device=device
+    )
+
+    opcheck(
+        torch.ops._C_cache_ops.concat_and_cache_nvfp4_mla,
+        (kv_c, k_pe, kv_cache, slot_mapping, scale),
+        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
+    )
+
+    # Route through the public entry point: concat_and_cache_mla dispatches
+    # to the nvfp4 op on kv_cache_dtype == "nvfp4_ds_mla".
+    ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale)
+
+    for i in range(num_tokens):
+        slot = slot_mapping_lst[i]
+        block_idx = slot // block_size
+        block_offset = slot % block_size
+        record = kv_cache[block_idx, block_offset]
+
+        # Group scales: E4M3(group_amax / 6.0). Round-to-nearest E4M3 stays
+        # within half a mantissa step (<= 6.25% relative); the slack also
+        # covers the kernel's approximate reciprocal.
+        kv_scales = (
+            record[nope_bytes : nope_bytes + num_groups]
+            .view(torch.float8_e4m3fn)
+            .float()
+        )
+        group_amax = kv_c[i].float().abs().reshape(num_groups, group_size).amax(dim=-1)
+        torch.testing.assert_close(kv_scales, group_amax / 6.0, atol=2**-9, rtol=0.08)
+
+        # NoPE payload: dequantize E2M1 nibbles x stored group scale. The
+        # E2M1 grid's largest half-gap is 1.0 (between 4 and 6), so the
+        # element error is bounded by ~1x the group scale.
+        fp4_vals = break_fp4_bytes(
+            record[:nope_bytes].unsqueeze(0), torch.float32
+        ).reshape(num_groups, group_size)
+        dequant = fp4_vals * kv_scales[:, None]
+        err = (dequant - kv_c[i].float().reshape(num_groups, group_size)).abs()
+        bound = 1.25 * kv_scales[:, None] + 2**-9
+        assert (err <= bound).all(), (
+            f"nvfp4 dequant error {err.max().item():.4f} exceeds the "
+            f"e2m1 grid bound at token {i}"
+        )
+        torch.testing.assert_close(
+            dequant.flatten(), kv_c[i].float(), atol=1.5, rtol=0.5
+        )
+
+        # The 16-byte alignment pad is zero-filled.
+        assert (record[nope_bytes + num_groups : rope_offset] == 0).all()
+
+        # RoPE lane is a verbatim 16-bit copy.
+        kv_rope = record[rope_offset:].view(dtype)
+        torch.testing.assert_close(kv_rope, k_pe[i], atol=0.0, rtol=0.0)
+
+    # Slots outside the mapping stay untouched (indexing/stride isolation).
+    written = torch.zeros(total_slots, dtype=torch.bool, device=device)
+    written[slot_mapping] = True
+    untouched = kv_cache.reshape(total_slots, entry_size)[~written]
+    assert (untouched == 0).all()
+
+
 @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
 @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
 @pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)