vllm icon indicating copy to clipboard operation
vllm copied to clipboard

[CI/Build] Per file CUDA Archs (improve wheel size and dev build times)

Open LucasWilkinson opened this issue 1 year ago • 3 comments

Improve dev build times and reduce wheel size by setting which compute capabilities we compile for on a per file basis instead of globally.

Also moved some of the torch library impl registrations into associated files (out of torch_bindings.cpp) this allows us to completely leave the implementation file out of the build process if there are no target architectures that the support that kernel. This means that we can rely on the torch dispatcher to raise a NotImplementedError instead of manually written compile guards.

This should also make it easier to keep the wheel size in check as newer Nvidia GPUs get released.

Results:

Wheel Size

217.69 MB (main) -> 179.79 MB (this PR)

Clean Cache Build Time (for a single target device)

Main:

==============================================
==   TORCH_CUDA_ARCH_LIST=7.5 MAX_JOBS=12   ==
==============================================
    Longest build steps for .cpp.o:
           1.0 weighted s to build CMakeFiles/_core_C.dir/csrc/core/torch_bindings.cpp.o (12.4 s elapsed time)
           2.0 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (23.5 s elapsed time)
           2.6 weighted s to build CMakeFiles/_moe_C.dir/csrc/moe/torch_bindings.cpp.o (31.5 s elapsed time)
           3.2 weighted s to build CMakeFiles/_C.dir/csrc/torch_bindings.cpp.o (38.5 s elapsed time)
    Longest build steps for .so (linking):
           0.1 weighted s to build _core_C.abi3.so (0.7 s elapsed time)
           0.1 weighted s to build _moe_C.abi3.so (1.0 s elapsed time)
           0.5 weighted s to build _deps/vllm-flash-attn-build/vllm_flash_attn_c.abi3.so (1.1 s elapsed time)
           6.2 weighted s to build _C.abi3.so (6.2 s elapsed time)
    Longest build steps for .cu.o:
          15.3 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (183.5 s elapsed time)
          15.3 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (183.5 s elapsed time)
          15.3 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (183.6 s elapsed time)
          15.3 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (183.7 s elapsed time)
          15.5 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (185.6 s elapsed time)
          15.5 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (185.9 s elapsed time)
          15.5 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (186.2 s elapsed time)
          37.4 weighted s to build CMakeFiles/_C.dir/csrc/quantization/cutlass_w8a8/scaled_mm_c3x.cu... (449.0 s elapsed time)
          43.9 weighted s to build CMakeFiles/_C.dir/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu... (527.4 s elapsed time)
         344.8 weighted s to build CMakeFiles/_C.dir/csrc/attention/attention_kernels.cu.o (1087.2 s elapsed time)
    1110.0 s weighted time (10120.4 s elapsed time sum, 9.1x parallelism)
    134 build steps completed, average of 0.12/s

==============================================
==   TORCH_CUDA_ARCH_LIST=8.0 MAX_JOBS=12   ==
==============================================
    Longest build steps for .cu.o:
          24.6 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (295.1 s elapsed time)
          30.6 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (222.9 s elapsed time)
          37.3 weighted s to build CMakeFiles/_C.dir/csrc/quantization/cutlass_w8a8/scaled_mm_c3x.cu... (447.5 s elapsed time)
          39.6 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (287.2 s elapsed time)
          42.7 weighted s to build CMakeFiles/_C.dir/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu... (512.4 s elapsed time)
          47.5 weighted s to build CMakeFiles/_moe_C.dir/csrc/moe/marlin_moe_ops.cu.o (570.3 s elapsed time)
          48.6 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (222.2 s elapsed time)
          55.4 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (301.1 s elapsed time)
         102.0 weighted s to build CMakeFiles/_C.dir/csrc/attention/attention_kernels.cu.o (1223.4 s elapsed time)
         119.9 weighted s to build CMakeFiles/_C.dir/csrc/quantization/gptq_marlin/gptq_marlin.cu.o (1435.4 s elapsed time)
    Longest build steps for .so (linking):
           0.1 weighted s to build _moe_C.abi3.so (1.0 s elapsed time)
           0.8 weighted s to build _C.abi3.so (8.2 s elapsed time)
           1.4 weighted s to build _deps/vllm-flash-attn-build/vllm_flash_attn_c.abi3.so (1.4 s elapsed time)
    1798.6 s weighted time (20368.4 s elapsed time sum, 11.3x parallelism)
    129 build steps completed, average of 0.07/s

==============================================
==   TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=12   ==
==============================================
    Longest build steps for .cu.o:
          22.9 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (274.5 s elapsed time)
          24.6 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (295.6 s elapsed time)
          28.2 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (225.5 s elapsed time)
          36.0 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (284.5 s elapsed time)
          37.3 weighted s to build CMakeFiles/_C.dir/csrc/quantization/cutlass_w8a8/scaled_mm_c3x.cu... (447.3 s elapsed time)
          49.7 weighted s to build CMakeFiles/_moe_C.dir/csrc/moe/marlin_moe_ops.cu.o (596.5 s elapsed time)
          52.2 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (299.2 s elapsed time)
          54.6 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (221.7 s elapsed time)
          56.1 weighted s to build CMakeFiles/_C.dir/csrc/attention/attention_kernels.cu.o (673.5 s elapsed time)
         125.1 weighted s to build CMakeFiles/_C.dir/csrc/quantization/gptq_marlin/gptq_marlin.cu.o (1441.4 s elapsed time)
    Longest build steps for .so (linking):
           0.1 weighted s to build _moe_C.abi3.so (1.0 s elapsed time)
           1.5 weighted s to build _C.abi3.so (7.3 s elapsed time)
           1.5 weighted s to build _deps/vllm-flash-attn-build/vllm_flash_attn_c.abi3.so (1.5 s elapsed time)
    1727.7 s weighted time (19516.2 s elapsed time sum, 11.3x parallelism)
    129 build steps completed, average of 0.07/s

This PR

==============================================
==   TORCH_CUDA_ARCH_LIST=7.5 MAX_JOBS=12   ==
==============================================
    Longest build steps for .cpp.o:
           1.0 weighted s to build CMakeFiles/_core_C.dir/csrc/core/torch_bindings.cpp.o (11.8 s elapsed time)
           1.8 weighted s to build _deps/vllm-flash-attn-build/CMakeFiles/vllm_flash_attn_c.dir/csrc... (22.1 s elapsed time)
           2.5 weighted s to build CMakeFiles/_moe_C.dir/csrc/moe/torch_bindings.cpp.o (29.5 s elapsed time)
           2.9 weighted s to build CMakeFiles/_C.dir/csrc/torch_bindings.cpp.o (34.3 s elapsed time)
    Longest build steps for .so (linking):
           0.1 weighted s to build _core_C.abi3.so (0.7 s elapsed time)
           0.1 weighted s to build _moe_C.abi3.so (1.1 s elapsed time)
           0.4 weighted s to build _deps/vllm-flash-attn-build/vllm_flash_attn_c.abi3.so (1.1 s elapsed time)
           2.8 weighted s to build _C.abi3.so (2.8 s elapsed time)
    Longest build steps for .cu.o:
           5.9 weighted s to build CMakeFiles/_C.dir/csrc/cache_kernels.cu.o (71.4 s elapsed time)
           6.0 weighted s to build CMakeFiles/_C.dir/csrc/custom_all_reduce.cu.o (71.6 s elapsed time)
           6.0 weighted s to build CMakeFiles/_moe_C.dir/csrc/moe/marlin_moe_ops.cu.o (72.1 s elapsed time)
           6.1 weighted s to build CMakeFiles/_moe_C.dir/csrc/moe/topk_softmax_kernels.cu.o (73.8 s elapsed time)
           6.2 weighted s to build CMakeFiles/_C.dir/csrc/mamba/causal_conv1d/causal_conv1d.cu.o (74.2 s elapsed time)
           6.2 weighted s to build CMakeFiles/_C.dir/csrc/quantization/gguf/gguf_kernel.cu.o (74.6 s elapsed time)
           7.3 weighted s to build CMakeFiles/_C.dir/csrc/quantization/gptq/q_gemm.cu.o (87.2 s elapsed time)
           8.0 weighted s to build CMakeFiles/_C.dir/csrc/mamba/mamba_ssm/selective_scan_fwd.cu.o (95.5 s elapsed time)
         155.0 weighted s to build CMakeFiles/_C.dir/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu... (445.5 s elapsed time)
         563.4 weighted s to build CMakeFiles/_C.dir/csrc/attention/attention_kernels.cu.o (937.2 s elapsed time)
    955.7 s weighted time (4139.9 s elapsed time sum, 4.3x parallelism)
    96 build steps completed, average of 0.10/s

==============================================
==   TORCH_CUDA_ARCH_LIST=8.0 MAX_JOBS=12   ==
==============================================
    Longest build steps for .cu.o:
           8.0 weighted s to build CMakeFiles/_C.dir/csrc/mamba/mamba_ssm/selective_scan_fwd.cu.o (96.4 s elapsed time)
           8.7 weighted s to build CMakeFiles/_C.dir/csrc/quantization/gptq_marlin/awq_marlin_repack... (69.0 s elapsed time)
           8.8 weighted s to build CMakeFiles/_C.dir/csrc/quantization/gptq_marlin/gptq_marlin_repac... (69.6 s elapsed time)
           9.9 weighted s to build CMakeFiles/_C.dir/csrc/quantization/marlin/dense/marlin_cuda_kern... (105.4 s elapsed time)
          10.3 weighted s to build CMakeFiles/_C.dir/csrc/quantization/marlin/sparse/marlin_24_cuda_... (104.3 s elapsed time)
          11.1 weighted s to build CMakeFiles/_C.dir/csrc/quantization/fp8/fp8_marlin.cu.o (116.0 s elapsed time)
          11.1 weighted s to build CMakeFiles/_C.dir/csrc/quantization/marlin/qqq/marlin_qqq_gemm_ke... (104.9 s elapsed time)
         113.3 weighted s to build CMakeFiles/_C.dir/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu... (432.6 s elapsed time)
         374.2 weighted s to build CMakeFiles/_C.dir/csrc/attention/attention_kernels.cu.o (1013.2 s elapsed time)
         521.9 weighted s to build CMakeFiles/_C.dir/csrc/quantization/gptq_marlin/gptq_marlin.cu.o (1036.0 s elapsed time)
    Longest build steps for .so (linking):
           3.6 weighted s to build _C.abi3.so (3.6 s elapsed time)
    1175.8 s weighted time (4289.0 s elapsed time sum, 3.6x parallelism)
    28 build steps completed, average of 0.02/s

==============================================
==   TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=12   ==
==============================================
    Longest build steps for .cu.o:
          20.9 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (250.7 s elapsed time)
          21.0 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (251.6 s elapsed time)
          21.0 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (251.9 s elapsed time)
          21.6 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (259.2 s elapsed time)
          21.6 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (259.7 s elapsed time)
          21.9 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (262.7 s elapsed time)
          22.1 weighted s to build CMakeFiles/_C.dir/csrc/quantization/machete/generated/machete_mm_... (265.0 s elapsed time)
          32.5 weighted s to build CMakeFiles/_C.dir/csrc/quantization/cutlass_w8a8/scaled_mm_c3x.cu... (389.9 s elapsed time)
          52.2 weighted s to build CMakeFiles/_C.dir/csrc/attention/attention_kernels.cu.o (626.7 s elapsed time)
         606.2 weighted s to build CMakeFiles/_C.dir/csrc/quantization/gptq_marlin/gptq_marlin.cu.o (1108.9 s elapsed time)
    Longest build steps for .so (linking):
           0.1 weighted s to build _moe_C.abi3.so (1.1 s elapsed time)
           0.5 weighted s to build _deps/vllm-flash-attn-build/vllm_flash_attn_c.abi3.so (1.0 s elapsed time)
           6.4 weighted s to build _C.abi3.so (6.4 s elapsed time)
    1267.9 s weighted time (8918.5 s elapsed time sum, 7.0x parallelism)
    124 build steps completed, average of 0.10/s

LucasWilkinson avatar Sep 26 '24 08:09 LucasWilkinson

👋 Hi! Thank you for contributing to the vLLM project. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

  • Add ready label to the PR
  • Enable auto-merge.

🚀

github-actions[bot] avatar Sep 26 '24 08:09 github-actions[bot]

/ready

LucasWilkinson avatar Oct 01 '24 20:10 LucasWilkinson

LGTM! This is a nice refactor ! Thanks @LucasWilkinson