[torch.compile] Add compile-only mode

Summary
=======

This PR is on the way to overlapping torch.compile and weight loading. See [design doc]([https://docs.google.com/document/d/1hssZeQv_lJlKqOr0vfpoqEUNn4ASGHVRB9a49yhcY6E/edit?tab=t.0](https://docs.google.com/document/d/1hssZeQv_lJlKqOr0vfpoqEUNn4ASGHVRB9a49yhcY6E/edit?tab=t.0)) and [proof-of-concept PR](https://github.com/vllm-project/vllm/pull/36072) and [RFC](https://github.com/vllm-project/vllm/issues/34956)

- Adds `vllm compile <model> [options]` CLI command and `vllm.compile_model()` Python API
- These APIs populate vLLM's torch.compile cache and do nothing else. A subsequent `vllm serve` call can read from the cache and perform a warm start.
- They also use minimal GPU memory. This is accomplished through a combination of using FakeTensors (tensors with no storage that report device correctly) and Meta tensors (tensors with no storage that report device="meta"). Note that there is still some minimal GPU memory allocation (< 10 Mb, from GPUModelRunner runtime buffers), I did not go track down all of it, but I'm also not sure it matters.
- In the future we can extend "vllm compile" to more than just torch.compile; for example, if vLLM uses triton kernels, or JIT'ed flashinfer kernels, `vllm compile` may also just compile those and saved the compiled artifacts somewhere.

How it works
============
- `FakeModelLoader` wraps the real model loader. It initializes weights on `meta` device and runs weight post-processing on meta tensors.
- Before torch.compile tracing, `swap_meta_params_to_fake()` converts meta params to FakeTensors. This is required so that torch.compile sees Tensors with the correct devices.
- In theory we should also get the FakeModelLoader to give us FakeTensors, but FakeTensors do not yet support the type of Tensor subclasses that vLLM uses.
- We raise the `CompilationDone` exception after cache artifacts are saved to avoid executing with fake tensors. (calling torch.compile performs both the compilation and an initial run of invoking the compiled artifact with the inputs)
- `EngineCore` early-returns after `compile_or_warm_up_model`, skipping KV cache allocation, scheduler, and sampler setup

Test plan
=========
- Added tests for compile_only cold start followed by a warm start. The tests verify that the compile_only cold start uses no GPU memory, and the warm start does end up reading from the cache.

Future work
===========
In the following order:
- add an option to overlap torch.compile and weight loading. The main process will do weight loading while spawning a new process to do compile-only work (that does not use gpu memory)
- Extend this design to more weight loading schemes. For example, we currently support no weight processing. This will involve getting the additional weight processing to support meta tensors.

Signed-off-by: Richard Zou <zou3519@gmail.com>

.buildkite/scripts/hardware_ci/run-xpu-test.sh

@@ -42,6 +42,7 @@ docker run \
     python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager --max-model-len 8192
     python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2
     python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
+    python3 examples/basic/offline_inference/generate.py --model OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc --block-size 64 --enforce-eager --max-model-len 8192
     cd tests
     pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
     pytest -v -s v1/engine

.pre-commit-config.yaml

@@ -39,7 +39,7 @@ repos:
   rev: 0.11.1
   hooks:
     - id: pip-compile
-      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu129, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
+      args: [requirements/test.in, -c, requirements/common.txt, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu129, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
       files: ^requirements/test\.(in|txt)$
     - id: pip-compile
       alias: pip-compile-rocm

CMakeLists.txt

@@ -340,8 +340,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
 
   list(APPEND VLLM_EXT_SRC
     "csrc/quantization/awq/gemm_kernels.cu"
-    "csrc/quantization/fp4/nvfp4_quant_entry.cu"
-    "csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
     "csrc/cutlass_extensions/common.cpp")
 
   set_gencode_flags_for_srcs(
@@ -489,59 +487,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                    " in CUDA target architectures")
   endif()
 
-  # The nvfp4_scaled_mm_sm120 kernels for Blackwell SM12x require
-  # CUDA 12.8 or later
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(FP4_ARCHS "12.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(FP4_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
-    set(SRCS
-      "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")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${FP4_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
-    message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
-  else()
-    message(STATUS "Not building NVFP4 as no compatible archs were found.")
-    # clear FP4_ARCHS
-    set(FP4_ARCHS)
-  endif()
-
-  # FP4 Archs and flags
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(FP4_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(FP4_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
-    set(SRCS
-      "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")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${FP4_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
-    message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
-  else()
-    message(STATUS "Not building NVFP4 as no compatible archs were found.")
-    # clear FP4_ARCHS
-    set(FP4_ARCHS)
-  endif()
-
   # CUTLASS MLA Archs and flags
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
     cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
@@ -681,34 +626,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     endif()
   endif()
 
-  # Only build W4A8 kernels if we are building for something compatible with sm90a
-  cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS)
-    set(SRCS
-       "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(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${W4A8_ARCHS}")
-
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
-
-    message(STATUS "Building W4A8 kernels for archs: ${W4A8_ARCHS}")
-  else()
-    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
-        AND W4A8_ARCHS)
-      message(STATUS "Not building W4A8 kernels as CUDA Compiler version is "
-                     "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
-                     "later if you intend on running w4a16 quantized models on "
-                     "Hopper.")
-    else()
-      message(STATUS "Not building W4A8 kernels as no compatible archs "
-                     "found in CUDA target architectures")
-    endif()
-  endif()
 
   # Hadacore kernels
   cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
@@ -760,7 +677,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   set(VLLM_STABLE_EXT_SRC
     "csrc/libtorch_stable/torch_bindings.cpp"
     "csrc/cutlass_extensions/common.cpp"
-    "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu")
+    "csrc/cuda_utils_kernels.cu"
+    "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")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_STABLE_EXT_SRC
@@ -978,6 +898,96 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     endif()
   endif()
 
+  #
+  # FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch)
+  #
+
+  # The nvfp4_scaled_mm_sm120 kernels for Blackwell SM12x require
+  # CUDA 12.8 or later
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(FP4_ARCHS "12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(FP4_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
+    set(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")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${FP4_ARCHS}")
+    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
+    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
+    message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
+  else()
+    message(STATUS "Not building NVFP4 as no compatible archs were found.")
+    # clear FP4_ARCHS
+    set(FP4_ARCHS)
+  endif()
+
+  # FP4 Archs and flags
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(FP4_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(FP4_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
+    set(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")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${FP4_ARCHS}")
+    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
+    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
+    message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
+  else()
+    message(STATUS "Not building NVFP4 as no compatible archs were found.")
+    # clear FP4_ARCHS
+    set(FP4_ARCHS)
+  endif()
+
+  #
+  # W4A8 kernels (moved from _C to _C_stable_libtorch)
+  #
+
+  # Only build W4A8 kernels if we are building for something compatible with sm90a
+  cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS)
+    set(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"
+       )
+
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${W4A8_ARCHS}")
+
+    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+
+    message(STATUS "Building W4A8 kernels for archs: ${W4A8_ARCHS}")
+  else()
+    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
+        AND W4A8_ARCHS)
+      message(STATUS "Not building W4A8 kernels as CUDA Compiler version is "
+                     "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
+                     "later if you intend on running w4a16 quantized models on "
+                     "Hopper.")
+    else()
+      message(STATUS "Not building W4A8 kernels as no compatible archs "
+                     "found in CUDA target architectures")
+    endif()
+  endif()
+
   message(STATUS "Enabling C_stable extension.")
   define_extension_target(
     _C_stable_libtorch

csrc/cuda_vec_utils.cuh

@@ -3,8 +3,8 @@
 
 #pragma once
 
-#include <c10/util/BFloat16.h>
-#include <c10/util/Half.h>
+#include <torch/headeronly/util/BFloat16.h>
+#include <torch/headeronly/util/Half.h>
 #include <cassert>
 
 #ifdef USE_ROCM

csrc/cutlass_extensions/cute_utils.cuh

@@ -1,7 +1,6 @@
 #pragma once
 
 #include <cute/tensor.hpp>
-#include <torch/all.h>
 namespace cute {
 
 ////////////////////////////////////////////////////////////////////

csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp

@@ -189,9 +189,9 @@ struct Sm90RowOrScalarBroadcastArray {
       }
 
       auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
-      Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
-      Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
-      Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
+      cute::Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
+      cute::Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
+      cute::Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
 
       for (int i = 0; i < size(tGS_gRow_flt); ++i) {
         if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) {
@@ -211,8 +211,8 @@ struct Sm90RowOrScalarBroadcastArray {
     begin_loop(int epi_m, int epi_n) {
       if (epi_m == 0) { // Assumes M-major subtile loop
         if (!params.row_broadcast) return; // Do not issue LDS when row is scalar 
-        Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
-        Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
+        cute::Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
+        cute::Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
         copy(tSR_sRow_flt, tSR_rRow_flt);
       }
     }
@@ -241,9 +241,9 @@ struct Sm90RowOrScalarBroadcastArray {
     auto [m, n, k, l] = args.tile_coord_mnkl;
     using ThreadCount = decltype(size(args.tiled_copy));
 
-    Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row_array[l]), make_shape(M,N,1), params.dRow);
-    Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n));          // (CTA_M, CTA_N)
-    Tensor sRow = make_tensor(make_smem_ptr(smem), 
+    cute::Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row_array[l]), make_shape(M,N,1), params.dRow);
+    cute::Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n));          // (CTA_M, CTA_N)
+    cute::Tensor sRow = make_tensor(make_smem_ptr(smem), 
         make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{}));  // (CTA_M, CTA_N)
     //// G2S: Gmem to Smem
     auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},
@@ -251,16 +251,16 @@ struct Sm90RowOrScalarBroadcastArray {
                                             Stride<_0,          _1>>{}, 
                                      Layout<_1>{});   
     auto thr_g2s = tiled_g2s.get_slice(args.thread_idx);
-    Tensor tGS_gRow = thr_g2s.partition_S(gRow);
-    Tensor tGS_sRow = thr_g2s.partition_D(sRow);
+    cute::Tensor tGS_gRow = thr_g2s.partition_S(gRow);
+    cute::Tensor tGS_sRow = thr_g2s.partition_D(sRow);
 
     //// G2S: Coord 
     auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})));
-    Tensor tGS_cRow = thr_g2s.partition_S(cRow);
+    cute::Tensor tGS_cRow = thr_g2s.partition_S(cRow);
 
     //// S2R: Smem to Reg
-    Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
-    Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow));                                           // (CPY,CPY_M,CPY_N)
+    cute::Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
+    cute::Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow));                                           // (CPY,CPY_M,CPY_N)
 
     return ConsumerStoreCallbacks<decltype(tGS_gRow), decltype(tGS_sRow), decltype(tGS_cRow), decltype(tiled_g2s), decltype(tSR_sRow), decltype(tSR_rRow), decltype(args.tCcD), decltype(args.residue_cD), ThreadCount>(
       tGS_gRow, 
@@ -389,7 +389,7 @@ struct Sm90ColOrScalarBroadcastArray {
 
     CUTLASS_DEVICE void
     begin() {
-      Tensor pred = make_tensor<bool>(shape(tCgCol));
+      cute::Tensor pred = make_tensor<bool>(shape(tCgCol));
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < size(pred); ++i) {
         pred(i) = get<0>(tCcCol(i)) < m;
@@ -409,7 +409,7 @@ struct Sm90ColOrScalarBroadcastArray {
     CUTLASS_DEVICE Array<Element, FragmentSize>
     visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
       Array<Element, FragmentSize> frg_col;
-      Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
+      cute::Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
 
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < FragmentSize; ++i) {
@@ -431,16 +431,16 @@ struct Sm90ColOrScalarBroadcastArray {
     auto [M, N, K, L] = args.problem_shape_mnkl;
     auto [m, n, k, l] = args.tile_coord_mnkl;
 
-    Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col_array[l]), make_shape(M,N,1), params.dCol);
-    Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>(                         // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    cute::Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col_array[l]), make_shape(M,N,1), params.dCol);
+    cute::Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>(                         // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
       mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
-    Tensor tCrCol = make_tensor_like(tCgCol);                                          // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    cute::Tensor tCrCol = make_tensor_like(tCgCol);                                          // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
 
     // Generate an identity tensor matching the shape of the global tensor and 
     //  partition the same way, this will be used to generate the predicate
     //  tensor for loading
-    Tensor cCol = make_identity_tensor(mCol.shape());
-    Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>(                         // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    cute::Tensor cCol = make_identity_tensor(mCol.shape());
+    cute::Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>(                         // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
       cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
 
     return ConsumerStoreCallbacks(

csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp

@@ -186,9 +186,9 @@ struct Sm90RowOrScalarBroadcast {
       }
 
       auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
-      Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
-      Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
-      Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
+      cute::Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
+      cute::Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
+      cute::Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
 
       for (int i = 0; i < size(tGS_gRow_flt); ++i) {
         if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) {
@@ -208,8 +208,8 @@ struct Sm90RowOrScalarBroadcast {
     begin_loop(int epi_m, int epi_n) {
       if (epi_m == 0) { // Assumes M-major subtile loop
         if (!params.row_broadcast) return; // Do not issue LDS when row is scalar 
-        Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
-        Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
+        cute::Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
+        cute::Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
         copy(tSR_sRow_flt, tSR_rRow_flt);
       }
     }
@@ -238,9 +238,9 @@ struct Sm90RowOrScalarBroadcast {
     auto [m, n, k, l] = args.tile_coord_mnkl;
     using ThreadCount = decltype(size(args.tiled_copy));
 
-    Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow);
-    Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n));          // (CTA_M, CTA_N)
-    Tensor sRow = make_tensor(make_smem_ptr(smem), 
+    cute::Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow);
+    cute::Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n));          // (CTA_M, CTA_N)
+    cute::Tensor sRow = make_tensor(make_smem_ptr(smem), 
         make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{}));  // (CTA_M, CTA_N)
     //// G2S: Gmem to Smem
     auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},
@@ -248,16 +248,16 @@ struct Sm90RowOrScalarBroadcast {
                                             Stride<_0,          _1>>{}, 
                                      Layout<_1>{});   
     auto thr_g2s = tiled_g2s.get_slice(args.thread_idx);
-    Tensor tGS_gRow = thr_g2s.partition_S(gRow);
-    Tensor tGS_sRow = thr_g2s.partition_D(sRow);
+    cute::Tensor tGS_gRow = thr_g2s.partition_S(gRow);
+    cute::Tensor tGS_sRow = thr_g2s.partition_D(sRow);
 
     //// G2S: Coord 
     auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})));
-    Tensor tGS_cRow = thr_g2s.partition_S(cRow);
+    cute::Tensor tGS_cRow = thr_g2s.partition_S(cRow);
 
     //// S2R: Smem to Reg
-    Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
-    Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow));                                           // (CPY,CPY_M,CPY_N)
+    cute::Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
+    cute::Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow));                                           // (CPY,CPY_M,CPY_N)
 
     return ConsumerStoreCallbacks<decltype(tGS_gRow), decltype(tGS_sRow), decltype(tGS_cRow), decltype(tiled_g2s), decltype(tSR_sRow), decltype(tSR_rRow), decltype(args.tCcD), decltype(args.residue_cD), ThreadCount>(
       tGS_gRow, 
@@ -382,7 +382,7 @@ struct Sm90ColOrScalarBroadcast {
 
     CUTLASS_DEVICE void
     begin() {
-      Tensor pred = make_tensor<bool>(shape(tCgCol));
+      cute::Tensor pred = make_tensor<bool>(shape(tCgCol));
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < size(pred); ++i) {
         pred(i) = get<0>(tCcCol(i)) < m;
@@ -402,7 +402,7 @@ struct Sm90ColOrScalarBroadcast {
     CUTLASS_DEVICE Array<Element, FragmentSize>
     visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
       Array<Element, FragmentSize> frg_col;
-      Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
+      cute::Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
 
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < FragmentSize; ++i) {
@@ -422,16 +422,16 @@ struct Sm90ColOrScalarBroadcast {
   get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) {
 
     auto [M, N, K, L] = args.problem_shape_mnkl;
-    Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol);
-    Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>(                         // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    cute::Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol);
+    cute::Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>(                         // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
       mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
-    Tensor tCrCol = make_tensor_like(tCgCol);                                          // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    cute::Tensor tCrCol = make_tensor_like(tCgCol);                                          // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
 
     // Generate an identity tensor matching the shape of the global tensor and 
     //  partition the same way, this will be used to generate the predicate
     //  tensor for loading
-    Tensor cCol = make_identity_tensor(mCol.shape());
-    Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>(                         // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    cute::Tensor cCol = make_identity_tensor(mCol.shape());
+    cute::Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>(                         // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
       cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
 
     return ConsumerStoreCallbacks(

csrc/cutlass_extensions/torch_utils.hpp

@@ -1,6 +1,21 @@
 #pragma once
 
-#include <torch/all.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>
+  #include <torch/headeronly/util/shim_utils.h>  // for STD_TORCH_CHECK
+using TorchTensor = torch::stable::Tensor;
+  #define TORCH_UTILS_CHECK STD_TORCH_CHECK
+#else
+  #include <torch/all.h>
+using TorchTensor = torch::Tensor;
+  #define TORCH_UTILS_CHECK TORCH_CHECK
+#endif
 
 #include "cute/layout.hpp"
 #include "cutlass/layout/matrix.h"
@@ -55,35 +70,35 @@ CUTE_HOST_DEVICE constexpr auto make_shape_from_idx(F&& f) {
 // If `tensor.dim() < rank(Stride{})`, the shape is padded with 1s and the extra
 // strides are set to be 0 or 1.
 template <typename Stride>
-static inline auto make_cute_layout(torch::Tensor const& tensor,
+static inline auto make_cute_layout(TorchTensor const& tensor,
                                     std::string_view name = "tensor") {
-  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)>;
+  TORCH_UTILS_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_CHECK(StrideEle::value == tensor.stride(idx), "Expected ",
-                        name, ".stride(", idx, ") to be ", StrideEle::value);
-            return StrideEle{};
-          } else {
-            if (tensor.size(idx) == 1) {
-              // use 0 stride for dim with size 1, this is easier for
-              // cute/cutlass to optimize (helps the TMA code flatten dims)
-              return StrideEle{0};
-            } else {
-              return tensor.stride(idx);
-            }
-          }
+    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);
+        return StrideEle{};
+      } else {
+        if (tensor.size(idx) == 1) {
+          // use 0 stride for dim with size 1, this is easier for
+          // cute/cutlass to optimize (helps the TMA code flatten dims)
+          return StrideEle{0};
         } else {
-          // Extra strides are assumed to be 0 or 1
-          if constexpr (cute::is_static_v<StrideEle>) {
-            static_assert(StrideEle::value == 0 || StrideEle::value == 1);
-          }
-          return StrideEle{};
+          return tensor.stride(idx);
         }
-      });
+      }
+    } else {
+      // Extra strides are assumed to be 0 or 1
+      if constexpr (cute::is_static_v<StrideEle>) {
+        static_assert(StrideEle::value == 0 || StrideEle::value == 1);
+      }
+      return StrideEle{};
+    }
+  });
 
   auto shape = cute::make_shape_from_idx<rank(Stride{})>([&](auto const& idx) {
     if (idx < tensor.dim())
@@ -97,7 +112,7 @@ static inline auto make_cute_layout(torch::Tensor const& tensor,
 
 template <typename Stride>
 static inline auto maybe_make_cute_layout(
-    std::optional<torch::Tensor> const& tensor,
+    std::optional<TorchTensor> const& tensor,
     std::string_view name = "tensor") {
   using Layout = decltype(make_cute_layout<Stride>(*tensor));
 
@@ -121,12 +136,12 @@ template <typename T>
 using equivalent_cutlass_type_t = typename equivalent_cutlass_type<T>::type;
 
 template <>
-struct equivalent_cutlass_type<c10::Half> {
+struct equivalent_cutlass_type<torch::headeronly::Half> {
   using type = cutlass::half_t;
 };
 
 template <>
-struct equivalent_cutlass_type<c10::BFloat16> {
+struct equivalent_cutlass_type<torch::headeronly::BFloat16> {
   using type = cutlass::bfloat16_t;
 };
 
@@ -134,8 +149,8 @@ struct equivalent_cutlass_type<c10::BFloat16> {
 // equivalent_scalar_t (basically inverse of equivalent_cutlass_type)
 //
 
-// Return a `c10::CppTypeToScalarType<T>` compatible type, i.e. get the C++ from
-// c10 that is equivalent to T, e.g.: `cutlass::half_t -> c10::Half`
+// Return a `torch::headeronly::CppTypeToScalarType<T>` compatible type, i.e.
+// get the C++ type equivalent to T, e.g.: `cutlass::half_t -> Half`
 template <typename T>
 struct equivalent_scalar_type {
   using type = T;
@@ -146,15 +161,15 @@ using equivalent_scalar_type_t = typename equivalent_scalar_type<T>::type;
 
 template <>
 struct equivalent_scalar_type<cutlass::half_t> {
-  using type = c10::Half;
+  using type = torch::headeronly::Half;
 };
 
 template <>
 struct equivalent_scalar_type<cutlass::bfloat16_t> {
-  using type = c10::BFloat16;
+  using type = torch::headeronly::BFloat16;
 };
 
-// get equivalent c10::ScalarType tag from compile time type
+// get equivalent torch::headeronly::ScalarType tag from compile time type
 template <typename T>
-static inline constexpr c10::ScalarType equivalent_scalar_type_v =
-    c10::CppTypeToScalarType<equivalent_scalar_type_t<T>>::value;
+static inline constexpr torch::headeronly::ScalarType equivalent_scalar_type_v =
+    torch::headeronly::CppTypeToScalarType<equivalent_scalar_type_t<T>>::value;

csrc/libtorch_stable/dispatch_utils.h

@@ -49,6 +49,15 @@
   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__))
+
 // Boolean dispatch
 #define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \
   if (expr) {                                            \

csrc/libtorch_stable/ops.h

@@ -84,4 +84,54 @@ void get_cutlass_batched_moe_mm_data(
     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);
+
 #endif

csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh

@@ -2,10 +2,9 @@
 #pragma once
 
 #include <cuda.h>
-#include <torch/all.h>
-#include <c10/cuda/CUDAStream.h>
+#include <torch/csrc/stable/tensor.h>
+#include "libtorch_stable/torch_utils.h"
 
-#include "core/scalar_type.hpp"
 #include "cutlass/bfloat16.h"
 #include "cutlass/float8.h"
 
@@ -41,7 +40,7 @@ __global__ void get_group_gemm_starts(
 }
 
 #define __CALL_GET_STARTS_KERNEL(TENSOR_C_TYPE, C_TYPE)                  \
-  else if (out_tensors.dtype() == TENSOR_C_TYPE) {                       \
+  else if (out_tensors.scalar_type() == TENSOR_C_TYPE) {                 \
     get_group_gemm_starts<cutlass::float_e4m3_t, int32_t, C_TYPE, float, \
                           cutlass::Array<cutlass::float_e4m3_t, 8>>      \
         <<<1, num_experts, 0, stream>>>(                                 \
@@ -66,23 +65,34 @@ __global__ void get_group_gemm_starts(
 namespace {
 
 void run_get_group_gemm_starts(
-    torch::Tensor const& expert_offsets, torch::Tensor& a_ptrs,
-    torch::Tensor& b_ptrs, torch::Tensor& out_ptrs,
-    torch::Tensor& a_scales_ptrs, torch::Tensor& b_scales_ptrs,
-    torch::Tensor& b_group_scales_ptrs, torch::Tensor const& a_tensors,
-    torch::Tensor const& b_tensors, torch::Tensor& out_tensors,
-    torch::Tensor const& a_scales, torch::Tensor const& b_scales,
-    torch::Tensor const& b_group_scales, const int64_t b_group_size) {
-  TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn);
-  TORCH_CHECK(b_tensors.dtype() == torch::kInt32);  // int4 8x packed into int32
-  TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
-  TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
-  TORCH_CHECK(b_group_scales.dtype() ==
-              torch::kFloat8_e4m3fn);  // the underlying torch type is e4m3
-  TORCH_CHECK(out_tensors.dtype() ==
-              torch::kBFloat16);  // only support bf16 for now
+    torch::stable::Tensor const& expert_offsets, torch::stable::Tensor& a_ptrs,
+    torch::stable::Tensor& b_ptrs, torch::stable::Tensor& out_ptrs,
+    torch::stable::Tensor& a_scales_ptrs, torch::stable::Tensor& b_scales_ptrs,
+    torch::stable::Tensor& b_group_scales_ptrs,
+    torch::stable::Tensor const& a_tensors,
+    torch::stable::Tensor const& b_tensors, torch::stable::Tensor& out_tensors,
+    torch::stable::Tensor const& a_scales,
+    torch::stable::Tensor const& b_scales,
+    torch::stable::Tensor const& b_group_scales, const int64_t b_group_size) {
+  STD_TORCH_CHECK(a_tensors.scalar_type() ==
+                  torch::headeronly::ScalarType::Float8_e4m3fn);
+  STD_TORCH_CHECK(
+      b_tensors.scalar_type() ==
+      torch::headeronly::ScalarType::Int);  // int4 8x packed into int32
+  STD_TORCH_CHECK(a_scales.scalar_type() ==
+                  torch::headeronly::ScalarType::Float);
+  STD_TORCH_CHECK(b_scales.scalar_type() ==
+                  torch::headeronly::ScalarType::Float);
+  STD_TORCH_CHECK(
+      b_group_scales.scalar_type() ==
+      torch::headeronly::ScalarType::Float8_e4m3fn);  // the underlying torch
+                                                      // type is e4m3
+  STD_TORCH_CHECK(
+      out_tensors.scalar_type() ==
+      torch::headeronly::ScalarType::BFloat16);  // only support bf16 for now
   // expect int64_t to avoid overflow during offset calculations
-  TORCH_CHECK(expert_offsets.dtype() == torch::kInt64);
+  STD_TORCH_CHECK(expert_offsets.scalar_type() ==
+                  torch::headeronly::ScalarType::Long);
 
   int num_experts = static_cast<int>(expert_offsets.size(0));
   // logical k, n
@@ -90,15 +100,16 @@ void run_get_group_gemm_starts(
   int64_t k = a_tensors.size(1);
   int64_t scale_k = cutlass::ceil_div(k, b_group_size);
 
-  auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
+  auto stream = get_current_cuda_stream(a_tensors.get_device_index());
 
   if (false) {
   }
-  __CALL_GET_STARTS_KERNEL(torch::kBFloat16, cutlass::bfloat16_t)
-  __CALL_GET_STARTS_KERNEL(torch::kFloat16, half)
+  __CALL_GET_STARTS_KERNEL(torch::headeronly::ScalarType::BFloat16,
+                           cutlass::bfloat16_t)
+  __CALL_GET_STARTS_KERNEL(torch::headeronly::ScalarType::Half, half)
   else {
-    TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
+    STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
   }
 }
 
-}  // namespace
+}  // namespace

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu

@@ -14,13 +14,12 @@
 #include "cutlass/util/mixed_dtype_utils.hpp"
 
 // vllm includes
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <torch/all.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include "libtorch_stable/torch_utils.h"
 #include "cutlass_extensions/torch_utils.hpp"
 #include "cutlass_extensions/common.hpp"
 
-#include "core/registration.h"
 #include "get_group_starts.cuh"
 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 #include "w4a8_utils.cuh"
@@ -168,31 +167,40 @@ struct W4A8GroupedGemmKernel {
   static_assert(sizeof(LayoutB_Reordered) % sizeof(int32_t) == 0,
                 "LayoutB_Reordered size must be divisible by 4 bytes");
 
-  static void grouped_mm(
-      torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
-      const torch::Tensor& b_tensors, const torch::Tensor& a_scales,
-      const torch::Tensor& b_scales, const torch::Tensor& b_group_scales,
-      const int64_t b_group_size, const torch::Tensor& expert_offsets,
-      const torch::Tensor& problem_sizes_torch, const torch::Tensor& a_strides,
-      const torch::Tensor& b_strides, const torch::Tensor& c_strides,
-      const torch::Tensor& group_scale_strides) {
+  static void grouped_mm(torch::stable::Tensor& out_tensors,
+                         const torch::stable::Tensor& a_tensors,
+                         const torch::stable::Tensor& b_tensors,
+                         const torch::stable::Tensor& a_scales,
+                         const torch::stable::Tensor& b_scales,
+                         const torch::stable::Tensor& b_group_scales,
+                         const int64_t b_group_size,
+                         const torch::stable::Tensor& expert_offsets,
+                         const torch::stable::Tensor& problem_sizes_torch,
+                         const torch::stable::Tensor& a_strides,
+                         const torch::stable::Tensor& b_strides,
+                         const torch::stable::Tensor& c_strides,
+                         const torch::stable::Tensor& group_scale_strides) {
     auto device = a_tensors.device();
     auto device_id = device.index();
-    const at::cuda::OptionalCUDAGuard device_guard(device);
-    auto stream = at::cuda::getCurrentCUDAStream(device_id);
+    const torch::stable::accelerator::DeviceGuard device_guard(device_id);
+    auto stream = get_current_cuda_stream(device_id);
 
     int num_experts = static_cast<int>(expert_offsets.size(0));
     int n = static_cast<int>(b_tensors.size(1));
     int k = static_cast<int>(b_tensors.size(2)) * PackFactor;
 
-    auto options_int =
-        torch::TensorOptions().dtype(torch::kInt64).device(device);
-    torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
-    torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
-    torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
-    torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
-    torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
-    torch::Tensor b_group_scales_ptrs = torch::empty(num_experts, options_int);
+    torch::stable::Tensor a_ptrs = torch::stable::empty(
+        num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
+    torch::stable::Tensor b_ptrs = torch::stable::empty(
+        num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
+    torch::stable::Tensor out_ptrs = torch::stable::empty(
+        num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
+    torch::stable::Tensor a_scales_ptrs = torch::stable::empty(
+        num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
+    torch::stable::Tensor b_scales_ptrs = torch::stable::empty(
+        num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
+    torch::stable::Tensor b_group_scales_ptrs = torch::stable::empty(
+        num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
 
     // get the correct offsets to pass to gemm
     run_get_group_gemm_starts(expert_offsets, a_ptrs, b_ptrs, out_ptrs,
@@ -247,9 +255,9 @@ struct W4A8GroupedGemmKernel {
 
     // Allocate workspace
     size_t workspace_size = GemmShuffled::get_workspace_size(arguments);
-    torch::Tensor workspace =
-        torch::empty(workspace_size,
-                     torch::TensorOptions().dtype(torch::kU8).device(device));
+    torch::stable::Tensor workspace = torch::stable::empty(
+        workspace_size, torch::headeronly::ScalarType::Byte, std::nullopt,
+        device);
 
     // Run GEMM
     GemmShuffled gemm;
@@ -294,14 +302,20 @@ using Kernel_256x128_2x1x1_Coop =
 using Kernel_128x256_2x1x1_Coop =
     W4A8GroupedGemmKernel<Shape<_128, _256>, Shape<_2, _1, _1>, Coop, CoopEpi>;
 
-void mm_dispatch(
-    torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
-    const torch::Tensor& b_tensors, const torch::Tensor& a_scales,
-    const torch::Tensor& b_scales, const torch::Tensor& b_group_scales,
-    const int64_t b_group_size, const torch::Tensor& expert_offsets,
-    const torch::Tensor& problem_sizes, const torch::Tensor& a_strides,
-    const torch::Tensor& b_strides, const torch::Tensor& c_strides,
-    const torch::Tensor& group_scale_strides, const std::string& schedule) {
+void mm_dispatch(torch::stable::Tensor& out_tensors,
+                 const torch::stable::Tensor& a_tensors,
+                 const torch::stable::Tensor& b_tensors,
+                 const torch::stable::Tensor& a_scales,
+                 const torch::stable::Tensor& b_scales,
+                 const torch::stable::Tensor& b_group_scales,
+                 const int64_t b_group_size,
+                 const torch::stable::Tensor& expert_offsets,
+                 const torch::stable::Tensor& problem_sizes,
+                 const torch::stable::Tensor& a_strides,
+                 const torch::stable::Tensor& b_strides,
+                 const torch::stable::Tensor& c_strides,
+                 const torch::stable::Tensor& group_scale_strides,
+                 const std::string& schedule) {
   if (schedule == "Kernel_128x16_1x1x1_Coop") {
     Kernel_128x16_1x1x1_Coop::grouped_mm(
         out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
@@ -358,18 +372,23 @@ void mm_dispatch(
         b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
         c_strides, group_scale_strides);
   } else {
-    TORCH_CHECK(false,
-                "cutlass_w4a8_moe_mm: unknown schedule string: ", schedule);
+    STD_TORCH_CHECK(false,
+                    "cutlass_w4a8_moe_mm: unknown schedule string: ", schedule);
   }
 }
 
-void mm(torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
-        const torch::Tensor& b_tensors, const torch::Tensor& a_scales,
-        const torch::Tensor& b_scales, const torch::Tensor& b_group_scales,
-        const int64_t b_group_size, const torch::Tensor& expert_offsets,
-        const torch::Tensor& problem_sizes, const torch::Tensor& a_strides,
-        const torch::Tensor& b_strides, const torch::Tensor& c_strides,
-        const torch::Tensor& group_scale_strides,
+void mm(torch::stable::Tensor& out_tensors,
+        const torch::stable::Tensor& a_tensors,
+        const torch::stable::Tensor& b_tensors,
+        const torch::stable::Tensor& a_scales,
+        const torch::stable::Tensor& b_scales,
+        const torch::stable::Tensor& b_group_scales, const int64_t b_group_size,
+        const torch::stable::Tensor& expert_offsets,
+        const torch::stable::Tensor& problem_sizes,
+        const torch::stable::Tensor& a_strides,
+        const torch::stable::Tensor& b_strides,
+        const torch::stable::Tensor& c_strides,
+        const torch::stable::Tensor& group_scale_strides,
         std::optional<std::string> maybe_schedule) {
   // user has specified a schedule
   if (maybe_schedule) {
@@ -406,26 +425,27 @@ void mm(torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
               a_strides, b_strides, c_strides, group_scale_strides, schedule);
 }
 
-std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b(
-    torch::Tensor const& b_tensors) {
-  TORCH_CHECK(b_tensors.dtype() == torch::kInt32);
-  TORCH_CHECK(b_tensors.dim() == 3);  // (experts, n, k)
-  TORCH_CHECK(b_tensors.is_contiguous());
-  TORCH_CHECK(b_tensors.is_cuda());
+std::tuple<torch::stable::Tensor, torch::stable::Tensor>
+encode_and_reorder_int4b(torch::stable::Tensor const& b_tensors) {
+  STD_TORCH_CHECK(b_tensors.scalar_type() ==
+                  torch::headeronly::ScalarType::Int);
+  STD_TORCH_CHECK(b_tensors.dim() == 3);  // (experts, n, k)
+  STD_TORCH_CHECK(b_tensors.is_contiguous());
+  STD_TORCH_CHECK(b_tensors.is_cuda());
 
   int n = static_cast<int>(b_tensors.size(1));
   int k = static_cast<int>(b_tensors.size(2)) * PackFactor;  // logical k
 
   // CUTLASS reorder_tensor requires k % 256 == 0 and n % 16 == 0.
   // These misalignments cause silent OOB unless run under Compute Sanitizer.
-  TORCH_CHECK(k % 256 == 0, "logical k must be divisible by 256");
-  TORCH_CHECK(n % 16 == 0, "n must be divisible by 16");
+  STD_TORCH_CHECK(k % 256 == 0, "logical k must be divisible by 256");
+  STD_TORCH_CHECK(n % 16 == 0, "n must be divisible by 16");
 
   // we will store the layout to an int32 tensor;
   // this is the number of elements we need per layout
   constexpr size_t layout_width = sizeof(LayoutB_Reordered) / sizeof(int32_t);
 
-  torch::Tensor b_tensors_packed = torch::empty_like(b_tensors);
+  torch::stable::Tensor b_tensors_packed = torch::stable::empty_like(b_tensors);
   int num_experts = static_cast<int>(b_tensors.size(0));
 
   auto b_ptr = static_cast<QuantType const*>(b_tensors.const_data_ptr());
@@ -435,7 +455,7 @@ std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b(
   size_t num_int4_elems = 1ull * num_experts * n * k;
   bool ok = vllm::cutlass_w4a8_utils::unified_encode_int4b(b_ptr, b_packed_ptr,
                                                            num_int4_elems);
-  TORCH_CHECK(ok, "unified_encode_int4b failed");
+  STD_TORCH_CHECK(ok, "unified_encode_int4b failed");
 
   // construct the layout once; assumes each expert has the same layout
   using LayoutType = LayoutB_Reordered;
@@ -456,28 +476,28 @@ std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b(
   }
 
   // save the packed layout to torch tensor so we can re-use it
-  auto cpu_opts =
-      torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU);
-  torch::Tensor layout_cpu =
-      torch::empty({num_experts, layout_width}, cpu_opts);
+  torch::stable::Tensor layout_cpu = torch::stable::empty(
+      {num_experts, layout_width}, torch::headeronly::ScalarType::Int,
+      std::nullopt, torch::stable::Device(torch::stable::DeviceType::CPU));
 
-  int32_t* layout_data = layout_cpu.data_ptr<int32_t>();
+  int32_t* layout_data = layout_cpu.mutable_data_ptr<int32_t>();
   for (int i = 0; i < num_experts; ++i) {
     std::memcpy(layout_data + i * layout_width,  // dst (int32*)
                 &layout_B_reordered,             // src (LayoutType*)
                 sizeof(LayoutType));             // number of bytes
   }
 
-  torch::Tensor packed_layout =
-      layout_cpu.to(b_tensors.device(), /*non_blocking=*/false);
+  torch::stable::Tensor packed_layout =
+      torch::stable::to(layout_cpu, b_tensors.device(),
+                        /*non_blocking=*/false);
 
   return {b_tensors_packed, packed_layout};
 }
 
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("cutlass_w4a8_moe_mm", &mm);
-  m.impl("cutlass_encode_and_reorder_int4b_grouped", &encode_and_reorder_int4b);
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("cutlass_w4a8_moe_mm", TORCH_BOX(&mm));
+  m.impl("cutlass_encode_and_reorder_int4b_grouped",
+         TORCH_BOX(&encode_and_reorder_int4b));
 }
 
 }  // namespace vllm::cutlass_w4a8_moe
-/////////////////////////////////////////////////////////////////////////////////////////////////

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu

@@ -3,14 +3,12 @@
 //   https://github.com/NVIDIA/cutlass/blob/main/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu
 //
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <torch/all.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include "libtorch_stable/torch_utils.h"
 #include "cutlass_extensions/torch_utils.hpp"
 #include "w4a8_utils.cuh"
 
-#include "core/registration.h"
-
 #include "cutlass/cutlass.h"
 #include <limits>
 
@@ -161,31 +159,31 @@ struct W4A8GemmKernel {
   using StrideD = typename GemmKernelShuffled::StrideD;
   using StrideS = typename CollectiveMainloopShuffled::StrideScale;
 
-  static torch::Tensor mm(torch::Tensor const& A,
-                          torch::Tensor const& B,             // already packed
-                          torch::Tensor const& group_scales,  // already packed
-                          int64_t group_size,
-                          torch::Tensor const& channel_scales,
-                          torch::Tensor const& token_scales,
-                          std::optional<at::ScalarType> const& maybe_out_type) {
+  static torch::stable::Tensor mm(
+      torch::stable::Tensor const& A,
+      torch::stable::Tensor const& B,             // already packed
+      torch::stable::Tensor const& group_scales,  // already packed
+      int64_t group_size, torch::stable::Tensor const& channel_scales,
+      torch::stable::Tensor const& token_scales,
+      std::optional<torch::headeronly::ScalarType> const& maybe_out_type) {
     // TODO: param validation
     int m = A.size(0);
     int k = A.size(1);
     int n = B.size(1);
 
     // safely cast group_size to int
-    TORCH_CHECK(group_size > 0 && group_size <= std::numeric_limits<int>::max(),
-                "group_size out of supported range for int: ", group_size);
+    STD_TORCH_CHECK(
+        group_size > 0 && group_size <= std::numeric_limits<int>::max(),
+        "group_size out of supported range for int: ", group_size);
     int const group_size_int = static_cast<int>(group_size);
 
     // Allocate output
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
+    const torch::stable::accelerator::DeviceGuard device_guard(
+        A.get_device_index());
     auto device = A.device();
-    auto stream = at::cuda::getCurrentCUDAStream(device.index());
-    torch::Tensor D =
-        torch::empty({m, n}, torch::TensorOptions()
-                                 .dtype(equivalent_scalar_type_v<ElementD>)
-                                 .device(device));
+    auto stream = get_current_cuda_stream(device.index());
+    torch::stable::Tensor D = torch::stable::empty(
+        {m, n}, equivalent_scalar_type_v<ElementD>, std::nullopt, device);
     // prepare arg pointers
     auto A_ptr = static_cast<MmaType const*>(A.const_data_ptr());
     auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr());
@@ -237,9 +235,9 @@ struct W4A8GemmKernel {
 
     // Workspace
     size_t workspace_size = GemmShuffled::get_workspace_size(arguments);
-    torch::Tensor workspace =
-        torch::empty(workspace_size,
-                     torch::TensorOptions().dtype(torch::kU8).device(device));
+    torch::stable::Tensor workspace = torch::stable::empty(
+        workspace_size, torch::headeronly::ScalarType::Byte, std::nullopt,
+        device);
 
     // Run GEMM
     GemmShuffled gemm;
@@ -269,14 +267,14 @@ using Kernel_128x64_1x1x1 = W4A8GemmKernel<Shape<_128, _64>, Shape<_1, _1, _1>>;
 using Kernel_128x32_1x1x1 = W4A8GemmKernel<Shape<_128, _32>, Shape<_1, _1, _1>>;
 using Kernel_128x16_1x1x1 = W4A8GemmKernel<Shape<_128, _16>, Shape<_1, _1, _1>>;
 
-torch::Tensor mm_dispatch(torch::Tensor const& A,
-                          torch::Tensor const& B,             // already packed
-                          torch::Tensor const& group_scales,  // already packed
-                          int64_t group_size,
-                          torch::Tensor const& channel_scales,
-                          torch::Tensor const& token_scales,
-                          std::optional<at::ScalarType> const& maybe_out_type,
-                          const std::string& schedule) {
+torch::stable::Tensor mm_dispatch(
+    torch::stable::Tensor const& A,
+    torch::stable::Tensor const& B,             // already packed
+    torch::stable::Tensor const& group_scales,  // already packed
+    int64_t group_size, torch::stable::Tensor const& channel_scales,
+    torch::stable::Tensor const& token_scales,
+    std::optional<torch::headeronly::ScalarType> const& maybe_out_type,
+    const std::string& schedule) {
   if (schedule == "256x128_1x1x1") {
     return Kernel_256x128_1x1x1::mm(A, B, group_scales, group_size,
                                     channel_scales, token_scales,
@@ -318,17 +316,18 @@ torch::Tensor mm_dispatch(torch::Tensor const& A,
                                    channel_scales, token_scales,
                                    maybe_out_type);
   }
-  TORCH_CHECK(false, "Unknown W4A8 schedule: ", schedule);
+  STD_TORCH_CHECK(false, "Unknown W4A8 schedule: ", schedule);
   return {};
 }
 
-torch::Tensor mm(torch::Tensor const& A,
-                 torch::Tensor const& B,             // already packed
-                 torch::Tensor const& group_scales,  // already packed
-                 int64_t group_size, torch::Tensor const& channel_scales,
-                 torch::Tensor const& token_scales,
-                 std::optional<at::ScalarType> const& maybe_out_type,
-                 std::optional<std::string> maybe_schedule) {
+torch::stable::Tensor mm(
+    torch::stable::Tensor const& A,
+    torch::stable::Tensor const& B,             // already packed
+    torch::stable::Tensor const& group_scales,  // already packed
+    int64_t group_size, torch::stable::Tensor const& channel_scales,
+    torch::stable::Tensor const& token_scales,
+    std::optional<torch::headeronly::ScalarType> const& maybe_out_type,
+    std::optional<std::string> maybe_schedule) {
   // requested a specific schedule
   if (maybe_schedule) {
     return mm_dispatch(A, B, group_scales, group_size, channel_scales,
@@ -378,14 +377,15 @@ torch::Tensor mm(torch::Tensor const& A,
 // ----------------------------------------------------------------------------
 // Pre-processing utils
 // ----------------------------------------------------------------------------
-torch::Tensor pack_scale_fp8(torch::Tensor const& scales) {
-  TORCH_CHECK(scales.dtype() == torch::kFloat8_e4m3fn);
-  TORCH_CHECK(scales.is_contiguous());
-  TORCH_CHECK(scales.is_cuda());
+torch::stable::Tensor pack_scale_fp8(torch::stable::Tensor const& scales) {
+  STD_TORCH_CHECK(scales.scalar_type() ==
+                  torch::headeronly::ScalarType::Float8_e4m3fn);
+  STD_TORCH_CHECK(scales.is_contiguous());
+  STD_TORCH_CHECK(scales.is_cuda());
 
-  auto packed_scales = torch::empty(
-      {scales.numel() * ScalePackSize},
-      torch::TensorOptions().dtype(scales.dtype()).device(scales.device()));
+  auto packed_scales =
+      torch::stable::empty({scales.numel() * ScalePackSize},
+                           scales.scalar_type(), std::nullopt, scales.device());
   auto scales_ptr = static_cast<MmaType const*>(scales.const_data_ptr());
   auto packed_scales_ptr =
       static_cast<cutlass::Array<ElementScale, ScalePackSize>*>(
@@ -396,15 +396,16 @@ torch::Tensor pack_scale_fp8(torch::Tensor const& scales) {
   return packed_scales;
 }
 
-torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) {
-  TORCH_CHECK(B.dtype() == torch::kInt32);
-  TORCH_CHECK(B.dim() == 2);
+torch::stable::Tensor encode_and_reorder_int4b(torch::stable::Tensor const& B) {
+  STD_TORCH_CHECK(B.scalar_type() == torch::headeronly::ScalarType::Int);
+  STD_TORCH_CHECK(B.dim() == 2);
 
-  torch::Tensor B_packed = torch::empty_like(B);
+  torch::stable::Tensor B_packed = torch::stable::empty_like(B);
 
   int k = B.size(0) * PackFactor;  // logical k
   int n = B.size(1);
-  TORCH_CHECK((n * k) % 32 == 0, "need multiples of 32 int4s for 16B chunks");
+  STD_TORCH_CHECK((n * k) % 32 == 0,
+                  "need multiples of 32 int4s for 16B chunks");
 
   auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr());
   auto B_packed_ptr = static_cast<QuantType*>(B_packed.data_ptr());
@@ -415,16 +416,17 @@ torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) {
 
   bool ok = vllm::cutlass_w4a8_utils::unified_encode_int4b(B_ptr, B_packed_ptr,
                                                            n * k);
-  TORCH_CHECK(ok, "unified_encode_int4b failed");
+  STD_TORCH_CHECK(ok, "unified_encode_int4b failed");
   cutlass::reorder_tensor(B_packed_ptr, layout_B, layout_B_reordered);
 
   return B_packed;
 }
 
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("cutlass_w4a8_mm", &mm);
-  m.impl("cutlass_pack_scale_fp8", &pack_scale_fp8);
-  m.impl("cutlass_encode_and_reorder_int4b", &encode_and_reorder_int4b);
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("cutlass_w4a8_mm", TORCH_BOX(&mm));
+  m.impl("cutlass_pack_scale_fp8", TORCH_BOX(&pack_scale_fp8));
+  m.impl("cutlass_encode_and_reorder_int4b",
+         TORCH_BOX(&encode_and_reorder_int4b));
 }
 
-}  // namespace vllm::cutlass_w4a8
+}  // namespace vllm::cutlass_w4a8

csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu

@@ -14,16 +14,15 @@
  * limitations under the License.
  */
 
-#include <torch/all.h>
+#include <torch/csrc/stable/tensor.h>
+#include "libtorch_stable/torch_utils.h"
+#include "libtorch_stable/dispatch_utils.h"
+#include "cuda_vec_utils.cuh"
 
 #include <cuda_runtime_api.h>
 #include <cuda_runtime.h>
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-
 #include <cuda_fp8.h>
-#include "dispatch_utils.h"
 
 #include "cuda_utils.h"
 #include "launch_bounds_utils.h"
@@ -118,17 +117,19 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
 
 }  // namespace vllm
 
-void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output,  // [..., d]
-                                     torch::Tensor& output_sf,
-                                     torch::Tensor& input,  // [..., 2 * d]
-                                     torch::Tensor& input_sf) {
+void silu_and_mul_nvfp4_quant_sm1xxa(
+    torch::stable::Tensor& output,  // [..., d]
+    torch::stable::Tensor& output_sf,
+    torch::stable::Tensor& input,  // [..., 2 * d]
+    torch::stable::Tensor& input_sf) {
   int32_t m = input.size(0);
   int32_t n = input.size(1) / 2;
 
-  TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
-  TORCH_CHECK(input.scalar_type() == at::ScalarType::Half ||
-                  input.scalar_type() == at::ScalarType::BFloat16,
-              "Unsupported input data type for quantize_to_fp4.");
+  STD_TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
+  STD_TORCH_CHECK(
+      input.scalar_type() == torch::headeronly::ScalarType::Half ||
+          input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "Unsupported input data type for quantize_to_fp4.");
 
   int multiProcessorCount =
       get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
@@ -136,8 +137,9 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output,  // [..., d]
   auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
   auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
   auto output_ptr = static_cast<int64_t*>(output.data_ptr());
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  auto stream = get_current_cuda_stream(input.get_device_index());
   dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
   int const numBlocksPerSM =
       vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
@@ -149,7 +151,7 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output,  // [..., d]
       int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
   dim3 grid(grid_x, grid_y);
 
-  VLLM_DISPATCH_HALF_TYPES(
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
       input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
         using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
         auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());

csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu

@@ -14,14 +14,12 @@
  * limitations under the License.
  */
 
-#include "core/registration.h"
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include "libtorch_stable/torch_utils.h"
 
-#include <torch/all.h>
 #include <cutlass/arch/arch.h>
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
 #include "cutlass_extensions/common.hpp"
 
 #include "cute/tensor.hpp"
@@ -122,7 +120,7 @@ __global__ void __get_group_gemm_starts(
 #define __CALL_GET_STARTS_KERNEL_BLOCKSCALE(ELEMENT_AB_TYPE, SF_TYPE,         \
                                             TENSOR_C_TYPE, C_TYPE, LayoutSFA, \
                                             LayoutSFB, ScaleConfig)           \
-  else if (out_tensors.dtype() == TENSOR_C_TYPE) {                            \
+  else if (out_tensors.scalar_type() == TENSOR_C_TYPE) {                      \
     __get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, float,          \
                             LayoutSFA, LayoutSFB, ScaleConfig>                \
         <<<1, num_experts, 0, stream>>>(                                      \
@@ -150,50 +148,64 @@ __global__ void __get_group_gemm_starts(
   }
 
 template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
-void run_get_group_gemm_starts(
-    const torch::Tensor& a_starts, const torch::Tensor& b_starts,
-    const torch::Tensor& out_starts, const torch::Tensor& a_scales_starts,
-    const torch::Tensor& b_scales_starts, const torch::Tensor& alpha_starts,
-    const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
-    const torch::Tensor& a_strides, const torch::Tensor& b_strides,
-    const torch::Tensor& c_strides, int64_t a_stride_val, int64_t b_stride_val,
-    int64_t c_stride_val,
-    /*these are used for their base addresses*/
-    torch::Tensor const& a_tensors, torch::Tensor const& b_tensors,
-    torch::Tensor const& out_tensors, torch::Tensor const& a_scales,
-    torch::Tensor const& b_scales, torch::Tensor const& alphas,
-    torch::Tensor const& expert_offsets, torch::Tensor const& sf_offsets,
-    torch::Tensor const& problem_sizes, int M, int N, int K) {
+void run_get_group_gemm_starts(const torch::stable::Tensor& a_starts,
+                               const torch::stable::Tensor& b_starts,
+                               const torch::stable::Tensor& out_starts,
+                               const torch::stable::Tensor& a_scales_starts,
+                               const torch::stable::Tensor& b_scales_starts,
+                               const torch::stable::Tensor& alpha_starts,
+                               const torch::stable::Tensor& layout_sfa,
+                               const torch::stable::Tensor& layout_sfb,
+                               const torch::stable::Tensor& a_strides,
+                               const torch::stable::Tensor& b_strides,
+                               const torch::stable::Tensor& c_strides,
+                               int64_t a_stride_val, int64_t b_stride_val,
+                               int64_t c_stride_val,
+                               /*these are used for their base addresses*/
+                               torch::stable::Tensor const& a_tensors,
+                               torch::stable::Tensor const& b_tensors,
+                               torch::stable::Tensor const& out_tensors,
+                               torch::stable::Tensor const& a_scales,
+                               torch::stable::Tensor const& b_scales,
+                               torch::stable::Tensor const& alphas,
+                               torch::stable::Tensor const& expert_offsets,
+                               torch::stable::Tensor const& sf_offsets,
+                               torch::stable::Tensor const& problem_sizes,
+                               int M, int N, int K) {
   int num_experts = (int)expert_offsets.size(0);
-  auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
+  auto stream = get_current_cuda_stream(a_tensors.get_device_index());
 
-  TORCH_CHECK(out_tensors.size(1) == N,
-              "Output tensor shape doesn't match expected shape");
-  TORCH_CHECK(K / 2 == b_tensors.size(2),
-              "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
-              " dimension must match");
+  STD_TORCH_CHECK(out_tensors.size(1) == N,
+                  "Output tensor shape doesn't match expected shape");
+  STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
+                  "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
+                  " dimension must match");
   if (false) {
   }
   //(ELEMENT_AB_TYPE, BS_TYPE, TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB,
   // ScaleConfig)
   __CALL_GET_STARTS_KERNEL_BLOCKSCALE(
-      cutlass::float_e2m1_t, cutlass::float_ue4m3_t, torch::kBFloat16,
-      cutlass::bfloat16_t, LayoutSFA, LayoutSFB, ScaleConfig)
+      cutlass::float_e2m1_t, cutlass::float_ue4m3_t,
+      torch::headeronly::ScalarType::BFloat16, cutlass::bfloat16_t, LayoutSFA,
+      LayoutSFB, ScaleConfig)
   __CALL_GET_STARTS_KERNEL_BLOCKSCALE(cutlass::float_e2m1_t,
-                                      cutlass::float_ue4m3_t, torch::kFloat16,
-                                      half, LayoutSFA, LayoutSFB, ScaleConfig)
+                                      cutlass::float_ue4m3_t,
+                                      torch::headeronly::ScalarType::Half, half,
+                                      LayoutSFA, LayoutSFB, ScaleConfig)
   else {
-    TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
+    STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
   }
 }
 
 template <typename OutType>
 void run_fp4_blockwise_scaled_group_mm_sm100(
-    torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
-    const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
-    const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
-    const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
-    int N, int K) {
+    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, int M, int N, int K) {
   using ProblemShape =
       cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
   using ElementType = cutlass::float_e2m1_t;
@@ -272,20 +284,40 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
 
   using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
   int num_experts = static_cast<int>(expert_offsets.size(0));
-  auto options_int =
-      torch::TensorOptions().dtype(torch::kInt64).device(a.device());
 
-  torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
-  torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
-  torch::Tensor a_strides1 = torch::empty(num_experts, options_int);
-  torch::Tensor b_strides1 = torch::empty(num_experts, options_int);
-  torch::Tensor c_strides1 = torch::empty(num_experts, options_int);
+  torch::stable::Tensor a_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor b_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor out_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor a_scales_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor b_scales_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor alpha_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor layout_sfa = torch::stable::empty(
+      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
+      a.device());
+  torch::stable::Tensor layout_sfb = torch::stable::empty(
+      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
+      a.device());
+  torch::stable::Tensor a_strides1 =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor b_strides1 =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor c_strides1 =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
 
   run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
       a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
@@ -308,7 +340,7 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
           typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
   typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
   scheduler.raster_order = RasterOrderOptions::AlongM;
-  hw_info.device_id = a.get_device();
+  hw_info.device_id = a.get_device_index();
   static std::unordered_map<int, int> cached_sm_counts;
   if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
     cached_sm_counts[hw_info.device_id] =
@@ -350,32 +382,35 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
       scheduler};
 
   size_t workspace_size = Gemm::get_workspace_size(args);
-  auto const workspace_options =
-      torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
-  auto workspace = torch::empty(workspace_size, workspace_options);
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device());
+  auto workspace =
+      torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
+                           std::nullopt, a.device());
+  const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
 
   auto can_implement_status = gemm_op.can_implement(args);
-  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
-              "Failed to implement GEMM: status=", (int)can_implement_status);
+  STD_TORCH_CHECK(
+      can_implement_status == cutlass::Status::kSuccess,
+      "Failed to implement GEMM: status=", (int)can_implement_status);
 
   // Run the GEMM
   auto status = gemm_op.initialize(args, workspace.data_ptr());
-  TORCH_CHECK(status == cutlass::Status::kSuccess,
-              "Failed to initialize GEMM: status=", (int)status,
-              " workspace_size=", workspace_size, " num_experts=", num_experts,
-              " M=", M, " N=", N, " K=", K);
+  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
+                  "Failed to initialize GEMM: status=", (int)status,
+                  " workspace_size=", workspace_size,
+                  " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
 
   status = gemm_op.run(args, workspace.data_ptr(), stream);
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
+  STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
 }
 
 void run_fp4_blockwise_scaled_group_mm_sm120(
-    torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
-    const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
-    const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
-    const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
-    int N, int K) {
+    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, int M, int N, int K) {
   using ProblemShape =
       cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
   using ElementType = cutlass::float_e2m1_t;
@@ -446,20 +481,40 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
 
   using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
   int num_experts = static_cast<int>(expert_offsets.size(0));
-  auto options_int =
-      torch::TensorOptions().dtype(torch::kInt64).device(a.device());
 
-  torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
-  torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
-  torch::Tensor a_strides1 = torch::empty(num_experts, options_int);
-  torch::Tensor b_strides1 = torch::empty(num_experts, options_int);
-  torch::Tensor c_strides1 = torch::empty(num_experts, options_int);
+  torch::stable::Tensor a_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor b_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor out_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor a_scales_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor b_scales_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor alpha_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor layout_sfa = torch::stable::empty(
+      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
+      a.device());
+  torch::stable::Tensor layout_sfb = torch::stable::empty(
+      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
+      a.device());
+  torch::stable::Tensor a_strides1 =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor b_strides1 =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor c_strides1 =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
 
   run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
       a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
@@ -480,7 +535,7 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
   using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions;
   typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
   scheduler.raster_order = RasterOrderOptions::AlongM;
-  hw_info.device_id = a.get_device();
+  hw_info.device_id = a.get_device_index();
   static std::unordered_map<int, int> cached_sm_counts;
   if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
     cached_sm_counts[hw_info.device_id] =
@@ -523,33 +578,36 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
       scheduler};
 
   size_t workspace_size = Gemm::get_workspace_size(args);
-  auto const workspace_options =
-      torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
-  auto workspace = torch::empty(workspace_size, workspace_options);
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device());
+  auto workspace =
+      torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
+                           std::nullopt, a.device());
+  const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
 
   auto can_implement_status = gemm_op.can_implement(args);
-  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
-              "Failed to implement GEMM: status=", (int)can_implement_status);
+  STD_TORCH_CHECK(
+      can_implement_status == cutlass::Status::kSuccess,
+      "Failed to implement GEMM: status=", (int)can_implement_status);
 
   // Run the GEMM
   auto status = gemm_op.initialize(args, workspace.data_ptr());
-  TORCH_CHECK(status == cutlass::Status::kSuccess,
-              "Failed to initialize GEMM: status=", (int)status,
-              " workspace_size=", workspace_size, " num_experts=", num_experts,
-              " M=", M, " N=", N, " K=", K);
+  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
+                  "Failed to initialize GEMM: status=", (int)status,
+                  " workspace_size=", workspace_size,
+                  " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
 
   status = gemm_op.run(args, workspace.data_ptr(), stream);
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
+  STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
 }
 
 template <typename OutType>
 void run_fp4_blockwise_scaled_group_mm(
-    torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
-    const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
-    const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
-    const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
-    int N, int K) {
+    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, int M, int N, int K) {
   int32_t version_num = get_sm_version_num();
 #if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
   if (version_num >= 120 && version_num < 130) {
@@ -567,7 +625,7 @@ void run_fp4_blockwise_scaled_group_mm(
     return;
   }
 #endif
-  TORCH_CHECK_NOT_IMPLEMENTED(
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
       false,
       "No compiled cutlass_fp4_group_mm kernel for CUDA device capability: ",
       version_num, ". Required capability: 100 or 120");
@@ -575,26 +633,31 @@ void run_fp4_blockwise_scaled_group_mm(
 
 #if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \
     (defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120)
-constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
-constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
+constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
+constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn;
 #endif
 
-#define CHECK_TYPE(x, st, m) \
-  TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)
+#define CHECK_TYPE(x, st, m)             \
+  STD_TORCH_CHECK(x.scalar_type() == st, \
+                  ": Inconsistency of torch::stable::Tensor type:", m)
 #define CHECK_TH_CUDA(x, m) \
-  TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
+  STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
 #define CHECK_CONTIGUOUS(x, m) \
-  TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
+  STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
 #define CHECK_INPUT(x, st, m) \
   CHECK_TH_CUDA(x, m);        \
   CHECK_CONTIGUOUS(x, m);     \
   CHECK_TYPE(x, st, m)
 
-void cutlass_fp4_group_mm(
-    torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
-    const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
-    const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
-    const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets) {
+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) {
 #if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \
     (defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120)
   // Input validation
@@ -602,30 +665,34 @@ void cutlass_fp4_group_mm(
   CHECK_INPUT(b, FLOAT4_E2M1X2, "b");
   CHECK_INPUT(a_blockscale, SF_DTYPE, "a_blockscale");
   CHECK_INPUT(b_blockscales, SF_DTYPE, "b_blockscales");
-  CHECK_INPUT(alphas, at::ScalarType::Float, "alphas");
+  CHECK_INPUT(alphas, torch::headeronly::ScalarType::Float, "alphas");
 
-  TORCH_CHECK(a_blockscale.dim() == 2,
-              "expected a_blockscale to be of shape [num_experts, rounded_m,"
-              " k // group_size], observed rank: ",
-              a_blockscale.dim())
-  TORCH_CHECK(b_blockscales.dim() == 3,
-              "expected b_blockscale to be of shape: "
-              " [num_experts, n, k // group_size], observed rank: ",
-              b_blockscales.dim())
-  TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be  a 2D tensor");
-  TORCH_CHECK(problem_sizes.size(1) == 3,
-              "problem_sizes must have the shape (num_experts, 3)");
-  TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
-              "Number of experts in problem_sizes must match expert_offsets");
-  TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
-              "problem_sizes must be int32.");
+  STD_TORCH_CHECK(
+      a_blockscale.dim() == 2,
+      "expected a_blockscale to be of shape [num_experts, rounded_m,"
+      " k // group_size], observed rank: ",
+      a_blockscale.dim())
+  STD_TORCH_CHECK(b_blockscales.dim() == 3,
+                  "expected b_blockscale to be of shape: "
+                  " [num_experts, n, k // group_size], observed rank: ",
+                  b_blockscales.dim())
+  STD_TORCH_CHECK(problem_sizes.dim() == 2,
+                  "problem_sizes must be  a 2D tensor");
+  STD_TORCH_CHECK(problem_sizes.size(1) == 3,
+                  "problem_sizes must have the shape (num_experts, 3)");
+  STD_TORCH_CHECK(
+      problem_sizes.size(0) == expert_offsets.size(0),
+      "Number of experts in problem_sizes must match expert_offsets");
+  STD_TORCH_CHECK(
+      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
+      "problem_sizes must be int32.");
 
   int M = static_cast<int>(a.size(0));
   int N = static_cast<int>(b.size(1));
   int E = static_cast<int>(b.size(0));
   int K = static_cast<int>(2 * b.size(2));
 
-  if (output.scalar_type() == torch::kBFloat16) {
+  if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
     run_fp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
         output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
         expert_offsets, sf_offsets, M, N, K);
@@ -633,7 +700,7 @@ void cutlass_fp4_group_mm(
   #if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
     int32_t version_num = get_sm_version_num();
     if (version_num >= 120 && version_num < 130) {
-      TORCH_CHECK_NOT_IMPLEMENTED(
+      STD_TORCH_CHECK_NOT_IMPLEMENTED(
           false, "SM120 NVFP4 MOE only supports bfloat16 output, got: ",
           output.scalar_type());
     }
@@ -643,7 +710,7 @@ void cutlass_fp4_group_mm(
         expert_offsets, sf_offsets, M, N, K);
   }
 #else
-  TORCH_CHECK_NOT_IMPLEMENTED(
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
       false,
       "No compiled cutlass_fp4_group_mm kernel, vLLM must "
       "be compiled with ENABLE_NVFP4_SM100 or ENABLE_NVFP4_SM120 for SM100/120 "
@@ -651,6 +718,6 @@ void cutlass_fp4_group_mm(
 #endif
 }
 
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("cutlass_fp4_group_mm", &cutlass_fp4_group_mm);
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("cutlass_fp4_group_mm", TORCH_BOX(&cutlass_fp4_group_mm));
 }

csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu

@@ -14,16 +14,15 @@
  * limitations under the License.
  */
 
-#include <torch/all.h>
+#include <torch/csrc/stable/tensor.h>
+#include "libtorch_stable/torch_utils.h"
+#include "libtorch_stable/dispatch_utils.h"
+#include "cuda_vec_utils.cuh"
 
 #include <cuda_runtime_api.h>
 #include <cuda_runtime.h>
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-
 #include <cuda_fp8.h>
-#include "dispatch_utils.h"
 
 #include "cuda_utils.h"
 #include "nvfp4_utils.cuh"
@@ -327,25 +326,28 @@ void quant_impl(void* output, void* output_scale, void* input,
 }  // namespace vllm
 
 /*Quantization entry for fp4 experts quantization*/
-#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
+#define CHECK_TH_CUDA(x, m) \
+  STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
 #define CHECK_CONTIGUOUS(x, m) \
-  TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
+  STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
 #define CHECK_INPUT(x, m) \
   CHECK_TH_CUDA(x, m);    \
   CHECK_CONTIGUOUS(x, m);
 
-constexpr auto HALF = at::ScalarType::Half;
-constexpr auto BF16 = at::ScalarType::BFloat16;
-constexpr auto FLOAT = at::ScalarType::Float;
-constexpr auto INT = at::ScalarType::Int;
-constexpr auto UINT8 = at::ScalarType::Byte;
+constexpr auto HALF = torch::headeronly::ScalarType::Half;
+constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
+constexpr auto FLOAT = torch::headeronly::ScalarType::Float;
+constexpr auto INT = torch::headeronly::ScalarType::Int;
+constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
 
 // Common validation for fp4 experts quantization entry points.
 static void validate_fp4_experts_quant_inputs(
-    torch::Tensor const& output, torch::Tensor const& output_scale,
-    torch::Tensor const& input, torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts, int64_t m_topk,
+    torch::stable::Tensor const& output,
+    torch::stable::Tensor const& 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, int64_t m_topk,
     int64_t k) {
   CHECK_INPUT(output, "output");
   CHECK_INPUT(output_scale, "output_scale");
@@ -354,41 +356,42 @@ static void validate_fp4_experts_quant_inputs(
   CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
   CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
 
-  TORCH_CHECK(output.dim() == 2);
-  TORCH_CHECK(output_scale.dim() == 2);
-  TORCH_CHECK(input.dim() == 2);
-  TORCH_CHECK(input_global_scale.dim() == 1);
-  TORCH_CHECK(input_offset_by_experts.dim() == 1);
-  TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
+  STD_TORCH_CHECK(output.dim() == 2);
+  STD_TORCH_CHECK(output_scale.dim() == 2);
+  STD_TORCH_CHECK(input.dim() == 2);
+  STD_TORCH_CHECK(input_global_scale.dim() == 1);
+  STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
+  STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
 
-  TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
-  TORCH_CHECK(input_global_scale.scalar_type() == FLOAT);
-  TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
-  TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
+  STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
+  STD_TORCH_CHECK(input_global_scale.scalar_type() == FLOAT);
+  STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
+  STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
   // output is uint8 (two nvfp4 values are packed into one uint8)
   // output_scale is int32 (four fp8 values are packed into one int32)
-  TORCH_CHECK(output.scalar_type() == UINT8);
-  TORCH_CHECK(output_scale.scalar_type() == INT);
+  STD_TORCH_CHECK(output.scalar_type() == UINT8);
+  STD_TORCH_CHECK(output_scale.scalar_type() == INT);
 
   const int BLOCK_SIZE = 16;
-  TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
+  STD_TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
   auto n_experts = input_global_scale.size(0);
-  TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
-  TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
-  TORCH_CHECK(output.size(0) == m_topk);
-  TORCH_CHECK(output.size(1) == k / 2);
+  STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
+  STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
+  STD_TORCH_CHECK(output.size(0) == m_topk);
+  STD_TORCH_CHECK(output.size(1) == k / 2);
   int scales_k = k / BLOCK_SIZE;
   // 4 means the swizzle requirement by nvidia nvfp4.
   int padded_k = (scales_k + (4 - 1)) / 4 * 4;
   // 4 means 4 fp8 values are packed into one int32
-  TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
+  STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
 }
 
 void scaled_fp4_experts_quant_sm1xxa(
-    torch::Tensor& output, torch::Tensor& output_scale,
-    torch::Tensor const& input, torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts) {
+    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) {
   auto m_topk = input.size(0);
   auto k = input.size(1);
 
@@ -397,11 +400,11 @@ void scaled_fp4_experts_quant_sm1xxa(
                                     output_scale_offset_by_experts, m_topk, k);
 
   auto n_experts = input_global_scale.size(0);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream =
-      at::cuda::getCurrentCUDAStream(input.get_device());
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
 
-  VLLM_DISPATCH_HALF_TYPES(
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
       input.scalar_type(), "nvfp4_experts_quant_kernel", [&] {
         using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
         vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
@@ -413,14 +416,15 @@ void scaled_fp4_experts_quant_sm1xxa(
 }
 
 void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
-    torch::Tensor& output, torch::Tensor& output_scale,
-    torch::Tensor const& input, torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts) {
+    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) {
   auto m_topk = input.size(0);
   // Input has gate || up layout, so k = input.size(1) / 2
   auto k_times_2 = input.size(1);
-  TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
+  STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
   auto k = k_times_2 / 2;
 
   validate_fp4_experts_quant_inputs(output, output_scale, input,
@@ -428,11 +432,11 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
                                     output_scale_offset_by_experts, m_topk, k);
 
   auto n_experts = input_global_scale.size(0);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  const cudaStream_t stream =
-      at::cuda::getCurrentCUDAStream(input.get_device());
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
 
-  VLLM_DISPATCH_HALF_TYPES(
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
       input.scalar_type(), "silu_mul_nvfp4_experts_quant_kernel", [&] {
         using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
         vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(

csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu

@@ -0,0 +1,172 @@
+/*
+ * Copyright (c) 2025, 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.
+ */
+
+#include <torch/csrc/stable/tensor.h>
+
+#include "libtorch_stable/torch_utils.h"
+
+#include "cutlass_extensions/common.hpp"
+#include "nvfp4_utils.cuh"
+
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
+    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
+void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
+                             torch::stable::Tensor const& input,
+                             torch::stable::Tensor const& output_sf,
+                             torch::stable::Tensor const& input_sf,
+                             bool is_sf_swizzled_layout);
+#endif
+
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
+    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
+void scaled_fp4_experts_quant_sm1xxa(
+    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);
+#endif
+
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
+    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
+void silu_and_mul_nvfp4_quant_sm1xxa(torch::stable::Tensor& output,
+                                     torch::stable::Tensor& output_sf,
+                                     torch::stable::Tensor& input,
+                                     torch::stable::Tensor& input_sf);
+#endif
+
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
+    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
+void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
+    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);
+#endif
+
+static bool nvfp4_quant_sm_supported() {
+  const int32_t sm = get_sm_version_num();
+#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
+  if (sm >= 100 && sm < 120) return true;
+#endif
+#if defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120
+  if (sm >= 120 && sm < 130) return true;
+#endif
+  return false;
+}
+
+void scaled_fp4_quant_out(torch::stable::Tensor const& input,
+                          torch::stable::Tensor const& input_sf,
+                          bool is_sf_swizzled_layout,
+                          torch::stable::Tensor& output,
+                          torch::stable::Tensor& output_sf) {
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
+    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
+  STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
+                  "No compiled nvfp4 quantization kernel for SM ",
+                  get_sm_version_num(),
+                  ". Recompile with the appropriate CUDA arch.");
+  return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
+                                 is_sf_swizzled_layout);
+#endif
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
+                                  "No compiled nvfp4 quantization kernel");
+}
+
+std::tuple<torch::stable::Tensor, torch::stable::Tensor> scaled_fp4_quant_func(
+    torch::stable::Tensor const& input, torch::stable::Tensor const& input_sf,
+    bool is_sf_swizzled_layout) {
+  int64_t n = input.size(-1);
+  int64_t m = input.numel() / n;
+  auto device = input.device();
+
+  // Two fp4 values packed into a uint8
+  auto output = torch::stable::empty(
+      {m, n / 2}, torch::headeronly::ScalarType::Byte, std::nullopt, device);
+
+  torch::stable::Tensor output_sf;
+  if (is_sf_swizzled_layout) {
+    auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
+    output_sf = torch::stable::empty(
+        {sf_m, sf_n}, torch::headeronly::ScalarType::Int, std::nullopt, device);
+  } else {
+    output_sf = torch::stable::empty({m, n / CVT_FP4_SF_VEC_SIZE},
+                                     torch::headeronly::ScalarType::Byte,
+                                     std::nullopt, device);
+  }
+
+  scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
+                       output_sf);
+  return {output, output_sf};
+}
+
+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) {
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
+    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
+  STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
+                  "No compiled nvfp4 experts quantization kernel for SM ",
+                  get_sm_version_num(),
+                  ". Recompile with the appropriate CUDA arch.");
+  return scaled_fp4_experts_quant_sm1xxa(
+      output, output_scale, input, input_global_scale, input_offset_by_experts,
+      output_scale_offset_by_experts);
+#endif
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
+      false, "No compiled nvfp4 experts quantization kernel");
+}
+
+void silu_and_mul_nvfp4_quant(torch::stable::Tensor& output,
+                              torch::stable::Tensor& output_sf,
+                              torch::stable::Tensor& input,
+                              torch::stable::Tensor& input_sf) {
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
+    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
+  STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
+                  "No compiled silu_and_mul nvfp4 quantization kernel for SM ",
+                  get_sm_version_num(),
+                  ". Recompile with the appropriate CUDA arch.");
+  return silu_and_mul_nvfp4_quant_sm1xxa(output, output_sf, input, input_sf);
+#endif
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
+      false, "No compiled silu_and_mul nvfp4 quantization kernel");
+}
+
+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) {
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
+    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
+  STD_TORCH_CHECK(nvfp4_quant_sm_supported(),
+                  "No compiled silu_and_mul nvfp4 experts quantization kernel "
+                  "for SM ",
+                  get_sm_version_num(),
+                  ". Recompile with the appropriate CUDA arch.");
+  return silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
+      output, output_scale, input, input_global_scale, input_offset_by_experts,
+      output_scale_offset_by_experts);
+#endif
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
+      false, "No compiled silu_and_mul nvfp4 experts quantization kernel");
+}

csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu

@@ -14,16 +14,16 @@
  * limitations under the License.
  */
 
-#include <torch/all.h>
+#include <torch/csrc/stable/tensor.h>
 
 #include <cuda_runtime_api.h>
 #include <cuda_runtime.h>
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-
 #include <cuda_fp8.h>
-#include "dispatch_utils.h"
+
+#include "libtorch_stable/torch_utils.h"
+#include "libtorch_stable/dispatch_utils.h"
+#include "cuda_vec_utils.cuh"
 
 #include "cuda_utils.h"
 #include "launch_bounds_utils.h"
@@ -173,18 +173,19 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
 
 }  // namespace vllm
 
-void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
-                             torch::Tensor const& input,
-                             torch::Tensor const& output_sf,
-                             torch::Tensor const& input_sf,
+void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output,
+                             torch::stable::Tensor const& input,
+                             torch::stable::Tensor const& output_sf,
+                             torch::stable::Tensor const& input_sf,
                              bool is_sf_swizzled_layout) {
   int32_t m = input.size(0);
   int32_t n = input.size(1);
 
-  TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
-  TORCH_CHECK(input.scalar_type() == at::ScalarType::Half ||
-                  input.scalar_type() == at::ScalarType::BFloat16,
-              "Unsupported input data type for quantize_to_fp4.");
+  STD_TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
+  STD_TORCH_CHECK(
+      input.scalar_type() == torch::headeronly::ScalarType::Half ||
+          input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "Unsupported input data type for quantize_to_fp4.");
 
   int multiProcessorCount =
       get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
@@ -192,8 +193,9 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
   auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
   auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
   auto output_ptr = static_cast<int64_t*>(output.data_ptr());
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  auto stream = get_current_cuda_stream(input.get_device_index());
 
   int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
 
@@ -213,15 +215,15 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
                  std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
     dim3 grid(grid_x, grid_y);
 
-    VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
-      using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
-      auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
-      // NOTE: We don't support e8m0 scales at this moment.
-      vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
-          m, n, num_padded_cols, input_ptr, input_sf_ptr,
-          reinterpret_cast<uint32_t*>(output_ptr),
-          reinterpret_cast<uint32_t*>(sf_out));
-    });
+    VLLM_STABLE_DISPATCH_HALF_TYPES(
+        input.scalar_type(), "nvfp4_quant_kernel", [&] {
+          using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
+          auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
+          vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
+              m, n, num_padded_cols, input_ptr, input_sf_ptr,
+              reinterpret_cast<uint32_t*>(output_ptr),
+              reinterpret_cast<uint32_t*>(sf_out));
+        });
   } else {
     int num_packed_cols = n / CVT_FP4_ELTS_PER_THREAD;
     int grid_y = vllm::div_round_up(num_packed_cols, static_cast<int>(block.x));
@@ -229,15 +231,15 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
         m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
     dim3 grid(grid_x, grid_y);
 
-    VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
-      using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
-      auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
-      // NOTE: We don't support e8m0 scales at this moment.
-      vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
-          <<<grid, block, 0, stream>>>(m, n, sf_n_unpadded, num_packed_cols,
-                                       input_ptr, input_sf_ptr,
-                                       reinterpret_cast<uint32_t*>(output_ptr),
-                                       reinterpret_cast<uint32_t*>(sf_out));
-    });
+    VLLM_STABLE_DISPATCH_HALF_TYPES(
+        input.scalar_type(), "nvfp4_quant_kernel", [&] {
+          using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
+          auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
+          vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
+              <<<grid, block, 0, stream>>>(
+                  m, n, sf_n_unpadded, num_packed_cols, input_ptr, input_sf_ptr,
+                  reinterpret_cast<uint32_t*>(output_ptr),
+                  reinterpret_cast<uint32_t*>(sf_out));
+        });
   }
 }

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu

@@ -14,32 +14,39 @@
  * limitations under the License.
  */
 
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
+#include <torch/csrc/stable/tensor.h>
+
+#include "libtorch_stable/torch_utils.h"
+
 #include "cutlass_extensions/common.hpp"
 
 #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
-                                  torch::Tensor const& B,
-                                  torch::Tensor const& A_sf,
-                                  torch::Tensor const& B_sf,
-                                  torch::Tensor const& alpha);
+void cutlass_scaled_fp4_mm_sm100a(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);
 #endif
 
 #if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
-void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A,
-                                  torch::Tensor const& B,
-                                  torch::Tensor const& A_sf,
-                                  torch::Tensor const& B_sf,
-                                  torch::Tensor const& alpha);
+void cutlass_scaled_fp4_mm_sm120a(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);
 #endif
 
-void cutlass_scaled_fp4_mm(torch::Tensor& D, const torch::Tensor& A,
-                           const torch::Tensor& B, const torch::Tensor& A_sf,
-                           const torch::Tensor& B_sf,
-                           const torch::Tensor& alpha) {
-  // Make sure we’re on A’s device.
-  const c10::cuda::OptionalCUDAGuard device_guard(device_of(A));
+void cutlass_scaled_fp4_mm(torch::stable::Tensor& D,
+                           const torch::stable::Tensor& A,
+                           const torch::stable::Tensor& B,
+                           const torch::stable::Tensor& A_sf,
+                           const torch::stable::Tensor& B_sf,
+                           const torch::stable::Tensor& alpha) {
+  // Make sure we're on A's device.
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      A.get_device_index());
   const int32_t sm = get_sm_version_num();
 
 #if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
@@ -56,8 +63,9 @@ void cutlass_scaled_fp4_mm(torch::Tensor& D, const torch::Tensor& A,
   }
 #endif
 
-  TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 mm kernel for SM ", sm,
-                              ". Recompile with CUDA >= 12.8 and CC >= 100.");
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
+      false, "No compiled nvfp4 mm kernel for SM ", sm,
+      ". Recompile with CUDA >= 12.8 and CC >= 100.");
 }
 
 bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability) {

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu

@@ -14,10 +14,9 @@
  * limitations under the License.
  */
 
-#include <torch/all.h>
+#include <torch/csrc/stable/tensor.h>
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "libtorch_stable/torch_utils.h"
 
 #include "cutlass_extensions/common.hpp"
 
@@ -127,8 +126,9 @@ struct Fp4GemmSm100 {
 
 template <typename Config>
 typename Config::Gemm::Arguments args_from_options(
-    at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
-    at::Tensor const& A_sf, at::Tensor const& B_sf, at::Tensor const& alpha,
+    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,
     int64_t M, int64_t N, int64_t K) {
   using ElementA = typename Config::Gemm::ElementA;
   using ElementB = typename Config::Gemm::ElementB;
@@ -174,19 +174,20 @@ typename Config::Gemm::Arguments args_from_options(
 }
 
 template <typename Config>
-void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
-             at::Tensor const& A_sf, at::Tensor const& B_sf,
-             at::Tensor const& alpha, int64_t m, int64_t n, int64_t k,
-             cudaStream_t stream) {
+void runGemm(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, int64_t m, int64_t n,
+             int64_t k, cudaStream_t stream) {
   typename Config::Gemm gemm;
 
   auto arguments =
       args_from_options<Config>(D, A, B, A_sf, B_sf, alpha, m, n, k);
 
   size_t workspace_size = Config::Gemm::get_workspace_size(arguments);
-  auto const workspace_options =
-      torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
-  auto workspace = torch::empty(workspace_size, workspace_options);
+  auto workspace =
+      torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
+                           std::nullopt, A.device());
 
   CUTLASS_CHECK(gemm.can_implement(arguments));
 
@@ -197,12 +198,13 @@ void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
 
 // Dispatch function to select appropriate config based on M
 template <typename OutType>
-void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
-                               torch::Tensor const& B,
-                               torch::Tensor const& A_sf,
-                               torch::Tensor const& B_sf,
-                               torch::Tensor const& alpha, int64_t m, int64_t n,
-                               int64_t k, cudaStream_t stream) {
+void cutlass_fp4_gemm_dispatch(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, int64_t m,
+                               int64_t n, int64_t k, cudaStream_t stream) {
   uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
 
   if (mp2 <= 16) {
@@ -222,61 +224,65 @@ void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
 
 #else
 template <typename OutType>
-void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
-                               torch::Tensor const& B,
-                               torch::Tensor const& A_sf,
-                               torch::Tensor const& B_sf,
-                               torch::Tensor const& alpha, int64_t m, int64_t n,
-                               int64_t k, cudaStream_t stream) {
-  TORCH_CHECK(false,
-              "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
-              "a CUTLASS 3.8 source directory to enable support.");
+void cutlass_fp4_gemm_dispatch(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, int64_t m,
+                               int64_t n, int64_t k, cudaStream_t stream) {
+  STD_TORCH_CHECK(false,
+                  "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
+                  "a CUTLASS 3.8 source directory to enable support.");
 }
 #endif  // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
 
-#define CHECK_TYPE(x, st, m) \
-  TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)
+#define CHECK_TYPE(x, st, m)             \
+  STD_TORCH_CHECK(x.scalar_type() == st, \
+                  ": Inconsistency of torch::stable::Tensor type:", m)
 #define CHECK_TH_CUDA(x, m) \
-  TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
+  STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
 #define CHECK_CONTIGUOUS(x, m) \
-  TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
+  STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
 #define CHECK_INPUT(x, st, m) \
   CHECK_TH_CUDA(x, m);        \
   CHECK_CONTIGUOUS(x, m);     \
   CHECK_TYPE(x, st, m)
 
-constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
-constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
+constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
+constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn;
 
-void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
-                                  torch::Tensor const& B,
-                                  torch::Tensor const& A_sf,
-                                  torch::Tensor const& B_sf,
-                                  torch::Tensor const& alpha) {
+void cutlass_scaled_fp4_mm_sm100a(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) {
   CHECK_INPUT(A, FLOAT4_E2M1X2, "a");
   CHECK_INPUT(B, FLOAT4_E2M1X2, "b");
 
   CHECK_INPUT(A_sf, SF_DTYPE, "scale_a");
   CHECK_INPUT(B_sf, SF_DTYPE, "scale_b");
 
-  CHECK_INPUT(alpha, at::ScalarType::Float, "alpha");
+  CHECK_INPUT(alpha, torch::headeronly::ScalarType::Float, "alpha");
 
-  TORCH_CHECK(A.dim() == 2, "a must be a matrix");
-  TORCH_CHECK(B.dim() == 2, "b must be a matrix");
-  TORCH_CHECK(A.sizes()[1] == B.sizes()[1],
-              "a and b shapes cannot be multiplied (", A.sizes()[0], "x",
-              A.sizes()[1], " and ", B.sizes()[0], "x", B.sizes()[1], ")");
+  STD_TORCH_CHECK(A.dim() == 2, "a must be a matrix");
+  STD_TORCH_CHECK(B.dim() == 2, "b must be a matrix");
+  STD_TORCH_CHECK(A.size(1) == B.size(1),
+                  "a and b shapes cannot be multiplied (", A.size(0), "x",
+                  A.size(1), " and ", B.size(0), "x", B.size(1), ")");
 
-  auto const m = A.sizes()[0];
-  auto const n = B.sizes()[0];
-  auto const k = A.sizes()[1] * 2;
+  auto const m = A.size(0);
+  auto const n = B.size(0);
+  auto const k = A.size(1) * 2;
 
   constexpr int alignment = 32;
-  TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ", alignment,
-              ", but got a shape: (", A.sizes()[0], "x", A.sizes()[1],
-              "), k: ", k, ".");
-  TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ", alignment,
-              ", but got b shape: (", B.sizes()[0], "x", B.sizes()[1], ").");
+  STD_TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ",
+                  alignment, ", but got a shape: (", A.size(0), "x", A.size(1),
+                  "), k: ", k, ".");
+  STD_TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ",
+                  alignment, ", but got b shape: (", B.size(0), "x", B.size(1),
+                  ").");
 
   auto round_up = [](int x, int y) { return (x + y - 1) / y * y; };
   int rounded_m = round_up(m, 128);
@@ -285,33 +291,34 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
   // integer.
   int rounded_k = round_up(k / 16, 4);
 
-  TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
-  TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
-  TORCH_CHECK(A_sf.sizes()[1] == B_sf.sizes()[1],
-              "scale_a and scale_b shapes cannot be multiplied (",
-              A_sf.sizes()[0], "x", A_sf.sizes()[1], " and ", B_sf.sizes()[0],
-              "x", B_sf.sizes()[1], ")");
-  TORCH_CHECK(A_sf.sizes()[0] == rounded_m && A_sf.sizes()[1] == rounded_k,
-              "scale_a must be padded and swizzled to a shape (", rounded_m,
-              "x", rounded_k, "), but got a shape (", A_sf.sizes()[0], "x",
-              A_sf.sizes()[1], ")");
-  TORCH_CHECK(B_sf.sizes()[0] == rounded_n && B_sf.sizes()[1] == rounded_k,
-              "scale_b must be padded and swizzled to a shape (", rounded_n,
-              "x", rounded_k, "), but got a shape (", B_sf.sizes()[0], "x",
-              B_sf.sizes()[1], ")");
+  STD_TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
+  STD_TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
+  STD_TORCH_CHECK(A_sf.size(1) == B_sf.size(1),
+                  "scale_a and scale_b shapes cannot be multiplied (",
+                  A_sf.size(0), "x", A_sf.size(1), " and ", B_sf.size(0), "x",
+                  B_sf.size(1), ")");
+  STD_TORCH_CHECK(A_sf.size(0) == rounded_m && A_sf.size(1) == rounded_k,
+                  "scale_a must be padded and swizzled to a shape (", rounded_m,
+                  "x", rounded_k, "), but got a shape (", A_sf.size(0), "x",
+                  A_sf.size(1), ")");
+  STD_TORCH_CHECK(B_sf.size(0) == rounded_n && B_sf.size(1) == rounded_k,
+                  "scale_b must be padded and swizzled to a shape (", rounded_n,
+                  "x", rounded_k, "), but got a shape (", B_sf.size(0), "x",
+                  B_sf.size(1), ")");
 
-  auto out_dtype = D.dtype();
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());
+  auto out_dtype = D.scalar_type();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      A.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(A.get_device_index());
 
-  if (out_dtype == at::ScalarType::Half) {
+  if (out_dtype == torch::headeronly::ScalarType::Half) {
     cutlass_fp4_gemm_dispatch<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n,
                                                k, stream);
-  } else if (out_dtype == at::ScalarType::BFloat16) {
+  } else if (out_dtype == torch::headeronly::ScalarType::BFloat16) {
     cutlass_fp4_gemm_dispatch<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha,
                                                    m, n, k, stream);
   } else {
-    TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm (", out_dtype,
-                ")");
+    STD_TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm (",
+                    out_dtype, ")");
   }
 }

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu

@@ -14,10 +14,9 @@
  * limitations under the License.
  */
 
-#include <torch/all.h>
+#include <torch/csrc/stable/tensor.h>
 
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "libtorch_stable/torch_utils.h"
 
 #include "cutlass_extensions/common.hpp"
 
@@ -34,19 +33,20 @@
 
 using namespace cute;
 
-#define CHECK_TYPE(x, st, m) \
-  TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)
+#define CHECK_TYPE(x, st, m)             \
+  STD_TORCH_CHECK(x.scalar_type() == st, \
+                  ": Inconsistency of torch::stable::Tensor type:", m)
 #define CHECK_TH_CUDA(x, m) \
-  TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
+  STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
 #define CHECK_CONTIGUOUS(x, m) \
-  TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
+  STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
 #define CHECK_INPUT(x, st, m) \
   CHECK_TH_CUDA(x, m);        \
   CHECK_CONTIGUOUS(x, m);     \
   CHECK_TYPE(x, st, m)
 
-constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
-constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
+constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
+constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn;
 
 struct sm120_fp4_config_M256 {
   using ClusterShape = Shape<_1, _1, _1>;
@@ -109,12 +109,13 @@ struct Fp4GemmSm120 {
 };
 
 template <typename Gemm>
-typename Gemm::Arguments args_from_options(at::Tensor& D, at::Tensor const& A,
-                                           at::Tensor const& B,
-                                           at::Tensor const& A_sf,
-                                           at::Tensor const& B_sf,
-                                           torch::Tensor const& alpha, int M,
-                                           int N, int K) {
+typename Gemm::Arguments args_from_options(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,
+                                           int M, int N, int K) {
   using ElementA = typename Gemm::ElementA;
   using ElementB = typename Gemm::ElementB;
   using ElementD = typename Gemm::ElementD;
@@ -158,18 +159,19 @@ typename Gemm::Arguments args_from_options(at::Tensor& D, at::Tensor const& A,
 }
 
 template <typename Gemm>
-void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
-             at::Tensor const& A_sf, at::Tensor const& B_sf,
-             torch::Tensor const& alpha, int M, int N, int K,
+void runGemm(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, int M, int N, int K,
              cudaStream_t stream) {
   Gemm gemm;
 
   auto arguments = args_from_options<Gemm>(D, A, B, A_sf, B_sf, alpha, M, N, K);
 
   size_t workspace_size = Gemm::get_workspace_size(arguments);
-  auto const workspace_options =
-      torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
-  auto workspace = torch::empty(workspace_size, workspace_options);
+  auto workspace =
+      torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
+                           std::nullopt, A.device());
 
   CUTLASS_CHECK(gemm.can_implement(arguments));
 
@@ -178,12 +180,13 @@ void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
   CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
 }
 
-void cutlass_fp4_bf16_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
-                                    torch::Tensor const& B,
-                                    torch::Tensor const& A_sf,
-                                    torch::Tensor const& B_sf,
-                                    torch::Tensor const& alpha, int m, int n,
-                                    int k, cudaStream_t stream) {
+void cutlass_fp4_bf16_gemm_dispatch(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, int m,
+                                    int n, int k, cudaStream_t stream) {
   uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
   if (mp2 <= 256) {
     runGemm<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::bfloat16_t>::Gemm>(
@@ -194,12 +197,13 @@ void cutlass_fp4_bf16_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
   }
 }
 
-void cutlass_fp4_f16_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
-                                   torch::Tensor const& B,
-                                   torch::Tensor const& A_sf,
-                                   torch::Tensor const& B_sf,
-                                   torch::Tensor const& alpha, int m, int n,
-                                   int k, cudaStream_t stream) {
+void cutlass_fp4_f16_gemm_dispatch(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, int m,
+                                   int n, int k, cudaStream_t stream) {
   uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
   if (mp2 <= 256) {
     runGemm<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::half_t>::Gemm>(
@@ -210,11 +214,12 @@ void cutlass_fp4_f16_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
   }
 }
 
-void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A,
-                                  torch::Tensor const& B,
-                                  torch::Tensor const& A_sf,
-                                  torch::Tensor const& B_sf,
-                                  torch::Tensor const& alpha) {
+void cutlass_scaled_fp4_mm_sm120a(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) {
 #if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED)
   CHECK_INPUT(A, FLOAT4_E2M1X2, "a");
   CHECK_INPUT(B, FLOAT4_E2M1X2, "b");
@@ -222,24 +227,25 @@ void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A,
   CHECK_INPUT(A_sf, SF_DTYPE, "scale_a");
   CHECK_INPUT(B_sf, SF_DTYPE, "scale_b");
 
-  CHECK_INPUT(alpha, at::ScalarType::Float, "alpha");
+  CHECK_INPUT(alpha, torch::headeronly::ScalarType::Float, "alpha");
 
-  TORCH_CHECK(A.dim() == 2, "a must be a matrix");
-  TORCH_CHECK(B.dim() == 2, "b must be a matrix");
-  TORCH_CHECK(A.sizes()[1] == B.sizes()[1],
-              "a and b shapes cannot be multiplied (", A.sizes()[0], "x",
-              A.sizes()[1], " and ", B.sizes()[0], "x", B.sizes()[1], ")");
+  STD_TORCH_CHECK(A.dim() == 2, "a must be a matrix");
+  STD_TORCH_CHECK(B.dim() == 2, "b must be a matrix");
+  STD_TORCH_CHECK(A.size(1) == B.size(1),
+                  "a and b shapes cannot be multiplied (", A.size(0), "x",
+                  A.size(1), " and ", B.size(0), "x", B.size(1), ")");
 
-  auto const m = A.sizes()[0];
-  auto const n = B.sizes()[0];
-  auto const k = A.sizes()[1] * 2;
+  auto const m = A.size(0);
+  auto const n = B.size(0);
+  auto const k = A.size(1) * 2;
 
   constexpr int alignment = 32;
-  TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ", alignment,
-              ", but got a shape: (", A.sizes()[0], "x", A.sizes()[1],
-              "), k: ", k, ".");
-  TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ", alignment,
-              ", but got b shape: (", B.sizes()[0], "x", B.sizes()[1], ").");
+  STD_TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ",
+                  alignment, ", but got a shape: (", A.size(0), "x", A.size(1),
+                  "), k: ", k, ".");
+  STD_TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ",
+                  alignment, ", but got b shape: (", B.size(0), "x", B.size(1),
+                  ").");
 
   auto round_up = [](int x, int y) { return (x + y - 1) / y * y; };
   int rounded_m = round_up(m, 128);
@@ -248,38 +254,39 @@ void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A,
   // integer.
   int rounded_k = round_up(k / 16, 4);
 
-  TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
-  TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
-  TORCH_CHECK(A_sf.sizes()[1] == B_sf.sizes()[1],
-              "scale_a and scale_b shapes cannot be multiplied (",
-              A_sf.sizes()[0], "x", A_sf.sizes()[1], " and ", B_sf.sizes()[0],
-              "x", B_sf.sizes()[1], ")");
-  TORCH_CHECK(A_sf.sizes()[0] == rounded_m && A_sf.sizes()[1] == rounded_k,
-              "scale_a must be padded and swizzled to a shape (", rounded_m,
-              "x", rounded_k, "), but got a shape (", A_sf.sizes()[0], "x",
-              A_sf.sizes()[1], ")");
-  TORCH_CHECK(B_sf.sizes()[0] == rounded_n && B_sf.sizes()[1] == rounded_k,
-              "scale_b must be padded and swizzled to a shape (", rounded_n,
-              "x", rounded_k, "), but got a shape (", B_sf.sizes()[0], "x",
-              B_sf.sizes()[1], ")");
+  STD_TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
+  STD_TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
+  STD_TORCH_CHECK(A_sf.size(1) == B_sf.size(1),
+                  "scale_a and scale_b shapes cannot be multiplied (",
+                  A_sf.size(0), "x", A_sf.size(1), " and ", B_sf.size(0), "x",
+                  B_sf.size(1), ")");
+  STD_TORCH_CHECK(A_sf.size(0) == rounded_m && A_sf.size(1) == rounded_k,
+                  "scale_a must be padded and swizzled to a shape (", rounded_m,
+                  "x", rounded_k, "), but got a shape (", A_sf.size(0), "x",
+                  A_sf.size(1), ")");
+  STD_TORCH_CHECK(B_sf.size(0) == rounded_n && B_sf.size(1) == rounded_k,
+                  "scale_b must be padded and swizzled to a shape (", rounded_n,
+                  "x", rounded_k, "), but got a shape (", B_sf.size(0), "x",
+                  B_sf.size(1), ")");
 
-  auto out_dtype = D.dtype();
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());
+  auto out_dtype = D.scalar_type();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      A.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(A.get_device_index());
 
-  if (out_dtype == at::ScalarType::BFloat16) {
+  if (out_dtype == torch::headeronly::ScalarType::BFloat16) {
     return cutlass_fp4_bf16_gemm_dispatch(D, A, B, A_sf, B_sf, alpha, m, n, k,
                                           stream);
-  } else if (out_dtype == at::ScalarType::Half) {
+  } else if (out_dtype == torch::headeronly::ScalarType::Half) {
     return cutlass_fp4_f16_gemm_dispatch(D, A, B, A_sf, B_sf, alpha, m, n, k,
                                          stream);
   } else {
-    TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm sm120 (",
-                out_dtype, ")");
+    STD_TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm sm120 (",
+                    out_dtype, ")");
   }
 #else
-  TORCH_CHECK(false,
-              "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
-              "a CUTLASS 3.8 source directory to enable support.");
+  STD_TORCH_CHECK(false,
+                  "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
+                  "a CUTLASS 3.8 source directory to enable support.");
 #endif  // defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED)
-}
+}

csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh

@@ -20,7 +20,7 @@
 #include <cuda_fp8.h>
 #include <utility>
 
-#include "../../cuda_vec_utils.cuh"
+#include "cuda_vec_utils.cuh"
 
 #if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \
     CUDA_VERSION >= 12090

csrc/libtorch_stable/torch_bindings.cpp

@@ -103,6 +103,102 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def(
       "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> "
       "bool");
+
+  // CUTLASS nvfp4 block scaled GEMM
+  ops.def(
+      "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
+      "                      Tensor block_scale_a, Tensor block_scale_b,"
+      "                      Tensor alpha) -> ()");
+
+  // cutlass nvfp4 block scaled group GEMM
+  ops.def(
+      "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
+      " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
+      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
+
+  // Compute NVFP4 block quantized tensor.
+  ops.def(
+      "scaled_fp4_quant(Tensor input,"
+      "                 Tensor input_scale, bool "
+      "is_sf_swizzled_layout) -> (Tensor, Tensor)");
+
+  // Out variant
+  // TODO: Add out_variant tag once PyTorch supports it (added in 2.11)
+  // This registration is now migrated to stable ABI
+  // at::Tag::out_variant is not available in the stable ABI (enum_tag.h is not
+  // yet in torch/headeronly), the tag should be applied from Python
+  // via torch.library.Library.define(..., tags=(torch.Tag.out_variant,))
+  // with the .impl remaining in C++.
+  // See pytorch/pytorch#176117.
+  ops.def(
+      "scaled_fp4_quant.out(Tensor input,"
+      "                     Tensor input_scale, bool "
+      "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) "
+      "-> ()");
+
+  // Compute NVFP4 experts quantization.
+  ops.def(
+      "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale,"
+      "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
+      "Tensor output_scale_offset_by_experts) -> ()");
+
+  // Fused SiLU+Mul+NVFP4 experts quantization.
+  ops.def(
+      "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! "
+      "output_scale,"
+      "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
+      "Tensor output_scale_offset_by_experts) -> ()");
+
+  // Fused SiLU+Mul+NVFP4 quantization.
+  ops.def(
+      "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
+      "Tensor input, Tensor input_global_scale) -> ()");
+
+  // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
+  // of the given capability
+  ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
+
+  // CUTLASS w4a8 GEMM
+  ops.def(
+      "cutlass_w4a8_mm("
+      "   Tensor A,"
+      "   Tensor B,"
+      "   Tensor group_scales,"
+      "   int    group_size,"
+      "   Tensor channel_scales,"
+      "   Tensor token_scales,"
+      "   ScalarType? out_type,"
+      "   str?   maybe_schedule"
+      ") -> Tensor");
+
+  // pack scales
+  ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor");
+
+  // encode and reorder weight matrix
+  ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor");
+
+  // CUTLASS w4a8 grouped GEMM
+  ops.def(
+      "cutlass_w4a8_moe_mm("
+      "   Tensor! out_tensors,"
+      "   Tensor a_tensors,"
+      "   Tensor b_tensors,"
+      "   Tensor a_scales,"
+      "   Tensor b_scales,"
+      "   Tensor b_group_scales,"
+      "   int b_group_size,"
+      "   Tensor expert_offsets,"
+      "   Tensor problem_sizes,"
+      "   Tensor a_strides,"
+      "   Tensor b_strides,"
+      "   Tensor c_strides,"
+      "   Tensor group_scale_strides,"
+      "   str? maybe_schedule"
+      ") -> ()");
+
+  ops.def(
+      "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, "
+      "Tensor)");
 #endif
 }
 
@@ -128,6 +224,18 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
            TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets));
   ops.impl("get_cutlass_batched_moe_mm_data",
            TORCH_BOX(&get_cutlass_batched_moe_mm_data));
+
+  // FP4/NVFP4 ops
+  ops.impl("cutlass_scaled_fp4_mm", TORCH_BOX(&cutlass_scaled_fp4_mm));
+  ops.impl("scaled_fp4_quant", TORCH_BOX(&scaled_fp4_quant_func));
+  ops.impl("scaled_fp4_quant.out", TORCH_BOX(&scaled_fp4_quant_out));
+  ops.impl("scaled_fp4_experts_quant", TORCH_BOX(&scaled_fp4_experts_quant));
+  ops.impl("silu_and_mul_scaled_fp4_experts_quant",
+           TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
+  ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
+
+  // W4A8 ops: impl registrations are in the source files
+  // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)
 #endif
 }
 
@@ -143,6 +251,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
            TORCH_BOX(&cutlass_group_gemm_supported));
   ops.impl("cutlass_scaled_mm_supports_block_fp8",
            TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8));
+  ops.impl("cutlass_scaled_mm_supports_fp4",
+           TORCH_BOX(&cutlass_scaled_mm_supports_fp4));
 #endif
 }
 

csrc/libtorch_stable/torch_utils.h

@@ -2,6 +2,7 @@
 
 #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>
 

csrc/ops.h

@@ -152,12 +152,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                         torch::Tensor& scale);
 
-#ifndef USE_ROCM
-void silu_and_mul_nvfp4_quant(torch::Tensor& out,
-                              torch::Tensor& output_block_scale,
-                              torch::Tensor& input,
-                              torch::Tensor& input_global_scale);
-#endif
 void persistent_masked_m_silu_mul_quant(
     const at::Tensor& input,   // (E, T, 2*H)
     const at::Tensor& counts,  // (E)
@@ -225,44 +219,6 @@ torch::Tensor ggml_moe_a8_vec(torch::Tensor X, torch::Tensor W,
 
 int64_t ggml_moe_get_block_size(int64_t type);
 
-#ifndef USE_ROCM
-
-bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability);
-
-void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
-                           torch::Tensor const& B, torch::Tensor const& A_sf,
-                           torch::Tensor const& B_sf,
-                           torch::Tensor const& alpha);
-
-void cutlass_fp4_group_mm(
-    torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
-    const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
-    const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
-    const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets);
-
-std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
-    torch::Tensor const& input, torch::Tensor const& input_scale,
-    bool is_sf_swizzled_layout);
-
-void scaled_fp4_quant_out(torch::Tensor const& input,
-                          torch::Tensor const& input_scale,
-                          bool is_sf_swizzled_layout, torch::Tensor& output,
-                          torch::Tensor& output_scale);
-
-void scaled_fp4_experts_quant(
-    torch::Tensor& output, torch::Tensor& output_scale,
-    torch::Tensor const& input, torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts);
-
-void silu_and_mul_scaled_fp4_experts_quant(
-    torch::Tensor& output, torch::Tensor& output_scale,
-    torch::Tensor const& input, torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts);
-
-#endif
-
 void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                               torch::Tensor const& scale,
                               std::optional<torch::Tensor> const& azp);

csrc/quantization/fp4/nvfp4_quant_entry.cu

@@ -1,163 +0,0 @@
-/*
- * Copyright (c) 2025, 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.
- */
-
-#include <torch/all.h>
-
-#include "cutlass_extensions/common.hpp"
-#include "nvfp4_utils.cuh"
-
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
-    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
-void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
-                             torch::Tensor const& input,
-                             torch::Tensor const& output_sf,
-                             torch::Tensor const& input_sf,
-                             bool is_sf_swizzled_layout);
-#endif
-
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
-    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
-void scaled_fp4_experts_quant_sm1xxa(
-    torch::Tensor& output, torch::Tensor& output_scale,
-    torch::Tensor const& input, torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts);
-#endif
-
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
-    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
-void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output,
-                                     torch::Tensor& output_sf,
-                                     torch::Tensor& input,
-                                     torch::Tensor& input_sf);
-#endif
-
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
-    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
-void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
-    torch::Tensor& output, torch::Tensor& output_scale,
-    torch::Tensor const& input, torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts);
-#endif
-
-static bool nvfp4_quant_sm_supported() {
-  const int32_t sm = get_sm_version_num();
-#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
-  if (sm >= 100 && sm < 120) return true;
-#endif
-#if defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120
-  if (sm >= 120 && sm < 130) return true;
-#endif
-  return false;
-}
-
-void scaled_fp4_quant_out(torch::Tensor const& input,
-                          torch::Tensor const& input_sf,
-                          bool is_sf_swizzled_layout, torch::Tensor& output,
-                          torch::Tensor& output_sf) {
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
-    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
-  TORCH_CHECK(nvfp4_quant_sm_supported(),
-              "No compiled nvfp4 quantization kernel for SM ",
-              get_sm_version_num(),
-              ". Recompile with the appropriate CUDA arch.");
-  return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
-                                 is_sf_swizzled_layout);
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
-}
-
-std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
-    torch::Tensor const& input, torch::Tensor const& input_sf,
-    bool is_sf_swizzled_layout) {
-  int64_t n = input.size(-1);
-  int64_t m = input.numel() / n;
-  auto device = input.device();
-
-  // Two fp4 values packed into a uint8
-  auto output = torch::empty(
-      {m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
-
-  torch::Tensor output_sf;
-  if (is_sf_swizzled_layout) {
-    auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
-    output_sf = torch::empty(
-        {sf_m, sf_n},
-        torch::TensorOptions().device(device).dtype(torch::kInt32));
-  } else {
-    output_sf = torch::empty(
-        {m, n / CVT_FP4_SF_VEC_SIZE},
-        torch::TensorOptions().device(device).dtype(torch::kUInt8));
-  }
-
-  scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
-                       output_sf);
-  return {output, output_sf};
-}
-
-void scaled_fp4_experts_quant(
-    torch::Tensor& output, torch::Tensor& output_scale,
-    torch::Tensor const& input, torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts) {
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
-    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
-  TORCH_CHECK(nvfp4_quant_sm_supported(),
-              "No compiled nvfp4 experts quantization kernel for SM ",
-              get_sm_version_num(),
-              ". Recompile with the appropriate CUDA arch.");
-  return scaled_fp4_experts_quant_sm1xxa(
-      output, output_scale, input, input_global_scale, input_offset_by_experts,
-      output_scale_offset_by_experts);
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(false,
-                              "No compiled nvfp4 experts quantization kernel");
-}
-
-void silu_and_mul_nvfp4_quant(torch::Tensor& output, torch::Tensor& output_sf,
-                              torch::Tensor& input, torch::Tensor& input_sf) {
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
-    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
-  TORCH_CHECK(nvfp4_quant_sm_supported(),
-              "No compiled silu_and_mul nvfp4 quantization kernel for SM ",
-              get_sm_version_num(),
-              ". Recompile with the appropriate CUDA arch.");
-  return silu_and_mul_nvfp4_quant_sm1xxa(output, output_sf, input, input_sf);
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(
-      false, "No compiled silu_and_mul nvfp4 quantization kernel");
-}
-
-void silu_and_mul_scaled_fp4_experts_quant(
-    torch::Tensor& output, torch::Tensor& output_scale,
-    torch::Tensor const& input, torch::Tensor const& input_global_scale,
-    torch::Tensor const& input_offset_by_experts,
-    torch::Tensor const& output_scale_offset_by_experts) {
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
-    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
-  TORCH_CHECK(nvfp4_quant_sm_supported(),
-              "No compiled silu_and_mul nvfp4 experts quantization kernel "
-              "for SM ",
-              get_sm_version_num(),
-              ". Recompile with the appropriate CUDA arch.");
-  return silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
-      output, output_scale, input, input_global_scale, input_offset_by_experts,
-      output_scale_offset_by_experts);
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(
-      false, "No compiled silu_and_mul nvfp4 experts quantization kernel");
-}

csrc/torch_bindings.cpp

@@ -109,13 +109,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
   ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
 
-#ifndef USE_ROCM
-  ops.def(
-      "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
-      "Tensor input, Tensor input_global_scale) -> ()");
-  ops.impl("silu_and_mul_nvfp4_quant", torch::kCUDA, &silu_and_mul_nvfp4_quant);
-#endif
-
   ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
   ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu);
 
@@ -332,47 +325,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "Tensor? qzeros_or_none, bool inplace) -> Tensor");
   // conditionally compiled so impl registrations are in source file
 
-  // CUTLASS w4a8 GEMM
-  ops.def(
-      "cutlass_w4a8_mm("
-      "   Tensor A,"
-      "   Tensor B,"
-      "   Tensor group_scales,"
-      "   int    group_size,"
-      "   Tensor channel_scales,"
-      "   Tensor token_scales,"
-      "   ScalarType? out_type,"
-      "   str?   maybe_schedule"
-      ") -> Tensor");
-  // pack scales
-  ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor");
-  // encode and reorder weight matrix
-  ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor");
-  // conditionally compiled so impl registration is in source file
-
-  // CUTLASS w4a8 grouped GEMM
-  ops.def(
-      "cutlass_w4a8_moe_mm("
-      "   Tensor! out_tensors,"
-      "   Tensor a_tensors,"
-      "   Tensor b_tensors,"
-      "   Tensor a_scales,"
-      "   Tensor b_scales,"
-      "   Tensor b_group_scales,"
-      "   int b_group_size,"
-      "   Tensor expert_offsets,"
-      "   Tensor problem_sizes,"
-      "   Tensor a_strides,"
-      "   Tensor b_strides,"
-      "   Tensor c_strides,"
-      "   Tensor group_scale_strides,"
-      "   str? maybe_schedule"
-      ") -> ()");
-  ops.def(
-      "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, "
-      "Tensor)");
-  // conditionally compiled so impl registration is in source file
-
 #endif
 
   // Dequantization for GGML.
@@ -409,20 +361,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("ggml_moe_get_block_size", &ggml_moe_get_block_size);
 
 #ifndef USE_ROCM
-  // CUTLASS nvfp4 block scaled GEMM
-  ops.def(
-      "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
-      "                      Tensor block_scale_a, Tensor block_scale_b,"
-      "                      Tensor alpha) -> ()");
-  ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);
-
-  // cutlass nvfp4 block scaled group GEMM
-  ops.def(
-      "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
-      " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
-      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
-  // conditionally compiled so impl registration is in source file
-
   // Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
   ops.def(
       "mxfp8_experts_quant("
@@ -455,44 +393,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "-> int");
   // conditionally compiled so impl in source file
 
-  // Compute NVFP4 block quantized tensor.
-  ops.def(
-      "scaled_fp4_quant(Tensor input,"
-      "                 Tensor input_scale, bool "
-      "is_sf_swizzled_layout) -> (Tensor, Tensor)");
-  ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant_func);
-
-  // Out variant
-  // TODO: Add {at::Tag::out_variant} tag and update all call sites
-  // to use the functional variant once vLLM upgrades PyTorch.
-  // See pytorch/pytorch#176117.
-  ops.def(
-      "scaled_fp4_quant.out(Tensor input,"
-      "                     Tensor input_scale, bool "
-      "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) "
-      "-> ()");
-  ops.impl("scaled_fp4_quant.out", torch::kCUDA, &scaled_fp4_quant_out);
-
-  // Compute NVFP4 experts quantization.
-  ops.def(
-      "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale,"
-      "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts) -> ()");
-  ops.impl("scaled_fp4_experts_quant", torch::kCUDA, &scaled_fp4_experts_quant);
-
-  // Fused SiLU+Mul+NVFP4 experts quantization.
-  ops.def(
-      "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! "
-      "output_scale,"
-      "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts) -> ()");
-  ops.impl("silu_and_mul_scaled_fp4_experts_quant", torch::kCUDA,
-           &silu_and_mul_scaled_fp4_experts_quant);
-
-  // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
-  // of the given capability
-  ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
-  ops.impl("cutlass_scaled_mm_supports_fp4", &cutlass_scaled_mm_supports_fp4);
 #endif
 
   // Quantized GEMM for GPTQ.

docs/design/cuda_graphs.md

@@ -225,7 +225,7 @@ outputs = model.generate(
 
 ### Piecewise compilation and full graph custom passes (attention fusion, sequence parallelism)
 
-Unfortunately, some custom compile passes have to see the whole graph to be effective and hence aren't compatible with piecewise compilation. This includes `AttnFusionPass` and `SequenceParallelismPass`. As a short-term solution, we automatically disable piecewise compilation (by setting `splitting_ops=[]`) when attention fusion is enabled. We use CUDA Graph modes `FULL` or `FULL_DECODE_ONLY` (depending on backend support). However, this leads to another optimization incompatibility and confusing performance tradeoffs.
+Unfortunately, some custom compile passes have to see the whole graph to be effective and hence aren't compatible with piecewise compilation. This includes `AttnQuantFusionPass` and `SequenceParallelismPass`. As a short-term solution, we automatically disable piecewise compilation (by setting `splitting_ops=[]`) when attention fusion is enabled. We use CUDA Graph modes `FULL` or `FULL_DECODE_ONLY` (depending on backend support). However, this leads to another optimization incompatibility and confusing performance tradeoffs.
 
 Long term, we've added the ability to partition the graph in Inductor instead of right after Dynamo. It can be enabled with `CompilationConfig.use_inductor_graph_partition=True` but is currently experimental and only available with `torch>=2.9`. This also increases compilation time as it has to compile the whole graph and cannot reuse piecewise compilation artifacts. Once vLLM supports 2.9, we plan to make this the default approach as it will also speed up piecewise cudagraph capture.
 

requirements/test.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu129 --python-platform x86_64-manylinux_2_28 --python-version 3.12
+#    uv pip compile requirements/test.in -c requirements/common.txt -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu129 --python-platform x86_64-manylinux_2_28 --python-version 3.12
 absl-py==2.1.0
     # via
     #   rouge-score
@@ -14,6 +14,7 @@ aiohappyeyeballs==2.6.1
     # via aiohttp
 aiohttp==3.13.3
     # via
+    #   -c requirements/common.txt
     #   aiohttp-cors
     #   datasets
     #   fsspec
@@ -225,7 +226,9 @@ et-xmlfile==2.0.0
 evaluate==0.4.3
     # via lm-eval
 fastapi==0.128.0
-    # via gpt-oss
+    # via
+    #   -c requirements/common.txt
+    #   gpt-oss
 fastparquet==2024.11.0
     # via genai-perf
 fastrlock==0.8.2
@@ -234,6 +237,7 @@ fastsafetensors==0.2.2
     # via -r requirements/test.in
 filelock==3.16.1
     # via
+    #   -c requirements/common.txt
     #   blobfile
     #   datasets
     #   diffusers
@@ -505,7 +509,9 @@ mbstrdecoder==1.1.3
 mdurl==0.1.2
     # via markdown-it-py
 mistral-common==1.10.0
-    # via -r requirements/test.in
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/test.in
 more-itertools==10.5.0
     # via lm-eval
 mpmath==1.3.0
@@ -655,13 +661,16 @@ omegaconf==2.3.0
 open-clip-torch==2.32.0
     # via -r requirements/test.in
 openai-harmony==0.0.4
-    # via gpt-oss
+    # via
+    #   -c requirements/common.txt
+    #   gpt-oss
 opencensus==0.11.4
     # via ray
 opencensus-context==0.1.3
     # via opencensus
 opencv-python-headless==4.13.0.90
     # via
+    #   -c requirements/common.txt
     #   -r requirements/test.in
     #   albucore
     #   albumentations
@@ -670,15 +679,17 @@ openpyxl==3.1.5
     # via -r requirements/test.in
 opentelemetry-api==1.35.0
     # via
+    #   -c requirements/common.txt
     #   opentelemetry-exporter-prometheus
     #   opentelemetry-sdk
     #   opentelemetry-semantic-conventions
 opentelemetry-exporter-prometheus==0.56b0
     # via ray
-opentelemetry-proto==1.36.0
+opentelemetry-proto==1.35.0
     # via ray
 opentelemetry-sdk==1.35.0
     # via
+    #   -c requirements/common.txt
     #   opentelemetry-exporter-prometheus
     #   ray
 opentelemetry-semantic-conventions==0.56b0
@@ -785,6 +796,7 @@ pqdm==0.2.0
     # via -r requirements/test.in
 prometheus-client==0.22.0
     # via
+    #   -c requirements/common.txt
     #   opentelemetry-exporter-prometheus
     #   ray
 propcache==0.2.0
@@ -793,8 +805,9 @@ propcache==0.2.0
     #   yarl
 proto-plus==1.26.1
     # via google-api-core
-protobuf==6.33.2
+protobuf==6.33.6
     # via
+    #   -c requirements/common.txt
     #   google-api-core
     #   googleapis-common-protos
     #   grpcio-reflection
@@ -836,6 +849,7 @@ pycryptodomex==3.22.0
     # via blobfile
 pydantic==2.12.0
     # via
+    #   -c requirements/common.txt
     #   -r requirements/test.in
     #   albumentations
     #   datamodel-code-generator
@@ -973,6 +987,7 @@ regex==2024.9.11
     #   transformers
 requests==2.32.3
     # via
+    #   -c requirements/common.txt
     #   azure-core
     #   buildkite-test-collector
     #   datasets
@@ -1085,6 +1100,7 @@ sentry-sdk==2.52.0
     # via wandb
 setuptools==77.0.3
     # via
+    #   -c requirements/common.txt
     #   lightning-utilities
     #   pytablewriter
     #   tensorboard
@@ -1099,6 +1115,7 @@ shellingham==1.5.4
     #   typer
 six==1.16.0
     # via
+    #   -c requirements/common.txt
     #   junit-xml
     #   lightly
     #   opencensus
@@ -1183,6 +1200,7 @@ tifffile==2025.3.30
     #   terratorch
 tiktoken==0.12.0
     # via
+    #   -c requirements/common.txt
     #   gpt-oss
     #   lm-eval
     #   mistral-common
@@ -1195,6 +1213,7 @@ timm==1.0.17
     #   torchgeo
 tokenizers==0.22.0
     # via
+    #   -c requirements/common.txt
     #   -r requirements/test.in
     #   transformers
 tomli==2.2.1
@@ -1271,6 +1290,7 @@ tqdm==4.67.3
     #   transformers
 transformers==4.57.5
     # via
+    #   -c requirements/common.txt
     #   -r requirements/test.in
     #   genai-perf
     #   peft
@@ -1297,6 +1317,7 @@ typeshed-client==2.8.2
     # via jsonargparse
 typing-extensions==4.15.0
     # via
+    #   -c requirements/common.txt
     #   aiosignal
     #   albumentations
     #   alembic

tests/compile/fusions_e2e/conftest.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import logging
+from collections import defaultdict
 
 import pytest
 import regex as re
@@ -52,6 +53,16 @@ def run_model(compile_config: int | CompilationConfig, model: str, **model_kwarg
         llm.llm_engine.vllm_config.compilation_config.compile_ranges_endpoints
     )
 
+    # Fetch match table from each worker via RPC and sum across workers.
+    worker_tables = llm.llm_engine.engine_core.collective_rpc(
+        "get_compilation_match_table"
+    )
+    combined: defaultdict[str, int] = defaultdict(int)
+    for table in worker_tables:
+        for k, v in table.items():
+            combined[k] += v
+    return dict(combined)
+
 
 @pytest.fixture
 def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
@@ -113,7 +124,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
         )
 
         with caplog_mp_spawn(logging.DEBUG) as log_holder:
-            run_model(full_compilation_config, model_name, **model_kwargs)
+            match_table = run_model(full_compilation_config, model_name, **model_kwargs)
 
         num_compile_ranges = len(full_compilation_config.get_compile_ranges())
         assert num_compile_ranges in [1, 2, 3]
@@ -155,11 +166,14 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
             else:
                 num_ranges_activated = num_compile_ranges
 
+            # TODO: Remove log counting in unit tests
+            # once all matchers implement VllmFusionPatternMatcherPass
             n_expected = tp_size * num_ranges_activated
-            assert len(log_matches) == n_expected, (
-                f"Could not find {n_expected} {match_name} "
-                f"(found {len(log_matches)}) in:\n {log_holder.text}"
-            )
+            if match_name != "attn_quant_fusion":
+                assert len(log_matches) == n_expected, (
+                    f"Could not find {n_expected} {match_name} "
+                    f"(found {len(log_matches)}) in:\n {log_holder.text}"
+                )
 
             expected_matches = getattr(matches, match_name)
 
@@ -215,6 +229,13 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
                     f"{tp_size * (num_ranges_activated - 1)} large-range "
                     f"entries (SP took precedence), found: {log_matches}"
                 )
+
+            elif match_name == "attn_quant_fusion":
+                actual_match = match_table.get(match_name, 0)
+                assert actual_match == expected_matches * n_expected, (
+                    f"Could not find {expected_matches * n_expected} "
+                    f"{match_name} (found {actual_match})."
+                )
             else:
                 expected_matches_list = [expected_matches] * n_expected
                 assert sorted(log_matches) == expected_matches_list, (

tests/compile/h100/test_startup.py

@@ -247,3 +247,92 @@ def test_model_startup(monkeypatch, vllm_runner, fresh_vllm_cache, spec):
 
     # Warm start — compiled artifacts loaded from disk cache.
     _check_model_run(vllm_runner, spec, is_cold_start=False)
+
+
+# ---------------------------------------------------------------------------
+# compile_model (compile-only) cold start tests
+# ---------------------------------------------------------------------------
+
+COMPILE_ONLY_SPECS = [
+    pytest.param(
+        ModelStartupSpec(
+            model="microsoft/Phi-tiny-MoE-instruct",
+            hf_overrides={},
+            cold_artifacts_saved=3,
+            warm_artifacts_saved=0,
+            warm_artifacts_loaded=3,
+        ),
+        id="phi_tiny_moe",
+    ),
+    pytest.param(
+        ModelStartupSpec(
+            model="openai/gpt-oss-120b",
+            hf_overrides={
+                "num_hidden_layers": 8,
+                "hidden_size": 256,
+                "intermediate_size": 512,
+                "num_attention_heads": 8,
+                "num_key_value_heads": 1,
+                "num_local_experts": 8,
+            },
+            cold_artifacts_saved=3,
+            warm_artifacts_saved=0,
+            warm_artifacts_loaded=3,
+        ),
+        id="gpt_oss_120b",
+    ),
+    pytest.param(
+        ModelStartupSpec(
+            model="zai-org/GLM-4.5",
+            hf_overrides=_SMALL_MOE_OVERRIDES,
+            cold_artifacts_saved=4,
+            warm_artifacts_saved=0,
+            warm_artifacts_loaded=4,
+        ),
+        id="glm_4.5",
+    ),
+]
+
+
+def _compile_only_cold_start(spec: ModelStartupSpec):
+    """Cold start using compile_model (fake weights, no GPU memory)."""
+    from vllm.compile_only import compile_model
+
+    old = compilation_counter.clone()
+    compile_model(
+        spec.model,
+        trust_remote_code=True,
+        max_model_len=256,
+        max_num_batched_tokens=1024,
+        block_size=64,
+        hf_overrides=spec.hf_overrides,
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            cudagraph_mode=CUDAGraphMode.NONE,
+            pass_config=PassConfig(fuse_allreduce_rms=False),
+        ),
+    )
+
+    saved = (
+        compilation_counter.num_compiled_artifacts_saved
+        - old.num_compiled_artifacts_saved
+    )
+    print(f"\n=== COMPILE-ONLY COLD START for {spec.model} ===")
+    print(f"  num_compiled_artifacts_saved={saved}")
+    assert saved == spec.cold_artifacts_saved, f"cold_artifacts_saved: got {saved}"
+
+
+@pytest.mark.parametrize("spec", COMPILE_ONLY_SPECS)
+@fork_new_process_for_each_test
+def test_compile_only_startup(monkeypatch, vllm_runner, fresh_vllm_cache, spec):
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+    # Cold start: compile-only in a forked child (fork before CUDA init).
+    ctx = mp.get_context("fork")
+    p = ctx.Process(target=_compile_only_cold_start, args=(spec,))
+    p.start()
+    p.join()
+    assert p.exitcode == 0, "Compile-only cold start failed"
+
+    # Warm start — compiled artifacts loaded from disk cache.
+    _check_model_run(vllm_runner, spec, is_cold_start=False)

tests/compile/passes/test_fusion_attn.py

@@ -9,7 +9,10 @@ from tests.compile.backend import LazyInitPass, TestBackend
 from tests.utils import TestFP8Layer, flat_product
 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
-from vllm.compilation.passes.fusion.attn_quant_fusion import ATTN_OP, AttnFusionPass
+from vllm.compilation.passes.fusion.attn_quant_fusion import (
+    ATTN_OP,
+    AttnQuantFusionPass,
+)
 from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS
 from vllm.compilation.passes.fx_utils import find_op_nodes
 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
@@ -384,7 +387,7 @@ def test_attention_quant_pattern(
 
         # Create test backend with fusion passes enabled
         noop_pass = NoOpEliminationPass(vllm_config)
-        attn_pass = LazyInitPass(AttnFusionPass, vllm_config)
+        attn_pass = LazyInitPass(AttnQuantFusionPass, vllm_config)
         cleanup_pass = PostCleanupPass(vllm_config)
 
         test_backend = TestBackend(noop_pass, attn_pass, cleanup_pass)
@@ -434,7 +437,7 @@ def test_attention_quant_pattern(
     # Only output quant ops are fused into attention.
     test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Dynamic)
 
-    # access the underlying `AttnFusionPass` on the `LazyInitPass`
+    # access the underlying `AttnQuantFusionPass` on the `LazyInitPass`
     assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)
 
     # Check attention ops in the graph before and after fusion

tests/kernels/moe/test_moe_weight_loading_padded.py

@@ -0,0 +1,292 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for FusedMoE weight loading with padded hidden dimensions.
+
+When using DeepEP backends or NIXL EP with models like nemotron_h,
+hidden_size may be rounded up (e.g., 2688 -> 3072) for backend requirements.
+Weight parameters are created with the padded size, but checkpoint weights
+have the original unpadded size. These tests verify that weight loading
+correctly handles this mismatch.
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+
+
+class TestGetHiddenDim:
+    """Unit tests for _get_hidden_dim."""
+
+    def test_2d_non_transposed_w2(self):
+        # w2: shard_dim=1 (intermediate), hidden=0
+        assert FusedMoE._get_hidden_dim(shard_dim=1, ndim=2) == 0
+
+    def test_2d_non_transposed_w13(self):
+        # w1/w3: shard_dim=0 (intermediate), hidden=1
+        assert FusedMoE._get_hidden_dim(shard_dim=0, ndim=2) == 1
+
+    def test_2d_transposed_w2(self):
+        # transposed w2: shard_dim=0, hidden=1
+        assert FusedMoE._get_hidden_dim(shard_dim=0, ndim=2) == 1
+
+    def test_2d_transposed_w13(self):
+        # transposed w1/w3: shard_dim=1, hidden=0
+        assert FusedMoE._get_hidden_dim(shard_dim=1, ndim=2) == 0
+
+    def test_3d_non_transposed_w2(self):
+        # 3D w2: shard_dim=2, hidden=1
+        assert FusedMoE._get_hidden_dim(shard_dim=2, ndim=3) == 1
+
+    def test_3d_non_transposed_w13(self):
+        # 3D w1/w3: shard_dim=1, hidden=2
+        assert FusedMoE._get_hidden_dim(shard_dim=1, ndim=3) == 2
+
+    def test_3d_transposed_w2(self):
+        # transposed 3D w2: shard_dim=1, hidden=2
+        assert FusedMoE._get_hidden_dim(shard_dim=1, ndim=3) == 2
+
+    def test_3d_transposed_w13(self):
+        # transposed 3D w1/w3: shard_dim=2, hidden=1
+        assert FusedMoE._get_hidden_dim(shard_dim=2, ndim=3) == 1
+
+    def test_1d_returns_zero(self):
+        # 1D per-channel scales: always returns 0
+        assert FusedMoE._get_hidden_dim(shard_dim=0, ndim=1) == 0
+        assert FusedMoE._get_hidden_dim(shard_dim=1, ndim=1) == 0
+
+    def test_invalid_shard_dim_raises(self):
+        # shard_dim outside the data dimensions should raise
+        with pytest.raises(ValueError, match="not a valid data dimension"):
+            FusedMoE._get_hidden_dim(shard_dim=0, ndim=3)
+
+
+class TestNarrowExpertDataForPadding:
+    """Unit tests for _narrow_expert_data_for_padding."""
+
+    def test_no_narrowing_when_shapes_match(self):
+        expert_data = torch.zeros(1024, 1024)
+        loaded_weight = torch.randn(1024, 1024)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=0
+        )
+        assert result.shape == loaded_weight.shape
+        assert result.data_ptr() == expert_data.data_ptr()
+
+    def test_narrow_w2_hidden_dim(self):
+        # w2: (hidden_size, intermediate_size) - hidden_size padded at dim 0
+        expert_data = torch.zeros(3072, 1024)
+        loaded_weight = torch.randn(2688, 1024)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=0
+        )
+        assert result.shape == (2688, 1024)
+
+    def test_narrow_w13_hidden_dim(self):
+        # w1/w3: (intermediate_size, hidden_size) - hidden_size padded at dim 1
+        expert_data = torch.zeros(2048, 3072)
+        loaded_weight = torch.randn(2048, 2688)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=1
+        )
+        assert result.shape == (2048, 2688)
+
+    def test_narrow_transposed_w2(self):
+        # transposed w2: (intermediate_size, hidden_size) - hidden at dim 1
+        expert_data = torch.zeros(1024, 3072)
+        loaded_weight = torch.randn(1024, 2688)
+        hidden_dim = FusedMoE._get_hidden_dim(shard_dim=0, ndim=2)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=hidden_dim
+        )
+        assert result.shape == (1024, 2688)
+
+    def test_narrow_3d_full_load(self):
+        # 3D tensor for full_load path: w2 (num_experts, hidden_size, intermediate)
+        expert_data = torch.zeros(8, 3072, 1024)
+        loaded_weight = torch.randn(8, 2688, 1024)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=1
+        )
+        assert result.shape == (8, 2688, 1024)
+
+    def test_narrow_1d_scale(self):
+        # 1D scale tensor: per-channel w2 scale (hidden_size,)
+        expert_data = torch.zeros(3072)
+        loaded_weight = torch.randn(2688)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=0
+        )
+        assert result.shape == (2688,)
+
+    def test_scalar_weight_no_op(self):
+        # 0-dim tensor should be a no-op
+        expert_data = torch.zeros(3072)
+        loaded_weight = torch.tensor(1.0)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=0
+        )
+        # ndim == 0, so no narrowing
+        assert result.shape == (3072,)
+
+    def test_no_narrowing_when_loaded_weight_larger(self):
+        # Guard: don't narrow if loaded_weight is larger than expert_data
+        expert_data = torch.zeros(2688, 1024)
+        loaded_weight = torch.randn(3072, 1024)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=0
+        )
+        assert result.shape == (2688, 1024)
+        assert result.data_ptr() == expert_data.data_ptr()
+
+    def test_negative_hidden_dim_is_noop(self):
+        # Negative hidden_dim should be a safe no-op (0 <= check)
+        expert_data = torch.zeros(3072, 1024)
+        loaded_weight = torch.randn(2688, 1024)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=-1
+        )
+        # -1 fails the 0 <= check, so no narrowing
+        assert result.shape == (3072, 1024)
+        assert result.data_ptr() == expert_data.data_ptr()
+
+    def test_only_narrows_hidden_dim(self):
+        # Verify that only the specified hidden_dim is narrowed,
+        # even when other dimensions also differ
+        expert_data = torch.zeros(3072, 2048)
+        loaded_weight = torch.randn(2688, 1024)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=0
+        )
+        # Only dim 0 (hidden) should be narrowed; dim 1 stays at 2048
+        assert result.shape == (2688, 2048)
+
+    def test_narrowed_data_shares_storage(self):
+        # Verify narrowing returns a view (writes go to original tensor)
+        expert_data = torch.zeros(3072, 1024)
+        loaded_weight = torch.randn(2688, 1024)
+        result = FusedMoE._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=0
+        )
+        result.copy_(loaded_weight)
+        # The first 2688 rows of expert_data should now have loaded_weight
+        assert torch.equal(expert_data[:2688, :], loaded_weight)
+        # Padded region should remain zero
+        assert torch.equal(expert_data[2688:, :], torch.zeros(3072 - 2688, 1024))
+
+
+class TestWeightLoadingWithPaddedHiddenSize:
+    """Integration-style tests that simulate padded weight loading."""
+
+    def test_load_w2_with_padding(self):
+        """Simulate loading w2 weights when hidden_size is padded."""
+        padded_hidden = 3072
+        original_hidden = 2688
+        intermediate = 1024
+
+        expert_data_full = torch.zeros(padded_hidden, intermediate)
+        loaded_weight = torch.randn(original_hidden, intermediate)
+
+        # w2 non-transposed: shard_dim=1, hidden_dim=0
+        hidden_dim = FusedMoE._get_hidden_dim(shard_dim=1, ndim=2)
+        expert_data = FusedMoE._narrow_expert_data_for_padding(
+            expert_data_full, loaded_weight, hidden_dim=hidden_dim
+        )
+        expert_data.copy_(loaded_weight)
+
+        assert torch.equal(expert_data_full[:original_hidden, :], loaded_weight)
+        assert torch.equal(
+            expert_data_full[original_hidden:, :],
+            torch.zeros(padded_hidden - original_hidden, intermediate),
+        )
+
+    def test_load_w13_with_padding(self):
+        """Simulate loading w1/w3 weights when hidden_size is padded."""
+        padded_hidden = 3072
+        original_hidden = 2688
+        intermediate = 1024
+
+        # w1/w3: (intermediate_size, hidden_size)
+        expert_data_full = torch.zeros(intermediate, padded_hidden)
+        loaded_weight = torch.randn(intermediate, original_hidden)
+
+        # w1 non-transposed: shard_dim=0, hidden_dim=1
+        hidden_dim = FusedMoE._get_hidden_dim(shard_dim=0, ndim=2)
+        expert_data = FusedMoE._narrow_expert_data_for_padding(
+            expert_data_full, loaded_weight, hidden_dim=hidden_dim
+        )
+        expert_data.copy_(loaded_weight)
+
+        assert torch.equal(expert_data_full[:, :original_hidden], loaded_weight)
+        assert torch.equal(
+            expert_data_full[:, original_hidden:],
+            torch.zeros(intermediate, padded_hidden - original_hidden),
+        )
+
+    def test_load_transposed_w2_with_padding(self):
+        """Simulate loading transposed w2 (GPTQ) with padded hidden_size."""
+        padded_hidden = 3072
+        original_hidden = 2688
+        intermediate = 1024
+
+        # transposed w2: (intermediate_size, hidden_size), shard_dim=0
+        expert_data_full = torch.zeros(intermediate, padded_hidden)
+        loaded_weight = torch.randn(intermediate, original_hidden)
+
+        hidden_dim = FusedMoE._get_hidden_dim(shard_dim=0, ndim=2)
+        expert_data = FusedMoE._narrow_expert_data_for_padding(
+            expert_data_full, loaded_weight, hidden_dim=hidden_dim
+        )
+        expert_data.copy_(loaded_weight)
+
+        assert torch.equal(expert_data_full[:, :original_hidden], loaded_weight)
+
+    def test_no_padding_is_noop(self):
+        """Verify that when sizes match, behavior is unchanged."""
+        hidden = 2048
+        intermediate = 1024
+
+        expert_data_full = torch.zeros(hidden, intermediate)
+        loaded_weight = torch.randn(hidden, intermediate)
+
+        hidden_dim = FusedMoE._get_hidden_dim(shard_dim=1, ndim=2)
+        expert_data = FusedMoE._narrow_expert_data_for_padding(
+            expert_data_full, loaded_weight, hidden_dim=hidden_dim
+        )
+        expert_data.copy_(loaded_weight)
+
+        assert torch.equal(expert_data_full, loaded_weight)
+
+    def test_bnb_shape_mismatch_raises(self):
+        """BnB + padded hidden_size should raise via weight_loader."""
+        from unittest.mock import MagicMock
+
+        num_experts = 1
+        padded_packed = 3072  # padded packed size
+        original_packed = 2688  # original packed size
+
+        # Build a param that looks like a BnB 4-bit MoE weight.
+        param_data = torch.zeros(num_experts, padded_packed, 1, dtype=torch.uint8)
+        param = torch.nn.Parameter(param_data, requires_grad=False)
+        param.use_bitsandbytes_4bit = True
+
+        loaded_weight = torch.randint(0, 255, (original_packed, 1), dtype=torch.uint8)
+
+        # Minimal FusedMoE mock so weight_loader reaches the BnB path.
+        moe = MagicMock(spec=FusedMoE)
+        moe.quant_config = None
+        moe.quant_method = MagicMock()
+        moe.quant_method.__class__.__name__ = "BitsAndBytesMethod"
+        moe._expert_map = None
+        moe.tp_rank = 0
+
+        # Call the real weight_loader (unbound) with our mock as self.
+        with pytest.raises(ValueError, match="BitsAndBytes"):
+            FusedMoE.weight_loader(
+                moe,
+                param,
+                loaded_weight,
+                weight_name="w2",
+                shard_id="w2",
+                expert_id=0,
+            )

vllm/__init__.py

@@ -14,6 +14,7 @@ import typing
 import vllm.env_override  # noqa: F401
 
 MODULE_ATTRS = {
+    "compile_model": ".compile_only:compile_model",
     "AsyncEngineArgs": ".engine.arg_utils:AsyncEngineArgs",
     "EngineArgs": ".engine.arg_utils:EngineArgs",
     "AsyncLLMEngine": ".engine.async_llm_engine:AsyncLLMEngine",
@@ -39,6 +40,7 @@ MODULE_ATTRS = {
 }
 
 if typing.TYPE_CHECKING:
+    from vllm.compile_only import compile_model as compile_model
     from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
     from vllm.engine.async_llm_engine import AsyncLLMEngine
     from vllm.engine.llm_engine import LLMEngine

vllm/compilation/backends.py

@@ -50,6 +50,16 @@ from .passes.pass_manager import PostGradPassManager
 logger = init_logger(__name__)
 
 
+class CompilationDone(Exception):
+    """Raised in compile-only mode after compilation is complete.
+
+    This signals that the vLLM-compile cache has been populated and
+    there is no need to actually execute the compiled code.
+    """
+
+    pass
+
+
 def make_copy_and_call(
     sym_tensor_indices: list[int],
     input_buffers: list[torch.Tensor | None],
@@ -990,7 +1000,14 @@ class VllmBackend:
         # Compute config/compiler/code hashes once and reuse
         config_hash = vllm_config.compute_hash()
         compiler_hash = self.compiler_manager.compute_hash(vllm_config)
-        forward_code_files = list(sorted(self.compilation_config.traced_files))
+        # Filter out PyTorch internal files — they are already covered
+        # by the torch version in env_factors.
+        torch_root = os.path.dirname(torch.__file__) + os.sep
+        forward_code_files = [
+            f
+            for f in sorted(self.compilation_config.traced_files)
+            if not f.startswith(torch_root)
+        ]
 
         logger.debug(
             "Traced files (to be considered for compilation cache):\n%s",

vllm/compilation/decorators.py

@@ -599,9 +599,23 @@ def _support_torch_compile(
                     # AOT artifact.
                     self.save_aot_compiled_function()
 
+                # In compile-only mode, raise CompilationDone after all
+                # piecewise graphs are compiled and cache artifacts saved.
+                # This is caught in gpu_worker.compile_or_warm_up_model()
+                # to skip execution with fake tensors.
+                if self.compilation_config.compile_only:
+                    from .backends import CompilationDone
+
+                    raise CompilationDone
+
                 with monitor_profiling_run():
                     output = self.aot_compiled_fn(self, *args, **kwargs)
             else:
+                # Same as above for non-AOT path.
+                if self.compilation_config.compile_only:
+                    from .backends import CompilationDone
+
+                    raise CompilationDone
                 with monitor_torch_compile(
                     self.vllm_config,
                     "torch.compile and initial profiling/warmup "

vllm/compilation/passes/fusion/attn_quant_fusion.py

@@ -1,15 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from abc import ABC, abstractmethod
+
 from collections.abc import Callable
-from typing import Any, ParamSpec
 
 import torch
-import torch._inductor.pattern_matcher as pm
-from torch import fx
 from torch._higher_order_ops.auto_functionalize import auto_functionalized
-from torch._inductor.pattern_matcher import PatternMatcherPass
 
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.logger import init_logger
@@ -22,14 +18,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import round_up
 
-from ..fx_utils import is_func
-from ..inductor_pass import enable_fake_mode
-from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
+from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement
 from .matcher_utils import MatcherQuantFP8
-from .rms_quant_fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
+from .rms_quant_fusion import QUANT_OPS
 
 logger = init_logger(__name__)
-P = ParamSpec("P")
+
 FP8_DTYPE = current_platform.fp8_dtype()
 FP4_DTYPE = torch.uint8
 
@@ -37,83 +31,10 @@ ATTN_OP = torch.ops.vllm.unified_attention_with_output.default
 RESHAPE_OP = torch.ops.aten.reshape.default
 
 
-class AttentionQuantPattern(ABC):
-    """
-    The base class for Attn+Quant fusions.
-    Should not be used directly.
-    """
-
-    def __init__(
-        self,
-        layer: Attention,
-        quant_key: QuantKey,
-        dtype: torch.dtype,
-    ) -> None:
-        self.layer = layer
-        self.layer_name = layer.layer_name
-        self.num_heads = layer.num_heads
-        self.head_size = layer.head_size
-        self.quant_key = quant_key
-        self.quant_dtype = quant_key.dtype
-        self.dtype = dtype
-
-        assert self.quant_key in QUANT_OPS, (
-            f"unsupported quantization scheme {self.quant_key}"
-        )
-        self.QUANT_OP = QUANT_OPS[self.quant_key]
-
-    def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor:
-        kwargs = {"dtype": self.dtype, "device": "cuda", **kwargs}
-        return torch.empty(*args, **kwargs)
-
-    def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor:
-        kwargs = {"dtype": self.quant_dtype, "device": "cuda", **kwargs}
-        return torch.empty(*args, **kwargs)
-
-    @staticmethod
-    def wrap_trace_fn(
-        trace_fn: Callable[P, fx.GraphModule],
-        *process_fx_fns: Callable[[fx.GraphModule], None],
-    ) -> Callable[P, fx.GraphModule]:
-        def wrapped(*args: P.args, **kwargs: P.kwargs) -> fx.GraphModule:
-            gm = trace_fn(*args, **kwargs)
-            for process_fx in process_fx_fns:
-                process_fx(gm)
-
-            return gm
-
-        return wrapped
-
-    @staticmethod
-    def fx_view_to_reshape(gm: torch.fx.GraphModule) -> None:
-        from torch._inductor.fx_passes.post_grad import view_to_reshape
-
-        view_to_reshape(gm)
-
-    @staticmethod
-    def remove_noop_permutes(gm: torch.fx.GraphModule) -> None:
-        for node in gm.graph.nodes:
-            if not is_func(node, torch.ops.aten.permute.default):
-                continue
-
-            dims = node.args[1]
-            if any(dim != i for i, dim in enumerate(dims)):
-                continue
-
-            # this is now an identity op, remove
-            node.replace_all_uses_with(node.args[0])
-            gm.graph.erase_node(node)
-
-    def register_if_supported(self, pm_pass: PatternMatcherPass) -> None:
-        if self.layer.impl.fused_output_quant_supported(self.quant_key):
-            self._register(pm_pass)
-
-    @abstractmethod
-    def _register(self, pm_pass: PatternMatcherPass) -> None:
-        raise NotImplementedError
+_FP8_QUANT_KEY = QuantKey(dtype=FP8_DTYPE, scale=kStaticTensorScale, symmetric=True)
 
 
-class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
+class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
     """
     Fusion for Attention+Fp8StaticQuant.
 
@@ -123,20 +44,16 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
     will be passed into Attention op as the `output_scale` argument.
     """
 
-    def __init__(
-        self,
-        layer: Attention,
-        dtype: torch.dtype,
-        symmetric: bool = True,
-    ) -> None:
-        quant_key = QuantKey(
-            dtype=FP8_DTYPE, scale=kStaticTensorScale, symmetric=symmetric
-        )
-        super().__init__(layer, quant_key, dtype)
-        self.quant_matcher = MatcherQuantFP8(quant_key)
+    def __init__(self, layer: Attention, dtype: torch.dtype):
+        self._layer_name = layer.layer_name
+        self._num_heads = layer.num_heads
+        self._head_size = layer.head_size
+        self._dtype = dtype
+        self._quant_matcher = MatcherQuantFP8(_FP8_QUANT_KEY)
 
-    def _register(self, pm_pass: PatternMatcherPass) -> None:
-        def pattern(
+    @property
+    def pattern(self) -> Callable[..., torch.Tensor]:
+        def _pattern(
             q: torch.Tensor,
             k: torch.Tensor,
             v: torch.Tensor,
@@ -150,18 +67,21 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
                 key=k,
                 value=v,
                 output=output_attn,
-                layer_name=self.layer_name,
+                layer_name=self._layer_name,
                 output_scale=None,
                 output_block_scale=None,
                 kv_cache_dummy_dep=kv_cache_dummy_dep,
             )
             attn_out_view = RESHAPE_OP(
-                at1[1], [q.shape[0], self.num_heads * self.head_size]
+                at1[1], [q.shape[0], self._num_heads * self._head_size]
             )
+            return self._quant_matcher(attn_out_view, scale)[0]
 
-            return self.quant_matcher(attn_out_view, scale)[0]
+        return _pattern
 
-        def replacement(
+    @property
+    def replacement(self) -> Callable[..., torch.Tensor]:
+        def _replacement(
             q: torch.Tensor,
             k: torch.Tensor,
             v: torch.Tensor,
@@ -169,10 +89,9 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
             scale: torch.Tensor,
             kv_cache_dummy_dep: torch.Tensor,
         ) -> torch.Tensor:
-            # attn output in quant_dtype
             output_attn = torch.empty(
-                [q.shape[0], self.num_heads, self.head_size],
-                dtype=self.quant_dtype,
+                [q.shape[0], self._num_heads, self._head_size],
+                dtype=FP8_DTYPE,
                 device=q.device,
             )
             at1 = auto_functionalized(
@@ -181,36 +100,32 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
                 key=k,
                 value=v,
                 output=output_attn,
-                layer_name=self.layer_name,
+                layer_name=self._layer_name,
                 output_scale=scale,
                 output_block_scale=None,
                 kv_cache_dummy_dep=kv_cache_dummy_dep,
             )
-            return RESHAPE_OP(at1[1], [-1, self.num_heads * self.head_size])
+            return RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size])
 
-        inputs = [
-            self.empty(5, self.num_heads, self.head_size),  # q
-            self.empty(5, self.num_heads, self.head_size),  # k
-            self.empty(5, self.num_heads, self.head_size),  # v
-            self.empty(5, self.num_heads, self.head_size),  # attn_output
-            empty_fp32(1, 1),  # scale
-            self.empty(0),  # kv_cache_dummy_dep
+        return _replacement
+
+    def get_inputs(self):
+        dtype = self._dtype
+        num_heads = self._num_heads
+        head_size = self._head_size
+        return [
+            self.empty(5, num_heads, head_size, dtype=dtype),  # q
+            self.empty(5, num_heads, head_size, dtype=dtype),  # k
+            self.empty(5, num_heads, head_size, dtype=dtype),  # v
+            self.empty(5, num_heads, head_size, dtype=dtype),  # attn_output
+            self.empty_fp32(1, 1),  # scale
+            self.empty(0, dtype=dtype),  # kv_cache_dummy_dep
         ]
 
-        pm.register_replacement(
-            pattern,
-            replacement,
-            inputs,
-            AttentionQuantPattern.wrap_trace_fn(
-                pm.fwd_only,
-                AttentionQuantPattern.fx_view_to_reshape,
-                AttentionQuantPattern.remove_noop_permutes,
-            ),
-            pm_pass,
-        )
 
-
-class AttentionNvfp4QuantPattern(AttentionQuantPattern):
+class AttnNvfp4QuantPattern(
+    VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]]
+):
     """
     Fusion for Attention+Nvfp4Quant.
 
@@ -220,11 +135,16 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
     will be passed into Attention op as the `output_scale` argument.
     """
 
-    def __init__(self, layer: Attention, dtype: torch.dtype) -> None:
-        super().__init__(layer, kNvfp4Dynamic, dtype)
+    def __init__(self, layer: Attention, dtype: torch.dtype):
+        self._layer_name = layer.layer_name
+        self._num_heads = layer.num_heads
+        self._head_size = layer.head_size
+        self._dtype = dtype
+        self._QUANT_OP = QUANT_OPS[kNvfp4Dynamic]
 
-    def _register(self, pm_pass: PatternMatcherPass) -> None:
-        def pattern(
+    @property
+    def pattern(self) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
+        def _pattern(
             q: torch.Tensor,
             k: torch.Tensor,
             v: torch.Tensor,
@@ -240,16 +160,16 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
                 key=k,
                 value=v,
                 output=output_attn,
-                layer_name=self.layer_name,
+                layer_name=self._layer_name,
                 output_scale=None,
                 output_block_scale=None,
                 kv_cache_dummy_dep=kv_cache_dummy_dep,
             )
             attn_out_view = RESHAPE_OP(
-                at1[1], [q.shape[0], self.num_heads * self.head_size]
+                at1[1], [q.shape[0], self._num_heads * self._head_size]
             )
             at2 = auto_functionalized(
-                self.QUANT_OP,
+                self._QUANT_OP,
                 input=attn_out_view,
                 input_scale=input_scale,
                 is_sf_swizzled_layout=True,
@@ -259,23 +179,25 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
             output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
             return at2[1], output_scale_view
 
-        def replacement(
+        return _pattern
+
+    @property
+    def replacement(self) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
+        def _replacement(
             q: torch.Tensor,
             k: torch.Tensor,
             v: torch.Tensor,
             output_attn: torch.Tensor,
-            output_quant: torch.Tensor,
+            _output_quant: torch.Tensor,
             output_scale: torch.Tensor,
             input_scale: torch.Tensor,
             kv_cache_dummy_dep: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor]:
-            # attention output in quant_dtype
             output_attn = torch.empty(
-                [q.shape[0], self.num_heads, self.head_size // 2],
-                dtype=self.quant_dtype,
+                [q.shape[0], self._num_heads, self._head_size // 2],
+                dtype=FP4_DTYPE,
                 device=q.device,
             )
-            # attention output block scale
             output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE)
             at2 = auto_functionalized(
                 ATTN_OP,
@@ -283,41 +205,35 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
                 key=k,
                 value=v,
                 output=output_attn,
-                layer_name=self.layer_name,
+                layer_name=self._layer_name,
                 output_scale=input_scale,
                 output_block_scale=output_scale_view,
                 kv_cache_dummy_dep=kv_cache_dummy_dep,
             )
-            output = RESHAPE_OP(at2[1], [-1, self.num_heads * self.head_size // 2])
+            output = RESHAPE_OP(at2[1], [-1, self._num_heads * self._head_size // 2])
             return output, at2[2]
 
-        inputs = [
-            empty_bf16(5, self.num_heads, self.head_size),  # q
-            empty_bf16(5, self.num_heads, self.head_size),  # k
-            empty_bf16(5, self.num_heads, self.head_size),  # v
-            empty_bf16(5, self.num_heads, self.head_size),  # output_attn
-            self.empty_quant(5, self.num_heads * self.head_size // 2),  # output_quant
-            empty_i32(
-                128, round_up(self.num_heads * self.head_size // 16, 4)
+        return _replacement
+
+    def get_inputs(self):
+        dtype = self._dtype
+        num_heads = self._num_heads
+        head_size = self._head_size
+        return [
+            self.empty_bf16(5, num_heads, head_size),  # q
+            self.empty_bf16(5, num_heads, head_size),  # k
+            self.empty_bf16(5, num_heads, head_size),  # v
+            self.empty_bf16(5, num_heads, head_size),  # output_attn
+            self.empty(5, num_heads * head_size // 2, dtype=FP4_DTYPE),  # output_quant
+            self.empty_i32(
+                128, round_up(num_heads * head_size // 16, 4)
             ),  # output_scale
-            empty_fp32(1, 1),  # input_scale
-            self.empty(0),  # kv_cache_dummy_dep
+            self.empty_fp32(1, 1),  # input_scale
+            self.empty(0, dtype=dtype),  # kv_cache_dummy_dep
         ]
 
-        pm.register_replacement(
-            pattern,
-            replacement,
-            inputs,
-            AttentionQuantPattern.wrap_trace_fn(
-                pm.fwd_only,
-                AttentionQuantPattern.fx_view_to_reshape,
-                AttentionQuantPattern.remove_noop_permutes,
-            ),
-            pm_pass,
-        )
 
-
-class AttnFusionPass(VllmPatternMatcherPass):
+class AttnQuantFusionPass(VllmFusionPatternMatcherPass):
     """
     This pass fuses post-attention quantization onto attention if supported.
 
@@ -330,43 +246,26 @@ class AttnFusionPass(VllmPatternMatcherPass):
     support are attention kernels, which need to support fusing output quant.
     """
 
-    @enable_fake_mode
     def __init__(self, config: VllmConfig) -> None:
-        super().__init__(config)
+        super().__init__(config, "attn_quant_fusion")
 
-        self.patterns = PatternMatcherPass(pass_name="attn_fusion_pass")
+        dtype = config.model_config.dtype
+        layers = list(get_layers_from_vllm_config(config, Attention).values())
 
-        attn_layers = get_layers_from_vllm_config(config, Attention)
-        for layer_name, layer in attn_layers.items():
-            pattern_fp8 = AttentionFp8StaticQuantPattern(
-                layer, config.model_config.dtype
-            )
-            pattern_fp8.register_if_supported(self.patterns)
-
-            if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
-                pattern_nvfp4 = AttentionNvfp4QuantPattern(
-                    layer, config.model_config.dtype
-                )
-                pattern_nvfp4.register_if_supported(self.patterns)
-
-        if len(attn_layers) == 0:
+        if len(layers) == 0:
             logger.warning(
                 "Attention + quant fusion is enabled, but no attention layers "
                 "were found in CompilationConfig.static_forward_context "
                 "so no fusion patterns were registered."
             )
 
-        self.dump_patterns(config, self.patterns)
+        for layer in layers:
+            if layer.impl.fused_output_quant_supported(_FP8_QUANT_KEY):
+                self.register(AttnFp8StaticQuantPattern(layer, dtype))
 
-    @VllmInductorPass.time_and_log
-    def __call__(self, graph: torch.fx.graph.Graph) -> None:
-        self.matched_count = self.patterns.apply(graph)
-        logger.debug("Fused quant onto %s attention nodes", self.matched_count)
+        if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
+            for layer in layers:
+                if layer.impl.fused_output_quant_supported(kNvfp4Dynamic):
+                    self.register(AttnNvfp4QuantPattern(layer, dtype))
 
-    def uuid(self) -> str:
-        return VllmInductorPass.hash_source(
-            self,
-            AttentionQuantPattern,
-            AttentionFp8StaticQuantPattern,
-            AttentionNvfp4QuantPattern,
-        )
+        self.dump_patterns(config, self.pm_pass)

vllm/compilation/passes/pass_manager.py

@@ -14,7 +14,7 @@ from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils.system_utils import set_env_var
 
-from .vllm_inductor_pass import VllmInductorPass
+from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
 
 if rocm_aiter_ops.is_enabled():
     from .fusion.rocm_aiter_fusion import (
@@ -25,7 +25,7 @@ if rocm_aiter_ops.is_enabled():
 
 if current_platform.is_cuda_alike():
     from .fusion.act_quant_fusion import ActivationQuantFusionPass
-    from .fusion.attn_quant_fusion import AttnFusionPass
+    from .fusion.attn_quant_fusion import AttnQuantFusionPass
     from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
     from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
     from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
@@ -108,6 +108,8 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
         self.fix_functionalization(graph)
         VllmInductorPass.dump_prefix = None  # Cleanup index
 
+        VllmPatternMatcherPass.log_match_summary()
+
     def configure(self, config: VllmConfig) -> None:
         self.pass_config = config.compilation_config.pass_config
 
@@ -144,7 +146,7 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
                 self.passes += [RopeKVCacheFusionPass(config)]
 
             if self.pass_config.fuse_attn_quant:
-                self.passes += [AttnFusionPass(config)]
+                self.passes += [AttnQuantFusionPass(config)]
 
             if self.pass_config.enable_qk_norm_rope_fusion:
                 self.passes += [SplitCoalescingPass(config)]

vllm/compilation/passes/vllm_inductor_pass.py

@@ -3,19 +3,24 @@
 import functools
 import operator
 import time
+from abc import ABC, abstractmethod
+from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import ClassVar
+from typing import Any, ClassVar, Generic, ParamSpec, TypeVar
 
 import regex as re
 import torch
+import torch._inductor.pattern_matcher as pm
+from torch import fx
 from torch._dynamo.utils import lazy_format_graph_code
 from torch._inductor.pattern_matcher import PatternMatcherPass, PatternPrettyPrinter
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 
-from .inductor_pass import InductorPass
+from .fx_utils import is_func
+from .inductor_pass import InductorPass, enable_fake_mode
 
 logger = init_logger(__name__)
 
@@ -79,18 +84,23 @@ class VllmInductorPass(InductorPass):
         logger.debug("%s completed in %.1f ms", self.pass_name, duration_ms)
 
 
+def get_match_table() -> dict[str, int]:
+    """Return a snapshot of the match table."""
+    return dict(VllmPatternMatcherPass.match_table)
+
+
 class VllmPatternMatcherPass(VllmInductorPass):
     """
     A VllmInductorPass that uses the Inductor pattern matcher.
-    Its main use is providing the dump_patterns utility that dumps the
-    Inductor pattern matcher patterns into a file, which greatly aids debugging.
-
-    TODO(luka) move more utilities to this pass.
+    Provides pattern registration with match counting, debug dumping, and logging.
     """
 
     matched_count: int = 0
     """The number of matched patterns in the pass."""
 
+    match_table: ClassVar[defaultdict[str, int]] = defaultdict(int)
+    """Global table mapping pass name to its total match count."""
+
     _OP_OVERLOAD_PATTERN: ClassVar[re.Pattern] = re.compile(
         r"<OpOverload\(op='([^']*)', overload='([^']*)'\)>"
     )
@@ -104,6 +114,11 @@ class VllmPatternMatcherPass(VllmInductorPass):
             )
         )
 
+    @classmethod
+    def log_match_summary(cls) -> None:
+        if cls.match_table:
+            logger.debug("fusion pass matches: %s", dict(cls.match_table))
+
     def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass) -> None:
         """
         If debug dumping is enabled, dump the Inductor pattern-matcher patterns
@@ -171,6 +186,124 @@ class VllmPatternMatcherPass(VllmInductorPass):
                     print(f"{pattern_repr}\n", file=f)
 
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class VllmPatternReplacement(ABC, Generic[P, R]):
+    """
+    A pattern/replacement pair for FX graph fusion.
+
+    Implement the three abstract members below, then pass
+    instances to VllmFusionPatternMatcherPass.register(). The pass will
+    find every occurrence of `pattern` in the graph and substitute it
+    with `replacement`.
+    """
+
+    # TODO(Badr): bound methods work for pattern registration since
+    # PyTorch 2.10. Once vLLM requires torch>=2.11, replace these properties
+    # with plain methods and drop the closure indirection.
+    @property
+    @abstractmethod
+    def pattern(self) -> Callable[P, R]:
+        """Returns a closure defining the FX subgraph to search for."""
+        ...
+
+    @property
+    @abstractmethod
+    def replacement(self) -> Callable[P, R]:
+        """
+        Returns a closure defining the FX subgraph to
+        substitute in place of each match.
+        """
+        ...
+
+    @abstractmethod
+    def get_inputs(self) -> list[torch.Tensor]:
+        """Example tensors used to trace pattern and replacement."""
+        ...
+
+    # Helpers for get_inputs: uninitialized tensors of common dtypes.
+    @staticmethod
+    def empty(*args, **kwargs) -> torch.Tensor:
+        return torch.empty(*args, device="cuda", **kwargs)
+
+    @staticmethod
+    def empty_bf16(*args, **kwargs) -> torch.Tensor:
+        return torch.empty(*args, dtype=torch.bfloat16, device="cuda", **kwargs)
+
+    @staticmethod
+    def empty_fp16(*args, **kwargs) -> torch.Tensor:
+        return torch.empty(*args, dtype=torch.float16, device="cuda", **kwargs)
+
+    @staticmethod
+    def empty_fp32(*args, **kwargs) -> torch.Tensor:
+        return torch.empty(*args, dtype=torch.float32, device="cuda", **kwargs)
+
+    @staticmethod
+    def empty_i32(*args, **kwargs) -> torch.Tensor:
+        return torch.empty(*args, dtype=torch.int32, device="cuda", **kwargs)
+
+
+def _fx_view_to_reshape(gm: fx.GraphModule) -> None:
+    from torch._inductor.fx_passes.post_grad import view_to_reshape
+
+    view_to_reshape(gm)
+
+
+def _remove_noop_permutes(gm: fx.GraphModule) -> None:
+    for node in gm.graph.nodes:
+        if not is_func(node, torch.ops.aten.permute.default):
+            continue
+        dims = node.args[1]
+        if any(dim != i for i, dim in enumerate(dims)):
+            continue
+        node.replace_all_uses_with(node.args[0])
+        gm.graph.erase_node(node)
+
+
+class VllmFusionPatternMatcherPass(VllmPatternMatcherPass):
+    """
+    A VllmPatternMatcherPass for passes that use VllmPatternReplacement objects.
+    Subclasses register patterns via self.register() in their own __init__.
+    """
+
+    def __init__(self, config: VllmConfig, pass_name: str) -> None:
+        super().__init__(config)
+        self.pass_name = pass_name
+        self.pm_pass = PatternMatcherPass(pass_name=pass_name)
+        self._pattern_replacements: list[VllmPatternReplacement] = []
+
+    @enable_fake_mode
+    def register(self, pr: VllmPatternReplacement) -> None:
+        pm.register_replacement(
+            pr.pattern,
+            pr.replacement,
+            pr.get_inputs(),
+            self._trace_fn,
+            self.pm_pass,
+        )
+        self._pattern_replacements.append(pr)
+
+    def uuid(self) -> str:
+        return VllmInductorPass.hash_source(
+            type(self),
+            *[type(pr) for pr in self._pattern_replacements],
+        )
+
+    @staticmethod
+    def _trace_fn(*args: Any, **kwargs: Any) -> fx.GraphModule:
+        gm = pm.fwd_only(*args, **kwargs)
+        _fx_view_to_reshape(gm)
+        _remove_noop_permutes(gm)
+        return gm
+
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: torch.fx.Graph) -> None:
+        self.matched_count = self.pm_pass.apply(graph)
+        VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count
+
+
 class PrinterInductorPass(VllmInductorPass):
     def __init__(self, name: str, config: VllmConfig) -> None:
         super().__init__(config)

vllm/compile_only.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Compile-only mode: populate vLLM's torch.compile cache
+without loading real model weights or allocating KV caches.
+
+The compile-only flag causes the model loader to be wrapped with
+FakeTensorMode (see ``fake_loader.wrap_loader_with_fake``), so the
+user's original ``load_format`` is preserved and the real loader's
+full pipeline runs — just with fake tensors instead of real weights.
+"""
+
+import argparse
+
+from vllm.logger import init_logger
+from vllm.usage.usage_lib import UsageContext
+
+logger = init_logger(__name__)
+
+
+def compile_model(
+    model: str,
+    *,
+    tensor_parallel_size: int = 1,
+    pipeline_parallel_size: int = 1,
+    quantization: str | None = None,
+    dtype: str = "auto",
+    trust_remote_code: bool = False,
+    **kwargs,
+) -> None:
+    """Pre-populate vLLM's torch.compile cache for a model.
+
+    Runs compilation using fake weights (zero GPU memory)
+    so that vLLM's torch.compile cache is populated.  Subsequent
+    ``vllm serve`` or ``LLM(...)`` calls for the same model
+    configuration will hit the warm cache and skip compilation.
+
+    Args:
+        model: HuggingFace model name or path.
+        tensor_parallel_size: Number of tensor parallel GPUs.
+        pipeline_parallel_size: Number of pipeline parallel stages.
+        quantization: Quantization method (e.g. "fp8").
+        dtype: Model dtype.
+        trust_remote_code: Trust remote code from HuggingFace.
+        **kwargs: Additional arguments passed to ``EngineArgs``.
+    """
+    from vllm.engine.arg_utils import EngineArgs
+
+    engine_args = EngineArgs(
+        model=model,
+        tensor_parallel_size=tensor_parallel_size,
+        pipeline_parallel_size=pipeline_parallel_size,
+        quantization=quantization,
+        dtype=dtype,
+        trust_remote_code=trust_remote_code,
+        enforce_eager=False,
+        **kwargs,
+    )
+
+    vllm_config = engine_args.create_engine_config(usage_context=UsageContext.LLM_CLASS)
+    vllm_config.compilation_config.compile_only = True
+
+    _run_compile_with_config(vllm_config)
+
+
+def run_compile_only(args: argparse.Namespace) -> None:
+    """Run compile-only mode from CLI arguments."""
+    from vllm.engine.arg_utils import EngineArgs
+
+    engine_args = EngineArgs.from_cli_args(args)
+    engine_args.enforce_eager = False
+
+    vllm_config = engine_args.create_engine_config(usage_context=UsageContext.LLM_CLASS)
+    vllm_config.compilation_config.compile_only = True
+
+    _run_compile_with_config(vllm_config)
+
+
+def _run_compile_with_config(vllm_config) -> None:
+    """Shared compile-only logic."""
+    from vllm.plugins import load_general_plugins
+    from vllm.v1.executor import Executor
+
+    load_general_plugins()
+
+    executor_class = Executor.get_class(vllm_config)
+    executor = executor_class(vllm_config)
+    executor.collective_rpc("compile_or_warm_up_model")
+
+    logger.info("Compile-only mode complete. Cache populated.")
+    executor.shutdown()

vllm/config/compilation.py

@@ -655,6 +655,12 @@ class CompilationConfig:
     local_cache_dir: str = field(default=None, init=False)  # type: ignore
     """local cache dir for each rank"""
 
+    compile_only: bool = False
+    """If True, run in compile-only mode: only torch.compile
+    compilation, skip CUDA graph capture, kernel warmup, and sampler
+    warmup. Used to pre-populate vLLM's torch.compile cache without
+    allocating KV caches or setting up the full engine."""
+
     fast_moe_cold_start: bool | None = None
     """Optimization for fast MOE cold start.
 
@@ -739,6 +745,8 @@ class CompilationConfig:
             "static_forward_context",
             "pass_config",  # handled separately below
             "dynamic_shapes_config",  # handled separately below
+            # compile_only doesn't affect the compiled graph
+            "compile_only",
         }
 
         from vllm.config.utils import get_hash_factors, hash_factors

vllm/entrypoints/cli/compile.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+
+from vllm.entrypoints.cli.types import CLISubcommand
+from vllm.entrypoints.openai.cli_args import make_arg_parser
+from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+DESCRIPTION = """[Experimental] Populate vLLM's torch.compile cache for a model.
+
+This command is experimental and a work in progress. Not all models and
+configurations are supported yet.
+
+This runs compilation using fake weights (zero GPU memory) so that
+vLLM's torch.compile cache is populated.  Subsequent ``vllm serve`` or
+``LLM(...)`` calls for the same model will hit the warm cache and skip
+compilation.
+"""
+
+
+class CompileSubcommand(CLISubcommand):
+    """The ``compile`` subcommand for the vLLM CLI."""
+
+    name = "compile"
+
+    @staticmethod
+    def cmd(args: argparse.Namespace) -> None:
+        from vllm.compile_only import run_compile_only
+
+        if hasattr(args, "model_tag") and args.model_tag is not None:
+            args.model = args.model_tag
+        run_compile_only(args)
+
+    def subparser_init(
+        self, subparsers: argparse._SubParsersAction
+    ) -> FlexibleArgumentParser:
+        compile_parser = subparsers.add_parser(
+            self.name,
+            help="[Experimental] Populate vLLM's torch.compile cache for a model.",
+            description=DESCRIPTION,
+            usage="vllm compile [model_tag] [options]",
+        )
+        compile_parser = make_arg_parser(compile_parser)
+        compile_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name)
+        return compile_parser
+
+
+def cmd_init() -> list[CLISubcommand]:
+    return [CompileSubcommand()]

vllm/entrypoints/cli/main.py

@@ -16,6 +16,7 @@ logger = init_logger(__name__)
 def main():
     import vllm.entrypoints.cli.benchmark.main
     import vllm.entrypoints.cli.collect_env
+    import vllm.entrypoints.cli.compile
     import vllm.entrypoints.cli.launch
     import vllm.entrypoints.cli.openai
     import vllm.entrypoints.cli.run_batch
@@ -26,6 +27,7 @@ def main():
     CMD_MODULES = [
         vllm.entrypoints.cli.openai,
         vllm.entrypoints.cli.serve,
+        vllm.entrypoints.cli.compile,
         vllm.entrypoints.cli.launch,
         vllm.entrypoints.cli.benchmark.main,
         vllm.entrypoints.cli.collect_env,

vllm/model_executor/layers/fused_moe/layer.py

@@ -860,6 +860,10 @@ class FusedMoE(CustomOp):
     ):
         # for per channel weight quantization
         if shard_id == "w2":
+            hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
+            expert_data = self._narrow_expert_data_for_padding(
+                expert_data, loaded_weight, hidden_dim=hidden_dim
+            )
             expert_data.copy_(loaded_weight)
         elif shard_id in ("w1", "w3"):
             self._load_w13(
@@ -870,6 +874,59 @@ class FusedMoE(CustomOp):
                 tp_rank=tp_rank,
             )
 
+    @staticmethod
+    def _get_hidden_dim(shard_dim: int, ndim: int) -> int:
+        """Compute the hidden dimension index from the shard (intermediate)
+        dimension and tensor rank.
+
+        For 2D weight tensors the two data dims are (0, 1). For 3D tensors
+        with an expert dimension at dim 0, they are (1, 2). ``shard_dim``
+        occupies one of these; the hidden dimension is the other.
+        For 1D tensors (e.g. per-channel scales) returns 0.
+        """
+        if ndim < 2:
+            return 0
+        dim_a = ndim - 2
+        dim_b = ndim - 1
+        if shard_dim == dim_a:
+            return dim_b
+        if shard_dim == dim_b:
+            return dim_a
+        raise ValueError(
+            f"shard_dim={shard_dim} is not a valid data dimension "
+            f"for a {ndim}D tensor (expected {dim_a} or {dim_b})"
+        )
+
+    @staticmethod
+    def _narrow_expert_data_for_padding(
+        expert_data: torch.Tensor,
+        loaded_weight: torch.Tensor,
+        hidden_dim: int,
+    ) -> torch.Tensor:
+        """Narrow expert_data hidden dim to match loaded_weight for padded
+        hidden_size.
+
+        When backends (e.g., DeepEP) round up hidden_size, weight parameters
+        are larger than checkpoint weights. Narrow the padded hidden dimension
+        before copying.
+
+        Args:
+            expert_data: The (possibly padded) parameter tensor to narrow.
+            loaded_weight: The checkpoint weight tensor with original size.
+            hidden_dim: The dimension index corresponding to hidden_size.
+                Must be non-negative.
+        """
+        if (
+            loaded_weight.ndim > 0
+            and 0 <= hidden_dim < expert_data.ndim
+            and hidden_dim < loaded_weight.ndim
+            and expert_data.shape[hidden_dim] > loaded_weight.shape[hidden_dim]
+        ):
+            expert_data = expert_data.narrow(
+                hidden_dim, 0, loaded_weight.shape[hidden_dim]
+            )
+        return expert_data
+
     def _load_w13(
         self,
         expert_data: torch.Tensor,
@@ -907,13 +964,10 @@ class FusedMoE(CustomOp):
         else:
             assert shard_id == "w3"
             expert_data = expert_data.narrow(shard_dim, shard_size, shard_size)
-
-        # Handle padding: if loaded_weight is smaller than expert_data (can happen
-        # on last TP shard with padding), copy to top-left corner
-        if expert_data.shape != loaded_weight.shape:
-            expert_data = expert_data[
-                : loaded_weight.shape[0], : loaded_weight.shape[1]
-            ]
+        hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
+        expert_data = self._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=hidden_dim
+        )
         expert_data.copy_(loaded_weight)
 
     def _load_w2(
@@ -943,12 +997,10 @@ class FusedMoE(CustomOp):
             narrow_size = min(shard_size, available)
             loaded_weight = loaded_weight.narrow(shard_dim, start_offset, narrow_size)
         # w2, down_proj: Load into only logical weight of w2.
-        # Handle padding: if loaded_weight is smaller than expert_data (can happen
-        # on last TP shard with padding), copy to top-left corner
-        if expert_data.shape != loaded_weight.shape:
-            expert_data = expert_data[
-                : loaded_weight.shape[0], : loaded_weight.shape[1]
-            ]
+        hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
+        expert_data = self._narrow_expert_data_for_padding(
+            expert_data, loaded_weight, hidden_dim=hidden_dim
+        )
         expert_data.copy_(loaded_weight)
 
     def _load_single_value(
@@ -1095,9 +1147,25 @@ class FusedMoE(CustomOp):
 
             expert_data = param.data[expert_id]
             if shard_id == "w2":
+                # BnB params are stored as flat packed tensors (e.g.
+                # (packed_size, 1)), not in the logical weight layout.
+                # Narrowing packed data for hidden-dim padding is not
+                # meaningful, so require an exact shape match.
+                if expert_data.shape != loaded_weight.shape:
+                    raise ValueError(
+                        "BitsAndBytes quantization with padded hidden_size "
+                        "(e.g., from DeepEP) is not supported. "
+                        f"Parameter shape {tuple(expert_data.shape)} != "
+                        f"checkpoint shape {tuple(loaded_weight.shape)}"
+                    )
                 expert_data.copy_(loaded_weight)
             elif shard_id in ("w1", "w3"):
-                # BNB inflight quantization has already sharded the weights
+                # BnB stores weights as flat packed tensors.  _load_w13 is
+                # still used to split the w1/w3 portions along shard_dim.
+                # _narrow_expert_data_for_padding will be a no-op since
+                # packed sizes should already match; if DeepEP padding
+                # causes a mismatch the copy_() will fail with a clear
+                # shape error.
                 full_load = True
                 self._load_w13(
                     shard_id=shard_id,

vllm/model_executor/layers/quantization/inc.py

@@ -6,14 +6,24 @@ from typing import TYPE_CHECKING, Any
 
 import regex as re
 import torch
+from torch.nn.parameter import Parameter
 
 from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+from vllm.model_executor.layers.linear import (
+    LinearBase,
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
 from vllm.model_executor.layers.quantization import (
     QuantizationConfig,
     QuantizationMethods,
 )
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.parameter import (
+    GroupQuantScaleParameter,
+    PackedvLLMParameter,
+    RowvLLMParameter,
+)
 from vllm.platforms import current_platform
 from vllm.scalar_type import scalar_types
 
@@ -402,16 +412,30 @@ class INCConfig(QuantizationConfig):
 
         return None
 
-    def apply_ipex_quant_layer(self, layer, prefix: str):
+    def apply_xpu_w4a16_quant_layer(self, layer, prefix: str):
         weight_bits, group_size, sym = self.get_layer_config(layer, prefix)
         if not self.check_quantized(weight_bits):
             if isinstance(layer, (LinearBase, ParallelLMHead)):
                 return UnquantizedLinearMethod()
             else:
                 return None
-        raise NotImplementedError(
-            "INC quantization is not supported during xpu kernel migration."
-        )
+
+        if weight_bits != 4:
+            raise NotImplementedError(
+                f"INC on XPU only supports 4-bit quantization, "
+                f"got weight_bits={weight_bits}."
+            )
+        if not sym:
+            raise NotImplementedError(
+                "INC W4A16 on XPU only supports symmetric quantization for now."
+            )
+        if isinstance(layer, (LinearBase, ParallelLMHead)):
+            return INCXPULinearMethod(
+                weight_bits=weight_bits,
+                group_size=group_size,
+                sym=sym,
+            )
+        return None
 
     def get_quant_method(self, layer: torch.nn.Module, prefix: str):
         if prefix and self.extra_config:
@@ -420,12 +444,8 @@ class INCConfig(QuantizationConfig):
                     layer_name == prefix or layer_name == f"model.{prefix}"
                 ) and self.extra_config[layer_name].get("bits", 16) >= 16:
                     return UnquantizedLinearMethod()
-        if (
-            current_platform.is_cpu()
-            or current_platform.is_xpu()
-            or self.backend == "ipex"
-        ):
-            return self.apply_ipex_quant_layer(layer, prefix)
+        if current_platform.is_xpu():
+            return self.apply_xpu_w4a16_quant_layer(layer, prefix)
         if "gptq" in self.packing_format or "gptq" in self.backend:
             return self.apply_gptq_quant_layer(layer, prefix)
         if "awq" in self.packing_format or "awq" in self.backend:
@@ -440,3 +460,136 @@ class INCConfig(QuantizationConfig):
         if is_auto_round_format:
             return cls.get_name()
         return None
+
+
+class INCXPULinearMethod(LinearMethodBase):
+    """XPU linear method for INC w4a16 GPTQ quantization (symmetric only).
+
+    Repacks GPTQ weights from [in_packed, out] to oneDNN [out, in_packed]
+    layout and calls torch.ops._xpu_C.int4_gemm_w4a16.
+
+    GPTQ format: qweight [in_packed, out] with sequential nibble order.
+
+    Note: Asymmetric quantization (sym=false) is not for now.
+
+    FIXME(yiliu30): Refine the implementation to reuse XPUwNa16LinearKernel.
+    """
+
+    def __init__(self, weight_bits: int, group_size: int, sym: bool):
+        self.weight_bits = weight_bits
+        self.group_size = group_size
+        self.sym = sym
+        self.pack_factor = 32 // weight_bits
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        del output_size  # Unused.
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        scales_and_zp_size = input_size_per_partition // self.group_size
+
+        # GPTQ: qweight [in // pack_factor, out] packed along input dim
+        qweight = PackedvLLMParameter(
+            data=torch.empty(
+                input_size_per_partition // self.pack_factor,
+                output_size_per_partition,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=0,
+            packed_factor=self.pack_factor,
+            weight_loader=weight_loader,
+        )
+        # scales: [num_groups, out] params_dtype
+        scales = GroupQuantScaleParameter(
+            data=torch.empty(
+                scales_and_zp_size,
+                output_size_per_partition,
+                dtype=params_dtype,
+            ),
+            input_dim=0,
+            output_dim=1,
+            weight_loader=weight_loader,
+        )
+        # qzeros: [num_groups, out // pack_factor] int32
+        qzeros = PackedvLLMParameter(
+            data=torch.empty(
+                scales_and_zp_size,
+                output_size_per_partition // self.pack_factor,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.pack_factor,
+            weight_loader=weight_loader,
+        )
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("scales", scales)
+        layer.register_parameter("qzeros", qzeros)
+
+        # GPTQ checkpoints may include g_idx for activation reordering.
+        # Register it so the weight loader doesn't error on unexpected keys.
+        g_idx = RowvLLMParameter(
+            data=torch.tensor(
+                [i // self.group_size for i in range(input_size_per_partition)],
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("g_idx", g_idx)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Repack GPTQ weights into kernel-ready NT layout."""
+        device = layer.qweight.data.device
+
+        # oneDNN int4 kernel requires strides[0]==1 ("NT format"), but GPTQ
+        # checkpoint is [K_packed, N] contiguous with strides (N, 1).
+        # Two transposes are needed — neither alone can achieve this:
+        #   1. .t().contiguous() → [N, K_packed] contiguous in memory
+        #   2. .t()              → [K_packed, N] view with strides (1, K_packed)
+        # The result has the same logical shape but strides[0]==1 as required.
+        qweight_ct = layer.qweight.data.t().contiguous()
+        layer.qweight = Parameter(qweight_ct.t(), requires_grad=False)
+
+        # Scales: [num_groups, out] — no change needed
+        layer.scales = Parameter(layer.scales.data, requires_grad=False)
+
+        # Symmetric: GPTQ v1 stores qzeros=7, effective zp = 7+1 = 8
+        # Kernel expects int8 scalar = 8
+        layer.qzeros = Parameter(
+            torch.tensor([8], dtype=torch.int8, device=device),
+            requires_grad=False,
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # qweight is already in NT layout [K_packed, N] (strides (1, K_packed))
+        # from process_weights_after_loading — pass directly to kernel.
+        out_shape = x.shape[:-1] + (layer.qweight.shape[1],)
+        reshaped_x = x.reshape(-1, x.shape[-1])
+        out = torch.ops._xpu_C.int4_gemm_w4a16(
+            reshaped_x,
+            layer.qweight,
+            bias,
+            layer.scales,
+            layer.qzeros,
+            self.group_size,
+            None,  # g_idx not needed: desc_act is always False for INC models
+        )
+        return out.reshape(out_shape)

vllm/model_executor/model_loader/fake_loader.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fake weight loader for compile-only mode.
+
+Initializes the model on the meta device (preserving Parameter subclasses
+like ModelWeightParameter), and runs process_weights_after_loading on
+meta tensors.
+
+This is used for compile-only mode where we want to run torch.compile
+without actually allocating any GPU memory for the model.
+"""
+
+import torch
+import torch.nn as nn
+
+from vllm.config import ModelConfig, VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
+from vllm.model_executor.model_loader.utils import (
+    initialize_model,
+    process_weights_after_loading,
+)
+from vllm.utils.torch_utils import set_default_torch_dtype
+
+logger = init_logger(__name__)
+
+
+class FakeModelLoader(BaseModelLoader):
+    """Model loader that initializes on meta device.
+
+    Model initialization runs on ``meta`` device because FakeTensorMode
+    doesn't preserve Parameter subclasses (e.g. ``ModelWeightParameter``
+    becomes a plain ``FakeTensor``).  No GPU memory is allocated.
+    """
+
+    def download_model(self, model_config: ModelConfig) -> None:
+        pass
+
+    def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
+        # No-op: all parameters are already on meta from init.
+        pass
+
+    def load_model(
+        self,
+        vllm_config: VllmConfig,
+        model_config: ModelConfig,
+        prefix: str = "",
+    ) -> nn.Module:
+        device_config = vllm_config.device_config
+        load_config = vllm_config.load_config
+        load_device = (
+            device_config.device if load_config.device is None else load_config.device
+        )
+        target_device = torch.device(load_device)
+
+        with set_default_torch_dtype(model_config.dtype):
+            # Initialize model on meta device — no GPU memory, and
+            # Parameter subclasses are preserved (unlike FakeTensorMode
+            # which converts them to FakeTensor).
+            with torch.device("meta"):
+                model = initialize_model(
+                    vllm_config=vllm_config,
+                    model_config=model_config,
+                    prefix=prefix,
+                )
+
+            # Run weight post-processing on meta tensors.
+            from vllm.model_executor.model_loader.base_loader import (
+                _has_online_quant,
+            )
+            from vllm.model_executor.model_loader.reload import (
+                finalize_layerwise_processing,
+            )
+
+            if _has_online_quant(model):
+                finalize_layerwise_processing(model, model_config)
+            process_weights_after_loading(model, model_config, target_device)
+
+            # Set fake_device on all parameters and buffers so
+            # swap_meta_params_to_fake knows the intended target device.
+            # Must happen AFTER process_weights_after_loading which may
+            # create new parameters.
+            for param in model.parameters():
+                param.fake_device = target_device
+            for buf in model.buffers():
+                buf.fake_device = target_device
+
+        return model.eval()
+
+
+def swap_meta_params_to_fake(model: nn.Module) -> None:
+    """Replace all meta-device parameters and buffers with FakeTensors.
+
+    Called before torch.compile so that Dynamo sees cuda-device tensors
+    during tracing.  Parameter subclasses are replaced with plain
+    FakeTensors at this point — isinstance checks are no longer needed
+    after weight processing is complete.
+    """
+    from torch._subclasses.fake_tensor import FakeTensorMode
+
+    fake_mode = FakeTensorMode()
+    with fake_mode:
+        for name, param in list(model.named_parameters()):
+            device = getattr(param, "fake_device", None)
+            assert device is not None and device != torch.device("meta"), (
+                f"Parameter {name} missing fake_device or has meta device"
+            )
+            fake_data = torch.empty(
+                param.shape,
+                dtype=param.dtype,
+                device=device,
+            )
+            # Navigate to the owning module and replace the parameter.
+            *path, attr = name.split(".")
+            parent = model.get_submodule(".".join(path)) if path else model
+            parent.register_parameter(
+                attr,
+                nn.Parameter(fake_data, requires_grad=param.requires_grad),
+            )
+        for name, buf in list(model.named_buffers()):
+            device = getattr(buf, "fake_device", buf.device)
+            if device == torch.device("meta"):
+                device = torch.device("cuda")
+            fake_buf = torch.empty(
+                buf.shape,
+                dtype=buf.dtype,
+                device=device,
+            )
+            *path, attr = name.split(".")
+            parent = model.get_submodule(".".join(path)) if path else model
+            parent.register_buffer(attr, fake_buf)

vllm/v1/worker/gpu_model_runner.py

@@ -4763,6 +4763,12 @@ class GPUModelRunner(
                 if load_dummy_weights:
                     self.load_config.load_format = "dummy"
                 model_loader = get_model_loader(self.load_config)
+                if self.vllm_config.compilation_config.compile_only:
+                    from vllm.model_executor.model_loader.fake_loader import (
+                        FakeModelLoader,
+                    )
+
+                    model_loader = FakeModelLoader()
                 self.model = model_loader.load_model(
                     vllm_config=self.vllm_config, model_config=self.model_config
                 )

vllm/v1/worker/gpu_worker.py

@@ -5,7 +5,7 @@
 import gc
 import os
 from collections.abc import Callable
-from contextlib import AbstractContextManager, nullcontext
+from contextlib import AbstractContextManager, nullcontext, suppress
 from datetime import timedelta
 from types import NoneType
 from typing import TYPE_CHECKING, Any
@@ -280,7 +280,12 @@ class Worker(WorkerBase):
 
             # take current memory snapshot
             self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
-            self.requested_memory = request_memory(init_snapshot, self.cache_config)
+            if self.compilation_config.compile_only:
+                # In compile-only mode with fake weights, skip memory
+                # validation since we don't allocate real GPU memory.
+                self.requested_memory = 0
+            else:
+                self.requested_memory = request_memory(init_snapshot, self.cache_config)
             logger.debug("worker init memory snapshot: %r", self.init_snapshot)
             logger.debug(
                 "worker requested memory: %sGiB", format_gib(self.requested_memory)
@@ -573,10 +578,59 @@ class Worker(WorkerBase):
                 if not any(x in compile_range for x in all_sizes):
                     warmup_sizes.append(compile_range.end)
 
+        if self.compilation_config.compile_only:
+            from vllm.compilation.backends import CompilationDone
+            from vllm.model_executor.model_loader.fake_loader import (
+                swap_meta_params_to_fake,
+            )
+
+            # Swap meta-device parameters to FakeTensors so that
+            # torch.compile sees cuda-device tensors during tracing.
+            swap_meta_params_to_fake(self.model_runner.model)
+
+            # Verify that no significant GPU memory was allocated for
+            # model weights. A small amount (< 64 MiB) may come from
+            # CUDA runtime or library initialization.
+            _COMPILE_ONLY_MEM_THRESHOLD = 64 * 1024 * 1024  # 64 MiB
+            mem_used = torch.accelerator.memory_allocated(self.device)
+            assert mem_used < _COMPILE_ONLY_MEM_THRESHOLD, (
+                f"compile-only mode should use minimal GPU memory after "
+                f"model loading, but {format_gib(mem_used)} GiB is "
+                f"allocated (threshold: "
+                f"{format_gib(_COMPILE_ONLY_MEM_THRESHOLD)} GiB)"
+            )
+
+            # In the normal path, the first torch.compile is triggered
+            # by profile_run() which calls _dummy_run(max_num_tokens).
+            # In compile-only mode we skip _initialize_kv_caches (which
+            # calls profile_run), so call it here to trigger compilation.
+            # CompilationDone is raised after vLLM's torch.compile cache
+            # and AOT artifact are saved, to prevent execution with fake
+            # tensors.
+            with suppress(CompilationDone):
+                self.model_runner.profile_run()
+
+            # Verify compilation didn't allocate significant GPU memory.
+            mem_used = torch.accelerator.memory_allocated(self.device)
+            assert mem_used < _COMPILE_ONLY_MEM_THRESHOLD, (
+                f"compile-only mode should use minimal GPU memory after "
+                f"compilation, but {format_gib(mem_used)} GiB is "
+                f"allocated (threshold: "
+                f"{format_gib(_COMPILE_ONLY_MEM_THRESHOLD)} GiB)"
+            )
+
+            logger.info(
+                "Compile-only mode: compilation complete. "
+                "Skipping kernel warmup, CUDA graphs, and "
+                "sampler warmup."
+            )
+            return self.compilation_config.compilation_time
+
         # We skip EPLB here since we don't want to record dummy metrics
         for size in sorted(warmup_sizes, reverse=True):
             logger.info("Compile and warming up model for size %d", size)
             self.model_runner._dummy_run(size, skip_eplb=True, remove_lora=False)
+
         self.model_runner.maybe_remove_all_loras(self.model_runner.lora_config)
 
         # Warmup and tune the kernels used during model execution before
@@ -703,6 +757,11 @@ class Worker(WorkerBase):
     def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
         return self.model_runner.get_supported_tasks()
 
+    def get_compilation_match_table(self) -> dict[str, int]:
+        from vllm.compilation.passes.vllm_inductor_pass import get_match_table
+
+        return get_match_table()
+
     def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
         """Get encoder timing stats from model runner."""
         return self.model_runner.get_encoder_timing_stats()