minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Sam (Kesen Li)
2026-09-18 14:50:46 -07:00
committed by GitHub
parent f5a1434700
commit d346b214fb
19 changed files with 2106 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/test_trtllm_mha_graph_metadata.py
+63
View File
@@ -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 = []
test/registered/e2e/models/test_qwen35_9b_nvfp4_kv_cache_sm100.py
+102
View File
@@ -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()
test/registered/kernel/quant/test_nvfp4_native_kv_layout.py
+358
View File
@@ -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"]))
test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py
+201 -7
View File
@@ -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."""
test/registered/unit/server_args/test_server_args.py
+229
View File
@@ -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)
test/registered/unit/test_model_overrides.py
+1
View File
@@ -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",