feat: add FP8 blockwise dequantize/GEMM operators and FP8(E4M3) KV cache support in paged attention

include/infinicore/ops.hpp

@@ -33,6 +33,8 @@
 #include "ops/flash_attention.hpp"
 #include "ops/fmin.hpp"
 #include "ops/fmod.hpp"
+#include "ops/fp8_blockwise_dequantize.hpp"
+#include "ops/fp8_blockwise_gemm.hpp"
 #include "ops/fp8_indexer_logits.hpp"
 #include "ops/fp8_indexer_quant.hpp"
 #include "ops/fp8_mla_rmsnorm_cache.hpp"

include/infinicore/ops/fp8_blockwise_dequantize.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(Fp8BlockwiseDequantize, Tensor, const Tensor &, const Tensor &);
+
+Tensor fp8_blockwise_dequantize(const Tensor &q,
+                                const Tensor &scales,
+                                const DataType &output_dtype);
+void fp8_blockwise_dequantize_(Tensor output,
+                               const Tensor &q,
+                               const Tensor &scales);
+
+} // namespace infinicore::op

include/infinicore/ops/fp8_blockwise_gemm.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(Fp8BlockwiseGemm, Tensor, const Tensor &, const Tensor &, const Tensor &);
+
+Tensor fp8_blockwise_gemm(const Tensor &a,
+                          const Tensor &q,
+                          const Tensor &scales);
+void fp8_blockwise_gemm_(Tensor output,
+                         const Tensor &a,
+                         const Tensor &q,
+                         const Tensor &scales);
+
+} // namespace infinicore::op

include/infinicore/ops/paged_attention.hpp

@@ -7,14 +7,16 @@
 
 namespace infinicore::op {
 
-INFINICORE_GRAPH_OP_CLASS(PagedAttention, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, std::optional<Tensor>, float);
+INFINICORE_GRAPH_OP_CLASS(PagedAttention, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, std::optional<Tensor>, float, std::optional<Tensor>, std::optional<Tensor>);
 
 Tensor paged_attention(const Tensor &q, const Tensor &k_cache, const Tensor &v_cache,
                        const Tensor &block_tables, const Tensor &kv_lens,
-                       std::optional<Tensor> alibi_slopes, float scale);
+                       std::optional<Tensor> alibi_slopes, float scale,
+                       std::optional<Tensor> k_scale = std::nullopt, std::optional<Tensor> v_scale = std::nullopt);
 
 void paged_attention_(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache,
                       const Tensor &block_tables, const Tensor &kv_lens,
-                      std::optional<Tensor> alibi_slopes, float scale);
+                      std::optional<Tensor> alibi_slopes, float scale,
+                      std::optional<Tensor> k_scale = std::nullopt, std::optional<Tensor> v_scale = std::nullopt);
 
 } // namespace infinicore::op

include/infinicore/ops/paged_attention_prefill.hpp

@@ -20,12 +20,15 @@ class PagedAttentionPrefill {
      * 7. cu_seqlens_q: Cumulative sequence lengths of Query (prefix sum for variable-length batch)
      * 8. alibi_slopes: ALiBi bias slopes (optional)
      * 9. scale: Scaling factor (typically 1/sqrt(head_size))
+     * 10. k_scale: Per-token dequant scales for FP8 K cache (optional)
+     * 11. v_scale: Per-token dequant scales for FP8 V cache (optional)
      */
-    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, std::optional<Tensor>, float);
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, std::optional<Tensor>, float, std::optional<Tensor>, std::optional<Tensor>);
 
     static void execute(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache,
                         Tensor block_tables, Tensor total_kv_lens, Tensor cum_seqlens_q,
-                        std::optional<Tensor> alibi_slopes, float scale);
+                        std::optional<Tensor> alibi_slopes, float scale,
+                        std::optional<Tensor> k_scale, std::optional<Tensor> v_scale);
 
     static common::OpDispatcher<schema> &dispatcher();
 };
@@ -37,7 +40,9 @@ Tensor paged_attention_prefill(Tensor q,
                                Tensor total_kv_lens,
                                Tensor cum_seqlens_q,
                                std::optional<Tensor> alibi_slopes,
-                               float scale);
+                               float scale,
+                               std::optional<Tensor> k_scale = std::nullopt,
+                               std::optional<Tensor> v_scale = std::nullopt);
 
 void paged_attention_prefill_(Tensor out,
                               Tensor q,
@@ -47,6 +52,8 @@ void paged_attention_prefill_(Tensor out,
                               Tensor total_kv_lens,
                               Tensor cum_seqlens_q,
                               std::optional<Tensor> alibi_slopes,
-                              float scale);
+                              float scale,
+                              std::optional<Tensor> k_scale = std::nullopt,
+                              std::optional<Tensor> v_scale = std::nullopt);
 
 } // namespace infinicore::op

include/infinicore/ops/paged_caching.hpp

@@ -3,11 +3,13 @@
 #include "../device.hpp"
 #include "../graph/graph.hpp"
 #include "common/op.hpp"
+#include <optional>
 
 namespace infinicore::op {
 
-INFINICORE_GRAPH_OP_CLASS(PagedCaching, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &);
+INFINICORE_GRAPH_OP_CLASS(PagedCaching, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &, std::optional<Tensor>, std::optional<Tensor>);
 
-void paged_caching_(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping);
+void paged_caching_(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping,
+                    std::optional<Tensor> k_scale = std::nullopt, std::optional<Tensor> v_scale = std::nullopt);
 
 } // namespace infinicore::op

include/infiniop.h

@@ -56,6 +56,8 @@
 #include "infiniop/ops/floor_divide.h"
 #include "infiniop/ops/fmin.h"
 #include "infiniop/ops/fmod.h"
+#include "infiniop/ops/fp8_blockwise_dequantize.h"
+#include "infiniop/ops/fp8_blockwise_gemm.h"
 #include "infiniop/ops/fp8_indexer_logits.h"
 #include "infiniop/ops/fp8_indexer_quant.h"
 #include "infiniop/ops/fp8_mla_rmsnorm_cache.h"

include/infiniop/ops/fp8_blockwise_dequantize.h

@@ -0,0 +1,42 @@
+#ifndef __INFINIOP_FP8_BLOCKWISE_DEQUANTIZE_API_H__
+#define __INFINIOP_FP8_BLOCKWISE_DEQUANTIZE_API_H__
+
+#include "../operator_descriptor.h"
+
+/**
+ * Dequantize a 2D FP8 E4M3FN blockwise-quantized weight tensor.
+ *
+ * The quantized input has shape [M, N] and dtype F8 (E4M3FN, stored as raw
+ * bytes). Scales have shape [M / BM, N / BN] and dtype F32, where BM and BN
+ * are the block sizes inferred from the shapes (typically 128 x 128). The
+ * output has shape [M, N] and may be FP16, BF16, or FP32. All tensors must be
+ * contiguous, M must be divisible by BM, and N must be divisible by BN.
+ *
+ * out[i, j] = fp8_e4m3_decode(q[i, j]) * scales[i / BM, j / BN]
+ */
+typedef struct InfiniopDescriptor *infiniopFp8BlockwiseDequantizeDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateFp8BlockwiseDequantizeDescriptor(
+    infiniopHandle_t handle,
+    infiniopFp8BlockwiseDequantizeDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t q_desc,
+    infiniopTensorDescriptor_t scales_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetFp8BlockwiseDequantizeWorkspaceSize(
+    infiniopFp8BlockwiseDequantizeDescriptor_t desc,
+    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopFp8BlockwiseDequantize(
+    infiniopFp8BlockwiseDequantizeDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *out,
+    const void *q,
+    const void *scales,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyFp8BlockwiseDequantizeDescriptor(
+    infiniopFp8BlockwiseDequantizeDescriptor_t desc);
+
+#endif

include/infiniop/ops/fp8_blockwise_gemm.h

@@ -0,0 +1,50 @@
+#ifndef __INFINIOP_FP8_BLOCKWISE_GEMM_API_H__
+#define __INFINIOP_FP8_BLOCKWISE_GEMM_API_H__
+
+#include "../operator_descriptor.h"
+
+/**
+ * Fused GEMM for FP8 E4M3FN blockwise-quantized weights (decode-oriented).
+ *
+ * Computes out = a @ dequantize(q, scales)^T without materializing the
+ * dequantized weight:
+ *
+ *   out[m, n] = sum_k a[m, k] * fp8_e4m3_decode(q[n, k]) * scales[n / BN, k / BK]
+ *
+ * - out: [M, N], dtype F16/BF16/F32, contiguous
+ * - a:   [M, K], same dtype as out, contiguous
+ * - q:   [N, K], dtype F8 (E4M3FN raw bytes), contiguous
+ * - scales: [N / BN, K / BK], dtype F32, contiguous (typically BN = BK = 128)
+ *
+ * N must be divisible by the scale row count and K by the scale col count.
+ * The operator is optimized for small M (decode); large M still works but is
+ * not the target use case.
+ */
+typedef struct InfiniopDescriptor *infiniopFp8BlockwiseGemmDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateFp8BlockwiseGemmDescriptor(
+    infiniopHandle_t handle,
+    infiniopFp8BlockwiseGemmDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t a_desc,
+    infiniopTensorDescriptor_t q_desc,
+    infiniopTensorDescriptor_t scales_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetFp8BlockwiseGemmWorkspaceSize(
+    infiniopFp8BlockwiseGemmDescriptor_t desc,
+    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopFp8BlockwiseGemm(
+    infiniopFp8BlockwiseGemmDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *out,
+    const void *a,
+    const void *q,
+    const void *scales,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyFp8BlockwiseGemmDescriptor(
+    infiniopFp8BlockwiseGemmDescriptor_t desc);
+
+#endif

include/infiniop/ops/paged_attention.h

@@ -27,6 +27,10 @@ typedef struct InfiniopDescriptor *infiniopPagedAttentionDescriptor_t;
  * Expected DType: int64_t (I64).
  * @param alibi_slopes_desc [Optional] Shape: (num_heads,).
  * Slopes for ALiBi (Attention with Linear Biases). Can be NULL.
+ * @param k_scale_desc [Optional] Shape: (num_blocks, num_kv_heads, block_size).
+ * Per-token dequant scales for the key cache. DType: F32.
+ * Required (non-NULL) iff k_cache/v_cache are F8; must be NULL otherwise.
+ * @param v_scale_desc [Optional] Same layout and rules as k_scale_desc.
  * @param scale     The attention scaling factor (typically 1/sqrt(head_size)).
  * @return infiniStatus_t Status code.
  */
@@ -40,6 +44,8 @@ __INFINI_C __export infiniStatus_t infiniopCreatePagedAttentionDescriptor(
     infiniopTensorDescriptor_t block_tables_desc,
     infiniopTensorDescriptor_t seq_lens_desc,
     infiniopTensorDescriptor_t alibi_slopes_desc,
+    infiniopTensorDescriptor_t k_scale_desc,
+    infiniopTensorDescriptor_t v_scale_desc,
     float scale);
 
 /**
@@ -65,6 +71,8 @@ __INFINI_C __export infiniStatus_t infiniopGetPagedAttentionWorkspaceSize(
  * @param block_tables Pointer to the block tables data.
  * @param seq_lens Pointer to the sequence lengths data.
  * @param alibi_slopes Pointer to the ALiBi slopes data. Can be NULL.
+ * @param k_scale Pointer to the per-token key dequant scales (F8 caches only). Can be NULL.
+ * @param v_scale Pointer to the per-token value dequant scales (F8 caches only). Can be NULL.
  * @param stream The CUDA stream for the operation. Can be NULL.
  * @return infiniStatus_t Status code of the operation.
  */
@@ -79,6 +87,8 @@ __INFINI_C __export infiniStatus_t infiniopPagedAttention(
     const void *block_tables,
     const void *seq_lens,
     const void *alibi_slopes,
+    const void *k_scale,
+    const void *v_scale,
     void *stream);
 
 /**

include/infiniop/ops/paged_attention_prefill.h

@@ -26,6 +26,11 @@ typedef struct InfiniopDescriptor *infiniopPagedAttentionPrefillDescriptor_t;
  * Shape: [batch_size + 1]
  * @param alibi_slopes_desc Optional descriptor for the ALiBi slopes tensor. Can be NULL.
  * Shape: [num_heads]
+ * @param k_scale_desc Optional descriptor for the per-token key dequant scales.
+ * Shape: [max_num_blocks, num_kv_heads, block_size], DType: F32.
+ * Required (non-NULL) iff k_cache/v_cache are F8; must be NULL otherwise.
+ * @param v_scale_desc Optional descriptor for the per-token value dequant scales.
+ * Same layout and rules as k_scale_desc.
  * @param scale The attention scaling factor (typically 1.0 / sqrt(head_size)).
  * @return infiniStatus_t Status code of the operation.
  */
@@ -40,6 +45,8 @@ __INFINI_C __export infiniStatus_t infiniopCreatePagedAttentionPrefillDescriptor
     infiniopTensorDescriptor_t seq_lens_desc,
     infiniopTensorDescriptor_t cum_seq_lens_q_desc,
     infiniopTensorDescriptor_t alibi_slopes_desc,
+    infiniopTensorDescriptor_t k_scale_desc,
+    infiniopTensorDescriptor_t v_scale_desc,
     float scale);
 
 /**
@@ -61,6 +68,8 @@ __INFINI_C __export infiniStatus_t infiniopGetPagedAttentionPrefillWorkspaceSize
  * @param seq_lens Pointer to the KV lengths data.
  * @param cum_seq_lens_q Pointer to the Q cumulative sequence lengths data (prefix sum).
  * @param alibi_slopes Pointer to the ALiBi slopes data. Can be NULL.
+ * @param k_scale Pointer to the per-token key dequant scales (F8 caches only). Can be NULL.
+ * @param v_scale Pointer to the per-token value dequant scales (F8 caches only). Can be NULL.
  * @param stream The device stream (e.g., cudaStream_t) for the operation.
  * @return infiniStatus_t Status code of the operation.
  */
@@ -76,6 +85,8 @@ __INFINI_C __export infiniStatus_t infiniopPagedAttentionPrefill(
     const void *seq_lens,
     const void *cum_seq_lens_q,
     const void *alibi_slopes,
+    const void *k_scale,
+    const void *v_scale,
     void *stream);
 
 /**

include/infiniop/ops/paged_caching.h

@@ -19,6 +19,11 @@ typedef struct InfiniopDescriptor *infiniopPagedCachingDescriptor_t;
  * @param k_desc Descriptor for the source key tensor.
  * @param v_desc Descriptor for the source value tensor.
  * @param slot_mapping_desc Descriptor for the slot mapping tensor.
+ * @param k_scale_desc [Optional] Descriptor for the per-token key dequant scales.
+ * Shape: [num_blocks, num_kv_heads, block_size], DType: F32.
+ * Required (non-NULL) iff the caches are F8; must be NULL otherwise.
+ * @param v_scale_desc [Optional] Descriptor for the per-token value dequant scales.
+ * Same layout and rules as k_scale_desc.
  * @return infiniStatus_t Status code of the operation.
  */
 __INFINI_C __export infiniStatus_t infiniopCreatePagedCachingDescriptor(
@@ -28,7 +33,9 @@ __INFINI_C __export infiniStatus_t infiniopCreatePagedCachingDescriptor(
     infiniopTensorDescriptor_t v_cache_desc,
     infiniopTensorDescriptor_t k_desc,
     infiniopTensorDescriptor_t v_desc,
-    infiniopTensorDescriptor_t slot_mapping_desc);
+    infiniopTensorDescriptor_t slot_mapping_desc,
+    infiniopTensorDescriptor_t k_scale_desc,
+    infiniopTensorDescriptor_t v_scale_desc);
 
 /**
  * @brief Retrieves the workspace size required for the Paged Caching operation.
@@ -51,6 +58,10 @@ __INFINI_C __export infiniStatus_t infiniopGetPagedCachingWorkspaceSize(
  * @param k Pointer to the source key tensor data.
  * @param v Pointer to the source value tensor data.
  * @param slot_mapping Pointer to the slot mapping data.
+ * @param k_scale [Optional] Pointer to the per-token key dequant scales.
+ * Written by this operator when the caches are F8 (quantization happens here).
+ * Must be NULL when the caches are not F8.
+ * @param v_scale [Optional] Pointer to the per-token value dequant scales.
  * @param stream The CUDA stream for the operation. Can be NULL.
  * @return infiniStatus_t Status code of the operation.
  */
@@ -63,6 +74,8 @@ __INFINI_C __export infiniStatus_t infiniopPagedCaching(
     const void *k,
     const void *v,
     const void *slot_mapping,
+    void *k_scale,
+    void *v_scale,
     void *stream);
 
 /**

python/infinicore/ops/paged_attention.py

@@ -10,6 +10,8 @@ def paged_attention(
     cache_lens: Tensor,
     alibi_slopes: Tensor | None = None,
     scale: float = 1.0,
+    k_scale: Tensor | None = None,
+    v_scale: Tensor | None = None,
     *,
     out: Tensor | None = None,
 ):
@@ -23,6 +25,8 @@ def paged_attention(
                 cache_lens._underlying,
                 alibi_slopes._underlying if alibi_slopes is not None else None,
                 scale,
+                k_scale._underlying if k_scale is not None else None,
+                v_scale._underlying if v_scale is not None else None,
             )
         )
 
@@ -35,6 +39,8 @@ def paged_attention(
         cache_lens._underlying,
         alibi_slopes._underlying if alibi_slopes is not None else None,
         scale,
+        k_scale._underlying if k_scale is not None else None,
+        v_scale._underlying if v_scale is not None else None,
     )
 
     return out