diff --git a/.ci/scripts/test_model_e2e.sh b/.ci/scripts/test_model_e2e.sh index 5cee37b19cf..9643c286d6d 100755 --- a/.ci/scripts/test_model_e2e.sh +++ b/.ci/scripts/test_model_e2e.sh @@ -354,7 +354,9 @@ EOF fi ;; qwen3_5_moe) - RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt 'What is the capital of France?' --max_new_tokens 128 --temperature 0" + RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt 'What is the capital of France?' --max_new_tokens 128 --temperature 0 --cuda_graph" + # CUDA graph capture requires cudaMallocAsync backend for stream-ordered allocations + export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync ;; voxtral_realtime) RUNNER_ARGS="--model_path ${MODEL_DIR}/model.pte --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --preprocessor_path ${MODEL_DIR}/$PREPROCESSOR --audio_path ${MODEL_DIR}/$AUDIO_FILE --temperature 0" diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 8a5ad285599..c6c270e6528 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -68,6 +68,7 @@ using executorch::runtime::Span; using executorch::runtime::etensor::Tensor; // SlimTensor type aliases +using cuda::CudaGraphPhase; using slim::CPU_DEVICE; using slim::DEFAULT_CUDA_DEVICE; using slim::DeviceTraits; @@ -80,6 +81,8 @@ namespace { constexpr char kSkipCopyOutputToCpuForMethod[] = "skip_copy_output_to_cpu_for_method"; constexpr char kUseSharedCudaStream[] = "use_shared_cuda_stream"; +constexpr char kEnableCudaGraphForMethod[] = "enable_cuda_graph_for_method"; +constexpr int kCudaGraphWarmupSteps = 3; } // anonymous namespace class ET_EXPERIMENTAL CudaBackend final @@ -146,6 +149,20 @@ class ET_EXPERIMENTAL CudaBackend final return method_in_csv(method_name, skip_copy_method_); } + void set_cuda_graph_method( + const std::array& raw) { + std::lock_guard guard(cuda_graph_method_mutex_); + cuda_graph_method_ = std::string(raw.data()); + } + + bool should_use_cuda_graph_for_method(const std::string& method_name) const { + if (method_name.empty()) { + return false; + } + std::lock_guard guard(cuda_graph_method_mutex_); + return method_in_csv(method_name, cuda_graph_method_); + } + // Create the shared CUDA stream. Called when use_shared_cuda_stream option // is set to true. The presence of shared_cuda_stream_ indicates shared mode. void create_shared_cuda_stream() { @@ -264,6 +281,17 @@ class ET_EXPERIMENTAL CudaBackend final ET_LOG(Error, "Option %s must be a boolean.", kUseSharedCudaStream); return Error::InvalidArgument; } + } else if (std::strcmp(option.key, kEnableCudaGraphForMethod) == 0) { + if (auto* val = std::get_if>( + &option.value)) { + set_cuda_graph_method(*val); + } else { + ET_LOG( + Error, + "Option %s must be a method name string.", + kEnableCudaGraphForMethod); + return Error::InvalidArgument; + } } } return Error::Ok; @@ -512,6 +540,17 @@ class ET_EXPERIMENTAL CudaBackend final method_name.c_str()); } + // Initialize CUDA graph state if enabled for this method. + if (should_use_cuda_graph_for_method(method_name)) { + handle->cuda_graph_state.phase = CudaGraphPhase::Warmup; + handle->cuda_graph_state.warmup_remaining = kCudaGraphWarmupSteps; + ET_LOG( + Info, + "CUDA graph enabled for method '%s' (warmup=%d)", + method_name.c_str(), + kCudaGraphWarmupSteps); + } + return (DelegateHandle*)handle; // Return the handle post-processing } @@ -538,6 +577,60 @@ class ET_EXPERIMENTAL CudaBackend final n_outputs, args.size()) + // --------------------------------------------------------------- + // CUDA graph REPLAY path — skip all tensor setup and just replay + // --------------------------------------------------------------- + if (handle->cuda_graph_state.phase == CudaGraphPhase::Replay) { + Result csr = getCurrentCUDAStream(0); + cudaStream_t cs = csr.get(); + ET_CHECK_OK_OR_RETURN_ERROR(csr.error()); + + // Copy new input data into static input buffers + for (size_t i = 0; i < n_inputs; i++) { + auto* cpu_tensor = &(args[i]->toTensor()); + cudaMemcpyAsync( + handle->cuda_graph_state.static_input_ptrs[i], + cpu_tensor->const_data_ptr(), + handle->cuda_graph_state.static_input_nbytes[i], + cudaMemcpyHostToDevice, + cs); + } + + // Replay the captured graph + cudaError_t gerr = + cudaGraphLaunch(handle->cuda_graph_state.graph_exec, cs); + ET_CHECK_OR_RETURN_ERROR( + gerr == cudaSuccess, + Internal, + "cudaGraphLaunch failed: %s", + cudaGetErrorString(gerr)); + + // Copy outputs back to CPU + const bool copy_outputs = + !should_skip_copy_for_method(handle->method_name); + if (copy_outputs) { + for (size_t i = 0; i < n_outputs; i++) { + auto* cpu_out = &(args[i + n_inputs]->toTensor()); + cudaMemcpyAsync( + cpu_out->mutable_data_ptr(), + handle->cuda_graph_state.static_output_ptrs[i], + handle->cuda_graph_state.static_output_nbytes[i], + cudaMemcpyDeviceToHost, + cs); + } + cudaStreamSynchronize(cs); + } + + return Error::Ok; + } + + // --------------------------------------------------------------- + // Normal path (also used for WARMUP and CAPTURE phases) + // --------------------------------------------------------------- + bool is_capture_step = + (handle->cuda_graph_state.phase == CudaGraphPhase::Warmup && + handle->cuda_graph_state.warmup_remaining == 0); + // NOTE: ExecuTorch tensors may be on CPU or GPU due to the skip-copy // optimization. We need to create GPU copies for CUDA kernel execution // using SlimTensor. @@ -548,6 +641,46 @@ class ET_EXPERIMENTAL CudaBackend final for (size_t i = 0; i < n_inputs; i++) { auto* cpu_tensor = &(args[i]->toTensor()); + // CAPTURE step: allocate persistent static GPU buffers + if (is_capture_step) { + auto sizes = cpu_tensor->sizes(); + auto strides = cpu_tensor->strides(); + std::vector sizes_vec(sizes.begin(), sizes.end()); + std::vector strides_vec(strides.begin(), strides.end()); + size_t nbytes = cpu_tensor->nbytes(); + + void* static_ptr = nullptr; + cudaError_t merr = cudaMalloc(&static_ptr, nbytes); + ET_CHECK_OR_RETURN_ERROR( + merr == cudaSuccess, + Internal, + "cudaMalloc for static input %zu failed: %s", + i, + cudaGetErrorString(merr)); + + cudaMemcpy( + static_ptr, + cpu_tensor->const_data_ptr(), + nbytes, + cudaMemcpyHostToDevice); + + handle->cuda_graph_state.static_input_ptrs.push_back(static_ptr); + handle->cuda_graph_state.static_input_sizes.push_back(sizes_vec); + handle->cuda_graph_state.static_input_strides.push_back(strides_vec); + handle->cuda_graph_state.static_input_scalar_types.push_back( + static_cast(cpu_tensor->scalar_type())); + handle->cuda_graph_state.static_input_nbytes.push_back(nbytes); + + gpu_inputs[i] = new SlimTensor(slim::from_blob( + static_ptr, + slim::makeArrayRef(sizes_vec), + slim::makeArrayRef(strides_vec), + static_cast(cpu_tensor->scalar_type()), + DEFAULT_CUDA_DEVICE, + 0)); + continue; + } + // Check if input data is already on GPU (skip-copy optimization for // inputs) This can happen when the caller has pre-staged data on GPU cudaPointerAttributes attributes{}; @@ -622,6 +755,23 @@ class ET_EXPERIMENTAL CudaBackend final Result cuda_stream_ret = getCurrentCUDAStream(0); cudaStream_t cuda_stream = cuda_stream_ret.get(); ET_CHECK_OK_OR_RETURN_ERROR(cuda_stream_ret.error()); + + if (is_capture_step) { + // ----- CUDA graph CAPTURE ----- + ET_LOG( + Info, + "CUDA graph: beginning stream capture for '%s'", + handle->method_name.c_str()); + + cudaError_t cerr = + cudaStreamBeginCapture(cuda_stream, cudaStreamCaptureModeRelaxed); + ET_CHECK_OR_RETURN_ERROR( + cerr == cudaSuccess, + Internal, + "cudaStreamBeginCapture failed: %s", + cudaGetErrorString(cerr)); + } + AOTIRuntimeError error = handle->run( handle->container_handle, reinterpret_cast(gpu_inputs.data()), @@ -647,6 +797,89 @@ class ET_EXPERIMENTAL CudaBackend final "AOTInductorModelContainerRun failed with error code %d", error); + if (is_capture_step) { + // End capture → instantiate graph + cudaError_t gerr = + cudaStreamEndCapture(cuda_stream, &handle->cuda_graph_state.graph); + ET_CHECK_OR_RETURN_ERROR( + gerr == cudaSuccess, + Internal, + "cudaStreamEndCapture failed: %s", + cudaGetErrorString(gerr)); + + gerr = cudaGraphInstantiate( + &handle->cuda_graph_state.graph_exec, + handle->cuda_graph_state.graph, + cudaGraphInstantiateFlagAutoFreeOnLaunch); + ET_CHECK_OR_RETURN_ERROR( + gerr == cudaSuccess, + Internal, + "cudaGraphInstantiate failed: %s", + cudaGetErrorString(gerr)); + + // Record static output pointers (stable under graph replay) + for (size_t i = 0; i < n_outputs; i++) { + SlimTensor* out = gpu_outputs[i]; + handle->cuda_graph_state.static_output_ptrs.push_back(out->data_ptr()); + + auto out_sizes = out->sizes(); + auto out_strides = out->strides(); + handle->cuda_graph_state.static_output_sizes.push_back( + std::vector(out_sizes.begin(), out_sizes.end())); + handle->cuda_graph_state.static_output_strides.push_back( + std::vector(out_strides.begin(), out_strides.end())); + handle->cuda_graph_state.static_output_scalar_types.push_back( + static_cast(out->dtype())); + handle->cuda_graph_state.static_output_nbytes.push_back(out->nbytes()); + } + + handle->cuda_graph_state.phase = CudaGraphPhase::Replay; + ET_LOG( + Info, + "CUDA graph: captured and instantiated for '%s'", + handle->method_name.c_str()); + + // Replay once to actually produce output (capture doesn't execute) + gerr = cudaGraphLaunch(handle->cuda_graph_state.graph_exec, cuda_stream); + ET_CHECK_OR_RETURN_ERROR( + gerr == cudaSuccess, + Internal, + "cudaGraphLaunch (first replay) failed: %s", + cudaGetErrorString(gerr)); + + // Copy capture-step outputs to CPU + const bool copy_outputs = + !should_skip_copy_for_method(handle->method_name); + if (copy_outputs) { + for (size_t i = 0; i < n_outputs; i++) { + auto* cpu_out = &(args[i + n_inputs]->toTensor()); + cudaMemcpyAsync( + cpu_out->mutable_data_ptr(), + handle->cuda_graph_state.static_output_ptrs[i], + handle->cuda_graph_state.static_output_nbytes[i], + cudaMemcpyDeviceToHost, + cuda_stream); + // Don't delete — static buffers are owned by the handle + gpu_outputs[i] = nullptr; + } + } + + return Error::Ok; + } + + // ----- Normal / WARMUP execution continues here ----- + + // Decrement warmup counter if in warmup phase + if (handle->cuda_graph_state.phase == CudaGraphPhase::Warmup && + handle->cuda_graph_state.warmup_remaining > 0) { + handle->cuda_graph_state.warmup_remaining--; + ET_LOG( + Info, + "CUDA graph warmup: %d steps remaining for '%s'", + handle->cuda_graph_state.warmup_remaining, + handle->method_name.c_str()); + } + const bool copy_outputs = !should_skip_copy_for_method(handle->method_name); if (copy_outputs) { @@ -741,6 +974,9 @@ class ET_EXPERIMENTAL CudaBackend final mutable std::mutex skip_copy_method_mutex_; std::string skip_copy_method_; + mutable std::mutex cuda_graph_method_mutex_; + std::string cuda_graph_method_; + // Shared CUDA stream for all methods. When set (non-null), all methods use // the same stream to ensure proper ordering (critical for skip-copy // optimization). Created when use_shared_cuda_stream option is set to true. diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index 02d3356379f..511293c4627 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -11,6 +11,7 @@ #include #include #include +#include namespace executorch { namespace backends { @@ -38,6 +39,52 @@ inline std::shared_ptr create_cuda_stream() { return std::shared_ptr( new cudaStream_t(stream), CudaStreamDeleter()); } + +enum class CudaGraphPhase { + Disabled = 0, + Warmup = 1, + Replay = 2, +}; + +// All CUDA graph related state grouped into a single struct. +struct CudaGraphState { + CudaGraphPhase phase = CudaGraphPhase::Disabled; + int warmup_remaining = 0; + + // Captured graph and executable instance + cudaGraph_t graph = nullptr; + cudaGraphExec_t graph_exec = nullptr; + + // Static input/output GPU buffers pinned during capture. + // These hold the tensor metadata; the underlying data pointers are fixed + // addresses that CUDA graph replay will write to / read from. + std::vector static_input_ptrs; + std::vector static_output_ptrs; + std::vector> static_input_sizes; + std::vector> static_input_strides; + std::vector> static_output_sizes; + std::vector> static_output_strides; + std::vector static_input_scalar_types; + std::vector static_output_scalar_types; + std::vector static_input_nbytes; + std::vector static_output_nbytes; + + ~CudaGraphState() { + if (graph_exec) { + cudaGraphExecDestroy(graph_exec); + } + if (graph) { + cudaGraphDestroy(graph); + } + // Only free input buffers — output buffers are owned by the AOTI runtime + // (allocated during graph capture via the caching allocator). + for (auto* ptr : static_input_ptrs) { + if (ptr) + cudaFree(ptr); + } + } +}; + // CUDA-specific delegate handle that extends AOTIDelegateHandle. // This consolidates CUDA stream management into a single location. struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { @@ -58,6 +105,9 @@ struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { bool has_cuda_stream() const { return cuda_stream != nullptr && *cuda_stream != nullptr; } + + // CUDA graph state (warmup, capture, replay, static buffers) + CudaGraphState cuda_graph_state; }; } // namespace cuda diff --git a/examples/models/qwen3_5_moe/main.cpp b/examples/models/qwen3_5_moe/main.cpp index 7f4e60596be..e7cd83dddc2 100644 --- a/examples/models/qwen3_5_moe/main.cpp +++ b/examples/models/qwen3_5_moe/main.cpp @@ -13,6 +13,8 @@ #include #include #include +#include +#include #include #include @@ -28,6 +30,7 @@ DEFINE_string(tokenizer_path, "", "HuggingFace tokenizer.json path."); DEFINE_string(prompt, "Hello", "Prompt text."); DEFINE_double(temperature, 0.8, "Sampling temperature (0 = greedy)."); DEFINE_int32(max_new_tokens, 128, "Maximum tokens to generate."); +DEFINE_bool(cuda_graph, false, "Enable CUDA graph for decode method."); namespace llm = ::executorch::extension::llm; using ::executorch::extension::from_blob; @@ -84,29 +87,25 @@ int main(int argc, char** argv) { } auto metadata = metadata_result.get(); + // Set CUDA graph option if requested (must be before load_method) + if (FLAGS_cuda_graph) { + executorch::runtime::BackendOptions<2> cuda_opts; + cuda_opts.set_option("enable_cuda_graph_for_method", "decode"); + executorch::runtime::set_option("CudaBackend", cuda_opts.view()); + printf("CUDA graph enabled for decode method\n"); + } + printf("Loading methods...\n"); - // Try loading both methods; fall back to single "forward" method - bool dual_method = true; - std::string prefill_method = "prefill"; auto err = module->load_method("prefill"); if (err != Error::Ok) { - // Try "forward" for single-method export - err = module->load_method("forward"); - if (err != Error::Ok) { - ET_LOG(Error, "Failed to load prefill/forward method"); - return 1; - } - prefill_method = "forward"; - dual_method = false; - printf("Using single-method mode (forward)\n"); + ET_LOG(Error, "Failed to load prefill method"); + return 1; } - if (dual_method) { - err = module->load_method("decode"); - if (err != Error::Ok) { - ET_LOG(Error, "Failed to load decode method"); - return 1; - } + err = module->load_method("decode"); + if (err != Error::Ok) { + ET_LOG(Error, "Failed to load decode method"); + return 1; } // Get EOS ids @@ -149,7 +148,7 @@ int main(int argc, char** argv) { prefill_inputs.push_back(tokens_tensor); prefill_inputs.push_back(pos_tensor); - auto prefill_result = module->execute(prefill_method, prefill_inputs); + auto prefill_result = module->execute("prefill", prefill_inputs); if (prefill_result.error() != Error::Ok) { ET_LOG(Error, "Prefill failed"); return 1; @@ -176,11 +175,6 @@ int main(int argc, char** argv) { // decode method, which may run on a different CUDA stream. cudaDeviceSynchronize(); - if (!dual_method) { - printf("Single-method mode: skipping decode\n"); - return 0; - } - // --------------------------------------------------------------- // Decode — generate tokens one at a time // ---------------------------------------------------------------