[Refactor] Remove dead minimax allreduce rms kernel (#46842)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 30, 2026 at 11:29 UTC
9e84ec86486df4aa5872c43a813e70772ab15f3c
4 files changed
-52
csrc/libtorch_stable/minimax_reduce_rms_kernel.cu
-29
@@ -804,35 +804,6 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
804
} // namespace tensorrt_llm
805
} // namespace vllm
806
807
-torch::stable::Tensor minimax_allreduce_rms(
808
- torch::stable::Tensor const& input,
809
- torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
810
- int64_t const rank, int64_t const nranks, double const eps) {
811
- const torch::stable::accelerator::DeviceGuard device_guard(
812
- input.get_device_index());
813
- auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
814
-
815
- allreduce_params.nranks = static_cast<int>(nranks);
816
- allreduce_params.rank = static_cast<int>(rank);
817
- allreduce_params.dtype = input.scalar_type();
818
- allreduce_params.size_q = static_cast<int>(input.numel());
819
- allreduce_params.hidden_dim = static_cast<int>(input.size(-1));
820
- allreduce_params.stride_q = allreduce_params.hidden_dim;
821
- allreduce_params.workspace =
822
- reinterpret_cast<void**>(workspace.mutable_data_ptr());
823
- allreduce_params.allreduce_in = const_cast<void*>(input.const_data_ptr());
824
- allreduce_params.rms_gamma = const_cast<void*>(norm_weight.const_data_ptr());
825
- allreduce_params.rms_eps = static_cast<float>(eps);
826
- allreduce_params.stream = get_current_cuda_stream(input.get_device_index());
827
-
828
- torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input);
829
- allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr();
830
-
831
- vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params);
832
-
833
- return rms_norm_out;
834
-}
835
-
807
std::tuple<torch::stable::Tensor, torch::stable::Tensor>
808
minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
809
torch::stable::Tensor const& norm_weight_q,
csrc/libtorch_stable/ops.h
-4
@@ -288,10 +288,6 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
288
int64_t cache_block_size);
289
290
#ifndef USE_ROCM
291
-torch::stable::Tensor minimax_allreduce_rms(
292
- torch::stable::Tensor const& input,
293
- torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
294
- int64_t const rank, int64_t const nranks, double const eps);
291
std::tuple<torch::stable::Tensor, torch::stable::Tensor>
292
minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
293
torch::stable::Tensor const& norm_weight_q,
csrc/libtorch_stable/torch_bindings.cpp
-5
@@ -449,10 +449,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
449
"int cache_block_size) -> ()");
450
451
#ifndef USE_ROCM
452
- ops.def(
453
- "minimax_allreduce_rms("
454
- "Tensor input, Tensor norm_weight, Tensor workspace, "
455
- "int rank, int nranks, float eps) -> Tensor");
452
ops.def(
453
"minimax_allreduce_rms_qk("
454
"Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, "
@@ -705,7 +701,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
701
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
702
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
703
#ifndef USE_ROCM
708
- ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
704
ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
705
#endif
706
ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
vllm/_custom_ops.py
-14
@@ -3785,20 +3785,6 @@ if hasattr(torch.ops._C, "hadacore_transform"):
3785
return torch.empty_like(x) if not inplace else x
3786
3787
3788
-if hasattr(torch.ops._C, "minimax_allreduce_rms"):
3789
-
3790
- @register_fake("_C::minimax_allreduce_rms")
3791
- def _minimax_allreduce_rms_fake(
3792
- input: torch.Tensor,
3793
- norm_weight: torch.Tensor,
3794
- workspace: torch.Tensor,
3795
- rank: int,
3796
- nranks: int,
3797
- eps: float,
3798
- ) -> torch.Tensor:
3799
- return torch.empty_like(input)
3800
-
3801
-
3788
if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
3789
3790
@register_fake("_C::minimax_allreduce_rms_qk")