onnxruntime
b0924f47 - Speed up CUDA CI build: split into per-arch OBJECT libraries, add --flash_nvcc_threads, and enable quick build mode (#28645)

Commit
65 days ago
Speed up CUDA CI build: split into per-arch OBJECT libraries, add --flash_nvcc_threads, and enable quick build mode (#28645) ## Description Speed up CUDA CI build times by splitting the monolithic CUDA provider into architecture-specific OBJECT libraries with independent `nvcc --threads` control, and introducing a quick build mode (`onnxruntime_QUICK_BUILD`) that reduces kernel instantiations for CI validation. ## Motivation and Context CUDA builds were bottlenecked by `--nvcc_threads 1` across all targets because flash attention (48 .cu files, SM80+) requires ~4GB per nvcc thread and caused OOM when compiled with higher thread counts. The old heuristic in `build.py` used `psutil` to auto-detect memory but was unreliable and always conservative. By splitting flash attention into its own OBJECT library, the rest of the build can safely use `--threads 4` while flash attention stays at `--threads 2`. Combined with quick build mode (fewer kernel variants), this significantly reduces CI wall-clock time. ## CI Time Saving * N1F1: `--nvcc_threads 1`. CI time is from checks of [PR 28607](https://github.com/microsoft/onnxruntime/pull/28607/checks). * N4F2: `--nvcc_threads 4 --flash_nvcc_threads 2`: CI time is from this PR. * N8F4: `--nvcc_threads 8 --flash_nvcc_threads 4`: CI time is from this PR. * N4F4: `--nvcc_threads 4 --flash_nvcc_threads 4`: CI time is from this PR. This is the final candidate. * Saving = N1F1 - N4F4 * Saving Ratio = (N1F1 - N4F4) / N1F1 Here is CI time (Build + Test time in minutes) saving: CI | N1F1 | N4F2 | N8F4 | N4F4 | Saved Minutes | Saving Ratio -- | -- | -- | -- | -- | -- | -- Linux CI | 35 + 38 | 35 + 32 | 35 + 32 | 36 + 27 | 10 | 14% Windows CI | 58 + 36 | 53 + 38 | 54 + 38 | 48 + 36 | 10 | 11% Plugin Linux CI | 53 + 26 | 38 + 17 | 39 + 39 | 39 + 15 | 25 | 32% Plugin Windows CI | 77 + 16 | 57 + 14 | 54 + 14 | 53 + 12 | 28 | 30% Windows TRT CI | 54 + 43 | 38 + 38 | 42 + 43 | 41 + 37 | 19 | 20% Note that this is only one time comparison. Cache might take effect with more runs, and might change the statistics. The CI time is reduced in the range of 11% to 32%. Total CI time saving is more than 90 minutes. ## Key Changes ### 1. CMake: Architecture-specific OBJECT Libraries | File | Change | |------|--------| | `cmake/onnxruntime_cuda_source_filters.cmake` | New macros: `onnxruntime_extract_flash_attention_sources()`, `onnxruntime_extract_llm_sources()`, `onnxruntime_extract_sm_specific_cuda_sources()` to partition sources by SM arch | | `cmake/onnxruntime_providers_cuda.cmake` | Create `flash_attention` (SM80+), `llm` (SM75+), `sm90_tma`, and `sm120_tma` OBJECT libraries with per-target `--threads`; merge fpA_intB SM90 launchers into SM90 TMA lib | | `cmake/onnxruntime_providers_cuda_plugin.cmake` | Mirror OBJECT library pattern for plugin EP build; consolidate shared compile options into a variable; fix `-Xcudafe --diag_suppress=550,2810` and `--std c++20` for CUDA 12.8 compatibility | | `cmake/onnxruntime_unittests.cmake` | Link new OBJECT libraries into test target | ### 2. Build Script: `--flash_nvcc_threads` and Default 4 | File | Change | |------|--------| | `tools/ci_build/build.py` | Remove `psutil`-based memory heuristic; add `--flash_nvcc_threads` forwarding; default `nvcc_threads` to 4 | | `tools/ci_build/build_args.py` | Add `--flash_nvcc_threads` CLI argument (default: same as `--nvcc_threads`) | ### 3. Quick Build Mode (`onnxruntime_QUICK_BUILD`) - Reduces flash attention kernels to hdim128 fp16 only (skips hdim32/64/96/192/256) - Guards some MoE SM90 generated launchers with `#ifndef ORT_QUICK_BUILD` - Restricts CUTLASS SM80 tile configs to 3 instantiations - Skips test cases that depend on excluded kernel variants (e.g., `test_gqa_fp8_fallback_unsupported_head_size` needs hdim64) - Applied to all CI pipelines **except** Linux CUDA CI (full build) and packaging pipelines ### 4. CI and Packaging Pipeline Updates All CUDA CI pipelines updated from `--nvcc_threads 1` to `--nvcc_threads 4 --flash_nvcc_threads 4`: - `.github/workflows/linux_cuda_ci.yml` - `.github/workflows/linux_cuda_plugin_ci.yml` (+ `QUICK_BUILD=ON`) - `.github/workflows/linux_tensorrt_ci.yml` (+ `QUICK_BUILD=ON`) - `.github/workflows/windows_cuda.yml` (+ `QUICK_BUILD=ON`) - `.github/workflows/windows_cuda_plugin.yml` (+ `QUICK_BUILD=ON`) - `.github/workflows/windows_tensorrt.yml` (+ `QUICK_BUILD=ON`) Packaging pipeline updated to use `--nvcc_threads 4 --flash_nvcc_threads 2`, except `--nvcc_threads 2 --flash_nvcc_threads 1` for cuda plugin: - Azure Pipelines: `custom-nuget-packaging-pipeline.yml`, `nuget-win-cuda-packaging-stage.yml`, `plugin-win-cuda-stage.yml`, `py-win-gpu-stage.yml` - Linux scripts: `build_cuda_plugin_package.sh`, `build_linux_python_package.sh` ### 5. Bug Fix: CUTLASS Heuristic for SIMT Kernels - `onnxruntime/contrib_ops/cuda/llm/cutlass_heuristic.cc`: Fixed `ORT_QUICK_BUILD` path to return proper tile config for SIMT (float) gemm type instead of discarding the type info ## Architecture Mapping | OBJECT Library | Min SM | Sources | Threads | |---|---|---|---| | `*_flash_attention` | SM80+ | `bert/flash_attention/*.cu` (48 files) | `onnxruntime_FLASH_NVCC_THREADS` (default: same as nvcc_threads) | | `*_llm` | SM75+ | `contrib_ops/cuda/llm/*.cu` (excl. SM90/SM120 launchers) | `onnxruntime_NVCC_THREADS` (default 4) | | `*_sm90_tma` | 90a-real | MoE TMA + fpA_intB SM90 launchers | `onnxruntime_NVCC_THREADS` | | `*_sm120_tma` | SM120+ | MoE SM120 TMA generated files | `onnxruntime_NVCC_THREADS` | | Parent target | All archs | Everything else | `onnxruntime_NVCC_THREADS` | ## New Build Options - `--nvcc_threads N` (default 4) — threads for all CUDA targets except flash attention - `--flash_nvcc_threads N` (default: same as `--nvcc_threads`) — threads specifically for flash attention compilation CMake cache variables: `onnxruntime_NVCC_THREADS`, `onnxruntime_FLASH_NVCC_THREADS` ## Testing - Built locally with `CMAKE_CUDA_ARCHITECTURES="75;80;86;89;90;100;120"`, `--nvcc_threads 4 --flash_nvcc_threads 2` - Verified flash attention .cu files compile only for SM80+ (checked `build.ninja` / VS project) - Verified LLM .cu files compile for SM75+ - Ran `onnxruntime_provider_test` — all CUDA EP tests pass - Ran `python test_qmoe_cuda.py` (MoE kernels), flash attention / GQA tests - No link errors in both in-tree provider and plugin EP builds - No nvcc warnings about duplicate `--threads` flags - Plugin CI compile options verified: `--std c++20`, `-Xcudafe --diag_suppress=550,2810`, MSVC `/bigobj` all applied to OBJECT libraries
Author
Parents
Loading