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]