fbc9ba6d30 · tduncan/Nvidia-Optimization-Stack-VLLM

csrc/libtorch_stable/cub_helpers.h renamed

csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp renamed

csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp renamed

csrc/libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh renamed

csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh renamed

csrc/libtorch_stable/dsv3_fused_a_gemm.cu

+2

@@ -733,6 +733,8 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output,

733 output.scalar_type() == torch::headeronly::ScalarType::BFloat16,

734 "Only BFloat16 output dtype is supported");

735

736 + const torch::stable::accelerator::DeviceGuard device_guard(

737 + mat_a.get_device_index());

738 STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");

739

740 auto stream = get_current_cuda_stream(mat_a.get_device_index());

csrc/libtorch_stable/fp32_router_gemm_entry.cu

+2

@@ -127,6 +127,8 @@ void fp32_router_gemm(

127 return;

128 }

129

130 + const torch::stable::accelerator::DeviceGuard device_guard(

131 + mat_a.get_device_index());

132 STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+");

133

134 auto stream = get_current_cuda_stream(mat_a.get_device_index());

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

+1 -1

@@ -57,7 +57,7 @@

57 #include "torch_utils.h"

58

59 #include "../cuda_compat.h"

60 -#include "../type_convert.cuh"

60 +#include "type_convert.cuh"

61 #include "../attention/dtype_fp8.cuh"

62 #include "dispatch_utils.h"

63

csrc/libtorch_stable/fused_qknorm_rope_kernel.cu

+1 -1

@@ -22,7 +22,7 @@

22

23 #include "async_util.cuh"

24 #include "../cuda_compat.h"

25 -#include "../type_convert.cuh"

25 +#include "type_convert.cuh"

26 #include "dispatch_utils.h"

27

28 #define CHECK_TYPE(x, st) \

csrc/libtorch_stable/layernorm_kernels.cu

+2 -2

@@ -2,9 +2,9 @@

2

3 #include "torch_utils.h"

4

5 -#include "../cub_helpers.h"

5 +#include "cub_helpers.h"

6 #include "../core/batch_invariant.hpp"

7 -#include "../type_convert.cuh"

7 +#include "type_convert.cuh"

8 #include "dispatch_utils.h"

9 #include "quantization/vectorization_utils.cuh"

10

csrc/libtorch_stable/layernorm_quant_kernels.cu

+2 -2

@@ -9,10 +9,10 @@

9

10 #include "torch_utils.h"

11

12 -#include "../cub_helpers.h"

12 +#include "cub_helpers.h"

13 #include "../core/batch_invariant.hpp"

14 #include "../quantization/w8a8/fp8/common.cuh"

15 -#include "../type_convert.cuh"

15 +#include "type_convert.cuh"

16 #include "dispatch_utils.h"

17 #include "quantization/vectorization_utils.cuh"

18

csrc/libtorch_stable/minimax_reduce_rms_kernel.h renamed

csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu

+2

@@ -151,6 +151,8 @@ void dsv3_router_gemm(

151 output.scalar_type() == torch::headeronly::ScalarType::BFloat16,

152 "output must be float32 or bf16");

153

154 + const torch::stable::accelerator::DeviceGuard device_guard(

155 + mat_a.get_device_index());

156 const int sm = getSMVersion();

157 STD_TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90");

158

csrc/libtorch_stable/moe/grouped_topk_kernels.cu

+3

@@ -1040,6 +1040,9 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(

1040 auto topk_indices = torch::stable::new_empty(

1041 scores, {num_tokens, topk}, torch::headeronly::ScalarType::Int);

1042 const bool pdl_flag = num_tokens <= vllm::moe::PDLEnableTokens;

1043 +

1044 + const torch::stable::accelerator::DeviceGuard device_guard(

1045 + scores.get_device_index());

1046 const cudaStream_t stream =

1047 get_current_cuda_stream(scores.get_device_index());

1048 auto const sf = static_cast<vllm::moe::ScoringFunc>(scoring_func);

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

+5

@@ -584,6 +584,8 @@ void moe_align_block_size(

584 torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,

585 torch::stable::Tensor num_tokens_post_pad,

586 std::optional<torch::stable::Tensor> maybe_expert_map) {

587 + const torch::stable::accelerator::DeviceGuard device_guard(

588 + topk_ids.get_device_index());

589 const cudaStream_t stream =

590 get_current_cuda_stream(topk_ids.get_device_index());

591

@@ -685,6 +687,8 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,

687 torch::stable::Tensor num_tokens_post_pad) {

688 namespace batched_kernel = vllm::moe::batched_moe_align_block_size;

689

690 + const torch::stable::accelerator::DeviceGuard device_guard(

691 + batch_num_tokens.get_device_index());

692 const cudaStream_t stream =

693 get_current_cuda_stream(batch_num_tokens.get_device_index());

694 int32_t const B = batch_num_tokens.size(0);

@@ -802,6 +806,7 @@ void moe_lora_align_block_size(

806

807 int device_max_shared_mem;

808 int dev = topk_ids.get_device_index();

809 + const torch::stable::accelerator::DeviceGuard device_guard(dev);

810 cudaDeviceGetAttribute(&device_max_shared_mem,

811 cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);

812 const cudaStream_t stream = get_current_cuda_stream(dev);

csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu

+6

@@ -87,6 +87,8 @@ void moe_permute_impl(

87 inv_permuted_idx.sizes().equals(token_expert_indices.sizes()),

88 "token_expert_indices shape must be same as inv_permuted_idx");

89

90 + const torch::stable::accelerator::DeviceGuard device_guard(

91 + input.get_device_index());

92 auto device = input.device();

93 auto n_token = input.sizes()[0];

94 auto n_hidden = input.sizes()[1];

@@ -182,6 +184,8 @@ void moe_unpermute(

184 permuted_hidden_states.scalar_type() == hidden_states.scalar_type(),

185 "permuted_hidden_states dtype must be same as hidden_states");

186

187 + const torch::stable::accelerator::DeviceGuard device_guard(

188 + hidden_states.get_device_index());

189 auto n_token = hidden_states.size(0);

190 auto n_hidden = hidden_states.size(1);

191 auto stream = get_current_cuda_stream(hidden_states.get_device_index());

@@ -238,6 +242,8 @@ void shuffle_rows(const torch::stable::Tensor& input_tensor,

242 STD_TORCH_CHECK(input_tensor.scalar_type() == output_tensor.scalar_type(),

243 "Input and output tensors must have the same data type");

244

245 + const torch::stable::accelerator::DeviceGuard device_guard(

246 + output_tensor.get_device_index());

247 auto stream = get_current_cuda_stream(output_tensor.get_device_index());

248 const int64_t blocks = output_tensor.size(0);

249 const int64_t threads = 256;

csrc/libtorch_stable/moe/topk_softmax_kernels.cu

+1 -1

@@ -25,7 +25,7 @@

25 #include <torch/headeronly/util/Exception.h>

26

27 #include "../../cuda_compat.h"

28 -#include "../../cub_helpers.h"

28 +#include "../cub_helpers.h"

29 #include "libtorch_stable/torch_utils.h"

30

31 #ifndef USE_ROCM

csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu

+1 -1

@@ -26,7 +26,7 @@

26 #include <torch/headeronly/util/Exception.h>

27

28 #include "../../cuda_compat.h"

29 -#include "../../cub_helpers.h"

29 +#include "../cub_helpers.h"

30 #include "libtorch_stable/torch_utils.h"

31 #ifndef USE_ROCM

32 #include <cuda_bf16.h>

csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh

+2

@@ -100,6 +100,8 @@ void run_get_group_gemm_starts(

100 int64_t k = a_tensors.size(1);

101 int64_t scale_k = cutlass::ceil_div(k, b_group_size);

102

103 + const torch::stable::accelerator::DeviceGuard device_guard(

104 + a_tensors.get_device_index());

105 auto stream = get_current_cuda_stream(a_tensors.get_device_index());

106

107 if (false) {

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu

+2 -2

@@ -17,11 +17,11 @@

17 #include <torch/csrc/stable/library.h>

18 #include <torch/csrc/stable/tensor.h>

19 #include "libtorch_stable/torch_utils.h"

20 -#include "cutlass_extensions/torch_utils.hpp"

20 +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"

21 #include "libtorch_stable/cutlass_extensions/common.hpp"

22

23 #include "get_group_starts.cuh"

24 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

24 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

25 #include "w4a8_utils.cuh"

26

27 namespace vllm::cutlass_w4a8_moe {

csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu

+2 -2

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

6 #include <torch/csrc/stable/library.h>

7 #include <torch/csrc/stable/tensor.h>

8 #include "libtorch_stable/torch_utils.h"

9 -#include "cutlass_extensions/torch_utils.hpp"

9 +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"

10 #include "w4a8_utils.cuh"

11

12 #include "cutlass/cutlass.h"

@@ -22,7 +22,7 @@

22 #include "cutlass/util/mixed_dtype_utils.hpp"

23

24 #include "libtorch_stable/cutlass_extensions/common.hpp"

25 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

25 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

26

27 #include <cuda_runtime.h>

28

csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu

+4

@@ -142,6 +142,8 @@ void mxfp4_run_get_group_gemm_starts(

142 torch::stable::Tensor const& sf_offsets,

143 torch::stable::Tensor const& problem_sizes, int M, int N, int K) {

144 int num_experts = (int)expert_offsets.size(0);

145 + const torch::stable::accelerator::DeviceGuard device_guard(

146 + a_tensors.get_device_index());

147 auto stream = get_current_cuda_stream(a_tensors.get_device_index());

148

149 STD_TORCH_CHECK(out_tensors.size(1) == N,

@@ -172,6 +174,8 @@ void run_mxfp4_blockwise_scaled_group_mm_sm100(

174 const torch::stable::Tensor& problem_sizes,

175 const torch::stable::Tensor& expert_offsets,

176 const torch::stable::Tensor& sf_offsets, int M, int N, int K) {

177 + const torch::stable::accelerator::DeviceGuard device_guard(

178 + a.get_device_index());

179 using ProblemShape =

180 cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;

181 using ElementType = cutlass::float_e2m1_t;

csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu

+6

@@ -173,6 +173,8 @@ void run_get_group_gemm_starts(const torch::stable::Tensor& a_starts,

173 torch::stable::Tensor const& problem_sizes,

174 int M, int N, int K) {

175 int num_experts = (int)expert_offsets.size(0);

176 + const torch::stable::accelerator::DeviceGuard device_guard(

177 + a_tensors.get_device_index());

178 auto stream = get_current_cuda_stream(a_tensors.get_device_index());

179

180 STD_TORCH_CHECK(out_tensors.size(1) == N,

@@ -206,6 +208,8 @@ void run_fp4_blockwise_scaled_group_mm_sm100(

208 const torch::stable::Tensor& problem_sizes,

209 const torch::stable::Tensor& expert_offsets,

210 const torch::stable::Tensor& sf_offsets, int M, int N, int K) {

211 + const torch::stable::accelerator::DeviceGuard device_guard(

212 + a.get_device_index());

213 using ProblemShape =

214 cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;

215 using ElementType = cutlass::float_e2m1_t;

@@ -411,6 +415,8 @@ void run_fp4_blockwise_scaled_group_mm_sm120(

415 const torch::stable::Tensor& problem_sizes,

416 const torch::stable::Tensor& expert_offsets,

417 const torch::stable::Tensor& sf_offsets, int M, int N, int K) {

418 + const torch::stable::accelerator::DeviceGuard device_guard(

419 + a.get_device_index());

420 using ProblemShape =

421 cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;

422 using ElementType = cutlass::float_e2m1_t;

csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh

+1 -1

@@ -8,7 +8,7 @@

8 #include "quantization/utils.cuh"

9 #include "quant_conversions.cuh"

10

11 -#include "../../../cub_helpers.h"

11 +#include "../../cub_helpers.h"

12 #include "../../../cuda_compat.h"

13

14 namespace vllm {

csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu

+2

@@ -150,6 +150,8 @@ void rearrange_kn_weight_as_n32k16_order(

150 void* b_zero_reorder =

151 has_zp ? b_zeros_reorder.value().mutable_data_ptr() : nullptr;

152

153 + const torch::stable::accelerator::DeviceGuard device_guard(

154 + b_qweight.get_device_index());

155 cudaStream_t stream = get_current_cuda_stream();

156 if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {

157 allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__half>(

csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh

+1 -1

@@ -1,6 +1,6 @@

1 #pragma once

2

3 -#include "cutlass_extensions/vllm_collective_builder.cuh"

3 +#include "libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh"

4 #include "machete_mainloop.cuh"

5

6 namespace cutlass::gemm::collective {

csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh

+3 -3

@@ -18,9 +18,9 @@

18 // clang-format on

19

20 #include "cutlass_extensions/cute_utils.cuh"

21 -#include "cutlass_extensions/vllm_numeric_conversion.cuh"

22 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

23 -#include "cutlass_extensions/torch_utils.hpp"

21 +#include "libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh"

22 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

23 +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"

24 #include "machete_collective_builder.cuh"

25 #include "machete_prepacked_layout.cuh"

26 #include "machete_interleaving_utils.cuh"

csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh

+1 -1

@@ -1,7 +1,7 @@

1 #pragma once

2

3 #include "machete_mm_kernel.cuh"

4 -#include "cutlass_extensions/torch_utils.hpp"

4 +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"

5 #include "core/scalar_type.hpp"

6 #include "libtorch_stable/torch_utils.h"

7

csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh

+1 -1

@@ -2,7 +2,7 @@

2

3 #include "machete_mm_kernel.cuh"

4 #include "cutlass_extensions/cute_utils.cuh"

5 -#include "cutlass_extensions/torch_utils.hpp"

5 +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"

6 #include <torch/headeronly/util/Exception.h>

7

8 namespace machete {

csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh

+1 -1

@@ -1,7 +1,7 @@

1 #pragma once

2

3 #include "machete_prepack_kernel.cuh"

4 -#include "cutlass_extensions/torch_utils.hpp"

4 +#include "libtorch_stable/cutlass_extensions/torch_utils.hpp"

5 #include "core/scalar_type.hpp"

6 #include "libtorch_stable/torch_utils.h"

7

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh

+1

@@ -37,6 +37,7 @@ void cutlass_gemm_caller(

37 typename GemmKernel::MainloopArguments mainloop_args,

38 typename GemmKernel::EpilogueArguments epilogue_args,

39 typename GemmKernel::TileSchedulerArguments scheduler = {}) {

40 + const torch::stable::accelerator::DeviceGuard device_guard(device.index());

41 cutlass::KernelHardwareInfo hw_info;

42 typename GemmKernel::Arguments args{cutlass::gemm::GemmUniversalMode::kGemm,

43 prob_shape,

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu

+1 -1

@@ -1,6 +1,6 @@

1 #include "scaled_mm_kernels.hpp"

2 #include "scaled_mm_sm90_int8_dispatch.cuh"

3 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

3 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

4

5 namespace vllm {

6

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu

+1 -1

@@ -1,6 +1,6 @@

1 #include "scaled_mm_kernels.hpp"

2 #include "scaled_mm_blockwise_sm100_fp8_dispatch.cuh"

3 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

3 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

4

5 namespace vllm {

6

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu

+1 -1

@@ -1,6 +1,6 @@

1 #include "scaled_mm_kernels.hpp"

2 #include "scaled_mm_blockwise_sm120_fp8_dispatch.cuh"

3 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

3 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

4

5 namespace vllm {

6

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu

+1 -1

@@ -1,7 +1,7 @@

1

2 #include "scaled_mm_kernels.hpp"

3 #include "scaled_mm_blockwise_sm90_fp8_dispatch.cuh"

4 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

4 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

5

6 namespace vllm {

7

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh

+1 -1

@@ -4,7 +4,7 @@

4

5 #include "scaled_mm.cuh"

6 #include "cutlass_gemm_caller.cuh"

7 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

7 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

8

9 /**

10 * This file defines Gemm kernel configurations for SM100 (fp8) based on the

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu

+1 -1

@@ -1,7 +1,7 @@

1 #include "scaled_mm_kernels.hpp"

2 #include "scaled_mm_sm120_fp8_dispatch.cuh"

3 #include "core/batch_invariant.hpp"

4 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

4 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

5

6 namespace vllm {

7

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh

+1 -1

@@ -4,7 +4,7 @@

4

5 #include "scaled_mm.cuh"

6 #include "cutlass_gemm_caller.cuh"

7 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

7 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

8

9 /**

10 * This file defines Gemm kernel configurations for SM90 (fp8) based on the Gemm

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu

+1 -1

@@ -1,6 +1,6 @@

1 #include "scaled_mm_kernels.hpp"

2 #include "scaled_mm_sm90_int8_dispatch.cuh"

3 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

3 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

4

5 namespace vllm {

6

csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh

+2

@@ -74,6 +74,8 @@ void run_get_group_gemm_starts(

74 bool per_act_token = a_scales.numel() != 1;

75 bool per_out_ch = b_scales.numel() != num_experts;

76

77 + const torch::stable::accelerator::DeviceGuard device_guard(

78 + a_tensors.get_device_index());

79 auto stream = get_current_cuda_stream(a_tensors.get_device_index());

80

81 if (false) {

csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh

+4 -1

@@ -7,8 +7,9 @@

7 #include "cutlass/gemm/device/gemm_universal_adapter.h"

8

9 #include <torch/csrc/stable/ops.h>

10 -#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

10 +#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

11 #include "libtorch_stable/cutlass_extensions/common.hpp"

12 +#include "libtorch_stable/torch_utils.h"

13 #include "get_group_starts.cuh"

14

15 using namespace cute;

@@ -103,6 +104,8 @@ void cutlass_group_gemm_caller(torch::stable::Tensor& out_tensors,

104

105 int num_experts = static_cast<int>(expert_offsets.size(0));

106

107 + const torch::stable::accelerator::DeviceGuard device_guard(

108 + a_tensors.get_device_index());

109 auto stream = get_current_cuda_stream(a_tensors.get_device_index());

110

111 auto device = a_tensors.device();

csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu

+5

@@ -212,6 +212,8 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(

212 "n and k must fit in int32");

213

214 int const num_experts = static_cast<int>(num_experts64);

215 + const torch::stable::accelerator::DeviceGuard device_guard(

216 + expert_first_token_offset.get_device_index());

217 auto stream =

218 get_current_cuda_stream(expert_first_token_offset.get_device_index());

219

@@ -241,6 +243,7 @@ void get_cutlass_moe_mm_data_caller(

243 const std::optional<torch::stable::Tensor>& blockscale_offsets,

244 const bool is_gated) {

245 auto device = topk_ids.device();

246 + const torch::stable::accelerator::DeviceGuard device_guard(device.index());

247 auto stream = get_current_cuda_stream(device.index());

248 torch::stable::Tensor atomic_buffer = torch::stable::new_zeros(

249 topk_ids, {num_experts}, torch::headeronly::ScalarType::Int);

@@ -311,6 +314,8 @@ void get_cutlass_batched_moe_mm_data_caller(

314 const torch::stable::Tensor& expert_num_tokens,

315 const int64_t num_local_experts, const int64_t padded_m, const int64_t n,

316 const int64_t k) {

317 + const torch::stable::accelerator::DeviceGuard device_guard(

318 + expert_offsets.get_device_index());

319 auto stream = get_current_cuda_stream(expert_offsets.get_device_index());

320

321 if (num_local_experts * padded_m > SWAP_AB_THRESHOLD) {

csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh

+1

@@ -156,6 +156,7 @@ inline void cutlass_gemm_caller(torch::stable::Tensor& out,

156 torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,

157 std::nullopt, device);

158

159 + const torch::stable::accelerator::DeviceGuard device_guard(device.index());

160 auto stream = get_current_cuda_stream(device.index());

161

162 CUTLASS_CHECK(gemm_op.can_implement(args));

csrc/libtorch_stable/quantization/w8a8/fp8/common.cu

+1 -1

@@ -1,6 +1,6 @@

1 #include "../../../../quantization/w8a8/fp8/common.cuh"

2 #include "../../../dispatch_utils.h"

3 -#include "../../../../cub_helpers.h"

3 +#include "../../../cub_helpers.h"

4 #include "../../vectorization_utils.cuh"

5 #include "../../../torch_utils.h"

6 #include <torch/csrc/stable/macros.h>

csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu

+4

@@ -203,6 +203,8 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,

203 STD_TORCH_CHECK(input.numel() % group_size == 0);

204 STD_TORCH_CHECK(output_s.dim() == 2);

205

206 + const torch::stable::accelerator::DeviceGuard device_guard(

207 + input.get_device_index());

208 cudaStream_t stream = get_current_cuda_stream();

209

210 constexpr int THREADS_PER_GROUP = 16;

@@ -506,6 +508,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,

508 "]; got [", output_s_packed.stride(0), ", ",

509 output_s_packed.stride(1), "].");

510

511 + const torch::stable::accelerator::DeviceGuard device_guard(

512 + input.get_device_index());

513 cudaStream_t stream = get_current_cuda_stream();

514

515 constexpr int THREADS_PER_GROUP = 8;

csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu

+1 -1

@@ -5,7 +5,7 @@

5 #include "../../../dispatch_utils.h"

6 #include "../../../torch_utils.h"

7 #include "../../vectorization_utils.cuh"

8 -#include "../../../../cub_helpers.h"

8 +#include "../../../cub_helpers.h"

9

10 static inline __device__ int8_t float_to_int8_rn(float x) {

11 #ifdef USE_ROCM

csrc/libtorch_stable/sampler.cu

+4

@@ -665,6 +665,8 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,

665 constexpr int kSortingAlgorithmThreshold = 12288;

666 constexpr int kSplitWorkThreshold = 200 * 1000;

667 constexpr int kNumThreadsPerBlock = 512;

668 + const torch::stable::accelerator::DeviceGuard device_guard(

669 + logits.get_device_index());

670 const cudaStream_t stream = get_current_cuda_stream();

671 const auto numColumns = logits.size(1);

672

@@ -727,6 +729,8 @@ void top_k_per_row_prefill(const torch::stable::Tensor& logits,

729 int64_t stride0, int64_t stride1, int64_t topK) {

730 constexpr int kSortingAlgorithmThreshold = 12288;

731 constexpr int kNumThreadsPerBlock = 512;

732 + const torch::stable::accelerator::DeviceGuard device_guard(

733 + logits.get_device_index());

734 const cudaStream_t stream = get_current_cuda_stream();

735

736 int numInsertionBlocks =

csrc/libtorch_stable/topk.cu

+5

@@ -21,6 +21,8 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,

21 int64_t max_seq_len) {

22 namespace P = vllm::persistent;

23

24 + const torch::stable::accelerator::DeviceGuard device_guard(

25 + logits.get_device_index());

26 const int64_t num_rows = logits.size(0);

27 const int64_t stride = logits.stride(0);

28 const cudaStream_t stream = get_current_cuda_stream();

@@ -260,6 +262,9 @@ void persistent_topk(const torch::stable::Tensor& logits,

262 k == 512 || k == 1024 || k == 2048,

263 "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);

264

265 + const torch::stable::accelerator::DeviceGuard device_guard(

266 + logits.get_device_index());

267 +

268 if (k == 512) {

269 launch_persistent_topk<512>(logits, lengths, output, workspace,

270 max_seq_len);

csrc/libtorch_stable/type_convert.cuh renamed

vllm/platforms/cuda.py

+6 -13

@@ -6,6 +6,7 @@ pynvml. However, it should not initialize cuda context.

6

7 from __future__ import annotations

8

9 +import contextlib

10 import os

11 import platform

12 from collections.abc import Callable

@@ -20,6 +21,9 @@ from typing_extensions import ParamSpec

21

22 # import custom ops, trigger op registration

23 import vllm._C_stable_libtorch # noqa

24 +

25 +with contextlib.suppress(ImportError):

26 + import vllm._qutlass_C # noqa

27 import vllm.envs as envs

28 from vllm.logger import init_logger

29 from vllm.utils.import_utils import import_pynvml

@@ -39,11 +43,6 @@ else:

43

44 logger = init_logger(__name__)

45

42 -try:

43 - import vllm._qutlass_C # noqa: F401

44 -except ImportError as e:

45 - logger.warning("Failed to import from vllm._qutlass_C: %r", e)

46 -

46 _P = ParamSpec("_P")

47 _R = TypeVar("_R")

48

@@ -221,16 +220,10 @@ class CudaPlatformBase(Platform):

220 import vllm._C_stable_libtorch # noqa: F401

221 except ImportError as e:

222 logger.warning_once("Failed to import from vllm._C_stable_libtorch: %r", e)

224 - try:

223 + with contextlib.suppress(ImportError):

224 import vllm._moe_C_stable_libtorch # noqa: F401

226 - except ImportError as e:

227 - logger.warning_once(

228 - "Failed to import from vllm._moe_C_stable_libtorch: %r", e

229 - )

230 - try:

225 + with contextlib.suppress(ImportError):

226 import vllm._qutlass_C # noqa: F401

232 - except ImportError as e:

233 - logger.warning_once("Failed to import from vllm._qutlass_C: %r", e)

227

228 @property

229 def supported_dtypes(self) -> list[torch.dtype]:

New stable abi cleanup (#46656)