NixlKVManager optimizations (#17654)
Signed-off-by: Ovidiu Mara <ovidium@nvidia.com>
This commit is contained in:
ovidiusm
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user