From d0359f3e0401f3ccb3b8fb66658908c08d265a44 Mon Sep 17 00:00:00 2001 From: mysterious hhhh <69612673+ultranationalism@users.noreply.github.com> Date: Sun, 19 Apr 2026 04:58:46 +0800 Subject: [PATCH] [Bugfix] Guard mxfp4_experts_quant bindings on ENABLE_NVFP4_SM100 (#40191) Signed-off-by: ultranationalism Signed-off-by: mgoin Co-authored-by: mgoin Co-authored-by: Claude Opus 4.7 (1M context) --- csrc/libtorch_stable/ops.h | 14 -------------- .../quantization/fp4/mxfp4_experts_quant.cu | 10 ++++++++++ csrc/libtorch_stable/torch_bindings.cpp | 8 ++------ 3 files changed, 12 insertions(+), 20 deletions(-) diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index fdf628d6fd9..176cd500633 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -134,20 +134,6 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, torch::stable::Tensor& input, torch::stable::Tensor& input_global_scale); -void mxfp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts, - int64_t n_experts); - -void silu_and_mul_mxfp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts, - int64_t n_experts); - void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, const torch::stable::Tensor& a, const torch::stable::Tensor& b, diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu index 9119f28a750..78e4eda0c01 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu @@ -23,6 +23,7 @@ #include #include +#include #include #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" @@ -420,3 +421,12 @@ void silu_and_mul_mxfp4_experts_quant( stream); }); } + +// Registered here (not torch_bindings.cpp) because VLLM_GPU_FLAGS is applied +// only under COMPILE_LANGUAGE:CUDA, so ENABLE_NVFP4_SM100 is invisible to +// .cpp files and cannot gate the registration from there. +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant)); + m.impl("silu_and_mul_mxfp4_experts_quant", + TORCH_BOX(&silu_and_mul_mxfp4_experts_quant)); +} diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 95ed6b44f10..124512e8162 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -252,12 +252,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { 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)); - ops.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant)); - ops.impl("silu_and_mul_mxfp4_experts_quant", - TORCH_BOX(&silu_and_mul_mxfp4_experts_quant)); - - // W4A8 ops: impl registrations are in the source files - // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu) + // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). + // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. #endif }