[NPU] Support mamba cache transfer for NPU (#20364)

2026-03-12 12:49:21 +08:00
parent abc672e717
commit dd82678b2d
1 changed files with 5 additions and 0 deletions
--- a/python/sglang/srt/disaggregation/ascend/conn.py
+++ b/python/sglang/srt/disaggregation/ascend/conn.py
@@ -34,6 +34,11 @@ class AscendKVManager(MooncakeKVManager):
        self.engine.batch_register(
            self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
        )
+        # Batch register state/extra pool data buffers
+        if self.kv_args.state_data_ptrs and self.kv_args.state_data_lens:
+            self.engine.batch_register(
+                self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
+            )

    def send_kvcache(
        self,