[Misc] Consolidate csrc shared headers and prune dead code (#60721)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
2026-10-12 09:46:22 -07:00
parent e0317b28ce
commit d79b52b182
242 changed files with 2783 additions and 3638 deletions
@@ -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:
@@ -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
@@ -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
@@ -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/
@@ -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:
@@ -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/

@@ -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
@@ -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}")
@@ -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}")
@@ -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=["."],
@@ -10,7 +10,7 @@
 #include <cmath>
 #include <type_traits>

-#include "../cuda_compat.h"
+#include "core/cuda_compat.h"
 #include "async_util.cuh"
 #include "cuda_vec_utils.cuh"
 #include "dispatch_utils.h"
@@ -6,9 +6,8 @@
 #include "../dispatch_utils.h"
 #include <torch/headeronly/core/ScalarType.h>

-#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 {

@@ -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 <torch/csrc/stable/library.h>

@@ -1,93 +0,0 @@
-#pragma once
-
-#include <torch/all.h>
-#include <c10/util/Optional.h>
-
-#include <map>
-#include <vector>
-
-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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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]
@@ -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 <algorithm>
@@ -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<int32_t>(num_tokens), block_size) +
-      2 * num_reqs;
+      div_ceil(static_cast<int32_t>(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<int32_t>(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<int32_t>()
                             : nullptr;
-  const int required_blocks = cuda_utils::ceil_div(total_tokens, block_size) +
-                              2 * static_cast<int32_t>(batch_size);
+  const int required_blocks =
+      div_ceil(total_tokens, block_size) + 2 * static_cast<int32_t>(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()),
@@ -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
@@ -1,18 +0,0 @@
-#pragma once
-#include <cstdlib>
-#include <string>
-
-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
@@ -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)
@@ -1,3 +0,0 @@
-#pragma once
-
-#define VLLM_IMPLIES(p, q) (!(p) || (q))
@@ -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)                                               \
@@ -1,7 +1,8 @@
 #pragma once

 #include <climits>
-#include <iostream>
+#include <cstdint>
+#include <cstdlib>

 inline constexpr uint32_t next_pow_2(uint32_t const num) {
  if (num <= 1) return num;
@@ -26,3 +27,16 @@ template <typename T>
 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
@@ -20,7 +20,7 @@
 */
 #pragma once

-#include "attention_generic.cuh"
+#include "dtype_generic.cuh"
 #include "dtype_float32.cuh"

 #ifndef USE_ROCM
@@ -20,7 +20,7 @@
 */
 #pragma once

-#include "attention_generic.cuh"
+#include "dtype_generic.cuh"
 #include "dtype_float32.cuh"

 #ifdef USE_ROCM
@@ -20,7 +20,7 @@
 */
 #pragma once

-#include "attention_generic.cuh"
+#include "dtype_generic.cuh"

 #include <stdint.h>

@@ -1,7 +1,8 @@
 #pragma once

-#include "attention_generic.cuh"
-#include "torch_utils.h"
+#include "dtype_generic.cuh"
+
+#include <torch/headeronly/util/Exception.h>

 #include <stdint.h>
 #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
@@ -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"
@@ -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,
@@ -0,0 +1,46 @@
+#pragma once
+
+#include <optional>
+#include <string>
+#include <torch/library.h>
+#include <tuple>
+#include <vector>
+
+void rms_norm(torch::Tensor& out, torch::Tensor& input,
+              std::optional<torch::Tensor> weight, double epsilon);
+
+void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
+                        std::optional<torch::Tensor> weight, double epsilon);
+
+void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
+                      std::optional<torch::Tensor> 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<torch::Tensor> const& azp);
+
+void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
+                               torch::Tensor& scales,
+                               std::optional<torch::Tensor> 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);
@@ -1,4 +1,4 @@
-#include "cache.h"
+#include <torch/all.h>
 #include "ops.h"
 #include "core/registration.h"

@@ -1,21 +1,8 @@
 #pragma once

+#include <cstdint>
 #include <stdio.h>

-#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 <typename T>
-HOST_DEVICE_INLINE constexpr std::enable_if_t<std::is_integral_v<T>, T>
-ceil_div(T a, T b) {
-  return (a + b - 1) / b;
-}
-
-};  // namespace cuda_utils
@@ -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 <int ngpus>
-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 <int ngpus>
-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 <int ngpus, bool final_sync = false>
 DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
  __syncthreads();
@@ -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 <typename Kernel>
-struct enable_sm120_only : Kernel {
-  template <typename... Args>
-  CUTLASS_DEVICE void operator()(Args&&... args) {
-#if defined __CUDA_ARCH__
-  #if __CUDA_ARCH__ == 1200
-    Kernel::operator()(std::forward<Args>(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 <typename Kernel>
 struct enable_sm120_family : Kernel {
@@ -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 <torch/csrc/stable/tensor.h>
-#endif
+#include <torch/csrc/stable/tensor.h>

 /*
   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 <typename T>
-struct identity {
-  CUTLASS_HOST_DEVICE
-  T operator()(T lhs) const { return lhs; }
-};
-
-template <typename ElementAcc, typename ElementD, typename TileShape>
-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<Compute, Accum>;
-  using ArgumentType = typename EVTCompute::Arguments;
-
-  template <typename... Args>
-  static ArgumentType prepare_args(Args... args) {
-    return {};
-  }
-};
-
 /*
 * This class provides the common load descriptors for the
 * ScaledEpilogue[...] classes
@@ -1,19 +1,12 @@
 #pragma once

-#include "torch_utils.h"
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Exception.h>
+#include <torch/headeronly/util/Half.h>

-// 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 <torch/csrc/stable/tensor.h>
-  #include <torch/headeronly/util/BFloat16.h>
-  #include <torch/headeronly/util/Half.h>
+// 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 <typename Stride>
 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<decltype(stride_ele)>;

    if (idx < tensor.dim()) {
      if constexpr (cute::is_static_v<StrideEle>) {
-        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 <typename T>
-struct equivalent_cutlass_type {
-  using type = T;
-};
-
-template <typename T>
-using equivalent_cutlass_type_t = typename equivalent_cutlass_type<T>::type;
-
-template <>
-struct equivalent_cutlass_type<torch::headeronly::Half> {
-  using type = cutlass::half_t;
-};
-
-template <>
-struct equivalent_cutlass_type<torch::headeronly::BFloat16> {
-  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<T>` compatible type, i.e.
@@ -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
@@ -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 {

@@ -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 <torch/all.h>
+#include <torch/headeronly/core/Dispatch.h>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/util/Exception.h>

 // 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"); \
  }
@@ -25,7 +25,7 @@
 #include <torch/headeronly/core/ScalarType.h>

 #include "core/registration.h"
-#include "libtorch_stable/torch_utils.h"
+#include "torch_utils.h"

 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
@@ -6,7 +6,7 @@
 #include <torch/headeronly/core/ScalarType.h>

 #include "core/registration.h"
-#include "libtorch_stable/torch_utils.h"
+#include "torch_utils.h"

 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
@@ -47,7 +47,7 @@
 #include <torch/csrc/stable/device.h>

 #include <cmath>
-#include "cuda_compat.h"
+#include "core/cuda_compat.h"
 #include "dispatch_utils.h"
 #include "type_convert.cuh"

@@ -57,16 +57,16 @@
 #include <torch/headeronly/core/ScalarType.h>
 #include <torch/csrc/stable/device.h>

-#include "cuda_compat.h"
+#include "core/cuda_compat.h"
 #include "dispatch_utils.h"
 #include "type_convert.cuh"

 #ifndef USE_ROCM
  #include <cuda_fp8.h>
-  #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
+  #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #else
  #include <hip/hip_fp8.h>
-  #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
+  #include "quantization/w8a8/fp8/amd/quant_utils.cuh"
 #endif
 #include <cuda_runtime.h>
 #include <type_traits>
@@ -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<<<grid, kBlockSize, 0, stream>>>(args...);
  // clang-format on
 #endif
@@ -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.
@@ -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"

@@ -10,7 +10,7 @@
 #include <cuda_runtime.h>

 #include "../torch_utils.h"
-#include "../../cuda_compat.h"
+#include "core/cuda_compat.h"

 namespace {

@@ -17,7 +17,6 @@

 #include "torch_utils.h"
 #include "ops.h"
-#include "../cuda_utils.h"

 #include <algorithm>
 #include <cstdint>
@@ -9,7 +9,7 @@
 #include <optional>

 #include "../torch_utils.h"
-#include "../../cuda_compat.h"
+#include "core/cuda_compat.h"

 namespace {

@@ -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"
@@ -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"
@@ -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 <float.h>
-#include <type_traits>
-
-namespace vllm {
-
-// Q*K^T operation.
-template <int THREAD_GROUP_SIZE, typename Vec, int N>
-inline __device__ float qk_dot_(const Vec (&q)[N], const Vec (&k)[N]) {
-  using A_vec = typename FloatVec<Vec>::Type;
-  // Compute the parallel products for Q*K^T (treat vector lanes separately).
-  A_vec qk_vec = mul<A_vec, Vec, Vec>(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 <typename T, int THREAD_GROUP_SIZE>
-struct Qk_dot {
-  template <typename Vec, int N>
-  static inline __device__ float dot(const Vec (&q)[N], const Vec (&k)[N]) {
-    return qk_dot_<THREAD_GROUP_SIZE>(q, k);
-  }
-};
-
-}  // namespace vllm
@@ -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;
-}
@@ -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 <torch/headeronly/core/Dispatch.h>
-#include <torch/headeronly/core/ScalarType.h>
-#include <torch/headeronly/util/Exception.h>
-
-// 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"); \
-  }
@@ -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<BoolConst>(...);
-/// });
-/// ```
-#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__();                \
-    }                                      \
-  }()
@@ -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()
@@ -1,822 +0,0 @@
-#pragma once
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/util/Exception.h>
-
-#include <optional>
-#include <string>
-#include <vector>
-
-#include <torch/csrc/stable/ops.h>
-
-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<torch::stable::Tensor> 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<torch::stable::Tensor> const& azp,
-                           std::optional<torch::stable::Tensor> 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<torch::stable::Tensor>& 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<torch::stable::Tensor, torch::stable::Tensor> 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<torch::stable::Tensor> 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<int64_t> prefill_tokens_with_context,
-    const std::optional<torch::stable::Tensor>& 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<torch::stable::Tensor> weight, double epsilon);
-
-void fused_add_rms_norm(torch::stable::Tensor& input,
-                        torch::stable::Tensor& residual,
-                        std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> scale_ub,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> scale_ub,
-                              std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> 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<torch::stable::Tensor> 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<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> position_ids,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor, torch::stable::Tensor>
-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<torch::stable::Tensor> index_q_norm_weight,
-    std::optional<torch::stable::Tensor> index_k_norm_weight,
-    int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
-    std::optional<torch::stable::Tensor> index_slot_mapping,
-    std::optional<torch::stable::Tensor> kv_cache,
-    std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
-    std::optional<torch::stable::Tensor> q_out,
-    std::optional<torch::stable::Tensor> index_q_out,
-    const std::string& kv_cache_dtype, bool skip_index_branch,
-    std::optional<torch::stable::Tensor> q_fp8_out, double q_fp8_scale,
-    std::optional<torch::stable::Tensor> kv_k_scale,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> 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<double> lower_bound,
-    std::optional<torch::stable::Tensor> output_gate,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> conv_weight,
-    std::optional<torch::stable::Tensor> conv_state,
-    std::optional<torch::stable::Tensor> conv_state_indices,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> initial_state,
-    std::optional<torch::stable::Tensor> final_state,
-    torch::stable::Tensor& out, torch::stable::Tensor const& cu_seqlens,
-    torch::stable::Tensor const& chunk_offsets, double scale,
-    std::optional<torch::stable::Tensor> group_state, int64_t groups,
-    std::optional<torch::stable::Tensor> checkpoint_state,
-    std::optional<torch::stable::Tensor> checkpoint_offsets,
-    std::optional<torch::stable::Tensor> checkpoint_state_indices,
-    std::optional<torch::stable::Tensor> state_cache,
-    std::optional<torch::stable::Tensor> state_indices,
-    std::optional<torch::stable::Tensor> has_initial_state);
-#endif
-
-#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
-void kimi_k3_attn_res(torch::stable::Tensor& prefix,
-                      std::optional<torch::stable::Tensor> delta,
-                      torch::stable::Tensor& blocks,
-                      torch::stable::Tensor const& norm_weight,
-                      torch::stable::Tensor const& qk_weight,
-                      std::optional<torch::stable::Tensor> 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<torch::stable::Tensor>& D_,
-    const std::optional<torch::stable::Tensor>& z_,
-    const std::optional<torch::stable::Tensor>& delta_bias_,
-    bool delta_softplus,
-    const std::optional<torch::stable::Tensor>& query_start_loc,
-    const std::optional<torch::stable::Tensor>& cache_indices,
-    const std::optional<torch::stable::Tensor>& has_initial_state,
-    const torch::stable::Tensor& ssm_states, int64_t null_block_id,
-    int64_t block_size,
-    const std::optional<torch::stable::Tensor>& block_idx_first_scheduled_token,
-    const std::optional<torch::stable::Tensor>& block_idx_last_scheduled_token,
-    const std::optional<torch::stable::Tensor>& initial_state_idx,
-    const std::optional<torch::stable::Tensor>& cu_chunk_seqlen,
-    const std::optional<torch::stable::Tensor>& last_chunk_indices);
-
-using fptr_t = int64_t;
-fptr_t init_custom_ar(const std::vector<int64_t>& 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<int64_t>& fake_ipc_ptrs);
-std::tuple<std::vector<int64_t>, std::vector<int64_t>>
-get_graph_buffer_ipc_meta(fptr_t _fa);
-void register_graph_buffers(fptr_t _fa,
-                            const std::vector<std::vector<int64_t>>& handles,
-                            const std::vector<std::vector<int64_t>>& offsets);
-std::tuple<int64_t, torch::stable::Tensor> 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<torch::stable::Tensor> 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<torch::stable::Tensor> const& azp);
-
-void dynamic_scaled_int8_quant(torch::stable::Tensor& out,
-                               torch::stable::Tensor const& input,
-                               torch::stable::Tensor& scales,
-                               std::optional<torch::stable::Tensor> 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<torch::headeronly::IntHeaderOnlyArrayRef> 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<torch::stable::Tensor> 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<torch::stable::Tensor> 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<torch::stable::Tensor> const& request_state_indices,
-    int64_t region_stride, int64_t max_union_rows,
-    std::optional<torch::stable::Tensor> const& miss_mask,
-    std::optional<torch::stable::Tensor> const& stats,
-    std::optional<torch::stable::Tensor> const& attention_indices,
-    int64_t attention_block_stride,
-    std::optional<torch::stable::Tensor> const& request_ids,
-    std::optional<torch::stable::Tensor> const& source_block_table,
-    int64_t source_block_size,
-    std::optional<torch::stable::Tensor> const& resolved_global_indices,
-    std::optional<torch::stable::Tensor> const& valid_counts,
-    std::optional<torch::stable::Tensor> const& swap_host_physical_rows,
-    std::optional<torch::stable::Tensor> const& swap_device_physical_rows,
-    std::optional<torch::stable::Tensor> const& swap_counts,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> const& request_state_indices,
-    std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> 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<torch::stable::Tensor> 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<torch::stable::Tensor> seq_starts = std::nullopt,
-    std::optional<torch::stable::Tensor> host_cache = std::nullopt,
-    std::optional<torch::stable::Tensor> host_row_ids = std::nullopt,
-    std::optional<torch::stable::Tensor> 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);
@@ -1,12 +0,0 @@
-#include <torch/csrc/stable/tensor.h>
-
-#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);
-}
@@ -1,20 +0,0 @@
-#pragma once
-
-#include <torch/csrc/stable/tensor.h>
-
-// 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);
@@ -1,88 +0,0 @@
-#pragma once
-
-#include <torch/csrc/inductor/aoti_torch/c/shim.h>
-#include <torch/csrc/stable/accelerator.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/util/shim_utils.h>
-
-#include <cuda_runtime.h>
-#include <cublas_v2.h>
-
-#include <deque>
-#include <mutex>
-#include <string>
-#include <vector>
-
-// 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<std::once_flag> device_flags;
-inline std::vector<cudaDeviceProp> 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<size_t>(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<cudaStream_t>(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<cublasHandle_t>(blas_handle_ptr);
-}
@@ -14,7 +14,7 @@
 #endif

 #include "selective_scan.h"
-#include "static_switch.h"
+#include "../dispatch_utils.h"

 template<int kNThreads_, int kNItems_, int kNRows_, bool kIsEvenLen_,
         bool kIsVariableB_, bool kIsVariableC_,
@@ -398,9 +398,9 @@ void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {
    // kIsVariableB, kIsVariableC and kHasZ are all set to True to reduce binary size
    constexpr bool kIsVariableB = true;
    constexpr bool kIsVariableC = true;
-    BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
-        BOOL_SWITCH(params.z_ptr != nullptr , kHasZ, [&] {
-            BOOL_SWITCH(params.query_start_loc_ptr != nullptr , kVarlen, [&] {
+    VLLM_STABLE_DISPATCH_BOOL(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
+        VLLM_STABLE_DISPATCH_BOOL(params.z_ptr != nullptr , kHasZ, [&] {
+            VLLM_STABLE_DISPATCH_BOOL(params.query_start_loc_ptr != nullptr , kVarlen, [&] {
                using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kNRows, kIsEvenLen, kIsVariableB, kIsVariableC, kHasZ,  kVarlen, input_t, weight_t, state_t>;
                constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t);
                dim3 grid(params.batch, params.dim / kNRows);
@@ -27,7 +27,7 @@
 #include <cooperative_groups.h>
 #include <cuda_runtime.h>

-#include "cuda_compat.h"
+#include "core/cuda_compat.h"
 #include "minimax_reduce_rms_kernel.h"

 #include <algorithm>
@@ -22,7 +22,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include <torch/headeronly/core/ScalarType.h>

-#include "libtorch_stable/torch_utils.h"
+#include "torch_utils.h"

 #include <cmath>
 #include <cstdint>
@@ -1,3 +1,2 @@
 sm*_kernel_*.cu
 kernel_selector.h
-kernel_*.cu
@@ -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                                                 \
@@ -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)               \
@@ -32,7 +32,7 @@
 #include <torch/headeronly/core/ScalarType.h>
 #include <torch/headeronly/util/Exception.h>

-#include "libtorch_stable/torch_utils.h"
+#include "torch_utils.h"

 #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t)               \
  static_assert(std::is_same<scalar_t, half>::value ||          \
@@ -8,11 +8,11 @@
 #include <torch/csrc/stable/tensor.h>
 #include <torch/headeronly/core/ScalarType.h>

-#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))

@@ -7,8 +7,8 @@
 #include <torch/headeronly/util/Exception.h>

 #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 <torch/csrc/stable/library.h>

--- a/Show More
+++ b/Show More