From fe52b49827e6560692d0c6b7dff2a719f5731fde Mon Sep 17 00:00:00 2001 From: Dayananda V Date: Fri, 7 Aug 2026 16:28:04 +0530 Subject: [PATCH] Fix IndexError in Triton backend with pipeline parallelism (#30340) Co-authored-by: Claude Sonnet 4.5 Co-authored-by: Shangming Cai --- python/sglang/srt/layers/attention/triton_backend.py | 8 +++++--- python/sglang/srt/mem_cache/memory_pool.py | 10 ++++++++-- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index c1f921ce1..2daae2f82 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -205,9 +205,11 @@ class TritonAttnBackend(AttentionBackend): self.v_head_dim = model_runner.token_to_kv_pool.get_v_head_dim() self.swa_v_head_dim = None else: - self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[ - -1 - ] + # Use start_layer instead of 0 to handle pipeline parallelism. + # In PP, start_layer may be > 0, so layer 0 isn't in this stage's buffer. + self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer( + model_runner.token_to_kv_pool.start_layer + ).shape[-1] self.swa_v_head_dim = None self.max_context_len = model_runner.model_config.context_len self.device = model_runner.device diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 0aa942076..3da4022d9 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -3869,7 +3869,11 @@ class HybridLinearKVPool(KVCache): ) def get_v_head_dim(self): - return self.full_kv_pool.get_value_buffer(0).shape[-1] + # Use start_layer to handle pipeline parallelism where layer 0 + # may not be present in this stage's buffer. + return self.full_kv_pool.get_value_buffer(self.full_kv_pool.start_layer).shape[ + -1 + ] def set_mla_kv_buffer( self, @@ -4991,4 +4995,6 @@ class MiniMaxSparseKVPool(KVCache): ) def get_v_head_dim(self): - return self.main_pool.get_value_buffer(0).shape[-1] + # Use start_layer to handle pipeline parallelism where layer 0 + # may not be present in this stage's buffer. + return self.main_pool.get_value_buffer(self.main_pool.start_layer).shape[-1]