feat: add musa_graph support metadata.

xllm/core/kernels/musa/llm_decode_metadata_update.cu

@@ -17,6 +17,7 @@ limitations under the License.
 #include <glog/logging.h>
 
 #include <algorithm>
+#include <vector>
 
 #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<size_t>(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<size_t>(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<size_t>(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<int32_t> kv_delta(static_cast<size_t>(actual_batch_size));
+      for (int64_t i = 0; i < actual_batch_size; ++i) {
+        kv_delta[static_cast<size_t>(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<size_t>(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<size_t>(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<size_t>(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<size_t>(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;

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