From a045a445395f2ed52c78293f4456a204bfe531cc Mon Sep 17 00:00:00 2001 From: wooway777 Date: Tue, 28 Jul 2026 08:10:02 +0000 Subject: [PATCH 1/8] feat-infiniop-add-fp8-dsa-operators --- include/infiniop.h | 4 + include/infiniop/ops/fp8_indexer_logits.h | 33 ++ include/infiniop/ops/fp8_indexer_quant.h | 64 +++ include/infiniop/ops/fp8_mla_rmsnorm_cache.h | 32 ++ include/infiniop/ops/fp8_sparse_mla.h | 36 ++ .../fp8_indexer_logits/fp8_indexer_logits.h | 67 +++ .../nvidia/fp8_indexer_logits_nvidia.cu | 211 ++++++++ .../nvidia/fp8_indexer_logits_nvidia.cuh | 8 + .../ops/fp8_indexer_logits/operator.cc | 102 ++++ .../ops/fp8_indexer_quant/fp8_indexer_quant.h | 46 ++ .../nvidia/fp8_indexer_quant_nvidia.cu | 473 ++++++++++++++++++ .../nvidia/fp8_indexer_quant_nvidia.cuh | 76 +++ .../ops/fp8_indexer_quant/operator.cc | 200 ++++++++ .../fp8_mla_rmsnorm_cache.h | 54 ++ .../nvidia/fp8_mla_rmsnorm_cache_nvidia.cu | 187 +++++++ .../nvidia/fp8_mla_rmsnorm_cache_nvidia.cuh | 8 + .../ops/fp8_mla_rmsnorm_cache/operator.cc | 101 ++++ .../ops/fp8_sparse_mla/fp8_sparse_mla.h | 71 +++ .../nvidia/fp8_sparse_mla_nvidia.cu | 320 ++++++++++++ .../nvidia/fp8_sparse_mla_nvidia.cuh | 8 + src/infiniop/ops/fp8_sparse_mla/operator.cc | 128 +++++ 21 files changed, 2229 insertions(+) create mode 100644 include/infiniop/ops/fp8_indexer_logits.h create mode 100644 include/infiniop/ops/fp8_indexer_quant.h create mode 100644 include/infiniop/ops/fp8_mla_rmsnorm_cache.h create mode 100644 include/infiniop/ops/fp8_sparse_mla.h create mode 100644 src/infiniop/ops/fp8_indexer_logits/fp8_indexer_logits.h create mode 100644 src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu create mode 100644 src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cuh create mode 100644 src/infiniop/ops/fp8_indexer_logits/operator.cc create mode 100644 src/infiniop/ops/fp8_indexer_quant/fp8_indexer_quant.h create mode 100644 src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cu create mode 100644 src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cuh create mode 100644 src/infiniop/ops/fp8_indexer_quant/operator.cc create mode 100644 src/infiniop/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache.h create mode 100644 src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cu create mode 100644 src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cuh create mode 100644 src/infiniop/ops/fp8_mla_rmsnorm_cache/operator.cc create mode 100644 src/infiniop/ops/fp8_sparse_mla/fp8_sparse_mla.h create mode 100644 src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu create mode 100644 src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cuh create mode 100644 src/infiniop/ops/fp8_sparse_mla/operator.cc diff --git a/include/infiniop.h b/include/infiniop.h index 67c8f9bb1..70cb3cd3f 100644 --- a/include/infiniop.h +++ b/include/infiniop.h @@ -56,6 +56,10 @@ #include "infiniop/ops/floor_divide.h" #include "infiniop/ops/fmin.h" #include "infiniop/ops/fmod.h" +#include "infiniop/ops/fp8_indexer_logits.h" +#include "infiniop/ops/fp8_indexer_quant.h" +#include "infiniop/ops/fp8_mla_rmsnorm_cache.h" +#include "infiniop/ops/fp8_sparse_mla.h" #include "infiniop/ops/fused_gated_delta_net_gating.h" #include "infiniop/ops/fused_moe.h" #include "infiniop/ops/gelu.h" diff --git a/include/infiniop/ops/fp8_indexer_logits.h b/include/infiniop/ops/fp8_indexer_logits.h new file mode 100644 index 000000000..3808b0292 --- /dev/null +++ b/include/infiniop/ops/fp8_indexer_logits.h @@ -0,0 +1,33 @@ +#ifndef __INFINIOP_FP8_INDEXER_LOGITS_API_H__ +#define __INFINIOP_FP8_INDEXER_LOGITS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFp8IndexerLogitsDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFp8IndexerLogitsDescriptor( + infiniopHandle_t handle, + infiniopFp8IndexerLogitsDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t logits_desc, + infiniopTensorDescriptor_t q_fp8_desc, + infiniopTensorDescriptor_t kv_cache_desc, + infiniopTensorDescriptor_t block_tables_desc, + infiniopTensorDescriptor_t weights_fp32_desc, + infiniopTensorDescriptor_t positions_desc, + infiniopTensorDescriptor_t request_ids_desc); + +__INFINI_C __export infiniStatus_t infiniopFp8IndexerLogits( + infiniopFp8IndexerLogitsDescriptor_t desc, + void *logits, + const void *q_fp8, + const void *kv_cache, + const void *block_tables, + const void *weights_fp32, + const void *positions, + const void *request_ids, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFp8IndexerLogitsDescriptor( + infiniopFp8IndexerLogitsDescriptor_t desc); + +#endif diff --git a/include/infiniop/ops/fp8_indexer_quant.h b/include/infiniop/ops/fp8_indexer_quant.h new file mode 100644 index 000000000..101a325a9 --- /dev/null +++ b/include/infiniop/ops/fp8_indexer_quant.h @@ -0,0 +1,64 @@ +#ifndef __INFINIOP_FP8_INDEXER_QUANT_API_H__ +#define __INFINIOP_FP8_INDEXER_QUANT_API_H__ + +#include "../operator_descriptor.h" + +#include + +typedef struct InfiniopDescriptor *infiniopFp8IndexerQuantDescriptor_t; +typedef struct InfiniopDescriptor *infiniopFusedFp8IndexerDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFp8IndexerQuantDescriptor( + infiniopHandle_t handle, + infiniopFp8IndexerQuantDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t q_fp8_desc, + infiniopTensorDescriptor_t weights_fp32_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t weights_desc); + +__INFINI_C __export infiniStatus_t infiniopFp8IndexerQuant( + infiniopFp8IndexerQuantDescriptor_t desc, + void *q_fp8, + void *weights_fp32, + const void *q, + const void *weights, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFp8IndexerQuantDescriptor( + infiniopFp8IndexerQuantDescriptor_t desc); + +__INFINI_C __export infiniStatus_t infiniopCreateFusedFp8IndexerDescriptor( + infiniopHandle_t handle, + infiniopFusedFp8IndexerDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t q_fp8_desc, + infiniopTensorDescriptor_t weights_fp32_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t q_raw_desc, + infiniopTensorDescriptor_t k_weights_desc, + infiniopTensorDescriptor_t norm_weight_desc, + infiniopTensorDescriptor_t norm_bias_desc, + infiniopTensorDescriptor_t positions_desc, + infiniopTensorDescriptor_t cos_sin_cache_desc, + infiniopTensorDescriptor_t slot_mapping_desc, + uint64_t rope_dim, + double eps, + double weights_scale); + +__INFINI_C __export infiniStatus_t infiniopFusedFp8Indexer( + infiniopFusedFp8IndexerDescriptor_t desc, + void *q_fp8, + void *weights_fp32, + void *k_cache, + const void *q_raw, + const void *k_weights, + const void *norm_weight, + const void *norm_bias, + const void *positions, + const void *cos_sin_cache, + const void *slot_mapping, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFusedFp8IndexerDescriptor( + infiniopFusedFp8IndexerDescriptor_t desc); + +#endif diff --git a/include/infiniop/ops/fp8_mla_rmsnorm_cache.h b/include/infiniop/ops/fp8_mla_rmsnorm_cache.h new file mode 100644 index 000000000..158e3ef94 --- /dev/null +++ b/include/infiniop/ops/fp8_mla_rmsnorm_cache.h @@ -0,0 +1,32 @@ +#ifndef __INFINIOP_FP8_MLA_RMSNORM_CACHE_API_H__ +#define __INFINIOP_FP8_MLA_RMSNORM_CACHE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFp8MlaRmsnormCacheDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFp8MlaRmsnormCacheDescriptor( + infiniopHandle_t handle, + infiniopFp8MlaRmsnormCacheDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t cache_desc, + infiniopTensorDescriptor_t vendor_cache_desc, + infiniopTensorDescriptor_t compressed_kv_desc, + infiniopTensorDescriptor_t norm_weight_desc, + infiniopTensorDescriptor_t rope_desc, + infiniopTensorDescriptor_t slot_mapping_desc, + double eps); + +__INFINI_C __export infiniStatus_t infiniopFp8MlaRmsnormCache( + infiniopFp8MlaRmsnormCacheDescriptor_t desc, + void *cache, + void *vendor_cache, + const void *compressed_kv, + const void *norm_weight, + const void *rope, + const void *slot_mapping, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFp8MlaRmsnormCacheDescriptor( + infiniopFp8MlaRmsnormCacheDescriptor_t desc); + +#endif diff --git a/include/infiniop/ops/fp8_sparse_mla.h b/include/infiniop/ops/fp8_sparse_mla.h new file mode 100644 index 000000000..3d199a614 --- /dev/null +++ b/include/infiniop/ops/fp8_sparse_mla.h @@ -0,0 +1,36 @@ +#ifndef __INFINIOP_FP8_SPARSE_MLA_API_H__ +#define __INFINIOP_FP8_SPARSE_MLA_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFp8SparseMlaDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFp8SparseMlaDescriptor( + infiniopHandle_t handle, + infiniopFp8SparseMlaDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t query_desc, + infiniopTensorDescriptor_t kv_cache_desc, + infiniopTensorDescriptor_t indices_desc, + infiniopTensorDescriptor_t topk_lens_desc, + float scale); + +__INFINI_C __export infiniStatus_t infiniopGetFp8SparseMlaWorkspaceSize( + infiniopFp8SparseMlaDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopFp8SparseMla( + infiniopFp8SparseMlaDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *query, + const void *kv_cache, + const void *indices, + const void *topk_lens, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFp8SparseMlaDescriptor( + infiniopFp8SparseMlaDescriptor_t desc); + +#endif diff --git a/src/infiniop/ops/fp8_indexer_logits/fp8_indexer_logits.h b/src/infiniop/ops/fp8_indexer_logits/fp8_indexer_logits.h new file mode 100644 index 000000000..b8f34edf2 --- /dev/null +++ b/src/infiniop/ops/fp8_indexer_logits/fp8_indexer_logits.h @@ -0,0 +1,67 @@ +#ifndef __FP8_INDEXER_LOGITS_H__ +#define __FP8_INDEXER_LOGITS_H__ + +#include "../../../utils.h" +#include "../../operator.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::fp8_indexer_logits::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + size_t _num_tokens; \ + size_t _num_heads; \ + size_t _head_dim; \ + size_t _num_cache_blocks; \ + size_t _block_size; \ + size_t _cache_stride; \ + size_t _num_requests; \ + size_t _max_blocks_per_request; \ + size_t _max_context_len; \ + \ + Descriptor( \ + size_t num_tokens, \ + size_t num_heads, \ + size_t head_dim, \ + size_t num_cache_blocks, \ + size_t block_size, \ + size_t cache_stride, \ + size_t num_requests, \ + size_t max_blocks_per_request, \ + size_t max_context_len, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _num_tokens(num_tokens), \ + _num_heads(num_heads), \ + _head_dim(head_dim), \ + _num_cache_blocks(num_cache_blocks), \ + _block_size(block_size), \ + _cache_stride(cache_stride), \ + _num_requests(num_requests), \ + _max_blocks_per_request(max_blocks_per_request), \ + _max_context_len(max_context_len) {} \ + \ + public: \ + static infiniStatus_t create( \ + infiniopHandle_t handle, \ + Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t logits_desc, \ + infiniopTensorDescriptor_t q_fp8_desc, \ + infiniopTensorDescriptor_t kv_cache_desc, \ + infiniopTensorDescriptor_t block_tables_desc, \ + infiniopTensorDescriptor_t weights_fp32_desc, \ + infiniopTensorDescriptor_t positions_desc, \ + infiniopTensorDescriptor_t request_ids_desc); \ + \ + infiniStatus_t calculate( \ + void *logits, \ + const void *q_fp8, \ + const void *kv_cache, \ + const void *block_tables, \ + const void *weights_fp32, \ + const void *positions, \ + const void *request_ids, \ + void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu b/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu new file mode 100644 index 000000000..740e9d887 --- /dev/null +++ b/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu @@ -0,0 +1,211 @@ +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "../../../tensor.h" +#include "fp8_indexer_logits_nvidia.cuh" + +#include +#include + +namespace { +constexpr size_t THREADS = 256; +template +INFINIOP_CUDA_KERNEL fp8IndexerLogitsKernel( + float *__restrict__ logits, + const cuda_fp8_e4m3 *__restrict__ q_fp8, + const uint8_t *__restrict__ kv_cache, + const int32_t *__restrict__ block_tables, + const float *__restrict__ weights_fp32, + const int64_t *__restrict__ positions, + const int32_t *__restrict__ request_ids, + size_t num_heads, + size_t head_dim, + size_t num_cache_blocks, + size_t block_size, + size_t cache_stride, + size_t num_requests, + size_t max_blocks_per_request, + size_t max_context_len) { + const size_t token = blockIdx.x; + constexpr size_t KEYS_PER_TILE = THREADS / LANES_PER_KEY; + const size_t tiles_per_block = (block_size + KEYS_PER_TILE - 1) / KEYS_PER_TILE; + const size_t tile_in_block = blockIdx.y % tiles_per_block; + const size_t logical_block = blockIdx.y / tiles_per_block; + const size_t lane = threadIdx.x % LANES_PER_KEY; + const size_t key_in_block = tile_in_block * KEYS_PER_TILE + + threadIdx.x / LANES_PER_KEY; + + extern __shared__ uint8_t shared_bytes[]; + auto *shared_q = reinterpret_cast(shared_bytes); + auto *shared_weights = reinterpret_cast( + shared_bytes + num_heads * head_dim); + const size_t q_base = token * num_heads * head_dim; + for (size_t i = threadIdx.x; i < num_heads * head_dim; i += blockDim.x) { + shared_q[i] = q_fp8[q_base + i]; + } + for (size_t i = threadIdx.x; i < num_heads; i += blockDim.x) { + shared_weights[i] = weights_fp32[token * num_heads + i]; + } + __syncthreads(); + + const int64_t position = positions[token]; + const int32_t request = request_ids[token]; + const size_t key_position = logical_block * block_size + key_in_block; + bool valid = key_in_block < block_size + && key_position < max_context_len + && position >= 0 + && key_position <= static_cast(position) + && request >= 0 + && static_cast(request) < num_requests + && logical_block < max_blocks_per_request; + int32_t physical_block = -1; + if (valid) { + physical_block = block_tables[static_cast(request) * max_blocks_per_request + logical_block]; + valid = physical_block >= 0 + && static_cast(physical_block) < num_cache_blocks; + } + + const uint8_t *cache_block = valid + ? kv_cache + + static_cast(physical_block) * block_size * cache_stride + : nullptr; + const auto *key = valid + ? reinterpret_cast( + cache_block + key_in_block * head_dim) + : nullptr; + const float key_scale = valid + ? *reinterpret_cast( + cache_block + block_size * head_dim + + key_in_block * sizeof(float)) + : 0.0f; + float acc = 0.0f; + for (size_t head = 0; head < num_heads; ++head) { + float dot = 0.0f; + if (valid) { + const auto *query = shared_q + head * head_dim; + for (size_t column = lane; column < head_dim; column += LANES_PER_KEY) { + dot += static_cast(query[column]) + * static_cast(key[column]); + } + } + // Every thread named by the full-warp mask must execute the shuffle, + // including invalid keys in a partial tile and graph-padding requests. + for (size_t offset = 1; offset < LANES_PER_KEY; offset <<= 1) { + dot += __shfl_xor_sync(0xffffffffu, dot, offset, LANES_PER_KEY); + } + if (valid && lane == 0) { + acc += fmaxf(dot * key_scale, 0.0f) * shared_weights[head]; + } + } + if (lane == 0 && key_position < max_context_len) { + logits[token * max_context_len + key_position] = valid + ? acc + : -CUDART_INF_F; + } +} +} // namespace + +namespace op::fp8_indexer_logits::nvidia { + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t logits_desc, + infiniopTensorDescriptor_t q_fp8_desc, + infiniopTensorDescriptor_t kv_cache_desc, + infiniopTensorDescriptor_t block_tables_desc, + infiniopTensorDescriptor_t weights_fp32_desc, + infiniopTensorDescriptor_t positions_desc, + infiniopTensorDescriptor_t request_ids_desc) { + const auto logits_shape = logits_desc->shape(); + const auto q_shape = q_fp8_desc->shape(); + const auto cache_shape = kv_cache_desc->shape(); + const auto blocks_shape = block_tables_desc->shape(); + const auto weights_shape = weights_fp32_desc->shape(); + CHECK_OR_RETURN(logits_shape.size() == 2 && q_shape.size() == 3 + && cache_shape.size() == 3 && blocks_shape.size() == 2 + && weights_shape.size() == 2 + && positions_desc->shape().size() == 1 + && request_ids_desc->shape().size() == 1, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(logits_shape[0] == q_shape[0] + && weights_shape[0] == q_shape[0] + && weights_shape[1] == q_shape[1] + && positions_desc->shape()[0] == q_shape[0] + && request_ids_desc->shape()[0] == q_shape[0] + && cache_shape[1] == 64 + && q_shape[2] == 128 + && cache_shape[2] == q_shape[2] + sizeof(float), + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(logits_desc->isContiguous() && q_fp8_desc->isContiguous() + && kv_cache_desc->isContiguous() + && block_tables_desc->isContiguous() + && weights_fp32_desc->isContiguous() + && positions_desc->isContiguous() + && request_ids_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + CHECK_OR_RETURN(logits_desc->dtype() == INFINI_DTYPE_F32 + && q_fp8_desc->dtype() == INFINI_DTYPE_F8 + && kv_cache_desc->dtype() == INFINI_DTYPE_U8 + && block_tables_desc->dtype() == INFINI_DTYPE_I32 + && weights_fp32_desc->dtype() == INFINI_DTYPE_F32 + && positions_desc->dtype() == INFINI_DTYPE_I64 + && request_ids_desc->dtype() == INFINI_DTYPE_I32, + INFINI_STATUS_BAD_TENSOR_DTYPE); + *desc_ptr = new Descriptor( + q_shape[0], q_shape[1], q_shape[2], cache_shape[0], cache_shape[1], + cache_shape[2], blocks_shape[0], blocks_shape[1], logits_shape[1], + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *logits, + const void *q_fp8, + const void *kv_cache, + const void *block_tables, + const void *weights_fp32, + const void *positions, + const void *request_ids, + void *stream) const { + // Pure decode has one token per active request. Use wide key tiling for + // decode batches up to eight while keeping prefill and mixed batches narrow. + const bool use_wide_decode = _num_tokens == _num_requests + && _num_tokens <= 8; + const size_t lanes_per_key = use_wide_decode ? 8 : 4; + const size_t keys_per_tile = THREADS / lanes_per_key; + const size_t tiles_per_block = (_block_size + keys_per_tile - 1) / keys_per_tile; + const dim3 grid( + static_cast(_num_tokens), + static_cast((_max_context_len + _block_size - 1) / _block_size * tiles_per_block)); + const size_t smem = _num_heads * _head_dim + _num_heads * sizeof(float); + if (use_wide_decode) { + fp8IndexerLogitsKernel<8><<< + grid, THREADS, smem, reinterpret_cast(stream)>>>( + reinterpret_cast(logits), + reinterpret_cast(q_fp8), + reinterpret_cast(kv_cache), + reinterpret_cast(block_tables), + reinterpret_cast(weights_fp32), + reinterpret_cast(positions), + reinterpret_cast(request_ids), + _num_heads, _head_dim, _num_cache_blocks, _block_size, _cache_stride, + _num_requests, _max_blocks_per_request, _max_context_len); + } else { + fp8IndexerLogitsKernel<4><<< + grid, THREADS, smem, reinterpret_cast(stream)>>>( + reinterpret_cast(logits), + reinterpret_cast(q_fp8), + reinterpret_cast(kv_cache), + reinterpret_cast(block_tables), + reinterpret_cast(weights_fp32), + reinterpret_cast(positions), + reinterpret_cast(request_ids), + _num_heads, _head_dim, _num_cache_blocks, _block_size, _cache_stride, + _num_requests, _max_blocks_per_request, _max_context_len); + } + return cudaGetLastError() == cudaSuccess + ? INFINI_STATUS_SUCCESS + : INFINI_STATUS_INTERNAL_ERROR; +} + +} // namespace op::fp8_indexer_logits::nvidia diff --git a/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cuh b/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cuh new file mode 100644 index 000000000..e11b79ce2 --- /dev/null +++ b/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __FP8_INDEXER_LOGITS_NVIDIA_H__ +#define __FP8_INDEXER_LOGITS_NVIDIA_H__ + +#include "../fp8_indexer_logits.h" + +DESCRIPTOR(nvidia) + +#endif diff --git a/src/infiniop/ops/fp8_indexer_logits/operator.cc b/src/infiniop/ops/fp8_indexer_logits/operator.cc new file mode 100644 index 000000000..2448832dd --- /dev/null +++ b/src/infiniop/ops/fp8_indexer_logits/operator.cc @@ -0,0 +1,102 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/fp8_indexer_logits.h" + +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) +#include "nvidia/fp8_indexer_logits_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateFp8IndexerLogitsDescriptor( + infiniopHandle_t handle, + infiniopFp8IndexerLogitsDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t logits_desc, + infiniopTensorDescriptor_t q_fp8_desc, + infiniopTensorDescriptor_t kv_cache_desc, + infiniopTensorDescriptor_t block_tables_desc, + infiniopTensorDescriptor_t weights_fp32_desc, + infiniopTensorDescriptor_t positions_desc, + infiniopTensorDescriptor_t request_ids_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::fp8_indexer_logits::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + logits_desc, q_fp8_desc, kv_cache_desc, block_tables_desc, \ + weights_fp32_desc, positions_desc, request_ids_desc) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopFp8IndexerLogits( + infiniopFp8IndexerLogitsDescriptor_t desc, + void *logits, + const void *q_fp8, + const void *kv_cache, + const void *block_tables, + const void *weights_fp32, + const void *positions, + const void *request_ids, + void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(logits, q_fp8, kv_cache, block_tables, weights_fp32, \ + positions, request_ids, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyFp8IndexerLogitsDescriptor( + infiniopFp8IndexerLogitsDescriptor_t desc) { +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/fp8_indexer_quant/fp8_indexer_quant.h b/src/infiniop/ops/fp8_indexer_quant/fp8_indexer_quant.h new file mode 100644 index 000000000..c7f9d9462 --- /dev/null +++ b/src/infiniop/ops/fp8_indexer_quant/fp8_indexer_quant.h @@ -0,0 +1,46 @@ +#ifndef __FP8_INDEXER_QUANT_H__ +#define __FP8_INDEXER_QUANT_H__ + +#include "../../../utils.h" +#include "../../operator.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::fp8_indexer_quant::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + size_t _num_groups; \ + size_t _head_dim; \ + size_t _threads; \ + infiniDtype_t _input_dtype; \ + \ + Descriptor( \ + size_t num_groups, \ + size_t head_dim, \ + size_t threads, \ + infiniDtype_t input_dtype, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _num_groups(num_groups), \ + _head_dim(head_dim), \ + _threads(threads), \ + _input_dtype(input_dtype) {} \ + \ + public: \ + static infiniStatus_t create( \ + infiniopHandle_t handle, \ + Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t q_fp8_desc, \ + infiniopTensorDescriptor_t weights_fp32_desc, \ + infiniopTensorDescriptor_t q_desc, \ + infiniopTensorDescriptor_t weights_desc); \ + \ + infiniStatus_t calculate( \ + void *q_fp8, \ + void *weights_fp32, \ + const void *q, \ + const void *weights, \ + void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cu b/src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cu new file mode 100644 index 000000000..a6a7bad74 --- /dev/null +++ b/src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cu @@ -0,0 +1,473 @@ +#include "../../../../utils.h" +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "../../../tensor.h" +#include "fp8_indexer_quant_nvidia.cuh" + +#include +#include +#include + +namespace { +template +__device__ __forceinline__ float to_float(T value); + +template <> +__device__ __forceinline__ float to_float(half value) { + return __half2float(value); +} + +template <> +__device__ __forceinline__ float to_float(cuda_bfloat16 value) { + return __bfloat162float(value); +} + +template +__device__ __forceinline__ T from_float(float value); + +template <> +__device__ __forceinline__ half from_float(float value) { + return __float2half_rn(value); +} + +template <> +__device__ __forceinline__ cuda_bfloat16 from_float(float value) { + return __float2bfloat16_rn(value); +} + +__device__ __forceinline__ float vendor_fp8_quotient(float value, float scale) { + const float reciprocal = __frcp_rn(scale); + const float quotient = value * reciprocal; + const float residual = fmaf(-scale, quotient, value); + return fmaf(residual, reciprocal, quotient); +} + +__device__ __forceinline__ uint8_t vendor_fp8_e4m3(float value) { + const uint32_t bits = __float_as_uint(value); + const uint32_t abs_bits = bits & 0x7fffffffU; + uint32_t magnitude = 0x7fU; + if (abs_bits < 0x43f00000U) { + if (abs_bits > 0x3c7fffffU) { + const uint32_t tie = (bits >> 20U) & 1U; + magnitude = (bits + tie + 0x0407ffffU) >> 20U; + } else { + const float absolute = __uint_as_float(abs_bits); + magnitude = __float_as_uint(absolute + 16384.0f); + } + } + const uint32_t sign = (bits >> 24U) & 0x80U; + return static_cast((magnitude & 0x7fU) | sign); +} + +template +INFINIOP_CUDA_KERNEL fp8IndexerQuantKernel( + cuda_fp8_e4m3 *__restrict__ q_fp8, + float *__restrict__ weights_fp32, + const T *__restrict__ q, + const T *__restrict__ weights, + size_t head_dim) { + const size_t group = blockIdx.x; + const size_t column = threadIdx.x; + const float value = column < head_dim + ? to_float(q[group * head_dim + column]) + : 0.0f; + + extern __shared__ float reduction[]; + reduction[column] = fabsf(value); + __syncthreads(); + for (size_t stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (column < stride) { + reduction[column] = fmaxf(reduction[column], reduction[column + stride]); + } + __syncthreads(); + } + + if (column == 0) { + const float abs_max = fmaxf(reduction[0], 1.0e-4f); + const float scale_raw = abs_max / 448.0f; + reduction[0] = exp2f(ceilf(log2f(scale_raw))); + weights_fp32[group] = to_float(weights[group]) * reduction[0]; + } + __syncthreads(); + + if (column < head_dim) { + const float quantized = fminf(448.0f, fmaxf(-448.0f, value / reduction[0])); + q_fp8[group * head_dim + column] = cuda_fp8_e4m3(quantized); + } +} + +template +INFINIOP_CUDA_KERNEL fusedFp8IndexerKernel( + cuda_fp8_e4m3 *__restrict__ q_fp8, + float *__restrict__ weights_fp32, + uint8_t *__restrict__ k_cache, + const T *__restrict__ q_raw, + const T *__restrict__ k_weights, + const T *__restrict__ norm_weight, + const T *__restrict__ norm_bias, + const int64_t *__restrict__ positions, + const T *__restrict__ cos_sin_cache, + const int64_t *__restrict__ slot_mapping, + size_t num_heads, + size_t head_dim, + size_t rope_dim, + size_t num_cache_blocks, + size_t block_size, + size_t cache_stride, + size_t max_positions, + float eps, + float weights_scale) { + const size_t work_per_token = num_heads + 1; + const size_t token = blockIdx.x / work_per_token; + const size_t work = blockIdx.x % work_per_token; + const size_t column = threadIdx.x; + const int64_t position_raw = positions[token]; + const size_t position = position_raw >= 0 + && static_cast(position_raw) < max_positions + ? static_cast(position_raw) + : 0; + const size_t half_rope_dim = rope_dim / 2; + + extern __shared__ float shared[]; + float *values = shared; + float *scratch = shared + head_dim; + + if (work < num_heads) { + const size_t group = token * num_heads + work; + const size_t q_base = group * head_dim; + float value = to_float(q_raw[q_base + column]); + if (column < rope_dim) { + const size_t pair = column & ~size_t(1); + const float x0 = to_float(q_raw[q_base + pair]); + const float x1 = to_float(q_raw[q_base + pair + 1]); + const size_t angle = column / 2; + const size_t cache_base = position * rope_dim; + const float cosine = to_float(cos_sin_cache[cache_base + angle]); + const float sine = to_float( + cos_sin_cache[cache_base + half_rope_dim + angle]); + value = (column & 1) != 0 + ? x0 * sine + x1 * cosine + : x0 * cosine - x1 * sine; + } + value = to_float(from_float(value)); + scratch[column] = fabsf(value); + __syncthreads(); + for (size_t stride = head_dim / 2; stride > 0; stride >>= 1) { + if (column < stride) { + scratch[column] = fmaxf(scratch[column], scratch[column + stride]); + } + __syncthreads(); + } + if (column == 0) { + const float abs_max = fmaxf(scratch[0], 1.0e-4f); + const float scale_raw = abs_max / 448.0f; + scratch[0] = exp2f(ceilf(log2f(scale_raw))); + const T scaled_weight = from_float( + to_float(k_weights[token * (head_dim + num_heads) + + head_dim + work]) + * weights_scale); + weights_fp32[group] = to_float(scaled_weight) * scratch[0]; + } + __syncthreads(); + const float quantized = fminf( + 448.0f, fmaxf(-448.0f, value / scratch[0])); + q_fp8[q_base + column] = cuda_fp8_e4m3(quantized); + return; + } + + const int64_t slot_raw = slot_mapping[token]; + if (slot_raw < 0 + || static_cast(slot_raw) >= num_cache_blocks * block_size) { + return; + } + const size_t kw_base = token * (head_dim + num_heads); + const size_t half_head_dim = head_dim / 2; + const bool active_lane = column < half_head_dim; + const size_t column0 = column * 2; + const size_t column1 = column0 + 1; + float input0 = 0.0f; + float input1 = 0.0f; + if (active_lane) { + input0 = to_float(k_weights[kw_base + column0]); + input1 = to_float(k_weights[kw_base + column1]); + // Match the vendor fused kernel: one 64-lane wave owns an adjacent + // BF16 pair and accumulates x0^2 with a fused multiply-add. + values[column] = input0 + input1; + scratch[column] = fmaf(input0, input0, input1 * input1); + } + __syncthreads(); + for (size_t stride = half_head_dim / 2; stride > 0; stride >>= 1) { + if (column < stride) { + values[column] += values[column + stride]; + scratch[column] += scratch[column + stride]; + } + __syncthreads(); + } + const float inv_head_dim = 1.0f / static_cast(head_dim); + const float mean = values[0] * inv_head_dim; + const float variance = fmaf(scratch[0], inv_head_dim, -mean * mean); + const float inv_std = rsqrtf(variance + eps); + if (active_lane) { + const float normalized0 = fmaf( + (input0 - mean) * inv_std, + to_float(norm_weight[column0]), + to_float(norm_bias[column0])); + const float normalized1 = fmaf( + (input1 - mean) * inv_std, + to_float(norm_weight[column1]), + to_float(norm_bias[column1])); + values[column0] = normalized0; + values[column1] = normalized1; + } + __syncthreads(); + + float value0 = 0.0f; + float value1 = 0.0f; + if (active_lane) { + value0 = values[column0]; + value1 = values[column1]; + } + if (active_lane && column0 < rope_dim) { + const float x0 = value0; + const float x1 = value1; + const size_t angle = column; + const size_t cache_base = position * rope_dim; + const float cosine = to_float(cos_sin_cache[cache_base + angle]); + const float sine = to_float( + cos_sin_cache[cache_base + half_rope_dim + angle]); + const float x1_sine = x1 * sine; + const float x0_sine = x0 * sine; + value0 = fmaf(x0, cosine, -x1_sine); + value1 = fmaf(x1, cosine, x0_sine); + } + if (active_lane) { + value0 = to_float(from_float(value0)); + value1 = to_float(from_float(value1)); + } + if (active_lane) { + scratch[column] = fmaxf(fabsf(value0), fabsf(value1)); + } + __syncthreads(); + for (size_t stride = half_head_dim / 2; stride > 0; stride >>= 1) { + if (column < stride) { + scratch[column] = fmaxf(scratch[column], scratch[column + stride]); + } + __syncthreads(); + } + if (column == 0) { + const float abs_max = fmaxf(scratch[0], 1.0e-4f); + const float scale_raw = abs_max / 448.0f; + scratch[0] = exp2f(ceilf(log2f(scale_raw))); + } + __syncthreads(); + + const size_t slot = static_cast(slot_raw); + const size_t cache_block = slot / block_size; + const size_t cache_offset = slot % block_size; + uint8_t *cache_base = k_cache + + cache_block * block_size * cache_stride; + if (active_lane) { + uint8_t *cache_values = cache_base + cache_offset * head_dim; + const float quantized0 = fminf( + 448.0f, fmaxf(-448.0f, vendor_fp8_quotient(value0, scratch[0]))); + const float quantized1 = fminf( + 448.0f, fmaxf(-448.0f, vendor_fp8_quotient(value1, scratch[0]))); + cache_values[column0] = vendor_fp8_e4m3(quantized0); + cache_values[column1] = vendor_fp8_e4m3(quantized1); + } + if (column == 0) { + *reinterpret_cast( + cache_base + block_size * head_dim + + cache_offset * sizeof(float)) + = scratch[0]; + } +} +} // namespace + +namespace op::fp8_indexer_quant::nvidia { + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t q_fp8_desc, + infiniopTensorDescriptor_t weights_fp32_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t weights_desc) { + const auto q_shape = q_desc->shape(); + const auto weights_shape = weights_desc->shape(); + CHECK_OR_RETURN(q_shape.size() == 3 && weights_shape.size() == 2, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(q_fp8_desc->shape() == q_shape + && weights_fp32_desc->shape() == weights_shape + && weights_shape[0] == q_shape[0] + && weights_shape[1] == q_shape[1], + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(q_desc->isContiguous() && weights_desc->isContiguous() + && q_fp8_desc->isContiguous() + && weights_fp32_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + const auto input_dtype = q_desc->dtype(); + CHECK_OR_RETURN(input_dtype == INFINI_DTYPE_F16 + || input_dtype == INFINI_DTYPE_BF16, + INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(weights_desc->dtype() == input_dtype + && q_fp8_desc->dtype() == INFINI_DTYPE_F8 + && weights_fp32_desc->dtype() == INFINI_DTYPE_F32, + INFINI_STATUS_BAD_TENSOR_DTYPE); + const size_t head_dim = q_shape[2]; + CHECK_OR_RETURN(head_dim > 0 && head_dim <= 1024, + INFINI_STATUS_BAD_TENSOR_SHAPE); + size_t threads = 1; + while (threads < head_dim) { + threads <<= 1; + } + *desc_ptr = new Descriptor( + q_shape[0] * q_shape[1], head_dim, threads, input_dtype, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *q_fp8, + void *weights_fp32, + const void *q, + const void *weights, + void *stream) const { + const auto cuda_stream = reinterpret_cast(stream); + const size_t smem = _threads * sizeof(float); + if (_input_dtype == INFINI_DTYPE_F16) { + fp8IndexerQuantKernel<<<_num_groups, _threads, smem, cuda_stream>>>( + reinterpret_cast(q_fp8), + reinterpret_cast(weights_fp32), + reinterpret_cast(q), + reinterpret_cast(weights), + _head_dim); + } else { + fp8IndexerQuantKernel<<<_num_groups, _threads, smem, cuda_stream>>>( + reinterpret_cast(q_fp8), + reinterpret_cast(weights_fp32), + reinterpret_cast(q), + reinterpret_cast(weights), + _head_dim); + } + return cudaGetLastError() == cudaSuccess + ? INFINI_STATUS_SUCCESS + : INFINI_STATUS_INTERNAL_ERROR; +} + +infiniStatus_t FusedDescriptor::create( + infiniopHandle_t handle, + FusedDescriptor **desc_ptr, + infiniopTensorDescriptor_t q_fp8_desc, + infiniopTensorDescriptor_t weights_fp32_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t q_raw_desc, + infiniopTensorDescriptor_t k_weights_desc, + infiniopTensorDescriptor_t norm_weight_desc, + infiniopTensorDescriptor_t norm_bias_desc, + infiniopTensorDescriptor_t positions_desc, + infiniopTensorDescriptor_t cos_sin_cache_desc, + infiniopTensorDescriptor_t slot_mapping_desc, + uint64_t rope_dim, + double eps, + double weights_scale) { + const auto q_shape = q_raw_desc->shape(); + const auto weights_shape = weights_fp32_desc->shape(); + const auto kw_shape = k_weights_desc->shape(); + const auto cache_shape = k_cache_desc->shape(); + const auto cos_sin_shape = cos_sin_cache_desc->shape(); + CHECK_OR_RETURN(q_shape.size() == 3 && weights_shape.size() == 2 + && kw_shape.size() == 2 && cache_shape.size() == 3 + && norm_weight_desc->shape().size() == 1 + && norm_bias_desc->shape().size() == 1 + && positions_desc->shape().size() == 1 + && cos_sin_shape.size() == 2 + && slot_mapping_desc->shape().size() == 1, + INFINI_STATUS_BAD_TENSOR_SHAPE); + const size_t num_tokens = q_shape[0]; + const size_t num_heads = q_shape[1]; + const size_t head_dim = q_shape[2]; + CHECK_OR_RETURN(head_dim == 128 && rope_dim > 0 && rope_dim <= head_dim + && rope_dim % 2 == 0 + && q_fp8_desc->shape() == q_shape + && weights_shape[0] == num_tokens + && weights_shape[1] == num_heads + && kw_shape[0] == num_tokens + && kw_shape[1] == head_dim + num_heads + && norm_weight_desc->shape()[0] == head_dim + && norm_bias_desc->shape()[0] == head_dim + && positions_desc->shape()[0] == num_tokens + && slot_mapping_desc->shape()[0] == num_tokens + && cos_sin_shape[1] == rope_dim + && cache_shape[1] > 0 + && cache_shape[2] == head_dim + sizeof(float), + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(q_fp8_desc->isContiguous() + && weights_fp32_desc->isContiguous() + && k_cache_desc->isContiguous() + && q_raw_desc->isContiguous() + && k_weights_desc->isContiguous() + && norm_weight_desc->isContiguous() + && norm_bias_desc->isContiguous() + && positions_desc->isContiguous() + && cos_sin_cache_desc->isContiguous() + && slot_mapping_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + const auto input_dtype = q_raw_desc->dtype(); + CHECK_OR_RETURN(input_dtype == INFINI_DTYPE_F16 + || input_dtype == INFINI_DTYPE_BF16, + INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(k_weights_desc->dtype() == input_dtype + && norm_weight_desc->dtype() == input_dtype + && norm_bias_desc->dtype() == input_dtype + && cos_sin_cache_desc->dtype() == input_dtype + && q_fp8_desc->dtype() == INFINI_DTYPE_F8 + && weights_fp32_desc->dtype() == INFINI_DTYPE_F32 + && k_cache_desc->dtype() == INFINI_DTYPE_U8 + && positions_desc->dtype() == INFINI_DTYPE_I64 + && slot_mapping_desc->dtype() == INFINI_DTYPE_I64, + INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(eps > 0.0 && weights_scale > 0.0, + INFINI_STATUS_BAD_PARAM); + *desc_ptr = new FusedDescriptor( + num_tokens, num_heads, head_dim, rope_dim, + cache_shape[0], cache_shape[1], cache_shape[2], cos_sin_shape[0], + input_dtype, eps, weights_scale, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t FusedDescriptor::calculate( + void *q_fp8, void *weights_fp32, void *k_cache, + const void *q_raw, const void *k_weights, + const void *norm_weight, const void *norm_bias, + const void *positions, const void *cos_sin_cache, + const void *slot_mapping, void *stream) const { + const size_t blocks = _num_tokens * (_num_heads + 1); + const size_t smem = 2 * _head_dim * sizeof(float); +#define LAUNCH_FUSED(T) \ + fusedFp8IndexerKernel<<(stream)>>>( \ + reinterpret_cast(q_fp8), \ + reinterpret_cast(weights_fp32), \ + reinterpret_cast(k_cache), \ + reinterpret_cast(q_raw), \ + reinterpret_cast(k_weights), \ + reinterpret_cast(norm_weight), \ + reinterpret_cast(norm_bias), \ + reinterpret_cast(positions), \ + reinterpret_cast(cos_sin_cache), \ + reinterpret_cast(slot_mapping), \ + _num_heads, _head_dim, _rope_dim, _num_cache_blocks, _block_size, \ + _cache_stride, _max_positions, _eps, _weights_scale) + if (_input_dtype == INFINI_DTYPE_F16) { + LAUNCH_FUSED(half); + } else { + LAUNCH_FUSED(cuda_bfloat16); + } +#undef LAUNCH_FUSED + return cudaGetLastError() == cudaSuccess + ? INFINI_STATUS_SUCCESS + : INFINI_STATUS_INTERNAL_ERROR; +} + +} // namespace op::fp8_indexer_quant::nvidia diff --git a/src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cuh b/src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cuh new file mode 100644 index 000000000..9212ccd9a --- /dev/null +++ b/src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cuh @@ -0,0 +1,76 @@ +#ifndef __FP8_INDEXER_QUANT_NVIDIA_H__ +#define __FP8_INDEXER_QUANT_NVIDIA_H__ + +#include "../fp8_indexer_quant.h" + +DESCRIPTOR(nvidia) + +namespace op::fp8_indexer_quant::nvidia { +class FusedDescriptor final : public InfiniopDescriptor { + size_t _num_tokens; + size_t _num_heads; + size_t _head_dim; + size_t _rope_dim; + size_t _num_cache_blocks; + size_t _block_size; + size_t _cache_stride; + size_t _max_positions; + infiniDtype_t _input_dtype; + float _eps; + float _weights_scale; + + FusedDescriptor( + size_t num_tokens, + size_t num_heads, + size_t head_dim, + size_t rope_dim, + size_t num_cache_blocks, + size_t block_size, + size_t cache_stride, + size_t max_positions, + infiniDtype_t input_dtype, + double eps, + double weights_scale, + infiniDevice_t device_type, + int device_id) + : InfiniopDescriptor{device_type, device_id}, + _num_tokens(num_tokens), + _num_heads(num_heads), + _head_dim(head_dim), + _rope_dim(rope_dim), + _num_cache_blocks(num_cache_blocks), + _block_size(block_size), + _cache_stride(cache_stride), + _max_positions(max_positions), + _input_dtype(input_dtype), + _eps(static_cast(eps)), + _weights_scale(static_cast(weights_scale)) {} + +public: + static infiniStatus_t create( + infiniopHandle_t handle, + FusedDescriptor **desc_ptr, + infiniopTensorDescriptor_t q_fp8_desc, + infiniopTensorDescriptor_t weights_fp32_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t q_raw_desc, + infiniopTensorDescriptor_t k_weights_desc, + infiniopTensorDescriptor_t norm_weight_desc, + infiniopTensorDescriptor_t norm_bias_desc, + infiniopTensorDescriptor_t positions_desc, + infiniopTensorDescriptor_t cos_sin_cache_desc, + infiniopTensorDescriptor_t slot_mapping_desc, + uint64_t rope_dim, + double eps, + double weights_scale); + + infiniStatus_t calculate( + void *q_fp8, void *weights_fp32, void *k_cache, + const void *q_raw, const void *k_weights, + const void *norm_weight, const void *norm_bias, + const void *positions, const void *cos_sin_cache, + const void *slot_mapping, void *stream) const; +}; +} // namespace op::fp8_indexer_quant::nvidia + +#endif diff --git a/src/infiniop/ops/fp8_indexer_quant/operator.cc b/src/infiniop/ops/fp8_indexer_quant/operator.cc new file mode 100644 index 000000000..75ec9301d --- /dev/null +++ b/src/infiniop/ops/fp8_indexer_quant/operator.cc @@ -0,0 +1,200 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/fp8_indexer_quant.h" + +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) +#include "nvidia/fp8_indexer_quant_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateFp8IndexerQuantDescriptor( + infiniopHandle_t handle, + infiniopFp8IndexerQuantDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t q_fp8_desc, + infiniopTensorDescriptor_t weights_fp32_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t weights_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::fp8_indexer_quant::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + q_fp8_desc, weights_fp32_desc, q_desc, weights_desc) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopFp8IndexerQuant( + infiniopFp8IndexerQuantDescriptor_t desc, + void *q_fp8, + void *weights_fp32, + const void *q, + const void *weights, + void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(q_fp8, weights_fp32, q, weights, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyFp8IndexerQuantDescriptor( + infiniopFp8IndexerQuantDescriptor_t desc) { +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} + +__INFINI_C infiniStatus_t infiniopCreateFusedFp8IndexerDescriptor( + infiniopHandle_t handle, + infiniopFusedFp8IndexerDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t q_fp8_desc, + infiniopTensorDescriptor_t weights_fp32_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t q_raw_desc, + infiniopTensorDescriptor_t k_weights_desc, + infiniopTensorDescriptor_t norm_weight_desc, + infiniopTensorDescriptor_t norm_bias_desc, + infiniopTensorDescriptor_t positions_desc, + infiniopTensorDescriptor_t cos_sin_cache_desc, + infiniopTensorDescriptor_t slot_mapping_desc, + uint64_t rope_dim, + double eps, + double weights_scale) { +#define CREATE_FUSED(CASE, NAMESPACE) \ + case CASE: \ + return op::fp8_indexer_quant::NAMESPACE::FusedDescriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + q_fp8_desc, weights_fp32_desc, k_cache_desc, q_raw_desc, \ + k_weights_desc, norm_weight_desc, norm_bias_desc, positions_desc, \ + cos_sin_cache_desc, slot_mapping_desc, rope_dim, eps, weights_scale) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE_FUSED(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE_FUSED(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE_FUSED(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE_FUSED(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE_FUSED +} + +__INFINI_C infiniStatus_t infiniopFusedFp8Indexer( + infiniopFusedFp8IndexerDescriptor_t desc, + void *q_fp8, + void *weights_fp32, + void *k_cache, + const void *q_raw, + const void *k_weights, + const void *norm_weight, + const void *norm_bias, + const void *positions, + const void *cos_sin_cache, + const void *slot_mapping, + void *stream) { +#define CALCULATE_FUSED(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(q_fp8, weights_fp32, k_cache, q_raw, k_weights, \ + norm_weight, norm_bias, positions, cos_sin_cache, \ + slot_mapping, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE_FUSED(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE_FUSED(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE_FUSED(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE_FUSED(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE_FUSED +} + +__INFINI_C infiniStatus_t infiniopDestroyFusedFp8IndexerDescriptor( + infiniopFusedFp8IndexerDescriptor_t desc) { +#define DESTROY_FUSED(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DESTROY_FUSED(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY_FUSED(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY_FUSED(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY_FUSED(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY_FUSED +} diff --git a/src/infiniop/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache.h b/src/infiniop/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache.h new file mode 100644 index 000000000..3549825d4 --- /dev/null +++ b/src/infiniop/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache.h @@ -0,0 +1,54 @@ +#ifndef __FP8_MLA_RMSNORM_CACHE_H__ +#define __FP8_MLA_RMSNORM_CACHE_H__ + +#include "../../../utils.h" +#include "../../operator.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::fp8_mla_rmsnorm_cache::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + size_t _num_tokens; \ + size_t _num_cache_blocks; \ + size_t _block_size; \ + bool _write_vendor_cache; \ + float _eps; \ + \ + Descriptor( \ + size_t num_tokens, \ + size_t num_cache_blocks, \ + size_t block_size, \ + bool write_vendor_cache, \ + double eps, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _num_tokens(num_tokens), \ + _num_cache_blocks(num_cache_blocks), \ + _block_size(block_size), \ + _write_vendor_cache(write_vendor_cache), \ + _eps(static_cast(eps)) {} \ + \ + public: \ + static infiniStatus_t create( \ + infiniopHandle_t handle, \ + Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t cache_desc, \ + infiniopTensorDescriptor_t vendor_cache_desc, \ + infiniopTensorDescriptor_t compressed_kv_desc, \ + infiniopTensorDescriptor_t norm_weight_desc, \ + infiniopTensorDescriptor_t rope_desc, \ + infiniopTensorDescriptor_t slot_mapping_desc, \ + double eps); \ + \ + infiniStatus_t calculate( \ + void *cache, \ + void *vendor_cache, \ + const void *compressed_kv, \ + const void *norm_weight, \ + const void *rope, \ + const void *slot_mapping, \ + void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cu b/src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cu new file mode 100644 index 000000000..f244c5f53 --- /dev/null +++ b/src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cu @@ -0,0 +1,187 @@ +#include "../../../../utils.h" +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "../../../tensor.h" +#include "fp8_mla_rmsnorm_cache_nvidia.cuh" + +#include +#include +#include + +namespace { +constexpr size_t LATENT_DIM = 512; +constexpr size_t GROUP_SIZE = 128; +constexpr size_t NUM_GROUPS = LATENT_DIM / GROUP_SIZE; +constexpr size_t ROPE_DIM = 64; +constexpr size_t VENDOR_CACHE_STRIDE = LATENT_DIM + ROPE_DIM; +constexpr size_t CACHE_STRIDE = LATENT_DIM + NUM_GROUPS * sizeof(float) + + ROPE_DIM * sizeof(cuda_bfloat16); + +INFINIOP_CUDA_KERNEL fp8MlaRmsnormCacheKernel( + uint8_t *__restrict__ cache, + cuda_bfloat16 *__restrict__ vendor_cache, + const cuda_bfloat16 *__restrict__ compressed_kv, + const cuda_bfloat16 *__restrict__ norm_weight, + const cuda_bfloat16 *__restrict__ rope, + const int64_t *__restrict__ slot_mapping, + size_t num_cache_tokens, + float eps) { + const size_t token = blockIdx.x; + const size_t column = threadIdx.x; + const int64_t slot_raw = slot_mapping[token]; + if (slot_raw < 0 || static_cast(slot_raw) >= num_cache_tokens) { + return; + } + + extern __shared__ float scratch[]; + const float input = __bfloat162float( + compressed_kv[token * LATENT_DIM + column]); + scratch[column] = input * input; + __syncthreads(); + for (size_t stride = LATENT_DIM / 2; stride > 0; stride >>= 1) { + if (column < stride) { + scratch[column] += scratch[column + stride]; + } + __syncthreads(); + } + const float inv_rms = rsqrtf(scratch[0] / static_cast(LATENT_DIM) + eps); + const cuda_bfloat16 normalized = __float2bfloat16_rn( + input * inv_rms * __bfloat162float(norm_weight[column])); + const float value = __bfloat162float(normalized); + scratch[column] = fabsf(value); + __syncthreads(); + + const size_t group = column / GROUP_SIZE; + const size_t lane = column % GROUP_SIZE; + const size_t group_base = group * GROUP_SIZE; + for (size_t stride = GROUP_SIZE / 2; stride > 0; stride >>= 1) { + if (lane < stride) { + scratch[group_base + lane] = fmaxf( + scratch[group_base + lane], + scratch[group_base + lane + stride]); + } + __syncthreads(); + } + const float abs_max = scratch[group_base]; + const float scale = abs_max > 0.0f ? abs_max / 448.0f : 0.0f; + + const size_t slot = static_cast(slot_raw); + uint8_t *cache_entry = cache + slot * CACHE_STRIDE; + const float quantized = scale > 0.0f + ? fminf(448.0f, fmaxf(-448.0f, value / scale)) + : 0.0f; + const cuda_fp8_e4m3 fp8_value(quantized); + reinterpret_cast(cache_entry)[column] = fp8_value; + if (lane == 0) { + reinterpret_cast(cache_entry + LATENT_DIM)[group] = scale; + } + const cuda_bfloat16 rope_value = column < ROPE_DIM + ? rope[token * ROPE_DIM + column] + : __float2bfloat16_rn(0.0f); + if (column < ROPE_DIM) { + reinterpret_cast( + cache_entry + LATENT_DIM + NUM_GROUPS * sizeof(float))[column] + = rope_value; + } + if (vendor_cache != nullptr) { + auto *vendor_entry = vendor_cache + slot * VENDOR_CACHE_STRIDE; + vendor_entry[column] = __float2bfloat16_rn( + static_cast(fp8_value) * scale); + if (column < ROPE_DIM) { + vendor_entry[LATENT_DIM + column] = rope_value; + } + } +} +} // namespace + +namespace op::fp8_mla_rmsnorm_cache::nvidia { + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t cache_desc, + infiniopTensorDescriptor_t vendor_cache_desc, + infiniopTensorDescriptor_t compressed_kv_desc, + infiniopTensorDescriptor_t norm_weight_desc, + infiniopTensorDescriptor_t rope_desc, + infiniopTensorDescriptor_t slot_mapping_desc, + double eps) { + const auto cache_shape = cache_desc->shape(); + const auto kv_shape = compressed_kv_desc->shape(); + const auto weight_shape = norm_weight_desc->shape(); + const auto rope_shape = rope_desc->shape(); + const auto slots_shape = slot_mapping_desc->shape(); + CHECK_OR_RETURN(cache_shape.size() == 3 && kv_shape.size() == 2 + && weight_shape.size() == 1 && rope_shape.size() == 2 + && slots_shape.size() == 1, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(kv_shape[1] == LATENT_DIM + && weight_shape[0] == LATENT_DIM + && rope_shape[0] == kv_shape[0] + && rope_shape[1] == ROPE_DIM + && slots_shape[0] == kv_shape[0] + && cache_shape[1] > 0 + && cache_shape[2] == CACHE_STRIDE, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(cache_desc->isContiguous() + && compressed_kv_desc->isContiguous() + && norm_weight_desc->isContiguous() + && rope_desc->isContiguous() + && slot_mapping_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + CHECK_OR_RETURN(cache_desc->dtype() == INFINI_DTYPE_U8 + && compressed_kv_desc->dtype() == INFINI_DTYPE_BF16 + && norm_weight_desc->dtype() == INFINI_DTYPE_BF16 + && rope_desc->dtype() == INFINI_DTYPE_BF16 + && slot_mapping_desc->dtype() == INFINI_DTYPE_I64, + INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(eps > 0.0, INFINI_STATUS_BAD_PARAM); + if (vendor_cache_desc != nullptr) { + const auto vendor_shape = vendor_cache_desc->shape(); + CHECK_OR_RETURN( + vendor_shape.size() == 3 + && vendor_shape[0] == cache_shape[0] + && vendor_shape[1] == cache_shape[1] + && vendor_shape[2] == VENDOR_CACHE_STRIDE, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN( + vendor_cache_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + CHECK_OR_RETURN( + vendor_cache_desc->dtype() == INFINI_DTYPE_BF16, + INFINI_STATUS_BAD_TENSOR_DTYPE); + } + *desc_ptr = new Descriptor( + kv_shape[0], cache_shape[0], cache_shape[1], + vendor_cache_desc != nullptr, eps, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *cache, + void *vendor_cache, + const void *compressed_kv, + const void *norm_weight, + const void *rope, + const void *slot_mapping, + void *stream) const { + fp8MlaRmsnormCacheKernel<<< + _num_tokens, LATENT_DIM, LATENT_DIM * sizeof(float), + reinterpret_cast(stream)>>>( + reinterpret_cast(cache), + _write_vendor_cache + ? reinterpret_cast(vendor_cache) + : nullptr, + reinterpret_cast(compressed_kv), + reinterpret_cast(norm_weight), + reinterpret_cast(rope), + reinterpret_cast(slot_mapping), + _num_cache_blocks * _block_size, + _eps); + return cudaGetLastError() == cudaSuccess + ? INFINI_STATUS_SUCCESS + : INFINI_STATUS_INTERNAL_ERROR; +} + +} // namespace op::fp8_mla_rmsnorm_cache::nvidia diff --git a/src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cuh b/src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cuh new file mode 100644 index 000000000..a408a331e --- /dev/null +++ b/src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __FP8_MLA_RMSNORM_CACHE_NVIDIA_H__ +#define __FP8_MLA_RMSNORM_CACHE_NVIDIA_H__ + +#include "../fp8_mla_rmsnorm_cache.h" + +DESCRIPTOR(nvidia) + +#endif diff --git a/src/infiniop/ops/fp8_mla_rmsnorm_cache/operator.cc b/src/infiniop/ops/fp8_mla_rmsnorm_cache/operator.cc new file mode 100644 index 000000000..112522dc1 --- /dev/null +++ b/src/infiniop/ops/fp8_mla_rmsnorm_cache/operator.cc @@ -0,0 +1,101 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/fp8_mla_rmsnorm_cache.h" + +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) +#include "nvidia/fp8_mla_rmsnorm_cache_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateFp8MlaRmsnormCacheDescriptor( + infiniopHandle_t handle, + infiniopFp8MlaRmsnormCacheDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t cache_desc, + infiniopTensorDescriptor_t vendor_cache_desc, + infiniopTensorDescriptor_t compressed_kv_desc, + infiniopTensorDescriptor_t norm_weight_desc, + infiniopTensorDescriptor_t rope_desc, + infiniopTensorDescriptor_t slot_mapping_desc, + double eps) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::fp8_mla_rmsnorm_cache::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + cache_desc, vendor_cache_desc, compressed_kv_desc, \ + norm_weight_desc, rope_desc, slot_mapping_desc, eps) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopFp8MlaRmsnormCache( + infiniopFp8MlaRmsnormCacheDescriptor_t desc, + void *cache, + void *vendor_cache, + const void *compressed_kv, + const void *norm_weight, + const void *rope, + const void *slot_mapping, + void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(cache, vendor_cache, compressed_kv, norm_weight, \ + rope, slot_mapping, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyFp8MlaRmsnormCacheDescriptor( + infiniopFp8MlaRmsnormCacheDescriptor_t desc) { +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/fp8_sparse_mla/fp8_sparse_mla.h b/src/infiniop/ops/fp8_sparse_mla/fp8_sparse_mla.h new file mode 100644 index 000000000..801edbd22 --- /dev/null +++ b/src/infiniop/ops/fp8_sparse_mla/fp8_sparse_mla.h @@ -0,0 +1,71 @@ +#ifndef __FP8_SPARSE_MLA_H__ +#define __FP8_SPARSE_MLA_H__ + +#include "../../../utils.h" +#include "../../operator.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::fp8_sparse_mla::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + size_t _num_tokens; \ + size_t _num_heads; \ + size_t _head_dim; \ + size_t _value_dim; \ + size_t _num_cache_tokens; \ + size_t _cache_stride; \ + size_t _topk; \ + size_t _groups; \ + size_t _workspace_size; \ + float _scale; \ + \ + Descriptor( \ + size_t num_tokens, \ + size_t num_heads, \ + size_t head_dim, \ + size_t value_dim, \ + size_t num_cache_tokens, \ + size_t cache_stride, \ + size_t topk, \ + size_t groups, \ + size_t workspace_size, \ + float scale, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _num_tokens(num_tokens), \ + _num_heads(num_heads), \ + _head_dim(head_dim), \ + _value_dim(value_dim), \ + _num_cache_tokens(num_cache_tokens), \ + _cache_stride(cache_stride), \ + _topk(topk), \ + _groups(groups), \ + _workspace_size(workspace_size), \ + _scale(scale) {} \ + \ + public: \ + static infiniStatus_t create( \ + infiniopHandle_t handle, \ + Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t output_desc, \ + infiniopTensorDescriptor_t query_desc, \ + infiniopTensorDescriptor_t kv_cache_desc, \ + infiniopTensorDescriptor_t indices_desc, \ + infiniopTensorDescriptor_t topk_lens_desc, \ + float scale); \ + \ + size_t workspaceSize() const { return _workspace_size; } \ + \ + infiniStatus_t calculate( \ + void *workspace, \ + size_t workspace_size, \ + void *output, \ + const void *query, \ + const void *kv_cache, \ + const void *indices, \ + const void *topk_lens, \ + void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu b/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu new file mode 100644 index 000000000..ad485065a --- /dev/null +++ b/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu @@ -0,0 +1,320 @@ +#include "../../../../utils.h" +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "../../../tensor.h" +#include "fp8_sparse_mla_nvidia.cuh" + +#include +#include +#include + +namespace { +constexpr size_t THREADS = 256; +constexpr size_t GROUP_SIZE = 64; +constexpr size_t LANES_PER_KEY = 4; +constexpr size_t LATENT_DIM = 512; +constexpr size_t ROPE_DIM = 64; +constexpr size_t HEAD_DIM = LATENT_DIM + ROPE_DIM; +constexpr size_t VALUE_DIM = LATENT_DIM; +constexpr size_t CACHE_STRIDE = LATENT_DIM + 4 * sizeof(float) + + ROPE_DIM * sizeof(cuda_bfloat16); +constexpr size_t MAX_GROUPS = 64; + +INFINIOP_CUDA_KERNEL fp8SparseMlaGroupKernel( + float *__restrict__ partial, + float *__restrict__ group_maxs, + float *__restrict__ group_sums, + const cuda_bfloat16 *__restrict__ query, + const uint8_t *__restrict__ kv_cache, + const int32_t *__restrict__ indices, + const int32_t *__restrict__ topk_lens, + size_t num_heads, + size_t num_cache_tokens, + size_t topk, + size_t groups, + float softmax_scale) { + const size_t group = blockIdx.x; + const size_t head = blockIdx.y; + const size_t token = blockIdx.z; + const size_t thread = threadIdx.x; + const size_t lane = thread % LANES_PER_KEY; + const size_t key_in_group = thread / LANES_PER_KEY; + const size_t key_slot = group * GROUP_SIZE + key_in_group; + + extern __shared__ float shared[]; + float *logits = shared; + float *scratch = shared + GROUP_SIZE; + + const int32_t valid_count_raw = topk_lens[token]; + const size_t valid_count = valid_count_raw > 0 + ? min(static_cast(valid_count_raw), topk) + : 0; + int32_t cache_index = -1; + bool valid = key_slot < valid_count; + if (valid) { + cache_index = indices[token * topk + key_slot]; + valid = cache_index >= 0 + && static_cast(cache_index) < num_cache_tokens; + } + + float dot = 0.0f; + if (valid) { + const auto *cache_entry = kv_cache + + static_cast(cache_index) * CACHE_STRIDE; + const auto *latent = reinterpret_cast(cache_entry); + const auto *scales = reinterpret_cast(cache_entry + LATENT_DIM); + const auto *rope = reinterpret_cast( + cache_entry + LATENT_DIM + 4 * sizeof(float)); + const auto *q = query + (token * num_heads + head) * HEAD_DIM; + for (size_t column = lane; column < LATENT_DIM; column += LANES_PER_KEY) { + dot += __bfloat162float(q[column]) + * static_cast(latent[column]) + * scales[column / 128]; + } + for (size_t column = lane; column < ROPE_DIM; column += LANES_PER_KEY) { + dot += __bfloat162float(q[LATENT_DIM + column]) + * __bfloat162float(rope[column]); + } + } + dot += __shfl_xor_sync(0xffffffffu, dot, 1, LANES_PER_KEY); + dot += __shfl_xor_sync(0xffffffffu, dot, 2, LANES_PER_KEY); + if (lane == 0) { + logits[key_in_group] = valid ? dot * softmax_scale : -CUDART_INF_F; + } + __syncthreads(); + + if (thread < GROUP_SIZE) { + scratch[thread] = logits[thread]; + } + __syncthreads(); + for (size_t stride = GROUP_SIZE / 2; stride > 0; stride >>= 1) { + if (thread < stride) { + scratch[thread] = fmaxf(scratch[thread], scratch[thread + stride]); + } + __syncthreads(); + } + const float group_max = scratch[0]; + + if (thread < GROUP_SIZE) { + const float value = logits[thread]; + const float probability = isfinite(value) && isfinite(group_max) + ? expf(value - group_max) + : 0.0f; + logits[thread] = probability; + scratch[thread] = probability; + } + __syncthreads(); + for (size_t stride = GROUP_SIZE / 2; stride > 0; stride >>= 1) { + if (thread < stride) { + scratch[thread] += scratch[thread + stride]; + } + __syncthreads(); + } + const float group_sum = scratch[0]; + const size_t stats_offset = (token * num_heads + head) * groups + group; + if (thread == 0) { + group_maxs[stats_offset] = group_max; + group_sums[stats_offset] = group_sum; + } + + const size_t partial_base = stats_offset * VALUE_DIM; + for (size_t value_column = thread; value_column < VALUE_DIM; + value_column += blockDim.x) { + float acc = 0.0f; + if (group_sum > 0.0f) { + for (size_t key = 0; key < GROUP_SIZE; ++key) { + const size_t slot = group * GROUP_SIZE + key; + if (slot >= valid_count) { + break; + } + const int32_t index = indices[token * topk + slot]; + if (index < 0 || static_cast(index) >= num_cache_tokens) { + continue; + } + const auto *cache_entry = kv_cache + + static_cast(index) * CACHE_STRIDE; + const auto *latent = reinterpret_cast(cache_entry); + const auto *scales = reinterpret_cast(cache_entry + LATENT_DIM); + acc += (logits[key] / group_sum) + * static_cast(latent[value_column]) + * scales[value_column / 128]; + } + } + partial[partial_base + value_column] = acc; + } +} + +INFINIOP_CUDA_KERNEL fp8SparseMlaReduceKernel( + cuda_bfloat16 *__restrict__ output, + const float *__restrict__ partial, + const float *__restrict__ group_maxs, + const float *__restrict__ group_sums, + size_t num_heads, + size_t groups) { + const size_t value_tile = blockIdx.x; + const size_t head = blockIdx.y; + const size_t token = blockIdx.z; + const size_t thread = threadIdx.x; + + extern __shared__ float shared[]; + float *max_values = shared; + float *weights = shared + MAX_GROUPS; + float *scratch = shared + 2 * MAX_GROUPS; + const size_t stats_base = (token * num_heads + head) * groups; + + if (thread < MAX_GROUPS) { + const float value = thread < groups + ? group_maxs[stats_base + thread] + : -CUDART_INF_F; + max_values[thread] = value; + scratch[thread] = value; + } + __syncthreads(); + for (size_t stride = MAX_GROUPS / 2; stride > 0; stride >>= 1) { + if (thread < stride) { + scratch[thread] = fmaxf(scratch[thread], scratch[thread + stride]); + } + __syncthreads(); + } + const float final_max = scratch[0]; + + if (thread < MAX_GROUPS) { + const float weight = thread < groups + && isfinite(max_values[thread]) + && isfinite(final_max) + ? expf(max_values[thread] - final_max) + * group_sums[stats_base + thread] + : 0.0f; + weights[thread] = weight; + scratch[thread] = weight; + } + __syncthreads(); + for (size_t stride = MAX_GROUPS / 2; stride > 0; stride >>= 1) { + if (thread < stride) { + scratch[thread] += scratch[thread + stride]; + } + __syncthreads(); + } + const float denominator = scratch[0]; + + const size_t value_column = value_tile * blockDim.x + thread; + if (value_column < VALUE_DIM) { + float acc = 0.0f; + for (size_t group = 0; group < groups; ++group) { + const size_t partial_offset = (stats_base + group) * VALUE_DIM; + acc += partial[partial_offset + value_column] * weights[group]; + } + const float value = denominator > 0.0f ? acc / denominator : 0.0f; + output[(token * num_heads + head) * VALUE_DIM + value_column] + = cuda_bfloat16(value); + } +} +} // namespace + +namespace op::fp8_sparse_mla::nvidia { + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t query_desc, + infiniopTensorDescriptor_t kv_cache_desc, + infiniopTensorDescriptor_t indices_desc, + infiniopTensorDescriptor_t topk_lens_desc, + float scale) { + const auto output_shape = output_desc->shape(); + const auto query_shape = query_desc->shape(); + const auto cache_shape = kv_cache_desc->shape(); + const auto indices_shape = indices_desc->shape(); + CHECK_OR_RETURN(output_shape.size() == 3 && query_shape.size() == 3 + && cache_shape.size() == 3 && indices_shape.size() == 3 + && topk_lens_desc->shape().size() == 1, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(output_shape[0] == query_shape[0] + && output_shape[1] == query_shape[1] + && query_shape[2] == HEAD_DIM + && output_shape[2] == VALUE_DIM + && cache_shape[1] == 1 + && cache_shape[2] == CACHE_STRIDE + && indices_shape[0] == query_shape[0] + && indices_shape[1] == 1 + && topk_lens_desc->shape()[0] == query_shape[0], + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(query_desc->dtype() == INFINI_DTYPE_BF16 + && output_desc->dtype() == INFINI_DTYPE_BF16 + && kv_cache_desc->dtype() == INFINI_DTYPE_U8 + && indices_desc->dtype() == INFINI_DTYPE_I32 + && topk_lens_desc->dtype() == INFINI_DTYPE_I32, + INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(output_desc->isContiguous() && query_desc->isContiguous() + && kv_cache_desc->isContiguous() + && indices_desc->isContiguous() + && topk_lens_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + const size_t topk = indices_shape[2]; + const size_t groups = (topk + GROUP_SIZE - 1) / GROUP_SIZE; + CHECK_OR_RETURN(topk > 0 && groups <= MAX_GROUPS, + INFINI_STATUS_BAD_TENSOR_SHAPE); + const size_t stats_elems = query_shape[0] * query_shape[1] * groups; + const size_t workspace_size = (stats_elems * VALUE_DIM + 2 * stats_elems) + * sizeof(float); + *desc_ptr = new Descriptor( + query_shape[0], query_shape[1], query_shape[2], output_shape[2], + cache_shape[0], cache_shape[2], topk, groups, workspace_size, scale, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *query, + const void *kv_cache, + const void *indices, + const void *topk_lens, + void *stream) const { + CHECK_OR_RETURN(workspace != nullptr && workspace_size >= _workspace_size, + INFINI_STATUS_INSUFFICIENT_WORKSPACE); + auto *partial = reinterpret_cast(workspace); + const size_t stats_elems = _num_tokens * _num_heads * _groups; + auto *group_maxs = partial + stats_elems * _value_dim; + auto *group_sums = group_maxs + stats_elems; + const auto cuda_stream = reinterpret_cast(stream); + const dim3 group_grid( + static_cast(_groups), + static_cast(_num_heads), + static_cast(_num_tokens)); + fp8SparseMlaGroupKernel<<>>( + partial, + group_maxs, + group_sums, + reinterpret_cast(query), + reinterpret_cast(kv_cache), + reinterpret_cast(indices), + reinterpret_cast(topk_lens), + _num_heads, + _num_cache_tokens, + _topk, + _groups, + _scale); + if (cudaGetLastError() != cudaSuccess) { + return INFINI_STATUS_INTERNAL_ERROR; + } + const dim3 reduce_grid( + static_cast((_value_dim + THREADS - 1) / THREADS), + static_cast(_num_heads), + static_cast(_num_tokens)); + fp8SparseMlaReduceKernel<<>>( + reinterpret_cast(output), + partial, + group_maxs, + group_sums, + _num_heads, + _groups); + return cudaGetLastError() == cudaSuccess + ? INFINI_STATUS_SUCCESS + : INFINI_STATUS_INTERNAL_ERROR; +} + +} // namespace op::fp8_sparse_mla::nvidia diff --git a/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cuh b/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cuh new file mode 100644 index 000000000..fcd78b3d0 --- /dev/null +++ b/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __FP8_SPARSE_MLA_NVIDIA_H__ +#define __FP8_SPARSE_MLA_NVIDIA_H__ + +#include "../fp8_sparse_mla.h" + +DESCRIPTOR(nvidia) + +#endif diff --git a/src/infiniop/ops/fp8_sparse_mla/operator.cc b/src/infiniop/ops/fp8_sparse_mla/operator.cc new file mode 100644 index 000000000..6668f90fe --- /dev/null +++ b/src/infiniop/ops/fp8_sparse_mla/operator.cc @@ -0,0 +1,128 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/fp8_sparse_mla.h" + +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) +#include "nvidia/fp8_sparse_mla_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateFp8SparseMlaDescriptor( + infiniopHandle_t handle, + infiniopFp8SparseMlaDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t query_desc, + infiniopTensorDescriptor_t kv_cache_desc, + infiniopTensorDescriptor_t indices_desc, + infiniopTensorDescriptor_t topk_lens_desc, + float scale) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::fp8_sparse_mla::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + output_desc, query_desc, kv_cache_desc, indices_desc, \ + topk_lens_desc, scale) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetFp8SparseMlaWorkspaceSize( + infiniopFp8SparseMlaDescriptor_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_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C infiniStatus_t infiniopFp8SparseMla( + infiniopFp8SparseMlaDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *query, + const void *kv_cache, + const void *indices, + const void *topk_lens, + void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, output, query, \ + kv_cache, indices, topk_lens, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyFp8SparseMlaDescriptor( + infiniopFp8SparseMlaDescriptor_t desc) { +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} From 9f1bb2d568ab31c909ded87736f9a2ab3052c313 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Tue, 28 Jul 2026 08:10:02 +0000 Subject: [PATCH 2/8] feat-infiniop-add-select-last-token-hidden --- include/infiniop.h | 1 + .../infiniop/ops/select_last_token_hidden.h | 25 ++++ .../cpu/select_last_token_hidden_cpu.cc | 71 ++++++++++ .../cpu/select_last_token_hidden_cpu.h | 8 ++ .../nvidia/select_last_token_hidden_nvidia.cu | 97 ++++++++++++++ .../select_last_token_hidden_nvidia.cuh | 8 ++ .../ops/select_last_token_hidden/operator.cc | 122 ++++++++++++++++++ .../select_last_token_hidden.h | 41 ++++++ test/infiniop/select_last_token_hidden.py | 108 ++++++++++++++++ 9 files changed, 481 insertions(+) create mode 100644 include/infiniop/ops/select_last_token_hidden.h create mode 100644 src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.cc create mode 100644 src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.h create mode 100644 src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cu create mode 100644 src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cuh create mode 100644 src/infiniop/ops/select_last_token_hidden/operator.cc create mode 100644 src/infiniop/ops/select_last_token_hidden/select_last_token_hidden.h create mode 100644 test/infiniop/select_last_token_hidden.py diff --git a/include/infiniop.h b/include/infiniop.h index 70cb3cd3f..9eb10a870 100644 --- a/include/infiniop.h +++ b/include/infiniop.h @@ -131,6 +131,7 @@ #include "infiniop/ops/rwkv5_wkv.h" #include "infiniop/ops/scal.h" #include "infiniop/ops/scatter.h" +#include "infiniop/ops/select_last_token_hidden.h" #include "infiniop/ops/selu.h" #include "infiniop/ops/sigmoid.h" #include "infiniop/ops/silu.h" diff --git a/include/infiniop/ops/select_last_token_hidden.h b/include/infiniop/ops/select_last_token_hidden.h new file mode 100644 index 000000000..ac86fa8b7 --- /dev/null +++ b/include/infiniop/ops/select_last_token_hidden.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_SELECT_LAST_TOKEN_HIDDEN_API_H__ +#define __INFINIOP_SELECT_LAST_TOKEN_HIDDEN_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSelectLastTokenHiddenDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSelectLastTokenHiddenDescriptor( + infiniopHandle_t handle, + infiniopSelectLastTokenHiddenDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t hidden_states_desc, + infiniopTensorDescriptor_t input_offsets_desc); + +__INFINI_C __export infiniStatus_t infiniopSelectLastTokenHidden( + infiniopSelectLastTokenHiddenDescriptor_t desc, + void *output, + const void *hidden_states, + const void *input_offsets, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySelectLastTokenHiddenDescriptor( + infiniopSelectLastTokenHiddenDescriptor_t desc); + +#endif diff --git a/src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.cc b/src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.cc new file mode 100644 index 000000000..d8820f812 --- /dev/null +++ b/src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.cc @@ -0,0 +1,71 @@ +#include "select_last_token_hidden_cpu.h" +#include "../../../../utils.h" +#include "../../../handle.h" +#include "../../../tensor.h" +#include +#include +#include + +namespace op::select_last_token_hidden::cpu { + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t hidden_states_desc, + infiniopTensorDescriptor_t input_offsets_desc) { + const auto output_shape = output_desc->shape(); + const auto hidden_shape = hidden_states_desc->shape(); + const auto offsets_shape = input_offsets_desc->shape(); + + CHECK_OR_RETURN(output_shape.size() == 3 && hidden_shape.size() == 3 && offsets_shape.size() == 1, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(offsets_shape[0] >= 2, INFINI_STATUS_BAD_TENSOR_SHAPE); + const size_t num_requests = offsets_shape[0] - 1; + CHECK_OR_RETURN(output_shape[0] == 1 && output_shape[1] == num_requests + && output_shape[2] == hidden_shape[2], + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(output_desc->isContiguous() && hidden_states_desc->isContiguous() + && input_offsets_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + CHECK_OR_RETURN(input_offsets_desc->dtype() == INFINI_DTYPE_I32, + INFINI_STATUS_BAD_TENSOR_DTYPE); + const auto hidden_dtype = hidden_states_desc->dtype(); + CHECK_OR_RETURN(hidden_dtype == INFINI_DTYPE_F16 || hidden_dtype == INFINI_DTYPE_BF16 + || hidden_dtype == INFINI_DTYPE_F32, + INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(output_desc->dtype() == hidden_dtype, INFINI_STATUS_BAD_TENSOR_DTYPE); + + const size_t total_tokens = hidden_shape[0] * hidden_shape[1]; + CHECK_OR_RETURN(total_tokens > 0 && hidden_shape[2] > 0, INFINI_STATUS_BAD_TENSOR_SHAPE); + *desc_ptr = new Descriptor( + num_requests, + total_tokens, + hidden_shape[2] * infiniSizeOf(hidden_dtype), + handle->device, + handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *output, + const void *hidden_states, + const void *input_offsets, + void *stream) const { + (void)stream; + auto *out = reinterpret_cast(output); + const auto *hidden = reinterpret_cast(hidden_states); + const auto *offsets = reinterpret_cast(input_offsets); + for (size_t request = 0; request < _num_requests; ++request) { + const int32_t row = offsets[request + 1] - 1; + CHECK_OR_RETURN(row >= 0 && static_cast(row) < _total_tokens, + INFINI_STATUS_BAD_PARAM); + std::memcpy( + out + request * _row_bytes, + hidden + static_cast(row) * _row_bytes, + _row_bytes); + } + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::select_last_token_hidden::cpu diff --git a/src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.h b/src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.h new file mode 100644 index 000000000..f03f80045 --- /dev/null +++ b/src/infiniop/ops/select_last_token_hidden/cpu/select_last_token_hidden_cpu.h @@ -0,0 +1,8 @@ +#ifndef __SELECT_LAST_TOKEN_HIDDEN_CPU_H__ +#define __SELECT_LAST_TOKEN_HIDDEN_CPU_H__ + +#include "../select_last_token_hidden.h" + +DESCRIPTOR(cpu) + +#endif diff --git a/src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cu b/src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cu new file mode 100644 index 000000000..9f9739209 --- /dev/null +++ b/src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cu @@ -0,0 +1,97 @@ +#include "../../../../utils.h" +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "../../../tensor.h" +#include "select_last_token_hidden_nvidia.cuh" +#include +#include + +template +INFINIOP_CUDA_KERNEL selectLastTokenHiddenKernel( + CopyT *__restrict__ output, + const CopyT *__restrict__ hidden_states, + const int32_t *__restrict__ input_offsets, + size_t row_width, + size_t total_tokens) { + const size_t request = blockIdx.x; + const int32_t row = input_offsets[request + 1] - 1; + if (row < 0 || static_cast(row) >= total_tokens) { + return; + } + const CopyT *src = hidden_states + static_cast(row) * row_width; + CopyT *dst = output + request * row_width; + for (size_t column = threadIdx.x; column < row_width; column += blockDim.x) { + dst[column] = src[column]; + } +} + +namespace op::select_last_token_hidden::nvidia { + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t hidden_states_desc, + infiniopTensorDescriptor_t input_offsets_desc) { + const auto output_shape = output_desc->shape(); + const auto hidden_shape = hidden_states_desc->shape(); + const auto offsets_shape = input_offsets_desc->shape(); + + CHECK_OR_RETURN(output_shape.size() == 3 && hidden_shape.size() == 3 && offsets_shape.size() == 1, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(offsets_shape[0] >= 2, INFINI_STATUS_BAD_TENSOR_SHAPE); + const size_t num_requests = offsets_shape[0] - 1; + CHECK_OR_RETURN(output_shape[0] == 1 && output_shape[1] == num_requests + && output_shape[2] == hidden_shape[2], + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(output_desc->isContiguous() && hidden_states_desc->isContiguous() + && input_offsets_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + CHECK_OR_RETURN(input_offsets_desc->dtype() == INFINI_DTYPE_I32, + INFINI_STATUS_BAD_TENSOR_DTYPE); + const auto hidden_dtype = hidden_states_desc->dtype(); + CHECK_OR_RETURN(hidden_dtype == INFINI_DTYPE_F16 || hidden_dtype == INFINI_DTYPE_BF16 + || hidden_dtype == INFINI_DTYPE_F32, + INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(output_desc->dtype() == hidden_dtype, INFINI_STATUS_BAD_TENSOR_DTYPE); + + const size_t total_tokens = hidden_shape[0] * hidden_shape[1]; + CHECK_OR_RETURN(total_tokens > 0 && hidden_shape[2] > 0, INFINI_STATUS_BAD_TENSOR_SHAPE); + *desc_ptr = new Descriptor( + num_requests, + total_tokens, + hidden_shape[2] * infiniSizeOf(hidden_dtype), + handle->device, + handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *output, + const void *hidden_states, + const void *input_offsets, + void *stream) const { + if (_num_requests == 0 || _row_bytes == 0) { + return INFINI_STATUS_SUCCESS; + } + constexpr size_t block_size = 256; + const auto cuda_stream = reinterpret_cast(stream); + if (_row_bytes % sizeof(uint4) == 0) { + selectLastTokenHiddenKernel<<<_num_requests, block_size, 0, cuda_stream>>>( + reinterpret_cast(output), + reinterpret_cast(hidden_states), + reinterpret_cast(input_offsets), + _row_bytes / sizeof(uint4), + _total_tokens); + } else { + selectLastTokenHiddenKernel<<<_num_requests, block_size, 0, cuda_stream>>>( + reinterpret_cast(output), + reinterpret_cast(hidden_states), + reinterpret_cast(input_offsets), + _row_bytes, + _total_tokens); + } + return cudaGetLastError() == cudaSuccess ? INFINI_STATUS_SUCCESS : INFINI_STATUS_INTERNAL_ERROR; +} + +} // namespace op::select_last_token_hidden::nvidia diff --git a/src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cuh b/src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cuh new file mode 100644 index 000000000..50caf0639 --- /dev/null +++ b/src/infiniop/ops/select_last_token_hidden/nvidia/select_last_token_hidden_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __SELECT_LAST_TOKEN_HIDDEN_NVIDIA_H__ +#define __SELECT_LAST_TOKEN_HIDDEN_NVIDIA_H__ + +#include "../select_last_token_hidden.h" + +DESCRIPTOR(nvidia) + +#endif diff --git a/src/infiniop/ops/select_last_token_hidden/operator.cc b/src/infiniop/ops/select_last_token_hidden/operator.cc new file mode 100644 index 000000000..8972c79fa --- /dev/null +++ b/src/infiniop/ops/select_last_token_hidden/operator.cc @@ -0,0 +1,122 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/select_last_token_hidden.h" + +#ifdef ENABLE_CPU_API +#include "cpu/select_last_token_hidden_cpu.h" +#endif +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/select_last_token_hidden_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateSelectLastTokenHiddenDescriptor( + infiniopHandle_t handle, + infiniopSelectLastTokenHiddenDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t hidden_states_desc, + infiniopTensorDescriptor_t input_offsets_desc) { + +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::select_last_token_hidden::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + output_desc, hidden_states_desc, input_offsets_desc) + + switch (handle->device) { +#ifdef ENABLE_CPU_API + CREATE(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopSelectLastTokenHidden( + infiniopSelectLastTokenHiddenDescriptor_t desc, + void *output, + const void *hidden_states, + const void *input_offsets, + void *stream) { + +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(output, hidden_states, input_offsets, 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 +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroySelectLastTokenHiddenDescriptor( + infiniopSelectLastTokenHiddenDescriptor_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 +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef DESTROY +} diff --git a/src/infiniop/ops/select_last_token_hidden/select_last_token_hidden.h b/src/infiniop/ops/select_last_token_hidden/select_last_token_hidden.h new file mode 100644 index 000000000..b6958691c --- /dev/null +++ b/src/infiniop/ops/select_last_token_hidden/select_last_token_hidden.h @@ -0,0 +1,41 @@ +#ifndef __SELECT_LAST_TOKEN_HIDDEN_H__ +#define __SELECT_LAST_TOKEN_HIDDEN_H__ + +#include "../../../utils.h" +#include "../../operator.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::select_last_token_hidden::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + size_t _num_requests; \ + size_t _total_tokens; \ + size_t _row_bytes; \ + \ + Descriptor( \ + size_t num_requests, \ + size_t total_tokens, \ + size_t row_bytes, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _num_requests(num_requests), \ + _total_tokens(total_tokens), \ + _row_bytes(row_bytes) {} \ + \ + public: \ + static infiniStatus_t create( \ + infiniopHandle_t handle, \ + Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t output_desc, \ + infiniopTensorDescriptor_t hidden_states_desc, \ + infiniopTensorDescriptor_t input_offsets_desc); \ + \ + infiniStatus_t calculate( \ + void *output, \ + const void *hidden_states, \ + const void *input_offsets, \ + void *stream) const; \ + }; \ + } + +#endif diff --git a/test/infiniop/select_last_token_hidden.py b/test/infiniop/select_last_token_hidden.py new file mode 100644 index 000000000..edeafd1df --- /dev/null +++ b/test/infiniop/select_last_token_hidden.py @@ -0,0 +1,108 @@ +import ctypes +from ctypes import POINTER, c_int32, c_void_p + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + InfiniDtypeNames, + TestTensor, + check_error, + get_args, + get_test_devices, + infiniopHandle_t, + infiniopOperatorDescriptor_t, + infiniopTensorDescriptor_t, + test_operator, +) + +LIBINFINIOP.infiniopCreateSelectLastTokenHiddenDescriptor.restype = c_int32 +LIBINFINIOP.infiniopCreateSelectLastTokenHiddenDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, +] +LIBINFINIOP.infiniopSelectLastTokenHidden.restype = c_int32 +LIBINFINIOP.infiniopSelectLastTokenHidden.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, +] +LIBINFINIOP.infiniopDestroySelectLastTokenHiddenDescriptor.restype = c_int32 +LIBINFINIOP.infiniopDestroySelectLastTokenHiddenDescriptor.argtypes = [ + infiniopOperatorDescriptor_t, +] + + +_TEST_CASES = [ + (4, 25, 64, (0, 3, 11, 18, 25)), + (4, 2048, 6144, (0, 512, 1024, 1536, 2048)), +] +_TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16, InfiniDtype.F32] + + +def test(handle, device, num_requests, total_tokens, hidden_size, offsets, dtype, sync): + print( + f"Testing SelectLastTokenHidden on {InfiniDeviceNames[device]} with " + f"requests={num_requests}, tokens={total_tokens}, hidden={hidden_size}, " + f"dtype={InfiniDtypeNames[dtype]}" + ) + hidden = TestTensor((1, total_tokens, hidden_size), None, dtype, device) + output = TestTensor( + (1, num_requests, hidden_size), None, dtype, device, mode="zeros" + ) + input_offsets = TestTensor.from_torch( + torch.tensor(offsets, dtype=torch.int32), InfiniDtype.I32, device + ) + expected_rows = torch.tensor( + [value - 1 for value in offsets[1:]], device=hidden.torch_tensor().device + ) + expected = ( + hidden.torch_tensor() + .view(total_tokens, hidden_size) + .index_select(0, expected_rows) + ) + output.update_torch_tensor(expected.view(1, num_requests, hidden_size)) + + if sync is not None: + sync() + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateSelectLastTokenHiddenDescriptor( + handle, + ctypes.byref(descriptor), + output.descriptor, + hidden.descriptor, + input_offsets.descriptor, + ) + ) + output.destroy_desc() + hidden.destroy_desc() + input_offsets.destroy_desc() + + check_error( + LIBINFINIOP.infiniopSelectLastTokenHidden( + descriptor, + output.data(), + hidden.data(), + input_offsets.data(), + None, + ) + ) + if sync is not None: + sync() + assert torch.equal(output.actual_tensor(), output.torch_tensor()) + check_error(LIBINFINIOP.infiniopDestroySelectLastTokenHiddenDescriptor(descriptor)) + + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + print("\033[92mTest passed!\033[0m") From f74c0d72900e5c5830a82dbc178d4f1a40500040 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Tue, 28 Jul 2026 08:10:02 +0000 Subject: [PATCH 3/8] feat-infiniop-add-cuda-style-silu-and-mul-backend --- .../nvidia/silu_and_mul_nvidia.cu | 123 ++++++++++++++++++ .../nvidia/silu_and_mul_nvidia.cuh | 8 ++ src/infiniop/ops/silu_and_mul/operator.cc | 64 +++++++++ 3 files changed, 195 insertions(+) create mode 100644 src/infiniop/ops/silu_and_mul/nvidia/silu_and_mul_nvidia.cu create mode 100644 src/infiniop/ops/silu_and_mul/nvidia/silu_and_mul_nvidia.cuh diff --git a/src/infiniop/ops/silu_and_mul/nvidia/silu_and_mul_nvidia.cu b/src/infiniop/ops/silu_and_mul/nvidia/silu_and_mul_nvidia.cu new file mode 100644 index 000000000..4bc785050 --- /dev/null +++ b/src/infiniop/ops/silu_and_mul/nvidia/silu_and_mul_nvidia.cu @@ -0,0 +1,123 @@ +#include "silu_and_mul_nvidia.cuh" + +#include "../../../devices/nvidia/nvidia_handle.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" + +#include +#include +#include +#include + +namespace op::silu_and_mul::nvidia { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +namespace { + +template +__device__ float to_float(T v) { + if constexpr (std::is_same_v) { + return __half2float(v); + } else if constexpr (std::is_same_v) { + return __bfloat162float(v); + } else { + return static_cast(v); + } +} + +template +__device__ T from_float(float v) { + if constexpr (std::is_same_v) { + return __float2half(v); + } else if constexpr (std::is_same_v) { + return __float2bfloat16(v); + } else { + return static_cast(v); + } +} + +template +INFINIOP_CUDA_KERNEL siluAndMulKernel(T *y, const T *x, size_t n, size_t hidden) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < n; idx += stride) { + size_t row = idx / hidden; + size_t col = idx - row * hidden; + const T *row_x = x + row * hidden * 2; + float gate = to_float(row_x[col]); + float up = to_float(row_x[col + hidden]); + float silu = gate / (1.0f + expf(-gate)); + y[idx] = from_float(silu * up); + } +} + +template +infiniStatus_t launch(const SiluAndMulInfo &info, void *y, const void *x, void *stream) { + size_t n = info.batch_size * info.out_hidden_dim; + if (n == 0) { + return INFINI_STATUS_SUCCESS; + } + constexpr int block = 256; + int grid = static_cast((n + block - 1) / block); + grid = grid > 65535 ? 65535 : grid; + siluAndMulKernel<<(stream)>>>( + reinterpret_cast(y), reinterpret_cast(x), n, info.out_hidden_dim); + CHECK_CUDA(cudaGetLastError()); + return INFINI_STATUS_SUCCESS; +} + +} // namespace + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc) { + + if (!desc_ptr) { + return INFINI_STATUS_BAD_PARAM; + } + + auto handle = reinterpret_cast(handle_); + auto dtype = y_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); + if (x_desc->dtype() != dtype) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + auto result = SiluAndMulInfo::create(y_desc, x_desc); + CHECK_RESULT(result); + + *desc_ptr = new Descriptor( + new Opaque{handle->internal()}, + result.take(), + 0, + handle->device, + handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *y, + const void *x, + void *stream) const { + switch (_info.dtype) { + case INFINI_DTYPE_F16: + return launch(_info, y, x, stream); + case INFINI_DTYPE_F32: + return launch(_info, y, x, stream); + case INFINI_DTYPE_BF16: + return launch(_info, y, x, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::silu_and_mul::nvidia diff --git a/src/infiniop/ops/silu_and_mul/nvidia/silu_and_mul_nvidia.cuh b/src/infiniop/ops/silu_and_mul/nvidia/silu_and_mul_nvidia.cuh new file mode 100644 index 000000000..f94a90e47 --- /dev/null +++ b/src/infiniop/ops/silu_and_mul/nvidia/silu_and_mul_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __SILU_AND_MUL_NVIDIA_API_H__ +#define __SILU_AND_MUL_NVIDIA_API_H__ + +#include "../silu_and_mul.h" + +DESCRIPTOR(nvidia) + +#endif // __SILU_AND_MUL_NVIDIA_API_H__ diff --git a/src/infiniop/ops/silu_and_mul/operator.cc b/src/infiniop/ops/silu_and_mul/operator.cc index 8040ddb71..ede1c59eb 100644 --- a/src/infiniop/ops/silu_and_mul/operator.cc +++ b/src/infiniop/ops/silu_and_mul/operator.cc @@ -2,6 +2,10 @@ #include "../../handle.h" #include "infiniop/ops/silu_and_mul.h" +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/silu_and_mul_nvidia.cuh" +#endif + #ifdef ENABLE_MOORE_API #include "moore/silu_and_mul_moore.h" #endif @@ -21,6 +25,21 @@ __INFINI_C infiniStatus_t infiniopCreateSiluAndMulDescriptor( x_desc); switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif #ifdef ENABLE_MOORE_API CREATE(INFINI_DEVICE_MOORE, moore); #endif @@ -39,6 +58,21 @@ __INFINI_C infiniStatus_t infiniopGetSiluAndMulWorkspaceSize(infiniopSiluAndMulD return INFINI_STATUS_SUCCESS; switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif #ifdef ENABLE_MOORE_API GET(INFINI_DEVICE_MOORE, moore); #endif @@ -62,6 +96,21 @@ __INFINI_C infiniStatus_t infiniopSiluAndMul( workspace, workspace_size, y, x, stream); switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif #ifdef ENABLE_MOORE_API CALCULATE(INFINI_DEVICE_MOORE, moore); #endif @@ -80,6 +129,21 @@ __INFINI_C infiniStatus_t infiniopDestroySiluAndMulDescriptor(infiniopSiluAndMul return INFINI_STATUS_SUCCESS; switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif #ifdef ENABLE_MOORE_API DESTROY(INFINI_DEVICE_MOORE, moore); #endif From 6de6f665b8addaf74cbacbca61ae16c50a92e823 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Tue, 28 Jul 2026 09:34:49 +0000 Subject: [PATCH 4/8] fix-infiniop-use-portable-negative-infinity --- .../fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu | 2 +- .../ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu b/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu index 740e9d887..c84148b5b 100644 --- a/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu +++ b/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu @@ -99,7 +99,7 @@ INFINIOP_CUDA_KERNEL fp8IndexerLogitsKernel( if (lane == 0 && key_position < max_context_len) { logits[token * max_context_len + key_position] = valid ? acc - : -CUDART_INF_F; + : -INFINITY; } } } // namespace diff --git a/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu b/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu index ad485065a..0f98cde14 100644 --- a/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu +++ b/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu @@ -79,7 +79,7 @@ INFINIOP_CUDA_KERNEL fp8SparseMlaGroupKernel( dot += __shfl_xor_sync(0xffffffffu, dot, 1, LANES_PER_KEY); dot += __shfl_xor_sync(0xffffffffu, dot, 2, LANES_PER_KEY); if (lane == 0) { - logits[key_in_group] = valid ? dot * softmax_scale : -CUDART_INF_F; + logits[key_in_group] = valid ? dot * softmax_scale : -INFINITY; } __syncthreads(); @@ -165,7 +165,7 @@ INFINIOP_CUDA_KERNEL fp8SparseMlaReduceKernel( if (thread < MAX_GROUPS) { const float value = thread < groups ? group_maxs[stats_base + thread] - : -CUDART_INF_F; + : -INFINITY; max_values[thread] = value; scratch[thread] = value; } From f6c87d406c51a4a66b93481ab19243bdcd814cdf Mon Sep 17 00:00:00 2001 From: wooway777 Date: Tue, 28 Jul 2026 11:04:52 +0000 Subject: [PATCH 5/8] fix(infiniop): make GLM FP8 ops portable across CUDA backends --- .../devices/nvidia/nvidia_kernel_common.cuh | 38 ++++ .../nvidia/fp8_indexer_logits_nvidia.cu | 15 +- .../ops/fp8_indexer_logits/operator.cc | 11 +- .../nvidia/fp8_indexer_quant_nvidia.cu | 35 +--- .../ops/fp8_indexer_quant/operator.cc | 20 +- .../nvidia/fp8_mla_rmsnorm_cache_nvidia.cu | 6 +- .../ops/fp8_mla_rmsnorm_cache/operator.cc | 11 +- .../nvidia/fp8_sparse_mla_nvidia.cu | 8 +- src/infiniop/ops/fp8_sparse_mla/operator.cc | 14 +- test/infiniop/fp8_indexer_logits.py | 186 ++++++++++++++++++ test/infiniop/fp8_indexer_quant.py | 114 +++++++++++ test/infiniop/fp8_mla_rmsnorm_cache.py | 186 ++++++++++++++++++ test/infiniop/fp8_sparse_mla.py | 185 +++++++++++++++++ 13 files changed, 784 insertions(+), 45 deletions(-) create mode 100644 test/infiniop/fp8_indexer_logits.py create mode 100644 test/infiniop/fp8_indexer_quant.py create mode 100644 test/infiniop/fp8_mla_rmsnorm_cache.py create mode 100644 test/infiniop/fp8_sparse_mla.py diff --git a/src/infiniop/devices/nvidia/nvidia_kernel_common.cuh b/src/infiniop/devices/nvidia/nvidia_kernel_common.cuh index 2d1b50148..658bc441c 100644 --- a/src/infiniop/devices/nvidia/nvidia_kernel_common.cuh +++ b/src/infiniop/devices/nvidia/nvidia_kernel_common.cuh @@ -7,6 +7,8 @@ #define INFINIOP_CUDA_KERNEL __global__ void #endif +#include + #include #include #ifndef ENABLE_HYGON_API @@ -32,6 +34,42 @@ using cuda_bfloat162 = nv_bfloat162; using cuda_fp8_e4m3 = __nv_fp8_e4m3; #endif +// Store FP8 E4M3FN values as raw bytes in portable CUDA-style kernels. +// Several CUDA-compatible backends either expose a different FP8 type name or +// do not ship cuda_fp8.h at all, while the tensor ABI is still one byte. +__forceinline__ __device__ uint8_t +infiniopFp8E4m3Encode(float value) { + const uint32_t bits = __float_as_uint(value); + const uint32_t abs_bits = bits & 0x7fffffffU; + uint32_t magnitude = 0x7fU; + if (abs_bits < 0x43f00000U) { + if (abs_bits > 0x3c7fffffU) { + const uint32_t tie = (bits >> 20U) & 1U; + magnitude = (bits + tie + 0x0407ffffU) >> 20U; + } else { + const float absolute = __uint_as_float(abs_bits); + magnitude = __float_as_uint(absolute + 16384.0f); + } + } + const uint32_t sign = (bits >> 24U) & 0x80U; + return static_cast((magnitude & 0x7fU) | sign); +} + +__forceinline__ __device__ float +infiniopFp8E4m3Decode(uint8_t value) { + const uint32_t magnitude = value & 0x7fU; + const uint32_t exponent = magnitude >> 3U; + const uint32_t mantissa = magnitude & 0x7U; + if (exponent == 0xfU && mantissa == 0x7U) { + return __uint_as_float(0x7fffffffU); + } + const float decoded = exponent == 0U + ? static_cast(mantissa) * 0.001953125f + : __uint_as_float((exponent + 120U) << 23U) + * (1.0f + static_cast(mantissa) * 0.125f); + return (value & 0x80U) != 0U ? -decoded : decoded; +} + namespace device::nvidia { // get the memory offset of the given element in a tensor given its flat index diff --git a/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu b/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu index c84148b5b..68d0481c0 100644 --- a/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu +++ b/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu @@ -11,7 +11,7 @@ constexpr size_t THREADS = 256; template INFINIOP_CUDA_KERNEL fp8IndexerLogitsKernel( float *__restrict__ logits, - const cuda_fp8_e4m3 *__restrict__ q_fp8, + const uint8_t *__restrict__ q_fp8, const uint8_t *__restrict__ kv_cache, const int32_t *__restrict__ block_tables, const float *__restrict__ weights_fp32, @@ -35,7 +35,7 @@ INFINIOP_CUDA_KERNEL fp8IndexerLogitsKernel( + threadIdx.x / LANES_PER_KEY; extern __shared__ uint8_t shared_bytes[]; - auto *shared_q = reinterpret_cast(shared_bytes); + auto *shared_q = shared_bytes; auto *shared_weights = reinterpret_cast( shared_bytes + num_heads * head_dim); const size_t q_base = token * num_heads * head_dim; @@ -69,8 +69,7 @@ INFINIOP_CUDA_KERNEL fp8IndexerLogitsKernel( + static_cast(physical_block) * block_size * cache_stride : nullptr; const auto *key = valid - ? reinterpret_cast( - cache_block + key_in_block * head_dim) + ? cache_block + key_in_block * head_dim : nullptr; const float key_scale = valid ? *reinterpret_cast( @@ -83,8 +82,8 @@ INFINIOP_CUDA_KERNEL fp8IndexerLogitsKernel( if (valid) { const auto *query = shared_q + head * head_dim; for (size_t column = lane; column < head_dim; column += LANES_PER_KEY) { - dot += static_cast(query[column]) - * static_cast(key[column]); + dot += infiniopFp8E4m3Decode(query[column]) + * infiniopFp8E4m3Decode(key[column]); } } // Every thread named by the full-warp mask must execute the shuffle, @@ -182,7 +181,7 @@ infiniStatus_t Descriptor::calculate( fp8IndexerLogitsKernel<8><<< grid, THREADS, smem, reinterpret_cast(stream)>>>( reinterpret_cast(logits), - reinterpret_cast(q_fp8), + reinterpret_cast(q_fp8), reinterpret_cast(kv_cache), reinterpret_cast(block_tables), reinterpret_cast(weights_fp32), @@ -194,7 +193,7 @@ infiniStatus_t Descriptor::calculate( fp8IndexerLogitsKernel<4><<< grid, THREADS, smem, reinterpret_cast(stream)>>>( reinterpret_cast(logits), - reinterpret_cast(q_fp8), + reinterpret_cast(q_fp8), reinterpret_cast(kv_cache), reinterpret_cast(block_tables), reinterpret_cast(weights_fp32), diff --git a/src/infiniop/ops/fp8_indexer_logits/operator.cc b/src/infiniop/ops/fp8_indexer_logits/operator.cc index 2448832dd..348397519 100644 --- a/src/infiniop/ops/fp8_indexer_logits/operator.cc +++ b/src/infiniop/ops/fp8_indexer_logits/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/fp8_indexer_logits.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) #include "nvidia/fp8_indexer_logits_nvidia.cuh" #endif @@ -33,6 +33,9 @@ __INFINI_C infiniStatus_t infiniopCreateFp8IndexerLogitsDescriptor( #ifdef ENABLE_QY_API CREATE(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia); #endif @@ -67,6 +70,9 @@ __INFINI_C infiniStatus_t infiniopFp8IndexerLogits( #ifdef ENABLE_QY_API CALCULATE(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia); #endif @@ -92,6 +98,9 @@ __INFINI_C infiniStatus_t infiniopDestroyFp8IndexerLogitsDescriptor( #ifdef ENABLE_QY_API DESTROY(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API DESTROY(INFINI_DEVICE_ALI, nvidia); #endif diff --git a/src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cu b/src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cu index a6a7bad74..379029028 100644 --- a/src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cu +++ b/src/infiniop/ops/fp8_indexer_quant/nvidia/fp8_indexer_quant_nvidia.cu @@ -42,26 +42,9 @@ __device__ __forceinline__ float vendor_fp8_quotient(float value, float scale) { return fmaf(residual, reciprocal, quotient); } -__device__ __forceinline__ uint8_t vendor_fp8_e4m3(float value) { - const uint32_t bits = __float_as_uint(value); - const uint32_t abs_bits = bits & 0x7fffffffU; - uint32_t magnitude = 0x7fU; - if (abs_bits < 0x43f00000U) { - if (abs_bits > 0x3c7fffffU) { - const uint32_t tie = (bits >> 20U) & 1U; - magnitude = (bits + tie + 0x0407ffffU) >> 20U; - } else { - const float absolute = __uint_as_float(abs_bits); - magnitude = __float_as_uint(absolute + 16384.0f); - } - } - const uint32_t sign = (bits >> 24U) & 0x80U; - return static_cast((magnitude & 0x7fU) | sign); -} - template INFINIOP_CUDA_KERNEL fp8IndexerQuantKernel( - cuda_fp8_e4m3 *__restrict__ q_fp8, + uint8_t *__restrict__ q_fp8, float *__restrict__ weights_fp32, const T *__restrict__ q, const T *__restrict__ weights, @@ -92,13 +75,13 @@ INFINIOP_CUDA_KERNEL fp8IndexerQuantKernel( if (column < head_dim) { const float quantized = fminf(448.0f, fmaxf(-448.0f, value / reduction[0])); - q_fp8[group * head_dim + column] = cuda_fp8_e4m3(quantized); + q_fp8[group * head_dim + column] = infiniopFp8E4m3Encode(quantized); } } template INFINIOP_CUDA_KERNEL fusedFp8IndexerKernel( - cuda_fp8_e4m3 *__restrict__ q_fp8, + uint8_t *__restrict__ q_fp8, float *__restrict__ weights_fp32, uint8_t *__restrict__ k_cache, const T *__restrict__ q_raw, @@ -171,7 +154,7 @@ INFINIOP_CUDA_KERNEL fusedFp8IndexerKernel( __syncthreads(); const float quantized = fminf( 448.0f, fmaxf(-448.0f, value / scratch[0])); - q_fp8[q_base + column] = cuda_fp8_e4m3(quantized); + q_fp8[q_base + column] = infiniopFp8E4m3Encode(quantized); return; } @@ -272,8 +255,8 @@ INFINIOP_CUDA_KERNEL fusedFp8IndexerKernel( 448.0f, fmaxf(-448.0f, vendor_fp8_quotient(value0, scratch[0]))); const float quantized1 = fminf( 448.0f, fmaxf(-448.0f, vendor_fp8_quotient(value1, scratch[0]))); - cache_values[column0] = vendor_fp8_e4m3(quantized0); - cache_values[column1] = vendor_fp8_e4m3(quantized1); + cache_values[column0] = infiniopFp8E4m3Encode(quantized0); + cache_values[column1] = infiniopFp8E4m3Encode(quantized1); } if (column == 0) { *reinterpret_cast( @@ -337,14 +320,14 @@ infiniStatus_t Descriptor::calculate( const size_t smem = _threads * sizeof(float); if (_input_dtype == INFINI_DTYPE_F16) { fp8IndexerQuantKernel<<<_num_groups, _threads, smem, cuda_stream>>>( - reinterpret_cast(q_fp8), + reinterpret_cast(q_fp8), reinterpret_cast(weights_fp32), reinterpret_cast(q), reinterpret_cast(weights), _head_dim); } else { fp8IndexerQuantKernel<<<_num_groups, _threads, smem, cuda_stream>>>( - reinterpret_cast(q_fp8), + reinterpret_cast(q_fp8), reinterpret_cast(weights_fp32), reinterpret_cast(q), reinterpret_cast(weights), @@ -447,7 +430,7 @@ infiniStatus_t FusedDescriptor::calculate( #define LAUNCH_FUSED(T) \ fusedFp8IndexerKernel<<(stream)>>>( \ - reinterpret_cast(q_fp8), \ + reinterpret_cast(q_fp8), \ reinterpret_cast(weights_fp32), \ reinterpret_cast(k_cache), \ reinterpret_cast(q_raw), \ diff --git a/src/infiniop/ops/fp8_indexer_quant/operator.cc b/src/infiniop/ops/fp8_indexer_quant/operator.cc index 75ec9301d..e241b9b00 100644 --- a/src/infiniop/ops/fp8_indexer_quant/operator.cc +++ b/src/infiniop/ops/fp8_indexer_quant/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/fp8_indexer_quant.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) #include "nvidia/fp8_indexer_quant_nvidia.cuh" #endif @@ -29,6 +29,9 @@ __INFINI_C infiniStatus_t infiniopCreateFp8IndexerQuantDescriptor( #ifdef ENABLE_QY_API CREATE(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia); #endif @@ -59,6 +62,9 @@ __INFINI_C infiniStatus_t infiniopFp8IndexerQuant( #ifdef ENABLE_QY_API CALCULATE(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia); #endif @@ -84,6 +90,9 @@ __INFINI_C infiniStatus_t infiniopDestroyFp8IndexerQuantDescriptor( #ifdef ENABLE_QY_API DESTROY(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API DESTROY(INFINI_DEVICE_ALI, nvidia); #endif @@ -127,6 +136,9 @@ __INFINI_C infiniStatus_t infiniopCreateFusedFp8IndexerDescriptor( #ifdef ENABLE_QY_API CREATE_FUSED(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + CREATE_FUSED(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API CREATE_FUSED(INFINI_DEVICE_ALI, nvidia); #endif @@ -165,6 +177,9 @@ __INFINI_C infiniStatus_t infiniopFusedFp8Indexer( #ifdef ENABLE_QY_API CALCULATE_FUSED(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + CALCULATE_FUSED(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API CALCULATE_FUSED(INFINI_DEVICE_ALI, nvidia); #endif @@ -190,6 +205,9 @@ __INFINI_C infiniStatus_t infiniopDestroyFusedFp8IndexerDescriptor( #ifdef ENABLE_QY_API DESTROY_FUSED(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + DESTROY_FUSED(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API DESTROY_FUSED(INFINI_DEVICE_ALI, nvidia); #endif diff --git a/src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cu b/src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cu index f244c5f53..b4e75b592 100644 --- a/src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cu +++ b/src/infiniop/ops/fp8_mla_rmsnorm_cache/nvidia/fp8_mla_rmsnorm_cache_nvidia.cu @@ -70,8 +70,8 @@ INFINIOP_CUDA_KERNEL fp8MlaRmsnormCacheKernel( const float quantized = scale > 0.0f ? fminf(448.0f, fmaxf(-448.0f, value / scale)) : 0.0f; - const cuda_fp8_e4m3 fp8_value(quantized); - reinterpret_cast(cache_entry)[column] = fp8_value; + const uint8_t fp8_value = infiniopFp8E4m3Encode(quantized); + cache_entry[column] = fp8_value; if (lane == 0) { reinterpret_cast(cache_entry + LATENT_DIM)[group] = scale; } @@ -86,7 +86,7 @@ INFINIOP_CUDA_KERNEL fp8MlaRmsnormCacheKernel( if (vendor_cache != nullptr) { auto *vendor_entry = vendor_cache + slot * VENDOR_CACHE_STRIDE; vendor_entry[column] = __float2bfloat16_rn( - static_cast(fp8_value) * scale); + infiniopFp8E4m3Decode(fp8_value) * scale); if (column < ROPE_DIM) { vendor_entry[LATENT_DIM + column] = rope_value; } diff --git a/src/infiniop/ops/fp8_mla_rmsnorm_cache/operator.cc b/src/infiniop/ops/fp8_mla_rmsnorm_cache/operator.cc index 112522dc1..b40fa1c20 100644 --- a/src/infiniop/ops/fp8_mla_rmsnorm_cache/operator.cc +++ b/src/infiniop/ops/fp8_mla_rmsnorm_cache/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/fp8_mla_rmsnorm_cache.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) #include "nvidia/fp8_mla_rmsnorm_cache_nvidia.cuh" #endif @@ -33,6 +33,9 @@ __INFINI_C infiniStatus_t infiniopCreateFp8MlaRmsnormCacheDescriptor( #ifdef ENABLE_QY_API CREATE(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia); #endif @@ -66,6 +69,9 @@ __INFINI_C infiniStatus_t infiniopFp8MlaRmsnormCache( #ifdef ENABLE_QY_API CALCULATE(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia); #endif @@ -91,6 +97,9 @@ __INFINI_C infiniStatus_t infiniopDestroyFp8MlaRmsnormCacheDescriptor( #ifdef ENABLE_QY_API DESTROY(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API DESTROY(INFINI_DEVICE_ALI, nvidia); #endif diff --git a/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu b/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu index 0f98cde14..987754c53 100644 --- a/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu +++ b/src/infiniop/ops/fp8_sparse_mla/nvidia/fp8_sparse_mla_nvidia.cu @@ -61,14 +61,14 @@ INFINIOP_CUDA_KERNEL fp8SparseMlaGroupKernel( if (valid) { const auto *cache_entry = kv_cache + static_cast(cache_index) * CACHE_STRIDE; - const auto *latent = reinterpret_cast(cache_entry); + const auto *latent = cache_entry; const auto *scales = reinterpret_cast(cache_entry + LATENT_DIM); const auto *rope = reinterpret_cast( cache_entry + LATENT_DIM + 4 * sizeof(float)); const auto *q = query + (token * num_heads + head) * HEAD_DIM; for (size_t column = lane; column < LATENT_DIM; column += LANES_PER_KEY) { dot += __bfloat162float(q[column]) - * static_cast(latent[column]) + * infiniopFp8E4m3Decode(latent[column]) * scales[column / 128]; } for (size_t column = lane; column < ROPE_DIM; column += LANES_PER_KEY) { @@ -133,10 +133,10 @@ INFINIOP_CUDA_KERNEL fp8SparseMlaGroupKernel( } const auto *cache_entry = kv_cache + static_cast(index) * CACHE_STRIDE; - const auto *latent = reinterpret_cast(cache_entry); + const auto *latent = cache_entry; const auto *scales = reinterpret_cast(cache_entry + LATENT_DIM); acc += (logits[key] / group_sum) - * static_cast(latent[value_column]) + * infiniopFp8E4m3Decode(latent[value_column]) * scales[value_column / 128]; } } diff --git a/src/infiniop/ops/fp8_sparse_mla/operator.cc b/src/infiniop/ops/fp8_sparse_mla/operator.cc index 6668f90fe..e563ef799 100644 --- a/src/infiniop/ops/fp8_sparse_mla/operator.cc +++ b/src/infiniop/ops/fp8_sparse_mla/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/fp8_sparse_mla.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) #include "nvidia/fp8_sparse_mla_nvidia.cuh" #endif @@ -32,6 +32,9 @@ __INFINI_C infiniStatus_t infiniopCreateFp8SparseMlaDescriptor( #ifdef ENABLE_QY_API CREATE(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API CREATE(INFINI_DEVICE_ALI, nvidia); #endif @@ -59,6 +62,9 @@ __INFINI_C infiniStatus_t infiniopGetFp8SparseMlaWorkspaceSize( #ifdef ENABLE_QY_API GET(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API GET(INFINI_DEVICE_ALI, nvidia); #endif @@ -93,6 +99,9 @@ __INFINI_C infiniStatus_t infiniopFp8SparseMla( #ifdef ENABLE_QY_API CALCULATE(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API CALCULATE(INFINI_DEVICE_ALI, nvidia); #endif @@ -118,6 +127,9 @@ __INFINI_C infiniStatus_t infiniopDestroyFp8SparseMlaDescriptor( #ifdef ENABLE_QY_API DESTROY(INFINI_DEVICE_QY, nvidia); #endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif #ifdef ENABLE_ALI_API DESTROY(INFINI_DEVICE_ALI, nvidia); #endif diff --git a/test/infiniop/fp8_indexer_logits.py b/test/infiniop/fp8_indexer_logits.py new file mode 100644 index 000000000..4ae933aa4 --- /dev/null +++ b/test/infiniop/fp8_indexer_logits.py @@ -0,0 +1,186 @@ +import ctypes +from ctypes import POINTER, c_int32, c_void_p + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + check_error, + get_args, + get_test_devices, + infiniopHandle_t, + infiniopOperatorDescriptor_t, + infiniopTensorDescriptor_t, + test_operator, +) + +LIBINFINIOP.infiniopCreateFp8IndexerLogitsDescriptor.restype = c_int32 +LIBINFINIOP.infiniopCreateFp8IndexerLogitsDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, +] +LIBINFINIOP.infiniopFp8IndexerLogits.restype = c_int32 +LIBINFINIOP.infiniopFp8IndexerLogits.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, +] +LIBINFINIOP.infiniopDestroyFp8IndexerLogitsDescriptor.restype = c_int32 +LIBINFINIOP.infiniopDestroyFp8IndexerLogitsDescriptor.argtypes = [ + infiniopOperatorDescriptor_t +] + +_TEST_CASES = [(2, 2), (3, 2)] +_TENSOR_DTYPES = [None] +_BLOCK_SIZE = 64 +_HEAD_DIM = 128 +_CACHE_STRIDE = _HEAD_DIM + 4 +_MAX_CONTEXT = 128 +_NUM_HEADS = 3 + + +def _make_cache(num_blocks): + keys = (torch.rand(num_blocks, _BLOCK_SIZE, _HEAD_DIM) * 8.0 - 4.0).to( + torch.float8_e4m3fn + ) + scales = torch.rand(num_blocks, _BLOCK_SIZE, dtype=torch.float32) * 0.02 + 0.001 + raw = torch.zeros( + (num_blocks, _BLOCK_SIZE * _CACHE_STRIDE), dtype=torch.uint8 + ) + raw[:, : _BLOCK_SIZE * _HEAD_DIM] = keys.view(torch.uint8).reshape( + num_blocks, -1 + ) + raw[:, _BLOCK_SIZE * _HEAD_DIM :] = scales.contiguous().view( + torch.uint8 + ).reshape(num_blocks, -1) + return raw.reshape(num_blocks, _BLOCK_SIZE, _CACHE_STRIDE), keys, scales + + +def test(handle, device, num_tokens, num_requests, _dtype, sync): + print( + f"Testing Fp8IndexerLogits on {InfiniDeviceNames[device]} with " + f"tokens={num_tokens}, requests={num_requests}" + ) + num_blocks = num_requests * 2 + q_src = ( + torch.rand(num_tokens, _NUM_HEADS, _HEAD_DIM) * 8.0 - 4.0 + ).to(torch.float8_e4m3fn) + weights_src = torch.rand(num_tokens, _NUM_HEADS, dtype=torch.float32) + cache_src, keys, key_scales = _make_cache(num_blocks) + block_tables_src = torch.arange( + num_blocks, dtype=torch.int32 + ).reshape(num_requests, 2) + positions_src = torch.tensor( + [70, 45] + [12] * (num_tokens - 2), dtype=torch.int64 + ) + request_ids_src = torch.tensor( + list(range(num_requests)) + [-1] * (num_tokens - num_requests), + dtype=torch.int32, + ) + + expected = torch.full( + (num_tokens, _MAX_CONTEXT), -torch.inf, dtype=torch.float32 + ) + for token in range(num_tokens): + request = int(request_ids_src[token]) + if request < 0 or request >= num_requests: + continue + for key_position in range(int(positions_src[token]) + 1): + logical_block = key_position // _BLOCK_SIZE + key_offset = key_position % _BLOCK_SIZE + physical_block = int(block_tables_src[request, logical_block]) + key = keys[physical_block, key_offset].float() + scale = key_scales[physical_block, key_offset] + acc = 0.0 + for head in range(_NUM_HEADS): + dot = torch.dot(q_src[token, head].float(), key) + acc += float(torch.relu(dot * scale) * weights_src[token, head]) + expected[token, key_position] = acc + + logits = TestTensor( + (num_tokens, _MAX_CONTEXT), + None, + InfiniDtype.F32, + device, + mode="zeros", + ) + q_fp8 = TestTensor.from_torch(q_src, InfiniDtype.F8, device) + kv_cache = TestTensor.from_torch(cache_src, InfiniDtype.U8, device) + block_tables = TestTensor.from_torch( + block_tables_src, InfiniDtype.I32, device + ) + weights = TestTensor.from_torch(weights_src, InfiniDtype.F32, device) + positions = TestTensor.from_torch(positions_src, InfiniDtype.I64, device) + request_ids = TestTensor.from_torch(request_ids_src, InfiniDtype.I32, device) + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateFp8IndexerLogitsDescriptor( + handle, + ctypes.byref(descriptor), + logits.descriptor, + q_fp8.descriptor, + kv_cache.descriptor, + block_tables.descriptor, + weights.descriptor, + positions.descriptor, + request_ids.descriptor, + ) + ) + for tensor in ( + logits, + q_fp8, + kv_cache, + block_tables, + weights, + positions, + request_ids, + ): + tensor.destroy_desc() + + check_error( + LIBINFINIOP.infiniopFp8IndexerLogits( + descriptor, + logits.data(), + q_fp8.data(), + kv_cache.data(), + block_tables.data(), + weights.data(), + positions.data(), + request_ids.data(), + None, + ) + ) + if sync is not None: + sync() + + actual = logits.actual_tensor().cpu() + assert torch.equal(torch.isneginf(actual), torch.isneginf(expected)) + finite = torch.isfinite(expected) + assert torch.allclose( + actual[finite], expected[finite], atol=2.0e-3, rtol=2.0e-4 + ) + check_error(LIBINFINIOP.infiniopDestroyFp8IndexerLogitsDescriptor(descriptor)) + + +if __name__ == "__main__": + args = get_args() + torch.manual_seed(1) + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + print("\033[92mFp8IndexerLogits test passed!\033[0m") diff --git a/test/infiniop/fp8_indexer_quant.py b/test/infiniop/fp8_indexer_quant.py new file mode 100644 index 000000000..889094701 --- /dev/null +++ b/test/infiniop/fp8_indexer_quant.py @@ -0,0 +1,114 @@ +import ctypes +from ctypes import POINTER, c_int32, c_void_p + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + InfiniDtypeNames, + TestTensor, + check_error, + get_args, + get_test_devices, + infiniopHandle_t, + infiniopOperatorDescriptor_t, + infiniopTensorDescriptor_t, + test_operator, +) + +LIBINFINIOP.infiniopCreateFp8IndexerQuantDescriptor.restype = c_int32 +LIBINFINIOP.infiniopCreateFp8IndexerQuantDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, +] +LIBINFINIOP.infiniopFp8IndexerQuant.restype = c_int32 +LIBINFINIOP.infiniopFp8IndexerQuant.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, +] +LIBINFINIOP.infiniopDestroyFp8IndexerQuantDescriptor.restype = c_int32 +LIBINFINIOP.infiniopDestroyFp8IndexerQuantDescriptor.argtypes = [ + infiniopOperatorDescriptor_t +] + +_TEST_CASES = [ + ((2, 3, 128),), + ((1, 5, 96),), + ((3, 2, 257),), +] +_TENSOR_DTYPES = [InfiniDtype.F16, InfiniDtype.BF16] + + +def test(handle, device, shape, dtype, sync): + print( + f"Testing Fp8IndexerQuant on {InfiniDeviceNames[device]} with " + f"shape={shape}, dtype={InfiniDtypeNames[dtype]}" + ) + q = TestTensor(shape, None, dtype, device, scale=8.0, bias=-4.0) + weights = TestTensor(shape[:2], None, dtype, device, scale=2.0, bias=-1.0) + q_fp8 = TestTensor(shape, None, InfiniDtype.F8, device, mode="zeros") + weights_fp32 = TestTensor( + shape[:2], None, InfiniDtype.F32, device, mode="zeros" + ) + + q_float = q.torch_tensor().float() + scales = torch.exp2( + torch.ceil(torch.log2(q_float.abs().amax(dim=-1).clamp_min(1.0e-4) / 448.0)) + ) + expected_q = (q_float / scales.unsqueeze(-1)).clamp(-448.0, 448.0).to( + torch.float8_e4m3fn + ) + expected_weights = weights.torch_tensor().float() * scales + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateFp8IndexerQuantDescriptor( + handle, + ctypes.byref(descriptor), + q_fp8.descriptor, + weights_fp32.descriptor, + q.descriptor, + weights.descriptor, + ) + ) + for tensor in (q_fp8, weights_fp32, q, weights): + tensor.destroy_desc() + + check_error( + LIBINFINIOP.infiniopFp8IndexerQuant( + descriptor, + q_fp8.data(), + weights_fp32.data(), + q.data(), + weights.data(), + None, + ) + ) + if sync is not None: + sync() + + assert torch.equal( + q_fp8.actual_tensor().view(torch.uint8), + expected_q.view(torch.uint8), + ) + assert torch.allclose( + weights_fp32.actual_tensor(), expected_weights, atol=1.0e-7, rtol=1.0e-6 + ) + check_error(LIBINFINIOP.infiniopDestroyFp8IndexerQuantDescriptor(descriptor)) + + +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[92mFp8IndexerQuant test passed!\033[0m") diff --git a/test/infiniop/fp8_mla_rmsnorm_cache.py b/test/infiniop/fp8_mla_rmsnorm_cache.py new file mode 100644 index 000000000..4d8ecd8c9 --- /dev/null +++ b/test/infiniop/fp8_mla_rmsnorm_cache.py @@ -0,0 +1,186 @@ +import ctypes +from ctypes import POINTER, c_double, c_int32, c_void_p + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + check_error, + get_args, + get_test_devices, + infiniopHandle_t, + infiniopOperatorDescriptor_t, + infiniopTensorDescriptor_t, + test_operator, +) + +LIBINFINIOP.infiniopCreateFp8MlaRmsnormCacheDescriptor.restype = c_int32 +LIBINFINIOP.infiniopCreateFp8MlaRmsnormCacheDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + c_double, +] +LIBINFINIOP.infiniopFp8MlaRmsnormCache.restype = c_int32 +LIBINFINIOP.infiniopFp8MlaRmsnormCache.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, +] +LIBINFINIOP.infiniopDestroyFp8MlaRmsnormCacheDescriptor.restype = c_int32 +LIBINFINIOP.infiniopDestroyFp8MlaRmsnormCacheDescriptor.argtypes = [ + infiniopOperatorDescriptor_t +] + +_LATENT_DIM = 512 +_GROUP_SIZE = 128 +_NUM_GROUPS = 4 +_ROPE_DIM = 64 +_CACHE_STRIDE = _LATENT_DIM + _NUM_GROUPS * 4 + _ROPE_DIM * 2 +_EPS = 1.0e-6 +_TEST_CASES = [(3, True), (2, False)] +_TENSOR_DTYPES = [None] + + +def test(handle, device, num_tokens, write_vendor_cache, _dtype, sync): + print( + f"Testing Fp8MlaRmsnormCache on {InfiniDeviceNames[device]} with " + f"tokens={num_tokens}, vendor_cache={write_vendor_cache}" + ) + num_blocks, block_size = 2, 4 + compressed_src = (torch.rand(num_tokens, _LATENT_DIM) * 4.0 - 2.0).to( + torch.bfloat16 + ) + weight_src = (torch.rand(_LATENT_DIM) + 0.5).to(torch.bfloat16) + rope_src = (torch.rand(num_tokens, _ROPE_DIM) * 2.0 - 1.0).to( + torch.bfloat16 + ) + slots_src = torch.tensor( + [0, block_size + 1] + ([-1] if num_tokens == 3 else []), + dtype=torch.int64, + ) + + cache = TestTensor( + (num_blocks, block_size, _CACHE_STRIDE), + None, + InfiniDtype.U8, + device, + mode="zeros", + ) + vendor_cache = ( + TestTensor( + (num_blocks, block_size, _LATENT_DIM + _ROPE_DIM), + None, + InfiniDtype.BF16, + device, + mode="zeros", + ) + if write_vendor_cache + else None + ) + compressed = TestTensor.from_torch( + compressed_src, InfiniDtype.BF16, device + ) + weight = TestTensor.from_torch(weight_src, InfiniDtype.BF16, device) + rope = TestTensor.from_torch(rope_src, InfiniDtype.BF16, device) + slots = TestTensor.from_torch(slots_src, InfiniDtype.I64, device) + + expected_cache = torch.zeros( + (num_blocks * block_size, _CACHE_STRIDE), dtype=torch.uint8 + ) + expected_vendor = torch.zeros( + (num_blocks * block_size, _LATENT_DIM + _ROPE_DIM), + dtype=torch.bfloat16, + ) + for token, slot in enumerate(slots_src.tolist()): + if slot < 0: + continue + value = compressed_src[token].float() + inv_rms = torch.rsqrt(value.square().mean() + _EPS) + normalized = (value * inv_rms * weight_src.float()).to(torch.bfloat16) + grouped = normalized.float().reshape(_NUM_GROUPS, _GROUP_SIZE) + scales = grouped.abs().amax(dim=-1) / 448.0 + quantized = torch.where( + scales[:, None] > 0, + (grouped / scales[:, None]).clamp(-448.0, 448.0), + torch.zeros_like(grouped), + ).to(torch.float8_e4m3fn) + expected_cache[slot, :_LATENT_DIM] = quantized.view(torch.uint8).reshape(-1) + scale_offset = _LATENT_DIM + expected_cache[slot, scale_offset : scale_offset + 16] = ( + scales.contiguous().view(torch.uint8) + ) + rope_offset = _LATENT_DIM + 16 + expected_cache[slot, rope_offset:] = rope_src[token].view(torch.uint8) + expected_vendor[slot, :_LATENT_DIM] = ( + quantized.float() * scales[:, None] + ).reshape(-1).to(torch.bfloat16) + expected_vendor[slot, _LATENT_DIM:] = rope_src[token] + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateFp8MlaRmsnormCacheDescriptor( + handle, + ctypes.byref(descriptor), + cache.descriptor, + vendor_cache.descriptor if vendor_cache is not None else None, + compressed.descriptor, + weight.descriptor, + rope.descriptor, + slots.descriptor, + _EPS, + ) + ) + tensors = [cache, compressed, weight, rope, slots] + if vendor_cache is not None: + tensors.append(vendor_cache) + for tensor in tensors: + tensor.destroy_desc() + + check_error( + LIBINFINIOP.infiniopFp8MlaRmsnormCache( + descriptor, + cache.data(), + vendor_cache.data() if vendor_cache is not None else None, + compressed.data(), + weight.data(), + rope.data(), + slots.data(), + None, + ) + ) + if sync is not None: + sync() + + actual_cache = cache.actual_tensor().cpu().reshape(-1, _CACHE_STRIDE) + assert torch.equal(actual_cache, expected_cache) + if vendor_cache is not None: + assert torch.equal( + vendor_cache.actual_tensor().cpu().reshape( + -1, _LATENT_DIM + _ROPE_DIM + ), + expected_vendor, + ) + check_error( + LIBINFINIOP.infiniopDestroyFp8MlaRmsnormCacheDescriptor(descriptor) + ) + + +if __name__ == "__main__": + args = get_args() + torch.manual_seed(2) + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + print("\033[92mFp8MlaRmsnormCache test passed!\033[0m") diff --git a/test/infiniop/fp8_sparse_mla.py b/test/infiniop/fp8_sparse_mla.py new file mode 100644 index 000000000..4d799acf1 --- /dev/null +++ b/test/infiniop/fp8_sparse_mla.py @@ -0,0 +1,185 @@ +import ctypes +from ctypes import POINTER, c_float, c_int32, c_uint64, c_void_p + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopHandle_t, + infiniopOperatorDescriptor_t, + infiniopTensorDescriptor_t, + test_operator, +) + +LIBINFINIOP.infiniopCreateFp8SparseMlaDescriptor.restype = c_int32 +LIBINFINIOP.infiniopCreateFp8SparseMlaDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + c_float, +] +LIBINFINIOP.infiniopGetFp8SparseMlaWorkspaceSize.restype = c_int32 +LIBINFINIOP.infiniopGetFp8SparseMlaWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_uint64), +] +LIBINFINIOP.infiniopFp8SparseMla.restype = c_int32 +LIBINFINIOP.infiniopFp8SparseMla.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_uint64, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, +] +LIBINFINIOP.infiniopDestroyFp8SparseMlaDescriptor.restype = c_int32 +LIBINFINIOP.infiniopDestroyFp8SparseMlaDescriptor.argtypes = [ + infiniopOperatorDescriptor_t +] + +_LATENT_DIM = 512 +_ROPE_DIM = 64 +_HEAD_DIM = _LATENT_DIM + _ROPE_DIM +_CACHE_STRIDE = _LATENT_DIM + 16 + _ROPE_DIM * 2 +_SCALE = 0.04 +_TEST_CASES = [(2, 2, 70)] +_TENSOR_DTYPES = [None] + + +def test(handle, device, num_tokens, num_heads, topk, _dtype, sync): + print( + f"Testing Fp8SparseMla on {InfiniDeviceNames[device]} with " + f"tokens={num_tokens}, heads={num_heads}, topk={topk}" + ) + num_cache_tokens = 80 + latent = (torch.rand(num_cache_tokens, _LATENT_DIM) * 8.0 - 4.0).to( + torch.float8_e4m3fn + ) + scales = torch.rand(num_cache_tokens, 4, dtype=torch.float32) * 0.02 + 0.002 + rope = (torch.rand(num_cache_tokens, _ROPE_DIM) * 2.0 - 1.0).to( + torch.bfloat16 + ) + cache_src = torch.zeros( + (num_cache_tokens, 1, _CACHE_STRIDE), dtype=torch.uint8 + ) + flat_cache = cache_src[:, 0] + flat_cache[:, :_LATENT_DIM] = latent.view(torch.uint8).reshape( + num_cache_tokens, -1 + ) + flat_cache[:, _LATENT_DIM : _LATENT_DIM + 16] = scales.contiguous().view( + torch.uint8 + ).reshape(num_cache_tokens, -1) + flat_cache[:, _LATENT_DIM + 16 :] = rope.view(torch.uint8).reshape( + num_cache_tokens, -1 + ) + query_src = (torch.rand(num_tokens, num_heads, _HEAD_DIM) - 0.5).to( + torch.bfloat16 + ) + indices_src = torch.arange(topk, dtype=torch.int32).repeat(num_tokens, 1, 1) + indices_src[0, 0, 2] = -1 + indices_src[1, 0, 67] = num_cache_tokens + topk_lens_src = torch.tensor([5, topk], dtype=torch.int32) + + dequantized = latent.float().reshape(num_cache_tokens, 4, 128) + dequantized = (dequantized * scales[:, :, None]).reshape( + num_cache_tokens, _LATENT_DIM + ) + keys = torch.cat([dequantized, rope.float()], dim=-1) + expected = torch.zeros( + (num_tokens, num_heads, _LATENT_DIM), dtype=torch.float32 + ) + for token in range(num_tokens): + valid_indices = [ + int(index) + for index in indices_src[token, 0, : int(topk_lens_src[token])] + if 0 <= int(index) < num_cache_tokens + ] + for head in range(num_heads): + logits = ( + keys[valid_indices] @ query_src[token, head].float() + ) * _SCALE + probabilities = torch.softmax(logits, dim=0) + expected[token, head] = ( + probabilities[:, None] * dequantized[valid_indices] + ).sum(dim=0) + + output = TestTensor( + (num_tokens, num_heads, _LATENT_DIM), + None, + InfiniDtype.BF16, + device, + mode="zeros", + ) + query = TestTensor.from_torch(query_src, InfiniDtype.BF16, device) + cache = TestTensor.from_torch(cache_src, InfiniDtype.U8, device) + indices = TestTensor.from_torch(indices_src, InfiniDtype.I32, device) + topk_lens = TestTensor.from_torch( + topk_lens_src, InfiniDtype.I32, device + ) + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateFp8SparseMlaDescriptor( + handle, + ctypes.byref(descriptor), + output.descriptor, + query.descriptor, + cache.descriptor, + indices.descriptor, + topk_lens.descriptor, + _SCALE, + ) + ) + for tensor in (output, query, cache, indices, topk_lens): + tensor.destroy_desc() + workspace_size = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetFp8SparseMlaWorkspaceSize( + descriptor, ctypes.byref(workspace_size) + ) + ) + workspace = TestWorkspace(workspace_size.value, output.device) + check_error( + LIBINFINIOP.infiniopFp8SparseMla( + descriptor, + workspace.data(), + workspace_size.value, + output.data(), + query.data(), + cache.data(), + indices.data(), + topk_lens.data(), + None, + ) + ) + if sync is not None: + sync() + + assert torch.allclose( + output.actual_tensor().cpu().float(), + expected, + atol=3.0e-2, + rtol=2.0e-2, + ) + check_error(LIBINFINIOP.infiniopDestroyFp8SparseMlaDescriptor(descriptor)) + + +if __name__ == "__main__": + args = get_args() + torch.manual_seed(3) + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + print("\033[92mFp8SparseMla test passed!\033[0m") From f90a1dce2233ef9c82b9a6649f805e73c487ac84 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Wed, 29 Jul 2026 01:37:06 +0000 Subject: [PATCH 6/8] feat(infinicore): wrap GLM FP8 operators --- include/infinicore/ops.hpp | 5 + include/infinicore/ops/fp8_indexer_logits.hpp | 22 ++ include/infinicore/ops/fp8_indexer_quant.hpp | 39 ++++ .../infinicore/ops/fp8_mla_rmsnorm_cache.hpp | 37 ++++ include/infinicore/ops/fp8_sparse_mla.hpp | 20 ++ .../ops/select_last_token_hidden.hpp | 13 ++ python/infinicore/__init__.py | 1 + python/infinicore/dtype.py | 1 + python/infinicore/utils.py | 4 + .../fp8_indexer_logits/fp8_indexer_logits.cc | 78 +++++++ .../fp8_indexer_logits_infiniop.cc | 72 ++++++ .../fp8_indexer_quant/fp8_indexer_quant.cc | 55 +++++ .../fp8_indexer_quant_infiniop.cc | 58 +++++ .../fp8_indexer_quant/fused_fp8_indexer.cc | 99 +++++++++ .../fused_fp8_indexer_infiniop.cc | 72 ++++++ .../fp8_mla_rmsnorm_cache.cc | 155 +++++++++++++ .../fp8_mla_rmsnorm_cache_infiniop.cc | 112 ++++++++++ .../ops/fp8_sparse_mla/fp8_sparse_mla.cc | 67 ++++++ .../fp8_sparse_mla/fp8_sparse_mla_infiniop.cc | 71 ++++++ .../select_last_token_hidden.cc | 54 +++++ .../select_last_token_hidden_infiniop.cc | 53 +++++ src/infinicore/pybind11/ops.hpp | 10 + .../pybind11/ops/fp8_indexer_logits.hpp | 20 ++ .../pybind11/ops/fp8_indexer_quant.hpp | 32 +++ .../pybind11/ops/fp8_mla_rmsnorm_cache.hpp | 28 +++ .../pybind11/ops/fp8_sparse_mla.hpp | 19 ++ .../pybind11/ops/select_last_token_hidden.hpp | 16 ++ test/infinicore/ops/glm_fp8_wrappers.py | 208 ++++++++++++++++++ 28 files changed, 1421 insertions(+) create mode 100644 include/infinicore/ops/fp8_indexer_logits.hpp create mode 100644 include/infinicore/ops/fp8_indexer_quant.hpp create mode 100644 include/infinicore/ops/fp8_mla_rmsnorm_cache.hpp create mode 100644 include/infinicore/ops/fp8_sparse_mla.hpp create mode 100644 include/infinicore/ops/select_last_token_hidden.hpp create mode 100644 src/infinicore/ops/fp8_indexer_logits/fp8_indexer_logits.cc create mode 100644 src/infinicore/ops/fp8_indexer_logits/fp8_indexer_logits_infiniop.cc create mode 100644 src/infinicore/ops/fp8_indexer_quant/fp8_indexer_quant.cc create mode 100644 src/infinicore/ops/fp8_indexer_quant/fp8_indexer_quant_infiniop.cc create mode 100644 src/infinicore/ops/fp8_indexer_quant/fused_fp8_indexer.cc create mode 100644 src/infinicore/ops/fp8_indexer_quant/fused_fp8_indexer_infiniop.cc create mode 100644 src/infinicore/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache.cc create mode 100644 src/infinicore/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache_infiniop.cc create mode 100644 src/infinicore/ops/fp8_sparse_mla/fp8_sparse_mla.cc create mode 100644 src/infinicore/ops/fp8_sparse_mla/fp8_sparse_mla_infiniop.cc create mode 100644 src/infinicore/ops/select_last_token_hidden/select_last_token_hidden.cc create mode 100644 src/infinicore/ops/select_last_token_hidden/select_last_token_hidden_infiniop.cc create mode 100644 src/infinicore/pybind11/ops/fp8_indexer_logits.hpp create mode 100644 src/infinicore/pybind11/ops/fp8_indexer_quant.hpp create mode 100644 src/infinicore/pybind11/ops/fp8_mla_rmsnorm_cache.hpp create mode 100644 src/infinicore/pybind11/ops/fp8_sparse_mla.hpp create mode 100644 src/infinicore/pybind11/ops/select_last_token_hidden.hpp create mode 100644 test/infinicore/ops/glm_fp8_wrappers.py diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index ed3abe502..7a4ec572a 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -29,6 +29,10 @@ #include "ops/flash_attention.hpp" #include "ops/fmin.hpp" #include "ops/fmod.hpp" +#include "ops/fp8_indexer_logits.hpp" +#include "ops/fp8_indexer_quant.hpp" +#include "ops/fp8_mla_rmsnorm_cache.hpp" +#include "ops/fp8_sparse_mla.hpp" #include "ops/fused_gated_delta_net_gating.hpp" #include "ops/fused_moe.hpp" #include "ops/gelu.hpp" @@ -68,6 +72,7 @@ #include "ops/rotmg.hpp" #include "ops/rwkv5_wkv.hpp" #include "ops/scal.hpp" +#include "ops/select_last_token_hidden.hpp" #include "ops/sigmoid.hpp" #include "ops/silu.hpp" #include "ops/silu_and_mul.hpp" diff --git a/include/infinicore/ops/fp8_indexer_logits.hpp b/include/infinicore/ops/fp8_indexer_logits.hpp new file mode 100644 index 000000000..523eb00c1 --- /dev/null +++ b/include/infinicore/ops/fp8_indexer_logits.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS( + Fp8IndexerLogits, + Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, + const Tensor &, const Tensor &); + +void fp8_indexer_logits_( + Tensor logits, + const Tensor &q_fp8, + const Tensor &kv_cache, + const Tensor &block_tables, + const Tensor &weights_fp32, + const Tensor &positions, + const Tensor &request_ids); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/fp8_indexer_quant.hpp b/include/infinicore/ops/fp8_indexer_quant.hpp new file mode 100644 index 000000000..8af946005 --- /dev/null +++ b/include/infinicore/ops/fp8_indexer_quant.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS( + Fp8IndexerQuant, Tensor, Tensor, const Tensor &, const Tensor &); + +void fp8_indexer_quant_( + Tensor q_fp8, + Tensor weights_fp32, + const Tensor &q, + const Tensor &weights); + +INFINICORE_GRAPH_OP_CLASS( + FusedFp8Indexer, + Tensor, Tensor, Tensor, + const Tensor &, const Tensor &, const Tensor &, const Tensor &, + const Tensor &, const Tensor &, const Tensor &, + size_t, double, double); + +void fused_fp8_indexer_( + Tensor q_fp8, + Tensor weights_fp32, + Tensor k_cache, + const Tensor &q_raw, + const Tensor &k_weights, + const Tensor &norm_weight, + const Tensor &norm_bias, + const Tensor &positions, + const Tensor &cos_sin_cache, + const Tensor &slot_mapping, + size_t rope_dim, + double eps, + double weights_scale); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/fp8_mla_rmsnorm_cache.hpp b/include/infinicore/ops/fp8_mla_rmsnorm_cache.hpp new file mode 100644 index 000000000..ebb55bcb1 --- /dev/null +++ b/include/infinicore/ops/fp8_mla_rmsnorm_cache.hpp @@ -0,0 +1,37 @@ +#pragma once + +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS( + Fp8MlaRmsnormCache, + Tensor, + const Tensor &, const Tensor &, const Tensor &, const Tensor &, + double); + +INFINICORE_GRAPH_OP_CLASS( + Fp8MlaRmsnormDualCache, + Tensor, Tensor, + const Tensor &, const Tensor &, const Tensor &, const Tensor &, + double); + +void fp8_mla_rmsnorm_cache_( + Tensor cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps); + +void fp8_mla_rmsnorm_dual_cache_( + Tensor cache, + Tensor vendor_cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/fp8_sparse_mla.hpp b/include/infinicore/ops/fp8_sparse_mla.hpp new file mode 100644 index 000000000..6ad6fa79d --- /dev/null +++ b/include/infinicore/ops/fp8_sparse_mla.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS( + Fp8SparseMla, + Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, float); + +void fp8_sparse_mla_( + Tensor output, + const Tensor &query, + const Tensor &kv_cache, + const Tensor &indices, + const Tensor &topk_lens, + float scale); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/select_last_token_hidden.hpp b/include/infinicore/ops/select_last_token_hidden.hpp new file mode 100644 index 000000000..9b9907051 --- /dev/null +++ b/include/infinicore/ops/select_last_token_hidden.hpp @@ -0,0 +1,13 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(SelectLastTokenHidden, Tensor, const Tensor &, const Tensor &); + +void select_last_token_hidden_(Tensor output, const Tensor &hidden_states, const Tensor &input_offsets); + +} // namespace infinicore::op diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index a74dad116..43682bc54 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -38,6 +38,7 @@ float16, float32, float64, + float8, half, int, int8, diff --git a/python/infinicore/dtype.py b/python/infinicore/dtype.py index a323471c2..69769bceb 100644 --- a/python/infinicore/dtype.py +++ b/python/infinicore/dtype.py @@ -66,6 +66,7 @@ def __hash__(self): cdouble = complex128 float16 = dtype(_infinicore.DataType.F16) half = float16 +float8 = dtype(_infinicore.DataType.F8) bfloat16 = dtype(_infinicore.DataType.BF16) uint8 = dtype(_infinicore.DataType.U8) int8 = dtype(_infinicore.DataType.I8) diff --git a/python/infinicore/utils.py b/python/infinicore/utils.py index fd39db0e1..6ef8c1ef1 100644 --- a/python/infinicore/utils.py +++ b/python/infinicore/utils.py @@ -19,6 +19,8 @@ def to_torch_dtype(infini_dtype): return torch.float64 elif infini_dtype == infinicore.bfloat16: return torch.bfloat16 + elif infini_dtype == infinicore.float8: + return torch.float8_e4m3fn elif infini_dtype == infinicore.int8: return torch.int8 elif infini_dtype == infinicore.int16: @@ -43,6 +45,8 @@ def to_infinicore_dtype(torch_dtype): return infinicore.float16 elif torch_dtype == torch.bfloat16: return infinicore.bfloat16 + elif torch_dtype == torch.float8_e4m3fn: + return infinicore.float8 elif torch_dtype == torch.int8: return infinicore.int8 elif torch_dtype == torch.int16: diff --git a/src/infinicore/ops/fp8_indexer_logits/fp8_indexer_logits.cc b/src/infinicore/ops/fp8_indexer_logits/fp8_indexer_logits.cc new file mode 100644 index 000000000..159ee38c9 --- /dev/null +++ b/src/infinicore/ops/fp8_indexer_logits/fp8_indexer_logits.cc @@ -0,0 +1,78 @@ +#include "infinicore/ops/fp8_indexer_logits.hpp" + +#include "../../utils.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Fp8IndexerLogits); + +Fp8IndexerLogits::Fp8IndexerLogits( + Tensor logits, + const Tensor &q_fp8, + const Tensor &kv_cache, + const Tensor &block_tables, + const Tensor &weights_fp32, + const Tensor &positions, + const Tensor &request_ids) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + logits, q_fp8, kv_cache, block_tables, weights_fp32, positions, request_ids); + INFINICORE_GRAPH_OP_DISPATCH( + q_fp8->device().getType(), logits, q_fp8, kv_cache, block_tables, + weights_fp32, positions, request_ids); +} + +void Fp8IndexerLogits::execute( + Tensor logits, + const Tensor &q_fp8, + const Tensor &kv_cache, + const Tensor &block_tables, + const Tensor &weights_fp32, + const Tensor &positions, + const Tensor &request_ids) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + Fp8IndexerLogits, logits, q_fp8, kv_cache, block_tables, + weights_fp32, positions, request_ids); +} + +void fp8_indexer_logits_( + Tensor logits, + const Tensor &q_fp8, + const Tensor &kv_cache, + const Tensor &block_tables, + const Tensor &weights_fp32, + const Tensor &positions, + const Tensor &request_ids) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + logits, q_fp8, kv_cache, block_tables, weights_fp32, positions, request_ids); + if (logits->ndim() != 2 || q_fp8->ndim() != 3 || kv_cache->ndim() != 3 + || block_tables->ndim() != 2 || weights_fp32->ndim() != 2 + || positions->ndim() != 1 || request_ids->ndim() != 1 + || logits->size(0) != q_fp8->size(0) + || weights_fp32->size(0) != q_fp8->size(0) + || weights_fp32->size(1) != q_fp8->size(1) + || positions->numel() != q_fp8->size(0) + || request_ids->numel() != q_fp8->size(0) + || kv_cache->size(1) != 64 || q_fp8->size(2) != 128 + || kv_cache->size(2) != q_fp8->size(2) + sizeof(float)) { + throw std::runtime_error("fp8_indexer_logits shape mismatch"); + } + if (logits->dtype() != DataType::F32 || q_fp8->dtype() != DataType::F8 + || kv_cache->dtype() != DataType::U8 + || block_tables->dtype() != DataType::I32 + || weights_fp32->dtype() != DataType::F32 + || positions->dtype() != DataType::I64 + || request_ids->dtype() != DataType::I32) { + throw std::runtime_error("fp8_indexer_logits dtype mismatch"); + } + for (const auto &tensor : {logits, q_fp8, kv_cache, block_tables, + weights_fp32, positions, request_ids}) { + if (!tensor->is_contiguous()) { + throw std::runtime_error("fp8_indexer_logits expects contiguous tensors"); + } + } + Fp8IndexerLogits::execute( + logits, q_fp8, kv_cache, block_tables, weights_fp32, positions, request_ids); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/fp8_indexer_logits/fp8_indexer_logits_infiniop.cc b/src/infinicore/ops/fp8_indexer_logits/fp8_indexer_logits_infiniop.cc new file mode 100644 index 000000000..46cfa7eef --- /dev/null +++ b/src/infinicore/ops/fp8_indexer_logits/fp8_indexer_logits_infiniop.cc @@ -0,0 +1,72 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/fp8_indexer_logits.hpp" + +namespace infinicore::op::fp8_indexer_logits_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Fp8IndexerLogits, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor logits, q_fp8, kv_cache, block_tables, weights_fp32; + graph::GraphTensor positions, request_ids; +}; + +void *plan( + Tensor logits, + const Tensor &q_fp8, + const Tensor &kv_cache, + const Tensor &block_tables, + const Tensor &weights_fp32, + const Tensor &positions, + const Tensor &request_ids) { + const size_t seed = hash_combine( + logits, q_fp8, kv_cache, block_tables, weights_fp32, positions, request_ids); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, + descriptor, + Fp8IndexerLogits, + seed, + logits->desc(), + q_fp8->desc(), + kv_cache->desc(), + block_tables->desc(), + weights_fp32->desc(), + positions->desc(), + request_ids->desc()); + return new PlannedMeta{ + descriptor, + graph::GraphTensor(logits), + graph::GraphTensor(q_fp8), + graph::GraphTensor(kv_cache), + graph::GraphTensor(block_tables), + graph::GraphTensor(weights_fp32), + graph::GraphTensor(positions), + graph::GraphTensor(request_ids)}; +} + +void run(void *planned_meta) { + auto *planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopFp8IndexerLogits( + planned->descriptor->desc, + planned->logits->data(), + planned->q_fp8->data(), + planned->kv_cache->data(), + planned->block_tables->data(), + planned->weights_fp32->data(), + planned->positions->data(), + planned->request_ids->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE( + Fp8IndexerLogits, + &plan, + &run, + &cleanup); + +} // namespace infinicore::op::fp8_indexer_logits_impl::infiniop diff --git a/src/infinicore/ops/fp8_indexer_quant/fp8_indexer_quant.cc b/src/infinicore/ops/fp8_indexer_quant/fp8_indexer_quant.cc new file mode 100644 index 000000000..a73f24cf8 --- /dev/null +++ b/src/infinicore/ops/fp8_indexer_quant/fp8_indexer_quant.cc @@ -0,0 +1,55 @@ +#include "infinicore/ops/fp8_indexer_quant.hpp" + +#include "../../utils.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Fp8IndexerQuant); + +Fp8IndexerQuant::Fp8IndexerQuant( + Tensor q_fp8, + Tensor weights_fp32, + const Tensor &q, + const Tensor &weights) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q_fp8, weights_fp32, q, weights); + INFINICORE_GRAPH_OP_DISPATCH( + q->device().getType(), q_fp8, weights_fp32, q, weights); +} + +void Fp8IndexerQuant::execute( + Tensor q_fp8, + Tensor weights_fp32, + const Tensor &q, + const Tensor &weights) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + Fp8IndexerQuant, q_fp8, weights_fp32, q, weights); +} + +void fp8_indexer_quant_( + Tensor q_fp8, + Tensor weights_fp32, + const Tensor &q, + const Tensor &weights) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q_fp8, weights_fp32, q, weights); + if (q->ndim() != 3 || weights->ndim() != 2 + || q_fp8->shape() != q->shape() + || weights_fp32->shape() != weights->shape() + || q->size(0) != weights->size(0) + || q->size(1) != weights->size(1)) { + throw std::runtime_error("fp8_indexer_quant shape mismatch"); + } + if ((q->dtype() != DataType::F16 && q->dtype() != DataType::BF16) + || weights->dtype() != q->dtype() + || q_fp8->dtype() != DataType::F8 + || weights_fp32->dtype() != DataType::F32) { + throw std::runtime_error("fp8_indexer_quant dtype mismatch"); + } + if (!q->is_contiguous() || !weights->is_contiguous() + || !q_fp8->is_contiguous() || !weights_fp32->is_contiguous()) { + throw std::runtime_error("fp8_indexer_quant expects contiguous tensors"); + } + Fp8IndexerQuant::execute(q_fp8, weights_fp32, q, weights); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/fp8_indexer_quant/fp8_indexer_quant_infiniop.cc b/src/infinicore/ops/fp8_indexer_quant/fp8_indexer_quant_infiniop.cc new file mode 100644 index 000000000..4a8b32a12 --- /dev/null +++ b/src/infinicore/ops/fp8_indexer_quant/fp8_indexer_quant_infiniop.cc @@ -0,0 +1,58 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/fp8_indexer_quant.hpp" + +namespace infinicore::op::fp8_indexer_quant_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Fp8IndexerQuant, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor q_fp8, weights_fp32, q, weights; +}; + +void *plan( + Tensor q_fp8, + Tensor weights_fp32, + const Tensor &q, + const Tensor &weights) { + const size_t seed = hash_combine(q_fp8, weights_fp32, q, weights); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, + descriptor, + Fp8IndexerQuant, + seed, + q_fp8->desc(), + weights_fp32->desc(), + q->desc(), + weights->desc()); + return new PlannedMeta{ + descriptor, + graph::GraphTensor(q_fp8), + graph::GraphTensor(weights_fp32), + graph::GraphTensor(q), + graph::GraphTensor(weights)}; +} + +void run(void *planned_meta) { + auto *planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopFp8IndexerQuant( + planned->descriptor->desc, + planned->q_fp8->data(), + planned->weights_fp32->data(), + planned->q->data(), + planned->weights->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE( + Fp8IndexerQuant, + &plan, + &run, + &cleanup); + +} // namespace infinicore::op::fp8_indexer_quant_impl::infiniop diff --git a/src/infinicore/ops/fp8_indexer_quant/fused_fp8_indexer.cc b/src/infinicore/ops/fp8_indexer_quant/fused_fp8_indexer.cc new file mode 100644 index 000000000..4f9ce7032 --- /dev/null +++ b/src/infinicore/ops/fp8_indexer_quant/fused_fp8_indexer.cc @@ -0,0 +1,99 @@ +#include "infinicore/ops/fp8_indexer_quant.hpp" + +#include "../../utils.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(FusedFp8Indexer); + +FusedFp8Indexer::FusedFp8Indexer( + Tensor q_fp8, Tensor weights_fp32, Tensor k_cache, + const Tensor &q_raw, const Tensor &k_weights, + const Tensor &norm_weight, const Tensor &norm_bias, + const Tensor &positions, const Tensor &cos_sin_cache, + const Tensor &slot_mapping, size_t rope_dim, + double eps, double weights_scale) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + q_fp8, weights_fp32, k_cache, q_raw, k_weights, norm_weight, + norm_bias, positions, cos_sin_cache, slot_mapping); + INFINICORE_GRAPH_OP_DISPATCH( + q_raw->device().getType(), q_fp8, weights_fp32, k_cache, + q_raw, k_weights, norm_weight, norm_bias, positions, + cos_sin_cache, slot_mapping, rope_dim, eps, weights_scale); +} + +void FusedFp8Indexer::execute( + Tensor q_fp8, Tensor weights_fp32, Tensor k_cache, + const Tensor &q_raw, const Tensor &k_weights, + const Tensor &norm_weight, const Tensor &norm_bias, + const Tensor &positions, const Tensor &cos_sin_cache, + const Tensor &slot_mapping, size_t rope_dim, + double eps, double weights_scale) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + FusedFp8Indexer, q_fp8, weights_fp32, k_cache, + q_raw, k_weights, norm_weight, norm_bias, positions, + cos_sin_cache, slot_mapping, rope_dim, eps, weights_scale); +} + +void fused_fp8_indexer_( + Tensor q_fp8, Tensor weights_fp32, Tensor k_cache, + const Tensor &q_raw, const Tensor &k_weights, + const Tensor &norm_weight, const Tensor &norm_bias, + const Tensor &positions, const Tensor &cos_sin_cache, + const Tensor &slot_mapping, size_t rope_dim, + double eps, double weights_scale) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + q_fp8, weights_fp32, k_cache, q_raw, k_weights, norm_weight, + norm_bias, positions, cos_sin_cache, slot_mapping); + if (q_raw->ndim() != 3 || k_weights->ndim() != 2 + || weights_fp32->ndim() != 2 || k_cache->ndim() != 3 + || norm_weight->ndim() != 1 || norm_bias->ndim() != 1 + || positions->ndim() != 1 || cos_sin_cache->ndim() != 2 + || slot_mapping->ndim() != 1 + || q_fp8->shape() != q_raw->shape() + || weights_fp32->size(0) != q_raw->size(0) + || weights_fp32->size(1) != q_raw->size(1) + || k_weights->size(0) != q_raw->size(0) + || k_weights->size(1) != q_raw->size(2) + q_raw->size(1) + || norm_weight->numel() != q_raw->size(2) + || norm_bias->numel() != q_raw->size(2) + || positions->numel() != q_raw->size(0) + || slot_mapping->numel() != q_raw->size(0) + || cos_sin_cache->size(1) != rope_dim + || k_cache->size(2) != q_raw->size(2) + sizeof(float)) { + throw std::runtime_error("fused_fp8_indexer shape mismatch"); + } + if ((q_raw->dtype() != DataType::F16 + && q_raw->dtype() != DataType::BF16) + || k_weights->dtype() != q_raw->dtype() + || norm_weight->dtype() != q_raw->dtype() + || norm_bias->dtype() != q_raw->dtype() + || cos_sin_cache->dtype() != q_raw->dtype() + || q_fp8->dtype() != DataType::F8 + || weights_fp32->dtype() != DataType::F32 + || k_cache->dtype() != DataType::U8 + || positions->dtype() != DataType::I64 + || slot_mapping->dtype() != DataType::I64) { + throw std::runtime_error("fused_fp8_indexer dtype mismatch"); + } + for (const auto &tensor : { + q_fp8, weights_fp32, k_cache, q_raw, k_weights, + norm_weight, norm_bias, positions, cos_sin_cache, slot_mapping}) { + if (!tensor->is_contiguous()) { + throw std::runtime_error( + "fused_fp8_indexer expects contiguous tensors"); + } + } + if (q_raw->size(2) != 128 || rope_dim == 0 + || rope_dim > q_raw->size(2) || rope_dim % 2 != 0 + || eps <= 0.0 || weights_scale <= 0.0) { + throw std::runtime_error("fused_fp8_indexer invalid parameters"); + } + FusedFp8Indexer::execute( + q_fp8, weights_fp32, k_cache, q_raw, k_weights, + norm_weight, norm_bias, positions, cos_sin_cache, + slot_mapping, rope_dim, eps, weights_scale); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/fp8_indexer_quant/fused_fp8_indexer_infiniop.cc b/src/infinicore/ops/fp8_indexer_quant/fused_fp8_indexer_infiniop.cc new file mode 100644 index 000000000..f7cfba105 --- /dev/null +++ b/src/infinicore/ops/fp8_indexer_quant/fused_fp8_indexer_infiniop.cc @@ -0,0 +1,72 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/fp8_indexer_quant.hpp" + +namespace infinicore::op::fused_fp8_indexer_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, FusedFp8Indexer, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor q_fp8, weights_fp32, k_cache; + graph::GraphTensor q_raw, k_weights, norm_weight, norm_bias; + graph::GraphTensor positions, cos_sin_cache, slot_mapping; +}; + +void *plan( + Tensor q_fp8, Tensor weights_fp32, Tensor k_cache, + const Tensor &q_raw, const Tensor &k_weights, + const Tensor &norm_weight, const Tensor &norm_bias, + const Tensor &positions, const Tensor &cos_sin_cache, + const Tensor &slot_mapping, size_t rope_dim, + double eps, double weights_scale) { + const size_t seed = hash_combine( + q_fp8, weights_fp32, k_cache, q_raw, k_weights, + norm_weight, norm_bias, positions, cos_sin_cache, slot_mapping, + rope_dim, eps, weights_scale); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, FusedFp8Indexer, seed, + q_fp8->desc(), weights_fp32->desc(), k_cache->desc(), + q_raw->desc(), k_weights->desc(), norm_weight->desc(), + norm_bias->desc(), positions->desc(), cos_sin_cache->desc(), + slot_mapping->desc(), static_cast(rope_dim), + eps, weights_scale); + return new PlannedMeta{ + descriptor, + graph::GraphTensor(q_fp8), + graph::GraphTensor(weights_fp32), + graph::GraphTensor(k_cache), + graph::GraphTensor(q_raw), + graph::GraphTensor(k_weights), + graph::GraphTensor(norm_weight), + graph::GraphTensor(norm_bias), + graph::GraphTensor(positions), + graph::GraphTensor(cos_sin_cache), + graph::GraphTensor(slot_mapping)}; +} + +void run(void *planned_meta) { + auto *planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopFusedFp8Indexer( + planned->descriptor->desc, + planned->q_fp8->data(), + planned->weights_fp32->data(), + planned->k_cache->data(), + planned->q_raw->data(), + planned->k_weights->data(), + planned->norm_weight->data(), + planned->norm_bias->data(), + planned->positions->data(), + planned->cos_sin_cache->data(), + planned->slot_mapping->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE( + FusedFp8Indexer, &plan, &run, &cleanup); + +} // namespace infinicore::op::fused_fp8_indexer_impl::infiniop diff --git a/src/infinicore/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache.cc b/src/infinicore/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache.cc new file mode 100644 index 000000000..708d868b0 --- /dev/null +++ b/src/infinicore/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache.cc @@ -0,0 +1,155 @@ +#include "infinicore/ops/fp8_mla_rmsnorm_cache.hpp" + +#include "../../utils.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Fp8MlaRmsnormCache); +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Fp8MlaRmsnormDualCache); + +Fp8MlaRmsnormCache::Fp8MlaRmsnormCache( + Tensor cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + cache, compressed_kv, norm_weight, rope, slot_mapping); + INFINICORE_GRAPH_OP_DISPATCH( + cache->device().getType(), cache, compressed_kv, norm_weight, + rope, slot_mapping, eps); +} + +void Fp8MlaRmsnormCache::execute( + Tensor cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + Fp8MlaRmsnormCache, cache, compressed_kv, norm_weight, + rope, slot_mapping, eps); +} + +Fp8MlaRmsnormDualCache::Fp8MlaRmsnormDualCache( + Tensor cache, + Tensor vendor_cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + cache, vendor_cache, compressed_kv, norm_weight, rope, slot_mapping); + INFINICORE_GRAPH_OP_DISPATCH( + cache->device().getType(), cache, vendor_cache, compressed_kv, + norm_weight, rope, slot_mapping, eps); +} + +void Fp8MlaRmsnormDualCache::execute( + Tensor cache, + Tensor vendor_cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + Fp8MlaRmsnormDualCache, cache, vendor_cache, compressed_kv, + norm_weight, rope, slot_mapping, eps); +} + +void fp8_mla_rmsnorm_cache_( + Tensor cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + cache, compressed_kv, norm_weight, rope, slot_mapping); + if (cache->ndim() != 3 || compressed_kv->ndim() != 2 + || norm_weight->ndim() != 1 || rope->ndim() != 2 + || slot_mapping->ndim() != 1 + || compressed_kv->size(1) != 512 + || norm_weight->numel() != 512 + || rope->size(0) != compressed_kv->size(0) + || rope->size(1) != 64 + || slot_mapping->numel() != compressed_kv->size(0) + || cache->size(2) != 656) { + throw std::runtime_error("fp8_mla_rmsnorm_cache shape mismatch"); + } + if (cache->dtype() != DataType::U8 + || compressed_kv->dtype() != DataType::BF16 + || norm_weight->dtype() != DataType::BF16 + || rope->dtype() != DataType::BF16 + || slot_mapping->dtype() != DataType::I64) { + throw std::runtime_error("fp8_mla_rmsnorm_cache dtype mismatch"); + } + for (const auto &tensor : { + cache, compressed_kv, norm_weight, rope, slot_mapping}) { + if (!tensor->is_contiguous()) { + throw std::runtime_error( + "fp8_mla_rmsnorm_cache expects contiguous tensors"); + } + } + if (eps <= 0.0) { + throw std::runtime_error("fp8_mla_rmsnorm_cache requires eps > 0"); + } + Fp8MlaRmsnormCache::execute( + cache, compressed_kv, norm_weight, rope, slot_mapping, eps); +} + +void fp8_mla_rmsnorm_dual_cache_( + Tensor cache, + Tensor vendor_cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + cache, vendor_cache, compressed_kv, norm_weight, rope, slot_mapping); + if (cache->ndim() != 3 || vendor_cache->ndim() != 3 + || compressed_kv->ndim() != 2 || norm_weight->ndim() != 1 + || rope->ndim() != 2 || slot_mapping->ndim() != 1 + || compressed_kv->size(1) != 512 + || norm_weight->numel() != 512 + || rope->size(0) != compressed_kv->size(0) + || rope->size(1) != 64 + || slot_mapping->numel() != compressed_kv->size(0) + || cache->size(2) != 656 + || vendor_cache->size(0) != cache->size(0) + || vendor_cache->size(1) != cache->size(1) + || vendor_cache->size(2) != 576) { + throw std::runtime_error("fp8_mla_rmsnorm_dual_cache shape mismatch"); + } + if (cache->dtype() != DataType::U8 + || vendor_cache->dtype() != DataType::BF16 + || compressed_kv->dtype() != DataType::BF16 + || norm_weight->dtype() != DataType::BF16 + || rope->dtype() != DataType::BF16 + || slot_mapping->dtype() != DataType::I64) { + throw std::runtime_error("fp8_mla_rmsnorm_dual_cache dtype mismatch"); + } + for (const auto &tensor : { + cache, vendor_cache, compressed_kv, norm_weight, rope, + slot_mapping}) { + if (!tensor->is_contiguous()) { + throw std::runtime_error( + "fp8_mla_rmsnorm_dual_cache expects contiguous tensors"); + } + } + if (eps <= 0.0) { + throw std::runtime_error( + "fp8_mla_rmsnorm_dual_cache requires eps > 0"); + } + Fp8MlaRmsnormDualCache::execute( + cache, vendor_cache, compressed_kv, norm_weight, rope, + slot_mapping, eps); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache_infiniop.cc b/src/infinicore/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache_infiniop.cc new file mode 100644 index 000000000..a8a30862f --- /dev/null +++ b/src/infinicore/ops/fp8_mla_rmsnorm_cache/fp8_mla_rmsnorm_cache_infiniop.cc @@ -0,0 +1,112 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/fp8_mla_rmsnorm_cache.hpp" + +namespace infinicore::op::fp8_mla_rmsnorm_cache_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Fp8MlaRmsnormCache, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor cache, compressed_kv, norm_weight, rope, slot_mapping; +}; + +void *plan( + Tensor cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps) { + const size_t seed = hash_combine( + cache, compressed_kv, norm_weight, rope, slot_mapping, eps); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, Fp8MlaRmsnormCache, seed, + cache->desc(), nullptr, compressed_kv->desc(), + norm_weight->desc(), rope->desc(), slot_mapping->desc(), eps); + return new PlannedMeta{ + descriptor, + graph::GraphTensor(cache), + graph::GraphTensor(compressed_kv), + graph::GraphTensor(norm_weight), + graph::GraphTensor(rope), + graph::GraphTensor(slot_mapping)}; +} + +void run(void *planned_meta) { + auto *planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopFp8MlaRmsnormCache( + planned->descriptor->desc, + planned->cache->data(), + nullptr, + planned->compressed_kv->data(), + planned->norm_weight->data(), + planned->rope->data(), + planned->slot_mapping->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE( + Fp8MlaRmsnormCache, &plan, &run, &cleanup); + +namespace dual { + +struct DualPlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor cache, vendor_cache, compressed_kv, norm_weight, rope, + slot_mapping; +}; + +void *plan_dual( + Tensor cache, + Tensor vendor_cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps) { + const size_t seed = hash_combine( + cache, vendor_cache, compressed_kv, norm_weight, rope, + slot_mapping, eps); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, Fp8MlaRmsnormCache, seed, + cache->desc(), vendor_cache->desc(), compressed_kv->desc(), + norm_weight->desc(), rope->desc(), slot_mapping->desc(), eps); + return new DualPlannedMeta{ + descriptor, + graph::GraphTensor(cache), + graph::GraphTensor(vendor_cache), + graph::GraphTensor(compressed_kv), + graph::GraphTensor(norm_weight), + graph::GraphTensor(rope), + graph::GraphTensor(slot_mapping)}; +} + +void run_dual(void *planned_meta) { + auto *planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopFp8MlaRmsnormCache( + planned->descriptor->desc, + planned->cache->data(), + planned->vendor_cache->data(), + planned->compressed_kv->data(), + planned->norm_weight->data(), + planned->rope->data(), + planned->slot_mapping->data(), + context::getStream())); +} + +void cleanup_dual(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE( + Fp8MlaRmsnormDualCache, &plan_dual, &run_dual, &cleanup_dual); + +} // namespace dual + +} // namespace infinicore::op::fp8_mla_rmsnorm_cache_impl::infiniop diff --git a/src/infinicore/ops/fp8_sparse_mla/fp8_sparse_mla.cc b/src/infinicore/ops/fp8_sparse_mla/fp8_sparse_mla.cc new file mode 100644 index 000000000..89c85db06 --- /dev/null +++ b/src/infinicore/ops/fp8_sparse_mla/fp8_sparse_mla.cc @@ -0,0 +1,67 @@ +#include "infinicore/ops/fp8_sparse_mla.hpp" + +#include "../../utils.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Fp8SparseMla); + +Fp8SparseMla::Fp8SparseMla( + Tensor output, + const Tensor &query, + const Tensor &kv_cache, + const Tensor &indices, + const Tensor &topk_lens, + float scale) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + output, query, kv_cache, indices, topk_lens); + INFINICORE_GRAPH_OP_DISPATCH( + query->device().getType(), output, query, kv_cache, indices, topk_lens, scale); +} + +void Fp8SparseMla::execute( + Tensor output, + const Tensor &query, + const Tensor &kv_cache, + const Tensor &indices, + const Tensor &topk_lens, + float scale) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + Fp8SparseMla, output, query, kv_cache, indices, topk_lens, scale); +} + +void fp8_sparse_mla_( + Tensor output, + const Tensor &query, + const Tensor &kv_cache, + const Tensor &indices, + const Tensor &topk_lens, + float scale) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + output, query, kv_cache, indices, topk_lens); + if (output->ndim() != 3 || query->ndim() != 3 || kv_cache->ndim() != 3 + || indices->ndim() != 3 || topk_lens->ndim() != 1 + || output->size(0) != query->size(0) + || output->size(1) != query->size(1) + || output->size(2) != 512 || query->size(2) != 576 + || kv_cache->size(1) != 1 || kv_cache->size(2) != 656 + || indices->size(0) != query->size(0) || indices->size(1) != 1 + || topk_lens->numel() != query->size(0)) { + throw std::runtime_error("fp8_sparse_mla shape mismatch"); + } + if (output->dtype() != DataType::BF16 || query->dtype() != DataType::BF16 + || kv_cache->dtype() != DataType::U8 + || indices->dtype() != DataType::I32 + || topk_lens->dtype() != DataType::I32) { + throw std::runtime_error("fp8_sparse_mla dtype mismatch"); + } + for (const auto &tensor : {output, query, kv_cache, indices, topk_lens}) { + if (!tensor->is_contiguous()) { + throw std::runtime_error("fp8_sparse_mla expects contiguous tensors"); + } + } + Fp8SparseMla::execute(output, query, kv_cache, indices, topk_lens, scale); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/fp8_sparse_mla/fp8_sparse_mla_infiniop.cc b/src/infinicore/ops/fp8_sparse_mla/fp8_sparse_mla_infiniop.cc new file mode 100644 index 000000000..b7bafd853 --- /dev/null +++ b/src/infinicore/ops/fp8_sparse_mla/fp8_sparse_mla_infiniop.cc @@ -0,0 +1,71 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/fp8_sparse_mla.hpp" + +namespace infinicore::op::fp8_sparse_mla_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Fp8SparseMla, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, output, query, kv_cache, indices, topk_lens; + float scale; +}; + +void *plan( + Tensor output, + const Tensor &query, + const Tensor &kv_cache, + const Tensor &indices, + const Tensor &topk_lens, + float scale) { + const size_t seed = hash_combine( + output, query, kv_cache, indices, topk_lens, scale); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, + descriptor, + Fp8SparseMla, + seed, + output->desc(), + query->desc(), + kv_cache->desc(), + indices->desc(), + topk_lens->desc(), + scale); + INFINIOP_WORKSPACE_TENSOR(workspace, Fp8SparseMla, descriptor); + return new PlannedMeta{ + descriptor, + graph::GraphTensor(workspace), + graph::GraphTensor(output), + graph::GraphTensor(query), + graph::GraphTensor(kv_cache), + graph::GraphTensor(indices), + graph::GraphTensor(topk_lens), + scale}; +} + +void run(void *planned_meta) { + auto *planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopFp8SparseMla( + planned->descriptor->desc, + planned->workspace->data(), + planned->workspace->numel(), + planned->output->data(), + planned->query->data(), + planned->kv_cache->data(), + planned->indices->data(), + planned->topk_lens->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE( + Fp8SparseMla, + &plan, + &run, + &cleanup); + +} // namespace infinicore::op::fp8_sparse_mla_impl::infiniop diff --git a/src/infinicore/ops/select_last_token_hidden/select_last_token_hidden.cc b/src/infinicore/ops/select_last_token_hidden/select_last_token_hidden.cc new file mode 100644 index 000000000..9074097db --- /dev/null +++ b/src/infinicore/ops/select_last_token_hidden/select_last_token_hidden.cc @@ -0,0 +1,54 @@ +#include "infinicore/ops/select_last_token_hidden.hpp" + +#include "../../utils.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(SelectLastTokenHidden); + +SelectLastTokenHidden::SelectLastTokenHidden( + Tensor output, + const Tensor &hidden_states, + const Tensor &input_offsets) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, hidden_states, input_offsets); + INFINICORE_GRAPH_OP_DISPATCH(output->device().getType(), output, hidden_states, input_offsets); +} + +void SelectLastTokenHidden::execute( + Tensor output, + const Tensor &hidden_states, + const Tensor &input_offsets) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(SelectLastTokenHidden, output, hidden_states, input_offsets); +} + +void select_last_token_hidden_(Tensor output, + const Tensor &hidden_states, + const Tensor &input_offsets) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, hidden_states, input_offsets); + if (hidden_states->ndim() != 3 || output->ndim() != 3 || input_offsets->ndim() != 1) { + throw std::runtime_error( + "select_last_token_hidden expects 3D hidden/output and 1D offsets"); + } + if (input_offsets->dtype() != DataType::I32) { + throw std::runtime_error("select_last_token_hidden expects int32 offsets"); + } + if (output->dtype() != hidden_states->dtype()) { + throw std::runtime_error("select_last_token_hidden output dtype mismatch"); + } + if (input_offsets->numel() < 2) { + throw std::runtime_error("select_last_token_hidden requires at least one request"); + } + const size_t num_requests = input_offsets->numel() - 1; + if (output->size(0) != 1 + || output->size(1) != num_requests + || output->size(2) != hidden_states->size(2)) { + throw std::runtime_error("select_last_token_hidden shape mismatch"); + } + if (!output->is_contiguous() || !hidden_states->is_contiguous() || !input_offsets->is_contiguous()) { + throw std::runtime_error("select_last_token_hidden expects contiguous tensors"); + } + SelectLastTokenHidden::execute(output, hidden_states, input_offsets); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/select_last_token_hidden/select_last_token_hidden_infiniop.cc b/src/infinicore/ops/select_last_token_hidden/select_last_token_hidden_infiniop.cc new file mode 100644 index 000000000..6b28a2eea --- /dev/null +++ b/src/infinicore/ops/select_last_token_hidden/select_last_token_hidden_infiniop.cc @@ -0,0 +1,53 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/select_last_token_hidden.hpp" + +namespace infinicore::op::select_last_token_hidden_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, SelectLastTokenHidden, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor output, hidden_states, input_offsets; +}; + +void *plan(Tensor output, const Tensor &hidden_states, const Tensor &input_offsets) { + const size_t seed = hash_combine(output, hidden_states, input_offsets); + + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, + descriptor, + SelectLastTokenHidden, + seed, + output->desc(), + hidden_states->desc(), + input_offsets->desc()); + + return new PlannedMeta{ + descriptor, + graph::GraphTensor(output), + graph::GraphTensor(hidden_states), + graph::GraphTensor(input_offsets)}; +} + +void run(void *planned_meta) { + auto *planned = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopSelectLastTokenHidden( + planned->descriptor->desc, + planned->output->data(), + planned->hidden_states->data(), + planned->input_offsets->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE( + SelectLastTokenHidden, + &plan, + &run, + &cleanup); + +} // namespace infinicore::op::select_last_token_hidden_impl::infiniop diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp index 7b0a1a9af..cb39473ab 100644 --- a/src/infinicore/pybind11/ops.hpp +++ b/src/infinicore/pybind11/ops.hpp @@ -50,6 +50,10 @@ #include "ops/floor_divide.hpp" #include "ops/fmin.hpp" #include "ops/fmod.hpp" +#include "ops/fp8_indexer_logits.hpp" +#include "ops/fp8_indexer_quant.hpp" +#include "ops/fp8_mla_rmsnorm_cache.hpp" +#include "ops/fp8_sparse_mla.hpp" #include "ops/fused_gated_delta_net_gating.hpp" #include "ops/fused_moe.hpp" #include "ops/gaussian_nll_loss.hpp" @@ -110,6 +114,7 @@ #include "ops/scal.hpp" #include "ops/scatter.hpp" #include "ops/selu.hpp" +#include "ops/select_last_token_hidden.hpp" #include "ops/sigmoid.hpp" #include "ops/silu.hpp" #include "ops/silu_and_mul.hpp" @@ -173,6 +178,10 @@ inline void bind(py::module &m) { bind_hinge_embedding_loss(m); bind_kv_caching(m); bind_fmod(m); + bind_fp8_indexer_logits(m); + bind_fp8_indexer_quant(m); + bind_fp8_mla_rmsnorm_cache(m); + bind_fp8_sparse_mla(m); bind_fused_gated_delta_net_gating(m); bind_fused_moe(m); bind_fmin(m); @@ -270,6 +279,7 @@ inline void bind(py::module &m) { bind_lerp(m); bind_triplet_margin_loss(m); bind_selu(m); + bind_select_last_token_hidden(m); bind_swap(m); bind_sinh(m); bind_layer_norm(m); diff --git a/src/infinicore/pybind11/ops/fp8_indexer_logits.hpp b/src/infinicore/pybind11/ops/fp8_indexer_logits.hpp new file mode 100644 index 000000000..d0779b672 --- /dev/null +++ b/src/infinicore/pybind11/ops/fp8_indexer_logits.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "infinicore/ops/fp8_indexer_logits.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { +inline void bind_fp8_indexer_logits(py::module &m) { + m.def("fp8_indexer_logits_", + &op::fp8_indexer_logits_, + py::arg("logits"), + py::arg("q_fp8"), + py::arg("kv_cache"), + py::arg("block_tables"), + py::arg("weights_fp32"), + py::arg("positions"), + py::arg("request_ids")); +} +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/fp8_indexer_quant.hpp b/src/infinicore/pybind11/ops/fp8_indexer_quant.hpp new file mode 100644 index 000000000..bfa241aae --- /dev/null +++ b/src/infinicore/pybind11/ops/fp8_indexer_quant.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include "infinicore/ops/fp8_indexer_quant.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { +inline void bind_fp8_indexer_quant(py::module &m) { + m.def("fp8_indexer_quant_", + &op::fp8_indexer_quant_, + py::arg("q_fp8"), + py::arg("weights_fp32"), + py::arg("q"), + py::arg("weights")); + m.def("fused_fp8_indexer_", + &op::fused_fp8_indexer_, + py::arg("q_fp8"), + py::arg("weights_fp32"), + py::arg("k_cache"), + py::arg("q_raw"), + py::arg("k_weights"), + py::arg("norm_weight"), + py::arg("norm_bias"), + py::arg("positions"), + py::arg("cos_sin_cache"), + py::arg("slot_mapping"), + py::arg("rope_dim"), + py::arg("eps"), + py::arg("weights_scale")); +} +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/fp8_mla_rmsnorm_cache.hpp b/src/infinicore/pybind11/ops/fp8_mla_rmsnorm_cache.hpp new file mode 100644 index 000000000..d7e2e4d43 --- /dev/null +++ b/src/infinicore/pybind11/ops/fp8_mla_rmsnorm_cache.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include "infinicore/ops/fp8_mla_rmsnorm_cache.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { +inline void bind_fp8_mla_rmsnorm_cache(py::module &m) { + m.def("fp8_mla_rmsnorm_cache_", + &op::fp8_mla_rmsnorm_cache_, + py::arg("cache"), + py::arg("compressed_kv"), + py::arg("norm_weight"), + py::arg("rope"), + py::arg("slot_mapping"), + py::arg("eps")); + m.def("fp8_mla_rmsnorm_dual_cache_", + &op::fp8_mla_rmsnorm_dual_cache_, + py::arg("cache"), + py::arg("vendor_cache"), + py::arg("compressed_kv"), + py::arg("norm_weight"), + py::arg("rope"), + py::arg("slot_mapping"), + py::arg("eps")); +} +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/fp8_sparse_mla.hpp b/src/infinicore/pybind11/ops/fp8_sparse_mla.hpp new file mode 100644 index 000000000..9999044d9 --- /dev/null +++ b/src/infinicore/pybind11/ops/fp8_sparse_mla.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "infinicore/ops/fp8_sparse_mla.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { +inline void bind_fp8_sparse_mla(py::module &m) { + m.def("fp8_sparse_mla_", + &op::fp8_sparse_mla_, + py::arg("output"), + py::arg("query"), + py::arg("kv_cache"), + py::arg("indices"), + py::arg("topk_lens"), + py::arg("scale")); +} +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/select_last_token_hidden.hpp b/src/infinicore/pybind11/ops/select_last_token_hidden.hpp new file mode 100644 index 000000000..7d131ffc5 --- /dev/null +++ b/src/infinicore/pybind11/ops/select_last_token_hidden.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "infinicore/ops/select_last_token_hidden.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { +inline void bind_select_last_token_hidden(py::module &m) { + m.def("select_last_token_hidden_", + &op::select_last_token_hidden_, + py::arg("output"), + py::arg("hidden_states"), + py::arg("input_offsets")); +} +} // namespace infinicore::ops diff --git a/test/infinicore/ops/glm_fp8_wrappers.py b/test/infinicore/ops/glm_fp8_wrappers.py new file mode 100644 index 000000000..bfd89737e --- /dev/null +++ b/test/infinicore/ops/glm_fp8_wrappers.py @@ -0,0 +1,208 @@ +import math + +import torch + +import infinicore +from infinicore.lib import _infinicore + + +def wrap(tensor): + return infinicore.from_torch(tensor) + + +def sync(): + infinicore.sync_device() + + +def test_indexer_quant(): + torch.manual_seed(20) + q = (torch.rand(2, 3, 128, device="cuda") * 8 - 4).to(torch.bfloat16) + weights = (torch.rand(2, 3, device="cuda") * 2 - 1).to(torch.bfloat16) + q_fp8 = torch.zeros_like(q, dtype=torch.float8_e4m3fn) + weights_fp32 = torch.zeros_like(weights, dtype=torch.float32) + tensors = [wrap(x) for x in (q_fp8, weights_fp32, q, weights)] + _infinicore.fp8_indexer_quant_(*(x._underlying for x in tensors)) + sync() + + scales = torch.exp2( + torch.ceil( + torch.log2(q.float().abs().amax(dim=-1).clamp_min(1.0e-4) / 448.0) + ) + ) + expected_q = (q.float() / scales[..., None]).clamp(-448, 448).to( + torch.float8_e4m3fn + ) + torch.testing.assert_close(q_fp8.float(), expected_q.float(), rtol=0, atol=0) + torch.testing.assert_close( + weights_fp32, weights.float() * scales, rtol=1e-6, atol=1e-7 + ) + + +def test_fused_indexer(): + torch.manual_seed(21) + tokens, heads, head_dim, rope_dim = 2, 3, 128, 64 + q_raw = torch.randn( + tokens, heads, head_dim, device="cuda", dtype=torch.bfloat16 + ) + k_weights = torch.randn( + tokens, head_dim + heads, device="cuda", dtype=torch.bfloat16 + ) + norm_weight = torch.ones(head_dim, device="cuda", dtype=torch.bfloat16) + norm_bias = torch.zeros(head_dim, device="cuda", dtype=torch.bfloat16) + positions = torch.tensor([0, 1], device="cuda", dtype=torch.int64) + cos_sin = torch.zeros(2, rope_dim, device="cuda", dtype=torch.bfloat16) + cos_sin[:, : rope_dim // 2] = 1 + slots = torch.tensor([0, 1], device="cuda", dtype=torch.int64) + q_fp8 = torch.zeros_like(q_raw, dtype=torch.float8_e4m3fn) + weights_fp32 = torch.zeros(tokens, heads, device="cuda", dtype=torch.float32) + k_cache = torch.zeros(1, 64, 132, device="cuda", dtype=torch.uint8) + tensors = [ + wrap(x) + for x in ( + q_fp8, + weights_fp32, + k_cache, + q_raw, + k_weights, + norm_weight, + norm_bias, + positions, + cos_sin, + slots, + ) + ] + _infinicore.fused_fp8_indexer_( + *(x._underlying for x in tensors), rope_dim, 1.0e-5, 1.0 + ) + sync() + assert torch.isfinite(q_fp8.float()).all() + assert torch.isfinite(weights_fp32).all() + assert torch.count_nonzero(q_fp8.float()) > 0 + assert torch.count_nonzero(k_cache[:, :2]) > 0 + + +def test_mla_cache_and_sparse(): + torch.manual_seed(22) + tokens = 70 + compressed = torch.randn(tokens, 512, device="cuda", dtype=torch.bfloat16) + norm_weight = torch.randn(512, device="cuda", dtype=torch.bfloat16) + rope = torch.randn(tokens, 64, device="cuda", dtype=torch.bfloat16) + slots = torch.arange(tokens, device="cuda", dtype=torch.int64) + cache = torch.zeros(2, 64, 656, device="cuda", dtype=torch.uint8) + vendor = torch.zeros(2, 64, 576, device="cuda", dtype=torch.bfloat16) + tensors = [wrap(x) for x in (cache, vendor, compressed, norm_weight, rope, slots)] + _infinicore.fp8_mla_rmsnorm_dual_cache_( + *(x._underlying for x in tensors), 1.0e-5 + ) + sync() + + entries = cache.view(-1, 656)[:tokens] + latent_fp8 = entries[:, :512].contiguous().view(torch.float8_e4m3fn) + scales = entries[:, 512:528].contiguous().view(torch.float32) + latent = ( + latent_fp8.float().view(tokens, 4, 128) * scales.view(tokens, 4, 1) + ).reshape(tokens, 512) + cached_rope = entries[:, 528:].contiguous().view(torch.bfloat16) + torch.testing.assert_close( + vendor.view(-1, 576)[:tokens, :512], + latent.to(torch.bfloat16), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + vendor.view(-1, 576)[:tokens, 512:], cached_rope, rtol=0, atol=0 + ) + + query = torch.randn(1, 2, 576, device="cuda", dtype=torch.bfloat16) + indices = torch.arange(tokens, device="cuda", dtype=torch.int32).view( + 1, 1, tokens + ) + lens = torch.tensor([tokens], device="cuda", dtype=torch.int32) + output = torch.zeros(1, 2, 512, device="cuda", dtype=torch.bfloat16) + sparse_tensors = [wrap(x) for x in (output, query, cache.view(-1, 1, 656), indices, lens)] + scale = float(576**-0.5) + _infinicore.fp8_sparse_mla_( + *(x._underlying for x in sparse_tensors), scale + ) + sync() + + keys = torch.cat([latent, cached_rope.float()], dim=-1) + expected = [] + for head in range(query.shape[1]): + logits = keys @ query[0, head].float() * scale + expected.append((torch.softmax(logits, dim=0)[:, None] * latent).sum(0)) + expected = torch.stack(expected).unsqueeze(0) + torch.testing.assert_close(output.float(), expected, rtol=0.02, atol=0.03) + + +def test_indexer_logits(): + torch.manual_seed(23) + tokens, heads, blocks = 2, 3, 4 + q = (torch.rand(tokens, heads, 128, device="cuda") * 8 - 4).to( + torch.float8_e4m3fn + ) + keys = (torch.rand(blocks, 64, 128, device="cuda") * 8 - 4).to( + torch.float8_e4m3fn + ) + scales = torch.rand(blocks, 64, device="cuda", dtype=torch.float32) * 0.02 + cache = torch.zeros(blocks, 64, 132, device="cuda", dtype=torch.uint8) + raw = cache.view(blocks, -1) + raw[:, : 64 * 128] = keys.view(torch.uint8).reshape(blocks, -1) + raw[:, 64 * 128 :] = scales.view(torch.uint8).reshape(blocks, -1) + block_tables = torch.arange(blocks, device="cuda", dtype=torch.int32).view(2, 2) + weights = torch.rand(tokens, heads, device="cuda", dtype=torch.float32) + positions = torch.tensor([70, 45], device="cuda", dtype=torch.int64) + request_ids = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + logits = torch.zeros(tokens, 128, device="cuda", dtype=torch.float32) + tensors = [ + wrap(x) + for x in ( + logits, + q, + cache, + block_tables, + weights, + positions, + request_ids, + ) + ] + _infinicore.fp8_indexer_logits_(*(x._underlying for x in tensors)) + sync() + + expected = torch.full_like(logits, -math.inf) + for token in range(tokens): + request = int(request_ids[token]) + for pos in range(int(positions[token]) + 1): + block = int(block_tables[request, pos // 64]) + offset = pos % 64 + dots = ( + q[token].float() * keys[block, offset].float()[None, :] + ).sum(-1) + expected[token, pos] = ( + torch.relu(dots * scales[block, offset]) * weights[token] + ).sum() + assert torch.equal(torch.isneginf(logits), torch.isneginf(expected)) + finite = torch.isfinite(expected) + torch.testing.assert_close( + logits[finite], expected[finite], rtol=2.0e-4, atol=2.0e-3 + ) + + +def test_select_last_token_hidden(): + torch.manual_seed(24) + hidden = torch.randn(1, 9, 16, device="cuda", dtype=torch.bfloat16) + offsets = torch.tensor([0, 2, 5, 9], device="cuda", dtype=torch.int32) + output = torch.zeros(1, 3, 16, device="cuda", dtype=torch.bfloat16) + tensors = [wrap(x) for x in (output, hidden, offsets)] + _infinicore.select_last_token_hidden_(*(x._underlying for x in tensors)) + sync() + torch.testing.assert_close(output, hidden[:, [1, 4, 8]], rtol=0, atol=0) + + +if __name__ == "__main__": + test_indexer_quant() + test_fused_indexer() + test_mla_cache_and_sparse() + test_indexer_logits() + test_select_last_token_hidden() + print("GLM FP8 InfiniCore wrapper tests passed") From 4123fb23f592d790c4aecac4c3799eedc6ee4a61 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Wed, 29 Jul 2026 02:17:24 +0000 Subject: [PATCH 7/8] test(infiniop): preserve fp8 inputs on device --- test/infiniop/fp8_indexer_logits.py | 7 ++++++- test/infiniop/libinfiniop/utils.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/test/infiniop/fp8_indexer_logits.py b/test/infiniop/fp8_indexer_logits.py index 4ae933aa4..7a544dec1 100644 --- a/test/infiniop/fp8_indexer_logits.py +++ b/test/infiniop/fp8_indexer_logits.py @@ -119,7 +119,12 @@ def test(handle, device, num_tokens, num_requests, _dtype, sync): device, mode="zeros", ) - q_fp8 = TestTensor.from_torch(q_src, InfiniDtype.F8, device) + # q_src is contiguous. Avoid from_torch's generic strided rearrangement: + # some device runtimes cannot round-trip FP8 through the float64 scratch + # tensor used by that path and silently replace the input with zeros. + q_fp8 = TestTensor( + q_src.shape, None, InfiniDtype.F8, device, mode="manual", set_tensor=q_src + ) kv_cache = TestTensor.from_torch(cache_src, InfiniDtype.U8, device) block_tables = TestTensor.from_torch( block_tables_src, InfiniDtype.I32, device diff --git a/test/infiniop/libinfiniop/utils.py b/test/infiniop/libinfiniop/utils.py index c0560dbde..8eaeb01ec 100644 --- a/test/infiniop/libinfiniop/utils.py +++ b/test/infiniop/libinfiniop/utils.py @@ -127,7 +127,7 @@ def __init__( elif mode == "manual": assert set_tensor is not None assert torch_shape == list(set_tensor.shape) - assert torch_strides == list(set_tensor.stride()) + assert torch_strides is None or torch_strides == list(set_tensor.stride()) self._torch_tensor = set_tensor.to(to_torch_dtype(dt)).to( torch_device_map[device] ) From 08cf5ced0484ce1b76567fdb762022c0261d754c Mon Sep 17 00:00:00 2001 From: wooway777 Date: Wed, 29 Jul 2026 02:47:01 +0000 Subject: [PATCH 8/8] chore: format fp8 ops and select op --- python/infinicore/__init__.py | 3 +- test/infinicore/ops/glm_fp8_wrappers.py | 42 +++++++++---------------- test/infiniop/fp8_indexer_logits.py | 38 ++++++++-------------- test/infiniop/fp8_indexer_quant.py | 8 ++--- test/infiniop/fp8_mla_rmsnorm_cache.py | 20 ++++-------- test/infiniop/fp8_sparse_mla.py | 32 ++++++------------- test/infiniop/libinfiniop/utils.py | 17 +++++----- 7 files changed, 54 insertions(+), 106 deletions(-) diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index 43682bc54..777c1be12 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -35,10 +35,10 @@ double, dtype, float, + float8, float16, float32, float64, - float8, half, int, int8, @@ -189,6 +189,7 @@ "complex128", "double", "float", + "float8", "float16", "float32", "float64", diff --git a/test/infinicore/ops/glm_fp8_wrappers.py b/test/infinicore/ops/glm_fp8_wrappers.py index bfd89737e..500471270 100644 --- a/test/infinicore/ops/glm_fp8_wrappers.py +++ b/test/infinicore/ops/glm_fp8_wrappers.py @@ -1,9 +1,9 @@ import math import torch +from infinicore.lib import _infinicore import infinicore -from infinicore.lib import _infinicore def wrap(tensor): @@ -25,12 +25,10 @@ def test_indexer_quant(): sync() scales = torch.exp2( - torch.ceil( - torch.log2(q.float().abs().amax(dim=-1).clamp_min(1.0e-4) / 448.0) - ) + torch.ceil(torch.log2(q.float().abs().amax(dim=-1).clamp_min(1.0e-4) / 448.0)) ) - expected_q = (q.float() / scales[..., None]).clamp(-448, 448).to( - torch.float8_e4m3fn + expected_q = ( + (q.float() / scales[..., None]).clamp(-448, 448).to(torch.float8_e4m3fn) ) torch.testing.assert_close(q_fp8.float(), expected_q.float(), rtol=0, atol=0) torch.testing.assert_close( @@ -41,9 +39,7 @@ def test_indexer_quant(): def test_fused_indexer(): torch.manual_seed(21) tokens, heads, head_dim, rope_dim = 2, 3, 128, 64 - q_raw = torch.randn( - tokens, heads, head_dim, device="cuda", dtype=torch.bfloat16 - ) + q_raw = torch.randn(tokens, heads, head_dim, device="cuda", dtype=torch.bfloat16) k_weights = torch.randn( tokens, head_dim + heads, device="cuda", dtype=torch.bfloat16 ) @@ -91,9 +87,7 @@ def test_mla_cache_and_sparse(): cache = torch.zeros(2, 64, 656, device="cuda", dtype=torch.uint8) vendor = torch.zeros(2, 64, 576, device="cuda", dtype=torch.bfloat16) tensors = [wrap(x) for x in (cache, vendor, compressed, norm_weight, rope, slots)] - _infinicore.fp8_mla_rmsnorm_dual_cache_( - *(x._underlying for x in tensors), 1.0e-5 - ) + _infinicore.fp8_mla_rmsnorm_dual_cache_(*(x._underlying for x in tensors), 1.0e-5) sync() entries = cache.view(-1, 656)[:tokens] @@ -114,16 +108,14 @@ def test_mla_cache_and_sparse(): ) query = torch.randn(1, 2, 576, device="cuda", dtype=torch.bfloat16) - indices = torch.arange(tokens, device="cuda", dtype=torch.int32).view( - 1, 1, tokens - ) + indices = torch.arange(tokens, device="cuda", dtype=torch.int32).view(1, 1, tokens) lens = torch.tensor([tokens], device="cuda", dtype=torch.int32) output = torch.zeros(1, 2, 512, device="cuda", dtype=torch.bfloat16) - sparse_tensors = [wrap(x) for x in (output, query, cache.view(-1, 1, 656), indices, lens)] + sparse_tensors = [ + wrap(x) for x in (output, query, cache.view(-1, 1, 656), indices, lens) + ] scale = float(576**-0.5) - _infinicore.fp8_sparse_mla_( - *(x._underlying for x in sparse_tensors), scale - ) + _infinicore.fp8_sparse_mla_(*(x._underlying for x in sparse_tensors), scale) sync() keys = torch.cat([latent, cached_rope.float()], dim=-1) @@ -138,12 +130,8 @@ def test_mla_cache_and_sparse(): def test_indexer_logits(): torch.manual_seed(23) tokens, heads, blocks = 2, 3, 4 - q = (torch.rand(tokens, heads, 128, device="cuda") * 8 - 4).to( - torch.float8_e4m3fn - ) - keys = (torch.rand(blocks, 64, 128, device="cuda") * 8 - 4).to( - torch.float8_e4m3fn - ) + q = (torch.rand(tokens, heads, 128, device="cuda") * 8 - 4).to(torch.float8_e4m3fn) + keys = (torch.rand(blocks, 64, 128, device="cuda") * 8 - 4).to(torch.float8_e4m3fn) scales = torch.rand(blocks, 64, device="cuda", dtype=torch.float32) * 0.02 cache = torch.zeros(blocks, 64, 132, device="cuda", dtype=torch.uint8) raw = cache.view(blocks, -1) @@ -175,9 +163,7 @@ def test_indexer_logits(): for pos in range(int(positions[token]) + 1): block = int(block_tables[request, pos // 64]) offset = pos % 64 - dots = ( - q[token].float() * keys[block, offset].float()[None, :] - ).sum(-1) + dots = (q[token].float() * keys[block, offset].float()[None, :]).sum(-1) expected[token, pos] = ( torch.relu(dots * scales[block, offset]) * weights[token] ).sum() diff --git a/test/infiniop/fp8_indexer_logits.py b/test/infiniop/fp8_indexer_logits.py index 7a544dec1..117cd2674 100644 --- a/test/infiniop/fp8_indexer_logits.py +++ b/test/infiniop/fp8_indexer_logits.py @@ -59,15 +59,11 @@ def _make_cache(num_blocks): torch.float8_e4m3fn ) scales = torch.rand(num_blocks, _BLOCK_SIZE, dtype=torch.float32) * 0.02 + 0.001 - raw = torch.zeros( - (num_blocks, _BLOCK_SIZE * _CACHE_STRIDE), dtype=torch.uint8 + raw = torch.zeros((num_blocks, _BLOCK_SIZE * _CACHE_STRIDE), dtype=torch.uint8) + raw[:, : _BLOCK_SIZE * _HEAD_DIM] = keys.view(torch.uint8).reshape(num_blocks, -1) + raw[:, _BLOCK_SIZE * _HEAD_DIM :] = ( + scales.contiguous().view(torch.uint8).reshape(num_blocks, -1) ) - raw[:, : _BLOCK_SIZE * _HEAD_DIM] = keys.view(torch.uint8).reshape( - num_blocks, -1 - ) - raw[:, _BLOCK_SIZE * _HEAD_DIM :] = scales.contiguous().view( - torch.uint8 - ).reshape(num_blocks, -1) return raw.reshape(num_blocks, _BLOCK_SIZE, _CACHE_STRIDE), keys, scales @@ -77,25 +73,21 @@ def test(handle, device, num_tokens, num_requests, _dtype, sync): f"tokens={num_tokens}, requests={num_requests}" ) num_blocks = num_requests * 2 - q_src = ( - torch.rand(num_tokens, _NUM_HEADS, _HEAD_DIM) * 8.0 - 4.0 - ).to(torch.float8_e4m3fn) + q_src = (torch.rand(num_tokens, _NUM_HEADS, _HEAD_DIM) * 8.0 - 4.0).to( + torch.float8_e4m3fn + ) weights_src = torch.rand(num_tokens, _NUM_HEADS, dtype=torch.float32) cache_src, keys, key_scales = _make_cache(num_blocks) - block_tables_src = torch.arange( - num_blocks, dtype=torch.int32 - ).reshape(num_requests, 2) - positions_src = torch.tensor( - [70, 45] + [12] * (num_tokens - 2), dtype=torch.int64 + block_tables_src = torch.arange(num_blocks, dtype=torch.int32).reshape( + num_requests, 2 ) + positions_src = torch.tensor([70, 45] + [12] * (num_tokens - 2), dtype=torch.int64) request_ids_src = torch.tensor( list(range(num_requests)) + [-1] * (num_tokens - num_requests), dtype=torch.int32, ) - expected = torch.full( - (num_tokens, _MAX_CONTEXT), -torch.inf, dtype=torch.float32 - ) + expected = torch.full((num_tokens, _MAX_CONTEXT), -torch.inf, dtype=torch.float32) for token in range(num_tokens): request = int(request_ids_src[token]) if request < 0 or request >= num_requests: @@ -126,9 +118,7 @@ def test(handle, device, num_tokens, num_requests, _dtype, sync): q_src.shape, None, InfiniDtype.F8, device, mode="manual", set_tensor=q_src ) kv_cache = TestTensor.from_torch(cache_src, InfiniDtype.U8, device) - block_tables = TestTensor.from_torch( - block_tables_src, InfiniDtype.I32, device - ) + block_tables = TestTensor.from_torch(block_tables_src, InfiniDtype.I32, device) weights = TestTensor.from_torch(weights_src, InfiniDtype.F32, device) positions = TestTensor.from_torch(positions_src, InfiniDtype.I64, device) request_ids = TestTensor.from_torch(request_ids_src, InfiniDtype.I32, device) @@ -177,9 +167,7 @@ def test(handle, device, num_tokens, num_requests, _dtype, sync): actual = logits.actual_tensor().cpu() assert torch.equal(torch.isneginf(actual), torch.isneginf(expected)) finite = torch.isfinite(expected) - assert torch.allclose( - actual[finite], expected[finite], atol=2.0e-3, rtol=2.0e-4 - ) + assert torch.allclose(actual[finite], expected[finite], atol=2.0e-3, rtol=2.0e-4) check_error(LIBINFINIOP.infiniopDestroyFp8IndexerLogitsDescriptor(descriptor)) diff --git a/test/infiniop/fp8_indexer_quant.py b/test/infiniop/fp8_indexer_quant.py index 889094701..482e1a70c 100644 --- a/test/infiniop/fp8_indexer_quant.py +++ b/test/infiniop/fp8_indexer_quant.py @@ -56,16 +56,14 @@ def test(handle, device, shape, dtype, sync): q = TestTensor(shape, None, dtype, device, scale=8.0, bias=-4.0) weights = TestTensor(shape[:2], None, dtype, device, scale=2.0, bias=-1.0) q_fp8 = TestTensor(shape, None, InfiniDtype.F8, device, mode="zeros") - weights_fp32 = TestTensor( - shape[:2], None, InfiniDtype.F32, device, mode="zeros" - ) + weights_fp32 = TestTensor(shape[:2], None, InfiniDtype.F32, device, mode="zeros") q_float = q.torch_tensor().float() scales = torch.exp2( torch.ceil(torch.log2(q_float.abs().amax(dim=-1).clamp_min(1.0e-4) / 448.0)) ) - expected_q = (q_float / scales.unsqueeze(-1)).clamp(-448.0, 448.0).to( - torch.float8_e4m3fn + expected_q = ( + (q_float / scales.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) ) expected_weights = weights.torch_tensor().float() * scales diff --git a/test/infiniop/fp8_mla_rmsnorm_cache.py b/test/infiniop/fp8_mla_rmsnorm_cache.py index 4d8ecd8c9..5f841b00b 100644 --- a/test/infiniop/fp8_mla_rmsnorm_cache.py +++ b/test/infiniop/fp8_mla_rmsnorm_cache.py @@ -64,9 +64,7 @@ def test(handle, device, num_tokens, write_vendor_cache, _dtype, sync): torch.bfloat16 ) weight_src = (torch.rand(_LATENT_DIM) + 0.5).to(torch.bfloat16) - rope_src = (torch.rand(num_tokens, _ROPE_DIM) * 2.0 - 1.0).to( - torch.bfloat16 - ) + rope_src = (torch.rand(num_tokens, _ROPE_DIM) * 2.0 - 1.0).to(torch.bfloat16) slots_src = torch.tensor( [0, block_size + 1] + ([-1] if num_tokens == 3 else []), dtype=torch.int64, @@ -90,9 +88,7 @@ def test(handle, device, num_tokens, write_vendor_cache, _dtype, sync): if write_vendor_cache else None ) - compressed = TestTensor.from_torch( - compressed_src, InfiniDtype.BF16, device - ) + compressed = TestTensor.from_torch(compressed_src, InfiniDtype.BF16, device) weight = TestTensor.from_torch(weight_src, InfiniDtype.BF16, device) rope = TestTensor.from_torch(rope_src, InfiniDtype.BF16, device) slots = TestTensor.from_torch(slots_src, InfiniDtype.I64, device) @@ -125,8 +121,8 @@ def test(handle, device, num_tokens, write_vendor_cache, _dtype, sync): rope_offset = _LATENT_DIM + 16 expected_cache[slot, rope_offset:] = rope_src[token].view(torch.uint8) expected_vendor[slot, :_LATENT_DIM] = ( - quantized.float() * scales[:, None] - ).reshape(-1).to(torch.bfloat16) + (quantized.float() * scales[:, None]).reshape(-1).to(torch.bfloat16) + ) expected_vendor[slot, _LATENT_DIM:] = rope_src[token] descriptor = infiniopOperatorDescriptor_t() @@ -168,14 +164,10 @@ def test(handle, device, num_tokens, write_vendor_cache, _dtype, sync): assert torch.equal(actual_cache, expected_cache) if vendor_cache is not None: assert torch.equal( - vendor_cache.actual_tensor().cpu().reshape( - -1, _LATENT_DIM + _ROPE_DIM - ), + vendor_cache.actual_tensor().cpu().reshape(-1, _LATENT_DIM + _ROPE_DIM), expected_vendor, ) - check_error( - LIBINFINIOP.infiniopDestroyFp8MlaRmsnormCacheDescriptor(descriptor) - ) + check_error(LIBINFINIOP.infiniopDestroyFp8MlaRmsnormCacheDescriptor(descriptor)) if __name__ == "__main__": diff --git a/test/infiniop/fp8_sparse_mla.py b/test/infiniop/fp8_sparse_mla.py index 4d799acf1..20f06b651 100644 --- a/test/infiniop/fp8_sparse_mla.py +++ b/test/infiniop/fp8_sparse_mla.py @@ -69,25 +69,17 @@ def test(handle, device, num_tokens, num_heads, topk, _dtype, sync): torch.float8_e4m3fn ) scales = torch.rand(num_cache_tokens, 4, dtype=torch.float32) * 0.02 + 0.002 - rope = (torch.rand(num_cache_tokens, _ROPE_DIM) * 2.0 - 1.0).to( - torch.bfloat16 - ) - cache_src = torch.zeros( - (num_cache_tokens, 1, _CACHE_STRIDE), dtype=torch.uint8 - ) + rope = (torch.rand(num_cache_tokens, _ROPE_DIM) * 2.0 - 1.0).to(torch.bfloat16) + cache_src = torch.zeros((num_cache_tokens, 1, _CACHE_STRIDE), dtype=torch.uint8) flat_cache = cache_src[:, 0] - flat_cache[:, :_LATENT_DIM] = latent.view(torch.uint8).reshape( - num_cache_tokens, -1 + flat_cache[:, :_LATENT_DIM] = latent.view(torch.uint8).reshape(num_cache_tokens, -1) + flat_cache[:, _LATENT_DIM : _LATENT_DIM + 16] = ( + scales.contiguous().view(torch.uint8).reshape(num_cache_tokens, -1) ) - flat_cache[:, _LATENT_DIM : _LATENT_DIM + 16] = scales.contiguous().view( - torch.uint8 - ).reshape(num_cache_tokens, -1) flat_cache[:, _LATENT_DIM + 16 :] = rope.view(torch.uint8).reshape( num_cache_tokens, -1 ) - query_src = (torch.rand(num_tokens, num_heads, _HEAD_DIM) - 0.5).to( - torch.bfloat16 - ) + query_src = (torch.rand(num_tokens, num_heads, _HEAD_DIM) - 0.5).to(torch.bfloat16) indices_src = torch.arange(topk, dtype=torch.int32).repeat(num_tokens, 1, 1) indices_src[0, 0, 2] = -1 indices_src[1, 0, 67] = num_cache_tokens @@ -98,9 +90,7 @@ def test(handle, device, num_tokens, num_heads, topk, _dtype, sync): num_cache_tokens, _LATENT_DIM ) keys = torch.cat([dequantized, rope.float()], dim=-1) - expected = torch.zeros( - (num_tokens, num_heads, _LATENT_DIM), dtype=torch.float32 - ) + expected = torch.zeros((num_tokens, num_heads, _LATENT_DIM), dtype=torch.float32) for token in range(num_tokens): valid_indices = [ int(index) @@ -108,9 +98,7 @@ def test(handle, device, num_tokens, num_heads, topk, _dtype, sync): if 0 <= int(index) < num_cache_tokens ] for head in range(num_heads): - logits = ( - keys[valid_indices] @ query_src[token, head].float() - ) * _SCALE + logits = (keys[valid_indices] @ query_src[token, head].float()) * _SCALE probabilities = torch.softmax(logits, dim=0) expected[token, head] = ( probabilities[:, None] * dequantized[valid_indices] @@ -126,9 +114,7 @@ def test(handle, device, num_tokens, num_heads, topk, _dtype, sync): query = TestTensor.from_torch(query_src, InfiniDtype.BF16, device) cache = TestTensor.from_torch(cache_src, InfiniDtype.U8, device) indices = TestTensor.from_torch(indices_src, InfiniDtype.I32, device) - topk_lens = TestTensor.from_torch( - topk_lens_src, InfiniDtype.I32, device - ) + topk_lens = TestTensor.from_torch(topk_lens_src, InfiniDtype.I32, device) descriptor = infiniopOperatorDescriptor_t() check_error( diff --git a/test/infiniop/libinfiniop/utils.py b/test/infiniop/libinfiniop/utils.py index 8eaeb01ec..7cad84219 100644 --- a/test/infiniop/libinfiniop/utils.py +++ b/test/infiniop/libinfiniop/utils.py @@ -1,10 +1,12 @@ -from typing import Sequence -import torch import ctypes +from typing import Sequence + import numpy as np +import torch + from .datatypes import * from .devices import * -from .liboperators import infiniopTensorDescriptor_t, LIBINFINIOP, infiniopHandle_t +from .liboperators import LIBINFINIOP, infiniopHandle_t, infiniopTensorDescriptor_t def check_error(status): @@ -602,9 +604,10 @@ def print_discrepancy( if actual.shape != expected.shape: raise ValueError("Tensors must have the same shape to compare.") - import torch import sys + import torch + is_terminal = sys.stdout.isatty() actual = actual.to("cpu") @@ -768,12 +771,9 @@ def get_test_devices(args): if args.qy: devices_to_test.append(InfiniDeviceEnum.QY) if args.cambricon: - import torch_mlu - devices_to_test.append(InfiniDeviceEnum.CAMBRICON) if args.ascend: import torch - import torch_npu torch.npu.set_device(0) # Ascend NPU needs explicit device initialization devices_to_test.append(InfiniDeviceEnum.ASCEND) @@ -783,12 +783,9 @@ def get_test_devices(args): devices_to_test.append(InfiniDeviceEnum.METAX) if args.moore: import torch - import torch_musa devices_to_test.append(InfiniDeviceEnum.MOORE) if args.kunlun: - import torch_xmlir - devices_to_test.append(InfiniDeviceEnum.KUNLUN) if args.hygon: import torch