NixlKVManager optimizations (#17654)

Signed-off-by: Ovidiu Mara <ovidium@nvidia.com>

python/sglang/srt/disaggregation/nixl/conn.py

@@ -144,10 +144,12 @@ class NixlKVManager(CommonKVManager):
                 "to run SGLang with NixlTransferEngine."
             ) from e
 
-        agent_config = nixl_agent_config(backends=[])
-        self.agent = nixl_agent(str(uuid.uuid4()), agent_config)
-
         backend = envs.SGLANG_DISAGGREGATION_NIXL_BACKEND.get()
+        agent_config = nixl_agent_config(
+            backends=[backend],
+            num_threads=(8 if disaggregation_mode == DisaggregationMode.PREFILL else 0),
+        )
+        self.agent = nixl_agent(str(uuid.uuid4()), agent_config)
 
         available_plugins = self.agent.get_plugin_list()
         if backend not in available_plugins:
@@ -155,9 +157,6 @@ class NixlKVManager(CommonKVManager):
                 f"NIXL backend '{backend}' not found. Available: {available_plugins}. "
                 f"Please install the required NIXL plugin or choose from: {available_plugins}"
             )
-
-        self.agent.create_backend(backend)
-        self.nixl_backend = backend
         logger.info(f"NIXL KVManager initialized with backend: {backend}")
 
         self.register_buffer_to_engine()
@@ -377,20 +376,47 @@ class NixlKVManager(CommonKVManager):
             ]
 
         src_addrs = []
+        src_lens = []
         dst_addrs = []
+        dst_lens = []
+
+        # Precompute block starts/lengths to reduce Python-level loops.
+        prefill_starts = np.fromiter(
+            (block[0] for block in prefill_kv_blocks), dtype=np.int64
+        )
+        dst_starts = np.fromiter((block[0] for block in dst_kv_blocks), dtype=np.int64)
+        block_lens = np.fromiter(
+            (len(block) for block in prefill_kv_blocks), dtype=np.int64
+        )
+
         for src_ptr, dst_ptr, item_len in layers_params:
-            for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
+            lengths = item_len * block_lens
-                src_addr = src_ptr + int(prefill_index[0]) * item_len
+            src_addrs.append(src_ptr + prefill_starts * item_len)
-                dst_addr = dst_ptr + int(decode_index[0]) * item_len
+            src_lens.append(lengths)
-                length = item_len * len(prefill_index)
+            dst_addrs.append(dst_ptr + dst_starts * item_len)
-                src_addrs.append((src_addr, length, self.kv_args.gpu_id))
+            dst_lens.append(lengths)
-                dst_addrs.append((dst_addr, length, dst_gpu_id))
+
+        def make_req_array(addr_chunks, len_chunks, gpu):
+            if not addr_chunks:
+                return np.empty((0, 3), dtype=np.int64)
+            flat_addrs = np.concatenate(addr_chunks)
+            flat_lens = np.concatenate(len_chunks)
+            return np.column_stack(
+                (
+                    flat_addrs,
+                    flat_lens,
+                    np.full_like(flat_addrs, gpu),
+                )
+            )
+
+        src_reqs = make_req_array(src_addrs, src_lens, self.kv_args.gpu_id)
+        dst_reqs = make_req_array(dst_addrs, dst_lens, dst_gpu_id)
 
         logger.debug(
             f"len(src_addrs): before group: {len(prefill_kv_indices)}, after group: {len(src_addrs)}"
         )
-        src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM")
+        src_descs = self.agent.get_xfer_descs(src_reqs, "VRAM")
-        dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM")
+        dst_descs = self.agent.get_xfer_descs(dst_reqs, "VRAM")
         # Transfer data
         xfer_handle = self.agent.initialize_xfer(
             "WRITE",
@@ -452,13 +478,6 @@ class NixlKVManager(CommonKVManager):
         src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
             self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
         )
-        # Create transfer descriptors
-        src_addrs = []
-        dst_addrs = []
-
-        bytes_per_token_on_prefill = src_kv_item_len // page_size
-        bytes_per_token_on_decode = dst_kv_item_len // page_size
-
         # Calculate precise byte offset and length for the sub-slice within the token
         src_head_slice_offset = src_head_start_offset * bytes_per_head_slice_to_send
         dst_head_slice_offset = dst_head_start_offset * bytes_per_head_slice_to_send
@@ -478,52 +497,53 @@ class NixlKVManager(CommonKVManager):
             for layer_id in range(layers_current_pp_stage)
         ]
 
+        prefill_indices = np.asarray(prefill_kv_indices, dtype=np.int64)
+        dst_indices = np.asarray(dst_kv_indices, dtype=np.int64)
+        bytes_per_token_prefill = src_kv_item_len // page_size
+        bytes_per_token_decode = dst_kv_item_len // page_size
+        token_offsets = np.arange(page_size, dtype=np.int64)
+
         src_addrs = []
         dst_addrs = []
 
-        # Calculate strides for a single token slot
-        bytes_per_token_on_prefill = src_kv_item_len // page_size
-        bytes_per_token_on_decode = dst_kv_item_len // page_size
-
         for src_ptr, dst_ptr in src_dst_ptr_pairs:
-            for i in range(len(prefill_kv_indices)):
+            src_page_bases = src_ptr + prefill_indices * src_kv_item_len
-                prefill_page_idx = int(prefill_kv_indices[i])
+            dst_page_bases = dst_ptr + dst_indices * dst_kv_item_len
-                decode_page_idx = int(dst_kv_indices[i])
 
-                # Get the starting addresses for the current src and dst pages
+            src_all = (
-                src_page_start_addr = src_ptr + prefill_page_idx * src_kv_item_len
+                src_page_bases[:, None]
-                dst_page_start_addr = dst_ptr + decode_page_idx * dst_kv_item_len
+                + token_offsets[None, :] * bytes_per_token_prefill
+                + src_head_slice_offset
+            ).ravel()
+            dst_all = (
+                dst_page_bases[:, None]
+                + token_offsets[None, :] * bytes_per_token_decode
+                + dst_head_slice_offset
+            ).ravel()
 
-                # Iterate through each valid token slot within the current page
+            src_addrs.append(src_all)
-                for token_slot_in_page in range(page_size):
+            dst_addrs.append(dst_all)
-                    # Calculate the start address of the current token slot
-                    src_token_slot_start_addr = (
-                        src_page_start_addr
-                        + token_slot_in_page * bytes_per_token_on_prefill
-                    )
-                    dst_token_slot_start_addr = (
-                        dst_page_start_addr
-                        + token_slot_in_page * bytes_per_token_on_decode
-                    )
 
-                    # Calculate final src and dst addresses by applying head-slice offsets
+        def make_req_array(addr_chunks, size, gpu):
-                    src_slice_addr = src_token_slot_start_addr + src_head_slice_offset
+            if not addr_chunks:
-                    dst_slice_addr = dst_token_slot_start_addr + dst_head_slice_offset
+                return np.empty((0, 3), dtype=np.int64)
-
+            flat_addrs = np.concatenate(addr_chunks)
-                    src_addrs.append(
+            return np.column_stack(
                 (
-                            src_slice_addr,
+                    flat_addrs,
-                            heads_bytes_per_token_to_send,
+                    np.full_like(flat_addrs, size),
-                            self.kv_args.gpu_id,
+                    np.full_like(flat_addrs, gpu),
                 )
             )
-                    dst_addrs.append(
+
-                        (dst_slice_addr, heads_bytes_per_token_to_send, dst_gpu_id)
+        src_reqs = make_req_array(
+            src_addrs, heads_bytes_per_token_to_send, self.kv_args.gpu_id
         )
+        dst_reqs = make_req_array(dst_addrs, heads_bytes_per_token_to_send, dst_gpu_id)
 
         # Use NIXL agent for transfer
-        src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM")
+        src_descs = self.agent.get_xfer_descs(src_reqs, "VRAM")
-        dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM")
+        dst_descs = self.agent.get_xfer_descs(dst_reqs, "VRAM")
 
         xfer_handle = self.agent.initialize_xfer(
             "WRITE", src_descs, dst_descs, peer_name, notif.encode("ascii")