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

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)

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

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