diff --git a/CMakeLists.txt b/CMakeLists.txt index 4d9691f81c6..3ab68b31584 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -111,6 +111,47 @@ set( pybind11_add_module(vllm_ascend_C ${VLLM_ASCEND_SRC}) +# Detect aclrtMemcpyBatchAsync availability (CANN 8.5+) +# Can be overridden via VLLM_ASCEND_ENABLE_BATCH_MEMCPY env var (registered +# in vllm_ascend/envs.py, forwarded by setup.py as a CMake variable): +# VLLM_ASCEND_ENABLE_BATCH_MEMCPY=1 -> force enable +# VLLM_ASCEND_ENABLE_BATCH_MEMCPY=0 -> force disable +# unset -> auto-detect from CANN headers +include(CheckCXXSourceCompiles) +set(CMAKE_REQUIRED_INCLUDES ${ASCEND_HOME_PATH}/include) +set(CMAKE_REQUIRED_LIBRARIES ascendcl) +set(CMAKE_REQUIRED_LINK_OPTIONS "-L${ASCEND_HOME_PATH}/lib64") + +if(DEFINED VLLM_ASCEND_ENABLE_BATCH_MEMCPY) + if("${VLLM_ASCEND_ENABLE_BATCH_MEMCPY}" STREQUAL "1") + message(STATUS "aclrtMemcpyBatchAsync: force enabled via VLLM_ASCEND_ENABLE_BATCH_MEMCPY=1") + target_compile_definitions(vllm_ascend_C PRIVATE CANN_MEMCPY_BATCH_ASYNC) + else() + message(STATUS "aclrtMemcpyBatchAsync: force disabled via VLLM_ASCEND_ENABLE_BATCH_MEMCPY=0") + endif() +else() + # Test the full code pattern we actually use, including struct member access. + # This ensures the macro is only defined when the API is fully compatible. + check_cxx_source_compiles(" + #include + int main() { + aclrtMemLocation loc = {}; + loc.type = ACL_MEM_LOCATION_TYPE_HOST; + loc.id = 0; + aclrtMemcpyBatchAttr attr = {}; + attr.srcLoc = loc; + attr.dstLoc = loc; + (void)aclrtMemcpyBatchAsync; + return 0; + } + " HAVE_ACLRT_MEMCPY_BATCH_ASYNC) + if(HAVE_ACLRT_MEMCPY_BATCH_ASYNC) + message(STATUS "aclrtMemcpyBatchAsync: detected in CANN headers, enabling batch memcpy path") + target_compile_definitions(vllm_ascend_C PRIVATE CANN_MEMCPY_BATCH_ASYNC) + else() + message(STATUS "aclrtMemcpyBatchAsync: not found in CANN headers, using fallback aclrtMemcpyAsync loop") + endif() +endif() # Prefer the CANN ACL headers over torch_npu's bundled third_party ACL copy. # torch_npu 2.9.0 ships an older acl_rt.h that does not declare # aclrtLaunchHostFunc, which breaks host-print compilation. diff --git a/csrc/torch_binding.cpp b/csrc/torch_binding.cpp index 6efce622bf6..a1aed1c6567 100644 --- a/csrc/torch_binding.cpp +++ b/csrc/torch_binding.cpp @@ -168,6 +168,128 @@ void swap_blocks(torch::Tensor &x, torch::Tensor &y, const torch::Tensor &z) return; } +void swap_blocks_batch(const torch::Tensor& src_ptrs, + const torch::Tensor& dst_ptrs, + const torch::Tensor& sizes, + int64_t direction) { + + TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); + TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); + TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU"); + TORCH_CHECK(src_ptrs.dtype() == torch::kInt64, "src_ptrs must be int64"); + TORCH_CHECK(dst_ptrs.dtype() == torch::kInt64, "dst_ptrs must be int64"); + TORCH_CHECK(sizes.dtype() == torch::kInt64, "sizes must be int64"); + + const int64_t n = src_ptrs.size(0); + TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs"); + TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs"); + + if (n == 0) return; + + const int64_t* src_data = src_ptrs.data_ptr(); + const int64_t* dst_data = dst_ptrs.data_ptr(); + const int64_t* size_data = sizes.data_ptr(); + + aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); + + aclrtMemcpyKind memcpy_kind; + switch (direction) { + case 0: + memcpy_kind = ACL_MEMCPY_HOST_TO_DEVICE; + break; + case 1: + memcpy_kind = ACL_MEMCPY_DEVICE_TO_HOST; + break; + case 2: + memcpy_kind = ACL_MEMCPY_DEVICE_TO_DEVICE; + break; + default: + TORCH_CHECK(false, + "swap_blocks_batch: invalid direction ", direction, + " (expected 0=H2D, 1=D2H, 2=D2D)"); + } + + // ========================================================================= + // path 1: aclrtMemcpyBatchAsync (CANN 8.5+) + // ========================================================================= +#if defined(CANN_MEMCPY_BATCH_ASYNC) + if (memcpy_kind != ACL_MEMCPY_DEVICE_TO_DEVICE) { + static_assert(sizeof(void*) == sizeof(int64_t), + "void* and int64_t must be the same size"); + static_assert(sizeof(size_t) == sizeof(int64_t), + "size_t and int64_t must be the same size"); + + void** dst_arr = reinterpret_cast( + const_cast(dst_data)); + void** src_arr = reinterpret_cast( + const_cast(src_data)); + size_t* size_arr = reinterpret_cast( + const_cast(size_data)); + size_t* dest_maxs = size_arr; + + // aclrtMemcpyBatchAttr uses srcLoc/dstLoc (aclrtMemLocation) + // to specify memory locations, not aclrtMemcpyKind. + int32_t device_id = 0; + aclrtGetDevice(&device_id); + + aclrtMemLocation host_loc = {}; + host_loc.type = ACL_MEM_LOCATION_TYPE_HOST; + host_loc.id = 0; + + aclrtMemLocation device_loc = {}; + device_loc.type = ACL_MEM_LOCATION_TYPE_DEVICE; + device_loc.id = device_id; + + aclrtMemcpyBatchAttr attr = {}; + if (memcpy_kind == ACL_MEMCPY_HOST_TO_DEVICE) { + attr.srcLoc = host_loc; + attr.dstLoc = device_loc; + } else { // ACL_MEMCPY_DEVICE_TO_HOST + attr.srcLoc = device_loc; + attr.dstLoc = host_loc; + } + + size_t attrs_index = 0; + size_t fail_index = 0; + + aclError result = aclrtMemcpyBatchAsync( + dst_arr, dest_maxs, src_arr, size_arr, + static_cast(n), + &attr, &attrs_index, 1, + &fail_index, stream); + + TORCH_CHECK(result == ACL_SUCCESS, + "aclrtMemcpyBatchAsync failed at index ", fail_index, + " with error code ", result); + return; + } +#endif + + // ========================================================================= + // path 2: aclrtMemcpyAsync + // ========================================================================= + for (int64_t i = 0; i < n; i++) { + void* dst = reinterpret_cast(dst_data[i]); + const void* src = reinterpret_cast(src_data[i]); + size_t copy_size = static_cast(size_data[i]); + + aclError ret = aclrtMemcpyAsync( + dst, + copy_size, + src, + copy_size, + memcpy_kind, + stream); + + TORCH_CHECK(ret == ACL_SUCCESS, + "aclrtMemcpyAsync failed at index ", i, + " with error code ", ret, + ", src=", src_data[i], + ", dst=", dst_data[i], + ", size=", size_data[i]); + } +} + AscendType get_dtype_from_torch(at::ScalarType scalarType) { if (scalarType == at::ScalarType::Float) { @@ -901,6 +1023,11 @@ TORCH_LIBRARY_EXPAND(CONCAT(_C, _ascend), ops) ops.def("swap_blocks(Tensor! x, Tensor! y, Tensor z) -> ()"); ops.impl("swap_blocks", torch::kPrivateUse1, &vllm_ascend::swap_blocks); + // swap_blocks_batch takes CPU tensors (int64 pointer/size arrays), not NPU + // tensors, so dispatch must be registered on the CPU backend. The function + // internally submits async memcpy on the current NPU stream. + ops.def("swap_blocks_batch(Tensor x, Tensor y, Tensor z, int direction) -> ()"); + ops.impl("swap_blocks_batch", torch::kCPU, &vllm_ascend::swap_blocks_batch); ops.def("device_print(str msg) -> ()"); ops.impl("device_print", c10::DispatchKey::CompositeExplicitAutograd, static_cast(&vllm_ascend::device_print)); diff --git a/setup.py b/setup.py index e2b899f38b6..2480cca779a 100644 --- a/setup.py +++ b/setup.py @@ -338,6 +338,11 @@ def configure(self, ext: CMakeExtension) -> None: # add TORCH_NPU_PATH cmake_args += [f"-DTORCH_NPU_PATH={torch_npu_path}"] + # Pass VLLM_ASCEND_ENABLE_BATCH_MEMCPY to CMake if explicitly set. + # When unset (None), CMake will auto-detect from CANN headers. + if envs.VLLM_ASCEND_ENABLE_BATCH_MEMCPY is not None: + cmake_args += [f"-DVLLM_ASCEND_ENABLE_BATCH_MEMCPY={envs.VLLM_ASCEND_ENABLE_BATCH_MEMCPY}"] + build_tool = [] # TODO(ganyi): ninja and ccache support for ascend c auto codegen. now we can only use make build # if which('ninja') is not None: diff --git a/vllm_ascend/envs.py b/vllm_ascend/envs.py index de408d361d8..9fc659c135b 100644 --- a/vllm_ascend/envs.py +++ b/vllm_ascend/envs.py @@ -111,6 +111,9 @@ "VLLM_ASCEND_FUSION_OP_TRANSPOSE_KV_CACHE_BY_BLOCK": lambda: bool( int(os.getenv("VLLM_ASCEND_FUSION_OP_TRANSPOSE_KV_CACHE_BY_BLOCK", "1")) ), + # Control the aclrtMemcpyBatchAsync compile path for KV cache offloading. + # "1": force enable, "0": force disable, None: auto-detect from CANN headers. + "VLLM_ASCEND_ENABLE_BATCH_MEMCPY": lambda: os.getenv("VLLM_ASCEND_ENABLE_BATCH_MEMCPY", None), } # end-env-vars-definition diff --git a/vllm_ascend/kv_offload/cpu_npu.py b/vllm_ascend/kv_offload/cpu_npu.py index 1f326a6c924..1f9f8391ac1 100644 --- a/vllm_ascend/kv_offload/cpu_npu.py +++ b/vllm_ascend/kv_offload/cpu_npu.py @@ -1,3 +1,6 @@ +from collections import deque +from dataclasses import dataclass + import numpy as np import torch from vllm.logger import logger @@ -7,6 +10,15 @@ from vllm.v1.kv_offload.worker.worker import OffloadingHandler, TransferResult, TransferSpec +@dataclass +class Transfer: + job_id: int + stream: torch.npu.Stream + start_event: torch.npu.Event + end_event: torch.npu.Event + num_bytes: int + + def expand_block_ids( block_ids: np.ndarray, block_size_factor: int, @@ -28,16 +40,15 @@ def expand_block_ids( """ assert skip_count < block_size_factor - first_range = np.arange(skip_count, block_size_factor) - full_range = np.arange(0, block_size_factor) - - output_idx = 0 - for i, block_id in enumerate(block_ids): - base_block_id = block_id * block_size_factor - indices = first_range if i == 0 else full_range - output_end_idx = output_idx + len(indices) - output[output_idx:output_end_idx] = base_block_id + indices - output_idx = output_end_idx + # Vectorized: compute all sub-block IDs at once + bases = block_ids * block_size_factor + offsets = np.arange(block_size_factor) + # shape: (num_blocks, block_size_factor) -> ravel to 1D + all_ids = (bases[:, None] + offsets[None, :]).ravel() + # Skip the first skip_count elements (only affects first block) + if skip_count > 0: + all_ids = all_ids[skip_count:] + output[: len(all_ids)] = all_ids class CpuNpuOffloadingHandler(OffloadingHandler): @@ -56,10 +67,12 @@ def __init__( self.d2h_stream = torch.npu.Stream() self.h2d_stream = torch.npu.Stream() - # job_id -> transfer npu event - self.transfer_events: dict[int, torch.npu.Event] = {} - # list of npu events available for reuse - self.events_pool: list[torch.npu.Event] = [] + # Ordered queue of in-flight transfers per direction + self._d2h_transfers: deque[Transfer] = deque() + self._h2d_transfers: deque[Transfer] = deque() + + # Reusable event pool to avoid allocation overhead + self._event_pool: list[torch.npu.Event] = [] pin_memory = is_pin_memory_available() @@ -94,24 +107,59 @@ def __init__( ) ) + # Pre-compute base pointers and block sizes for batch copies. + # In vllm-ascend, each layer's KV cache is stored as a tuple + # (key_cache, value_cache), so we flatten them into individual + # sub-tensors for batching: [layer0_key, layer0_value, + # layer1_key, layer1_value, ...]. + npu_base_ptrs = [] + cpu_base_ptrs = [] + block_sizes_in_bytes = [] + + for npu_tensor, cpu_tensor in zip(self.npu_tensors, self.cpu_tensors): + for kv_idx in range(2): # 0=key, 1=value + npu_t = npu_tensor[kv_idx] + cpu_t = cpu_tensor[kv_idx] + npu_base_ptrs.append(npu_t.data_ptr()) + cpu_base_ptrs.append(cpu_t.data_ptr()) + # block size in bytes = stride of dim 0 (elements) * element size + block_sizes_in_bytes.append(npu_t.stride(0) * npu_t.element_size()) + + self._npu_base_ptrs = np.array(npu_base_ptrs, dtype=np.int64) + self._cpu_base_ptrs = np.array(cpu_base_ptrs, dtype=np.int64) + self._block_size_in_bytes_arr = np.array(block_sizes_in_bytes, dtype=np.int64) + # Total bytes per block across all sub-tensors (for transfer stats) + self._total_bytes_per_block = int(self._block_size_in_bytes_arr.sum()) + + def _get_event(self) -> torch.npu.Event: + if self._event_pool: + return self._event_pool.pop() + return torch.npu.Event(enable_timing=True) + + def _recycle_event(self, event: torch.npu.Event) -> None: + self._event_pool.append(event) + def transfer_async(self, job_id: int, spec: TransferSpec) -> bool: - logger.info("start transfer_async...") src_spec, dst_spec = spec if isinstance(src_spec, CPULoadStoreSpec): assert isinstance(dst_spec, GPULoadStoreSpec) stream = self.h2d_stream - src_tensors = self.cpu_tensors - dst_tensors = self.npu_tensors + src_base_ptrs = self._cpu_base_ptrs + dst_base_ptrs = self._npu_base_ptrs src_block_size_factor = self.block_size_factor dst_block_size_factor = 1 + is_d2h = False + transfers = self._h2d_transfers else: assert isinstance(src_spec, GPULoadStoreSpec) assert isinstance(dst_spec, CPULoadStoreSpec) stream = self.d2h_stream - src_tensors = self.npu_tensors - dst_tensors = self.cpu_tensors + src_base_ptrs = self._npu_base_ptrs + dst_base_ptrs = self._cpu_base_ptrs src_block_size_factor = 1 dst_block_size_factor = self.block_size_factor + is_d2h = True + transfers = self._d2h_transfers src_blocks = src_spec.block_ids dst_blocks = dst_spec.block_ids @@ -123,58 +171,91 @@ def transfer_async(self, job_id: int, spec: TransferSpec) -> bool: assert src_sub_block_count == dst_blocks.size * dst_block_size_factor - dst_sub_blocks_to_skip - src_to_dst = np.empty((src_sub_block_count, 2), dtype=np.int64) - expand_block_ids(src_blocks, src_block_size_factor, src_to_dst[:, 0]) + # Expand block IDs into sub-block IDs + src_block_ids = np.empty(src_sub_block_count, dtype=np.int64) + dst_block_ids = np.empty(src_sub_block_count, dtype=np.int64) + expand_block_ids(src_blocks, src_block_size_factor, src_block_ids) expand_block_ids( dst_blocks, dst_block_size_factor, - src_to_dst[:, 1], + dst_block_ids, skip_count=dst_sub_blocks_to_skip, ) - src_to_dst_tensor = torch.from_numpy(src_to_dst) - event = self.events_pool.pop() if self.events_pool else torch.npu.Event() - with torch.npu.stream(stream): - for src_tensor, dst_tensor in zip(src_tensors, dst_tensors): - src_key_cache, src_value_cache = src_tensor[0], src_tensor[1] - dst_key_cache, dst_value_cache = dst_tensor[0], dst_tensor[1] + # Build flat pointer arrays for all sub-tensors × all block pairs. + # sub-tensors = [layer0_key, layer0_value, layer1_key, layer1_value, ...] + # Fully vectorized via numpy broadcasting (no Python loop). + num_pairs = src_sub_block_count + num_sub_tensors = len(self._block_size_in_bytes_arr) + total = num_pairs * num_sub_tensors + + # (num_sub_tensors, 1) + (1, num_pairs) * (num_sub_tensors, 1) -> (num_sub_tensors, num_pairs) + bsz_col = self._block_size_in_bytes_arr[:, None] # (T, 1) + all_src = (src_base_ptrs[:, None] + src_block_ids[None, :] * bsz_col).ravel() + all_dst = (dst_base_ptrs[:, None] + dst_block_ids[None, :] * bsz_col).ravel() + all_sizes = np.broadcast_to(bsz_col, (num_sub_tensors, num_pairs)).ravel().copy() + + batch_src = torch.from_numpy(all_src) + batch_dst = torch.from_numpy(all_dst) + batch_sizes = torch.from_numpy(all_sizes) - torch.ops._C_ascend.swap_blocks(src_key_cache, dst_key_cache, src_to_dst_tensor) - torch.ops._C_ascend.swap_blocks(src_value_cache, dst_value_cache, src_to_dst_tensor) + start_event = self._get_event() + end_event = self._get_event() - event.record(stream) + if is_d2h: + # Wait for model computation to finish before reading NPU data + stream.wait_stream(torch.npu.current_stream()) + if transfers: + # Ensure this transfer starts only after the previous one completes + last_transfer = transfers[-1] + stream.wait_event(last_transfer.end_event) - self.transfer_events[job_id] = event + with torch.npu.stream(stream): + start_event.record(stream) + if total > 0: + direction = 0 if not is_d2h else 1 + torch.ops._C_ascend.swap_blocks_batch(batch_src, batch_dst, batch_sizes, direction) + end_event.record(stream) + + transfers.append( + Transfer( + job_id=job_id, + stream=stream, + start_event=start_event, + end_event=end_event, + num_bytes=src_sub_block_count * self._total_bytes_per_block, + ) + ) - # success return True def get_finished(self) -> list[TransferResult]: results: list[TransferResult] = [] - finished_job_ids = [] - for job_id, event in self.transfer_events.items(): - if event.query(): + for transfers, transfer_type in [ + (self._d2h_transfers, ("NPU", "CPU")), + (self._h2d_transfers, ("CPU", "NPU")), + ]: + while transfers and transfers[0].end_event.query(): + transfer = transfers.popleft() + transfer_time = transfer.start_event.elapsed_time(transfer.end_event) * 1e-3 results.append( TransferResult( - job_id=job_id, + job_id=transfer.job_id, success=True, - transfer_size=None, - transfer_time=None, - transfer_type=None, + transfer_size=transfer.num_bytes, + transfer_time=transfer_time, + transfer_type=transfer_type, ) ) - finished_job_ids.append(job_id) - self.events_pool.append(event) - for job_id in finished_job_ids: - del self.transfer_events[job_id] + self._recycle_event(transfer.start_event) + self._recycle_event(transfer.end_event) return results def wait(self, job_ids: set[int]) -> None: """ Wait (block) until all specified transfer jobs are completed. """ - for job_id in job_ids: - event = self.transfer_events.get(job_id) - if event is not None: - # This will block until the NPU event is complete - event.synchronize() + for transfers in (self._d2h_transfers, self._h2d_transfers): + for transfer in transfers: + if transfer.job_id in job_ids: + transfer.end_event.synchronize()