minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[env] Make max KV chunk capacity configurable via SGLANG_MAX_KV_CHUNK_CAPACITY (#25120)

Browse Source
This commit is contained in:
Khoa Pham
2026-05-12 22:37:45 -07:00
committed by GitHub
parent b0018ad015
commit c665edec6e
5 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/references/environment_variables.md
+1
View File
@@ -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` |
docs_new/docs/references/environment_variables.mdx
+5
View File
@@ -137,6 +137,11 @@ SGLang supports various environment variables that can be used to configure its
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Sets the threshold for enabling chunked prefix caching</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`8192`</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`SGLANG_MAX_KV_CHUNK_CAPACITY`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Maximum number of tokens in each KV chunk for DeepSeek MHA chunked prefix cache</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`131072`</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable RoPE fusion in Fused Multi-Layer Attention</td>
python/sglang/srt/environ.py
+1
View File
@@ -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)
python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py
+2 -3
View File
@@ -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
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
+1 -1
View File
@@ -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:
#