[GDN][KDA][mem_cache] int8 checkpoint pool for the linear-attn prefix cache (#28185)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
@@ -0,0 +1,93 @@
|
||||
"""
|
||||
End-to-end test for the int8 mamba checkpoint pool on a real GDN-hybrid model.
|
||||
|
||||
Launches Qwen3-Next-80B-A3B (a gated-delta-net / linear-attention hybrid) with
|
||||
``--enable-int8-mamba-checkpoint`` and checks two things while the int8 dequant
|
||||
path is exercised:
|
||||
|
||||
* KLDivergenceMixin — on a prefix/decode cache HIT the generated logprobs are
|
||||
compared (KL) against a full recompute. This is the *sensitive* precision
|
||||
guard: it directly bounds how far the int8-reused state moves the output
|
||||
distribution from the exact-recompute distribution.
|
||||
* test_gsm8k — end-to-end task accuracy holds.
|
||||
|
||||
NOTE: the int8 checkpoint is only engaged when a cached prefix is reused FROM the
|
||||
int8 pool, which requires ``--mamba-scheduler-strategy extra_buffer`` — the default
|
||||
``no_buffer`` only snapshots the recurrent state at the full-sequence leaf, so a
|
||||
fixed-prefix / divergent-question workload reuses ~0 mamba state and the int8 path
|
||||
would never fire.
|
||||
|
||||
Usage:
|
||||
python3 -m unittest test_int8_mamba_checkpoint_e2e
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin
|
||||
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
|
||||
from sglang.test.test_utils import DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST
|
||||
|
||||
register_cuda_ci(est_time=400, stage="base-c", runner_config="4-gpu-h100")
|
||||
|
||||
|
||||
class TestInt8MambaCheckpointE2E(KLDivergenceMixin, DefaultServerBase):
|
||||
"""int8 mamba checkpoint pool on Qwen3-Next-80B-A3B (GDN-hybrid)."""
|
||||
|
||||
model = DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST
|
||||
|
||||
# Cache-hit KL: int8 is a lossy codec, so its cache-hit divergence is
|
||||
# inherently larger than the bf16/fp8 reuse the other KL tests bound (~0.005),
|
||||
# and it grows with context length (a longer prefix = a fuller state = larger
|
||||
# absolute rounding error in the logits). Measured on a Qwen3.5-35B stand-in
|
||||
# over LongBench-V2 prompts: prefill ~0.044, decode ~0.024. Thresholds are set
|
||||
# to ~2x that, to cover model differences (80B) and the reuse path's
|
||||
# run-to-run noise while still catching a real int8 regression.
|
||||
kl_div_thres = 0.06
|
||||
kl_div_thres_prefill = 0.10
|
||||
kl_div_thres_decode = 0.06
|
||||
kl_div_max_samples = 16
|
||||
kl_div_prefill_max_new_tokens = 512
|
||||
kl_div_decode_max_new_tokens = 512
|
||||
|
||||
gsm8k_threshold = 0.90
|
||||
num_gsm8k_questions = 100
|
||||
num_shots = 8
|
||||
parallel = 8
|
||||
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--tp-size",
|
||||
"4",
|
||||
"--mem-fraction-static",
|
||||
"0.7",
|
||||
"--enable-int8-mamba-checkpoint",
|
||||
"--mamba-scheduler-strategy",
|
||||
"extra_buffer",
|
||||
]
|
||||
|
||||
def test_gsm8k(self):
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k
|
||||
|
||||
url = urlparse(self.base_url)
|
||||
args = SimpleNamespace(
|
||||
num_shots=self.num_shots,
|
||||
data_path=None,
|
||||
num_questions=self.num_gsm8k_questions,
|
||||
max_new_tokens=512,
|
||||
parallel=self.parallel,
|
||||
host=f"http://{url.hostname}",
|
||||
port=int(url.port),
|
||||
)
|
||||
metrics = run_few_shot_gsm8k(args)
|
||||
print(
|
||||
f"[{self.__class__.__name__}] GSM8K accuracy: {metrics['accuracy']:.3f} "
|
||||
f"(threshold: {self.gsm8k_threshold})"
|
||||
)
|
||||
self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,178 @@
|
||||
"""Tests for Int8CheckpointStore (int8-compressed cached linear-attn states).
|
||||
|
||||
CPU tests cover the codec error bound, store/load round-trip, and the active-pool
|
||||
copy-on-write helpers. The CUDA test reproduces the validated decode-output error
|
||||
(int8 checkpoint loaded once then decoded continues bf16) ~ 0.5%, far below the
|
||||
bf16-baseline-relative threshold that GSM8K showed is quality-safe.
|
||||
|
||||
python -m pytest test/srt/mem_cache/test_int8_checkpoint_store.py -v
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.mamba_checkpoint_pool import (
|
||||
Int8CheckpointStore,
|
||||
MambaCheckpointPool,
|
||||
)
|
||||
|
||||
H, V, K = 32, 128, 128
|
||||
L = 4
|
||||
|
||||
|
||||
def _rand_state(n, device="cpu"):
|
||||
# KDA-like state magnitudes (see fp8_checkpoint_probe: |S| mean ~6e-2)
|
||||
return torch.randn(L, n, H, V, K, device=device) * 6e-2
|
||||
|
||||
|
||||
class TestInt8CheckpointCodec(unittest.TestCase):
|
||||
def test_quantize_dequantize_error_bound(self):
|
||||
s = _rand_state(8)
|
||||
q, scale = Int8CheckpointStore.quantize(s)
|
||||
self.assertEqual(q.dtype, torch.int8)
|
||||
self.assertEqual(scale.shape, (L, 8, H, 1, K)) # per (layer,slot,head,k-chan)
|
||||
deq = Int8CheckpointStore.dequantize(q, scale, torch.float32)
|
||||
rel = (deq - s).norm() / s.norm()
|
||||
# uniform int8 per-channel on a ~uniform state: well under 1%
|
||||
self.assertLess(rel.item(), 1e-2, f"int8 codec rel err too high: {rel}")
|
||||
|
||||
def test_symmetric_and_zero(self):
|
||||
s = torch.zeros(L, 1, H, V, K)
|
||||
q, scale = Int8CheckpointStore.quantize(s)
|
||||
self.assertTrue(torch.equal(q, torch.zeros_like(q)))
|
||||
deq = Int8CheckpointStore.dequantize(q, scale, torch.float32)
|
||||
self.assertTrue(torch.equal(deq, s))
|
||||
|
||||
def test_store_load_roundtrip(self):
|
||||
store = Int8CheckpointStore(
|
||||
num_layers=L,
|
||||
num_slots=16,
|
||||
num_heads=H,
|
||||
head_v_dim=V,
|
||||
head_k_dim=K,
|
||||
device="cpu",
|
||||
)
|
||||
s = _rand_state(4)
|
||||
slots = torch.tensor([1, 3, 5, 7])
|
||||
store.store(slots, s)
|
||||
out = store.load(slots, torch.float32)
|
||||
# load == dequant of stored
|
||||
q, scale = Int8CheckpointStore.quantize(s)
|
||||
ref = Int8CheckpointStore.dequantize(
|
||||
q, scale.to(store.scale.dtype), torch.float32
|
||||
)
|
||||
self.assertLess((out - ref).abs().max().item(), 1e-3)
|
||||
|
||||
def test_cow_helpers(self):
|
||||
store = Int8CheckpointStore(
|
||||
num_layers=L,
|
||||
num_slots=16,
|
||||
num_heads=H,
|
||||
head_v_dim=V,
|
||||
head_k_dim=K,
|
||||
device="cpu",
|
||||
)
|
||||
active = torch.zeros(L, 10, H, V, K) # bf16/fp32 active pool
|
||||
active[:, 2] = _rand_state(1).squeeze(1)
|
||||
# store active slot 2 -> ckpt slot 4
|
||||
store.store_from_pool(active, torch.tensor([2]), torch.tensor([4]))
|
||||
# load ckpt slot 4 -> active slot 6 (cache-hit COW)
|
||||
store.copy_to_pool(active, torch.tensor([4]), torch.tensor([6]))
|
||||
rel = (active[:, 6] - active[:, 2]).norm() / active[:, 2].norm()
|
||||
self.assertLess(rel.item(), 1e-2)
|
||||
|
||||
def test_memory_is_half_of_bf16(self):
|
||||
store = Int8CheckpointStore(
|
||||
num_layers=L,
|
||||
num_slots=100,
|
||||
num_heads=H,
|
||||
head_v_dim=V,
|
||||
head_k_dim=K,
|
||||
device="cpu",
|
||||
)
|
||||
bf16_per_slot = L * H * V * K * 2
|
||||
# int8 data (1B) + small per-(head,k) bf16 scale -> well under bf16; ~2x slots
|
||||
self.assertLess(store.bytes_per_slot(), bf16_per_slot * 0.6)
|
||||
|
||||
def test_estimate_matches_actual_mem(self):
|
||||
# the pre-allocation estimate (used to fit-check HBM before building the
|
||||
# pool) must equal the real allocated footprint, for any temporal dtype
|
||||
for tdt in (torch.bfloat16, torch.float32):
|
||||
kw = dict(
|
||||
num_layers=L,
|
||||
num_slots=64,
|
||||
num_heads=H,
|
||||
head_v_dim=V,
|
||||
head_k_dim=K,
|
||||
conv_shapes=[(4, K)],
|
||||
conv_dtype=torch.bfloat16,
|
||||
temporal_dtype=tdt,
|
||||
)
|
||||
est = MambaCheckpointPool.estimate_mem_usage_bytes(**kw)
|
||||
pool = MambaCheckpointPool(**kw, device="cpu")
|
||||
self.assertEqual(est["qdata"] + est["scale"] + est["conv"], est["total"])
|
||||
self.assertEqual(est["total"], pool.mem_usage_bytes())
|
||||
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "needs CUDA + fla kernels")
|
||||
class TestInt8CheckpointDecodeError(unittest.TestCase):
|
||||
def test_decode_error_within_bound(self):
|
||||
try:
|
||||
from sglang.srt.layers.attention.fla.kda import fused_recurrent_kda
|
||||
except (ImportError, ModuleNotFoundError) as e:
|
||||
self.skipTest(f"fla kernels unavailable: {e}")
|
||||
|
||||
dev = "cuda"
|
||||
torch.manual_seed(0)
|
||||
|
||||
def synth(T, s):
|
||||
torch.manual_seed(s)
|
||||
q = torch.randn(1, T, H, K, device=dev, dtype=torch.bfloat16) * 0.5
|
||||
k = torch.randn(1, T, H, K, device=dev, dtype=torch.bfloat16) * 0.5
|
||||
v = (torch.randn(1, T, H, V, device=dev) * 0.5).bfloat16()
|
||||
beta = torch.rand(1, T, H, device=dev, dtype=torch.bfloat16)
|
||||
g = -torch.rand(1, T, H, K, device=dev, dtype=torch.float32) * 0.1 - 0.005
|
||||
return q, k, v, g, beta
|
||||
|
||||
def decode(state, inp):
|
||||
st = state.clone()
|
||||
o, _ = fused_recurrent_kda(
|
||||
q=inp[0],
|
||||
k=inp[1],
|
||||
v=inp[2],
|
||||
g=inp[3],
|
||||
beta=inp[4],
|
||||
scale=K**-0.5,
|
||||
initial_state=st,
|
||||
inplace_final_state=True,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
cu_seqlens=None,
|
||||
)
|
||||
return o.float()
|
||||
|
||||
S = torch.zeros(1, H, V, K, device=dev, dtype=torch.float32)
|
||||
pre = synth(512, 0)
|
||||
fused_recurrent_kda(
|
||||
q=pre[0],
|
||||
k=pre[1],
|
||||
v=pre[2],
|
||||
g=pre[3],
|
||||
beta=pre[4],
|
||||
scale=K**-0.5,
|
||||
initial_state=S,
|
||||
inplace_final_state=True,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
cu_seqlens=None,
|
||||
)
|
||||
dec = synth(128, 1)
|
||||
o_ref = decode(S, dec)
|
||||
q, scale = Int8CheckpointStore.quantize(S) # [1,H,V,K]
|
||||
S_int8 = Int8CheckpointStore.dequantize(q, scale, torch.float32)
|
||||
o_int8 = decode(S_int8, dec)
|
||||
rel = (o_int8 - o_ref).norm() / o_ref.norm()
|
||||
self.assertLess(rel.item(), 1.5e-2, f"int8 decode err {rel} too high")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user