feat(kv-cache): support SM100 NVFP4 GenMHA and speculative decoding (#36340)
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user