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

This commit is contained in:
Sam (Kesen Li)
2026-09-18 14:50:46 -07:00
committed by GitHub
parent f5a1434700
commit d346b214fb
19 changed files with 2106 additions and 112 deletions
@@ -30,9 +30,10 @@ SGLang supports the following quantized KV cache formats:
FP4 quantization is currently experimental.
</Warning>
[OCP (Open Compute Project)](https://www.opencompute.org) specifies MXFP4 (Microscaling FP4), a 4-bit floating-point format:
[OCP (Open Compute Project)](https://www.opencompute.org) specifies MXFP4 (Microscaling FP4), a 4-bit floating-point format. SGLang exposes two experimental E2M1 KV-cache recipes:
- **E2M1** (1 sign bit, 2 exponent bits, 1 mantissa bit): Uses block-based microscaling where tensors are divided into blocks of consecutive elements, with each block sharing a single 8-bit exponential scaling factor. While OCP specifies blocks of 32 elements, SGLang's current implementation uses blocks of 16 elements for KV cache quantization.
- **`nvfp4`**: NVIDIA FP4 with 16-value blocks, E4M3 block scales, and a per-tensor global scale.
- **`fp4_mx_block16`**: An E2M1 block-size-16 compatibility recipe. It is distinct from the standard OCP MXFP4 block-size-32 format.
## Usage
@@ -62,6 +63,38 @@ python3 -m sglang.launch_server \
--kv-cache-dtype fp4_mx_block16 \
```
### SM100 native NVFP4 recipes
On SM100, prefill can either consume packed NVFP4 directly or dequantize it into an FP8 E4M3 workspace. Select the online dequantization dtype with `--prefill-kv-cache-dequant-dtype`; SGLang chooses the corresponding attention implementation.
For native NVFP4 prefill and decode:
```bash Command
python3 -m sglang.launch_server \
--model-path Qwen/Qwen3.5-35B-A3B-FP8 \
--tp-size 1 \
--kv-cache-dtype nvfp4 \
--prefill-kv-cache-dequant-dtype nvfp4 \
--page-size 16
```
For FP8 prefill backed by a temporary dequantization workspace, with native NVFP4 decode:
```bash Command
python3 -m sglang.launch_server \
--model-path Qwen/Qwen3.5-35B-A3B-FP8 \
--tp-size 1 \
--kv-cache-dtype nvfp4 \
--prefill-kv-cache-dequant-dtype fp8_e4m3 \
--page-size 16
```
`nvfp4` means that prefill consumes the packed native NVFP4 cache directly without any additional dequantization. The default value, `auto`, selects this native NVFP4 mode on SM100 and FP8 E4M3 dequantization on SM120. Decode consumes native NVFP4 in both recipes. The native recipe avoids the FP8 workspace and its token-linear scale copy; the FP8 recipe retains both the linear scales used during prefill and the physical scale layout used during decode. KV data remains stored as packed FP4 either way.
<Note>
Native NVFP4 prefill requires SM100, a page size divisible by 4, and an attention head dimension divisible by 64. TRT-LLM GenMHA uses FP8 query and output buffers internally; SGLang converts the result back to the model activation dtype. Top-k-1 EAGLE/EAGLE3/NEXTN and breadth-1 NGRAM speculative decoding are supported: an EAGLE-family draft worker uses `trtllm_mha`, and target verification consumes the physical NVFP4 cache directly in eager execution and CUDA Graphs. Set `--speculative-ngram-max-bfs-breadth=1` for NGRAM. With the mixed FlashInfer-prefill/TRT-LLM-decode recipe, SGLang resolves `--speculative-attention-mode` to `decode` so verification does not depend on FlashInfer's transient dequantization workspace. Other speculative algorithms, PD disaggregation, hierarchical KV cache, and LMCache are not currently supported by the SM100 native NVFP4 path. SM120 XQA continues to use its architecture-specific linear scale layout and BF16 query/output path.
</Note>
### Scaling Factors
FP8 quantization requires scaling factors to properly quantize and dequantize the KV cache.
@@ -114,6 +147,8 @@ FP4 and FP8 quantization require additional memory for block-based scaling facto
This enables longer context lengths or more concurrent requests within the same memory budget.
For native `nvfp4`, each logical scalar costs 0.5 bytes of packed FP4 data plus 1/16 byte of block-scale storage, compared with 2 bytes for BF16. The resulting theoretical KV-token capacity multiplier is `2 / 0.5625 = 3.56×`. The mixed SM100 recipe owns a second scale layout and a shared one-layer FP8 prefill workspace, so its exact capacity depends on the number of full-attention layers. SGLang includes those auxiliary buffers in its KV-pool sizing calculation.
### Accuracy Impact
#### FP8 Accuracy
@@ -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(
+216 -7
View File
@@ -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
+2
View File
@@ -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.
@@ -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",
@@ -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, (
@@ -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
@@ -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)
@@ -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),
@@ -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,
)
@@ -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
@@ -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):
@@ -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 += (
@@ -29,6 +29,45 @@ DEVICE = "cuda"
PAGE_SIZE = 128
@pytest.mark.parametrize(
"max_running_requests,max_draft_tokens,max_cuda_graph_bs,expected",
[
(32, None, None, 32),
(32, 0, 16, 32),
(32, 4, 64, 256),
(7, 16, 4, 112),
],
)
def test_native_nvfp4_output_capacity_includes_verify_width(
max_running_requests, max_draft_tokens, max_cuda_graph_bs, expected
):
assert (
trtllm_mha_backend._native_fp4_decode_output_capacity(
max_running_requests, max_draft_tokens, max_cuda_graph_bs
)
== expected
)
@pytest.mark.parametrize(
"max_context_len,max_prefill_tokens,chunked_prefill_limit,expected",
[
(4096, 8192, 0, 8192),
(8192, 4096, 0, 8192),
(8192, 16384, 2048, 2048),
],
)
def test_native_nvfp4_output_capacity_includes_unchunked_batch(
max_context_len, max_prefill_tokens, chunked_prefill_limit, expected
):
assert (
trtllm_mha_backend._native_fp4_prefill_output_capacity(
max_context_len, max_prefill_tokens, chunked_prefill_limit
)
== expected
)
def _make_backend_for_hook_test(speculative_num_draft_tokens=None):
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
@@ -60,6 +99,30 @@ def _make_backend_for_hook_test(speculative_num_draft_tokens=None):
return backend
@pytest.mark.parametrize(
"uses_genmha,prefill_native,decode_native,forward_mode,expected",
[
(True, True, True, ForwardMode.EXTEND, True),
(True, True, True, ForwardMode.TARGET_VERIFY, True),
# Hybrid mode=decode routes target verification into the decode child.
(True, False, True, ForwardMode.TARGET_VERIFY, True),
(True, False, True, ForwardMode.EXTEND, False),
# SM120 XQA has native access metadata but not the physical GenMHA layout.
(False, False, True, ForwardMode.TARGET_VERIFY, False),
],
)
def test_extend_selects_native_nvfp4_layout_per_call(
uses_genmha, prefill_native, decode_native, forward_mode, expected
):
backend = TRTLLMHAAttnBackend.__new__(TRTLLMHAAttnBackend)
backend.uses_trtllm_gen_native_fp4 = uses_genmha
backend.prefill_uses_native_fp4 = prefill_native
backend.decode_uses_native_fp4 = decode_native
forward_batch = SimpleNamespace(forward_mode=forward_mode)
assert backend._forward_extend_uses_native_fp4(forward_batch) is expected
def test_cuda_graph_metadata_launch_runs_in_graph_hook(monkeypatch):
calls = []
@@ -0,0 +1,102 @@
import unittest
import torch
from sglang.srt.utils.common import is_sm100_supported
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
# The SM100 CI pool has no single-GPU runner. Use all four GPUs in the B200
# runner so the declared resource and the server topology stay aligned.
register_cuda_ci(
est_time=900,
stage="extra-b",
runner_config="4-gpu-b200",
)
MODEL = "Qwen/Qwen3.5-9B"
GSM8K_ACCURACY_THRESHOLD = 0.70
GSM8K_NUM_QUESTIONS = 200
GSM8K_NUM_SHOTS = 8
COMMON_ARGS = [
"--tp-size",
"4",
"--kv-cache-dtype",
"nvfp4",
"--page-size",
"16",
"--max-total-tokens",
"131072",
"--max-running-requests",
"64",
]
MTP_ARGS = [
"--speculative-algorithm",
"NEXTN",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
]
HAS_FOUR_SM100_GPUS = is_sm100_supported() and torch.cuda.device_count() >= 4
@unittest.skipUnless(HAS_FOUR_SM100_GPUS, "requires 4 SM100 GPUs with CUDA 12.8+")
class TestQwen35NVFP4KVNativePrefillSM100(GSM8KMixin, DefaultServerBase):
"""Native NVFP4 prefill and decode without MTP."""
model = MODEL
gsm8k_score_threshold = GSM8K_ACCURACY_THRESHOLD
gsm8k_num_examples = GSM8K_NUM_QUESTIONS
gsm8k_num_threads = 128
gsm8k_num_shots = GSM8K_NUM_SHOTS
other_args = COMMON_ARGS + ["--prefill-kv-cache-dequant-dtype", "nvfp4"]
@unittest.skipUnless(HAS_FOUR_SM100_GPUS, "requires 4 SM100 GPUs with CUDA 12.8+")
class TestQwen35NVFP4KVDQPrefillSM100(GSM8KMixin, DefaultServerBase):
"""FP8-dequantized prefill and native NVFP4 decode without MTP."""
model = MODEL
gsm8k_score_threshold = GSM8K_ACCURACY_THRESHOLD
gsm8k_num_examples = GSM8K_NUM_QUESTIONS
gsm8k_num_threads = 128
gsm8k_num_shots = GSM8K_NUM_SHOTS
other_args = COMMON_ARGS + ["--prefill-kv-cache-dequant-dtype", "fp8_e4m3"]
@unittest.skipUnless(HAS_FOUR_SM100_GPUS, "requires 4 SM100 GPUs with CUDA 12.8+")
class TestQwen35NVFP4KVNativePrefillMTPSM100(GSM8KMixin, DefaultServerBase):
"""Native NVFP4 prefill and decode with NEXTN MTP."""
model = MODEL
gsm8k_score_threshold = GSM8K_ACCURACY_THRESHOLD
gsm8k_num_examples = GSM8K_NUM_QUESTIONS
gsm8k_num_threads = 128
gsm8k_num_shots = GSM8K_NUM_SHOTS
gsm8k_accept_length_thres = 1.2
other_args = COMMON_ARGS + ["--prefill-kv-cache-dequant-dtype", "nvfp4"] + MTP_ARGS
@unittest.skipUnless(HAS_FOUR_SM100_GPUS, "requires 4 SM100 GPUs with CUDA 12.8+")
class TestQwen35NVFP4KVDQPrefillMTPSM100(GSM8KMixin, DefaultServerBase):
"""FP8-dequantized prefill and native NVFP4 decode with NEXTN MTP."""
model = MODEL
gsm8k_score_threshold = GSM8K_ACCURACY_THRESHOLD
gsm8k_num_examples = GSM8K_NUM_QUESTIONS
gsm8k_num_threads = 128
gsm8k_num_shots = GSM8K_NUM_SHOTS
gsm8k_accept_length_thres = 1.2
other_args = (
COMMON_ARGS + ["--prefill-kv-cache-dequant-dtype", "fp8_e4m3"] + MTP_ARGS
)
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,358 @@
"""SM100 parity tests for SGLang's TRT-LLM-native NVFP4 KV layout."""
import math
import sys
import pytest
import torch
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
NVFP4KVCacheMethod,
)
from sglang.srt.utils import is_sm100_supported
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(
est_time=30,
stage="base-b-kernel-unit",
runner_config="4-gpu-b200",
)
pytestmark = pytest.mark.skipif(
not is_sm100_supported(), reason="TRT-LLM native NVFP4 layout requires SM100"
)
@torch.inference_mode()
def test_nvfp4_native_layout_matches_flashinfer_reference():
from flashinfer.fp4_quantization import nvfp4_quantize_paged_kv_cache
torch.manual_seed(7)
pages, heads, page_size, head_dim = 4, 4, 16, 128
total_tokens = pages * page_size
k_global_scale = torch.tensor([0.025], dtype=torch.float32, device="cuda")
v_global_scale = torch.tensor([0.03125], dtype=torch.float32, device="cuda")
k_nhd = torch.randn(
pages, page_size, heads, head_dim, dtype=torch.bfloat16, device="cuda"
)
v_nhd = torch.randn_like(k_nhd)
k_hnd = k_nhd.permute(0, 2, 1, 3).contiguous()
v_hnd = v_nhd.permute(0, 2, 1, 3).contiguous()
method = NVFP4KVCacheMethod(num_layers=1, device="cuda", page_size=page_size)
method.configure_attention_backends("trtllm_mha", "trtllm_mha")
buffers = method.create_buffers(
total_tokens, heads, head_dim, layer_num=1, device="cuda"
)
# Match multi-step EAGLE's per-step view while exercising page boundaries
# and every token mod-4 position used by TRT-LLM's V-scale interleave.
loc_storage = torch.empty((total_tokens, 3), dtype=torch.int64, device="cuda")
loc_storage[:, 0] = torch.randperm(total_tokens, device="cuda")
loc = loc_storage[:, 0]
assert loc.stride() == (3,)
method.quantize_and_store(
buffers["k_buffer"][0],
buffers["v_buffer"][0],
buffers["k_scale_buffer"],
buffers["v_scale_buffer"],
loc,
k_nhd.reshape(total_tokens, heads, head_dim)[loc],
v_nhd.reshape(total_tokens, heads, head_dim)[loc],
k_scale=k_global_scale,
v_scale=v_global_scale,
native_k_scale_buffer=buffers["native_k_scale_buffer"][0],
native_v_scale_buffer=buffers["native_v_scale_buffer"][0],
)
torch.cuda.synchronize()
(ref_k, ref_v), (ref_ks, ref_vs), _, _ = nvfp4_quantize_paged_kv_cache(
k_hnd,
v_hnd,
kv_layout="HND",
k_global_sf=1.0 / k_global_scale,
v_global_sf=1.0 / v_global_scale,
)
got_k = (
buffers["k_buffer"][0]
.view(pages, page_size, heads, head_dim // 2)
.permute(0, 2, 1, 3)
)
got_v = (
buffers["v_buffer"][0]
.view(pages, page_size, heads, head_dim // 2)
.permute(0, 2, 1, 3)
)
got_ks = buffers["native_k_scale_buffer"][0].view(torch.float8_e4m3fn)
got_vs = buffers["native_v_scale_buffer"][0].view(torch.float8_e4m3fn)
torch.testing.assert_close(got_k, ref_k, rtol=0, atol=0)
torch.testing.assert_close(got_v, ref_v, rtol=0, atol=0)
torch.testing.assert_close(got_ks.float(), ref_ks.float(), rtol=0, atol=0)
torch.testing.assert_close(got_vs.float(), ref_vs.float(), rtol=0, atol=0)
@pytest.mark.parametrize(
"total_tokens,max_kv_len,page_table_width",
[
(64, 64, 1),
# Mirror the Qwen3.5 server's short-prompt launch: the active sequence
# occupies only part of one page, while the kernel receives the model's
# full context limit and a correspondingly wide page-table stride.
(26, 262144, 4096),
],
)
@torch.inference_mode()
def test_nvfp4_native_prefill_attention_matches_bf16_reference(
total_tokens: int, max_kv_len: int, page_table_width: int
):
"""Exercise SGLang's writer and FlashInfer's context kernel together."""
import flashinfer
torch.manual_seed(11)
pages, page_size = 1, 64
q_heads, kv_heads, head_dim = 16, 2, 256
global_scale = torch.ones(1, dtype=torch.float32, device="cuda")
q = torch.randn(
total_tokens, q_heads, head_dim, dtype=torch.bfloat16, device="cuda"
)
k = torch.randn(
total_tokens, kv_heads, head_dim, dtype=torch.bfloat16, device="cuda"
)
v = torch.randn_like(k)
method = NVFP4KVCacheMethod(num_layers=1, device="cuda", page_size=page_size)
method.configure_attention_backends("trtllm_mha", "trtllm_mha")
buffers = method.create_buffers(
pages * page_size, kv_heads, head_dim, layer_num=1, device="cuda"
)
method.quantize_and_store(
buffers["k_buffer"][0],
buffers["v_buffer"][0],
None,
None,
torch.arange(total_tokens, device="cuda"),
k,
v,
k_scale=global_scale,
v_scale=global_scale,
native_k_scale_buffer=buffers["native_k_scale_buffer"][0],
native_v_scale_buffer=buffers["native_v_scale_buffer"][0],
)
k_cache = (
buffers["k_buffer"][0]
.view(pages, page_size, kv_heads, head_dim // 2)
.permute(0, 2, 1, 3)
)
v_cache = (
buffers["v_buffer"][0]
.view(pages, page_size, kv_heads, head_dim // 2)
.permute(0, 2, 1, 3)
)
block_scales = (
buffers["native_k_scale_buffer"][0].view(torch.float8_e4m3fn),
buffers["native_v_scale_buffer"][0].view(torch.float8_e4m3fn),
)
q_fp8 = q.to(torch.float8_e4m3fn)
out = torch.empty_like(q_fp8)
flashinfer.prefill.trtllm_batch_context_with_kv_cache(
query=q_fp8,
kv_cache=(k_cache, v_cache),
workspace_buffer=torch.zeros(
256 * 1024 * 1024, dtype=torch.uint8, device="cuda"
),
block_tables=torch.zeros(
(1, page_table_width), dtype=torch.int32, device="cuda"
),
seq_lens=torch.tensor([total_tokens], dtype=torch.int32, device="cuda"),
max_q_len=total_tokens,
max_kv_len=max_kv_len,
bmm1_scale=1.0 / math.sqrt(head_dim),
bmm2_scale=1.0,
batch_size=1,
cum_seq_lens_q=torch.tensor(
[0, total_tokens], dtype=torch.int32, device="cuda"
),
cum_seq_lens_kv=torch.tensor(
[0, total_tokens], dtype=torch.int32, device="cuda"
),
out=out,
kv_cache_sf=block_scales,
causal=True,
)
# Compare against the same FP8 query and BF16 K/V before KV quantization.
# The threshold mirrors FlashInfer's native NVFP4 attention regression.
repeat = q_heads // kv_heads
reference = (
torch.nn.functional.scaled_dot_product_attention(
q_fp8.bfloat16().transpose(0, 1).unsqueeze(0),
k.repeat_interleave(repeat, dim=1).transpose(0, 1).unsqueeze(0),
v.repeat_interleave(repeat, dim=1).transpose(0, 1).unsqueeze(0),
is_causal=True,
)
.squeeze(0)
.transpose(0, 1)
)
cosine = torch.nn.functional.cosine_similarity(
out.float().reshape(-1), reference.float().reshape(-1), dim=0
)
assert cosine.item() > 0.86, f"native NVFP4 prefill cosine={cosine.item():.4f}"
@torch.inference_mode()
def test_nvfp4_native_target_verify_matches_bf16_reference():
"""Exercise the multi-query-token GenMHA path used by TARGET_VERIFY."""
import flashinfer
torch.manual_seed(17)
page_size, prefix, verify_len = 32, 40, 4
seq_len = prefix + verify_len
pages = math.ceil(seq_len / page_size)
q_heads, kv_heads, head_dim = 8, 2, 128
global_scale = torch.ones(1, dtype=torch.float32, device="cuda")
q = torch.randn(verify_len, q_heads, head_dim, dtype=torch.bfloat16, device="cuda")
k = torch.randn(seq_len, kv_heads, head_dim, dtype=torch.bfloat16, device="cuda")
v = torch.randn_like(k)
method = NVFP4KVCacheMethod(num_layers=1, device="cuda", page_size=page_size)
method.configure_attention_backends("trtllm_mha", "trtllm_mha")
buffers = method.create_buffers(
pages * page_size, kv_heads, head_dim, layer_num=1, device="cuda"
)
method.quantize_and_store(
buffers["k_buffer"][0],
buffers["v_buffer"][0],
None,
None,
torch.arange(seq_len, device="cuda"),
k,
v,
k_scale=global_scale,
v_scale=global_scale,
native_k_scale_buffer=buffers["native_k_scale_buffer"][0],
native_v_scale_buffer=buffers["native_v_scale_buffer"][0],
)
kv_cache = (
buffers["k_buffer"][0]
.view(pages, page_size, kv_heads, head_dim // 2)
.permute(0, 2, 1, 3),
buffers["v_buffer"][0]
.view(pages, page_size, kv_heads, head_dim // 2)
.permute(0, 2, 1, 3),
)
block_scales = (
buffers["native_k_scale_buffer"][0].view(torch.float8_e4m3fn),
buffers["native_v_scale_buffer"][0].view(torch.float8_e4m3fn),
)
q_fp8 = q.to(torch.float8_e4m3fn)
out = torch.empty_like(q_fp8)
flashinfer.decode.trtllm_batch_decode_with_kv_cache(
query=q_fp8,
kv_cache=kv_cache,
workspace_buffer=torch.zeros(
256 * 1024 * 1024, dtype=torch.uint8, device="cuda"
),
block_tables=torch.arange(pages, dtype=torch.int32, device="cuda").view(
1, pages
),
seq_lens=torch.tensor([seq_len], dtype=torch.int32, device="cuda"),
max_seq_len=seq_len,
bmm1_scale=1.0 / math.sqrt(head_dim),
bmm2_scale=1.0,
out=out,
kv_cache_sf=block_scales,
q_len_per_req=verify_len,
)
repeat = q_heads // kv_heads
k_ref = k.repeat_interleave(repeat, dim=1).permute(1, 0, 2).float()
v_ref = v.repeat_interleave(repeat, dim=1).permute(1, 0, 2).float()
q_ref = q_fp8.bfloat16().permute(1, 0, 2).float()
scores = torch.einsum("hqd,hkd->hqk", q_ref, k_ref) / math.sqrt(head_dim)
key_positions = torch.arange(seq_len, device="cuda").view(1, 1, -1)
query_positions = (prefix + torch.arange(verify_len, device="cuda")).view(1, -1, 1)
scores.masked_fill_(key_positions > query_positions, float("-inf"))
reference = torch.einsum(
"hqk,hkd->hqd", torch.softmax(scores, dim=-1), v_ref
).permute(1, 0, 2)
cosine = torch.nn.functional.cosine_similarity(
out.float().reshape(-1), reference.reshape(-1), dim=0
)
assert cosine.item() > 0.86, (
f"native NVFP4 target-verify cosine={cosine.item():.4f}"
)
@torch.inference_mode()
def test_nvfp4_native_scale_move_preserves_logical_rows():
from sglang.srt.layers.quantization.nvfp4_kv_cache import (
move_nvfp4_native_scales,
nvfp4_v_scale_swizzle_indices,
)
pages, heads, page_size, scale_dim = 3, 2, 16, 8
k_scale = (
torch.arange(
pages * heads * page_size * scale_dim,
dtype=torch.int64,
device="cuda",
)
.remainder(251)
.to(torch.uint8)
.view(pages, heads, page_size, scale_dim)
)
v_scale = torch.zeros_like(k_scale)
# Seed V through the inverse logical mapping so its token rows have a clear
# identity even though physical storage is interleaved.
logical_v = (
torch.arange(
pages * page_size * heads * scale_dim,
dtype=torch.int64,
device="cuda",
)
.remainder(251)
.to(torch.uint8)
.view(pages * page_size, heads, scale_dim)
)
tokens = torch.arange(page_size, device="cuda")[:, None]
scales = torch.arange(scale_dim, device="cuda")[None, :]
sw_t, sw_s = nvfp4_v_scale_swizzle_indices(tokens, scales, scale_dim)
for page in range(pages):
for head in range(heads):
v_scale[page, head, sw_t, sw_s] = logical_v[
page * page_size : (page + 1) * page_size, head
]
src = torch.tensor([1, 15, 16, 35], dtype=torch.int64, device="cuda")
tgt = torch.tensor([46, 32, 31, 4], dtype=torch.int64, device="cuda")
expected_k = k_scale.clone()
expected_v = logical_v.clone()
expected_k[tgt // page_size, :, tgt % page_size, :] = expected_k[
src // page_size, :, src % page_size, :
]
expected_v[tgt] = expected_v[src]
move_nvfp4_native_scales(k_scale, v_scale, tgt, src)
torch.cuda.synchronize()
torch.testing.assert_close(k_scale, expected_k, rtol=0, atol=0)
got_v = torch.empty_like(logical_v)
for page in range(pages):
for head in range(heads):
got_v[page * page_size : (page + 1) * page_size, head] = v_scale[
page, head, sw_t, sw_s
]
torch.testing.assert_close(got_v, expected_v, rtol=0, atol=0)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
@@ -5,6 +5,8 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=35, suite="base-a-test-cpu")
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
@@ -70,8 +72,6 @@ class TestKVCacheQuantRegistry(CustomTestCase):
resolve_kv_cache_quant("fp4_e2m1")
def test_model_runner_rejects_legacy_fp4_alias(self):
from types import SimpleNamespace
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_context
@@ -125,6 +125,8 @@ class TestCPUFP8KVCacheMethod(CustomTestCase):
cache_v,
k_scale=0.5,
v_scale=0.25,
native_k_scale_buffer=None,
native_v_scale_buffer=None,
)
torch.testing.assert_close(
@@ -170,18 +172,32 @@ class TestNVFP4KVCacheMethod(CustomTestCase):
NVFP4KVCacheMethod,
)
m = NVFP4KVCacheMethod(num_layers=4, device="cpu")
m = NVFP4KVCacheMethod(num_layers=4, device="cpu", native_scale_layout=True)
self.assertEqual(m.name, "nvfp4")
self.assertEqual(m.SCALE_BLOCK_SIZE, 16)
self.assertTrue(m.needs_dequant_workspace())
self.assertTrue(m.needs_native_fp4_scales())
self.assertTrue(m.needs_global_scale())
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
KVCacheAttentionAccessKind,
)
self.assertEqual(
m.resolve_attention_access("prefill", "trtllm_mha").kind,
KVCacheAttentionAccessKind.NATIVE_FP4,
)
self.assertEqual(
m.resolve_attention_access("prefill", "flashinfer").kind,
KVCacheAttentionAccessKind.DEQUANT_WORKSPACE,
)
def test_create_buffers_shapes(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
NVFP4KVCacheMethod,
)
m = NVFP4KVCacheMethod(num_layers=4, device="cpu")
m = NVFP4KVCacheMethod(num_layers=4, device="cpu", native_scale_layout=True)
size, heads, dim, layers = 64, 8, 128, 4
bufs = m.create_buffers(size, heads, dim, layers, "cpu")
@@ -189,11 +205,17 @@ class TestNVFP4KVCacheMethod(CustomTestCase):
self.assertEqual(len(bufs["v_buffer"]), layers)
self.assertEqual(len(bufs["k_scale_buffer"]), layers)
self.assertEqual(len(bufs["v_scale_buffer"]), layers)
self.assertEqual(len(bufs["native_k_scale_buffer"]), layers)
self.assertEqual(len(bufs["native_v_scale_buffer"]), layers)
# FP4 packed: (size, heads, dim//2)
self.assertEqual(bufs["k_buffer"][0].shape, (size, heads, dim // 2))
# Block scales: (size, heads, dim//16)
self.assertEqual(bufs["k_scale_buffer"][0].shape, (size, heads, dim // 16))
self.assertEqual(
bufs["native_k_scale_buffer"][0].shape,
(size // 16, heads, 16, dim // 16),
)
# Dequant workspace: (size, heads, dim), FP8
self.assertEqual(bufs["dq_k_buffer"].shape, (size, heads, dim))
self.assertEqual(bufs["dq_k_buffer"].dtype, torch.float8_e4m3fn)
@@ -204,10 +226,134 @@ class TestNVFP4KVCacheMethod(CustomTestCase):
NVFP4KVCacheMethod,
)
m = NVFP4KVCacheMethod(num_layers=4, device="cpu")
m = NVFP4KVCacheMethod(num_layers=4, device="cpu", native_scale_layout=True)
cell = m.compute_cell_size(head_num=8, head_dim=128, num_layers=4, kv_size=1)
# FP4: 8*64*4*2 = 4096, scales: 8*8*4*2 = 512, dq: 8*128*2 = 2048
self.assertEqual(cell, 4096 + 512 + 2048)
# FP4: 4096, linear scales: 512, native scales: 512, shared DQ: 2048.
self.assertEqual(cell, 4096 + 512 + 512 + 2048)
def test_active_prefill_recipe_controls_auxiliary_memory(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
NVFP4KVCacheMethod,
)
size, heads, dim, layers = 64, 8, 128, 4
native = NVFP4KVCacheMethod(
num_layers=layers,
device="cpu",
page_size=16,
native_scale_layout=True,
)
native.configure_attention_backends("trtllm_mha", "trtllm_mha")
native_bufs = native.create_buffers(size, heads, dim, layers, "cpu")
self.assertIsNone(native_bufs["k_scale_buffer"])
self.assertIsNone(native_bufs["v_scale_buffer"])
self.assertIsNone(native_bufs["dq_k_buffer"])
self.assertIsNotNone(native_bufs["native_k_scale_buffer"])
self.assertEqual(native.compute_cell_size(heads, dim, layers, 1), 4096 + 512)
mixed = NVFP4KVCacheMethod(
num_layers=layers,
device="cpu",
page_size=16,
native_scale_layout=True,
)
mixed.configure_attention_backends("flashinfer", "trtllm_mha")
mixed_bufs = mixed.create_buffers(size, heads, dim, layers, "cpu")
self.assertIsNotNone(mixed_bufs["k_scale_buffer"])
self.assertIsNotNone(mixed_bufs["dq_k_buffer"])
self.assertIsNotNone(mixed_bufs["native_k_scale_buffer"])
self.assertEqual(
mixed.compute_cell_size(heads, dim, layers, 1),
4096 + 512 + 512 + 2048,
)
def test_server_args_backend_selection_uses_resolution_projection(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
KVCacheAttentionAccessKind,
NVFP4KVCacheMethod,
)
# Model/backend hooks declare overrides without mutating the raw
# ServerArgs fields. Pool sizing and allocation must observe the same
# resolved pair, and must not depend on a private ServerArgs method.
server_args = SimpleNamespace(
attention_backend="triton",
prefill_attention_backend=None,
decode_attention_backend=None,
_resolved_overrides=[
(
"test_model_override",
{
"prefill_attention_backend": "flashinfer",
"decode_attention_backend": "trtllm_mha",
},
)
],
)
method = NVFP4KVCacheMethod(
num_layers=1,
device="cpu",
page_size=16,
native_scale_layout=True,
)
method.configure_attention_backends_from_server_args(server_args)
accesses = method.active_attention_accesses()
self.assertEqual(
[access.kind for access in accesses],
[
KVCacheAttentionAccessKind.DEQUANT_WORKSPACE,
KVCacheAttentionAccessKind.NATIVE_FP4,
],
)
def test_xqa_recipe_retains_linear_scales(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
NVFP4KVCacheMethod,
)
size, heads, dim, layers = 64, 8, 128, 4
xqa = NVFP4KVCacheMethod(
num_layers=layers,
device="cpu",
page_size=16,
native_scale_layout=False,
)
xqa.configure_attention_backends("flashinfer", "trtllm_mha")
bufs = xqa.create_buffers(size, heads, dim, layers, "cpu")
self.assertIsNotNone(bufs["k_scale_buffer"])
self.assertIsNotNone(bufs["dq_k_buffer"])
self.assertIsNone(bufs["native_k_scale_buffer"])
self.assertFalse(xqa.needs_native_fp4_scales())
self.assertEqual(
xqa.compute_cell_size(heads, dim, layers, 1),
4096 + 512 + 2048,
)
def test_native_v_scale_swizzle_reference(self):
from sglang.srt.layers.quantization.nvfp4_kv_cache import (
nvfp4_v_scale_swizzle_indices,
)
token = torch.arange(16)[:, None]
scale = torch.arange(8)[None, :]
swizzled_token, swizzled_scale = nvfp4_v_scale_swizzle_indices(
token, scale, scale_dim=8
)
# Every logical (token, scale) pair maps bijectively inside each
# four-token group and agrees with FlashInfer/TRT-LLM's published map.
flat = (swizzled_token * 8 + swizzled_scale).flatten()
self.assertEqual(torch.unique(flat).numel(), 16 * 8)
self.assertEqual(
(swizzled_token[3, 7].item(), swizzled_scale[3, 7].item()), (3, 7)
)
self.assertEqual(
(swizzled_token[1, 4].item(), swizzled_scale[1, 4].item()), (2, 1)
)
def test_scales_init(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
@@ -220,6 +366,54 @@ class TestNVFP4KVCacheMethod(CustomTestCase):
self.assertTrue(torch.all(m.v_scales_gpu == 1.0))
self.assertEqual(len(m.k_scales_gpu), 4)
def test_sm100_scale_loading_preserves_uncalibrated_fallback(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
NVFP4KVCacheMethod,
)
attention = SimpleNamespace(
layer_id=0,
k_scale=torch.tensor(1.0),
v_scale=torch.tensor(1.0),
)
model = SimpleNamespace(
layers=[SimpleNamespace(self_attn=SimpleNamespace(attn=attention))]
)
method = NVFP4KVCacheMethod(num_layers=1, device="cpu")
with patch(
"sglang.srt.layers.quantization.fp4_kv_cache_quant_method.get_platform",
return_value=SimpleNamespace(is_sm100=True),
):
method.load_scales_from_model(model)
self.assertEqual(method.get_bmm_scales(0), (1.0, 1.0))
def test_sm100_scale_loading_converts_calibrated_checkpoint_scales(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
NVFP4KVCacheMethod,
)
attention = SimpleNamespace(
layer_id=0,
k_scale=torch.tensor(0.002),
v_scale=torch.tensor(0.003),
)
model = SimpleNamespace(
layers=[SimpleNamespace(self_attn=SimpleNamespace(attn=attention))]
)
method = NVFP4KVCacheMethod(num_layers=1, device="cpu")
with patch(
"sglang.srt.layers.quantization.fp4_kv_cache_quant_method.get_platform",
return_value=SimpleNamespace(is_sm100=True),
):
method.load_scales_from_model(model)
k_scale, v_scale = method.get_bmm_scales(0)
self.assertAlmostEqual(k_scale, 0.012)
self.assertAlmostEqual(v_scale, 0.018)
@skip_if_no_blackwell_nvfp4
def test_quantize_dequantize_roundtrip(self):
"""Test NVFP4 quantize->dequantize roundtrip on CUDA."""
@@ -34,6 +34,8 @@ from sglang.srt.arg_groups.hisparse_hook import (
)
from sglang.srt.arg_groups.kv_cache_hook import (
handle_cache_compatibility,
handle_kv4_compatibility,
handle_nvfp4_prefill_kv_dequant_dtype,
validate_prefill_only_disable_kv_cache_args,
)
from sglang.srt.arg_groups.mamba_hook import handle_mamba_backend
@@ -746,6 +748,233 @@ class TestMambaCacheStochasticRounding(unittest.TestCase):
handle_mamba_backend(server_args)
class TestKV4Compatibility(unittest.TestCase):
def setUp(self):
self._use_mla_backend_patcher = patch(
"sglang.srt.arg_groups.kv_cache_hook.use_mla_backend",
return_value=False,
)
self._use_mla_backend_patcher.start()
self.addCleanup(self._use_mla_backend_patcher.stop)
@staticmethod
def _make_nvfp4_args(**overrides):
return ServerArgs(
model_path="dummy",
kv_cache_dtype="nvfp4",
attention_backend="trtllm_mha",
**overrides,
)
@staticmethod
def _make_unrouted_nvfp4_args(**overrides):
return ServerArgs(model_path="dummy", kv_cache_dtype="nvfp4", **overrides)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_prefill_kv_dequant_dtype_selects_native_backends_on_sm100(self):
args = self._make_unrouted_nvfp4_args(prefill_kv_cache_dequant_dtype="nvfp4")
handle_nvfp4_prefill_kv_dequant_dtype(args)
self.assertEqual(
resolution_result(args, "prefill_attention_backend"), "trtllm_mha"
)
self.assertEqual(
resolution_result(args, "decode_attention_backend"), "trtllm_mha"
)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_prefill_kv_dequant_dtype_selects_fp8_prefill_on_sm100(self):
args = self._make_unrouted_nvfp4_args(prefill_kv_cache_dequant_dtype="fp8_e4m3")
handle_nvfp4_prefill_kv_dequant_dtype(args)
self.assertEqual(
resolution_result(args, "prefill_attention_backend"), "flashinfer"
)
self.assertEqual(
resolution_result(args, "decode_attention_backend"), "trtllm_mha"
)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_prefill_kv_dequant_dtype_auto_defaults_to_native_on_sm100(self):
args = self._make_unrouted_nvfp4_args()
handle_nvfp4_prefill_kv_dequant_dtype(args)
self.assertEqual(
resolution_result(args, "prefill_kv_cache_dequant_dtype"), "nvfp4"
)
self.assertEqual(
resolution_result(args, "prefill_attention_backend"), "trtllm_mha"
)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_prefill_kv_dequant_dtype_auto_preserves_fp8_prefill_recipe(self):
args = self._make_unrouted_nvfp4_args(
prefill_attention_backend="flashinfer",
decode_attention_backend="trtllm_mha",
)
handle_nvfp4_prefill_kv_dequant_dtype(args)
self.assertEqual(
resolution_result(args, "prefill_kv_cache_dequant_dtype"), "fp8_e4m3"
)
@override_platform(is_cuda=True, is_sm100=False, is_sm120=True)
def test_prefill_kv_dequant_dtype_auto_defaults_to_fp8_on_sm120(self):
args = self._make_unrouted_nvfp4_args()
handle_nvfp4_prefill_kv_dequant_dtype(args)
self.assertEqual(
resolution_result(args, "prefill_kv_cache_dequant_dtype"), "fp8_e4m3"
)
self.assertEqual(
resolution_result(args, "prefill_attention_backend"), "flashinfer"
)
@override_platform(is_cuda=True, is_sm100=False, is_sm120=True)
def test_prefill_kv_dequant_dtype_rejects_native_prefill_off_sm100(self):
args = self._make_unrouted_nvfp4_args(prefill_kv_cache_dequant_dtype="nvfp4")
with self.assertRaisesRegex(ValueError, "requires SM100"):
handle_nvfp4_prefill_kv_dequant_dtype(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_prefill_kv_dequant_dtype_rejects_conflicting_prefill_backend(self):
args = self._make_unrouted_nvfp4_args(
prefill_kv_cache_dequant_dtype="nvfp4",
prefill_attention_backend="flashinfer",
)
with self.assertRaisesRegex(ValueError, "Remove the backend option"):
handle_nvfp4_prefill_kv_dequant_dtype(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_prefill_kv_dequant_dtype_rejects_conflicting_decode_backend(self):
args = self._make_unrouted_nvfp4_args(
prefill_kv_cache_dequant_dtype="fp8_e4m3",
decode_attention_backend="flashinfer",
)
with self.assertRaisesRegex(ValueError, "NVFP4 decode requires"):
handle_nvfp4_prefill_kv_dequant_dtype(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_prefill_kv_dequant_dtype_rejects_non_nvfp4_storage(self):
args = ServerArgs(
model_path="dummy",
kv_cache_dtype="fp8_e4m3",
prefill_kv_cache_dequant_dtype="nvfp4",
)
with self.assertRaisesRegex(ValueError, "applies only"):
handle_nvfp4_prefill_kv_dequant_dtype(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_sm100_native_nvfp4_allows_topk_one_speculative_decoding(self):
for algorithm in ("EAGLE", "EAGLE3", "NEXTN"):
for prefill_backend, decode_backend in (
(None, None),
("flashinfer", "trtllm_mha"),
):
for speculative_attention_mode in ("prefill", "decode"):
with self.subTest(
algorithm=algorithm,
prefill_backend=prefill_backend,
decode_backend=decode_backend,
speculative_attention_mode=speculative_attention_mode,
):
args = self._make_nvfp4_args(
prefill_attention_backend=prefill_backend,
decode_attention_backend=decode_backend,
speculative_algorithm=algorithm,
speculative_eagle_topk=1,
speculative_attention_mode=speculative_attention_mode,
)
handle_kv4_compatibility(args)
expected_mode = (
"decode"
if prefill_backend == "flashinfer"
and speculative_attention_mode == "prefill"
else speculative_attention_mode
)
self.assertEqual(
resolution_result(args, "speculative_attention_mode"),
expected_mode,
)
self.assertEqual(
resolution_result(
args, "speculative_draft_attention_backend"
),
"trtllm_mha",
)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_sm100_native_nvfp4_rejects_unvalidated_spec_algorithms(self):
for algorithm in ("STANDALONE", "FROZEN_KV_MTP", "CUSTOM_SPEC"):
with self.subTest(algorithm=algorithm):
args = self._make_nvfp4_args(speculative_algorithm=algorithm)
with self.assertRaisesRegex(ValueError, "supports EAGLE"):
handle_kv4_compatibility(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_sm100_native_nvfp4_rejects_non_native_draft_backend(self):
args = self._make_nvfp4_args(
speculative_algorithm="EAGLE",
speculative_eagle_topk=1,
speculative_draft_attention_backend="flashinfer",
)
with self.assertRaisesRegex(ValueError, "physical NVFP4 KV layout"):
handle_kv4_compatibility(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_sm100_native_nvfp4_ngram_needs_no_draft_backend(self):
args = self._make_nvfp4_args(
speculative_algorithm="NGRAM",
speculative_ngram_max_bfs_breadth=1,
)
handle_kv4_compatibility(args)
self.assertIsNone(
resolution_result(args, "speculative_draft_attention_backend")
)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_sm100_native_nvfp4_rejects_branched_ngram(self):
args = self._make_nvfp4_args(
speculative_algorithm="NGRAM",
speculative_ngram_max_bfs_breadth=2,
)
with self.assertRaisesRegex(ValueError, "speculative-ngram-max-bfs-breadth=1"):
handle_kv4_compatibility(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_sm100_native_nvfp4_rejects_prefix_commit_spec_algorithms(self):
for algorithm in ("DFLASH", "DSPARK"):
with self.subTest(algorithm=algorithm):
args = self._make_nvfp4_args(speculative_algorithm=algorithm)
with self.assertRaisesRegex(ValueError, algorithm):
handle_kv4_compatibility(args)
@override_platform(is_cuda=True, is_sm100=False, is_sm120=True)
def test_sm120_xqa_keeps_existing_speculative_support(self):
args = self._make_nvfp4_args(speculative_algorithm="EAGLE")
handle_kv4_compatibility(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_sm100_native_nvfp4_allows_monolithic_non_speculative_inference(self):
args = self._make_nvfp4_args()
handle_kv4_compatibility(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_sm100_native_nvfp4_rejects_host_tiered_cache(self):
for option in ("enable_hierarchical_cache", "enable_lmcache"):
with self.subTest(option=option):
args = self._make_nvfp4_args(**{option: True})
with self.assertRaisesRegex(ValueError, "host pools"):
handle_kv4_compatibility(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_sm100_native_nvfp4_rejects_pd_disaggregation(self):
args = self._make_nvfp4_args(disaggregation_mode="decode")
with self.assertRaisesRegex(ValueError, "PD disaggregation"):
handle_kv4_compatibility(args)
@override_platform(is_cuda=True, is_sm100=True, is_sm120=False)
def test_nvfp4_rejects_unified_memory(self):
args = self._make_nvfp4_args(enable_unified_memory=True)
with self.assertRaisesRegex(ValueError, "enable-unified-memory"):
handle_kv4_compatibility(args)
class TestLoadBalanceMethod(unittest.TestCase):
def _load_balance_args(self, **kwargs):
server_args = ServerArgs(model_path="dummy", **kwargs)
@@ -73,6 +73,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
"disable_hybrid_swa_memory",
"sampling_backend",
"attention_backend",
"prefill_kv_cache_dequant_dtype",
"page_size",
"moe_runner_backend",
"quantization",