From c665edec6e8d636c387fa3f348a13bdfa497a14a Mon Sep 17 00:00:00 2001 From: Khoa Pham Date: Tue, 12 May 2026 22:37:45 -0700 Subject: [PATCH] [env] Make max KV chunk capacity configurable via `SGLANG_MAX_KV_CHUNK_CAPACITY` (#25120) --- docs/references/environment_variables.md | 1 + docs_new/docs/references/environment_variables.mdx | 5 +++++ python/sglang/srt/environ.py | 1 + .../srt/model_executor/forward_batch_deepseek_mha_mixin.py | 5 ++--- .../deepseek_common/attention_forward_methods/forward_mha.py | 2 +- 5 files changed, 10 insertions(+), 4 deletions(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 45e51b9ab..4b6be1a21 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -32,6 +32,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_IS_FLASHINFER_AVAILABLE` | Control FlashInfer availability check | `true` | | `SGLANG_SKIP_P2P_CHECK` | Skip P2P (peer-to-peer) access check | `false` | | `SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD` | Sets the threshold for enabling chunked prefix caching | `8192` | +| `SGLANG_MAX_KV_CHUNK_CAPACITY` | Maximum number of tokens in each KV chunk for DeepSeek MHA chunked prefix cache | `131072` | | `SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION` | Enable RoPE fusion in Fused Multi-Layer Attention | `1` | | `SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP` | Disable overlap schedule for consecutive prefill batches | `false` | | `SGLANG_SCHEDULER_MAX_RECV_PER_POLL` | Set the maximum number of requests per poll, with a negative value indicating no limit | `-1` | diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index 26791cbdd..4ab44d01d 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -137,6 +137,11 @@ SGLang supports various environment variables that can be used to configure its Sets the threshold for enabling chunked prefix caching `8192` + + `SGLANG_MAX_KV_CHUNK_CAPACITY` + Maximum number of tokens in each KV chunk for DeepSeek MHA chunked prefix cache + `131072` + `SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION` Enable RoPE fusion in Fused Multi-Layer Attention diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 57edcdd80..e40a24c0a 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -406,6 +406,7 @@ class Envs: # DeepSeek MHA Optimization SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD = EnvInt(8192) + SGLANG_MAX_KV_CHUNK_CAPACITY = EnvInt(128 * 1024) # DeepEP SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False) diff --git a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py index 2840f9f23..769d30235 100644 --- a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py +++ b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py @@ -7,6 +7,7 @@ import torch import triton import triton.language as tl +from sglang.srt.environ import envs from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton @@ -44,9 +45,7 @@ class ForwardBatchDeepSeekMHAMixin: mha_one_shot_kv_indices: Optional[torch.Tensor] = None def get_max_chunk_capacity(self): - # Maximum number of tokens in each chunk - # TODO: Should be changed to a better value, maybe passed through server args - return 128 * 1024 + return envs.SGLANG_MAX_KV_CHUNK_CAPACITY.get() def set_prefix_chunk_idx(self, idx: int): self.prefix_chunk_idx = idx diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index d710c9018..a2846434e 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -53,7 +53,7 @@ if _use_aiter_gfx95: # The minimum sum_prefix_length to enable mha with kv chunking, 8192 by default (can be changed with SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD) # For batches with smaller sum_prefix_length > 0, MLA kernel with absorption will be used instead. # max_kv_chunk_capacity: -# The maximum number of tokens in each kv chunk, 128 * 1024 by default (can be get with forward_batch.get_max_chunk_capacity()) +# The maximum number of tokens in each kv chunk, 128 * 1024 by default (can be changed with SGLANG_MAX_KV_CHUNK_CAPACITY, or get with forward_batch.get_max_chunk_capacity()) # The forward methods for MHA in DeepSeek models: #