Deepseek v4: support mixed dtype compression states (#27277)

Co-authored-by: zhujunyu <zhujunyu.666@bytedance.com>

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -58,6 +58,24 @@ MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP = 1
 
 logger = logging.getLogger(__name__)
 
+
+def _get_dsv4_compress_state_dtypes() -> tuple[torch.dtype, torch.dtype]:
+    dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower()
+    if dtype_name in ("float32", "fp32"):
+        return torch.float32, torch.float32
+    if dtype_name in ("bfloat16", "bf16"):
+        if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+            raise ValueError(
+                "SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16 is not supported when "
+                "SGLANG_OPT_USE_ONLINE_COMPRESS=1; online c128 state must stay float32."
+            )
+        return torch.bfloat16, torch.bfloat16
+    raise ValueError(
+        "Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE="
+        f"{dtype_name!r}. Expected one of: float32, fp32, bfloat16, bf16."
+    )
+
+
 _is_npu = is_npu()
 _is_hip = is_hip()
 
@@ -418,7 +436,8 @@ class ModelRunnerKVCacheMixin:
                 swa_page_size=swa_page_size,
                 sliding_window=self.model_config.window_size,
                 dtype=self.kv_cache_dtype,
-                state_dtype=self.state_dtype,
+                c4_state_dtype=self.c4_state_dtype,
+                c128_state_dtype=self.c128_state_dtype,
                 qk_nope_head_dim=self.model_config.qk_nope_head_dim,
                 qk_rope_head_dim=self.model_config.qk_rope_head_dim,
                 indexer_head_dim=self.model_config.index_head_dim,
@@ -947,12 +966,12 @@ class ModelRunnerKVCacheMixin:
             self.c4_state_pool_size = config.c4_state_pool_size
             self.c128_state_pool_size = config.c128_state_pool_size
 
-        # state_dtype is a DSV4 architectural constant (fp32 for c4/c128
-        # state buffers); set unconditionally so draft workers have it before
-        # _init_pools reads it (target path also overwrites this in the
-        # configurator's resolve() for parity, harmless here).
+        # Draft worker does not own the compression-state pools, but keep the
+        # dtype attributes initialized so _init_pools can share one code path.
         if is_deepseek_v4(self.model_config.hf_config):
-            self.state_dtype = torch.float32
+            self.c4_state_dtype, self.c128_state_dtype = (
+                _get_dsv4_compress_state_dtypes()
+            )
 
         self._init_pools()
 

python/sglang/srt/model_executor/pool_configurator.py

@@ -62,6 +62,23 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _get_dsv4_compress_state_dtype_sizes() -> tuple[int, int]:
+    dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower()
+    if dtype_name in ("float32", "fp32"):
+        return 4, 4
+    if dtype_name in ("bfloat16", "bf16"):
+        if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+            raise ValueError(
+                "SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16 is not supported when "
+                "SGLANG_OPT_USE_ONLINE_COMPRESS=1; online c128 state must stay float32."
+            )
+        return 2, 2
+    raise ValueError(
+        "Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE="
+        f"{dtype_name!r}. Expected one of: float32, fp32, bfloat16, bf16."
+    )
+
+
 class MemoryPoolConfigurator:
     """Base class for memory pool configurators.
 
@@ -419,16 +436,18 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
         )
 
         attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
-        state_dtype_size = 4
-        c4_state_bytes = 2 * 2 * attn_head_dim * state_dtype_size
+        c4_state_dtype_size, c128_state_dtype_size = (
+            _get_dsv4_compress_state_dtype_sizes()
+        )
+        c4_state_bytes = 2 * 2 * attn_head_dim * c4_state_dtype_size
         # Online c128 stores (max, sum, kv) per slot (3*head_dim) instead of
         # raw (kv, score) (2*head_dim). Combined with ring_size=1 this still
         # nets a large reduction (~3/256x) but the per-slot bytes go up.
         c128_online = envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
         c128_state_bytes = (
-            (3 if c128_online else 2 * 1) * attn_head_dim * state_dtype_size
+            (3 if c128_online else 2 * 1) * attn_head_dim * c128_state_dtype_size
         )
-        c4_indexer_state_bytes = 2 * 2 * self.indexer_head_dim * state_dtype_size
+        c4_indexer_state_bytes = 2 * 2 * self.indexer_head_dim * c4_state_dtype_size
 
         c4_state_ratio = self.c4_ring_size / self.swa_page_size
         c128_state_ratio = self.c128_ring_size / self.swa_page_size