From f3f8c19c0b83f2c57939cde4a6db830da4102a54 Mon Sep 17 00:00:00 2001 From: David Young Date: Sun, 5 Jul 2026 01:00:28 +0100 Subject: [PATCH 1/2] =?UTF-8?q?kv:=20nvfp4=5Fds=5Fmla=20=E2=80=94=204-bit?= =?UTF-8?q?=20NVFP4=20KV=20cache=20for=20B12X=20sparse=20MLA=20(SM120)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an opt-in "nvfp4_ds_mla" KV cache dtype for the B12X sparse-MLA backend: the 512-dim MLA latent is stored as packed NVFP4 (E2M1 data + per-16-group E4M3 scales) instead of fp8, shrinking the per-token record from 656 B to 432 B per layer (256 B FP4 NoPE + 32 B E4M3 scales + 16 B alignment pad + 128 B BF16 RoPE) for +39-48% KV pool at equal budget. Behavior is unchanged unless opted in: every change is gated on kv_cache_dtype == "nvfp4_ds_mla", and fp8_ds_mla serving takes byte-identical code paths — including the b12x call signatures. The scale_format / caps kwargs are forwarded to b12x ONLY for the FP4 record, so fp8 serving keeps working on a b12x tree without the nvfp4 read-path port. Write side: csrc concat_and_cache_nvfp4_mla, in-tree in libtorch_stable/cache_kernels.cu (+ ops.h decl, _C_cache_ops schema), guarded by ENABLE_NVFP4_SM100/SM120 with a clear error on pre-Blackwell builds. _custom_ops falls back to loading a companion vllm/_nvfp4_mla_cache_C.so iff the main build lacks the op, so the feature can also ship as an overlay on an existing image. Read side: requires the b12x ScaleFormat.NVFP4_E4M3 (== 2) decode/extend path (companion b12x PR to follow); until that lands, requesting nvfp4_ds_mla fails loudly at plan construction with an unexpected-kwarg error. B12X_MLA_SPARSE only; FLASHMLA_SPARSE still canonicalizes to fp8_ds_mla. Validated on GLM-5.2 753B @ TP4/DCP4 on 4x RTX PRO 6000 (SM120): KV pool 454,510 vs 307,547 tokens (+47.8%) at util 0.96; GPQA-Diamond 174/198 vs 175/198 for fp8 KV on the same checkpoint (statistically tied); NIAH 30/30 from 4k to 360k; needle retrieved at 460k depth; decode speed within noise of fp8 at matched context; zero OOMs. Signed-off-by: David Young Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JzPRoS8j7b78iivwSmFv4y --- csrc/libtorch_stable/cache_kernels.cu | 162 ++++++++++++++++++ csrc/libtorch_stable/ops.h | 6 + csrc/libtorch_stable/torch_bindings.cpp | 8 + vllm/_custom_ops.py | 36 ++++ vllm/config/cache.py | 1 + .../layers/attention/mla_attention.py | 7 +- vllm/utils/torch_utils.py | 3 +- .../attention/backends/mla/b12x_mla_sparse.py | 64 ++++++- vllm/v1/kv_cache_interface.py | 10 +- 9 files changed, 285 insertions(+), 12 deletions(-) diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index a1ac81cb10a4..c83f68bce531 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/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(result); } +#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120) +template +__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::Type; + using PVec = PackedVec; + + 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(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::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(in_vec, 1.0f, + scale_out); + +#if CVT_FP4_PACK16 + uint8_t* data_dst = token_dst + group * 8; + reinterpret_cast(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(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(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 __global__ void indexer_k_quant_and_cache_kernel( const scalar_t* __restrict__ k, // [num_tokens, head_dim] @@ -839,6 +935,11 @@ void reshape_and_cache_flash( kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ reinterpret_cast(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] @@ -846,6 +947,11 @@ void concat_and_cache_mla( // + 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 + <<>>( + reinterpret_cast(kv_c.data_ptr()), + reinterpret_cast(k_pe.data_ptr()), + reinterpret_cast(kv_cache.data_ptr()), + slot_mapping.const_data_ptr(), 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 diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index d60b68a5868d..c0c51a20816e 100644 --- a/csrc/libtorch_stable/ops.h +++ b/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, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 1be7217ce789..cc9666299551 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/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)); diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 16e0df0df648..dc953325fe61 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2625,11 +2625,47 @@ def concat_and_cache_mla( kv_cache_dtype: str, scale: torch.Tensor, ) -> None: + if kv_cache_dtype == "nvfp4_ds_mla": + concat_and_cache_nvfp4_mla(kv_c, k_pe, kv_cache, slot_mapping, scale) + return torch.ops._C_cache_ops.concat_and_cache_mla( kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale ) +_NVFP4_MLA_CACHE_EXT_LOADED = False + + +def _ensure_nvfp4_mla_cache_ext() -> None: + """Load the companion _C_cache_ops fragment that registers + ``concat_and_cache_nvfp4_mla`` when the main ``_C_stable_libtorch`` + build does not carry it (out-of-tree companion extension build).""" + global _NVFP4_MLA_CACHE_EXT_LOADED + if _NVFP4_MLA_CACHE_EXT_LOADED: + return + if not hasattr(torch.ops._C_cache_ops, "concat_and_cache_nvfp4_mla"): + import os as _os + + _ext = _os.path.join( + _os.path.dirname(__file__), "_nvfp4_mla_cache_C.so" + ) + torch.ops.load_library(_ext) + _NVFP4_MLA_CACHE_EXT_LOADED = True + + +def concat_and_cache_nvfp4_mla( + kv_c: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + scale: torch.Tensor, +) -> None: + _ensure_nvfp4_mla_cache_ext() + torch.ops._C_cache_ops.concat_and_cache_nvfp4_mla( + kv_c, k_pe, kv_cache, slot_mapping, scale + ) + + def concat_and_cache_mla_rope_fused( positions: torch.Tensor, q_pe: torch.Tensor, diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 9b96c64513b1..b1967df29241 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -25,6 +25,7 @@ "fp8_e5m2", "fp8_inc", "fp8_ds_mla", + "nvfp4_ds_mla", "turboquant_k8v4", "turboquant_4bit_nc", "turboquant_k3v4_nc", diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 38c911bd5971..6cbe8a30e5c2 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -354,6 +354,11 @@ def _canonicalize_sparse_mla_kv_cache_dtype( kv_cache_dtype: CacheDType, ) -> CacheDType: backend_name = attn_backend.get_name() + if backend_name == "B12X_MLA_SPARSE" and kv_cache_dtype == "nvfp4_ds_mla": + # B12X keeps the fp8_ds_mla default, but natively accepts the + # experimental nvfp4_ds_mla record (432 B/token: 256B FP4 NoPE + + # 32B E4M3 group-16 scales + 16B pad + 128B BF16 RoPE). + return kv_cache_dtype if backend_name in ( "FLASHMLA_SPARSE", "B12X_MLA_SPARSE", @@ -755,7 +760,7 @@ def forward_impl( k_c_normed = k_c_normed[:num_actual_toks, ...] k_pe = k_pe[:num_actual_toks, ...] - if fp8_attention and self.kv_cache_dtype != "fp8_ds_mla": + if fp8_attention and self.kv_cache_dtype not in ("fp8_ds_mla", "nvfp4_ds_mla"): kv_cache = kv_cache.view(current_platform.fp8_dtype()) # Sparse MLA impls only support forward_mqa (decode-style attention) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 9269fbb44d72..de291d3cf9c0 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -43,6 +43,7 @@ "fp8_per_token_head": torch.uint8, "fp8_inc": torch.float8_e4m3fn, "fp8_ds_mla": torch.uint8, + "nvfp4_ds_mla": torch.uint8, "turboquant_k8v4": torch.uint8, "turboquant_4bit_nc": torch.uint8, "turboquant_k3v4_nc": torch.uint8, @@ -75,7 +76,7 @@ def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: return ( kv_cache_dtype.startswith("fp8") or kv_cache_dtype.endswith("per_token_head") - or kv_cache_dtype == "nvfp4" + or kv_cache_dtype in ("nvfp4", "nvfp4_ds_mla") ) diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 691f7c0d8e7f..50a8bc8bc9cc 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -143,6 +143,7 @@ class B12xMLASparseBackend(AttentionBackend): "auto", "bfloat16", "fp8_ds_mla", + "nvfp4_ds_mla", "fp8", # aliases for fp8_ds_mla on this backend "fp8_e4m3", ] @@ -222,6 +223,10 @@ def get_kv_cache_shape( # scales + 128 BF16 RoPE). Mirrors the FlashMLA / SPARSE_MLA_SM120 # layout; b12x's GLM_NSA decode reads the same record. return (num_blocks, block_size, 656) + if cache_dtype_str == "nvfp4_ds_mla": + # NVFP4 MLA latent: 256 B NoPE data + 32 B E4M3 scales + + # 16 B alignment pad + 128 B BF16 RoPE. + return (num_blocks, block_size, 432) return (num_blocks, block_size, head_size) @@ -560,6 +565,22 @@ def __init__( # DCP backend. The kernel must therefore plan for, and return, the full # gathered head set; the outer layer reduces/scatters it back afterward. self._input_num_heads = self.num_heads * self.dcp_world_size + # b12x ScaleFormat.NVFP4_E4M3 == 2 selects the + # 432 B/token FP4 latent record in the unified SM120 decode/extend + # kernels; None keeps the dtype-inferred format (ARBITRARY_FP32 for + # the 656 B fp8_ds_mla record). + self._b12x_scale_format = ( + 2 if self.kv_cache_dtype == "nvfp4_ds_mla" else None + ) + # Forwarded into every plan/decode/extend b12x + # call ONLY for the FP4 record, so fp8_ds_mla serving keeps the stock + # b12x call signature (works on a b12x tree without the nvfp4-ds-mla + # read-path port; the port is required only to serve nvfp4_ds_mla). + self._b12x_nvfp4_kwargs: dict[str, Any] = ( + {} + if self._b12x_scale_format is None + else {"scale_format": self._b12x_scale_format} + ) # Split-K cap: ceil(topk / tile). Bounds the borrowed mid_out/mid_lse # chunk dim and the workspace max_chunks_per_row. @@ -616,6 +637,18 @@ def __init__( def _make_plan( mode: str, max_q_rows: int, num_q_heads: int, max_batch: int ) -> Any: + # The FP4 record needs the caps to carry the + # cache dtype + b12x ScaleFormat so the scratch planner sizes for + # the 432 B record; omit both for fp8_ds_mla so the caps stay + # constructible on a stock (pre-nvfp4-port) b12x tree. + caps_kwargs: dict[str, Any] = ( + {} + if self._b12x_scale_format is None + else { + "kv_cache_dtype": self.kv_cache_dtype, + "scale_format": self._b12x_scale_format, + } + ) return plan_sparse_mla_scratch( B12XSparseMLAScratchCaps( device=self.device, @@ -630,6 +663,7 @@ def _make_plan( max_batch=int(max_batch), max_chunks_per_row=self._num_splits_cap, page_size=self.block_size, + **caps_kwargs, ) ) @@ -708,18 +742,24 @@ def _prewarm_extend_kernels_once(self, max_batched: int) -> None: rows_to_warm = (1, 2, 4, max(1, int(max_batched))) seen_rows: set[int] = set() - # GLM fp8_ds_mla cache records are 656 B/token; the real KV cache is - # laid out (num_blocks, block_size, 656) (see the allocator at the - # block-shape branch above), so a page's stride(0) = block_size*656. - # The prewarm dummy must match that layout -- (1, block_size, 656) -- - # so _cache_block_stride_bytes sees stride >= page_size*656. The prior - # (block_size, 1, 656) shape put block_size in dim 0, giving stride(0) - # = 656 < page_size*656, which tripped the SM120 stride assertion + # GLM cache records are 656 B/token (fp8_ds_mla) or 432 B/token + # (nvfp4_ds_mla); the real KV cache is laid out + # (num_blocks, block_size, record_bytes) (see the allocator at the + # block-shape branch above), so a page's stride(0) = + # block_size*record_bytes. The prewarm dummy must match that layout -- + # (1, block_size, record_bytes) -- so _cache_block_stride_bytes sees + # stride >= page_size*record_bytes. The prior (block_size, 1, ...) + # shape put block_size in dim 0, giving stride(0) = record_bytes < + # page_size*record_bytes, which tripped the SM120 stride assertion # whenever this prewarm ran (i.e. spec + cudagraphs, the first config # to reach here; verifier-only and eager-snap both skipped it). # One page is enough: prewarm top-k indices all point at slot zero. + # Record width follows the cache dtype. + record_bytes = 432 if self.kv_cache_dtype == "nvfp4_ds_mla" else 656 kv_cache = torch.zeros( - (1, self.block_size, 656), dtype=torch.uint8, device=self.device + (1, self.block_size, record_bytes), + dtype=torch.uint8, + device=self.device, ) for rows in rows_to_warm: rows = int(rows) @@ -758,6 +798,7 @@ def _prewarm_extend_kernels_once(self, max_batched: int) -> None: v_head_dim=self.kv_lora_rank, return_lse=True, lse_scale="natural", + **self._b12x_nvfp4_kwargs, ) else: self._sparse_mla_extend_forward( @@ -765,6 +806,7 @@ def _prewarm_extend_kernels_once(self, max_batched: int) -> None: kv_cache=kv_cache, sm_scale=self.scale, v_head_dim=self.kv_lora_rank, + **self._b12x_nvfp4_kwargs, ) self._sync_warmup() @@ -887,7 +929,7 @@ def forward_mqa( kv_cache = kv_u8.reshape(-1, self.block_size, kv_u8.shape[-1]) else: raise ValueError( - "B12X_MLA_SPARSE expected fp8_ds_mla KV cache as " + "B12X_MLA_SPARSE expected fp8_ds_mla/nvfp4_ds_mla KV cache as " f"(blocks,{self.block_size},bytes) or (slots,1,bytes), got " f"{tuple(kv_u8.shape)}" ) @@ -934,6 +976,7 @@ def forward_mqa( forced_num_splits=self._num_splits_cap, return_lse=True, lse_scale="natural", + **self._b12x_nvfp4_kwargs, ), ) if self._pad_heads: @@ -951,6 +994,7 @@ def forward_mqa( sm_scale=self.scale, v_head_dim=self.kv_lora_rank, forced_num_splits=self._num_splits_cap, + **self._b12x_nvfp4_kwargs, ), ) if self._pad_heads: @@ -989,6 +1033,7 @@ def forward_mqa( v_head_dim=self.kv_lora_rank, return_lse=True, lse_scale="natural", + **self._b12x_nvfp4_kwargs, ), ) else: @@ -999,6 +1044,7 @@ def forward_mqa( kv_cache=kv_cache, sm_scale=self.scale, v_head_dim=self.kv_lora_rank, + **self._b12x_nvfp4_kwargs, ), ) if self._pad_heads: diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index d12ed8480ae6..a9d8543b820a 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -63,7 +63,7 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: return KVQuantMode.INT8_PER_TOKEN_HEAD if kv_cache_dtype == "fp8_per_token_head": return KVQuantMode.FP8_PER_TOKEN_HEAD - if kv_cache_dtype == "nvfp4": + if kv_cache_dtype in ("nvfp4", "nvfp4_ds_mla"): return KVQuantMode.NVFP4 if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"): return KVQuantMode.FP8_PER_TENSOR @@ -395,6 +395,10 @@ def storage_block_size(self) -> int: @property def real_page_size_bytes(self) -> int: + if self.cache_dtype_str == "nvfp4_ds_mla": + # 432 B/token: 256B FP4 NoPE + 32B E4M3 scales + + # 16B alignment pad + 128B BF16 RoPE. + return self.block_size * 432 if self.cache_dtype_str == "fp8_ds_mla": if self.model_version == "deepseek_v4": # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token. @@ -592,6 +596,10 @@ def storage_block_size(self) -> int: @property def real_page_size_bytes(self) -> int: + if self.cache_dtype_str == "nvfp4_ds_mla": + # 432 B/token: 256B FP4 NoPE + 32B E4M3 scales + + # 16B alignment pad + 128B BF16 RoPE. + return self.storage_block_size * 432 if self.model_version == "deepseek_v4" and self.cache_dtype_str == "fp8_ds_mla": # DeepseekV4 FlashMLA: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B # per token. FlashInfer's contiguous bf16/fp8 cache falls through to From d122ed33874127a19ee0a3df9149f8692ef5a168 Mon Sep 17 00:00:00 2001 From: David Young Date: Tue, 7 Jul 2026 11:58:48 +0100 Subject: [PATCH 2/2] =?UTF-8?q?tests:=20concat=5Fand=5Fcache=5Fnvfp4=5Fmla?= =?UTF-8?q?=20=E2=80=94=20dequant=20reference=20vs=20E2M1=20grid=20+=20gro?= =?UTF-8?q?up=20scales?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Quantizes a random MLA latent through the op and dequantizes the cache record with a torch reference (E2M1 nibble table x per-group E4M3 scales): asserts the stored scales match E4M3(group_amax/6) within half a mantissa step, bounds the per-element NoPE error by the E2M1 grid half-gap (1.25x group scale), and checks the 16-byte pad is zeroed, the 16-bit RoPE lane is copied verbatim, and unmapped slots stay untouched. Skips cleanly without CUDA, on ROCm, and below SM100. Signed-off-by: David Young Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JzPRoS8j7b78iivwSmFv4y --- tests/kernels/attention/test_cache.py | 118 ++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 4cbeb7a0b97c..2cb3d89f5a2b 100644 --- a/tests/kernels/attention/test_cache.py +++ b/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)