From bb27d7b89fba88108af10e7cd29cbb1e30de5958 Mon Sep 17 00:00:00 2001 From: Ali Chen Date: Wed, 9 Sep 2026 00:53:26 -0700 Subject: [PATCH 1/2] [SM120] Add opt-in NVFP4 sparse-MLA KV cache format for DeepSeek-V4 FlashInfer PR #4955 added native NVFP4 sparse-MLA kernels for SM120/SM121 alongside the existing FP8 ones. This wires them into the DSv4 attention path behind SGLANG_SM120_KV_CACHE_FORMAT, which defaults to "fp8" and leaves every existing code path untouched. The NVFP4 cache ABI is 384 B/token (packed E2M1 nope + BF16 rope + E4M3 group scales) against FP8's 584 B, and it has no per-page padding, so the pool groups it into the page sizes the kernels accept (64 primary, 2 or 64 extra) and the flat slot space stays identity-mapped. That also removes the 256->64 page-split copy the FP8 path performs on every attention call. Three fused kernels only emit the FP8 ABI, so under NVFP4: - the SWA store takes the existing bf16-intermediate path (SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE) and quantizes on store; - the compressed c4/c128 store lands in a BF16 staging buffer via the existing bf16_store path and is quantized afterwards; - the pre-quantized FP8 write path asserts instead of corrupting silently. Gains are shape-dependent and only materialize when attention runs all 64 heads per rank (attn_tp=1, i.e. DP-attention, PP, or CP). Under attention TP the NVFP4 prefill kernel is slower than FP8, so this stays opt-in. Co-Authored-By: Claude Opus 5 (1M context) --- .../kernels/ops/attention/flash_mla_sm120.py | 129 ++++++++++++++++++ python/sglang/srt/environ.py | 2 + .../layers/attention/deepseek_v4_backend.py | 40 +++--- .../layers/attention/dsv4/compressor_v2.py | 49 +++++++ .../srt/mem_cache/deepseek_v4_memory_pool.py | 67 ++++++++- python/sglang/srt/models/deepseek_v4.py | 5 +- 6 files changed, 274 insertions(+), 18 deletions(-) diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py index 6f39fa05b252..309d3fb3f166 100644 --- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py +++ b/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, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a81b8c0695a9..714312774662 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 2f1f1f2f3536..16a7b651e0c3 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index b9268e8c6eeb..fbafbdb6c228 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 54eabcb7e878..4ce8b6ae34bf 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -2,6 +2,7 @@ import logging from contextlib import nullcontext +from functools import lru_cache from typing import List, Literal, NamedTuple, Optional, Tuple import torch @@ -22,7 +23,7 @@ from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool from sglang.srt.mem_cache.memory_pool import KVCache from sglang.srt.runtime_context import get_exec, get_spec -from sglang.srt.utils import ceil_div, is_hip +from sglang.srt.utils import ceil_div, is_hip, is_sm120_supported logger = logging.getLogger(__name__) @@ -30,6 +31,23 @@ ONLINE_C128 = not _is_hip and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get() +# FlashInfer SM120 NVFP4 sparse-MLA cache ABI: 224 B packed E2M1 nope + +# 128 B BF16 rope + 32 B E4M3 group scales. +NVFP4_BYTES_PER_TOKEN = 384 + + +@lru_cache(maxsize=1) +def is_nvfp4_kv_cache() -> bool: + """Whether the DSv4 KV pools use FlashInfer's NVFP4 sparse-MLA cache ABI.""" + if envs.SGLANG_SM120_KV_CACHE_FORMAT.get() != "nvfp4": + return False + if not is_sm120_supported(): + raise ValueError( + "SGLANG_SM120_KV_CACHE_FORMAT=nvfp4 requires SM120/SM121; FlashInfer " + "builds the NVFP4 sparse-MLA kernels for consumer Blackwell only." + ) + return True + def get_compress_state_ring_size( compress_ratio: int, is_speculative: bool = False @@ -87,6 +105,12 @@ def __init__( self.quantize_block_size = 64 self.rope_storage_dtype = torch.bfloat16 self.k_with_scale_buffer_dtype = torch.int8 + self.is_nvfp4 = is_nvfp4_kv_cache() + # FlashInfer's NVFP4 sparse-MLA kernels accept a primary page size of + # 64 and an extra-cache page size of 2 or 64. The NVFP4 ABI has no + # per-page padding, so regrouping a flat slot space into 64-token + # pages is pure reindexing and leaves slot ids unchanged. + self.nvfp4_page_size = min(self.page_size, 64) self._create_buffers() def _create_buffers(self): @@ -104,6 +128,8 @@ def _create_buffers(self): ] def get_bytes_per_token(self) -> int: + if self.is_nvfp4: + return NVFP4_BYTES_PER_TOKEN dim_per_token = ( self.qk_nope_head_dim + self.qk_rope_head_dim * self.rope_storage_dtype.itemsize @@ -112,8 +138,32 @@ def get_bytes_per_token(self) -> int: ) return dim_per_token + def nvfp4_cache_view(self, buf: torch.Tensor) -> torch.Tensor: + """[num_nvfp4_pages, nvfp4_page_size, 384] uint8 view of a raw buffer. + + Both the append helper and the attention kernel derive the page split + from this shape, so writers and readers must use this single view. + ``get_key_buffer`` hands out an ``fp8_e4m3`` view of the raw storage + under ``--kv-cache-dtype fp8_e4m3``; the NVFP4 ABI is opaque bytes. + """ + if buf.dtype != torch.uint8: + buf = buf.view(torch.uint8) + return buf.view(-1, self.nvfp4_page_size, NVFP4_BYTES_PER_TOKEN) + def create_buffer(self, *, num_pages: int): bytes_per_token = self.get_bytes_per_token() + if self.is_nvfp4: + # The NVFP4 ABI packs each page exactly, so the slot space is + # regrouped into nvfp4_page_size-token rows with no padding. + assert self.store_dtype == torch.uint8 + self.kv_cache_total_dim = bytes_per_token + self.bytes_per_page_padded = self.nvfp4_page_size * bytes_per_token + return torch.zeros( + ceil_div(num_pages * self.page_size, self.nvfp4_page_size), + self.bytes_per_page_padded, + dtype=self.store_dtype, + device=self.device, + ) self.kv_cache_total_dim = bytes_per_token bytes_per_page_non_padded = self.page_size * bytes_per_token self.bytes_per_page_padded = ceil_div(bytes_per_page_non_padded, 576) * 576 @@ -137,6 +187,9 @@ def set_key_buffer( loc: torch.Tensor, cache_nope_fp8_rope_bf16_pack: NopeFp8RopeBf16Pack, ): + assert ( + not self.is_nvfp4 + ), "pre-quantized FP8 KV writes are not supported by the NVFP4 cache ABI" dsv4_index_buf_accessor.SetKAndS.execute( pool=self, buf=self.kv_buffer[layer_id], @@ -150,6 +203,14 @@ def set_key_buffer_fused( loc: torch.Tensor, cache_k: torch.Tensor, ) -> None: + if self.is_nvfp4: + from flashinfer.mla import nvfp4_quantize_append_sparse_mla_cache + + return nvfp4_quantize_append_sparse_mla_cache( + cache_k, + loc, + self.nvfp4_cache_view(self.kv_buffer[layer_id]), + ) return fused_store_cache( input=cache_k, cache=self.kv_buffer[layer_id], @@ -1187,6 +1248,10 @@ def set_swa_key_buffer_radix_fused_norm_rope( freqs_cis: torch.Tensor, positions: torch.Tensor, ) -> None: + assert not self.swa_kv_pool.is_nvfp4, ( + "the fused norm+rope+store kernel writes the FP8 cache ABI; " + "SGLANG_SM120_KV_CACHE_FORMAT=nvfp4 needs the unfused store path" + ) fused_k_norm_rope_flashmla( kv=kv, kv_weight=kv_weight, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 3c883cd2f3ba..e9f1439d34da 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -110,6 +110,7 @@ prepare_context_parallel_metadata, ) from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.mem_cache.deepseek_v4_memory_pool import is_nvfp4_kv_cache from sglang.srt.mem_cache.memory_pool import RadixAttention from sglang.srt.model_executor.cuda_graph_config import ( Backend, @@ -968,7 +969,9 @@ def _compute_kv_to_cache( Replaces the bf16-kv-intermediate path. Used everywhere except the DSA prefill-CP case (which needs bf16 kv for the cross-rank all-gather). """ - if envs.SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE.get(): + # The fused kernel below only emits the FP8 cache ABI, so an NVFP4 KV + # cache has to take the bf16-intermediate path and quantize on store. + if envs.SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE.get() or is_nvfp4_kv_cache(): # Quantize the nope payload from bf16-rounded values (the fused # kernel quantizes from fp32 registers; the bf16 rounding moves # values across fp8 bins relative to bf16-sourced consumers). From b165cd6aaab271b55925204d1fa7782266fd47e8 Mon Sep 17 00:00:00 2001 From: Ali Chen Date: Thu, 10 Sep 2026 22:16:39 -0700 Subject: [PATCH 2/2] [SM120] Add regression tests for the NVFP4 sparse-MLA cache path Every case corresponds to a defect that reached a running server while building this path, so these are regression guards rather than coverage: staging-buffer pointer stability under CUDA graph capture, bucketing so that cache does not grow without bound, per-call zeroing, int64 out_loc, uint8 coercion in nvfp4_cache_view, and the platform gate. All of that runs on CPU. The slot-addressing round-trip needs the kernels and is skipped off SM120. sglang CI has no SM120 runner, so the file is not registered in run_suite.py. Co-Authored-By: Claude Opus 5 (1M context) --- test/srt/mem_cache/test_dsv4_nvfp4_cache.py | 252 ++++++++++++++++++++ 1 file changed, 252 insertions(+) create mode 100644 test/srt/mem_cache/test_dsv4_nvfp4_cache.py diff --git a/test/srt/mem_cache/test_dsv4_nvfp4_cache.py b/test/srt/mem_cache/test_dsv4_nvfp4_cache.py new file mode 100644 index 000000000000..a38d2bb9cdf8 --- /dev/null +++ b/test/srt/mem_cache/test_dsv4_nvfp4_cache.py @@ -0,0 +1,252 @@ +"""Unit tests for the DeepSeek-V4 NVFP4 sparse-MLA KV cache path (SM120/SM121). + +Most of this runs on CPU and pins the invariants that the integration silently +depends on. Each one corresponds to a bug that reached a running server during +development, so they are regression tests rather than documentation: + + - ``_nvfp4_stage`` must never move a bucket's storage. A decode CUDA graph + bakes in the staging address at capture time (captured at the largest + ``cuda_graph_max_bs``), and a later prefill needs far more rows; a grow-in- + place scheme frees the tensor the replayed graph still writes through. + - ``_nvfp4_stage`` must bucket rather than key on the exact row count, or the + cache grows without bound (prefill row counts follow each chunk's token + count). + - Its rows must be zeroed per call: entries the compress kernel skips must + not be quantized into the pool as stale data. + - ``nvfp4_cache_view`` must force uint8. ``get_key_buffer`` hands out an + ``fp8_e4m3`` view under ``--kv-cache-dtype fp8_e4m3`` and the NVFP4 op + rejects anything but uint8. + - The platform gate must reject the flag on non-SM120 hardware instead of + letting the pool allocate an NVFP4 layout no kernel can read. + +The slot-addressing round-trip needs the kernels and therefore real SM120 +hardware; it is skipped elsewhere. sglang CI has no SM120 runner, so this file +is not registered in run_suite.py. + +Run: python3 test/srt/mem_cache/test_dsv4_nvfp4_cache.py +""" + +import types +import unittest +from unittest import mock + +import torch + +from sglang.srt.layers.attention.dsv4.compressor_v2 import CompressorBackendMixin +from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( + NVFP4_BYTES_PER_TOKEN, + DeepSeekV4SingleKVPool, + is_nvfp4_kv_cache, +) + + +def _is_sm120() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 12 + + +class TestNvfp4StageBuffer(unittest.TestCase): + """The compressed-cache staging buffer, on CPU.""" + + def setUp(self): + self.mixin = CompressorBackendMixin() + self.dev = torch.device("cpu") + + def test_bucket_storage_is_stable_when_a_larger_request_arrives(self): + """A decode-sized bucket must survive a later prefill-sized request. + + This is the CUDA-graph hazard: the graph captured at bs=32 keeps writing + through whatever pointer it saw, so growing that allocation corrupts it. + """ + rows_small, _ = self.mixin._nvfp4_stage(4, 32, self.dev) + ptr_small = rows_small.data_ptr() + + # A prefill chunk of 32768 tokens at compress_ratio 4 emits 8192 rows. + self.mixin._nvfp4_stage(4, 8192, self.dev) + + rows_again, _ = self.mixin._nvfp4_stage(4, 32, self.dev) + self.assertEqual( + rows_again.data_ptr(), + ptr_small, + "staging storage moved; a replayed CUDA graph would write through a " + "freed pointer", + ) + + def test_row_counts_are_bucketed_to_powers_of_two(self): + """Distinct row counts must not each get their own tensor.""" + seen = set() + for n in (33, 40, 50, 64): + rows, _ = self.mixin._nvfp4_stage(4, n, self.dev) + self.assertEqual(rows.shape, (n, 512)) + seen.add(rows.untyped_storage().data_ptr()) + # 33..64 all round up to 64, so one backing allocation serves them all. + self.assertEqual(len(seen), 1) + + # ... and a count past the bucket does allocate a new one. + rows, _ = self.mixin._nvfp4_stage(4, 65, self.dev) + self.assertNotIn(rows.untyped_storage().data_ptr(), seen) + + def test_compress_ratios_do_not_share_a_buffer(self): + rows_c4, _ = self.mixin._nvfp4_stage(4, 64, self.dev) + rows_c128, _ = self.mixin._nvfp4_stage(128, 64, self.dev) + self.assertNotEqual(rows_c4.data_ptr(), rows_c128.data_ptr()) + + def test_rows_are_zeroed_on_every_call(self): + rows, _ = self.mixin._nvfp4_stage(4, 16, self.dev) + rows.fill_(1.0) + rows_again, _ = self.mixin._nvfp4_stage(4, 16, self.dev) + self.assertTrue(torch.equal(rows_again, torch.zeros_like(rows_again))) + + def test_out_loc_is_int64_identity(self): + """The compress store kernel rejects int32 out_loc.""" + rows, loc = self.mixin._nvfp4_stage(4, 12, self.dev) + self.assertEqual(loc.dtype, torch.int64) + self.assertEqual(rows.shape[0], loc.shape[0]) + self.assertTrue(torch.equal(loc, torch.arange(12, dtype=torch.int64))) + + +class TestNvfp4CacheView(unittest.TestCase): + """The single view both the append helper and the attention kernel use.""" + + def _view(self, buf, page_size): + pool = types.SimpleNamespace(nvfp4_page_size=page_size) + return DeepSeekV4SingleKVPool.nvfp4_cache_view(pool, buf) + + def test_abi_is_384_bytes_per_token(self): + self.assertEqual(NVFP4_BYTES_PER_TOKEN, 384) + + def test_view_shape_and_contiguity(self): + num_rows, page_size = 8, 64 + buf = torch.zeros( + num_rows, page_size * NVFP4_BYTES_PER_TOKEN, dtype=torch.uint8 + ) + view = self._view(buf, page_size) + self.assertEqual(view.shape, (num_rows, page_size, NVFP4_BYTES_PER_TOKEN)) + # The op derives the page split from these strides. + self.assertEqual( + view.stride(), (page_size * NVFP4_BYTES_PER_TOKEN, NVFP4_BYTES_PER_TOKEN, 1) + ) + + def test_view_forces_uint8_from_an_fp8_buffer(self): + """--kv-cache-dtype fp8_e4m3 makes get_key_buffer hand out an fp8 view.""" + page_size = 64 + buf = torch.zeros(4, page_size * NVFP4_BYTES_PER_TOKEN, dtype=torch.uint8) + view = self._view(buf.view(torch.float8_e4m3fn), page_size) + self.assertEqual(view.dtype, torch.uint8) + self.assertEqual(view.shape, (4, page_size, NVFP4_BYTES_PER_TOKEN)) + + def test_extra_cache_page_sizes(self): + """c4 pools carry page 64 and c128 pools page 2; the op accepts both.""" + for page_size in (2, 64): + buf = torch.zeros(6, page_size * NVFP4_BYTES_PER_TOKEN, dtype=torch.uint8) + view = self._view(buf, page_size) + self.assertEqual(view.shape[1], page_size) + self.assertEqual(view.numel(), buf.numel()) + + +class TestPlatformGate(unittest.TestCase): + def tearDown(self): + is_nvfp4_kv_cache.cache_clear() + + def _gate(self, fmt, sm120): + is_nvfp4_kv_cache.cache_clear() + with mock.patch( + "sglang.srt.mem_cache.deepseek_v4_memory_pool.envs" + ) as envs, mock.patch( + "sglang.srt.mem_cache.deepseek_v4_memory_pool.is_sm120_supported", + return_value=sm120, + ): + envs.SGLANG_SM120_KV_CACHE_FORMAT.get.return_value = fmt + return is_nvfp4_kv_cache() + + def test_default_is_fp8(self): + self.assertFalse(self._gate("fp8", sm120=True)) + + def test_fp8_is_allowed_off_sm120(self): + self.assertFalse(self._gate("fp8", sm120=False)) + + def test_nvfp4_on_sm120(self): + self.assertTrue(self._gate("nvfp4", sm120=True)) + + def test_nvfp4_off_sm120_raises(self): + with self.assertRaisesRegex(ValueError, "SM120"): + self._gate("nvfp4", sm120=False) + + +@unittest.skipUnless(_is_sm120(), "NVFP4 sparse MLA requires SM120/SM121") +class TestSlotAddressingRoundTrip(unittest.TestCase): + """Append at arbitrary flat slots, read them back through the attention op. + + Writers regroup the flat slot space into 64-token pages and readers index it + by the same flat ids; this pins that the two agree, which is what makes the + pool's page regrouping safe. + """ + + def test_round_trip_matches_bf16_reference(self): + from flashinfer.mla import nvfp4_quantize_append_sparse_mla_cache + from flashinfer.mla._sparse_mla_nvfp4_sm120 import ( + _sparse_mla_nvfp4_sm120_paged_attention, + ) + + torch.manual_seed(0) + dev, page, d = "cuda", 64, 512 + num_pages, heads, topk = 64, 64, 128 + num_slots = num_pages * page + batch = 4 + + kv = (torch.randn(num_slots, d, dtype=torch.bfloat16, device=dev) / 10).clamp( + -1, 1 + ) + cache = torch.zeros( + num_pages, page, NVFP4_BYTES_PER_TOKEN, dtype=torch.uint8, device=dev + ) + + # Two shuffled chunks, mirroring incremental prefill then decode stores. + perm = torch.randperm(num_slots, device=dev).to(torch.int32) + for chunk in (perm[: num_slots // 2], perm[num_slots // 2 :]): + nvfp4_quantize_append_sparse_mla_cache( + kv[chunk.long()].contiguous(), chunk, cache + ) + + q = (torch.randn(batch, heads, d, dtype=torch.bfloat16, device=dev) / 10).clamp( + -1, 1 + ) + idx = torch.randint(0, num_slots, (batch, topk), dtype=torch.int32, device=dev) + sm_scale = d**-0.5 + sink = torch.zeros(heads, dtype=torch.float32, device=dev) + + out = torch.empty(batch, heads, d, dtype=torch.bfloat16, device=dev) + lse = torch.empty(batch, heads, dtype=torch.float32, device=dev) + splits = (topk + 63) // 64 + _sparse_mla_nvfp4_sm120_paged_attention( + q, + cache, + idx, + out, + lse, + sm_scale, + topk_length=torch.full((batch,), topk, dtype=torch.int32, device=dev), + attn_sink=sink, + mid_out=torch.empty( + batch, heads, splits, d, dtype=torch.bfloat16, device=dev + ), + mid_lse=torch.empty(batch, heads, splits, dtype=torch.float32, device=dev), + use_prefill=False, + ) + + gathered = kv[idx.long()].float() + scores = torch.einsum("bhd,bkd->bhk", q.float(), gathered) * sm_scale + scores = torch.cat( + [scores, sink.view(1, heads, 1).expand(batch, heads, 1)], dim=-1 + ) + ref = torch.einsum("bhk,bkd->bhd", scores.softmax(dim=-1)[..., :topk], gathered) + + cos = torch.nn.functional.cosine_similarity( + out.float().flatten(1), ref.flatten(1), dim=-1 + ) + # NVFP4 is a 4-bit format; this bound catches addressing errors, not + # quantization error (a slot mismatch collapses cosine to ~0). + self.assertGreater(cos.min().item(), 0.97) + + +if __name__ == "__main__": + unittest.main()