diff --git a/xllm/core/kernels/musa/llm_decode_metadata_update.cu b/xllm/core/kernels/musa/llm_decode_metadata_update.cu index 5d74fb9e60..852d6a53ec 100644 --- a/xllm/core/kernels/musa/llm_decode_metadata_update.cu +++ b/xllm/core/kernels/musa/llm_decode_metadata_update.cu @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include "core/kernels/musa/llm_decode_metadata_update.h" @@ -67,6 +68,109 @@ __global__ void llm_decode_metadata_update_kernel( } // namespace +void memcpy_async(void* dst, + const void* src, + size_t bytes, + cudaMemcpyKind kind, + LlmDecodeMetadataUpdateStream stream) { + if (bytes == 0 || dst == nullptr || src == nullptr) { + return; + } + const cudaError_t error = cudaMemcpyAsync(dst, src, bytes, kind, stream); + CHECK_EQ(error, cudaSuccess) + << "llm_decode_metadata memcpy failed: " << cudaGetErrorString(error); +} + +void update_llm_decode_metadata_from_host( + const LlmDecodeMetadataHostUpdateParams& params, + LlmDecodeMetadataUpdateStream stream) { + const int64_t actual_num_tokens = params.actual_num_tokens; + const int64_t padded_num_tokens = params.padded_num_tokens; + const int64_t actual_batch_size = params.actual_batch_size; + const int64_t actual_indices_size = params.actual_indices_size; + if (actual_num_tokens <= 0 && actual_batch_size <= 0) { + return; + } + + const size_t token_bytes = + static_cast(actual_num_tokens) * sizeof(int32_t); + memcpy_async(params.dst_tokens, + params.src_tokens, + token_bytes, + cudaMemcpyDeviceToDevice, + stream); + memcpy_async(params.dst_positions, + params.src_positions, + token_bytes, + cudaMemcpyDeviceToDevice, + stream); + memcpy_async(params.dst_new_cache_slots, + params.src_new_cache_slots, + token_bytes, + cudaMemcpyDeviceToDevice, + stream); + + if (padded_num_tokens > actual_num_tokens) { + const size_t pad_bytes = + static_cast(padded_num_tokens - actual_num_tokens) * + sizeof(int32_t); + const cudaError_t zero_tokens = cudaMemsetAsync( + params.dst_tokens + actual_num_tokens, 0, pad_bytes, stream); + CHECK_EQ(zero_tokens, cudaSuccess) + << "llm_decode_metadata token padding memset failed: " + << cudaGetErrorString(zero_tokens); + const cudaError_t zero_slots = cudaMemsetAsync( + params.dst_new_cache_slots + actual_num_tokens, 0, pad_bytes, stream); + CHECK_EQ(zero_slots, cudaSuccess) + << "llm_decode_metadata slot padding memset failed: " + << cudaGetErrorString(zero_slots); + } + + if (actual_batch_size >= 0 && params.host_kv_seq_lens != nullptr) { + const size_t kv_cu_bytes = + static_cast(actual_batch_size + 1) * sizeof(int32_t); + memcpy_async(params.dst_kv_seq_lens, + params.host_kv_seq_lens, + kv_cu_bytes, + cudaMemcpyHostToDevice, + stream); + if (actual_batch_size > 0 && params.dst_kv_seq_lens_delta != nullptr) { + std::vector kv_delta(static_cast(actual_batch_size)); + for (int64_t i = 0; i < actual_batch_size; ++i) { + kv_delta[static_cast(i)] = + params.host_kv_seq_lens[i + 1] - params.host_kv_seq_lens[i]; + } + memcpy_async(params.dst_kv_seq_lens_delta, + kv_delta.data(), + static_cast(actual_batch_size) * sizeof(int32_t), + cudaMemcpyHostToDevice, + stream); + } + } + + if (params.host_paged_kv_indptr != nullptr && actual_batch_size >= 0) { + memcpy_async(params.dst_paged_kv_indptr, + params.host_paged_kv_indptr, + static_cast(actual_batch_size + 1) * sizeof(int32_t), + cudaMemcpyHostToDevice, + stream); + } + if (params.host_paged_kv_indices != nullptr && actual_indices_size > 0) { + memcpy_async(params.dst_paged_kv_indices, + params.host_paged_kv_indices, + static_cast(actual_indices_size) * sizeof(int32_t), + cudaMemcpyHostToDevice, + stream); + } + if (params.host_paged_kv_last_page_len != nullptr && actual_batch_size > 0) { + memcpy_async(params.dst_paged_kv_last_page_len, + params.host_paged_kv_last_page_len, + static_cast(actual_batch_size) * sizeof(int32_t), + cudaMemcpyHostToDevice, + stream); + } +} + void update_llm_decode_metadata(const LlmDecodeMetadataUpdateParams& params, LlmDecodeMetadataUpdateStream stream) { int64_t max_work_size = params.actual_num_tokens; diff --git a/xllm/core/kernels/musa/llm_decode_metadata_update.h b/xllm/core/kernels/musa/llm_decode_metadata_update.h index a887fd033a..6decb6e840 100644 --- a/xllm/core/kernels/musa/llm_decode_metadata_update.h +++ b/xllm/core/kernels/musa/llm_decode_metadata_update.h @@ -49,4 +49,32 @@ struct LlmDecodeMetadataUpdateParams { void update_llm_decode_metadata(const LlmDecodeMetadataUpdateParams& params, LlmDecodeMetadataUpdateStream stream); +// CPU-direct graph replay path: H2D paged-KV / kv_seq_lens from host mirrors, +// D2D tokens / positions / new_cache_slots from device inputs. +struct LlmDecodeMetadataHostUpdateParams { + const int32_t* src_tokens; + const int32_t* src_positions; + const int32_t* src_new_cache_slots; + const int32_t* host_kv_seq_lens; + const int32_t* host_paged_kv_indptr; + const int32_t* host_paged_kv_indices; + const int32_t* host_paged_kv_last_page_len; + int32_t* dst_tokens; + int32_t* dst_positions; + int32_t* dst_new_cache_slots; + int32_t* dst_kv_seq_lens; + int32_t* dst_kv_seq_lens_delta; + int32_t* dst_paged_kv_indptr; + int32_t* dst_paged_kv_indices; + int32_t* dst_paged_kv_last_page_len; + int64_t actual_num_tokens; + int64_t padded_num_tokens; + int64_t actual_batch_size; + int64_t actual_indices_size; +}; + +void update_llm_decode_metadata_from_host( + const LlmDecodeMetadataHostUpdateParams& params, + LlmDecodeMetadataUpdateStream stream); + } // namespace xllm::kernel::cuda