[Feature] Add optional FP8 (float8_e4m3fn) KV cache pool

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.

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,
         )
 

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):

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

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,

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

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

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"]
 

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()