@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a201ad72d8
Fork 0

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 7, 14:00 UTC

[Refactor][Kernel] Add global helper to deduplicate vectorized memory ops (#35105)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: LopezCastroRoberto <roberto.lopez.castro@udc.es> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>

Roberto L. Castro committed Feb 28, 2026 at 01:28 UTC a201ad72d87eaaa1fe20e2f42378be4ddbc867f4
6 files changed +474 -372
csrc/activation_kernels.cu
+90 -205
@@ -5,117 +5,11 @@
5 #include <cmath>
6
7 #include "cuda_compat.h"
8 +#include "cuda_vec_utils.cuh"
9 #include "dispatch_utils.h"
10
11 namespace vllm {
12
12 -struct alignas(32) u32x8_t {
13 - uint32_t u0, u1, u2, u3, u4, u5, u6, u7;
14 -};
15 -
16 -__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
17 -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
18 - defined(CUDA_VERSION) && CUDA_VERSION >= 12090
19 - asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
20 - : "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
21 - "=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
22 - : "l"(ptr));
23 -#else
24 - const uint4* uint_ptr = reinterpret_cast<const uint4*>(ptr);
25 - uint4 top_half = __ldg(&uint_ptr[0]);
26 - uint4 bottom_half = __ldg(&uint_ptr[1]);
27 - val.u0 = top_half.x;
28 - val.u1 = top_half.y;
29 - val.u2 = top_half.z;
30 - val.u3 = top_half.w;
31 - val.u4 = bottom_half.x;
32 - val.u5 = bottom_half.y;
33 - val.u6 = bottom_half.z;
34 - val.u7 = bottom_half.w;
35 -#endif
36 -}
37 -
38 -__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
39 -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
40 - defined(CUDA_VERSION) && CUDA_VERSION >= 12090
41 - asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
42 - :
43 - : "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
44 - "r"(val.u4), "r"(val.u5), "r"(val.u6), "r"(val.u7)
45 - : "memory");
46 -#else
47 - uint4* uint_ptr = reinterpret_cast<uint4*>(ptr);
48 - uint_ptr[0] = make_uint4(val.u0, val.u1, val.u2, val.u3);
49 - uint_ptr[1] = make_uint4(val.u4, val.u5, val.u6, val.u7);
50 -#endif
51 -}
52 -
53 -template <bool support_256>
54 -struct VecTraits;
55 -
56 -template <>
57 -struct VecTraits<true> {
58 - static constexpr int ARCH_MAX_VEC_SIZE = 32;
59 - using vec_t = u32x8_t;
60 -};
61 -
62 -template <>
63 -struct VecTraits<false> {
64 - static constexpr int ARCH_MAX_VEC_SIZE = 16;
65 - using vec_t = int4;
66 -};
67 -
68 -template <typename T>
69 -struct PackedTraits;
70 -
71 -template <>
72 -struct PackedTraits<c10::BFloat16> {
73 - using packed_t = __nv_bfloat162;
74 -};
75 -
76 -template <>
77 -struct PackedTraits<c10::Half> {
78 - using packed_t = __half2;
79 -};
80 -
81 -template <>
82 -struct PackedTraits<float> {
83 - using packed_t = float2;
84 -};
85 -
86 -template <typename packed_t>
87 -__device__ __forceinline__ float2 cast_to_float2(const packed_t& val) {
88 - if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
89 - return __bfloat1622float2(val);
90 - } else if constexpr (std::is_same_v<packed_t, __half2>) {
91 - return __half22float2(val);
92 - } else if constexpr (std::is_same_v<packed_t, float2>) {
93 - return float2(val);
94 - }
95 -}
96 -
97 -template <typename packed_t>
98 -__device__ __forceinline__ packed_t cast_to_packed(const float2& val) {
99 - if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
100 - return __float22bfloat162_rn(val);
101 - } else if constexpr (std::is_same_v<packed_t, __half2>) {
102 - return __float22half2_rn(val);
103 - } else if constexpr (std::is_same_v<packed_t, float2>) {
104 - return float2(val);
105 - }
106 -}
107 -
108 -template <typename packed_t>
109 -__device__ __forceinline__ packed_t packed_mul(const packed_t& x,
110 - const packed_t& y) {
111 - if constexpr (std::is_same_v<packed_t, __nv_bfloat162> ||
112 - std::is_same_v<packed_t, __half2>) {
113 - return __hmul2(x, y);
114 - } else if constexpr (std::is_same_v<packed_t, float2>) {
115 - return make_float2(x.x * y.x, x.y * y.y);
116 - }
117 -}
118 -
13 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
14 bool act_first>
15 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
@@ -131,16 +25,6 @@ __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
25 : packed_mul(x, PACKED_ACT_FN(y));
26 }
27
134 -// Check if all pointers are 16-byte aligned for int4 vectorized access
135 -__host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
136 - return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
137 -}
138 -
139 -// Check if all pointers are 16-byte aligned for longlong4_32a vectorized access
140 -__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
141 - return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
142 -}
143 -
28 // Activation and gating kernel template.
29 template <typename scalar_t, typename packed_t,
30 scalar_t (*ACT_FN)(const scalar_t&),
@@ -155,36 +39,32 @@ __global__ void act_and_mul_kernel(
39 scalar_t* out_ptr = out + blockIdx.x * d;
40
41 if constexpr (use_vec) {
158 - // Fast path: 128-bit/256-bit vectorized loop
159 - using vec_t = typename VecTraits<use_256b>::vec_t;
160 - constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
161 - constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
42 + using cuda_t = typename CUDATypeConverter<scalar_t>::Type;
43 + using pvec_t = PackedVec<cuda_t, use_256b>;
44
163 - const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
164 - const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
165 - vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
166 - const int num_vecs = d / 2 / VEC_SIZE;
45 + const pvec_t* x_vec = reinterpret_cast<const pvec_t*>(x_ptr);
46 + const pvec_t* y_vec = reinterpret_cast<const pvec_t*>(y_ptr);
47 + pvec_t* out_vec = reinterpret_cast<pvec_t*>(out_ptr);
48 + const int num_vecs = d / 2 / pvec_t::NUM_ELTS;
49
50 for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
169 - vec_t x, y;
51 + pvec_t x, y;
52 if constexpr (use_256b) {
53 ld256(x, &x_vec[i]);
54 ld256(y, &y_vec[i]);
55 } else {
174 - x = VLLM_LDG(&x_vec[i]);
175 - y = VLLM_LDG(&y_vec[i]);
56 + ld128(x, &x_vec[i]);
57 + ld128(y, &y_vec[i]);
58 }
177 - auto* xp = reinterpret_cast<packed_t*>(&x);
178 - auto* yp = reinterpret_cast<packed_t*>(&y);
59 #pragma unroll
180 - for (int j = 0; j < VEC_SIZE; j++) {
181 - xp[j] =
182 - packed_compute<packed_t, PACKED_ACT_FN, act_first>(xp[j], yp[j]);
60 + for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
61 + x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
62 + x.elts[j], y.elts[j]);
63 }
64 if constexpr (use_256b) {
65 st256(x, &out_vec[i]);
66 } else {
187 - out_vec[i] = x;
67 + st128(x, &out_vec[i]);
68 }
69 }
70 } else {
@@ -272,51 +152,54 @@ packed_gelu_tanh_kernel(const packed_t& val) {
152 // Launch activation and gating kernel.
153 // Use ACT_FIRST (bool) indicating whether to apply the activation function
154 // first.
275 -#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
276 - auto dtype = input.scalar_type(); \
277 - int d = input.size(-1) / 2; \
278 - int64_t num_tokens = input.numel() / input.size(-1); \
279 - if (num_tokens == 0) { \
280 - return; \
281 - } \
282 - dim3 grid(num_tokens); \
283 - int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
284 - int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16; \
285 - int vec_size = support_vec / at::elementSize(dtype); \
286 - const bool use_vec = (d % vec_size == 0); \
287 - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
288 - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
289 - if (use_vec) { \
290 - dim3 block(std::min(d / vec_size, 1024)); \
291 - if (cc_major >= 10 && num_tokens > 128) { \
292 - VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
293 - vllm::act_and_mul_kernel< \
294 - scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
295 - KERNEL<scalar_t>, \
296 - PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
297 - ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
298 - out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
299 - }); \
300 - } else { \
301 - VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
302 - vllm::act_and_mul_kernel< \
303 - scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
304 - KERNEL<scalar_t>, \
305 - PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
306 - ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
307 - out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
308 - }); \
309 - } \
310 - } else { \
311 - dim3 block(std::min(d, 1024)); \
312 - VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
313 - vllm::act_and_mul_kernel< \
314 - scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
315 - KERNEL<scalar_t>, \
316 - PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
317 - ACT_FIRST, false><<<grid, block, 0, stream>>>( \
318 - out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
319 - }); \
155 +#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
156 + auto dtype = input.scalar_type(); \
157 + int d = input.size(-1) / 2; \
158 + int64_t num_tokens = input.numel() / input.size(-1); \
159 + if (num_tokens == 0) { \
160 + return; \
161 + } \
162 + dim3 grid(num_tokens); \
163 + int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
164 + int support_vec = \
165 + (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \
166 + ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE \
167 + : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE; \
168 + int vec_size = support_vec / at::elementSize(dtype); \
169 + const bool use_vec = (d % vec_size == 0); \
170 + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
171 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
172 + if (use_vec) { \
173 + dim3 block(std::min(d / vec_size, 1024)); \
174 + if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \
175 + VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
176 + vllm::act_and_mul_kernel< \
177 + scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
178 + KERNEL<scalar_t>, \
179 + PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
180 + ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
181 + out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
182 + }); \
183 + } else { \
184 + VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
185 + vllm::act_and_mul_kernel< \
186 + scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
187 + KERNEL<scalar_t>, \
188 + PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
189 + ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
190 + out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
191 + }); \
192 + } \
193 + } else { \
194 + dim3 block(std::min(d, 1024)); \
195 + VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
196 + vllm::act_and_mul_kernel< \
197 + scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
198 + KERNEL<scalar_t>, \
199 + PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
200 + ACT_FIRST, false><<<grid, block, 0, stream>>>( \
201 + out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
202 + }); \
203 }
204
205 void silu_and_mul(torch::Tensor& out, // [..., d]
@@ -378,35 +261,31 @@ __global__ void act_and_mul_kernel_with_param(
261 scalar_t* out_ptr = out + blockIdx.x * d;
262
263 if constexpr (use_vec) {
381 - // Fast path: 128-bit/256-bit vectorized loop
382 - using vec_t = typename VecTraits<use_256b>::vec_t;
383 - constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
384 - constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
264 + using cuda_t = typename CUDATypeConverter<scalar_t>::Type;
265 + using pvec_t = PackedVec<cuda_t, use_256b>;
266
386 - const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
387 - const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
388 - vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
389 - const int num_vecs = d / 2 / VEC_SIZE;
267 + const pvec_t* x_vec = reinterpret_cast<const pvec_t*>(x_ptr);
268 + const pvec_t* y_vec = reinterpret_cast<const pvec_t*>(y_ptr);
269 + pvec_t* out_vec = reinterpret_cast<pvec_t*>(out_ptr);
270 + const int num_vecs = d / 2 / pvec_t::NUM_ELTS;
271
272 for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
392 - vec_t x, y;
273 + pvec_t x, y;
274 if constexpr (use_256b) {
275 ld256(x, &x_vec[i]);
276 ld256(y, &y_vec[i]);
277 } else {
397 - x = VLLM_LDG(&x_vec[i]);
398 - y = VLLM_LDG(&y_vec[i]);
278 + ld128(x, &x_vec[i]);
279 + ld128(y, &y_vec[i]);
280 }
400 - auto* xp = reinterpret_cast<packed_t*>(&x);
401 - auto* yp = reinterpret_cast<packed_t*>(&y);
281 #pragma unroll
403 - for (int j = 0; j < VEC_SIZE; j++) {
404 - xp[j] = packed_mul(PACKED_ACT_FN(xp[j], param), yp[j]);
282 + for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
283 + x.elts[j] = packed_mul(PACKED_ACT_FN(x.elts[j], param), y.elts[j]);
284 }
285 if constexpr (use_256b) {
286 st256(x, &out_vec[i]);
287 } else {
409 - out_vec[i] = x;
288 + st128(x, &out_vec[i]);
289 }
290 }
291 } else {
@@ -499,21 +378,24 @@ __global__ void swigluoai_and_mul_kernel(
378 } \
379 dim3 grid(num_tokens); \
380 int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
502 - int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16; \
381 + int support_vec = \
382 + (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \
383 + ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE \
384 + : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE; \
385 int vec_size = support_vec / at::elementSize(dtype); \
386 const bool use_vec = (d % vec_size == 0); \
387 const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
388 const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
389 if (use_vec) { \
390 dim3 block(std::min(d / vec_size, 1024)); \
509 - if (cc_major >= 10 && num_tokens > 128) { \
391 + if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \
392 VLLM_DISPATCH_FLOATING_TYPES( \
393 dtype, "act_and_mul_kernel_with_param", [&] { \
394 vllm::act_and_mul_kernel_with_param< \
513 - scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
395 + scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
396 KERNEL<scalar_t>, \
397 PACKED_KERNEL< \
516 - typename vllm::PackedTraits<scalar_t>::packed_t>, \
398 + typename vllm::PackedTypeConverter<scalar_t>::Type>, \
399 true, true><<<grid, block, 0, stream>>>( \
400 out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, \
401 PARAM); \
@@ -522,10 +404,10 @@ __global__ void swigluoai_and_mul_kernel(
404 VLLM_DISPATCH_FLOATING_TYPES( \
405 dtype, "act_and_mul_kernel_with_param", [&] { \
406 vllm::act_and_mul_kernel_with_param< \
525 - scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
407 + scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
408 KERNEL<scalar_t>, \
409 PACKED_KERNEL< \
528 - typename vllm::PackedTraits<scalar_t>::packed_t>, \
410 + typename vllm::PackedTypeConverter<scalar_t>::Type>, \
411 true, false><<<grid, block, 0, stream>>>( \
412 out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, \
413 PARAM); \
@@ -535,9 +417,9 @@ __global__ void swigluoai_and_mul_kernel(
417 dim3 block(std::min(d, 1024)); \
418 VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
419 vllm::act_and_mul_kernel_with_param< \
538 - scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
420 + scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
421 KERNEL<scalar_t>, \
540 - PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
422 + PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
423 false><<<grid, block, 0, stream>>>( \
424 out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM); \
425 }); \
@@ -629,14 +511,17 @@ __global__ void activation_kernel(
511 } \
512 dim3 grid(num_tokens); \
513 int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
632 - int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16; \
514 + int support_vec = \
515 + (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \
516 + ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE \
517 + : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE; \
518 int vec_size = support_vec / at::elementSize(dtype); \
519 const bool use_vec = (d % vec_size == 0); \
520 const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
521 const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
522 if (use_vec) { \
523 dim3 block(std::min(d / vec_size, 1024)); \
639 - if (cc_major >= 10 && num_tokens > 128) { \
524 + if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \
525 VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
526 vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true> \
527 <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
csrc/cuda_vec_utils.cuh new
+334
@@ -0,0 +1,334 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +#pragma once
5 +
6 +#include <c10/util/BFloat16.h>
7 +#include <c10/util/Half.h>
8 +#include <cassert>
9 +
10 +#ifdef USE_ROCM
11 + #include <hip/hip_runtime.h>
12 +#else
13 + #include <cuda_bf16.h>
14 + #include <cuda_fp16.h>
15 + #include <cuda_runtime.h>
16 +#endif
17 +
18 +// Device-side: SM100+ architecture with CUDA 12.9+ toolkit, which
19 +// together enable 256-bit (v8.u32) PTX load/store instructions.
20 +// Use for PTX instruction selection with architecture fallback paths.
21 +#if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
22 + defined(CUDA_VERSION) && CUDA_VERSION >= 12090
23 + #define VLLM_256B_PTX_ENABLED 1
24 +#else
25 + #define VLLM_256B_PTX_ENABLED 0
26 +#endif
27 +
28 +namespace vllm {
29 +
30 +// ============================================================
31 +// Types and traits
32 +// ============================================================
33 +
34 +// 256-bit (32-byte) aligned vector type: 8 x uint32_t
35 +struct alignas(32) u32x8_t {
36 + uint32_t d[8];
37 +};
38 +
39 +// VecTraits — select between 128-bit (int4) and 256-bit
40 +// (u32x8_t) vector types at compile time.
41 +template <bool support_256>
42 +struct VecTraits;
43 +
44 +template <>
45 +struct VecTraits<true> {
46 + static constexpr int ARCH_MAX_VEC_SIZE = 32;
47 + using vec_t = u32x8_t;
48 +};
49 +
50 +template <>
51 +struct VecTraits<false> {
52 + static constexpr int ARCH_MAX_VEC_SIZE = 16;
53 + using vec_t = int4;
54 +};
55 +
56 +// PackedTypeConverter — map between CUDA scalar and packed types
57 +// half <-> half2, __nv_bfloat16 <-> __nv_bfloat162, etc.
58 +template <typename T>
59 +struct PackedTypeConverter {
60 + static_assert(sizeof(T) == 0,
61 + "PackedTypeConverter is not specialized for this type.");
62 +};
63 +
64 +template <>
65 +struct PackedTypeConverter<half2> {
66 + using Type = half;
67 +};
68 +
69 +template <>
70 +struct PackedTypeConverter<half> {
71 + using Type = half2;
72 +};
73 +
74 +template <>
75 +struct PackedTypeConverter<__nv_bfloat162> {
76 + using Type = __nv_bfloat16;
77 +};
78 +
79 +template <>
80 +struct PackedTypeConverter<__nv_bfloat16> {
81 + using Type = __nv_bfloat162;
82 +};
83 +
84 +template <>
85 +struct PackedTypeConverter<float> {
86 + using Type = float2;
87 +};
88 +
89 +template <>
90 +struct PackedTypeConverter<float2> {
91 + using Type = float;
92 +};
93 +
94 +template <>
95 +struct PackedTypeConverter<c10::Half> {
96 + using Type = half2;
97 +};
98 +
99 +template <>
100 +struct PackedTypeConverter<c10::BFloat16> {
101 + using Type = __nv_bfloat162;
102 +};
103 +
104 +// CUDATypeConverter — map PyTorch scalar types to CUDA scalar
105 +// c10::Half -> half, c10::BFloat16 -> __nv_bfloat16
106 +template <typename T>
107 +struct CUDATypeConverter {
108 + using Type = T;
109 +};
110 +
111 +template <>
112 +struct CUDATypeConverter<c10::Half> {
113 + using Type = half;
114 +};
115 +
116 +template <>
117 +struct CUDATypeConverter<c10::BFloat16> {
118 + using Type = __nv_bfloat16;
119 +};
120 +
121 +// PackedVec — typed vector container for packed element access.
122 +// Derives alignment and element count from VecTraits.
123 +// Type is the CUDA scalar type (e.g. half, __nv_bfloat16).
124 +template <class Type, bool use_256b>
125 +struct alignas(VecTraits<use_256b>::ARCH_MAX_VEC_SIZE) PackedVec {
126 + static constexpr int NUM_ELTS =
127 + VecTraits<use_256b>::ARCH_MAX_VEC_SIZE /
128 + sizeof(typename PackedTypeConverter<Type>::Type);
129 + typename PackedTypeConverter<Type>::Type elts[NUM_ELTS];
130 +};
131 +
132 +// ============================================================
133 +// Load / store primitives
134 +// ============================================================
135 +
136 +// 256-bit load / store — SM100+ only (PTX v8 instructions).
137 +__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
138 +#if VLLM_256B_PTX_ENABLED
139 + asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
140 + : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]),
141 + "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7])
142 + : "l"(ptr));
143 +#else
144 + assert(false && "ld256 requires SM100+ with CUDA 12.9+");
145 +#endif
146 +}
147 +
148 +__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
149 +#if VLLM_256B_PTX_ENABLED
150 + asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
151 + :
152 + : "l"(ptr), "r"(val.d[0]), "r"(val.d[1]), "r"(val.d[2]),
153 + "r"(val.d[3]), "r"(val.d[4]), "r"(val.d[5]), "r"(val.d[6]),
154 + "r"(val.d[7])
155 + : "memory");
156 +#else
157 + assert(false && "st256 requires SM100+ with CUDA 12.9+");
158 +#endif
159 +}
160 +
161 +// Generic ld256 / st256 for any 32-byte aligned type (e.g. PackedVec).
162 +// Non-template overloads above are preferred for u32x8_t.
163 +template <typename T>
164 +__device__ __forceinline__ void ld256(T& val, const T* ptr) {
165 + static_assert(sizeof(T) == 32, "ld256 requires a 32-byte type");
166 + ld256(reinterpret_cast<u32x8_t&>(val), reinterpret_cast<const u32x8_t*>(ptr));
167 +}
168 +
169 +template <typename T>
170 +__device__ __forceinline__ void st256(T& val, T* ptr) {
171 + static_assert(sizeof(T) == 32, "st256 requires a 32-byte type");
172 + st256(reinterpret_cast<u32x8_t&>(val), reinterpret_cast<u32x8_t*>(ptr));
173 +}
174 +
175 +// 128-bit load / store via __ldg (read-only cache hint).
176 +template <typename T>
177 +__device__ __forceinline__ void ld128(T& val, const T* ptr) {
178 + static_assert(sizeof(T) == 16, "ld128 requires a 16-byte type");
179 + *reinterpret_cast<int4*>(&val) = __ldg(reinterpret_cast<const int4*>(ptr));
180 +}
181 +
182 +template <typename T>
183 +__device__ __forceinline__ void st128(T& val, T* ptr) {
184 + static_assert(sizeof(T) == 16, "st128 requires a 16-byte type");
185 + *reinterpret_cast<int4*>(ptr) = *reinterpret_cast<int4*>(&val);
186 +}
187 +
188 +// 256-bit cache-streaming (.cs) load / store — SM100+ only.
189 +__forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) {
190 +#if VLLM_256B_PTX_ENABLED
191 + u32x8_t val;
192 + asm volatile("ld.global.cs.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];"
193 + : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]),
194 + "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7])
195 + : "l"(addr));
196 + return val;
197 +#else
198 + assert(false && "ld256_cs requires SM100+ with CUDA 12.9+");
199 + return {};
200 +#endif
201 +}
202 +
203 +__forceinline__ __device__ void st256_cs(u32x8_t* addr, u32x8_t val) {
204 +#if VLLM_256B_PTX_ENABLED
205 + asm volatile(
206 + "st.global.cs.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" ::"l"(addr),
207 + "r"(val.d[0]), "r"(val.d[1]), "r"(val.d[2]), "r"(val.d[3]), "r"(val.d[4]),
208 + "r"(val.d[5]), "r"(val.d[6]), "r"(val.d[7]));
209 +#else
210 + assert(false && "st256_cs requires SM100+ with CUDA 12.9+");
211 +#endif
212 +}
213 +
214 +// 32-bit cache-streaming (.cs) load / store — SM100+ only.
215 +__forceinline__ __device__ int ld32_cs(const int* addr) {
216 +#if VLLM_256B_PTX_ENABLED
217 + int val;
218 + asm volatile("ld.global.cs.b32 %0, [%1];" : "=r"(val) : "l"(addr));
219 + return val;
220 +#else
221 + assert(false && "ld32_cs requires SM100+ with CUDA 12.9+");
222 + return 0;
223 +#endif
224 +}
225 +
226 +__forceinline__ __device__ void st32_cs(int* addr, int val) {
227 +#if VLLM_256B_PTX_ENABLED
228 + asm volatile("st.global.cs.b32 [%0], %1;" ::"l"(addr), "r"(val));
229 +#else
230 + assert(false && "st32_cs requires SM100+ with CUDA 12.9+");
231 +#endif
232 +}
233 +
234 +// Predicated 256-bit / 128-bit cache-global (.cg) loads.
235 +// Returns zero if pred is false. SM100+ only.
236 +__device__ __forceinline__ void ld256_cg_or_zero(u32x8_t& val, const void* ptr,
237 + bool pred) {
238 +#if VLLM_256B_PTX_ENABLED
239 + asm volatile(
240 + "{\n"
241 + " .reg .pred pr;\n"
242 + " setp.ne.u32 pr, %8, 0;\n"
243 + " mov.u32 %0, 0;\n"
244 + " mov.u32 %1, 0;\n"
245 + " mov.u32 %2, 0;\n"
246 + " mov.u32 %3, 0;\n"
247 + " mov.u32 %4, 0;\n"
248 + " mov.u32 %5, 0;\n"
249 + " mov.u32 %6, 0;\n"
250 + " mov.u32 %7, 0;\n"
251 + " @pr ld.global.cg.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%9];\n"
252 + "}\n"
253 + : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]),
254 + "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7])
255 + : "r"((int)pred), "l"(ptr));
256 +#else
257 + assert(false && "ld256_cg_or_zero requires SM100+ with CUDA 12.9+");
258 +#endif
259 +}
260 +
261 +__device__ __forceinline__ void ld128_cg_or_zero(uint4& val, const void* ptr,
262 + bool pred) {
263 +#if VLLM_256B_PTX_ENABLED
264 + uint32_t r0, r1, r2, r3;
265 +
266 + asm volatile(
267 + "{\n"
268 + " .reg .pred pr;\n"
269 + " setp.ne.u32 pr, %4, 0;\n"
270 + " mov.u32 %0, 0;\n"
271 + " mov.u32 %1, 0;\n"
272 + " mov.u32 %2, 0;\n"
273 + " mov.u32 %3, 0;\n"
274 + " @pr ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%5];\n"
275 + "}\n"
276 + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
277 + : "r"((int)pred), "l"(ptr));
278 +
279 + val = uint4{r0, r1, r2, r3};
280 +#else
281 + assert(false && "ld128_cg_or_zero requires SM100+ with CUDA 12.9+");
282 +#endif
283 +}
284 +
285 +// ============================================================
286 +// Alignment helpers
287 +// ============================================================
288 +
289 +__host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
290 + return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
291 +}
292 +
293 +__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
294 + return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
295 +}
296 +
297 +// ============================================================
298 +// Packed type conversion and arithmetic
299 +// ============================================================
300 +
301 +template <typename packed_t>
302 +__device__ __forceinline__ float2 cast_to_float2(const packed_t& val) {
303 + if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
304 + return __bfloat1622float2(val);
305 + } else if constexpr (std::is_same_v<packed_t, __half2>) {
306 + return __half22float2(val);
307 + } else if constexpr (std::is_same_v<packed_t, float2>) {
308 + return float2(val);
309 + }
310 +}
311 +
312 +template <typename packed_t>
313 +__device__ __forceinline__ packed_t cast_to_packed(const float2& val) {
314 + if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
315 + return __float22bfloat162_rn(val);
316 + } else if constexpr (std::is_same_v<packed_t, __half2>) {
317 + return __float22half2_rn(val);
318 + } else if constexpr (std::is_same_v<packed_t, float2>) {
319 + return float2(val);
320 + }
321 +}
322 +
323 +template <typename packed_t>
324 +__device__ __forceinline__ packed_t packed_mul(const packed_t& x,
325 + const packed_t& y) {
326 + if constexpr (std::is_same_v<packed_t, __nv_bfloat162> ||
327 + std::is_same_v<packed_t, __half2>) {
328 + return __hmul2(x, y);
329 + } else if constexpr (std::is_same_v<packed_t, float2>) {
330 + return make_float2(x.x * y.x, x.y * y.y);
331 + }
332 +}
333 +
334 +} // namespace vllm
csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
+18 -18
@@ -39,12 +39,12 @@ namespace vllm {
39 template <class Type, bool UE8M0_SF = false>
40 __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
41 silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols,
42 - int32_t num_padded_cols,
42 + int32_t num_packed_cols,
43 Type const* __restrict__ in,
44 float const* __restrict__ SFScale,
45 uint32_t* __restrict__ out,
46 uint32_t* __restrict__ SFout) {
47 - using PackedVec = vllm::PackedVec<Type>;
47 + using PackedVec = vllm::PackedVec<Type, CVT_FP4_PACK16>;
48 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
49 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
50 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
@@ -63,7 +63,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
63
64 // Input tensor row/col loops.
65 for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
66 - if (colIdx < num_padded_cols) {
66 + if (colIdx < num_packed_cols) {
67 PackedVec in_vec;
68 PackedVec in_vec2;
69 int64_t inOffset =
@@ -73,19 +73,19 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
73
74 bool valid = (rowIdx < numRows) && (elem_idx < numCols);
75 if constexpr (CVT_FP4_PACK16) {
76 - ld256_or_zero_cg_u32<Type>(
77 - in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
78 - valid);
79 - ld256_or_zero_cg_u32<Type>(
80 - in_vec2, &reinterpret_cast<const uint32_t*>(in)[inOffset2 * 8],
81 - valid);
76 + ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),
77 + &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
78 + valid);
79 + ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec2),
80 + &reinterpret_cast<const uint32_t*>(in)[inOffset2 * 8],
81 + valid);
82 } else {
83 - ld128_or_zero_cg_u32<Type>(
84 - in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
85 - valid);
86 - ld128_or_zero_cg_u32<Type>(
87 - in_vec2, &reinterpret_cast<const uint32_t*>(in)[inOffset2 * 4],
88 - valid);
83 + ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec),
84 + &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
85 + valid);
86 + ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec2),
87 + &reinterpret_cast<const uint32_t*>(in)[inOffset2 * 4],
88 + valid);
89 }
90
91 // Compute silu and mul
@@ -142,9 +142,9 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
142 int const numBlocksPerSM =
143 vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
144
145 - int sf_n_unpadded = int(n / CVT_FP4_ELTS_PER_THREAD);
145 + int num_packed_cols = int(n / CVT_FP4_ELTS_PER_THREAD);
146
147 - int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast<int>(block.x));
147 + int grid_y = vllm::div_round_up(num_packed_cols, static_cast<int>(block.x));
148 int grid_x = std::min(
149 int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
150 dim3 grid(grid_x, grid_y);
@@ -154,7 +154,7 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
154 using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
155 auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
156 vllm::silu_mul_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
157 - m, n, sf_n_unpadded, input_ptr, input_sf_ptr,
157 + m, n, num_packed_cols, input_ptr, input_sf_ptr,
158 reinterpret_cast<uint32_t*>(output_ptr),
159 reinterpret_cast<uint32_t*>(sf_out));
160 });
csrc/quantization/fp4/nvfp4_experts_quant.cu
+2 -2
@@ -43,7 +43,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
43 uint32_t* input_offset_by_experts,
44 uint32_t* output_scale_offset_by_experts, int n_experts,
45 bool low_latency) {
46 - using PackedVec = PackedVec<Type>;
46 + using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
47 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
48 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
49 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
@@ -155,7 +155,7 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
155 float const* SFScale, uint32_t* out, uint32_t* SFout,
156 uint32_t* input_offset_by_experts,
157 uint32_t* output_scale_offset_by_experts, int n_experts) {
158 - using PackedVec = PackedVec<Type>;
158 + using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
159 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
160 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
161 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
csrc/quantization/fp4/nvfp4_quant_kernels.cu
+14 -14
@@ -42,7 +42,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
42 Type const* __restrict__ in,
43 float const* __restrict__ SFScale,
44 uint32_t* __restrict__ out, uint32_t* __restrict__ SFout) {
45 - using PackedVec = vllm::PackedVec<Type>;
45 + using PackedVec = vllm::PackedVec<Type, CVT_FP4_PACK16>;
46
47 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
48 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
@@ -71,13 +71,13 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
71 // If we are outside valid rows OR outside valid columns -> Use Zeros
72 bool valid = (rowIdx < numRows) && (elem_idx < numCols);
73 if constexpr (CVT_FP4_PACK16) {
74 - ld256_or_zero_cg_u32<Type>(
75 - in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
76 - valid);
74 + ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),
75 + &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
76 + valid);
77 } else {
78 - ld128_or_zero_cg_u32<Type>(
79 - in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
80 - valid);
78 + ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec),
79 + &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
80 + valid);
81 }
82
83 auto sf_out =
@@ -114,7 +114,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
114 float const* __restrict__ SFScale,
115 uint32_t* __restrict__ out,
116 uint32_t* __restrict__ SFout) {
117 - using PackedVec = PackedVec<Type>;
117 + using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
118
119 static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
120 (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
@@ -139,13 +139,13 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
139 // If we are outside valid rows OR outside valid columns -> Use Zeros
140 bool valid = (rowIdx < numRows) && (elem_idx < numCols);
141 if constexpr (CVT_FP4_PACK16) {
142 - ld256_or_zero_cg_u32<Type>(
143 - in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
144 - valid);
142 + ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),
143 + &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
144 + valid);
145 } else {
146 - ld128_or_zero_cg_u32<Type>(
147 - in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
148 - valid);
146 + ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec),
147 + &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
148 + valid);
149 }
150
151 auto sf_out =
csrc/quantization/fp4/nvfp4_utils.cuh
+16 -133
@@ -19,8 +19,10 @@
19 #include <cuda_runtime.h>
20 #include <cuda_fp8.h>
21
22 -#if (defined(NVFP4_ENABLE_ELTS16) && (CUDART_VERSION >= 12090) && \
23 - defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100)
22 +#include "../../cuda_vec_utils.cuh"
23 +
24 +#if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \
25 + CUDA_VERSION >= 12090
26 #define ELTS_PER_THREAD 16
27 constexpr int CVT_FP4_ELTS_PER_THREAD = 16;
28 constexpr bool CVT_FP4_PACK16 = true;
@@ -34,68 +36,6 @@ constexpr int CVT_FP4_SF_VEC_SIZE = 16;
36
37 namespace vllm {
38
37 -// Convert PyTorch cpp type to CUDA type
38 -template <typename T>
39 -struct CUDATypeConverter {
40 - using Type = T;
41 -};
42 -
43 -template <>
44 -struct CUDATypeConverter<at::Half> {
45 - using Type = half;
46 -};
47 -
48 -template <>
49 -struct CUDATypeConverter<at::BFloat16> {
50 - using Type = __nv_bfloat16;
51 -};
52 -
53 -// Get type2 from type or vice versa (applied to half and bfloat16)
54 -template <typename T>
55 -struct TypeConverter {
56 - using Type = half2;
57 -}; // keep for generality
58 -
59 -template <>
60 -struct TypeConverter<half2> {
61 - using Type = half;
62 -};
63 -
64 -template <>
65 -struct TypeConverter<half> {
66 - using Type = half2;
67 -};
68 -
69 -template <>
70 -struct TypeConverter<__nv_bfloat162> {
71 - using Type = __nv_bfloat16;
72 -};
73 -
74 -template <>
75 -struct TypeConverter<__nv_bfloat16> {
76 - using Type = __nv_bfloat162;
77 -};
78 -
79 -#if (defined(NVFP4_ENABLE_ELTS16) && (CUDART_VERSION >= 12090) && \
80 - defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100)
81 -// Define a 32 bytes packed data type.
82 -template <class Type>
83 -struct alignas(32) PackedVec {
84 - typename TypeConverter<Type>::Type elts[8];
85 -};
86 -#else
87 -// Define a 16 bytes packed data type.
88 -template <class Type>
89 -struct alignas(16) PackedVec {
90 - typename TypeConverter<Type>::Type elts[4];
91 -};
92 -#endif
93 -
94 -template <>
95 -struct PackedVec<__nv_fp8_e4m3> {
96 - __nv_fp8x2_e4m3 elts[8];
97 -};
98 -
39 template <typename Int>
40 __host__ __device__ inline Int round_up(Int x, Int y) {
41 static_assert(std::is_integral_v<Int>,
@@ -208,56 +148,6 @@ __device__ __forceinline__ float reciprocal_approximate_ftz(float a) {
148 return b;
149 }
150
211 -template <class Type>
212 -__device__ __forceinline__ void ld128_or_zero_cg_u32(PackedVec<Type>& out,
213 - const void* ptr,
214 - bool pred) {
215 - uint32_t r0, r1, r2, r3;
216 -
217 - asm volatile(
218 - "{\n"
219 - " .reg .pred pr;\n"
220 - " setp.ne.u32 pr, %4, 0;\n"
221 - " mov.u32 %0, 0;\n"
222 - " mov.u32 %1, 0;\n"
223 - " mov.u32 %2, 0;\n"
224 - " mov.u32 %3, 0;\n"
225 - " @pr ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%5];\n"
226 - "}\n"
227 - : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
228 - : "r"((int)pred), "l"(ptr));
229 -
230 - *reinterpret_cast<uint4*>(&out) = uint4{r0, r1, r2, r3};
231 -}
232 -
233 -template <class Type>
234 -__device__ __forceinline__ void ld256_or_zero_cg_u32(PackedVec<Type>& out,
235 - const void* ptr,
236 - bool pred) {
237 - uint32_t r0, r1, r2, r3, r4, r5, r6, r7;
238 -
239 - asm volatile(
240 - "{\n"
241 - " .reg .pred pr;\n"
242 - " setp.ne.u32 pr, %8, 0;\n"
243 - " mov.u32 %0, 0;\n"
244 - " mov.u32 %1, 0;\n"
245 - " mov.u32 %2, 0;\n"
246 - " mov.u32 %3, 0;\n"
247 - " mov.u32 %4, 0;\n"
248 - " mov.u32 %5, 0;\n"
249 - " mov.u32 %6, 0;\n"
250 - " mov.u32 %7, 0;\n"
251 - " @pr ld.global.cg.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%9];\n"
252 - "}\n"
253 - : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3), "=r"(r4), "=r"(r5), "=r"(r6),
254 - "=r"(r7)
255 - : "r"((int)pred), "l"(ptr));
256 -
257 - reinterpret_cast<uint4*>(&out)[0] = uint4{r0, r1, r2, r3};
258 - reinterpret_cast<uint4*>(&out)[1] = uint4{r4, r5, r6, r7};
259 -}
260 -
151 // Compute SF output offset for swizzled tensor core layout.
152 // SF layout: [numMTiles, numKTiles, 32, 4, 4]
153 // Caller must precompute: numKTiles = (numCols + 63) / 64
@@ -315,8 +205,8 @@ __device__ __forceinline__ uint8_t* sf_out_rowmajor_u8(int row, int pack,
205
206 // Quantizes the provided PackedVec into the uint32_t output
207 template <class Type, int CVT_FP4_NUM_THREADS_PER_SF, bool UE8M0_SF = false>
318 -__device__ __forceinline__ fp4_packed_t
319 -cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
208 +__device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4(
209 + PackedVec<Type, CVT_FP4_PACK16>& vec, float SFScaleVal, uint8_t* SFout) {
210 // Get absolute maximum values among the local 8 values.
211 auto localMax = __habs2(vec.elts[0]);
212
@@ -372,11 +262,7 @@ cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
262
263 #pragma unroll
264 for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
375 - if constexpr (std::is_same_v<Type, half>) {
376 - fp2Vals[i] = __half22float2(vec.elts[i]);
377 - } else {
378 - fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
379 - }
265 + fp2Vals[i] = cast_to_float2(vec.elts[i]);
266 fp2Vals[i].x *= outputScale;
267 fp2Vals[i].y *= outputScale;
268 }
@@ -395,22 +281,19 @@ __device__ __forceinline__ float2 silu2(float2 x) {
281 }
282
283 template <class Type>
398 -__inline__ __device__ PackedVec<Type> compute_silu_mul(
399 - const PackedVec<Type>& x_vec, const PackedVec<Type>& y_vec) {
400 - PackedVec<Type> result;
284 +__inline__ __device__ PackedVec<Type, CVT_FP4_PACK16> compute_silu_mul(
285 + const PackedVec<Type, CVT_FP4_PACK16>& x_vec,
286 + const PackedVec<Type, CVT_FP4_PACK16>& y_vec) {
287 + PackedVec<Type, CVT_FP4_PACK16> result;
288
289 #pragma unroll
290 for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
291 // silu_mul in float32
405 - if constexpr (std::is_same_v<Type, half>) {
406 - float2 silu_vec = silu2(__half22float2(x_vec.elts[i]));
407 - result.elts[i] = __float22half2_rn(
408 - __fmul2_rn(silu_vec, __half22float2(y_vec.elts[i])));
409 - } else {
410 - float2 silu_vec = silu2(__bfloat1622float2(x_vec.elts[i]));
411 - result.elts[i] = __float22bfloat162_rn(
412 - __fmul2_rn(silu_vec, __bfloat1622float2(y_vec.elts[i])));
413 - }
292 + using packed_t = typename PackedTypeConverter<Type>::Type;
293 + float2 silu_vec = silu2(cast_to_float2(x_vec.elts[i]));
294 + float2 y_f2 = cast_to_float2(y_vec.elts[i]);
295 + result.elts[i] = cast_to_packed<packed_t>(
296 + make_float2(silu_vec.x * y_f2.x, silu_vec.y * y_f2.y));
297 }
298 return result;
299 }