From dd82678b2d9b44d3c7854b43900b5e224d9cc7a0 Mon Sep 17 00:00:00 2001
From: Hexq0210 <893781835@qq.com>
Date: Thu, 12 Mar 2026 12:49:21 +0800
Subject: [PATCH] [NPU] Support mamba cache transfer for NPU (#20364)

---
 python/sglang/srt/disaggregation/ascend/conn.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/sglang/srt/disaggregation/ascend/conn.py b/python/sglang/srt/disaggregation/ascend/conn.py
index 1eff99404..d8ef16afa 100644
--- 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,