[Feature] Unified memory: support decode context parallelism for Kimi-Linear (#36890)

This commit is contained in:
Cheng Wan
2026-09-01 12:44:26 -07:00
committed by GitHub
parent 3315356cc0
commit 0b1ce3d140
20 changed files with 671 additions and 105 deletions
@@ -91,10 +91,8 @@ class ForwardBatchDeepSeekMHAMixin:
self.prefix_chunk_starts_cpu[idx],
self.prefix_chunk_seq_lens_cpu[idx],
)
# None on a backend that never bound a translator.
src = get_attn_backend().kv_index_translator
if src is not None:
chunk_kv_indices = src.translate_full_attn_ids(chunk_kv_indices)
translator = get_attn_backend().kv_index_translator
chunk_kv_indices = translator.translate_dcp_read_ids(chunk_kv_indices)
self.prefix_chunk_kv_indices.append(chunk_kv_indices)
# Here we suppose the length of each chunk is equal
@@ -1020,7 +1020,7 @@ class ModelRunner:
self.attn_backend = backends.attn_backend
self.decode_attn_backend = backends.decode_attn_backend
self.decode_attn_backend_group = backends.decode_attn_backend_group
self.kv_index_translator.assert_backends_carry_translator(
self.kv_index_translator.bind_and_verify_backends(
[self.attn_backend, self.decode_attn_backend]
)