[SM120] Add opt-in NVFP4 sparse-MLA KV cache format for DeepSeek-V4

python/sglang/kernels/ops/attention/flash_mla_sm120.py

@@ -208,6 +208,7 @@ def _sm120_sparse_decode_fwd(
 # SM120 FlashMLA: default FlashInfer (CUTLASS SM120 sparse MLA decode).
 # Override with SGLANG_SM120_FLASHMLA_BACKEND=triton|torch to force fallback.
 _sm120_default_backend = envs.SGLANG_SM120_FLASHMLA_BACKEND.get()
+_sm120_kv_cache_format = envs.SGLANG_SM120_KV_CACHE_FORMAT.get()
 
 
 def flash_mla_with_kvcache_sm120(**kwargs):
@@ -229,6 +230,19 @@ def flash_mla_with_kvcache_sm120(**kwargs):
     extra_topk_length = kwargs.get("extra_topk_length")
 
     if _sm120_default_backend == "flashinfer":
+        if _sm120_kv_cache_format == "nvfp4":
+            return _flash_mla_flashinfer_nvfp4(
+                q,
+                k_cache,
+                indices,
+                topk_length,
+                attn_sink,
+                head_dim_v,
+                softmax_scale,
+                extra_k_cache,
+                extra_indices,
+                extra_topk_length,
+            )
         return _flash_mla_flashinfer(
             q,
             k_cache,
@@ -474,6 +488,121 @@ def _split_kv_pages_to_64(
     )
 
 
+_nvfp4_plan_cache: dict = {}
+
+
+def _nvfp4_plan(num_tokens, num_heads, topk, extra_topk, extra_page_size, device):
+    """Phase + chunks-per-block decision for the NVFP4 sparse-MLA kernels.
+
+    FlashInfer's planner is CUDA-graph safe (it skips calibration while the
+    stream is capturing and falls back to its CPB heuristic), but it is pure
+    Python, so memoize the decision per shape to keep it off the hot path.
+    """
+    key = (num_tokens, num_heads, topk, extra_topk, extra_page_size)
+    cached = _nvfp4_plan_cache.get(key)
+    if cached is not None:
+        return cached
+
+    from flashinfer.mla._sparse_mla_nvfp4_sm120_plan import (
+        NVFP4KernelVariant,
+        plan_nvfp4_sparse_mla_sm120,
+    )
+
+    planned = plan_nvfp4_sparse_mla_sm120(
+        num_tokens,
+        num_heads,
+        topk,
+        _PBS_DST,
+        device,
+        extra_topk=extra_topk,
+        extra_page_size=extra_page_size,
+        has_topk_length=True,
+        has_extra_topk_length=extra_topk > 0,
+        has_attn_sink=True,
+    )
+    if planned is None:
+        raise ValueError(
+            "no NVFP4 sparse MLA kernel serves "
+            f"T={num_tokens}, H={num_heads}, topk={topk}, extra_topk={extra_topk}"
+        )
+    decision = (planned.variant is NVFP4KernelVariant.PREFILL_STREAMING, planned.cpb)
+    _nvfp4_plan_cache[key] = decision
+    return decision
+
+
+def _flash_mla_flashinfer_nvfp4(
+    q,
+    k_cache,
+    indices,
+    topk_length,
+    attn_sink,
+    head_dim_v,
+    softmax_scale,
+    extra_k_cache,
+    extra_indices,
+    extra_topk_length,
+):
+    """FlashInfer SM120 NVFP4 sparse MLA (384 B/token cache ABI).
+
+    The pool hands over caches already grouped into the page sizes the NVFP4
+    kernels accept (64 primary, 2 or 64 extra) with no per-page padding, so
+    unlike the FP8 path there is no page-split copy here.
+    """
+    from flashinfer.mla._sparse_mla_nvfp4_sm120 import (
+        _sparse_mla_nvfp4_sm120_paged_attention,
+    )
+
+    B, _, H, _ = q.shape
+    dev = q.device
+    idx = indices.squeeze(1) if indices.dim() == 3 else indices
+    extra_idx = (
+        extra_indices.squeeze(1)
+        if extra_indices is not None and extra_indices.dim() == 3
+        else extra_indices
+    )
+
+    topk = idx.shape[-1]
+    extra_topk = extra_idx.shape[-1] if extra_idx is not None else 0
+    extra_page_size = extra_k_cache.shape[1] if extra_k_cache is not None else 0
+    use_prefill, cpb = _nvfp4_plan(B, H, topk, extra_topk, extra_page_size, dev)
+
+    output = torch.empty(B, H, head_dim_v, dtype=torch.bfloat16, device=dev)
+    out_lse = torch.empty(B, H, dtype=torch.float32, device=dev)
+
+    if use_prefill:
+        mid_out = None
+        mid_lse = None
+    else:
+        _BI = 64
+        num_splits = (topk + _BI - 1) // _BI + (
+            (extra_topk + _BI - 1) // _BI if extra_topk > 0 else 0
+        )
+        mid_out = torch.empty(
+            B, H, num_splits, head_dim_v, dtype=torch.bfloat16, device=dev
+        )
+        mid_lse = torch.empty(B, H, num_splits, dtype=torch.float32, device=dev)
+
+    _sparse_mla_nvfp4_sm120_paged_attention(
+        q.squeeze(1) if q.ndim == 4 else q,
+        k_cache,
+        idx,
+        output,
+        out_lse,
+        softmax_scale,
+        topk_length=topk_length,
+        attn_sink=attn_sink,
+        extra_kv_cache=extra_k_cache,
+        extra_indices=extra_idx,
+        extra_topk_length=extra_topk_length,
+        mid_out=mid_out,
+        mid_lse=mid_lse,
+        use_prefill=use_prefill,
+        chunks_per_block_override=cpb,
+    )
+
+    return (output.unsqueeze(1), None)
+
+
 def _flash_mla_flashinfer(
     q,
     k_cache,

python/sglang/srt/environ.py

@@ -966,6 +966,8 @@ class Envs:
     SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS = EnvInt(1)
     # SM120 FlashMLA decode backend: "flashinfer" (default), "triton", or "torch".
     SGLANG_SM120_FLASHMLA_BACKEND = EnvStr("flashinfer")
+    # SM120 DSv4 sparse-MLA KV cache ABI: "fp8" (default) or "nvfp4".
+    SGLANG_SM120_KV_CACHE_FORMAT = EnvStr("fp8")
     SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096)
     SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048)
     SGLANG_FLASHINFER_AUTOTUNE_CACHE = EnvBool(True)

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -1705,24 +1705,32 @@ def forward(
 
             swa_window_size = token_to_kv_pool.swa_window_size
             assert swa_k_cache.ndim == 2
-            k_cache_total_dim = token_to_kv_pool.swa_kv_pool.kv_cache_total_dim
-            swa_k_cache = swa_k_cache[:, : swa_window_size * k_cache_total_dim].view(
-                swa_k_cache.shape[0], swa_window_size, 1, k_cache_total_dim
-            )
+            is_nvfp4_kv = token_to_kv_pool.swa_kv_pool.is_nvfp4
+            if is_nvfp4_kv:
+                swa_k_cache = token_to_kv_pool.swa_kv_pool.nvfp4_cache_view(swa_k_cache)
+            else:
+                k_cache_total_dim = token_to_kv_pool.swa_kv_pool.kv_cache_total_dim
+                swa_k_cache = swa_k_cache[
+                    :, : swa_window_size * k_cache_total_dim
+                ].view(swa_k_cache.shape[0], swa_window_size, 1, k_cache_total_dim)
 
             if extra_k_cache is not None:
-                page_sizes = {
-                    4: token_to_kv_pool.page_size // 4,
-                    128: token_to_kv_pool.page_size // 128,
-                }
-                extra_k_cache = extra_k_cache[
-                    :, : page_sizes[compress_ratio] * k_cache_total_dim
-                ].view(
-                    extra_k_cache.shape[0],
-                    page_sizes[compress_ratio],
-                    1,
-                    k_cache_total_dim,
-                )
+                if is_nvfp4_kv:
+                    _, _, extra_pool = token_to_kv_pool.layer_mapping[layer_id]
+                    extra_k_cache = extra_pool.nvfp4_cache_view(extra_k_cache)
+                else:
+                    page_sizes = {
+                        4: token_to_kv_pool.page_size // 4,
+                        128: token_to_kv_pool.page_size // 128,
+                    }
+                    extra_k_cache = extra_k_cache[
+                        :, : page_sizes[compress_ratio] * k_cache_total_dim
+                    ].view(
+                        extra_k_cache.shape[0],
+                        page_sizes[compress_ratio],
+                        1,
+                        k_cache_total_dim,
+                    )
             swa_page_indices = core_attn_metadata.swa_page_indices
             swa_topk_lengths = core_attn_metadata.swa_topk_lengths
 

python/sglang/srt/layers/attention/dsv4/compressor_v2.py

@@ -139,6 +139,36 @@ def _get_out_loc(self, compress_ratio: int) -> torch.Tensor:
         attr_name = f"c{compress_ratio}_out_loc"
         return getattr(self.forward_metadata.core_metadata, attr_name)
 
+    def _nvfp4_stage(
+        self, compress_ratio: int, num_rows: int, device: torch.device
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """BF16 staging rows plus their identity out_loc, one row per entry.
+
+        Bucketed to a power of two and never evicted, then sliced. A decode
+        CUDA graph bakes in the buffer address at capture time, so the backing
+        tensor for a bucket must never be reallocated; bucketing also bounds
+        the buffer count, which keying on the raw row count would not (prefill
+        row counts follow the chunk's token count and vary per request).
+        Rows are zeroed each step because entries the compress kernel skips
+        must not be quantized into the pool as stale data. out_loc is int64;
+        the compress store kernel rejects int32.
+        """
+        cache = getattr(self, "_nvfp4_stage_cache", None)
+        if cache is None:
+            cache = self._nvfp4_stage_cache = {}
+        capacity = 1 << max(0, num_rows - 1).bit_length()
+        key = (compress_ratio, capacity)
+        entry = cache.get(key)
+        if entry is None:
+            entry = (
+                torch.zeros(capacity, 512, dtype=torch.bfloat16, device=device),
+                torch.arange(capacity, dtype=torch.int64, device=device),
+            )
+            cache[key] = entry
+        rows = entry[0][:num_rows]
+        rows.zero_()
+        return rows, entry[1][:num_rows]
+
     def _forward_compress_all_in_one(
         self,
         *,
@@ -241,6 +271,17 @@ def forward_unified(
             page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
             if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"):
                 out_loc = compress_kv_pool._translate_loc_to_hisparse_device(out_loc)
+            if compress_kv_pool.is_nvfp4:
+                # The all-in-one kernel only emits the FP8 cache ABI. Land its
+                # normed+roped output in a compacted BF16 staging buffer and
+                # quantize into the NVFP4 pool afterwards.
+                nvfp4_out_loc = out_loc
+                nvfp4_cache = compress_kv_pool.nvfp4_cache_view(kv_cache)
+                kv_cache, out_loc = self._nvfp4_stage(
+                    compressor.ratio, out_loc.shape[0], out_loc.device
+                )
+                page_size = 1
+                bf16_store = True
         self._forward_compress_all_in_one(
             kv_score_buffer=state_pool.kv_score_buffer.kv_score,
             kv_score_input=kv_score_input,
@@ -257,6 +298,14 @@ def forward_unified(
             use_fp4_indexer=use_fp4_indexer,
             bf16_store=bf16_store,
         )
+        if not compressor.is_in_indexer and not is_unified_kv_triton():
+            _, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id]
+            if compress_kv_pool.is_nvfp4:
+                from flashinfer.mla import nvfp4_quantize_append_sparse_mla_cache
+
+                nvfp4_quantize_append_sparse_mla_cache(
+                    kv_cache, nvfp4_out_loc, nvfp4_cache
+                )
         online_c128_mtp = getattr(self, "online_c128_mtp", None)
         if online_c128_mtp is not None:
             online_c128_mtp.write_prefix_states(