diff --git a/docs/features.md b/docs/features.md index 9ebb9e963..f24545c71 100644 --- a/docs/features.md +++ b/docs/features.md @@ -51,6 +51,10 @@ Adopting the original design from [SGLang](https://github.com/sgl-project/sglang ![radix](https://lmsys.org/images/blog/sglang/radix_attn.jpg) *Illustration of Radix Attention from [LMSYS Blog](https://lmsys.org/blog/2024-01-17-sglang/).* +## FP8 KV Cache + +The KV cache can be stored in FP8 (E4M3) to halve its VRAM footprint, doubling the maximum context length or concurrency at the same memory budget. Enable with `--kv-dtype float8`; the default `auto` inherits the compute dtype (no quantisation). K/V are clamped to ±448 before cast (calibrated `k_scale`/`v_scale` from W8A8 checkpoints are ignored in this version). FP8 KV with the FlashAttention backend requires sm_90+ (Hopper/Blackwell); on older GPUs use `--attn fi`. + ## Overlap Scheduling To further reduce CPU overhead, Mini-SGLang employs overlap scheduling, a technique proposed in [NanoFlow](https://arxiv.org/abs/2408.12757). This approach overlaps the CPU scheduling overhead with GPU computation, improving overall system throughput. diff --git a/python/minisgl/attention/fi.py b/python/minisgl/attention/fi.py index f390137a7..c99f0c0c8 100644 --- a/python/minisgl/attention/fi.py +++ b/python/minisgl/attention/fi.py @@ -57,7 +57,8 @@ class FIMetadata(BaseAttnMetadata): page_size: Literal[1] # currently only support page_size=1 pos_encoding_mode: str seq_lens_cpu: torch.Tensor # on cpu - dtype: torch.dtype + q_dtype: torch.dtype # compute (Q) dtype + kv_dtype: torch.dtype # K/V storage dtype; differs from `q_dtype` for quantised pools wrapper: BatchPrefillWithPagedKVCacheWrapper | BatchDecodeWithPagedKVCacheWrapper initialized: bool = False # fmt: on @@ -141,9 +142,9 @@ def _initialize_metadata_once(self, metadata: FIMetadata) -> None: page_size=metadata.page_size, pos_encoding_mode=metadata.pos_encoding_mode, seq_lens=metadata.seq_lens_cpu, - data_type=metadata.dtype, - q_data_type=metadata.dtype, - kv_data_type=metadata.dtype, + data_type=metadata.q_dtype, + q_data_type=metadata.q_dtype, + kv_data_type=metadata.kv_dtype, non_blocking=True, ) else: @@ -158,8 +159,8 @@ def _initialize_metadata_once(self, metadata: FIMetadata) -> None: page_size=metadata.page_size, pos_encoding_mode=metadata.pos_encoding_mode, seq_lens=metadata.seq_lens_cpu, - q_data_type=metadata.dtype, - kv_data_type=metadata.dtype, + q_data_type=metadata.q_dtype, + kv_data_type=metadata.kv_dtype, non_blocking=True, causal=True, ) @@ -220,7 +221,8 @@ def prepare_metadata(self, batch: Batch) -> None: page_size=1, pos_encoding_mode="NONE", seq_lens_cpu=seq_len_cpu, - dtype=self.kvcache.dtype, + q_dtype=self.kvcache.dtype, + kv_dtype=self.kvcache.store_dtype, wrapper=self.decode_wrappers if batch.is_decode else self.prefill_wrapper, ) diff --git a/python/minisgl/engine/config.py b/python/minisgl/engine/config.py index bbee54e00..73c6bd7c9 100644 --- a/python/minisgl/engine/config.py +++ b/python/minisgl/engine/config.py @@ -29,6 +29,12 @@ class EngineConfig: use_pynccl: bool = True max_seq_len_override: int | None = None num_page_override: int | None = None # if not None, will override the number of pages + # resolves to compute `dtype` when unset; set to torch.float8_e4m3fn for FP8 KV cache. + kv_dtype: torch.dtype = None + + def __post_init__(self) -> None: + if self.kv_dtype is None: + object.__setattr__(self, "kv_dtype", self.dtype) @cached_property def hf_config(self): diff --git a/python/minisgl/engine/engine.py b/python/minisgl/engine/engine.py index ea29a96b0..c238f4ff6 100644 --- a/python/minisgl/engine/engine.py +++ b/python/minisgl/engine/engine.py @@ -60,6 +60,8 @@ def __init__(self, config: EngineConfig): page_size=config.page_size, device=self.device, dtype=self.dtype, + kv_dtype=config.kv_dtype, + attention_backend=config.attention_backend, ) # ======================= Page table initialization ======================== @@ -152,7 +154,7 @@ def _determine_num_pages(self, old_free_memory: int, config: EngineConfig) -> in * config.model_config.head_dim * div_even(config.model_config.num_kv_heads, config.tp_info.size, allow_replicate=True) * config.page_size - * self.dtype.itemsize + * config.kv_dtype.itemsize * config.model_config.num_layers ) num_pages = config.num_page_override diff --git a/python/minisgl/kvcache/__init__.py b/python/minisgl/kvcache/__init__.py index 3f5390c98..19e18ce54 100644 --- a/python/minisgl/kvcache/__init__.py +++ b/python/minisgl/kvcache/__init__.py @@ -2,10 +2,11 @@ from typing import TYPE_CHECKING, Protocol -from minisgl.utils import Registry +import torch + +from minisgl.utils import Registry, init_logger, is_sm90_supported if TYPE_CHECKING: - import torch from minisgl.models import ModelConfig from .base import ( @@ -16,6 +17,8 @@ SizeInfo, ) +logger = init_logger(__name__) + class CacheManagerCreator(Protocol): def __call__(self, device: torch.device) -> BasePrefixCache: ... @@ -30,10 +33,41 @@ def create_kvcache_pool( page_size: int, dtype: torch.dtype, device: torch.device, + kv_dtype: torch.dtype, + attention_backend: str, ) -> BaseKVCachePool: - from .mha_pool import MHAKVCache # TODO: support other variants (e.g. MLA) - - return MHAKVCache( + if kv_dtype == dtype: + from .mha_pool import MHAKVCache # TODO: support other variants (e.g. MLA) + + cls: type[BaseKVCachePool] = MHAKVCache + elif kv_dtype == torch.float8_e4m3fn: + # FA fp8 KV requires sm_90+ (FA3 on Hopper, FA4 on Blackwell). Refuse + # the combination on pre-Hopper hardware -- the kernel accepts fp8 + # tensors at the Python boundary but the underlying SASS path is + # missing, so the failure mode is silent corruption. + if "fa" in attention_backend.split(",") and not is_sm90_supported(): + major, minor = torch.cuda.get_device_capability(device) + raise ValueError( + f"FP8 KV cache with the FlashAttention backend requires sm_90 " + f"(Hopper) or sm_100 (Blackwell). Detected sm_{major}{minor}. " + f"Use --attention-backend fi (FlashInfer) instead, or run on a " + f"supported GPU." + ) + logger.warning_rank0( + "FP8 KV cache enabled with scale=1.0. K/V are clamped to +/-448 before " + "cast; outliers beyond +/-448 saturate. Expect quality regression on " + "long-context or outlier-heavy workloads. Plumbing-only; calibrated " + "k_scale/v_scale in checkpoints are ignored in this version." + ) + from .quantized_mha_pool import QuantizedMHAKVCache + + cls = QuantizedMHAKVCache + else: + raise ValueError( + f"Unsupported kv_dtype {kv_dtype}; only torch.float8_e4m3fn is supported." + ) + + return cls( num_kv_heads=model_config.num_kv_heads, num_pages=num_pages, page_size=page_size, diff --git a/python/minisgl/kvcache/base.py b/python/minisgl/kvcache/base.py index 1328c0294..e3e40d3c9 100644 --- a/python/minisgl/kvcache/base.py +++ b/python/minisgl/kvcache/base.py @@ -30,7 +30,13 @@ def device(self) -> torch.device: ... @property @abstractmethod - def dtype(self) -> torch.dtype: ... + def dtype(self) -> torch.dtype: + """Compute dtype: the dtype downstream attention / projection layers see.""" + + @property + def store_dtype(self) -> torch.dtype: + """K/V buffer storage dtype; override for pools that compute and store at different precisions.""" + return self.dtype @property @abstractmethod diff --git a/python/minisgl/kvcache/quantized_mha_pool.py b/python/minisgl/kvcache/quantized_mha_pool.py new file mode 100644 index 000000000..d6f80af64 --- /dev/null +++ b/python/minisgl/kvcache/quantized_mha_pool.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import torch + +from .mha_pool import MHAKVCache + +_FP8_E4M3FN_MAX = 448.0 + + +class QuantizedMHAKVCache(MHAKVCache): + """MHA KV cache stored as float8_e4m3fn with implicit scale=1.0. + + Calibrated k_scale/v_scale from W8A8 checkpoints are silently ignored. + """ + + def __init__( + self, + num_kv_heads: int, + num_layers: int, + head_dim: int, + num_pages: int, + page_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> None: + super().__init__( + num_kv_heads=num_kv_heads, + num_layers=num_layers, + head_dim=head_dim, + num_pages=num_pages, + page_size=page_size, + dtype=torch.float8_e4m3fn, + device=device, + ) + self._compute_dtype = dtype + + def store_kv( + self, k: torch.Tensor, v: torch.Tensor, out_loc: torch.Tensor, layer_id: int + ) -> None: + # torch.to(float8_e4m3fn) does NOT saturate -- out-of-range fp16 lands as NaN. + k_q = k.clamp(-_FP8_E4M3FN_MAX, _FP8_E4M3FN_MAX).to(torch.float8_e4m3fn) + v_q = v.clamp(-_FP8_E4M3FN_MAX, _FP8_E4M3FN_MAX).to(torch.float8_e4m3fn) + super().store_kv(k_q, v_q, out_loc, layer_id) + + @property + def dtype(self) -> torch.dtype: + return self._compute_dtype + + @property + def store_dtype(self) -> torch.dtype: + return self._kv_buffer.dtype diff --git a/python/minisgl/server/args.py b/python/minisgl/server/args.py index 3ec88f8d2..58ac29fed 100644 --- a/python/minisgl/server/args.py +++ b/python/minisgl/server/args.py @@ -83,6 +83,14 @@ def parse_args(args: List[str], run_shell: bool = False) -> Tuple[ServerArgs, bo help="Data type for model weights and activations. 'auto' will use FP16 for FP32/FP16 models and BF16 for BF16 models.", ) + parser.add_argument( + "--kv-dtype", + type=str, + default="auto", + choices=["auto", "float8"], + help="KV cache storage dtype. 'auto' uses the compute dtype (no quantisation); 'float8' enables the FP8 (E4M3) KV cache.", + ) + parser.add_argument( "--tensor-parallel-size", "--tp-size", @@ -257,8 +265,14 @@ def parse_args(args: List[str], run_shell: bool = False) -> Tuple[ServerArgs, bo "float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32, + "float8": torch.float8_e4m3fn, } kwargs["dtype"] = DTYPE_MAP[dtype_str] if isinstance(dtype_str, str) else dtype_str + + # "auto" inherits the compute dtype (no KV-cache quantisation). + kv_dtype_str = kwargs["kv_dtype"] + kwargs["kv_dtype"] = kwargs["dtype"] if kv_dtype_str == "auto" else DTYPE_MAP[kv_dtype_str] + kwargs["tp_info"] = DistributedInfo(0, kwargs["tensor_parallel_size"]) del kwargs["tensor_parallel_size"] diff --git a/tests/kernel/test_kvcache_quantized.py b/tests/kernel/test_kvcache_quantized.py new file mode 100644 index 000000000..64494398d --- /dev/null +++ b/tests/kernel/test_kvcache_quantized.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import torch + +import minisgl.distributed.info as _dist_info +from minisgl.kvcache.quantized_mha_pool import QuantizedMHAKVCache, _FP8_E4M3FN_MAX +from minisgl.utils import call_if_main + +_FP8 = torch.float8_e4m3fn + + +def _make_pool(num_pages=16, page_size=1, num_heads=4, head_dim=64, num_layers=1): + _dist_info._TP_INFO = None + _dist_info.set_tp_info(rank=0, size=1) + return QuantizedMHAKVCache( + num_kv_heads=num_heads, + num_layers=num_layers, + head_dim=head_dim, + num_pages=num_pages, + page_size=page_size, + dtype=torch.float16, + device=torch.device("cuda:0"), + ) + + +def _flat(t): + return t.reshape(t.shape[0], -1).contiguous() + + +def _read(pool, indices, layer=0, side="k"): + buf = pool.k_cache(layer) if side == "k" else pool.v_cache(layer) + return buf.reshape(-1, buf.shape[-2], buf.shape[-1])[indices] + + +@call_if_main(__name__) +def test_quantized_kvcache(): + pool = _make_pool() + assert pool.dtype == torch.float16 + assert pool.store_dtype == _FP8 + + indices = torch.tensor([1, 5, 11], device="cuda", dtype=torch.int64) + torch.manual_seed(0) + + # in-range round-trip + clamp saturation + NaN/Inf propagation, byte-exact + k = torch.randn(3, 4, 64, dtype=torch.float16, device="cuda") + v = torch.randn(3, 4, 64, dtype=torch.float16, device="cuda") + k[0, 0, 0] = 1000.0 # saturates to +448 + v[0, 0, 0] = -1000.0 # saturates to -448 + k[1, 0, 0] = float("nan") # NaN -> NaN + v[1, 0, 0] = float("inf") # +inf -> +448 (via clamp) + + expected_k = k.clamp(-_FP8_E4M3FN_MAX, _FP8_E4M3FN_MAX).to(_FP8) + expected_v = v.clamp(-_FP8_E4M3FN_MAX, _FP8_E4M3FN_MAX).to(_FP8) + pool.store_kv(_flat(k), _flat(v), indices, layer_id=0) + + got_k = _read(pool, indices, side="k") + got_v = _read(pool, indices, side="v") + assert torch.equal(got_k.view(torch.uint8), expected_k.view(torch.uint8)) + assert torch.equal(got_v.view(torch.uint8), expected_v.view(torch.uint8)) + + k_fp16 = got_k.to(torch.float16) + v_fp16 = got_v.to(torch.float16) + assert k_fp16[0, 0, 0].item() == _FP8_E4M3FN_MAX + assert v_fp16[0, 0, 0].item() == -_FP8_E4M3FN_MAX + assert torch.isnan(k_fp16[1, 0, 0]).item() + assert v_fp16[1, 0, 0].item() == _FP8_E4M3FN_MAX # inf clamped + + +@call_if_main(__name__) +def test_default_cast_to_fp8_produces_nan(): + """Pins Phase 0 (A1): plain .to(fp8) does NOT saturate -- out-of-range -> NaN. + + If a future torch makes .to(fp8) saturating, remove the clamp in + QuantizedMHAKVCache.store_kv. + """ + x = torch.tensor([500.0], dtype=torch.float16, device="cuda") + assert torch.isnan(x.to(_FP8).to(torch.float16)).item()