[Feature] Unified memory: support decode context parallelism for Kimi-Linear (#36890)
This commit is contained in:
@@ -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]
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user