diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 7905aae65..7eef4a9d1 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/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): - src_addr = src_ptr + int(prefill_index[0]) * item_len - dst_addr = dst_ptr + int(decode_index[0]) * item_len - length = item_len * len(prefill_index) - src_addrs.append((src_addr, length, self.kv_args.gpu_id)) - dst_addrs.append((dst_addr, length, dst_gpu_id)) + lengths = item_len * block_lens + src_addrs.append(src_ptr + prefill_starts * item_len) + src_lens.append(lengths) + dst_addrs.append(dst_ptr + dst_starts * item_len) + dst_lens.append(lengths) + + 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") - dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM") + src_descs = self.agent.get_xfer_descs(src_reqs, "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)): - prefill_page_idx = int(prefill_kv_indices[i]) - decode_page_idx = int(dst_kv_indices[i]) + src_page_bases = src_ptr + prefill_indices * src_kv_item_len + dst_page_bases = dst_ptr + dst_indices * dst_kv_item_len - # Get the starting addresses for the current src and dst pages - src_page_start_addr = src_ptr + prefill_page_idx * src_kv_item_len - dst_page_start_addr = dst_ptr + decode_page_idx * dst_kv_item_len + src_all = ( + src_page_bases[:, None] + + 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 - for token_slot_in_page in range(page_size): - # 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 - ) + src_addrs.append(src_all) + dst_addrs.append(dst_all) - # Calculate final src and dst addresses by applying head-slice offsets - src_slice_addr = src_token_slot_start_addr + src_head_slice_offset - dst_slice_addr = dst_token_slot_start_addr + dst_head_slice_offset + def make_req_array(addr_chunks, size, gpu): + if not addr_chunks: + return np.empty((0, 3), dtype=np.int64) + flat_addrs = np.concatenate(addr_chunks) + return np.column_stack( + ( + flat_addrs, + np.full_like(flat_addrs, size), + np.full_like(flat_addrs, gpu), + ) + ) - src_addrs.append( - ( - src_slice_addr, - heads_bytes_per_token_to_send, - self.kv_args.gpu_id, - ) - ) - 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") - dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM") + src_descs = self.agent.get_xfer_descs(src_reqs, "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")