A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
[Build/CI] Fix CUDA 11.8 build (#17679)
Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Signed-off-by: Tyler Michael Smith <tysmith@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Tyler Michael Smith committed
May 22, 2025 at 15:13 UTC
6e588da0f4b90e695a20779c3d5a079e56ad3a7b
9 files changed
+78
-15
CMakeLists.txt
+5
-1
@@ -30,7 +30,11 @@ set(ignoreMe "${VLLM_PYTHON_PATH}")
30
set(PYTHON_SUPPORTED_VERSIONS "3.9" "3.10" "3.11" "3.12")
31
32
# Supported NVIDIA architectures.
33
-set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0")
33
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL)
34
+ set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0")
35
+else()
36
+ set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0")
37
+endif()
38
39
# Supported AMD GPU architectures.
40
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201")
csrc/moe/moe_ops.h
+3
-1
@@ -28,4 +28,6 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
28
torch::Tensor num_tokens_post_pad, int64_t top_k,
29
int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
30
int64_t BLOCK_SIZE_K, int64_t bit);
31
-#endif
\ No newline at end of file
31
+#endif
32
+
33
+bool moe_permute_unpermute_supported();
\ No newline at end of file
csrc/moe/moe_permute_unpermute_op.cu
+42
-1
@@ -5,6 +5,9 @@
5
#include "permute_unpermute_kernels/dispatch.h"
6
#include "core/registration.h"
7
8
+// moe_permute kernels require at least CUDA 12.0
9
+#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
10
+
11
void moe_permute(
12
const torch::Tensor& input, // [n_token, hidden]
13
const torch::Tensor& topk_weights, //[n_token, topk]
@@ -127,7 +130,45 @@ void moe_unpermute(
130
});
131
}
132
133
+#else
134
+
135
+void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_weights,
136
+ torch::Tensor& topk_ids,
137
+ const torch::Tensor& token_expert_indicies,
138
+ const std::optional<torch::Tensor>& expert_map,
139
+ int64_t n_expert, int64_t n_local_expert, int64_t topk,
140
+ const std::optional<int64_t>& align_block_size,
141
+ torch::Tensor& permuted_input,
142
+ torch::Tensor& expert_first_token_offset,
143
+ torch::Tensor& src_row_id2dst_row_id_map,
144
+ torch::Tensor& m_indices) {
145
+ TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0");
146
+}
147
+
148
+void moe_unpermute(const torch::Tensor& input,
149
+ const torch::Tensor& topk_weights, torch::Tensor& topk_ids,
150
+ const torch::Tensor& token_expert_indicies,
151
+ const std::optional<torch::Tensor>& expert_map,
152
+ int64_t n_expert, int64_t n_local_expert, int64_t topk,
153
+ const std::optional<int64_t>& align_block_size,
154
+ torch::Tensor& permuted_input,
155
+ torch::Tensor& expert_first_token_offset,
156
+ torch::Tensor& src_row_id2dst_row_id_map,
157
+ torch::Tensor& m_indices) {
158
+ TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0");
159
+}
160
+
161
+#endif
162
+
163
+bool moe_permute_unpermute_supported() {
164
+#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
165
+ return true;
166
+#else
167
+ return false;
168
+#endif
169
+}
170
+
171
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
172
m.impl("moe_permute", &moe_permute);
173
m.impl("moe_unpermute", &moe_unpermute);
133
-}
\ No newline at end of file
174
+}
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
+6
-4
@@ -1,6 +1,9 @@
1
2
#include "moe_permute_unpermute_kernel.h"
3
4
+// moe_permute kernels require at least CUDA 12.0
5
+#if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
6
+
7
// CubKeyValueSorter definition begin
8
CubKeyValueSorter::CubKeyValueSorter()
9
: num_experts_(0), num_bits_(sizeof(int) * 8) {}
@@ -131,9 +134,6 @@ __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
134
int num_experts) {
135
auto tidx = threadIdx.x;
136
auto bidx = blockIdx.x;
134
- auto lidx = tidx & 31;
135
- auto widx = tidx >> 5;
136
- auto warp_count = (blockDim.x + 31) >> 5;
137
auto offset = bidx * blockDim.x;
138
auto bound = min(offset + blockDim.x, size);
139
extern __shared__ int smem_expert_map[];
@@ -226,4 +226,6 @@ void getMIndices(int64_t* expert_first_token_offset,
226
expert_first_token_offset, align_expert_first_token_offset, m_indices,
227
num_local_expert, align_block_size);
228
}
229
-}
\ No newline at end of file
229
+}
230
+
231
+#endif
csrc/moe/torch_bindings.cpp
+3
-1
@@ -77,7 +77,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
77
"Tensor topk_ids,Tensor src_row_id2dst_row_id_map, Tensor "
78
"expert_first_token_offset, int n_expert, int n_local_expert,int "
79
"topk, Tensor! hidden_states)->()");
80
- // conditionally compiled so impl registration is in source file
80
+
81
+ m.def("moe_permute_unpermute_supported() -> bool");
82
+ m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported);
83
84
#endif
85
}
csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+1
-1
@@ -123,7 +123,7 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
123
}
124
125
bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
126
- // CUTLASS groped FP8 kernels need at least CUDA 12.3
126
+ // CUTLASS grouped FP8 kernels need at least CUDA 12.3
127
// and SM90 (Hopper)
128
129
#if defined CUDA_VERSION
docker/Dockerfile
+11
-5
@@ -263,8 +263,11 @@ if [ "$TARGETPLATFORM" != "linux/arm64" ]; then \
263
export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0 10.0+PTX'; \
264
else \
265
export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0+PTX'; \
266
- fi && \
267
- export FLASHINFER_ENABLE_AOT=1; \
266
+ fi; \
267
+ CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
268
+ if [ "$CUDA_MAJOR" -lt 12 ]; then \
269
+ export FLASHINFER_ENABLE_SM90=0; \
270
+ fi; \
271
uv pip install --system --no-build-isolation "git+https://github.com/flashinfer-ai/flashinfer@21ea1d2545f74782b91eb8c08fd503ac4c0743fc" ; \
272
fi
273
COPY examples examples
@@ -275,7 +278,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
278
. /etc/environment && \
279
uv pip list
280
278
-# Although we build Flashinfer with AOT mode, there's still
281
+# Even when we build Flashinfer with AOT mode, there's still
282
# some issues w.r.t. JIT compilation. Therefore we need to
283
# install build dependencies for JIT compilation.
284
# TODO: Remove this once FlashInfer AOT wheel is fixed
@@ -303,8 +306,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \
306
uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"
307
308
# install development dependencies (for testing)
306
-RUN --mount=type=cache,target=/root/.cache/uv \
307
- uv pip install --system -r requirements/dev.txt
309
+RUN --mount=type=cache,target=/root/.cache/uv \
310
+ CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
311
+ if [ "$CUDA_MAJOR" -ge 12 ]; then \
312
+ uv pip install --system -r requirements/dev.txt; \
313
+ fi
314
315
# install development dependencies (for testing)
316
RUN --mount=type=cache,target=/root/.cache/uv \
tests/kernels/moe/test_moe_permute_unpermute.py
+3
-1
@@ -13,7 +13,7 @@ import torch
13
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
14
from vllm.model_executor.layers.fused_moe.layer import determine_expert_map
15
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
16
- moe_permute, moe_unpermute)
16
+ moe_permute, moe_permute_unpermute_supported, moe_unpermute)
17
from vllm.platforms import current_platform
18
19
NUM_EXPERTS = [16, 64]
@@ -167,6 +167,8 @@ def torch_unpermute(permuted_hidden_states: torch.Tensor,
167
def test_moe_permute_unpermute(n_token: int, n_hidden: int, topk: int,
168
n_expert: int, ep_size: int, dtype: torch.dtype,
169
align_block_size: Optional[int]):
170
+ if not moe_permute_unpermute_supported():
171
+ pytest.skip("moe_permute_unpermute is not supported on this platform.")
172
fill_invalid_expert = 0
173
ep_rank = np.random.randint(0, ep_size)
174
expert_map = None
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+4
@@ -182,3 +182,7 @@ def moe_unpermute(
182
expert_first_token_offset, n_expert,
183
n_local_expert, topk, hidden_states)
184
return hidden_states
185
+
186
+
187
+def moe_permute_unpermute_supported():
188
+ return torch.ops._moe_C.moe_permute_unpermute_supported()