| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | // nvfp4_ds_mla KV cache writer and gather/upconvert reader. |
| 5 | // |
| 6 | // Separate TU because its e2m1 conversions are only accepted by ptxas for a |
| 7 | // suffixed Blackwell target (sm_100a / sm_100f); only this file is compiled |
| 8 | // with the suffixed arch list. |
| 9 | |
| 10 | #include "nvfp4_ds_mla_cache.h" |
| 11 | |
| 12 | #include "quantization/fp4/nvfp4_utils.cuh" // vllm::fp32_vec16_to_e2m1() |
| 13 | |
| 14 | #include <cuda_bf16.h> |
| 15 | #include <cuda_fp8.h> |
| 16 | |
| 17 | #if defined(CUDART_VERSION) && CUDART_VERSION >= 12080 |
| 18 | #include <cuda_fp4.h> |
| 19 | #define VLLM_HAS_CUDA_FP4 1 |
| 20 | #else |
| 21 | #define VLLM_HAS_CUDA_FP4 0 |
| 22 | #endif |
| 23 | |
| 24 | // __CUDA_ARCH_FAMILY_SPECIFIC__ is defined only for the suffixed targets. The |
| 25 | // non-native paths trap rather than emulate; the host entry points reject this |
| 26 | // format off SM100, so they are unreachable. |
| 27 | #if VLLM_HAS_CUDA_FP4 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \ |
| 28 | defined(__CUDA_ARCH_FAMILY_SPECIFIC__) |
| 29 | #define VLLM_NVFP4_NATIVE_CVT 1 |
| 30 | #else |
| 31 | #define VLLM_NVFP4_NATIVE_CVT 0 |
| 32 | #endif |
| 33 | |
| 34 | namespace vllm { |
| 35 | |
| 36 | // nvfp4_ds_mla cache format, 352 B/token: |
| 37 | // [0, 256) 512 x e2m1 NoPE, packed 2/byte (low nibble = even element) |
| 38 | // [256, 320) 64 x e4m3 RoPE (unscaled: e4m3's 4 exponent bits span the |
| 39 | // RoPE range unaided) |
| 40 | // [320, 352) 32 x e4m3 NoPE scale factors (one per 16 elements), permuted |
| 41 | // Quantization: sf = e4m3(max(amax_16 / 6, 2^-9)), q = round(x / float(sf)) |
| 42 | // Dequantization: x = float(q) * float(sf) |
| 43 | // |
| 44 | // The scale-factor region is an 8x4 -> 4x8 transpose: element block s -- NoPE |
| 45 | // dims [16s, 16s + 16) -- has its scale at byte kSfNopeOff + nvfp4_sf_byte(s), |
| 46 | // and byte p holds the scale of element block 4 * (p & 7) + (p >> 3). Keep |
| 47 | // nvfp4_sf_byte() in lockstep with its counterpart in FlashMLA |
| 48 | // (csrc/sm100/decode/head64/config.h and tests/quant.py). |
| 49 | constexpr int kSfNopeOff = 256 + 64; |
| 50 | |
| 51 | __host__ __device__ constexpr int nvfp4_sf_byte(int s) { |
| 52 | return 8 * (s & 3) + (s >> 2); |
| 53 | } |
| 54 | |
| 55 | // Thread q's scale groups {4c + q} must land on the contiguous bytes 8q + c. |
| 56 | // This also makes it a permutation of [0, 32), since both sides cover that |
| 57 | // range exactly once. |
| 58 | constexpr bool nvfp4_sf_perm_ok() { |
| 59 | for (int q = 0; q < 4; ++q) { |
| 60 | for (int c = 0; c < 8; ++c) { |
| 61 | if (nvfp4_sf_byte(4 * c + q) != 8 * q + c) { |
| 62 | return false; |
| 63 | } |
| 64 | } |
| 65 | } |
| 66 | return true; |
| 67 | } |
| 68 | static_assert(nvfp4_sf_perm_ok(), |
| 69 | "NVFP4 SF byte order must match FlashMLA's nvfp4_sf_byte()"); |
| 70 | |
| 71 | // `cvt.rn.satfinite.e2m1x2.f32 d, a, b` puts a in d's HIGH nibble and b in its |
| 72 | // LOW nibble, so the odd element of a pair is the first source. |
| 73 | __device__ __forceinline__ uint64_t |
| 74 | nvfp4_pack16_e2m1_rn(const float (&vals)[16], float inv_sf) { |
| 75 | #if VLLM_NVFP4_NATIVE_CVT |
| 76 | float2 pairs[8]; |
| 77 | #pragma unroll |
| 78 | for (int i = 0; i < 8; i++) { |
| 79 | pairs[i] = make_float2(vals[2 * i] * inv_sf, vals[2 * i + 1] * inv_sf); |
| 80 | } |
| 81 | const u32x2 packed = fp32_vec16_to_e2m1(pairs); |
| 82 | return (static_cast<uint64_t>(packed.hi) << 32) | packed.lo; |
| 83 | #else |
| 84 | __trap(); |
| 85 | return 0; |
| 86 | #endif |
| 87 | } |
| 88 | |
| 89 | // The e2m1 data and its e4m3 scale both convert to f16 and multiply there; the |
| 90 | // product then goes f16 -> f32 -> bf16, as no f16 -> bf16 instruction exists. |
| 91 | __device__ __forceinline__ void nvfp4_unpack16_e2m1(uint64_t raw, |
| 92 | uint8_t sf_byte, |
| 93 | __nv_bfloat16 (&out)[16]) { |
| 94 | #if VLLM_NVFP4_NATIVE_CVT |
| 95 | const __half2_raw sf2_raw = __nv_cvt_fp8x2_to_halfraw2( |
| 96 | static_cast<__nv_fp8x2_storage_t>(sf_byte * 0x0101u), __NV_E4M3); |
| 97 | const __half2 sf2 = *reinterpret_cast<const __half2*>(&sf2_raw); |
| 98 | __nv_bfloat162* out2 = reinterpret_cast<__nv_bfloat162*>(out); |
| 99 | #pragma unroll |
| 100 | for (int i = 0; i < 8; i++) { |
| 101 | // low nibble -> .x, so .x = element 2i, matching the writer. |
| 102 | const __half2_raw v_raw = __nv_cvt_fp4x2_to_halfraw2( |
| 103 | static_cast<__nv_fp4x2_storage_t>(raw >> (8 * i)), __NV_E2M1); |
| 104 | const __half2 v = __hmul2(*reinterpret_cast<const __half2*>(&v_raw), sf2); |
| 105 | out2[i] = __float22bfloat162_rn(__half22float2(v)); |
| 106 | } |
| 107 | #else |
| 108 | __trap(); |
| 109 | #endif |
| 110 | } |
| 111 | |
| 112 | __device__ __forceinline__ float nvfp4_tile_scale(const float* vals, |
| 113 | float max_val, |
| 114 | uint8_t* sf_out) { |
| 115 | float amax = 0.0f; |
| 116 | #pragma unroll |
| 117 | for (int i = 0; i < 16; i++) { |
| 118 | amax = fmaxf(amax, fabsf(vals[i])); |
| 119 | } |
| 120 | const float sf_f = fmaxf(amax / max_val, 0.001953125f); // 2^-9 |
| 121 | const __nv_fp8_e4m3 sf8(sf_f); |
| 122 | *sf_out = *reinterpret_cast<const uint8_t*>(&sf8); |
| 123 | return float(sf8); |
| 124 | } |
| 125 | |
| 126 | // One CUDA block (64 threads) per token. |
| 127 | // Warp 0: 32 threads, one 16-element NoPE tile each. |
| 128 | // Warp 1: 32 threads, two RoPE elements each. |
| 129 | __global__ void concat_and_cache_nvfp4_ds_mla_kernel( |
| 130 | const __nv_bfloat16* __restrict__ kv_c, // [num_tokens, 512] |
| 131 | const __nv_bfloat16* __restrict__ k_pe, // [num_tokens, 64] |
| 132 | uint8_t* __restrict__ kv_cache, // [num_blocks, block_size, 352] |
| 133 | const int64_t* __restrict__ slot_mapping, // [num_tokens] |
| 134 | const int block_stride, // |
| 135 | const int entry_stride, // |
| 136 | const int kv_c_stride, // |
| 137 | const int k_pe_stride, // |
| 138 | const int block_size) { |
| 139 | const int64_t token_idx = blockIdx.x; |
| 140 | const int64_t slot_idx = slot_mapping[token_idx]; |
| 141 | // NOTE: slot_idx can be -1 if the token is padded |
| 142 | if (slot_idx < 0) { |
| 143 | return; |
| 144 | } |
| 145 | const int64_t block_idx = slot_idx / block_size; |
| 146 | const int64_t block_offset = slot_idx % block_size; |
| 147 | uint8_t* token_ptr = |
| 148 | kv_cache + block_idx * block_stride + block_offset * entry_stride; |
| 149 | |
| 150 | if (threadIdx.x < 32) { |
| 151 | const int tile = threadIdx.x; |
| 152 | const __nv_bfloat16* src = kv_c + token_idx * kv_c_stride + tile * 16; |
| 153 | float vals[16]; |
| 154 | #pragma unroll |
| 155 | for (int i = 0; i < 16; i++) { |
| 156 | vals[i] = __bfloat162float(src[i]); |
| 157 | } |
| 158 | // `tile` is the element-block index; its scale byte is permuted. |
| 159 | const float sf = nvfp4_tile_scale( |
| 160 | vals, 6.0f, token_ptr + kSfNopeOff + nvfp4_sf_byte(tile)); |
| 161 | const float inv_sf = 1.0f / sf; |
| 162 | *reinterpret_cast<uint64_t*>(token_ptr + tile * 8) = |
| 163 | nvfp4_pack16_e2m1_rn(vals, inv_sf); |
| 164 | } else { |
| 165 | // Plain e4m3, unscaled. All 32 lanes take 2 elements each, matching the |
| 166 | // reader and concat_and_cache_ds_mla_kernel. |
| 167 | const int lane = threadIdx.x - 32; |
| 168 | const __nv_bfloat16* src = k_pe + token_idx * k_pe_stride + lane * 2; |
| 169 | uint8_t q[2]; |
| 170 | #pragma unroll |
| 171 | for (int i = 0; i < 2; i++) { |
| 172 | const __nv_fp8_e4m3 v8(__bfloat162float(src[i])); |
| 173 | q[i] = *reinterpret_cast<const uint8_t*>(&v8); |
| 174 | } |
| 175 | *reinterpret_cast<uint16_t*>(token_ptr + 256 + lane * 2) = |
| 176 | *reinterpret_cast<const uint16_t*>(q); |
| 177 | } |
| 178 | } |
| 179 | |
| 180 | // One warp per token (mirrors cp_gather_and_upconvert_fp8_kv_cache). |
| 181 | __global__ void cp_gather_and_upconvert_nvfp4_kv_cache_kernel( |
| 182 | const uint8_t* __restrict__ src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 352] |
| 183 | __nv_bfloat16* __restrict__ dst, // [total_tokens, 576] |
| 184 | const int32_t* __restrict__ block_table, // [num_reqs, BLOCK_INDICES] |
| 185 | const int32_t* __restrict__ workspace_starts, // [num_reqs] |
| 186 | const int32_t num_reqs, const int32_t block_size, |
| 187 | const int32_t total_tokens, const int64_t block_table_stride, |
| 188 | const int64_t cache_block_stride, const int64_t cache_entry_stride, |
| 189 | const int64_t dst_entry_stride) { |
| 190 | const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; |
| 191 | if (flat_warp_id >= total_tokens) return; |
| 192 | const int lane_id = threadIdx.x & 31; |
| 193 | |
| 194 | // Binary search to find which request owns this output token |
| 195 | int lo = 0, hi = num_reqs - 1; |
| 196 | while (lo < hi) { |
| 197 | int mid = (lo + hi + 1) >> 1; |
| 198 | if (workspace_starts[mid] <= flat_warp_id) |
| 199 | lo = mid; |
| 200 | else |
| 201 | hi = mid - 1; |
| 202 | } |
| 203 | const int req_id = lo; |
| 204 | |
| 205 | const int out_token_id = flat_warp_id; |
| 206 | const int token_offset = out_token_id - workspace_starts[req_id]; |
| 207 | const int cache_block_idx = token_offset / block_size; |
| 208 | const int offset_in_block = token_offset % block_size; |
| 209 | const int physical_block = |
| 210 | block_table[req_id * block_table_stride + cache_block_idx]; |
| 211 | |
| 212 | const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride + |
| 213 | offset_in_block * cache_entry_stride; |
| 214 | __nv_bfloat16* dst_ptr = dst + out_token_id * dst_entry_stride; |
| 215 | |
| 216 | // `lane_id` is the element-block index, so its scale byte is permuted. |
| 217 | { |
| 218 | const uint64_t raw = |
| 219 | *reinterpret_cast<const uint64_t*>(token_ptr + lane_id * 8); |
| 220 | const uint8_t sf_byte = token_ptr[kSfNopeOff + nvfp4_sf_byte(lane_id)]; |
| 221 | __nv_bfloat16 out[16]; |
| 222 | nvfp4_unpack16_e2m1(raw, sf_byte, out); |
| 223 | int4* nope_dst = reinterpret_cast<int4*>(dst_ptr) + lane_id * 2; |
| 224 | nope_dst[0] = *reinterpret_cast<const int4*>(&out[0]); |
| 225 | nope_dst[1] = *reinterpret_cast<const int4*>(&out[8]); |
| 226 | } |
| 227 | |
| 228 | // RoPE is unscaled e4m3. |
| 229 | { |
| 230 | const __nv_fp8_e4m3* raw = |
| 231 | reinterpret_cast<const __nv_fp8_e4m3*>(token_ptr + 256 + lane_id * 2); |
| 232 | const __nv_bfloat16 out[2] = {__float2bfloat16_rn(float(raw[0])), |
| 233 | __float2bfloat16_rn(float(raw[1]))}; |
| 234 | *reinterpret_cast<uint32_t*>(dst_ptr + 512 + lane_id * 2) = |
| 235 | *reinterpret_cast<const uint32_t*>(out); |
| 236 | } |
| 237 | } |
| 238 | |
| 239 | void launch_concat_and_cache_nvfp4_ds_mla(const void* kv_c, const void* k_pe, |
| 240 | void* kv_cache, |
| 241 | const int64_t* slot_mapping, |
| 242 | int block_stride, int entry_stride, |
| 243 | int kv_c_stride, int k_pe_stride, |
| 244 | int block_size, int num_tokens, |
| 245 | cudaStream_t stream) { |
| 246 | dim3 grid(num_tokens); |
| 247 | dim3 block(64); // warp 0: 32 NoPE tiles; warp 1: 4 RoPE tiles |
| 248 | concat_and_cache_nvfp4_ds_mla_kernel<<<grid, block, 0, stream>>>( |
| 249 | reinterpret_cast<const __nv_bfloat16*>(kv_c), |
| 250 | reinterpret_cast<const __nv_bfloat16*>(k_pe), |
| 251 | reinterpret_cast<uint8_t*>(kv_cache), slot_mapping, block_stride, |
| 252 | entry_stride, kv_c_stride, k_pe_stride, block_size); |
| 253 | } |
| 254 | |
| 255 | void launch_cp_gather_and_upconvert_nvfp4_kv_cache( |
| 256 | const uint8_t* src_cache, void* dst, const int32_t* block_table, |
| 257 | const int32_t* workspace_starts, int32_t num_reqs, int32_t block_size, |
| 258 | int32_t total_tokens, int64_t block_table_stride, |
| 259 | int64_t cache_block_stride, int64_t cache_entry_stride, |
| 260 | int64_t dst_entry_stride, cudaStream_t stream) { |
| 261 | constexpr int warps_per_block = 8; |
| 262 | const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block; |
| 263 | const int block_size_threads = warps_per_block * 32; // 256 threads |
| 264 | cp_gather_and_upconvert_nvfp4_kv_cache_kernel<<<grid_size, block_size_threads, |
| 265 | 0, stream>>>( |
| 266 | src_cache, reinterpret_cast<__nv_bfloat16*>(dst), block_table, |
| 267 | workspace_starts, num_reqs, block_size, total_tokens, block_table_stride, |
| 268 | cache_block_stride, cache_entry_stride, dst_entry_stride); |
| 269 | } |
| 270 | |
| 271 | } // namespace vllm |