introduce kernel for converting e4m3fn kv_cache to e4m3fnuz (#4113)

Summary:

X-link: facebookresearch/FBGEMM#1197

As titled. NV uses e4m3fn for fp8 KV cache, and AMD uses e4m3fnuz. If we want to do HH serving with NV prefill and AMD decode, we need a way to convert between the two.

In practice this can be done with two operations: overwriting -0 (which are valid in e4m3fn but not in uz) with +0, and multiplying the scale qparam by 2, since the exponent bias differs by 1 between the two formats. The scale qparam is packed as the low bits of a __half2 (fp16x2).

Differential Revision: D74502334

Author: bradleyhd

fbgemm_gpu/experimental/gen_ai/src/kv_cache/kv_cache.cpp

@@ -49,6 +49,8 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
           DEFAULT_PAGE_SIZE) ") -> (Tensor, Tensor)");
   m.def(
       "quantize_qkv_per_head(Tensor amax, Tensor XQKV, Tensor varseq_seqpos, Tensor? varseq_batch, Tensor q_seqstarts, Tensor cache_K, Tensor cache_V, Tensor XQ_O, int max_seq_len, Tensor? qparam_k=None, Tensor? qparam_v=None) -> Tensor");
+  m.def(
+      "convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace(Tensor cache_K, Tensor cache_V, Tensor qparam_K, Tensor qparam_V) -> ()");
 }
 
 TORCH_LIBRARY_IMPL(fbgemm, CPU, m) {
@@ -73,6 +75,9 @@ TORCH_LIBRARY_IMPL(fbgemm, CUDA, m) {
   m.impl("dequantize_int4_cache", dequantize_int4_cache);
   m.impl("dequantize_fp8_cache", dequantize_fp8_cache);
   m.impl("quantize_qkv_per_head", quantize_qkv_per_head);
+  m.impl(
+      "convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace",
+      fbgemm_gpu::convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace);
 }
 
 at::Tensor rope_qkv_varseq_prefill_meta(
@@ -295,6 +300,13 @@ at::Tensor quantize_qkv_per_head_meta(
       at::empty_symint({B_KV, N_KVH}, cache_K.options().dtype(at::kFloat));
   return at::empty_like(XQKV);
 }
+
+void convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace_meta(
+    at::Tensor cache_K,
+    at::Tensor cache_V,
+    at::Tensor qparam_K,
+    at::Tensor qparam_v) {};
+
 TORCH_LIBRARY_IMPL(fbgemm, Meta, m) {
   m.impl("rope_qkv_varseq_prefill", rope_qkv_varseq_prefill_meta);
   m.impl("rope_qkv_decoding", rope_qkv_decoding_meta);
@@ -304,6 +316,9 @@ TORCH_LIBRARY_IMPL(fbgemm, Meta, m) {
   m.impl("xpos_qkv_decoding", xpos_qkv_decoding_meta);
   m.impl("dequantize_int4_cache", dequantize_int4_cache_meta);
   m.impl("dequantize_fp8_cache", dequantize_fp8_cache_meta);
+  m.impl(
+      "convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace",
+      convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace_meta);
 }
 
 } // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/kv_cache/kv_cache.cu

@@ -72,7 +72,6 @@ __global__ void dequantize_int4_cache_kernel(
         cache_V_dq // [B][MAX_T][N_KVH][D_H]
 ) {
   auto N_KVH = cache_K.size(2);
-  auto MAX_T = cache_K.size(1);
   auto D_H = cache_K_dq.size(3);
 
   auto b = blockIdx.x;
@@ -2331,6 +2330,112 @@ at::Tensor xpos_qkv_decoding(
   return XQ_O;
 }
 
+#if (defined(USE_ROCM))
+/**
+ * Converts the contents of a FP8 KV cache from e4m3fn (NV) to e4m3fnuz (AMD).
+ * These formats differ in their support for negative zero, and in their
+ * exponent bias. Negative zeros are replaced with positive zero, and the scale
+ * qparam is multiplied by 2.0, because we know that the scale will be applied
+ * to the k/v value and is equivalent to recomputing the exponent bias.
+ *
+ * This in an inplace operation.
+ *
+ * It is assumed that inputs will have been generated with scale_ub = max(fp16)
+ * / 2 to avoid overflow. Some debug mode assertions are in place, but there are
+ * no runtime guarantees.
+ *
+ * As written, this kernel is only valid on AMD, because it relies on threads
+ * 32-63 to convert the V tensors. NV only has threads 0-31 per warp.
+ */
+__global__ void convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace_kernel(
+    at::PackedTensorAccessor64<uint8_t, 4, at::RestrictPtrTraits>
+        cache_K, // [B][MAX_T][N_KVH][D_H]
+    at::PackedTensorAccessor64<uint8_t, 4, at::RestrictPtrTraits>
+        cache_V, // [B][MAX_T][N_KVH][D_H]
+    at::PackedTensorAccessor64<uint32_t, 4, at::RestrictPtrTraits> qparam_K,
+    at::PackedTensorAccessor64<uint32_t, 4, at::RestrictPtrTraits> qparam_V) {
+  auto N_KVH = cache_K.size(2);
+  auto MAX_T = cache_K.size(1);
+  auto D_H = cache_K.size(3);
+  CUDA_KERNEL_ASSERT(D_H == 128);
+
+  auto b = blockIdx.x;
+  int h = 0, t = 0;
+  uint8_t* head;
+  __half* qparam_scale_ptr;
+
+  for (auto t_h = threadIdx.y + blockIdx.y * blockDim.y; t_h < MAX_T * N_KVH;
+       t_h += blockDim.y * gridDim.y) {
+    h = t_h % N_KVH;
+    t = t_h / N_KVH;
+
+    auto tidx = threadIdx.x;
+    if (threadIdx.x < 32) {
+      head = &cache_K[b][t][h][0];
+      qparam_scale_ptr = reinterpret_cast<__half*>(&qparam_K[b][t][h][0]);
+    } else {
+      head = &cache_V[b][t][h][0];
+      qparam_scale_ptr = reinterpret_cast<__half*>(&qparam_V[b][t][h][0]);
+      tidx -= 32;
+    }
+    auto D_H_idx = tidx * 4; // Reading 4 bytes at once.
+
+    // Our only goal here is to detect negative zeros that are valid
+    // in e4m3fn, but not valid in e4m3fnuz, and overwrite them with zeros.
+    // E.g. 0x80708070u ^ 0x80808080u = 0xf000f0u & 0x80708070u = 0x00700070u
+    uint32_t packed_fp8x4_vals = *reinterpret_cast<uint32_t*>(&head[D_H_idx]);
+    *reinterpret_cast<uint32_t*>(&head[D_H_idx]) =
+        (packed_fp8x4_vals ^ 0x80808080) & packed_fp8x4_vals;
+
+    // Multiply qparam scale (member x) by 2 to compensate for the exponent
+    // bias difference (1) between e4m3fn and e4m3fnuz. We only need to do this
+    // once per row. In debug mode, assert that 2.0*scale as a float would not
+    // exceed the max value of __half.
+    if (tidx == 0) {
+      CUDA_KERNEL_ASSERT(
+          __half2float(*qparam_scale_ptr) * 2.0f <=
+          __half2float(std::numeric_limits<__half>::max()));
+      *qparam_scale_ptr = __float2half(__half2float(*qparam_scale_ptr) * 2.0f);
+    }
+  }
+}
+
+void convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace(
+    at::Tensor cache_K,
+    at::Tensor cache_V,
+    at::Tensor qparam_K,
+    at::Tensor qparam_V) {
+  TORCH_CHECK(cache_K.is_cuda());
+  TORCH_CHECK(cache_V.is_cuda());
+
+  auto B = cache_K.size(0);
+
+  constexpr int32_t kMaxBlocks = 1;
+  dim3 blocks(B, std::max<int32_t>(1, kMaxBlocks / B));
+  dim3 threads(kThreadsPerWarp, kWarpsPerBlock);
+
+  convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace_kernel<<<
+      blocks,
+      threads,
+      0,
+      at::cuda::getCurrentCUDAStream()>>>(
+      cache_K.packed_accessor64<uint8_t, 4, at::RestrictPtrTraits>(),
+      cache_V.packed_accessor64<uint8_t, 4, at::RestrictPtrTraits>(),
+      qparam_K.packed_accessor64<uint32_t, 4, at::RestrictPtrTraits>(),
+      qparam_V.packed_accessor64<uint32_t, 4, at::RestrictPtrTraits>());
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+#else
+void convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace(
+    at::Tensor cache_K,
+    at::Tensor cache_V,
+    at::Tensor qparam_K,
+    at::Tensor qparam_V) {
+  throw std::runtime_error(
+      "convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace is only supported on AMD");
+}
+#endif
+
 #if (defined(USE_ROCM) && ROCM_VERSION >= 60200) || \
     (defined(CUDA_VERSION) && CUDA_VERSION >= 12000)
 
@@ -2550,7 +2655,6 @@ __global__ void dequantize_fp8_cache_kernel_paged(
     int32_t block_tables_b_stride,
     int32_t page_size) {
   auto N_KVH = cache_K.size(2);
-  auto MAX_T = cache_K.size(1);
   auto D_H = cache_K_dq.size(3);
   auto D_H_q = cache_K.size(3);
   CUDA_KERNEL_ASSERT(D_H == 128);
@@ -2647,20 +2751,20 @@ std::tuple<at::Tensor, at::Tensor> dequantize_fp8_cache(
   auto D_H = (D_HQ - fp8_qparam_offset);
 
   // TODO:
-  // The below allocates Tensors that have the same shape as cache_K and cache_V
-  // to store their dequantize results. For paged KV cache, this can be a bit
-  // inefficient because it has the shape of [1 x (MAX_PAGES * PAGE_SIZE) x
-  // N_KVH x D_H] to accommodate pages globally across batch instances, and
-  // if we have very large MAX_PAGES then we are essentially allocating a very
-  // huge Tensor here. The benefit is that the following users of this
-  // dequantized results can reuse the existing block_tables to access their
-  // elements. If we want to be more efficient, there are two possible
-  // approaches: (1) Allocate a shorter Tensor here and store the dequantize
-  // results in a more compact manner, but that requires creating a new
-  // block_tables here and making sure the following users all use the
-  // correct block_tables. (2) From outside, keep a persistent buffer that has a
-  // matching shape with the original paged KV and feed the same buffer
-  // into this function at every layer to reuse it and prevent allocation.
+  // The below allocates Tensors that have the same shape as cache_K and
+  // cache_V to store their dequantize results. For paged KV cache, this can
+  // be a bit inefficient because it has the shape of [1 x (MAX_PAGES *
+  // PAGE_SIZE) x N_KVH x D_H] to accommodate pages globally across batch
+  // instances, and if we have very large MAX_PAGES then we are essentially
+  // allocating a very huge Tensor here. The benefit is that the following
+  // users of this dequantized results can reuse the existing block_tables to
+  // access their elements. If we want to be more efficient, there are two
+  // possible approaches: (1) Allocate a shorter Tensor here and store the
+  // dequantize results in a more compact manner, but that requires creating a
+  // new block_tables here and making sure the following users all use the
+  // correct block_tables. (2) From outside, keep a persistent buffer that has
+  // a matching shape with the original paged KV and feed the same buffer into
+  // this function at every layer to reuse it and prevent allocation.
 
   auto cache_K_dq = at::empty(
       {B_KV, MAX_T, N_KVH, D_H}, cache_K.options().dtype(at::kBFloat16));
@@ -2889,8 +2993,8 @@ at::Tensor quantize_qkv_per_head(
   float* const scale_q_ptr = scale_q.data_ptr<float>();
   // Launch the kernel
   // TODO: Launch the kernel with B_T * N_H_L blocks only in case of decode.
-  // Currently, we are launching the kernel with B_T * HH blocks for decode and
-  // KV blocks just return when qparam_k is passed as nullptr.
+  // Currently, we are launching the kernel with B_T * HH blocks for decode
+  // and KV blocks just return when qparam_k is passed as nullptr.
   quantizeQKVPerHead<<<
       grid_size,
       block_size,

fbgemm_gpu/experimental/gen_ai/src/kv_cache/kv_cache.h

@@ -193,4 +193,10 @@ at::Tensor mqa_attn(
     std::optional<at::Tensor> qparam_k,
     std::optional<at::Tensor> qparam_v);
 
+void convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace(
+    at::Tensor cache_K,
+    at::Tensor cache_V,
+    at::Tensor qparam_K,
+    at::Tensor qparam_v);
+
 } // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/test/kv_cache/kv_cache_test.py

@@ -70,6 +70,7 @@ def _get_varseq_batch_seqpos(
     return varseq_batch, varseq_seqpos
 
 
+@unittest.skipIf(not torch.cuda.is_available(), "Skipping because no GPU is available.")
 class KVCacheTests(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
@@ -353,6 +354,85 @@ def test_fp8_kv_cache(self, MAX_T: int, N_KVH_L: int) -> None:
             cache_v[:, :T], cache_v_bf16[:, :T], atol=1.0e-2, rtol=5.0e-2
         )
 
+    @settings(deadline=None)
+    @given(
+        MAX_T=st.sampled_from([8000, 16384]),
+        N_KVH_L=st.sampled_from([1, 2]),
+    )
+    @unittest.skipIf(
+        not torch.cuda.is_available() or not HAS_XFORMERS or not torch.version.hip,
+        "Skip when no AMD GPU or xformers is not available",
+    )
+    def test_fp8_kv_e4m3fn_convert_to_e4m3fnuz(self, MAX_T: int, N_KVH_L: int) -> None:
+        B = 2
+        D_H = 128
+
+        qparam_offset = 0
+
+        # Init K to all 1s
+        cache_fp8x4_k = torch.full(
+            size=(B, MAX_T, N_KVH_L, int(D_H) + qparam_offset),
+            fill_value=0x01,
+            dtype=torch.uint8,
+            device=self.device,
+        )
+
+        # Choose random elements to set to negative zero, 0x80
+        cache_fp8x4_k_flat = cache_fp8x4_k.flatten()
+        random_indices = torch.randperm(cache_fp8x4_k_flat.size(0))[
+            : cache_fp8x4_k_flat.size(0) // 2
+        ]
+        cache_fp8x4_k_flat[random_indices] = 0x80
+        cache_fp8x4_k = cache_fp8x4_k_flat.reshape(cache_fp8x4_k.shape)
+
+        # Expected K has +0 in place of -0
+        cache_fp8x4_k_expected_flat = cache_fp8x4_k_flat.clone()
+        cache_fp8x4_k_expected_flat[random_indices] = 0x00
+        cache_fp8x4_k_expected = cache_fp8x4_k_expected_flat.reshape(
+            cache_fp8x4_k.shape
+        )
+
+        # Repeat for V
+        cache_fp8x4_v = torch.full(
+            size=(B, MAX_T, N_KVH_L, int(D_H) + qparam_offset),
+            fill_value=0x01,
+            dtype=torch.uint8,
+            device=self.device,
+        )
+        cache_fp8x4_v_flat = cache_fp8x4_v.flatten()
+        random_indices = torch.randperm(cache_fp8x4_v_flat.size(0))[
+            : cache_fp8x4_v_flat.size(0) // 2
+        ]
+        cache_fp8x4_v_flat[random_indices] = 0x80
+        cache_fp8x4_v = cache_fp8x4_v_flat.reshape(cache_fp8x4_v.shape)
+        cache_fp8x4_v_expected_flat = cache_fp8x4_v_flat.clone()
+        cache_fp8x4_v_expected_flat[random_indices] = 0x00
+        cache_fp8x4_v_expected = cache_fp8x4_v_expected_flat.reshape(
+            cache_fp8x4_v.shape
+        )
+
+        qparam_fp16x2_k = torch.full(
+            size=(B, MAX_T, N_KVH_L, 1), fill_value=0x3C003C00, dtype=torch.uint32
+        ).cuda()
+        qparam_fp16x2_k_expected = torch.full(
+            size=(B, MAX_T, N_KVH_L, 1), fill_value=0x3C004000, dtype=torch.uint32
+        ).cuda()
+        qparam_fp16x2_v = torch.full(
+            size=(B, MAX_T, N_KVH_L, 1), fill_value=0x3C003C00, dtype=torch.uint32
+        ).cuda()
+        qparam_fp16x2_v_expected = torch.full(
+            size=(B, MAX_T, N_KVH_L, 1), fill_value=0x3C004000, dtype=torch.uint32
+        ).cuda()
+
+        torch.ops.fbgemm.convert_e4m3fn_kv_cache_to_e4m3fnuz_inplace(
+            cache_fp8x4_k, cache_fp8x4_v, qparam_fp16x2_k, qparam_fp16x2_v
+        )
+
+        assert torch.equal(cache_fp8x4_k, cache_fp8x4_k_expected)
+        assert torch.equal(cache_fp8x4_v, cache_fp8x4_v_expected)
+        assert torch.equal(qparam_fp16x2_k, qparam_fp16x2_k_expected)
+        assert torch.equal(qparam_fp16x2_v, qparam_fp16x2_v_expected)
+
     @settings(deadline=None)
     @given(
         prefill=st.booleans(),