Fix IndexError in Triton backend with pipeline parallelism (#30340)

Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Dayananda V
2026-08-07 18:58:04 +08:00
committed by GitHub
co-authored by Claude Sonnet 4.5 Shangming Cai
parent 7af3d000f2
commit fe52b49827
2 changed files with 13 additions and 5 deletions
+8 -2
View File
@@ -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]