From 434d3f652b49151fe38e90252eba995035b1821e Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 3 Aug 2026 06:04:54 +0000 Subject: [PATCH] fix(nvidia): support graph replay, MoE decode, and diff builds --- .../elementwise/nvidia/elementwise_nvidia.cuh | 39 ++++++++++++++++++- src/infiniop/ops/diff/nvidia/diff_nvidia.cuh | 3 ++ .../nvidia/moe_fused_dense_nvidia.cu | 19 ++++++--- 3 files changed, 54 insertions(+), 7 deletions(-) diff --git a/src/infiniop/elementwise/nvidia/elementwise_nvidia.cuh b/src/infiniop/elementwise/nvidia/elementwise_nvidia.cuh index adebf758f..180f65a2d 100644 --- a/src/infiniop/elementwise/nvidia/elementwise_nvidia.cuh +++ b/src/infiniop/elementwise/nvidia/elementwise_nvidia.cuh @@ -21,6 +21,28 @@ __device__ __forceinline__ const T *typedInputPtr(const void *ptr) { return reinterpret_cast(ptr); } +template +struct InputPointerArray { + const void *values[N]; +}; + +/** + * @brief Stores elementwise input pointers in device workspace. + * + * The pointer array is passed by value as a kernel argument. This is required + * for CUDA Graph capture: a captured cudaMemcpyAsync from inputs.data() would + * retain a pointer to a temporary host std::vector that is destroyed before + * graph replay. + */ +template +INFINIOP_CUDA_KERNEL storeInputPointers( + const void **output, + InputPointerArray inputs) { + for (size_t i = threadIdx.x; i < N; i += blockDim.x) { + output[i] = inputs.values[i]; + } +} + /** * @brief Computes the output index in memory, accounting for strides if non-contiguous. * @@ -297,8 +319,21 @@ private: const int8_t *info_meta_start = info.getMetaStart(); const int8_t *d_meta_start = reinterpret_cast(workspace) + input_arr_size; - // copy the input pointer array and meta to device - CHECK_CUDA(cudaMemcpyAsync(workspace, h_inputs_arr, input_arr_size, cudaMemcpyHostToDevice, stream)); + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + CHECK_CUDA(cudaStreamIsCapturing(stream, &capture_status)); + if (capture_status == cudaStreamCaptureStatusNone) { + CHECK_CUDA(cudaMemcpyAsync(workspace, h_inputs_arr, input_arr_size, cudaMemcpyHostToDevice, stream)); + } else { + // A captured H2D copy from a temporary std::vector would retain an + // invalid host pointer for replay. Kernel arguments are stored by + // value in the graph node instead. + InputPointerArray input_pointers{}; + for (size_t i = 0; i < N; ++i) { + input_pointers.values[i] = h_inputs_arr[i]; + } + storeInputPointers<<<1, N, 0, stream>>>( + reinterpret_cast(workspace), input_pointers); + } CHECK_CUDA(cudaMemcpyAsync((void *)d_meta_start, info_meta_start, info.getMetaMemSize(), cudaMemcpyHostToDevice, stream)); // offset/assign the pointers diff --git a/src/infiniop/ops/diff/nvidia/diff_nvidia.cuh b/src/infiniop/ops/diff/nvidia/diff_nvidia.cuh index 83772d853..24f5a958d 100644 --- a/src/infiniop/ops/diff/nvidia/diff_nvidia.cuh +++ b/src/infiniop/ops/diff/nvidia/diff_nvidia.cuh @@ -1,8 +1,11 @@ #ifndef __DIFF_NVIDIA_H__ #define __DIFF_NVIDIA_H__ +#include "../../../../utils.h" #include "../../../operator.h" #include +#include +#include namespace op::diff::nvidia { diff --git a/src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu b/src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu index 0b6357f9f..2df6a74f0 100644 --- a/src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu +++ b/src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu @@ -403,6 +403,7 @@ infiniStatus_t launch_cutlass_gemm_grouped_device_meta(int problem_count, int64_t *d_ldb, int64_t *d_ldc, int64_t *d_ldd, + bool full_occupancy_wave, cudaStream_t stream) { if (problem_count == 0) { return INFINI_STATUS_SUCCESS; @@ -436,7 +437,15 @@ infiniStatus_t launch_cutlass_gemm_grouped_device_meta(int problem_count, cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly>::GemmKernel; using Gemm = cutlass::gemm::device::GemmGrouped; - const int threadblock_count = std::min(Gemm::sufficient(), problem_count); + const int occupancy_wave = Gemm::sufficient(); + // Decode routes to only top-k experts, but each expert still contains many + // N-dimension tiles. A full persistent wave lets CUTLASS distribute those + // tiles across the GPU instead of leaving most SMs idle. Prefill retains + // the previous problem-count cap to avoid scheduler overhead from excess + // threadblocks when most experts receive no tokens. + const int threadblock_count = full_occupancy_wave + ? occupancy_wave + : std::min(occupancy_wave, problem_count); if (threadblock_count <= 0) { return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; } @@ -547,7 +556,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info, // A D2H count copy and stream sync cannot be captured for replay. CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta( topk, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, - grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream)); + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, true, stream)); swiglu_kernel<<<(topk * intermediate_size + 255) / 256, 256, 0, stream>>>(gate_up, activated, topk, intermediate_size); @@ -560,7 +569,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info, output_permutation, pairs, num_experts, hidden_size, intermediate_size, block_size); CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta( topk, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, - grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream)); + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, true, stream)); apply_shuffle_mul_sum_kernel<<<1, std::min(hidden_size, 1024), 0, stream>>>( expert_out, reinterpret_cast(output), output_permutation, @@ -604,7 +613,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info, // avoids making the runtime problem count depend on device routing data. CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta( num_experts, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, - grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream)); + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, false, stream)); swiglu_kernel<<<(max_num_tokens_padded * intermediate_size + 255) / 256, 256, 0, stream>>>( gate_up, activated, max_num_tokens_padded, intermediate_size); @@ -615,7 +624,7 @@ infiniStatus_t calculate_typed(const MoeFusedDenseInfo &info, w2_t, expert_out, num_experts, hidden_size, intermediate_size); CHECK_STATUS(launch_cutlass_gemm_grouped_device_meta( num_experts, grouped_problems, grouped_ptr_a, grouped_ptr_b, grouped_ptr_c, grouped_ptr_d, - grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, stream)); + grouped_lda, grouped_ldb, grouped_ldc, grouped_ldd, false, stream)); apply_shuffle_mul_sum_kernel<<>>( expert_out, reinterpret_cast(output), output_permutation,