[Feature] Add FP4 KV Cache Design and support SM120 GPUs (#21601)

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

@@ -291,7 +291,11 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
         from sglang.srt.layers.attention.linear.utils import (
             initialize_linear_attn_config,
         )
-        from sglang.srt.utils import is_blackwell, is_npu
+        from sglang.srt.utils import (
+            is_blackwell,
+            is_npu,
+            is_sm120_supported,
+        )
 
         if not is_npu():
             from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
@@ -320,13 +324,25 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
         hybrid_backend_cls = HybridLinearAttnBackend
         if hybrid_gdn_config(runner.model_config) is not None:
             if is_blackwell():
-                assert (
-                    runner.server_args.attention_backend == "triton"
-                    or runner.server_args.attention_backend == "trtllm_mha"
-                    or runner.server_args.attention_backend == "fa4"
-                    or runner.server_args.attention_backend == "flashinfer"
-                ), "triton, trtllm_mha, fa4, or flashinfer backend are the only supported backends on Blackwell GPUs for hybrid GDN models, use --attention-backend to specify the backend."
-            if is_npu():
+                if is_sm120_supported():
+                    allowed = {"triton", "trtllm_mha", "flashinfer"}
+                else:
+                    allowed = {"triton", "trtllm_mha", "fa4"}
+                attn_be = runner.server_args.attention_backend
+                prefill_be = runner.server_args.prefill_attention_backend
+                decode_be = runner.server_args.decode_attention_backend
+                # When using split prefill/decode backends, check each individually
+                if prefill_be and decode_be:
+                    assert prefill_be in allowed and decode_be in allowed, (
+                        f"Only {allowed} backends are supported on Blackwell GPUs for hybrid GDN models. "
+                        f"Got prefill={prefill_be}, decode={decode_be}."
+                    )
+                else:
+                    assert attn_be in allowed, (
+                        f"Only {allowed} backends are supported on Blackwell GPUs for hybrid GDN models. "
+                        f"Got attention_backend={attn_be}."
+                    )
+            elif is_npu():
                 assert (
                     runner.server_args.attention_backend == "ascend"
                 ), "ascend backend is the only supported backend on NPU for hybrid GDN models, use --attention-backend ascend to specify the backend."

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

@@ -26,6 +26,9 @@ from sglang.kernels.ops.attention.utils import (
 from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
+    KVCacheAttentionAccessKind,
+)
 from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
@@ -321,9 +324,52 @@ class FlashInferAttnBackend(AttentionBackend):
         self.dllm_config = DllmConfig.from_server_args(model_runner.server_args)
         self.is_dllm_model = self.dllm_config is not None
 
+        self.kv_cache_quant_method = self.token_to_kv_pool.get_kv_cache_quant_method()
+        self.prefill_kv_access = self.kv_cache_quant_method.resolve_attention_access(
+            "prefill", "flashinfer"
+        )
+        self.decode_kv_access = self.kv_cache_quant_method.resolve_attention_access(
+            "decode", "flashinfer"
+        )
+        prefill_backend, decode_backend = (
+            model_runner.server_args.get_attention_backends()
+        )
+        if self.__class__ is FlashInferAttnBackend:
+            if prefill_backend == "flashinfer":
+                self._check_kv_attention_access("prefill", self.prefill_kv_access)
+            if decode_backend == "flashinfer":
+                self._check_kv_attention_access("decode", self.decode_kv_access)
+
+        self.prefill_uses_dequant_workspace = (
+            self.prefill_kv_access is not None
+            and self.prefill_kv_access.kind
+            == KVCacheAttentionAccessKind.DEQUANT_WORKSPACE
+        )
+        self.decode_uses_dequant_workspace = (
+            self.decode_kv_access is not None
+            and self.decode_kv_access.kind
+            == KVCacheAttentionAccessKind.DEQUANT_WORKSPACE
+        )
+        self.is_nvfp4_kvcache = any(
+            access is not None and access.scale_recipe == "nvfp4"
+            for access in (self.prefill_kv_access, self.decode_kv_access)
+        )
+        self.dq_page_table = None
+        self.dq_paged_kernel_lens = None
+        self.cpu_req_pool_indices = None
+        # FP4 fake-quant prefill/decode exposes an FP8 workspace to FlashInfer.
+        self.flashinfer_kv_cache_dtype = (
+            torch.float8_e4m3fn
+            if (
+                self.prefill_uses_dequant_workspace
+                or self.decode_uses_dequant_workspace
+            )
+            else model_runner.kv_cache_dtype
+        )
+
         # Parse constants
         self.decode_use_tensor_cores = should_use_tensor_core(
-            kv_cache_dtype=model_runner.kv_cache_dtype,
+            kv_cache_dtype=self.flashinfer_kv_cache_dtype,
             num_attention_heads=model_runner.model_config.num_attention_heads
             // get_parallel().attn_tp_size,
             num_kv_heads=model_runner.model_config.get_num_kv_heads(
@@ -331,6 +377,7 @@ class FlashInferAttnBackend(AttentionBackend):
             ),
         )
         self.max_context_len = model_runner.model_config.context_len
+        self.page_size = model_runner.page_size
         self.skip_prefill = skip_prefill
         self.is_multimodal = model_runner.model_config.is_multimodal
         assert not (
@@ -490,6 +537,16 @@ class FlashInferAttnBackend(AttentionBackend):
             List[BatchPrefillWithPagedKVCacheWrapper]
         ] = None
 
+    def _check_kv_attention_access(self, phase: str, access) -> None:
+        if access is not None:
+            return
+        method_name = getattr(self.kv_cache_quant_method, "name", "unknown")
+        available = self.kv_cache_quant_method.describe_attention_accesses(phase)
+        raise ValueError(
+            f"KV cache method {method_name!r} does not support {phase} with "
+            f"flashinfer attention backend. Available {phase} accesses: {available}."
+        )
+
     @staticmethod
     def _resolve_swa_kv_pool(model_runner: ModelRunner) -> Optional[BaseSWAKVPool]:
         """Return the SWA KV pool to translate against, or None for non-SWA models.
@@ -776,6 +833,86 @@ class FlashInferAttnBackend(AttentionBackend):
                     self.cuda_graph_swa_out_cache_loc[:n]
                 )
 
+    def _prepare_dequant_workspace_metadata_for_extend(
+        self, forward_batch: ForwardBatch, use_ragged: bool = False
+    ):
+        """Prepare FlashInfer metadata for an FP4 dequant workspace.
+
+        Some FP4 recipes store packed KV but expose an FP8 workspace to
+        FlashInfer prefill. This builds the workspace page table, exact paged
+        lengths, and CPU request ids needed to populate that workspace before
+        the prefill kernel runs.
+        """
+        self.dq_page_table = None
+        self.dq_paged_kernel_lens = None
+        self.cpu_req_pool_indices = None
+        if not (
+            self.prefill_uses_dequant_workspace
+            and forward_batch.forward_mode.is_extend_without_speculative()
+        ):
+            return
+
+        # Ragged prefill handles current-chunk K/V with raw tensors, so the
+        # paged side only contains cached prefix lengths. Non-ragged prefill
+        # uses the dequant workspace for prefix + current chunk, so it needs
+        # full sequence lengths. These CPU length containers may arrive as
+        # Python lists or CPU tensors depending on the metadata builder.
+        paged_seq_lens_cpu = (
+            forward_batch.extend_prefix_lens_cpu
+            if use_ragged
+            else forward_batch.seq_lens_cpu
+        )
+        raw_paged_seq_lens = (
+            paged_seq_lens_cpu
+            if isinstance(paged_seq_lens_cpu, list)
+            else paged_seq_lens_cpu.tolist()
+        )
+        paged_seq_lens = [
+            int(seq_len.item()) if isinstance(seq_len, torch.Tensor) else int(seq_len)
+            for seq_len in raw_paged_seq_lens
+        ]
+        if sum(paged_seq_lens) <= 0:
+            self.cpu_req_pool_indices = forward_batch.req_pool_indices.to(
+                "cpu", non_blocking=True
+            )
+            return
+
+        # dq_buffer layout is page-aligned: each request occupies
+        # ceil(seq_len/page_size)*page_size slots, starting after a page_size
+        # dummy prefix. dq_page_table maps only actual token positions and skips
+        # padding gaps; dq_paged_kernel_lens stores real lengths so FlashInfer
+        # causal offsets use seq_len - q_len, not page_align(seq_len) - q_len.
+        seq_lens_with_scratch = paged_seq_lens + [256]
+        starts = []
+        next_start = self.page_size
+        for seq_len in seq_lens_with_scratch:
+            starts.append(next_start)
+            padded_len = (
+                (seq_len + self.page_size - 1) // self.page_size
+            ) * self.page_size
+            next_start += padded_len
+
+        device = forward_batch.req_pool_indices.device
+        indices = [
+            torch.arange(start, start + seq_len, device=device, dtype=torch.int32)
+            for start, seq_len in zip(starts, seq_lens_with_scratch)
+            if seq_len > 0
+        ]
+        self.dq_page_table = torch.cat(indices) if indices else None
+        self.dq_paged_kernel_lens = torch.tensor(
+            paged_seq_lens,
+            dtype=torch.int32,
+            device=device,
+        )
+        self.cpu_req_pool_indices = forward_batch.req_pool_indices.to(
+            "cpu", non_blocking=True
+        )
+
+    def _kv_write_scales(self, layer: RadixAttention):
+        if self.kv_cache_quant_method.needs_global_scale():
+            return None, None
+        return layer.k_scale, layer.v_scale
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         swa_out_cache_loc = None
         if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
@@ -844,6 +981,10 @@ class FlashInferAttnBackend(AttentionBackend):
                 # Use new backend-specific implementation
                 multi_item_params = self._process_multi_item_scoring(forward_batch)
 
+            self._prepare_dequant_workspace_metadata_for_extend(
+                forward_batch, use_ragged
+            )
+
             self.indices_updater_prefill.update(
                 forward_batch.req_pool_indices,
                 forward_batch.seq_lens,
@@ -858,6 +999,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 multi_item_params=multi_item_params,
                 cross_attention_custom_mask=forward_batch.cross_attention_custom_mask,
                 extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu,
+                custom_kv_indices=self.dq_page_table,
             )
             self.forward_metadata = PrefillMetadata(
                 self.prefill_wrappers_paged,
@@ -1128,18 +1270,40 @@ class FlashInferAttnBackend(AttentionBackend):
         logits_soft_cap = layer.logit_cap
 
         q = q.contiguous()
+
+        assert not (
+            self.prefill_uses_dequant_workspace and layer.is_cross_attention
+        ), "FP4 dequant KV cache is not supported for cross-attention"
+
+        # We perform dequant for chunk prefill/cache reuse.
+        pool = self.token_to_kv_pool
+        if self.prefill_uses_dequant_workspace:
+            kv_cache = pool.get_flashinfer_dequant_workspace_kv_buffer(
+                layer,
+                self.req_to_token_pool.req_to_token,
+                self.cpu_req_pool_indices,
+                forward_batch.extend_prefix_lens_cpu,
+                forward_batch.extend_seq_lens_cpu,
+                self.page_size,
+                prepare_workspace=self.dq_page_table is not None,
+                use_ragged=self.forward_metadata.use_ragged,
+                k_cur=k,
+                v_cur=v,
+            )
+        else:
+            kv_cache = pool.get_kv_buffer(layer.layer_id)
+
+        # use paged attention
         if not self.forward_metadata.use_ragged:
-            if k is not None:
+            if k is not None and save_kv_cache:
                 assert v is not None
-                if save_kv_cache:
-                    self.token_to_kv_pool.set_kv_buffer(
-                        layer,
-                        KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc),
-                        k,
-                        v,
-                        layer.k_scale,
-                        layer.v_scale,
-                    )
+                self.token_to_kv_pool.set_kv_buffer(
+                    layer,
+                    KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc),
+                    k,
+                    v,
+                    *self._kv_write_scales(layer),
+                )
 
             causal = (
                 not layer.is_cross_attention
@@ -1147,7 +1311,7 @@ class FlashInferAttnBackend(AttentionBackend):
             )
             o = prefill_wrapper_paged.forward(
                 q.view(-1, layer.tp_q_head_num, layer.head_dim),
-                self.token_to_kv_pool.get_kv_buffer(layer.layer_id),
+                kv_cache,
                 causal=causal,
                 sm_scale=layer.scaling,
                 # Disable sliding window attention for multi-item scoring:
@@ -1175,6 +1339,9 @@ class FlashInferAttnBackend(AttentionBackend):
             # previously cached context without re-materializing KV tensors (e.g., the
             # IQuestLoopCoder path uses token_to_kv_pool as the KV source).
             if k is None and v is None:
+                assert (
+                    not self.prefill_uses_dequant_workspace
+                ), "KV cache must be provided for ragged attention when using FP4 dequant KV cache"
                 k = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)[0]
                 v = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)[1]
             causal = True
@@ -1219,7 +1386,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 )
                 o2, s2 = prefill_wrapper_paged.forward_return_lse(
                     q.view(-1, layer.tp_q_head_num, layer.head_dim),
-                    self.token_to_kv_pool.get_kv_buffer(layer.layer_id),
+                    kv_cache,
                     causal=False,
                     sm_scale=layer.scaling,
                     window_left=swa_window_left,
@@ -1234,8 +1401,7 @@ class FlashInferAttnBackend(AttentionBackend):
                     KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc),
                     k,
                     v,
-                    layer.k_scale,
-                    layer.v_scale,
+                    *self._kv_write_scales(layer),
                 )
 
         return o.view(-1, layer.tp_q_head_num * layer.head_dim)
@@ -1267,14 +1433,29 @@ class FlashInferAttnBackend(AttentionBackend):
                     KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc),
                     k,
                     v,
-                    layer.k_scale,
-                    layer.v_scale,
+                    *self._kv_write_scales(layer),
                 )
 
+        if self.decode_uses_dequant_workspace:
+            kv_cache = (
+                self.token_to_kv_pool.get_flashinfer_decode_dequant_workspace_kv_buffer(
+                    layer,
+                    self.req_to_token_pool.req_to_token,
+                    forward_batch.req_pool_indices,
+                    (
+                        forward_batch.seq_lens_cpu
+                        if forward_batch.seq_lens_cpu is not None
+                        else forward_batch.seq_lens
+                    ),
+                )
+            )
+        else:
+            kv_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
+
         # Call the wrapped function
         o = decode_wrapper.forward(
             q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
-            self.token_to_kv_pool.get_kv_buffer(layer.layer_id),
+            kv_cache,
             sm_scale=layer.scaling,
             logits_soft_cap=layer.logit_cap,
             # Must use _float to avoid device-to-host copy that breaks cuda graph capture.
@@ -1306,7 +1487,7 @@ class FlashInferIndicesUpdaterDecode:
             get_parallel().attn_tp_size
         )
         self.head_dim = model_runner.model_config.head_dim
-        self.data_type = model_runner.kv_cache_dtype
+        self.data_type = attn_backend.flashinfer_kv_cache_dtype
         self.q_data_type = model_runner.dtype
         self.sliding_window_size = model_runner.sliding_window_size
         self.attn_backend = attn_backend
@@ -1574,7 +1755,7 @@ class FlashInferIndicesUpdaterPrefill:
             get_parallel().attn_tp_size
         )
         self.head_dim = model_runner.model_config.head_dim
-        self.data_type = model_runner.kv_cache_dtype
+        self.data_type = attn_backend.flashinfer_kv_cache_dtype
         self.q_data_type = model_runner.dtype
         self.sliding_window_size = model_runner.sliding_window_size
         self.attn_backend = attn_backend
@@ -1610,6 +1791,7 @@ class FlashInferIndicesUpdaterPrefill:
         multi_item_params: Optional[MultiItemScoringParams] = None,
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         extend_prefix_lens_cpu: Optional[List[int]] = None,
+        custom_kv_indices: Optional[torch.Tensor] = None,
     ):
         # Keep the signature for type checking. It will be assigned during runtime.
         raise NotImplementedError()
@@ -1629,6 +1811,7 @@ class FlashInferIndicesUpdaterPrefill:
         multi_item_params: Optional[MultiItemScoringParams] = None,
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         extend_prefix_lens_cpu: Optional[List[int]] = None,
+        custom_kv_indices: Optional[torch.Tensor] = None,
     ):
         if use_ragged:
             assert prefix_lens is not None
@@ -1658,6 +1841,7 @@ class FlashInferIndicesUpdaterPrefill:
             fixed_split_size=fixed_split_size,
             multi_item_params=multi_item_params,
             seq_lens_cpu=seq_lens_cpu,
+            custom_kv_indices=custom_kv_indices,
         )
 
     def update_sliding_window(
@@ -1675,7 +1859,12 @@ class FlashInferIndicesUpdaterPrefill:
         multi_item_params: Optional[MultiItemScoringParams] = None,
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         extend_prefix_lens_cpu: Optional[List[int]] = None,
+        custom_kv_indices: Optional[torch.Tensor] = None,
     ):
+        if custom_kv_indices is not None:
+            raise RuntimeError(
+                "NVFP4 custom KV indices are only supported by the single-wrapper FlashInfer path."
+            )
         if prefix_lens is None:
             num_accept_tokens = getattr(spec_info, "num_accept_tokens", None)
             prefix_lens = (
@@ -1796,7 +1985,12 @@ class FlashInferIndicesUpdaterPrefill:
         multi_item_params: Optional[MultiItemScoringParams] = None,
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         extend_prefix_lens_cpu: Optional[List[int]] = None,
+        custom_kv_indices: Optional[torch.Tensor] = None,
     ):
+        if custom_kv_indices is not None:
+            raise RuntimeError(
+                "NVFP4 custom KV indices are not supported for cross-attention."
+            )
         for wrapper_id in range(2):
             if wrapper_id == 0:
                 # normal attention
@@ -1848,28 +2042,43 @@ class FlashInferIndicesUpdaterPrefill:
         multi_item_params: Optional[MultiItemScoringParams] = None,
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         seq_lens_cpu: Optional[torch.Tensor] = None,
+        custom_kv_indices: Optional[torch.Tensor] = None,
     ):
         bs = len(seq_lens)
         if spec_info is None:
             assert prefix_lens is not None
             assert len(seq_lens) == len(req_pool_indices)
             # Normal extend
-            kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
+            # custom_kv_indices uses exact dq_paged_kernel_lens so FlashInfer causal
+            # offsets are based on real token counts, not page-aligned padding.
+            if (
+                custom_kv_indices is not None
+                and self.attn_backend.dq_paged_kernel_lens is not None
+            ):
+                kv_indptr[1 : bs + 1] = torch.cumsum(
+                    self.attn_backend.dq_paged_kernel_lens, dim=0
+                )
+            else:
+                kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
             kv_indptr = kv_indptr[: bs + 1]
-            kv_indices = torch.empty(
-                paged_kernel_lens_sum + 256,
-                dtype=torch.int32,
-                device=req_pool_indices.device,
-            )
-            create_flashinfer_kv_indices_triton[(bs,)](
-                self.req_to_token,
-                req_pool_indices,
-                paged_kernel_lens,
-                kv_indptr,
-                kv_start_idx,
-                kv_indices,
-                self.req_to_token.shape[1],
-            )
+
+            if custom_kv_indices is not None:
+                kv_indices = custom_kv_indices
+            else:
+                kv_indices = torch.empty(
+                    paged_kernel_lens_sum + 256,
+                    dtype=torch.int32,
+                    device=req_pool_indices.device,
+                )
+                create_flashinfer_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    req_pool_indices,
+                    paged_kernel_lens,
+                    kv_indptr,
+                    kv_start_idx,
+                    kv_indices,
+                    self.req_to_token.shape[1],
+                )
             qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
             qo_indptr = qo_indptr[: bs + 1]
 

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

@@ -25,6 +25,9 @@ from sglang.srt.layers.attention.flashinfer_backend import (
     FlashInferAttnBackend,
     FlashInferMultiStepDraftBackend,
 )
+from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
+    KVCacheAttentionAccessKind,
+)
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
@@ -104,6 +107,17 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         super().__init__(
             model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf
         )
+        self.decode_kv_access = self.kv_cache_quant_method.resolve_attention_access(
+            "decode", "trtllm_mha"
+        )
+        self._check_decode_kv_access()
+        self.decode_uses_native_fp4 = (
+            self.decode_kv_access.kind == KVCacheAttentionAccessKind.NATIVE_FP4
+        )
+        self.is_nvfp4_kvcache = (
+            self.decode_uses_native_fp4
+            and self.decode_kv_access.scale_recipe == "nvfp4"
+        )
 
         config = model_runner.model_config
 
@@ -177,6 +191,24 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         #   KV fp8: q_type = fp8, out_type=model_runner.dtype
         self.is_xqa_impl = is_sm90_supported() or is_sm120_supported()
 
+    def _check_decode_kv_access(self) -> None:
+        supported_kinds = {
+            KVCacheAttentionAccessKind.PLAIN,
+            KVCacheAttentionAccessKind.NATIVE_FP4,
+        }
+        if (
+            self.decode_kv_access is not None
+            and self.decode_kv_access.kind in supported_kinds
+        ):
+            return
+
+        method_name = getattr(self.kv_cache_quant_method, "name", "unknown")
+        available = self.kv_cache_quant_method.describe_attention_accesses("decode")
+        raise ValueError(
+            f"KV cache method {method_name!r} does not support decode with "
+            f"trtllm_mha. Available decode accesses: {available}."
+        )
+
     @staticmethod
     def _resolve_swa_kv_pool(model_runner: ModelRunner) -> Optional[SWAKVPool]:
         """Return the SWAKVPool to translate against, or None for non-SWA models.
@@ -847,6 +879,45 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
 
         self.forward_metadata = metadata
 
+    def _reshape_paged_kv_cache(
+        self,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        layer: RadixAttention,
+        head_dim: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        k_cache = k_cache.view(
+            -1, self.page_size, layer.tp_k_head_num, head_dim
+        ).permute(0, 2, 1, 3)
+        v_cache = v_cache.view(
+            -1, self.page_size, layer.tp_v_head_num, head_dim
+        ).permute(0, 2, 1, 3)
+        if layer.tp_k_head_num == 1:
+            k_cache = canonicalize_stride(k_cache)
+        if layer.tp_v_head_num == 1:
+            v_cache = canonicalize_stride(v_cache)
+        return k_cache, v_cache
+
+    def _get_nvfp4_bmm_scales(self, layer: RadixAttention) -> tuple[float, float]:
+        assert self.is_nvfp4_kvcache
+        return self.kv_cache_quant_method.get_bmm_scales(layer.layer_id)
+
+    def _get_nvfp4_decode_kv_cache(self, layer: RadixAttention) -> tuple[
+        tuple[torch.Tensor, torch.Tensor],
+        tuple[torch.Tensor, torch.Tensor],
+    ]:
+        assert self.is_nvfp4_kvcache
+        k_fp4, v_fp4, k_scale, v_scale = self.token_to_kv_pool.get_raw_kv_buffer(
+            layer.layer_id
+        )
+        kv_cache = self._reshape_paged_kv_cache(
+            k_fp4, v_fp4, layer, layer.head_dim // 2
+        )
+        kv_cache_block_scales = self._reshape_paged_kv_cache(
+            k_scale, v_scale, layer, layer.head_dim // 16
+        )
+        return kv_cache, kv_cache_block_scales
+
     def forward_decode(
         self,
         q: torch.Tensor,
@@ -862,6 +933,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
 
         use_fused_fp8_path = self._should_use_fused_fp8_path(save_kv_cache, k)
         use_fused_qkv = use_fused_fp8_path and not self.is_xqa_impl
+        pool = self.token_to_kv_pool
 
         if use_fused_fp8_path:
             fused_q = self._fused_fp8_qkv_kv_cache(
@@ -872,15 +944,13 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             k = None
             v = None
         else:
-            # Use original set_kv_buffer path
             if save_kv_cache and k is not None:
                 self.token_to_kv_pool.set_kv_buffer(
                     layer,
                     KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc),
                     k,
                     v,
-                    layer.k_scale,
-                    layer.v_scale,
+                    *self._kv_write_scales(layer),
                 )
 
         # For XQA, q_dtype should be bf16. For trtllm-gen,
@@ -893,31 +963,26 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         ):
             q = q.to(torch.float8_e4m3fn)
         q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim)
-        k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
-        # shape conversion:
-        # [num_pages, page_size, num_kv_heads, head_dim] -> [num_pages, num_kv_heads, page_size, head_dim]
-        k_cache = k_cache.view(
-            -1, self.page_size, layer.tp_k_head_num, layer.head_dim
-        ).permute(0, 2, 1, 3)
-        v_cache = v_cache.view(
-            -1, self.page_size, layer.tp_v_head_num, layer.head_dim
-        ).permute(0, 2, 1, 3)
 
-        if layer.tp_k_head_num == 1:
-            k_cache = canonicalize_stride(k_cache)
-        if layer.tp_v_head_num == 1:
-            v_cache = canonicalize_stride(v_cache)
+        if self.is_nvfp4_kvcache:
+            kv_cache, kv_cache_block_scales = self._get_nvfp4_decode_kv_cache(layer)
+        else:
+            k_cache, v_cache = pool.get_kv_buffer(layer.layer_id)
+            kv_cache = self._reshape_paged_kv_cache(
+                k_cache, v_cache, layer, layer.head_dim
+            )
+            kv_cache_block_scales = None
 
-        kv_cache = (k_cache, v_cache)
-
-        bmm1_scale, bmm2_scale = self._get_bmm_scales(layer, q_scale)
-        # sink: additional value per head in the denominator of the softmax.
+        if self.is_nvfp4_kvcache:
+            k_scale, v_scale = self._get_nvfp4_bmm_scales(layer)
+            bmm1_scale = q_scale * k_scale * layer.scaling
+            bmm2_scale = v_scale
+        else:
+            bmm1_scale, bmm2_scale = self._get_bmm_scales(layer, q_scale)
         attention_sink = kwargs.get("sinks", None)
 
         page_table = self._get_layer_page_table(layer, forward_batch)
 
-        # Call TRT-LLM kernel
-        # raw_out: like q, [bs, acc_q_len, num_q_heads, head_dim] but with output dtype
         o = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
             query=q,
             kv_cache=kv_cache,
@@ -931,7 +996,10 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             sinks=attention_sink,
             skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
             out_dtype=self.q_data_type,  # model_runner.dtype
+            kv_cache_sf=kv_cache_block_scales,
         )
+        if self.is_nvfp4_kvcache and o.dtype != self.q_data_type:
+            o = o.to(self.q_data_type)
 
         return o.view(-1, layer.tp_q_head_num * layer.head_dim)
 
@@ -945,6 +1013,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         save_kv_cache=True,
         **kwargs,
     ):
+        if self.decode_uses_native_fp4:
+            raise RuntimeError(
+                "TRTLLM MHA with native FP4 KV cache supports decode only; "
+                "use a separate prefill backend such as flashinfer or triton."
+            )
+
         cache_loc = forward_batch.out_cache_loc
 
         use_fused_fp8_path = self._should_use_fused_fp8_path(save_kv_cache, k)

python/sglang/srt/layers/quantization/fp4_kv_cache_quant_method.py

@@ -16,38 +16,211 @@ KV cache quantization strategy pattern.
 
 Three-player design:
   quant_method (pure compute)  ►  Pool (buffer + batch dequant)  ►  Backend (view adaptation)
+
+Why do we need attention access rules?
+- Problem: torch.float4_e2m1fn_x2 only describes packed FP4 storage. It does
+  not say whether the recipe is NVFP4 or fp4_mx_block16, nor how scales are
+  interpreted.
+- Problem: prefill and decode may use different KV views for the same recipe.
+  For example, NVFP4 uses a FlashInfer dequant workspace for prefill, but TRTLLM
+  MHA consumes native packed FP4 plus scales for decode.
+- Problem: putting these recipe/backend combinations directly into each backend
+  as dtype checks makes unsupported paths hard to spot and future recipes hard
+  to add safely.
+- Approach: the bottom registry declares KVCacheAttentionAccess entries for
+  each recipe. A backend resolves one entry by (phase, backend_name, tags),
+  then either uses the declared access pattern or fails fast if that
+  combination is unsupported.
 """
 
 from abc import ABC, abstractmethod
-from typing import Optional
+from dataclasses import dataclass
+from enum import Enum
+from typing import Iterable, Optional
 
 import torch
 from torch import Tensor
 
 from sglang.srt.layers.quantization.kvfp4_tensor import E2M1_MAX
+from sglang.srt.utils.common import is_sm100_supported
 
 
-class FP4KVCacheQuantMethod(ABC):
-    """Abstract base for FP4 KV cache quantization strategies.
+class KVCacheAttentionPhase(str, Enum):
+    PREFILL = "prefill"
+    DECODE = "decode"
+
+
+class KVCacheAttentionAccessKind(str, Enum):
+    # KV cache is already in the dtype/layout expected by the attention backend.
+    PLAIN = "plain"
+    # KV cache is stored quantized, then dequantized and decompressed into a
+    # temporary workspace before attention.
+    DEQUANT_WORKSPACE = "dequant_workspace"
+    # Attention backend directly consumes FP4 KV cache storage and scales.
+    NATIVE_FP4 = "native_fp4"
+
+
+@dataclass(frozen=True)
+class KVCacheBackendMatcher:
+    exact: frozenset[str] = frozenset()
+    tags: frozenset[str] = frozenset()
+    any_backend: bool = False
+
+    def matches(self, backend_name: str, backend_tags: Iterable[str]) -> bool:
+        backend_tags = frozenset(backend_tags)
+        return (
+            self.any_backend
+            or backend_name in self.exact
+            or (bool(self.tags) and self.tags.issubset(backend_tags))
+        )
+
+
+@dataclass(frozen=True)
+class KVCacheAttentionAccess:
+    """Describes how one attention backend reads KV cache for one phase.
+
+    Fields:
+    - phase: prefill or decode stage where this rule applies.
+    - kind: access mode, such as plain KV, dequant workspace, or native FP4.
+    - backend_matcher: backend names/tags that select this rule, e.g.
+      exact={"trtllm_mha"} or a FlashInfer dequant-workspace tag.
+    - storage_dtype: dtype stored in the KV pool, e.g. torch.float4_e2m1fn_x2.
+    - attention_kv_dtype: dtype consumed by attention after any conversion, e.g.
+      FP8 workspace for FlashInfer prefill or packed FP4 for TRTLLM decode.
+    - scale_recipe: scale semantics for this FP4 recipe, e.g. "nvfp4" or
+      "fp4_mx_block16".
+    - workspace_dtype: temporary workspace dtype for dequant/decompress paths;
+      None means no temporary workspace is needed.
+    """
+
+    phase: KVCacheAttentionPhase
+    kind: KVCacheAttentionAccessKind
+    backend_matcher: KVCacheBackendMatcher
+    storage_dtype: Optional[torch.dtype] = None
+    attention_kv_dtype: Optional[torch.dtype] = None
+    scale_recipe: Optional[str] = None
+    workspace_dtype: Optional[torch.dtype] = None
+
+    def matches(self, phase, backend_name: str, backend_tags: Iterable[str]) -> bool:
+        return self.phase == KVCacheAttentionPhase(
+            phase
+        ) and self.backend_matcher.matches(backend_name, backend_tags)
+
+
+class KVCacheQuantMethodBase(ABC):
+    """Abstract base for KV cache quantization strategies.
 
     Owns the quantize/dequantize computation.  The Pool owns the buffers and
     orchestrates the batch dequant loop.  Backends only do view/reshape.
-
-    All operations (quantize_and_store, dequantize_prev_kv) use FlashInfer
-    kernels or pure tensor ops, so they are CUDA-graph compatible.
     """
 
     name: str
     SCALE_BLOCK_SIZE: int = 1
 
+    def attention_accesses(self) -> tuple[KVCacheAttentionAccess, ...]:
+        return KV_CACHE_ATTENTION_ACCESS_REGISTRY.get(self.name, ())
+
+    def resolve_attention_access(
+        self, phase, backend_name: str, backend_tags: Iterable[str] = ()
+    ) -> Optional[KVCacheAttentionAccess]:
+        for access in self.attention_accesses():
+            if access.matches(phase, backend_name, backend_tags):
+                return access
+        return None
+
+    def describe_attention_accesses(self, phase=None) -> str:
+        accesses = self.attention_accesses()
+        if phase is not None:
+            phase = KVCacheAttentionPhase(phase)
+            accesses = tuple(access for access in accesses if access.phase == phase)
+        if not accesses:
+            return "none"
+        return "; ".join(
+            f"{access.phase.value}:{access.kind.value}:"
+            f"exact={sorted(access.backend_matcher.exact)}:"
+            f"tags={sorted(access.backend_matcher.tags)}"
+            for access in accesses
+        )
+
     def needs_dequant_workspace(self) -> bool:
-        """Whether the pool should allocate dq_k_buffer / dq_v_buffer for prefill."""
-        return False
+        """Whether the pool should allocate dq_k_buffer / dq_v_buffer."""
+        return any(
+            access.kind == KVCacheAttentionAccessKind.DEQUANT_WORKSPACE
+            for access in self.attention_accesses()
+        )
+
+    def needs_plain_kv_dequant_read(self) -> bool:
+        """Whether plain attention reads require dequantizing packed KV first."""
+        return any(
+            access.kind == KVCacheAttentionAccessKind.PLAIN
+            and access.storage_dtype is not None
+            for access in self.attention_accesses()
+        )
+
+    def dequant_workspace_dtype(self) -> Optional[torch.dtype]:
+        """Workspace dtype required by DEQUANT_WORKSPACE access rules."""
+        workspace_dtypes = set()
+        for access in self.attention_accesses():
+            if access.kind != KVCacheAttentionAccessKind.DEQUANT_WORKSPACE:
+                continue
+            if access.workspace_dtype is None:
+                raise ValueError(
+                    f"KV cache method {self.name!r} declares DEQUANT_WORKSPACE "
+                    "without workspace_dtype."
+                )
+            workspace_dtypes.add(access.workspace_dtype)
+
+        if not workspace_dtypes:
+            return None
+        if len(workspace_dtypes) != 1:
+            raise ValueError(
+                f"KV cache method {self.name!r} declares multiple dequant "
+                f"workspace dtypes: {sorted(str(dtype) for dtype in workspace_dtypes)}."
+            )
+        return next(iter(workspace_dtypes))
+
+    def kv_storage_dtype(self) -> torch.dtype:
+        """Packed KV storage dtype declared by attention access rules."""
+        storage_dtypes = {
+            access.storage_dtype
+            for access in self.attention_accesses()
+            if access.storage_dtype is not None
+        }
+        if not storage_dtypes:
+            return torch.uint8
+        if len(storage_dtypes) != 1:
+            raise ValueError(
+                f"KV cache method {self.name!r} declares multiple storage "
+                f"dtypes: {sorted(str(dtype) for dtype in storage_dtypes)}."
+            )
+        return next(iter(storage_dtypes))
+
+    def plain_attention_kv_dtype(self) -> Optional[torch.dtype]:
+        """Dtype produced when packed KV is dequantized for PLAIN attention."""
+        attention_dtypes = {
+            access.attention_kv_dtype
+            for access in self.attention_accesses()
+            if access.kind == KVCacheAttentionAccessKind.PLAIN
+            and access.storage_dtype is not None
+            and access.attention_kv_dtype is not None
+        }
+        if not attention_dtypes:
+            return None
+        if len(attention_dtypes) != 1:
+            raise ValueError(
+                f"KV cache method {self.name!r} declares multiple plain attention "
+                f"KV dtypes: {sorted(str(dtype) for dtype in attention_dtypes)}."
+            )
+        return next(iter(attention_dtypes))
 
     def needs_global_scale(self) -> bool:
         """Whether this method uses a per-layer global FP32 scale."""
         return False
 
+    def scale_buffer_view_dtype(self) -> Optional[torch.dtype]:
+        """Optional dtype view for stored per-block scales."""
+        return None
+
     @abstractmethod
     def create_buffers(
         self, size: int, head_num: int, head_dim: int, layer_num: int, device: str
@@ -91,23 +264,73 @@ class FP4KVCacheQuantMethod(ABC):
         """Dequantize stored FP4 KV (selected token indices already applied).
 
         Returns:
-            (k_fp8, v_fp8): Both in torch.float8_e4m3fn dtype with shape
-            matching the input (after unpacking). These are written into the
-            shared dequant workspace buffer for the FlashInfer FP8 prefill kernel.
+            Dequantized K/V tensors with shape matching the input after unpacking.
         """
 
+    def dequantize_kv_tensor(
+        self,
+        fp4_tensor: Tensor,
+        scales: Tensor,
+        layer_id: int,
+        dtype: Optional[torch.dtype] = None,
+    ) -> Tensor:
+        """Dequantize one packed FP4 KV tensor for plain attention reads."""
+        raise NotImplementedError(
+            f"KV cache method {self.name!r} does not support plain KV dequant reads."
+        )
+
     @abstractmethod
     def compute_cell_size(
         self, head_num: int, head_dim: int, num_layers: int, kv_size: int
     ) -> int:
         """Per-token memory footprint in bytes (for capacity estimation)."""
 
-    def load_scales_from_model(self, model_runner, sm_version: int = None) -> None:
+    def load_scales_from_model(self, model_runner) -> None:
         """Load per-layer global scales from model weights (no-op by default)."""
         pass
 
 
-class NVFP4KVMethod(FP4KVCacheQuantMethod):
+class UnquantizedKVCacheMethod(KVCacheQuantMethodBase):
+    """Identity method for BF16 / FP8 KV cache: no extra quantization."""
+
+    name = "unquantized"
+    SCALE_BLOCK_SIZE = 1
+
+    def create_buffers(self, size, head_num, head_dim, layer_num, device) -> dict:
+        pass
+
+    def quantize_and_store(
+        self,
+        k_buffer,
+        v_buffer,
+        k_scale_buffer,
+        v_scale_buffer,
+        loc,
+        cache_k,
+        cache_v,
+        k_scale=None,
+        v_scale=None,
+    ) -> None:
+        raise RuntimeError(
+            "Unquantized KV cache writes are handled by MHATokenToKVPool.set_kv_buffer."
+        )
+
+    def dequantize_prev_kv(
+        self, k_fp4, k_scales, v_fp4, v_scales, layer_id
+    ) -> tuple[Tensor, Tensor]:
+        raise NotImplementedError(
+            "Unquantized KV cache does not support FP4 KV dequantization."
+        )
+
+    def compute_cell_size(
+        self, head_num: int, head_dim: int, num_layers: int, kv_size: int
+    ) -> int:
+        raise NotImplementedError(
+            "Unquantized KV cache capacity is computed by the default pool configurator."
+        )
+
+
+class NVFP4KVCacheMethod(KVCacheQuantMethodBase):
     """NVFP4 two-level scaling: global FP32 + per-block FP8 E4M3.
 
     Supported on SM100 and SM120.
@@ -116,29 +339,28 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod):
     name = "nvfp4"
     SCALE_BLOCK_SIZE = 16
 
-    def __init__(self, num_layers: int, device: str, sm_version: int = 120):
+    def __init__(self, num_layers: int, device: str):
         self.num_layers = num_layers
         self.device = device
-        self.sm_version = sm_version
         # Per-layer global FP32 scales; filled by load_scales_from_model()
         self.k_scales_gpu = torch.ones(num_layers, dtype=torch.float32, device=device)
         self.v_scales_gpu = torch.ones(num_layers, dtype=torch.float32, device=device)
-
-    def needs_dequant_workspace(self) -> bool:
-        return (
-            True  # prefill uses FP8 dequant workspace; future native FP4 kernel → False
-        )
+        self.k_scales_float = [1.0] * num_layers
+        self.v_scales_float = [1.0] * num_layers
 
     def needs_global_scale(self) -> bool:
         return True
 
-    def load_scales_from_model(self, model_runner, sm_version: int = None) -> None:
-        if sm_version is not None:
-            self.sm_version = sm_version
+    def scale_buffer_view_dtype(self) -> Optional[torch.dtype]:
+        return torch.float8_e4m3fn
 
+    def load_scales_from_model(self, model_runner) -> None:
         from sglang.srt.model_loader.utils import resolve_language_model
 
-        language_model = resolve_language_model(model_runner.model)
+        model = getattr(model_runner, "model", model_runner)
+        language_model = (
+            model if hasattr(model, "layers") else resolve_language_model(model)
+        )
 
         attention_layers = []
         for layer in language_model.layers:
@@ -158,17 +380,19 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod):
 
         # k_scales_gpu is indexed by global (absolute) layer_id.  Resize if the model
         # has layers with global IDs larger than what was pre-allocated.
-        # This happens in hybrid models (e.g., GDN) where only a subset of layers
-        # are full-attention, but their layer_ids are non-contiguous.
         max_global_id = max(layer.layer_id for layer in attention_layers)
         required_size = max_global_id + 1
         if required_size > len(self.k_scales_gpu):
+            old_size = len(self.k_scales_gpu)
             self.k_scales_gpu = torch.ones(
                 required_size, dtype=torch.float32, device=self.device
             )
             self.v_scales_gpu = torch.ones(
                 required_size, dtype=torch.float32, device=self.device
             )
+            extra_layers = required_size - old_size
+            self.k_scales_float.extend([1.0] * extra_layers)
+            self.v_scales_float.extend([1.0] * extra_layers)
 
         k_scales_cpu = self.k_scales_gpu.cpu().clone()
         v_scales_cpu = self.v_scales_gpu.cpu().clone()
@@ -192,23 +416,28 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod):
             # The FP4 data type itself is identical on both architectures.
             # Reference: TRT-LLM FP8QDQLinearMethod.process_weights_after_loading_fused_qkv_linear
             # https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/modules/linear.py
-            if self.sm_version == 100:
+            if is_sm100_supported():
                 k_scale *= E2M1_MAX
                 v_scale *= E2M1_MAX
             k_scales_cpu[layer_id] = k_scale
             v_scales_cpu[layer_id] = v_scale
+            self.k_scales_float[layer_id] = k_scale
+            self.v_scales_float[layer_id] = v_scale
 
         self.k_scales_gpu.copy_(k_scales_cpu, non_blocking=True)
         self.v_scales_gpu.copy_(v_scales_cpu, non_blocking=True)
 
+    def get_bmm_scales(self, layer_id: int) -> tuple[float, float]:
+        return self.k_scales_float[layer_id], self.v_scales_float[layer_id]
+
     def create_buffers(
         self, size: int, head_num: int, head_dim: int, layer_num: int, device: str
     ) -> dict:
         m = size
         n = head_num
         k = head_dim
-        store_dtype = torch.uint8
-        dq_dtype = torch.float8_e4m3fn
+        store_dtype = self.kv_storage_dtype()
+        dq_dtype = self.dequant_workspace_dtype()
 
         k_buffer = [
             torch.zeros((m, n, k // 2), dtype=store_dtype, device=device)
@@ -230,9 +459,17 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod):
             )
             for _ in range(layer_num)
         ]
-        # Shared dequant workspace — one copy, reused per layer during prefill
-        dq_k_buffer = torch.zeros((m, n, k), dtype=dq_dtype, device=device)
-        dq_v_buffer = torch.zeros((m, n, k), dtype=dq_dtype, device=device)
+        # Shared dequant workspace: one copy, reused per layer during prefill.
+        dq_k_buffer = (
+            torch.zeros((m, n, k), dtype=dq_dtype, device=device)
+            if dq_dtype is not None
+            else None
+        )
+        dq_v_buffer = (
+            torch.zeros((m, n, k), dtype=dq_dtype, device=device)
+            if dq_dtype is not None
+            else None
+        )
 
         return {
             "k_buffer": k_buffer,
@@ -265,10 +502,15 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod):
             cache_v.contiguous(), v_scale
         )
 
-        k_buffer[loc] = cache_k.view(torch.uint8)
-        v_buffer[loc] = cache_v.view(torch.uint8)
-        k_scale_buffer[loc] = cache_k_fp4_sf.view(torch.uint8)
-        v_scale_buffer[loc] = cache_v_fp4_sf.view(torch.uint8)
+        cache_k = cache_k.view(torch.uint8)
+        cache_v = cache_v.view(torch.uint8)
+        cache_k_fp4_sf = cache_k_fp4_sf.view(torch.uint8)
+        cache_v_fp4_sf = cache_v_fp4_sf.view(torch.uint8)
+
+        k_buffer[loc] = cache_k
+        v_buffer[loc] = cache_v
+        k_scale_buffer[loc] = cache_k_fp4_sf
+        v_scale_buffer[loc] = cache_v_fp4_sf
 
     def dequantize_prev_kv(
         self,
@@ -300,26 +542,43 @@ class NVFP4KVMethod(FP4KVCacheQuantMethod):
         scale_size = (
             head_num * (head_dim // self.SCALE_BLOCK_SIZE) * num_layers * 2 * kv_size
         )
-        # Dequant workspace: shared across layers (not multiplied by num_layers), FP8
-        dq_size = head_num * head_dim * 2 * kv_size
+        # Dequant workspace is shared across layers, not multiplied by num_layers.
+        dq_dtype = self.dequant_workspace_dtype()
+        dq_size = (
+            head_num
+            * head_dim
+            * 2
+            * kv_size
+            * torch.empty((), dtype=dq_dtype).element_size()
+            if dq_dtype is not None
+            else 0
+        )
         return fp4_size + scale_size + dq_size
 
 
-class BlockFP4KVMethod(FP4KVCacheQuantMethod):
-    """Block-wise FP4 single-level scaling (similar to MXFP4 but block_size=16)."""
+class FP4MXBlock16KVCacheMethod(KVCacheQuantMethodBase):
+    """Block-16 FP4 E2M1 single-level scaling.
 
-    name = "blockfp4"
+    This is intentionally not called MXFP4: standard MXFP4 uses a block size of
+    32, while this KV cache recipe stores one scale per 16 FP4 values.
+    """
+
+    name = "fp4_mx_block16"
     SCALE_BLOCK_SIZE = 16
 
-    def needs_dequant_workspace(self) -> bool:
-        return True
+    def __init__(
+        self,
+        num_layers: Optional[int] = None,
+        device: Optional[str] = None,
+    ):
+        pass
 
     def create_buffers(
         self, size: int, head_num: int, head_dim: int, layer_num: int, device: str
     ) -> dict:
         m = size
-        store_dtype = torch.uint8
-        dq_dtype = torch.float8_e4m3fn
+        store_dtype = self.kv_storage_dtype()
+        dq_dtype = self.dequant_workspace_dtype()
 
         k_buffer = [
             torch.zeros((m, head_num, head_dim // 2), dtype=store_dtype, device=device)
@@ -329,7 +588,7 @@ class BlockFP4KVMethod(FP4KVCacheQuantMethod):
             torch.zeros((m, head_num, head_dim // 2), dtype=store_dtype, device=device)
             for _ in range(layer_num)
         ]
-        # MXFP4 flattens head dimensions for scale storage
+        # Block-16 FP4 flattens head dimensions for scale storage
         k_scale_buffer = [
             torch.zeros(
                 (m, (head_num * head_dim) // self.SCALE_BLOCK_SIZE),
@@ -346,11 +605,15 @@ class BlockFP4KVMethod(FP4KVCacheQuantMethod):
             )
             for _ in range(layer_num)
         ]
-        dq_k_buffer = torch.zeros(
-            (m, head_num, head_dim), dtype=dq_dtype, device=device
+        dq_k_buffer = (
+            torch.zeros((m, head_num, head_dim), dtype=dq_dtype, device=device)
+            if dq_dtype is not None
+            else None
         )
-        dq_v_buffer = torch.zeros(
-            (m, head_num, head_dim), dtype=dq_dtype, device=device
+        dq_v_buffer = (
+            torch.zeros((m, head_num, head_dim), dtype=dq_dtype, device=device)
+            if dq_dtype is not None
+            else None
         )
 
         return {
@@ -375,15 +638,34 @@ class BlockFP4KVMethod(FP4KVCacheQuantMethod):
         k_scale=None,
         v_scale=None,
     ) -> None:
-        from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil
+        from sglang.srt.layers.quantization.kvfp4_tensor import (
+            FP4MXBlock16KVQuantizeUtil,
+        )
+
+        cache_k_fp4, cache_k_sf = FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k)
+        cache_v_fp4, cache_v_sf = FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_v)
 
-        cache_k_fp4, cache_k_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_k)
-        cache_v_fp4, cache_v_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_v)
         k_buffer[loc] = cache_k_fp4
         v_buffer[loc] = cache_v_fp4
         k_scale_buffer[loc] = cache_k_sf
         v_scale_buffer[loc] = cache_v_sf
 
+    def dequantize_kv_tensor(
+        self,
+        fp4_tensor: Tensor,
+        scales: Tensor,
+        layer_id: int,
+        dtype: Optional[torch.dtype] = None,
+    ) -> Tensor:
+        from sglang.srt.layers.quantization.kvfp4_tensor import (
+            FP4MXBlock16KVQuantizeUtil,
+        )
+
+        target_dtype = dtype or self.plain_attention_kv_dtype() or torch.bfloat16
+        return FP4MXBlock16KVQuantizeUtil.batched_dequantize(
+            fp4_tensor, scales, dtype=target_dtype
+        )
+
     def dequantize_prev_kv(
         self,
         k_fp4: Tensor,
@@ -392,11 +674,10 @@ class BlockFP4KVMethod(FP4KVCacheQuantMethod):
         v_scales: Tensor,
         layer_id: int,
     ) -> tuple[Tensor, Tensor]:
-        from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil
-
-        k_bf16 = BlockFP4KVQuantizeUtil.batched_dequantize(k_fp4, k_scales)
-        v_bf16 = BlockFP4KVQuantizeUtil.batched_dequantize(v_fp4, v_scales)
-        return k_bf16.to(torch.float8_e4m3fn), v_bf16.to(torch.float8_e4m3fn)
+        return (
+            self.dequantize_kv_tensor(k_fp4, k_scales, layer_id),
+            self.dequantize_kv_tensor(v_fp4, v_scales, layer_id),
+        )
 
     def compute_cell_size(
         self, head_num: int, head_dim: int, num_layers: int, kv_size: int
@@ -405,22 +686,152 @@ class BlockFP4KVMethod(FP4KVCacheQuantMethod):
         scale_size = (
             (head_num * head_dim // self.SCALE_BLOCK_SIZE) * num_layers * 2 * kv_size
         )
-        dq_size = head_num * head_dim * 2 * kv_size
+        dq_dtype = self.dequant_workspace_dtype()
+        dq_size = (
+            head_num
+            * head_dim
+            * 2
+            * kv_size
+            * torch.empty((), dtype=dq_dtype).element_size()
+            if dq_dtype is not None
+            else 0
+        )
         return fp4_size + scale_size + dq_size
 
 
-# Registry: name → class.  Only classes for fp4_e2m1 dtype need to be listed.
-FP4_KV_CACHE_QUANT_REGISTRY: dict[str, type[FP4KVCacheQuantMethod]] = {
-    "nvfp4": NVFP4KVMethod,
-    "blockfp4": BlockFP4KVMethod,
+# Registry: method name -> attention access rules.
+_PREFILL = KVCacheAttentionPhase.PREFILL
+_DECODE = KVCacheAttentionPhase.DECODE
+_PLAIN_KIND = KVCacheAttentionAccessKind.PLAIN
+_DQ_WORKSPACE_KIND = KVCacheAttentionAccessKind.DEQUANT_WORKSPACE
+_NATIVE_FP4_KIND = KVCacheAttentionAccessKind.NATIVE_FP4
+_ANY_BACKEND = KVCacheBackendMatcher(any_backend=True)
+_NVFP4_SCALE = "nvfp4"
+_FP4_MX_SCALE = "fp4_mx_block16"
+_FP8_E4M3 = torch.float8_e4m3fn
+_TORCH_FP4 = getattr(torch, "float4_e2m1fn_x2", None)
+_BF16 = torch.bfloat16
+_NVFP4_PREFILL_BACKENDS = frozenset({"flashinfer"})
+_NVFP4_DECODE_BACKENDS = frozenset({"trtllm_mha"})
+_FP4_MX_MHA_BACKENDS = frozenset(
+    {"triton", "torch_native", "flex_attention", "trtllm_mha"}
+)
+_FP4_MX_PREFILL_BACKENDS = _FP4_MX_MHA_BACKENDS | frozenset({"fa4"})
+
+
+def _backend_matcher(backends) -> KVCacheBackendMatcher:
+    if isinstance(backends, KVCacheBackendMatcher):
+        return backends
+    return KVCacheBackendMatcher(exact=backends)
+
+
+def _plain(
+    phase: KVCacheAttentionPhase,
+    backends,
+    scale: Optional[str] = None,
+    attention_dtype: Optional[torch.dtype] = None,
+) -> KVCacheAttentionAccess:
+    return KVCacheAttentionAccess(
+        phase,
+        _PLAIN_KIND,
+        _backend_matcher(backends),
+        storage_dtype=torch.uint8 if scale is not None else None,
+        attention_kv_dtype=attention_dtype,
+        scale_recipe=scale,
+    )
+
+
+def _dq_workspace(
+    phase: KVCacheAttentionPhase,
+    backends,
+    scale: str,
+    attention_dtype: torch.dtype,
+) -> KVCacheAttentionAccess:
+    return KVCacheAttentionAccess(
+        phase,
+        _DQ_WORKSPACE_KIND,
+        _backend_matcher(backends),
+        storage_dtype=torch.uint8,
+        attention_kv_dtype=attention_dtype,
+        scale_recipe=scale,
+        workspace_dtype=attention_dtype,
+    )
+
+
+def _native_fp4(
+    phase: KVCacheAttentionPhase,
+    backends,
+    scale: str,
+    attention_dtype: Optional[torch.dtype],
+) -> KVCacheAttentionAccess:
+    return KVCacheAttentionAccess(
+        phase,
+        _NATIVE_FP4_KIND,
+        _backend_matcher(backends),
+        storage_dtype=torch.uint8,
+        attention_kv_dtype=attention_dtype,
+        scale_recipe=scale,
+    )
+
+
+KV_CACHE_ATTENTION_ACCESS_REGISTRY: dict[str, tuple[KVCacheAttentionAccess, ...]] = {
+    UnquantizedKVCacheMethod.name: (
+        _plain(_PREFILL, _ANY_BACKEND),
+        _plain(_DECODE, _ANY_BACKEND),
+    ),
+    NVFP4KVCacheMethod.name: (
+        _dq_workspace(_PREFILL, _NVFP4_PREFILL_BACKENDS, _NVFP4_SCALE, _FP8_E4M3),
+        _native_fp4(_DECODE, _NVFP4_DECODE_BACKENDS, _NVFP4_SCALE, _TORCH_FP4),
+    ),
+    FP4MXBlock16KVCacheMethod.name: (
+        _plain(_PREFILL, _FP4_MX_PREFILL_BACKENDS, _FP4_MX_SCALE, _BF16),
+        _plain(_DECODE, _FP4_MX_MHA_BACKENDS, _FP4_MX_SCALE, _BF16),
+    ),
 }
 
 
-def get_fp4_kv_cache_quant_method(name: str, **kwargs) -> FP4KVCacheQuantMethod:
-    """Instantiate a FP4KVCacheQuantMethod by recipe name."""
-    if name not in FP4_KV_CACHE_QUANT_REGISTRY:
+# Registry: explicit --kv-cache-dtype value -> method class.
+KV_CACHE_QUANT_REGISTRY: dict[str, type[KVCacheQuantMethodBase]] = {
+    "nvfp4": NVFP4KVCacheMethod,
+    "fp4_mx_block16": FP4MXBlock16KVCacheMethod,
+}
+
+
+def resolve_kv_cache_quant(kv_cache_dtype) -> Optional[str]:
+    """Resolve the explicit FP4 KV cache recipe from ``--kv-cache-dtype``."""
+
+    if not isinstance(kv_cache_dtype, str):
+        if (
+            hasattr(torch, "float4_e2m1fn_x2")
+            and kv_cache_dtype == torch.float4_e2m1fn_x2
+        ):
+            raise ValueError(
+                "FP4 KV cache storage dtype does not identify the recipe. "
+                "Pass the explicit --kv-cache-dtype value: 'nvfp4' or 'fp4_mx_block16'."
+            )
+        return None
+
+    if kv_cache_dtype == "fp4_e2m1":
         raise ValueError(
-            f"Unknown fp4_kv_cache_recipe: '{name}'. "
-            f"Available: {list(FP4_KV_CACHE_QUANT_REGISTRY)}"
+            "--kv-cache-dtype=fp4_e2m1 is deprecated. "
+            "Use --kv-cache-dtype=fp4_mx_block16."
         )
-    return FP4_KV_CACHE_QUANT_REGISTRY[name](**kwargs)
+    if kv_cache_dtype == "mxfp4":
+        raise ValueError(
+            "--kv-cache-dtype=mxfp4 is reserved for true MXFP4 block-size-32 "
+            "semantics. Use --kv-cache-dtype=fp4_mx_block16 for the current "
+            "block-size-16 FP4 KV recipe."
+        )
+    if kv_cache_dtype in KV_CACHE_QUANT_REGISTRY:
+        return kv_cache_dtype
+    return None
+
+
+def get_kv_cache_quant_method(name: str, **kwargs) -> KVCacheQuantMethodBase:
+    """Instantiate a KVCacheQuantMethodBase by internal method name."""
+    if name not in KV_CACHE_QUANT_REGISTRY:
+        raise ValueError(
+            f"Unknown KV cache quantization method: '{name}'. "
+            f"Available: {list(KV_CACHE_QUANT_REGISTRY)}"
+        )
+    return KV_CACHE_QUANT_REGISTRY[name](**kwargs)

python/sglang/srt/layers/quantization/kvfp4_tensor.py

@@ -25,40 +25,35 @@ class FP4KVCacheRecipe(Enum):
 
 E2M1_MAX = 6.0
 MAX_BLOCK_SCALE_FP8 = 448.0  # Maximum FP8 E4M3 value
-# Put constants directly on CUDA if available
-_device = "cuda" if torch.cuda.is_available() else "cpu"
 # E2M1 format: 1 sign bit + 2 exponent bits + 1 mantissa bit = 4 bits
 # 16 possible values: 0x0-0xF
 # Negative values: 0x8-0xF (sign bit = 1)
 # Positive values: 0x0-0x7 (sign bit = 0)
-E2M1_VALUES = torch.tensor(
-    [
-        0,
-        0.5,
-        1,
-        1.5,
-        2,
-        3,
-        4,
-        6,  # 0x0-0x7: positive values
-        -0,
-        -0.5,
-        -1,
-        -1.5,
-        -2,
-        -3,
-        -4,
-        -6,
-    ],  # 0x8-0xF: negative values
-    dtype=torch.float32,
-    device=_device,
-)
-E2M1_BOUNDS = torch.tensor(
-    [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5], dtype=torch.float32, device=_device
+# Keep constants as Python literals. Compiled helpers materialize them with
+# input.new_tensor(), so they follow the caller device without a global GPU tensor
+# or a CPU tensor .to(device) in the hot path.
+E2M1_VALUES = (
+    0.0,
+    0.5,
+    1.0,
+    1.5,
+    2.0,
+    3.0,
+    4.0,
+    6.0,
+    -0.0,
+    -0.5,
+    -1.0,
+    -1.5,
+    -2.0,
+    -3.0,
+    -4.0,
+    -6.0,
 )
+E2M1_BOUNDS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
 
 
-class BlockFP4KVQuantizeUtil:
+class FP4MXBlock16KVQuantizeUtil:
     """Block-wise FP4 (E2M1) quantization for KV cache.
 
     Similar to MXFP4 but uses block_size=16 (MXFP4 spec defines block_size=32).
@@ -69,6 +64,7 @@ class BlockFP4KVQuantizeUtil:
     @torch.compile
     def batched_quantize(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         """
+
         Quantize tensor to KVFP4 format
         Args:
             tensor: Input tensor of shape [B, M, N]
@@ -95,7 +91,8 @@ class BlockFP4KVQuantizeUtil:
         abs_vals = scaled.abs()
 
         # Pure tensor version (CUDA Graph safe)
-        magnitude_bits = torch.sum(abs_vals.unsqueeze(-1) >= E2M1_BOUNDS, dim=-1)
+        bounds = tensor.new_tensor(E2M1_BOUNDS, dtype=torch.float32)
+        magnitude_bits = torch.sum(abs_vals.unsqueeze(-1) >= bounds, dim=-1)
 
         # Combine sign and magnitude
         fp4_vals = sign_bits + magnitude_bits.to(torch.uint8)
@@ -136,7 +133,8 @@ class BlockFP4KVQuantizeUtil:
         magnitude_idx = fp4_vals & 0x07
 
         # Convert to float values
-        float_vals = E2M1_VALUES[magnitude_idx.long()]
+        values = quant_tensor.new_tensor(E2M1_VALUES[:8], dtype=torch.float32)
+        float_vals = values[magnitude_idx.long()]
         float_vals = torch.where(sign_mask, -float_vals, float_vals)
 
         # Reshape for block-wise scaling
@@ -177,21 +175,34 @@ class NVFP4KVQuantizeUtil:
                 block_scales: shape [B, M, N/16], dtype float8_e4m3fn
                 global_scale: passthrough
         """
-        from sglang.srt.utils import is_sm90_supported, is_sm100_supported
+        from sglang.srt.utils import (
+            is_sm90_supported,
+            is_sm100_supported,
+            is_sm120_supported,
+        )
 
-        assert is_sm90_supported(), "NVFP4 KV cache quantize requires SM90+ GPU"
+        assert (
+            is_sm100_supported() or is_sm120_supported() or is_sm90_supported()
+        ), "NVFP4 KV cache quantize requires SM100/SM120 or SM90 fallback GPU"
 
         b, m, n = tensor.shape
         tensor_2d = tensor.reshape(b * m, n)
 
+        # The KV cache path passes preloaded per-layer scales already on device.
+        # Keep scalar/0-d support for tests and future fallback paths, but do not
+        # silently move tensor scales here.
         if isinstance(global_scale, (int, float)):
             global_scale = torch.tensor(
                 [global_scale], dtype=torch.float32, device=tensor.device
             )
         elif global_scale.dim() == 0:
             global_scale = global_scale.unsqueeze(0)
+        elif global_scale.device != tensor.device:
+            raise ValueError(
+                "NVFP4 global scale tensor must already be on the KV tensor device."
+            )
 
-        if is_sm100_supported():
+        if is_sm100_supported() or is_sm120_supported():
             from flashinfer import nvfp4_kv_quantize
 
             # nvfp4_kv_quantize takes global_scale directly (not inverted)
@@ -238,18 +249,29 @@ class NVFP4KVQuantizeUtil:
         Returns:
             Dequantized tensor of shape [B, M, N]
         """
-        from sglang.srt.utils import is_sm100_supported
+        from sglang.srt.utils import (
+            is_sm90_supported,
+            is_sm100_supported,
+            is_sm120_supported,
+        )
 
         b, m, n_half = quant_tensor.shape
 
+        # The KV cache path passes preloaded per-layer scales already on device.
+        # Keep scalar/0-d support for tests and future fallback paths, but do not
+        # silently move tensor scales here.
         if isinstance(global_scale, (int, float)):
             global_scale = torch.tensor(
                 [global_scale], dtype=torch.float32, device=quant_tensor.device
             )
         elif global_scale.dim() == 0:
             global_scale = global_scale.unsqueeze(0)
+        elif global_scale.device != quant_tensor.device:
+            raise ValueError(
+                "NVFP4 global scale tensor must already be on the KV tensor device."
+            )
 
-        if is_sm100_supported():
+        if is_sm100_supported() or is_sm120_supported():
             from flashinfer import nvfp4_kv_dequantize
 
             quant_2d = quant_tensor.view(torch.uint8).reshape(b * m, n_half)
@@ -259,6 +281,9 @@ class NVFP4KVQuantizeUtil:
             )
             return output_2d.reshape(b, m, -1)
         else:
+            assert (
+                is_sm90_supported()
+            ), "NVFP4 KV cache dequantize requires SM100/SM120 or SM90 fallback GPU"
             # Pure PyTorch fallback for SM90
             n = n_half * 2
             fp4_vals = torch.empty(
@@ -266,7 +291,8 @@ class NVFP4KVQuantizeUtil:
             )
             fp4_vals[..., 0::2] = quant_tensor & 0x0F
             fp4_vals[..., 1::2] = (quant_tensor >> 4) & 0x0F
-            float_vals = E2M1_VALUES[fp4_vals.long()]
+            values = quant_tensor.new_tensor(E2M1_VALUES, dtype=torch.float32)
+            float_vals = values[fp4_vals.long()]
             reshaped = float_vals.view(b, m * n // 16, 16)
             block_scales_float = block_scales.float().unsqueeze(-1)
             scaled = reshaped * block_scales_float

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -21,6 +21,10 @@ from sglang.srt.configs.model_config import (
 )
 from sglang.srt.distributed.parallel_state import get_world_group
 from sglang.srt.environ import envs
+from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
+    get_kv_cache_quant_method,
+    resolve_kv_cache_quant,
+)
 from sglang.srt.mem_cache.allocation_sizing import get_req_to_token_extra_context_len
 from sglang.srt.mem_cache.allocator import (
     BaseTokenToKVPoolAllocator,
@@ -153,6 +157,7 @@ class KVCacheConfigurator:
     gpu_id: int
     ps: ParallelState
     pp_group: Any
+    model: Any
     model_config: ModelConfig
     server_args: ServerArgs
     kv_cache_dtype: torch.dtype
@@ -178,6 +183,20 @@ class KVCacheConfigurator:
         self.mambaish_config = mambaish_config(self.model_config)
         self.hybrid_gdn_config = hybrid_gdn_config(self.model_config)
 
+    def _build_fp4_quant_method(self, *, num_layers: int):
+        if not is_float4_e2m1fn_x2(self.kv_cache_dtype):
+            return None
+        quant_name = resolve_kv_cache_quant(self.server_args.kv_cache_dtype)
+        if quant_name is None:
+            return None
+        quant_method = get_kv_cache_quant_method(
+            quant_name,
+            num_layers=num_layers,
+            device=self.device,
+        )
+        quant_method.load_scales_from_model(self.model)
+        return quant_method
+
     def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult:
         """Apply a resolved MemoryPoolConfig and initialize pools."""
         if not self.spec_algorithm.is_none() and self.is_draft_worker:
@@ -349,7 +368,7 @@ class KVCacheConfigurator:
                 "Supported configurations today: plain MHA models on CUDA with the FA "
                 "(fa3/fa4) prefill backend, --is-embedding, --chunked-prefill-size=-1, "
                 "--disable-radix-cache, no context-parallel attention, no HiSparse, "
-                "and --kv-cache-dtype != fp4_e2m1."
+                "and --kv-cache-dtype not in {nvfp4, fp4_mx_block16}."
             )
         return _InitializedPools(
             req_to_token_pool=req_to_token_pool,
@@ -552,7 +571,7 @@ class KVCacheConfigurator:
                 f"{unsupported_pool_family}. Supported configurations today: plain MHA "
                 "models on CUDA with the FA (fa3/fa4) prefill backend, --is-embedding, "
                 "--chunked-prefill-size=-1, --disable-radix-cache, no context-parallel "
-                "attention, no HiSparse, and --kv-cache-dtype != fp4_e2m1."
+                "attention, no HiSparse, and --kv-cache-dtype not in {nvfp4, fp4_mx_block16}."
             )
 
     def _build_req_to_token_pool(self, *, max_num_reqs: int) -> ReqToTokenPool:
@@ -786,18 +805,19 @@ class KVCacheConfigurator:
                     mha_pool_class=mha_pool_class,
                 )
             else:
+                quant_method = None
                 if is_float4_e2m1fn_x2(self.kv_cache_dtype):
                     assert (
                         not enable_page_major
                     ), "page-major KV layout is not supported with fp4 KV cache"
-                    token_to_kv_pool = self._build_mha_fp4_kv_pool(
-                        max_total_num_tokens=sizes.max_total_num_tokens,
-                    )
-                else:
-                    token_to_kv_pool = self._build_mha_kv_pool(
-                        max_total_num_tokens=sizes.max_total_num_tokens,
-                        mha_pool_class=mha_pool_class,
+                    quant_method = self._build_fp4_quant_method(
+                        num_layers=self.layer_info.num_effective_layers
                     )
+                token_to_kv_pool = self._build_mha_kv_pool(
+                    max_total_num_tokens=sizes.max_total_num_tokens,
+                    mha_pool_class=mha_pool_class,
+                    quant_method=quant_method,
+                )
         return token_to_kv_pool
 
     def _build_dsv4_kv_pool(
@@ -1179,6 +1199,18 @@ class KVCacheConfigurator:
                 "kv_lora_rank": self.model_config.kv_lora_rank,
                 "qk_rope_head_dim": self.model_config.qk_rope_head_dim,
             }
+        full_attention_layer_ids = (
+            [0]
+            if self.is_draft_worker
+            else [
+                i
+                for i in self.mambaish_config.full_attention_layer_ids
+                if self.layer_info.start_layer <= i < self.layer_info.end_layer
+            ]
+        )
+        quant_method = self._build_fp4_quant_method(
+            num_layers=len(full_attention_layer_ids)
+        )
         token_to_kv_pool = HybridLinearKVPool(
             page_size=self.server_args.page_size,
             size=max_total_num_tokens,
@@ -1186,15 +1218,7 @@ class KVCacheConfigurator:
             head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
             head_dim=self.model_config.head_dim,
             # if draft worker, we only need 1 attention layer's kv pool
-            full_attention_layer_ids=(
-                [0]
-                if self.is_draft_worker
-                else [
-                    i
-                    for i in self.mambaish_config.full_attention_layer_ids
-                    if self.layer_info.start_layer <= i < self.layer_info.end_layer
-                ]
-            ),
+            full_attention_layer_ids=full_attention_layer_ids,
             device=self.device,
             mamba_pool=req_to_token_pool.mamba_pool,
             enable_memory_saver=self.server_args.enable_memory_saver,
@@ -1202,7 +1226,8 @@ class KVCacheConfigurator:
             use_mla=self.use_mla_backend,
             start_layer=self.layer_info.start_layer,
             full_kv_pool_class=mha_pool_class,
-            post_capture_active=self.post_capture_kv_active,
+            quant_method=quant_method,
+            post_capture_active=self.post_capture_kv_active and quant_method is None,
             **extra_args,
         )
         return token_to_kv_pool
@@ -1226,13 +1251,18 @@ class KVCacheConfigurator:
         return token_to_kv_pool
 
     def _build_mha_kv_pool(
-        self, *, max_total_num_tokens: int, mha_pool_class: type
+        self, *, max_total_num_tokens: int, mha_pool_class: type, quant_method=None
     ) -> KVCache:
         pool_cls = (
             NoOpMHATokenToKVPool
             if self.server_args.prefill_only_disable_kv_cache
             else mha_pool_class
         )
+        pool_kwargs = {}
+        if quant_method is not None:
+            pool_kwargs["quant_method"] = quant_method
+        else:
+            pool_kwargs["post_capture_active"] = self.post_capture_kv_active
         token_to_kv_pool = pool_cls(
             max_total_num_tokens,
             page_size=self.server_args.page_size,
@@ -1247,7 +1277,7 @@ class KVCacheConfigurator:
             end_layer=self.layer_info.end_layer,
             enable_alt_stream=not self.server_args.enable_pdmux,
             enable_kv_cache_copy=(self.server_args.speculative_algorithm is not None),
-            post_capture_active=self.post_capture_kv_active,
+            **pool_kwargs,
         )
         return token_to_kv_pool
 

python/sglang/srt/mem_cache/kv_cache_dtype.py

@@ -53,14 +53,23 @@ def configure_kv_cache_dtype(
     elif server_args_kv_cache_dtype in ("bf16", "bfloat16"):
         kv_cache_dtype = torch.bfloat16
     elif server_args_kv_cache_dtype == "fp4_e2m1":
+        raise ValueError(
+            "--kv-cache-dtype=fp4_e2m1 is deprecated. "
+            "Use --kv-cache-dtype=fp4_mx_block16."
+        )
+    elif server_args_kv_cache_dtype in ("nvfp4", "fp4_mx_block16"):
         if hasattr(torch, "float4_e2m1fn_x2"):
             kv_cache_dtype = torch.float4_e2m1fn_x2
-            logger.warning(f"FP4 (E2M1) KV Cache might lead to a accuracy drop!")
-        else:
             logger.warning(
-                f"--kv-cache-dtype falls back to 'auto' because this torch version does not support torch.float4_e2m1fn_x2"
+                "%s KV Cache might lead to an accuracy drop!",
+                server_args_kv_cache_dtype.upper(),
+            )
+        else:
+            raise ValueError(
+                f"--kv-cache-dtype={server_args_kv_cache_dtype} requires "
+                "torch.float4_e2m1fn_x2 support. Please use PyTorch 2.8.0+ "
+                "with CUDA 12.8+."
             )
-            kv_cache_dtype = model_dtype
     else:
         raise ValueError(f"Unsupported kv_cache_dtype: {server_args_kv_cache_dtype}.")
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -51,6 +51,9 @@ from sglang.srt.configs.mamba_utils import BaseLinearStateParams
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
+from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
+    UnquantizedKVCacheMethod,
+)
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator
 from sglang.srt.mem_cache.kv_vmm_backing import KvVmmBufferOwner
@@ -73,6 +76,7 @@ from sglang.srt.utils import (
     cpu_has_amx_support,
     is_cpu,
     is_cuda,
+    is_float4_e2m1fn_x2,
     is_hip,
     is_npu,
     next_power_of_2,
@@ -1387,6 +1391,24 @@ class KVCache(abc.ABC):
     def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
         raise NotImplementedError()
 
+    def get_kv_cache_quant_method(self) -> Any:
+        """Return the concrete KV quant method, unwrapping composite KV pools."""
+        fallback = None
+        for pool in (
+            self,
+            getattr(self, "full_kv_pool", None),
+            getattr(self, "swa_kv_pool", None),
+        ):
+            if pool is None:
+                continue
+            quant_method = getattr(pool, "quant_method", None)
+            if quant_method is None:
+                continue
+            if getattr(quant_method, "name", None) != "unquantized":
+                return quant_method
+            fallback = quant_method
+        return fallback
+
     def maybe_get_custom_mem_pool(self):
         return self.custom_mem_pool
 
@@ -1411,6 +1433,7 @@ class MHATokenToKVPool(KVCache):
         enable_alt_stream: bool = True,
         enable_kv_cache_copy: bool = False,
         kv_cache_layout: Optional[str] = None,
+        quant_method=None,
         post_capture_active: bool = False,
     ):
         if post_capture_active:
@@ -1482,6 +1505,10 @@ class MHATokenToKVPool(KVCache):
                     assert self.head_dim % self._kv_vector_x == 0
                     assert self.v_head_dim % self._kv_vector_x == 0
 
+        self.quant_method = (
+            quant_method if quant_method is not None else UnquantizedKVCacheMethod()
+        )
+
         self._create_buffers()
 
         self.device_module = torch.get_device_module(self.device)
@@ -1554,12 +1581,100 @@ class MHATokenToKVPool(KVCache):
             self._kv_copy_config,
         )
 
+    @property
+    def is_quantized_kv_cache(self) -> bool:
+        return not isinstance(self.quant_method, UnquantizedKVCacheMethod)
+
     def _create_buffers(self):
-        if self.post_capture_active:
-            self._alloc_post_capture_buffers()
+        if self.is_quantized_kv_cache:
+            if self.post_capture_active:
+                raise NotImplementedError(
+                    "Post-capture KV backing is not supported for quantized KV cache."
+                )
+            self._create_quantized_buffers()
         else:
-            self._create_buffers_normal()
+            self.k_scale_buffer = None
+            self.v_scale_buffer = None
+            self.dq_k_buffer = None
+            self.dq_v_buffer = None
+            if self.post_capture_active:
+                self._alloc_post_capture_buffers()
+            else:
+                self._create_buffers_normal()
         self._kv_buffer_descs = self._build_kv_buffer_descs()
+        self._init_data_ptrs_and_strides()
+
+    def _create_quantized_buffers(self):
+        # Quantized recipes own packed-data, scale, and workspace shapes.
+        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            with (
+                torch.cuda.use_mem_pool(self.custom_mem_pool)
+                if self.enable_custom_mem_pool
+                else nullcontext()
+            ):
+                buf = self.quant_method.create_buffers(
+                    self.size + self.page_size,
+                    self.head_num,
+                    self.head_dim,
+                    self.layer_num,
+                    self.device,
+                )
+        self.k_buffer = buf["k_buffer"]
+        self.v_buffer = buf["v_buffer"]
+        self.k_scale_buffer = buf.get("k_scale_buffer")
+        self.v_scale_buffer = buf.get("v_scale_buffer")
+        self.dq_k_buffer = buf.get("dq_k_buffer")
+        self.dq_v_buffer = buf.get("dq_v_buffer")
+        self.store_dtype = buf.get("store_dtype", torch.uint8)
+        self._check_quantized_buffer_access_requirements()
+
+    def _check_quantized_buffer_access_requirements(self):
+        expected_workspace_dtype = self.quant_method.dequant_workspace_dtype()
+        has_k_workspace = self.dq_k_buffer is not None
+        has_v_workspace = self.dq_v_buffer is not None
+        if has_k_workspace != has_v_workspace:
+            raise RuntimeError(
+                f"KV cache method {self.quant_method.name!r} created only one "
+                "dequant workspace buffer."
+            )
+
+        if expected_workspace_dtype is None:
+            if has_k_workspace:
+                raise RuntimeError(
+                    f"KV cache method {self.quant_method.name!r} does not declare "
+                    "DEQUANT_WORKSPACE access but created dequant buffers."
+                )
+            return
+
+        if not has_k_workspace:
+            raise RuntimeError(
+                f"KV cache method {self.quant_method.name!r} declares "
+                "DEQUANT_WORKSPACE access but did not create dequant buffers."
+            )
+
+        if (
+            self.dq_k_buffer.dtype != expected_workspace_dtype
+            or self.dq_v_buffer.dtype != expected_workspace_dtype
+        ):
+            raise RuntimeError(
+                f"KV cache method {self.quant_method.name!r} declares dequant "
+                f"workspace dtype {expected_workspace_dtype}, but created "
+                f"{self.dq_k_buffer.dtype}/{self.dq_v_buffer.dtype}."
+            )
+
+    def _slot_move_pointer_buffers(self):
+        """Buffers whose pointers/strides are used when KV slots are remapped.
+
+        FP4 KV cache stores data and per-block scales separately, so slot moves
+        must update both. This list feeds data_ptrs/data_strides; it does not
+        copy tensor contents by itself.
+        """
+        buffers = [*self.k_buffer, *self.v_buffer]
+        if getattr(self, "k_scale_buffer", None) is not None:
+            buffers.extend([*self.k_scale_buffer, *self.v_scale_buffer])
+        return buffers
+
+    def _init_data_ptrs_and_strides(self):
         self.k_data_ptrs = torch.tensor(
             [x.data_ptr() for x in self.k_buffer],
             dtype=torch.uint64,
@@ -1570,11 +1685,16 @@ class MHATokenToKVPool(KVCache):
             dtype=torch.uint64,
             device=self.device,
         )
-        self.data_ptrs = torch.cat([self.k_data_ptrs, self.v_data_ptrs], dim=0)
+        slot_move_pointer_buffers = self._slot_move_pointer_buffers()
+        self.data_ptrs = torch.tensor(
+            [x.data_ptr() for x in slot_move_pointer_buffers],
+            dtype=torch.uint64,
+            device=self.device,
+        )
         self.data_strides = torch.tensor(
             [
                 np.prod(x.shape[1:]) * x.dtype.itemsize
-                for x in self.k_buffer + self.v_buffer
+                for x in slot_move_pointer_buffers
             ],
             device=self.device,
         )
@@ -1716,6 +1836,14 @@ class MHATokenToKVPool(KVCache):
     def _clear_buffers(self):
         del self.k_buffer
         del self.v_buffer
+        if hasattr(self, "k_scale_buffer") and self.k_scale_buffer is not None:
+            del self.k_scale_buffer
+        if hasattr(self, "v_scale_buffer") and self.v_scale_buffer is not None:
+            del self.v_scale_buffer
+        if hasattr(self, "dq_k_buffer") and self.dq_k_buffer is not None:
+            del self.dq_k_buffer
+        if hasattr(self, "dq_v_buffer") and self.dq_v_buffer is not None:
+            del self.dq_v_buffer
         if self._post_capture_owner is not None:
             self._post_capture_owner.close()
             self._post_capture_owner = None
@@ -1723,12 +1851,14 @@ class MHATokenToKVPool(KVCache):
     def get_kv_size_bytes(self):
         assert hasattr(self, "k_buffer")
         assert hasattr(self, "v_buffer")
-        k_size_bytes = 0
-        for k_cache in self.k_buffer:
-            k_size_bytes += get_tensor_size_bytes(k_cache)
-        v_size_bytes = 0
-        for v_cache in self.v_buffer:
-            v_size_bytes += get_tensor_size_bytes(v_cache)
+        k_size_bytes = get_tensor_size_bytes(self.k_buffer)
+        v_size_bytes = get_tensor_size_bytes(self.v_buffer)
+        if getattr(self, "k_scale_buffer", None) is not None:
+            k_size_bytes += get_tensor_size_bytes(self.k_scale_buffer)
+            v_size_bytes += get_tensor_size_bytes(self.v_scale_buffer)
+        if getattr(self, "dq_k_buffer", None) is not None:
+            k_size_bytes += get_tensor_size_bytes(self.dq_k_buffer)
+            v_size_bytes += get_tensor_size_bytes(self.dq_v_buffer)
         return k_size_bytes, v_size_bytes
 
     # for disagg
@@ -1798,9 +1928,19 @@ class MHATokenToKVPool(KVCache):
 
     def _get_key_buffer(self, layer_id: int):
         # for internal use of referencing
+        local_layer_id = layer_id - self.start_layer
+        if (
+            self.is_quantized_kv_cache
+            and self.quant_method.needs_plain_kv_dequant_read()
+        ):
+            return self.quant_method.dequantize_kv_tensor(
+                self.k_buffer[local_layer_id],
+                self.k_scale_buffer[local_layer_id],
+                layer_id,
+            )
         if self.store_dtype != self.dtype:
-            return self.k_buffer[layer_id - self.start_layer].view(self.dtype)
-        return self.k_buffer[layer_id - self.start_layer]
+            return self.k_buffer[local_layer_id].view(self.dtype)
+        return self.k_buffer[local_layer_id]
 
     def get_key_buffer(self, layer_id: int):
         # note: get_key_buffer is hooked with synchronization for layer-wise KV cache loading
@@ -1812,9 +1952,19 @@ class MHATokenToKVPool(KVCache):
 
     def _get_value_buffer(self, layer_id: int):
         # for internal use of referencing
+        local_layer_id = layer_id - self.start_layer
+        if (
+            self.is_quantized_kv_cache
+            and self.quant_method.needs_plain_kv_dequant_read()
+        ):
+            return self.quant_method.dequantize_kv_tensor(
+                self.v_buffer[local_layer_id],
+                self.v_scale_buffer[local_layer_id],
+                layer_id,
+            )
         if self.store_dtype != self.dtype:
-            return self.v_buffer[layer_id - self.start_layer].view(self.dtype)
-        return self.v_buffer[layer_id - self.start_layer]
+            return self.v_buffer[local_layer_id].view(self.dtype)
+        return self.v_buffer[local_layer_id]
 
     def get_value_buffer(self, layer_id: int):
         if self.layer_transfer_counter is not None:
@@ -1839,10 +1989,25 @@ class MHATokenToKVPool(KVCache):
         # Catch stale slot ids here instead of as illegal-addr / silent KV
         # corruption in the store_kvcache write (gated on SGLANG_ENABLE_ASYNC_ASSERT).
         maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA)")
-        if layer_id_override is not None:
-            layer_id = layer_id_override
-        else:
-            layer_id = layer.layer_id
+        layer_id = (
+            layer_id_override if layer_id_override is not None else layer.layer_id
+        )
+        global_layer_id = layer.layer_id if layer is not None else layer_id
+
+        if self.is_quantized_kv_cache:
+            if dcp_kv_mask is not None:
+                raise RuntimeError("dcp_kv_mask is not supported for FP4 KV cache.")
+            self._set_quantized_kv_buffer(
+                layer_id,
+                global_layer_id,
+                loc,
+                cache_k,
+                cache_v,
+                k_scale,
+                v_scale,
+            )
+            return
+
         if cache_k.dtype != self.dtype:
             if k_scale is not None:
                 cache_k.div_(k_scale)
@@ -1936,6 +2101,255 @@ class MHATokenToKVPool(KVCache):
             same_kv_dim=self.same_kv_dim,
         )
 
+    def _quantized_scales(self, global_layer_id: int, k_scale, v_scale):
+        if k_scale is None and hasattr(self.quant_method, "k_scales_gpu"):
+            k_scale = self.quant_method.k_scales_gpu[
+                global_layer_id : global_layer_id + 1
+            ]
+            v_scale = self.quant_method.v_scales_gpu[
+                global_layer_id : global_layer_id + 1
+            ]
+        return k_scale, v_scale
+
+    def _set_quantized_kv_buffer(
+        self,
+        layer_id: int,
+        global_layer_id: int,
+        loc_info,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+        k_scale=None,
+        v_scale=None,
+    ) -> None:
+        loc, _, _ = unwrap_write_loc(loc_info)
+        local_layer_id = layer_id - self.start_layer
+        k_scale, v_scale = self._quantized_scales(global_layer_id, k_scale, v_scale)
+        self.quant_method.quantize_and_store(
+            self.k_buffer[local_layer_id],
+            self.v_buffer[local_layer_id],
+            (
+                self.k_scale_buffer[local_layer_id]
+                if self.k_scale_buffer is not None
+                else None
+            ),
+            (
+                self.v_scale_buffer[local_layer_id]
+                if self.v_scale_buffer is not None
+                else None
+            ),
+            loc,
+            cache_k,
+            cache_v,
+            k_scale,
+            v_scale,
+        )
+
+    def get_raw_kv_buffer(
+        self, layer_id: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        local_layer_id = layer_id - self.start_layer
+        if self.k_scale_buffer is None or self.v_scale_buffer is None:
+            raise RuntimeError("Raw FP4 KV cache requested from a non-FP4 KV pool.")
+        k_scale = self.k_scale_buffer[local_layer_id]
+        v_scale = self.v_scale_buffer[local_layer_id]
+        scale_view_dtype = self.quant_method.scale_buffer_view_dtype()
+        if scale_view_dtype is not None:
+            k_scale = k_scale.view(scale_view_dtype)
+            v_scale = v_scale.view(scale_view_dtype)
+        return (
+            self.k_buffer[local_layer_id],
+            self.v_buffer[local_layer_id],
+            k_scale,
+            v_scale,
+        )
+
+    def get_dequant_workspace(self) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.dq_k_buffer is None or self.dq_v_buffer is None:
+            raise RuntimeError(
+                "Dequant workspace requested from a KV pool without FP4 dequant buffers."
+            )
+        return self.dq_k_buffer, self.dq_v_buffer
+
+    def get_flashinfer_dequant_workspace_kv_buffer(
+        self,
+        layer: RadixAttention,
+        req_to_token: torch.Tensor,
+        req_pool_indices_cpu,
+        extend_prefix_lens_cpu,
+        extend_seq_lens_cpu,
+        page_size: int,
+        *,
+        prepare_workspace: bool,
+        use_ragged: bool,
+        k_cur: Optional[torch.Tensor] = None,
+        v_cur: Optional[torch.Tensor] = None,
+        layer_id_override: Optional[int] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return the FlashInfer FP8 KV view for a quantized KV cache.
+
+        FlashInfer prefill consumes FP8 KV. Quantized pools store packed FP4 plus
+        per-block scales, so the pool owns the dequant workspace and returns the
+        view shape expected by FlashInfer.
+        """
+        if not self.is_quantized_kv_cache:
+            raise RuntimeError(
+                "FlashInfer quantized KV buffer requested from a non-quantized KV pool."
+            )
+
+        if prepare_workspace:
+            transfer_cur_kv = not use_ragged
+            k_cur_fp8 = (
+                k_cur.to(torch.float8_e4m3fn)
+                if k_cur is not None and transfer_cur_kv
+                else None
+            )
+            v_cur_fp8 = (
+                v_cur.to(torch.float8_e4m3fn)
+                if v_cur is not None and transfer_cur_kv
+                else None
+            )
+            self._prepare_dequant_extend_workspace(
+                layer.layer_id if layer_id_override is None else layer_id_override,
+                layer.layer_id,
+                req_to_token,
+                req_pool_indices_cpu,
+                extend_prefix_lens_cpu,
+                extend_seq_lens_cpu,
+                page_size,
+                k_cur_fp8=k_cur_fp8,
+                v_cur_fp8=v_cur_fp8,
+            )
+
+        k_buffer_dq, v_buffer_dq = self.get_dequant_workspace()
+        return (
+            k_buffer_dq.view(-1, layer.tp_k_head_num, layer.head_dim),
+            v_buffer_dq.view(-1, layer.tp_v_head_num, layer.head_dim),
+        )
+
+    def get_flashinfer_decode_dequant_workspace_kv_buffer(
+        self,
+        layer: RadixAttention,
+        req_to_token: torch.Tensor,
+        req_pool_indices,
+        seq_lens,
+        *,
+        layer_id_override: Optional[int] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if not self.is_quantized_kv_cache:
+            raise RuntimeError(
+                "FlashInfer dequant workspace requested from a non-quantized KV pool."
+            )
+
+        self._prepare_dequant_decode_workspace(
+            layer.layer_id if layer_id_override is None else layer_id_override,
+            layer.layer_id,
+            req_to_token,
+            req_pool_indices,
+            seq_lens,
+        )
+        k_buffer_dq, v_buffer_dq = self.get_dequant_workspace()
+        return (
+            k_buffer_dq.view(-1, layer.tp_k_head_num, layer.head_dim),
+            v_buffer_dq.view(-1, layer.tp_v_head_num, layer.head_dim),
+        )
+
+    @staticmethod
+    def _to_cpu_int_list(values) -> list[int]:
+        if isinstance(values, list):
+            return [int(value) for value in values]
+        if isinstance(values, torch.Tensor):
+            return [int(value) for value in values.cpu().tolist()]
+        return [int(value) for value in values]
+
+    def _prepare_dequant_extend_workspace(
+        self,
+        layer_id: int,
+        global_layer_id: int,
+        req_to_token: torch.Tensor,
+        req_pool_indices_cpu,
+        extend_prefix_lens_cpu,
+        extend_seq_lens_cpu,
+        page_size: int,
+        k_cur_fp8: Optional[torch.Tensor] = None,
+        v_cur_fp8: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Build the shared FP8 workspace used by FlashInfer extend attention.
+
+        Cached prefix tokens are stored as packed FP4 plus per-block scales, so
+        paged prefill dequantizes those prefix tokens into the FP8 workspace.
+        The current extend chunk can already be FP8 and is copied into the same
+        workspace after the prefix region.
+        """
+        k_fp4, v_fp4, k_scales, v_scales = self.get_raw_kv_buffer(layer_id)
+        dq_k, dq_v = self.get_dequant_workspace()
+
+        cur_batch_start_loc_cpu = 0
+        cur_token_idx_dq = page_size
+
+        for i in range(len(req_pool_indices_cpu)):
+            req_idx = int(req_pool_indices_cpu[i])
+            prev_len = int(extend_prefix_lens_cpu[i])
+            extend_len = int(extend_seq_lens_cpu[i])
+
+            if prev_len > 0:
+                prev_indices = req_to_token[req_idx, :prev_len]
+                k_prev_fp8, v_prev_fp8 = self.quant_method.dequantize_prev_kv(
+                    k_fp4[prev_indices],
+                    k_scales[prev_indices],
+                    v_fp4[prev_indices],
+                    v_scales[prev_indices],
+                    global_layer_id,
+                )
+                dq_k[cur_token_idx_dq : cur_token_idx_dq + prev_len] = k_prev_fp8
+                dq_v[cur_token_idx_dq : cur_token_idx_dq + prev_len] = v_prev_fp8
+
+            if k_cur_fp8 is not None:
+                cur_end = cur_batch_start_loc_cpu + extend_len
+                dst_start = cur_token_idx_dq + prev_len
+                dst_end = dst_start + extend_len
+                dq_k[dst_start:dst_end] = k_cur_fp8[cur_batch_start_loc_cpu:cur_end]
+                dq_v[dst_start:dst_end] = v_cur_fp8[cur_batch_start_loc_cpu:cur_end]
+                cur_batch_start_loc_cpu = cur_end
+
+            workspace_len = prev_len + (extend_len if k_cur_fp8 is not None else 0)
+            cur_token_idx_dq = (
+                (cur_token_idx_dq + workspace_len + page_size - 1)
+                // page_size
+                * page_size
+            )
+
+        return dq_k, dq_v
+
+    def _prepare_dequant_decode_workspace(
+        self,
+        layer_id: int,
+        global_layer_id: int,
+        req_to_token: torch.Tensor,
+        req_pool_indices,
+        seq_lens,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        k_fp4, v_fp4, k_scales, v_scales = self.get_raw_kv_buffer(layer_id)
+        dq_k, dq_v = self.get_dequant_workspace()
+
+        req_pool_indices_cpu = self._to_cpu_int_list(req_pool_indices)
+        seq_lens_cpu = self._to_cpu_int_list(seq_lens)
+
+        for req_idx, seq_len in zip(req_pool_indices_cpu, seq_lens_cpu):
+            if seq_len <= 0:
+                continue
+            kv_indices = req_to_token[req_idx, :seq_len]
+            k_prev_fp8, v_prev_fp8 = self.quant_method.dequantize_prev_kv(
+                k_fp4[kv_indices],
+                k_scales[kv_indices],
+                v_fp4[kv_indices],
+                v_scales[kv_indices],
+                global_layer_id,
+            )
+            dq_k[kv_indices] = k_prev_fp8
+            dq_v[kv_indices] = v_prev_fp8
+
+        return dq_k, dq_v
+
     def set_kv_buffer_prefix_valid(
         self,
         layer: RadixAttention,
@@ -2047,6 +2461,10 @@ class MHATokenToKVPool(KVCache):
         # per-layer buffers here ignore page_size in move_kv_cache_native.
         if self.use_native_move_kv_cache:
             move_kv_cache_native(self.k_buffer, self.v_buffer, tgt_loc, src_loc)
+            if getattr(self, "k_scale_buffer", None) is not None:
+                move_kv_cache_native(
+                    self.k_scale_buffer, self.v_scale_buffer, tgt_loc, src_loc
+                )
             return
 
         N = tgt_loc.numel()
@@ -2266,10 +2684,10 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool):
             cache_k_nope_fp4_sf = self.k_scale_buffer[layer_id - self.start_layer]
 
             from sglang.srt.layers.quantization.kvfp4_tensor import (
-                BlockFP4KVQuantizeUtil,
+                FP4MXBlock16KVQuantizeUtil,
             )
 
-            cache_k_nope_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize(
+            cache_k_nope_fp4_dequant = FP4MXBlock16KVQuantizeUtil.batched_dequantize(
                 cache_k_nope_fp4, cache_k_nope_fp4_sf
             )
             return cache_k_nope_fp4_dequant
@@ -2284,10 +2702,10 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool):
             cache_v_nope_fp4_sf = self.v_scale_buffer[layer_id - self.start_layer]
 
             from sglang.srt.layers.quantization.kvfp4_tensor import (
-                BlockFP4KVQuantizeUtil,
+                FP4MXBlock16KVQuantizeUtil,
             )
 
-            cache_v_nope_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize(
+            cache_v_nope_fp4_dequant = FP4MXBlock16KVQuantizeUtil.batched_dequantize(
                 cache_v_nope_fp4, cache_v_nope_fp4_sf
             )
             return cache_v_nope_fp4_dequant
@@ -2318,11 +2736,15 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool):
                 cache_v.div_(v_scale)
 
             from sglang.srt.layers.quantization.kvfp4_tensor import (
-                BlockFP4KVQuantizeUtil,
+                FP4MXBlock16KVQuantizeUtil,
             )
 
-            cache_k, cache_k_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_k)
-            cache_v, cache_v_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_v)
+            cache_k, cache_k_fp4_sf = FP4MXBlock16KVQuantizeUtil.batched_quantize(
+                cache_k
+            )
+            cache_v, cache_v_fp4_sf = FP4MXBlock16KVQuantizeUtil.batched_quantize(
+                cache_v
+            )
 
         if self.store_dtype != self.dtype:
             cache_k = cache_k.view(self.store_dtype)
@@ -2526,6 +2948,7 @@ class HybridLinearKVPool(KVCache):
         qk_rope_head_dim: int = None,
         start_layer: Optional[int] = None,
         full_kv_pool_class: Optional[type] = None,
+        quant_method=None,
         # When provided (shared-KV-pool path), use this pool for the
         # full-attention layers instead of constructing one internally.
         full_kv_pool: Optional[KVCache] = None,
@@ -2551,20 +2974,28 @@ class HybridLinearKVPool(KVCache):
             self.full_kv_pool = full_kv_pool
         elif not use_mla:
             TokenToKVPoolClass = MHATokenToKVPool
+            quant_method_kwarg = {"quant_method": quant_method}
 
             if current_platform.is_out_of_tree():
                 TokenToKVPoolClass = current_platform.get_mha_kv_pool_cls()
+                quant_method_kwarg = {}
             elif _is_npu:
+                assert not is_float4_e2m1fn_x2(
+                    dtype
+                ), "FP4 is not supported on NPU yet."
                 from sglang.srt.hardware_backend.npu.memory_pool_npu import (
                     NPUMHATokenToKVPool,
                 )
 
                 TokenToKVPoolClass = NPUMHATokenToKVPool
+                quant_method_kwarg = {}
             elif full_kv_pool_class is not None:
                 # Caller-selected MHA layout variant (e.g. the page-major
                 # PageMajorMHATokenToKVPool). NPU / out-of-tree classes keep
                 # priority since they don't understand alternate layouts.
                 TokenToKVPoolClass = full_kv_pool_class
+            else:
+                TokenToKVPoolClass = MHATokenToKVPool
 
             post_capture_kwargs = (
                 {"post_capture_active": True} if post_capture_active else {}
@@ -2579,6 +3010,7 @@ class HybridLinearKVPool(KVCache):
                 device=device,
                 enable_memory_saver=enable_memory_saver,
                 enable_kv_cache_copy=enable_kv_cache_copy,
+                **quant_method_kwarg,
                 **post_capture_kwargs,
             )
         else:
@@ -2665,14 +3097,18 @@ class HybridLinearKVPool(KVCache):
         if self.layer_transfer_counter is not None:
             self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
 
-    def get_key_buffer(self, layer_id: int):
+    def get_key_buffer(self, layer_id: int, scale: Optional[float] = None):
         self._wait_for_layer(layer_id)
         layer_id = self._transfer_full_attention_id(layer_id)
+        if scale is not None:
+            return self.full_kv_pool.get_key_buffer(layer_id, scale)
         return self.full_kv_pool.get_key_buffer(layer_id)
 
-    def get_value_buffer(self, layer_id: int):
+    def get_value_buffer(self, layer_id: int, scale: Optional[float] = None):
         self._wait_for_layer(layer_id)
         layer_id = self._transfer_full_attention_id(layer_id)
+        if scale is not None:
+            return self.full_kv_pool.get_value_buffer(layer_id, scale)
         return self.full_kv_pool.get_value_buffer(layer_id)
 
     def get_kv_buffer(self, layer_id: int):
@@ -2680,6 +3116,30 @@ class HybridLinearKVPool(KVCache):
         layer_id = self._transfer_full_attention_id(layer_id)
         return self.full_kv_pool.get_kv_buffer(layer_id)
 
+    def get_raw_kv_buffer(
+        self, layer_id: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        self._wait_for_layer(layer_id)
+        layer_id = self._transfer_full_attention_id(layer_id)
+        return self.full_kv_pool.get_raw_kv_buffer(layer_id)
+
+    def get_dequant_workspace(self) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.full_kv_pool.get_dequant_workspace()
+
+    def get_flashinfer_dequant_workspace_kv_buffer(self, layer, *args, **kwargs):
+        self._wait_for_layer(layer.layer_id)
+        local_layer_id = self._transfer_full_attention_id(layer.layer_id)
+        return self.full_kv_pool.get_flashinfer_dequant_workspace_kv_buffer(
+            layer, *args, layer_id_override=local_layer_id, **kwargs
+        )
+
+    def get_flashinfer_decode_dequant_workspace_kv_buffer(self, layer, *args, **kwargs):
+        self._wait_for_layer(layer.layer_id)
+        local_layer_id = self._transfer_full_attention_id(layer.layer_id)
+        return self.full_kv_pool.get_flashinfer_decode_dequant_workspace_kv_buffer(
+            layer, *args, layer_id_override=local_layer_id, **kwargs
+        )
+
     @contextmanager
     def _transfer_id_context(self, layer: RadixAttention):
         @contextmanager
@@ -2712,7 +3172,7 @@ class HybridLinearKVPool(KVCache):
         if not self.use_mla:
             write_loc = full_loc if full_loc is not None else loc
             self.full_kv_pool.set_kv_buffer(
-                None,
+                layer,
                 write_loc,
                 cache_k,
                 cache_v,
@@ -3095,10 +3555,10 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
             cache_k_nope_fp4_sf = self.kv_scale_buffer[layer_id - self.start_layer]
 
             from sglang.srt.layers.quantization.kvfp4_tensor import (
-                BlockFP4KVQuantizeUtil,
+                FP4MXBlock16KVQuantizeUtil,
             )
 
-            cache_k_nope_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize(
+            cache_k_nope_fp4_dequant = FP4MXBlock16KVQuantizeUtil.batched_dequantize(
                 cache_k_nope_fp4, cache_k_nope_fp4_sf
             )
             return cache_k_nope_fp4_dequant
@@ -3119,10 +3579,10 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
         assert not self.dsa_kv_cache_store_fp8
         if cache_k.dtype != self.dtype:
             from sglang.srt.layers.quantization.kvfp4_tensor import (
-                BlockFP4KVQuantizeUtil,
+                FP4MXBlock16KVQuantizeUtil,
             )
 
-            cache_k_fp4, cache_k_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize(
+            cache_k_fp4, cache_k_fp4_sf = FP4MXBlock16KVQuantizeUtil.batched_quantize(
                 cache_k
             )
 
@@ -3158,14 +3618,14 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
         else:
             if cache_k_nope.dtype != self.dtype:
                 from sglang.srt.layers.quantization.kvfp4_tensor import (
-                    BlockFP4KVQuantizeUtil,
+                    FP4MXBlock16KVQuantizeUtil,
                 )
 
                 cache_k_nope_fp4, cache_k_nope_fp4_sf = (
-                    BlockFP4KVQuantizeUtil.batched_quantize(cache_k_nope)
+                    FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k_nope)
                 )
                 cache_k_rope_fp4, cache_k_rope_fp4_sf = (
-                    BlockFP4KVQuantizeUtil.batched_quantize(cache_k_rope)
+                    FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k_rope)
                 )
 
             if self.store_dtype != self.dtype:

python/sglang/srt/model_executor/model_runner.py

@@ -516,6 +516,7 @@ class ModelRunner:
             gpu_id=self.gpu_id,
             ps=self.ps,
             pp_group=self.pp_group,
+            model=self.model,
             model_config=self.model_config,
             server_args=self.server_args,
             kv_cache_dtype=self.kv_cache_dtype,
@@ -1073,14 +1074,19 @@ class ModelRunner:
             )
 
     def configure_kv_cache_dtype(self):
+        spec_algorithm = getattr(self, "spec_algorithm", None)
         resolved_kv_cache_dtype, self.kv_cache_dtype = (
             kv_cache_dtype.configure_kv_cache_dtype(
                 server_args_kv_cache_dtype=self.server_args.kv_cache_dtype,
-                model=self.model,
-                model_dtype=self.dtype,
-                is_draft_worker=self.is_draft_worker,
-                is_dflash=self.spec_algorithm.is_dflash(),
-                speculative_draft_attention_backend=self.server_args.speculative_draft_attention_backend,
+                model=getattr(self, "model", None),
+                model_dtype=getattr(self, "dtype", torch.bfloat16),
+                is_draft_worker=getattr(self, "is_draft_worker", False),
+                is_dflash=(
+                    spec_algorithm.is_dflash() if spec_algorithm is not None else False
+                ),
+                speculative_draft_attention_backend=getattr(
+                    self.server_args, "speculative_draft_attention_backend", None
+                ),
             )
         )
         if resolved_kv_cache_dtype is not None:

python/sglang/srt/model_executor/pool_configurator.py

@@ -273,6 +273,8 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
                 cell_size = (cell_size // 2) + (
                     (n * k * effective_num_layers * 2 * kv_size) // scale_block_size
                 )
+                # FP4 prefill uses one shared FP8 dequant workspace across layers.
+                cell_size += n * k * 2 * kv_size
 
         return cell_size
 

python/sglang/srt/server_args.py

@@ -605,10 +605,21 @@ class ServerArgs:
             help=(
                 'Data type for kv cache storage. "auto" will use model data type. '
                 '"bf16" or "bfloat16" for BF16 KV cache. "fp8_e5m2" and '
-                '"fp8_e4m3" are supported for CUDA 11.8+. "fp4_e2m1" (only '
-                "mxfp4) is supported for CUDA 12.8+ and PyTorch 2.8.0+"
+                '"fp8_e4m3" are supported for CUDA 11.8+. "nvfp4" selects '
+                'the NVFP4 FP4 E2M1 KV cache recipe; "fp4_mx_block16" '
+                "selects the MX-style block-size-16 FP4 E2M1 KV cache "
+                "recipe. Both require CUDA 12.8+ and PyTorch 2.8.0+"
             ),
-            choices=["auto", "fp8_e5m2", "fp8_e4m3", "bf16", "bfloat16", "fp4_e2m1"],
+            choices=[
+                "auto",
+                "fp8_e5m2",
+                "fp8_e4m3",
+                "bf16",
+                "bfloat16",
+                "nvfp4",
+                "fp4_mx_block16",
+                "fp4_e2m1",
+            ],
             resolvable=True,
         ),
     ] = "auto"
@@ -4964,7 +4975,7 @@ class ServerArgs:
         """Check FP4 KV cache compatibility with the attention backend"""
         from sglang.srt.arg_groups.overrides import resolved_view
 
-        if self.kv_cache_dtype != "fp4_e2m1":
+        if self.kv_cache_dtype not in ("nvfp4", "fp4_mx_block16"):
             return
 
         use_mla_backend = self.use_mla_backend()
@@ -4972,6 +4983,13 @@ class ServerArgs:
         attention_backend = resolved_view(self).attention_backend
 
         if is_cuda():
+            if self.kv_cache_dtype == "nvfp4" and not (
+                is_sm100_supported() or is_sm120_supported()
+            ):
+                raise RuntimeError(
+                    "--kv-cache-dtype=nvfp4 requires Blackwell SM100 or SM120. "
+                    "Use --kv-cache-dtype=fp4_mx_block16 for the block-size-16 FP4 recipe."
+                )
             if (
                 prefill_backend != decode_backend and prefill_backend != "fa4"
             ):  # Take care of prefill=fa4 later
@@ -5008,7 +5026,6 @@ class ServerArgs:
                             "cutlass_mla",
                             "flashinfer",
                             "trtllm_mla",
-                            "flashmla",
                         ]
                         assert attention_backend in KV4_ATTENTION_MLA_BACKEND_CHOICES, (
                             f"KV4 MLA expects attention_backend to be one of "
@@ -5895,11 +5912,11 @@ class ServerArgs:
                 "Other prefill-only workloads may be supported in a future change once "
                 "their attention paths stop reading or writing the paged KV cache."
             )
-        if self.kv_cache_dtype == "fp4_e2m1":
+        if self.kv_cache_dtype in ("nvfp4", "fp4_mx_block16"):
             raise ValueError(
                 "--prefill-only-disable-kv-cache does not currently support "
-                "--kv-cache-dtype=fp4_e2m1 because the FP4 pool uses a separate "
-                "allocation path."
+                "--kv-cache-dtype=nvfp4 or --kv-cache-dtype=fp4_mx_block16 because "
+                "the FP4 pool uses a separate allocation path."
             )
 
         # Structural preconditions for the FA backend's fa_skip_kv_cache path,