diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index 5e93e1457..664766b61 100644 --- a/include/infinicore/ops.hpp +++ b/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" diff --git a/include/infinicore/ops/fp8_blockwise_dequantize.hpp b/include/infinicore/ops/fp8_blockwise_dequantize.hpp new file mode 100644 index 000000000..db77385ca --- /dev/null +++ b/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 diff --git a/include/infinicore/ops/fp8_blockwise_gemm.hpp b/include/infinicore/ops/fp8_blockwise_gemm.hpp new file mode 100644 index 000000000..080748a73 --- /dev/null +++ b/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 diff --git a/include/infinicore/ops/paged_attention.hpp b/include/infinicore/ops/paged_attention.hpp index 8c906c95e..ba8e91e6c 100644 --- a/include/infinicore/ops/paged_attention.hpp +++ b/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, float); +INFINICORE_GRAPH_OP_CLASS(PagedAttention, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, std::optional, float, std::optional, std::optional); Tensor paged_attention(const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, const Tensor &block_tables, const Tensor &kv_lens, - std::optional alibi_slopes, float scale); + std::optional alibi_slopes, float scale, + std::optional k_scale = std::nullopt, std::optional 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 alibi_slopes, float scale); + std::optional alibi_slopes, float scale, + std::optional k_scale = std::nullopt, std::optional v_scale = std::nullopt); } // namespace infinicore::op diff --git a/include/infinicore/ops/paged_attention_prefill.hpp b/include/infinicore/ops/paged_attention_prefill.hpp index 952924528..5ff649a76 100644 --- a/include/infinicore/ops/paged_attention_prefill.hpp +++ b/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, float); + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, std::optional, float, std::optional, std::optional); 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 alibi_slopes, float scale); + std::optional alibi_slopes, float scale, + std::optional k_scale, std::optional v_scale); static common::OpDispatcher &dispatcher(); }; @@ -37,7 +40,9 @@ Tensor paged_attention_prefill(Tensor q, Tensor total_kv_lens, Tensor cum_seqlens_q, std::optional alibi_slopes, - float scale); + float scale, + std::optional k_scale = std::nullopt, + std::optional 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 alibi_slopes, - float scale); + float scale, + std::optional k_scale = std::nullopt, + std::optional v_scale = std::nullopt); } // namespace infinicore::op diff --git a/include/infinicore/ops/paged_caching.hpp b/include/infinicore/ops/paged_caching.hpp index 403b4b738..ec5e9c23d 100644 --- a/include/infinicore/ops/paged_caching.hpp +++ b/include/infinicore/ops/paged_caching.hpp @@ -3,11 +3,13 @@ #include "../device.hpp" #include "../graph/graph.hpp" #include "common/op.hpp" +#include 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, std::optional); -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 k_scale = std::nullopt, std::optional v_scale = std::nullopt); } // namespace infinicore::op diff --git a/include/infiniop.h b/include/infiniop.h index 9f632e27f..df334c286 100644 --- a/include/infiniop.h +++ b/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" diff --git a/include/infiniop/ops/fp8_blockwise_dequantize.h b/include/infiniop/ops/fp8_blockwise_dequantize.h new file mode 100644 index 000000000..cf9a50398 --- /dev/null +++ b/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 diff --git a/include/infiniop/ops/fp8_blockwise_gemm.h b/include/infiniop/ops/fp8_blockwise_gemm.h new file mode 100644 index 000000000..11817d158 --- /dev/null +++ b/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 diff --git a/include/infiniop/ops/paged_attention.h b/include/infiniop/ops/paged_attention.h index 7f1656ef3..f980ec533 100644 --- a/include/infiniop/ops/paged_attention.h +++ b/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); /** diff --git a/include/infiniop/ops/paged_attention_prefill.h b/include/infiniop/ops/paged_attention_prefill.h index e2e93076b..0629c1ce9 100644 --- a/include/infiniop/ops/paged_attention_prefill.h +++ b/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); /** diff --git a/include/infiniop/ops/paged_caching.h b/include/infiniop/ops/paged_caching.h index d85125c30..ee8892e32 100644 --- a/include/infiniop/ops/paged_caching.h +++ b/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); /** diff --git a/python/infinicore/ops/paged_attention.py b/python/infinicore/ops/paged_attention.py index dfefa6a76..084b92127 100644 --- a/python/infinicore/ops/paged_attention.py +++ b/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 diff --git a/python/infinicore/ops/paged_attention_prefill.py b/python/infinicore/ops/paged_attention_prefill.py index 848f74abf..3eed335ab 100644 --- a/python/infinicore/ops/paged_attention_prefill.py +++ b/python/infinicore/ops/paged_attention_prefill.py @@ -11,10 +11,14 @@ def paged_attention_prefill( cu_seqlens_q: Tensor, alibi_slopes: Tensor | None = None, scale: float = 1.0, + k_scale: Tensor | None = None, + v_scale: Tensor | None = None, *, out: Tensor | None = None, ): alibi_ptr = alibi_slopes._underlying if alibi_slopes is not None else None + k_scale_ptr = k_scale._underlying if k_scale is not None else None + v_scale_ptr = v_scale._underlying if v_scale is not None else None if out is None: return Tensor( @@ -27,6 +31,8 @@ def paged_attention_prefill( cu_seqlens_q._underlying, alibi_ptr, scale, + k_scale_ptr, + v_scale_ptr, ) ) @@ -40,6 +46,8 @@ def paged_attention_prefill( cu_seqlens_q._underlying, alibi_ptr, scale, + k_scale_ptr, + v_scale_ptr, ) return out diff --git a/python/infinicore/ops/paged_caching.py b/python/infinicore/ops/paged_caching.py index e3b8d63fb..92596d5b6 100644 --- a/python/infinicore/ops/paged_caching.py +++ b/python/infinicore/ops/paged_caching.py @@ -8,6 +8,8 @@ def paged_caching( k: Tensor, v: Tensor, slot_mapping: Tensor, + k_scale: Tensor | None = None, + v_scale: Tensor | None = None, ): Tensor( _infinicore.paged_caching_( @@ -16,6 +18,8 @@ def paged_caching( k._underlying, v._underlying, slot_mapping._underlying, + k_scale._underlying if k_scale is not None else None, + v_scale._underlying if v_scale is not None else None, ) ) return (k_cache, v_cache) diff --git a/src/infinicore/ops/fp8_blockwise_dequantize/fp8_blockwise_dequantize.cc b/src/infinicore/ops/fp8_blockwise_dequantize/fp8_blockwise_dequantize.cc new file mode 100644 index 000000000..d9d872c69 --- /dev/null +++ b/src/infinicore/ops/fp8_blockwise_dequantize/fp8_blockwise_dequantize.cc @@ -0,0 +1,36 @@ +#include "infinicore/ops/fp8_blockwise_dequantize.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Fp8BlockwiseDequantize); + +Fp8BlockwiseDequantize::Fp8BlockwiseDequantize(Tensor output, + const Tensor &q, + const Tensor &scales) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, q, scales); + INFINICORE_GRAPH_OP_DISPATCH(output->device().getType(), output, q, scales); +} + +void Fp8BlockwiseDequantize::execute(Tensor output, + const Tensor &q, + const Tensor &scales) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Fp8BlockwiseDequantize, output, q, scales); +} + +Tensor fp8_blockwise_dequantize(const Tensor &q, + const Tensor &scales, + const DataType &output_dtype) { + auto output = Tensor::empty(q->shape(), output_dtype, q->device()); + Fp8BlockwiseDequantize::execute(output, q, scales); + return output; +} + +void fp8_blockwise_dequantize_(Tensor output, + const Tensor &q, + const Tensor &scales) { + Fp8BlockwiseDequantize::execute(output, q, scales); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/fp8_blockwise_dequantize/fp8_blockwise_dequantize_infiniop.cc b/src/infinicore/ops/fp8_blockwise_dequantize/fp8_blockwise_dequantize_infiniop.cc new file mode 100644 index 000000000..6fd1ba80b --- /dev/null +++ b/src/infinicore/ops/fp8_blockwise_dequantize/fp8_blockwise_dequantize_infiniop.cc @@ -0,0 +1,48 @@ +#include "../../utils.hpp" +#include "../infiniop_impl.hpp" +#include "infinicore/common/hash.hpp" +#include "infinicore/ops/common/cache.hpp" +#include "infinicore/ops/fp8_blockwise_dequantize.hpp" + +#include + +namespace infinicore::op::fp8_blockwise_dequantize_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Fp8BlockwiseDequantize, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor output, q, scales; +}; + +void *plan(Tensor output, const Tensor &q, const Tensor &scales) { + const size_t seed = hash_combine(output, q, scales); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, Fp8BlockwiseDequantize, + seed, output->desc(), q->desc(), scales->desc()); + return new PlannedMeta{ + descriptor, + graph::GraphTensor(output), + graph::GraphTensor(q), + graph::GraphTensor(scales)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopFp8BlockwiseDequantize( + planned->descriptor->desc, + nullptr, 0, + planned->output->data(), + planned->q->data(), + planned->scales->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Fp8BlockwiseDequantize, &plan, &run, &cleanup); + +} // namespace infinicore::op::fp8_blockwise_dequantize_impl::infiniop diff --git a/src/infinicore/ops/fp8_blockwise_gemm/fp8_blockwise_gemm.cc b/src/infinicore/ops/fp8_blockwise_gemm/fp8_blockwise_gemm.cc new file mode 100644 index 000000000..08753f83b --- /dev/null +++ b/src/infinicore/ops/fp8_blockwise_gemm/fp8_blockwise_gemm.cc @@ -0,0 +1,41 @@ +#include "infinicore/ops/fp8_blockwise_gemm.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Fp8BlockwiseGemm); + +Fp8BlockwiseGemm::Fp8BlockwiseGemm(Tensor output, + const Tensor &a, + const Tensor &q, + const Tensor &scales) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, a, q, scales); + INFINICORE_GRAPH_OP_DISPATCH(output->device().getType(), output, a, q, scales); +} + +void Fp8BlockwiseGemm::execute(Tensor output, + const Tensor &a, + const Tensor &q, + const Tensor &scales) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Fp8BlockwiseGemm, output, a, q, scales); +} + +Tensor fp8_blockwise_gemm(const Tensor &a, + const Tensor &q, + const Tensor &scales) { + const auto M = a->size(0); + const auto N = q->size(0); + auto output = Tensor::empty({M, N}, a->dtype(), a->device()); + Fp8BlockwiseGemm::execute(output, a, q, scales); + return output; +} + +void fp8_blockwise_gemm_(Tensor output, + const Tensor &a, + const Tensor &q, + const Tensor &scales) { + Fp8BlockwiseGemm::execute(output, a, q, scales); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/fp8_blockwise_gemm/fp8_blockwise_gemm_infiniop.cc b/src/infinicore/ops/fp8_blockwise_gemm/fp8_blockwise_gemm_infiniop.cc new file mode 100644 index 000000000..2f8ae7bf2 --- /dev/null +++ b/src/infinicore/ops/fp8_blockwise_gemm/fp8_blockwise_gemm_infiniop.cc @@ -0,0 +1,50 @@ +#include "../../utils.hpp" +#include "../infiniop_impl.hpp" +#include "infinicore/common/hash.hpp" +#include "infinicore/ops/common/cache.hpp" +#include "infinicore/ops/fp8_blockwise_gemm.hpp" + +#include + +namespace infinicore::op::fp8_blockwise_gemm_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Fp8BlockwiseGemm, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor output, a, q, scales; +}; + +void *plan(Tensor output, const Tensor &a, const Tensor &q, const Tensor &scales) { + const size_t seed = hash_combine(output, a, q, scales); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, Fp8BlockwiseGemm, + seed, output->desc(), a->desc(), q->desc(), scales->desc()); + return new PlannedMeta{ + descriptor, + graph::GraphTensor(output), + graph::GraphTensor(a), + graph::GraphTensor(q), + graph::GraphTensor(scales)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopFp8BlockwiseGemm( + planned->descriptor->desc, + nullptr, 0, + planned->output->data(), + planned->a->data(), + planned->q->data(), + planned->scales->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Fp8BlockwiseGemm, &plan, &run, &cleanup); + +} // namespace infinicore::op::fp8_blockwise_gemm_impl::infiniop diff --git a/src/infinicore/ops/paged_attention/paged_attention.cc b/src/infinicore/ops/paged_attention/paged_attention.cc index 60de2ae66..e8270e9c7 100644 --- a/src/infinicore/ops/paged_attention/paged_attention.cc +++ b/src/infinicore/ops/paged_attention/paged_attention.cc @@ -7,32 +7,36 @@ INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PagedAttention); PagedAttention::PagedAttention(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, const Tensor &block_tables, const Tensor &kv_lens, - std::optional alibi_slopes, float scale) { + std::optional alibi_slopes, float scale, + std::optional k_scale, std::optional v_scale) { INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, kv_lens); INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); + out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale, k_scale, v_scale); } void PagedAttention::execute(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, const Tensor &block_tables, const Tensor &kv_lens, - std::optional alibi_slopes, float scale) { + std::optional alibi_slopes, float scale, + std::optional k_scale, std::optional v_scale) { INFINICORE_GRAPH_OP_RECORD_OR_RUN( PagedAttention, - out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); + out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale, k_scale, v_scale); } Tensor paged_attention(const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, const Tensor &block_tables, const Tensor &kv_lens, - std::optional alibi_slopes, float scale) { + std::optional alibi_slopes, float scale, + std::optional k_scale, std::optional v_scale) { auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); - paged_attention_(out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); + paged_attention_(out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale, k_scale, v_scale); return out; } 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 alibi_slopes, float scale) { - PagedAttention::execute(out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); + std::optional alibi_slopes, float scale, + std::optional k_scale, std::optional v_scale) { + PagedAttention::execute(out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale, k_scale, v_scale); } } // namespace infinicore::op diff --git a/src/infinicore/ops/paged_attention/paged_attention_infiniop.cc b/src/infinicore/ops/paged_attention/paged_attention_infiniop.cc index 733733a6b..8e7af764b 100644 --- a/src/infinicore/ops/paged_attention/paged_attention_infiniop.cc +++ b/src/infinicore/ops/paged_attention/paged_attention_infiniop.cc @@ -11,18 +11,22 @@ struct PlannedMeta { graph::GraphTensor workspace, out, q, k_cache, v_cache, block_tables, cache_lens; std::optional alibi_slopes; float scale; + std::optional k_scale, v_scale; }; void *plan(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, const Tensor &block_tables, const Tensor &cache_lens, - std::optional alibi_slopes, float scale) { - size_t seed = hash_combine(out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes); + std::optional alibi_slopes, float scale, + std::optional k_scale, std::optional v_scale) { + size_t seed = hash_combine(out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes, k_scale, v_scale); INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( Descriptor, descriptor, PagedAttention, seed, out->desc(), q->desc(), k_cache->desc(), v_cache->desc(), block_tables->desc(), cache_lens->desc(), alibi_slopes ? alibi_slopes.value()->desc() : nullptr, + k_scale ? k_scale.value()->desc() : nullptr, + v_scale ? v_scale.value()->desc() : nullptr, scale); INFINIOP_WORKSPACE_TENSOR(workspace, PagedAttention, descriptor); @@ -37,7 +41,9 @@ void *plan(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_c graph::GraphTensor(block_tables), graph::GraphTensor(cache_lens), alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale}; + scale, + k_scale ? std::optional(graph::GraphTensor(*k_scale)) : std::nullopt, + v_scale ? std::optional(graph::GraphTensor(*v_scale)) : std::nullopt}; } void run(void *planned_meta) { @@ -55,6 +61,8 @@ void run(void *planned_meta) { p->block_tables->data(), p->cache_lens->data(), p->alibi_slopes.has_value() ? p->alibi_slopes.value()->data() : nullptr, + p->k_scale.has_value() ? p->k_scale.value()->data() : nullptr, + p->v_scale.has_value() ? p->v_scale.value()->data() : nullptr, context::getStream())); } diff --git a/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc b/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc index 8e37a3ac7..e36ff53d1 100644 --- a/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc +++ b/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc @@ -17,7 +17,9 @@ void *plan(Tensor out, const Tensor &block_tables, const Tensor &cache_lens, std::optional alibi_slopes, - float scale); + float scale, + std::optional k_scale, + std::optional v_scale); void run(void *planned_meta); void cleanup(void **planned_meta_ptr); } // namespace infinicore::op::paged_attention_impl::infiniop @@ -96,14 +98,20 @@ void *plan(Tensor out, const Tensor &block_tables, const Tensor &cache_lens, std::optional alibi_slopes, - float scale) { + float scale, + std::optional k_scale, + std::optional v_scale) { INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().getType())); INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, cache_lens); if (alibi_slopes) { INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, *alibi_slopes); } - const bool use_flash_attention = canUseFlashAttention(out, q, k_cache, v_cache, block_tables, cache_lens); + // FP8 KV caches (k_scale/v_scale present) never match the flash path's + // dtype checks, so canUseFlashAttention already rejects them; the scales + // are forwarded to the InfiniOP fallback. + const bool use_flash_attention = canUseFlashAttention(out, q, k_cache, v_cache, block_tables, cache_lens) + && !k_scale.has_value() && !v_scale.has_value(); auto flash_out = out->unsqueeze(1); auto flash_q = q->unsqueeze(1); auto flash_k_cache = k_cache->permute({0, 2, 1, 3}); @@ -111,7 +119,7 @@ void *plan(Tensor out, void *fallback_meta = use_flash_attention ? nullptr : paged_attention_impl::infiniop::plan( - out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes, scale); + out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes, scale, k_scale, v_scale); return new PlannedMeta{ TensorMeta(flash_out), TensorMeta(flash_q), TensorMeta(flash_k_cache), TensorMeta(flash_v_cache), diff --git a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill.cc b/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill.cc index d223198b4..a193a3094 100644 --- a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill.cc +++ b/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill.cc @@ -11,29 +11,32 @@ common::OpDispatcher &PagedAttentionPrefill::disp void PagedAttentionPrefill::execute(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, - std::optional alibi_slopes, float scale) { + std::optional alibi_slopes, float scale, + std::optional k_scale, std::optional v_scale) { INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q); infinicore::context::setDevice(out->device()); dispatcher().lookup(out->device().getType())(out, q, k_cache, v_cache, block_tables, - kv_lens, cum_seqlens_q, alibi_slopes, scale); + kv_lens, cum_seqlens_q, alibi_slopes, scale, k_scale, v_scale); } Tensor paged_attention_prefill(Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, - std::optional alibi_slopes, float scale) { + std::optional alibi_slopes, float scale, + std::optional k_scale, std::optional v_scale) { auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); - paged_attention_prefill_(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale); + paged_attention_prefill_(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale, k_scale, v_scale); return out; } void paged_attention_prefill_(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, - std::optional alibi_slopes, float scale) { + std::optional alibi_slopes, float scale, + std::optional k_scale, std::optional v_scale) { - PagedAttentionPrefill::execute(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale); + PagedAttentionPrefill::execute(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale, k_scale, v_scale); } } // namespace infinicore::op diff --git a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniop.cc b/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniop.cc index 05de90cb8..daaf38aec 100644 --- a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniop.cc +++ b/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniop.cc @@ -17,8 +17,9 @@ thread_local common::OpCache void calculate(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, - std::optional alibi_slopes, float scale) { - size_t seed = hash_combine(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale); + std::optional alibi_slopes, float scale, + std::optional k_scale, std::optional v_scale) { + size_t seed = hash_combine(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale, k_scale, v_scale); auto device = context::getDevice(); auto &cache = caches.getCache(device); @@ -37,6 +38,8 @@ void calculate(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, kv_lens->desc(), cum_seqlens_q->desc(), alibi_slopes.has_value() ? alibi_slopes.value()->desc() : nullptr, + k_scale.has_value() ? k_scale.value()->desc() : nullptr, + v_scale.has_value() ? v_scale.value()->desc() : nullptr, scale)); cache.put(seed, desc); } else { @@ -59,6 +62,8 @@ void calculate(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, kv_lens->data(), cum_seqlens_q->data(), alibi_slopes.has_value() ? alibi_slopes.value()->data() : nullptr, + k_scale.has_value() ? k_scale.value()->data() : nullptr, + v_scale.has_value() ? v_scale.value()->data() : nullptr, context::getStream())); } diff --git a/src/infinicore/ops/paged_caching/paged_caching.cc b/src/infinicore/ops/paged_caching/paged_caching.cc index 7eecc10f1..b9d20732e 100644 --- a/src/infinicore/ops/paged_caching/paged_caching.cc +++ b/src/infinicore/ops/paged_caching/paged_caching.cc @@ -5,16 +5,19 @@ namespace infinicore::op { INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PagedCaching); -PagedCaching::PagedCaching(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { +PagedCaching::PagedCaching(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping, + std::optional k_scale, std::optional v_scale) { INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, slot_mapping); - INFINICORE_GRAPH_OP_DISPATCH(k->device().getType(), k_cache, v_cache, k, v, slot_mapping); + INFINICORE_GRAPH_OP_DISPATCH(k->device().getType(), k_cache, v_cache, k, v, slot_mapping, k_scale, v_scale); } -void PagedCaching::execute(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(PagedCaching, k_cache, v_cache, k, v, slot_mapping); +void PagedCaching::execute(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping, + std::optional k_scale, std::optional v_scale) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(PagedCaching, k_cache, v_cache, k, v, slot_mapping, k_scale, v_scale); } -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 k_scale, std::optional v_scale) { constexpr Size MAX_TOKENS_PER_LAUNCH = 32768; const Size num_tokens = k->size(0); @@ -24,7 +27,8 @@ void paged_caching_(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tenso PagedCaching::execute(k_cache, v_cache, k->narrow({{0, start, chunk_size}}), v->narrow({{0, start, chunk_size}}), - slot_mapping->narrow({{0, start, chunk_size}})); + slot_mapping->narrow({{0, start, chunk_size}}), + k_scale, v_scale); } } diff --git a/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc b/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc index 5e8be049a..4d2eb987d 100644 --- a/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc +++ b/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc @@ -10,14 +10,18 @@ struct PlannedMeta { std::shared_ptr descriptor; graph::GraphTensor workspace, k_cache, v_cache, k, v, slot_mapping; + std::optional k_scale, v_scale; }; -void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { - size_t key = hash_combine(k_cache, v_cache, k, v, slot_mapping); +void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping, + std::optional k_scale, std::optional v_scale) { + size_t key = hash_combine(k_cache, v_cache, k, v, slot_mapping, k_scale, v_scale); INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( Descriptor, descriptor, PagedCaching, - key, k_cache->desc(), v_cache->desc(), k->desc(), v->desc(), slot_mapping->desc()); + key, k_cache->desc(), v_cache->desc(), k->desc(), v->desc(), slot_mapping->desc(), + k_scale ? k_scale.value()->desc() : nullptr, + v_scale ? v_scale.value()->desc() : nullptr); INFINIOP_WORKSPACE_TENSOR(workspace, PagedCaching, descriptor); @@ -28,7 +32,9 @@ void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, con graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), - graph::GraphTensor(slot_mapping)}; + graph::GraphTensor(slot_mapping), + k_scale ? std::optional(graph::GraphTensor(*k_scale)) : std::nullopt, + v_scale ? std::optional(graph::GraphTensor(*v_scale)) : std::nullopt}; } void run(void *planned_meta) { @@ -44,6 +50,8 @@ void run(void *planned_meta) { p->k->data(), p->v->data(), p->slot_mapping->data(), + p->k_scale.has_value() ? p->k_scale.value()->data() : nullptr, + p->v_scale.has_value() ? p->v_scale.value()->data() : nullptr, context::getStream())); } diff --git a/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc b/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc index 0c5c6ece4..10fde2ea8 100644 --- a/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc +++ b/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc @@ -7,19 +7,39 @@ #include "base/reshape_and_cache_flash.h" +namespace infinicore::op::paged_caching_impl::infiniop { +void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping, + std::optional k_scale, std::optional v_scale); +void run(void *planned_meta); +void cleanup(void **planned_meta_ptr); +} // namespace infinicore::op::paged_caching_impl::infiniop + namespace infinicore::op::paged_caching_impl::infiniops { namespace { using TensorMeta = ::infinicore::op::infiniops::TensorMeta; struct PlannedMeta { - TensorMeta k, v, slot_mapping, scale, k_cache, v_cache; - graph::GraphTensor k_tensor, v_tensor, slot_mapping_tensor, scale_tensor, k_cache_tensor, v_cache_tensor; + // Populated on the flash path only. + std::optional k, v, slot_mapping, scale, k_cache, v_cache; + std::optional k_tensor, v_tensor, slot_mapping_tensor, scale_tensor, k_cache_tensor, v_cache_tensor; + void *fallback_meta; }; } // namespace -void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { +void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping, + std::optional k_scale, std::optional v_scale) { INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(k_cache->device().getType())); INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, slot_mapping); + // FP8 KV caches need on-write quantization, which the flash path does not + // implement; fall back to the InfiniOP implementation. + if (k_scale.has_value() || v_scale.has_value() || k_cache->dtype() == DataType::F8 + || v_cache->dtype() == DataType::F8) { + return new PlannedMeta{ + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + paged_caching_impl::infiniop::plan(k_cache, v_cache, k, v, slot_mapping, k_scale, v_scale)}; + } + // The "auto" cache path ignores scales, but the canonical API requires them. auto scale = Tensor::empty({1}, DataType::F32, k_cache->device()); auto k_cache_view = k_cache->permute({0, 2, 1, 3}); @@ -27,29 +47,38 @@ void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, con return new PlannedMeta{ TensorMeta(k), TensorMeta(v), TensorMeta(slot_mapping), TensorMeta(scale), TensorMeta(k_cache_view), TensorMeta(v_cache_view), - graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(slot_mapping), graph::GraphTensor(scale), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache)}; + graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(slot_mapping), graph::GraphTensor(scale), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), + nullptr}; } void run(void *planned_meta) { auto planned = reinterpret_cast(planned_meta); + if (planned->fallback_meta != nullptr) { + paged_caching_impl::infiniop::run(planned->fallback_meta); + return; + } infini::ops::Handle handle; handle.set_stream(context::getStream()); infini::ops::Config config; infini::ops::ReshapeAndCacheFlash::Call( handle, config, - planned->k.tensor(planned->k_tensor), - planned->v.tensor(planned->v_tensor), - planned->slot_mapping.tensor(planned->slot_mapping_tensor), - planned->scale.tensor(planned->scale_tensor), - planned->scale.tensor(planned->scale_tensor), + planned->k->tensor(*planned->k_tensor), + planned->v->tensor(*planned->v_tensor), + planned->slot_mapping->tensor(*planned->slot_mapping_tensor), + planned->scale->tensor(*planned->scale_tensor), + planned->scale->tensor(*planned->scale_tensor), std::string{"auto"}, - planned->k_cache.tensor(planned->k_cache_tensor), - planned->v_cache.tensor(planned->v_cache_tensor)); + planned->k_cache->tensor(*planned->k_cache_tensor), + planned->v_cache->tensor(*planned->v_cache_tensor)); } void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); + auto planned = *reinterpret_cast(planned_meta_ptr); + if (planned->fallback_meta != nullptr) { + paged_caching_impl::infiniop::cleanup(&planned->fallback_meta); + } + delete planned; *planned_meta_ptr = nullptr; } diff --git a/src/infinicore/pybind11/ops/paged_attention.hpp b/src/infinicore/pybind11/ops/paged_attention.hpp index ab77c87a4..c3ef04646 100644 --- a/src/infinicore/pybind11/ops/paged_attention.hpp +++ b/src/infinicore/pybind11/ops/paged_attention.hpp @@ -8,21 +8,37 @@ namespace py = pybind11; namespace infinicore::ops { -Tensor py_paged_attention(Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, pybind11::object alibi_slopes, float scale) { +Tensor py_paged_attention(Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, pybind11::object alibi_slopes, float scale, pybind11::object k_scale, pybind11::object v_scale) { std::optional alibi_slopes_tensor = std::nullopt; if (!alibi_slopes.is_none()) { alibi_slopes_tensor = alibi_slopes.cast(); } - return op::paged_attention(q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes_tensor, scale); + std::optional k_scale_tensor = std::nullopt; + if (!k_scale.is_none()) { + k_scale_tensor = k_scale.cast(); + } + std::optional v_scale_tensor = std::nullopt; + if (!v_scale.is_none()) { + v_scale_tensor = v_scale.cast(); + } + return op::paged_attention(q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes_tensor, scale, k_scale_tensor, v_scale_tensor); } -void py_paged_attention_(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, pybind11::object alibi_slopes, float scale) { +void py_paged_attention_(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, pybind11::object alibi_slopes, float scale, pybind11::object k_scale, pybind11::object v_scale) { std::optional alibi_slopes_tensor = std::nullopt; if (!alibi_slopes.is_none()) { alibi_slopes_tensor = alibi_slopes.cast(); } + std::optional k_scale_tensor = std::nullopt; + if (!k_scale.is_none()) { + k_scale_tensor = k_scale.cast(); + } + std::optional v_scale_tensor = std::nullopt; + if (!v_scale.is_none()) { + v_scale_tensor = v_scale.cast(); + } - op::paged_attention_(out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes_tensor, scale); + op::paged_attention_(out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes_tensor, scale, k_scale_tensor, v_scale_tensor); } inline void bind_paged_attention(py::module &m) { @@ -35,6 +51,8 @@ inline void bind_paged_attention(py::module &m) { py::arg("cache_lens"), py::arg("alibi_slopes"), py::arg("scale"), + py::arg("k_scale") = py::none(), + py::arg("v_scale") = py::none(), R"doc(Paged attention of query and key cache tensors.)doc"); m.def("paged_attention_", @@ -47,6 +65,8 @@ inline void bind_paged_attention(py::module &m) { py::arg("cache_lens"), py::arg("alibi_slopes"), py::arg("scale"), + py::arg("k_scale") = py::none(), + py::arg("v_scale") = py::none(), R"doc(In-place paged attention of query and key cache tensors.)doc"); } diff --git a/src/infinicore/pybind11/ops/paged_attention_prefill.hpp b/src/infinicore/pybind11/ops/paged_attention_prefill.hpp index 13b2c2683..7937f7258 100644 --- a/src/infinicore/pybind11/ops/paged_attention_prefill.hpp +++ b/src/infinicore/pybind11/ops/paged_attention_prefill.hpp @@ -14,13 +14,23 @@ Tensor py_paged_attention_prefill(Tensor q, Tensor history_lens, Tensor cu_seqlens_q, py::object alibi_slopes, - float scale) { + float scale, + py::object k_scale, + py::object v_scale) { std::optional alibi_slopes_tensor = std::nullopt; if (!alibi_slopes.is_none()) { alibi_slopes_tensor = alibi_slopes.cast(); } + std::optional k_scale_tensor = std::nullopt; + if (!k_scale.is_none()) { + k_scale_tensor = k_scale.cast(); + } + std::optional v_scale_tensor = std::nullopt; + if (!v_scale.is_none()) { + v_scale_tensor = v_scale.cast(); + } return op::paged_attention_prefill( - q, k_cache, v_cache, block_tables, history_lens, cu_seqlens_q, alibi_slopes_tensor, scale); + q, k_cache, v_cache, block_tables, history_lens, cu_seqlens_q, alibi_slopes_tensor, scale, k_scale_tensor, v_scale_tensor); } void py_paged_attention_prefill_(Tensor out, @@ -31,12 +41,22 @@ void py_paged_attention_prefill_(Tensor out, Tensor history_lens, Tensor cu_seqlens_q, py::object alibi_slopes, - float scale) { + float scale, + py::object k_scale, + py::object v_scale) { std::optional alibi_slopes_tensor = std::nullopt; if (!alibi_slopes.is_none()) { alibi_slopes_tensor = alibi_slopes.cast(); } - op::paged_attention_prefill_(out, q, k_cache, v_cache, block_tables, history_lens, cu_seqlens_q, alibi_slopes_tensor, scale); + std::optional k_scale_tensor = std::nullopt; + if (!k_scale.is_none()) { + k_scale_tensor = k_scale.cast(); + } + std::optional v_scale_tensor = std::nullopt; + if (!v_scale.is_none()) { + v_scale_tensor = v_scale.cast(); + } + op::paged_attention_prefill_(out, q, k_cache, v_cache, block_tables, history_lens, cu_seqlens_q, alibi_slopes_tensor, scale, k_scale_tensor, v_scale_tensor); } inline void bind_paged_attention_prefill(py::module &m) { @@ -50,6 +70,8 @@ inline void bind_paged_attention_prefill(py::module &m) { py::arg("cu_seqlens_q"), py::arg("alibi_slopes") = py::none(), py::arg("scale") = 1.0, + py::arg("k_scale") = py::none(), + py::arg("v_scale") = py::none(), R"doc(Paged attention prefill for packed variable-length queries.)doc"); m.def("paged_attention_prefill_", @@ -63,6 +85,8 @@ inline void bind_paged_attention_prefill(py::module &m) { py::arg("cu_seqlens_q"), py::arg("alibi_slopes") = py::none(), py::arg("scale") = 1.0, + py::arg("k_scale") = py::none(), + py::arg("v_scale") = py::none(), R"doc(In-place paged attention prefill for packed variable-length queries.)doc"); } diff --git a/src/infinicore/pybind11/ops/paged_caching.hpp b/src/infinicore/pybind11/ops/paged_caching.hpp index 4320b4eef..0103a85fa 100644 --- a/src/infinicore/pybind11/ops/paged_caching.hpp +++ b/src/infinicore/pybind11/ops/paged_caching.hpp @@ -8,14 +8,28 @@ namespace py = pybind11; namespace infinicore::ops { +inline void py_paged_caching_(Tensor k_cache, Tensor v_cache, Tensor k, Tensor v, Tensor slot_mapping, py::object k_scale, py::object v_scale) { + std::optional k_scale_tensor = std::nullopt; + if (!k_scale.is_none()) { + k_scale_tensor = k_scale.cast(); + } + std::optional v_scale_tensor = std::nullopt; + if (!v_scale.is_none()) { + v_scale_tensor = v_scale.cast(); + } + op::paged_caching_(k_cache, v_cache, k, v, slot_mapping, k_scale_tensor, v_scale_tensor); +} + inline void bind_paged_caching(py::module &m) { m.def("paged_caching_", - &op::paged_caching_, + &ops::py_paged_caching_, py::arg("k_cache"), py::arg("v_cache"), py::arg("k"), py::arg("v"), py::arg("slot_mapping"), + py::arg("k_scale") = py::none(), + py::arg("v_scale") = py::none(), R"doc(Paged caching of key and value tensors.)doc"); } diff --git a/src/infiniop/ops/avg_pool3d/operator.cc b/src/infiniop/ops/avg_pool3d/operator.cc index df58e1282..02bc7dca8 100644 --- a/src/infiniop/ops/avg_pool3d/operator.cc +++ b/src/infiniop/ops/avg_pool3d/operator.cc @@ -5,7 +5,7 @@ #ifdef ENABLE_CPU_API #include "cpu/avg_pool3d_cpu.h" #endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) +#if (defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API)) && defined(ENABLE_CUDNN_API) #include "nvidia/avg_pool3d_nvidia.cuh" #endif #ifdef ENABLE_METAX_API @@ -40,10 +40,10 @@ __INFINI_C infiniStatus_t infiniopCreateAvgPool3dDescriptor( #ifdef ENABLE_CPU_API CREATE(INFINI_DEVICE_CPU, cpu); #endif -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) && defined(ENABLE_CUDNN_API) CREATE(INFINI_DEVICE_NVIDIA, nvidia); #endif -#ifdef ENABLE_QY_API +#if defined(ENABLE_QY_API) && defined(ENABLE_CUDNN_API) CREATE(INFINI_DEVICE_QY, nvidia); #endif #ifdef ENABLE_METAX_API @@ -71,10 +71,10 @@ __INFINI_C infiniStatus_t infiniopGetAvgPool3dWorkspaceSize(infiniopAvgPool3dDes #ifdef ENABLE_CPU_API GET(INFINI_DEVICE_CPU, cpu) #endif -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) && defined(ENABLE_CUDNN_API) GET(INFINI_DEVICE_NVIDIA, nvidia) #endif -#ifdef ENABLE_QY_API +#if defined(ENABLE_QY_API) && defined(ENABLE_CUDNN_API) GET(INFINI_DEVICE_QY, nvidia) #endif #ifdef ENABLE_METAX_API @@ -109,10 +109,10 @@ __INFINI_C infiniStatus_t infiniopAvgPool3d( #ifdef ENABLE_CPU_API CALCULATE(INFINI_DEVICE_CPU, cpu); #endif -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) && defined(ENABLE_CUDNN_API) CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); #endif -#ifdef ENABLE_QY_API +#if defined(ENABLE_QY_API) && defined(ENABLE_CUDNN_API) CALCULATE(INFINI_DEVICE_QY, nvidia); #endif #ifdef ENABLE_METAX_API @@ -142,10 +142,10 @@ infiniopDestroyAvgPool3dDescriptor(infiniopAvgPool3dDescriptor_t desc) { #ifdef ENABLE_CPU_API DELETE(INFINI_DEVICE_CPU, cpu); #endif -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) && defined(ENABLE_CUDNN_API) DELETE(INFINI_DEVICE_NVIDIA, nvidia); #endif -#ifdef ENABLE_QY_API +#if defined(ENABLE_QY_API) && defined(ENABLE_CUDNN_API) DELETE(INFINI_DEVICE_QY, nvidia); #endif #ifdef ENABLE_METAX_API diff --git a/src/infiniop/ops/fp8_blockwise_dequantize/cpu/fp8_blockwise_dequantize_cpu.cc b/src/infiniop/ops/fp8_blockwise_dequantize/cpu/fp8_blockwise_dequantize_cpu.cc new file mode 100644 index 000000000..9cd54ead6 --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_dequantize/cpu/fp8_blockwise_dequantize_cpu.cc @@ -0,0 +1,83 @@ +#include "fp8_blockwise_dequantize_cpu.h" + +#include "../../../../utils/custom_types.h" +#include "../../../devices/cpu/cpu_handle.h" + +#include +#include +#include + +namespace op::fp8_blockwise_dequantize::cpu { +namespace { + +// Decode one FP8 E4M3FN byte: 1 sign bit, 4 exponent bits (bias 7), 3 mantissa +// bits. E4M3FN has no infinity; exponent 15 with mantissa 7 is NaN. +inline float decode_e4m3(uint8_t value) { + const int exponent = (value >> 3) & 0xf; + const int mantissa = value & 0x7; + if (exponent == 0xf && mantissa == 0x7) { + return std::numeric_limits::quiet_NaN(); + } + const float decoded = exponent == 0 + ? std::ldexp(static_cast(mantissa), -9) + : std::ldexp(1.0f + static_cast(mantissa) * 0.125f, exponent - 7); + return value & 0x80 ? -decoded : decoded; +} + +template +void dequantize(T *out, + const uint8_t *q, + const float *scales, + const Fp8BlockwiseDequantizeInfo &info) { +#ifdef ENABLE_OMP +#pragma omp parallel for +#endif + for (ptrdiff_t row = 0; row < static_cast(info.rows); ++row) { + const size_t scale_row = row / info.block_rows; + for (size_t col = 0; col < info.cols; ++col) { + const size_t index = row * info.cols + col; + const float scale = scales[scale_row * info.scales_cols + col / info.block_cols]; + out[index] = utils::cast(decode_e4m3(q[index]) * scale); + } + } +} + +} // namespace + +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { delete _opaque; } + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t scales_desc) { + auto info = Fp8BlockwiseDequantizeInfo::create(out_desc, q_desc, scales_desc); + CHECK_RESULT(info); + *desc_ptr = new Descriptor(new Opaque{}, info.take(), handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *, size_t, void *out, + const void *q, const void *scales, void *) const { + auto q_ptr = reinterpret_cast(q); + auto scales_ptr = reinterpret_cast(scales); + switch (_info.output_dtype) { + case INFINI_DTYPE_F16: + dequantize(reinterpret_cast(out), q_ptr, scales_ptr, _info); + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_BF16: + dequantize(reinterpret_cast(out), q_ptr, scales_ptr, _info); + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_F32: + dequantize(reinterpret_cast(out), q_ptr, scales_ptr, _info); + return INFINI_STATUS_SUCCESS; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::fp8_blockwise_dequantize::cpu diff --git a/src/infiniop/ops/fp8_blockwise_dequantize/cpu/fp8_blockwise_dequantize_cpu.h b/src/infiniop/ops/fp8_blockwise_dequantize/cpu/fp8_blockwise_dequantize_cpu.h new file mode 100644 index 000000000..b163b8cd8 --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_dequantize/cpu/fp8_blockwise_dequantize_cpu.h @@ -0,0 +1,8 @@ +#ifndef __FP8_BLOCKWISE_DEQUANTIZE_CPU_H__ +#define __FP8_BLOCKWISE_DEQUANTIZE_CPU_H__ + +#include "../fp8_blockwise_dequantize.h" + +DESCRIPTOR(cpu) + +#endif diff --git a/src/infiniop/ops/fp8_blockwise_dequantize/fp8_blockwise_dequantize.h b/src/infiniop/ops/fp8_blockwise_dequantize/fp8_blockwise_dequantize.h new file mode 100644 index 000000000..92612f85a --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_dequantize/fp8_blockwise_dequantize.h @@ -0,0 +1,35 @@ +#ifndef __FP8_BLOCKWISE_DEQUANTIZE_H__ +#define __FP8_BLOCKWISE_DEQUANTIZE_H__ + +#include "../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::fp8_blockwise_dequantize::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + Fp8BlockwiseDequantizeInfo _info; \ + \ + Descriptor(Opaque *opaque, Fp8BlockwiseDequantizeInfo info, \ + infiniDevice_t device_type, int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), _info(info) {} \ + \ + public: \ + ~Descriptor(); \ + size_t workspaceSize() const { return 0; } \ + \ + static infiniStatus_t create( \ + infiniopHandle_t handle, Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t out_desc, \ + infiniopTensorDescriptor_t q_desc, \ + infiniopTensorDescriptor_t scales_desc); \ + \ + infiniStatus_t calculate( \ + void *workspace, size_t workspace_size, void *out, \ + const void *q, const void *scales, void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/fp8_blockwise_dequantize/info.h b/src/infiniop/ops/fp8_blockwise_dequantize/info.h new file mode 100644 index 000000000..d366397fa --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_dequantize/info.h @@ -0,0 +1,51 @@ +#ifndef __FP8_BLOCKWISE_DEQUANTIZE_INFO_H__ +#define __FP8_BLOCKWISE_DEQUANTIZE_INFO_H__ + +#include "../../../utils.h" +#include "../../tensor.h" + +namespace op::fp8_blockwise_dequantize { + +class Fp8BlockwiseDequantizeInfo { + Fp8BlockwiseDequantizeInfo() = default; + +public: + infiniDtype_t output_dtype; + size_t rows; + size_t cols; + size_t block_rows; + size_t block_cols; + size_t scales_cols; + + static utils::Result create( + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t scales_desc) { + CHECK_OR_RETURN(out_desc != nullptr && q_desc != nullptr && scales_desc != nullptr, + INFINI_STATUS_NULL_POINTER); + CHECK_DTYPE(out_desc->dtype(), INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); + CHECK_DTYPE(q_desc->dtype(), INFINI_DTYPE_F8); + CHECK_DTYPE(scales_desc->dtype(), INFINI_DTYPE_F32); + CHECK_OR_RETURN(out_desc->ndim() == 2 && q_desc->ndim() == 2 && scales_desc->ndim() == 2, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(out_desc->isContiguous() && q_desc->isContiguous() && scales_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + + const size_t rows = q_desc->dim(0); + const size_t cols = q_desc->dim(1); + CHECK_OR_RETURN(rows > 0 && cols > 0 && out_desc->dim(0) == rows && out_desc->dim(1) == cols, + INFINI_STATUS_BAD_TENSOR_SHAPE); + + const size_t scales_rows = scales_desc->dim(0); + const size_t scales_cols = scales_desc->dim(1); + CHECK_OR_RETURN(scales_rows > 0 && scales_cols > 0 && rows % scales_rows == 0 && cols % scales_cols == 0, + INFINI_STATUS_BAD_TENSOR_SHAPE); + + return utils::Result(Fp8BlockwiseDequantizeInfo{ + out_desc->dtype(), rows, cols, rows / scales_rows, cols / scales_cols, scales_cols}); + } +}; + +} // namespace op::fp8_blockwise_dequantize + +#endif diff --git a/src/infiniop/ops/fp8_blockwise_dequantize/nvidia/fp8_blockwise_dequantize_nvidia.cu b/src/infiniop/ops/fp8_blockwise_dequantize/nvidia/fp8_blockwise_dequantize_nvidia.cu new file mode 100644 index 000000000..fe73ded16 --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_dequantize/nvidia/fp8_blockwise_dequantize_nvidia.cu @@ -0,0 +1,111 @@ +#include "fp8_blockwise_dequantize_nvidia.cuh" + +#include "../../../devices/nvidia/nvidia_handle.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" + +#include +#include +#include + +#include + +namespace op::fp8_blockwise_dequantize::nvidia { +namespace { + +template +__device__ __forceinline__ T cast_output(float value); + +template <> +__device__ __forceinline__ half cast_output(float value) { + return __float2half_rn(value); +} + +template <> +__device__ __forceinline__ __nv_bfloat16 cast_output(float value) { + return __float2bfloat16_rn(value); +} + +template <> +__device__ __forceinline__ float cast_output(float value) { + return value; +} + +template +INFINIOP_CUDA_KERNEL dequantize_kernel( + T *out, + const uint8_t *q, + const float *scales, + size_t numel, + size_t cols, + size_t block_rows, + size_t block_cols, + size_t scales_cols) { + const size_t index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= numel) { + return; + } + + const size_t row = index / cols; + const size_t col = index - row * cols; + const size_t scale_index = (row / block_rows) * scales_cols + col / block_cols; + out[index] = cast_output(infiniopFp8E4m3Decode(q[index]) * scales[scale_index]); +} + +template +void launch(T *out, + const uint8_t *q, + const float *scales, + const Fp8BlockwiseDequantizeInfo &info, + cudaStream_t stream) { + constexpr size_t block_size = 256; + const size_t numel = info.rows * info.cols; + const size_t grid_size = (numel + block_size - 1) / block_size; + dequantize_kernel<<>>( + out, q, scales, numel, info.cols, + info.block_rows, info.block_cols, info.scales_cols); +} + +} // namespace + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { delete _opaque; } + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t scales_desc) { + auto info = Fp8BlockwiseDequantizeInfo::create(out_desc, q_desc, scales_desc); + CHECK_RESULT(info); + auto nvidia_handle = reinterpret_cast(handle); + *desc_ptr = new Descriptor( + new Opaque{nvidia_handle->internal()}, info.take(), handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *, size_t, void *out, + const void *q, const void *scales, void *stream) const { + auto cuda_stream = reinterpret_cast(stream); + auto q_ptr = reinterpret_cast(q); + auto scales_ptr = reinterpret_cast(scales); + switch (_info.output_dtype) { + case INFINI_DTYPE_F16: + launch(reinterpret_cast(out), q_ptr, scales_ptr, _info, cuda_stream); + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_BF16: + launch(reinterpret_cast<__nv_bfloat16 *>(out), q_ptr, scales_ptr, _info, cuda_stream); + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_F32: + launch(reinterpret_cast(out), q_ptr, scales_ptr, _info, cuda_stream); + return INFINI_STATUS_SUCCESS; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::fp8_blockwise_dequantize::nvidia diff --git a/src/infiniop/ops/fp8_blockwise_dequantize/nvidia/fp8_blockwise_dequantize_nvidia.cuh b/src/infiniop/ops/fp8_blockwise_dequantize/nvidia/fp8_blockwise_dequantize_nvidia.cuh new file mode 100644 index 000000000..e9e92bbfc --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_dequantize/nvidia/fp8_blockwise_dequantize_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __FP8_BLOCKWISE_DEQUANTIZE_NVIDIA_CUH__ +#define __FP8_BLOCKWISE_DEQUANTIZE_NVIDIA_CUH__ + +#include "../fp8_blockwise_dequantize.h" + +DESCRIPTOR(nvidia) + +#endif diff --git a/src/infiniop/ops/fp8_blockwise_dequantize/operator.cc b/src/infiniop/ops/fp8_blockwise_dequantize/operator.cc new file mode 100644 index 000000000..9375cbd98 --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_dequantize/operator.cc @@ -0,0 +1,100 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/fp8_blockwise_dequantize.h" + +#ifdef ENABLE_CPU_API +#include "cpu/fp8_blockwise_dequantize_cpu.h" +#endif +#ifdef ENABLE_NVIDIA_API +#include "nvidia/fp8_blockwise_dequantize_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateFp8BlockwiseDequantizeDescriptor( + infiniopHandle_t handle, + infiniopFp8BlockwiseDequantizeDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t scales_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::fp8_blockwise_dequantize::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + out_desc, q_desc, scales_desc) + switch (handle->device) { +#ifdef ENABLE_CPU_API + CREATE(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetFp8BlockwiseDequantizeWorkspaceSize( + infiniopFp8BlockwiseDequantizeDescriptor_t desc, + size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc) \ + ->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_CPU_API + GET(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C infiniStatus_t infiniopFp8BlockwiseDequantize( + infiniopFp8BlockwiseDequantizeDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *q, + const void *scales, + void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, out, q, scales, stream) + switch (desc->device_type) { +#ifdef ENABLE_CPU_API + CALCULATE(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyFp8BlockwiseDequantizeDescriptor( + infiniopFp8BlockwiseDequantizeDescriptor_t desc) { +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_CPU_API + DESTROY(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/fp8_blockwise_gemm/cpu/fp8_blockwise_gemm_cpu.cc b/src/infiniop/ops/fp8_blockwise_gemm/cpu/fp8_blockwise_gemm_cpu.cc new file mode 100644 index 000000000..8c68d9de9 --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_gemm/cpu/fp8_blockwise_gemm_cpu.cc @@ -0,0 +1,92 @@ +#include "fp8_blockwise_gemm_cpu.h" + +#include "../../../../utils/custom_types.h" +#include "../../../devices/cpu/cpu_handle.h" + +#include +#include +#include + +namespace op::fp8_blockwise_gemm::cpu { +namespace { + +// Decode one FP8 E4M3FN byte: 1 sign bit, 4 exponent bits (bias 7), 3 mantissa +// bits. E4M3FN has no infinity; exponent 15 with mantissa 7 is NaN. +inline float decode_e4m3(uint8_t value) { + const int exponent = (value >> 3) & 0xf; + const int mantissa = value & 0x7; + if (exponent == 0xf && mantissa == 0x7) { + return std::numeric_limits::quiet_NaN(); + } + const float decoded = exponent == 0 + ? std::ldexp(static_cast(mantissa), -9) + : std::ldexp(1.0f + static_cast(mantissa) * 0.125f, exponent - 7); + return value & 0x80 ? -decoded : decoded; +} + +template +void gemm(T *out, + const T *a, + const uint8_t *q, + const float *scales, + const Fp8BlockwiseGemmInfo &info) { +#ifdef ENABLE_OMP +#pragma omp parallel for collapse(2) +#endif + for (ptrdiff_t n = 0; n < static_cast(info.N); ++n) { + for (ptrdiff_t m = 0; m < static_cast(info.M); ++m) { + float acc = 0.0f; + const size_t scale_row = n / info.block_n; + for (size_t kb = 0; kb < info.K / info.block_k; ++kb) { + const float scale = scales[scale_row * info.scales_cols + kb]; + float chunk = 0.0f; + for (size_t k = kb * info.block_k; k < (kb + 1) * info.block_k; ++k) { + chunk += utils::cast(a[m * info.K + k]) * decode_e4m3(q[n * info.K + k]); + } + acc += scale * chunk; + } + out[m * info.N + n] = utils::cast(acc); + } + } +} + +} // namespace + +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { delete _opaque; } + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t scales_desc) { + auto info = Fp8BlockwiseGemmInfo::create(out_desc, a_desc, q_desc, scales_desc); + CHECK_RESULT(info); + *desc_ptr = new Descriptor(new Opaque{}, info.take(), handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *, size_t, void *out, + const void *a, const void *q, const void *scales, void *) const { + auto q_ptr = reinterpret_cast(q); + auto scales_ptr = reinterpret_cast(scales); + switch (_info.dtype) { + case INFINI_DTYPE_F16: + gemm(reinterpret_cast(out), reinterpret_cast(a), q_ptr, scales_ptr, _info); + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_BF16: + gemm(reinterpret_cast(out), reinterpret_cast(a), q_ptr, scales_ptr, _info); + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_F32: + gemm(reinterpret_cast(out), reinterpret_cast(a), q_ptr, scales_ptr, _info); + return INFINI_STATUS_SUCCESS; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::fp8_blockwise_gemm::cpu diff --git a/src/infiniop/ops/fp8_blockwise_gemm/cpu/fp8_blockwise_gemm_cpu.h b/src/infiniop/ops/fp8_blockwise_gemm/cpu/fp8_blockwise_gemm_cpu.h new file mode 100644 index 000000000..37cf0cc42 --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_gemm/cpu/fp8_blockwise_gemm_cpu.h @@ -0,0 +1,8 @@ +#ifndef __FP8_BLOCKWISE_GEMM_CPU_H__ +#define __FP8_BLOCKWISE_GEMM_CPU_H__ + +#include "../fp8_blockwise_gemm.h" + +DESCRIPTOR(cpu) + +#endif diff --git a/src/infiniop/ops/fp8_blockwise_gemm/fp8_blockwise_gemm.h b/src/infiniop/ops/fp8_blockwise_gemm/fp8_blockwise_gemm.h new file mode 100644 index 000000000..a0d80d294 --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_gemm/fp8_blockwise_gemm.h @@ -0,0 +1,37 @@ +#ifndef __FP8_BLOCKWISE_GEMM_H__ +#define __FP8_BLOCKWISE_GEMM_H__ + +#include "../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::fp8_blockwise_gemm::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + Fp8BlockwiseGemmInfo _info; \ + \ + Descriptor(Opaque *opaque, Fp8BlockwiseGemmInfo info, \ + infiniDevice_t device_type, int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), _info(info) {} \ + \ + public: \ + ~Descriptor(); \ + size_t workspaceSize() const { return 0; } \ + \ + static infiniStatus_t create( \ + infiniopHandle_t handle, Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t out_desc, \ + infiniopTensorDescriptor_t a_desc, \ + infiniopTensorDescriptor_t q_desc, \ + infiniopTensorDescriptor_t scales_desc); \ + \ + infiniStatus_t calculate( \ + void *workspace, size_t workspace_size, void *out, \ + const void *a, const void *q, const void *scales, \ + void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/fp8_blockwise_gemm/info.h b/src/infiniop/ops/fp8_blockwise_gemm/info.h new file mode 100644 index 000000000..3ef26e453 --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_gemm/info.h @@ -0,0 +1,61 @@ +#ifndef __FP8_BLOCKWISE_GEMM_INFO_H__ +#define __FP8_BLOCKWISE_GEMM_INFO_H__ + +#include "../../../utils.h" +#include "../../tensor.h" + +namespace op::fp8_blockwise_gemm { + +class Fp8BlockwiseGemmInfo { + Fp8BlockwiseGemmInfo() = default; + +public: + infiniDtype_t dtype; + size_t M, N, K; + size_t block_n; + size_t block_k; + size_t scales_cols; + + static utils::Result create( + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t scales_desc) { + CHECK_OR_RETURN(out_desc != nullptr && a_desc != nullptr && q_desc != nullptr && scales_desc != nullptr, + INFINI_STATUS_NULL_POINTER); + const infiniDtype_t dtype = out_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); + CHECK_OR_RETURN(a_desc->dtype() == dtype, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_DTYPE(q_desc->dtype(), INFINI_DTYPE_F8); + CHECK_DTYPE(scales_desc->dtype(), INFINI_DTYPE_F32); + CHECK_OR_RETURN(out_desc->ndim() == 2 && a_desc->ndim() == 2 && q_desc->ndim() == 2 && scales_desc->ndim() == 2, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(out_desc->isContiguous() && a_desc->isContiguous() && q_desc->isContiguous() && scales_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + + const size_t M = out_desc->dim(0); + const size_t N = out_desc->dim(1); + const size_t K = a_desc->dim(1); + CHECK_OR_RETURN(M > 0 && N > 0 && K > 0, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(a_desc->dim(0) == M && q_desc->dim(0) == N && q_desc->dim(1) == K, + INFINI_STATUS_BAD_TENSOR_SHAPE); + + const size_t scales_rows = scales_desc->dim(0); + const size_t scales_cols = scales_desc->dim(1); + CHECK_OR_RETURN(scales_rows > 0 && scales_cols > 0 && N % scales_rows == 0 && K % scales_cols == 0, + INFINI_STATUS_BAD_TENSOR_SHAPE); + const size_t block_n = N / scales_rows; + const size_t block_k = K / scales_cols; + // The kernel addresses scales per 128-wide K chunk and requires + // 4-byte vectorizable rows. + CHECK_OR_RETURN(block_k % 128 == 0 && K % 128 == 0, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(block_n % 16 == 0 && K % 4 == 0, INFINI_STATUS_BAD_TENSOR_SHAPE); + + return utils::Result(Fp8BlockwiseGemmInfo{ + dtype, M, N, K, block_n, block_k, scales_cols}); + } +}; + +} // namespace op::fp8_blockwise_gemm + +#endif diff --git a/src/infiniop/ops/fp8_blockwise_gemm/nvidia/fp8_blockwise_gemm_nvidia.cu b/src/infiniop/ops/fp8_blockwise_gemm/nvidia/fp8_blockwise_gemm_nvidia.cu new file mode 100644 index 000000000..d8a829d1f --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_gemm/nvidia/fp8_blockwise_gemm_nvidia.cu @@ -0,0 +1,571 @@ +#include "fp8_blockwise_gemm_nvidia.cuh" + +#include "../../../devices/nvidia/nvidia_handle.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" + +#include +#include +#include + +#include +#include +#include + +namespace op::fp8_blockwise_gemm::nvidia { +namespace { + +// --------------------------------------------------------------------------- +// dtype helpers +// --------------------------------------------------------------------------- + +template +__device__ __forceinline__ T from_float(float value); + +template <> +__device__ __forceinline__ half from_float(float value) { + return __float2half_rn(value); +} + +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float value) { + return __float2bfloat16_rn(value); +} + +template <> +__device__ __forceinline__ float from_float(float value) { + return value; +} + +// Load 4 consecutive elements (8B for half/bf16, 16B for float) and convert to +// float. The address is always 4-element aligned (K % 4 == 0 and lanes walk in +// steps of 4). +template +__device__ __forceinline__ void load4(const T *p, float *v); + +template <> +__device__ __forceinline__ void load4(const half *p, float *v) { + const half2 h01 = *reinterpret_cast(p); + const half2 h23 = *reinterpret_cast(p + 2); + const float2 f01 = __half22float2(h01); + const float2 f23 = __half22float2(h23); + v[0] = f01.x; + v[1] = f01.y; + v[2] = f23.x; + v[3] = f23.y; +} + +template <> +__device__ __forceinline__ void load4<__nv_bfloat16>(const __nv_bfloat16 *p, float *v) { + const __nv_bfloat162 b01 = *reinterpret_cast(p); + const __nv_bfloat162 b23 = *reinterpret_cast(p + 2); + const float2 f01 = __bfloat1622float2(b01); + const float2 f23 = __bfloat1622float2(b23); + v[0] = f01.x; + v[1] = f01.y; + v[2] = f23.x; + v[3] = f23.y; +} + +template <> +__device__ __forceinline__ void load4(const float *p, float *v) { + const float4 f = *reinterpret_cast(p); + v[0] = f.x; + v[1] = f.y; + v[2] = f.z; + v[3] = f.w; +} + +// --------------------------------------------------------------------------- +// Fused FP8 blockwise GEMM: out[m, n] = sum_k a[m,k] * w[n,k] * s[n/BN, k/BK] +// +// Decode-oriented GEMV shape: one warp computes the M_TILE outputs of one +// weight row (warp-per-row, 4 warps per block). The 32 lanes of a warp cover +// a 512-byte K group (16 FP8 each), so each FP8 weight byte is read exactly +// once and reused for all M_TILE activation rows in registers; the next group +// is prefetched into registers while the current one is consumed. Per-128- +// chunk partial dots are scaled and accumulated in FP32. K tails that are not +// a multiple of 512 bytes fall back to a 4-byte path. +// --------------------------------------------------------------------------- + +constexpr int TN = 4; // weight rows per thread block (4 warps x 1 row) + +template +__device__ __forceinline__ void load16(const T *p, float *v); + +template <> +__device__ __forceinline__ void load16(const half *p, float *v) { + load4(p, v); + load4(p + 4, v + 4); + load4(p + 8, v + 8); + load4(p + 12, v + 12); +} + +template <> +__device__ __forceinline__ void load16<__nv_bfloat16>(const __nv_bfloat16 *p, float *v) { + load4(p, v); + load4(p + 4, v + 4); + load4(p + 8, v + 8); + load4(p + 12, v + 12); +} + +template <> +__device__ __forceinline__ void load16(const float *p, float *v) { + load4(p, v); + load4(p + 4, v + 4); + load4(p + 8, v + 8); + load4(p + 12, v + 12); +} + +__device__ __forceinline__ void decode16(uint4 q16, float *w) { + w[0] = infiniopFp8E4m3Decode(q16.x & 0xffU); + w[1] = infiniopFp8E4m3Decode((q16.x >> 8) & 0xffU); + w[2] = infiniopFp8E4m3Decode((q16.x >> 16) & 0xffU); + w[3] = infiniopFp8E4m3Decode(q16.x >> 24); + w[4] = infiniopFp8E4m3Decode(q16.y & 0xffU); + w[5] = infiniopFp8E4m3Decode((q16.y >> 8) & 0xffU); + w[6] = infiniopFp8E4m3Decode((q16.y >> 16) & 0xffU); + w[7] = infiniopFp8E4m3Decode(q16.y >> 24); + w[8] = infiniopFp8E4m3Decode(q16.z & 0xffU); + w[9] = infiniopFp8E4m3Decode((q16.z >> 8) & 0xffU); + w[10] = infiniopFp8E4m3Decode((q16.z >> 16) & 0xffU); + w[11] = infiniopFp8E4m3Decode(q16.z >> 24); + w[12] = infiniopFp8E4m3Decode(q16.w & 0xffU); + w[13] = infiniopFp8E4m3Decode((q16.w >> 8) & 0xffU); + w[14] = infiniopFp8E4m3Decode((q16.w >> 16) & 0xffU); + w[15] = infiniopFp8E4m3Decode(q16.w >> 24); +} + +template +INFINIOP_CUDA_KERNEL fp8_blockwise_gemm_kernel( + T *__restrict__ out, + const T *__restrict__ a, + const uint8_t *__restrict__ q, + const float *__restrict__ scales, + size_t M, size_t N, size_t K, + size_t block_n, size_t block_k, size_t scales_cols) { + const size_t row = static_cast(blockIdx.x) * TN + (threadIdx.x >> 5); + if (row >= N) { + return; + } + const size_t m0 = static_cast(blockIdx.y) * M_TILE; + const int lane = threadIdx.x & 31; + + { + const uint8_t *q_row = q + row * K; + const size_t scale_row = row / block_n; + + float acc[M_TILE]; +#pragma unroll + for (int m = 0; m < M_TILE; ++m) { + acc[m] = 0.0f; + } + + // Main loop: 512-byte groups (lane*16 within group; lanes 0-7 cover the + // first 128-wide sub-chunk, 8-15 the second, etc.). + const size_t k_groups = K / 512; + const size_t lane_off = static_cast(lane) * 16; + const size_t sub_chunk = static_cast(lane) / 8; + + uint4 q_next = (k_groups > 0) + ? *reinterpret_cast(q_row + lane_off) + : make_uint4(0, 0, 0, 0); + for (size_t kg = 0; kg < k_groups; ++kg) { + const uint4 q_cur = q_next; + const size_t k = (kg + 1) * 512 + lane_off; + if (kg + 1 < k_groups) { + q_next = *reinterpret_cast(q_row + k); + } + + float w[16]; + decode16(q_cur, w); + + const size_t k_base = kg * 512 + sub_chunk * 128; + const float scale = scales[scale_row * scales_cols + k_base / block_k]; + + float cacc[M_TILE]; +#pragma unroll + for (int m = 0; m < M_TILE; ++m) { + cacc[m] = 0.0f; + } +#pragma unroll + for (int m = 0; m < M_TILE; ++m) { + if (m0 + m >= M) { + break; + } + float av[16]; + load16(a + (m0 + m) * K + kg * 512 + lane_off, av); +#pragma unroll + for (int j = 0; j < 16; ++j) { + cacc[m] = fmaf(w[j], av[j], cacc[m]); + } + } +#pragma unroll + for (int m = 0; m < M_TILE; ++m) { + acc[m] = fmaf(scale, cacc[m], acc[m]); + } + } + + // Tail: remaining 128-wide chunks (K % 512 != 0), 4 bytes per lane. + for (size_t kc = k_groups * 4; kc < K / 128; ++kc) { + const size_t k = kc * 128 + static_cast(lane) * 4; + const uint32_t q4 = *reinterpret_cast(q_row + k); + float w[4]; + w[0] = infiniopFp8E4m3Decode(q4 & 0xffU); + w[1] = infiniopFp8E4m3Decode((q4 >> 8) & 0xffU); + w[2] = infiniopFp8E4m3Decode((q4 >> 16) & 0xffU); + w[3] = infiniopFp8E4m3Decode(q4 >> 24); + + const float scale = scales[scale_row * scales_cols + (kc * 128) / block_k]; + float cacc[M_TILE]; +#pragma unroll + for (int m = 0; m < M_TILE; ++m) { + cacc[m] = 0.0f; + } +#pragma unroll + for (int m = 0; m < M_TILE; ++m) { + if (m0 + m >= M) { + break; + } + float av[4]; + load4(a + (m0 + m) * K + k, av); + cacc[m] = w[0] * av[0] + w[1] * av[1] + w[2] * av[2] + w[3] * av[3]; + } +#pragma unroll + for (int m = 0; m < M_TILE; ++m) { + acc[m] = fmaf(scale, cacc[m], acc[m]); + } + } + + // Warp reduce and store. +#pragma unroll + for (int m = 0; m < M_TILE; ++m) { + if (m0 + m >= M) { + break; + } + float v = acc[m]; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + v += __shfl_xor_sync(0xffffffffu, v, offset); + } + if (lane == 0) { + out[(m0 + m) * N + row] = from_float(v); + } + } + } +} + +// --------------------------------------------------------------------------- +// Tensor-core path (mma.m16n8k16) for 9 <= M <= 32 with F16/BF16 activations. +// +// The SIMT warp-per-row kernel above re-reads the activation row per weight +// element and becomes instruction-throughput bound once M grows (W4: ~5-6 +// TFLOP/s at M=16). This path instead treats decode as a skinny GEMM: +// CTA tile = M_BLOCKS*16 x 32 (N), K streamed in 128-wide chunks +// warp = one n8 block; mma.m16n8k16.row.col accumulates each 128-K +// chunk into a partial C, which is then promoted with the +// (n-block, k-chunk) scale: c_fin += scale * c_part. +// FP8 codes are decoded in registers with a bit-placement trick whose result +// is the true value times 2^-120 (BF16) / 2^-8 (F16); that power-of-two +// factor is folded into the block scale at promote time, so the mma inputs +// are exact and no per-element multiply is needed. (The NaN code 0x7F is not +// special-cased: the encoder saturates at 448, so quantized weights never +// contain it.) +// +// K % 128 == 0 and block_k % 128 == 0 are guaranteed by Fp8BlockwiseGemmInfo, +// so every 128-wide K chunk maps to exactly one scale column. Rows beyond +// M/N are zero-filled on load and discarded on store, so any M/N tail works. +// --------------------------------------------------------------------------- + +constexpr int MMA_N_TILE = 32; // weight rows per CTA (4 warps x n8) +constexpr int MMA_K_CHUNK = 128; // K per pipeline stage (one scale sub-chunk) +constexpr int MMA_THREADS = 128; // 4 warps +constexpr int MMA_A_STRIDE = 136; // sA row stride in elements (128 + 8 pad) +constexpr int MMA_W_STRIDE = 144; // sW row stride in bytes (128 + 16 pad) + +template +struct MmaTraits; + +template <> +struct MmaTraits<__nv_bfloat16> { + // Two e4m3 codes (low byte first) -> bf16x2, each the true value * 2^-120. + static __device__ __forceinline__ uint32_t decodePair(uint32_t two) { + const uint32_t lo = ((two & 0x7fU) << 4) | ((two & 0x80U) << 8); + const uint32_t hi = ((two & 0x7f00U) >> 4) | (two & 0x8000U); + return lo | (hi << 16); + } + static constexpr float kDecodeScale = 0x1p+120f; + static __device__ __forceinline__ void mma(float c[4], const uint32_t a[4], const uint32_t b[2]) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + } +}; + +template <> +struct MmaTraits { + // Two e4m3 codes (low byte first) -> half2, each the true value * 2^-8. + static __device__ __forceinline__ uint32_t decodePair(uint32_t two) { + const uint32_t lo = ((two & 0x7fU) << 7) | ((two & 0x80U) << 8); + const uint32_t hi = ((two & 0x7f00U) >> 1) | (two & 0x8000U); + return lo | (hi << 16); + } + static constexpr float kDecodeScale = 256.0f; + static __device__ __forceinline__ void mma(float c[4], const uint32_t a[4], const uint32_t b[2]) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + } +}; + +template +INFINIOP_CUDA_KERNEL fp8_blockwise_gemm_mma_kernel( + T *__restrict__ out, + const T *__restrict__ a, + const uint8_t *__restrict__ q, + const float *__restrict__ scales, + size_t M, size_t N, size_t K, + size_t block_n, size_t block_k, size_t scales_cols) { + + constexpr int A_ROWS = M_BLOCKS * 16; + __shared__ T sA[2][A_ROWS][MMA_A_STRIDE]; + __shared__ uint8_t sW[2][MMA_N_TILE][MMA_W_STRIDE]; + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int g = lane >> 2; // mma group id (row within m16 / column within n8) + const int t = lane & 3; // mma thread id in group + + const size_t n_base = static_cast(blockIdx.x) * MMA_N_TILE; + const size_t m_base = static_cast(blockIdx.y) * A_ROWS; + + const int nchunks = static_cast(K / MMA_K_CHUNK); + const int kb_per_scale = static_cast(block_k / MMA_K_CHUNK); + + // Global -> register staging. A: A_ROWS*256B over 128 threads (16B each, + // row = idx/16, seg = idx%16). W: 32 rows x 128B (row = idx/8, seg = idx%8). + // Out-of-range rows are zero-filled (they never contribute to the output). + uint4 a_stage[A_ROWS / 8]; + uint4 w_stage[2]; + auto stage_chunk = [&](int c) { + const size_t k0 = static_cast(c) * MMA_K_CHUNK; +#pragma unroll + for (int i = 0; i < A_ROWS / 8; ++i) { + const int idx = tid + i * MMA_THREADS; + const size_t m = m_base + (idx >> 4); + a_stage[i] = (m < M) + ? *reinterpret_cast(a + m * K + k0 + (idx & 15) * 8) + : make_uint4(0, 0, 0, 0); + } +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int idx = tid + i * MMA_THREADS; + const size_t n = n_base + (idx >> 3); + w_stage[i] = (n < N) + ? *reinterpret_cast(q + n * K + k0 + (idx & 7) * 16) + : make_uint4(0, 0, 0, 0); + } + }; + auto store_chunk = [&](int buf) { +#pragma unroll + for (int i = 0; i < A_ROWS / 8; ++i) { + const int idx = tid + i * MMA_THREADS; + *reinterpret_cast(&sA[buf][idx >> 4][(idx & 15) * 8]) = a_stage[i]; + } +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int idx = tid + i * MMA_THREADS; + *reinterpret_cast(&sW[buf][idx >> 3][(idx & 7) * 16]) = w_stage[i]; + } + }; + + // Accumulators: c_fin holds the scale-promoted sum over all K chunks; + // c_part is the raw mma result of the current 128-wide chunk. + float c_fin[M_BLOCKS][4]; +#pragma unroll + for (int mb = 0; mb < M_BLOCKS; ++mb) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + c_fin[mb][i] = 0.0f; + } + } + + stage_chunk(0); + store_chunk(0); + __syncthreads(); + + for (int c = 0; c < nchunks; ++c) { + const int buf = c & 1; + const bool has_next = (c + 1 < nchunks); + if (has_next) { + stage_chunk(c + 1); // loads in flight during the compute below + } + + float c_part[M_BLOCKS][4]; +#pragma unroll + for (int mb = 0; mb < M_BLOCKS; ++mb) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + c_part[mb][i] = 0.0f; + } + } + + // B fragments for this warp's n8 block, decoded on the fly. + const uint8_t *wrow = &sW[buf][warp * 8 + g][0]; +#pragma unroll + for (int step = 0; step < MMA_K_CHUNK / 16; ++step) { + const int kk = step * 16; + const uint32_t b[2] = { + MmaTraits::decodePair(*reinterpret_cast(wrow + kk + t * 2)), + MmaTraits::decodePair(*reinterpret_cast(wrow + kk + t * 2 + 8)), + }; +#pragma unroll + for (int mb = 0; mb < M_BLOCKS; ++mb) { + const T *arow0 = &sA[buf][mb * 16 + g][0]; + const T *arow1 = &sA[buf][mb * 16 + g + 8][0]; + const uint32_t a_frag[4] = { + *reinterpret_cast(arow0 + kk + t * 2), + *reinterpret_cast(arow1 + kk + t * 2), + *reinterpret_cast(arow0 + kk + t * 2 + 8), + *reinterpret_cast(arow1 + kk + t * 2 + 8), + }; + MmaTraits::mma(c_part[mb], a_frag, b); + } + } + + // Promote the chunk partials with the block scale (the decode + // power-of-two factor folded in). The two columns a thread holds + // (2t, 2t+1) always sit in one scale block (block_n % 16 == 0). + const size_t n0 = n_base + warp * 8 + t * 2; + float s = 0.0f; + if (n0 < N) { + s = scales[(n0 / block_n) * scales_cols + c / kb_per_scale] * MmaTraits::kDecodeScale; + } +#pragma unroll + for (int mb = 0; mb < M_BLOCKS; ++mb) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + c_fin[mb][i] = fmaf(s, c_part[mb][i], c_fin[mb][i]); + } + } + + if (has_next) { + store_chunk(buf ^ 1); + } + __syncthreads(); + } + + // Epilogue: thread (g, t) owns C rows g / g+8 and columns 2t / 2t+1. +#pragma unroll + for (int mb = 0; mb < M_BLOCKS; ++mb) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + const size_t m = m_base + mb * 16 + g + (i >> 1) * 8; + const size_t n = n_base + warp * 8 + t * 2 + (i & 1); + if (m < M && n < N) { + out[m * N + n] = from_float(c_fin[mb][i]); + } + } + } +} + +template +void launch_mma_mblocks(T *out, const T *a, const uint8_t *q, const float *scales, + const Fp8BlockwiseGemmInfo &info, cudaStream_t stream) { + dim3 grid((info.N + MMA_N_TILE - 1) / MMA_N_TILE, + (info.M + M_BLOCKS * 16 - 1) / (M_BLOCKS * 16)); + fp8_blockwise_gemm_mma_kernel<<>>( + out, a, q, scales, info.M, info.N, info.K, + info.block_n, info.block_k, info.scales_cols); +} + +template +void launch_mtile(T *out, const T *a, const uint8_t *q, const float *scales, + const Fp8BlockwiseGemmInfo &info, cudaStream_t stream) { + dim3 grid((info.N + TN - 1) / TN, (info.M + M_TILE - 1) / M_TILE); + fp8_blockwise_gemm_kernel<<>>( + out, a, q, scales, info.M, info.N, info.K, + info.block_n, info.block_k, info.scales_cols); +} + +template +void launch(T *out, const T *a, const uint8_t *q, const float *scales, + const Fp8BlockwiseGemmInfo &info, cudaStream_t stream) { + const size_t m = info.M; + // Tensor-core path for the decode range where the SIMT kernel goes + // instruction-throughput bound (W4: fused loses to naive from M >= 16). + // INFINIOP_FP8_GEMM_MMA=0 forces the SIMT kernels (A/B debugging). + if constexpr (std::is_same_v || std::is_same_v) { + static const bool mma_enabled = [] { + const char *env = std::getenv("INFINIOP_FP8_GEMM_MMA"); + return env == nullptr || env[0] != '0'; + }(); + if (mma_enabled && m > 8 && m <= 32) { + if (m <= 16) { + launch_mma_mblocks(out, a, q, scales, info, stream); + } else { + launch_mma_mblocks(out, a, q, scales, info, stream); + } + return; + } + } + if (m <= 1) { + launch_mtile(out, a, q, scales, info, stream); + } else if (m <= 2) { + launch_mtile(out, a, q, scales, info, stream); + } else if (m <= 4) { + launch_mtile(out, a, q, scales, info, stream); + } else if (m <= 8) { + launch_mtile(out, a, q, scales, info, stream); + } else { + launch_mtile(out, a, q, scales, info, stream); + } +} + +} // namespace + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { delete _opaque; } + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t scales_desc) { + auto info = Fp8BlockwiseGemmInfo::create(out_desc, a_desc, q_desc, scales_desc); + CHECK_RESULT(info); + auto nvidia_handle = reinterpret_cast(handle); + *desc_ptr = new Descriptor( + new Opaque{nvidia_handle->internal()}, info.take(), handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *, size_t, void *out, + const void *a, const void *q, const void *scales, void *stream) const { + auto cuda_stream = reinterpret_cast(stream); + auto q_ptr = reinterpret_cast(q); + auto scales_ptr = reinterpret_cast(scales); + switch (_info.dtype) { + case INFINI_DTYPE_F16: + launch(reinterpret_cast(out), reinterpret_cast(a), q_ptr, scales_ptr, _info, cuda_stream); + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_BF16: + launch(reinterpret_cast<__nv_bfloat16 *>(out), reinterpret_cast(a), q_ptr, scales_ptr, _info, cuda_stream); + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_F32: + launch(reinterpret_cast(out), reinterpret_cast(a), q_ptr, scales_ptr, _info, cuda_stream); + return INFINI_STATUS_SUCCESS; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::fp8_blockwise_gemm::nvidia diff --git a/src/infiniop/ops/fp8_blockwise_gemm/nvidia/fp8_blockwise_gemm_nvidia.cuh b/src/infiniop/ops/fp8_blockwise_gemm/nvidia/fp8_blockwise_gemm_nvidia.cuh new file mode 100644 index 000000000..f70f872f7 --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_gemm/nvidia/fp8_blockwise_gemm_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __FP8_BLOCKWISE_GEMM_NVIDIA_CUH__ +#define __FP8_BLOCKWISE_GEMM_NVIDIA_CUH__ + +#include "../fp8_blockwise_gemm.h" + +DESCRIPTOR(nvidia) + +#endif diff --git a/src/infiniop/ops/fp8_blockwise_gemm/operator.cc b/src/infiniop/ops/fp8_blockwise_gemm/operator.cc new file mode 100644 index 000000000..6c8aadca5 --- /dev/null +++ b/src/infiniop/ops/fp8_blockwise_gemm/operator.cc @@ -0,0 +1,102 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/fp8_blockwise_gemm.h" + +#ifdef ENABLE_CPU_API +#include "cpu/fp8_blockwise_gemm_cpu.h" +#endif +#ifdef ENABLE_NVIDIA_API +#include "nvidia/fp8_blockwise_gemm_nvidia.cuh" +#endif + +__INFINI_C 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) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::fp8_blockwise_gemm::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + out_desc, a_desc, q_desc, scales_desc) + switch (handle->device) { +#ifdef ENABLE_CPU_API + CREATE(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetFp8BlockwiseGemmWorkspaceSize( + infiniopFp8BlockwiseGemmDescriptor_t desc, + size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc) \ + ->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_CPU_API + GET(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C 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) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, out, a, q, scales, stream) + switch (desc->device_type) { +#ifdef ENABLE_CPU_API + CALCULATE(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyFp8BlockwiseGemmDescriptor( + infiniopFp8BlockwiseGemmDescriptor_t desc) { +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_CPU_API + DESTROY(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/paged_attention/ascend/paged_attention_ascend.cc b/src/infiniop/ops/paged_attention/ascend/paged_attention_ascend.cc index ed0b10854..506ee75a2 100644 --- a/src/infiniop/ops/paged_attention/ascend/paged_attention_ascend.cc +++ b/src/infiniop/ops/paged_attention/ascend/paged_attention_ascend.cc @@ -19,11 +19,16 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t block_tables_desc, infiniopTensorDescriptor_t cache_lens_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { auto info = PagedAttentionInfo::create( out_desc, q_desc, k_cache_desc, v_cache_desc, - block_tables_desc, cache_lens_desc, alibi_slopes_desc, scale); + block_tables_desc, cache_lens_desc, alibi_slopes_desc, k_scale_desc, v_scale_desc, scale); CHECK_RESULT(info); + if (info->cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } auto handle_ascend = reinterpret_cast(handle); *desc_ptr = new Descriptor( @@ -40,6 +45,7 @@ infiniStatus_t Descriptor::calculate( void *workspace, size_t workspace_size, void *out, const void *q, const void *k_cache, const void *v_cache, const void *block_tables, const void *cache_lens, const void *alibi_slopes, + const void *k_scale, const void *v_scale, void *stream) const { (void)workspace; (void)workspace_size; diff --git a/src/infiniop/ops/paged_attention/bang/paged_attention_bang.mlu b/src/infiniop/ops/paged_attention/bang/paged_attention_bang.mlu index 18734d583..7a3da7727 100644 --- a/src/infiniop/ops/paged_attention/bang/paged_attention_bang.mlu +++ b/src/infiniop/ops/paged_attention/bang/paged_attention_bang.mlu @@ -297,11 +297,16 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t block_tables_desc, infiniopTensorDescriptor_t seq_lens_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { auto handle = reinterpret_cast(handle_); - auto info = PagedAttentionInfo::create(out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, seq_lens_desc, alibi_slopes_desc, scale); + auto info = PagedAttentionInfo::create(out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, seq_lens_desc, alibi_slopes_desc, k_scale_desc, v_scale_desc, scale); CHECK_RESULT(info); + if (info->cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } *desc_ptr = new Descriptor( new Opaque{static_cast(handle)->internal()}, @@ -320,6 +325,8 @@ infiniStatus_t Descriptor::calculate( const void *block_tables, const void *seq_lens, const void *alibi_slopes, + const void *k_scale, + const void *v_scale, void *stream) const { (void)workspace; diff --git a/src/infiniop/ops/paged_attention/cuda/kernel.cuh b/src/infiniop/ops/paged_attention/cuda/kernel.cuh index c6ccd049b..c19f18372 100644 --- a/src/infiniop/ops/paged_attention/cuda/kernel.cuh +++ b/src/infiniop/ops/paged_attention/cuda/kernel.cuh @@ -32,6 +32,12 @@ __device__ void pagedAttentionKernel( const int num_heads = gridDim.x; const int64_t seq_len = seq_lens_[seq_idx]; if (seq_len == 0) { + // Zero-length sequence: write defined zeros instead of leaving the + // caller's output buffer untouched. + Tdata *out_ptr = out_ + seq_idx * o_stride + head_idx * HEAD_SIZE; + for (size_t h_dim = threadIdx.x; h_dim < HEAD_SIZE; h_dim += NUM_THREADS) { + out_ptr[h_dim] = static_cast(0.0f); + } return; } diff --git a/src/infiniop/ops/paged_attention/cuda/kernel_fp8.cuh b/src/infiniop/ops/paged_attention/cuda/kernel_fp8.cuh new file mode 100644 index 000000000..c8158b94b --- /dev/null +++ b/src/infiniop/ops/paged_attention/cuda/kernel_fp8.cuh @@ -0,0 +1,415 @@ +#ifndef __PAGED_ATTENTION_FP8_KERNEL_CUH__ +#define __PAGED_ATTENTION_FP8_KERNEL_CUH__ + +//================================================================================ +// Paged Attention Decode Kernel for FP8(E4M3) KV Caches (clean-room) +// +// v1.6: deeper memory-level parallelism on top of the v1.5 warp-parallel scan. +// One CTA of NUM_WARPS*32 threads (default 8 warps = 256 threads) per +// (sequence, query head); warp `w` visits tokens w, w+NUM_WARPS, ... of every +// referenced page (per-page striding; the global token index t = t_base + tb +// keeps ALiBi correct for any page size). Each lane owns DL = HEAD_SIZE/32 +// consecutive head dims (4 for hd128, 2 for hd64), so a single uint32/uint16 +// load fetches all of the lane's E4M3 codes for one token and the warp covers +// HEAD_SIZE contiguous bytes per K/V row. +// +// The token stream is walked with a (logical block, token-in-block) cursor and +// register double buffering: while token i is being processed, the packed K/V +// words and both per-token scales of token i+1 are already in flight (loads +// issued one iteration ahead, across page boundaries). The qk dot product is +// a shuffle-only warp reduction; every warp keeps its own online-softmax state +// (m, l, acc[DL]) in registers; the partial states are merged once at the end +// through shared memory (the only __syncthreads in the kernel), rescaling by +// exp2(m_w - m_total) in the standard online-softmax fashion. +// +// Semantics are unchanged from v1: dequant-on-load +// x = e4m3_decode(code) * scale[physical_block, kv_head, slot] +// with per-token-per-kv-head F32 scales written by paged_caching, online +// softmax in the log2 domain (scale * log2e, optional ALiBi, final division +// by l + 1e-6), F16/BF16 output, HEAD_SIZE in {64, 128}. +// +// v2: optional cross-CTA split-kv (flash-decoding). When num_splits > 1 the +// grid gains a z dimension (split_idx = blockIdx.z) and each CTA scans only +// the contiguous token shard [split_idx*shard, min(seq_len, +shard)) of its +// (sequence, query head); instead of the final output it writes the merged +// per-CTA state (m, l, unnormalized acc) to workspace partials, laid out as +// partial_acc [num_splits, num_seqs, num_heads, HEAD_SIZE] +// partial_m/l [num_splits, num_seqs, num_heads] +// (same convention as the F16/BF16 family in kernel_v2.cuh). A second +// combine kernel (one CTA of HEAD_SIZE threads per (seq, head)) then merges +// the shards in the log2 domain and divides by l + 1e-6. Shards with no +// token produce the neutral element (m = -inf, l = 0, acc = 0). +// +// Alignment note: the packed code loads assume each K/V row is DL-byte +// aligned, i.e. the cache base pointer and the batch/head/row strides are +// multiples of DL bytes. This holds for any contiguous cache pool (row stride +// == HEAD_SIZE) and for block/head slices of one. +//================================================================================ + +#include +#include + +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" + +namespace op::paged_attention::cuda { + +// Number of warps per CTA in the FP8 decode kernel. Tunable (4/8/16); the +// launcher sizes the block from this constant, so the two stay in sync. +constexpr int kFp8DecodeNumWarps = 32; + +// Upper bound on cross-CTA split-kv shards; the launcher clamps num_splits to +// this and sizes the workspace for it. +constexpr int kFp8DecodeMaxSplits = 8; + +template +__device__ void flashAttentionDecodeFp8Kernel( + Tdata *out_, + float *partial_acc_, // [num_splits, num_seqs, num_heads, HEAD_SIZE]; nullptr => no split-kv + float *partial_m_, // [num_splits, num_seqs, num_heads] + float *partial_l_, // [num_splits, num_seqs, num_heads] + const Tdata *q_, + const uint8_t *k_cache_, + const uint8_t *v_cache_, + const float *k_scale_, + const float *v_scale_, + const Tindex *block_tables_, + const Tindex *cache_lens_, + const float *alibi_slopes_, + size_t num_kv_heads, + float scale, + size_t max_num_blocks_per_seq, + size_t page_block_size, + ptrdiff_t q_stride, + ptrdiff_t q_head_stride, + ptrdiff_t k_batch_stride, + ptrdiff_t k_row_stride, + ptrdiff_t k_head_stride, + ptrdiff_t v_batch_stride, + ptrdiff_t v_row_stride, + ptrdiff_t v_head_stride, + ptrdiff_t o_stride, + ptrdiff_t o_head_stride, + ptrdiff_t k_scale_block_stride, + ptrdiff_t k_scale_head_stride, + ptrdiff_t k_scale_slot_stride, + ptrdiff_t v_scale_block_stride, + ptrdiff_t v_scale_head_stride, + ptrdiff_t v_scale_slot_stride, + int num_splits) { + + static_assert(HEAD_SIZE == 64 || HEAD_SIZE == 128, + "FP8 decode kernel supports head_size 64/128 only."); + + constexpr int DL = HEAD_SIZE / 32; // head dims per lane: 4 (hd128) or 2 (hd64) + using PackT = std::conditional_t
; + constexpr int NUM_WARPS = kFp8DecodeNumWarps; + + const size_t seq_idx = blockIdx.y; + const size_t head_idx = blockIdx.x; + const int split_idx = static_cast(blockIdx.z); + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + const int seq_len = static_cast(cache_lens_[seq_idx]); + if (seq_len <= 0) { + // A zero-length sequence has no tokens to attend to. Publish the + // neutral element (split-kv partials) or zeros (direct output) so the + // row is well-defined: combine would otherwise read stale workspace + // and the direct path would leave the caller's buffer untouched. + const int tid0 = threadIdx.x; + if (partial_m_ != nullptr) { + const size_t n0 = gridDim.y * gridDim.x; // num_seqs * num_heads + const size_t idx0 = static_cast(split_idx) * n0 + seq_idx * gridDim.x + head_idx; + if (tid0 == 0) { + partial_m_[idx0] = -INFINITY; + partial_l_[idx0] = 0.0f; + } + if (tid0 < HEAD_SIZE) { + partial_acc_[idx0 * HEAD_SIZE + tid0] = 0.0f; + } + } else if (tid0 < HEAD_SIZE) { + Tdata *out_ptr = out_ + seq_idx * o_stride + head_idx * o_head_stride + tid0; + if constexpr (std::is_same_v) { + *out_ptr = __float2half_rn(0.0f); + } else if constexpr (std::is_same_v) { + *out_ptr = __float2bfloat16_rn(0.0f); + } else { + *out_ptr = static_cast(0.0f); + } + } + return; + } + + // This CTA's contiguous token shard [tok_lo, tok_hi). Without split-kv + // (num_splits == 1) the shard is the whole sequence. + const int shard = (seq_len + num_splits - 1) / num_splits; + const int tok_lo = min(split_idx * shard, seq_len); + const int tok_hi = min(seq_len, tok_lo + shard); + + const size_t num_heads = gridDim.x; + const size_t num_queries_per_kv = num_heads / num_kv_heads; + const size_t kv_head_idx = head_idx / num_queries_per_kv; + + const float alibi_slope = (alibi_slopes_ == nullptr) ? 0.0f : alibi_slopes_[head_idx]; + constexpr float kLog2e = 1.4426950408889634f; + const float scale_log2 = scale * kLog2e; + + const Tindex *block_table = block_tables_ + seq_idx * max_num_blocks_per_seq; + + // This lane's head dims are d0 .. d0+DL-1 (contiguous => packed byte loads). + const int d0 = lane * DL; + const Tdata *q_ptr = q_ + seq_idx * q_stride + head_idx * q_head_stride + d0; + float q_reg[DL]; +#pragma unroll + for (int j = 0; j < DL; ++j) { + q_reg[j] = static_cast(q_ptr[j]); + } + + // Per-warp online softmax state; acc holds this lane's DL dims. + float acc[DL]; +#pragma unroll + for (int j = 0; j < DL; ++j) { + acc[j] = 0.0f; + } + float m = -INFINITY; + float l = 0.0f; + + const int pbs = static_cast(page_block_size); + + // This CTA's token window within logical block `lb`, intersected with the + // split shard: [tokenBegin, tokenEnd). tokenBegin is nonzero only on the + // shard's first page; the last page of the shard may be partial. + auto tokenBegin = [&](int lb) -> int { + return max(tok_lo - lb * pbs, 0); + }; + auto tokenEnd = [&](int lb) -> int { + return min(pbs, tok_hi - lb * pbs); + }; + + // One token's prefetched payload: the lane's packed K/V code words plus + // the (warp-uniform) per-token dequant scales. + struct KvPack { + PackT k, v; + float ks, vs; + }; + + // Load token `tb` of logical block `lb` into `pack`. All four loads are + // independent; the block_table read is L1-cached and shared by all warps. + auto loadToken = [&](int lb, int tb, KvPack &pack) { + const ptrdiff_t physical_block = static_cast(block_table[lb]); + const uint8_t *k_base = k_cache_ + physical_block * k_batch_stride + static_cast(kv_head_idx) * k_head_stride; + const uint8_t *v_base = v_cache_ + physical_block * v_batch_stride + static_cast(kv_head_idx) * v_head_stride; + const float *k_scale_base = k_scale_ + physical_block * k_scale_block_stride + static_cast(kv_head_idx) * k_scale_head_stride; + const float *v_scale_base = v_scale_ + physical_block * v_scale_block_stride + static_cast(kv_head_idx) * v_scale_head_stride; + pack.k = *reinterpret_cast(k_base + tb * k_row_stride + d0); + pack.v = *reinterpret_cast(v_base + tb * v_row_stride + d0); + pack.ks = k_scale_base[tb * k_scale_slot_stride]; + pack.vs = v_scale_base[tb * v_scale_slot_stride]; + }; + + // Per-warp token cursor: (logical block, token-in-block). Within each + // page's window warp `warp` owns tokens tokenBegin+warp, + // tokenBegin+warp+NUM_WARPS, ...; pages with no token for this warp are + // skipped. The shard bounds make this equivalent to the v1 full-sequence + // walk when num_splits == 1. + int cur_lb = tok_lo / pbs; + int cur_tb = tokenBegin(cur_lb) + warp; + while (cur_lb * pbs < tok_hi && cur_tb >= tokenEnd(cur_lb)) { + ++cur_lb; + cur_tb = tokenBegin(cur_lb) + warp; + } + bool has_cur = cur_lb * pbs < tok_hi; + + KvPack cur{}, nxt{}; + if (has_cur) { + loadToken(cur_lb, cur_tb, cur); + } + + // Software-pipelined scan: process `cur` while `nxt` is being fetched. + while (has_cur) { + // Advance the cursor (possibly across page boundaries) and issue the + // next token's loads before touching the current payload. + int nxt_lb = cur_lb; + int nxt_tb = cur_tb + NUM_WARPS; + if (nxt_tb >= tokenEnd(nxt_lb)) { + ++nxt_lb; + nxt_tb = tokenBegin(nxt_lb) + warp; + while (nxt_lb * pbs < tok_hi && nxt_tb >= tokenEnd(nxt_lb)) { + ++nxt_lb; + nxt_tb = tokenBegin(nxt_lb) + warp; + } + } + const bool has_nxt = nxt_lb * pbs < tok_hi; + if (has_nxt) { + loadToken(nxt_lb, nxt_tb, nxt); + } + + const int t = cur_lb * pbs + cur_tb; + + float partial = 0.0f; +#pragma unroll + for (int j = 0; j < DL; ++j) { + const uint8_t code = static_cast((cur.k >> (8 * j)) & 0xFF); + partial += q_reg[j] * (infiniopFp8E4m3Decode(code) * cur.ks); + } + // Shuffle-only warp reduction (no smem, no __syncthreads). +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + partial += __shfl_xor_sync(0xffffffff, partial, offset); + } + const float qk = partial; + + float score = qk * scale_log2; + if (alibi_slope != 0.0f) { + score += (alibi_slope * static_cast(t - (seq_len - 1))) * kLog2e; + } + const float m_new = fmaxf(m, score); + const float alpha = exp2f(m - m_new); + const float beta = exp2f(score - m_new); + l = l * alpha + beta; + m = m_new; + +#pragma unroll + for (int j = 0; j < DL; ++j) { + const uint8_t code = static_cast((cur.v >> (8 * j)) & 0xFF); + acc[j] = acc[j] * alpha + beta * (infiniopFp8E4m3Decode(code) * cur.vs); + } + + cur = nxt; + cur_lb = nxt_lb; + cur_tb = nxt_tb; + has_cur = has_nxt; + } + + // ---- Cross-warp merge (the only block-wide synchronization) ---- + __shared__ float m_part[NUM_WARPS]; + __shared__ float l_part[NUM_WARPS]; + __shared__ float acc_part[NUM_WARPS][HEAD_SIZE]; + + if (lane == 0) { + m_part[warp] = m; + l_part[warp] = l; + } +#pragma unroll + for (int j = 0; j < DL; ++j) { + acc_part[warp][d0 + j] = acc[j]; + } + __syncthreads(); + + // Scalar (m, l) merge, computed redundantly by all threads. A warp that + // saw no token has m = -inf / l = 0 and contributes weight 0. The explicit + // -inf guard (instead of relying on exp2f(-inf - m_total) == 0) also keeps + // an all-empty split shard — where m_total itself is -inf and + // exp2f(-inf - -inf) would be NaN — at the neutral element (0, 0, 0). + float m_total = m_part[0]; +#pragma unroll + for (int w = 1; w < NUM_WARPS; ++w) { + m_total = fmaxf(m_total, m_part[w]); + } + float wgt[NUM_WARPS]; + float l_total = 0.0f; +#pragma unroll + for (int w = 0; w < NUM_WARPS; ++w) { + wgt[w] = (m_part[w] == -INFINITY) ? 0.0f : exp2f(m_part[w] - m_total); + l_total += l_part[w] * wgt[w]; + } + + // HEAD_SIZE output dims over the CTA: thread `tid` writes dim `tid` + // (threads with tid >= HEAD_SIZE idle out). + const int tid = threadIdx.x; + + // Split-kv: publish this shard's merged (m, l, unnormalized acc) and let + // the combine kernel produce the final output. + if (partial_m_ != nullptr) { + const size_t n = gridDim.y * gridDim.x; // num_seqs * num_heads + const size_t idx = static_cast(split_idx) * n + seq_idx * gridDim.x + head_idx; + if (tid == 0) { + partial_m_[idx] = (l_total > 0.0f) ? m_total : -INFINITY; + partial_l_[idx] = l_total; + } + if (tid < HEAD_SIZE) { + float o = 0.0f; +#pragma unroll + for (int w = 0; w < NUM_WARPS; ++w) { + o += acc_part[w][tid] * wgt[w]; + } + partial_acc_[idx * HEAD_SIZE + tid] = o; + } + return; + } + + const float inv_l = 1.0f / (l_total + 1e-6f); + if (tid < HEAD_SIZE) { + float o = 0.0f; +#pragma unroll + for (int w = 0; w < NUM_WARPS; ++w) { + o += acc_part[w][tid] * wgt[w]; + } + o *= inv_l; + Tdata *out_ptr = out_ + seq_idx * o_stride + head_idx * o_head_stride + tid; + if constexpr (std::is_same_v) { + *out_ptr = __float2half_rn(o); + } else if constexpr (std::is_same_v) { + *out_ptr = __float2bfloat16_rn(o); + } else { + *out_ptr = static_cast(o); + } + } +} + +// Cross-CTA split-kv combine: one CTA of HEAD_SIZE threads per (sequence, +// query head) merges the per-shard partials in the log2 domain and writes the +// final output. Mirrors the FP8 kernel's own cross-warp merge; empty shards +// carry m = -inf / l = 0 and contribute weight 0 (split 0 is never empty for +// seq_len > 0, so m_total stays finite). +template +__device__ void flashAttentionDecodeFp8SplitKvCombineKernel( + Tdata *out_, + const float *partial_acc_, // [num_splits, num_seqs, num_heads, HEAD_SIZE] + const float *partial_m_, // [num_splits, num_seqs, num_heads] + const float *partial_l_, // [num_splits, num_seqs, num_heads] + int num_splits, + ptrdiff_t o_stride, + ptrdiff_t o_head_stride) { + + const size_t seq_idx = blockIdx.y; + const size_t head_idx = blockIdx.x; + const int tid = threadIdx.x; + + const size_t n = gridDim.y * gridDim.x; // num_seqs * num_heads + const size_t base = seq_idx * gridDim.x + head_idx; + + // Scalar (m, l) merge, computed redundantly by all threads. + float m_total = -INFINITY; + for (int s = 0; s < num_splits; ++s) { + m_total = fmaxf(m_total, partial_m_[s * n + base]); + } + float wgt[kFp8DecodeMaxSplits]; + float l_total = 0.0f; + for (int s = 0; s < num_splits; ++s) { + const float ms = partial_m_[s * n + base]; + wgt[s] = (ms == -INFINITY) ? 0.0f : exp2f(ms - m_total); + l_total += partial_l_[s * n + base] * wgt[s]; + } + const float inv_l = 1.0f / (l_total + 1e-6f); + + if (tid < HEAD_SIZE) { + float o = 0.0f; + for (int s = 0; s < num_splits; ++s) { + o += partial_acc_[(s * n + base) * HEAD_SIZE + tid] * wgt[s]; + } + o *= inv_l; + Tdata *out_ptr = out_ + seq_idx * o_stride + head_idx * o_head_stride + tid; + if constexpr (std::is_same_v) { + *out_ptr = __float2half_rn(o); + } else if constexpr (std::is_same_v) { + *out_ptr = __float2bfloat16_rn(o); + } else { + *out_ptr = static_cast(o); + } + } +} + +} // namespace op::paged_attention::cuda + +#endif // __PAGED_ATTENTION_FP8_KERNEL_CUH__ diff --git a/src/infiniop/ops/paged_attention/cuda/kernel_v2.cuh b/src/infiniop/ops/paged_attention/cuda/kernel_v2.cuh index 489c117b5..b63cc1eb3 100644 --- a/src/infiniop/ops/paged_attention/cuda/kernel_v2.cuh +++ b/src/infiniop/ops/paged_attention/cuda/kernel_v2.cuh @@ -163,6 +163,20 @@ __device__ void flashAttentionDecodeWarpKernel( const int seq_len = static_cast(cache_lens_[seq_idx]); if (seq_len <= 0) { + // Zero-length sequence: write defined zeros instead of leaving the + // caller's output buffer untouched. + Tdata *out_ptr = out_ + seq_idx * o_stride + head_idx * HEAD_SIZE; +#pragma unroll + for (int i = 0; i < DIMS_PER_THREAD; ++i) { + const int dim = lane * DIMS_PER_THREAD + i; + if constexpr (std::is_same_v) { + out_ptr[dim] = __float2half_rn(0.0f); + } else if constexpr (std::is_same_v) { + out_ptr[dim] = __float2bfloat16_rn(0.0f); + } else { + out_ptr[dim] = static_cast(0.0f); + } + } return; } @@ -371,7 +385,11 @@ __device__ void flashAttentionDecodeSplitKvWarpKernel( constexpr int DIMS_PER_THREAD = HEAD_SIZE / kWarpSize; const int seq_len = static_cast(cache_lens_[seq_idx]); - if (seq_len <= 0 || num_splits <= 0) { + // No early return for seq_len <= 0: shard becomes 0, so every split takes + // the empty-shard path below and publishes the neutral element (m=-inf, + // l=0, acc=0). Combine then emits a defined zero row instead of reading + // stale workspace. + if (num_splits <= 0) { return; } @@ -601,7 +619,7 @@ __device__ void flashAttentionDecodeSplitKvCombineWarpKernel( float acc = 0.0f; for (int s = 0; s < num_splits; ++s) { const float ms = partial_m[s * n + base]; - const float w = exp2f(ms - m); + const float w = (ms == -INFINITY) ? 0.0f : exp2f(ms - m); acc += partial_acc[(s * n + base) * HEAD_SIZE + dim] * w; } const float o = acc * inv_l; @@ -661,7 +679,11 @@ __device__ void flashAttentionDecodeSplitKvCtaKernel( const int warp_id = tid / kWarpSize; const int seq_len = static_cast(cache_lens_[seq_idx]); - if (seq_len <= 0 || num_splits <= 0) { + // No early return for seq_len <= 0: shard becomes 0, so every split takes + // the empty-shard path below and publishes the neutral element (m=-inf, + // l=0, acc=0). Combine then emits a defined zero row instead of reading + // stale workspace. + if (num_splits <= 0) { return; } @@ -1108,6 +1130,16 @@ __device__ void flashAttentionDecodeCtaPipelinedKernel( const int seq_len = static_cast(cache_lens_[seq_idx]); if (seq_len <= 0) { + // Zero-length sequence: write defined zeros instead of leaving the + // caller's output buffer untouched. + Tdata *out_ptr = out_ + seq_idx * o_stride + head_idx * HEAD_SIZE; + if constexpr (std::is_same_v) { + out_ptr[tid] = __float2half_rn(0.0f); + } else if constexpr (std::is_same_v) { + out_ptr[tid] = __float2bfloat16_rn(0.0f); + } else { + out_ptr[tid] = static_cast(0.0f); + } return; } @@ -1285,6 +1317,19 @@ __device__ void flashAttentionDecodeCtaKernel( const int seq_len = static_cast(cache_lens_[seq_idx]); if (seq_len <= 0) { + // Zero-length sequence: write defined zeros instead of leaving the + // caller's output buffer untouched. + Tdata *out_ptr = out_ + seq_idx * o_stride + head_idx * HEAD_SIZE; +#pragma unroll + for (int i = 0; i < kPack; ++i) { + if constexpr (std::is_same_v) { + out_ptr[dim + i] = __float2half_rn(0.0f); + } else if constexpr (std::is_same_v) { + out_ptr[dim + i] = __float2bfloat16_rn(0.0f); + } else { + out_ptr[dim + i] = static_cast(0.0f); + } + } return; } @@ -1778,6 +1823,23 @@ __device__ void flashAttentionDecodeCtaGqaKernel( const int seq_len = static_cast(cache_lens_[seq_idx]); if (seq_len <= 0) { + // Zero-length sequence: write defined zeros for all NGROUPS query + // heads instead of leaving the caller's output buffer untouched. +#pragma unroll + for (int g = 0; g < NGROUPS; ++g) { + const int q_head = kv_head_idx * NGROUPS + g; + Tdata *out_ptr = out_ + seq_idx * o_stride + q_head * HEAD_SIZE; +#pragma unroll + for (int i = 0; i < kPack; ++i) { + if constexpr (std::is_same_v) { + out_ptr[dim + i] = __float2half_rn(0.0f); + } else if constexpr (std::is_same_v) { + out_ptr[dim + i] = __float2bfloat16_rn(0.0f); + } else { + out_ptr[dim + i] = static_cast(0.0f); + } + } + } return; } diff --git a/src/infiniop/ops/paged_attention/info.h b/src/infiniop/ops/paged_attention/info.h index 3b7f0364e..d474a0e46 100644 --- a/src/infiniop/ops/paged_attention/info.h +++ b/src/infiniop/ops/paged_attention/info.h @@ -14,6 +14,7 @@ class PagedAttentionInfo { public: infiniDtype_t dtype; + infiniDtype_t cache_dtype; infiniDtype_t index_dtype; float scale; @@ -38,6 +39,16 @@ class PagedAttentionInfo { ptrdiff_t block_table_batch_stride; ptrdiff_t cache_lens_stride; + // --- FP8(E4M3) KV cache: per-token dequant scales [num_blocks, num_kv_heads, block_size] --- + // Only meaningful when cache_dtype == INFINI_DTYPE_F8. + ptrdiff_t q_head_stride; + ptrdiff_t k_scale_block_stride; + ptrdiff_t k_scale_head_stride; + ptrdiff_t k_scale_slot_stride; + ptrdiff_t v_scale_block_stride; + ptrdiff_t v_scale_head_stride; + ptrdiff_t v_scale_slot_stride; + static utils::Result create( infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t q_desc, @@ -46,14 +57,38 @@ class PagedAttentionInfo { infiniopTensorDescriptor_t block_tables_desc, infiniopTensorDescriptor_t cache_lens_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { auto dtype = q_desc->dtype(); CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); - if (out_desc->dtype() != dtype || k_cache_desc->dtype() != dtype || v_cache_desc->dtype() != dtype) { + if (out_desc->dtype() != dtype) { return INFINI_STATUS_BAD_TENSOR_DTYPE; } + // The caches either keep the compute dtype or store FP8(E4M3) codes + // plus per-token F32 scales produced by paged_caching. + auto cache_dtype = k_cache_desc->dtype(); + const bool cache_fp8 = (cache_dtype == INFINI_DTYPE_F8); + const bool has_k_scale = k_scale_desc.has_value() && k_scale_desc.value() != nullptr; + const bool has_v_scale = v_scale_desc.has_value() && v_scale_desc.value() != nullptr; + if (cache_fp8) { + if (v_cache_desc->dtype() != INFINI_DTYPE_F8) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (!has_k_scale || !has_v_scale) { + return INFINI_STATUS_BAD_PARAM; + } + } else { + if (cache_dtype != dtype || v_cache_desc->dtype() != dtype) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (has_k_scale || has_v_scale) { + return INFINI_STATUS_BAD_PARAM; + } + } + if (q_desc->ndim() != 3 || out_desc->ndim() != 3) { return INFINI_STATUS_BAD_TENSOR_SHAPE; } @@ -150,6 +185,22 @@ class PagedAttentionInfo { return INFINI_STATUS_BAD_TENSOR_SHAPE; } + // Per-token dequant scales for the FP8 path: [num_blocks, num_kv_heads, page_block_size]. + if (cache_fp8) { + for (const auto &scale_desc : {k_scale_desc.value(), v_scale_desc.value()}) { + if (scale_desc->dtype() != INFINI_DTYPE_F32) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (scale_desc->ndim() != 3) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + const auto scale_shape = scale_desc->shape(); + if (scale_shape[0] != num_blocks || scale_shape[1] != num_kv_heads || scale_shape[2] != page_block_size) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + } + } + const size_t max_num_blocks_per_seq = block_tables_desc->shape()[1]; // Strides (in elements) @@ -168,8 +219,17 @@ class PagedAttentionInfo { const ptrdiff_t block_table_batch_stride = block_tables_desc->stride(0); const ptrdiff_t cache_lens_stride = cache_lens_desc->stride(0); + const ptrdiff_t q_head_stride = q_desc->stride(1); + const ptrdiff_t k_scale_block_stride = cache_fp8 ? k_scale_desc.value()->stride(0) : 0; + const ptrdiff_t k_scale_head_stride = cache_fp8 ? k_scale_desc.value()->stride(1) : 0; + const ptrdiff_t k_scale_slot_stride = cache_fp8 ? k_scale_desc.value()->stride(2) : 0; + const ptrdiff_t v_scale_block_stride = cache_fp8 ? v_scale_desc.value()->stride(0) : 0; + const ptrdiff_t v_scale_head_stride = cache_fp8 ? v_scale_desc.value()->stride(1) : 0; + const ptrdiff_t v_scale_slot_stride = cache_fp8 ? v_scale_desc.value()->stride(2) : 0; + return utils::Result(PagedAttentionInfo{ dtype, + cache_dtype, block_tables_dt, scale, num_seqs, @@ -190,6 +250,13 @@ class PagedAttentionInfo { o_head_stride, block_table_batch_stride, cache_lens_stride, + q_head_stride, + k_scale_block_stride, + k_scale_head_stride, + k_scale_slot_stride, + v_scale_block_stride, + v_scale_head_stride, + v_scale_slot_stride, }); } }; diff --git a/src/infiniop/ops/paged_attention/metax/paged_attention_metax.maca b/src/infiniop/ops/paged_attention/metax/paged_attention_metax.maca index 2b9d64b17..ac99e6d5d 100644 --- a/src/infiniop/ops/paged_attention/metax/paged_attention_metax.maca +++ b/src/infiniop/ops/paged_attention/metax/paged_attention_metax.maca @@ -192,11 +192,16 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t block_tables_desc, infiniopTensorDescriptor_t cache_lens_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { - auto info_res = PagedAttentionInfo::create(out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, cache_lens_desc, alibi_slopes_desc, scale); + auto info_res = PagedAttentionInfo::create(out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, cache_lens_desc, alibi_slopes_desc, k_scale_desc, v_scale_desc, scale); CHECK_RESULT(info_res); auto info = info_res.take(); + if (info.cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } // Reserve workspace for optional split-kv decode (partial acc + m/l). // Workspace is independent of runtime env toggles; kernels will clamp num_splits <= kMaxSplits. constexpr size_t kMaxSplits = 8; @@ -214,6 +219,7 @@ infiniStatus_t Descriptor::calculate( void *workspace, size_t workspace_size, void *out, const void *q, const void *k_cache, const void *v_cache, const void *block_tables, const void *cache_lens, const void *alibi_slopes, + const void *k_scale, const void *v_scale, void *stream_) const { bool need_workspace = false; diff --git a/src/infiniop/ops/paged_attention/moore/paged_attention_moore.mu b/src/infiniop/ops/paged_attention/moore/paged_attention_moore.mu index 3227a1913..e48045a2b 100644 --- a/src/infiniop/ops/paged_attention/moore/paged_attention_moore.mu +++ b/src/infiniop/ops/paged_attention/moore/paged_attention_moore.mu @@ -188,11 +188,16 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t block_tables_desc, infiniopTensorDescriptor_t cache_lens_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { - auto info_res = PagedAttentionInfo::create(out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, cache_lens_desc, alibi_slopes_desc, scale); + auto info_res = PagedAttentionInfo::create(out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, cache_lens_desc, alibi_slopes_desc, k_scale_desc, v_scale_desc, scale); CHECK_RESULT(info_res); auto info = info_res.take(); + if (info.cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } // Reserve workspace for optional split-kv decode (partial acc + m/l). // Workspace is independent of runtime env toggles; kernels will clamp num_splits <= kMaxSplits. constexpr size_t kMaxSplits = 8; @@ -210,6 +215,7 @@ infiniStatus_t Descriptor::calculate( void *workspace, size_t workspace_size, void *out, const void *q, const void *k_cache, const void *v_cache, const void *block_tables, const void *cache_lens, const void *alibi_slopes, + const void *k_scale, const void *v_scale, void *stream_) const { bool need_workspace = false; diff --git a/src/infiniop/ops/paged_attention/nvidia/paged_attention_fp8.cu b/src/infiniop/ops/paged_attention/nvidia/paged_attention_fp8.cu new file mode 100644 index 000000000..1aa4bfa66 --- /dev/null +++ b/src/infiniop/ops/paged_attention/nvidia/paged_attention_fp8.cu @@ -0,0 +1,368 @@ +#include + +#include +#include +#include +#include + +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" + +#include "../cuda/kernel_fp8.cuh" + +// FP8(E4M3) KV-cache decode launchers. Independent of the F16/BF16 kernel +// family: k_cache/v_cache carry E4M3 codes and k_scale/v_scale the per-token +// F32 dequant scales produced by paged_caching. `dtype` is the q/out dtype +// (F16 or BF16); accumulation is always float. +// +// Split-kv (flash-decoding): when the waves heuristic (or an env override) +// picks num_splits > 1, the decode kernel runs with grid.z = num_splits and +// writes per-shard (m, l, acc) partials into the workspace, then a combine +// kernel merges them into `out`. The workspace layout matches the F16/BF16 +// family: partial_acc [kFp8DecodeMaxSplits, num_seqs, num_heads, head_size] +// followed by partial_m / partial_l [kFp8DecodeMaxSplits, num_seqs, num_heads]. + +namespace op::paged_attention::nvidia { + +namespace { + +constexpr size_t ceilDiv(size_t a, size_t b) { + return (a + b - 1) / b; +} + +inline int getSmCount() { + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { + return 0; + } + int sm_count = 0; + if (cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device) != cudaSuccess) { + return 0; + } + return sm_count; +} + +// Same FA2-style "waves" heuristic as the F16/BF16 split-kv launchers: shard +// the KV sequence so that base_blocks * num_splits CTAs fill the SMs without +// paying for more combine work than the split saves. seqlen_k is an upper +// bound (max pages * page size). +inline int chooseNumSplitsHeuristic(size_t num_heads, size_t num_seqs, size_t seqlen_k, int sm_count) { + if (sm_count <= 0) { + return 1; + } + if (num_heads == 0 || num_seqs == 0) { + return 1; + } + if (seqlen_k <= 256) { + return 1; + } + + const size_t base_blocks = num_heads * num_seqs; + int best_splits = 1; + // Baseline: one kernel, base_blocks CTAs, each scanning seqlen_k tokens. + size_t best_score = (ceilDiv(base_blocks, static_cast(sm_count)) * seqlen_k); + + size_t prev_work_per_block = seqlen_k; + for (int s = 2; s <= op::paged_attention::cuda::kFp8DecodeMaxSplits; ++s) { + const size_t blocks = base_blocks * static_cast(s); + const size_t waves_split = ceilDiv(blocks, static_cast(sm_count)); + const size_t work_per_block = ceilDiv(seqlen_k, static_cast(s)); + // If this split count doesn't reduce per-block work vs the previous split, it's effectively redundant. + if (work_per_block == prev_work_per_block) { + continue; + } + prev_work_per_block = work_per_block; + // Combine is one extra kernel with base_blocks blocks; approximate as one more wave unit. + const size_t waves_combine = ceilDiv(base_blocks, static_cast(sm_count)); + const size_t score = waves_split * work_per_block + waves_combine; + if (score < best_score) { + best_score = score; + best_splits = s; + } + } + return best_splits; +} + +template +INFINIOP_CUDA_KERNEL flashAttentionDecodeFp8( + Tdata *out, + float *partial_acc, + float *partial_m, + float *partial_l, + const Tdata *q, + const uint8_t *k_cache, + const uint8_t *v_cache, + const float *k_scale, + const float *v_scale, + const Tindex *block_tables, + const Tindex *cache_lens, + const float *alibi_slopes, + size_t num_kv_heads, + float scale, + size_t max_num_blocks_per_seq, + size_t page_block_size, + ptrdiff_t q_stride, + ptrdiff_t q_head_stride, + ptrdiff_t k_batch_stride, + ptrdiff_t k_row_stride, + ptrdiff_t k_head_stride, + ptrdiff_t v_batch_stride, + ptrdiff_t v_row_stride, + ptrdiff_t v_head_stride, + ptrdiff_t o_stride, + ptrdiff_t o_head_stride, + ptrdiff_t k_scale_block_stride, + ptrdiff_t k_scale_head_stride, + ptrdiff_t k_scale_slot_stride, + ptrdiff_t v_scale_block_stride, + ptrdiff_t v_scale_head_stride, + ptrdiff_t v_scale_slot_stride, + int num_splits) { + op::paged_attention::cuda::flashAttentionDecodeFp8Kernel( + out, partial_acc, partial_m, partial_l, + q, k_cache, v_cache, k_scale, v_scale, block_tables, cache_lens, alibi_slopes, + num_kv_heads, scale, max_num_blocks_per_seq, page_block_size, + q_stride, q_head_stride, + k_batch_stride, k_row_stride, k_head_stride, + v_batch_stride, v_row_stride, v_head_stride, + o_stride, o_head_stride, + k_scale_block_stride, k_scale_head_stride, k_scale_slot_stride, + v_scale_block_stride, v_scale_head_stride, v_scale_slot_stride, + num_splits); +} + +template +INFINIOP_CUDA_KERNEL flashAttentionDecodeFp8SplitKvCombine( + Tdata *out, + const float *partial_acc, + const float *partial_m, + const float *partial_l, + int num_splits, + ptrdiff_t o_stride, + ptrdiff_t o_head_stride) { + op::paged_attention::cuda::flashAttentionDecodeFp8SplitKvCombineKernel( + out, partial_acc, partial_m, partial_l, num_splits, o_stride, o_head_stride); +} + +// Split-kv policy for the FP8 decode kernel. Default (no env) is "auto": the +// waves heuristic splits only when it profits (small grids, long contexts). +// Env knobs mirror the F16/BF16 family: +// INFINIOP_FLASH_DECODE_SPLITKV = 0/false (never) | 1/true (force) | auto +// INFINIOP_FLASH_NUM_SPLITS = 1..8 (fixed, implies split) | auto +inline int chooseNumSplitsFp8(size_t num_heads, size_t num_seqs, + size_t max_num_blocks_per_seq, size_t page_block_size) { + const char *splitkv = std::getenv("INFINIOP_FLASH_DECODE_SPLITKV"); + if (splitkv && (std::strcmp(splitkv, "0") == 0 || std::strcmp(splitkv, "false") == 0)) { + return 1; + } + const bool forced = splitkv && (std::strcmp(splitkv, "1") == 0 || std::strcmp(splitkv, "true") == 0); + + int num_splits = 1; + const char *ns = std::getenv("INFINIOP_FLASH_NUM_SPLITS"); + const bool ns_fixed = ns && std::strcmp(ns, "auto") != 0 && std::atoi(ns) > 0; + if (ns_fixed) { + num_splits = std::atoi(ns); + } else if (forced) { + num_splits = 4; // fixed default, matching the F16/BF16 hd128 launcher + } else { + const size_t seqlen_k = max_num_blocks_per_seq * page_block_size; + num_splits = chooseNumSplitsHeuristic(num_heads, num_seqs, seqlen_k, getSmCount()); + } + if (num_splits < 1) { + num_splits = 1; + } + if (num_splits > op::paged_attention::cuda::kFp8DecodeMaxSplits) { + num_splits = op::paged_attention::cuda::kFp8DecodeMaxSplits; + } + + if (const char *dbg = std::getenv("INFINIOP_FLASH_DEBUG_SPLITS")) { + if (std::strcmp(dbg, "1") == 0 || std::strcmp(dbg, "true") == 0) { + static size_t last_seqs = ~static_cast(0); + static size_t last_heads = ~static_cast(0); + static size_t last_cap = ~static_cast(0); + static int last_splits = -1; + const size_t cap = max_num_blocks_per_seq * page_block_size; + if (num_seqs != last_seqs || num_heads != last_heads || cap != last_cap || num_splits != last_splits) { + last_seqs = num_seqs; + last_heads = num_heads; + last_cap = cap; + last_splits = num_splits; + std::fprintf(stderr, + "[INFINIOP][paged_attention][fp8] splitkv: heads=%zu seqs=%zu seqlen_k~%zu -> num_splits=%d\n", + num_heads, num_seqs, cap, num_splits); + } + } + } + return num_splits; +} + +template +infiniStatus_t launch_decode_fp8_impl( + void *workspace, + size_t workspace_size, + void *out, + const void *q, + const void *k_cache, + const void *v_cache, + const void *k_scale, + const void *v_scale, + infiniDtype_t dtype, + const Tindex *block_tables, + const Tindex *cache_lens, + const float *alibi_slopes, + size_t num_heads, + size_t num_seqs, + size_t num_kv_heads, + float scale, + size_t head_size, + size_t max_num_blocks_per_seq, + size_t page_block_size, + ptrdiff_t q_stride, + ptrdiff_t q_head_stride, + ptrdiff_t k_batch_stride, + ptrdiff_t k_row_stride, + ptrdiff_t k_head_stride, + ptrdiff_t v_batch_stride, + ptrdiff_t v_row_stride, + ptrdiff_t v_head_stride, + ptrdiff_t o_stride, + ptrdiff_t o_head_stride, + ptrdiff_t k_scale_block_stride, + ptrdiff_t k_scale_head_stride, + ptrdiff_t k_scale_slot_stride, + ptrdiff_t v_scale_block_stride, + ptrdiff_t v_scale_head_stride, + ptrdiff_t v_scale_slot_stride, + cudaStream_t stream) { + + if (num_heads == 0 || num_seqs == 0) { + return INFINI_STATUS_SUCCESS; + } + + const int num_splits = chooseNumSplitsFp8(num_heads, num_seqs, max_num_blocks_per_seq, page_block_size); + + float *partial_acc = nullptr; + float *partial_m = nullptr; + float *partial_l = nullptr; + if (num_splits > 1) { + const size_t n = num_seqs * num_heads; + const size_t acc_elems = static_cast(op::paged_attention::cuda::kFp8DecodeMaxSplits) * n * head_size; + const size_t ml_elems = static_cast(op::paged_attention::cuda::kFp8DecodeMaxSplits) * n; + const size_t needed_bytes = (acc_elems + 2 * ml_elems) * sizeof(float); + if (workspace == nullptr || workspace_size < needed_bytes) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + float *ws = static_cast(workspace); + partial_acc = ws; + partial_m = partial_acc + acc_elems; + partial_l = partial_m + ml_elems; + } + + // One CTA per (sequence, query head, split shard). The block size follows + // the kernel's warp count (kFp8DecodeNumWarps); each lane owns + // HEAD_SIZE/32 consecutive dims. num_splits == 1 takes the same path with + // grid.z == 1 and null partials, identical to the pre-split-kv kernel. + const dim3 grid(static_cast(num_heads), static_cast(num_seqs), + static_cast(num_splits)); + const dim3 grid_combine(static_cast(num_heads), static_cast(num_seqs), 1); + constexpr uint32_t kBlockThreads = op::paged_attention::cuda::kFp8DecodeNumWarps * 32; + +#define LAUNCH_FP8_DECODE(Tdata, HEAD_SIZE) \ + do { \ + flashAttentionDecodeFp8 \ + <<>>( \ + static_cast(out), \ + partial_acc, partial_m, partial_l, \ + static_cast(q), \ + static_cast(k_cache), \ + static_cast(v_cache), \ + static_cast(k_scale), \ + static_cast(v_scale), \ + block_tables, cache_lens, alibi_slopes, \ + num_kv_heads, scale, max_num_blocks_per_seq, page_block_size, \ + q_stride, q_head_stride, \ + k_batch_stride, k_row_stride, k_head_stride, \ + v_batch_stride, v_row_stride, v_head_stride, \ + o_stride, o_head_stride, \ + k_scale_block_stride, k_scale_head_stride, k_scale_slot_stride, \ + v_scale_block_stride, v_scale_head_stride, v_scale_slot_stride, \ + num_splits); \ + if (num_splits > 1) { \ + flashAttentionDecodeFp8SplitKvCombine \ + <<>>( \ + static_cast(out), \ + partial_acc, partial_m, partial_l, \ + num_splits, o_stride, o_head_stride); \ + } \ + return INFINI_STATUS_SUCCESS; \ + } while (false) + +#define DISPATCH_FP8_DECODE_HEAD_SIZE(HEAD_SIZE_) \ + do { \ + if (dtype == INFINI_DTYPE_F16) { \ + LAUNCH_FP8_DECODE(half, HEAD_SIZE_); \ + } \ + if (dtype == INFINI_DTYPE_BF16) { \ + LAUNCH_FP8_DECODE(__nv_bfloat16, HEAD_SIZE_); \ + } \ + return INFINI_STATUS_BAD_TENSOR_DTYPE; \ + } while (false) + + switch (head_size) { + case 64: + DISPATCH_FP8_DECODE_HEAD_SIZE(64); + case 128: + DISPATCH_FP8_DECODE_HEAD_SIZE(128); + default: + // FP8 decode implements head_size 64/128 only. + return INFINI_STATUS_NOT_IMPLEMENTED; + } + +#undef DISPATCH_FP8_DECODE_HEAD_SIZE +#undef LAUNCH_FP8_DECODE +} + +} // namespace + +#define DEFINE_LAUNCH_DECODE_FP8(SUFFIX, Tindex) \ + infiniStatus_t launch_decode_fp8_##SUFFIX( \ + void *workspace, size_t workspace_size, \ + void *out, const void *q, const void *k_cache, const void *v_cache, \ + const void *k_scale, const void *v_scale, \ + infiniDtype_t dtype, \ + const Tindex *block_tables, const Tindex *cache_lens, \ + const float *alibi_slopes, \ + size_t num_heads, size_t num_seqs, size_t num_kv_heads, \ + float scale, size_t head_size, \ + size_t max_num_blocks_per_seq, size_t page_block_size, \ + ptrdiff_t q_stride, ptrdiff_t q_head_stride, \ + ptrdiff_t k_batch_stride, ptrdiff_t k_row_stride, ptrdiff_t k_head_stride, \ + ptrdiff_t v_batch_stride, ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, \ + ptrdiff_t o_stride, ptrdiff_t o_head_stride, \ + ptrdiff_t k_scale_block_stride, ptrdiff_t k_scale_head_stride, \ + ptrdiff_t k_scale_slot_stride, \ + ptrdiff_t v_scale_block_stride, ptrdiff_t v_scale_head_stride, \ + ptrdiff_t v_scale_slot_stride, \ + cudaStream_t stream) { \ + return launch_decode_fp8_impl( \ + workspace, workspace_size, \ + out, q, k_cache, v_cache, k_scale, v_scale, dtype, \ + block_tables, cache_lens, alibi_slopes, \ + num_heads, num_seqs, num_kv_heads, scale, head_size, \ + max_num_blocks_per_seq, page_block_size, \ + q_stride, q_head_stride, \ + k_batch_stride, k_row_stride, k_head_stride, \ + v_batch_stride, v_row_stride, v_head_stride, \ + o_stride, o_head_stride, \ + k_scale_block_stride, k_scale_head_stride, k_scale_slot_stride, \ + v_scale_block_stride, v_scale_head_stride, v_scale_slot_stride, \ + stream); \ + } + +DEFINE_LAUNCH_DECODE_FP8(i64, int64_t) +DEFINE_LAUNCH_DECODE_FP8(i32, int32_t) +DEFINE_LAUNCH_DECODE_FP8(u32, uint32_t) + +#undef DEFINE_LAUNCH_DECODE_FP8 + +} // namespace op::paged_attention::nvidia diff --git a/src/infiniop/ops/paged_attention/nvidia/paged_attention_mla_hd576_v512.cu b/src/infiniop/ops/paged_attention/nvidia/paged_attention_mla_hd576_v512.cu index 3c449a168..3a3524690 100644 --- a/src/infiniop/ops/paged_attention/nvidia/paged_attention_mla_hd576_v512.cu +++ b/src/infiniop/ops/paged_attention/nvidia/paged_attention_mla_hd576_v512.cu @@ -49,6 +49,20 @@ INFINIOP_CUDA_KERNEL flashAttentionDecodeMlaHd576V512Warp( const int seq_len = static_cast(cache_lens_[seq_idx]); if (seq_len <= 0) { + // Zero-length sequence: write defined zeros instead of leaving the + // caller's output buffer untouched. + Tdata *out_ptr = out_ + seq_idx * o_stride + head_idx * o_head_stride; +#pragma unroll + for (int i = 0; i < kVDimsPerThread; ++i) { + const int dim = lane * kVDimsPerThread + i; + if constexpr (std::is_same_v) { + out_ptr[dim] = __float2half_rn(0.0f); + } else if constexpr (std::is_same_v) { + out_ptr[dim] = __float2bfloat16_rn(0.0f); + } else { + out_ptr[dim] = static_cast(0.0f); + } + } return; } @@ -244,7 +258,11 @@ INFINIOP_CUDA_KERNEL flashAttentionDecodeMlaHd576V512SplitKv( const int lane = threadIdx.x; const int seq_len = static_cast(cache_lens_[seq_idx]); - if (seq_len <= 0 || num_splits <= 0) { + // No early return for seq_len <= 0: shard becomes 0, so every split takes + // the empty-shard path below and publishes the neutral element (m=-inf, + // l=0, acc=0). Combine then emits a defined zero row instead of reading + // stale workspace. + if (num_splits <= 0) { return; } @@ -444,7 +462,7 @@ INFINIOP_CUDA_KERNEL flashAttentionDecodeMlaHd576V512SplitKvCombine( float acc = 0.0f; for (int s = 0; s < num_splits; ++s) { const float ms = partial_m[s * n + base]; - const float w = exp2f(ms - m); + const float w = (ms == -INFINITY) ? 0.0f : exp2f(ms - m); acc += partial_acc[(s * n + base) * kMlaValueSize + dim] * w; } const float o = acc * inv_l; diff --git a/src/infiniop/ops/paged_attention/nvidia/paged_attention_nvidia.cu b/src/infiniop/ops/paged_attention/nvidia/paged_attention_nvidia.cu index 45bf10524..9d64e18d2 100644 --- a/src/infiniop/ops/paged_attention/nvidia/paged_attention_nvidia.cu +++ b/src/infiniop/ops/paged_attention/nvidia/paged_attention_nvidia.cu @@ -9,6 +9,63 @@ namespace op::paged_attention::nvidia { +// FP8(E4M3) KV-cache decode launchers (paged_attention_fp8.cu). head_size +// 64/128; other head sizes return INFINI_STATUS_NOT_IMPLEMENTED. Split-kv +// partials go through the workspace (sized in Descriptor::create). +infiniStatus_t launch_decode_fp8_i64( + void *workspace, size_t workspace_size, + void *out, const void *q, const void *k_cache, const void *v_cache, + const void *k_scale, const void *v_scale, + infiniDtype_t dtype, + const int64_t *block_tables, const int64_t *cache_lens, + const float *alibi_slopes, + size_t num_heads, size_t num_seqs, size_t num_kv_heads, + float scale, size_t head_size, + size_t max_num_blocks_per_seq, size_t page_block_size, + ptrdiff_t q_stride, ptrdiff_t q_head_stride, + ptrdiff_t k_batch_stride, ptrdiff_t k_row_stride, ptrdiff_t k_head_stride, + ptrdiff_t v_batch_stride, ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, + ptrdiff_t o_stride, ptrdiff_t o_head_stride, + ptrdiff_t k_scale_block_stride, ptrdiff_t k_scale_head_stride, ptrdiff_t k_scale_slot_stride, + ptrdiff_t v_scale_block_stride, ptrdiff_t v_scale_head_stride, ptrdiff_t v_scale_slot_stride, + cudaStream_t stream); + +infiniStatus_t launch_decode_fp8_i32( + void *workspace, size_t workspace_size, + void *out, const void *q, const void *k_cache, const void *v_cache, + const void *k_scale, const void *v_scale, + infiniDtype_t dtype, + const int32_t *block_tables, const int32_t *cache_lens, + const float *alibi_slopes, + size_t num_heads, size_t num_seqs, size_t num_kv_heads, + float scale, size_t head_size, + size_t max_num_blocks_per_seq, size_t page_block_size, + ptrdiff_t q_stride, ptrdiff_t q_head_stride, + ptrdiff_t k_batch_stride, ptrdiff_t k_row_stride, ptrdiff_t k_head_stride, + ptrdiff_t v_batch_stride, ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, + ptrdiff_t o_stride, ptrdiff_t o_head_stride, + ptrdiff_t k_scale_block_stride, ptrdiff_t k_scale_head_stride, ptrdiff_t k_scale_slot_stride, + ptrdiff_t v_scale_block_stride, ptrdiff_t v_scale_head_stride, ptrdiff_t v_scale_slot_stride, + cudaStream_t stream); + +infiniStatus_t launch_decode_fp8_u32( + void *workspace, size_t workspace_size, + void *out, const void *q, const void *k_cache, const void *v_cache, + const void *k_scale, const void *v_scale, + infiniDtype_t dtype, + const uint32_t *block_tables, const uint32_t *cache_lens, + const float *alibi_slopes, + size_t num_heads, size_t num_seqs, size_t num_kv_heads, + float scale, size_t head_size, + size_t max_num_blocks_per_seq, size_t page_block_size, + ptrdiff_t q_stride, ptrdiff_t q_head_stride, + ptrdiff_t k_batch_stride, ptrdiff_t k_row_stride, ptrdiff_t k_head_stride, + ptrdiff_t v_batch_stride, ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, + ptrdiff_t o_stride, ptrdiff_t o_head_stride, + ptrdiff_t k_scale_block_stride, ptrdiff_t k_scale_head_stride, ptrdiff_t k_scale_slot_stride, + ptrdiff_t v_scale_block_stride, ptrdiff_t v_scale_head_stride, ptrdiff_t v_scale_slot_stride, + cudaStream_t stream); + infiniStatus_t launch_decode_hd64_i64( void *workspace, size_t workspace_size, void *out, const void *q, const void *k_cache, const void *v_cache, @@ -189,13 +246,17 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t block_tables_desc, infiniopTensorDescriptor_t cache_lens_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { - auto info_res = PagedAttentionInfo::create(out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, cache_lens_desc, alibi_slopes_desc, scale); + auto info_res = PagedAttentionInfo::create(out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, cache_lens_desc, alibi_slopes_desc, k_scale_desc, v_scale_desc, scale); CHECK_RESULT(info_res); auto info = info_res.take(); - // Reserve workspace for optional split-kv decode (partial acc + m/l). - // Workspace is independent of runtime env toggles; kernels will clamp num_splits <= kMaxSplits. + // Reserve workspace for optional split-kv decode (partial acc + m/l), + // shared by the F16/BF16 family and the FP8 decode kernel (v2). + // Workspace is independent of runtime env toggles; kernels clamp + // num_splits <= kMaxSplits, so the reservation covers every decision. constexpr size_t kMaxSplits = 8; const size_t per_split = info.num_seqs * info.num_heads * (info.value_size + 2) * sizeof(float); const size_t workspace_bytes = kMaxSplits * per_split; @@ -211,8 +272,46 @@ infiniStatus_t Descriptor::calculate( void *workspace, size_t workspace_size, void *out, const void *q, const void *k_cache, const void *v_cache, const void *block_tables, const void *cache_lens, const void *alibi_slopes, + const void *k_scale, const void *v_scale, void *stream_) const { + auto stream = static_cast(stream_); + + const float *alibi_ptr = (alibi_slopes == nullptr) ? nullptr : static_cast(alibi_slopes); + + // FP8(E4M3) KV caches dispatch to the dedicated clean-room decode kernel + // (split-kv partials, when enabled, live in the workspace). + if (_info.cache_dtype == INFINI_DTYPE_F8) { +#define CALCULATE_FP8(Tindex, SUFFIX) \ + return launch_decode_fp8_##SUFFIX( \ + workspace, workspace_size, \ + out, q, k_cache, v_cache, k_scale, v_scale, _info.dtype, \ + static_cast(block_tables), \ + static_cast(cache_lens), alibi_ptr, \ + _info.num_heads, _info.num_seqs, _info.num_kv_heads, \ + _info.scale, _info.head_size, \ + _info.max_num_blocks_per_seq, _info.page_block_size, \ + _info.q_stride, _info.q_head_stride, \ + _info.k_batch_stride, _info.k_row_stride, _info.k_head_stride, \ + _info.v_batch_stride, _info.v_row_stride, _info.v_head_stride, \ + _info.o_stride, _info.o_head_stride, \ + _info.k_scale_block_stride, _info.k_scale_head_stride, _info.k_scale_slot_stride, \ + _info.v_scale_block_stride, _info.v_scale_head_stride, _info.v_scale_slot_stride, \ + stream) + + if (_info.index_dtype == INFINI_DTYPE_I64) { + CALCULATE_FP8(int64_t, i64); + } + if (_info.index_dtype == INFINI_DTYPE_I32) { + CALCULATE_FP8(int32_t, i32); + } + if (_info.index_dtype == INFINI_DTYPE_U32) { + CALCULATE_FP8(uint32_t, u32); + } + return INFINI_STATUS_BAD_TENSOR_DTYPE; +#undef CALCULATE_FP8 + } + bool need_workspace = false; if (const char *env = std::getenv("INFINIOP_FLASH_DECODE_SPLITKV")) { // "auto" may enable split-kv depending on the runtime heuristic. @@ -225,10 +324,6 @@ infiniStatus_t Descriptor::calculate( return INFINI_STATUS_INSUFFICIENT_WORKSPACE; } - auto stream = static_cast(stream_); - - const float *alibi_ptr = (alibi_slopes == nullptr) ? nullptr : static_cast(alibi_slopes); - if (_info.index_dtype == INFINI_DTYPE_I64) { const auto *block_table_i64 = static_cast(block_tables); const auto *cache_lens_i64 = static_cast(cache_lens); diff --git a/src/infiniop/ops/paged_attention/operator.cc b/src/infiniop/ops/paged_attention/operator.cc index 3c3cce9f9..6be37d276 100644 --- a/src/infiniop/ops/paged_attention/operator.cc +++ b/src/infiniop/ops/paged_attention/operator.cc @@ -28,16 +28,21 @@ __INFINI_C 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) { infiniopTensorDescriptor_t alibi_opt = (alibi_slopes_desc == nullptr) ? nullptr : alibi_slopes_desc; + infiniopTensorDescriptor_t k_scale_opt = (k_scale_desc == nullptr) ? nullptr : k_scale_desc; + infiniopTensorDescriptor_t v_scale_opt = (v_scale_desc == nullptr) ? nullptr : v_scale_desc; #define CREATE(CASE, NAMESPACE) \ case CASE: \ return op::paged_attention::NAMESPACE::Descriptor::create( \ handle, \ reinterpret_cast(desc_ptr), \ - out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, seq_lens_desc, alibi_opt, scale); + out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, \ + seq_lens_desc, alibi_opt, k_scale_opt, v_scale_opt, scale); switch (handle->device) { #ifdef ENABLE_NVIDIA_API @@ -113,13 +118,14 @@ __INFINI_C infiniStatus_t infiniopPagedAttention( void *workspace, size_t workspace_size, void *out, const void *q, const void *k_cache, const void *v_cache, const void *block_tables, const void *seq_lens, const void *alibi_slopes, + const void *k_scale, const void *v_scale, void *stream) { #define CALCULATE(CASE, NAMESPACE) \ case CASE: \ return reinterpret_cast(desc)->calculate( \ workspace, workspace_size, out, q, k_cache, v_cache, block_tables, \ - seq_lens, alibi_slopes, stream); + seq_lens, alibi_slopes, k_scale, v_scale, stream); switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API diff --git a/src/infiniop/ops/paged_attention/paged_attention.h b/src/infiniop/ops/paged_attention/paged_attention.h index 662f39386..97b5b94d2 100644 --- a/src/infiniop/ops/paged_attention/paged_attention.h +++ b/src/infiniop/ops/paged_attention/paged_attention.h @@ -39,6 +39,8 @@ infiniopTensorDescriptor_t block_tables_desc, \ infiniopTensorDescriptor_t seq_lens_desc, \ const std::optional &alibi_slopes_desc, \ + const std::optional &k_scale_desc, \ + const std::optional &v_scale_desc, \ float scale); \ \ infiniStatus_t calculate( \ @@ -46,6 +48,7 @@ void *out, const void *q, const void *k_cache, const void *v_cache, \ const void *block_tables, const void *seq_lens, \ const void *alibi_slopes, \ + const void *k_scale, const void *v_scale, \ void *stream) const; \ }; \ } diff --git a/src/infiniop/ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend.cc b/src/infiniop/ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend.cc index 8cf9364d4..fffed642e 100644 --- a/src/infiniop/ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend.cc +++ b/src/infiniop/ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend.cc @@ -20,12 +20,17 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t seq_lens_desc, infiniopTensorDescriptor_t cum_seq_lens_q_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { auto info = PagedAttentionPrefillInfo::create( out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, seq_lens_desc, cum_seq_lens_q_desc, - alibi_slopes_desc, scale); + alibi_slopes_desc, k_scale_desc, v_scale_desc, scale); CHECK_RESULT(info); + if (info->cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } auto handle_ascend = reinterpret_cast(handle); *desc_ptr = new Descriptor( @@ -45,6 +50,7 @@ infiniStatus_t Descriptor::calculate( const void *seq_lens, const void *cum_seq_lens_q, const void *alibi_slopes, + const void *k_scale, const void *v_scale, void *stream) const { (void)workspace; (void)workspace_size; diff --git a/src/infiniop/ops/paged_attention_prefill/bang/paged_attention_prefill_bang.mlu b/src/infiniop/ops/paged_attention_prefill/bang/paged_attention_prefill_bang.mlu index dfdb781cd..5a76f0b80 100644 --- a/src/infiniop/ops/paged_attention_prefill/bang/paged_attention_prefill_bang.mlu +++ b/src/infiniop/ops/paged_attention_prefill/bang/paged_attention_prefill_bang.mlu @@ -328,14 +328,19 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t total_kv_lens_desc, infiniopTensorDescriptor_t cum_seqlens_q_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { auto handle = reinterpret_cast(handle_); auto info = PagedAttentionPrefillInfo::create( out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, total_kv_lens_desc, cum_seqlens_q_desc, - alibi_slopes_desc, scale); + alibi_slopes_desc, k_scale_desc, v_scale_desc, scale); CHECK_RESULT(info); + if (info->cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } *desc_ptr = new Descriptor( new Opaque{static_cast(handle)->internal()}, @@ -355,6 +360,8 @@ infiniStatus_t Descriptor::calculate( const void *total_kv_lens, const void *cum_seq_lens_q, const void *alibi_slopes, + const void *k_scale, + const void *v_scale, void *stream) const { (void)workspace; diff --git a/src/infiniop/ops/paged_attention_prefill/cuda/kernel_fp8.cuh b/src/infiniop/ops/paged_attention_prefill/cuda/kernel_fp8.cuh new file mode 100644 index 000000000..52b04788f --- /dev/null +++ b/src/infiniop/ops/paged_attention_prefill/cuda/kernel_fp8.cuh @@ -0,0 +1,122 @@ +#ifndef __PAGED_ATTENTION_PREFILL_FP8_KERNEL_CUH__ +#define __PAGED_ATTENTION_PREFILL_FP8_KERNEL_CUH__ + +//================================================================================ +// FP8(E4M3) KV-cache gather-dequant kernels for paged prefill (plan B). +// +// The existing F16/BF16 prefill kernels stay untouched. For F8 caches we +// first gather the blocks referenced by each sequence's block table and +// dequantize them (e4m3_decode(code) * scale[block, kv_head, slot]) into a +// compact BF16/F16 scratch cache laid out as +// [num_seqs * max_num_blocks_per_seq, num_kv_heads, page_block_size, head_size] +// together with identity block tables (scratch block of (seq, page) is +// seq * max_num_blocks_per_seq + page), then run the regular prefill kernel +// on the scratch. Scratch memory comes from the operator workspace. +//================================================================================ + +#include + +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" + +namespace op::paged_attention_prefill::cuda { + +namespace fp8 { + +template +__device__ __forceinline__ Tdata fromFloat(float value) { + if constexpr (std::is_same_v) { + return __float2half_rn(value); + } else if constexpr (std::is_same_v) { + return __float2bfloat16_rn(value); + } else { + return static_cast(value); + } +} + +} // namespace fp8 + +// block_tables_scratch[i] = i for i in [0, total) +template +__device__ void fillIdentityBlockTablesKernel( + Tindex *block_tables_scratch, + const size_t total) { + const size_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < total) { + block_tables_scratch[i] = static_cast(i); + } +} + +// One CTA per (logical page, sequence, kv head): dequantize the referenced +// page of K and V into the compact scratch cache. Pages past the sequence's +// kv length are skipped. +template +__device__ void gatherDequantFp8KvKernel( + Tdata *k_scratch_, + Tdata *v_scratch_, + const uint8_t *k_cache_, + const uint8_t *v_cache_, + const float *k_scale_, + const float *v_scale_, + const Tindex *block_tables_, + const Tindex *total_kv_lens_, + const size_t num_kv_heads, + const size_t head_size, + const size_t value_size, + const size_t page_block_size, + const size_t max_num_blocks_per_seq, + const ptrdiff_t block_table_batch_stride, + const ptrdiff_t k_batch_stride, + const ptrdiff_t k_row_stride, + const ptrdiff_t k_head_stride, + const ptrdiff_t v_batch_stride, + const ptrdiff_t v_row_stride, + const ptrdiff_t v_head_stride, + const ptrdiff_t k_scale_block_stride, + const ptrdiff_t k_scale_head_stride, + const ptrdiff_t k_scale_slot_stride, + const ptrdiff_t v_scale_block_stride, + const ptrdiff_t v_scale_head_stride, + const ptrdiff_t v_scale_slot_stride) { + + const size_t page = blockIdx.x; + const size_t seq = blockIdx.y; + const size_t head = blockIdx.z; + + const size_t kv_len = static_cast(total_kv_lens_[seq]); + if (page * page_block_size >= kv_len) { + return; // page not referenced by this sequence + } + + const ptrdiff_t phys = static_cast(block_tables_[seq * block_table_batch_stride + page]); + const size_t scratch_block = seq * max_num_blocks_per_seq + page; + + // K page + { + const uint8_t *k_src = k_cache_ + phys * k_batch_stride + static_cast(head) * k_head_stride; + const float *k_scl = k_scale_ + phys * k_scale_block_stride + static_cast(head) * k_scale_head_stride; + Tdata *k_dst = k_scratch_ + (scratch_block * num_kv_heads + head) * (page_block_size * head_size); + for (size_t idx = threadIdx.x; idx < page_block_size * head_size; idx += blockDim.x) { + const size_t slot = idx / head_size; + const size_t d = idx - slot * head_size; + k_dst[idx] = fp8::fromFloat( + infiniopFp8E4m3Decode(k_src[slot * k_row_stride + d]) * k_scl[slot * k_scale_slot_stride]); + } + } + + // V page (value_size may differ from head_size for MLA) + { + const uint8_t *v_src = v_cache_ + phys * v_batch_stride + static_cast(head) * v_head_stride; + const float *v_scl = v_scale_ + phys * v_scale_block_stride + static_cast(head) * v_scale_head_stride; + Tdata *v_dst = v_scratch_ + (scratch_block * num_kv_heads + head) * (page_block_size * value_size); + for (size_t idx = threadIdx.x; idx < page_block_size * value_size; idx += blockDim.x) { + const size_t slot = idx / value_size; + const size_t d = idx - slot * value_size; + v_dst[idx] = fp8::fromFloat( + infiniopFp8E4m3Decode(v_src[slot * v_row_stride + d]) * v_scl[slot * v_scale_slot_stride]); + } + } +} + +} // namespace op::paged_attention_prefill::cuda + +#endif // __PAGED_ATTENTION_PREFILL_FP8_KERNEL_CUH__ diff --git a/src/infiniop/ops/paged_attention_prefill/info.h b/src/infiniop/ops/paged_attention_prefill/info.h index 39926283e..a8c544aef 100644 --- a/src/infiniop/ops/paged_attention_prefill/info.h +++ b/src/infiniop/ops/paged_attention_prefill/info.h @@ -15,6 +15,7 @@ class PagedAttentionPrefillInfo { public: infiniDtype_t dtype; + infiniDtype_t cache_dtype; infiniDtype_t index_dtype; float scale; @@ -41,6 +42,15 @@ class PagedAttentionPrefillInfo { ptrdiff_t block_table_batch_stride; + // --- FP8(E4M3) KV cache: per-token dequant scales [num_blocks, num_kv_heads, block_size] --- + // Only meaningful when cache_dtype == INFINI_DTYPE_F8. + ptrdiff_t k_scale_block_stride; + ptrdiff_t k_scale_head_stride; + ptrdiff_t k_scale_slot_stride; + ptrdiff_t v_scale_block_stride; + ptrdiff_t v_scale_head_stride; + ptrdiff_t v_scale_slot_stride; + static utils::Result create( infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t q_desc, @@ -50,14 +60,38 @@ class PagedAttentionPrefillInfo { infiniopTensorDescriptor_t total_kv_lens_desc, infiniopTensorDescriptor_t cum_seqlens_q_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { auto dtype = q_desc->dtype(); CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); - if (out_desc->dtype() != dtype || k_cache_desc->dtype() != dtype || v_cache_desc->dtype() != dtype) { + if (out_desc->dtype() != dtype) { return INFINI_STATUS_BAD_TENSOR_DTYPE; } + // The caches either keep the compute dtype or store FP8(E4M3) codes + // plus per-token F32 scales produced by paged_caching. + auto cache_dtype = k_cache_desc->dtype(); + const bool cache_fp8 = (cache_dtype == INFINI_DTYPE_F8); + const bool has_k_scale = k_scale_desc.has_value() && k_scale_desc.value() != nullptr; + const bool has_v_scale = v_scale_desc.has_value() && v_scale_desc.value() != nullptr; + if (cache_fp8) { + if (v_cache_desc->dtype() != INFINI_DTYPE_F8) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (!has_k_scale || !has_v_scale) { + return INFINI_STATUS_BAD_PARAM; + } + } else { + if (cache_dtype != dtype || v_cache_desc->dtype() != dtype) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (has_k_scale || has_v_scale) { + return INFINI_STATUS_BAD_PARAM; + } + } + // q/out: [total_q, heads, head_dim] if (q_desc->ndim() != 3 || out_desc->ndim() != 3) { return INFINI_STATUS_BAD_TENSOR_SHAPE; @@ -139,6 +173,22 @@ class PagedAttentionPrefillInfo { return INFINI_STATUS_BAD_TENSOR_SHAPE; } + // Per-token dequant scales for the FP8 path: [num_blocks, num_kv_heads, page_block_size]. + if (cache_fp8) { + for (const auto &scale_desc : {k_scale_desc.value(), v_scale_desc.value()}) { + if (scale_desc->dtype() != INFINI_DTYPE_F32) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (scale_desc->ndim() != 3) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + const auto scale_shape = scale_desc->shape(); + if (scale_shape[0] != num_blocks || scale_shape[1] != num_kv_heads || scale_shape[2] != page_block_size) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + } + } + const size_t num_seqs = total_kv_lens_desc->shape()[0]; if (cum_seqlens_q_desc->shape()[0] != num_seqs + 1) { return INFINI_STATUS_BAD_PARAM; @@ -162,6 +212,13 @@ class PagedAttentionPrefillInfo { const ptrdiff_t block_table_batch_stride = block_tables_desc->stride(0); + const ptrdiff_t k_scale_block_stride = cache_fp8 ? k_scale_desc.value()->stride(0) : 0; + const ptrdiff_t k_scale_head_stride = cache_fp8 ? k_scale_desc.value()->stride(1) : 0; + const ptrdiff_t k_scale_slot_stride = cache_fp8 ? k_scale_desc.value()->stride(2) : 0; + const ptrdiff_t v_scale_block_stride = cache_fp8 ? v_scale_desc.value()->stride(0) : 0; + const ptrdiff_t v_scale_head_stride = cache_fp8 ? v_scale_desc.value()->stride(1) : 0; + const ptrdiff_t v_scale_slot_stride = cache_fp8 ? v_scale_desc.value()->stride(2) : 0; + if (const char *dbg = std::getenv("INFINIOP_DEBUG_PREFILL_INFO")) { static bool printed = false; if (!printed && std::strcmp(dbg, "1") == 0) { @@ -182,6 +239,7 @@ class PagedAttentionPrefillInfo { return utils::Result(PagedAttentionPrefillInfo{ dtype, + cache_dtype, block_tables_dt, scale, num_seqs, @@ -204,6 +262,12 @@ class PagedAttentionPrefillInfo { o_stride, o_head_stride, block_table_batch_stride, + k_scale_block_stride, + k_scale_head_stride, + k_scale_slot_stride, + v_scale_block_stride, + v_scale_head_stride, + v_scale_slot_stride, }); } }; diff --git a/src/infiniop/ops/paged_attention_prefill/metax/paged_attention_prefill_metax.maca b/src/infiniop/ops/paged_attention_prefill/metax/paged_attention_prefill_metax.maca index d26eca008..4e2f3733c 100644 --- a/src/infiniop/ops/paged_attention_prefill/metax/paged_attention_prefill_metax.maca +++ b/src/infiniop/ops/paged_attention_prefill/metax/paged_attention_prefill_metax.maca @@ -1456,13 +1456,18 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t total_kv_lens_desc, infiniopTensorDescriptor_t cum_seqlens_q_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { auto info = PagedAttentionPrefillInfo::create( out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, total_kv_lens_desc, cum_seqlens_q_desc, - alibi_slopes_desc, scale); + alibi_slopes_desc, k_scale_desc, v_scale_desc, scale); CHECK_RESULT(info); + if (info->cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } // Optional split-kv prefill requires workspace for partial (m, l, acc). // IMPORTANT: Unlike decode, prefill's total_q_tokens can be very large, so we must NOT reserve @@ -1506,6 +1511,7 @@ infiniStatus_t Descriptor::calculate( const void *total_kv_lens, const void *cum_seqlens_q, const void *alibi_slopes, + const void *k_scale, const void *v_scale, void *stream_) const { auto stream = static_cast(stream_); diff --git a/src/infiniop/ops/paged_attention_prefill/moore/paged_attention_prefill_moore.mu b/src/infiniop/ops/paged_attention_prefill/moore/paged_attention_prefill_moore.mu index 8c5afd0c5..6aba4eb62 100644 --- a/src/infiniop/ops/paged_attention_prefill/moore/paged_attention_prefill_moore.mu +++ b/src/infiniop/ops/paged_attention_prefill/moore/paged_attention_prefill_moore.mu @@ -199,15 +199,20 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t seq_lens_desc, infiniopTensorDescriptor_t cum_seq_lens_q_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { auto info = PagedAttentionPrefillInfo::create( out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, seq_lens_desc, cum_seq_lens_q_desc, - alibi_slopes_desc, scale); + alibi_slopes_desc, k_scale_desc, v_scale_desc, scale); CHECK_RESULT(info); + if (info->cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } *desc_ptr = new Descriptor( new Opaque{reinterpret_cast(handle)->internal()}, @@ -223,6 +228,7 @@ infiniStatus_t Descriptor::calculate( const void *seq_lens, const void *cum_seq_lens_q, const void *alibi_slopes, + const void *k_scale, const void *v_scale, void *stream_) const { musaStream_t stream = (musaStream_t)stream_; diff --git a/src/infiniop/ops/paged_attention_prefill/nvidia/paged_attention_prefill_nvidia.cu b/src/infiniop/ops/paged_attention_prefill/nvidia/paged_attention_prefill_nvidia.cu index c3183525b..8f9e6c5e1 100644 --- a/src/infiniop/ops/paged_attention_prefill/nvidia/paged_attention_prefill_nvidia.cu +++ b/src/infiniop/ops/paged_attention_prefill/nvidia/paged_attention_prefill_nvidia.cu @@ -13,6 +13,7 @@ // #include "paged_attention_prefill_fa2.cuh" #include "paged_attention_prefill_nvidia.cuh" +#include "../cuda/kernel_fp8.cuh" #include "../cuda/kernel_v2.cuh" namespace op::paged_attention_prefill::nvidia { @@ -1567,6 +1568,175 @@ infiniStatus_t launch_prefill_warpcta16( return INFINI_STATUS_BAD_TENSOR_SHAPE; } } + +// ============================================================================ +// FP8(E4M3) KV-cache prefill (plan B): gather-dequant the referenced pages +// into a compact BF16/F16 scratch cache with identity block tables, then run +// the regular prefill kernel on the scratch. The F16/BF16 kernels and their +// dispatch above are not modified. +// ============================================================================ + +template +INFINIOP_CUDA_KERNEL FillIdentityBlockTables( + Tindex *block_tables_scratch, size_t total) { + op::paged_attention_prefill::cuda::fillIdentityBlockTablesKernel( + block_tables_scratch, total); +} + +template +INFINIOP_CUDA_KERNEL GatherDequantFp8Kv( + Tdata *k_scratch, Tdata *v_scratch, + const uint8_t *k_cache, const uint8_t *v_cache, + const float *k_scale, const float *v_scale, + const Tindex *block_tables, const Tindex *total_kv_lens, + size_t num_kv_heads, size_t head_size, size_t value_size, + size_t page_block_size, size_t max_num_blocks_per_seq, + ptrdiff_t block_table_batch_stride, + ptrdiff_t k_batch_stride, ptrdiff_t k_row_stride, ptrdiff_t k_head_stride, + ptrdiff_t v_batch_stride, ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, + ptrdiff_t k_scale_block_stride, ptrdiff_t k_scale_head_stride, ptrdiff_t k_scale_slot_stride, + ptrdiff_t v_scale_block_stride, ptrdiff_t v_scale_head_stride, ptrdiff_t v_scale_slot_stride) { + op::paged_attention_prefill::cuda::gatherDequantFp8KvKernel( + k_scratch, v_scratch, k_cache, v_cache, k_scale, v_scale, + block_tables, total_kv_lens, + num_kv_heads, head_size, value_size, page_block_size, max_num_blocks_per_seq, + block_table_batch_stride, + k_batch_stride, k_row_stride, k_head_stride, + v_batch_stride, v_row_stride, v_head_stride, + k_scale_block_stride, k_scale_head_stride, k_scale_slot_stride, + v_scale_block_stride, v_scale_head_stride, v_scale_slot_stride); +} + +constexpr size_t alignUp256(size_t x) { + return (x + 255) / 256 * 256; +} + +struct Fp8ScratchLayout { + size_t v_offset; + size_t bt_offset; + size_t total; +}; + +// Workspace layout: [K scratch | V scratch | identity block tables], each +// segment 256-byte aligned. Tdata is F16/BF16 (2 bytes), guaranteed by info. +inline Fp8ScratchLayout fp8ScratchLayout(const PagedAttentionPrefillInfo &info) { + const size_t tindex_size = (info.index_dtype == INFINI_DTYPE_I64) ? sizeof(int64_t) : sizeof(int32_t); + const size_t scratch_blocks = info.num_seqs * info.max_num_blocks_per_seq; + const size_t k_bytes = scratch_blocks * info.num_kv_heads * info.page_block_size * info.head_size * sizeof(uint16_t); + const size_t v_bytes = scratch_blocks * info.num_kv_heads * info.page_block_size * info.value_size * sizeof(uint16_t); + const size_t bt_bytes = scratch_blocks * tindex_size; + Fp8ScratchLayout layout; + layout.v_offset = alignUp256(k_bytes); + layout.bt_offset = layout.v_offset + alignUp256(v_bytes); + layout.total = layout.bt_offset + alignUp256(bt_bytes); + return layout; +} + +template +infiniStatus_t launch_prefill_fp8( + void *workspace, size_t workspace_size, + Tdata *out, const Tdata *q, + const void *k_cache, const void *v_cache, + const void *k_scale, const void *v_scale, + const Tindex *block_tables, const Tindex *total_kv_lens, const Tindex *cu_seqlens_q, + const float *alibi_slopes, + const PagedAttentionPrefillInfo &info, + cudaStream_t stream) { + + const Fp8ScratchLayout layout = fp8ScratchLayout(info); + if (workspace == nullptr || workspace_size < layout.total) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + const size_t scratch_blocks = info.num_seqs * info.max_num_blocks_per_seq; + if (scratch_blocks == 0) { + return INFINI_STATUS_SUCCESS; + } + + Tdata *k_scratch = static_cast(workspace); + Tdata *v_scratch = reinterpret_cast(static_cast(workspace) + layout.v_offset); + Tindex *bt_scratch = reinterpret_cast(static_cast(workspace) + layout.bt_offset); + + // Identity block tables over the compact scratch: (seq, page) -> seq * mbps + page. + { + constexpr int threads = 256; + const size_t blocks = ceilDiv(scratch_blocks, static_cast(threads)); + FillIdentityBlockTables + <<(blocks), threads, 0, stream>>>(bt_scratch, scratch_blocks); + } + + // Gather + dequantize the referenced pages. + { + const dim3 grid(static_cast(info.max_num_blocks_per_seq), + static_cast(info.num_seqs), + static_cast(info.num_kv_heads)); + const dim3 block(256); + GatherDequantFp8Kv<<>>( + k_scratch, v_scratch, + static_cast(k_cache), static_cast(v_cache), + static_cast(k_scale), static_cast(v_scale), + block_tables, total_kv_lens, + info.num_kv_heads, info.head_size, info.value_size, + info.page_block_size, info.max_num_blocks_per_seq, + info.block_table_batch_stride, + info.k_batch_stride, info.k_row_stride, info.k_head_stride, + info.v_batch_stride, info.v_row_stride, info.v_head_stride, + info.k_scale_block_stride, info.k_scale_head_stride, info.k_scale_slot_stride, + info.v_scale_block_stride, info.v_scale_head_stride, info.v_scale_slot_stride); + } + + // Contiguous scratch strides. + const ptrdiff_t s_k_batch = static_cast(info.num_kv_heads * info.page_block_size * info.head_size); + const ptrdiff_t s_k_head = static_cast(info.page_block_size * info.head_size); + const ptrdiff_t s_k_row = static_cast(info.head_size); + const ptrdiff_t s_v_batch = static_cast(info.num_kv_heads * info.page_block_size * info.value_size); + const ptrdiff_t s_v_head = static_cast(info.page_block_size * info.value_size); + const ptrdiff_t s_v_row = static_cast(info.value_size); + const ptrdiff_t s_bt_batch = static_cast(info.max_num_blocks_per_seq); + +#define LAUNCH_PREFILL_ON_SCRATCH(LAUNCHER) \ + return LAUNCHER( \ + out, q, k_scratch, v_scratch, bt_scratch, total_kv_lens, cu_seqlens_q, alibi_slopes, \ + info.num_heads, info.num_seqs, info.num_kv_heads, info.total_q_tokens, \ + info.head_size, info.scale, info.max_num_blocks_per_seq, info.page_block_size, \ + s_bt_batch, \ + info.q_stride, info.q_head_stride, \ + s_k_batch, s_k_row, s_k_head, \ + s_v_batch, s_v_row, s_v_head, \ + info.o_stride, info.o_head_stride, stream) + + // Follow the same default kernel selection as the regular path. The + // split-kv/mma variants are not supported on the F8 path; fall back to + // the default tile kernel for head_size 64/128 and to "ref" otherwise. + const char *k = default_prefill_kernel(info); + if (std::strcmp(k, "warp") == 0) { + LAUNCH_PREFILL_ON_SCRATCH(launch_prefill_warp); + } + if (std::strcmp(k, "warpcta") == 0) { + LAUNCH_PREFILL_ON_SCRATCH(launch_prefill); + } + if (std::strcmp(k, "warpcta8pipe") == 0) { + LAUNCH_PREFILL_ON_SCRATCH(launch_prefill_warpcta8pipe); + } + if (std::strcmp(k, "warpcta16") == 0) { + LAUNCH_PREFILL_ON_SCRATCH(launch_prefill_warpcta16); + } + if (std::strcmp(k, "ref") == 0) { + return launch_prefill_ref( + out, q, k_scratch, v_scratch, bt_scratch, total_kv_lens, cu_seqlens_q, alibi_slopes, + info.num_heads, info.num_seqs, info.num_kv_heads, info.total_q_tokens, + info.head_size, info.value_size, info.scale, info.max_num_blocks_per_seq, info.page_block_size, + s_bt_batch, + info.q_stride, info.q_head_stride, + s_k_batch, s_k_row, s_k_head, + s_v_batch, s_v_row, s_v_head, + info.o_stride, info.o_head_stride, stream); + } + // "warpcta8" (and any remaining default) — head_size 64/128. + LAUNCH_PREFILL_ON_SCRATCH(launch_prefill_warpcta8); + +#undef LAUNCH_PREFILL_ON_SCRATCH +} } // namespace struct Descriptor::Opaque { @@ -1588,12 +1758,14 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t total_kv_lens_desc, infiniopTensorDescriptor_t cum_seqlens_q_desc, const std::optional &alibi_slopes_desc, + const std::optional &k_scale_desc, + const std::optional &v_scale_desc, float scale) { auto info = PagedAttentionPrefillInfo::create( out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, total_kv_lens_desc, cum_seqlens_q_desc, - alibi_slopes_desc, scale); + alibi_slopes_desc, k_scale_desc, v_scale_desc, scale); CHECK_RESULT(info); // Optional split-kv prefill requires workspace for partial (m, l, acc). @@ -1621,8 +1793,12 @@ infiniStatus_t Descriptor::create( const size_t n = info->total_q_tokens * info->num_heads; const size_t splitkv_workspace_bytes = use_splitkv ? (static_cast(num_splits) * n * (info->head_size + 2) * sizeof(float)) : 0; - const size_t workspace_bytes = splitkv_workspace_bytes; - // const size_t workspace_bytes = splitkv_workspace_bytes + fa2_workspace_bytes; + // FP8 caches (plan B) need scratch for the gather-dequantized K/V plus the + // compact identity block tables. Split-kv prefill is not supported on the + // F8 path, so the two never coexist. + const size_t fp8_workspace_bytes = (info->cache_dtype == INFINI_DTYPE_F8) ? fp8ScratchLayout(*info).total : 0; + + const size_t workspace_bytes = splitkv_workspace_bytes + fp8_workspace_bytes; *desc_ptr = new Descriptor( new Opaque{reinterpret_cast(handle)->internal()}, @@ -1638,6 +1814,7 @@ infiniStatus_t Descriptor::calculate( const void *total_kv_lens, const void *cum_seqlens_q, const void *alibi_slopes, + const void *k_scale, const void *v_scale, void *stream_) const { auto stream = static_cast(stream_); @@ -1645,6 +1822,45 @@ infiniStatus_t Descriptor::calculate( const void *total_kv_lens_ptr = total_kv_lens; const void *cu_seqlens_q_ptr = cum_seqlens_q; + // FP8(E4M3) KV caches: gather-dequant into scratch, then run the regular + // F16/BF16 prefill kernel on the scratch (plan B; split-kv unsupported). + if (_info.cache_dtype == INFINI_DTYPE_F8) { +#define CALCULATE_FP8_PREFILL(Tindex, Tdata) \ + return launch_prefill_fp8( \ + workspace, workspace_size, \ + static_cast(out), static_cast(q), \ + k_cache, v_cache, k_scale, v_scale, \ + static_cast(block_tables), \ + static_cast(total_kv_lens_ptr), \ + static_cast(cu_seqlens_q_ptr), \ + alibi_ptr, _info, stream) + + if (_info.index_dtype == INFINI_DTYPE_I64) { + if (_info.dtype == INFINI_DTYPE_F16) { + CALCULATE_FP8_PREFILL(int64_t, half); + } + if (_info.dtype == INFINI_DTYPE_BF16) { + CALCULATE_FP8_PREFILL(int64_t, __nv_bfloat16); + } + } else if (_info.index_dtype == INFINI_DTYPE_I32) { + if (_info.dtype == INFINI_DTYPE_F16) { + CALCULATE_FP8_PREFILL(int32_t, half); + } + if (_info.dtype == INFINI_DTYPE_BF16) { + CALCULATE_FP8_PREFILL(int32_t, __nv_bfloat16); + } + } else if (_info.index_dtype == INFINI_DTYPE_U32) { + if (_info.dtype == INFINI_DTYPE_F16) { + CALCULATE_FP8_PREFILL(uint32_t, half); + } + if (_info.dtype == INFINI_DTYPE_BF16) { + CALCULATE_FP8_PREFILL(uint32_t, __nv_bfloat16); + } + } + return INFINI_STATUS_BAD_TENSOR_DTYPE; +#undef CALCULATE_FP8_PREFILL + } + bool use_splitkv = false; if (const char *env = std::getenv("INFINIOP_FLASH_PREFILL_SPLITKV")) { use_splitkv = (std::strcmp(env, "1") == 0) || (std::strcmp(env, "true") == 0); diff --git a/src/infiniop/ops/paged_attention_prefill/operator.cc b/src/infiniop/ops/paged_attention_prefill/operator.cc index bd06d36ca..8fb3ca7ad 100644 --- a/src/infiniop/ops/paged_attention_prefill/operator.cc +++ b/src/infiniop/ops/paged_attention_prefill/operator.cc @@ -29,9 +29,13 @@ __INFINI_C 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) { infiniopTensorDescriptor_t alibi_opt = (alibi_slopes_desc == nullptr) ? nullptr : alibi_slopes_desc; + infiniopTensorDescriptor_t k_scale_opt = (k_scale_desc == nullptr) ? nullptr : k_scale_desc; + infiniopTensorDescriptor_t v_scale_opt = (v_scale_desc == nullptr) ? nullptr : v_scale_desc; #define CREATE(CASE, NAMESPACE) \ case CASE: \ @@ -39,7 +43,7 @@ __INFINI_C infiniStatus_t infiniopCreatePagedAttentionPrefillDescriptor( handle, \ reinterpret_cast(desc_ptr), \ out_desc, q_desc, k_cache_desc, v_cache_desc, block_tables_desc, \ - seq_lens_desc, cum_seq_lens_q_desc, alibi_opt, scale); + seq_lens_desc, cum_seq_lens_q_desc, alibi_opt, k_scale_opt, v_scale_opt, scale); switch (handle->device) { #ifdef ENABLE_NVIDIA_API @@ -118,13 +122,14 @@ __INFINI_C 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) { #define CALCULATE(CASE, NAMESPACE) \ case CASE: \ return reinterpret_cast(desc)->calculate( \ workspace, workspace_size, out, q, k_cache, v_cache, block_tables, \ - seq_lens, cum_seq_lens_q, alibi_slopes, stream); + seq_lens, cum_seq_lens_q, alibi_slopes, k_scale, v_scale, stream); switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API diff --git a/src/infiniop/ops/paged_attention_prefill/paged_attention_prefill.h b/src/infiniop/ops/paged_attention_prefill/paged_attention_prefill.h index 4c9222205..8098f4493 100644 --- a/src/infiniop/ops/paged_attention_prefill/paged_attention_prefill.h +++ b/src/infiniop/ops/paged_attention_prefill/paged_attention_prefill.h @@ -40,6 +40,8 @@ infiniopTensorDescriptor_t seq_lens_desc, \ infiniopTensorDescriptor_t cum_seq_lens_q_desc, \ const std::optional &alibi_slopes_desc, \ + const std::optional &k_scale_desc, \ + const std::optional &v_scale_desc, \ float scale); \ \ infiniStatus_t calculate( \ @@ -49,6 +51,7 @@ const void *seq_lens, \ const void *cum_seq_lens_q, \ const void *alibi_slopes, \ + const void *k_scale, const void *v_scale, \ void *stream) const; \ }; \ } diff --git a/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.cc b/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.cc index ff4e59faa..be8cad65a 100644 --- a/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.cc +++ b/src/infiniop/ops/paged_caching/ascend/paged_caching_ascend.cc @@ -16,10 +16,15 @@ infiniStatus_t Descriptor::create( 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) { - auto info = PagedCachingInfo::create(k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc); + auto info = PagedCachingInfo::create(k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc, k_scale_desc, v_scale_desc); CHECK_RESULT(info); + if (info->cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } auto handle_ascend = reinterpret_cast(handle); *desc_ptr = new Descriptor( @@ -37,6 +42,7 @@ infiniStatus_t Descriptor::calculate( void *k_cache, void *v_cache, const void *k, const void *v, const void *slot_mapping, + void *k_scale, void *v_scale, void *stream) const { (void)workspace; (void)workspace_size; diff --git a/src/infiniop/ops/paged_caching/bang/paged_caching_bang.mlu b/src/infiniop/ops/paged_caching/bang/paged_caching_bang.mlu index 92fc38dd8..c535e3053 100644 --- a/src/infiniop/ops/paged_caching/bang/paged_caching_bang.mlu +++ b/src/infiniop/ops/paged_caching/bang/paged_caching_bang.mlu @@ -123,11 +123,16 @@ infiniStatus_t Descriptor::create( 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) { auto handle = reinterpret_cast(handle_); - auto info = PagedCachingInfo::create(k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc); + auto info = PagedCachingInfo::create(k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc, k_scale_desc, v_scale_desc); CHECK_RESULT(info); + if (info->cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } *desc_ptr = new Descriptor( new Opaque{static_cast(handle)->internal()}, @@ -144,6 +149,8 @@ infiniStatus_t Descriptor::calculate( const void *k, const void *v, const void *slot_mapping, + void *k_scale, + void *v_scale, void *stream) const { (void)workspace; diff --git a/src/infiniop/ops/paged_caching/cuda/kernel_fp8.cuh b/src/infiniop/ops/paged_caching/cuda/kernel_fp8.cuh new file mode 100644 index 000000000..e626bf357 --- /dev/null +++ b/src/infiniop/ops/paged_caching/cuda/kernel_fp8.cuh @@ -0,0 +1,147 @@ +#ifndef __PAGED_CACHING_FP8_KERNEL_CUH__ +#define __PAGED_CACHING_FP8_KERNEL_CUH__ + +//================================================================================ +// Paged Caching FP8(E4M3) Quantizing Write Kernel +// +// Same grid/slot semantics as the plain copy kernel in kernel.cuh, but each +// block additionally performs dynamic per-token-per-head quantization: +// amax = max(|x[0:head_size]|) over the head_dim vector +// scale = amax / 448 (scale = 1 when amax == 0) +// q = e4m3_encode(x / scale) +// The scale is written to k_scale/v_scale at [block, kv_head, block_offset] +// and the encoded byte to the F8 cache. amax==0 yields scale=1 and all-zero +// codes (encode(0)==0). +//================================================================================ + +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" + +namespace op::paged_caching::cuda { + +namespace { + +// Block-wide max reduction. All NUM_THREADS threads must participate. +template +__device__ __forceinline__ float blockReduceMax(float value) { + constexpr int NUM_WARPS = NUM_THREADS / 32; + __shared__ float warp_max[NUM_WARPS]; + + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value = fmaxf(value, __shfl_xor_sync(0xffffffff, value, offset)); + } + if (lane == 0) { + warp_max[warp] = value; + } + __syncthreads(); + + float result = (threadIdx.x < NUM_WARPS) ? warp_max[threadIdx.x] : 0.0f; + if (warp == 0) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + result = fmaxf(result, __shfl_xor_sync(0xffffffff, result, offset)); + } + if (lane == 0) { + warp_max[0] = result; + } + } + __syncthreads(); + result = warp_max[0]; + __syncthreads(); + return result; +} + +// Quantize one head_dim vector (src -> dst) and write its dequant scale. +// All NUM_THREADS threads of the block must call this together. +template +__device__ __forceinline__ void quantizeHeadVector( + uint8_t *dst, + float *scale_out, + const Tdata *src, + const size_t num_dims) { + float amax = 0.0f; + for (int i = threadIdx.x; i < num_dims; i += NUM_THREADS) { + amax = fmaxf(amax, fabsf(static_cast(src[i]))); + } + amax = blockReduceMax(amax); + + const float scale = (amax > 0.0f) ? (amax / 448.0f) : 1.0f; + if (threadIdx.x == 0) { + *scale_out = scale; + } + const float inv_scale = 1.0f / scale; + for (int i = threadIdx.x; i < num_dims; i += NUM_THREADS) { + dst[i] = infiniopFp8E4m3Encode(static_cast(src[i]) * inv_scale); + } + // Keep the whole block in lockstep for the next vector's reduction. + __syncthreads(); +} + +} // namespace + +template < + typename Tdata, // Data type of the source K/V tensors (half, __nv_bfloat16) + int NUM_THREADS // Number of threads per block, configured at launch time + > +__device__ void pagedCachingFp8Kernel( + // ----- Output Tensors (F8 codes stored as raw bytes) ----- + uint8_t *k_cache_ptr, // [num_blocks, nkvh, block_size, dh] + uint8_t *v_cache_ptr, // [num_blocks, nkvh, block_size, dv] + float *k_scale_ptr, // [num_blocks, nkvh, block_size] + float *v_scale_ptr, // [num_blocks, nkvh, block_size] + // ----- Input Tensors ----- + const Tdata *k_ptr, // [ntok, nkvh, dh] + const Tdata *v_ptr, // [ntok, nkvh, dv] + const int64_t *slot_mapping_ptr, // [ntok] + // ----- Metadata ----- + const size_t head_size, // Dimension of each key head (dh_k) + const size_t v_head_size, // Dimension of each value head (dh_v) + const size_t block_size, // Number of tokens per block in the KV cache + // ----- Stride Information (identical semantics to the copy kernel) ----- + const ptrdiff_t k_src_stride, + const ptrdiff_t v_src_stride, + const ptrdiff_t k_src_head_stride, + const ptrdiff_t v_src_head_stride, + const ptrdiff_t k_cache_block_stride, + const ptrdiff_t v_cache_block_stride, + const ptrdiff_t k_cache_head_stride, + const ptrdiff_t v_cache_head_stride, + const ptrdiff_t k_cache_slot_stride, + const ptrdiff_t v_cache_slot_stride, + // ----- Scale strides ([num_blocks, nkvh, block_size]) ----- + const ptrdiff_t k_scale_block_stride, + const ptrdiff_t k_scale_head_stride, + const ptrdiff_t k_scale_slot_stride, + const ptrdiff_t v_scale_block_stride, + const ptrdiff_t v_scale_head_stride, + const ptrdiff_t v_scale_slot_stride) { + + const int token_idx = blockIdx.y; + const int head_idx = blockIdx.x; + + const int64_t slot_idx = slot_mapping_ptr[token_idx]; + if (slot_idx < 0) { + return; + } + const int64_t physical_block_idx = slot_idx / block_size; + const int64_t block_offset = slot_idx % block_size; + + const Tdata *k_src_head_ptr = k_ptr + token_idx * k_src_stride + head_idx * k_src_head_stride; + const Tdata *v_src_head_ptr = v_ptr + token_idx * v_src_stride + head_idx * v_src_head_stride; + + uint8_t *k_dst_head_ptr = k_cache_ptr + physical_block_idx * k_cache_block_stride + head_idx * k_cache_head_stride + block_offset * k_cache_slot_stride; + uint8_t *v_dst_head_ptr = v_cache_ptr + physical_block_idx * v_cache_block_stride + head_idx * v_cache_head_stride + block_offset * v_cache_slot_stride; + + float *k_scale_out = k_scale_ptr + physical_block_idx * k_scale_block_stride + head_idx * k_scale_head_stride + block_offset * k_scale_slot_stride; + float *v_scale_out = v_scale_ptr + physical_block_idx * v_scale_block_stride + head_idx * v_scale_head_stride + block_offset * v_scale_slot_stride; + + quantizeHeadVector(k_dst_head_ptr, k_scale_out, k_src_head_ptr, head_size); + quantizeHeadVector(v_dst_head_ptr, v_scale_out, v_src_head_ptr, v_head_size); +} + +} // namespace op::paged_caching::cuda + +#endif // __PAGED_CACHING_FP8_KERNEL_CUH__ diff --git a/src/infiniop/ops/paged_caching/info.h b/src/infiniop/ops/paged_caching/info.h index 8a6808390..6fe4cb4e9 100644 --- a/src/infiniop/ops/paged_caching/info.h +++ b/src/infiniop/ops/paged_caching/info.h @@ -14,6 +14,7 @@ class PagedCachingInfo { public: // --- Data Type --- infiniDtype_t dtype; + infiniDtype_t cache_dtype; // --- Shape Dimensions --- size_t num_tokens; @@ -34,18 +35,51 @@ class PagedCachingInfo { ptrdiff_t k_cache_slot_stride; ptrdiff_t v_cache_slot_stride; + // --- Strides for the per-token dequant scales ([num_blocks, num_kv_heads, block_size]) --- + // Only meaningful when cache_dtype == INFINI_DTYPE_F8. + ptrdiff_t k_scale_block_stride; + ptrdiff_t k_scale_head_stride; + ptrdiff_t k_scale_slot_stride; + ptrdiff_t v_scale_block_stride; + ptrdiff_t v_scale_head_stride; + ptrdiff_t v_scale_slot_stride; + static utils::Result create( infiniopTensorDescriptor_t k_cache_desc, 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) { auto dtype = k_desc->dtype(); CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - if (v_desc->dtype() != dtype || k_cache_desc->dtype() != dtype || v_cache_desc->dtype() != dtype) { + if (v_desc->dtype() != dtype) { return INFINI_STATUS_BAD_TENSOR_DTYPE; } + // The caches either keep the source dtype (plain copy) or store FP8(E4M3) + // codes plus per-token F32 scales (dynamic quantization on write). + auto cache_dtype = k_cache_desc->dtype(); + const bool cache_fp8 = (cache_dtype == INFINI_DTYPE_F8); + if (cache_fp8) { + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); + if (v_cache_desc->dtype() != INFINI_DTYPE_F8) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (k_scale_desc == nullptr || v_scale_desc == nullptr) { + printf("F8 paged_caching requires k_scale and v_scale.\n"); + return INFINI_STATUS_BAD_PARAM; + } + } else { + if (cache_dtype != dtype || v_cache_desc->dtype() != dtype) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (k_scale_desc != nullptr || v_scale_desc != nullptr) { + printf("k_scale/v_scale are only valid for F8 caches.\n"); + return INFINI_STATUS_BAD_PARAM; + } + } if (slot_mapping_desc->dtype() != INFINI_DTYPE_I64) { printf("slot_mapping must be int64_t.\n"); return INFINI_STATUS_BAD_TENSOR_DTYPE; @@ -83,6 +117,23 @@ class PagedCachingInfo { return INFINI_STATUS_BAD_TENSOR_SHAPE; } + // --- Validate per-token scale tensors for the FP8 path --- + if (cache_fp8) { + const size_t num_blocks = k_cache_shape[0]; + for (auto scale_desc : {k_scale_desc, v_scale_desc}) { + if (scale_desc->dtype() != INFINI_DTYPE_F32) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + if (scale_desc->ndim() != 3) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + const auto scale_shape = scale_desc->shape(); + if (scale_shape[0] != num_blocks || scale_shape[1] != num_kv_heads || scale_shape[2] != block_size) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + } + } + // --- Extract strides for memory access --- ptrdiff_t k_src_stride = k_desc->stride(0); ptrdiff_t v_src_stride = v_desc->stride(0); @@ -95,8 +146,16 @@ class PagedCachingInfo { ptrdiff_t k_cache_slot_stride = k_cache_desc->stride(2); ptrdiff_t v_cache_slot_stride = v_cache_desc->stride(2); + ptrdiff_t k_scale_block_stride = cache_fp8 ? k_scale_desc->stride(0) : 0; + ptrdiff_t k_scale_head_stride = cache_fp8 ? k_scale_desc->stride(1) : 0; + ptrdiff_t k_scale_slot_stride = cache_fp8 ? k_scale_desc->stride(2) : 0; + ptrdiff_t v_scale_block_stride = cache_fp8 ? v_scale_desc->stride(0) : 0; + ptrdiff_t v_scale_head_stride = cache_fp8 ? v_scale_desc->stride(1) : 0; + ptrdiff_t v_scale_slot_stride = cache_fp8 ? v_scale_desc->stride(2) : 0; + return utils::Result(PagedCachingInfo{ dtype, + cache_dtype, num_tokens, num_kv_heads, head_size, @@ -111,7 +170,13 @@ class PagedCachingInfo { k_cache_head_stride, v_cache_head_stride, k_cache_slot_stride, - v_cache_slot_stride}); + v_cache_slot_stride, + k_scale_block_stride, + k_scale_head_stride, + k_scale_slot_stride, + v_scale_block_stride, + v_scale_head_stride, + v_scale_slot_stride}); } }; diff --git a/src/infiniop/ops/paged_caching/metax/paged_caching_metax.maca b/src/infiniop/ops/paged_caching/metax/paged_caching_metax.maca index a14f5073a..240e72470 100644 --- a/src/infiniop/ops/paged_caching/metax/paged_caching_metax.maca +++ b/src/infiniop/ops/paged_caching/metax/paged_caching_metax.maca @@ -40,10 +40,15 @@ infiniStatus_t Descriptor::create( 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) { - auto info = PagedCachingInfo::create(k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc); + auto info = PagedCachingInfo::create(k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc, k_scale_desc, v_scale_desc); CHECK_RESULT(info); + if (info->cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } // Create and return the Descriptor instance. *desc_ptr = new Descriptor( @@ -153,6 +158,7 @@ infiniStatus_t Descriptor::calculate( void *k_cache, void *v_cache, const void *k, const void *v, const void *slot_mapping, + void *k_scale, void *v_scale, void *stream_) const { hcStream_t stream = (hcStream_t)stream_; diff --git a/src/infiniop/ops/paged_caching/moore/paged_caching_moore.mu b/src/infiniop/ops/paged_caching/moore/paged_caching_moore.mu index 579e7cf05..893b1ad22 100644 --- a/src/infiniop/ops/paged_caching/moore/paged_caching_moore.mu +++ b/src/infiniop/ops/paged_caching/moore/paged_caching_moore.mu @@ -40,10 +40,15 @@ infiniStatus_t Descriptor::create( 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) { - auto info = PagedCachingInfo::create(k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc); + auto info = PagedCachingInfo::create(k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc, k_scale_desc, v_scale_desc); CHECK_RESULT(info); + if (info->cache_dtype == INFINI_DTYPE_F8) { + return INFINI_STATUS_NOT_IMPLEMENTED; + } // Create and return the Descriptor instance. *desc_ptr = new Descriptor( @@ -153,6 +158,7 @@ infiniStatus_t Descriptor::calculate( void *k_cache, void *v_cache, const void *k, const void *v, const void *slot_mapping, + void *k_scale, void *v_scale, void *stream_) const { musaStream_t stream = (musaStream_t)stream_; diff --git a/src/infiniop/ops/paged_caching/nvidia/paged_caching_nvidia.cu b/src/infiniop/ops/paged_caching/nvidia/paged_caching_nvidia.cu index 202c02b23..5b2b106e0 100644 --- a/src/infiniop/ops/paged_caching/nvidia/paged_caching_nvidia.cu +++ b/src/infiniop/ops/paged_caching/nvidia/paged_caching_nvidia.cu @@ -1,6 +1,7 @@ #include "../../../devices/nvidia/nvidia_common.cuh" #include "../../../devices/nvidia/nvidia_kernel_common.cuh" #include "../cuda/kernel.cuh" +#include "../cuda/kernel_fp8.cuh" #include "paged_caching_nvidia.cuh" template @@ -21,6 +22,29 @@ INFINIOP_CUDA_KERNEL pagedCaching( k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, v_cache_head_stride, k_cache_slot_stride, v_cache_slot_stride); } +template +INFINIOP_CUDA_KERNEL pagedCachingFp8( + uint8_t *k_cache, uint8_t *v_cache, + float *k_scale, float *v_scale, + const Tdata *k, const Tdata *v, + const int64_t *slot_mapping, + const size_t head_size, const size_t v_head_size, const size_t block_size, + const ptrdiff_t k_src_stride, const ptrdiff_t v_src_stride, + const ptrdiff_t k_src_head_stride, const ptrdiff_t v_src_head_stride, + const ptrdiff_t k_cache_block_stride, const ptrdiff_t v_cache_block_stride, + const ptrdiff_t k_cache_head_stride, const ptrdiff_t v_cache_head_stride, + const ptrdiff_t k_cache_slot_stride, const ptrdiff_t v_cache_slot_stride, + const ptrdiff_t k_scale_block_stride, const ptrdiff_t k_scale_head_stride, const ptrdiff_t k_scale_slot_stride, + const ptrdiff_t v_scale_block_stride, const ptrdiff_t v_scale_head_stride, const ptrdiff_t v_scale_slot_stride) { + op::paged_caching::cuda::pagedCachingFp8Kernel( + k_cache, v_cache, k_scale, v_scale, k, v, slot_mapping, head_size, v_head_size, + block_size, k_src_stride, v_src_stride, + k_src_head_stride, v_src_head_stride, + k_cache_block_stride, v_cache_block_stride, k_cache_head_stride, v_cache_head_stride, k_cache_slot_stride, v_cache_slot_stride, + k_scale_block_stride, k_scale_head_stride, k_scale_slot_stride, + v_scale_block_stride, v_scale_head_stride, v_scale_slot_stride); +} + namespace op::paged_caching::nvidia { // PIMPL struct definition struct Descriptor::Opaque { @@ -40,9 +64,11 @@ infiniStatus_t Descriptor::create( 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) { - auto info = PagedCachingInfo::create(k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc); + auto info = PagedCachingInfo::create(k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc, k_scale_desc, v_scale_desc); CHECK_RESULT(info); // Create and return the Descriptor instance. @@ -147,16 +173,85 @@ infiniStatus_t launchKernel(const PagedCachingInfo &info, return INFINI_STATUS_SUCCESS; } +// FP8(E4M3) quantizing write path: same grid/block geometry as the copy +// kernel, but each block computes the per-(token, head) amax, writes the +// dequant scale, and stores E4M3 codes instead of raw values. +template +infiniStatus_t launchKernelFp8(const PagedCachingInfo &info, + void *k_cache, void *v_cache, + void *k_scale, void *v_scale, + const void *k, const void *v, + const void *slot_mapping, + cudaStream_t stream) { + + dim3 grid(uint64_t(info.num_kv_heads), uint64_t(info.num_tokens), 1); + dim3 block(NUM_THREADS); + size_t shared_mem_size = 0; + +#define LAUNCH_FP8(Tdata) \ + pagedCachingFp8 \ + <<>>( \ + (uint8_t *)k_cache, \ + (uint8_t *)v_cache, \ + (float *)k_scale, \ + (float *)v_scale, \ + (const Tdata *)k, \ + (const Tdata *)v, \ + (const int64_t *)slot_mapping, \ + info.head_size, \ + info.v_head_size, \ + info.block_size, \ + info.k_src_stride, \ + info.v_src_stride, \ + info.k_src_head_stride, \ + info.v_src_head_stride, \ + info.k_cache_block_stride, \ + info.v_cache_block_stride, \ + info.k_cache_head_stride, \ + info.v_cache_head_stride, \ + info.k_cache_slot_stride, \ + info.v_cache_slot_stride, \ + info.k_scale_block_stride, \ + info.k_scale_head_stride, \ + info.k_scale_slot_stride, \ + info.v_scale_block_stride, \ + info.v_scale_head_stride, \ + info.v_scale_slot_stride) + + if (info.dtype == INFINI_DTYPE_F16) { + LAUNCH_FP8(half); + } else if (info.dtype == INFINI_DTYPE_BF16) { + LAUNCH_FP8(__nv_bfloat16); + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +#undef LAUNCH_FP8 + return INFINI_STATUS_SUCCESS; +} + // Execution method implementation infiniStatus_t Descriptor::calculate( void *workspace, size_t workspace_size, void *k_cache, void *v_cache, const void *k, const void *v, const void *slot_mapping, + void *k_scale, void *v_scale, void *stream_) const { cudaStream_t stream = (cudaStream_t)stream_; + if (_info.cache_dtype == INFINI_DTYPE_F8) { + if (_opaque->internal->maxThreadsPerBlock() >= CUDA_BLOCK_SIZE_1024) { + return launchKernelFp8(_info, k_cache, v_cache, k_scale, v_scale, k, v, slot_mapping, stream); + } else if (_opaque->internal->maxThreadsPerBlock() >= CUDA_BLOCK_SIZE_512) { + return launchKernelFp8(_info, k_cache, v_cache, k_scale, v_scale, k, v, slot_mapping, stream); + } else if (_opaque->internal->maxThreadsPerBlock() >= CUDA_BLOCK_SIZE_4096) { + return launchKernelFp8(_info, k_cache, v_cache, k_scale, v_scale, k, v, slot_mapping, stream); + } else { + return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; + } + } + // Dispatch logic based on the GPU's maximum threads per block. // This allows selecting the largest, most efficient block size the hardware supports. if (_opaque->internal->maxThreadsPerBlock() >= CUDA_BLOCK_SIZE_1024) { diff --git a/src/infiniop/ops/paged_caching/operator.cc b/src/infiniop/ops/paged_caching/operator.cc index a32b97669..4d86f67bd 100644 --- a/src/infiniop/ops/paged_caching/operator.cc +++ b/src/infiniop/ops/paged_caching/operator.cc @@ -25,14 +25,17 @@ __INFINI_C 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) { #define CREATE(CASE, NAMESPACE) \ case CASE: \ return op::paged_caching::NAMESPACE::Descriptor::create( \ handle, \ reinterpret_cast(desc_ptr), \ - k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc); + k_cache_desc, v_cache_desc, k_desc, v_desc, slot_mapping_desc, \ + k_scale_desc, v_scale_desc); switch (handle->device) { #ifdef ENABLE_NVIDIA_API @@ -115,12 +118,14 @@ __INFINI_C infiniStatus_t infiniopPagedCaching( void *k_cache, void *v_cache, const void *k, const void *v, const void *slot_mapping, + void *k_scale, void *v_scale, void *stream) { #define CALCULATE(CASE, NAMESPACE) \ case CASE: \ return reinterpret_cast(desc)->calculate( \ - workspace, workspace_size, k_cache, v_cache, k, v, slot_mapping, stream); + workspace, workspace_size, k_cache, v_cache, k, v, slot_mapping, \ + k_scale, v_scale, stream); switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API diff --git a/src/infiniop/ops/paged_caching/paged_caching.h b/src/infiniop/ops/paged_caching/paged_caching.h index cbd7a2701..faf85fb8d 100644 --- a/src/infiniop/ops/paged_caching/paged_caching.h +++ b/src/infiniop/ops/paged_caching/paged_caching.h @@ -36,13 +36,16 @@ 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); \ \ infiniStatus_t calculate( \ void *workspace, size_t workspace_size, \ void *k_cache, void *v_cache, \ const void *k, const void *v, \ const void *slot_mapping, \ + void *k_scale, void *v_scale, \ void *stream) const; \ }; \ } diff --git a/test/infiniop/fp8_blockwise_dequantize.py b/test/infiniop/fp8_blockwise_dequantize.py new file mode 100644 index 000000000..c50100b4d --- /dev/null +++ b/test/infiniop/fp8_blockwise_dequantize.py @@ -0,0 +1,125 @@ +import ctypes +from ctypes import c_uint64 + +import torch + +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + InfiniDtypeNames, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopOperatorDescriptor_t, + test_operator, + to_torch_dtype, +) + +_TEST_CASES = [ + # (M, N, BM, BN) + (256, 256, 128, 128), + (1024, 1024, 128, 128), + (384, 1536, 128, 128), + (512, 768, 64, 128), + (256, 512, 128, 64), +] + +_TENSOR_DTYPES = [InfiniDtype.F16, InfiniDtype.BF16, InfiniDtype.F32] + + +def reference_dequantize(q, scales, block_m, block_n, out_dtype): + """Torch oracle: decode E4M3 bytes, apply the per-block scale, cast.""" + q_float = q.to(torch.float32) + scales_full = scales.repeat_interleave(block_m, dim=0).repeat_interleave(block_n, dim=1) + return (q_float * scales_full).to(out_dtype) + + +def run_dequantize(handle, device, q_data, scales_data, block_m, block_n, dtype, sync=None): + q = TestTensor.from_torch(q_data, InfiniDtype.F8, device) + scales = TestTensor.from_torch(scales_data, InfiniDtype.F32, device) + out = TestTensor(q_data.shape, None, dtype, device, mode="zeros") + expected = reference_dequantize( + q.torch_tensor(), scales.torch_tensor(), block_m, block_n, to_torch_dtype(dtype) + ) + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateFp8BlockwiseDequantizeDescriptor( + handle, + ctypes.byref(descriptor), + out.descriptor, + q.descriptor, + scales.descriptor, + ) + ) + for tensor in (out, q, scales): + tensor.destroy_desc() + + workspace_size = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetFp8BlockwiseDequantizeWorkspaceSize( + descriptor, ctypes.byref(workspace_size) + ) + ) + workspace = TestWorkspace(workspace_size.value, device) + check_error( + LIBINFINIOP.infiniopFp8BlockwiseDequantize( + descriptor, + workspace.data(), + workspace_size.value, + out.data(), + q.data(), + scales.data(), + None, + ) + ) + if sync is not None: + sync() + + if dtype == InfiniDtype.F16: + # CPU f32->f16 conversion truncates instead of rounding to nearest, so + # allow a 1-ulp fp16 gap; BF16 (round-to-nearest-even) and F32 are exact. + torch.testing.assert_close(out.actual_tensor(), expected, atol=1e-3, rtol=1e-3) + else: + torch.testing.assert_close(out.actual_tensor(), expected, atol=0, rtol=0) + check_error(LIBINFINIOP.infiniopDestroyFp8BlockwiseDequantizeDescriptor(descriptor)) + + +def test( + handle, + device, + m, + n, + block_m, + block_n, + dtype, + sync=None, +): + print( + f"Testing FP8 blockwise dequantize on {InfiniDeviceNames[device]} " + f"with q_shape=({m}, {n}), block=({block_m}, {block_n}), " + f"out_dtype={InfiniDtypeNames[dtype]}" + ) + + probe = TestTensor((1,), None, InfiniDtype.F32, device, mode="zeros") + torch_device = probe.actual_tensor().device + # Random floats cover signs, normals, subnormals and zero once quantized to + # E4M3; all bit patterns produced by .to(torch.float8_e4m3fn) are legal. + q_data = ((torch.rand(m, n, dtype=torch.float32, device=torch_device) * 2 - 1) * 8).to( + torch.float8_e4m3fn + ) + scales_data = ( + torch.rand(m // block_m, n // block_n, dtype=torch.float32, device=torch_device) * 2 + 0.5 + ) + run_dequantize(handle, device, q_data, scales_data, block_m, block_n, dtype, sync) + + +if __name__ == "__main__": + args = get_args() + torch.manual_seed(0) + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + print("\033[92mTest passed!\033[0m") diff --git a/test/infiniop/fp8_blockwise_gemm.py b/test/infiniop/fp8_blockwise_gemm.py new file mode 100644 index 000000000..7af05d660 --- /dev/null +++ b/test/infiniop/fp8_blockwise_gemm.py @@ -0,0 +1,114 @@ +import ctypes +from ctypes import c_uint64 + +import torch + +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + InfiniDtypeNames, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopOperatorDescriptor_t, + test_operator, + to_torch_dtype, +) + +_TEST_CASES = [ + # (M, N, K, BM, BK) + (1, 256, 256, 128, 128), + (2, 256, 512, 128, 128), + (4, 512, 384, 64, 128), + (8, 1024, 512, 128, 256), + (16, 4096, 4096, 128, 128), + (13, 2048, 1024, 128, 128), + (32, 768, 1280, 128, 128), +] + +_TENSOR_DTYPES = [InfiniDtype.F16, InfiniDtype.BF16, InfiniDtype.F32] + + +def reference_gemm(a, q, scales, block_n, block_k): + """Torch oracle in fp32: dequantize per block, then matmul.""" + w = q.to(torch.float32) * scales.repeat_interleave(block_n, dim=0).repeat_interleave(block_k, dim=1) + return a.to(torch.float32) @ w.t() + + +def run_case(handle, device, a_data, q_data, scales_data, block_n, block_k, dtype, sync=None): + a = TestTensor.from_torch(a_data, dtype, device) + q = TestTensor.from_torch(q_data, InfiniDtype.F8, device) + scales = TestTensor.from_torch(scales_data, InfiniDtype.F32, device) + out = TestTensor((a_data.shape[0], q_data.shape[0]), None, dtype, device, mode="zeros") + expected = reference_gemm(a.torch_tensor(), q.torch_tensor(), scales.torch_tensor(), block_n, block_k) + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateFp8BlockwiseGemmDescriptor( + handle, + ctypes.byref(descriptor), + out.descriptor, + a.descriptor, + q.descriptor, + scales.descriptor, + ) + ) + for tensor in (out, a, q, scales): + tensor.destroy_desc() + + workspace_size = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetFp8BlockwiseGemmWorkspaceSize(descriptor, ctypes.byref(workspace_size)) + ) + workspace = TestWorkspace(workspace_size.value, device) + check_error( + LIBINFINIOP.infiniopFp8BlockwiseGemm( + descriptor, + workspace.data(), + workspace_size.value, + out.data(), + a.data(), + q.data(), + scales.data(), + None, + ) + ) + if sync is not None: + sync() + + actual = out.actual_tensor().to(torch.float32) + rel = (actual - expected).abs().mean() / expected.abs().mean().clamp(min=1e-6) + limit = 1e-3 if dtype == InfiniDtype.F32 else 2e-2 + assert rel < limit, f"relative error {rel.item():.6f} >= {limit}" + check_error(LIBINFINIOP.infiniopDestroyFp8BlockwiseGemmDescriptor(descriptor)) + + +def test(handle, device, m, n, k, block_n, block_k, dtype, sync=None): + print( + f"Testing FP8 blockwise GEMM on {InfiniDeviceNames[device]} " + f"with M={m}, N={n}, K={k}, block=({block_n}, {block_k}), " + f"dtype={InfiniDtypeNames[dtype]}" + ) + torch_dtype = to_torch_dtype(dtype) + probe = TestTensor((1,), None, InfiniDtype.F32, device, mode="zeros") + torch_device = probe.actual_tensor().device + + a_data = torch.randn(m, k, dtype=torch_dtype, device=torch_device) * 0.5 + q_data = ((torch.rand(n, k, dtype=torch.float32, device=torch_device) * 2 - 1) * 4).to( + torch.float8_e4m3fn + ) + scales_data = ( + torch.rand(n // block_n, k // block_k, dtype=torch.float32, device=torch_device) * 0.02 + 0.005 + ) + run_case(handle, device, a_data, q_data, scales_data, block_n, block_k, dtype, sync) + + +if __name__ == "__main__": + args = get_args() + torch.manual_seed(0) + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + print("\033[92mTest passed!\033[0m") diff --git a/test/infiniop/libinfiniop/op_register.py b/test/infiniop/libinfiniop/op_register.py index ec9add01a..482c3c142 100644 --- a/test/infiniop/libinfiniop/op_register.py +++ b/test/infiniop/libinfiniop/op_register.py @@ -1394,6 +1394,70 @@ def mxfp4_dequantize_(lib): ] +@OpRegister.operator +def fp8_blockwise_dequantize_(lib): + lib.infiniopCreateFp8BlockwiseDequantizeDescriptor.restype = c_int32 + lib.infiniopCreateFp8BlockwiseDequantizeDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + ] + lib.infiniopGetFp8BlockwiseDequantizeWorkspaceSize.restype = c_int32 + lib.infiniopGetFp8BlockwiseDequantizeWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopFp8BlockwiseDequantize.restype = c_int32 + lib.infiniopFp8BlockwiseDequantize.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyFp8BlockwiseDequantizeDescriptor.restype = c_int32 + lib.infiniopDestroyFp8BlockwiseDequantizeDescriptor.argtypes = [ + infiniopOperatorDescriptor_t, + ] + + +@OpRegister.operator +def fp8_blockwise_gemm_(lib): + lib.infiniopCreateFp8BlockwiseGemmDescriptor.restype = c_int32 + lib.infiniopCreateFp8BlockwiseGemmDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + ] + lib.infiniopGetFp8BlockwiseGemmWorkspaceSize.restype = c_int32 + lib.infiniopGetFp8BlockwiseGemmWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopFp8BlockwiseGemm.restype = c_int32 + lib.infiniopFp8BlockwiseGemm.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyFp8BlockwiseGemmDescriptor.restype = c_int32 + lib.infiniopDestroyFp8BlockwiseGemmDescriptor.argtypes = [ + infiniopOperatorDescriptor_t, + ] + + @OpRegister.operator def per_channel_quant_int8_(lib): lib.infiniopCreatePerChannelQuantI8Descriptor.restype = c_int32 @@ -2197,6 +2261,8 @@ def paged_attention_(lib): infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, c_void_p, + infiniopTensorDescriptor_t, # k_scale_desc + infiniopTensorDescriptor_t, # v_scale_desc c_float, ] @@ -2218,7 +2284,9 @@ def paged_attention_(lib): c_void_p, c_void_p, c_void_p, - c_void_p, + c_void_p, # k_scale + c_void_p, # v_scale + c_void_p, # stream ] lib.infiniopDestroyPagedAttentionDescriptor.restype = c_int32 @@ -2238,6 +2306,8 @@ def paged_caching_(lib): infiniopTensorDescriptor_t, # k_desc infiniopTensorDescriptor_t, # v_desc infiniopTensorDescriptor_t, # slot_mapping_desc + infiniopTensorDescriptor_t, # k_scale_desc + infiniopTensorDescriptor_t, # v_scale_desc ] # infiniopGetPagedCachingWorkspaceSize @@ -2258,6 +2328,8 @@ def paged_caching_(lib): c_void_p, # k c_void_p, # v c_void_p, # slot_mapping + c_void_p, # k_scale + c_void_p, # v_scale c_void_p, # stream ] @@ -2281,7 +2353,9 @@ def paged_attention_prefill_(lib): infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, - infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, # alibi_slopes_desc + infiniopTensorDescriptor_t, # k_scale_desc + infiniopTensorDescriptor_t, # v_scale_desc c_float, ] @@ -2303,8 +2377,10 @@ def paged_attention_prefill_(lib): c_void_p, c_void_p, c_void_p, - c_void_p, - c_void_p, + c_void_p, # alibi_slopes + c_void_p, # k_scale + c_void_p, # v_scale + c_void_p, # stream ] lib.infiniopDestroyPagedAttentionPrefillDescriptor.restype = c_int32 diff --git a/test/infiniop/paged_attention.py b/test/infiniop/paged_attention.py index d7a6b9679..4e4937415 100644 --- a/test/infiniop/paged_attention.py +++ b/test/infiniop/paged_attention.py @@ -14,6 +14,7 @@ profile_operation, InfiniDtype, InfiniDtypeNames, + InfiniDeviceEnum, InfiniDeviceNames, infiniopOperatorDescriptor_t, TestWorkspace, @@ -209,6 +210,8 @@ def test( block_tables.descriptor, seq_lens.descriptor, alibi_slopes_desc, + None, + None, scale, ) ) @@ -247,6 +250,8 @@ def lib_paged_attention(): seq_lens.data(), alibi_slopes_data, None, + None, + None, ) ) @@ -277,6 +282,190 @@ def lib_paged_attention(): check_error(LIBINFINIOP.infiniopDestroyPagedAttentionDescriptor(descriptor)) +# ============================================================================== +# FP8 (E4M3) KV-Cache Decode Test +# ============================================================================== +# FP8 decode v1 supports head_size 64/128 with value_size == head_size. +# The caches carry E4M3 codes plus per-token F32 scales; q/out stay F16/BF16. +_TEST_CASES_FP8_ = [ + # (num_seqs, num_heads, num_kv_heads, head_size, block_size, max_seq_len, use_alibi) + (1, 1, 1, 128, 16, 1024, False), + (4, 40, 40, 128, 16, 1024, True), + # Small grid + long context: the split-kv heuristic picks num_splits > 1, + # exercising the cross-CTA shard scan + combine pass (with ALiBi here). + (1, 8, 8, 128, 16, 4096, True), + (8, 64, 8, 128, 16, 2048, False), + (3, 8, 8, 64, 16, 1024, False), +] + +_TENSOR_DTYPES_FP8_ = [InfiniDtype.F16, InfiniDtype.BF16] + + +def quantize_cache_ref(pool_float): + """ + Quantize a float32 cache pool [num_blocks, nkvh, block_size, d] the same way + paged_caching does: per-(block, head, slot) amax/448 scaling, then E4M3 RNE. + Returns (codes_uint8, scales_float32). + """ + amax = pool_float.abs().amax(dim=-1) + scales = torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)) + inv_scales = 1.0 / scales + codes = (pool_float * inv_scales.unsqueeze(-1)).to(torch.float8_e4m3fn) + return codes, scales + + +def test_fp8( + handle, + device, + num_seqs, + num_heads, + num_kv_heads, + head_size, + block_size, + max_seq_len, + use_alibi, + dtype, + sync, +): + print( + f"Testing PagedAttention FP8 on {InfiniDeviceNames[device]} with " + f"num_seqs={num_seqs}, num_heads={num_heads}, head_size={head_size}, " + f"block_size={block_size}, dtype={InfiniDtypeNames[dtype]}, use_alibi={use_alibi}" + ) + + scale = 1.0 / (head_size**0.5) + max_blocks_per_seq = (max_seq_len + block_size - 1) // block_size + num_blocks = num_seqs * max_blocks_per_seq + + q = TestTensor((num_seqs, num_heads, head_size), None, dtype, device) + out = TestTensor((num_seqs, num_heads, head_size), None, dtype, device) + + # Build FP8 caches by quantizing random float pools per token. + k_pool = TestTensor( + (num_blocks, num_kv_heads, block_size, head_size), + None, + InfiniDtype.F32, + device, + scale=4.0, + bias=-2.0, + ) + v_pool = TestTensor( + (num_blocks, num_kv_heads, block_size, head_size), + None, + InfiniDtype.F32, + device, + scale=4.0, + bias=-2.0, + ) + k_codes_torch, k_scales_torch = quantize_cache_ref(k_pool.torch_tensor()) + v_codes_torch, v_scales_torch = quantize_cache_ref(v_pool.torch_tensor()) + + k_cache = TestTensor.from_torch(k_codes_torch, InfiniDtype.F8, device) + v_cache = TestTensor.from_torch(v_codes_torch, InfiniDtype.F8, device) + k_scale_t = TestTensor.from_torch(k_scales_torch, InfiniDtype.F32, device) + v_scale_t = TestTensor.from_torch(v_scales_torch, InfiniDtype.F32, device) + + seq_lens_torch = torch.randint(1, max_seq_len, (num_seqs,), dtype=torch.int64) + seq_lens = TestTensor.from_torch(seq_lens_torch, InfiniDtype.I64, device) + + block_tables_py = torch.arange( + 0, num_seqs * max_blocks_per_seq, dtype=torch.int64 + ).view(num_seqs, max_blocks_per_seq) + block_tables = TestTensor.from_torch(block_tables_py, InfiniDtype.I64, device) + + alibi_slopes_desc = ctypes.c_void_p(0) + alibi_slopes_data = ctypes.c_void_p(0) + alibi_slopes_torch = None + if use_alibi: + alibi_slopes = TestTensor((num_heads,), None, InfiniDtype.F32, device) + alibi_slopes_desc = alibi_slopes.descriptor + alibi_slopes_data = alibi_slopes.data() + alibi_slopes_torch = alibi_slopes.torch_tensor() + + # Reference: dequantize the caches to float32 and run the attention reference + # entirely in float32 so only the kernel's own arithmetic is compared. + k_deq = k_codes_torch.float() * k_scales_torch.unsqueeze(-1) + v_deq = v_codes_torch.float() * v_scales_torch.unsqueeze(-1) + ans = ref_single_query_cached_kv_attention( + q.torch_tensor().float(), + k_deq, + v_deq, + block_tables.torch_tensor(), + seq_lens.torch_tensor(), + scale, + alibi_slopes_torch, + ) + + if sync: + sync() + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreatePagedAttentionDescriptor( + handle, + ctypes.byref(descriptor), + out.descriptor, + q.descriptor, + k_cache.descriptor, + v_cache.descriptor, + block_tables.descriptor, + seq_lens.descriptor, + alibi_slopes_desc, + k_scale_t.descriptor, + v_scale_t.descriptor, + scale, + ) + ) + + workspace_size = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetPagedAttentionWorkspaceSize( + descriptor, ctypes.byref(workspace_size) + ) + ) + workspace = TestWorkspace(workspace_size.value, q.device) + + q.destroy_desc() + out.destroy_desc() + k_cache.destroy_desc() + v_cache.destroy_desc() + k_scale_t.destroy_desc() + v_scale_t.destroy_desc() + block_tables.destroy_desc() + seq_lens.destroy_desc() + if use_alibi: + alibi_slopes.destroy_desc() + + check_error( + LIBINFINIOP.infiniopPagedAttention( + descriptor, + workspace.data(), + workspace_size.value, + out.data(), + q.data(), + k_cache.data(), + v_cache.data(), + block_tables.data(), + seq_lens.data(), + alibi_slopes_data, + k_scale_t.data(), + v_scale_t.data(), + None, + ) + ) + + if sync: + sync() + + atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype) + out_float = out.actual_tensor().float() + if DEBUG: + debug(out_float, ans, atol=atol, rtol=rtol) + assert torch.allclose(out_float, ans, atol=atol, rtol=rtol) + + check_error(LIBINFINIOP.infiniopDestroyPagedAttentionDescriptor(descriptor)) + + if __name__ == "__main__": args = get_args() @@ -288,5 +477,7 @@ def lib_paged_attention(): for device in get_test_devices(args): test_operator(device, test, _TEST_CASES_, _TENSOR_DTYPES) + if device == InfiniDeviceEnum.NVIDIA: + test_operator(device, test_fp8, _TEST_CASES_FP8_, _TENSOR_DTYPES_FP8_) print("\033[92mTest passed!\033[0m") diff --git a/test/infiniop/paged_attention_prefill.py b/test/infiniop/paged_attention_prefill.py index 6f9c83e41..52124edb3 100644 --- a/test/infiniop/paged_attention_prefill.py +++ b/test/infiniop/paged_attention_prefill.py @@ -4,6 +4,7 @@ import torch from libinfiniop import ( LIBINFINIOP, + InfiniDeviceEnum, InfiniDeviceNames, InfiniDtype, InfiniDtypeNames, @@ -269,6 +270,8 @@ def torch_paged_attention_multi_turn(): seq_lens.descriptor, cum_seq_lens_q.descriptor, None, + None, + None, scale, ) ) @@ -296,6 +299,8 @@ def lib_attn(): cum_seq_lens_q.data(), None, None, + None, + None, ) ) @@ -328,6 +333,222 @@ def lib_attn(): ) +# ============================================================================== +# FP8 (E4M3) KV-Cache Prefill Test +# ============================================================================== +# FP8 prefill (plan B) gather-dequantizes the paged F8 caches into a compact +# F16/BF16 scratch and runs the regular prefill kernels on it. The caches carry +# E4M3 codes plus per-token F32 scales; q/out stay F16/BF16. +_TEST_CASES_FP8 = [ + # num_seqs, num_heads, num_kv_heads, head_size, block_size, max_step_len, num_rounds, index_dtypes + (1, 4, 4, 128, 8, 16, 2, InfiniDtype.I64), + (2, 8, 8, 128, 16, 32, 2, InfiniDtype.I32), + (4, 16, 16, 128, 8, 64, 2, InfiniDtype.I64), + (2, 8, 8, 64, 16, 32, 1, InfiniDtype.I64), +] + +_TENSOR_DTYPES_FP8 = [InfiniDtype.F16, InfiniDtype.BF16] + + +def quantize_per_token_ref(x): + """ + Reference dynamic per-token-per-head FP8(E4M3) quantization with the same + float32 op order as the kernels: amax/448 scaling, then E4M3 RNE encoding. + """ + x_float = x.float() + amax = x_float.abs().amax(dim=-1) + scales = torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)) + inv_scales = 1.0 / scales + codes = (x_float * inv_scales.unsqueeze(-1)).to(torch.float8_e4m3fn) + return codes, scales + + +def test_fp8( + handle, + device, + num_seqs, + num_heads, + num_kv_heads, + head_size, + block_size, + max_step_len, + num_rounds, + index_dtype=InfiniDtype.I64, + *tail, +): + if len(tail) == 2: + dtype, sync = tail + value_size = head_size + elif len(tail) == 3: + value_size, dtype, sync = tail + else: + raise ValueError(f"Unexpected paged_attention_prefill FP8 test arguments: {tail}") + print( + f"Testing PagedAttentionPrefill FP8 on {InfiniDeviceNames[device]} with " + f"seqs:{num_seqs}, heads:{num_heads}, head_size:{head_size}, value_size:{value_size}, " + f"block:{block_size}, max_step_len:{max_step_len}, num_rounds:{num_rounds}, dtype:{InfiniDtypeNames[dtype]}, " + f"index_dtype:{InfiniDtypeNames[index_dtype]}" + ) + + num_blocks = 2048 + manager = SimpleCacheManager(num_blocks, block_size) + scale = head_size**-0.5 + + # Quantized cache pools: E4M3 codes plus per-token F32 scales. + k_pool = TestTensor( + (num_blocks, num_kv_heads, block_size, head_size), + None, + InfiniDtype.F32, + device, + scale=4.0, + bias=-2.0, + ) + v_pool = TestTensor( + (num_blocks, num_kv_heads, block_size, value_size), + None, + InfiniDtype.F32, + device, + scale=4.0, + bias=-2.0, + ) + k_codes, k_scales = quantize_per_token_ref(k_pool.torch_tensor()) + v_codes, v_scales = quantize_per_token_ref(v_pool.torch_tensor()) + + torch_idx_type = ( + torch.int32 if index_dtype in (InfiniDtype.I32, InfiniDtype.U32) else torch.int64 + ) + + for r in range(num_rounds): + query_lens_cpu = torch.randint(1, max_step_len + 1, (num_seqs,), dtype=torch.int64) + q_total_tokens = query_lens_cpu.sum().item() + q_packed_tensors = torch.zeros(q_total_tokens, num_heads, head_size) + + seq_lens_list = [] + all_block_tables = [] + cum_seq_lens_q_list = [] + cum_q_lens = 0 + for i in range(num_seqs): + cum_seq_lens_q_list.append(cum_q_lens) + + cur_q_len = query_lens_cpu[i].item() + table, total_len = manager.allocate_slots(i, cur_q_len) + cur_seq_lens = total_len - cur_q_len + seq_lens_list.append(total_len) + all_block_tables.append(table) + + # Simulated KV insertion: quantize the new tokens and write codes+scales. + k_new = torch.randn(cur_q_len, num_kv_heads, head_size) + v_new = torch.randn(cur_q_len, num_kv_heads, value_size) + q_val = torch.randn(cur_q_len, num_heads, head_size) + q_packed_tensors[cum_q_lens : cum_q_lens + cur_q_len] = q_val + cum_q_lens = cum_q_lens + cur_q_len + + k_new_codes, k_new_scales = quantize_per_token_ref(k_new) + v_new_codes, v_new_scales = quantize_per_token_ref(v_new) + for t in range(cur_q_len): + logical_pos = cur_seq_lens + t + b_id = table[logical_pos // block_size] + off = logical_pos % block_size + k_codes[b_id, :, off, :] = k_new_codes[t].to(k_codes.device) + v_codes[b_id, :, off, :] = v_new_codes[t].to(v_codes.device) + k_scales[b_id, :, off] = k_new_scales[t].to(k_scales.device) + v_scales[b_id, :, off] = v_new_scales[t].to(v_scales.device) + + cum_seq_lens_q_list.append(cum_q_lens) + + q_new = TestTensor.from_torch(q_packed_tensors, dtype, device) + out = TestTensor((q_total_tokens, num_heads, value_size), None, dtype, device) + out.actual_tensor().zero_() + + k_cache = TestTensor.from_torch(k_codes, InfiniDtype.F8, device) + v_cache = TestTensor.from_torch(v_codes, InfiniDtype.F8, device) + k_scale_t = TestTensor.from_torch(k_scales, InfiniDtype.F32, device) + v_scale_t = TestTensor.from_torch(v_scales, InfiniDtype.F32, device) + + seq_lens = TestTensor.from_torch( + torch.tensor(seq_lens_list, dtype=torch_idx_type), index_dtype, device + ) + cum_seq_lens_q = TestTensor.from_torch( + torch.tensor(cum_seq_lens_q_list, dtype=torch_idx_type), index_dtype, device + ) + max_blocks = max(len(t) for t in all_block_tables) + padded_tables = [t + [0] * (max_blocks - len(t)) for t in all_block_tables] + block_tables = TestTensor.from_torch( + torch.tensor(padded_tables, dtype=torch_idx_type), index_dtype, device + ) + + # Reference: dequantize the caches and run the attention in float32. + k_deq = k_codes.float() * k_scales.unsqueeze(-1) + v_deq = v_codes.float() * v_scales.unsqueeze(-1) + ans = ref_paged_attention_multi_turn( + q_new.torch_tensor().float(), + k_deq, + v_deq, + block_tables.torch_tensor(), + seq_lens.torch_tensor(), + cum_seq_lens_q.torch_tensor(), + scale, + ) + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreatePagedAttentionPrefillDescriptor( + handle, + ctypes.byref(descriptor), + out.descriptor, + q_new.descriptor, + k_cache.descriptor, + v_cache.descriptor, + block_tables.descriptor, + seq_lens.descriptor, + cum_seq_lens_q.descriptor, + None, + k_scale_t.descriptor, + v_scale_t.descriptor, + scale, + ) + ) + + workspace_size = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetPagedAttentionPrefillWorkspaceSize( + descriptor, ctypes.byref(workspace_size) + ) + ) + workspace = TestWorkspace(workspace_size.value, device) + + check_error( + LIBINFINIOP.infiniopPagedAttentionPrefill( + descriptor, + workspace.data(), + workspace_size.value, + out.data(), + q_new.data(), + k_cache.data(), + v_cache.data(), + block_tables.data(), + seq_lens.data(), + cum_seq_lens_q.data(), + None, + k_scale_t.data(), + v_scale_t.data(), + None, + ) + ) + if sync: + sync() + + atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype) + out_float = out.actual_tensor().float() + if DEBUG: + debug(out_float, ans, atol=atol, rtol=rtol) + assert torch.allclose(out_float, ans, atol=atol, rtol=rtol) + + check_error( + LIBINFINIOP.infiniopDestroyPagedAttentionPrefillDescriptor(descriptor) + ) + + # ============================================================================== # Main Execution # ============================================================================== @@ -341,5 +562,7 @@ def lib_attn(): for device in get_test_devices(args): test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + if device == InfiniDeviceEnum.NVIDIA: + test_operator(device, test_fp8, _TEST_CASES_FP8, _TENSOR_DTYPES_FP8) print("\033[92mTest passed!\033[0m") diff --git a/test/infiniop/paged_caching.py b/test/infiniop/paged_caching.py index 3d4405943..b374897f6 100644 --- a/test/infiniop/paged_caching.py +++ b/test/infiniop/paged_caching.py @@ -13,6 +13,7 @@ profile_operation, InfiniDtype, InfiniDtypeNames, + InfiniDeviceEnum, InfiniDeviceNames, infiniopOperatorDescriptor_t, TestWorkspace, @@ -175,6 +176,8 @@ def test( k.descriptor, v.descriptor, slot_mapping.descriptor, + None, + None, ) ) @@ -207,6 +210,8 @@ def lib_paged_caching(): v.data(), slot_mapping.data(), None, + None, + None, ) ) @@ -246,6 +251,227 @@ def lib_paged_caching(): check_error(LIBINFINIOP.infiniopDestroyPagedCachingDescriptor(descriptor)) +# ============================================================================== +# FP8 (E4M3) Cache Test +# ============================================================================== +# FP8 cases reuse the shape tuples of _TEST_CASES_; the cache pools are F8 while +# the source K/V stay in _TENSOR_DTYPES_FP8_ and per-token scales are F32. +_TENSOR_DTYPES_FP8_ = [InfiniDtype.F16, InfiniDtype.BF16] + + +def quantize_per_token_ref(x): + """ + Reference dynamic per-token-per-head FP8(E4M3) quantization, replicating the + kernel's exact float32 op order: + amax = max(|x|) over the head dim + scale = amax / 448 (scale = 1 when amax == 0) + q = e4m3_encode(x * (1 / scale)) + + NOTE: computed on CPU deliberately. On CUDA, torch lowers tensor/scalar + division to multiply-by-reciprocal, which differs from the kernel's IEEE + division by 1 ulp; at exact E4M3 rounding midpoints that flips the chosen + code and breaks bitwise comparison. CPU torch uses true division and + matches the kernel bit-for-bit. + """ + x = x.cpu() + x_float = x.float() + amax = x_float.abs().amax(dim=-1) + scale = torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)) + inv_scale = 1.0 / scale + q = (x_float * inv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) + return q, scale + + +def test_fp8( + handle, + device, + num_seqs, + max_seq_len, + num_kv_heads, + head_size, + block_size, + *tail, +): + if len(tail) == 2: + dtype, sync = tail + value_size = head_size + elif len(tail) == 3: + value_size, dtype, sync = tail + else: + raise ValueError(f"Unexpected paged_caching FP8 test arguments: {tail}") + print( + f"Testing PagedCaching FP8 on {InfiniDeviceNames[device]} with " + f"num_seqs={num_seqs}, max_seq_len={max_seq_len}, num_kv_heads={num_kv_heads}, " + f"head_size={head_size}, value_size={value_size}, block_size={block_size}, " + f"kv dtype={InfiniDtypeNames[dtype]}, cache dtype=F8" + ) + + num_blocks = 4096 + + context_lens_torch = torch.randint( + 1, max_seq_len + 1, (num_seqs,), dtype=torch.int64 + ) + ntok = torch.sum(context_lens_torch).item() + if ntok == 0: + print("Skipping test case with ntok=0") + return + + slot_mapping_list = [] + current_slot = 0 + for length in context_lens_torch: + start_slot = current_slot + slot_mapping_list.extend(range(start_slot, start_slot + length.item())) + current_slot += length.item() + assert current_slot <= num_blocks * block_size + slot_mapping_torch = torch.tensor(slot_mapping_list, dtype=torch.int64) + + # Source K/V stay in half precision; caches hold raw F8 bytes. + k = TestTensor( + (ntok, num_kv_heads, head_size), None, dtype, device, scale=4.0, bias=-2.0 + ) + v = TestTensor( + (ntok, num_kv_heads, value_size), None, dtype, device, scale=4.0, bias=-2.0 + ) + slot_mapping = TestTensor.from_torch(slot_mapping_torch, InfiniDtype.I64, device) + + k_cache_pool = TestTensor( + (num_blocks, num_kv_heads, block_size, head_size), + None, + InfiniDtype.F8, + device, + mode="float8_e4m3fn", + ) + v_cache_pool = TestTensor( + (num_blocks, num_kv_heads, block_size, value_size), + None, + InfiniDtype.F8, + device, + mode="float8_e4m3fn", + ) + k_scale_pool = TestTensor( + (num_blocks, num_kv_heads, block_size), None, InfiniDtype.F32, device, mode="zeros" + ) + v_scale_pool = TestTensor( + (num_blocks, num_kv_heads, block_size), None, InfiniDtype.F32, device, mode="zeros" + ) + + # Reference: quantize every token, then scatter codes and scales into the pools. + qk, sk = quantize_per_token_ref(k.torch_tensor()) + qv, sv = quantize_per_token_ref(v.torch_tensor()) + k_cache_ref = k_cache_pool.torch_tensor().view(torch.uint8).clone() + v_cache_ref = v_cache_pool.torch_tensor().view(torch.uint8).clone() + k_scale_ref = k_scale_pool.torch_tensor().clone() + v_scale_ref = v_scale_pool.torch_tensor().clone() + for i in range(ntok): + slot = slot_mapping_torch[i].item() + block_idx = slot // block_size + block_offset = slot % block_size + k_cache_ref[block_idx, :, block_offset, :] = qk[i].view(torch.uint8) + v_cache_ref[block_idx, :, block_offset, :] = qv[i].view(torch.uint8) + k_scale_ref[block_idx, :, block_offset] = sk[i] + v_scale_ref[block_idx, :, block_offset] = sv[i] + + if sync: + sync() + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreatePagedCachingDescriptor( + handle, + ctypes.byref(descriptor), + k_cache_pool.descriptor, + v_cache_pool.descriptor, + k.descriptor, + v.descriptor, + slot_mapping.descriptor, + k_scale_pool.descriptor, + v_scale_pool.descriptor, + ) + ) + + workspace_size = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetPagedCachingWorkspaceSize( + descriptor, ctypes.byref(workspace_size) + ) + ) + workspace = TestWorkspace(workspace_size.value, device) + + k.destroy_desc() + v.destroy_desc() + k_cache_pool.destroy_desc() + v_cache_pool.destroy_desc() + slot_mapping.destroy_desc() + k_scale_pool.destroy_desc() + v_scale_pool.destroy_desc() + + check_error( + LIBINFINIOP.infiniopPagedCaching( + descriptor, + workspace.data(), + workspace_size.value, + k_cache_pool.data(), + v_cache_pool.data(), + k.data(), + v.data(), + slot_mapping.data(), + k_scale_pool.data(), + v_scale_pool.data(), + None, + ) + ) + + if sync: + sync() + + if DEBUG: + debug( + k_cache_pool.actual_tensor().view(torch.uint8), + k_cache_ref, + atol=0, + rtol=0, + ) + debug(k_scale_pool.actual_tensor(), k_scale_ref, atol=0, rtol=1e-5) + + kc = k_cache_pool.actual_tensor().view(torch.uint8) + vc = v_cache_pool.actual_tensor().view(torch.uint8) + for name, actual, ref, src in ( + ("K", kc, k_cache_ref, k.torch_tensor()), + ("V", vc, v_cache_ref, v.torch_tensor()), + ): + m = actual != ref + if m.any(): + print(f"[DIAG] {name} cache mismatch {m.sum().item()}/{actual.numel()}") + inv = {s: t for t, s in enumerate(slot_mapping_torch.tolist())} + for idx in m.nonzero()[:6]: + b, h, off, d = [int(x) for x in idx.tolist()] + tok = inv.get(b * block_size + off, -1) + x = src[tok, h, d].float().item() if tok >= 0 else float("nan") + amax = src[tok, h].float().abs().max().item() if tok >= 0 else float("nan") + sc = amax / 448.0 if amax > 0 else 1.0 + y = x * (1.0 / sc) + print( + f"[DIAG] {name} blk{b} h{h} off{off} d{d} tok{tok} " + f"ref=0x{ref[b, h, off, d].item():02x} act=0x{actual[b, h, off, d].item():02x} " + f"x={x!r} amax={amax!r} x*inv_scale={y!r}" + ) + print( + f"[DIAG] scale maxdiff " + f"k={(k_scale_pool.actual_tensor() - k_scale_ref).abs().max().item()} " + f"v={(v_scale_pool.actual_tensor() - v_scale_ref).abs().max().item()}" + ) + assert torch.equal(k_cache_pool.actual_tensor().view(torch.uint8), k_cache_ref) + assert torch.equal(v_cache_pool.actual_tensor().view(torch.uint8), v_cache_ref) + assert torch.allclose( + k_scale_pool.actual_tensor(), k_scale_ref, atol=0, rtol=1e-5 + ) + assert torch.allclose( + v_scale_pool.actual_tensor(), v_scale_ref, atol=0, rtol=1e-5 + ) + + check_error(LIBINFINIOP.infiniopDestroyPagedCachingDescriptor(descriptor)) + + if __name__ == "__main__": args = get_args() @@ -257,5 +483,7 @@ def lib_paged_caching(): for device in get_test_devices(args): test_operator(device, test, _TEST_CASES_, _TENSOR_DTYPES) + if device == InfiniDeviceEnum.NVIDIA: + test_operator(device, test_fp8, _TEST_CASES_, _TENSOR_DTYPES_FP8_) print("\033[92mTest passed!\033[0m") diff --git a/test/infiniop/paged_caching_prefill.py b/test/infiniop/paged_caching_prefill.py index f39cd2afc..a133c07f0 100644 --- a/test/infiniop/paged_caching_prefill.py +++ b/test/infiniop/paged_caching_prefill.py @@ -173,6 +173,8 @@ def torch_caching(): k_in.descriptor, v_in.descriptor, slot_mapping.descriptor, + None, + None, ) ) @@ -196,6 +198,8 @@ def lib_caching(): v_in.data(), slot_mapping.data(), None, + None, + None, ) ) diff --git a/xmake.lua b/xmake.lua index 7ddfbf710..b181e664a 100644 --- a/xmake.lua +++ b/xmake.lua @@ -1,5 +1,4 @@ add_rules("mode.debug", "mode.release") -add_requires("boost", {configs = {stacktrace = true}}) add_requires("pybind11") -- Define color codes @@ -92,7 +91,7 @@ end option("cuda_arch") set_showmenu(true) set_description("Set CUDA GPU architecture (e.g. sm_90)") - set_values("sm_50", "sm_60", "sm_70", "sm_75", "sm_80", "sm_86", "sm_89", "sm_90", "sm_90a") + set_values("sm_50", "sm_60", "sm_70", "sm_75", "sm_80", "sm_86", "sm_89", "sm_90", "sm_90a", "sm_120") set_category("option") option_end() diff --git a/xmake/nvidia.lua b/xmake/nvidia.lua index 0d2fd0f81..1dd3ff554 100644 --- a/xmake/nvidia.lua +++ b/xmake/nvidia.lua @@ -13,9 +13,12 @@ local FLASH_ATTN_ROOT = get_config("flash-attn") local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") --- Apply -gencode from `xmake f --cuda_arch=sm_80` (comma-separated values supported). +-- Apply gencode from `xmake f --cuda_arch=sm_80` (comma-separated values supported). +-- Goes through cugencodes (not raw cuflags) so the device-link step inherits the +-- arch too: nvcc >= 13 device-links for sm_75 by default when no gencode is given, +-- which produces binaries without any SASS for the actual GPU. -- Returns true when explicit arch flags were added. -local function apply_cuda_arch_flags(add_fn) +local function apply_cuda_arch_flags(target) local arch_opt = get_config("cuda_arch") if not arch_opt or type(arch_opt) ~= "string" or arch_opt == "" then return false @@ -23,8 +26,7 @@ local function apply_cuda_arch_flags(add_fn) for _, arch in ipairs(arch_opt:split(",")) do arch = arch:trim() if arch ~= "" then - local compute = arch:gsub("sm_", "compute_") - add_fn("-gencode=arch=" .. compute .. ",code=" .. arch) + target:add("cugencodes", arch) end end return true @@ -118,30 +120,18 @@ target("infiniop-nvidia") end -- CUDA arch: explicit --cuda_arch > nvidia-smi auto-detect > native - if not apply_cuda_arch_flags(function(flag) target:add("cuflags", flag) end) then - local ok, sm_str = os.iorunv("nvidia-smi", {"--query-gpu=compute_cap", "--format=csv,noheader,nounits"}) - if ok and sm_str then - local major, minor = sm_str:match("(%d+)%.(%d+)") - if major then - local sm = tonumber(major) * 10 + tonumber(minor) - local archs = {} - if sm >= 75 then table.insert(archs, "sm_75") end - if sm >= 80 then table.insert(archs, "sm_80") end - if sm >= 86 then table.insert(archs, "sm_86") end - if sm >= 89 then table.insert(archs, "sm_89") end + if not apply_cuda_arch_flags(target) then + -- os.iorunv returns (stdout, stderr); take the first GPU's "major.minor" + local out = os.iorunv("nvidia-smi", {"--query-gpu=compute_cap", "--format=csv,noheader,nounits"}) + local sm_str = out and out:match("[^\r\n]+") + local major, minor = (sm_str or ""):match("(%d+)%.(%d+)") + if major then + local sm = tonumber(major) * 10 + tonumber(minor) + if sm == 90 then -- H100 (sm_90a): use sm_90a for cutlass 3.x - if sm == 90 then - target:add("cuflags", "-gencode=arch=compute_90a,code=sm_90a") - elseif sm > 90 then - table.insert(archs, "sm_90") - end - if #archs == 0 then - target:add("cugencodes", "native") - end - for _, arch in ipairs(archs) do - local compute = arch:gsub("sm_", "compute_") - target:add("cuflags", "-gencode=arch=" .. compute .. ",code=" .. arch) - end + target:add("cugencodes", "sm_90a") + elseif sm >= 75 then + target:add("cugencodes", "sm_" .. sm) else target:add("cugencodes", "native") end @@ -216,6 +206,12 @@ target("infiniop-nvidia") set_languages("cxx17") add_files("../src/infiniop/devices/nvidia/*.cu", "../src/infiniop/ops/*/nvidia/*.cu", "../src/infiniop/ops/*/*/nvidia/*.cu") + -- avg_pool3d's NVIDIA implementation requires cuDNN; exclude it when cudnn=n + -- to avoid mixing PyTorch's bundled cuDNN with the system CUDA runtime. + if not has_config("cudnn") then + remove_files("../src/infiniop/ops/avg_pool3d/nvidia/*.cu") + end + if has_config("ninetoothed") then add_files("../build/ninetoothed/*.c", "../build/ninetoothed/*.cpp") end @@ -226,6 +222,12 @@ target("infinirt-nvidia") add_deps("infini-utils") on_install(function (target) end) + on_load(function (target) + if not apply_cuda_arch_flags(target) then + target:add("cugencodes", "native") + end + end) + set_policy("build.cuda.devlink", true) set_toolchains("cuda") add_links("cudart") @@ -284,7 +286,7 @@ target("flash-attn-nvidia") add_links("cudart") on_load(function (target) - if not apply_cuda_arch_flags(function(flag) target:add("cuflags", flag) end) then + if not apply_cuda_arch_flags(target) then target:add("cugencodes", "native") end end)