| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | #pragma once |
| 5 | |
| 6 | #include <torch/headeronly/util/BFloat16.h> |
| 7 | #include <torch/headeronly/util/Half.h> |
| 8 | #include <cassert> |
| 9 | |
| 10 | #ifdef USE_ROCM |
| 11 | #include <hip/hip_runtime.h> |
| 12 | #include <hip/hip_bf16.h> |
| 13 | #include <hip/hip_fp16.h> |
| 14 | #else |
| 15 | #include <cuda_bf16.h> |
| 16 | #include <cuda_fp16.h> |
| 17 | #include <cuda_runtime.h> |
| 18 | #endif |
| 19 | |
| 20 | // Device-side: SM100+ architecture with CUDA 12.9+ toolkit, which |
| 21 | // together enable 256-bit (v8.u32) PTX load/store instructions. |
| 22 | // Use for PTX instruction selection with architecture fallback paths. |
| 23 | #if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \ |
| 24 | defined(CUDART_VERSION) && CUDART_VERSION >= 12090 |
| 25 | #define VLLM_256B_PTX_ENABLED 1 |
| 26 | #else |
| 27 | #define VLLM_256B_PTX_ENABLED 0 |
| 28 | #endif |
| 29 | |
| 30 | namespace vllm { |
| 31 | |
| 32 | // ============================================================ |
| 33 | // Types and traits |
| 34 | // ============================================================ |
| 35 | |
| 36 | // 256-bit (32-byte) aligned vector type: 8 x uint32_t |
| 37 | struct alignas(32) u32x8_t { |
| 38 | uint32_t d[8]; |
| 39 | }; |
| 40 | |
| 41 | // VecTraits — select between 128-bit (int4) and 256-bit |
| 42 | // (u32x8_t) vector types at compile time. |
| 43 | template <bool support_256> |
| 44 | struct VecTraits; |
| 45 | |
| 46 | template <> |
| 47 | struct VecTraits<true> { |
| 48 | static constexpr int ARCH_MAX_VEC_SIZE = 32; |
| 49 | using vec_t = u32x8_t; |
| 50 | }; |
| 51 | |
| 52 | template <> |
| 53 | struct VecTraits<false> { |
| 54 | static constexpr int ARCH_MAX_VEC_SIZE = 16; |
| 55 | using vec_t = int4; |
| 56 | }; |
| 57 | |
| 58 | // PackedTypeConverter — map between CUDA scalar and packed types |
| 59 | // half <-> half2, __nv_bfloat16 <-> __nv_bfloat162, etc. |
| 60 | template <typename T> |
| 61 | struct PackedTypeConverter { |
| 62 | static_assert(sizeof(T) == 0, |
| 63 | "PackedTypeConverter is not specialized for this type."); |
| 64 | }; |
| 65 | |
| 66 | template <> |
| 67 | struct PackedTypeConverter<half2> { |
| 68 | using Type = half; |
| 69 | }; |
| 70 | |
| 71 | template <> |
| 72 | struct PackedTypeConverter<half> { |
| 73 | using Type = half2; |
| 74 | }; |
| 75 | |
| 76 | template <> |
| 77 | struct PackedTypeConverter<__nv_bfloat162> { |
| 78 | using Type = __nv_bfloat16; |
| 79 | }; |
| 80 | |
| 81 | template <> |
| 82 | struct PackedTypeConverter<__nv_bfloat16> { |
| 83 | using Type = __nv_bfloat162; |
| 84 | }; |
| 85 | |
| 86 | template <> |
| 87 | struct PackedTypeConverter<float> { |
| 88 | using Type = float2; |
| 89 | }; |
| 90 | |
| 91 | template <> |
| 92 | struct PackedTypeConverter<float2> { |
| 93 | using Type = float; |
| 94 | }; |
| 95 | |
| 96 | template <> |
| 97 | struct PackedTypeConverter<c10::Half> { |
| 98 | using Type = half2; |
| 99 | }; |
| 100 | |
| 101 | template <> |
| 102 | struct PackedTypeConverter<c10::BFloat16> { |
| 103 | using Type = __nv_bfloat162; |
| 104 | }; |
| 105 | |
| 106 | // CUDATypeConverter — map PyTorch scalar types to CUDA scalar |
| 107 | // c10::Half -> half, c10::BFloat16 -> __nv_bfloat16 |
| 108 | template <typename T> |
| 109 | struct CUDATypeConverter { |
| 110 | using Type = T; |
| 111 | }; |
| 112 | |
| 113 | template <> |
| 114 | struct CUDATypeConverter<c10::Half> { |
| 115 | using Type = half; |
| 116 | }; |
| 117 | |
| 118 | template <> |
| 119 | struct CUDATypeConverter<c10::BFloat16> { |
| 120 | using Type = __nv_bfloat16; |
| 121 | }; |
| 122 | |
| 123 | // PackedVec — typed vector container for packed element access. |
| 124 | // Derives alignment and element count from VecTraits. |
| 125 | // Type is the CUDA scalar type (e.g. half, __nv_bfloat16). |
| 126 | template <class Type, bool use_256b> |
| 127 | struct alignas(VecTraits<use_256b>::ARCH_MAX_VEC_SIZE) PackedVec { |
| 128 | static constexpr int NUM_ELTS = |
| 129 | VecTraits<use_256b>::ARCH_MAX_VEC_SIZE / |
| 130 | sizeof(typename PackedTypeConverter<Type>::Type); |
| 131 | typename PackedTypeConverter<Type>::Type elts[NUM_ELTS]; |
| 132 | }; |
| 133 | |
| 134 | // ============================================================ |
| 135 | // Load / store primitives |
| 136 | // ============================================================ |
| 137 | |
| 138 | // 256-bit load / store — SM100+ only (PTX v8 instructions). |
| 139 | __device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) { |
| 140 | #if VLLM_256B_PTX_ENABLED |
| 141 | asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n" |
| 142 | : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]), |
| 143 | "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7]) |
| 144 | : "l"(ptr)); |
| 145 | #else |
| 146 | assert(false && "ld256 requires SM100+ with CUDA 12.9+"); |
| 147 | #endif |
| 148 | } |
| 149 | |
| 150 | __device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) { |
| 151 | #if VLLM_256B_PTX_ENABLED |
| 152 | asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n" |
| 153 | : |
| 154 | : "l"(ptr), "r"(val.d[0]), "r"(val.d[1]), "r"(val.d[2]), |
| 155 | "r"(val.d[3]), "r"(val.d[4]), "r"(val.d[5]), "r"(val.d[6]), |
| 156 | "r"(val.d[7]) |
| 157 | : "memory"); |
| 158 | #else |
| 159 | assert(false && "st256 requires SM100+ with CUDA 12.9+"); |
| 160 | #endif |
| 161 | } |
| 162 | |
| 163 | // Generic ld256 / st256 for any 32-byte aligned type (e.g. PackedVec). |
| 164 | // Non-template overloads above are preferred for u32x8_t. |
| 165 | template <typename T> |
| 166 | __device__ __forceinline__ void ld256(T& val, const T* ptr) { |
| 167 | static_assert(sizeof(T) == 32, "ld256 requires a 32-byte type"); |
| 168 | ld256(reinterpret_cast<u32x8_t&>(val), reinterpret_cast<const u32x8_t*>(ptr)); |
| 169 | } |
| 170 | |
| 171 | template <typename T> |
| 172 | __device__ __forceinline__ void st256(T& val, T* ptr) { |
| 173 | static_assert(sizeof(T) == 32, "st256 requires a 32-byte type"); |
| 174 | st256(reinterpret_cast<u32x8_t&>(val), reinterpret_cast<u32x8_t*>(ptr)); |
| 175 | } |
| 176 | |
| 177 | // 128-bit load / store via __ldg (read-only cache hint). |
| 178 | template <typename T> |
| 179 | __device__ __forceinline__ void ld128(T& val, const T* ptr) { |
| 180 | static_assert(sizeof(T) == 16, "ld128 requires a 16-byte type"); |
| 181 | *reinterpret_cast<int4*>(&val) = __ldg(reinterpret_cast<const int4*>(ptr)); |
| 182 | } |
| 183 | |
| 184 | template <typename T> |
| 185 | __device__ __forceinline__ void st128(T& val, T* ptr) { |
| 186 | static_assert(sizeof(T) == 16, "st128 requires a 16-byte type"); |
| 187 | *reinterpret_cast<int4*>(ptr) = *reinterpret_cast<int4*>(&val); |
| 188 | } |
| 189 | |
| 190 | // 256-bit cache-streaming (.cs) load / store — SM100+ only. |
| 191 | __forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) { |
| 192 | #if VLLM_256B_PTX_ENABLED |
| 193 | u32x8_t val; |
| 194 | asm volatile("ld.global.cs.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];" |
| 195 | : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]), |
| 196 | "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7]) |
| 197 | : "l"(addr)); |
| 198 | return val; |
| 199 | #else |
| 200 | assert(false && "ld256_cs requires SM100+ with CUDA 12.9+"); |
| 201 | return u32x8_t{}; |
| 202 | #endif |
| 203 | } |
| 204 | |
| 205 | __forceinline__ __device__ void st256_cs(u32x8_t* addr, u32x8_t val) { |
| 206 | #if VLLM_256B_PTX_ENABLED |
| 207 | asm volatile( |
| 208 | "st.global.cs.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" ::"l"(addr), |
| 209 | "r"(val.d[0]), "r"(val.d[1]), "r"(val.d[2]), "r"(val.d[3]), "r"(val.d[4]), |
| 210 | "r"(val.d[5]), "r"(val.d[6]), "r"(val.d[7])); |
| 211 | #else |
| 212 | assert(false && "st256_cs requires SM100+ with CUDA 12.9+"); |
| 213 | #endif |
| 214 | } |
| 215 | |
| 216 | // 32-bit load / store. |
| 217 | __device__ __forceinline__ int ld32(const int* addr) { return __ldg(addr); } |
| 218 | |
| 219 | __device__ __forceinline__ void st32(int* addr, int val) { *addr = val; } |
| 220 | |
| 221 | // 32-bit cache-streaming (.cs) load / store. |
| 222 | // Falls back to ld32/st32 on ROCm (no .cs hint). |
| 223 | __forceinline__ __device__ int ld32_cs(const int* addr) { |
| 224 | int val; |
| 225 | #ifndef USE_ROCM |
| 226 | asm volatile("ld.global.cs.b32 %0, [%1];" : "=r"(val) : "l"(addr)); |
| 227 | #else |
| 228 | val = ld32(addr); |
| 229 | #endif |
| 230 | return val; |
| 231 | } |
| 232 | |
| 233 | __forceinline__ __device__ void st32_cs(int* addr, int val) { |
| 234 | #ifndef USE_ROCM |
| 235 | asm volatile("st.global.cs.b32 [%0], %1;" ::"l"(addr), "r"(val)); |
| 236 | #else |
| 237 | st32(addr, val); |
| 238 | #endif |
| 239 | } |
| 240 | |
| 241 | // 128-bit cache-streaming (.cs) load / store. |
| 242 | // Falls back to ld128/st128 on ROCm (no .cs hint). |
| 243 | __forceinline__ __device__ int4 ld128_cs(const int4* addr) { |
| 244 | int4 val; |
| 245 | #ifndef USE_ROCM |
| 246 | asm volatile("ld.global.cs.v4.u32 {%0,%1,%2,%3}, [%4];" |
| 247 | : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) |
| 248 | : "l"(addr)); |
| 249 | #else |
| 250 | ld128(val, addr); |
| 251 | #endif |
| 252 | return val; |
| 253 | } |
| 254 | |
| 255 | __forceinline__ __device__ void st128_cs(int4* addr, int4 val) { |
| 256 | #ifndef USE_ROCM |
| 257 | asm volatile("st.global.cs.v4.u32 [%0], {%1,%2,%3,%4};" ::"l"(addr), |
| 258 | "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w)); |
| 259 | #else |
| 260 | st128(val, addr); |
| 261 | #endif |
| 262 | } |
| 263 | |
| 264 | // Predicated 256-bit / 128-bit cache-global (.cg) loads. |
| 265 | // Returns zero if pred is false. SM100+ only. |
| 266 | __device__ __forceinline__ void ld256_cg_or_zero(u32x8_t& val, const void* ptr, |
| 267 | bool pred) { |
| 268 | #if VLLM_256B_PTX_ENABLED |
| 269 | asm volatile( |
| 270 | "{\n" |
| 271 | " .reg .pred pr;\n" |
| 272 | " setp.ne.u32 pr, %8, 0;\n" |
| 273 | " mov.u32 %0, 0;\n" |
| 274 | " mov.u32 %1, 0;\n" |
| 275 | " mov.u32 %2, 0;\n" |
| 276 | " mov.u32 %3, 0;\n" |
| 277 | " mov.u32 %4, 0;\n" |
| 278 | " mov.u32 %5, 0;\n" |
| 279 | " mov.u32 %6, 0;\n" |
| 280 | " mov.u32 %7, 0;\n" |
| 281 | " @pr ld.global.cg.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%9];\n" |
| 282 | "}\n" |
| 283 | : "=r"(val.d[0]), "=r"(val.d[1]), "=r"(val.d[2]), "=r"(val.d[3]), |
| 284 | "=r"(val.d[4]), "=r"(val.d[5]), "=r"(val.d[6]), "=r"(val.d[7]) |
| 285 | : "r"((int)pred), "l"(ptr)); |
| 286 | #else |
| 287 | assert(false && "ld256_cg_or_zero requires SM100+ with CUDA 12.9+"); |
| 288 | #endif |
| 289 | } |
| 290 | |
| 291 | __device__ __forceinline__ void ld128_cg_or_zero(uint4& val, const void* ptr, |
| 292 | bool pred) { |
| 293 | #ifndef USE_ROCM |
| 294 | uint32_t r0, r1, r2, r3; |
| 295 | |
| 296 | asm volatile( |
| 297 | "{\n" |
| 298 | " .reg .pred pr;\n" |
| 299 | " setp.ne.u32 pr, %4, 0;\n" |
| 300 | " mov.u32 %0, 0;\n" |
| 301 | " mov.u32 %1, 0;\n" |
| 302 | " mov.u32 %2, 0;\n" |
| 303 | " mov.u32 %3, 0;\n" |
| 304 | " @pr ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%5];\n" |
| 305 | "}\n" |
| 306 | : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) |
| 307 | : "r"((int)pred), "l"(ptr)); |
| 308 | |
| 309 | val = uint4{r0, r1, r2, r3}; |
| 310 | #else |
| 311 | assert(false && "ld128_cg_or_zero is not supported on ROCm"); |
| 312 | #endif |
| 313 | } |
| 314 | |
| 315 | // ============================================================ |
| 316 | // Alignment helpers |
| 317 | // ============================================================ |
| 318 | |
| 319 | __host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) { |
| 320 | return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0; |
| 321 | } |
| 322 | |
| 323 | // ============================================================ |
| 324 | // Packed type conversion and arithmetic |
| 325 | // ============================================================ |
| 326 | |
| 327 | template <typename packed_t> |
| 328 | __device__ __forceinline__ float2 cast_to_float2(const packed_t& val) { |
| 329 | if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) { |
| 330 | return __bfloat1622float2(val); |
| 331 | } else if constexpr (std::is_same_v<packed_t, __half2>) { |
| 332 | return __half22float2(val); |
| 333 | } else if constexpr (std::is_same_v<packed_t, float2>) { |
| 334 | return float2(val); |
| 335 | } |
| 336 | } |
| 337 | |
| 338 | template <typename packed_t> |
| 339 | __device__ __forceinline__ packed_t cast_to_packed(const float2& val) { |
| 340 | if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) { |
| 341 | return __float22bfloat162_rn(val); |
| 342 | } else if constexpr (std::is_same_v<packed_t, __half2>) { |
| 343 | return __float22half2_rn(val); |
| 344 | } else if constexpr (std::is_same_v<packed_t, float2>) { |
| 345 | return float2(val); |
| 346 | } |
| 347 | } |
| 348 | |
| 349 | template <typename packed_t> |
| 350 | __device__ __forceinline__ packed_t packed_mul(const packed_t& x, |
| 351 | const packed_t& y) { |
| 352 | if constexpr (std::is_same_v<packed_t, __nv_bfloat162> || |
| 353 | std::is_same_v<packed_t, __half2>) { |
| 354 | return __hmul2(x, y); |
| 355 | } else if constexpr (std::is_same_v<packed_t, float2>) { |
| 356 | return make_float2(x.x * y.x, x.y * y.y); |
| 357 | } |
| 358 | } |
| 359 | |
| 360 | } // namespace vllm |