[CI Failure] Fix CUDA wheel build for the Kimi K3 fused MLA kernel (#52210)
Misha Goin committed
Aug 13, 2026 at 16:56 UTC
2e2ffd104bd0b12e36704c957d09cda9488d5773
1 file changed
+31
-37
csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu
+31
-37
@@ -123,14 +123,10 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src,
123
int rope_elem_base = 0) {
124
uint4 const v = *reinterpret_cast<const uint4*>(src);
125
if constexpr (FP8 || APPLY_ROPE) {
126
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
127
- // _typeConvert<BFloat16> is unavailable on pre-Ampere. Kimi K3 uses
128
- // bf16 inputs, so discard unsupported conversion paths in those builds.
129
- if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
126
+ using Converter = vllm::_typeConvert<scalar_t>;
127
+ if constexpr (!Converter::exists) {
128
return;
129
} else {
132
-#endif
133
- using Converter = vllm::_typeConvert<scalar_t>;
130
auto const* p =
131
reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
132
float f[kVecElems];
@@ -175,17 +171,15 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src,
171
}
172
*reinterpret_cast<uint2*>(dst) = out;
173
#else
178
- uint8_t out[kVecElems];
174
+ uint8_t out[kVecElems];
175
#pragma unroll
180
- for (int i = 0; i < kVecElems; i++) {
181
- float s = fminf(fmaxf(f[i] * scale_inv, -kFp8Max), kFp8Max);
182
- out[i] = rocm_cvt_float_to_fp8_e4m3(s);
183
- }
184
- *reinterpret_cast<uint2*>(dst) = *reinterpret_cast<uint2 const*>(out);
176
+ for (int i = 0; i < kVecElems; i++) {
177
+ float s = fminf(fmaxf(f[i] * scale_inv, -kFp8Max), kFp8Max);
178
+ out[i] = rocm_cvt_float_to_fp8_e4m3(s);
179
+ }
180
+ *reinterpret_cast<uint2*>(dst) = *reinterpret_cast<uint2 const*>(out);
181
#endif
186
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
182
}
188
-#endif
183
} else {
184
*reinterpret_cast<uint4*>(dst) = v;
185
}
@@ -197,28 +191,32 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src,
191
template <typename scalar_t>
192
__device__ __forceinline__ void copyChunk8UnitFp8(uint8_t* dst,
193
const scalar_t* src) {
200
-#ifndef USE_ROCM
201
- uint4 const input = *reinterpret_cast<const uint4*>(src);
194
using Converter = vllm::_typeConvert<scalar_t>;
203
- auto const* input2 =
204
- reinterpret_cast<typename Converter::packed_hip_type const*>(&input);
205
- uint2 output;
206
- auto* output2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&output);
195
+ if constexpr (!Converter::exists) {
196
+ return;
197
+ } else {
198
+#ifndef USE_ROCM
199
+ uint4 const input = *reinterpret_cast<const uint4*>(src);
200
+ auto const* input2 =
201
+ reinterpret_cast<typename Converter::packed_hip_type const*>(&input);
202
+ uint2 output;
203
+ auto* output2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&output);
204
#pragma unroll
208
- for (int i = 0; i < 4; ++i) {
209
- if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
210
- output2[i] = __nv_cvt_bfloat16raw2_to_fp8x2(
211
- static_cast<__nv_bfloat162_raw>(input2[i]), __NV_SATFINITE,
212
- __NV_E4M3);
213
- } else {
214
- output2[i] = __nv_cvt_halfraw2_to_fp8x2(
215
- static_cast<__half2_raw>(input2[i]), __NV_SATFINITE, __NV_E4M3);
205
+ for (int i = 0; i < 4; ++i) {
206
+ if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
207
+ output2[i] = __nv_cvt_bfloat16raw2_to_fp8x2(
208
+ static_cast<__nv_bfloat162_raw>(input2[i]), __NV_SATFINITE,
209
+ __NV_E4M3);
210
+ } else {
211
+ output2[i] = __nv_cvt_halfraw2_to_fp8x2(
212
+ static_cast<__half2_raw>(input2[i]), __NV_SATFINITE, __NV_E4M3);
213
+ }
214
}
217
- }
218
- *reinterpret_cast<uint2*>(dst) = output;
215
+ *reinterpret_cast<uint2*>(dst) = output;
216
#else
220
- copyChunk8<scalar_t, true>(dst, src, 1.0f);
217
+ copyChunk8<scalar_t, true>(dst, src, 1.0f);
218
#endif
219
+ }
220
}
221
222
// Concat + store one head's full key: dst[e] = [k_nope | k_pe], e in [0, 192).
@@ -360,12 +358,10 @@ __device__ __forceinline__ void writeDsMlaCache(
358
scalar_t* row16 = reinterpret_cast<scalar_t*>(row);
359
scalar_t* rope_dst = row16 + kKvLoraRank / 2 + 8 + laneId * 2;
360
if constexpr (APPLY_ROPE) {
363
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
364
- if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
361
+ using Converter = vllm::_typeConvert<scalar_t>;
362
+ if constexpr (!Converter::exists) {
363
return;
364
} else {
367
-#endif
368
- using Converter = vllm::_typeConvert<scalar_t>;
365
using packed_t = typename Converter::packed_hip_type;
366
packed_t const src = *reinterpret_cast<packed_t const*>(pe + laneId * 2);
367
float2 const xy = Converter::convert(src);
@@ -374,9 +370,7 @@ __device__ __forceinline__ void writeDsMlaCache(
370
static_cast<float>(cos_sin[laneId + kQkRopeHeadDim / 2]);
371
*reinterpret_cast<packed_t*>(rope_dst) = Converter::convert(
372
make_float2(xy.x * cos - xy.y * sin, xy.x * sin + xy.y * cos));
377
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
373
}
379
-#endif
374
} else {
375
*reinterpret_cast<int32_t*>(rope_dst) =
376
*reinterpret_cast<const int32_t*>(pe + laneId * 2);