From 0a212c611909bb106e980b8bdc1a1720b077f883 Mon Sep 17 00:00:00 2001 From: Zhihao Wang <101526713+xiuhu17@users.noreply.github.com> Date: Fri, 24 Jul 2026 15:52:00 -0700 Subject: [PATCH] [RL] DSV4: add env to quantize SWA KV cache from bf16-rounded values (#31086) Signed-off-by: zhihaow6 --- docs_new/docs/references/environment_variables.mdx | 5 +++++ python/sglang/srt/environ.py | 3 +++ python/sglang/srt/models/deepseek_v4.py | 9 +++++++++ 3 files changed, 17 insertions(+) diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index db244e8e1..4a847ea47 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -710,6 +710,11 @@ SGLang supports various environment variables that can be used to configure its Default reasoning_effort for the DeepSeek V4 chat encoder when a request does not set it (accepts max, high; empty means unset). "" + + SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE + For DeepSeek V4, quantize the SWA FP8 KV cache from BF16-rounded values instead of FP32 registers. This matches trainer-side QAT and the DSA prefill-CP path, at the cost of an extra BF16 KV materialization and separate cache-store kernels. + false + diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d4a741e89..506ba4cfe 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1018,6 +1018,9 @@ class Envs: # Default reasoning_effort for dsv4 chat encoder when request doesn't set it. # Accepts "", "max", "high" (empty string means unset); other values filtered to None. SGLANG_DSV4_REASONING_EFFORT = EnvStr("") + # Quantize the SWA fp8 KV cache from bf16-rounded values (matches + # trainer-side QAT and the DSA-CP path) instead of fp32 registers. + SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE = EnvBool(False) # CUDA kernels SGLANG_OPT_DEEPGEMM_HC_PRENORM = EnvBool(True) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index f054c162c..2c13f6f51 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -722,6 +722,15 @@ class MQALayer(MqaAttentionBase): Replaces the bf16-kv-intermediate path. Used everywhere except the DSA prefill-CP case (which needs bf16 kv for the cross-rank all-gather). """ + if envs.SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE.get(): + # Quantize the nope payload from bf16-rounded values (the fused + # kernel quantizes from fp32 registers; the bf16 rounding moves + # values across fp8 bins relative to bf16-sourced consumers). + kv = self._compute_kv_bf16(x, positions, qkv_a=qkv_a) + attn_backend.store_cache( + layer_id=self.layer_id, swa_k=kv, forward_batch=forward_batch + ) + return if qkv_a is not None: kv = qkv_a[..., self.q_lora_rank :] else: