[Build] Fix CUDA release wheel builds (#50243)
Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Kevin H. Luu committed
Jul 30, 2026 at 11:48 UTC
5f8f72866c4923fb06902cf60cf219a2340982e8
3 files changed
+90
-50
csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu
+85
-45
@@ -115,50 +115,57 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src,
115
int rope_elem_base = 0) {
116
uint4 const v = *reinterpret_cast<const uint4*>(src);
117
if constexpr (FP8 || APPLY_ROPE) {
118
- using Converter = vllm::_typeConvert<scalar_t>;
119
- auto const* p =
120
- reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
121
- float f[kVecElems];
118
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
119
+ // _typeConvert<BFloat16> is unavailable on pre-Ampere. Kimi K3 uses
120
+ // bf16 inputs, so discard unsupported conversion paths in those builds.
121
+ if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
122
+ return;
123
+ } else {
124
+#endif
125
+ using Converter = vllm::_typeConvert<scalar_t>;
126
+ auto const* p =
127
+ reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
128
+ float f[kVecElems];
129
#pragma unroll
123
- for (int i = 0; i < 4; i++) {
124
- float2 x = Converter::convert(p[i]);
125
- f[2 * i] = x.x;
126
- f[2 * i + 1] = x.y;
127
- }
128
- if constexpr (APPLY_ROPE) {
130
+ for (int i = 0; i < 4; i++) {
131
+ float2 x = Converter::convert(p[i]);
132
+ f[2 * i] = x.x;
133
+ f[2 * i + 1] = x.y;
134
+ }
135
+ if constexpr (APPLY_ROPE) {
136
#pragma unroll
130
- for (int i = 0; i < kVecElems / 2; i++) {
131
- int const pair_idx = rope_elem_base / 2 + i;
132
- float const cos = static_cast<float>(cos_sin[pair_idx]);
133
- float const sin =
134
- static_cast<float>(cos_sin[pair_idx + kQkRopeHeadDim / 2]);
135
- float const x = f[2 * i];
136
- float const y = f[2 * i + 1];
137
- f[2 * i] = x * cos - y * sin;
138
- f[2 * i + 1] = x * sin + y * cos;
137
+ for (int i = 0; i < kVecElems / 2; i++) {
138
+ int const pair_idx = rope_elem_base / 2 + i;
139
+ float const cos = static_cast<float>(cos_sin[pair_idx]);
140
+ float const sin =
141
+ static_cast<float>(cos_sin[pair_idx + kQkRopeHeadDim / 2]);
142
+ float const x = f[2 * i];
143
+ float const y = f[2 * i + 1];
144
+ f[2 * i] = x * cos - y * sin;
145
+ f[2 * i + 1] = x * sin + y * cos;
146
+ }
147
}
140
- }
141
- if constexpr (!FP8) {
142
- uint4 out;
143
- auto* o = reinterpret_cast<typename Converter::packed_hip_type*>(&out);
148
+ if constexpr (!FP8) {
149
+ uint4 out;
150
+ auto* o = reinterpret_cast<typename Converter::packed_hip_type*>(&out);
151
#pragma unroll
145
- for (int i = 0; i < kVecElems / 2; i++) {
146
- o[i] = Converter::convert(make_float2(f[2 * i], f[2 * i + 1]));
152
+ for (int i = 0; i < kVecElems / 2; i++) {
153
+ o[i] = Converter::convert(make_float2(f[2 * i], f[2 * i + 1]));
154
+ }
155
+ *reinterpret_cast<uint4*>(dst) = out;
156
+ return;
157
}
148
- *reinterpret_cast<uint4*>(dst) = out;
149
- return;
150
- }
158
#ifndef USE_ROCM
152
- uint2 out;
153
- auto* o2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out);
159
+ uint2 out;
160
+ auto* o2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out);
161
#pragma unroll
155
- for (int i = 0; i < 4; i++) {
156
- float2 s = make_float2(f[2 * i] * scale_inv, f[2 * i + 1] * scale_inv);
157
- s.x = fminf(fmaxf(s.x, -kFp8Max), kFp8Max);
158
- s.y = fminf(fmaxf(s.y, -kFp8Max), kFp8Max);
159
- o2[i] = __nv_cvt_float2_to_fp8x2(s, __NV_SATFINITE, __NV_E4M3);
160
- }
161
- *reinterpret_cast<uint2*>(dst) = out;
162
+ for (int i = 0; i < 4; i++) {
163
+ float2 s = make_float2(f[2 * i] * scale_inv, f[2 * i + 1] * scale_inv);
164
+ s.x = fminf(fmaxf(s.x, -kFp8Max), kFp8Max);
165
+ s.y = fminf(fmaxf(s.y, -kFp8Max), kFp8Max);
166
+ o2[i] = __nv_cvt_float2_to_fp8x2(s, __NV_SATFINITE, __NV_E4M3);
167
+ }
168
+ *reinterpret_cast<uint2*>(dst) = out;
169
#else
170
uint8_t out[kVecElems];
171
#pragma unroll
@@ -167,6 +174,9 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src,
174
out[i] = rocm_cvt_float_to_fp8_e4m3(s);
175
}
176
*reinterpret_cast<uint2*>(dst) = *reinterpret_cast<uint2 const*>(out);
177
+#endif
178
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
179
+ }
180
#endif
181
} else {
182
*reinterpret_cast<uint4*>(dst) = v;
@@ -280,14 +290,23 @@ __device__ __forceinline__ void writeDsMlaCache(
290
scalar_t* row16 = reinterpret_cast<scalar_t*>(row);
291
scalar_t* rope_dst = row16 + kKvLoraRank / 2 + 8 + laneId * 2;
292
if constexpr (APPLY_ROPE) {
283
- using Converter = vllm::_typeConvert<scalar_t>;
284
- using packed_t = typename Converter::packed_hip_type;
285
- packed_t const src = *reinterpret_cast<packed_t const*>(pe + laneId * 2);
286
- float2 const xy = Converter::convert(src);
287
- float const cos = static_cast<float>(cos_sin[laneId]);
288
- float const sin = static_cast<float>(cos_sin[laneId + kQkRopeHeadDim / 2]);
289
- *reinterpret_cast<packed_t*>(rope_dst) = Converter::convert(
290
- make_float2(xy.x * cos - xy.y * sin, xy.x * sin + xy.y * cos));
293
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
294
+ if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
295
+ return;
296
+ } else {
297
+#endif
298
+ using Converter = vllm::_typeConvert<scalar_t>;
299
+ using packed_t = typename Converter::packed_hip_type;
300
+ packed_t const src = *reinterpret_cast<packed_t const*>(pe + laneId * 2);
301
+ float2 const xy = Converter::convert(src);
302
+ float const cos = static_cast<float>(cos_sin[laneId]);
303
+ float const sin =
304
+ static_cast<float>(cos_sin[laneId + kQkRopeHeadDim / 2]);
305
+ *reinterpret_cast<packed_t*>(rope_dst) = Converter::convert(
306
+ make_float2(xy.x * cos - xy.y * sin, xy.x * sin + xy.y * cos));
307
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
308
+ }
309
+#endif
310
} else {
311
*reinterpret_cast<int32_t*>(rope_dst) =
312
*reinterpret_cast<const int32_t*>(pe + laneId * 2);
@@ -662,6 +681,21 @@ static void launchPdl(KernelT kernel, int num_tokens, int num_heads,
681
#endif
682
}
683
684
+void checkBfloat16Support(torch::headeronly::ScalarType dtype) {
685
+#ifndef USE_ROCM
686
+ if (dtype == torch::headeronly::ScalarType::BFloat16) {
687
+ static int const sm_version = getSMVersion();
688
+ STD_TORCH_CHECK(
689
+ sm_version >= 80,
690
+ "Kimi K3 fused MLA operations require sm_80+ (Ampere or newer); got "
691
+ "sm_",
692
+ sm_version);
693
+ }
694
+#else
695
+ (void)dtype;
696
+#endif
697
+}
698
+
699
} // namespace kimi_k3_fused_ops
700
} // namespace vllm
701
@@ -735,6 +769,7 @@ void fused_kimi_k3_mla_key_concat_kv_cache_insert(
769
kv_c_normed.scalar_type() == dt &&
770
k_out.scalar_type() == dt && k_cache.scalar_type() == dt,
771
"all tensors must share k_nope's (bf16/fp16) dtype");
772
+ kk3::checkBfloat16Support(dt);
773
774
int const num_tokens = static_cast<int>(k_nope.size(0));
775
int const num_heads = static_cast<int>(k_nope.size(1));
@@ -827,6 +862,7 @@ void fused_kimi_k3_mla_key_concat_ds_mla_insert(
862
STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
863
slot_mapping.scalar_type() == ScalarType::Long,
864
"slot_mapping must be int64 CUDA");
865
+ kk3::checkBfloat16Support(dt);
866
867
int const num_tokens = static_cast<int>(k_nope.size(0));
868
int const num_heads = static_cast<int>(k_nope.size(1));
@@ -940,6 +976,7 @@ void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
976
check_scale(k_scale_inv, "k_scale_inv must be scalar float32 CUDA");
977
check_scale(v_scale_inv, "v_scale_inv must be scalar float32 CUDA");
978
check_scale(cache_scale_inv, "cache_scale_inv must be scalar float32 CUDA");
979
+ kk3::checkBfloat16Support(dt);
980
981
int const num_tokens = static_cast<int>(k_nope.size(0));
982
int const num_heads = static_cast<int>(k_nope.size(1));
@@ -1053,6 +1090,7 @@ void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(
1090
k_cache.size(1) == cache_block_size &&
1091
k_cache.size(2) == 576 && k_cache.stride(2) == 1,
1092
"k_cache shape [nblk, block_size, 576] contiguous");
1093
+ kk3::checkBfloat16Support(dt);
1094
1095
int const num_tokens = static_cast<int>(ql_nope.size(0));
1096
int const num_heads = static_cast<int>(ql_nope.size(1));
@@ -1127,6 +1165,7 @@ void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(
1165
};
1166
check_scale(q_scale_inv, "q_scale_inv must be scalar float32 CUDA");
1167
check_scale(cache_scale_inv, "cache_scale_inv must be scalar float32 CUDA");
1168
+ kk3::checkBfloat16Support(dt);
1169
1170
int const num_tokens = static_cast<int>(ql_nope.size(0));
1171
int const num_heads = static_cast<int>(ql_nope.size(1));
@@ -1192,6 +1231,7 @@ void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(
1231
k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
1232
k_cache.size(2) == 656 && k_cache.stride(2) == 1,
1233
"k_cache shape [nblk, block_size, 656] uint8 contiguous");
1234
+ kk3::checkBfloat16Support(dt);
1235
1236
int const num_tokens = static_cast<int>(ql_nope.size(0));
1237
int const num_heads = static_cast<int>(ql_nope.size(1));
docker/Dockerfile
+5
-5
@@ -26,8 +26,8 @@ ARG CUDA_VERSION=13.0.3
26
ARG PYTHON_VERSION=3.12
27
ARG UBUNTU_VERSION=22.04
28
# DeepEPv2 requires NCCL >= 2.30.4 (GIN backend).
29
-# This version is only used for CUDA 13+ builds; CUDA 12 falls back to
30
-# the default NCCL version shipped with the base image.
29
+# This version is used for CUDA 12+ builds to keep the DeepEP build and
30
+# runtime NCCL versions compatible.
31
ARG NCCL_VERSION=2.30.7
32
33
# By parameterizing the base images, we allow third-party to use their own
@@ -487,7 +487,7 @@ ARG NCCL_VERSION
487
RUN --mount=type=cache,target=/opt/uv/cache \
488
mkdir -p /tmp/ep_kernels_workspace/dist && \
489
CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
490
- if [ "$CUDA_MAJOR" -ge 13 ] && [ -n "$NCCL_VERSION" ]; then \
490
+ if [ -n "$NCCL_VERSION" ]; then \
491
echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
492
> /tmp/nccl-override.txt && \
493
export UV_OVERRIDE=/tmp/nccl-override.txt; \
@@ -742,9 +742,9 @@ RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" \
742
743
# DeepEPv2 GIN requires NCCL >= 2.30.4 at both compile and runtime. torch pins
744
# an older version as a transitive dep; this override forces uv to use our
745
-# pinned version whenever nvidia-nccl-cu* is resolved. Empty on CUDA 12 (no-op).
745
+# pinned version whenever nvidia-nccl-cu* is resolved.
746
RUN CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \
747
- if [ "$CUDA_MAJOR" -ge 13 ]; then \
747
+ if [ -n "$NCCL_VERSION" ]; then \
748
echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \
749
> /etc/uv-overrides.txt; \
750
else \
docs/assets/contributing/dockerfile-stages-dependency.png
Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ