diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 49bc3c32b6..e4ae7f897e 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -829,9 +829,7 @@ steps: - examples/features/data_parallel/data_parallel_offline.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py - - csrc/custom_quickreduce.cu - - csrc/ops.h - - csrc/torch_bindings.cpp + - csrc/rocm/quickreduce/ - vllm/model_executor/layers/ - vllm/entrypoints/llm.py - vllm/config/parallel.py @@ -1514,7 +1512,7 @@ steps: parallelism: 5 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - csrc/rocm/ - tests/kernels/moe @@ -1560,6 +1558,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention - vllm/model_executor/layers/attention - tests/kernels/attention @@ -1585,10 +1584,10 @@ steps: - vllm/platforms/rocm.py - tests/kernels/core/test_minimax_reduce_rms.py - tests/utils.py - - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu - - csrc/libtorch_stable/minimax_reduce_rms_kernel.h - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - csrc/ops.h + - csrc/torch_bindings.cpp - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py optional: true commands: @@ -1610,9 +1609,9 @@ steps: - vllm/models/kimi_k3/nvidia/ops/third_party/kda/ - vllm/models/kimi_k3/amd/ - CMakeLists.txt - - csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu + - csrc/ops.h + - csrc/torch_bindings.cpp - tests/models/kimi_k3/test_amd_attn_res.py - tests/models/kimi_k3/test_amd_kda_decode.py - tests/models/kimi_k3/test_kda.py @@ -1633,7 +1632,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/models/kimi_k3/ - - csrc/libtorch_stable/kimi_k3/ + - csrc/kimi_k3/ - tests/models/kimi_k3/ - vllm/platforms/rocm.py commands: @@ -3512,9 +3511,7 @@ steps: optional: true working_dir: "/vllm-workspace/" source_file_dependencies: - - csrc/custom_quickreduce.cu - - csrc/ops.h - - csrc/torch_bindings.cpp + - csrc/rocm/quickreduce/ - vllm/distributed/ - vllm/model_executor/layers/ - vllm/entrypoints/llm.py @@ -4033,7 +4030,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ @@ -4604,10 +4601,10 @@ steps: - vllm/platforms/rocm.py - tests/kernels/core/test_minimax_reduce_rms.py - tests/utils.py - - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu - - csrc/libtorch_stable/minimax_reduce_rms_kernel.h - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - csrc/ops.h + - csrc/torch_bindings.cpp - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py optional: true commands: @@ -4622,7 +4619,7 @@ steps: source_file_dependencies: - csrc/quantization/fp4/ - csrc/attention/mla/ - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/model_executor/layers/fused_moe/cutlass_moe.py - vllm/v1/attention/backends/triton_attn.py - vllm/v1/attention/backends/rocm_attn.py @@ -4674,7 +4671,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu + - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - vllm/platforms/rocm.py @@ -4690,6 +4687,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention - vllm/model_executor/layers/attention - tests/kernels/attention @@ -4709,7 +4707,7 @@ steps: parallelism: 5 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - csrc/rocm/ - tests/kernels/moe @@ -4759,7 +4757,7 @@ steps: - vllm/utils/torch_utils.py - vllm/platforms/ - vllm/_aiter_ops.py - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/distributed/device_communicators/ - vllm/config commands: diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index f83946386a..0b1a73136b 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -510,9 +510,7 @@ steps: - examples/features/data_parallel/data_parallel_offline.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py - - csrc/custom_quickreduce.cu - - csrc/ops.h - - csrc/torch_bindings.cpp + - csrc/rocm/quickreduce/ - vllm/model_executor/layers/ - vllm/entrypoints/llm.py - vllm/config/parallel.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index aa0971b4c5..472c19e74a 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -72,10 +72,10 @@ steps: num_devices: 2 device: h100 source_file_dependencies: - - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu - - csrc/libtorch_stable/minimax_reduce_rms_kernel.h - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - csrc/ops.h + - csrc/torch_bindings.cpp - vllm/model_executor/layers/minimax_rms_norm/ - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py - tests/kernels/core/test_minimax_reduce_rms.py @@ -122,7 +122,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: - - csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu + - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - vllm/platforms/rocm.py @@ -316,6 +316,7 @@ steps: timeout_in_minutes: 30 source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention # TODO: remove this dependency (https://github.com/vllm-project/vllm/issues/32267) - vllm/model_executor/layers/attention @@ -335,6 +336,7 @@ steps: - image-build-amd source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention - vllm/model_executor/layers/attention - tests/kernels/attention @@ -467,7 +469,7 @@ steps: num_devices: 1 timeout_in_minutes: 50 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -485,7 +487,7 @@ steps: device: mi300_1 timeout_in_minutes: 55 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -576,11 +578,10 @@ steps: # optional: true source_file_dependencies: - csrc/quantization/fp4/ - - csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu - - csrc/libtorch_stable/quantization/fp4/ - - csrc/libtorch_stable/quantization/w8a8/cutlass/ + - csrc/nvfp4_kv_cache_kernels.cu + - csrc/quantization/w8a8/cutlass/ - csrc/attention/mla/ - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/model_executor/layers/fusion/ - vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py - vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -606,9 +607,9 @@ steps: - vllm/config/attention.py - vllm/v1/attention/backends/registry.py - vllm/_custom_ops.py - - csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu + - csrc/ops.h + - csrc/torch_bindings.cpp - tests/kernels/attention/test_flashinfer_mla_decode.py - tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py - tests/kernels/attention/test_cache.py @@ -853,7 +854,7 @@ steps: device: h100 num_devices: 2 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -896,7 +897,7 @@ steps: device: b200-k8s num_devices: 2 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -1009,7 +1010,7 @@ steps: source_file_dependencies: - csrc/quantization/fp4/ - csrc/attention/mla/ - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/model_executor/layers/fused_moe/cutlass_moe.py - vllm/v1/attention/backends/triton_attn.py - vllm/v1/attention/backends/rocm_attn.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 0660795105..f6a4a89f11 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -230,7 +230,7 @@ steps: - vllm/v1/attention/backends/mla/ - vllm/model_executor/layers/attention/sparse_mla_attention.py - vllm/v1/worker/gpu/ - - csrc/libtorch_stable/hisparse_kernels.cu + - csrc/hisparse_kernels.cu commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hisparse.txt @@ -1185,7 +1185,7 @@ steps: optional: true num_devices: 4 source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ @@ -1203,7 +1203,7 @@ steps: optional: true num_devices: 4 source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ @@ -1224,7 +1224,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 288e7621ca..552910e85a 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -132,7 +132,7 @@ steps: device: b200-k8s source_file_dependencies: - vllm/models/kimi_k3/ - - csrc/libtorch_stable/kimi_k3/ + - csrc/kimi_k3/ - tests/models/kimi_k3/ - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py commands: diff --git a/.gitignore b/.gitignore index 46418c0392..57d0d54fc8 100644 --- a/.gitignore +++ b/.gitignore @@ -237,9 +237,6 @@ benchmarks/**/*.json actionlint shellcheck*/ -# Ignore moe/marlin_moe gen code -csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_* - # Ignore ep_kernels_workspace folder ep_kernels_workspace/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8e257edccc..78a672c110 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,7 +21,7 @@ repos: rev: v21.1.2 hooks: - id: clang-format - exclude: 'csrc/libtorch_stable/moe/topk_softmax_kernels.cu|vllm/third_party/.*' + exclude: 'csrc/moe/topk_softmax_kernels.cu|vllm/third_party/.*' types_or: [c++, cuda] args: [--style=file, --verbose] - repo: https://github.com/DavidAnson/markdownlint-cli2 diff --git a/CMakeLists.txt b/CMakeLists.txt index b4dd9ada5e..3a71d7e1e5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -363,8 +363,8 @@ endif() if(VLLM_GPU_LANG STREQUAL "HIP") set(VLLM_EXT_SRC - "csrc/torch_bindings.cpp" - "csrc/custom_quickreduce.cu") + "csrc/rocm/quickreduce/torch_bindings.cpp" + "csrc/rocm/quickreduce/custom_quickreduce.cu") message(STATUS "Enabling C extension.") define_extension_target( @@ -374,16 +374,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP") SOURCES ${VLLM_EXT_SRC} COMPILE_FLAGS ${VLLM_GPU_FLAGS} ARCHITECTURES ${VLLM_GPU_ARCHES} - INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR} - INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR} USE_SABI 3 WITH_SOABI) - - # If CUTLASS is compiled on NVCC >= 12.5, it by default uses - # cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the - # driver API. This causes problems when linking with earlier versions of CUDA. - # Setting this variable sidesteps the issue by calling the driver directly. - target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) endif() # _C HIP endif if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") @@ -419,36 +411,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() set(VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/torch_bindings.cpp" - "csrc/libtorch_stable/cuda_view.cu" - "csrc/libtorch_stable/cuda_utils_kernels.cu" - "csrc/libtorch_stable/activation_kernels.cu" - "csrc/libtorch_stable/ngram_embedding_kernels.cu" - "csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu" - "csrc/libtorch_stable/quantization/activation_kernels.cu" - "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu" - "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu" - "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu" - "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu" - "csrc/libtorch_stable/quantization/gptq/q_gemm.cu" - "csrc/libtorch_stable/pos_encoding_kernels.cu" - "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu" - "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu" - "csrc/libtorch_stable/layernorm_kernels.cu" - "csrc/libtorch_stable/layernorm_quant_kernels.cu" - "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu" - "csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu" - "csrc/libtorch_stable/attention/merge_attn_states.cu" - "csrc/libtorch_stable/sampler.cu" - "csrc/libtorch_stable/topk.cu" - "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" - "csrc/libtorch_stable/cache_kernels.cu" - "csrc/libtorch_stable/cache_kernels_fused.cu" - "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu" - "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp" - "csrc/libtorch_stable/custom_all_reduce.cu" - "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu" - "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu") + "csrc/torch_bindings.cpp" + "csrc/cuda_view.cu" + "csrc/cuda_utils_kernels.cu" + "csrc/activation_kernels.cu" + "csrc/ngram_embedding_kernels.cu" + "csrc/vocab_parallel_embedding_kernels.cu" + "csrc/quantization/activation_kernels.cu" + "csrc/quantization/w8a8/int8/scaled_quant.cu" + "csrc/quantization/w8a8/fp8/common.cu" + "csrc/quantization/w8a8/fp8/per_token_group_quant.cu" + "csrc/quantization/gptq/q_gemm.cu" + "csrc/pos_encoding_kernels.cu" + "csrc/fused_qknorm_rope_kernel.cu" + "csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu" + "csrc/layernorm_kernels.cu" + "csrc/layernorm_quant_kernels.cu" + "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu" + "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu" + "csrc/attention/merge_attn_states.cu" + "csrc/sampler.cu" + "csrc/topk.cu" + "csrc/mamba/selective_scan_fwd.cu" + "csrc/cache_kernels.cu" + "csrc/cache_kernels_fused.cu" + "csrc/custom_all_gather_reduce_scatter.cu" + "csrc/custom_all_gather_reduce_scatter_ops.cpp" + "csrc/custom_all_reduce.cu" + "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu" + "csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu") if(VLLM_GPU_LANG STREQUAL "CUDA" AND DEFINED CMAKE_CUDA_COMPILER_VERSION AND @@ -470,10 +461,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/hisparse_kernels.cu" - "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu" - "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu" - "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu") + "csrc/hisparse_kernels.cu" + "csrc/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu" + "csrc/attention/dcp_utils/dcp_direct_kv_gather.cu" + "csrc/attention/dcp_utils/dcp_direct_q_gather.cu") SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") @@ -508,12 +499,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") FetchContent_MakeAvailable(cutlass) list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/cutlass_extensions/common.cpp" - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu" - "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu" - "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu") + "csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu" + "csrc/quantization/fp4/nvfp4_quant_entry.cu" + "csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu" + "csrc/quantization/awq/gemm_kernels.cu" + "csrc/minimax_reduce_rms_kernel.cu") # # Machete kernels @@ -527,7 +517,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # preselected input type pairs and schedules. # Generate sources: set(MACHETE_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/machete/generate.py) + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/machete/generate.py) file(MD5 ${MACHETE_GEN_SCRIPT} MACHETE_GEN_SCRIPT_HASH) message(STATUS "Machete generation script hash: ${MACHETE_GEN_SCRIPT_HASH}") @@ -560,7 +550,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() # Add machete generated sources - file(GLOB MACHETE_GEN_SOURCES "csrc/libtorch_stable/quantization/machete/generated/*.cu") + file(GLOB MACHETE_GEN_SOURCES "csrc/quantization/machete/generated/*.cu") list(APPEND VLLM_STABLE_EXT_SRC ${MACHETE_GEN_SOURCES}) # forward compatible @@ -569,7 +559,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") CUDA_ARCHS "${MACHETE_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC - csrc/libtorch_stable/quantization/machete/machete_pytorch.cu) + csrc/quantization/machete/machete_pytorch.cu) message(STATUS "Building Machete kernels for archs: ${MACHETE_ARCHS}") else() if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 @@ -587,7 +577,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # The fused MiniMax-M3 KV insert quantizes NVFP4 with SM100-family-only # cvt instructions, so its SM10x/11x code is built for the family targets. set(MINIMAX_M3_FUSED_SRC - "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu") + "csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu") list(REMOVE_ITEM VLLM_STABLE_EXT_SRC "${MINIMAX_M3_FUSED_SRC}") set_gencode_flags_for_srcs( SRCS "${VLLM_STABLE_EXT_SRC}" @@ -611,9 +601,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(COOPERATIVE_TOPK_ARCHS) list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/cooperative_topk.cu") + "csrc/cooperative_topk.cu") set_gencode_flags_for_srcs( - SRCS "csrc/libtorch_stable/cooperative_topk.cu" + SRCS "csrc/cooperative_topk.cu" CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}") endif() @@ -656,7 +646,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # preselected input type pairs and schedules. # Generate sources: set(MARLIN_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/marlin/generate_kernels.py) + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py) file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH) list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") @@ -691,7 +681,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if (MARLIN_ARCHS) - file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_float16.cu") + file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_float16.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_ARCHS}") @@ -701,7 +691,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC}) - file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_bfloat16.cu") + file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_bfloat16.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_BF16_ARCHS}") @@ -713,7 +703,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if (MARLIN_SM75_ARCHS) - file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm75_kernel_*.cu") + file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_SM75_ARCHS}") @@ -725,7 +715,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if (MARLIN_FP8_ARCHS) - file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm89_kernel_*.cu") + file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_FP8_ARCHS}") @@ -737,10 +727,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() set(MARLIN_SRCS - "csrc/libtorch_stable/quantization/marlin/marlin.cu" - "csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu" - "csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu" - "csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu") + "csrc/quantization/marlin/marlin.cu" + "csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu" + "csrc/quantization/marlin/gptq_marlin_repack.cu" + "csrc/quantization/marlin/awq_marlin_repack.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_SRCS}" CUDA_ARCHS "${MARLIN_OTHER_ARCHS}") @@ -763,7 +753,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS) - set(DSV3_FUSED_A_GEMM_SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu") + set(DSV3_FUSED_A_GEMM_SRCS "csrc/dsv3_fused_a_gemm.cu") set_gencode_flags_for_srcs( SRCS "${DSV3_FUSED_A_GEMM_SRCS}" CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}") @@ -778,8 +768,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS) set(FP32_ROUTER_GEMM_SRCS - "csrc/libtorch_stable/fp32_router_gemm_entry.cu" - "csrc/libtorch_stable/fp32_router_gemm.cu") + "csrc/fp32_router_gemm_entry.cu" + "csrc/fp32_router_gemm.cu") set_gencode_flags_for_srcs( SRCS "${FP32_ROUTER_GEMM_SRCS}" CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}") @@ -799,11 +789,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS) set(SCALED_MM_SM90_SRCS - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu") + "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu") set_gencode_flags_for_srcs( SRCS "${SCALED_MM_SM90_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") @@ -834,9 +824,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) set(SCALED_MM_SM120_SRCS - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu" + "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu" ) set_gencode_flags_for_srcs( SRCS "${SCALED_MM_SM120_SRCS}" @@ -868,9 +858,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) set(SCALED_MM_SM100_SRCS - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu" + "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu" ) set_gencode_flags_for_srcs( SRCS "${SCALED_MM_SM100_SRCS}" @@ -901,7 +891,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # subtract out the archs that are already built for 3x list(REMOVE_ITEM SCALED_MM_2X_ARCHS ${SCALED_MM_3X_ARCHS}) if (SCALED_MM_2X_ARCHS) - set(SCALED_MM_C2X_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu") + set(SCALED_MM_C2X_SRCS "csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu") set_gencode_flags_for_srcs( SRCS "${SCALED_MM_C2X_SRCS}" CUDA_ARCHS "${SCALED_MM_2X_ARCHS}") @@ -927,7 +917,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # if it's possible to compile MoE kernels that use its output. cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS) - set(CUTLASS_MOE_SM90_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu") + set(CUTLASS_MOE_SM90_SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MOE_SM90_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") @@ -951,7 +941,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) - set(CUTLASS_MOE_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu") + set(CUTLASS_MOE_SM100_SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MOE_SM100_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") @@ -976,7 +966,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS) - set(CUTLASS_MOE_DATA_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu") + set(CUTLASS_MOE_DATA_SRCS "csrc/quantization/w8a8/cutlass/moe/moe_data.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MOE_DATA_SRCS}" CUDA_ARCHS "${CUTLASS_MOE_DATA_ARCHS}") @@ -1008,12 +998,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS) set(FP4_SM120_SRCS - "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" - "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") + "csrc/quantization/fp4/nvfp4_quant_kernels.cu" + "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" + "csrc/quantization/fp4/nvfp4_experts_quant.cu" + "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu" + "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/nvfp4_kv_cache_kernels.cu") set_gencode_flags_for_srcs( SRCS "${FP4_SM120_SRCS}" CUDA_ARCHS "${FP4_SM120_ARCHS}") @@ -1034,15 +1024,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS) set(FP4_SM100_SRCS - "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" - "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu" - "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu" - "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu" - "csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu") + "csrc/quantization/fp4/nvfp4_quant_kernels.cu" + "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" + "csrc/quantization/fp4/nvfp4_experts_quant.cu" + "csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu" + "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/quantization/fp4/mxfp4_experts_quant.cu" + "csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu" + "csrc/nvfp4_kv_cache_kernels.cu" + "csrc/nvfp4_ds_mla_cache_kernels.cu") if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) message(STATUS "Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).") @@ -1066,9 +1056,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS) set(W4A8_SRCS - "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu" - "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu" - "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu" + "csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu" + "csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu" + "csrc/quantization/cutlass_w4a8/w4a8_utils.cu" ) set_gencode_flags_for_srcs( @@ -1101,7 +1091,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS) set(CUTLASS_MLA_SRCS - "csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu") + "csrc/attention/mla/sm100_cutlass_mla_kernel.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MLA_SRCS}" CUDA_ARCHS "${MLA_ARCHS}") @@ -1125,7 +1115,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(FUSED_KDA_DECODE_ARCHS) set(FUSED_KDA_DECODE_SRC - "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu") + "csrc/kimi_k3/fused_kda_decode_kernel.cu") set_gencode_flags_for_srcs( SRCS "${FUSED_KDA_DECODE_SRC}" CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}") @@ -1137,7 +1127,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(FUSED_GDN_DECODE_ARCHS) set(FUSED_GDN_DECODE_SRC - "csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu") + "csrc/gdn/fused_gdn_decode_kernel.cu") set_gencode_flags_for_srcs( SRCS "${FUSED_GDN_DECODE_SRC}" CUDA_ARCHS "${FUSED_GDN_DECODE_ARCHS}") @@ -1154,7 +1144,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(KIMI_K3_ATTN_RES_ARCHS) set(KIMI_K3_ATTN_RES_SRC - "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu") + "csrc/kimi_k3/attn_res_kernel.cu") set_gencode_flags_for_srcs( SRCS "${KIMI_K3_ATTN_RES_SRC}" CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}") @@ -1169,7 +1159,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # Hadacore kernels cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}") if(HADACORE_ARCHS) - set(HADACORE_SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu") + set(HADACORE_SRCS "csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu") set_gencode_flags_for_srcs( SRCS "${HADACORE_SRCS}" CUDA_ARCHS "${HADACORE_ARCHS}") @@ -1193,7 +1183,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") list(FILTER FUSED_KDA_DECODE_HIP_ARCHS INCLUDE REGEX "gfx942|gfx950") if(FUSED_KDA_DECODE_HIP_ARCHS) set(FUSED_KDA_DECODE_HIP_SRC - "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu") + "csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu") list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_HIP_SRC}") message(STATUS "Building fused KDA decode for HIP archs: ${FUSED_KDA_DECODE_HIP_ARCHS}") @@ -1203,7 +1193,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") list(FILTER FUSED_KDA_CHUNK_HIP_ARCHS INCLUDE REGEX "gfx950") if(FUSED_KDA_CHUNK_HIP_ARCHS) set(FUSED_KDA_CHUNK_HIP_SRC - "csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu") + "csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu") list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_CHUNK_HIP_SRC}") message(STATUS "Building fused KDA chunk for HIP archs: ${FUSED_KDA_CHUNK_HIP_ARCHS}") @@ -1289,21 +1279,21 @@ endif() # set(VLLM_MOE_EXT_SRC - "csrc/libtorch_stable/moe/torch_bindings.cpp" - "csrc/libtorch_stable/moe/moe_align_sum_kernels.cu" - "csrc/libtorch_stable/moe/topk_softmax_kernels.cu" - "csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu") + "csrc/moe/torch_bindings.cpp" + "csrc/moe/moe_align_sum_kernels.cu" + "csrc/moe/topk_softmax_kernels.cu" + "csrc/moe/topk_softplus_sqrt_kernels.cu") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC - "csrc/libtorch_stable/moe/moe_wna16.cu" - "csrc/libtorch_stable/moe/grouped_topk_kernels.cu") + "csrc/moe/moe_wna16.cu" + "csrc/moe/grouped_topk_kernels.cu") endif() if(VLLM_GPU_LANG STREQUAL "CUDA") set(MOE_PERMUTE_SRC - "csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu" - "csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu") + "csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu" + "csrc/moe/moe_permute_unpermute_op.cu") list(APPEND VLLM_MOE_EXT_SRC "${MOE_PERMUTE_SRC}") endif() @@ -1314,7 +1304,7 @@ set_gencode_flags_for_srcs( if(VLLM_GPU_LANG STREQUAL "CUDA") set(VLLM_MOE_WNA16_SRC - "csrc/libtorch_stable/moe/moe_wna16.cu") + "csrc/moe/moe_wna16.cu") set_gencode_flags_for_srcs( SRCS "${VLLM_MOE_WNA16_SRC}" @@ -1349,7 +1339,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # preselected input type pairs and schedules. # Generate sources: set(MOE_MARLIN_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py) + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py) file(MD5 ${MOE_MARLIN_GEN_SCRIPT} MOE_MARLIN_GEN_SCRIPT_HASH) list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MOE_MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") @@ -1363,6 +1353,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH=$ENV{PYTHONPATH} ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR} + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/moe/marlin_moe_wna16 RESULT_VARIABLE moe_marlin_generation_result OUTPUT_VARIABLE moe_marlin_generation_output OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log @@ -1384,7 +1375,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_ARCHS) - file(GLOB MARLIN_MOE_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_*.cu") + file(GLOB MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/sm80_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_SRC}" CUDA_ARCHS "${MARLIN_MOE_ARCHS}") @@ -1396,7 +1387,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_SM75_ARCHS) - file(GLOB MARLIN_MOE_SM75_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm75_kernel_*.cu") + file(GLOB MARLIN_MOE_SM75_SRC "csrc/moe/marlin_moe_wna16/sm75_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_SM75_SRC}" CUDA_ARCHS "${MARLIN_MOE_SM75_ARCHS}") @@ -1408,7 +1399,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_FP8_ARCHS) - file(GLOB MARLIN_MOE_FP8_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm89_kernel_*.cu") + file(GLOB MARLIN_MOE_FP8_SRC "csrc/moe/marlin_moe_wna16/sm89_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_FP8_SRC}" CUDA_ARCHS "${MARLIN_MOE_FP8_ARCHS}") @@ -1419,7 +1410,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_FP8_SRC}) endif() - set(MARLIN_MOE_OTHER_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu") + set(MARLIN_MOE_OTHER_SRC "csrc/moe/marlin_moe_wna16/ops.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_OTHER_SRC}" CUDA_ARCHS "${MARLIN_MOE_OTHER_ARCHS}") diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 730b963c23..c4d111e77b 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -497,7 +497,7 @@ set(VLLM_EXT_SRC "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" "csrc/cpu/mamba_cpu.cpp" - "csrc/moe/dynamic_4bit_int_moe_cpu.cpp" + "csrc/cpu/dynamic_4bit_int_moe_cpu.cpp" "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/cpu_isa.cpp" @@ -599,7 +599,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" - "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") + "csrc/cpu/dynamic_4bit_int_moe_cpu.cpp") set(VLLM_EXT_SRC_AVX2 "csrc/cpu/sgl-kernels/fla.cpp" @@ -617,7 +617,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" - "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") + "csrc/cpu/dynamic_4bit_int_moe_cpu.cpp") message(STATUS "CPU extension (AVX512F + BF16 + VNNI + AMX) source files: ${VLLM_EXT_SRC_AVX512} ${VLLM_EXT_SRC_SGL}") message(STATUS "CPU extension (AVX512F) source files: ${VLLM_EXT_SRC_AVX512}") diff --git a/cmake/hipify.py b/cmake/hipify.py index f4932260c9..2cf58a7b7d 100755 --- a/cmake/hipify.py +++ b/cmake/hipify.py @@ -65,7 +65,7 @@ if __name__ == "__main__": project_directory=args.project_dir, output_directory=args.output_dir, # Hipify resolves quoted includes next to the including file first; vLLM - # uses paths relative to csrc/ (e.g. "libtorch_stable/torch_utils.h" + # uses paths relative to csrc/ (e.g. "torch_utils.h" # from quantization/w8a8/fp8/*.cu). Without an include root here, those # headers are never found and are not hipified or rewritten in dependents. header_include_dirs=["."], diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/activation_kernels.cu similarity index 99% rename from csrc/libtorch_stable/activation_kernels.cu rename to csrc/activation_kernels.cu index a7ee87b35f..9e28c3343b 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/csrc/activation_kernels.cu @@ -10,7 +10,7 @@ #include #include -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "async_util.cuh" #include "cuda_vec_utils.cuh" #include "dispatch_utils.h" diff --git a/csrc/libtorch_stable/async_util.cuh b/csrc/async_util.cuh similarity index 100% rename from csrc/libtorch_stable/async_util.cuh rename to csrc/async_util.cuh diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu b/csrc/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu rename to csrc/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh b/csrc/attention/dcp_utils/dcp_direct_common.cuh similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh rename to csrc/attention/dcp_utils/dcp_direct_common.cuh diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu b/csrc/attention/dcp_utils/dcp_direct_kv_gather.cu similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu rename to csrc/attention/dcp_utils/dcp_direct_kv_gather.cu diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu b/csrc/attention/dcp_utils/dcp_direct_q_gather.cu similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu rename to csrc/attention/dcp_utils/dcp_direct_q_gather.cu diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/attention/merge_attn_states.cu similarity index 99% rename from csrc/libtorch_stable/attention/merge_attn_states.cu rename to csrc/attention/merge_attn_states.cu index cc89397f68..3be334fd64 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/attention/merge_attn_states.cu @@ -6,9 +6,8 @@ #include "../dispatch_utils.h" #include -#include "../../attention/attention_dtypes.h" -#include "attention_utils.cuh" -#include "../../quantization/w8a8/fp8/common.cuh" +#include "core/vec_dtypes/vec_dtypes.h" +#include "quantization/w8a8/fp8/common.cuh" namespace vllm { diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp b/csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp rename to csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp b/csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp rename to csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp b/csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp rename to csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp b/csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp rename to csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp diff --git a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu b/csrc/attention/mla/sm100_cutlass_mla_kernel.cu similarity index 99% rename from csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu rename to csrc/attention/mla/sm100_cutlass_mla_kernel.cu index 150e324628..b04c014b98 100644 --- a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu +++ b/csrc/attention/mla/sm100_cutlass_mla_kernel.cu @@ -18,7 +18,7 @@ limitations under the License. * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929 * by Alcanderian JieXin Liang */ -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include diff --git a/csrc/cache.h b/csrc/cache.h deleted file mode 100644 index dbc2675fc7..0000000000 --- a/csrc/cache.h +++ /dev/null @@ -1,93 +0,0 @@ -#pragma once - -#include -#include - -#include -#include - -void swap_blocks(torch::Tensor& src, torch::Tensor& dst, - int64_t block_size_in_bytes, - const torch::Tensor& block_mapping); - -void swap_blocks_batch(const torch::Tensor& src_ptrs, - const torch::Tensor& dst_ptrs, - const torch::Tensor& sizes, - bool is_src_access_order_any); - -void reshape_and_cache(torch::Tensor& key, torch::Tensor& value, - torch::Tensor& key_cache, torch::Tensor& value_cache, - torch::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::Tensor& k_scale, torch::Tensor& v_scale); - -void reshape_and_cache_flash(torch::Tensor& key, torch::Tensor& value, - torch::Tensor& key_cache, - torch::Tensor& value_cache, - torch::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::Tensor& k_scale, torch::Tensor& v_scale); - -void concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe, - torch::Tensor& kv_cache, torch::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::Tensor& scale); - -// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla -void concat_and_cache_mla_rope_fused( - torch::Tensor& positions, torch::Tensor& q_pe, torch::Tensor& k_pe, - torch::Tensor& kv_c, torch::Tensor& rope_cos_sin_cache, bool rope_is_neox, - torch::Tensor& kv_cache_slot_mapping, torch::Tensor& kv_cache, - const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale); - -// Just for unittest -void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache, - const double scale, const std::string& kv_cache_dtype); - -void gather_and_maybe_dequant_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& cu_seq_lens, // [BATCH+1] - torch::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] - int64_t num_tokens, const std::string& kv_cache_dtype, - torch::Tensor const& scale, - std::optional seq_starts = std::nullopt); - -// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. -void cp_gather_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& cu_seq_lens, // [BATCH+1] - int64_t batch_size, std::optional seq_starts = std::nullopt); - -// Gather and upconvert FP8 KV cache to BF16 workspace -void cp_gather_and_upconvert_fp8_kv_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] - torch::Tensor const& dst, // [TOT_TOKENS, 576] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size, std::optional seq_starts = std::nullopt); - -// Indexer K quantization and cache function -void indexer_k_quant_and_cache( - torch::Tensor& k, // [num_tokens, head_dim] - torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::Tensor& slot_mapping, // [num_tokens] - int64_t quant_block_size, // quantization block size - const std::string& scale_fmt); - -// Concatenate query nope and rope for MLA/DSA attention -void concat_mla_q( - torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] - torch::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] - torch::Tensor& q_out); // [num_tokens, num_heads, nope_dim + rope_dim] - -// Extract function to gather quantized K cache -void cp_gather_indexer_k_quant_cache( - const torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::Tensor& dst_k, // [num_tokens, head_dim] - torch::Tensor& dst_scale, // [num_tokens, head_dim / quant_block_size * 4] - const torch::Tensor& block_table, // [batch_size, num_blocks] - const torch::Tensor& cu_seq_lens); // [batch_size + 1] diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/cache_kernels.cu similarity index 99% rename from csrc/libtorch_stable/cache_kernels.cu rename to csrc/cache_kernels.cu index 9130adf87d..d1c1e9df98 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -1,8 +1,8 @@ #include "torch_utils.h" #include "dispatch_utils.h" -#include "../cuda_utils.h" -#include "../cuda_compat.h" +#include "core/utils.hpp" +#include "core/cuda_compat.h" #include "quantization/vectorization_utils.cuh" #include "concat_mla_q.cuh" @@ -12,9 +12,9 @@ #endif #ifdef USE_ROCM - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #include @@ -712,8 +712,7 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel( } __syncthreads(); - for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x)); - iter++) { + for (int iter = 0; iter < div_ceil(batch_size, int(blockDim.x)); iter++) { int tid = iter * blockDim.x + threadIdx.x; if (tid < batch_size) { const int seq_start = cu_seq_lens[tid]; @@ -1229,7 +1228,7 @@ __device__ __forceinline__ bool map_gather_page_task( const int32_t source_begin = seq_starts == nullptr ? 0 : seq_starts[req_id]; const int32_t first_block = source_begin / block_size; const int32_t num_pages = - cuda_utils::ceil_div(source_begin + seq_len, block_size) - first_block; + div_ceil(source_begin + seq_len, block_size) - first_block; if (relative_page < num_pages) { page.req_id = req_id; page.logical_block = first_block + relative_page; @@ -1408,8 +1407,7 @@ void gather_and_maybe_dequant_cache( const int64_t dst_entry_stride = dst.stride(0); const int32_t page_threads = num_tokens >= (1 << 20) ? 128 : 256; const int32_t required_blocks = - cuda_utils::ceil_div(static_cast(num_tokens), block_size) + - 2 * num_reqs; + div_ceil(static_cast(num_tokens), block_size) + 2 * num_reqs; const dim3 grid(required_blocks); const dim3 block(page_threads); @@ -1685,7 +1683,7 @@ void cp_gather_cache( const int32_t num_reqs = static_cast(batch_size); const int32_t required_blocks = - cuda_utils::ceil_div(total_tokens, block_size) + 2 * num_reqs; + div_ceil(total_tokens, block_size) + 2 * num_reqs; const dim3 grid(required_blocks); const dim3 block(256); @@ -1828,8 +1826,8 @@ void cp_gather_and_upconvert_fp8_kv_cache( const int32_t* seq_starts_ptr = seq_starts.has_value() ? seq_starts.value().const_data_ptr() : nullptr; - const int required_blocks = cuda_utils::ceil_div(total_tokens, block_size) + - 2 * static_cast(batch_size); + const int required_blocks = + div_ceil(total_tokens, block_size) + 2 * static_cast(batch_size); vllm::cp_gather_and_upconvert_fp8_kv_cache_page<<< required_blocks, block_size_threads, 0, stream>>>( src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()), diff --git a/csrc/libtorch_stable/cache_kernels_fused.cu b/csrc/cache_kernels_fused.cu similarity index 98% rename from csrc/libtorch_stable/cache_kernels_fused.cu rename to csrc/cache_kernels_fused.cu index 744e3a2dcd..4dc691636d 100644 --- a/csrc/libtorch_stable/cache_kernels_fused.cu +++ b/csrc/cache_kernels_fused.cu @@ -1,13 +1,13 @@ #include "torch_utils.h" #include "dispatch_utils.h" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" -#include "../quantization/w8a8/fp8/common.cuh" +#include "quantization/w8a8/fp8/common.cuh" #ifdef USE_ROCM - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #ifdef USE_ROCM diff --git a/csrc/libtorch_stable/concat_mla_q.cuh b/csrc/concat_mla_q.cuh similarity index 100% rename from csrc/libtorch_stable/concat_mla_q.cuh rename to csrc/concat_mla_q.cuh diff --git a/csrc/libtorch_stable/cooperative_topk.cu b/csrc/cooperative_topk.cu similarity index 100% rename from csrc/libtorch_stable/cooperative_topk.cu rename to csrc/cooperative_topk.cu diff --git a/csrc/libtorch_stable/cooperative_topk.cuh b/csrc/cooperative_topk.cuh similarity index 100% rename from csrc/libtorch_stable/cooperative_topk.cuh rename to csrc/cooperative_topk.cuh diff --git a/csrc/core/batch_invariant.hpp b/csrc/core/batch_invariant.hpp deleted file mode 100644 index 8273bc74b1..0000000000 --- a/csrc/core/batch_invariant.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once -#include -#include - -namespace vllm { - -// vllm_is_batch_invariant(); returns true -// if env VLLM_BATCH_INVARIANT=1 -inline bool vllm_is_batch_invariant() { - static bool cached = []() { - std::string env_key = "VLLM_BATCH_INVARIANT"; - const char* val = std::getenv(env_key.c_str()); - return (val && std::atoi(val) != 0) ? 1 : 0; - }(); - return cached; -} - -} // namespace vllm diff --git a/csrc/cuda_compat.h b/csrc/core/cuda_compat.h similarity index 90% rename from csrc/cuda_compat.h rename to csrc/core/cuda_compat.h index d7d589db62..0eafeff621 100644 --- a/csrc/cuda_compat.h +++ b/csrc/core/cuda_compat.h @@ -54,12 +54,6 @@ struct Utils { __shfl_xor(var, lane_mask, width) #endif -#ifndef USE_ROCM - #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane) -#else - #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane) -#endif - #ifndef USE_ROCM #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \ __shfl_down_sync(uint32_t(-1), var, lane_delta) diff --git a/csrc/core/exception.hpp b/csrc/core/exception.hpp deleted file mode 100644 index f3b2ffaef6..0000000000 --- a/csrc/core/exception.hpp +++ /dev/null @@ -1,3 +0,0 @@ -#pragma once - -#define VLLM_IMPLIES(p, q) (!(p) || (q)) diff --git a/csrc/core/registration.h b/csrc/core/registration.h index 4d0ce1c572..e5396e9a8b 100644 --- a/csrc/core/registration.h +++ b/csrc/core/registration.h @@ -12,11 +12,6 @@ // could be a macro instead of a literal token. #define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE) -// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME -// could be a macro instead of a literal token. -#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \ - TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE) - // REGISTER_EXTENSION allows the shared library to be loaded and initialized // via python's import statement. #define REGISTER_EXTENSION(NAME) \ diff --git a/csrc/libtorch_stable/core/math.hpp b/csrc/core/utils.hpp similarity index 70% rename from csrc/libtorch_stable/core/math.hpp rename to csrc/core/utils.hpp index 6764e1fd60..36b8c2e38e 100644 --- a/csrc/libtorch_stable/core/math.hpp +++ b/csrc/core/utils.hpp @@ -1,7 +1,8 @@ #pragma once #include -#include +#include +#include inline constexpr uint32_t next_pow_2(uint32_t const num) { if (num <= 1) return num; @@ -26,3 +27,16 @@ template inline constexpr T round_to_next_multiple_of(T a, T b) { return a % b == 0 ? a : ((a / b) + 1) * b; } + +namespace vllm { + +// Returns true if env VLLM_BATCH_INVARIANT=1 +inline bool vllm_is_batch_invariant() { + static bool cached = []() { + const char* val = std::getenv("VLLM_BATCH_INVARIANT"); + return val && std::atoi(val) != 0; + }(); + return cached; +} + +} // namespace vllm diff --git a/csrc/attention/dtype_bfloat16.cuh b/csrc/core/vec_dtypes/dtype_bfloat16.cuh similarity index 99% rename from csrc/attention/dtype_bfloat16.cuh rename to csrc/core/vec_dtypes/dtype_bfloat16.cuh index 97a25baa1f..5d01a51e76 100644 --- a/csrc/attention/dtype_bfloat16.cuh +++ b/csrc/core/vec_dtypes/dtype_bfloat16.cuh @@ -20,7 +20,7 @@ */ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include "dtype_float32.cuh" #ifndef USE_ROCM diff --git a/csrc/attention/dtype_float16.cuh b/csrc/core/vec_dtypes/dtype_float16.cuh similarity index 99% rename from csrc/attention/dtype_float16.cuh rename to csrc/core/vec_dtypes/dtype_float16.cuh index 3a1815f0ed..6883d521d2 100644 --- a/csrc/attention/dtype_float16.cuh +++ b/csrc/core/vec_dtypes/dtype_float16.cuh @@ -20,7 +20,7 @@ */ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include "dtype_float32.cuh" #ifdef USE_ROCM diff --git a/csrc/attention/dtype_float32.cuh b/csrc/core/vec_dtypes/dtype_float32.cuh similarity index 99% rename from csrc/attention/dtype_float32.cuh rename to csrc/core/vec_dtypes/dtype_float32.cuh index 7c6a686db3..33f78c492e 100644 --- a/csrc/attention/dtype_float32.cuh +++ b/csrc/core/vec_dtypes/dtype_float32.cuh @@ -20,7 +20,7 @@ */ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include diff --git a/csrc/attention/dtype_fp8.cuh b/csrc/core/vec_dtypes/dtype_fp8.cuh similarity index 88% rename from csrc/attention/dtype_fp8.cuh rename to csrc/core/vec_dtypes/dtype_fp8.cuh index 3d56859d8f..6d3c35e978 100644 --- a/csrc/attention/dtype_fp8.cuh +++ b/csrc/core/vec_dtypes/dtype_fp8.cuh @@ -1,7 +1,8 @@ #pragma once -#include "attention_generic.cuh" -#include "torch_utils.h" +#include "dtype_generic.cuh" + +#include #include #ifdef ENABLE_FP8 @@ -31,7 +32,7 @@ inline Fp8KVCacheDataType get_fp8_kv_cache_data_type( } else if (dtype_str == "fp8_e5m2") { return Fp8KVCacheDataType::kFp8E5M2; } - TORCH_UTILS_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str); + STD_TORCH_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str); } // fp8 vector types for quantization of kv cache diff --git a/csrc/attention/attention_generic.cuh b/csrc/core/vec_dtypes/dtype_generic.cuh similarity index 100% rename from csrc/attention/attention_generic.cuh rename to csrc/core/vec_dtypes/dtype_generic.cuh diff --git a/csrc/attention/attention_dtypes.h b/csrc/core/vec_dtypes/vec_dtypes.h similarity index 79% rename from csrc/attention/attention_dtypes.h rename to csrc/core/vec_dtypes/vec_dtypes.h index 64f86381d9..41a7b91ca4 100644 --- a/csrc/attention/attention_dtypes.h +++ b/csrc/core/vec_dtypes/vec_dtypes.h @@ -1,6 +1,6 @@ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include "dtype_float16.cuh" #include "dtype_float32.cuh" #include "dtype_bfloat16.cuh" diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index 0666ee413f..0290c7ab60 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -14,7 +14,7 @@ namespace cpu_attention { enum class ISA { AMX, VEC, VEC16, NEON, VXE, RVV, VSX, AMX_FP8 }; -// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly. +// Mirrors csrc/core/vec_dtypes/dtype_fp8.cuh Fp8KVCacheDataType exactly. enum class Fp8KVCacheDataType { kAuto = 0, kFp8E4M3 = 1, diff --git a/csrc/moe/dynamic_4bit_int_moe_cpu.cpp b/csrc/cpu/dynamic_4bit_int_moe_cpu.cpp similarity index 100% rename from csrc/moe/dynamic_4bit_int_moe_cpu.cpp rename to csrc/cpu/dynamic_4bit_int_moe_cpu.cpp diff --git a/csrc/cpu/ops.h b/csrc/cpu/ops.h new file mode 100644 index 0000000000..ead20b01ae --- /dev/null +++ b/csrc/cpu/ops.h @@ -0,0 +1,46 @@ +#pragma once + +#include +#include +#include +#include +#include + +void rms_norm(torch::Tensor& out, torch::Tensor& input, + std::optional weight, double epsilon); + +void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, + std::optional weight, double epsilon); + +void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, + std::optional key, int64_t head_size, + torch::Tensor& cos_sin_cache, bool is_neox, + int64_t rope_dim_offset, bool inverse); + +void silu_and_mul(torch::Tensor& out, torch::Tensor& input); + +void gelu_and_mul(torch::Tensor& out, torch::Tensor& input); + +void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input); + +void gelu_tanh(torch::Tensor& out, torch::Tensor& input); + +void gelu_new(torch::Tensor& out, torch::Tensor& input); + +void gelu_fast(torch::Tensor& out, torch::Tensor& input); + +void gelu_quick(torch::Tensor& out, torch::Tensor& input); + +void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, + torch::Tensor const& scale, + std::optional const& azp); + +void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, + torch::Tensor& scales, + std::optional const& azp); + +torch::Tensor dynamic_4bit_int_moe_cpu( + torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, + torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, + int64_t intermediate_size, int64_t group_size, + bool apply_router_weight_on_input, int64_t activation_kind); diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 2cb1be095d..f39cf31c1f 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -1,4 +1,4 @@ -#include "cache.h" +#include #include "ops.h" #include "core/registration.h" diff --git a/csrc/libtorch_stable/cub_helpers.h b/csrc/cub_helpers.h similarity index 100% rename from csrc/libtorch_stable/cub_helpers.h rename to csrc/cub_helpers.h diff --git a/csrc/cuda_utils.h b/csrc/cuda_utils.h index 6e62ea208d..bca7c75bc1 100644 --- a/csrc/cuda_utils.h +++ b/csrc/cuda_utils.h @@ -1,21 +1,8 @@ #pragma once +#include #include -#if defined(__HIPCC__) - #define HOST_DEVICE_INLINE __host__ __device__ - #define DEVICE_INLINE __device__ - #define HOST_INLINE __host__ -#elif defined(__CUDACC__) || defined(_NVHPC_CUDA) - #define HOST_DEVICE_INLINE __host__ __device__ __forceinline__ - #define DEVICE_INLINE __device__ __forceinline__ - #define HOST_INLINE __host__ __forceinline__ -#else - #define HOST_DEVICE_INLINE inline - #define DEVICE_INLINE inline - #define HOST_INLINE inline -#endif - #define CUDA_CHECK(cmd) \ do { \ cudaError_t e = cmd; \ @@ -29,13 +16,3 @@ int64_t get_device_attribute(int64_t attribute, int64_t device_id); int64_t get_max_shared_memory_per_block_device_attribute(int64_t device_id); - -namespace cuda_utils { - -template -HOST_DEVICE_INLINE constexpr std::enable_if_t, T> -ceil_div(T a, T b) { - return (a + b - 1) / b; -} - -}; // namespace cuda_utils \ No newline at end of file diff --git a/csrc/libtorch_stable/cuda_utils_kernels.cu b/csrc/cuda_utils_kernels.cu similarity index 100% rename from csrc/libtorch_stable/cuda_utils_kernels.cu rename to csrc/cuda_utils_kernels.cu diff --git a/csrc/libtorch_stable/cuda_vec_utils.cuh b/csrc/cuda_vec_utils.cuh similarity index 100% rename from csrc/libtorch_stable/cuda_vec_utils.cuh rename to csrc/cuda_vec_utils.cuh diff --git a/csrc/libtorch_stable/cuda_view.cu b/csrc/cuda_view.cu similarity index 100% rename from csrc/libtorch_stable/cuda_view.cu rename to csrc/cuda_view.cu diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu b/csrc/custom_all_gather_reduce_scatter.cu similarity index 100% rename from csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu rename to csrc/custom_all_gather_reduce_scatter.cu diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp b/csrc/custom_all_gather_reduce_scatter_ops.cpp similarity index 100% rename from csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp rename to csrc/custom_all_gather_reduce_scatter_ops.cpp diff --git a/csrc/libtorch_stable/custom_all_reduce.cu b/csrc/custom_all_reduce.cu similarity index 100% rename from csrc/libtorch_stable/custom_all_reduce.cu rename to csrc/custom_all_reduce.cu diff --git a/csrc/custom_collective_common.cuh b/csrc/custom_collective_common.cuh index 353dcd07e8..8c6a98149d 100644 --- a/csrc/custom_collective_common.cuh +++ b/csrc/custom_collective_common.cuh @@ -216,21 +216,6 @@ DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; } -template -DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, - int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; - st_flag_release(peer_counter_ptr, flag); - while (ld_flag_acquire(self_counter_ptr) != flag); - } - __syncthreads(); - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - // This function is meant to be used as the second or the final // synchronization barrier in the all reduce kernel. If it's the final // synchronization barrier, we don't need to make any visibility guarantees @@ -279,22 +264,6 @@ DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; } -template -DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, - int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], - flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM); - while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], - __ATOMIC_ACQUIRE, - __MEMORY_SCOPE_DEVICE) < flag); - } - __syncthreads(); - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - template DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { __syncthreads(); diff --git a/csrc/libtorch_stable/cutlass_extensions/common.hpp b/csrc/cutlass_extensions/common.hpp similarity index 89% rename from csrc/libtorch_stable/cutlass_extensions/common.hpp rename to csrc/cutlass_extensions/common.hpp index ca96b0ef3f..0a6fbd8384 100644 --- a/csrc/libtorch_stable/cutlass_extensions/common.hpp +++ b/csrc/cutlass_extensions/common.hpp @@ -25,7 +25,15 @@ inline int get_cuda_max_shared_memory_per_block_opt_in(int const device) { return max_shared_mem_per_block_opt_in; } -int32_t get_sm_version_num(); +inline int32_t get_sm_version_num() { + int32_t major_capability, minor_capability; + cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, + 0); + cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, + 0); + int32_t version_num = major_capability * 10 + minor_capability; + return version_num; +} /** * A wrapper for a kernel that is used to guard against compilation on @@ -110,21 +118,6 @@ struct enable_sm100_to_sm120 : Kernel { } }; -template -struct enable_sm120_only : Kernel { - template - CUTLASS_DEVICE void operator()(Args&&... args) { -#if defined __CUDA_ARCH__ - #if __CUDA_ARCH__ == 1200 - Kernel::operator()(std::forward(args)...); - #else - printf("This kernel only supports sm120a.\n"); - asm("trap;"); - #endif -#endif - } -}; - // SM12x family includes SM120 (RTX 5090) and SM121 (DGX Spark GB10) template struct enable_sm120_family : Kernel { diff --git a/csrc/libtorch_stable/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp similarity index 100% rename from csrc/libtorch_stable/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp rename to csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp diff --git a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp similarity index 100% rename from csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp rename to csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp diff --git a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp similarity index 93% rename from csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp rename to csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp index c2ddcea6d0..30f4a547f8 100644 --- a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp +++ b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp @@ -1,15 +1,9 @@ #pragma once -#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp" -#include "cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp" +#include "broadcast_load_epilogue_c3x.hpp" +#include "broadcast_load_epilogue_array_c3x.hpp" -// This header is shared by both _C (unstable ABI) and _C_stable_libtorch -// (stable ABI) targets. When compiled under the stable ABI target, -// TORCH_TARGET_VERSION is defined and Tensor is unavailable, so we -// use torch::stable::Tensor instead. -#ifdef TORCH_TARGET_VERSION - #include -#endif +#include /* This file defines custom epilogues for fusing channel scales, token scales, @@ -23,38 +17,10 @@ namespace vllm::c3x { -#ifdef TORCH_TARGET_VERSION using TensorType = torch::stable::Tensor; -#else -using TensorType = torch::Tensor; -#endif using namespace cute; -template -struct identity { - CUTLASS_HOST_DEVICE - T operator()(T lhs) const { return lhs; } -}; - -template -struct TrivialEpilogue { - private: - using Accum = cutlass::epilogue::fusion::Sm90AccFetch; - using Compute = cutlass::epilogue::fusion::Sm90Compute< - cutlass::epilogue::thread::Identity, ElementD, ElementAcc, - cutlass::FloatRoundStyle::round_to_nearest>; - - public: - using EVTCompute = cutlass::epilogue::fusion::Sm90EVT; - using ArgumentType = typename EVTCompute::Arguments; - - template - static ArgumentType prepare_args(Args... args) { - return {}; - } -}; - /* * This class provides the common load descriptors for the * ScaledEpilogue[...] classes diff --git a/csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp b/csrc/cutlass_extensions/torch_utils.hpp similarity index 76% rename from csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp rename to csrc/cutlass_extensions/torch_utils.hpp index 5f973033fd..c5b9d9591e 100644 --- a/csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp +++ b/csrc/cutlass_extensions/torch_utils.hpp @@ -1,19 +1,12 @@ #pragma once -#include "torch_utils.h" +#include +#include +#include +#include -// This header is shared between _C (unstable ABI, used by machete) and -// _C_stable_libtorch (stable ABI, used by W4A8/sparse). TORCH_TARGET_VERSION -// is defined only for the stable target, so we switch includes and types -// accordingly. TorchTensor (not Tensor) avoids ambiguity with cute::Tensor. -#ifdef TORCH_TARGET_VERSION - #include - #include - #include +// TorchTensor (not Tensor) avoids ambiguity with cute::Tensor. using TorchTensor = torch::stable::Tensor; -#else -using TorchTensor = torch::Tensor; -#endif #include "cute/layout.hpp" #include "cutlass/layout/matrix.h" @@ -70,15 +63,15 @@ CUTE_HOST_DEVICE constexpr auto make_shape_from_idx(F&& f) { template static inline auto make_cute_layout(TorchTensor const& tensor, std::string_view name = "tensor") { - TORCH_UTILS_CHECK(tensor.dim() <= rank(Stride{})); + STD_TORCH_CHECK(tensor.dim() <= rank(Stride{})); auto stride = cute::transform_with_idx(Stride{}, [&](auto const& stride_ele, auto const& idx) { using StrideEle = std::decay_t; if (idx < tensor.dim()) { if constexpr (cute::is_static_v) { - TORCH_UTILS_CHECK(StrideEle::value == tensor.stride(idx), "Expected ", - name, ".stride(", idx, ") to be ", StrideEle::value); + STD_TORCH_CHECK(StrideEle::value == tensor.stride(idx), "Expected ", + name, ".stride(", idx, ") to be ", StrideEle::value); return StrideEle{}; } else { if (tensor.size(idx) == 1) { @@ -122,29 +115,7 @@ static inline auto maybe_make_cute_layout( } // -// Torch Type to Cutlass Type (equivalent_cutlass_type) -// - -template -struct equivalent_cutlass_type { - using type = T; -}; - -template -using equivalent_cutlass_type_t = typename equivalent_cutlass_type::type; - -template <> -struct equivalent_cutlass_type { - using type = cutlass::half_t; -}; - -template <> -struct equivalent_cutlass_type { - using type = cutlass::bfloat16_t; -}; - -// -// equivalent_scalar_t (basically inverse of equivalent_cutlass_type) +// Cutlass Type to Torch Type (equivalent_scalar_type) // // Return a `torch::headeronly::CppTypeToScalarType` compatible type, i.e. diff --git a/csrc/libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh b/csrc/cutlass_extensions/vllm_collective_builder.cuh similarity index 100% rename from csrc/libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh rename to csrc/cutlass_extensions/vllm_collective_builder.cuh diff --git a/csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh b/csrc/cutlass_extensions/vllm_numeric_conversion.cuh similarity index 99% rename from csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh rename to csrc/cutlass_extensions/vllm_numeric_conversion.cuh index 90f226cf64..0f74053fd2 100644 --- a/csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh +++ b/csrc/cutlass_extensions/vllm_numeric_conversion.cuh @@ -1,9 +1,9 @@ #pragma once #include "cutlass/numeric_conversion.h" -#include "cutlass_extensions/vllm_custom_types.cuh" -#include "cutlass_extensions/cute_utils.cuh" -#include "cutlass_extensions/vllm_type_utils.cuh" +#include "vllm_custom_types.cuh" +#include "cute_utils.cuh" +#include "vllm_type_utils.cuh" // this file extends: // https://github.com/NVIDIA/cutlass/blob/cutlass-3.5.0/include/cutlass/numeric_conversion.h diff --git a/csrc/cutlass_extensions/vllm_type_utils.cuh b/csrc/cutlass_extensions/vllm_type_utils.cuh index 500ed508c8..ba79dddf6f 100644 --- a/csrc/cutlass_extensions/vllm_type_utils.cuh +++ b/csrc/cutlass_extensions/vllm_type_utils.cuh @@ -2,7 +2,7 @@ #include "cutlass/half.h" #include "cuda_bf16.h" -#include "cutlass_extensions/vllm_custom_types.cuh" +#include "vllm_custom_types.cuh" namespace cutlass { diff --git a/csrc/dispatch_utils.h b/csrc/dispatch_utils.h index de0c505b7a..aa125a08fd 100644 --- a/csrc/dispatch_utils.h +++ b/csrc/dispatch_utils.h @@ -1,158 +1,188 @@ /* - * Adapted from - * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + * Stable ABI compatible dispatch utilities for vLLM. + * Adapted from dispatch_utils.h to use PyTorch's header-only (THO_*) macros + * instead of the ATen (AT_*) macros. + * + * These macros use: + * - THO_DISPATCH_SWITCH instead of AT_DISPATCH_SWITCH + * - THO_DISPATCH_CASE instead of AT_DISPATCH_CASE + * - torch::headeronly::ScalarType instead of at::ScalarType + * + * Add more macros here as needed when migrating additional kernels. */ #pragma once -#include +#include +#include +#include // Need a special dispatch case macro since we will nest the FP8 dispatch. // Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'. -#define AT_DISPATCH_FP8_CASE(enum_type, ...) \ - AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) +#define VLLM_STABLE_DISPATCH_FP8_CASE(enum_type, ...) \ + THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) -#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +// Same idea, for dispatching on an int32/int64 index tensor (e.g. topk_ids) +// nested inside a value-type dispatch. Named 'idx_t' instead of 'scalar_t'. +#define VLLM_STABLE_DISPATCH_IDX_CASE(enum_type, ...) \ + THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, idx_t, __VA_ARGS__) -#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) +#define VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(...) \ + VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Int, \ + __VA_ARGS__) \ + VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Long, \ + __VA_ARGS__) -#define VLLM_DISPATCH_CASE_HALF_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +#define VLLM_STABLE_DISPATCH_IDX_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(__VA_ARGS__)) -#define VLLM_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) +#define VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) -// ROCm devices might use either fn or fnuz, so set up dispatch table for both. -// A host-based check at runtime will create a preferred FP8 type for ROCm -// such that the correct kernel is dispatched. +#define VLLM_STABLE_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) + +#define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH( \ + TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) + +// FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format #ifdef USE_ROCM - #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ - AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) - - #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) + #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ + VLLM_STABLE_DISPATCH_FP8_CASE( \ + torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ + VLLM_STABLE_DISPATCH_FP8_CASE( \ + torch::headeronly::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) #else - #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ - AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) - - #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) + #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ + VLLM_STABLE_DISPATCH_FP8_CASE( \ + torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) #endif // When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. -// See AT_DISPATCH_FP8_CASE above. -#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) +// See VLLM_STABLE_DISPATCH_FP8_CASE above. +#define VLLM_STABLE_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) -#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__)) +// Half types dispatch (Half + BFloat16) +#define VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) -#define VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) +#define VLLM_STABLE_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) -#define VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(__VA_ARGS__)) +// Quant type dispatch (FP8 + INT8) +#ifdef USE_ROCM + #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ + __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fnuz, \ + __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) +#else + #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ + __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) +#endif -#define VLLM_DISPATCH_CASE_INTEGRAL_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) +#define VLLM_STABLE_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__)) -#define VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::UInt16, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::UInt32, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::UInt64, __VA_ARGS__) - -#define VLLM_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH( \ - TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \ - switch (VEC_SIZE) { \ - case 16: { \ - constexpr int vec_size = 16; \ - __VA_ARGS__(); \ - break; \ - } \ - case 8: { \ - constexpr int vec_size = 8; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int vec_size = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - case 2: { \ - constexpr int vec_size = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - default: { \ - constexpr int vec_size = 1; \ - __VA_ARGS__(); \ - break; \ - } \ +// Group size dispatch (pure C++ if/else, no ATen dependency) +#define VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \ + if (group_size == 128) { \ + constexpr int const_group_size = 128; \ + __VA_ARGS__(); \ + } else if (group_size == 64) { \ + constexpr int const_group_size = 64; \ + __VA_ARGS__(); \ } -#define VLLM_DISPATCH_BOOL(expr, const_expr, ...) \ - if (expr) { \ - constexpr bool const_expr = true; \ - __VA_ARGS__(); \ - } else { \ - constexpr bool const_expr = false; \ - __VA_ARGS__(); \ +// Boolean dispatch +#define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \ + if (expr) { \ + constexpr bool const_expr = true; \ + __VA_ARGS__(); \ + } else { \ + constexpr bool const_expr = false; \ + __VA_ARGS__(); \ } -#define VLLM_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \ - if (group_size == 128) { \ - constexpr int const_group_size = 128; \ - __VA_ARGS__(); \ - } else if (group_size == 64) { \ - constexpr int const_group_size = 64; \ - __VA_ARGS__(); \ +// Vec size dispatch (pure C++ switch, no ATen dependency) +#define VLLM_STABLE_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \ + switch (VEC_SIZE) { \ + case 16: { \ + constexpr int vec_size = 16; \ + __VA_ARGS__(); \ + break; \ + } \ + case 8: { \ + constexpr int vec_size = 8; \ + __VA_ARGS__(); \ + break; \ + } \ + case 4: { \ + constexpr int vec_size = 4; \ + __VA_ARGS__(); \ + break; \ + } \ + case 2: { \ + constexpr int vec_size = 2; \ + __VA_ARGS__(); \ + break; \ + } \ + default: { \ + constexpr int vec_size = 1; \ + __VA_ARGS__(); \ + break; \ + } \ } -#define VLLM_DISPATCH_RANK234(NUM_DIMS, ...) \ - switch (NUM_DIMS) { \ - case 2: { \ - constexpr int tensor_rank = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - case 3: { \ - constexpr int tensor_rank = 3; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int tensor_rank = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - default: \ - TORCH_CHECK(false, "Expects rank 2, 3 or 4 tensors but got ", NUM_DIMS); \ +// Tensor rank dispatch (2D, 3D, 4D) +#define VLLM_STABLE_DISPATCH_RANK234(NUM_DIMS, ...) \ + switch (NUM_DIMS) { \ + case 2: { \ + constexpr int tensor_rank = 2; \ + __VA_ARGS__(); \ + break; \ + } \ + case 3: { \ + constexpr int tensor_rank = 3; \ + __VA_ARGS__(); \ + break; \ + } \ + case 4: { \ + constexpr int tensor_rank = 4; \ + __VA_ARGS__(); \ + break; \ + } \ + default: \ + STD_TORCH_CHECK( \ + false, "Expects rank 2, 3 or 4 tensors but got unsupported rank"); \ } diff --git a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu b/csrc/dsv3_fused_a_gemm.cu similarity index 99% rename from csrc/libtorch_stable/dsv3_fused_a_gemm.cu rename to csrc/dsv3_fused_a_gemm.cu index 4250683b91..c0989f8420 100644 --- a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu +++ b/csrc/dsv3_fused_a_gemm.cu @@ -25,7 +25,7 @@ #include #include "core/registration.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/fp32_router_gemm.cu b/csrc/fp32_router_gemm.cu similarity index 100% rename from csrc/libtorch_stable/fp32_router_gemm.cu rename to csrc/fp32_router_gemm.cu diff --git a/csrc/libtorch_stable/fp32_router_gemm_entry.cu b/csrc/fp32_router_gemm_entry.cu similarity index 99% rename from csrc/libtorch_stable/fp32_router_gemm_entry.cu rename to csrc/fp32_router_gemm_entry.cu index fc09193c86..e55ac7e721 100644 --- a/csrc/libtorch_stable/fp32_router_gemm_entry.cu +++ b/csrc/fp32_router_gemm_entry.cu @@ -6,7 +6,7 @@ #include #include "core/registration.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu rename to csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 5a5caacce7..4eb8d7756b 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -47,7 +47,7 @@ #include #include -#include "cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "type_convert.cuh" diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu rename to csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu index 3892790c40..dd88645904 100644 --- a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu +++ b/csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -57,16 +57,16 @@ #include #include -#include "cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "type_convert.cuh" #ifndef USE_ROCM #include - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #else #include - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #endif #include #include @@ -835,7 +835,7 @@ static void launchPdlSlots(KernelT kernel, int num_tokens, int slots_per_token, // clang-format off // hipify's CUDA->HIP regex catastrophically backtracks on "> > >"; keep the // launch closer as ">>>". clang-format would otherwise re-split it (it does - // not parse the CUDA launch syntax in this libtorch_stable file). + // not parse the CUDA launch syntax in this stable-ABI file). kernel<<>>(args...); // clang-format on #endif diff --git a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu rename to csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu index 0a45af8ade..09484daf33 100644 --- a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -54,15 +54,15 @@ #include "torch_utils.h" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "type_convert.cuh" -#include "../attention/dtype_fp8.cuh" +#include "core/vec_dtypes/dtype_fp8.cuh" #include "dispatch_utils.h" #ifdef USE_ROCM - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif // Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs. diff --git a/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu b/csrc/fused_qknorm_rope_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_qknorm_rope_kernel.cu rename to csrc/fused_qknorm_rope_kernel.cu index 01bad79cd8..a75ede86fc 100644 --- a/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu +++ b/csrc/fused_qknorm_rope_kernel.cu @@ -21,7 +21,7 @@ #include "torch_utils.h" #include "async_util.cuh" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "type_convert.cuh" #include "dispatch_utils.h" diff --git a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu b/csrc/gdn/fused_gdn_decode_kernel.cu similarity index 99% rename from csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu rename to csrc/gdn/fused_gdn_decode_kernel.cu index 1229404dc6..5e606ffd7f 100644 --- a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu +++ b/csrc/gdn/fused_gdn_decode_kernel.cu @@ -10,7 +10,7 @@ #include #include "../torch_utils.h" -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" namespace { diff --git a/csrc/libtorch_stable/hisparse_kernels.cu b/csrc/hisparse_kernels.cu similarity index 99% rename from csrc/libtorch_stable/hisparse_kernels.cu rename to csrc/hisparse_kernels.cu index 5c7621c2f5..da3857d1e9 100644 --- a/csrc/libtorch_stable/hisparse_kernels.cu +++ b/csrc/hisparse_kernels.cu @@ -17,7 +17,6 @@ #include "torch_utils.h" #include "ops.h" -#include "../cuda_utils.h" #include #include diff --git a/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu b/csrc/kimi_k3/attn_res_kernel.cu similarity index 100% rename from csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu rename to csrc/kimi_k3/attn_res_kernel.cu diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu similarity index 100% rename from csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu rename to csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu b/csrc/kimi_k3/fused_kda_decode_kernel.cu similarity index 99% rename from csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu rename to csrc/kimi_k3/fused_kda_decode_kernel.cu index 00e5ce794d..c5214ff486 100644 --- a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu +++ b/csrc/kimi_k3/fused_kda_decode_kernel.cu @@ -9,7 +9,7 @@ #include #include "../torch_utils.h" -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" namespace { diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu similarity index 100% rename from csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu rename to csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu diff --git a/csrc/libtorch_stable/launch_bounds_utils.h b/csrc/launch_bounds_utils.h similarity index 100% rename from csrc/libtorch_stable/launch_bounds_utils.h rename to csrc/launch_bounds_utils.h diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/layernorm_kernels.cu similarity index 99% rename from csrc/libtorch_stable/layernorm_kernels.cu rename to csrc/layernorm_kernels.cu index b342c59f18..8d6741ddac 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/layernorm_kernels.cu @@ -3,7 +3,7 @@ #include "torch_utils.h" #include "cub_helpers.h" -#include "../core/batch_invariant.hpp" +#include "core/utils.hpp" #include "type_convert.cuh" #include "dispatch_utils.h" #include "quantization/vectorization_utils.cuh" diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/layernorm_quant_kernels.cu similarity index 99% rename from csrc/libtorch_stable/layernorm_quant_kernels.cu rename to csrc/layernorm_quant_kernels.cu index f43be531de..7f8fb9e510 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/csrc/layernorm_quant_kernels.cu @@ -10,8 +10,8 @@ #include "torch_utils.h" #include "cub_helpers.h" -#include "../core/batch_invariant.hpp" -#include "../quantization/w8a8/fp8/common.cuh" +#include "core/utils.hpp" +#include "quantization/w8a8/fp8/common.cuh" #include "type_convert.cuh" #include "dispatch_utils.h" #include "quantization/vectorization_utils.cuh" diff --git a/csrc/libtorch_stable/attention/attention_utils.cuh b/csrc/libtorch_stable/attention/attention_utils.cuh deleted file mode 100644 index 29783cb0e5..0000000000 --- a/csrc/libtorch_stable/attention/attention_utils.cuh +++ /dev/null @@ -1,57 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp - * Copyright (c) 2023, The vLLM team. - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include "../../cuda_compat.h" -#include "../../attention/attention_dtypes.h" - -#include -#include - -namespace vllm { - -// Q*K^T operation. -template -inline __device__ float qk_dot_(const Vec (&q)[N], const Vec (&k)[N]) { - using A_vec = typename FloatVec::Type; - // Compute the parallel products for Q*K^T (treat vector lanes separately). - A_vec qk_vec = mul(q[0], k[0]); -#pragma unroll - for (int ii = 1; ii < N; ++ii) { - qk_vec = vllm::fma(q[ii], k[ii], qk_vec); - } - - // Finalize the reduction across lanes. - float qk = sum(qk_vec); -#pragma unroll - for (int mask = THREAD_GROUP_SIZE / 2; mask >= 1; mask /= 2) { - qk += VLLM_SHFL_XOR_SYNC(qk, mask); - } - return qk; -} - -template -struct Qk_dot { - template - static inline __device__ float dot(const Vec (&q)[N], const Vec (&k)[N]) { - return qk_dot_(q, k); - } -}; - -} // namespace vllm diff --git a/csrc/libtorch_stable/cutlass_extensions/common.cpp b/csrc/libtorch_stable/cutlass_extensions/common.cpp deleted file mode 100644 index 5bc9463bfa..0000000000 --- a/csrc/libtorch_stable/cutlass_extensions/common.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "common.hpp" - -int32_t get_sm_version_num() { - int32_t major_capability, minor_capability; - cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, - 0); - cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, - 0); - int32_t version_num = major_capability * 10 + minor_capability; - return version_num; -} \ No newline at end of file diff --git a/csrc/libtorch_stable/dispatch_utils.h b/csrc/libtorch_stable/dispatch_utils.h deleted file mode 100644 index aa125a08fd..0000000000 --- a/csrc/libtorch_stable/dispatch_utils.h +++ /dev/null @@ -1,188 +0,0 @@ -/* - * Stable ABI compatible dispatch utilities for vLLM. - * Adapted from dispatch_utils.h to use PyTorch's header-only (THO_*) macros - * instead of the ATen (AT_*) macros. - * - * These macros use: - * - THO_DISPATCH_SWITCH instead of AT_DISPATCH_SWITCH - * - THO_DISPATCH_CASE instead of AT_DISPATCH_CASE - * - torch::headeronly::ScalarType instead of at::ScalarType - * - * Add more macros here as needed when migrating additional kernels. - */ -#pragma once - -#include -#include -#include - -// Need a special dispatch case macro since we will nest the FP8 dispatch. -// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'. -#define VLLM_STABLE_DISPATCH_FP8_CASE(enum_type, ...) \ - THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) - -// Same idea, for dispatching on an int32/int64 index tensor (e.g. topk_ids) -// nested inside a value-type dispatch. Named 'idx_t' instead of 'scalar_t'. -#define VLLM_STABLE_DISPATCH_IDX_CASE(enum_type, ...) \ - THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, idx_t, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(...) \ - VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Int, \ - __VA_ARGS__) \ - VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Long, \ - __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_IDX_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(__VA_ARGS__)) - -#define VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) - -#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ - VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) - -#define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH( \ - TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) - -// FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format -#ifdef USE_ROCM - #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ - VLLM_STABLE_DISPATCH_FP8_CASE( \ - torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - VLLM_STABLE_DISPATCH_FP8_CASE( \ - torch::headeronly::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) -#else - #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ - VLLM_STABLE_DISPATCH_FP8_CASE( \ - torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) -#endif - -// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. -// See VLLM_STABLE_DISPATCH_FP8_CASE above. -#define VLLM_STABLE_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) - -// Half types dispatch (Half + BFloat16) -#define VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) - -// Quant type dispatch (FP8 + INT8) -#ifdef USE_ROCM - #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ - __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fnuz, \ - __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) -#else - #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ - __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) -#endif - -#define VLLM_STABLE_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__)) - -// Group size dispatch (pure C++ if/else, no ATen dependency) -#define VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \ - if (group_size == 128) { \ - constexpr int const_group_size = 128; \ - __VA_ARGS__(); \ - } else if (group_size == 64) { \ - constexpr int const_group_size = 64; \ - __VA_ARGS__(); \ - } - -// Boolean dispatch -#define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \ - if (expr) { \ - constexpr bool const_expr = true; \ - __VA_ARGS__(); \ - } else { \ - constexpr bool const_expr = false; \ - __VA_ARGS__(); \ - } - -// Vec size dispatch (pure C++ switch, no ATen dependency) -#define VLLM_STABLE_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \ - switch (VEC_SIZE) { \ - case 16: { \ - constexpr int vec_size = 16; \ - __VA_ARGS__(); \ - break; \ - } \ - case 8: { \ - constexpr int vec_size = 8; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int vec_size = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - case 2: { \ - constexpr int vec_size = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - default: { \ - constexpr int vec_size = 1; \ - __VA_ARGS__(); \ - break; \ - } \ - } - -// Tensor rank dispatch (2D, 3D, 4D) -#define VLLM_STABLE_DISPATCH_RANK234(NUM_DIMS, ...) \ - switch (NUM_DIMS) { \ - case 2: { \ - constexpr int tensor_rank = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - case 3: { \ - constexpr int tensor_rank = 3; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int tensor_rank = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - default: \ - STD_TORCH_CHECK( \ - false, "Expects rank 2, 3 or 4 tensors but got unsupported rank"); \ - } diff --git a/csrc/libtorch_stable/mamba/static_switch.h b/csrc/libtorch_stable/mamba/static_switch.h deleted file mode 100644 index 840cb2374a..0000000000 --- a/csrc/libtorch_stable/mamba/static_switch.h +++ /dev/null @@ -1,28 +0,0 @@ -// Inspired by -// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h -// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h - -// clang-format off -// adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/static_switch.h -#pragma once - -/// @param COND - a boolean expression to switch by -/// @param CONST_NAME - a name given for the constexpr bool variable. -/// @param ... - code to execute for true and false -/// -/// Usage: -/// ``` -/// BOOL_SWITCH(flag, BoolConst, [&] { -/// some_function(...); -/// }); -/// ``` -#define BOOL_SWITCH(COND, CONST_NAME, ...) \ - [&] { \ - if (COND) { \ - constexpr bool CONST_NAME = true; \ - return __VA_ARGS__(); \ - } else { \ - constexpr bool CONST_NAME = false; \ - return __VA_ARGS__(); \ - } \ - }() diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py b/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py deleted file mode 100644 index 6948a98a37..0000000000 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py +++ /dev/null @@ -1,315 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import glob -import itertools -import os -import subprocess -import sys - -import jinja2 - -ARCHS = [] -SUPPORT_FP8 = False -SUPPORT_SM75 = False -SUPPORT_SM80 = False -for arch in sys.argv[1].split(","): - arch = arch[: arch.index(".") + 2].replace(".", "") - arch = int(arch) - # SM89 and the SM12x family (SM120 RTX 5090, SM121 DGX Spark GB10) - # fully support mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32. - # SM90 and SM100 can use this PTX, but it’s simulated - # with FP16 MMA, so it cannot achieve any acceleration. - if arch == 89 or arch // 10 == 12: - SUPPORT_FP8 = True - if arch >= 80: - SUPPORT_SM80 = True - if arch == 75: - SUPPORT_SM75 = True - -FILE_HEAD_COMMENT = """ -// auto generated by generate_kernels.py -// clang-format off -""".lstrip() - -FILE_HEAD = ( - FILE_HEAD_COMMENT - + """ -#include "kernel.h" -#include "marlin_template.h" - -namespace MARLIN_NAMESPACE_NAME { -""" -) - -TEMPLATE = ( - "template __global__ void Marlin<" - "{{a_type_id}}, " - "{{b_type_id}}, " - "{{c_type_id}}, " - "{{s_type_id}}, " - "{{threads}}, " - "{{thread_m_blocks}}, " - "{{thread_n_blocks}}, " - "{{thread_k_blocks}}, " - "{{m_block_size_8}}, " - "{{stages}}, " - "{{group_blocks}}, " - "{{is_zp_float}}>" - "( MARLIN_KERNEL_PARAMS );" -) - -THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)] - -THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4] - -QUANT_CONFIGS = [ - # AWQ-INT4 - { - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 - { - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 2, 4, 8], - }, - # AWQ-INT8 - { - "b_type": "kU8B128", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 2, 4, 8], - }, - # FP8 - { - "b_type": "kFE4M3fn", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 8], - }, - # NVFP4 - { - "b_type": "kFE2M1f", - "s_type": "kFE4M3fn", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [1], - }, - # MXFP4 - { - "a_type": ["kBFloat16"], - "b_type": "kFE2M1f", - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [2], - }, - # MXFP8 - { - "a_type": ["kBFloat16"], - "b_type": "kFE4M3fn", - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [2], - }, - # AWQ-INT4 with INT8 activation - { - "a_type": ["kS8"], - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 with INT8 activation - { - "a_type": ["kS8"], - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # AWQ-INT4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # MXFP4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kFE2M1f", - "c_type": ["kBFloat16"], - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [2], - }, -] - - -def remove_old_kernels(): - for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"): - subprocess.call(["rm", "-f", filename]) - - filename = os.path.dirname(__file__) + "/kernel_selector.h" - subprocess.call(["rm", "-f", filename]) - - -def generate_new_kernels(): - result_dict = {} - sm_75_result_dict = {} - - for quant_config in QUANT_CONFIGS: - c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"]) - a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"]) - b_type = quant_config["b_type"] - all_group_blocks = quant_config["group_blocks"] - all_m_blocks = quant_config["thread_m_blocks"] - all_thread_configs = quant_config["thread_configs"] - - for a_type, c_type in itertools.product(a_types, c_types): - if not SUPPORT_FP8 and a_type == "kFE4M3fn": - continue - if "16" in a_type and "16" in c_type and a_type != c_type: - continue - s_type = quant_config.get("s_type", c_type) - if (a_type, b_type, c_type) not in result_dict: - result_dict[(a_type, b_type, c_type)] = [] - if a_type in ["kFloat16", "kS8"] and c_type == "kFloat16": - sm_75_result_dict[(a_type, b_type, c_type)] = [] - - for group_blocks, m_blocks, thread_configs in itertools.product( - all_group_blocks, all_m_blocks, all_thread_configs - ): - thread_k, thread_n, threads = thread_configs - - if threads == 256: - # for small batch (m_blocks == 1), - # we only need (128, 128, 256) - # for large batch (m_blocks > 1), - # we only need (64, 256, 256) - if m_blocks <= 1 and (thread_k, thread_n) != (128, 128): - continue - if m_blocks > 1 and (thread_k, thread_n) != (64, 256): - continue - - config = { - "threads": threads, - "s_type": s_type, - "thread_m_blocks": max(m_blocks, 1), - "thread_k_blocks": thread_k // 16, - "thread_n_blocks": thread_n // 16, - "m_block_size_8": "true" if m_blocks == 0.5 else "false", - "stages": 4, - "group_blocks": group_blocks, - "is_zp_float": "false", - } - - if SUPPORT_SM80: - result_dict[(a_type, b_type, c_type)].append(config) - if (a_type, b_type, c_type) in sm_75_result_dict and SUPPORT_SM75: - config_sm75 = config.copy() - config_sm75["stages"] = 2 - sm_75_result_dict[(a_type, b_type, c_type)].append(config_sm75) - - kernel_selector_str = FILE_HEAD_COMMENT - - for result_dict_tmp in [result_dict, sm_75_result_dict]: - for (a_type, b_type, c_type), config_list in result_dict_tmp.items(): - all_template_str_list = [] - if not config_list: - continue - for config in config_list: - s_type = config["s_type"] - template_str = jinja2.Template(TEMPLATE).render( - a_type_id=f"vllm::{a_type}.id()", - b_type_id=f"vllm::{b_type}.id()", - c_type_id=f"vllm::{c_type}.id()", - s_type_id=f"vllm::{s_type}.id()", - **config, - ) - all_template_str_list.append(template_str) - - conditions = [ - f"a_type == vllm::{a_type}", - f"b_type == vllm::{b_type}", - f"c_type == vllm::{c_type}", - f"s_type == vllm::{s_type}", - f"threads == {config['threads']}", - f"thread_m_blocks == {config['thread_m_blocks']}", - f"thread_n_blocks == {config['thread_n_blocks']}", - f"thread_k_blocks == {config['thread_k_blocks']}", - f"m_block_size_8 == {config['m_block_size_8']}", - f"stages == {config['stages']}", - f"group_blocks == {config['group_blocks']}", - f"is_zp_float == {config['is_zp_float']}", - ] - conditions = " && ".join(conditions) - - if kernel_selector_str == FILE_HEAD_COMMENT: - kernel_selector_str += f"if ({conditions})\n kernel = " - else: - kernel_selector_str += f"else if ({conditions})\n kernel = " - - kernel_template2 = ( - "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, " - "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, " - "{{thread_n_blocks}}, {{thread_k_blocks}}, " - "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, " - "{{is_zp_float}}>;" - ) - - kernel_selector_str += ( - jinja2.Template(kernel_template2).render( - a_type_id=f"vllm::{a_type}.id()", - b_type_id=f"vllm::{b_type}.id()", - c_type_id=f"vllm::{c_type}.id()", - s_type_id=f"vllm::{s_type}.id()", - **config, - ) - + "\n" - ) - - file_content = FILE_HEAD + "\n\n" - file_content += "\n\n".join(all_template_str_list) + "\n\n}\n" - if a_type == "kFE4M3fn": - filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - elif result_dict_tmp is sm_75_result_dict: - filename = f"sm75_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - else: - filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - - filename = filename.lower() - - with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: - f.write(file_content) - - if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: - kernel_selector_str += ( - "else if (a_type == vllm::kFE4M3fn)\n" - " STD_TORCH_CHECK(false, " - '"marlin kernel with fp8 activation is not built.");' - ) - - with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f: - f.write(kernel_selector_str) - - -if __name__ == "__main__": - remove_old_kernels() - generate_new_kernels() diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h deleted file mode 100644 index 3266d28b1e..0000000000 --- a/csrc/libtorch_stable/ops.h +++ /dev/null @@ -1,822 +0,0 @@ -#pragma once - -#include -#include -#include - -#include -#include -#include - -#include - -inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) { - // Ensure tensor is on CUDA - STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device"); - - // Get the raw data pointer - void* data_ptr = tensor.mutable_data_ptr(); - - /// Create a new tensor from the raw data pointer - return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(), - tensor.device(), tensor.scalar_type()); -} - -void per_token_group_quant_fp8(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double fp8_min, - double fp8_max, bool scale_ue8m0, - bool dummy_is_scale_transposed, - bool dummy_is_tma_aligned); - -// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales. -void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s_packed, - int64_t group_size, double eps, - double min_8bit, double max_8bit); - -void per_token_group_quant_int8(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double int8_min, - double int8_max); - -#ifndef USE_ROCM -bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); -bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability); -bool cutlass_group_gemm_supported(int64_t cuda_device_capability); - -void cutlass_scaled_mm(torch::stable::Tensor& out, - torch::stable::Tensor const& a, - torch::stable::Tensor const& b, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - std::optional const& bias); - -void cutlass_moe_mm(torch::stable::Tensor& out_tensors, - torch::stable::Tensor const& a_tensors, - torch::stable::Tensor const& b_tensors, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - torch::stable::Tensor const& expert_offsets, - torch::stable::Tensor const& problem_sizes, - torch::stable::Tensor const& a_strides, - torch::stable::Tensor const& b_strides, - torch::stable::Tensor const& c_strides, bool per_act_token, - bool per_out_ch); - -void cutlass_scaled_mm_azp(torch::stable::Tensor& out, - torch::stable::Tensor const& a, - torch::stable::Tensor const& b, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - torch::stable::Tensor const& azp_adj, - std::optional const& azp, - std::optional const& bias); - -void get_cutlass_moe_mm_data( - const torch::stable::Tensor& topk_ids, - torch::stable::Tensor& expert_offsets, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, - torch::stable::Tensor& input_permutation, - torch::stable::Tensor& output_permutation, const int64_t num_experts, - const int64_t n, const int64_t k, - const std::optional& blockscale_offsets, - const bool is_gated); - -void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( - const torch::stable::Tensor& expert_first_token_offset, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k, - const bool swap_ab); - -void get_cutlass_batched_moe_mm_data( - torch::stable::Tensor& expert_offsets, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, - const torch::stable::Tensor& expert_num_tokens, - const int64_t num_local_experts, const int64_t padded_m, const int64_t n, - const int64_t k); - -// FP4/NVFP4 ops -bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability); - -void cutlass_scaled_fp4_mm(torch::stable::Tensor& D, - torch::stable::Tensor const& A, - torch::stable::Tensor const& B, - torch::stable::Tensor const& A_sf, - torch::stable::Tensor const& B_sf, - torch::stable::Tensor const& alpha); - -void cutlass_fp4_group_mm(torch::stable::Tensor& output, - const torch::stable::Tensor& a, - const torch::stable::Tensor& b, - const torch::stable::Tensor& a_blockscale, - const torch::stable::Tensor& b_blockscales, - const torch::stable::Tensor& alphas, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& sf_offsets); - -std::tuple scaled_fp4_quant_func( - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout); - -void scaled_fp4_quant_out(torch::stable::Tensor const& input, - torch::stable::Tensor const& input_scale, - bool is_sf_swizzled_layout, - torch::stable::Tensor& output, - torch::stable::Tensor& output_scale); - -void scaled_fp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_global_scale, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts); - -void silu_and_mul_scaled_fp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_global_scale, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts); - -void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, - torch::stable::Tensor& output_block_scale, - torch::stable::Tensor& input, - torch::stable::Tensor& input_global_scale); - -void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, - const torch::stable::Tensor& a, - const torch::stable::Tensor& b, - const torch::stable::Tensor& a_blockscale, - const torch::stable::Tensor& b_blockscales, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& sf_offsets); - -// AWQ ops -torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats, - torch::stable::Tensor _kernel, - torch::stable::Tensor _scaling_factors, - torch::stable::Tensor _zeros, - int64_t split_k_iters); - -torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, - torch::stable::Tensor _scaling_factors, - torch::stable::Tensor _zeros, - int64_t split_k_iters, int64_t thx, - int64_t thy); - -// DSV3 fused A GEMM: conditionally compiled so declaration and impl -// registration are in the source file (dsv3_fused_a_gemm.cu) - -#endif - -// CPU tensor -> CUDA UVA view (shared CUDA/ROCm) -torch::stable::Tensor get_cuda_view_from_cpu_tensor( - torch::stable::Tensor& cpu_tensor); - -// Attention kernels (shared CUDA/ROCm) -void merge_attn_states( - torch::stable::Tensor& output, - std::optional output_lse, - const torch::stable::Tensor& prefix_output, - const torch::stable::Tensor& prefix_lse, - const torch::stable::Tensor& suffix_output, - const torch::stable::Tensor& suffix_lse, - const std::optional prefill_tokens_with_context, - const std::optional& output_scale = std::nullopt); - -torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, - bool inplace); - -// Layernorm kernels (shared CUDA/ROCm) -void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input, - std::optional weight, double epsilon); - -void fused_add_rms_norm(torch::stable::Tensor& input, - torch::stable::Tensor& residual, - std::optional weight, - double epsilon); - -// Layernorm-quant kernels (shared CUDA/ROCm) -void rms_norm_static_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& weight, - torch::stable::Tensor& scale, double epsilon); - -void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& residual, - torch::stable::Tensor& weight, - torch::stable::Tensor& scale, - double epsilon); - -// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) -void rms_norm_dynamic_per_token_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor const& weight, torch::stable::Tensor& scales, - double const var_epsilon, std::optional scale_ub, - std::optional residual); - -void rms_norm_per_block_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor const& weight, - torch::stable::Tensor& scales, - double const var_epsilon, - std::optional scale_ub, - std::optional residual, - int64_t group_size, bool is_scale_transposed); - -void silu_and_mul_per_block_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scales, - int64_t group_size, - std::optional scale_ub, - bool is_scale_transposed); - -// Positional encoding kernels (shared CUDA/ROCm) -void rotary_embedding(torch::stable::Tensor& positions, - torch::stable::Tensor& query, - std::optional key, - int64_t head_size, torch::stable::Tensor& cos_sin_cache, - bool is_neox, int64_t rope_dim_offset, bool inverse); - -void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, - int64_t num_heads_k, int64_t num_heads_v, - int64_t head_dim, double eps, - torch::stable::Tensor& q_weight, - torch::stable::Tensor& k_weight, - torch::stable::Tensor& cos_sin_cache, bool is_neox, - torch::stable::Tensor& position_ids, - int64_t forced_token_heads_per_warp); - -void fused_deepseek_v4_kv_rope_insert( - torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, - std::optional fp8_scale, bool kv_mxfp8); - -torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8, - bool apply_q_rope, bool is_q_interleaved); - -void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( - torch::stable::Tensor& q, torch::stable::Tensor const& kv, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, double eps, - int64_t cache_block_size, bool apply_q_norm); - -void fused_kimi_k3_mla_key_concat_kv_cache_insert( - torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, - torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_key_concat_ds_mla_insert( - torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, - torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, - torch::stable::Tensor& k_out); - -void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, - torch::stable::Tensor const& v, - torch::stable::Tensor& k_fp8, - torch::stable::Tensor& v_fp8); - -void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( - torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, - torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8, - torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& q_scale_inv, - torch::stable::Tensor const& k_scale_inv, - torch::stable::Tensor const& v_scale_inv, - torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( - torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, - torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, - torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( - torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, - torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, - torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& q_scale_inv, - torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( - torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, - torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, - torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( - torch::stable::Tensor const& q, torch::stable::Tensor const& kv, - torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, - torch::stable::Tensor const& fp8_scale, - torch::stable::Tensor const& q_fp8_scale_inv, double eps, - int64_t cache_block_size, bool apply_q_norm); - -#ifndef USE_ROCM -std::tuple -minimax_allreduce_rms_qk(torch::stable::Tensor qkv, - torch::stable::Tensor const& norm_weight_q, - torch::stable::Tensor const& norm_weight_k, - torch::stable::Tensor workspace, int64_t const q_size, - int64_t const kv_size, int64_t const rank, - int64_t const nranks, double const eps); -#endif - -// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV / -// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs -// the index branch and scatters k/v/index_k into their paged caches. -void fused_minimax_m3_qknorm_rope_kv_insert( - torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight, - torch::stable::Tensor const& k_norm_weight, - torch::stable::Tensor const& cos_sin_cache, - torch::stable::Tensor const& positions, int64_t num_heads, - int64_t num_kv_heads, int64_t rotary_dim, double eps, - std::optional index_q_norm_weight, - std::optional index_k_norm_weight, - int64_t num_index_heads, std::optional slot_mapping, - std::optional index_slot_mapping, - std::optional kv_cache, - std::optional index_cache, int64_t block_size, - std::optional q_out, - std::optional index_q_out, - const std::string& kv_cache_dtype, bool skip_index_branch, - std::optional q_fp8_out, double q_fp8_scale, - std::optional kv_k_scale, - std::optional kv_v_scale); - -#ifdef VLLM_ENABLE_FUSED_KDA_DECODE -void fused_kda_decode( - torch::stable::Tensor const& x, torch::stable::Tensor const& weight, - std::optional bias, - torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, - torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, - torch::stable::Tensor const& dt_bias, - torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, - torch::stable::Tensor& out, std::optional lower_bound, - std::optional output_gate, - std::optional norm_weight, double norm_eps); - -#endif - -#ifdef VLLM_ENABLE_FUSED_GDN_DECODE -void fused_gdn_decode_post_conv_mtp( - torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a, - torch::stable::Tensor const& b, torch::stable::Tensor const& a_log, - torch::stable::Tensor const& dt_bias, - torch::stable::Tensor const& state_indices, - torch::stable::Tensor const& cu_seqlens, - torch::stable::Tensor const& num_accepted_tokens, - torch::stable::Tensor& state, torch::stable::Tensor const& output_gate, - torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out, - double scale, double norm_eps, const std::string& output_gate_activation); - -#endif - -#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK -void fused_kda_prologue( - torch::stable::Tensor const& q, torch::stable::Tensor const& k, - torch::stable::Tensor const& v, torch::stable::Tensor const& raw_g, - torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, - torch::stable::Tensor const& dt_bias, torch::stable::Tensor& qg, - torch::stable::Tensor& w, torch::stable::Tensor& u, - torch::stable::Tensor& kg_t, torch::stable::Tensor& aqk, - torch::stable::Tensor& decay, torch::stable::Tensor const& cu_seqlens, - torch::stable::Tensor const& chunk_indices, - std::optional conv_weight, - std::optional conv_state, - std::optional conv_state_indices, - std::optional conv_has_initial_state, double scale, - double lower_bound); - -void fused_kda_chunk( - torch::stable::Tensor const& qg, torch::stable::Tensor const& w, - torch::stable::Tensor const& u, torch::stable::Tensor const& kg_t, - torch::stable::Tensor const& aqk, torch::stable::Tensor const& decay, - std::optional initial_state, - std::optional final_state, - torch::stable::Tensor& out, torch::stable::Tensor const& cu_seqlens, - torch::stable::Tensor const& chunk_offsets, double scale, - std::optional group_state, int64_t groups, - std::optional checkpoint_state, - std::optional checkpoint_offsets, - std::optional checkpoint_state_indices, - std::optional state_cache, - std::optional state_indices, - std::optional has_initial_state); -#endif - -#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES -void kimi_k3_attn_res(torch::stable::Tensor& prefix, - std::optional delta, - torch::stable::Tensor& blocks, - torch::stable::Tensor const& norm_weight, - torch::stable::Tensor const& qk_weight, - std::optional output_norm_weight, - torch::stable::Tensor& output, int64_t num_blocks, - int64_t block_write_idx, double eps, - double output_norm_eps); -#endif - -// Sampler kernels (shared CUDA/ROCm) -void apply_repetition_penalties_( - torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, - const torch::stable::Tensor& output_mask, - const torch::stable::Tensor& repetition_penalties); - -void top_k_per_row_prefill(const torch::stable::Tensor& logits, - const torch::stable::Tensor& rowStarts, - const torch::stable::Tensor& rowEnds, - torch::stable::Tensor& indices, int64_t numRows, - int64_t stride0, int64_t stride1, int64_t topK); - -void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, - const torch::stable::Tensor& seqLens, - torch::stable::Tensor& indices, int64_t numRows, - int64_t stride0, int64_t stride1, int64_t topK); - -void persistent_topk(const torch::stable::Tensor& logits, - const torch::stable::Tensor& lengths, - torch::stable::Tensor& output, - torch::stable::Tensor& workspace, int64_t k, - int64_t max_seq_len); - -#ifdef VLLM_ENABLE_COOPERATIVE_TOPK -void cooperative_topk(const torch::stable::Tensor& logits, - const torch::stable::Tensor& lengths, - torch::stable::Tensor& output, - torch::stable::Tensor& workspace, int64_t k, - int64_t max_seq_len); -#endif - -void selective_scan_fwd( - const torch::stable::Tensor& u, const torch::stable::Tensor& delta, - const torch::stable::Tensor& A, const torch::stable::Tensor& B, - const torch::stable::Tensor& C, - const std::optional& D_, - const std::optional& z_, - const std::optional& delta_bias_, - bool delta_softplus, - const std::optional& query_start_loc, - const std::optional& cache_indices, - const std::optional& has_initial_state, - const torch::stable::Tensor& ssm_states, int64_t null_block_id, - int64_t block_size, - const std::optional& block_idx_first_scheduled_token, - const std::optional& block_idx_last_scheduled_token, - const std::optional& initial_state_idx, - const std::optional& cu_chunk_seqlen, - const std::optional& last_chunk_indices); - -using fptr_t = int64_t; -fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, - torch::stable::Tensor& rank_data, int64_t rank, - bool fully_connected); -void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes); -void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes); -void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t local_buffer, - fptr_t multicast_buffer, fptr_t epoch_buffer, - int64_t stage_sz_bytes); -void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes); -void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, - fptr_t local_buffer, fptr_t epoch_buffer, - int64_t stage_sz_bytes); -void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, - fptr_t local_buffer, fptr_t multicast_buffer, - int64_t stage_sz_bytes, int64_t block_limit); -void dispose(fptr_t _fa); -int64_t meta_size(); -void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); -std::tuple, std::vector> -get_graph_buffer_ipc_meta(fptr_t _fa); -void register_graph_buffers(fptr_t _fa, - const std::vector>& handles, - const std::vector>& offsets); -std::tuple allocate_shared_buffer_and_handle( - int64_t size); -int64_t open_mem_handle(torch::stable::Tensor& mem_handle); -void free_shared_buffer(int64_t buffer); - -// Activation kernels (shared CUDA/ROCm) -void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); -void silu_and_mul_clamp(torch::stable::Tensor& out, - torch::stable::Tensor& input, double limit, - double alpha = 1.0, double beta = 0.0); - -void silu_and_mul_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& scale); - -void persistent_masked_m_silu_mul_quant( - const torch::stable::Tensor& input, // (E, T, 2*H) - const torch::stable::Tensor& tokens_per_expert, // (E) - torch::stable::Tensor& y_q, // (E, T, H) [OUT] - torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT] - bool use_ue8m0); - -void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_tanh_and_mul(torch::stable::Tensor& out, - torch::stable::Tensor& input); -void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double threshold); -void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double alpha = 1.702, double limit = 7.0); -void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double beta = 1.0, double linear_beta = -1.0); -void situ_and_mul_quant( - torch::stable::Tensor& out, torch::stable::Tensor& scale, - torch::stable::Tensor& input, double beta = 1.0, double linear_beta = -1.0, - int64_t group_size = 0, - std::optional num_valid_tokens = std::nullopt, - int64_t topk = 1); -void masked_situ_and_mul(torch::stable::Tensor& out, - torch::stable::Tensor& input, - const torch::stable::Tensor& expert_num_tokens, - double beta = 1.0, double linear_beta = -1.0); -void masked_moe_activation(torch::stable::Tensor& out, - torch::stable::Tensor& input, - const torch::stable::Tensor& valid_token_counts, - const std::string& activation, - double clamp_limit = 0.0, double alpha = 1.0, - double beta = 0.0, double situ_beta = 1.0, - double situ_linear_beta = -1.0); -void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); - -void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input); - -// INT8 quantization kernels (shared CUDA/ROCm) -void static_scaled_int8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor const& scale, - std::optional const& azp); - -void dynamic_scaled_int8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scales, - std::optional const& azp); - -// FP8 quantization kernels (shared CUDA/ROCm) -void static_scaled_fp8_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor const& scale, - std::optional group_shape = - std::nullopt); - -void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scale); - -void dynamic_per_token_scaled_fp8_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor& scale, - std::optional const& scale_ub); - -// GPTQ kernels (shared CUDA/ROCm) -torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, - torch::stable::Tensor b_q_weight, - torch::stable::Tensor b_gptq_qzeros, - torch::stable::Tensor b_gptq_scales, - bool use_exllama, bool use_v2_format, - int64_t bit); - -void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit); - -// Cache ops (shared CUDA/ROCm) -void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, - int64_t block_size_in_bytes, - const torch::stable::Tensor& block_mapping); - -// Batch swap: submit all block copies in a single driver call. -void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, - const torch::stable::Tensor& dst_ptrs, - const torch::stable::Tensor& sizes, - bool is_src_access_order_any); - -void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, - torch::stable::Tensor& key_cache, - torch::stable::Tensor& value_cache, - torch::stable::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale); - -void reshape_and_cache_flash( - torch::stable::Tensor& key, torch::stable::Tensor& value, - torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, - torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, - torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale); - -void concat_and_cache_mla(torch::stable::Tensor& kv_c, - torch::stable::Tensor& k_pe, - torch::stable::Tensor& kv_cache, - torch::stable::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::stable::Tensor& scale); - -void concat_and_cache_mla_grouped( - torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe, - torch::stable::Tensor& kv_cache_ptrs, torch::stable::Tensor& slot_mapping, - int64_t block_size, int64_t block_stride, int64_t entry_stride, - std::optional kv_scales, - const std::string& kv_cache_dtype); - -#ifndef USE_ROCM -// HiSparse kernels use raw PTX in the row copy; CUDA-only. -void hisparse_resolve_residency( - torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, - torch::stable::Tensor const& hot_block_table, - torch::stable::Tensor const& global_indices, - torch::stable::Tensor& hot_indices, - torch::stable::Tensor& device_global_indices, - torch::stable::Tensor& lru_slots, - std::optional const& request_state_indices, - int64_t region_stride, int64_t max_union_rows, - std::optional const& miss_mask, - std::optional const& stats, - std::optional const& attention_indices, - int64_t attention_block_stride, - std::optional const& request_ids, - std::optional const& source_block_table, - int64_t source_block_size, - std::optional const& resolved_global_indices, - std::optional const& valid_counts, - std::optional const& swap_host_physical_rows, - std::optional const& swap_device_physical_rows, - std::optional const& swap_counts, - std::optional const& resident_block_table, - int64_t resident_block_size, int64_t resident_null_block); - -void hisparse_invalidate_written_slots( - torch::stable::Tensor& device_global_indices, - torch::stable::Tensor const& request_state_indices, - torch::stable::Tensor const& req_id_per_token, - torch::stable::Tensor const& written_slots); - -void hisparse_gather_plan( - torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, - torch::stable::Tensor const& global_indices, - torch::stable::Tensor const& hot_indices, - torch::stable::Tensor const& miss_mask, - std::optional const& request_state_indices, - std::optional const& attention_indices, - int64_t attention_block_stride); - -void hisparse_gather_compact(torch::stable::Tensor const& host_cache, - torch::stable::Tensor& hot_cache, - torch::stable::Tensor const& miss_global_indices, - torch::stable::Tensor const& miss_hot_indices, - torch::stable::Tensor const& miss_counts); - -#endif // !USE_ROCM - -// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla -void concat_and_cache_mla_rope_fused( - torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, - torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c, - torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox, - torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache, - const std::string& kv_cache_dtype, - torch::stable::Tensor& kv_cache_quant_scale); - -// Just for unittest -void convert_fp8(torch::stable::Tensor& dst_cache, - torch::stable::Tensor& src_cache, const double scale, - const std::string& kv_cache_dtype); - -void gather_and_maybe_dequant_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRIES...] - torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] - torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] - int64_t num_tokens, const std::string& kv_cache_dtype, - torch::stable::Tensor const& scale, - std::optional seq_starts = std::nullopt); - -// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. -void cp_gather_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRIES...] - torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] - int64_t batch_size, - std::optional seq_starts = std::nullopt); - -// Gather and upconvert FP8 KV cache to BF16 workspace -void cp_gather_and_upconvert_fp8_kv_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // 656] - torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size, - std::optional seq_starts = std::nullopt, - std::optional host_cache = std::nullopt, - std::optional host_row_ids = std::nullopt, - std::optional device_row_ids = std::nullopt); - -// Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace -void cp_gather_and_upconvert_nvfp4_kv_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // 352] - torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size); - -// Indexer K quantization and cache function -void indexer_k_quant_and_cache( - torch::stable::Tensor& k, // [num_tokens, head_dim] - torch::stable::Tensor& kv_cache, // [num_blocks, block_size, - // cache_stride] - torch::stable::Tensor& slot_mapping, // [num_tokens] - int64_t quant_block_size, // quantization block size - const std::string& scale_fmt); - -// Concatenate query nope and rope for MLA/DSA attention -void concat_mla_q( - torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] - torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] - torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim + - // rope_dim] - -// Extract function to gather quantized K cache -void cp_gather_indexer_k_quant_cache( - const torch::stable::Tensor& kv_cache, // [num_blocks, block_size, - // cache_stride] - torch::stable::Tensor& dst_k, // [num_tokens, head_dim] - torch::stable::Tensor& dst_scale, // [num_tokens, head_dim / - // quant_block_size * 4] - const torch::stable::Tensor& block_table, // [batch_size, num_blocks] - const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] - -// Fused vocab-parallel embedding lookup (see -// vocab_parallel_embedding_kernels.cu). -void vocab_parallel_embedding( - torch::stable::Tensor& out, const torch::stable::Tensor& input_ids, - const torch::stable::Tensor& weight, int64_t org_vocab_start_index, - int64_t org_vocab_end_index, int64_t num_org_vocab_padding, - int64_t added_vocab_start_index, int64_t added_vocab_end_index); - -// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu). -void ngram_compute_n_gram_ids( - int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, - torch::stable::Tensor& ne_mods, - torch::stable::Tensor& exclusive_ne_embedder_size_sums, - torch::stable::Tensor& exclusive_req_len_sums, - torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, - torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); diff --git a/csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu deleted file mode 100644 index 2ffbee7ca2..0000000000 --- a/csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu +++ /dev/null @@ -1,12 +0,0 @@ -#include - -#include "libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h" - -void per_token_group_quant_int8(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double int8_min, - double int8_max) { - per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, - int8_min, int8_max); -} diff --git a/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h b/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h deleted file mode 100644 index 6630c0dece..0000000000 --- a/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include - -// 8-bit per-token-group quantization helper used by both FP8 and INT8 -void per_token_group_quant_8bit(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double min_8bit, - double max_8bit, bool scale_ue8m0 = false); - -// Public op: register-resident packed quant for the DeepGEMM Blackwell path. -// Restricted to group_size == 128 and bf16/fp16 input; other configurations -// raise STD_TORCH_CHECK. The legacy shared-memory fallback was removed because -// no production caller (deep_gemm_moe / input_quant_fp8) uses other shapes. -void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s_packed, - int64_t group_size, double eps, - double min_8bit, double max_8bit); diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp deleted file mode 100644 index 3d4f9ab273..0000000000 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ /dev/null @@ -1,1171 +0,0 @@ -#include "ops.h" -#include "cuda_utils.h" -#include "core/registration.h" - -#include - -// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility. -// Note: We register under namespace "_C" so ops are accessible as -// torch.ops._C. for compatibility with existing code. -STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { - // Compute per-token-group FP8 quantized tensor and scaling factor. - // The dummy arguments are here so we can correctly fuse with RMSNorm. - ops.def( - "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! " - "output_s, " - "int group_size, float eps, float fp8_min, float fp8_max, bool " - "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned " - ") -> ()"); - // Compute per-token-group 8-bit quantized tensor and UE8M0-packed, - // TMA-aligned scales for DeepGEMM. - ops.def( - "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, " - "Tensor! output_s_packed, int group_size, float eps, float fp8_min, " - "float fp8_max) -> ()"); - // Compute per-token-group INT8 quantized tensor and scaling factor. - ops.def( - "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! " - "output_s, int group_size, float eps, float int8_min, float int8_max) -> " - "()"); - ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); - -#ifndef USE_ROCM - - // Note about marlin kernel 'workspace' arguments: - // Technically these should be mutable since they are modified by the kernel. - // But since they are set back to zero once the kernel is finished we can - // hand wave and say that they have no net effect. - // - // The reason to mark 'workspace' as immutable is so that they don't interfere - // with using ScalarType arguments in the ops. If they are marked as mutable, - // pytorch throws an assert in - // 'torch._higher_order_ops._register_effectful_op' that prevents these - // kernels from being torch.compile'd. - // See the following document for more info on custom types and ops that use - // custom types: - // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA - - // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. - ops.def( - "machete_supported_schedules(" - " ScalarType a_type," - " int b_type," - " ScalarType? maybe_group_scales_type," - " ScalarType? maybe_group_zeros_type," - " ScalarType? maybe_channel_scales_type," - " ScalarType? maybe_token_scales_type," - " ScalarType? maybe_out_type" - ") -> str[]"); - ops.def( - "machete_mm(" - " Tensor A," - " Tensor B," - " int b_type," - " ScalarType? out_type," - " Tensor? group_scales," - " Tensor? group_zeros," - " int? group_size," - " Tensor? channel_scales," - " Tensor? token_scales," - " str? schedule" - ") -> Tensor"); - ops.def( - "machete_prepack_B(" - " Tensor B," - " ScalarType a_type," - " int b_type," - " ScalarType? group_scales_type" - ") -> Tensor"); - // conditionally compiled so impl registration is in source file - - // Marlin GEMM - ops.def( - "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " - "Tensor? b_bias_or_none,Tensor b_scales, " - "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " - "Tensor workspace, int b_type_id, " - "SymInt size_m, SymInt size_n, SymInt size_k, " - "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // gptq_marlin repack from GPTQ. - ops.def( - "gptq_marlin_repack(Tensor b_q_weight, " - "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // awq_marlin repack from AWQ. - ops.def( - "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " - "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // preprocess W-int4A-fp8 weight for marlin kernel - ops.def( - "marlin_int4_fp8_preprocess(Tensor qweight, " - "Tensor? qzeros_or_none, bool inplace) -> Tensor"); - // conditionally compiled so impl registrations are in source file -#endif - -#ifndef USE_ROCM - // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column - // quantization, as well as bias - ops.def( - "cutlass_scaled_mm(Tensor! out, Tensor a," - " Tensor b, Tensor a_scales," - " Tensor b_scales, Tensor? bias) -> ()"); - - // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column - // quantization. - ops.def( - "cutlass_scaled_mm_azp(Tensor! out, Tensor a," - " Tensor b, Tensor a_scales," - " Tensor b_scales, Tensor azp_adj," - " Tensor? azp, Tensor? bias) -> ()"); - - // Check if cutlass scaled_mm is supported for CUDA devices of the given - // capability - ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool"); - - // Check if cutlass grouped gemm is supported for CUDA devices of the given - // capability - ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool"); - - // CUTLASS w8a8 grouped GEMM - ops.def( - "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, " - " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, " - " Tensor problem_sizes, Tensor a_strides, " - " Tensor b_strides, Tensor c_strides, bool per_act_token, " - " bool per_out_ch) -> ()"); - - // A function that computes data required to run fused MoE with w8a8 grouped - // GEMM. It takes topk_ids as an input, and computes expert_offsets - // (token start indices of each expert). In addition to this, it computes - // problem sizes for each expert's multiplication used by the two mms called - // from fused MoE operation, and arrays with permutations required to shuffle - // and de-shuffle the input/output of the fused operation. - ops.def( - "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, " - " Tensor! problem_sizes1, Tensor! problem_sizes2, " - " Tensor! input_permutation, " - " Tensor! output_permutation, int num_experts, " - " int n, int k, Tensor? blockscale_offsets, " - " bool is_gated) -> ()"); - - // compute per-expert problem sizes from expert_first_token_offset - // produced by vLLM's moe_permute kernel - ops.def( - "get_cutlass_moe_mm_problem_sizes_from_expert_offsets(" - " Tensor expert_first_token_offset, " - " Tensor! problem_sizes1, " - " Tensor! problem_sizes2, " - " int n, int k, bool swap_ab) -> ()"); - - // A function that computes data required to run fused MoE with w8a8 grouped - // GEMM in batched expert format. It takes expert_num_tokens - // as an input, and computes expert_offsets (token start indices of each - // expert). In addition to this, it computes problem sizes for each expert's - // multiplication used by the two mms called from fused MoE operation. - ops.def( - "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, " - " Tensor! problem_sizes1, " - " Tensor! problem_sizes2, " - " Tensor expert_num_tokens, " - " int num_local_experts, int padded_m, " - " int n, int k) -> ()"); - - // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3) - ops.def( - "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> " - "bool"); - - // CUTLASS nvfp4 block scaled GEMM - ops.def( - "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b," - " Tensor block_scale_a, Tensor block_scale_b," - " Tensor alpha) -> ()"); - - // cutlass nvfp4 block scaled group GEMM - ops.def( - "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b," - " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas," - " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); - - // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE) - ops.def( - "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b," - " Tensor a_blockscale, Tensor b_blockscales," - " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); - - // Compute NVFP4 block quantized tensor. - ops.def( - "scaled_fp4_quant(Tensor input," - " Tensor input_scale, bool " - "is_sf_swizzled_layout) -> (Tensor, Tensor)"); - - // Out variant - // TODO: Add out_variant tag once PyTorch supports it (added in 2.11) - // This registration is now migrated to stable ABI - // at::Tag::out_variant is not available in the stable ABI (enum_tag.h is not - // yet in torch/headeronly), the tag should be applied from Python - // via torch.library.Library.define(..., tags=(torch.Tag.out_variant,)) - // with the .impl remaining in C++. - // See pytorch/pytorch#176117. - ops.def( - "scaled_fp4_quant.out(Tensor input," - " Tensor input_scale, bool " - "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) " - "-> ()"); - - // Compute NVFP4 experts quantization. - ops.def( - "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale," - "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts) -> ()"); - - // Fused SiLU+Mul+NVFP4 experts quantization. - ops.def( - "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! " - "output_scale," - "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts) -> ()"); - - // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs). - ops.def( - "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," - "Tensor input, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); - - // Fused SiLU+Mul+MXFP4 experts quantization. - ops.def( - "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " - "output_scale," - "Tensor input, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); - - // Fused SiLU+Mul+NVFP4 quantization. - ops.def( - "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, " - "Tensor input, Tensor input_global_scale) -> ()"); - - // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices - // of the given capability - ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool"); - - // CUTLASS w4a8 GEMM - ops.def( - "cutlass_w4a8_mm(" - " Tensor A," - " Tensor B," - " Tensor group_scales," - " int group_size," - " Tensor channel_scales," - " Tensor token_scales," - " ScalarType? out_type," - " str? maybe_schedule" - ") -> Tensor"); - - // pack scales - ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor"); - - // encode and reorder weight matrix - ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor"); - - // CUTLASS w4a8 grouped GEMM - ops.def( - "cutlass_w4a8_moe_mm(" - " Tensor! out_tensors," - " Tensor a_tensors," - " Tensor b_tensors," - " Tensor a_scales," - " Tensor b_scales," - " Tensor b_group_scales," - " int b_group_size," - " Tensor expert_offsets," - " Tensor problem_sizes," - " Tensor a_strides," - " Tensor b_strides," - " Tensor c_strides," - " Tensor group_scale_strides," - " str? maybe_schedule" - ") -> ()"); - - ops.def( - "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, " - "Tensor)"); - - // SM100 CUTLASS MLA decode - // conditionally compiled so impl registrations are in source file - ops.def( - "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope," - " Tensor q_pe, Tensor kv_c_and_k_pe_cache," - " Tensor seq_lens, Tensor page_table," - " Tensor workspace, float scale," - " int num_kv_splits) -> ()"); - - ops.def( - "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches," - " int sm_count, int num_kv_splits) " - "-> int"); - // Quantized GEMM for AWQ. - ops.def( - "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, " - "Tensor _zeros, SymInt split_k_iters) -> Tensor"); - - // Dequantization for AWQ. - ops.def( - "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, " - "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor"); - - // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). - // conditionally compiled so impl registration is in source file - ops.def( - "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, " - "bool enable_pdl=False) -> ()"); - - // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). - // conditionally compiled so impl registration is in source file - ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); -#endif - - // Merge attn states - // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 - // can be used to combine partial attention results (in the split-KV case) - ops.def( - "merge_attn_states(" - " Tensor! output," - " Tensor!? output_lse," - " Tensor prefix_output," - " Tensor prefix_lse," - " Tensor suffix_output," - " Tensor suffix_lse," - " int!? prefill_tokens_with_context," - " Tensor? output_scale=None) -> ()"); - - // Hadamard transforms - // conditionally compiled so impl registration is in source file - ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor"); - - // Apply Root Mean Square (RMS) Normalization to the input tensor. - ops.def( - "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) " - "-> " - "()"); - - // In-place fused Add and RMS Normalization. - ops.def( - "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " - "float epsilon) -> ()"); - - // Layernorm-quant - // Apply Root Mean Square (RMS) Normalization to the input tensor. - ops.def( - "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, " - "Tensor scale, float epsilon) -> " - "()"); - - // In-place fused Add and RMS Normalization. - ops.def( - "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, " - "Tensor! residual, Tensor weight, " - "Tensor scale, float epsilon) -> ()"); - - // Fused Layernorm + Quant kernels - ops.def( - "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, " - "Tensor weight, Tensor! scale, float epsilon, " - "Tensor? scale_ub, Tensor!? residual) -> ()"); - - // Fused Layernorm + Block quant kernels - ops.def( - "rms_norm_per_block_quant(Tensor! result, Tensor input, " - "Tensor weight, Tensor! scale, float epsilon, " - "Tensor? scale_ub, Tensor!? residual, int group_size, " - "bool is_scale_transposed) -> ()"); - - // Fused SiLU+Mul + per-block quantization - ops.def( - "silu_and_mul_per_block_quant(" - "Tensor! out, " - "Tensor input, " - "Tensor! scales, " - "int group_size, " - "Tensor? scale_ub=None, " - "bool is_scale_transposed=False) -> ()"); - - // Rotary embedding - // Apply GPT-NeoX or GPT-J style rotary embedding to query and key. - ops.def( - "rotary_embedding(Tensor positions, Tensor! query," - " Tensor!? key, int head_size," - " Tensor cos_sin_cache, bool is_neox, int " - "rope_dim_offset=0, bool inverse=False) -> ()"); - - // Function for fused QK Norm and RoPE - ops.def( - "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " - "int num_heads_k, int num_heads_v, int head_dim, float eps, " - "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " - "bool is_neox, Tensor position_ids, " - "int forced_token_heads_per_warp=-1) -> ()"); - - // q_head_padded is the padded Q head count of the returned tensor, or 0 to - // do the KV insert alone and return an empty tensor. The Q knobs are - // independent: apply_q_norm and apply_q_rope each drop that step for Q - // alone (KV is always rotated), and is_q_interleaved reads and writes Q in - // FlashMLA's mega-attention chunk-interleaved layout, which moves the - // padding heads to the tail of every head-dim chunk. - ops.def( - "fused_deepseek_v4_kv_rope_insert(" - "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " - "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, " - "bool kv_mxfp8=False) -> " - "()"); - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" - "Tensor q_in, Tensor kv, Tensor! k_cache, " - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "int q_head_padded, float eps, int cache_block_size, " - "bool apply_q_norm=True, bool kv_mxfp8=False, bool apply_q_rope=True, " - "bool is_q_interleaved=False) -> Tensor"); - - // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate - // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(" - "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, " - "Tensor position_ids, Tensor cos_sin_cache, float eps, " - "int cache_block_size, bool apply_q_norm=True) -> ()"); - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(" - "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, " - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " - "int cache_block_size, bool apply_q_norm=True) -> ()"); - - // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion. - ops.def( - "fused_kimi_k3_mla_key_concat_kv_cache_insert(" - "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " - "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_key_concat_ds_mla_insert(" - "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " - "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_kv_concat(Tensor k_nope, Tensor k_pe, Tensor! k_out) " - "-> ()"); - ops.def( - "fused_kimi_k3_mla_kv_concat_quant_fp8(" - "Tensor k_nope, Tensor k_pe, Tensor v, Tensor! k_fp8, Tensor! v_fp8) " - "-> ()"); - ops.def( - "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" - "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, " - "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, " - "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, " - "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, " - "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); - - // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the - // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla). - ops.def( - "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(" - "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " - "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(" - "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " - "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " - "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, " - "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(" - "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " - "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - -#ifndef USE_ROCM - ops.def( - "minimax_allreduce_rms_qk(" - "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " - "Tensor workspace, int q_size, int kv_size, int rank, int nranks, " - "float eps) -> (Tensor, Tensor)"); -#endif - - // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert. - ops.def( - "fused_minimax_m3_qknorm_rope_kv_insert(" - "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, " - "Tensor cos_sin_cache, Tensor positions, int num_heads, " - "int num_kv_heads, int rotary_dim, float eps, " - "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, " - "int num_index_heads, " - "Tensor? slot_mapping, Tensor? index_slot_mapping, " - "Tensor!? kv_cache, Tensor!? index_cache, " - "int block_size, Tensor!? q_out, Tensor!? index_q_out, " - "str kv_cache_dtype, bool skip_index_branch=False, " - "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, " - "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()"); - -#ifdef VLLM_ENABLE_FUSED_KDA_DECODE - ops.def( - "fused_kda_decode(" - "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, " - "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, " - "Tensor state_indices, Tensor! state, Tensor! out, " - "float? lower_bound=None, Tensor? output_gate=None, " - "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()"); -#endif - -#ifdef VLLM_ENABLE_FUSED_GDN_DECODE - ops.def( - "fused_gdn_decode_post_conv_mtp(" - "Tensor mixed_qkv, Tensor a, Tensor b, Tensor A_log, Tensor dt_bias, " - "Tensor state_indices, Tensor cu_seqlens, Tensor num_accepted_tokens, " - "Tensor! state, Tensor output_gate, Tensor norm_weight, Tensor! out, " - "float scale, float norm_eps=1e-5, " - "str output_gate_activation='silu') -> ()"); -#endif - -#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK - ops.def( - "fused_kda_prologue(" - "Tensor q, Tensor k, Tensor v, Tensor raw_g, Tensor raw_beta, " - "Tensor A_log, Tensor dt_bias, Tensor! qg, Tensor! w, Tensor! u, " - "Tensor! kg_t, Tensor! aqk, Tensor! decay, Tensor cu_seqlens, " - "Tensor chunk_indices, Tensor? conv_weight, Tensor(e!)? conv_state, " - "Tensor? conv_state_indices, Tensor? conv_has_initial_state, " - "float scale, float lower_bound) -> ()"); - ops.def( - "fused_kda_chunk(" - "Tensor qg, Tensor w, Tensor u, Tensor kg_t, Tensor aqk, Tensor decay, " - "Tensor? initial_state, Tensor(a!)? final_state, Tensor! out, " - "Tensor cu_seqlens, Tensor chunk_offsets, float scale, " - "Tensor(b!)? group_state, int groups, " - "Tensor(c!)? checkpoint_state=None, Tensor? checkpoint_offsets=None, " - "Tensor? checkpoint_state_indices=None, Tensor(d!)? state_cache=None, " - "Tensor? state_indices=None, Tensor? has_initial_state=None) -> ()"); -#endif - -#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES - ops.def( - "kimi_k3_attn_res(" - "Tensor! prefix, Tensor? delta, Tensor! blocks, Tensor norm_weight, " - "Tensor qk_weight, Tensor? output_norm_weight, Tensor! output, " - "int num_blocks, " - "int block_write_idx, float eps, " - "float output_norm_eps) -> ()"); -#endif - - // Apply repetition penalties to logits in-place. - ops.def( - "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " - "Tensor output_mask, Tensor repetition_penalties) -> ()"); - - // Optimized top-k per row operations. - ops.def( - "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, " - "Tensor! indices, int numRows, int stride0, " - "int stride1, int topK) -> ()"); - - ops.def( - "top_k_per_row_decode(Tensor logits, int next_n, " - "Tensor seq_lens, Tensor! indices, " - "int numRows, int stride0, int stride1, int topK) -> ()"); - - ops.def( - "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " - "Tensor workspace, int k, int max_seq_len) -> ()"); - -#ifdef VLLM_ENABLE_COOPERATIVE_TOPK - ops.def( - "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " - "Tensor workspace, int k, int max_seq_len) -> ()"); -#endif - - // Activation ops - ops.def( - "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " - "y_q, Tensor! y_s, bool use_ue8m0) -> ()"); - ops.def("weak_ref_tensor(Tensor input) -> Tensor"); - - // Activation function used in SwiGLU. - ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); - - ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); - - // SwiGLU activation with input clamping. - // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to - // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up. - ops.def( - "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, " - "float alpha=1.0, float beta=0.0) -> ()"); - - // SwiGLU activation with FP8 quantization. - ops.def( - "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); - - // Activation function used in GeGLU with `none` approximation. - ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); - - // Activation function used in GeGLU with `tanh` approximation. - ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); - - // FATReLU implementation. - ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()"); - - ops.def( - "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float " - "limit=7.0) " - "-> ()"); - - // SituGLU implementation used in Kimi models. - ops.def( - "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float " - "linear_beta=-1.0) -> ()"); - // Fused SituGLU activation + dynamic FP8 quantization for the Humming w2 path - // (writes the fp8 down input and its float32 scale). group_size=0 -> - // per-token scale [.., 1]; group_size=128 -> k-major block-FP8 scale [.., - // d/128]. - ops.def( - "situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, " - "float beta=1.0, float linear_beta=-1.0, int group_size=0, " - "Tensor? num_valid_tokens=None, int topk=1) -> ()"); - ops.def( - "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " - "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); - ops.def( - "masked_moe_activation(Tensor! out, Tensor input, Tensor " - "valid_token_counts, str activation, float clamp_limit=0.0, float " - "alpha=1.0, float beta=0.0, float situ_beta=1.0, float " - "situ_linear_beta=-1.0) -> ()"); - - // GELU implementation used in GPT-2. - ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); - - // Approximate GELU implementation. - ops.def("gelu_fast(Tensor! out, Tensor input) -> ()"); - - // Quick GELU implementation. - ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); - - // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2 - ops.def("relu_squared(Tensor! out, Tensor input) -> ()"); - - // Compute int8 quantized tensor for given scaling factor. - ops.def( - "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale," - "Tensor? azp) -> ()"); - - // Compute int8 quantized tensor and scaling factor - ops.def( - "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, " - "Tensor!? azp) -> ()"); - - // Compute FP8 quantized tensor for given scaling factor. - // Supports per-tensor, per-channel, per-token, and arbitrary 2D group - // scaling. Optional group_m/group_n specify the group shape explicitly; - // required for 1D scales to disambiguate per-channel vs per-token. - ops.def( - "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, " - "int[]? group_shape=None) -> ()"); - - // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor. - ops.def( - "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) " - "-> " - "()"); - - // Compute dynamic-per-token FP8 quantized tensor and scaling factor. - ops.def( - "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, " - "Tensor! scale, Tensor? scale_ub) -> " - "()"); - - // Quantized GEMM for GPTQ. - // Note: even though the C++ inferred schema is correct for this op, it seems - // to prevent the meta function registry. - ops.def( - "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, " - "Tensor b_gptq_scales, bool use_exllama, bool " - "use_v2_format, int bit) " - "-> Tensor"); - - // Post processing for GPTQ. - ops.def("gptq_shuffle(Tensor! q_weight, int bit) -> ()"); - - // Mamba selective scan kernel - ops.def( - "selective_scan_fwd(Tensor! u, Tensor! delta," - "Tensor! A, Tensor! B, Tensor! C," - "Tensor? D_, Tensor!? z_, Tensor? delta_bias_," - "bool delta_softplus," - "Tensor? query_start_loc," - "Tensor? cache_indices," - "Tensor? has_initial_state," - "Tensor! ssm_states," - "int null_block_id," - "int block_size," - "Tensor? block_idx_first_scheduled_token," - "Tensor? block_idx_last_scheduled_token," - "Tensor? initial_state_idx," - "Tensor? cu_chunk_seqlen," - "Tensor? last_chunk_indices) -> ()"); - - // LongCat n-gram embedding index kernel. All tensor args are marked mutable - // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table - // and n_gram_ids are actually written in place. - // Fused vocab-parallel embedding: gather the rows this rank owns and write - // zeros for the rest, so the following all-reduce reconstructs the full - // embedding. - ops.def( - "vocab_parallel_embedding(Tensor! out, Tensor input_ids, Tensor weight, " - "int org_vocab_start_index, int org_vocab_end_index, " - "int num_org_vocab_padding, int added_vocab_start_index, " - "int added_vocab_end_index) -> ()"); - ops.def( - "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, " - "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, " - "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, " - "Tensor(f!) row_indices, Tensor(g!) column_starts, " - "Tensor(h!) n_gram_ids) -> ()"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { - // LongCat n-gram embedding index kernel. - ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); - ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding)); - - // Per-token group quantization - ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); - ops.impl("per_token_group_fp8_quant_packed", - TORCH_BOX(&per_token_group_quant_8bit_packed)); - ops.impl("per_token_group_quant_int8", - TORCH_BOX(&per_token_group_quant_int8)); - -#ifndef USE_ROCM - // CUTLASS scaled_mm ops - ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm)); - ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp)); - ops.impl("cutlass_moe_mm", TORCH_BOX(&cutlass_moe_mm)); - ops.impl("get_cutlass_moe_mm_data", TORCH_BOX(&get_cutlass_moe_mm_data)); - ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets", - TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets)); - ops.impl("get_cutlass_batched_moe_mm_data", - TORCH_BOX(&get_cutlass_batched_moe_mm_data)); - - // FP4/NVFP4 ops - ops.impl("cutlass_scaled_fp4_mm", TORCH_BOX(&cutlass_scaled_fp4_mm)); - ops.impl("scaled_fp4_quant", TORCH_BOX(&scaled_fp4_quant_func)); - ops.impl("scaled_fp4_quant.out", TORCH_BOX(&scaled_fp4_quant_out)); - ops.impl("scaled_fp4_experts_quant", TORCH_BOX(&scaled_fp4_experts_quant)); - ops.impl("silu_and_mul_scaled_fp4_experts_quant", - TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); - ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); - // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). - // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. - - // AWQ ops - ops.impl("awq_gemm", TORCH_BOX(&awq_gemm)); - ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize)); - - // DSV3 fused A GEMM: conditionally compiled so impl registration is in - // source file (dsv3_fused_a_gemm.cu) -#endif - - ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states)); - - // Layernorm kernels (shared CUDA/ROCm) - ops.impl("rms_norm", TORCH_BOX(&rms_norm)); - ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm)); - - // Layernorm-quant kernels (shared CUDA/ROCm) - ops.impl("rms_norm_static_fp8_quant", TORCH_BOX(&rms_norm_static_fp8_quant)); - ops.impl("fused_add_rms_norm_static_fp8_quant", - TORCH_BOX(&fused_add_rms_norm_static_fp8_quant)); - - // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) - ops.impl("rms_norm_dynamic_per_token_quant", - TORCH_BOX(&rms_norm_dynamic_per_token_quant)); - ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant)); - ops.impl("silu_and_mul_per_block_quant", - TORCH_BOX(&silu_and_mul_per_block_quant)); - - // Positional encoding kernels (shared CUDA/ROCm) - ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); - ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); - ops.impl("fused_deepseek_v4_kv_rope_insert", - TORCH_BOX(&fused_deepseek_v4_kv_rope_insert)); - ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); - ops.impl( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert)); - ops.impl( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); - ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert", - TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); - ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", - TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); - ops.impl("fused_kimi_k3_mla_kv_concat", - TORCH_BOX(&fused_kimi_k3_mla_kv_concat)); - ops.impl("fused_kimi_k3_mla_kv_concat_quant_fp8", - TORCH_BOX(&fused_kimi_k3_mla_kv_concat_quant_fp8)); - ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", - TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); - ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", - TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert)); - ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", - TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert)); - ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", - TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert)); -#ifndef USE_ROCM - ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); -#endif - ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", - TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); -#ifdef VLLM_ENABLE_FUSED_KDA_DECODE - ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); -#endif - -#ifdef VLLM_ENABLE_FUSED_GDN_DECODE - ops.impl("fused_gdn_decode_post_conv_mtp", - TORCH_BOX(&fused_gdn_decode_post_conv_mtp)); -#endif - -#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK - ops.impl("fused_kda_prologue", TORCH_BOX(&fused_kda_prologue)); - ops.impl("fused_kda_chunk", TORCH_BOX(&fused_kda_chunk)); -#endif - -#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES - ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); -#endif - - // Sampler kernels (shared CUDA/ROCm) - ops.impl("apply_repetition_penalties_", - TORCH_BOX(&apply_repetition_penalties_)); - ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); - ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); - ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); -#ifdef VLLM_ENABLE_COOPERATIVE_TOPK - ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); -#endif - - // Activation kernels (shared CUDA/ROCm) - ops.impl("persistent_masked_m_silu_mul_quant", - TORCH_BOX(&persistent_masked_m_silu_mul_quant)); - ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor)); - ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant)); - ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul)); - ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu)); - ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul)); - ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); - ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); - ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); - ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); - ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant)); - ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); - ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation)); - ops.impl("gelu_new", TORCH_BOX(&gelu_new)); - ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); - ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); - ops.impl("relu_squared", TORCH_BOX(&relu_squared)); - ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp)); - - // INT8 quantization kernels - ops.impl("static_scaled_int8_quant", TORCH_BOX(&static_scaled_int8_quant)); - ops.impl("dynamic_scaled_int8_quant", TORCH_BOX(&dynamic_scaled_int8_quant)); - - // FP8 quantization kernels - ops.impl("static_scaled_fp8_quant", TORCH_BOX(&static_scaled_fp8_quant)); - ops.impl("dynamic_scaled_fp8_quant", TORCH_BOX(&dynamic_scaled_fp8_quant)); - ops.impl("dynamic_per_token_scaled_fp8_quant", - TORCH_BOX(&dynamic_per_token_scaled_fp8_quant)); - - // GPTQ kernels - ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm)); - ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle)); - - // Mamba kernels - ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { - ops.impl("get_cuda_view_from_cpu_tensor", - TORCH_BOX(&get_cuda_view_from_cpu_tensor)); -} - -STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) { - cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int"); - cuda_utils.def( - "get_max_shared_memory_per_block_device_attribute(int device_id) -> int"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd, - cuda_utils) { - cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute)); - cuda_utils.impl("get_max_shared_memory_per_block_device_attribute", - TORCH_BOX(&get_max_shared_memory_per_block_device_attribute)); -} - -// These capability-check functions take only primitive args (no tensors), so -// there is no device to dispatch on. CompositeExplicitAutograd makes them -// available for all backends. This is the stable ABI equivalent of calling -// ops.impl("op_name", &func) without a dispatch key in the non-stable API. -STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { -#ifndef USE_ROCM - ops.impl("cutlass_scaled_mm_supports_fp8", - TORCH_BOX(&cutlass_scaled_mm_supports_fp8)); - ops.impl("cutlass_group_gemm_supported", - TORCH_BOX(&cutlass_group_gemm_supported)); - ops.impl("cutlass_scaled_mm_supports_block_fp8", - TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8)); - ops.impl("cutlass_scaled_mm_supports_fp4", - TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); -#endif -} - -// Cache ops -STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { - // Swap in (out) the cache blocks from src to dst. - ops.def( - "swap_blocks(Tensor src, Tensor! dst," - " int block_size_in_bytes, Tensor block_mapping) -> ()"); - - // Batch swap: submit all block copies in a single driver call. - ops.def( - "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," - " Tensor sizes," - " bool is_src_access_order_any=False) -> ()"); - - // Reshape the key and value tensors and cache them. - ops.def( - "reshape_and_cache(Tensor key, Tensor value," - " Tensor! key_cache, Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - - // Reshape the key and value tensors and cache them. - ops.def( - "reshape_and_cache_flash(Tensor key, Tensor value," - " Tensor! key_cache," - " Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - - // Concat kv_c and k_pe and cache them. - ops.def( - "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," - " Tensor! kv_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor scale) -> ()"); - - // Grouped concat_and_cache_mla across all layers. Each layer's cache base - // pointer and optional plain-FP8 scale are read from device tensors. - ops.def( - "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe," - " Tensor kv_cache_ptrs," - " Tensor slot_mapping," - " int block_size, int block_stride," - " int entry_stride," - " Tensor? kv_scales=None," - " str kv_cache_dtype='auto') -> ()"); - -#ifndef USE_ROCM - ops.def( - "hisparse_resolve_residency(Tensor host_cache," - " Tensor! hot_cache," - " Tensor hot_block_table," - " Tensor global_indices," - " Tensor! hot_indices," - " Tensor! device_global_indices," - " Tensor! lru_slots," - " Tensor? request_state_indices," - " int region_stride," - " int max_union_rows," - " Tensor(a!)? miss_mask=None," - " Tensor(b!)? stats=None," - " Tensor(c!)? attention_indices=None," - " int attention_block_stride=0," - " Tensor? request_ids=None," - " Tensor? source_block_table=None," - " int source_block_size=0," - " Tensor(d!)? resolved_global_indices=None," - " Tensor(e!)? valid_counts=None," - " Tensor(f!)? swap_host_physical_rows=None," - " Tensor(g!)? swap_device_physical_rows=None," - " Tensor(h!)? swap_counts=None," - " Tensor? resident_block_table=None," - " int resident_block_size=0," - " int resident_null_block=0) -> ()"); - - ops.def( - "hisparse_invalidate_written_slots(Tensor! device_global_indices," - " Tensor request_state_indices," - " Tensor req_id_per_token," - " Tensor written_slots) -> ()"); - - ops.def( - "hisparse_gather_plan(Tensor host_cache," - " Tensor! hot_cache," - " Tensor global_indices," - " Tensor hot_indices," - " Tensor miss_mask," - " Tensor? request_state_indices," - " Tensor(a!)? attention_indices=None," - " int attention_block_stride=0) -> ()"); - - ops.def( - "hisparse_gather_compact(Tensor host_cache," - " Tensor! hot_cache," - " Tensor miss_global_indices," - " Tensor miss_hot_indices," - " Tensor miss_counts) -> ()"); - -#endif // !USE_ROCM - - // Rotate Q and K, then write to kv cache for MLA - ops.def( - "concat_and_cache_mla_rope_fused(" - " Tensor positions," - " Tensor! q_pe," - " Tensor! k_pe," - " Tensor kv_c," - " Tensor cos_sin_cache," - " bool is_neox," - " Tensor slot_mapping," - " Tensor! kv_cache," - " str kv_cache_dtype," - " Tensor kv_cache_scale) -> ()"); - - // Convert the key and value cache to fp8 data type. - ops.def( - "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " - "str kv_cache_dtype) -> ()"); - - // Gather cache blocks from src_cache to dst, dequantizing from - // src_cache's dtype to dst's dtype if necessary. - ops.def( - "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " - " Tensor block_table, Tensor cu_seq_lens, " - " Tensor token_to_seq, " - " int num_tokens, " - " str kv_cache_dtype, " - " Tensor scale, Tensor? seq_starts) -> ()"); - - ops.def( - "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " - "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); - - ops.def( - "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " - "seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? " - "device_row_ids=None) -> ()"); - - ops.def( - "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor workspace_starts, int batch_size) -> ()"); - - ops.def( - "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " - "slot_mapping, " - "int quant_block_size, str kv_cache_dtype) -> ()"); - - ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); - - ops.def( - "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " - "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); -} - -STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) { - custom_ar.def( - "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " - "int rank, bool fully_connected) -> int"); - custom_ar.def( - "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " - "int reg_buffer_sz_bytes) -> ()"); - custom_ar.def("dispose(int fa) -> ()"); - custom_ar.def("meta_size() -> int"); - custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()"); - custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])"); - custom_ar.def( - "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()"); - custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)"); - custom_ar.def("open_mem_handle(Tensor mem_handle) -> int"); - custom_ar.def("free_shared_buffer(int ptr) -> ()"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) { - custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar)); - custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) { - custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) { - custom_ar.impl("dispose", TORCH_BOX(&dispose)); - custom_ar.impl("meta_size", TORCH_BOX(&meta_size)); - custom_ar.impl("register_buffer", TORCH_BOX(®ister_buffer)); - custom_ar.impl("get_graph_buffer_ipc_meta", - TORCH_BOX(&get_graph_buffer_ipc_meta)); - custom_ar.impl("register_graph_buffers", TORCH_BOX(®ister_graph_buffers)); - custom_ar.impl("allocate_shared_buffer_and_handle", - TORCH_BOX(&allocate_shared_buffer_and_handle)); - custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { - ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { - ops.impl("swap_blocks", TORCH_BOX(&swap_blocks)); - ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); - ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); - ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); - ops.impl("concat_and_cache_mla_grouped", - TORCH_BOX(&concat_and_cache_mla_grouped)); - -#ifndef USE_ROCM - ops.impl("hisparse_resolve_residency", - TORCH_BOX(&hisparse_resolve_residency)); - ops.impl("hisparse_invalidate_written_slots", - TORCH_BOX(&hisparse_invalidate_written_slots)); - ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan)); - ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact)); -#endif // !USE_ROCM - ops.impl("concat_and_cache_mla_rope_fused", - TORCH_BOX(&concat_and_cache_mla_rope_fused)); - ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); - ops.impl("gather_and_maybe_dequant_cache", - TORCH_BOX(&gather_and_maybe_dequant_cache)); - ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache)); - ops.impl("cp_gather_and_upconvert_fp8_kv_cache", - TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache)); - ops.impl("cp_gather_and_upconvert_nvfp4_kv_cache", - TORCH_BOX(&cp_gather_and_upconvert_nvfp4_kv_cache)); - ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache)); - ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q)); - ops.impl("cp_gather_indexer_k_quant_cache", - TORCH_BOX(&cp_gather_indexer_k_quant_cache)); -} - -REGISTER_EXTENSION(_C_stable_libtorch) diff --git a/csrc/libtorch_stable/torch_utils.h b/csrc/libtorch_stable/torch_utils.h deleted file mode 100644 index f02346739e..0000000000 --- a/csrc/libtorch_stable/torch_utils.h +++ /dev/null @@ -1,88 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -#include -#include - -#include -#include -#include -#include - -// Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED. -#define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \ - STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__) - -// Device properties cache for stable ABI compatibility. -// Uses raw CUDA/HIP APIs instead of ATen functions. -// Using inline ensures a single instance across all translation units. -inline std::deque device_flags; -inline std::vector device_properties; -inline std::once_flag vectors_init_flag; - -inline void do_init_device_vectors() { - int device_count; - cudaError_t err = cudaGetDeviceCount(&device_count); - if (err != cudaSuccess) { - STD_TORCH_CHECK(false, "cudaGetDeviceCount failed: " + - std::string(cudaGetErrorString(err))); - } - device_flags.resize(device_count); - device_properties.resize(device_count); -} - -inline void initDeviceVectors() { - std::call_once(vectors_init_flag, do_init_device_vectors); -} - -inline void initDeviceProperty(int device_index) { - cudaDeviceProp device_prop{}; - cudaError_t err = cudaGetDeviceProperties(&device_prop, device_index); - if (err != cudaSuccess) { - STD_TORCH_CHECK(false, "cudaGetDeviceProperties failed: " + - std::string(cudaGetErrorString(err))); - } - device_properties[device_index] = device_prop; -} - -// Get device properties using raw CUDA/HIP APIs (stable ABI compatible). -// Caches results per device so cudaGetDeviceProperties is called at most once -// per device. -inline cudaDeviceProp* get_device_prop() { - initDeviceVectors(); - int device_index; - cudaError_t err = cudaGetDevice(&device_index); - if (err != cudaSuccess) { - STD_TORCH_CHECK( - false, "cudaGetDevice failed: " + std::string(cudaGetErrorString(err))); - } - STD_TORCH_CHECK(device_index >= 0 && static_cast(device_index) < - device_properties.size(), - "CUDA device index " + std::to_string(device_index) + - " out of range [0, " + - std::to_string(device_properties.size()) + ")"); - - std::call_once(device_flags[device_index], initDeviceProperty, device_index); - return &device_properties[device_index]; -} - -// Utility to get the current CUDA stream for a given device using stable APIs. -// Returns a cudaStream_t for use in kernel launches. -inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) { - void* stream_ptr = nullptr; - TORCH_ERROR_CODE_CHECK( - aoti_torch_get_current_cuda_stream(device_index, &stream_ptr)); - return reinterpret_cast(stream_ptr); -} - -// Utility to get the current cuBLAS handle using stable APIs. -inline cublasHandle_t get_current_cuda_blas_handle() { - void* blas_handle_ptr = nullptr; - TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&blas_handle_ptr)); - return reinterpret_cast(blas_handle_ptr); -} diff --git a/csrc/libtorch_stable/mamba/selective_scan.h b/csrc/mamba/selective_scan.h similarity index 100% rename from csrc/libtorch_stable/mamba/selective_scan.h rename to csrc/mamba/selective_scan.h diff --git a/csrc/libtorch_stable/mamba/selective_scan_fwd.cu b/csrc/mamba/selective_scan_fwd.cu similarity index 99% rename from csrc/libtorch_stable/mamba/selective_scan_fwd.cu rename to csrc/mamba/selective_scan_fwd.cu index c17b06a79b..cfc1f038f1 100644 --- a/csrc/libtorch_stable/mamba/selective_scan_fwd.cu +++ b/csrc/mamba/selective_scan_fwd.cu @@ -14,7 +14,7 @@ #endif #include "selective_scan.h" -#include "static_switch.h" +#include "../dispatch_utils.h" template; constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t); dim3 grid(params.batch, params.dim / kNRows); diff --git a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu b/csrc/minimax_reduce_rms_kernel.cu similarity index 99% rename from csrc/libtorch_stable/minimax_reduce_rms_kernel.cu rename to csrc/minimax_reduce_rms_kernel.cu index 58d61b353d..5a3da3dadf 100644 --- a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu +++ b/csrc/minimax_reduce_rms_kernel.cu @@ -27,7 +27,7 @@ #include #include -#include "cuda_compat.h" +#include "core/cuda_compat.h" #include "minimax_reduce_rms_kernel.h" #include diff --git a/csrc/libtorch_stable/minimax_reduce_rms_kernel.h b/csrc/minimax_reduce_rms_kernel.h similarity index 100% rename from csrc/libtorch_stable/minimax_reduce_rms_kernel.h rename to csrc/minimax_reduce_rms_kernel.h diff --git a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu b/csrc/moe/grouped_topk_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/grouped_topk_kernels.cu rename to csrc/moe/grouped_topk_kernels.cu index 852b2b468a..ea89aa38a1 100644 --- a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu +++ b/csrc/moe/grouped_topk_kernels.cu @@ -22,7 +22,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore b/csrc/moe/marlin_moe_wna16/.gitignore similarity index 73% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore rename to csrc/moe/marlin_moe_wna16/.gitignore index 7dc482a894..ba805f9250 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore +++ b/csrc/moe/marlin_moe_wna16/.gitignore @@ -1,3 +1,2 @@ sm*_kernel_*.cu kernel_selector.h -kernel_*.cu diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h b/csrc/moe/marlin_moe_wna16/kernel.h similarity index 95% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h rename to csrc/moe/marlin_moe_wna16/kernel.h index 25c57f3b2d..1fd00c12d1 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h +++ b/csrc/moe/marlin_moe_wna16/kernel.h @@ -3,8 +3,8 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "libtorch_stable/quantization/marlin/marlin.cuh" -#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" +#include "quantization/marlin/marlin.cuh" +#include "quantization/marlin/marlin_dtypes.cuh" #include "core/scalar_type.hpp" #define MARLIN_KERNEL_PARAMS \ diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h b/csrc/moe/marlin_moe_wna16/marlin_template.h similarity index 99% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h rename to csrc/moe/marlin_moe_wna16/marlin_template.h index aac9bd157f..e2d80ce6bf 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h +++ b/csrc/moe/marlin_moe_wna16/marlin_template.h @@ -23,10 +23,10 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "libtorch_stable/quantization/marlin/marlin.cuh" -#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" -#include "libtorch_stable/quantization/marlin/dequant.h" -#include "libtorch_stable/quantization/marlin/marlin_mma.h" +#include "quantization/marlin/marlin.cuh" +#include "quantization/marlin/marlin_dtypes.cuh" +#include "quantization/marlin/dequant.h" +#include "quantization/marlin/marlin_mma.h" #include "core/scalar_type.hpp" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu b/csrc/moe/marlin_moe_wna16/ops.cu similarity index 99% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu rename to csrc/moe/marlin_moe_wna16/ops.cu index 48de07dc75..8336b46e15 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu +++ b/csrc/moe/marlin_moe_wna16/ops.cu @@ -32,7 +32,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::value || \ diff --git a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh b/csrc/moe/moeTopKFuncs.cuh similarity index 100% rename from csrc/libtorch_stable/moe/moeTopKFuncs.cuh rename to csrc/moe/moeTopKFuncs.cuh diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/moe/moe_align_sum_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/moe_align_sum_kernels.cu rename to csrc/moe/moe_align_sum_kernels.cu index 5c9a522feb..a2424f8339 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/moe/moe_align_sum_kernels.cu @@ -8,11 +8,11 @@ #include #include -#include "../../cuda_compat.h" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/quantization/vectorization.cuh" -#include "libtorch_stable/torch_utils.h" +#include "core/cuda_compat.h" +#include "core/utils.hpp" +#include "dispatch_utils.h" +#include "quantization/vectorization.cuh" +#include "torch_utils.h" #define CEILDIV(x, y) (((x) + (y) - 1) / (y)) diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/moe/moe_ops.h similarity index 100% rename from csrc/libtorch_stable/moe/moe_ops.h rename to csrc/moe/moe_ops.h diff --git a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu b/csrc/moe/moe_permute_unpermute_op.cu similarity index 99% rename from csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu rename to csrc/moe/moe_permute_unpermute_op.cu index 957cd60e49..31c937b2a2 100644 --- a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu +++ b/csrc/moe/moe_permute_unpermute_op.cu @@ -7,8 +7,8 @@ #include #include "core/registration.h" -#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" -#include "libtorch_stable/torch_utils.h" +#include "moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" +#include "torch_utils.h" #include diff --git a/csrc/libtorch_stable/moe/moe_wna16.cu b/csrc/moe/moe_wna16.cu similarity index 99% rename from csrc/libtorch_stable/moe/moe_wna16.cu rename to csrc/moe/moe_wna16.cu index 9345a7c9f7..a8971041d0 100644 --- a/csrc/libtorch_stable/moe/moe_wna16.cu +++ b/csrc/moe/moe_wna16.cu @@ -8,7 +8,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "moe_wna16_utils.h" #define DIVIDE(x, size) (((x) + (size) - 1) / (size)) diff --git a/csrc/libtorch_stable/moe/moe_wna16_utils.h b/csrc/moe/moe_wna16_utils.h similarity index 100% rename from csrc/libtorch_stable/moe/moe_wna16_utils.h rename to csrc/moe/moe_wna16_utils.h diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h b/csrc/moe/permute_unpermute_kernels/dispatch.h similarity index 100% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h rename to csrc/moe/permute_unpermute_kernels/dispatch.h diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu similarity index 98% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu rename to csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu index 8d8bdf8baa..56c1e2bddf 100644 --- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu @@ -1,7 +1,7 @@ #include #include -#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" +#include "moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" // moe_permute kernels require at least CUDA 12.0 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h similarity index 97% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h rename to csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h index 53afad2f5d..1496eccb8d 100644 --- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h @@ -10,7 +10,7 @@ #include "cutlass/array.h" #include "cutlass/numeric_size.h" -#include "libtorch_stable/moe/permute_unpermute_kernels/dispatch.h" +#include "moe/permute_unpermute_kernels/dispatch.h" template inline T* get_ptr(torch::stable::Tensor& t) { diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl similarity index 100% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl rename to csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/topk_softmax_kernels.cu rename to csrc/moe/topk_softmax_kernels.cu index 098644a1b9..0e46277017 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/moe/topk_softmax_kernels.cu @@ -24,9 +24,9 @@ #include #include -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" #include "../cub_helpers.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #ifndef USE_ROCM #include diff --git a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu b/csrc/moe/topk_softplus_sqrt_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu rename to csrc/moe/topk_softplus_sqrt_kernels.cu index fbb07216d2..237f76334a 100644 --- a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/moe/topk_softplus_sqrt_kernels.cu @@ -25,9 +25,9 @@ #include #include -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" #include "../cub_helpers.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #ifndef USE_ROCM #include #include diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp similarity index 100% rename from csrc/libtorch_stable/moe/torch_bindings.cpp rename to csrc/moe/torch_bindings.cpp diff --git a/csrc/libtorch_stable/ngram_embedding_kernels.cu b/csrc/ngram_embedding_kernels.cu similarity index 100% rename from csrc/libtorch_stable/ngram_embedding_kernels.cu rename to csrc/ngram_embedding_kernels.cu diff --git a/csrc/libtorch_stable/nvfp4_ds_mla_cache.h b/csrc/nvfp4_ds_mla_cache.h similarity index 100% rename from csrc/libtorch_stable/nvfp4_ds_mla_cache.h rename to csrc/nvfp4_ds_mla_cache.h diff --git a/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu b/csrc/nvfp4_ds_mla_cache_kernels.cu similarity index 100% rename from csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu rename to csrc/nvfp4_ds_mla_cache_kernels.cu diff --git a/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu b/csrc/nvfp4_kv_cache_kernels.cu similarity index 99% rename from csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu rename to csrc/nvfp4_kv_cache_kernels.cu index bb11e474a6..341c4cdd93 100644 --- a/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu +++ b/csrc/nvfp4_kv_cache_kernels.cu @@ -15,10 +15,10 @@ // - reciprocal_approximate_ftz() for fast reciprocal #define NVFP4_ENABLE_ELTS16 1 -#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh" +#include "quantization/fp4/nvfp4_utils.cuh" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/torch_utils.h" +#include "dispatch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/ops.h b/csrc/ops.h index d7ee2d080c..3266d28b1e 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -1,72 +1,822 @@ #pragma once +#include +#include +#include + #include #include -#include -#include - -#include "core/scalar_type.hpp" - #include -// rms_norm and fused_add_rms_norm declarations also exist in -// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here -// because the CPU build still uses these torch::Tensor declarations. -void rms_norm(torch::Tensor& out, torch::Tensor& input, - std::optional weight, double epsilon); +#include -void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, - std::optional weight, double epsilon); +inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) { + // Ensure tensor is on CUDA + STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device"); -// rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable -// ABI for CUDA). It remains here because the CPU build still uses these -// torch::Tensor declarations. -void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, - std::optional key, int64_t head_size, - torch::Tensor& cos_sin_cache, bool is_neox, - int64_t rope_dim_offset, bool inverse); + // Get the raw data pointer + void* data_ptr = tensor.mutable_data_ptr(); -void silu_and_mul(torch::Tensor& out, torch::Tensor& input); + /// Create a new tensor from the raw data pointer + return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(), + tensor.device(), tensor.scalar_type()); +} -void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit, - double alpha = 1.0, double beta = 0.0); +void per_token_group_quant_fp8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double fp8_min, + double fp8_max, bool scale_ue8m0, + bool dummy_is_scale_transposed, + bool dummy_is_tma_aligned); -void gelu_and_mul(torch::Tensor& out, torch::Tensor& input); +// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales. +void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s_packed, + int64_t group_size, double eps, + double min_8bit, double max_8bit); -void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input); +void per_token_group_quant_int8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double int8_min, + double int8_max); -void gelu_tanh(torch::Tensor& out, torch::Tensor& input); +#ifndef USE_ROCM +bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); +bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability); +bool cutlass_group_gemm_supported(int64_t cuda_device_capability); -void gelu_new(torch::Tensor& out, torch::Tensor& input); +void cutlass_scaled_mm(torch::stable::Tensor& out, + torch::stable::Tensor const& a, + torch::stable::Tensor const& b, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + std::optional const& bias); -void gelu_fast(torch::Tensor& out, torch::Tensor& input); +void cutlass_moe_mm(torch::stable::Tensor& out_tensors, + torch::stable::Tensor const& a_tensors, + torch::stable::Tensor const& b_tensors, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + torch::stable::Tensor const& expert_offsets, + torch::stable::Tensor const& problem_sizes, + torch::stable::Tensor const& a_strides, + torch::stable::Tensor const& b_strides, + torch::stable::Tensor const& c_strides, bool per_act_token, + bool per_out_ch); -void gelu_quick(torch::Tensor& out, torch::Tensor& input); +void cutlass_scaled_mm_azp(torch::stable::Tensor& out, + torch::stable::Tensor const& a, + torch::stable::Tensor const& b, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + torch::stable::Tensor const& azp_adj, + std::optional const& azp, + std::optional const& bias); -void relu_squared(torch::Tensor& out, torch::Tensor& input); +void get_cutlass_moe_mm_data( + const torch::stable::Tensor& topk_ids, + torch::stable::Tensor& expert_offsets, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, + torch::stable::Tensor& input_permutation, + torch::stable::Tensor& output_permutation, const int64_t num_experts, + const int64_t n, const int64_t k, + const std::optional& blockscale_offsets, + const bool is_gated); -void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, - torch::Tensor const& scale, - std::optional const& azp); +void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( + const torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k, + const bool swap_ab); -void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, - torch::Tensor& scales, - std::optional const& azp); +void get_cutlass_batched_moe_mm_data( + torch::stable::Tensor& expert_offsets, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, + const torch::stable::Tensor& expert_num_tokens, + const int64_t num_local_experts, const int64_t padded_m, const int64_t n, + const int64_t k); -torch::Tensor dynamic_4bit_int_moe_cpu( - torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, - torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, - int64_t intermediate_size, int64_t group_size, - bool apply_router_weight_on_input, int64_t activation_kind); +// FP4/NVFP4 ops +bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability); + +void cutlass_scaled_fp4_mm(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha); + +void cutlass_fp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& alphas, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets); + +std::tuple scaled_fp4_quant_func( + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout); + +void scaled_fp4_quant_out(torch::stable::Tensor const& input, + torch::stable::Tensor const& input_scale, + bool is_sf_swizzled_layout, + torch::stable::Tensor& output, + torch::stable::Tensor& output_scale); + +void scaled_fp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_global_scale, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts); + +void silu_and_mul_scaled_fp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_global_scale, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts); + +void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, + torch::stable::Tensor& output_block_scale, + torch::stable::Tensor& input, + torch::stable::Tensor& input_global_scale); + +void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets); + +// AWQ ops +torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats, + torch::stable::Tensor _kernel, + torch::stable::Tensor _scaling_factors, + torch::stable::Tensor _zeros, + int64_t split_k_iters); + +torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, + torch::stable::Tensor _scaling_factors, + torch::stable::Tensor _zeros, + int64_t split_k_iters, int64_t thx, + int64_t thy); + +// DSV3 fused A GEMM: conditionally compiled so declaration and impl +// registration are in the source file (dsv3_fused_a_gemm.cu) + +#endif + +// CPU tensor -> CUDA UVA view (shared CUDA/ROCm) +torch::stable::Tensor get_cuda_view_from_cpu_tensor( + torch::stable::Tensor& cpu_tensor); + +// Attention kernels (shared CUDA/ROCm) +void merge_attn_states( + torch::stable::Tensor& output, + std::optional output_lse, + const torch::stable::Tensor& prefix_output, + const torch::stable::Tensor& prefix_lse, + const torch::stable::Tensor& suffix_output, + const torch::stable::Tensor& suffix_lse, + const std::optional prefill_tokens_with_context, + const std::optional& output_scale = std::nullopt); + +torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, + bool inplace); + +// Layernorm kernels (shared CUDA/ROCm) +void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input, + std::optional weight, double epsilon); + +void fused_add_rms_norm(torch::stable::Tensor& input, + torch::stable::Tensor& residual, + std::optional weight, + double epsilon); + +// Layernorm-quant kernels (shared CUDA/ROCm) +void rms_norm_static_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& weight, + torch::stable::Tensor& scale, double epsilon); + +void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& residual, + torch::stable::Tensor& weight, + torch::stable::Tensor& scale, + double epsilon); + +// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) +void rms_norm_dynamic_per_token_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, torch::stable::Tensor& scales, + double const var_epsilon, std::optional scale_ub, + std::optional residual); + +void rms_norm_per_block_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, + torch::stable::Tensor& scales, + double const var_epsilon, + std::optional scale_ub, + std::optional residual, + int64_t group_size, bool is_scale_transposed); + +void silu_and_mul_per_block_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scales, + int64_t group_size, + std::optional scale_ub, + bool is_scale_transposed); + +// Positional encoding kernels (shared CUDA/ROCm) +void rotary_embedding(torch::stable::Tensor& positions, + torch::stable::Tensor& query, + std::optional key, + int64_t head_size, torch::stable::Tensor& cos_sin_cache, + bool is_neox, int64_t rope_dim_offset, bool inverse); + +void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, + int64_t num_heads_k, int64_t num_heads_v, + int64_t head_dim, double eps, + torch::stable::Tensor& q_weight, + torch::stable::Tensor& k_weight, + torch::stable::Tensor& cos_sin_cache, bool is_neox, + torch::stable::Tensor& position_ids, + int64_t forced_token_heads_per_warp); + +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8); + +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8, + bool apply_q_rope, bool is_q_interleaved); + +void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& kv, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, double eps, + int64_t cache_block_size, bool apply_q_norm); + +void fused_kimi_k3_mla_key_concat_kv_cache_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_key_concat_ds_mla_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor& k_out); + +void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor const& v, + torch::stable::Tensor& k_fp8, + torch::stable::Tensor& v_fp8); + +void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8, + torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& k_scale_inv, + torch::stable::Tensor const& v_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + torch::stable::Tensor const& q, torch::stable::Tensor const& kv, + torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& fp8_scale, + torch::stable::Tensor const& q_fp8_scale_inv, double eps, + int64_t cache_block_size, bool apply_q_norm); + +#ifndef USE_ROCM +std::tuple +minimax_allreduce_rms_qk(torch::stable::Tensor qkv, + torch::stable::Tensor const& norm_weight_q, + torch::stable::Tensor const& norm_weight_k, + torch::stable::Tensor workspace, int64_t const q_size, + int64_t const kv_size, int64_t const rank, + int64_t const nranks, double const eps); +#endif + +// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV / +// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs +// the index branch and scatters k/v/index_k into their paged caches. +void fused_minimax_m3_qknorm_rope_kv_insert( + torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight, + torch::stable::Tensor const& k_norm_weight, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& positions, int64_t num_heads, + int64_t num_kv_heads, int64_t rotary_dim, double eps, + std::optional index_q_norm_weight, + std::optional index_k_norm_weight, + int64_t num_index_heads, std::optional slot_mapping, + std::optional index_slot_mapping, + std::optional kv_cache, + std::optional index_cache, int64_t block_size, + std::optional q_out, + std::optional index_q_out, + const std::string& kv_cache_dtype, bool skip_index_branch, + std::optional q_fp8_out, double q_fp8_scale, + std::optional kv_k_scale, + std::optional kv_v_scale); + +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE +void fused_kda_decode( + torch::stable::Tensor const& x, torch::stable::Tensor const& weight, + std::optional bias, + torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, + torch::stable::Tensor& out, std::optional lower_bound, + std::optional output_gate, + std::optional norm_weight, double norm_eps); + +#endif + +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE +void fused_gdn_decode_post_conv_mtp( + torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a, + torch::stable::Tensor const& b, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, + torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& num_accepted_tokens, + torch::stable::Tensor& state, torch::stable::Tensor const& output_gate, + torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out, + double scale, double norm_eps, const std::string& output_gate_activation); + +#endif + +#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK +void fused_kda_prologue( + torch::stable::Tensor const& q, torch::stable::Tensor const& k, + torch::stable::Tensor const& v, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, torch::stable::Tensor& qg, + torch::stable::Tensor& w, torch::stable::Tensor& u, + torch::stable::Tensor& kg_t, torch::stable::Tensor& aqk, + torch::stable::Tensor& decay, torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& chunk_indices, + std::optional conv_weight, + std::optional conv_state, + std::optional conv_state_indices, + std::optional conv_has_initial_state, double scale, + double lower_bound); + +void fused_kda_chunk( + torch::stable::Tensor const& qg, torch::stable::Tensor const& w, + torch::stable::Tensor const& u, torch::stable::Tensor const& kg_t, + torch::stable::Tensor const& aqk, torch::stable::Tensor const& decay, + std::optional initial_state, + std::optional final_state, + torch::stable::Tensor& out, torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& chunk_offsets, double scale, + std::optional group_state, int64_t groups, + std::optional checkpoint_state, + std::optional checkpoint_offsets, + std::optional checkpoint_state_indices, + std::optional state_cache, + std::optional state_indices, + std::optional has_initial_state); +#endif + +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES +void kimi_k3_attn_res(torch::stable::Tensor& prefix, + std::optional delta, + torch::stable::Tensor& blocks, + torch::stable::Tensor const& norm_weight, + torch::stable::Tensor const& qk_weight, + std::optional output_norm_weight, + torch::stable::Tensor& output, int64_t num_blocks, + int64_t block_write_idx, double eps, + double output_norm_eps); +#endif + +// Sampler kernels (shared CUDA/ROCm) +void apply_repetition_penalties_( + torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, + const torch::stable::Tensor& output_mask, + const torch::stable::Tensor& repetition_penalties); + +void top_k_per_row_prefill(const torch::stable::Tensor& logits, + const torch::stable::Tensor& rowStarts, + const torch::stable::Tensor& rowEnds, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK); + +void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, + const torch::stable::Tensor& seqLens, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK); + +void persistent_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); + +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK +void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); +#endif + +void selective_scan_fwd( + const torch::stable::Tensor& u, const torch::stable::Tensor& delta, + const torch::stable::Tensor& A, const torch::stable::Tensor& B, + const torch::stable::Tensor& C, + const std::optional& D_, + const std::optional& z_, + const std::optional& delta_bias_, + bool delta_softplus, + const std::optional& query_start_loc, + const std::optional& cache_indices, + const std::optional& has_initial_state, + const torch::stable::Tensor& ssm_states, int64_t null_block_id, + int64_t block_size, + const std::optional& block_idx_first_scheduled_token, + const std::optional& block_idx_last_scheduled_token, + const std::optional& initial_state_idx, + const std::optional& cu_chunk_seqlen, + const std::optional& last_chunk_indices); using fptr_t = int64_t; -#ifdef USE_ROCM -fptr_t init_custom_qr(int64_t rank, int64_t world_size, - std::optional qr_max_size = std::nullopt); -void qr_destroy(fptr_t _fa); -torch::Tensor qr_get_handle(fptr_t _fa); -void qr_open_handles(fptr_t _fa, const std::vector& handles); -void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, - int64_t quant_level, bool cast_bf2half = false); -int64_t qr_max_size(); -#endif +fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, + torch::stable::Tensor& rank_data, int64_t rank, + bool fully_connected); +void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t local_buffer, + fptr_t multicast_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); +void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t local_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); +void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t local_buffer, fptr_t multicast_buffer, + int64_t stage_sz_bytes, int64_t block_limit); +void dispose(fptr_t _fa); +int64_t meta_size(); +void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); +std::tuple, std::vector> +get_graph_buffer_ipc_meta(fptr_t _fa); +void register_graph_buffers(fptr_t _fa, + const std::vector>& handles, + const std::vector>& offsets); +std::tuple allocate_shared_buffer_and_handle( + int64_t size); +int64_t open_mem_handle(torch::stable::Tensor& mem_handle); +void free_shared_buffer(int64_t buffer); + +// Activation kernels (shared CUDA/ROCm) +void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); +void silu_and_mul_clamp(torch::stable::Tensor& out, + torch::stable::Tensor& input, double limit, + double alpha = 1.0, double beta = 0.0); + +void silu_and_mul_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& scale); + +void persistent_masked_m_silu_mul_quant( + const torch::stable::Tensor& input, // (E, T, 2*H) + const torch::stable::Tensor& tokens_per_expert, // (E) + torch::stable::Tensor& y_q, // (E, T, H) [OUT] + torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT] + bool use_ue8m0); + +void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_tanh_and_mul(torch::stable::Tensor& out, + torch::stable::Tensor& input); +void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double threshold); +void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double alpha = 1.702, double limit = 7.0); +void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double beta = 1.0, double linear_beta = -1.0); +void situ_and_mul_quant( + torch::stable::Tensor& out, torch::stable::Tensor& scale, + torch::stable::Tensor& input, double beta = 1.0, double linear_beta = -1.0, + int64_t group_size = 0, + std::optional num_valid_tokens = std::nullopt, + int64_t topk = 1); +void masked_situ_and_mul(torch::stable::Tensor& out, + torch::stable::Tensor& input, + const torch::stable::Tensor& expert_num_tokens, + double beta = 1.0, double linear_beta = -1.0); +void masked_moe_activation(torch::stable::Tensor& out, + torch::stable::Tensor& input, + const torch::stable::Tensor& valid_token_counts, + const std::string& activation, + double clamp_limit = 0.0, double alpha = 1.0, + double beta = 0.0, double situ_beta = 1.0, + double situ_linear_beta = -1.0); +void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); + +void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input); + +// INT8 quantization kernels (shared CUDA/ROCm) +void static_scaled_int8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor const& scale, + std::optional const& azp); + +void dynamic_scaled_int8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scales, + std::optional const& azp); + +// FP8 quantization kernels (shared CUDA/ROCm) +void static_scaled_fp8_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor const& scale, + std::optional group_shape = + std::nullopt); + +void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scale); + +void dynamic_per_token_scaled_fp8_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor& scale, + std::optional const& scale_ub); + +// GPTQ kernels (shared CUDA/ROCm) +torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, + torch::stable::Tensor b_q_weight, + torch::stable::Tensor b_gptq_qzeros, + torch::stable::Tensor b_gptq_scales, + bool use_exllama, bool use_v2_format, + int64_t bit); + +void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit); + +// Cache ops (shared CUDA/ROCm) +void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, + int64_t block_size_in_bytes, + const torch::stable::Tensor& block_mapping); + +// Batch swap: submit all block copies in a single driver call. +void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, + const torch::stable::Tensor& dst_ptrs, + const torch::stable::Tensor& sizes, + bool is_src_access_order_any); + +void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale); + +void reshape_and_cache_flash( + torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale); + +void concat_and_cache_mla(torch::stable::Tensor& kv_c, + torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& scale); + +void concat_and_cache_mla_grouped( + torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache_ptrs, torch::stable::Tensor& slot_mapping, + int64_t block_size, int64_t block_stride, int64_t entry_stride, + std::optional kv_scales, + const std::string& kv_cache_dtype); + +#ifndef USE_ROCM +// HiSparse kernels use raw PTX in the row copy; CUDA-only. +void hisparse_resolve_residency( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& hot_block_table, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor& hot_indices, + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor& lru_slots, + std::optional const& request_state_indices, + int64_t region_stride, int64_t max_union_rows, + std::optional const& miss_mask, + std::optional const& stats, + std::optional const& attention_indices, + int64_t attention_block_stride, + std::optional const& request_ids, + std::optional const& source_block_table, + int64_t source_block_size, + std::optional const& resolved_global_indices, + std::optional const& valid_counts, + std::optional const& swap_host_physical_rows, + std::optional const& swap_device_physical_rows, + std::optional const& swap_counts, + std::optional const& resident_block_table, + int64_t resident_block_size, int64_t resident_null_block); + +void hisparse_invalidate_written_slots( + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor const& request_state_indices, + torch::stable::Tensor const& req_id_per_token, + torch::stable::Tensor const& written_slots); + +void hisparse_gather_plan( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor const& hot_indices, + torch::stable::Tensor const& miss_mask, + std::optional const& request_state_indices, + std::optional const& attention_indices, + int64_t attention_block_stride); + +void hisparse_gather_compact(torch::stable::Tensor const& host_cache, + torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& miss_global_indices, + torch::stable::Tensor const& miss_hot_indices, + torch::stable::Tensor const& miss_counts); + +#endif // !USE_ROCM + +// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla +void concat_and_cache_mla_rope_fused( + torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, + torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c, + torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox, + torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache, + const std::string& kv_cache_dtype, + torch::stable::Tensor& kv_cache_quant_scale); + +// Just for unittest +void convert_fp8(torch::stable::Tensor& dst_cache, + torch::stable::Tensor& src_cache, const double scale, + const std::string& kv_cache_dtype); + +void gather_and_maybe_dequant_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] + int64_t num_tokens, const std::string& kv_cache_dtype, + torch::stable::Tensor const& scale, + std::optional seq_starts = std::nullopt); + +// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. +void cp_gather_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + int64_t batch_size, + std::optional seq_starts = std::nullopt); + +// Gather and upconvert FP8 KV cache to BF16 workspace +void cp_gather_and_upconvert_fp8_kv_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // 656] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& workspace_starts, // [BATCH] + int64_t batch_size, + std::optional seq_starts = std::nullopt, + std::optional host_cache = std::nullopt, + std::optional host_row_ids = std::nullopt, + std::optional device_row_ids = std::nullopt); + +// Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace +void cp_gather_and_upconvert_nvfp4_kv_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // 352] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& workspace_starts, // [BATCH] + int64_t batch_size); + +// Indexer K quantization and cache function +void indexer_k_quant_and_cache( + torch::stable::Tensor& k, // [num_tokens, head_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& slot_mapping, // [num_tokens] + int64_t quant_block_size, // quantization block size + const std::string& scale_fmt); + +// Concatenate query nope and rope for MLA/DSA attention +void concat_mla_q( + torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] + torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] + torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim + + // rope_dim] + +// Extract function to gather quantized K cache +void cp_gather_indexer_k_quant_cache( + const torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& dst_k, // [num_tokens, head_dim] + torch::stable::Tensor& dst_scale, // [num_tokens, head_dim / + // quant_block_size * 4] + const torch::stable::Tensor& block_table, // [batch_size, num_blocks] + const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] + +// Fused vocab-parallel embedding lookup (see +// vocab_parallel_embedding_kernels.cu). +void vocab_parallel_embedding( + torch::stable::Tensor& out, const torch::stable::Tensor& input_ids, + const torch::stable::Tensor& weight, int64_t org_vocab_start_index, + int64_t org_vocab_end_index, int64_t num_org_vocab_padding, + int64_t added_vocab_start_index, int64_t added_vocab_end_index); + +// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu). +void ngram_compute_n_gram_ids( + int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, + torch::stable::Tensor& ne_mods, + torch::stable::Tensor& exclusive_ne_embedder_size_sums, + torch::stable::Tensor& exclusive_req_len_sums, + torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, + torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/persistent_topk.cuh similarity index 100% rename from csrc/libtorch_stable/persistent_topk.cuh rename to csrc/persistent_topk.cuh diff --git a/csrc/libtorch_stable/pos_encoding_kernels.cu b/csrc/pos_encoding_kernels.cu similarity index 99% rename from csrc/libtorch_stable/pos_encoding_kernels.cu rename to csrc/pos_encoding_kernels.cu index 74af743b09..bad616fc87 100644 --- a/csrc/libtorch_stable/pos_encoding_kernels.cu +++ b/csrc/pos_encoding_kernels.cu @@ -1,6 +1,6 @@ #include "torch_utils.h" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/activation_kernels.cu b/csrc/quantization/activation_kernels.cu similarity index 99% rename from csrc/libtorch_stable/quantization/activation_kernels.cu rename to csrc/quantization/activation_kernels.cu index 822a41969e..9d7da85a22 100644 --- a/csrc/libtorch_stable/quantization/activation_kernels.cu +++ b/csrc/quantization/activation_kernels.cu @@ -1,10 +1,10 @@ -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include -#include "libtorch_stable/core/math.hpp" -#include "cuda_compat.h" -#include "libtorch_stable/dispatch_utils.h" +#include "core/utils.hpp" +#include "core/cuda_compat.h" +#include "dispatch_utils.h" #include "quantization/w8a8/fp8/common.cuh" #ifndef USE_ROCM diff --git a/csrc/libtorch_stable/quantization/awq/dequantize.cuh b/csrc/quantization/awq/dequantize.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/awq/dequantize.cuh rename to csrc/quantization/awq/dequantize.cuh diff --git a/csrc/libtorch_stable/quantization/awq/gemm_kernels.cu b/csrc/quantization/awq/gemm_kernels.cu similarity index 99% rename from csrc/libtorch_stable/quantization/awq/gemm_kernels.cu rename to csrc/quantization/awq/gemm_kernels.cu index 43301f0fe2..a0be657045 100644 --- a/csrc/libtorch_stable/quantization/awq/gemm_kernels.cu +++ b/csrc/quantization/awq/gemm_kernels.cu @@ -9,9 +9,9 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023} #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/quantization/awq/dequantize.cuh" +#include "quantization/awq/dequantize.cuh" #include diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh b/csrc/quantization/cutlass_w4a8/get_group_starts.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh rename to csrc/quantization/cutlass_w4a8/get_group_starts.cuh index 7217617912..aff0502550 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh +++ b/csrc/quantization/cutlass_w4a8/get_group_starts.cuh @@ -3,7 +3,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "cutlass/bfloat16.h" #include "cutlass/float8.h" diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu b/csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu similarity index 98% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu rename to csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu index 1d7eb04093..82d6b2ec8a 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu +++ b/csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu @@ -16,12 +16,12 @@ // vllm includes #include #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "torch_utils.h" +#include "cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/common.hpp" #include "get_group_starts.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" #include "w4a8_utils.cuh" namespace vllm::cutlass_w4a8_moe { diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu b/csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu similarity index 98% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu rename to csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu index 8f19394b39..c4e96e8aee 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu +++ b/csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu @@ -5,8 +5,8 @@ #include #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "torch_utils.h" +#include "cutlass_extensions/torch_utils.hpp" #include "w4a8_utils.cuh" #include "cutlass/cutlass.h" @@ -21,8 +21,8 @@ #include "cutlass/util/packed_stride.hpp" #include "cutlass/util/mixed_dtype_utils.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/common.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" #include diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu b/csrc/quantization/cutlass_w4a8/w4a8_utils.cu similarity index 100% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu rename to csrc/quantization/cutlass_w4a8/w4a8_utils.cu diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cuh b/csrc/quantization/cutlass_w4a8/w4a8_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cuh rename to csrc/quantization/cutlass_w4a8/w4a8_utils.cuh diff --git a/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu similarity index 97% rename from csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu rename to csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu index 6238a27191..2064775111 100644 --- a/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu +++ b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu @@ -15,8 +15,8 @@ */ #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include "../../cuda_vec_utils.cuh" #include @@ -25,7 +25,7 @@ #include #include "cuda_utils.h" -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" // Define before including nvfp4_utils.cuh so the header // can use this macro during compilation. diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu b/csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu rename to csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu index e4d2f22012..3234717f6b 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu +++ b/csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu @@ -8,11 +8,11 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cute/tensor.hpp" #include "cutlass/tensor_ref.h" diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/quantization/fp4/mxfp4_experts_quant.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu rename to csrc/quantization/fp4/mxfp4_experts_quant.cu index 086c07acf0..4c6875f062 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/quantization/fp4/mxfp4_experts_quant.cu @@ -25,9 +25,9 @@ #include #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "torch_utils.h" +#include "dispatch_utils.h" +#include "cutlass_extensions/common.hpp" #include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" @@ -41,7 +41,7 @@ static_assert(CVT_FP4_ELTS_PER_THREAD == 16, #define VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED 0 #endif -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" #if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu b/csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu rename to csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu index ea33444063..8e8a30e956 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu +++ b/csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu @@ -16,11 +16,11 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cute/tensor.hpp" #include "cutlass/tensor_ref.h" diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu b/csrc/quantization/fp4/nvfp4_experts_quant.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu rename to csrc/quantization/fp4/nvfp4_experts_quant.cu index 6296f0d901..55c722f1de 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu +++ b/csrc/quantization/fp4/nvfp4_experts_quant.cu @@ -15,8 +15,8 @@ */ #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include "../../cuda_vec_utils.cuh" #include @@ -26,7 +26,7 @@ #include "cuda_utils.h" #include "nvfp4_utils.cuh" -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu b/csrc/quantization/fp4/nvfp4_quant_entry.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu rename to csrc/quantization/fp4/nvfp4_quant_entry.cu index e1e7e7a74d..ace7bbbc10 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu +++ b/csrc/quantization/fp4/nvfp4_quant_entry.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "nvfp4_utils.cuh" #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu b/csrc/quantization/fp4/nvfp4_quant_kernels.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu rename to csrc/quantization/fp4/nvfp4_quant_kernels.cu index 7ec418ee38..13a3026cc8 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu +++ b/csrc/quantization/fp4/nvfp4_quant_kernels.cu @@ -21,12 +21,12 @@ #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" // Define before including nvfp4_utils.cuh so the header // can use this macro during compilation. diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu b/csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu similarity index 97% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu rename to csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu index bfb526fcd4..6a60fe6f60 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu +++ b/csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D, diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu b/csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu rename to csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu index af9f24a70e..f925a599db 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu +++ b/csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cutlass/cutlass.h" @@ -31,8 +31,7 @@ #include "cutlass/util/packed_stride.hpp" -#include "libtorch_stable/core/math.hpp" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu b/csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu rename to csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu index 3a45ede8df..02ded13c78 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu +++ b/csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cutlass/cutlass.h" @@ -31,8 +31,7 @@ #include "cutlass/util/packed_stride.hpp" -#include "libtorch_stable/core/math.hpp" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh b/csrc/quantization/fp4/nvfp4_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh rename to csrc/quantization/fp4/nvfp4_utils.cuh diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu b/csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu rename to csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu index 56dd470387..f864d6037e 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu +++ b/csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu @@ -2,7 +2,7 @@ #include "../../torch_utils.h" #include "../../dispatch_utils.h" -#include "../../../core/batch_invariant.hpp" +#include "core/utils.hpp" #include "layernorm_utils.cuh" #include "quant_conversions.cuh" diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu b/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu similarity index 100% rename from csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu rename to csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu diff --git a/csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh b/csrc/quantization/fused_kernels/layernorm_utils.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh rename to csrc/quantization/fused_kernels/layernorm_utils.cuh index 9a3d221562..d633cb9a6e 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh +++ b/csrc/quantization/fused_kernels/layernorm_utils.cuh @@ -4,12 +4,12 @@ * __device__ layernorm utilities. */ -#include "libtorch_stable/quantization/vectorization.cuh" +#include "quantization/vectorization.cuh" #include "quantization/utils.cuh" #include "quant_conversions.cuh" #include "../../cub_helpers.h" -#include "../../../cuda_compat.h" +#include "core/cuda_compat.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh b/csrc/quantization/fused_kernels/quant_conversions.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh rename to csrc/quantization/fused_kernels/quant_conversions.cuh index a81a0b9feb..1953e90f5d 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh +++ b/csrc/quantization/fused_kernels/quant_conversions.cuh @@ -4,9 +4,9 @@ * __device__ helper functions to deal with float -> quant datatype conversion */ -#include "libtorch_stable/quantization/vectorization.cuh" +#include "quantization/vectorization.cuh" // TODO(luka/varun):refactor common.cuh to use this file instead -#include "../../../quantization/w8a8/fp8/common.cuh" +#include "quantization/w8a8/fp8/common.cuh" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/gptq/compat.cuh b/csrc/quantization/gptq/compat.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/compat.cuh rename to csrc/quantization/gptq/compat.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh b/csrc/quantization/gptq/matrix_view.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/matrix_view.cuh rename to csrc/quantization/gptq/matrix_view.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/q_gemm.cu b/csrc/quantization/gptq/q_gemm.cu similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/q_gemm.cu rename to csrc/quantization/gptq/q_gemm.cu diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_2.cuh b/csrc/quantization/gptq/qdq_2.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_2.cuh rename to csrc/quantization/gptq/qdq_2.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_3.cuh b/csrc/quantization/gptq/qdq_3.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_3.cuh rename to csrc/quantization/gptq/qdq_3.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_4.cuh b/csrc/quantization/gptq/qdq_4.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_4.cuh rename to csrc/quantization/gptq/qdq_4.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_8.cuh b/csrc/quantization/gptq/qdq_8.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_8.cuh rename to csrc/quantization/gptq/qdq_8.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_util.cuh b/csrc/quantization/gptq/qdq_util.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_util.cuh rename to csrc/quantization/gptq/qdq_util.cuh diff --git a/csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu b/csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu similarity index 99% rename from csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu rename to csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu index a5d199db27..2575b20b8e 100644 --- a/csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu +++ b/csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu @@ -11,8 +11,8 @@ Redistribution and use in source and binary forms, with or without modification, THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ***********/ -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/machete/Readme.md b/csrc/quantization/machete/Readme.md similarity index 100% rename from csrc/libtorch_stable/quantization/machete/Readme.md rename to csrc/quantization/machete/Readme.md diff --git a/csrc/libtorch_stable/quantization/machete/generate.py b/csrc/quantization/machete/generate.py similarity index 99% rename from csrc/libtorch_stable/quantization/machete/generate.py rename to csrc/quantization/machete/generate.py index 11a5bbdd13..60a49c587f 100644 --- a/csrc/libtorch_stable/quantization/machete/generate.py +++ b/csrc/quantization/machete/generate.py @@ -496,7 +496,8 @@ def create_sources(impl_configs: list[ImplConfig], num_impl_files=8): def generate(): - # See csrc/quantization/machete/Readme.md, the Codegeneration for more info + # See csrc/quantization/machete/Readme.md, the + # Codegeneration for more info # about how this works SCRIPT_DIR = os.path.dirname(__file__) diff --git a/csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh b/csrc/quantization/machete/machete_collective_builder.cuh similarity index 94% rename from csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh rename to csrc/quantization/machete/machete_collective_builder.cuh index ba8da0af2c..ee825583de 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh +++ b/csrc/quantization/machete/machete_collective_builder.cuh @@ -1,6 +1,6 @@ #pragma once -#include "libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh" +#include "cutlass_extensions/vllm_collective_builder.cuh" #include "machete_mainloop.cuh" namespace cutlass::gemm::collective { diff --git a/csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh b/csrc/quantization/machete/machete_interleaving_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh rename to csrc/quantization/machete/machete_interleaving_utils.cuh diff --git a/csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh b/csrc/quantization/machete/machete_mainloop.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh rename to csrc/quantization/machete/machete_mainloop.cuh index a0d1ded887..9127a701a0 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh +++ b/csrc/quantization/machete/machete_mainloop.cuh @@ -17,8 +17,9 @@ // before feeding them into the tensor cores in registers, we need the weight // matrix to be A. To achieve this we compute the transpose of Y = XW^t as // Y^t = W^tX^t. This is mostly done outside of this file in -// csrc/quantization/machete/machete_mm_kernel.cuh, but this why A is the -// quantized/narrow type and has the prepacked layout despite the API being: +// csrc/quantization/machete/machete_mm_kernel.cuh, but this why +// A is the quantized/narrow type and has the prepacked layout despite the API +// being: // B_prepacked = machete_prepack_B(B) // Y = machete_mm(A, B_prepacked) // diff --git a/csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh b/csrc/quantization/machete/machete_mm_kernel.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh rename to csrc/quantization/machete/machete_mm_kernel.cuh index 57655f369c..db3321a39d 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh +++ b/csrc/quantization/machete/machete_mm_kernel.cuh @@ -18,9 +18,9 @@ // clang-format on #include "cutlass_extensions/cute_utils.cuh" -#include "libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/vllm_numeric_conversion.cuh" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include "machete_collective_builder.cuh" #include "machete_prepacked_layout.cuh" #include "machete_interleaving_utils.cuh" diff --git a/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh b/csrc/quantization/machete/machete_mm_launcher.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh rename to csrc/quantization/machete/machete_mm_launcher.cuh index 5ecf03daf7..14206f1d5c 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh +++ b/csrc/quantization/machete/machete_mm_launcher.cuh @@ -1,9 +1,9 @@ #pragma once #include "machete_mm_kernel.cuh" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include "core/scalar_type.hpp" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh b/csrc/quantization/machete/machete_prepack_kernel.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh rename to csrc/quantization/machete/machete_prepack_kernel.cuh index 6a28dceccc..e1e054e5a0 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh +++ b/csrc/quantization/machete/machete_prepack_kernel.cuh @@ -2,7 +2,7 @@ #include "machete_mm_kernel.cuh" #include "cutlass_extensions/cute_utils.cuh" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include namespace machete { diff --git a/csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh b/csrc/quantization/machete/machete_prepack_launcher.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh rename to csrc/quantization/machete/machete_prepack_launcher.cuh index 63af5e496b..26bf200732 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh +++ b/csrc/quantization/machete/machete_prepack_launcher.cuh @@ -1,9 +1,9 @@ #pragma once #include "machete_prepack_kernel.cuh" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include "core/scalar_type.hpp" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh b/csrc/quantization/machete/machete_prepacked_layout.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh rename to csrc/quantization/machete/machete_prepacked_layout.cuh diff --git a/csrc/libtorch_stable/quantization/machete/machete_pytorch.cu b/csrc/quantization/machete/machete_pytorch.cu similarity index 100% rename from csrc/libtorch_stable/quantization/machete/machete_pytorch.cu rename to csrc/quantization/machete/machete_pytorch.cu diff --git a/csrc/libtorch_stable/quantization/marlin/.gitignore b/csrc/quantization/marlin/.gitignore similarity index 73% rename from csrc/libtorch_stable/quantization/marlin/.gitignore rename to csrc/quantization/marlin/.gitignore index 7dc482a894..ba805f9250 100644 --- a/csrc/libtorch_stable/quantization/marlin/.gitignore +++ b/csrc/quantization/marlin/.gitignore @@ -1,3 +1,2 @@ sm*_kernel_*.cu kernel_selector.h -kernel_*.cu diff --git a/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu b/csrc/quantization/marlin/awq_marlin_repack.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu rename to csrc/quantization/marlin/awq_marlin_repack.cu index 55ce5b4e73..09b647a03a 100644 --- a/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu +++ b/csrc/quantization/marlin/awq_marlin_repack.cu @@ -7,7 +7,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" namespace marlin { diff --git a/csrc/libtorch_stable/quantization/marlin/dequant.h b/csrc/quantization/marlin/dequant.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/dequant.h rename to csrc/quantization/marlin/dequant.h diff --git a/csrc/libtorch_stable/quantization/marlin/generate_kernels.py b/csrc/quantization/marlin/generate_kernels.py similarity index 96% rename from csrc/libtorch_stable/quantization/marlin/generate_kernels.py rename to csrc/quantization/marlin/generate_kernels.py index 0b908d86d5..9e0af3c275 100644 --- a/csrc/libtorch_stable/quantization/marlin/generate_kernels.py +++ b/csrc/quantization/marlin/generate_kernels.py @@ -161,12 +161,15 @@ QUANT_CONFIGS = [ }, ] +# Kernels are written next to the kernel.h / marlin_template.h they include. +OUTPUT_DIR = sys.argv[2] if len(sys.argv) > 2 else os.path.dirname(__file__) + def remove_old_kernels(): - for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"): + for filename in glob.glob(OUTPUT_DIR + "/*kernel_*.cu"): subprocess.call(["rm", "-f", filename]) - filename = os.path.dirname(__file__) + "/kernel_selector.h" + filename = OUTPUT_DIR + "/kernel_selector.h" subprocess.call(["rm", "-f", filename]) @@ -297,7 +300,7 @@ def generate_new_kernels(): filename = filename.lower() - with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: + with open(os.path.join(OUTPUT_DIR, filename), "w") as f: f.write(file_content) if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: @@ -307,7 +310,7 @@ def generate_new_kernels(): '"marlin kernel with fp8 activation is not built.");' ) - with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f: + with open(os.path.join(OUTPUT_DIR, "kernel_selector.h"), "w") as f: f.write(kernel_selector_str) diff --git a/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu b/csrc/quantization/marlin/gptq_marlin_repack.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu rename to csrc/quantization/marlin/gptq_marlin_repack.cu index 5dab9fe932..3fcf3406eb 100644 --- a/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu +++ b/csrc/quantization/marlin/gptq_marlin_repack.cu @@ -7,7 +7,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" namespace marlin { diff --git a/csrc/libtorch_stable/quantization/marlin/kernel.h b/csrc/quantization/marlin/kernel.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/kernel.h rename to csrc/quantization/marlin/kernel.h diff --git a/csrc/libtorch_stable/quantization/marlin/marlin.cu b/csrc/quantization/marlin/marlin.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/marlin.cu rename to csrc/quantization/marlin/marlin.cu index cd52f3c736..d77ad692b2 100644 --- a/csrc/libtorch_stable/quantization/marlin/marlin.cu +++ b/csrc/quantization/marlin/marlin.cu @@ -32,7 +32,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::value || \ diff --git a/csrc/libtorch_stable/quantization/marlin/marlin.cuh b/csrc/quantization/marlin/marlin.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin.cuh rename to csrc/quantization/marlin/marlin.cuh diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh b/csrc/quantization/marlin/marlin_dtypes.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh rename to csrc/quantization/marlin/marlin_dtypes.cuh diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu b/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu rename to csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu index b48601435e..c446326909 100644 --- a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu +++ b/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu @@ -8,7 +8,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" // for only non-zp format (like gptq) __global__ void marlin_int4_fp8_preprocess_kernel_without_zp( diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_mma.h b/csrc/quantization/marlin/marlin_mma.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin_mma.h rename to csrc/quantization/marlin/marlin_mma.h diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_template.h b/csrc/quantization/marlin/marlin_template.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin_template.h rename to csrc/quantization/marlin/marlin_template.h diff --git a/csrc/libtorch_stable/quantization/vectorization.cuh b/csrc/quantization/vectorization.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/vectorization.cuh rename to csrc/quantization/vectorization.cuh diff --git a/csrc/libtorch_stable/quantization/vectorization_utils.cuh b/csrc/quantization/vectorization_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/vectorization_utils.cuh rename to csrc/quantization/vectorization_utils.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh b/csrc/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh rename to csrc/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh index a93a136bf4..0722101b38 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh @@ -5,7 +5,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "cutlass/cutlass.h" @@ -19,8 +19,8 @@ #include "cutlass/gemm/collective/collective_builder.hpp" #include "cutlass/util/packed_stride.hpp" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" // clang-format on namespace vllm::c3x { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm.cuh index cf18b63a88..b20c055fb8 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm.cuh @@ -14,8 +14,8 @@ #include "cutlass/epilogue/collective/collective_builder.hpp" #include "cutlass/gemm/collective/collective_builder.hpp" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" // clang-format on /* diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu similarity index 90% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu index 5178eed072..bc088cf633 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_int8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu similarity index 89% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu index b793001226..f3df69850e 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm100_fp8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh index fa2add9182..1e7a31f5ca 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh @@ -2,7 +2,7 @@ #include -#include "cuda_utils.h" +#include "core/utils.hpp" #include "cutlass/cutlass.h" #include "cutlass/numeric_types.h" @@ -242,10 +242,10 @@ void cutlass_gemm_blockwise_sm100_fp8_dispatch(torch::stable::Tensor& out, if (!swap_ab) { constexpr int TILE_N = 128; int tile_m = 256; - if (cuda_utils::ceil_div(n, TILE_N) * cuda_utils::ceil_div(m, 64) <= sms) { + if (div_ceil(n, TILE_N) * div_ceil(m, 64) <= sms) { tile_m = 64; } - else if (cuda_utils::ceil_div(n, TILE_N) * cuda_utils::ceil_div(m, 128) <= sms) { + else if (div_ceil(n, TILE_N) * div_ceil(m, 128) <= sms) { tile_m = 128; } if (tile_m == 64) { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu similarity index 96% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu index 021c16d668..9f3c8a196d 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm120_fp8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh index 9c921ea90a..34bb7ef3f2 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh @@ -2,7 +2,6 @@ #include -#include "cuda_utils.h" #include "cutlass/cutlass.h" #include "cutlass/numeric_types.h" diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu similarity index 89% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu index a97909d37a..d3318c4876 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu @@ -1,7 +1,7 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm90_fp8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp similarity index 82% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp index 913436186c..a15bc1da4a 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp @@ -1,7 +1,7 @@ #include #include -#include "cuda_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" template void dispatch_scaled_mm(torch::stable::Tensor& c, @@ -41,14 +41,12 @@ void dispatch_scaled_mm(torch::stable::Tensor& c, STD_TORCH_CHECK(b_scales.dim() == 2, "b scale must be 2d tensor."); int32_t version_num = get_sm_version_num(); if (version_num >= 90) { - STD_TORCH_CHECK( - a.size(0) == a_scales.size(0) && - cuda_utils::ceil_div(a.size(1), int64_t(128)) == a_scales.size(1), - "a_scale_group_shape must be [1, 128]."); - STD_TORCH_CHECK( - cuda_utils::ceil_div(b.size(0), int64_t(128)) == b_scales.size(0) && - cuda_utils::ceil_div(b.size(1), int64_t(128)) == b_scales.size(1), - "b_scale_group_shape must be [128, 128]."); + STD_TORCH_CHECK(a.size(0) == a_scales.size(0) && + div_ceil(a.size(1), int64_t(128)) == a_scales.size(1), + "a_scale_group_shape must be [1, 128]."); + STD_TORCH_CHECK(div_ceil(b.size(0), int64_t(128)) == b_scales.size(0) && + div_ceil(b.size(1), int64_t(128)) == b_scales.size(1), + "b_scale_group_shape must be [128, 128]."); } STD_TORCH_CHECK(!bias, "Bias not yet supported blockwise scaled_mm"); diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu index 84040a6a22..59022b041e 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm100_fp8_dispatch.cuh" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh index 42f9e4d5bc..f790b3653d 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh @@ -4,7 +4,7 @@ #include "scaled_mm.cuh" #include "cutlass_gemm_caller.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" /** * This file defines Gemm kernel configurations for SM100 (fp8) based on the diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu similarity index 91% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu index 1f709699eb..5cfe1d8d3a 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu @@ -1,7 +1,7 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm120_fp8_dispatch.cuh" -#include "core/batch_invariant.hpp" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu index e86c9bd48d..c94b9bd072 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_fp8_dispatch.cuh" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh index 2fae3016c3..f78b8daea5 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh @@ -4,7 +4,7 @@ #include "scaled_mm.cuh" #include "cutlass_gemm_caller.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" /** * This file defines Gemm kernel configurations for SM90 (fp8) based on the Gemm diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu similarity index 91% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu index bf2bd030e3..717a2a5883 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_int8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh b/csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh rename to csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh index 6942c4b48c..2189e5d0ab 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh +++ b/csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh @@ -3,7 +3,7 @@ #include #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "cutlass/bfloat16.h" #include "cutlass/float8.h" diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh rename to csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh index b4cd520e96..12241fb299 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh +++ b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh @@ -7,9 +7,9 @@ #include "cutlass/gemm/device/gemm_universal_adapter.h" #include -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" -#include "libtorch_stable/torch_utils.h" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/common.hpp" +#include "torch_utils.h" #include "get_group_starts.cuh" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu rename to csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu index 1f88bf140d..f03f2c0bc4 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu +++ b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu @@ -1,6 +1,6 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu rename to csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu index d494bfb9d0..14bb0586cd 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu +++ b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu @@ -1,6 +1,6 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu rename to csrc/quantization/w8a8/cutlass/moe/moe_data.cu index 8ebc3dd9e8..b6043e8033 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu +++ b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu @@ -1,11 +1,11 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include #include -#include "libtorch_stable/dispatch_utils.h" +#include "dispatch_utils.h" #include diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu similarity index 98% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu index 184a26b491..bd60058d73 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu @@ -9,8 +9,8 @@ #include "scaled_mm_c2x_sm89_fp8_dispatch.cuh" #include "scaled_mm_c2x_sm89_int8_dispatch.cuh" -#include "core/batch_invariant.hpp" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp" using namespace vllm; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cuh similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cuh index 9f9e12e982..30e587c63a 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cuh @@ -3,7 +3,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" // clang-format will break include orders // clang-format off @@ -22,8 +22,8 @@ #include "cutlass/epilogue/threadblock/fusion/visitors.hpp" #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" // clang-format on using namespace cute; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu index 51f84d2ffd..0c05f59d3f 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu @@ -3,9 +3,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" void cutlass_scaled_mm_sm75(torch::stable::Tensor& c, torch::stable::Tensor const& a, diff --git a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh index 32925c0c3c..cb050c9223 100644 --- a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh @@ -5,7 +5,7 @@ #include #include -#include "../../../../attention/attention_dtypes.h" +#include "core/vec_dtypes/vec_dtypes.h" #include namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/common.cu b/csrc/quantization/w8a8/fp8/common.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/fp8/common.cu rename to csrc/quantization/w8a8/fp8/common.cu index abebd6f02f..4714997869 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/common.cu +++ b/csrc/quantization/w8a8/fp8/common.cu @@ -1,5 +1,5 @@ -#include "../../../../quantization/w8a8/fp8/common.cuh" -#include "../../../../cuda_compat.h" +#include "common.cuh" +#include "core/cuda_compat.h" #include "../../../dispatch_utils.h" #include "../../../cub_helpers.h" #include "../../vectorization_utils.cuh" diff --git a/csrc/quantization/w8a8/fp8/common.cuh b/csrc/quantization/w8a8/fp8/common.cuh index 087f509916..143be698cc 100644 --- a/csrc/quantization/w8a8/fp8/common.cuh +++ b/csrc/quantization/w8a8/fp8/common.cuh @@ -1,17 +1,17 @@ #pragma once -#include "libtorch_stable/quantization/vectorization.cuh" -#include "../../utils.cuh" +#include "quantization/vectorization.cuh" +#include "quantization/utils.cuh" #include -// This header is shared between _C and _C_stable_libtorch targets. +// This header is shared between _rocm_C and _C_stable_libtorch targets. // torch_utils.h provides get_device_prop(). We need to pass USE_CUDA // to the .so to expose some of the shims used by torch_utils.h. For now -// this is only done for _C_stable_libtorch and not for _C, so we use the -// non stable at::cuda::getCurrentDeviceProperties for _C for now. +// this is only done for _C_stable_libtorch and not for _rocm_C, so we use +// the non stable at::cuda::getCurrentDeviceProperties for _rocm_C for now. #ifdef TORCH_TARGET_VERSION - #include "../../../libtorch_stable/torch_utils.h" + #include "torch_utils.h" #else #ifdef USE_ROCM #include @@ -19,9 +19,9 @@ #endif #ifndef USE_ROCM - #include "nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #else - #include "amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #endif // Determines the preferred FP8 type for the current platform. diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 7fcf541fd6..77aabe4651 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -1,6 +1,6 @@ #pragma once -#include "../../../../attention/attention_dtypes.h" +#include "core/vec_dtypes/vec_dtypes.h" #include #include #include @@ -243,7 +243,7 @@ __inline__ __device__ uint8_t scaled_vec_conversion( } // torch stable-ABI (headeronly) scalar types delegate to the CUDA-native -// conversions, so libtorch_stable kernels dispatched on c10::BFloat16 / +// conversions, so stable-ABI kernels dispatched on c10::BFloat16 / // c10::Half quantize correctly without manual casts. template <> __inline__ __device__ uint8_t scaled_vec_conversion( diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu rename to csrc/quantization/w8a8/fp8/per_token_group_quant.cu index a0aee0cd92..d2a8886438 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu @@ -3,8 +3,6 @@ #include #include -#include "libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h" - #include #ifdef USE_ROCM @@ -13,10 +11,10 @@ #include #endif -#include "libtorch_stable/quantization/vectorization.cuh" -#include "libtorch_stable/quantization/vectorization_utils.cuh" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/torch_utils.h" +#include "quantization/vectorization.cuh" +#include "quantization/vectorization_utils.cuh" +#include "dispatch_utils.h" +#include "torch_utils.h" __device__ __forceinline__ float GroupReduceMax(float val) { #ifdef USE_ROCM @@ -346,11 +344,13 @@ inline int GetGroupsPerBlockX(int64_t padded_groups_per_row) { return 4; } -void per_token_group_quant_8bit(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double min_8bit, - double max_8bit, bool scale_ue8m0) { +// 8-bit per-token-group quantization helper used by both FP8 and INT8 +static void per_token_group_quant_8bit(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, + double min_8bit, double max_8bit, + bool scale_ue8m0 = false) { STD_TORCH_CHECK(input.is_contiguous()); STD_TORCH_CHECK(output_q.is_contiguous()); @@ -870,3 +870,12 @@ void per_token_group_quant_fp8(const torch::stable::Tensor& input, per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0); } + +void per_token_group_quant_int8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double int8_min, + double int8_max) { + per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, + int8_min, int8_max); +} diff --git a/csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu b/csrc/quantization/w8a8/int8/scaled_quant.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu rename to csrc/quantization/w8a8/int8/scaled_quant.cu diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 2a11653a58..caf6ee9a14 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -19,11 +19,11 @@ #include #include #include -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include -#include "../attention/dtype_fp8.cuh" -#include "../quantization/w8a8/fp8/amd/quant_utils.cuh" +#include "core/vec_dtypes/dtype_fp8.cuh" +#include "quantization/w8a8/fp8/amd/quant_utils.cuh" // ROCm 6.2 compatibility: map OCP fp8 types to FNUZ variants if OCP is absent #if !defined(HIP_FP8_TYPE_OCP) diff --git a/csrc/rocm/dispatch_utils.h b/csrc/rocm/dispatch_utils.h new file mode 100644 index 0000000000..12026965f0 --- /dev/null +++ b/csrc/rocm/dispatch_utils.h @@ -0,0 +1,24 @@ +/* + * Adapted from + * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + */ +#pragma once + +#include + +// Need a special dispatch case macro since we will nest the FP8 dispatch. +// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'. +#define AT_DISPATCH_FP8_CASE(enum_type, ...) \ + AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) + +// ROCm devices might use either fn or fnuz, so set up dispatch table for both. +// A host-based check at runtime will create a preferred FP8 type for ROCm +// such that the correct kernel is dispatched. +#define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) + +// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. +// See AT_DISPATCH_FP8_CASE above. +#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) diff --git a/csrc/rocm/q_gemm_rdna3.cu b/csrc/rocm/q_gemm_rdna3.cu index e6f1c2fb17..8ba5b48f05 100644 --- a/csrc/rocm/q_gemm_rdna3.cu +++ b/csrc/rocm/q_gemm_rdna3.cu @@ -3,7 +3,8 @@ // // W4A16 GPTQ kernel for RDNA3 (gfx1100 / RX 7900 XTX class), templated on the // activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit -// kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes: +// kernel (csrc/quantization/gptq/q_gemm.cu) with the following +// changes: // // 1. Deterministic split-K epilogue. K is split across gridDim.z blocks; // each block stores its FP32 block partial to a scratch tensor, and the diff --git a/csrc/quickreduce/base.h b/csrc/rocm/quickreduce/base.h similarity index 100% rename from csrc/quickreduce/base.h rename to csrc/rocm/quickreduce/base.h diff --git a/csrc/custom_quickreduce.cu b/csrc/rocm/quickreduce/custom_quickreduce.cu similarity index 99% rename from csrc/custom_quickreduce.cu rename to csrc/rocm/quickreduce/custom_quickreduce.cu index d4e5d179a5..c1e75ae05d 100644 --- a/csrc/custom_quickreduce.cu +++ b/csrc/rocm/quickreduce/custom_quickreduce.cu @@ -5,7 +5,7 @@ #ifdef USE_ROCM - #include "quickreduce/quick_reduce.h" + #include "rocm/quickreduce/quick_reduce.h" quickreduce::fptr_t init_custom_qr(int64_t rank, int64_t world_size, std::optional qr_max_size) { diff --git a/csrc/quickreduce/quick_reduce.h b/csrc/rocm/quickreduce/quick_reduce.h similarity index 100% rename from csrc/quickreduce/quick_reduce.h rename to csrc/rocm/quickreduce/quick_reduce.h diff --git a/csrc/quickreduce/quick_reduce_impl.cuh b/csrc/rocm/quickreduce/quick_reduce_impl.cuh similarity index 100% rename from csrc/quickreduce/quick_reduce_impl.cuh rename to csrc/rocm/quickreduce/quick_reduce_impl.cuh diff --git a/csrc/rocm/quickreduce/torch_bindings.cpp b/csrc/rocm/quickreduce/torch_bindings.cpp new file mode 100644 index 0000000000..b7d9a695ac --- /dev/null +++ b/csrc/rocm/quickreduce/torch_bindings.cpp @@ -0,0 +1,51 @@ +#include +#include "core/registration.h" +#include +#include + +#include +#include + +#ifdef USE_ROCM +// Defined in custom_quickreduce.cu. +using fptr_t = int64_t; +fptr_t init_custom_qr(int64_t rank, int64_t world_size, + std::optional qr_max_size = std::nullopt); +void qr_destroy(fptr_t _fa); +torch::Tensor qr_get_handle(fptr_t _fa); +void qr_open_handles(fptr_t _fa, const std::vector& handles); +void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, + int64_t quant_level, bool cast_bf2half = false); +int64_t qr_max_size(); +#endif + +// Note on op signatures: +// The X_meta signatures are for the meta functions corresponding to op X. +// They must be kept in sync with the signature for X. Generally, only +// functions that return Tensors require a meta function. +// +// See the following links for detailed docs on op registration and function +// schemas. +// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9 +// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations + +#ifdef USE_ROCM +TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { + // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C). + custom_ar.def( + "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool " + "cast_bf2half) -> ()"); + custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce); + + custom_ar.def("init_custom_qr", &init_custom_qr); + custom_ar.def("qr_destroy", &qr_destroy); + custom_ar.def("qr_get_handle", &qr_get_handle); + + custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()"); + custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles); + + custom_ar.def("qr_max_size", &qr_max_size); +} +#endif + +REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index 79af081455..bd7c0daa1d 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -12,10 +12,10 @@ #include #include -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "quantization/w8a8/fp8/common.cuh" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" // Number of streams the pre-allocated split-K pool covers. ~7.5 MiB per slot. static constexpr int64_t kWvSlots = 8; diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index 4588b216fe..be3cda7be6 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -10,8 +10,7 @@ #include #include // std::in_range -#include "../cuda_compat.h" -#include "dispatch_utils.h" +#include "core/cuda_compat.h" // Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts #if defined(__GFX11__) || defined(__GFX12__) diff --git a/csrc/libtorch_stable/sampled_topk.cuh b/csrc/sampled_topk.cuh similarity index 100% rename from csrc/libtorch_stable/sampled_topk.cuh rename to csrc/sampled_topk.cuh diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/sampler.cu similarity index 99% rename from csrc/libtorch_stable/sampler.cu rename to csrc/sampler.cu index 09151004c1..e9145e3b99 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/sampler.cu @@ -1,4 +1,4 @@ -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "torch_utils.h" diff --git a/csrc/libtorch_stable/topk.cu b/csrc/topk.cu similarity index 100% rename from csrc/libtorch_stable/topk.cu rename to csrc/topk.cu diff --git a/csrc/libtorch_stable/topk_histogram_4096.cuh b/csrc/topk_histogram_4096.cuh similarity index 100% rename from csrc/libtorch_stable/topk_histogram_4096.cuh rename to csrc/topk_histogram_4096.cuh diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index bcf0ce3e1c..3d4f9ab273 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -1,39 +1,1171 @@ -// Provides torch::Tensor for ops.h (previously included transitively via -// cache.h, which is no longer included here after cache ops moved to -// _C_stable_libtorch). -#include #include "ops.h" +#include "cuda_utils.h" #include "core/registration.h" -#include -#include -// Note on op signatures: -// The X_meta signatures are for the meta functions corresponding to op X. -// They must be kept in sync with the signature for X. Generally, only -// functions that return Tensors require a meta function. -// -// See the following links for detailed docs on op registration and function -// schemas. -// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9 -// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations +#include -#ifdef USE_ROCM -TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { - // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C). - custom_ar.def( - "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool " - "cast_bf2half) -> ()"); - custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce); +// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility. +// Note: We register under namespace "_C" so ops are accessible as +// torch.ops._C. for compatibility with existing code. +STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { + // Compute per-token-group FP8 quantized tensor and scaling factor. + // The dummy arguments are here so we can correctly fuse with RMSNorm. + ops.def( + "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! " + "output_s, " + "int group_size, float eps, float fp8_min, float fp8_max, bool " + "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned " + ") -> ()"); + // Compute per-token-group 8-bit quantized tensor and UE8M0-packed, + // TMA-aligned scales for DeepGEMM. + ops.def( + "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, " + "Tensor! output_s_packed, int group_size, float eps, float fp8_min, " + "float fp8_max) -> ()"); + // Compute per-token-group INT8 quantized tensor and scaling factor. + ops.def( + "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! " + "output_s, int group_size, float eps, float int8_min, float int8_max) -> " + "()"); + ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); - custom_ar.def("init_custom_qr", &init_custom_qr); - custom_ar.def("qr_destroy", &qr_destroy); - custom_ar.def("qr_get_handle", &qr_get_handle); +#ifndef USE_ROCM - custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()"); - custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles); + // Note about marlin kernel 'workspace' arguments: + // Technically these should be mutable since they are modified by the kernel. + // But since they are set back to zero once the kernel is finished we can + // hand wave and say that they have no net effect. + // + // The reason to mark 'workspace' as immutable is so that they don't interfere + // with using ScalarType arguments in the ops. If they are marked as mutable, + // pytorch throws an assert in + // 'torch._higher_order_ops._register_effectful_op' that prevents these + // kernels from being torch.compile'd. + // See the following document for more info on custom types and ops that use + // custom types: + // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA - custom_ar.def("qr_max_size", &qr_max_size); -} + // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. + ops.def( + "machete_supported_schedules(" + " ScalarType a_type," + " int b_type," + " ScalarType? maybe_group_scales_type," + " ScalarType? maybe_group_zeros_type," + " ScalarType? maybe_channel_scales_type," + " ScalarType? maybe_token_scales_type," + " ScalarType? maybe_out_type" + ") -> str[]"); + ops.def( + "machete_mm(" + " Tensor A," + " Tensor B," + " int b_type," + " ScalarType? out_type," + " Tensor? group_scales," + " Tensor? group_zeros," + " int? group_size," + " Tensor? channel_scales," + " Tensor? token_scales," + " str? schedule" + ") -> Tensor"); + ops.def( + "machete_prepack_B(" + " Tensor B," + " ScalarType a_type," + " int b_type," + " ScalarType? group_scales_type" + ") -> Tensor"); + // conditionally compiled so impl registration is in source file + + // Marlin GEMM + ops.def( + "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " + "Tensor? b_bias_or_none,Tensor b_scales, " + "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " + "Tensor workspace, int b_type_id, " + "SymInt size_m, SymInt size_n, SymInt size_k, " + "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // gptq_marlin repack from GPTQ. + ops.def( + "gptq_marlin_repack(Tensor b_q_weight, " + "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // awq_marlin repack from AWQ. + ops.def( + "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " + "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // preprocess W-int4A-fp8 weight for marlin kernel + ops.def( + "marlin_int4_fp8_preprocess(Tensor qweight, " + "Tensor? qzeros_or_none, bool inplace) -> Tensor"); + // conditionally compiled so impl registrations are in source file #endif -REGISTER_EXTENSION(TORCH_EXTENSION_NAME) +#ifndef USE_ROCM + // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column + // quantization, as well as bias + ops.def( + "cutlass_scaled_mm(Tensor! out, Tensor a," + " Tensor b, Tensor a_scales," + " Tensor b_scales, Tensor? bias) -> ()"); + + // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column + // quantization. + ops.def( + "cutlass_scaled_mm_azp(Tensor! out, Tensor a," + " Tensor b, Tensor a_scales," + " Tensor b_scales, Tensor azp_adj," + " Tensor? azp, Tensor? bias) -> ()"); + + // Check if cutlass scaled_mm is supported for CUDA devices of the given + // capability + ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool"); + + // Check if cutlass grouped gemm is supported for CUDA devices of the given + // capability + ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool"); + + // CUTLASS w8a8 grouped GEMM + ops.def( + "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, " + " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, " + " Tensor problem_sizes, Tensor a_strides, " + " Tensor b_strides, Tensor c_strides, bool per_act_token, " + " bool per_out_ch) -> ()"); + + // A function that computes data required to run fused MoE with w8a8 grouped + // GEMM. It takes topk_ids as an input, and computes expert_offsets + // (token start indices of each expert). In addition to this, it computes + // problem sizes for each expert's multiplication used by the two mms called + // from fused MoE operation, and arrays with permutations required to shuffle + // and de-shuffle the input/output of the fused operation. + ops.def( + "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, " + " Tensor! problem_sizes1, Tensor! problem_sizes2, " + " Tensor! input_permutation, " + " Tensor! output_permutation, int num_experts, " + " int n, int k, Tensor? blockscale_offsets, " + " bool is_gated) -> ()"); + + // compute per-expert problem sizes from expert_first_token_offset + // produced by vLLM's moe_permute kernel + ops.def( + "get_cutlass_moe_mm_problem_sizes_from_expert_offsets(" + " Tensor expert_first_token_offset, " + " Tensor! problem_sizes1, " + " Tensor! problem_sizes2, " + " int n, int k, bool swap_ab) -> ()"); + + // A function that computes data required to run fused MoE with w8a8 grouped + // GEMM in batched expert format. It takes expert_num_tokens + // as an input, and computes expert_offsets (token start indices of each + // expert). In addition to this, it computes problem sizes for each expert's + // multiplication used by the two mms called from fused MoE operation. + ops.def( + "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, " + " Tensor! problem_sizes1, " + " Tensor! problem_sizes2, " + " Tensor expert_num_tokens, " + " int num_local_experts, int padded_m, " + " int n, int k) -> ()"); + + // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3) + ops.def( + "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> " + "bool"); + + // CUTLASS nvfp4 block scaled GEMM + ops.def( + "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b," + " Tensor block_scale_a, Tensor block_scale_b," + " Tensor alpha) -> ()"); + + // cutlass nvfp4 block scaled group GEMM + ops.def( + "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b," + " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas," + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + + // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE) + ops.def( + "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b," + " Tensor a_blockscale, Tensor b_blockscales," + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + + // Compute NVFP4 block quantized tensor. + ops.def( + "scaled_fp4_quant(Tensor input," + " Tensor input_scale, bool " + "is_sf_swizzled_layout) -> (Tensor, Tensor)"); + + // Out variant + // TODO: Add out_variant tag once PyTorch supports it (added in 2.11) + // This registration is now migrated to stable ABI + // at::Tag::out_variant is not available in the stable ABI (enum_tag.h is not + // yet in torch/headeronly), the tag should be applied from Python + // via torch.library.Library.define(..., tags=(torch.Tag.out_variant,)) + // with the .impl remaining in C++. + // See pytorch/pytorch#176117. + ops.def( + "scaled_fp4_quant.out(Tensor input," + " Tensor input_scale, bool " + "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) " + "-> ()"); + + // Compute NVFP4 experts quantization. + ops.def( + "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale," + "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts) -> ()"); + + // Fused SiLU+Mul+NVFP4 experts quantization. + ops.def( + "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! " + "output_scale," + "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts) -> ()"); + + // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs). + ops.def( + "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + + // Fused SiLU+Mul+MXFP4 experts quantization. + ops.def( + "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " + "output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + + // Fused SiLU+Mul+NVFP4 quantization. + ops.def( + "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, " + "Tensor input, Tensor input_global_scale) -> ()"); + + // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices + // of the given capability + ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool"); + + // CUTLASS w4a8 GEMM + ops.def( + "cutlass_w4a8_mm(" + " Tensor A," + " Tensor B," + " Tensor group_scales," + " int group_size," + " Tensor channel_scales," + " Tensor token_scales," + " ScalarType? out_type," + " str? maybe_schedule" + ") -> Tensor"); + + // pack scales + ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor"); + + // encode and reorder weight matrix + ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor"); + + // CUTLASS w4a8 grouped GEMM + ops.def( + "cutlass_w4a8_moe_mm(" + " Tensor! out_tensors," + " Tensor a_tensors," + " Tensor b_tensors," + " Tensor a_scales," + " Tensor b_scales," + " Tensor b_group_scales," + " int b_group_size," + " Tensor expert_offsets," + " Tensor problem_sizes," + " Tensor a_strides," + " Tensor b_strides," + " Tensor c_strides," + " Tensor group_scale_strides," + " str? maybe_schedule" + ") -> ()"); + + ops.def( + "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, " + "Tensor)"); + + // SM100 CUTLASS MLA decode + // conditionally compiled so impl registrations are in source file + ops.def( + "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope," + " Tensor q_pe, Tensor kv_c_and_k_pe_cache," + " Tensor seq_lens, Tensor page_table," + " Tensor workspace, float scale," + " int num_kv_splits) -> ()"); + + ops.def( + "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches," + " int sm_count, int num_kv_splits) " + "-> int"); + // Quantized GEMM for AWQ. + ops.def( + "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, " + "Tensor _zeros, SymInt split_k_iters) -> Tensor"); + + // Dequantization for AWQ. + ops.def( + "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, " + "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor"); + + // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). + // conditionally compiled so impl registration is in source file + ops.def( + "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, " + "bool enable_pdl=False) -> ()"); + + // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). + // conditionally compiled so impl registration is in source file + ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); +#endif + + // Merge attn states + // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 + // can be used to combine partial attention results (in the split-KV case) + ops.def( + "merge_attn_states(" + " Tensor! output," + " Tensor!? output_lse," + " Tensor prefix_output," + " Tensor prefix_lse," + " Tensor suffix_output," + " Tensor suffix_lse," + " int!? prefill_tokens_with_context," + " Tensor? output_scale=None) -> ()"); + + // Hadamard transforms + // conditionally compiled so impl registration is in source file + ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor"); + + // Apply Root Mean Square (RMS) Normalization to the input tensor. + ops.def( + "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) " + "-> " + "()"); + + // In-place fused Add and RMS Normalization. + ops.def( + "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " + "float epsilon) -> ()"); + + // Layernorm-quant + // Apply Root Mean Square (RMS) Normalization to the input tensor. + ops.def( + "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, " + "Tensor scale, float epsilon) -> " + "()"); + + // In-place fused Add and RMS Normalization. + ops.def( + "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, " + "Tensor! residual, Tensor weight, " + "Tensor scale, float epsilon) -> ()"); + + // Fused Layernorm + Quant kernels + ops.def( + "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, " + "Tensor weight, Tensor! scale, float epsilon, " + "Tensor? scale_ub, Tensor!? residual) -> ()"); + + // Fused Layernorm + Block quant kernels + ops.def( + "rms_norm_per_block_quant(Tensor! result, Tensor input, " + "Tensor weight, Tensor! scale, float epsilon, " + "Tensor? scale_ub, Tensor!? residual, int group_size, " + "bool is_scale_transposed) -> ()"); + + // Fused SiLU+Mul + per-block quantization + ops.def( + "silu_and_mul_per_block_quant(" + "Tensor! out, " + "Tensor input, " + "Tensor! scales, " + "int group_size, " + "Tensor? scale_ub=None, " + "bool is_scale_transposed=False) -> ()"); + + // Rotary embedding + // Apply GPT-NeoX or GPT-J style rotary embedding to query and key. + ops.def( + "rotary_embedding(Tensor positions, Tensor! query," + " Tensor!? key, int head_size," + " Tensor cos_sin_cache, bool is_neox, int " + "rope_dim_offset=0, bool inverse=False) -> ()"); + + // Function for fused QK Norm and RoPE + ops.def( + "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " + "int num_heads_k, int num_heads_v, int head_dim, float eps, " + "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " + "bool is_neox, Tensor position_ids, " + "int forced_token_heads_per_warp=-1) -> ()"); + + // q_head_padded is the padded Q head count of the returned tensor, or 0 to + // do the KV insert alone and return an empty tensor. The Q knobs are + // independent: apply_q_norm and apply_q_rope each drop that step for Q + // alone (KV is always rotated), and is_q_interleaved reads and writes Q in + // FlashMLA's mega-attention chunk-interleaved layout, which moves the + // padding heads to the tail of every head-dim chunk. + ops.def( + "fused_deepseek_v4_kv_rope_insert(" + "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " + "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, " + "bool kv_mxfp8=False) -> " + "()"); + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" + "Tensor q_in, Tensor kv, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "int q_head_padded, float eps, int cache_block_size, " + "bool apply_q_norm=True, bool kv_mxfp8=False, bool apply_q_rope=True, " + "bool is_q_interleaved=False) -> Tensor"); + + // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate + // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(" + "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, " + "Tensor position_ids, Tensor cos_sin_cache, float eps, " + "int cache_block_size, bool apply_q_norm=True) -> ()"); + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(" + "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " + "int cache_block_size, bool apply_q_norm=True) -> ()"); + + // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion. + ops.def( + "fused_kimi_k3_mla_key_concat_kv_cache_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_key_concat_ds_mla_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_kv_concat(Tensor k_nope, Tensor k_pe, Tensor! k_out) " + "-> ()"); + ops.def( + "fused_kimi_k3_mla_kv_concat_quant_fp8(" + "Tensor k_nope, Tensor k_pe, Tensor v, Tensor! k_fp8, Tensor! v_fp8) " + "-> ()"); + ops.def( + "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" + "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, " + "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, " + "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, " + "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + + // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the + // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla). + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + +#ifndef USE_ROCM + ops.def( + "minimax_allreduce_rms_qk(" + "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " + "Tensor workspace, int q_size, int kv_size, int rank, int nranks, " + "float eps) -> (Tensor, Tensor)"); +#endif + + // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert. + ops.def( + "fused_minimax_m3_qknorm_rope_kv_insert(" + "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, " + "Tensor cos_sin_cache, Tensor positions, int num_heads, " + "int num_kv_heads, int rotary_dim, float eps, " + "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, " + "int num_index_heads, " + "Tensor? slot_mapping, Tensor? index_slot_mapping, " + "Tensor!? kv_cache, Tensor!? index_cache, " + "int block_size, Tensor!? q_out, Tensor!? index_q_out, " + "str kv_cache_dtype, bool skip_index_branch=False, " + "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, " + "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()"); + +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE + ops.def( + "fused_kda_decode(" + "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, " + "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, " + "Tensor state_indices, Tensor! state, Tensor! out, " + "float? lower_bound=None, Tensor? output_gate=None, " + "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()"); +#endif + +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE + ops.def( + "fused_gdn_decode_post_conv_mtp(" + "Tensor mixed_qkv, Tensor a, Tensor b, Tensor A_log, Tensor dt_bias, " + "Tensor state_indices, Tensor cu_seqlens, Tensor num_accepted_tokens, " + "Tensor! state, Tensor output_gate, Tensor norm_weight, Tensor! out, " + "float scale, float norm_eps=1e-5, " + "str output_gate_activation='silu') -> ()"); +#endif + +#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK + ops.def( + "fused_kda_prologue(" + "Tensor q, Tensor k, Tensor v, Tensor raw_g, Tensor raw_beta, " + "Tensor A_log, Tensor dt_bias, Tensor! qg, Tensor! w, Tensor! u, " + "Tensor! kg_t, Tensor! aqk, Tensor! decay, Tensor cu_seqlens, " + "Tensor chunk_indices, Tensor? conv_weight, Tensor(e!)? conv_state, " + "Tensor? conv_state_indices, Tensor? conv_has_initial_state, " + "float scale, float lower_bound) -> ()"); + ops.def( + "fused_kda_chunk(" + "Tensor qg, Tensor w, Tensor u, Tensor kg_t, Tensor aqk, Tensor decay, " + "Tensor? initial_state, Tensor(a!)? final_state, Tensor! out, " + "Tensor cu_seqlens, Tensor chunk_offsets, float scale, " + "Tensor(b!)? group_state, int groups, " + "Tensor(c!)? checkpoint_state=None, Tensor? checkpoint_offsets=None, " + "Tensor? checkpoint_state_indices=None, Tensor(d!)? state_cache=None, " + "Tensor? state_indices=None, Tensor? has_initial_state=None) -> ()"); +#endif + +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.def( + "kimi_k3_attn_res(" + "Tensor! prefix, Tensor? delta, Tensor! blocks, Tensor norm_weight, " + "Tensor qk_weight, Tensor? output_norm_weight, Tensor! output, " + "int num_blocks, " + "int block_write_idx, float eps, " + "float output_norm_eps) -> ()"); +#endif + + // Apply repetition penalties to logits in-place. + ops.def( + "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " + "Tensor output_mask, Tensor repetition_penalties) -> ()"); + + // Optimized top-k per row operations. + ops.def( + "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, " + "Tensor! indices, int numRows, int stride0, " + "int stride1, int topK) -> ()"); + + ops.def( + "top_k_per_row_decode(Tensor logits, int next_n, " + "Tensor seq_lens, Tensor! indices, " + "int numRows, int stride0, int stride1, int topK) -> ()"); + + ops.def( + "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); + +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.def( + "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); +#endif + + // Activation ops + ops.def( + "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " + "y_q, Tensor! y_s, bool use_ue8m0) -> ()"); + ops.def("weak_ref_tensor(Tensor input) -> Tensor"); + + // Activation function used in SwiGLU. + ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); + + ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); + + // SwiGLU activation with input clamping. + // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to + // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up. + ops.def( + "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, " + "float alpha=1.0, float beta=0.0) -> ()"); + + // SwiGLU activation with FP8 quantization. + ops.def( + "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); + + // Activation function used in GeGLU with `none` approximation. + ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); + + // Activation function used in GeGLU with `tanh` approximation. + ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); + + // FATReLU implementation. + ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()"); + + ops.def( + "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float " + "limit=7.0) " + "-> ()"); + + // SituGLU implementation used in Kimi models. + ops.def( + "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float " + "linear_beta=-1.0) -> ()"); + // Fused SituGLU activation + dynamic FP8 quantization for the Humming w2 path + // (writes the fp8 down input and its float32 scale). group_size=0 -> + // per-token scale [.., 1]; group_size=128 -> k-major block-FP8 scale [.., + // d/128]. + ops.def( + "situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, " + "float beta=1.0, float linear_beta=-1.0, int group_size=0, " + "Tensor? num_valid_tokens=None, int topk=1) -> ()"); + ops.def( + "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " + "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); + ops.def( + "masked_moe_activation(Tensor! out, Tensor input, Tensor " + "valid_token_counts, str activation, float clamp_limit=0.0, float " + "alpha=1.0, float beta=0.0, float situ_beta=1.0, float " + "situ_linear_beta=-1.0) -> ()"); + + // GELU implementation used in GPT-2. + ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); + + // Approximate GELU implementation. + ops.def("gelu_fast(Tensor! out, Tensor input) -> ()"); + + // Quick GELU implementation. + ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); + + // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2 + ops.def("relu_squared(Tensor! out, Tensor input) -> ()"); + + // Compute int8 quantized tensor for given scaling factor. + ops.def( + "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale," + "Tensor? azp) -> ()"); + + // Compute int8 quantized tensor and scaling factor + ops.def( + "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, " + "Tensor!? azp) -> ()"); + + // Compute FP8 quantized tensor for given scaling factor. + // Supports per-tensor, per-channel, per-token, and arbitrary 2D group + // scaling. Optional group_m/group_n specify the group shape explicitly; + // required for 1D scales to disambiguate per-channel vs per-token. + ops.def( + "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, " + "int[]? group_shape=None) -> ()"); + + // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor. + ops.def( + "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) " + "-> " + "()"); + + // Compute dynamic-per-token FP8 quantized tensor and scaling factor. + ops.def( + "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, " + "Tensor! scale, Tensor? scale_ub) -> " + "()"); + + // Quantized GEMM for GPTQ. + // Note: even though the C++ inferred schema is correct for this op, it seems + // to prevent the meta function registry. + ops.def( + "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, " + "Tensor b_gptq_scales, bool use_exllama, bool " + "use_v2_format, int bit) " + "-> Tensor"); + + // Post processing for GPTQ. + ops.def("gptq_shuffle(Tensor! q_weight, int bit) -> ()"); + + // Mamba selective scan kernel + ops.def( + "selective_scan_fwd(Tensor! u, Tensor! delta," + "Tensor! A, Tensor! B, Tensor! C," + "Tensor? D_, Tensor!? z_, Tensor? delta_bias_," + "bool delta_softplus," + "Tensor? query_start_loc," + "Tensor? cache_indices," + "Tensor? has_initial_state," + "Tensor! ssm_states," + "int null_block_id," + "int block_size," + "Tensor? block_idx_first_scheduled_token," + "Tensor? block_idx_last_scheduled_token," + "Tensor? initial_state_idx," + "Tensor? cu_chunk_seqlen," + "Tensor? last_chunk_indices) -> ()"); + + // LongCat n-gram embedding index kernel. All tensor args are marked mutable + // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table + // and n_gram_ids are actually written in place. + // Fused vocab-parallel embedding: gather the rows this rank owns and write + // zeros for the rest, so the following all-reduce reconstructs the full + // embedding. + ops.def( + "vocab_parallel_embedding(Tensor! out, Tensor input_ids, Tensor weight, " + "int org_vocab_start_index, int org_vocab_end_index, " + "int num_org_vocab_padding, int added_vocab_start_index, " + "int added_vocab_end_index) -> ()"); + ops.def( + "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, " + "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, " + "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, " + "Tensor(f!) row_indices, Tensor(g!) column_starts, " + "Tensor(h!) n_gram_ids) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { + // LongCat n-gram embedding index kernel. + ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); + ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding)); + + // Per-token group quantization + ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); + ops.impl("per_token_group_fp8_quant_packed", + TORCH_BOX(&per_token_group_quant_8bit_packed)); + ops.impl("per_token_group_quant_int8", + TORCH_BOX(&per_token_group_quant_int8)); + +#ifndef USE_ROCM + // CUTLASS scaled_mm ops + ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm)); + ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp)); + ops.impl("cutlass_moe_mm", TORCH_BOX(&cutlass_moe_mm)); + ops.impl("get_cutlass_moe_mm_data", TORCH_BOX(&get_cutlass_moe_mm_data)); + ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets", + TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets)); + ops.impl("get_cutlass_batched_moe_mm_data", + TORCH_BOX(&get_cutlass_batched_moe_mm_data)); + + // FP4/NVFP4 ops + ops.impl("cutlass_scaled_fp4_mm", TORCH_BOX(&cutlass_scaled_fp4_mm)); + ops.impl("scaled_fp4_quant", TORCH_BOX(&scaled_fp4_quant_func)); + ops.impl("scaled_fp4_quant.out", TORCH_BOX(&scaled_fp4_quant_out)); + ops.impl("scaled_fp4_experts_quant", TORCH_BOX(&scaled_fp4_experts_quant)); + ops.impl("silu_and_mul_scaled_fp4_experts_quant", + TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); + ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); + // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). + // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. + + // AWQ ops + ops.impl("awq_gemm", TORCH_BOX(&awq_gemm)); + ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize)); + + // DSV3 fused A GEMM: conditionally compiled so impl registration is in + // source file (dsv3_fused_a_gemm.cu) +#endif + + ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states)); + + // Layernorm kernels (shared CUDA/ROCm) + ops.impl("rms_norm", TORCH_BOX(&rms_norm)); + ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm)); + + // Layernorm-quant kernels (shared CUDA/ROCm) + ops.impl("rms_norm_static_fp8_quant", TORCH_BOX(&rms_norm_static_fp8_quant)); + ops.impl("fused_add_rms_norm_static_fp8_quant", + TORCH_BOX(&fused_add_rms_norm_static_fp8_quant)); + + // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) + ops.impl("rms_norm_dynamic_per_token_quant", + TORCH_BOX(&rms_norm_dynamic_per_token_quant)); + ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant)); + ops.impl("silu_and_mul_per_block_quant", + TORCH_BOX(&silu_and_mul_per_block_quant)); + + // Positional encoding kernels (shared CUDA/ROCm) + ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); + ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); + ops.impl("fused_deepseek_v4_kv_rope_insert", + TORCH_BOX(&fused_deepseek_v4_kv_rope_insert)); + ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); + ops.impl( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert)); + ops.impl( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); + ops.impl("fused_kimi_k3_mla_kv_concat", + TORCH_BOX(&fused_kimi_k3_mla_kv_concat)); + ops.impl("fused_kimi_k3_mla_kv_concat_quant_fp8", + TORCH_BOX(&fused_kimi_k3_mla_kv_concat_quant_fp8)); + ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert)); +#ifndef USE_ROCM + ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); +#endif + ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", + TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE + ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); +#endif + +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE + ops.impl("fused_gdn_decode_post_conv_mtp", + TORCH_BOX(&fused_gdn_decode_post_conv_mtp)); +#endif + +#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK + ops.impl("fused_kda_prologue", TORCH_BOX(&fused_kda_prologue)); + ops.impl("fused_kda_chunk", TORCH_BOX(&fused_kda_chunk)); +#endif + +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); +#endif + + // Sampler kernels (shared CUDA/ROCm) + ops.impl("apply_repetition_penalties_", + TORCH_BOX(&apply_repetition_penalties_)); + ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); + ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); + ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); +#endif + + // Activation kernels (shared CUDA/ROCm) + ops.impl("persistent_masked_m_silu_mul_quant", + TORCH_BOX(&persistent_masked_m_silu_mul_quant)); + ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor)); + ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant)); + ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul)); + ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu)); + ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul)); + ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); + ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); + ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); + ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); + ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant)); + ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); + ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation)); + ops.impl("gelu_new", TORCH_BOX(&gelu_new)); + ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); + ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); + ops.impl("relu_squared", TORCH_BOX(&relu_squared)); + ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp)); + + // INT8 quantization kernels + ops.impl("static_scaled_int8_quant", TORCH_BOX(&static_scaled_int8_quant)); + ops.impl("dynamic_scaled_int8_quant", TORCH_BOX(&dynamic_scaled_int8_quant)); + + // FP8 quantization kernels + ops.impl("static_scaled_fp8_quant", TORCH_BOX(&static_scaled_fp8_quant)); + ops.impl("dynamic_scaled_fp8_quant", TORCH_BOX(&dynamic_scaled_fp8_quant)); + ops.impl("dynamic_per_token_scaled_fp8_quant", + TORCH_BOX(&dynamic_per_token_scaled_fp8_quant)); + + // GPTQ kernels + ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm)); + ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle)); + + // Mamba kernels + ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { + ops.impl("get_cuda_view_from_cpu_tensor", + TORCH_BOX(&get_cuda_view_from_cpu_tensor)); +} + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) { + cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int"); + cuda_utils.def( + "get_max_shared_memory_per_block_device_attribute(int device_id) -> int"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd, + cuda_utils) { + cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute)); + cuda_utils.impl("get_max_shared_memory_per_block_device_attribute", + TORCH_BOX(&get_max_shared_memory_per_block_device_attribute)); +} + +// These capability-check functions take only primitive args (no tensors), so +// there is no device to dispatch on. CompositeExplicitAutograd makes them +// available for all backends. This is the stable ABI equivalent of calling +// ops.impl("op_name", &func) without a dispatch key in the non-stable API. +STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { +#ifndef USE_ROCM + ops.impl("cutlass_scaled_mm_supports_fp8", + TORCH_BOX(&cutlass_scaled_mm_supports_fp8)); + ops.impl("cutlass_group_gemm_supported", + TORCH_BOX(&cutlass_group_gemm_supported)); + ops.impl("cutlass_scaled_mm_supports_block_fp8", + TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8)); + ops.impl("cutlass_scaled_mm_supports_fp4", + TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); +#endif +} + +// Cache ops +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { + // Swap in (out) the cache blocks from src to dst. + ops.def( + "swap_blocks(Tensor src, Tensor! dst," + " int block_size_in_bytes, Tensor block_mapping) -> ()"); + + // Batch swap: submit all block copies in a single driver call. + ops.def( + "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," + " Tensor sizes," + " bool is_src_access_order_any=False) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache(Tensor key, Tensor value," + " Tensor! key_cache, Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache_flash(Tensor key, Tensor value," + " Tensor! key_cache," + " Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Concat kv_c and k_pe and cache them. + ops.def( + "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," + " Tensor! kv_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor scale) -> ()"); + + // Grouped concat_and_cache_mla across all layers. Each layer's cache base + // pointer and optional plain-FP8 scale are read from device tensors. + ops.def( + "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe," + " Tensor kv_cache_ptrs," + " Tensor slot_mapping," + " int block_size, int block_stride," + " int entry_stride," + " Tensor? kv_scales=None," + " str kv_cache_dtype='auto') -> ()"); + +#ifndef USE_ROCM + ops.def( + "hisparse_resolve_residency(Tensor host_cache," + " Tensor! hot_cache," + " Tensor hot_block_table," + " Tensor global_indices," + " Tensor! hot_indices," + " Tensor! device_global_indices," + " Tensor! lru_slots," + " Tensor? request_state_indices," + " int region_stride," + " int max_union_rows," + " Tensor(a!)? miss_mask=None," + " Tensor(b!)? stats=None," + " Tensor(c!)? attention_indices=None," + " int attention_block_stride=0," + " Tensor? request_ids=None," + " Tensor? source_block_table=None," + " int source_block_size=0," + " Tensor(d!)? resolved_global_indices=None," + " Tensor(e!)? valid_counts=None," + " Tensor(f!)? swap_host_physical_rows=None," + " Tensor(g!)? swap_device_physical_rows=None," + " Tensor(h!)? swap_counts=None," + " Tensor? resident_block_table=None," + " int resident_block_size=0," + " int resident_null_block=0) -> ()"); + + ops.def( + "hisparse_invalidate_written_slots(Tensor! device_global_indices," + " Tensor request_state_indices," + " Tensor req_id_per_token," + " Tensor written_slots) -> ()"); + + ops.def( + "hisparse_gather_plan(Tensor host_cache," + " Tensor! hot_cache," + " Tensor global_indices," + " Tensor hot_indices," + " Tensor miss_mask," + " Tensor? request_state_indices," + " Tensor(a!)? attention_indices=None," + " int attention_block_stride=0) -> ()"); + + ops.def( + "hisparse_gather_compact(Tensor host_cache," + " Tensor! hot_cache," + " Tensor miss_global_indices," + " Tensor miss_hot_indices," + " Tensor miss_counts) -> ()"); + +#endif // !USE_ROCM + + // Rotate Q and K, then write to kv cache for MLA + ops.def( + "concat_and_cache_mla_rope_fused(" + " Tensor positions," + " Tensor! q_pe," + " Tensor! k_pe," + " Tensor kv_c," + " Tensor cos_sin_cache," + " bool is_neox," + " Tensor slot_mapping," + " Tensor! kv_cache," + " str kv_cache_dtype," + " Tensor kv_cache_scale) -> ()"); + + // Convert the key and value cache to fp8 data type. + ops.def( + "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " + "str kv_cache_dtype) -> ()"); + + // Gather cache blocks from src_cache to dst, dequantizing from + // src_cache's dtype to dst's dtype if necessary. + ops.def( + "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " + " Tensor block_table, Tensor cu_seq_lens, " + " Tensor token_to_seq, " + " int num_tokens, " + " str kv_cache_dtype, " + " Tensor scale, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " + "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " + "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " + "seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? " + "device_row_ids=None) -> ()"); + + ops.def( + "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, " + "Tensor block_table, Tensor workspace_starts, int batch_size) -> ()"); + + ops.def( + "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " + "slot_mapping, " + "int quant_block_size, str kv_cache_dtype) -> ()"); + + ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); + + ops.def( + "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " + "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); +} + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) { + custom_ar.def( + "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " + "int rank, bool fully_connected) -> int"); + custom_ar.def( + "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ar.def("dispose(int fa) -> ()"); + custom_ar.def("meta_size() -> int"); + custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()"); + custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])"); + custom_ar.def( + "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()"); + custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)"); + custom_ar.def("open_mem_handle(Tensor mem_handle) -> int"); + custom_ar.def("free_shared_buffer(int ptr) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) { + custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar)); + custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) { + custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) { + custom_ar.impl("dispose", TORCH_BOX(&dispose)); + custom_ar.impl("meta_size", TORCH_BOX(&meta_size)); + custom_ar.impl("register_buffer", TORCH_BOX(®ister_buffer)); + custom_ar.impl("get_graph_buffer_ipc_meta", + TORCH_BOX(&get_graph_buffer_ipc_meta)); + custom_ar.impl("register_graph_buffers", TORCH_BOX(®ister_graph_buffers)); + custom_ar.impl("allocate_shared_buffer_and_handle", + TORCH_BOX(&allocate_shared_buffer_and_handle)); + custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { + ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { + ops.impl("swap_blocks", TORCH_BOX(&swap_blocks)); + ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); + ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); + ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); + ops.impl("concat_and_cache_mla_grouped", + TORCH_BOX(&concat_and_cache_mla_grouped)); + +#ifndef USE_ROCM + ops.impl("hisparse_resolve_residency", + TORCH_BOX(&hisparse_resolve_residency)); + ops.impl("hisparse_invalidate_written_slots", + TORCH_BOX(&hisparse_invalidate_written_slots)); + ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan)); + ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact)); +#endif // !USE_ROCM + ops.impl("concat_and_cache_mla_rope_fused", + TORCH_BOX(&concat_and_cache_mla_rope_fused)); + ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); + ops.impl("gather_and_maybe_dequant_cache", + TORCH_BOX(&gather_and_maybe_dequant_cache)); + ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache)); + ops.impl("cp_gather_and_upconvert_fp8_kv_cache", + TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache)); + ops.impl("cp_gather_and_upconvert_nvfp4_kv_cache", + TORCH_BOX(&cp_gather_and_upconvert_nvfp4_kv_cache)); + ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache)); + ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q)); + ops.impl("cp_gather_indexer_k_quant_cache", + TORCH_BOX(&cp_gather_indexer_k_quant_cache)); +} + +REGISTER_EXTENSION(_C_stable_libtorch) diff --git a/csrc/torch_utils.h b/csrc/torch_utils.h index 898b9e113e..f02346739e 100644 --- a/csrc/torch_utils.h +++ b/csrc/torch_utils.h @@ -1,17 +1,88 @@ #pragma once -// Shared TORCH_UTILS_CHECK across both libtorch stable and unstable source -// files. Keep this header free of CUTLASS/CUTE so attention/quant headers can -// use it. -// -// If TORCH_TARGET_VERSION is defined, we are building _C_stable_libtorch.so so -// use STD_TORCH_CHECK via header-only. -// Otherwise, use TORCH_CHECK via torch/all.h. +#include +#include +#include +#include +#include -#ifdef TORCH_TARGET_VERSION - #include - #define TORCH_UTILS_CHECK STD_TORCH_CHECK -#else - #include - #define TORCH_UTILS_CHECK TORCH_CHECK -#endif +#include +#include + +#include +#include +#include +#include + +// Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED. +#define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \ + STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__) + +// Device properties cache for stable ABI compatibility. +// Uses raw CUDA/HIP APIs instead of ATen functions. +// Using inline ensures a single instance across all translation units. +inline std::deque device_flags; +inline std::vector device_properties; +inline std::once_flag vectors_init_flag; + +inline void do_init_device_vectors() { + int device_count; + cudaError_t err = cudaGetDeviceCount(&device_count); + if (err != cudaSuccess) { + STD_TORCH_CHECK(false, "cudaGetDeviceCount failed: " + + std::string(cudaGetErrorString(err))); + } + device_flags.resize(device_count); + device_properties.resize(device_count); +} + +inline void initDeviceVectors() { + std::call_once(vectors_init_flag, do_init_device_vectors); +} + +inline void initDeviceProperty(int device_index) { + cudaDeviceProp device_prop{}; + cudaError_t err = cudaGetDeviceProperties(&device_prop, device_index); + if (err != cudaSuccess) { + STD_TORCH_CHECK(false, "cudaGetDeviceProperties failed: " + + std::string(cudaGetErrorString(err))); + } + device_properties[device_index] = device_prop; +} + +// Get device properties using raw CUDA/HIP APIs (stable ABI compatible). +// Caches results per device so cudaGetDeviceProperties is called at most once +// per device. +inline cudaDeviceProp* get_device_prop() { + initDeviceVectors(); + int device_index; + cudaError_t err = cudaGetDevice(&device_index); + if (err != cudaSuccess) { + STD_TORCH_CHECK( + false, "cudaGetDevice failed: " + std::string(cudaGetErrorString(err))); + } + STD_TORCH_CHECK(device_index >= 0 && static_cast(device_index) < + device_properties.size(), + "CUDA device index " + std::to_string(device_index) + + " out of range [0, " + + std::to_string(device_properties.size()) + ")"); + + std::call_once(device_flags[device_index], initDeviceProperty, device_index); + return &device_properties[device_index]; +} + +// Utility to get the current CUDA stream for a given device using stable APIs. +// Returns a cudaStream_t for use in kernel launches. +inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) { + void* stream_ptr = nullptr; + TORCH_ERROR_CODE_CHECK( + aoti_torch_get_current_cuda_stream(device_index, &stream_ptr)); + return reinterpret_cast(stream_ptr); +} + +// Utility to get the current cuBLAS handle using stable APIs. +inline cublasHandle_t get_current_cuda_blas_handle() { + void* blas_handle_ptr = nullptr; + TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&blas_handle_ptr)); + return reinterpret_cast(blas_handle_ptr); +} diff --git a/csrc/libtorch_stable/type_convert.cuh b/csrc/type_convert.cuh similarity index 100% rename from csrc/libtorch_stable/type_convert.cuh rename to csrc/type_convert.cuh diff --git a/csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu b/csrc/vocab_parallel_embedding_kernels.cu similarity index 100% rename from csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu rename to csrc/vocab_parallel_embedding_kernels.cu diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 4095c41a7e..6f76838f55 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -232,7 +232,7 @@ Supported quantization scheme/hardware combinations: - Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) -- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu) +- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu) ### SiLU+Mul + Quantization (`fuse_act_quant`) @@ -258,7 +258,7 @@ Supported quantization scheme/hardware combinations: - Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) - CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/) -- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu) +- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu) ### RMSNorm + Padding (`fuse_act_padding`) diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index 7ca207c8e4..173a2ea344 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -185,10 +185,7 @@ def test_every_dsv3_routed_shape_is_instantiated() -> None: let (3216, 7168) ship without its DISPATCH_DSV3_SHAPE(7168, 3216). """ source = ( - Path(__file__).resolve().parents[2] - / "csrc" - / "libtorch_stable" - / "dsv3_fused_a_gemm.cu" + Path(__file__).resolve().parents[2] / "csrc" / "dsv3_fused_a_gemm.cu" ).read_text(encoding="utf-8") # Benchmark-only shapes live behind VLLM_K3_BENCH_SHAPES and are not built # by default, so they must not count as available. diff --git a/tools/ci/kernel_symbol_map.py b/tools/ci/kernel_symbol_map.py index c88f2d9ad5..3177a10bba 100755 --- a/tools/ci/kernel_symbol_map.py +++ b/tools/ci/kernel_symbol_map.py @@ -32,13 +32,13 @@ Output (gzipped JSON): "errors": [], # [{"object", "source", "error"}, ...] "objects": [ { - "source": "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu", + "source": "csrc/fused_qknorm_rope_kernel.cu", "target": "_C", "object": "CMakeFiles/_C.dir/csrc/.../fused_qknorm_rope_kernel.cu.o", "device": true, "symbols": ["_Z21fusedQKNormRopeKernelI...", ...], - "deps": ["csrc/libtorch_stable/fused_qknorm_rope_kernel.cu", - "csrc/cuda_compat.h", ...] + "deps": ["csrc/fused_qknorm_rope_kernel.cu", + "csrc/core/cuda_compat.h", ...] # "error": "..." # present when cuobjdump failed on it }, ... diff --git a/vllm/distributed/device_communicators/quick_all_reduce.py b/vllm/distributed/device_communicators/quick_all_reduce.py index 406869ea2c..0aa2a60902 100644 --- a/vllm/distributed/device_communicators/quick_all_reduce.py +++ b/vllm/distributed/device_communicators/quick_all_reduce.py @@ -28,7 +28,7 @@ from vllm.distributed.utils import is_weak_contiguous # noqa: E402, F401 class QuickReduceRegime(Enum): - # Keep integer ids aligned with csrc/quickreduce/quick_reduce.h + # Keep integer ids aligned with csrc/rocm/quickreduce/quick_reduce.h FP = 0 INT8 = 1 INT6 = 2 diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py index da8c69eed0..a541debfe9 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py @@ -104,7 +104,8 @@ class CutlassInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): # azp_adj is the AZP adjustment term, used to account for weights. # It does not depend on scales or azp, so it is the same for # static and dynamic quantization. - # For more details, see csrc/quantization/w8a8/cutlass/Epilogues.md + # For more details, see + # csrc/quantization/w8a8/cutlass/Epilogues.md # https://github.com/vllm-project/vllm/blob/main/csrc/quantization/w8a8/cutlass/Epilogues.md if not config.input_symmetric: weight = getattr(layer, w_q_name) diff --git a/vllm/models/common/ops/fused_qk_rmsnorm.py b/vllm/models/common/ops/fused_qk_rmsnorm.py index e64e818210..fec8719a89 100644 --- a/vllm/models/common/ops/fused_qk_rmsnorm.py +++ b/vllm/models/common/ops/fused_qk_rmsnorm.py @@ -68,7 +68,8 @@ class FusedQKVRMSNormKernel(VllmTritonJitKernel["FusedQKVRMSNormKernel.CompileKe weight_ptr = kv_weight_ptr row_out = kv_out_ptr + token_idx * kv_out_stride - # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's + # RMSNorm in fp32 throughout — matches + # csrc/layernorm_kernels.cu's # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which # keep x, rrms, and w all in fp32 and perform a single cast at store. block = tl.arange(0, BLOCK_SIZE) diff --git a/vllm/models/kimi_k3/amd/ops/kda_chunk.py b/vllm/models/kimi_k3/amd/ops/kda_chunk.py index 035a49f852..a3fbada908 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_chunk.py +++ b/vllm/models/kimi_k3/amd/ops/kda_chunk.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ROCm entry point for the fused Kimi-K3 KDA chunk kernel. -The kernel in ``csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu`` +The kernel in ``csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu`` replaces the chunk-state recurrence and the output GEMM of the Triton chunk path with a single launch that keeps the per-chunk state in registers, so the ``[chunks, H, V, K]`` state tensor and the recomputed values never reach HBM. diff --git a/vllm/models/kimi_k3/amd/ops/kda_decode.py b/vllm/models/kimi_k3/amd/ops/kda_decode.py index 132cd58cf8..682776a807 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_decode.py +++ b/vllm/models/kimi_k3/amd/ops/kda_decode.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ROCm entry points for the fused Kimi-K3 KDA decode kernel. -The kernel in ``csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu`` +The kernel in ``csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu`` replaces, for a pure non-speculative decode batch, the three Triton launches and two copies the AMD KDA layer otherwise runs per layer: the packed causal conv1d update, the recurrent delta-rule step, and the gated output RMSNorm. diff --git a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py index 79de3af0bc..a39ec787e0 100644 --- a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py +++ b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py @@ -3,7 +3,7 @@ """Fused MLA prefill and decode epilogues for Kimi-K3. Thin wrappers over the CUDA ops in -``csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which +``csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``. - ``fused_mla_key_concat_kv_cache_insert`` (bf16): optionally apply RoPE, diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 0ff34f8ac5..4fe0be0b1f 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -118,7 +118,7 @@ token's KV cache is 352 Bytes, structured as: elements, stored permuted (an 8x4 -> 4x8 transpose: the scale for element block `s` lives at byte `8 * (s & 3) + (s >> 2)`) so that the 8 scales one FlashMLA dequant thread needs are contiguous. See the layout comment in - `csrc/libtorch_stable/cache_kernels.cu`. + `csrc/cache_kernels.cu`. """ diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index f1b0fa688e..f7474b6bed 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -712,7 +712,7 @@ class HiSparseRuntime: "HiSparse requires its compiled _C_cache_ops CUDA kernels " "(host-resident decode has no Python fallback). Rebuild vLLM " "from source so " - "csrc/libtorch_stable/hisparse_kernels.cu is included." + "csrc/hisparse_kernels.cu is included." ) self.max_num_reqs = max_num_reqs self.row_width = row_width