feat(kv-cache): support SM100 NVFP4 GenMHA and speculative decoding (#36340)

python/sglang/srt/arg_groups/fields/exec_.py

@@ -158,6 +158,21 @@ class ExecKernel(msgspec.Struct):
             resolvable=True,
         ),
     ] = None
+    prefill_kv_cache_dequant_dtype: A[
+        str,
+        Arg(
+            help=(
+                "Online dequantization dtype used by prefill attention when "
+                "--kv-cache-dtype=nvfp4. 'nvfp4' reads the packed cache directly "
+                "without additional dequantization; 'fp8_e4m3' dequantizes it "
+                "into a temporary FP8 workspace. This does not change the stored "
+                "KV-cache dtype. 'auto' selects NVFP4 without additional "
+                "dequantization on SM100 and FP8 E4M3 otherwise."
+            ),
+            choices=["auto", "nvfp4", "fp8_e4m3"],
+            resolvable=True,
+        ),
+    ] = "auto"
     sampling_backend: A[
         Optional[str],
         Arg(

python/sglang/srt/arg_groups/kv_cache_hook.py

@@ -21,6 +21,80 @@ from sglang.srt.runtime_context import get_platform
 
 logger = logging.getLogger(__name__)
 
+_NVFP4_PREFILL_BACKEND = {
+    "fp8_e4m3": "flashinfer",
+    "nvfp4": "trtllm_mha",
+}
+_NVFP4_PREFILL_DEQUANT_DTYPE = {
+    backend: dtype for dtype, backend in _NVFP4_PREFILL_BACKEND.items()
+}
+
+
+def handle_nvfp4_prefill_kv_dequant_dtype(server_args: Any) -> None:
+    """Resolve the public prefill dequantization dtype to attention backends."""
+
+    cfg = resolving_view(server_args)
+    requested_dtype = cfg.prefill_kv_cache_dequant_dtype
+    if cfg.kv_cache_dtype != "nvfp4":
+        if requested_dtype != "auto":
+            raise ValueError(
+                "--prefill-kv-cache-dequant-dtype applies only with "
+                "--kv-cache-dtype=nvfp4."
+            )
+        return
+
+    if requested_dtype == "auto":
+        explicit_backend = cfg.prefill_attention_backend or cfg.attention_backend
+        if explicit_backend in _NVFP4_PREFILL_DEQUANT_DTYPE:
+            requested_dtype = _NVFP4_PREFILL_DEQUANT_DTYPE[explicit_backend]
+        else:
+            if explicit_backend is not None:
+                raise ValueError(
+                    "NVFP4 prefill supports an FP8 E4M3 workspace or native "
+                    f"NVFP4, but backend {explicit_backend!r} provides neither."
+                )
+            requested_dtype = "nvfp4" if get_platform().is_sm100 else "fp8_e4m3"
+
+    if requested_dtype == "nvfp4" and not get_platform().is_sm100:
+        raise ValueError(
+            "Native NVFP4 prefill currently requires SM100; use "
+            "--prefill-kv-cache-dequant-dtype=fp8_e4m3 on this platform."
+        )
+
+    target_prefill_backend = _NVFP4_PREFILL_BACKEND[requested_dtype]
+    explicit_prefill_backend = cfg.prefill_attention_backend
+    if (
+        explicit_prefill_backend is not None
+        and explicit_prefill_backend != target_prefill_backend
+    ):
+        raise ValueError(
+            f"--prefill-kv-cache-dequant-dtype={requested_dtype} requires prefill "
+            f"backend {target_prefill_backend!r}, but "
+            f"--prefill-attention-backend={explicit_prefill_backend!r} was set. "
+            "Remove the backend option and select the KV dtype only."
+        )
+
+    explicit_decode_backend = cfg.decode_attention_backend
+    if explicit_decode_backend not in (None, "trtllm_mha"):
+        raise ValueError(
+            "NVFP4 decode requires --decode-attention-backend=trtllm_mha; got "
+            f"{explicit_decode_backend!r}. Remove the backend option; NVFP4 "
+            "selects the supported decode implementation automatically."
+        )
+
+    updates = {
+        "prefill_attention_backend": target_prefill_backend,
+        "decode_attention_backend": "trtllm_mha",
+    }
+    if cfg.prefill_kv_cache_dequant_dtype == "auto":
+        updates["prefill_kv_cache_dequant_dtype"] = requested_dtype
+    declare_resolution(server_args, "_handle_nvfp4_prefill_kv_dequant_dtype", **updates)
+    logger.info(
+        "NVFP4 prefill dequant dtype: %s; prefill input: %s; decode input: nvfp4.",
+        requested_dtype,
+        requested_dtype,
+    )
+
 
 def handle_mxfp8_kv_cache_compatibility(server_args: Any) -> None:
     """MXFP8 KV cache uses operands available only on SM100+ (Blackwell)."""
@@ -54,14 +128,149 @@ def handle_kv4_compatibility(server_args: Any) -> None:
                 "--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
-            logger.warning(
-                f"Attention: Using KV4 with PREFILL = {prefill_backend} "
-                f"and DECODE = {decode_backend}. "
-                f"Compatibility issues are unlikely, but may occur in rare edge cases."
+        if cfg.enable_unified_memory:
+            raise ValueError(
+                "FP4 KV cache does not yet support --enable-unified-memory: "
+                "the unified MHA pool does not allocate FP4 block scales or "
+                "the prefill dequant workspace."
             )
+
+        # SM100 trtllm_mha owns physical, kernel-native NVFP4 scales. The
+        # transfer and host-tier pools do not preserve that layout yet. Keep
+        # these combinations fail-fast while allowing target verification to
+        # reuse the same monolithic cache and GenMHA kernels.
+        uses_sm100_trtllm_nvfp4 = (
+            cfg.kv_cache_dtype == "nvfp4"
+            and get_platform().is_sm100
+            and "trtllm_mha" in (prefill_backend, decode_backend)
+        )
+        uses_mixed_nvfp4 = (
+            cfg.kv_cache_dtype == "nvfp4"
+            and prefill_backend == "flashinfer"
+            and decode_backend == "trtllm_mha"
+        )
+        uses_sm100_mixed_nvfp4 = uses_sm100_trtllm_nvfp4 and uses_mixed_nvfp4
+        speculative_algorithm = (
+            cfg.speculative_algorithm.upper()
+            if cfg.speculative_algorithm is not None
+            else None
+        )
+        supported_native_spec_algorithms = {
+            "EAGLE",
+            "EAGLE3",
+            "NEXTN",
+            "NGRAM",
+        }
+        if uses_sm100_trtllm_nvfp4 and speculative_algorithm in (
+            "DFLASH",
+            "DSPARK",
+        ):
+            # These workers commit only a prefix of a dense candidate block
+            # through set_kv_buffer_prefix_valid(). That specialized writer
+            # does not produce GenMHA's physical K/V scale layout yet.
+            raise ValueError(
+                "SM100 native NVFP4 speculative decoding does not yet support "
+                f"{speculative_algorithm}; use EAGLE/NEXTN or NGRAM."
+            )
+        if (
+            uses_sm100_trtllm_nvfp4
+            and speculative_algorithm is not None
+            and speculative_algorithm not in supported_native_spec_algorithms
+        ):
+            # Do not silently treat STANDALONE, FROZEN_KV_MTP, or a custom
+            # plugin algorithm as EAGLE. Their draft/cache-commit contracts may
+            # differ, and none currently has native-layout coverage here.
+            raise ValueError(
+                "SM100 native NVFP4 speculative decoding supports EAGLE, "
+                "EAGLE3, NEXTN, and breadth-1 NGRAM; got "
+                f"{speculative_algorithm}."
+            )
+        if (
+            uses_sm100_trtllm_nvfp4
+            and speculative_algorithm == "NGRAM"
+            and cfg.speculative_ngram_max_bfs_breadth != 1
+        ):
+            # TRT-LLM MHA's target-verify metadata supports a linear chain
+            # only. NGRAM defaults to a breadth-10 tree, so reject that default
+            # explicitly rather than reaching the later generic paged-backend
+            # assertion with a misleading compatibility message.
+            raise ValueError(
+                "SM100 native NVFP4 NGRAM speculative decoding requires "
+                "--speculative-ngram-max-bfs-breadth=1 because trtllm_mha "
+                "supports linear target verification only; got "
+                f"{cfg.speculative_ngram_max_bfs_breadth}."
+            )
+        uses_draft_model = speculative_algorithm not in (None, "NGRAM")
+        if uses_sm100_trtllm_nvfp4 and uses_draft_model:
+            # A draft worker owns another physical KV pool. It cannot inherit a
+            # target-only hybrid pair: draft-extend would then select the
+            # prefill child even though the draft pool and its block scales use
+            # the native GenMHA layout. Give every draft phase one layout and
+            # one backend, including prefill-graph capture and multi-step
+            # decode. This also covers explicit hybrid target configurations,
+            # for which the model-default hook intentionally does not choose a
+            # draft backend.
+            draft_backend = cfg.speculative_draft_attention_backend
+            if draft_backend is None:
+                logger.warning(
+                    "SM100 native NVFP4 speculative decoding uses trtllm_mha "
+                    "for the draft worker."
+                )
+                declare_resolution(
+                    server_args,
+                    "_handle_kv4_compatibility",
+                    speculative_draft_attention_backend="trtllm_mha",
+                )
+            elif draft_backend != "trtllm_mha":
+                raise ValueError(
+                    "SM100 native NVFP4 speculative decoding requires "
+                    "--speculative-draft-attention-backend=trtllm_mha so the "
+                    "draft worker consumes its physical NVFP4 KV layout; got "
+                    f"{draft_backend!r}."
+                )
+        if (
+            uses_sm100_mixed_nvfp4
+            and cfg.speculative_algorithm is not None
+            and cfg.speculative_attention_mode == "prefill"
+        ):
+            # FlashInfer prefill reads a transient FP8 dequant workspace. Its
+            # host-built page layout cannot be refreshed inside a target-verify
+            # CUDA graph, whereas the decode child consumes the physical NVFP4
+            # cache directly for both eager and graph execution.
+            logger.warning(
+                "SM100 mixed NVFP4 speculative decoding routes target verify "
+                "to trtllm_mha; overriding --speculative-attention-mode=prefill "
+                "to decode."
+            )
+            declare_resolution(
+                server_args,
+                "_handle_kv4_compatibility",
+                speculative_attention_mode="decode",
+            )
+        if uses_sm100_trtllm_nvfp4 and cfg.disaggregation_mode != "null":
+            raise ValueError(
+                "SM100 native NVFP4 with trtllm_mha does not yet support PD "
+                "disaggregation because its physical block-scale layout is not "
+                "implemented by the KV transfer path."
+            )
+        if uses_sm100_trtllm_nvfp4 and (
+            cfg.enable_hierarchical_cache or cfg.enable_lmcache
+        ):
+            raise ValueError(
+                "SM100 native NVFP4 with trtllm_mha does not yet support "
+                "hierarchical KV cache or LMCache because their host pools do "
+                "not preserve the physical block-scale layout."
+            )
+
+        if prefill_backend != decode_backend and prefill_backend != "fa4":
+            # NVFP4 with FP8 prefill is a supported mixed-storage recipe.
+            if not uses_mixed_nvfp4:
+                logger.warning(
+                    f"Attention: Using KV4 with PREFILL = {prefill_backend} "
+                    f"and DECODE = {decode_backend}. "
+                    "Compatibility issues are unlikely, but may occur in rare "
+                    "edge cases."
+                )
         else:
             if prefill_backend == "fa4":
                 if uses_mla:  # FA4 + MLA

python/sglang/srt/arg_groups/pipeline.py

@@ -163,6 +163,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
         handle_cache_compatibility,
         handle_kv4_compatibility,
         handle_mxfp8_kv_cache_compatibility,
+        handle_nvfp4_prefill_kv_dequant_dtype,
         handle_page_major_kv_layout,
         handle_prefill_only_disable_kv_cache,
         handle_unified_memory_pool,
@@ -258,6 +259,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
     )
 
     run_hook(handle_deterministic_inference, server_args)
+    run_hook(handle_nvfp4_prefill_kv_dequant_dtype, server_args)
     run_hook(handle_attention_backend_compatibility, server_args)
     # Must run after the attention backend is resolved so the trtllm_mla
     # default (auto-selected for DeepseekV3ForCausalLM on sm100) is visible.

python/sglang/srt/arg_groups/resolution_hooks.py

@@ -79,6 +79,7 @@ _OVERRIDABLE_HOOKS: FrozenSet[str] = frozenset(
         "handle_gpu_memory_settings",
         "handle_model_specific_adjustments",
         "handle_deterministic_inference",
+        "handle_nvfp4_prefill_kv_dequant_dtype",
         "handle_attention_backend_compatibility",
         "disable_prefill_cuda_graph_for_deepseek_trtllm_mla",
         "handle_mamba_backend",

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

@@ -434,7 +434,11 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
                 if get_platform().is_sm120:
                     allowed = {"triton", "trtllm_mha", "flashinfer"}
                 else:
-                    allowed = {"triton", "trtllm_mha", "fa4"}
+                    # FlashInfer paged prefill is also valid for SM100 hybrid
+                    # GDN models. In particular, quantized KV recipes use it
+                    # to expose an FP8 dequant workspace while a different
+                    # backend (for example TRT-LLM GenMHA) owns decode.
+                    allowed = {"triton", "trtllm_mha", "fa4", "flashinfer"}
                 prefill_be = runner.prefill_attention_backend_str
                 decode_be = runner.decode_attention_backend_str
                 assert prefill_be in allowed and decode_be in allowed, (

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

@@ -472,9 +472,25 @@ class FlashInferAttnBackend(AttentionBackend):
 
         fmha_backend = "auto"
         if get_platform().is_sm100:
+            fmha_backend = "fa2"
             # Disable CUTLASS backend when piecewise cuda graph is enabled
-            # due to TMA descriptor initialization issues on SM100 GPUs.
-            if not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
+            # due to TMA descriptor initialization issues on SM100 GPUs. The
+            # current FlashInfer SM100 CUTLASS FMHA dispatch only instantiates
+            # 64x64, 128x128, and 192x128 head dimensions. Keep unsupported
+            # shapes (for example Qwen3.5's 256x256) on the FA2 fallback.
+            cutlass_supported_head_dims = {
+                (64, 64),
+                (128, 128),
+                (192, 128),
+            }
+            head_dims = (
+                model_runner.model_config.head_dim,
+                model_runner.model_config.v_head_dim,
+            )
+            if (
+                head_dims in cutlass_supported_head_dims
+                and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
+            ):
                 fmha_backend = "cutlass"
         self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
             self.workspace_buffer, "NHD", backend=fmha_backend

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

@@ -42,14 +42,18 @@ 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.memory_pool import KVWriteLoc
+from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import (
     get_buffer,
+    get_exec,
     get_parallel,
     get_platform,
+    get_schedule,
     get_spec,
+    max_prefill_buffer_tokens,
+    max_speculative_num_draft_tokens,
 )
 from sglang.srt.speculative.ragged_verify import (
     build_ragged_target_verify_geometry,
@@ -76,6 +80,56 @@ DEFAULT_WORKSPACE_SIZE_MB = 512
 # Reuse this workspace buffer across all TRTLLM MHA wrappers
 
 
+def _native_fp4_decode_output_capacity(
+    max_running_requests: int,
+    max_draft_tokens: Optional[int],
+    max_cuda_graph_bs: Optional[int] = None,
+) -> int:
+    """Maximum FP8 output rows for eager/graph decode and target verify."""
+    request_capacity = max(max_running_requests, max_cuda_graph_bs or 0)
+    return request_capacity * max(1, max_draft_tokens or 1)
+
+
+def _native_fp4_prefill_output_capacity(
+    max_context_len: int,
+    max_prefill_tokens: int,
+    chunked_prefill_limit: int,
+) -> int:
+    """Maximum FP8 output rows for one admitted prefill batch."""
+    if chunked_prefill_limit > 0:
+        return chunked_prefill_limit
+    return max(max_context_len, max_prefill_tokens)
+
+
+def _trtllm_native_nvfp4_kv_buffer(token_to_kv_pool, layer_id: int):
+    """Return the pool-owned buffers in TRT-LLM GenMHA's native layout."""
+    pool = token_to_kv_pool
+    if isinstance(pool, HybridLinearKVPool):
+        pool._wait_for_layer(layer_id)
+        layer_id = pool._transfer_full_attention_id(layer_id)
+        pool = pool.full_kv_pool
+    elif pool.layer_transfer_counter is not None:
+        pool.layer_transfer_counter.wait_until(layer_id - pool.start_layer)
+
+    local_layer_id = layer_id - pool.start_layer
+    if pool.native_k_scale_buffer is None or pool.native_v_scale_buffer is None:
+        raise RuntimeError(
+            "TRT-LLM native FP4 KV cache requested from a pool without native scales."
+        )
+    k_scale = pool.native_k_scale_buffer[local_layer_id]
+    v_scale = pool.native_v_scale_buffer[local_layer_id]
+    scale_view_dtype = pool.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 (
+        pool.k_buffer[local_layer_id],
+        pool.v_buffer[local_layer_id],
+        k_scale,
+        v_scale,
+    )
+
+
 @dataclass
 class TRTLLMMHAMetadata:
     # Sequence lengths for the forward batch
@@ -140,14 +194,32 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         super().__init__(
             model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf
         )
+        self.prefill_kv_access = self.kv_cache_quant_method.resolve_attention_access(
+            "prefill", "trtllm_mha"
+        )
         self.decode_kv_access = self.kv_cache_quant_method.resolve_attention_access(
             "decode", "trtllm_mha"
         )
-        self._check_decode_kv_access()
+        prefill_is_trtllm_mha = (
+            model_runner.prefill_attention_backend_str == "trtllm_mha"
+        )
+        decode_is_trtllm_mha = model_runner.decode_attention_backend_str == "trtllm_mha"
+        if prefill_is_trtllm_mha:
+            self._check_prefill_kv_access()
+        if decode_is_trtllm_mha:
+            self._check_decode_kv_access()
+        self.prefill_uses_native_fp4 = (
+            prefill_is_trtllm_mha
+            and self.prefill_kv_access.kind == KVCacheAttentionAccessKind.NATIVE_FP4
+        )
         self.decode_uses_native_fp4 = (
-            self.decode_kv_access.kind == KVCacheAttentionAccessKind.NATIVE_FP4
+            decode_is_trtllm_mha
+            and self.decode_kv_access.kind == KVCacheAttentionAccessKind.NATIVE_FP4
         )
         self.is_nvfp4_kvcache = (
+            self.prefill_uses_native_fp4
+            and self.prefill_kv_access.scale_recipe == "nvfp4"
+        ) or (
             self.decode_uses_native_fp4
             and self.decode_kv_access.scale_recipe == "nvfp4"
         )
@@ -165,6 +237,57 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self.device = model_runner.device
 
+        # XQA (SM90/SM120) consumes the legacy linear NVFP4 scale layout and
+        # BF16 Q/O. TRT-LLM GenMHA (SM100) consumes physical HND scales and
+        # requires FP8 Q/O.
+        self.is_xqa_impl = get_platform().is_sm90 or get_platform().is_sm120
+        if self.prefill_uses_native_fp4 and self.is_xqa_impl:
+            raise ValueError(
+                "Native NVFP4 prefill with trtllm_mha requires SM100 "
+                "TRT-LLM GenMHA. Use --prefill-attention-backend flashinfer "
+                "with XQA on SM90/SM120."
+            )
+        self.uses_trtllm_gen_native_fp4 = self.is_nvfp4_kvcache and not self.is_xqa_impl
+
+        # Speculative decoding
+        # Only support topk <= 1 for now.
+        self.topk = get_spec().speculative_eagle_topk or 0
+        self.speculative_step_id = speculative_step_id
+        self.target_verify_metadata = {}
+        self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens
+
+        self._nvfp4_fp8_output = None
+        if self.uses_trtllm_gen_native_fp4:
+            prefill_limit = 0
+            if self.prefill_uses_native_fp4 and not skip_prefill:
+                # Includes PP dynamic-chunk growth and piecewise capture bounds.
+                prefill_limit = _native_fp4_prefill_output_capacity(
+                    self.max_context_len,
+                    get_schedule().max_prefill_tokens or 0,
+                    max_prefill_buffer_tokens(),
+                )
+            decode_limit = 0
+            if self.decode_uses_native_fp4:
+                # TARGET_VERIFY submits one query row per draft token. Use the
+                # widest adaptive-spec candidate too: the output buffer is
+                # shared by eager execution and every captured CUDA graph.
+                decode_limit = _native_fp4_decode_output_capacity(
+                    model_runner.max_running_requests,
+                    max_speculative_num_draft_tokens(),
+                    get_exec().graph.cuda_graph_config.decode.max_bs,
+                )
+            max_native_tokens = max(prefill_limit, decode_limit)
+            num_q_heads = config.num_attention_heads // get_parallel().attn_tp_size
+            self._nvfp4_fp8_output = get_buffer(
+                f"trtllm_mha_nvfp4_output_{max_native_tokens}_"
+                f"{num_q_heads}_{config.head_dim}",
+                lambda: torch.empty(
+                    (max_native_tokens, num_q_heads, config.head_dim),
+                    dtype=torch.float8_e4m3fn,
+                    device=self.device,
+                ),
+            )
+
         # Workspace allocation
         self.workspace_size = workspace_size_bytes
         # Allocate buffers
@@ -180,13 +303,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         # CUDA graph state
         self.decode_cuda_graph_metadata = {}
 
-        # Speculative decoding
-        # Only support topk <= 1 for now.
-        self.topk = get_spec().speculative_eagle_topk or 0
-        self.speculative_step_id = speculative_step_id
-        self.target_verify_metadata = {}
-
-        self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens
         # True iff the model declares ENCODER_ONLY (bidirectional) layers, which
         # need the expanded TARGET_VERIFY metadata (TRTLLMMHAMetadata.encoder_*).
         self.expand_encoder_only_verify = any(
@@ -230,11 +346,8 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         # TRTLLM-GEN:
         #   KV bf16: q_type = bf16, out_type=model_runner.dtype
         #   KV fp8: q_type = fp8, out_type=model_runner.dtype
-        self.is_xqa_impl = get_platform().is_sm90 or get_platform().is_sm120
-
         # fmha_v2 prefill kernel supports SM90 and SM120
         self.use_fmha_v2 = get_platform().is_sm90 or get_platform().is_sm120
-
         # trtllm-gen serves page_size >= 128 only through its dynamic
         # tokens-per-page kernels, which exist solely for GQA with equal QK/V
         # head dims (power-of-2 pages). Mirror that precondition here so an
@@ -302,6 +415,65 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             f"trtllm_mha. Available decode accesses: {available}."
         )
 
+    def _check_prefill_kv_access(self) -> None:
+        supported_kinds = {
+            KVCacheAttentionAccessKind.PLAIN,
+            KVCacheAttentionAccessKind.NATIVE_FP4,
+        }
+        if (
+            self.prefill_kv_access is not None
+            and self.prefill_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("prefill")
+        raise ValueError(
+            f"KV cache method {method_name!r} does not support prefill with "
+            f"trtllm_mha. Available prefill accesses: {available}."
+        )
+
+    def _nvfp4_output_view(self, q: torch.Tensor) -> torch.Tensor:
+        if self._nvfp4_fp8_output is None:
+            raise RuntimeError("Native NVFP4 output buffer was not initialized.")
+        if q.shape[0] > self._nvfp4_fp8_output.shape[0]:
+            raise RuntimeError(
+                "TRT-LLM NVFP4 attention received more query tokens than its "
+                f"preallocated FP8 output buffer: {q.shape[0]} > "
+                f"{self._nvfp4_fp8_output.shape[0]}. Increase "
+                "--chunked-prefill-size, --max-running-requests, or the "
+                "speculative/CUDA-graph output capacity."
+            )
+        return self._nvfp4_fp8_output[: q.shape[0]].view_as(q)
+
+    def _forward_extend_uses_native_fp4(self, forward_batch: ForwardBatch) -> bool:
+        """Whether this extend-family call must consume physical NVFP4.
+
+        A hybrid backend can route TARGET_VERIFY to its decode child. That
+        child's normal extend role belongs to FlashInfer, so its prefill flag
+        is false even though this particular verify call must use the native
+        GenMHA decode layout.
+        """
+        return self.uses_trtllm_gen_native_fp4 and (
+            self.prefill_uses_native_fp4
+            or (
+                self.decode_uses_native_fp4
+                and forward_batch.forward_mode.is_target_verify()
+            )
+        )
+
+    def _finalize_nvfp4_output(
+        self, output: torch.Tensor, forward_batch: ForwardBatch
+    ) -> torch.Tensor:
+        if output.dtype == self.q_data_type:
+            return output
+        model_output = forward_batch._attn_output
+        if model_output is not None and model_output.numel() == output.numel():
+            model_output = model_output.view_as(output)
+            model_output.copy_(output)
+            return model_output
+        return output.to(self.q_data_type)
+
     @staticmethod
     def _resolve_swa_kv_pool(model_runner: ModelRunner) -> Optional[SWAKVPool]:
         """Return the SWAKVPool to translate against, or None for non-SWA models.
@@ -1159,10 +1331,18 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         sinks: Optional[torch.Tensor],
         q_len_per_req: int = 1,
         kv_cache_sf=None,
+        out: Optional[torch.Tensor] = None,
+        out_dtype: Optional[torch.dtype] = None,
     ) -> torch.Tensor:
         """Run decode, optionally sorting and splitting requests by KV length."""
 
-        def run_group(group_query, group_block_tables, group_seq_lens):
+        resolved_out_dtype = (
+            out_dtype
+            if out_dtype is not None
+            else (out.dtype if out is not None else self.q_data_type)
+        )
+
+        def run_group(group_query, group_block_tables, group_seq_lens, group_out=None):
             kwargs = {}
             if q_len_per_req != 1:
                 kwargs["q_len_per_req"] = q_len_per_req
@@ -1178,7 +1358,8 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                 window_left=window_left,
                 sinks=sinks,
                 skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
-                out_dtype=self.q_data_type,
+                out=group_out,
+                out_dtype=None if group_out is not None else resolved_out_dtype,
                 kv_cache_sf=kv_cache_sf,
                 multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer,
                 **kwargs,
@@ -1187,16 +1368,20 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         num_requests = seq_lens.shape[0]
         num_splits = min(self.decode_seq_len_splits, num_requests)
         if num_splits == 1:
-            return run_group(query, block_tables, seq_lens)
+            return run_group(query, block_tables, seq_lens, out)
 
         order = torch.argsort(seq_lens)
         query_by_request = query.view(
             num_requests, q_len_per_req, query.shape[-2], query.shape[-1]
         )
-        output_by_request = torch.empty(
-            query_by_request.shape,
-            dtype=self.q_data_type,
-            device=query.device,
+        output_by_request = (
+            out.view_as(query_by_request)
+            if out is not None
+            else torch.empty(
+                query_by_request.shape,
+                dtype=resolved_out_dtype,
+                device=query.device,
+            )
         )
         for indices in torch.tensor_split(order, num_splits):
             group_output = run_group(
@@ -1220,15 +1405,26 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         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
-        )
+        if self.is_xqa_impl:
+            k_fp4, v_fp4, k_scale, v_scale = self.token_to_kv_pool.get_raw_kv_buffer(
+                layer.layer_id
+            )
+        else:
+            k_fp4, v_fp4, k_scale, v_scale = _trtllm_native_nvfp4_kv_buffer(
+                self.token_to_kv_pool, 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
-        )
+        if self.is_xqa_impl:
+            kv_cache_block_scales = self._reshape_paged_kv_cache(
+                k_scale, v_scale, layer, layer.head_dim // 16
+            )
+        else:
+            # SM100 native scale buffers are already physical HND with
+            # contiguous [page-token, block-scale] dimensions. V is
+            # four-token interleaved.
+            kv_cache_block_scales = (k_scale, v_scale)
         return kv_cache, kv_cache_block_scales
 
     def forward_decode(
@@ -1271,7 +1467,10 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         # For XQA, q_dtype should be bf16. For trtllm-gen,
         # q_dtype should be FP8 when KV is in FP8.
         q_scale = 1.0
-        if (
+        if self.decode_uses_native_fp4 and not self.is_xqa_impl:
+            # SM100 TRT-LLM GenMHA requires FP8 Q for native NVFP4 KV.
+            q = q.to(torch.float8_e4m3fn)
+        elif (
             self.data_type == torch.float8_e4m3fn
             and not self.is_xqa_impl
             and not use_fused_qkv
@@ -1300,7 +1499,11 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         attention_sink = kwargs.get("sinks", None)
 
         page_table = self._get_layer_page_table(layer, forward_batch)
-
+        native_out = (
+            self._nvfp4_output_view(q)
+            if self.decode_uses_native_fp4 and not self.is_xqa_impl
+            else None
+        )
         o = self._run_fixed_q_len_decode(
             q,
             kv_cache,
@@ -1310,10 +1513,16 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             bmm2_scale=bmm2_scale,
             window_left=layer.sliding_window_size,
             sinks=attention_sink,
+            out=native_out,
+            out_dtype=(
+                None
+                if self.decode_uses_native_fp4 and not self.is_xqa_impl
+                else self.q_data_type
+            ),
             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)
+        if self.decode_uses_native_fp4 and not self.is_xqa_impl:
+            o = self._finalize_nvfp4_output(o, forward_batch)
 
         return o.view(-1, layer.tp_q_head_num * layer.head_dim)
 
@@ -1327,14 +1536,13 @@ 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
         cp_active = is_cp_active(forward_batch)
+        uses_native_fp4 = self._forward_extend_uses_native_fp4(forward_batch)
+        if uses_native_fp4 and cp_active:
+            raise NotImplementedError(
+                "Native NVFP4 TRT-LLM prefill does not yet support context parallelism."
+            )
 
         # The fused path writes rank-local K/V directly to cache. CP needs
         # the strategy to gather K/V into full logical token order first.
@@ -1369,12 +1577,13 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                         KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc),
                         k,
                         v,
-                        layer.k_scale,
-                        layer.v_scale,
+                        *self._kv_write_scales(layer),
                     )
 
         q_scale = 1.0
-        if (
+        if uses_native_fp4:
+            q = q.to(torch.float8_e4m3fn)
+        elif (
             self.data_type == torch.float8_e4m3fn
             and (
                 not self.is_xqa_impl
@@ -1383,39 +1592,51 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             and not use_fused_qkv
         ):
             q = q.to(torch.float8_e4m3fn)
-
         if self.use_fmha_v2:
             q = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
         else:
             q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim)
 
-        # NHD layout (native pool format): [num_pages, page_size, num_kv_heads, head_dim]
-        k_cache_raw, v_cache_raw = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
-
         is_decode_mode = (
             forward_batch.forward_mode.is_target_verify()
             or forward_batch.forward_mode.is_draft_extend_v2()
         )
 
-        if not self.use_fmha_v2 or is_decode_mode:
-            # Decode and SM100 batch_context kernels require HND layout.
-            k_cache, v_cache = self._reshape_paged_kv_cache(
-                k_cache_raw, v_cache_raw, layer, layer.head_dim
-            )
+        if uses_native_fp4:
+            kv_cache, kv_cache_block_scales = self._get_nvfp4_decode_kv_cache(layer)
+            k_cache, v_cache = kv_cache
         else:
-            k_cache = k_cache_raw.view(
-                -1, self.page_size, layer.tp_k_head_num, layer.head_dim
-            )
-            v_cache = v_cache_raw.view(
-                -1, self.page_size, layer.tp_v_head_num, layer.head_dim
+            # Native pool format is NHD:
+            # [num_pages, page_size, num_kv_heads, head_dim].
+            k_cache_raw, v_cache_raw = self.token_to_kv_pool.get_kv_buffer(
+                layer.layer_id
             )
+            if not self.use_fmha_v2 or is_decode_mode:
+                # Decode and SM100 batch_context kernels require HND layout.
+                k_cache, v_cache = self._reshape_paged_kv_cache(
+                    k_cache_raw, v_cache_raw, layer, layer.head_dim
+                )
+            else:
+                k_cache = k_cache_raw.view(
+                    -1, self.page_size, layer.tp_k_head_num, layer.head_dim
+                )
+                v_cache = v_cache_raw.view(
+                    -1, self.page_size, layer.tp_v_head_num, layer.head_dim
+                )
 
-        kv_cache = (k_cache, v_cache)
+            kv_cache = (k_cache, v_cache)
+            kv_cache_block_scales = None
         # sink: additional value per head in the denominator of the softmax.
         attention_sink = kwargs.get("sinks", None)
-        bmm1_scale, bmm2_scale = self._get_bmm_scales(layer, q_scale)
+        if uses_native_fp4:
+            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)
 
         page_table = self._get_layer_page_table(layer, forward_batch)
+        native_out = self._nvfp4_output_view(q) if uses_native_fp4 else None
 
         if is_decode_mode:
             if (
@@ -1445,7 +1666,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                     window_left=layer.sliding_window_size,
                     sinks=attention_sink,
                     skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
-                    out_dtype=self.q_data_type,
+                    out=native_out,
+                    out_dtype=(None if uses_native_fp4 else self.q_data_type),
+                    kv_cache_sf=kv_cache_block_scales,
                     q_len_per_req=1,
                     multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer,
                 )
@@ -1462,7 +1685,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                     window_left=layer.sliding_window_size,
                     sinks=attention_sink,
                     skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
-                    out_dtype=self.q_data_type,
+                    out=native_out,
+                    out_dtype=(None if uses_native_fp4 else self.q_data_type),
+                    kv_cache_sf=kv_cache_block_scales,
                     q_len_per_req=None,
                     max_q_len=self.forward_metadata.max_seq_len_q,
                     cum_seq_lens_q=self.forward_metadata.cu_seqlens_q,
@@ -1478,6 +1703,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                     bmm2_scale=bmm2_scale,
                     window_left=layer.sliding_window_size,
                     sinks=attention_sink,
+                    out=native_out,
+                    out_dtype=(None if uses_native_fp4 else self.q_data_type),
+                    kv_cache_sf=kv_cache_block_scales,
                     q_len_per_req=self.forward_metadata.max_seq_len_q,
                 )
         elif self.use_fmha_v2 and not cp_active:
@@ -1541,7 +1769,8 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                     sinks=attention_sink,
                     skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(),
                     out=out,
-                    out_dtype=self.q_data_type,
+                    out_dtype=(None if uses_native_fp4 else self.q_data_type),
+                    kv_cache_sf=kv_cache_block_scales,
                 )
 
             if cp_active:
@@ -1555,7 +1784,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                     attention_backend=CPAttentionBackendKind.TRTLLM_MHA,
                 )
             else:
-                out = forward_batch._attn_output
+                out = native_out if uses_native_fp4 else forward_batch._attn_output
                 if out is not None:
                     out = out.view_as(q)
                 o = _trtllm_context_attn(
@@ -1567,6 +1796,8 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                     out=out,
                 )
 
+        if uses_native_fp4:
+            o = self._finalize_nvfp4_output(o, forward_batch)
         return o.view(-1, layer.tp_q_head_num * layer.head_dim)
 
 

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

@@ -120,6 +120,47 @@ class KVCacheQuantMethodBase(ABC):
     def attention_accesses(self) -> tuple[KVCacheAttentionAccess, ...]:
         return KV_CACHE_ATTENTION_ACCESS_REGISTRY.get(self.name, ())
 
+    def configure_attention_backends(
+        self, prefill_backend: str, decode_backend: str
+    ) -> None:
+        """Select the accesses that this server will actually instantiate.
+
+        Keeping this selection on the recipe object lets memory allocation
+        omit compatibility workspaces/layouts that the chosen backend pair can
+        never read. Directly-constructed methods retain the complete registry,
+        which is useful for introspection and backwards-compatible unit tests.
+        """
+        selected = []
+        for phase, backend in (
+            (KVCacheAttentionPhase.PREFILL, prefill_backend),
+            (KVCacheAttentionPhase.DECODE, decode_backend),
+        ):
+            access = self.resolve_attention_access(phase, backend)
+            if access is not None:
+                selected.append(access)
+        self._active_attention_accesses = tuple(selected)
+
+    def configure_attention_backends_from_server_args(self, server_args) -> None:
+        """Select accesses from the resolved per-phase attention backends.
+
+        Attention backend resolution is declaration based: model hooks can
+        override a pristine ``ServerArgs`` without mutating its raw fields.
+        Always use the public resolution projection here instead of depending
+        on a private ``ServerArgs`` helper, and keep pool sizing/allocation on
+        the same selection path.
+        """
+        from sglang.srt.arg_groups.overrides import (
+            attention_backends_of,
+            resolved_view,
+        )
+
+        self.configure_attention_backends(
+            *attention_backends_of(resolved_view(server_args))
+        )
+
+    def active_attention_accesses(self) -> tuple[KVCacheAttentionAccess, ...]:
+        return getattr(self, "_active_attention_accesses", self.attention_accesses())
+
     def resolve_attention_access(
         self, phase, backend_name: str, backend_tags: Iterable[str] = ()
     ) -> Optional[KVCacheAttentionAccess]:
@@ -146,9 +187,20 @@ class KVCacheQuantMethodBase(ABC):
         """Whether the pool should allocate dq_k_buffer / dq_v_buffer."""
         return any(
             access.kind == KVCacheAttentionAccessKind.DEQUANT_WORKSPACE
-            for access in self.attention_accesses()
+            for access in self.active_attention_accesses()
         )
 
+    def has_native_fp4_access(self) -> bool:
+        """Whether a selected backend consumes native packed FP4 + scales."""
+        return any(
+            access.kind == KVCacheAttentionAccessKind.NATIVE_FP4
+            for access in self.active_attention_accesses()
+        )
+
+    def needs_native_fp4_scales(self) -> bool:
+        """Whether the pool needs a separate native FP4 scale layout."""
+        return self.has_native_fp4_access()
+
     def needs_plain_kv_dequant_read(self) -> bool:
         """Whether plain attention reads require dequantizing packed KV first."""
         return any(
@@ -160,7 +212,7 @@ class KVCacheQuantMethodBase(ABC):
     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():
+        for access in self.active_attention_accesses():
             if access.kind != KVCacheAttentionAccessKind.DEQUANT_WORKSPACE:
                 continue
             if access.workspace_dtype is None:
@@ -249,6 +301,8 @@ class KVCacheQuantMethodBase(ABC):
         cache_v: Tensor,
         k_scale=None,
         v_scale=None,
+        native_k_scale_buffer: Optional[Tensor] = None,
+        native_v_scale_buffer: Optional[Tensor] = None,
     ) -> None:
         """Quantize cache_k / cache_v and write into buffers at loc."""
 
@@ -310,6 +364,8 @@ class UnquantizedKVCacheMethod(KVCacheQuantMethodBase):
         cache_v,
         k_scale=None,
         v_scale=None,
+        native_k_scale_buffer=None,
+        native_v_scale_buffer=None,
     ) -> None:
         raise RuntimeError(
             "Unquantized KV cache writes are handled by MHATokenToKVPool.set_kv_buffer."
@@ -363,6 +419,8 @@ class CPUFP8KVCacheMethod(KVCacheQuantMethodBase):
         cache_v,
         k_scale=None,
         v_scale=None,
+        native_k_scale_buffer=None,
+        native_v_scale_buffer=None,
     ) -> None:
         k_scale = 1.0 if k_scale is None else k_scale
         v_scale = 1.0 if v_scale is None else v_scale
@@ -391,9 +449,21 @@ class NVFP4KVCacheMethod(KVCacheQuantMethodBase):
     name = "nvfp4"
     SCALE_BLOCK_SIZE = 16
 
-    def __init__(self, num_layers: int, device: str):
+    def __init__(
+        self,
+        num_layers: int,
+        device: str,
+        page_size: int = 16,
+        native_scale_layout: Optional[bool] = None,
+    ):
         self.num_layers = num_layers
         self.device = device
+        self.page_size = page_size
+        self.use_trtllm_gen_native_scale_layout = (
+            get_platform().is_sm100
+            if native_scale_layout is None
+            else native_scale_layout
+        )
         # 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)
@@ -403,6 +473,17 @@ class NVFP4KVCacheMethod(KVCacheQuantMethodBase):
     def needs_global_scale(self) -> bool:
         return True
 
+    def needs_native_fp4_scales(self) -> bool:
+        """Whether SM100 TRT-LLM GenMHA's physical HND scales are needed."""
+        return self.has_native_fp4_access() and self.use_trtllm_gen_native_scale_layout
+
+    def needs_linear_scale_buffer(self) -> bool:
+        # FlashInfer DQ prefill always needs token-linear scales. SM120 XQA also
+        # consumes the legacy linear layout instead of SM100's GenMHA layout.
+        return self.needs_dequant_workspace() or (
+            self.has_native_fp4_access() and not self.use_trtllm_gen_native_scale_layout
+        )
+
     def scale_buffer_view_dtype(self) -> Optional[torch.dtype]:
         return torch.float8_e4m3fn
 
@@ -461,15 +542,21 @@ class NVFP4KVCacheMethod(KVCacheQuantMethodBase):
                 if hasattr(layer, "v_scale") and layer.v_scale is not None
                 else 1.0
             )
-            # SM100 uses TRT-LLM XQA kernels that expect KV scales as
+            # SM100 uses TRT-LLM GenMHA kernels that expect KV scales as
             # amax / 448, but the calibrated checkpoint stores amax / (6 * 448).
             # We multiply by E2M1_MAX (6.0) to bridge the gap.  SM120 uses a
             # different kernel path where scales already include this factor.
             # 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 get_platform().is_sm100:
+            # BaseKVCacheMethod uses exactly 1.0 when the checkpoint did not
+            # provide calibrated KV scales. Keep that neutral fallback: turning
+            # it into 6.0 needlessly pushes the online block scales toward the
+            # low-precision end of E4M3. Real calibrated scales are tiny
+            # positive values and need the E2M1_MAX conversion below.
+            if get_platform().is_sm100 and k_scale != 1.0:
                 k_scale *= E2M1_MAX
+            if get_platform().is_sm100 and v_scale != 1.0:
                 v_scale *= E2M1_MAX
             k_scales_cpu[layer_id] = k_scale
             v_scales_cpu[layer_id] = v_scale
@@ -490,6 +577,24 @@ class NVFP4KVCacheMethod(KVCacheQuantMethodBase):
         k = head_dim
         store_dtype = self.kv_storage_dtype()
         dq_dtype = self.dequant_workspace_dtype()
+        needs_linear_scales = self.needs_linear_scale_buffer()
+        needs_native_scales = self.needs_native_fp4_scales()
+
+        if needs_native_scales:
+            if self.page_size % 4 != 0:
+                raise ValueError(
+                    "Native NVFP4 requires page_size divisible by 4, got "
+                    f"{self.page_size}."
+                )
+            if k % 64 != 0:
+                raise ValueError(
+                    f"Native NVFP4 requires head_dim divisible by 64, got {k}."
+                )
+            if m % self.page_size != 0:
+                raise ValueError(
+                    "NVFP4 pool rows must be page aligned, got "
+                    f"rows={m}, page_size={self.page_size}."
+                )
 
         k_buffer = [
             torch.zeros((m, n, k // 2), dtype=store_dtype, device=device)
@@ -499,18 +604,52 @@ class NVFP4KVCacheMethod(KVCacheQuantMethodBase):
             torch.zeros((m, n, k // 2), dtype=store_dtype, device=device)
             for _ in range(layer_num)
         ]
-        k_scale_buffer = [
-            torch.zeros(
-                (m, n, k // self.SCALE_BLOCK_SIZE), dtype=store_dtype, device=device
-            )
-            for _ in range(layer_num)
-        ]
-        v_scale_buffer = [
-            torch.zeros(
-                (m, n, k // self.SCALE_BLOCK_SIZE), dtype=store_dtype, device=device
-            )
-            for _ in range(layer_num)
-        ]
+        k_scale_buffer = (
+            [
+                torch.zeros(
+                    (m, n, k // self.SCALE_BLOCK_SIZE),
+                    dtype=store_dtype,
+                    device=device,
+                )
+                for _ in range(layer_num)
+            ]
+            if needs_linear_scales
+            else None
+        )
+        v_scale_buffer = (
+            [
+                torch.zeros(
+                    (m, n, k // self.SCALE_BLOCK_SIZE),
+                    dtype=store_dtype,
+                    device=device,
+                )
+                for _ in range(layer_num)
+            ]
+            if needs_linear_scales
+            else None
+        )
+        native_scale_shape = (
+            m // self.page_size,
+            n,
+            self.page_size,
+            k // self.SCALE_BLOCK_SIZE,
+        )
+        native_k_scale_buffer = (
+            [
+                torch.zeros(native_scale_shape, dtype=store_dtype, device=device)
+                for _ in range(layer_num)
+            ]
+            if needs_native_scales
+            else None
+        )
+        native_v_scale_buffer = (
+            [
+                torch.zeros(native_scale_shape, dtype=store_dtype, device=device)
+                for _ in range(layer_num)
+            ]
+            if needs_native_scales
+            else None
+        )
         # Shared dequant workspace: one copy, reused per layer during prefill.
         dq_k_buffer = (
             torch.zeros((m, n, k), dtype=dq_dtype, device=device)
@@ -528,6 +667,8 @@ class NVFP4KVCacheMethod(KVCacheQuantMethodBase):
             "v_buffer": v_buffer,
             "k_scale_buffer": k_scale_buffer,
             "v_scale_buffer": v_scale_buffer,
+            "native_k_scale_buffer": native_k_scale_buffer,
+            "native_v_scale_buffer": native_v_scale_buffer,
             "dq_k_buffer": dq_k_buffer,
             "dq_v_buffer": dq_v_buffer,
             "store_dtype": store_dtype,
@@ -544,8 +685,13 @@ class NVFP4KVCacheMethod(KVCacheQuantMethodBase):
         cache_v: Tensor,
         k_scale=None,
         v_scale=None,
+        native_k_scale_buffer: Optional[Tensor] = None,
+        native_v_scale_buffer: Optional[Tensor] = None,
     ) -> None:
         from sglang.srt.layers.quantization.kvfp4_tensor import NVFP4KVQuantizeUtil
+        from sglang.srt.layers.quantization.nvfp4_kv_cache import (
+            store_nvfp4_kv_cache,
+        )
 
         cache_k, cache_k_fp4_sf, _ = NVFP4KVQuantizeUtil.quantize(
             cache_k.contiguous(), k_scale
@@ -559,10 +705,20 @@ class NVFP4KVCacheMethod(KVCacheQuantMethodBase):
         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
+        store_nvfp4_kv_cache(
+            cache_k,
+            cache_v,
+            cache_k_fp4_sf,
+            cache_v_fp4_sf,
+            loc,
+            k_buffer,
+            v_buffer,
+            k_scale_buffer,
+            v_scale_buffer,
+            native_k_scale_buffer,
+            native_v_scale_buffer,
+            self.page_size,
+        )
 
     def dequantize_prev_kv(
         self,
@@ -590,10 +746,14 @@ class NVFP4KVCacheMethod(KVCacheQuantMethodBase):
     ) -> int:
         # FP4 data: per-layer, K+V
         fp4_size = head_num * (head_dim // 2) * num_layers * 2 * kv_size
-        # Block scales: per-layer, K+V (uint8)
-        scale_size = (
+        # Linear scales serve the FP8-prefill compatibility recipe; native HND
+        # scales serve TRT-LLM GenMHA. Mixed mode intentionally owns both.
+        one_scale_layout_size = (
             head_num * (head_dim // self.SCALE_BLOCK_SIZE) * num_layers * 2 * kv_size
         )
+        scale_size = one_scale_layout_size * (
+            int(self.needs_linear_scale_buffer()) + int(self.needs_native_fp4_scales())
+        )
         # Dequant workspace is shared across layers, not multiplied by num_layers.
         dq_dtype = self.dequant_workspace_dtype()
         dq_size = (
@@ -622,6 +782,7 @@ class FP4MXBlock16KVCacheMethod(KVCacheQuantMethodBase):
         self,
         num_layers: Optional[int] = None,
         device: Optional[str] = None,
+        page_size: Optional[int] = None,
     ):
         pass
 
@@ -689,6 +850,8 @@ class FP4MXBlock16KVCacheMethod(KVCacheQuantMethodBase):
         cache_v,
         k_scale=None,
         v_scale=None,
+        native_k_scale_buffer=None,
+        native_v_scale_buffer=None,
     ) -> None:
         from sglang.srt.layers.quantization.kvfp4_tensor import (
             FP4MXBlock16KVQuantizeUtil,
@@ -763,8 +926,9 @@ _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"})
+_NVFP4_DQ_KV_PREFILL_BACKENDS = frozenset({"flashinfer"})
+_NVFP4_KV_PREFILL_BACKENDS = frozenset({"trtllm_mha"})
+_NVFP4_KV_DECODE_BACKENDS = frozenset({"trtllm_mha"})
 _FP4_MX_MHA_BACKENDS = frozenset(
     {"triton", "torch_native", "flex_attention", "trtllm_mha"}
 )
@@ -837,8 +1001,9 @@ KV_CACHE_ATTENTION_ACCESS_REGISTRY: dict[str, tuple[KVCacheAttentionAccess, ...]
         _plain(_DECODE, _CPU_FP8_BACKENDS),
     ),
     NVFP4KVCacheMethod.name: (
-        _dq_workspace(_PREFILL, _NVFP4_PREFILL_BACKENDS, _NVFP4_SCALE, _FP8_E4M3),
-        _native_fp4(_DECODE, _NVFP4_DECODE_BACKENDS, _NVFP4_SCALE, _TORCH_FP4),
+        _dq_workspace(_PREFILL, _NVFP4_DQ_KV_PREFILL_BACKENDS, _NVFP4_SCALE, _FP8_E4M3),
+        _native_fp4(_PREFILL, _NVFP4_KV_PREFILL_BACKENDS, _NVFP4_SCALE, _TORCH_FP4),
+        _native_fp4(_DECODE, _NVFP4_KV_DECODE_BACKENDS, _NVFP4_SCALE, _TORCH_FP4),
     ),
     FP4MXBlock16KVCacheMethod.name: (
         _plain(_PREFILL, _FP4_MX_PREFILL_BACKENDS, _FP4_MX_SCALE, _BF16),

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

@@ -0,0 +1,290 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Storage helpers for TRT-LLM GenMHA's native NVFP4 KV layout.
+
+Packed K/V data stays in SGLang's slot-major NHD pool.  TRT-LLM permits the
+outer data strides to be non-contiguous, so the backend can expose an HND view
+without copying it.  Scale factors are stricter: the page/token and block-scale
+dimensions must be contiguous, and V scales use a four-token interleave.  The
+kernels below create and maintain that native scale view while preserving an
+optional linear scale view used by FlashInfer's FP8-prefill compatibility path.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+def nvfp4_v_scale_swizzle_indices(
+    token_offsets: torch.Tensor, scale_indices: torch.Tensor, scale_dim: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return TRT-LLM's four-token-interleaved V-scale coordinates.
+
+    This small torch reference is intentionally device agnostic so the layout
+    contract can be unit-tested without a GPU. ``scale_dim`` is head_dim / 16.
+
+    FlashInfer's ``nvfp4_block_scale_interleave`` is the 128x4 scale layout for
+    GEMM/MoE and is not this KV-cache layout. Its slot-mapping paged-KV append
+    writes linear V scales; its complete-cache conversion helper applies this
+    four-token permutation, but only while rewriting the whole cache. The kernel
+    below instead fuses the permutation with SGLang's incremental slot scatter
+    so it does not materialize another cache-sized scale tensor.
+    """
+    if scale_dim % 4 != 0:
+        raise ValueError(f"NVFP4 scale_dim must be divisible by 4, got {scale_dim}.")
+    scale_group = scale_dim // 4
+    swizzled_token = (token_offsets // 4) * 4 + scale_indices // scale_group
+    swizzled_scale = (scale_indices % scale_group) * 4 + token_offsets % 4
+    return swizzled_token, swizzled_scale
+
+
+@triton.jit
+def _store_nvfp4_kv_kernel(
+    k_src,
+    v_src,
+    k_scale_src,
+    v_scale_src,
+    loc,
+    k_dst,
+    v_dst,
+    k_scale_linear_dst,
+    v_scale_linear_dst,
+    k_scale_native_dst,
+    v_scale_native_dst,
+    loc_stride: tl.constexpr,
+    num_heads: tl.constexpr,
+    packed_dim: tl.constexpr,
+    scale_dim: tl.constexpr,
+    page_size: tl.constexpr,
+    BLOCK_PACKED: tl.constexpr,
+    BLOCK_SCALE: tl.constexpr,
+    STORE_LINEAR: tl.constexpr,
+    STORE_NATIVE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    head_idx = tl.program_id(1)
+    slot = tl.load(loc + token_idx * loc_stride).to(tl.int64)
+
+    packed_offsets = tl.arange(0, BLOCK_PACKED)
+    packed_mask = packed_offsets < packed_dim
+    src_packed_base = (token_idx * num_heads + head_idx) * packed_dim
+    dst_packed_base = (slot * num_heads + head_idx) * packed_dim
+    k_packed = tl.load(k_src + src_packed_base + packed_offsets, mask=packed_mask)
+    v_packed = tl.load(v_src + src_packed_base + packed_offsets, mask=packed_mask)
+    tl.store(k_dst + dst_packed_base + packed_offsets, k_packed, mask=packed_mask)
+    tl.store(v_dst + dst_packed_base + packed_offsets, v_packed, mask=packed_mask)
+
+    scale_offsets = tl.arange(0, BLOCK_SCALE)
+    scale_mask = scale_offsets < scale_dim
+    src_scale_base = (token_idx * num_heads + head_idx) * scale_dim
+    k_scale = tl.load(k_scale_src + src_scale_base + scale_offsets, mask=scale_mask)
+    v_scale = tl.load(v_scale_src + src_scale_base + scale_offsets, mask=scale_mask)
+
+    if STORE_LINEAR:
+        dst_scale_base = (slot * num_heads + head_idx) * scale_dim
+        tl.store(
+            k_scale_linear_dst + dst_scale_base + scale_offsets,
+            k_scale,
+            mask=scale_mask,
+        )
+        tl.store(
+            v_scale_linear_dst + dst_scale_base + scale_offsets,
+            v_scale,
+            mask=scale_mask,
+        )
+
+    if STORE_NATIVE:
+        page = slot // page_size
+        token_offset = slot % page_size
+        native_page_head_base = (page * num_heads + head_idx) * page_size * scale_dim
+
+        k_native_offset = (
+            native_page_head_base + token_offset * scale_dim + scale_offsets
+        )
+        tl.store(k_scale_native_dst + k_native_offset, k_scale, mask=scale_mask)
+
+        scale_group = scale_dim // 4
+        swizzled_token = (token_offset // 4) * 4 + scale_offsets // scale_group
+        swizzled_scale = (scale_offsets % scale_group) * 4 + token_offset % 4
+        v_native_offset = (
+            native_page_head_base + swizzled_token * scale_dim + swizzled_scale
+        )
+        tl.store(v_scale_native_dst + v_native_offset, v_scale, mask=scale_mask)
+
+
+def store_nvfp4_kv_cache(
+    k_src: torch.Tensor,
+    v_src: torch.Tensor,
+    k_scale_src: torch.Tensor,
+    v_scale_src: torch.Tensor,
+    loc: torch.Tensor,
+    k_dst: torch.Tensor,
+    v_dst: torch.Tensor,
+    k_scale_linear_dst: Optional[torch.Tensor],
+    v_scale_linear_dst: Optional[torch.Tensor],
+    k_scale_native_dst: Optional[torch.Tensor],
+    v_scale_native_dst: Optional[torch.Tensor],
+    page_size: int,
+) -> None:
+    """Scatter one layer of quantized K/V and both selected scale layouts."""
+    if (k_scale_linear_dst is None) != (v_scale_linear_dst is None):
+        raise ValueError("Linear NVFP4 K/V scale buffers must be provided together.")
+    if (k_scale_native_dst is None) != (v_scale_native_dst is None):
+        raise ValueError("Native NVFP4 K/V scale buffers must be provided together.")
+    store_linear = k_scale_linear_dst is not None
+    store_native = k_scale_native_dst is not None
+    if not (store_linear or store_native):
+        raise ValueError("At least one NVFP4 scale layout must be selected.")
+
+    num_tokens, num_heads, packed_dim = k_src.shape
+    scale_dim = k_scale_src.shape[-1]
+    if store_native:
+        if page_size % 4 != 0:
+            raise ValueError(
+                f"Native NVFP4 requires page_size divisible by 4, got {page_size}."
+            )
+        if scale_dim % 4 != 0:
+            raise ValueError(
+                "Native NVFP4 requires head_dim divisible by 64; "
+                f"got scale_dim={scale_dim} (head_dim={scale_dim * 16})."
+            )
+    expected_data_shape = (num_tokens, num_heads, packed_dim)
+    expected_scale_shape = (num_tokens, num_heads, scale_dim)
+    if v_src.shape != expected_data_shape:
+        raise ValueError(f"K/V packed shapes differ: {k_src.shape} vs {v_src.shape}.")
+    if (
+        k_scale_src.shape != expected_scale_shape
+        or v_scale_src.shape != expected_scale_shape
+    ):
+        raise ValueError(
+            "Unexpected NVFP4 scale shapes: "
+            f"K={k_scale_src.shape}, V={v_scale_src.shape}, expected={expected_scale_shape}."
+        )
+    if loc.numel() != num_tokens:
+        raise ValueError(f"loc has {loc.numel()} entries for {num_tokens} KV rows.")
+
+    # Compile-time-false branches do not dereference these placeholder pointers.
+    linear_k = k_scale_linear_dst if store_linear else k_dst
+    linear_v = v_scale_linear_dst if store_linear else v_dst
+    native_k = k_scale_native_dst if store_native else k_dst
+    native_v = v_scale_native_dst if store_native else v_dst
+    _store_nvfp4_kv_kernel[(num_tokens, num_heads)](
+        k_src,
+        v_src,
+        k_scale_src,
+        v_scale_src,
+        loc,
+        k_dst,
+        v_dst,
+        linear_k,
+        linear_v,
+        native_k,
+        native_v,
+        loc_stride=loc.stride(0),
+        num_heads=num_heads,
+        packed_dim=packed_dim,
+        scale_dim=scale_dim,
+        page_size=page_size,
+        BLOCK_PACKED=triton.next_power_of_2(packed_dim),
+        BLOCK_SCALE=triton.next_power_of_2(scale_dim),
+        STORE_LINEAR=store_linear,
+        STORE_NATIVE=store_native,
+        num_warps=4,
+    )
+
+
+@triton.jit
+def _move_nvfp4_native_scales_kernel(
+    k_scale,
+    v_scale,
+    tgt_loc,
+    src_loc,
+    num_heads: tl.constexpr,
+    scale_dim: tl.constexpr,
+    page_size: tl.constexpr,
+    BLOCK_SCALE: tl.constexpr,
+):
+    move_idx = tl.program_id(0)
+    head_idx = tl.program_id(1)
+    target = tl.load(tgt_loc + move_idx).to(tl.int64)
+    source = tl.load(src_loc + move_idx).to(tl.int64)
+    target_page, target_token = target // page_size, target % page_size
+    source_page, source_token = source // page_size, source % page_size
+
+    scale_offsets = tl.arange(0, BLOCK_SCALE)
+    mask = scale_offsets < scale_dim
+    source_base = (source_page * num_heads + head_idx) * page_size * scale_dim
+    target_base = (target_page * num_heads + head_idx) * page_size * scale_dim
+
+    k_values = tl.load(
+        k_scale + source_base + source_token * scale_dim + scale_offsets, mask=mask
+    )
+    tl.store(
+        k_scale + target_base + target_token * scale_dim + scale_offsets,
+        k_values,
+        mask=mask,
+    )
+
+    scale_group = scale_dim // 4
+    source_swizzled_token = (source_token // 4) * 4 + scale_offsets // scale_group
+    source_swizzled_scale = (scale_offsets % scale_group) * 4 + source_token % 4
+    target_swizzled_token = (target_token // 4) * 4 + scale_offsets // scale_group
+    target_swizzled_scale = (scale_offsets % scale_group) * 4 + target_token % 4
+    v_values = tl.load(
+        v_scale
+        + source_base
+        + source_swizzled_token * scale_dim
+        + source_swizzled_scale,
+        mask=mask,
+    )
+    tl.store(
+        v_scale
+        + target_base
+        + target_swizzled_token * scale_dim
+        + target_swizzled_scale,
+        v_values,
+        mask=mask,
+    )
+
+
+def move_nvfp4_native_scales(
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+    tgt_loc: torch.Tensor,
+    src_loc: torch.Tensor,
+) -> None:
+    """Move logical token scale rows between native HND/swizzled slots."""
+    if tgt_loc.numel() == 0:
+        return
+    if k_scale.shape != v_scale.shape or k_scale.ndim != 4:
+        raise ValueError(
+            f"Expected matching [pages, heads, page, scales] tensors, got "
+            f"{k_scale.shape} and {v_scale.shape}."
+        )
+    _, num_heads, page_size, scale_dim = k_scale.shape
+    _move_nvfp4_native_scales_kernel[(tgt_loc.numel(), num_heads)](
+        k_scale,
+        v_scale,
+        tgt_loc,
+        src_loc,
+        num_heads=num_heads,
+        scale_dim=scale_dim,
+        page_size=page_size,
+        BLOCK_SCALE=triton.next_power_of_2(scale_dim),
+        num_warps=1,
+    )

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -320,7 +320,9 @@ class KVCacheConfigurator:
             quant_name,
             num_layers=num_layers,
             device=self.device,
+            page_size=self.page_size,
         )
+        quant_method.configure_attention_backends_from_server_args(self.server_args)
         quant_method.load_scales_from_model(self.model)
         return quant_method
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -2145,6 +2145,8 @@ class MHATokenToKVPool(KVCache):
         else:
             self.k_scale_buffer = None
             self.v_scale_buffer = None
+            self.native_k_scale_buffer = None
+            self.native_v_scale_buffer = None
             self.dq_k_buffer = None
             self.dq_v_buffer = None
             if self.post_capture_active:
@@ -2173,6 +2175,8 @@ class MHATokenToKVPool(KVCache):
         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.native_k_scale_buffer = buf.get("native_k_scale_buffer")
+        self.native_v_scale_buffer = buf.get("native_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)
@@ -2182,6 +2186,24 @@ class MHATokenToKVPool(KVCache):
         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
+        has_k_native_scales = self.native_k_scale_buffer is not None
+        has_v_native_scales = self.native_v_scale_buffer is not None
+        requires_native_scales = (
+            self.quant_method.needs_native_fp4_scales()
+            if hasattr(self.quant_method, "needs_native_fp4_scales")
+            else False
+        )
+        if has_k_native_scales != has_v_native_scales:
+            raise RuntimeError(
+                f"KV cache method {self.quant_method.name!r} created only one "
+                "native FP4 scale buffer."
+            )
+        if requires_native_scales != has_k_native_scales:
+            expectation = "requires" if requires_native_scales else "does not require"
+            raise RuntimeError(
+                f"KV cache method {self.quant_method.name!r} {expectation} native "
+                f"FP4 scales, but buffer presence is {has_k_native_scales}."
+            )
         if has_k_workspace != has_v_workspace:
             raise RuntimeError(
                 f"KV cache method {self.quant_method.name!r} created only one "
@@ -2390,6 +2412,16 @@ class MHATokenToKVPool(KVCache):
             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, "native_k_scale_buffer")
+            and self.native_k_scale_buffer is not None
+        ):
+            del self.native_k_scale_buffer
+        if (
+            hasattr(self, "native_v_scale_buffer")
+            and self.native_v_scale_buffer is not None
+        ):
+            del self.native_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:
@@ -2406,6 +2438,9 @@ class MHATokenToKVPool(KVCache):
         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, "native_k_scale_buffer", None) is not None:
+            k_size_bytes += get_tensor_size_bytes(self.native_k_scale_buffer)
+            v_size_bytes += get_tensor_size_bytes(self.native_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)
@@ -2686,6 +2721,12 @@ class MHATokenToKVPool(KVCache):
         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)
+        native_scale_kwargs = {}
+        if self.native_k_scale_buffer is not None:
+            native_scale_kwargs = {
+                "native_k_scale_buffer": self.native_k_scale_buffer[local_layer_id],
+                "native_v_scale_buffer": self.native_v_scale_buffer[local_layer_id],
+            }
         self.quant_method.quantize_and_store(
             self.k_buffer[local_layer_id],
             self.v_buffer[local_layer_id],
@@ -2704,6 +2745,7 @@ class MHATokenToKVPool(KVCache):
             cache_v,
             k_scale,
             v_scale,
+            **native_scale_kwargs,
         )
 
     def get_raw_kv_buffer(
@@ -3022,6 +3064,7 @@ class MHATokenToKVPool(KVCache):
             for kb, vb in zip(self.k_buffer, self.v_buffer):
                 kb[pages_t, :, offs_t, :] = kb[pages_s, :, offs_s, :]
                 vb[pages_t, :, offs_t, :] = vb[pages_s, :, offs_s, :]
+            self._move_native_fp4_scales(tgt_loc, src_loc)
             return
 
         self._move_kv_cache_impl(tgt_loc, src_loc)
@@ -3035,6 +3078,7 @@ class MHATokenToKVPool(KVCache):
                 move_kv_cache_native(
                     self.k_scale_buffer, self.v_scale_buffer, tgt_loc, src_loc
                 )
+            self._move_native_fp4_scales(tgt_loc, src_loc)
             return
 
         N = tgt_loc.numel()
@@ -3058,6 +3102,7 @@ class MHATokenToKVPool(KVCache):
                 next_power_of_2(N),
                 cfg,
             )
+            self._move_native_fp4_scales(tgt_loc, src_loc)
             return
 
         # Huge N: chunk, but each chunk's upper is still pow2(<= cap)
@@ -3073,6 +3118,21 @@ class MHATokenToKVPool(KVCache):
                 next_power_of_2(chunk_len),
                 cfg,
             )
+        self._move_native_fp4_scales(tgt_loc, src_loc)
+
+    def _move_native_fp4_scales(
+        self, tgt_loc: torch.Tensor, src_loc: torch.Tensor
+    ) -> None:
+        if self.native_k_scale_buffer is None:
+            return
+        from sglang.srt.layers.quantization.nvfp4_kv_cache import (
+            move_nvfp4_native_scales,
+        )
+
+        for k_scale, v_scale in zip(
+            self.native_k_scale_buffer, self.native_v_scale_buffer
+        ):
+            move_nvfp4_native_scales(k_scale, v_scale, tgt_loc, src_loc)
 
 
 class NoOpMHATokenToKVPool(MHATokenToKVPool):

python/sglang/srt/model_executor/pool_configurator.py

@@ -410,14 +410,31 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
             )
 
             if is_float4_e2m1fn_x2(kv_cache_dtype):
-                # kv_scale_buffer
-                scale_block_size = 16
-                k = model_config.head_dim
-                cell_size = (cell_size // 2) + (
-                    (n * k * effective_num_layers * 2 * kv_size) // scale_block_size
+                from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
+                    get_kv_cache_quant_method,
+                    resolve_kv_cache_quant,
+                )
+
+                quant_name = resolve_kv_cache_quant(kvc.kv_cache_dtype_str)
+                if quant_name is None:
+                    raise ValueError(
+                        "FP4 storage dtype requires an explicit KV recipe name."
+                    )
+                quant_method = get_kv_cache_quant_method(
+                    quant_name,
+                    num_layers=effective_num_layers,
+                    device=kvc.device,
+                    page_size=kvc.page_size,
+                )
+                quant_method.configure_attention_backends_from_server_args(
+                    kvc.server_args
+                )
+                cell_size = quant_method.compute_cell_size(
+                    n,
+                    model_config.head_dim,
+                    effective_num_layers,
+                    kv_size,
                 )
-                # FP4 prefill uses one shared FP8 dequant workspace across layers.
-                cell_size += n * k * 2 * kv_size
             elif self.kv_cache_dtype_str == "mxfp8":
                 scale_block_size = 32
                 cell_size += (