From af38c67e284b8f39b10640fb3e9ecadf08b9e8b0 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Tue, 14 Jul 2026 05:36:15 +0000 Subject: [PATCH 1/4] pepe: glm 5.2 ops trial --- ...45\345\205\245\350\277\233\345\272\246.md" | 80 ++++ include/infinicore/adaptor/aten_adaptor.hpp | 15 +- .../adaptor/vllm_iluvatar_adaptor.hpp | 44 +++ include/infinicore/ops/cast.hpp | 5 + .../infinicore/ops/concat_and_cache_mla.hpp | 16 + .../ops/concat_and_cache_mla_int8.hpp | 16 + include/infinicore/ops/concat_mla_q.hpp | 11 + .../ops/dynamic_scaled_int8_quant.hpp | 11 + include/infinicore/ops/grouped_topk_vllm.hpp | 8 + .../infinicore/ops/moe_argsort_bincount.hpp | 7 + include/infinicore/ops/moe_expand_input.hpp | 9 + .../infinicore/ops/moe_silu_and_mul_quant.hpp | 9 + include/infinicore/ops/moe_sum_vllm.hpp | 7 + include/infinicore/ops/moe_topk_vllm.hpp | 7 + include/infinicore/ops/scaled_mm_w4a8.hpp | 9 + include/infinicore/ops/scaled_mm_w8a8.hpp | 12 + .../ops/vocab_parallel_embedding.hpp | 6 + include/infinicore/ops/w4a8_group_gemm.hpp | 6 + include/infinicore/ops/w8a8_group_gemm.hpp | 6 + python/infinicore/__init__.py | 23 +- python/infinicore/ops/add_rms_norm.py | 17 + python/infinicore/ops/concat_and_cache_mla.py | 20 + .../ops/concat_and_cache_mla_int8.py | 17 + python/infinicore/ops/concat_mla_q.py | 12 + .../ops/dynamic_scaled_int8_quant.py | 18 + python/infinicore/ops/grouped_topk_vllm.py | 41 ++ python/infinicore/ops/moe_argsort_bincount.py | 14 + python/infinicore/ops/moe_expand_input.py | 22 ++ .../infinicore/ops/moe_silu_and_mul_quant.py | 11 + python/infinicore/ops/moe_sum_vllm.py | 20 + python/infinicore/ops/moe_topk_vllm.py | 58 +++ python/infinicore/ops/scaled_mm_w4a8.py | 25 ++ python/infinicore/ops/w4a8_group_gemm.py | 28 ++ python/infinicore/ops/w8a8_group_gemm.py | 36 ++ src/infinicore/adaptor/aten_adaptor.cc | 19 + .../adaptor/vllm_iluvatar_adaptor.cc | 361 ++++++++++++++++++ .../allocators/pinnable_block_allocator.cc | 19 +- src/infinicore/nn/parameter.cc | 18 +- .../ops/add_rms_norm/add_rms_norm.cc | 15 + src/infinicore/ops/cast/cast.cc | 21 + .../concat_and_cache_mla.cc | 77 ++++ .../concat_and_cache_mla_int8.cc | 89 +++++ .../ops/concat_mla_q/concat_mla_q.cc | 83 ++++ .../grouped_topk_vllm/grouped_topk_vllm.cc | 82 ++++ .../ops/mha_kvcache/mha_kvcache_flashattn.cc | 2 + .../moe_argsort_bincount.cc | 44 +++ .../ops/moe_expand_input/moe_expand_input.cc | 91 +++++ .../moe_silu_and_mul_quant.cc | 82 ++++ .../ops/moe_sum_vllm/moe_sum_vllm.cc | 68 ++++ .../ops/moe_topk_vllm/moe_topk_vllm.cc | 94 +++++ .../mha_varlen_flashattn.cc | 2 + .../dynamic_scaled_int8_quant.cc | 73 ++++ .../ops/scaled_mm_w4a8/scaled_mm_w4a8.cc | 79 ++++ .../ops/scaled_mm_w8a8/scaled_mm_w8a8.cc | 80 ++++ .../vocab_parallel_embedding.cc | 26 ++ .../ops/w4a8_group_gemm/w4a8_group_gemm.cc | 89 +++++ .../ops/w8a8_group_gemm/w8a8_group_gemm.cc | 89 +++++ src/infinicore/pybind11/ops.hpp | 26 ++ src/infinicore/pybind11/ops/add_rms_norm.hpp | 7 + .../pybind11/ops/concat_and_cache_mla.hpp | 23 ++ .../ops/concat_and_cache_mla_int8.hpp | 24 ++ src/infinicore/pybind11/ops/concat_mla_q.hpp | 26 ++ .../ops/dynamic_scaled_int8_quant.hpp | 26 ++ .../pybind11/ops/grouped_topk_vllm.hpp | 17 + .../pybind11/ops/moe_argsort_bincount.hpp | 10 + .../pybind11/ops/moe_expand_input.hpp | 10 + .../pybind11/ops/moe_silu_and_mul_quant.hpp | 10 + src/infinicore/pybind11/ops/moe_sum_vllm.hpp | 10 + src/infinicore/pybind11/ops/moe_topk_vllm.hpp | 26 ++ .../pybind11/ops/scaled_mm_w4a8.hpp | 10 + .../pybind11/ops/w4a8_group_gemm.hpp | 9 + .../pybind11/ops/w8a8_group_gemm.hpp | 9 + src/infinicore/tensor/view.cc | 32 +- .../ops/moe_sum/nvidia/moe_sum_nvidia.cu | 4 +- src/infiniop/ops/moe_sum/operator.cc | 14 +- .../nvidia/moe_topk_sigmoid_nvidia.cu | 4 +- src/infiniop/ops/moe_topk_sigmoid/operator.cc | 14 +- .../nvidia/moe_topk_softmax_nvidia.cu | 4 +- src/infiniop/ops/moe_topk_softmax/operator.cc | 14 +- .../nvidia/silu_and_mul_nvidia.cu | 123 ++++++ .../nvidia/silu_and_mul_nvidia.cuh | 8 + src/infiniop/ops/silu_and_mul/operator.cc | 64 ++++ xmake.lua | 31 +- xmake/iluvatar.lua | 2 +- 84 files changed, 2770 insertions(+), 36 deletions(-) create mode 100644 "docs/GLM-5.2-W4A8-InfiniCore\347\256\227\345\255\220\346\216\245\345\205\245\350\277\233\345\272\246.md" create mode 100644 include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp create mode 100644 include/infinicore/ops/cast.hpp create mode 100644 include/infinicore/ops/concat_and_cache_mla.hpp create mode 100644 include/infinicore/ops/concat_and_cache_mla_int8.hpp create mode 100644 include/infinicore/ops/concat_mla_q.hpp create mode 100644 include/infinicore/ops/dynamic_scaled_int8_quant.hpp create mode 100644 include/infinicore/ops/grouped_topk_vllm.hpp create mode 100644 include/infinicore/ops/moe_argsort_bincount.hpp create mode 100644 include/infinicore/ops/moe_expand_input.hpp create mode 100644 include/infinicore/ops/moe_silu_and_mul_quant.hpp create mode 100644 include/infinicore/ops/moe_sum_vllm.hpp create mode 100644 include/infinicore/ops/moe_topk_vllm.hpp create mode 100644 include/infinicore/ops/scaled_mm_w4a8.hpp create mode 100644 include/infinicore/ops/scaled_mm_w8a8.hpp create mode 100644 include/infinicore/ops/vocab_parallel_embedding.hpp create mode 100644 include/infinicore/ops/w4a8_group_gemm.hpp create mode 100644 include/infinicore/ops/w8a8_group_gemm.hpp create mode 100644 python/infinicore/ops/concat_and_cache_mla.py create mode 100644 python/infinicore/ops/concat_and_cache_mla_int8.py create mode 100644 python/infinicore/ops/concat_mla_q.py create mode 100644 python/infinicore/ops/dynamic_scaled_int8_quant.py create mode 100644 python/infinicore/ops/grouped_topk_vllm.py create mode 100644 python/infinicore/ops/moe_argsort_bincount.py create mode 100644 python/infinicore/ops/moe_expand_input.py create mode 100644 python/infinicore/ops/moe_silu_and_mul_quant.py create mode 100644 python/infinicore/ops/moe_sum_vllm.py create mode 100644 python/infinicore/ops/moe_topk_vllm.py create mode 100644 python/infinicore/ops/scaled_mm_w4a8.py create mode 100644 python/infinicore/ops/w4a8_group_gemm.py create mode 100644 python/infinicore/ops/w8a8_group_gemm.py create mode 100644 src/infinicore/adaptor/vllm_iluvatar_adaptor.cc create mode 100644 src/infinicore/ops/cast/cast.cc create mode 100644 src/infinicore/ops/concat_and_cache_mla/concat_and_cache_mla.cc create mode 100644 src/infinicore/ops/concat_and_cache_mla_int8/concat_and_cache_mla_int8.cc create mode 100644 src/infinicore/ops/concat_mla_q/concat_mla_q.cc create mode 100644 src/infinicore/ops/grouped_topk_vllm/grouped_topk_vllm.cc create mode 100644 src/infinicore/ops/moe_argsort_bincount/moe_argsort_bincount.cc create mode 100644 src/infinicore/ops/moe_expand_input/moe_expand_input.cc create mode 100644 src/infinicore/ops/moe_silu_and_mul_quant/moe_silu_and_mul_quant.cc create mode 100644 src/infinicore/ops/moe_sum_vllm/moe_sum_vllm.cc create mode 100644 src/infinicore/ops/moe_topk_vllm/moe_topk_vllm.cc create mode 100644 src/infinicore/ops/quant/dynamic_scaled_int8_quant/dynamic_scaled_int8_quant.cc create mode 100644 src/infinicore/ops/scaled_mm_w4a8/scaled_mm_w4a8.cc create mode 100644 src/infinicore/ops/scaled_mm_w8a8/scaled_mm_w8a8.cc create mode 100644 src/infinicore/ops/vocab_parallel_embedding/vocab_parallel_embedding.cc create mode 100644 src/infinicore/ops/w4a8_group_gemm/w4a8_group_gemm.cc create mode 100644 src/infinicore/ops/w8a8_group_gemm/w8a8_group_gemm.cc create mode 100644 src/infinicore/pybind11/ops/concat_and_cache_mla.hpp create mode 100644 src/infinicore/pybind11/ops/concat_and_cache_mla_int8.hpp create mode 100644 src/infinicore/pybind11/ops/concat_mla_q.hpp create mode 100644 src/infinicore/pybind11/ops/dynamic_scaled_int8_quant.hpp create mode 100644 src/infinicore/pybind11/ops/grouped_topk_vllm.hpp create mode 100644 src/infinicore/pybind11/ops/moe_argsort_bincount.hpp create mode 100644 src/infinicore/pybind11/ops/moe_expand_input.hpp create mode 100644 src/infinicore/pybind11/ops/moe_silu_and_mul_quant.hpp create mode 100644 src/infinicore/pybind11/ops/moe_sum_vllm.hpp create mode 100644 src/infinicore/pybind11/ops/moe_topk_vllm.hpp create mode 100644 src/infinicore/pybind11/ops/scaled_mm_w4a8.hpp create mode 100644 src/infinicore/pybind11/ops/w4a8_group_gemm.hpp create mode 100644 src/infinicore/pybind11/ops/w8a8_group_gemm.hpp 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/docs/GLM-5.2-W4A8-InfiniCore\347\256\227\345\255\220\346\216\245\345\205\245\350\277\233\345\272\246.md" "b/docs/GLM-5.2-W4A8-InfiniCore\347\256\227\345\255\220\346\216\245\345\205\245\350\277\233\345\272\246.md" new file mode 100644 index 000000000..dcfb3924d --- /dev/null +++ "b/docs/GLM-5.2-W4A8-InfiniCore\347\256\227\345\255\220\346\216\245\345\205\245\350\277\233\345\272\246.md" @@ -0,0 +1,80 @@ +# GLM-5.2-W4A8 InfiniCore 算子接入进度 + +记录环境：TG150 .58 `/home/wuwei/InfiniCore`，容器 `pepe`，Iluvatar BI-V150 (`ivcore11`)。 + +## 已接入/验证 + +| Trace 算子类别 | 当前 InfiniCore 接入状态 | 验证 | +|---|---|---| +| `pyinfer::perf::dynamic_scaled_int8_quant` | 新增 `infinicore.dynamic_scaled_int8_quant`，通过 `vllm_iluvatar/_C` dlopen+dlsym 桥接 | F16/BF16 torch zero-copy 数值验证通过，scale/output 与参考一致 | +| `pyinfer::perf::fused_add_rms_norm` | 新增 `infinicore.add_rms_norm_inplace`，通过 `vllm_iluvatar/_C` 桥接 | F16/BF16 数值验证通过 | +| `silu_and_mul` | 新增/启用 InfiniOp Iluvatar 路径，复用 NVIDIA CUDA kernel 结构 | `python3 test/infiniop/silu_and_mul.py --iluvatar` 通过 | +| `moe_sum` | InfiniOp NVIDIA backend 增加 Iluvatar dispatch/编译 guard | `deepseek_moe.py --iluvatar` 覆盖通过 | +| `pyinfer::perf::topk_softmax` / `topk_sigmoid` | 新增 `infinicore.moe_topk_softmax_vllm` / `moe_topk_sigmoid_vllm`，通过 `vllm_iluvatar/_C` 桥接，三输出匹配 vLLM trace: weights、expert ids、source rows | F16/BF16/F32、renormalize true/false、correction_bias 验证通过 | +| `pyinfer::perf::moe_grouped_topk` / `_moe_C.grouped_topk` | 新增 `infinicore.grouped_topk_vllm`，通过 `vllm_iluvatar/_C` 桥接；仅开放 correction-bias 路径，no-bias 在 Python/C++ 层显式拒绝 | 与 vLLM 实际生产入口 `torch.ops._moe_C.grouped_topk` 在 F16/BF16、softmax/sigmoid、renormalize true/false 下逐位一致；F16 也匹配 `vllm_iluvatar` torch reference；BF16 与 torch reference 存在偏差，记录为复用生产 perf 语义的风险 | +| `pyinfer::cuinfer::scaled_mm_w4a8` / `_C.scaled_mm_w4a8` | 新增 `infinicore.scaled_mm_w4a8`，通过 `vllm_iluvatar/_C` 的 cuinfer 符号桥接；用于 dense W4A8 linear packed GEMM，当前不改权重布局 | 与 vLLM 实际生产入口 `torch.ops._C.scaled_mm_w4a8` 在 F16/BF16、有无 bias 下逐位一致 | +| `pyinfer::cuinfer::w4a8_group_gemm` / `_C.w4a8_group_gemm` | 新增 `infinicore.w4a8_group_gemm_`，通过 `vllm_iluvatar/_C` 的 cuinfer 符号桥接；用于 MoE W4A8 group GEMM | 与 vLLM 实际生产入口 `torch.ops._C.w4a8_group_gemm` 在 F16/BF16、有无 bias、有无 sorted_token_ids、prefill/decode 标志下逐位一致；未声称匹配 torch fallback，因为 fallback 对构造样本与 production 不一致 | +| `pyinfer::cuinfer::w8a8_group_gemm` / `_C.w8a8_group_gemm` | 新增 `infinicore.w8a8_group_gemm_`，通过 `vllm_iluvatar/_C` 的 cuinfer 符号桥接；仅开放 `trans_weight=True` 的 TN prefill 路径 | 与 vLLM 实际生产入口 `torch.ops._C.w8a8_group_gemm` 在 F16/BF16、有无 bias、有无 sorted_token_ids、多个 E/M/K/N shape 下逐位一致；`is_decode=True` 会触发 vLLM/CUINFER internal error，已在 Python/C++ 层显式 guard 禁用 | +| `moe_topk_softmax` | InfiniOp NVIDIA backend 增加 Iluvatar dispatch/编译 guard | `topksoftmax.py --iluvatar`、`deepseek_moe.py --iluvatar` 通过 | +| `moe_topk_sigmoid` | InfiniOp NVIDIA backend 增加 Iluvatar dispatch/编译 guard | `deepseek_moe.py --iluvatar` 覆盖通过 | +| `rms_norm` / `rope` / `mrope` / `kv_caching` / `topksoftmax` | 仓库已有 InfiniCore/InfiniOp 接口和 Iluvatar 可运行路径 | 既有测试曾通过；本轮重跑了 `topksoftmax` | +| `pyinfer::perf::rotary_embedding` | 暂不 bridge；继续使用已验证的 InfiniCore `rope/mrope` 路径 | 直接调用 vLLM perf 对 F16/BF16、2D/3D、Neox/GPT-J 均产生 NaN；已撤回 bridge，避免暴露错误接口 | + +## 已明确不接入的错误路径 + +- `scaled_mm_int8` 的 NVIDIA path 曾尝试打开 Iluvatar dispatch，但首个 BF16 case 数值失败；已撤回，当前 `scaled_mm` 下没有残留 `ENABLE_ILUVATAR` / `INFINI_DEVICE_ILUVATAR` 修改。 +- GLM 的 W4A8/GEMM 热路径应继续优先复用 cuinfer/vLLM 已用 so，而不是暴露未验证的 InfiniOp int8 GEMM 路径。 + +## 当前新增接口 + +- C++ header: `include/infinicore/ops/dynamic_scaled_int8_quant.hpp` +- Python: `infinicore.dynamic_scaled_int8_quant(input, input_scales, out=None)` + - `input`: F16/BF16 contiguous，last dim 为 hidden size + - `input_scales`: F32 contiguous，numel = `input.numel / input.shape[-1]` + - `out`: optional I8 contiguous，shape 与 input 相同 +- C++/Python: `add_rms_norm_inplace(input, residual, weight, epsilon=1e-5)` +- C++/Python: `concat_mla_q(ql_nope, q_pe, out=None)`，当前限制 GLM MLA 维度 `512 + 64 -> 576` +- C++/Python: `concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype="auto", scale=...)` +- C++/Python: `concat_and_cache_mla_int8(kv_c_int8, kv_c_scale, k_pe_int8, k_pe_scale, kv_cache, kv_cache_scale, slot_mapping)` +- C++/Python: `moe_topk_softmax_vllm(gating_output, topk, renormalize=False, correction_bias=None, out=None)` +- C++/Python: `moe_topk_sigmoid_vllm(gating_output, topk, renormalize=False, correction_bias=None, out=None)` +- C++/Python: `grouped_topk_vllm(scores, num_expert_group, topk_group, topk, renormalize, routed_scaling_factor=1.0, bias=..., scoring_func="softmax", out=None)` + - 当前必须传 `bias` / correction bias；no-bias 路径因 vLLM perf 与 reference 不一致而禁用。 + - 对 torch tensor 做 zero-copy 验证时使用 `infinicore.tensor.from_torch(...)` 包装输入和 `out` buffer。 +- C++/Python: `scaled_mm_w4a8(a, b, a_scales, b_scales, bias=None, trans_weight=False, out=None)` + - `a`: int8 activation `(M,K)`；`b`: packed int4-as-int8 weight `(K,N/2)` when `trans_weight=False`；`a_scales`: F32 `(M,1)`；`b_scales`: F32 `(N,1)`；`out`: F16/BF16 `(M,N)`。 + - 当前验证的是 dense W4A8 NN layout，与 vLLM `scaled_mm_w4a8_packed_fwd`/`torch.ops._C.scaled_mm_w4a8` 路径一致。 +- C++/Python: `w4a8_group_gemm_(out, input, weight, input_scale, weight_scale, tokens_per_experts, sorted_token_ids=None, bias=None, trans_weight=True, is_decode=False)` + - `input`: int8 `(M,K)`；`weight`: packed int4-as-int8 `(E,N,K/2)` for TN；`input_scale`: F32 `(M,1)`；`weight_scale`: F32 `(E,N,1)`；`tokens_per_experts`: int32 `(E,)`；`out`: F16/BF16 `(M,N)`。 + - 与 vLLM production op 逐位一致；prefill 构造样本中 `tokens_per_experts` 使用 CPU int32，decode 构造样本中使用 GPU int32，均已覆盖。 +- C++/Python: `w8a8_group_gemm_(out, input, weight, input_scale, weight_scale, tokens_per_experts, sorted_token_ids=None, bias=None, trans_weight=True, is_decode=False)` + - `input`: int8 `(M,K)`；`weight`: int8 `(E,N,K)` for TN；`input_scale`: F32 `(M,1)`；`weight_scale`: F32 `(E,N,1)`；`tokens_per_experts`: int32 `(E,)`；`out`: F16/BF16 `(M,N)`。 + - 当前只开放 `trans_weight=True` 且 `is_decode=False`。`trans_weight=False` 与 `is_decode=True` 均会在 Python/C++ 层报错；后者是因为 vLLM/CUINFER production kernel decode 分支在当前环境触发 `CUINFER_STATUS_INTERNAL_ERROR`。 + - 与 vLLM production prefill op 逐位一致；验证覆盖 F16/BF16、有无 bias、有无 sorted_token_ids，以及 `(E,M,K,N)=(4,16,128,256),(4,64,256,512),(8,128,256,512)`。 + +## 待接入/待确认 + +| Trace 符号 | 建议路径 | 备注 | +|---|---|---| +| `concat_and_cache_mla_rope_fused` | 暂不接入 | 直接对比 vLLM perf 与其 torch fallback：q_pe/k_pe 原地 RoPE 结果不一致，cache 近似一致；为避免暴露错误语义，已撤回 bridge | +| ixCCL send/recv/allreduce/allgather | 归入 distributed/ccl 层 | 不属于单 GPU kernel；需单独检查 InfiniCore CCL 接口 | + +## 本轮验证命令摘要 + +```bash +xmake f --iluvatar-gpu=true --ccl=false --aten=true --iluvatar_arch=ivcore11 --cuda=/usr/local/corex-4.5.0.20260619 -cv +xmake build infinicore_cpp_api +xmake build _infinicore +xmake install infinicore_cpp_api +xmake install _infinicore +xmake build infiniop-iluvatar +xmake build infiniop +xmake install infiniop +python3 test/infiniop/silu_and_mul.py --iluvatar +python3 test/infiniop/topksoftmax.py --iluvatar +python3 test/infiniop/deepseek_moe.py --iluvatar +# grouped_topk_vllm: Python zero-copy torch buffer 测试，对比 torch.ops._moe_C.grouped_topk，并验证 no-bias guard +# scaled_mm_w4a8: Python zero-copy torch buffer 测试，对比 torch.ops._C.scaled_mm_w4a8，F16/BF16、有无 bias 逐位一致 +# w4a8_group_gemm_: Python zero-copy torch buffer 测试，对比 torch.ops._C.w4a8_group_gemm，F16/BF16、有无 bias、有无 sorted_token_ids、prefill/decode 逐位一致 +# w8a8_group_gemm_: Python zero-copy torch buffer 测试，对比 torch.ops._C.w8a8_group_gemm，F16/BF16、有无 bias、有无 sorted_token_ids、prefill 逐位一致；decode/trans_weight guard 已验证 +``` diff --git a/include/infinicore/adaptor/aten_adaptor.hpp b/include/infinicore/adaptor/aten_adaptor.hpp index 74b053dc5..ded231482 100644 --- a/include/infinicore/adaptor/aten_adaptor.hpp +++ b/include/infinicore/adaptor/aten_adaptor.hpp @@ -1,5 +1,6 @@ #ifdef ENABLE_ATEN #pragma once + #include "../context/context.hpp" #include "../tensor.hpp" @@ -29,6 +30,8 @@ inline at::ScalarType to_at_dtype(DataType dtype) { return at::kHalf; case DataType::BF16: return at::kBFloat16; + case DataType::I8: + return at::kChar; case DataType::I32: return at::kInt; case DataType::I64: @@ -40,8 +43,12 @@ inline at::ScalarType to_at_dtype(DataType dtype) { inline at::Device to_at_device(const Device &device) { // PyTorch ATen only exposes standard device types (e.g. kCPU/kCUDA). - // Treat MetaX/QY devices as CUDA devices for ATen tensor interoperability. - if (device.getType() == Device::Type::NVIDIA || device.getType() == Device::Type::METAX || device.getType() == Device::Type::QY || device.getType() == Device::Type::HYGON) { + // Treat CUDA-compatible vendor devices as CUDA devices for ATen tensor interoperability. + if (device.getType() == Device::Type::NVIDIA + || device.getType() == Device::Type::METAX + || device.getType() == Device::Type::QY + || device.getType() == Device::Type::ILUVATAR + || device.getType() == Device::Type::HYGON) { return at::Device(at::kCUDA, device.getIndex()); } else if (device.getType() == Device::Type::CPU) { return at::Device(at::kCPU); @@ -64,6 +71,10 @@ c10::hip::HIPStream get_hip_stream(); c10::cuda::CUDAStream get_cuda_stream(); #endif +#if defined(ENABLE_ILUVATAR_API) +void set_aten_stream_to_infinicore(); +#endif + #if defined(ENABLE_MOORE_API) c10::musa::MUSAStream get_musa_stream(); #endif diff --git a/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp b/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp new file mode 100644 index 000000000..8f6a47c09 --- /dev/null +++ b/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp @@ -0,0 +1,44 @@ +#pragma once + +#ifdef ENABLE_ATEN +#include +#include +#include + +namespace infinicore::adaptor::vllm_iluvatar { + +bool available(); +bool dynamic_scaled_int8_quant_available(); +bool concat_mla_q_available(); +bool concat_and_cache_mla_available(); +bool concat_and_cache_mla_int8_available(); +bool topk_softmax_available(); +bool topk_sigmoid_available(); +bool grouped_topk_available(); +bool scaled_mm_w4a8_available(); +bool scaled_mm_w8a8_available(); +bool w4a8_group_gemm_available(); +bool w8a8_group_gemm_available(); +bool argsort_bincount_with_inv_pos_available(); +bool expand_moe_input_with_inv_pos_available(); +bool silu_and_mul_quant_available(); +bool moe_sum_vllm_available(); +void fused_add_rms_norm(at::Tensor &input, at::Tensor &residual, at::Tensor &weight, float epsilon); +void dynamic_scaled_int8_quant(at::Tensor &output, at::Tensor &input_scales, const at::Tensor &input); +void concat_mla_q(at::Tensor &ql_nope, at::Tensor &q_pe, at::Tensor &q_out); +void concat_and_cache_mla(at::Tensor &kv_c, at::Tensor &k_pe, at::Tensor &kv_cache, at::Tensor &slot_mapping, const std::string &kv_cache_dtype, at::Tensor &scale); +void concat_and_cache_mla_int8(at::Tensor &kv_c_int8, at::Tensor &kv_c_scale, at::Tensor &k_pe_int8, at::Tensor &k_pe_scale, at::Tensor &kv_cache, at::Tensor &kv_cache_scale, at::Tensor &slot_mapping); +void topk_softmax(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional correction_bias); +void topk_sigmoid(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional correction_bias); +void grouped_topk(at::Tensor &topk_weights, at::Tensor &topk_ids, const at::Tensor &scores, std::optional bias, int64_t num_expert_group, int64_t topk_group, const std::string &scoring_func, bool renormalize); +void scaled_mm_w4a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional bias, bool trans_weight); +void scaled_mm_w8a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional bias, bool trans_weight); +void w4a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode); +void w8a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode); +void argsort_bincount_with_inv_pos(const at::Tensor &topk_ids, at::Tensor &tokens_per_experts, at::Tensor &sorted_indices, at::Tensor &inv_pos, int64_t num_experts); +void expand_moe_input_with_inv_pos(at::Tensor &expand_states, std::optional expand_scales, const at::Tensor &hidden_states, const at::Tensor &inv_pos, int64_t top_k, int64_t group_size, int64_t format); +void silu_and_mul_quant(at::Tensor &output, std::optional output_scale, const at::Tensor &input, int64_t format); +void moe_sum_vllm(at::Tensor &output, const at::Tensor &input, std::optional topk_weights, std::optional extra_residual, double routed_scale, double residual_scale); + +} // namespace infinicore::adaptor::vllm_iluvatar +#endif // ENABLE_ATEN diff --git a/include/infinicore/ops/cast.hpp b/include/infinicore/ops/cast.hpp new file mode 100644 index 000000000..4aa907922 --- /dev/null +++ b/include/infinicore/ops/cast.hpp @@ -0,0 +1,5 @@ +#pragma once +#include "../tensor.hpp" +namespace infinicore::op { +void cast_(Tensor output, const Tensor &input); +} diff --git a/include/infinicore/ops/concat_and_cache_mla.hpp b/include/infinicore/ops/concat_and_cache_mla.hpp new file mode 100644 index 000000000..6acc70b21 --- /dev/null +++ b/include/infinicore/ops/concat_and_cache_mla.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "../tensor.hpp" +#include + +namespace infinicore::op { + +void concat_and_cache_mla_(const Tensor &kv_c, + const Tensor &k_pe, + Tensor kv_cache, + const Tensor &slot_mapping, + const std::string &kv_cache_dtype, + Tensor scale); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/concat_and_cache_mla_int8.hpp b/include/infinicore/ops/concat_and_cache_mla_int8.hpp new file mode 100644 index 000000000..3de0f4d13 --- /dev/null +++ b/include/infinicore/ops/concat_and_cache_mla_int8.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "../tensor.hpp" + +namespace infinicore::op { + +void concat_and_cache_mla_int8_(const Tensor &kv_c_int8, + const Tensor &kv_c_scale, + const Tensor &k_pe_int8, + const Tensor &k_pe_scale, + Tensor kv_cache, + Tensor kv_cache_scale, + const Tensor &slot_mapping); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/concat_mla_q.hpp b/include/infinicore/ops/concat_mla_q.hpp new file mode 100644 index 000000000..90d84c596 --- /dev/null +++ b/include/infinicore/ops/concat_mla_q.hpp @@ -0,0 +1,11 @@ +#pragma once + +#include "../device.hpp" +#include "../tensor.hpp" + +namespace infinicore::op { + +void concat_mla_q_(const Tensor &ql_nope, const Tensor &q_pe, Tensor q_out); +Tensor concat_mla_q(const Tensor &ql_nope, const Tensor &q_pe); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/dynamic_scaled_int8_quant.hpp b/include/infinicore/ops/dynamic_scaled_int8_quant.hpp new file mode 100644 index 000000000..f0a114821 --- /dev/null +++ b/include/infinicore/ops/dynamic_scaled_int8_quant.hpp @@ -0,0 +1,11 @@ +#pragma once + +#include "../device.hpp" +#include "../tensor.hpp" + +namespace infinicore::op { + +void dynamic_scaled_int8_quant_(Tensor output, const Tensor &input, Tensor input_scales); +Tensor dynamic_scaled_int8_quant(const Tensor &input, Tensor input_scales); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/grouped_topk_vllm.hpp b/include/infinicore/ops/grouped_topk_vllm.hpp new file mode 100644 index 000000000..11d30cdf4 --- /dev/null +++ b/include/infinicore/ops/grouped_topk_vllm.hpp @@ -0,0 +1,8 @@ +#pragma once +#include "../device.hpp" +#include "../tensor.hpp" +#include +#include +namespace infinicore::op { +void grouped_topk_vllm_(Tensor topk_weights, Tensor topk_ids, const Tensor &scores, int64_t num_expert_group, int64_t topk_group, bool renormalize, float routed_scaling_factor, const Tensor &bias = Tensor(), const std::string &scoring_func = "softmax"); +} diff --git a/include/infinicore/ops/moe_argsort_bincount.hpp b/include/infinicore/ops/moe_argsort_bincount.hpp new file mode 100644 index 000000000..b710383b9 --- /dev/null +++ b/include/infinicore/ops/moe_argsort_bincount.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../tensor.hpp" +#include + +namespace infinicore::op { +void moe_argsort_bincount_with_inv_pos_(Tensor tokens_per_experts, Tensor sorted_indices, Tensor inv_pos, const Tensor &topk_ids, int64_t num_experts); +} diff --git a/include/infinicore/ops/moe_expand_input.hpp b/include/infinicore/ops/moe_expand_input.hpp new file mode 100644 index 000000000..0f1f9cabd --- /dev/null +++ b/include/infinicore/ops/moe_expand_input.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "../tensor.hpp" +#include +#include + +namespace infinicore::op { +// format: 0=normal, 1=quant, 2=packed. group_size is used by quant/packed paths. +void moe_expand_input_with_inv_pos_(Tensor expand_states, std::optional expand_scales, const Tensor &hidden_states, const Tensor &inv_pos, int64_t top_k, int64_t group_size = 128, int64_t format = 0); +} // namespace infinicore::op diff --git a/include/infinicore/ops/moe_silu_and_mul_quant.hpp b/include/infinicore/ops/moe_silu_and_mul_quant.hpp new file mode 100644 index 000000000..496c9f3e9 --- /dev/null +++ b/include/infinicore/ops/moe_silu_and_mul_quant.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "../tensor.hpp" +#include +#include + +namespace infinicore::op { +// input shape [M, 2N]. format: 0=normal fp output, 1=quant int8+scale, 2=packed int8+scale. +void moe_silu_and_mul_quant_(Tensor output, std::optional output_scale, const Tensor &input, int64_t format = 0); +} // namespace infinicore::op diff --git a/include/infinicore/ops/moe_sum_vllm.hpp b/include/infinicore/ops/moe_sum_vllm.hpp new file mode 100644 index 000000000..93aaab35b --- /dev/null +++ b/include/infinicore/ops/moe_sum_vllm.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../tensor.hpp" +#include + +namespace infinicore::op { +void moe_sum_vllm_(Tensor output, const Tensor &input, std::optional topk_weights = std::nullopt, std::optional extra_residual = std::nullopt, double routed_scale = 1.0, double residual_scale = 1.0); +} diff --git a/include/infinicore/ops/moe_topk_vllm.hpp b/include/infinicore/ops/moe_topk_vllm.hpp new file mode 100644 index 000000000..cedba1746 --- /dev/null +++ b/include/infinicore/ops/moe_topk_vllm.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../device.hpp" +#include "../tensor.hpp" +namespace infinicore::op { +void moe_topk_softmax_vllm_(Tensor topk_weights, Tensor topk_ids, Tensor token_expert_indices, const Tensor &gating_output, bool renormalize = false, const Tensor &correction_bias = Tensor()); +void moe_topk_sigmoid_vllm_(Tensor topk_weights, Tensor topk_ids, Tensor token_expert_indices, const Tensor &gating_output, bool renormalize = false, const Tensor &correction_bias = Tensor()); +} // namespace infinicore::op diff --git a/include/infinicore/ops/scaled_mm_w4a8.hpp b/include/infinicore/ops/scaled_mm_w4a8.hpp new file mode 100644 index 000000000..c7d1f2f1f --- /dev/null +++ b/include/infinicore/ops/scaled_mm_w4a8.hpp @@ -0,0 +1,9 @@ +#pragma once + +#include "../tensor.hpp" +#include + +namespace infinicore::op { +void scaled_mm_w4a8_(Tensor out, const Tensor &a, const Tensor &b, const Tensor &a_scales, const Tensor &b_scales, std::optional bias = std::nullopt, bool trans_weight = false); +Tensor scaled_mm_w4a8(const Tensor &a, const Tensor &b, const Tensor &a_scales, const Tensor &b_scales, std::optional bias = std::nullopt, bool trans_weight = false); +} // namespace infinicore::op diff --git a/include/infinicore/ops/scaled_mm_w8a8.hpp b/include/infinicore/ops/scaled_mm_w8a8.hpp new file mode 100644 index 000000000..54b7a4ee9 --- /dev/null +++ b/include/infinicore/ops/scaled_mm_w8a8.hpp @@ -0,0 +1,12 @@ +#pragma once +#include "infinicore/tensor.hpp" +#include + +namespace infinicore::op { +Tensor scaled_mm_w8a8(const Tensor &a, const Tensor &b, + const Tensor &a_scales, const Tensor &b_scales, + std::optional bias = std::nullopt, bool trans_weight = true); +void scaled_mm_w8a8_(Tensor out, const Tensor &a, const Tensor &b, + const Tensor &a_scales, const Tensor &b_scales, + std::optional bias = std::nullopt, bool trans_weight = true); +} // namespace infinicore::op diff --git a/include/infinicore/ops/vocab_parallel_embedding.hpp b/include/infinicore/ops/vocab_parallel_embedding.hpp new file mode 100644 index 000000000..25390664a --- /dev/null +++ b/include/infinicore/ops/vocab_parallel_embedding.hpp @@ -0,0 +1,6 @@ +#pragma once +#include "../tensor.hpp" +#include +namespace infinicore::op { +void vocab_parallel_embedding_(Tensor output, const Tensor &indices, const Tensor &weight, int64_t vocab_start, int64_t vocab_end); +} diff --git a/include/infinicore/ops/w4a8_group_gemm.hpp b/include/infinicore/ops/w4a8_group_gemm.hpp new file mode 100644 index 000000000..60312c14c --- /dev/null +++ b/include/infinicore/ops/w4a8_group_gemm.hpp @@ -0,0 +1,6 @@ +#pragma once +#include "../tensor.hpp" +#include +namespace infinicore::op { +void w4a8_group_gemm_(Tensor out, const Tensor &input, const Tensor &weight, const Tensor &input_scale, const Tensor &weight_scale, const Tensor &tokens_per_experts, std::optional sorted_token_ids = std::nullopt, std::optional bias = std::nullopt, bool trans_weight = true, bool is_decode = false); +} diff --git a/include/infinicore/ops/w8a8_group_gemm.hpp b/include/infinicore/ops/w8a8_group_gemm.hpp new file mode 100644 index 000000000..68239471c --- /dev/null +++ b/include/infinicore/ops/w8a8_group_gemm.hpp @@ -0,0 +1,6 @@ +#pragma once +#include "../tensor.hpp" +#include +namespace infinicore::op { +void w8a8_group_gemm_(Tensor out, const Tensor &input, const Tensor &weight, const Tensor &input_scale, const Tensor &weight_scale, const Tensor &tokens_per_experts, std::optional sorted_token_ids = std::nullopt, std::optional bias = std::nullopt, bool trans_weight = true, bool is_decode = false); +} diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index 808995640..b01222408 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -50,7 +50,7 @@ ) from infinicore.ops.acos import acos from infinicore.ops.add import add -from infinicore.ops.add_rms_norm import add_rms_norm +from infinicore.ops.add_rms_norm import add_rms_norm, add_rms_norm_inplace from infinicore.ops.addbmm import addbmm from infinicore.ops.addcmul import addcmul from infinicore.ops.addr import addr @@ -76,11 +76,15 @@ from infinicore.ops.broadcast_to import broadcast_to from infinicore.ops.cat import cat from infinicore.ops.cdist import cdist +from infinicore.ops.concat_and_cache_mla import concat_and_cache_mla +from infinicore.ops.concat_and_cache_mla_int8 import concat_and_cache_mla_int8 +from infinicore.ops.concat_mla_q import concat_mla_q from infinicore.ops.conv2d import conv2d from infinicore.ops.cross_entropy import cross_entropy from infinicore.ops.diff import diff from infinicore.ops.digamma import digamma from infinicore.ops.dist import dist +from infinicore.ops.dynamic_scaled_int8_quant import dynamic_scaled_int8_quant from infinicore.ops.equal import equal from infinicore.ops.flipud import flipud from infinicore.ops.float_power import float_power @@ -88,6 +92,7 @@ from infinicore.ops.floor_divide import floor_divide from infinicore.ops.fmin import fmin from infinicore.ops.fmod import fmod +from infinicore.ops.grouped_topk_vllm import grouped_topk_vllm from infinicore.ops.hypot import hypot from infinicore.ops.index_add import index_add from infinicore.ops.index_copy import index_copy @@ -108,6 +113,11 @@ from infinicore.ops.mha import mha from infinicore.ops.mha_kvcache import mha_kvcache from infinicore.ops.mha_varlen import mha_varlen +from infinicore.ops.moe_argsort_bincount import moe_argsort_bincount_with_inv_pos_ +from infinicore.ops.moe_expand_input import moe_expand_input_with_inv_pos_ +from infinicore.ops.moe_silu_and_mul_quant import moe_silu_and_mul_quant_ +from infinicore.ops.moe_sum_vllm import moe_sum_vllm_ +from infinicore.ops.moe_topk_vllm import moe_topk_sigmoid_vllm, moe_topk_softmax_vllm from infinicore.ops.moore_mate_flash_attn import ( moore_mate_flash_attn_decode, moore_mate_flash_attn_prefill, @@ -127,6 +137,7 @@ from infinicore.ops.rotm import rotm from infinicore.ops.rotmg import rotmg from infinicore.ops.scal import scal +from infinicore.ops.scaled_mm_w4a8 import scaled_mm_w4a8 from infinicore.ops.scatter import scatter from infinicore.ops.sinh import sinh from infinicore.ops.squeeze import squeeze @@ -139,6 +150,8 @@ from infinicore.ops.vander import vander from infinicore.ops.var import var from infinicore.ops.var_mean import var_mean +from infinicore.ops.w4a8_group_gemm import w4a8_group_gemm_ +from infinicore.ops.w8a8_group_gemm import w8a8_group_gemm_ from infinicore.tensor import ( Tensor, empty, @@ -205,6 +218,7 @@ "add", "addr", "add_rms_norm", + "add_rms_norm_inplace", "argwhere", "asin", "asum", @@ -218,6 +232,8 @@ "floor", "attention", "mrope", + "moe_topk_sigmoid_vllm", + "moe_topk_softmax_vllm", "block_diag", "kron", "bitwise_right_shift", @@ -228,6 +244,9 @@ "fmod", "cat", "conv2d", + "concat_mla_q", + "concat_and_cache_mla", + "concat_and_cache_mla_int8", "inner", "masked_select", "logaddexp", @@ -237,6 +256,7 @@ "mul", "mul_scalar", "diff", + "dynamic_scaled_int8_quant", "digamma", "dist", "logdet", @@ -278,6 +298,7 @@ "paged_attention", "paged_attention_prefill", "hypot", + "grouped_topk_vllm", "index_copy", "index_add", "take", diff --git a/python/infinicore/ops/add_rms_norm.py b/python/infinicore/ops/add_rms_norm.py index a5de7bd92..bd0546b6b 100644 --- a/python/infinicore/ops/add_rms_norm.py +++ b/python/infinicore/ops/add_rms_norm.py @@ -32,3 +32,20 @@ def add_rms_norm(a, b, weight, epsilon=1e-5, *, out=None, residual=None): ) return out, residual + + +def add_rms_norm_inplace(input, residual, weight, epsilon=1e-5): + """ + vLLM-style in-place fused Add and RMSNorm. + + Mutates input and residual in place: + residual = input_old + residual + input = RMSNorm(residual, weight, epsilon) + """ + _infinicore.add_rms_norm_inplace( + input._underlying, + residual._underlying, + weight._underlying, + epsilon, + ) + return input, residual diff --git a/python/infinicore/ops/concat_and_cache_mla.py b/python/infinicore/ops/concat_and_cache_mla.py new file mode 100644 index 000000000..088002e64 --- /dev/null +++ b/python/infinicore/ops/concat_and_cache_mla.py @@ -0,0 +1,20 @@ +from infinicore.lib import _infinicore + + +def concat_and_cache_mla( + kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype="auto", scale=None +): + """Write concat([kv_c, k_pe], -1) into paged MLA KV cache at slot_mapping.""" + if scale is None: + raise ValueError( + "concat_and_cache_mla requires a float32 scale tensor; pass a scalar tensor for kv_cache_dtype='auto'" + ) + _infinicore.concat_and_cache_mla_( + kv_c._underlying, + k_pe._underlying, + kv_cache._underlying, + slot_mapping._underlying, + kv_cache_dtype, + scale._underlying, + ) + return kv_cache diff --git a/python/infinicore/ops/concat_and_cache_mla_int8.py b/python/infinicore/ops/concat_and_cache_mla_int8.py new file mode 100644 index 000000000..097c97d36 --- /dev/null +++ b/python/infinicore/ops/concat_and_cache_mla_int8.py @@ -0,0 +1,17 @@ +from infinicore.lib import _infinicore + + +def concat_and_cache_mla_int8( + kv_c_int8, kv_c_scale, k_pe_int8, k_pe_scale, kv_cache, kv_cache_scale, slot_mapping +): + """Write pre-quantized MLA latent/rope KV and scales into int8 paged cache.""" + _infinicore.concat_and_cache_mla_int8_( + kv_c_int8._underlying, + kv_c_scale._underlying, + k_pe_int8._underlying, + k_pe_scale._underlying, + kv_cache._underlying, + kv_cache_scale._underlying, + slot_mapping._underlying, + ) + return kv_cache, kv_cache_scale diff --git a/python/infinicore/ops/concat_mla_q.py b/python/infinicore/ops/concat_mla_q.py new file mode 100644 index 000000000..96b2cdef8 --- /dev/null +++ b/python/infinicore/ops/concat_mla_q.py @@ -0,0 +1,12 @@ +from infinicore.lib import _infinicore +from infinicore.tensor import empty + + +def concat_mla_q(ql_nope, q_pe, *, out=None): + """Concatenate MLA ql_nope and q_pe along the last dimension into out.""" + if out is None: + shape = list(ql_nope.shape) + shape[-1] = ql_nope.shape[-1] + q_pe.shape[-1] + out = empty(tuple(shape), dtype=ql_nope.dtype, device=ql_nope.device) + _infinicore.concat_mla_q_(ql_nope._underlying, q_pe._underlying, out._underlying) + return out diff --git a/python/infinicore/ops/dynamic_scaled_int8_quant.py b/python/infinicore/ops/dynamic_scaled_int8_quant.py new file mode 100644 index 000000000..0d7b1f4bd --- /dev/null +++ b/python/infinicore/ops/dynamic_scaled_int8_quant.py @@ -0,0 +1,18 @@ +from infinicore.dtype import int8 +from infinicore.lib import _infinicore +from infinicore.tensor import empty + + +def dynamic_scaled_int8_quant(input, input_scales, *, out=None): + """Per-token dynamic scaled int8 quantization. + + Mutates/writes: + out: int8 tensor with the same shape as input + input_scales: float32 tensor with numel == input.numel / input.shape[-1] + """ + if out is None: + out = empty(input.shape, dtype=int8, device=input.device) + _infinicore.dynamic_scaled_int8_quant_( + out._underlying, input._underlying, input_scales._underlying + ) + return out diff --git a/python/infinicore/ops/grouped_topk_vllm.py b/python/infinicore/ops/grouped_topk_vllm.py new file mode 100644 index 000000000..c9ef0f320 --- /dev/null +++ b/python/infinicore/ops/grouped_topk_vllm.py @@ -0,0 +1,41 @@ +from infinicore.dtype import float32, int32 +from infinicore.lib import _infinicore +from infinicore.tensor import empty + + +def grouped_topk_vllm( + scores, + num_expert_group, + topk_group, + topk, + renormalize, + routed_scaling_factor=1.0, + bias=None, + scoring_func="softmax", + *, + out=None, + topk_ids_dtype=int32, +): + if bias is None: + raise RuntimeError( + "grouped_topk_vllm currently requires correction bias; vllm_iluvatar perf no-bias path mismatches reference" + ) + if out is None: + shape = (scores.shape[0], topk) + out = ( + empty(shape, dtype=float32, device=scores.device), + empty(shape, dtype=topk_ids_dtype, device=scores.device), + ) + topk_weights, topk_ids = out + _infinicore.grouped_topk_vllm_( + topk_weights._underlying, + topk_ids._underlying, + scores._underlying, + num_expert_group, + topk_group, + renormalize, + float(routed_scaling_factor), + None if bias is None else bias._underlying, + scoring_func, + ) + return out diff --git a/python/infinicore/ops/moe_argsort_bincount.py b/python/infinicore/ops/moe_argsort_bincount.py new file mode 100644 index 000000000..96d1ccbe3 --- /dev/null +++ b/python/infinicore/ops/moe_argsort_bincount.py @@ -0,0 +1,14 @@ +from infinicore.lib import _infinicore + + +def moe_argsort_bincount_with_inv_pos_( + tokens_per_experts, sorted_indices, inv_pos, topk_ids, num_experts: int +): + _infinicore.moe_argsort_bincount_with_inv_pos_( + tokens_per_experts._underlying, + sorted_indices._underlying, + inv_pos._underlying, + topk_ids._underlying, + num_experts, + ) + return tokens_per_experts, sorted_indices, inv_pos diff --git a/python/infinicore/ops/moe_expand_input.py b/python/infinicore/ops/moe_expand_input.py new file mode 100644 index 000000000..fc3886374 --- /dev/null +++ b/python/infinicore/ops/moe_expand_input.py @@ -0,0 +1,22 @@ +from infinicore.lib import _infinicore + + +def moe_expand_input_with_inv_pos_( + expand_states, + expand_scales, + hidden_states, + inv_pos, + top_k: int, + group_size: int = 128, + format: int = 0, +): + _infinicore.moe_expand_input_with_inv_pos_( + expand_states._underlying, + None if expand_scales is None else expand_scales._underlying, + hidden_states._underlying, + inv_pos._underlying, + top_k, + group_size, + format, + ) + return expand_states if expand_scales is None else (expand_states, expand_scales) diff --git a/python/infinicore/ops/moe_silu_and_mul_quant.py b/python/infinicore/ops/moe_silu_and_mul_quant.py new file mode 100644 index 000000000..7ac6bbca2 --- /dev/null +++ b/python/infinicore/ops/moe_silu_and_mul_quant.py @@ -0,0 +1,11 @@ +from infinicore.lib import _infinicore + + +def moe_silu_and_mul_quant_(output, output_scale, input, format: int = 0): + _infinicore.moe_silu_and_mul_quant_( + output._underlying, + None if output_scale is None else output_scale._underlying, + input._underlying, + format, + ) + return output if output_scale is None else (output, output_scale) diff --git a/python/infinicore/ops/moe_sum_vllm.py b/python/infinicore/ops/moe_sum_vllm.py new file mode 100644 index 000000000..35fb544e6 --- /dev/null +++ b/python/infinicore/ops/moe_sum_vllm.py @@ -0,0 +1,20 @@ +from infinicore.lib import _infinicore + + +def moe_sum_vllm_( + output, + input, + topk_weights=None, + extra_residual=None, + routed_scale: float = 1.0, + residual_scale: float = 1.0, +): + _infinicore.moe_sum_vllm_( + output._underlying, + input._underlying, + None if topk_weights is None else topk_weights._underlying, + None if extra_residual is None else extra_residual._underlying, + routed_scale, + residual_scale, + ) + return output diff --git a/python/infinicore/ops/moe_topk_vllm.py b/python/infinicore/ops/moe_topk_vllm.py new file mode 100644 index 000000000..f28b7728d --- /dev/null +++ b/python/infinicore/ops/moe_topk_vllm.py @@ -0,0 +1,58 @@ +from infinicore.dtype import float32, int32 +from infinicore.lib import _infinicore +from infinicore.tensor import empty + + +def _alloc_outputs(gating_output, topk, topk_ids_dtype): + shape = (gating_output.shape[0], topk) + # vLLM perf allows fp32 weights for all gating dtypes; use fp32 by default for stable comparisons. + topk_weights = empty(shape, dtype=float32, device=gating_output.device) + topk_ids = empty(shape, dtype=topk_ids_dtype, device=gating_output.device) + token_expert_indices = empty(shape, dtype=int32, device=gating_output.device) + return topk_weights, topk_ids, token_expert_indices + + +def moe_topk_softmax_vllm( + gating_output, + topk, + renormalize=False, + correction_bias=None, + *, + out=None, + topk_ids_dtype=int32, +): + if out is None: + out = _alloc_outputs(gating_output, topk, topk_ids_dtype) + topk_weights, topk_ids, token_expert_indices = out + _infinicore.moe_topk_softmax_vllm_( + topk_weights._underlying, + topk_ids._underlying, + token_expert_indices._underlying, + gating_output._underlying, + renormalize, + None if correction_bias is None else correction_bias._underlying, + ) + return out + + +def moe_topk_sigmoid_vllm( + gating_output, + topk, + renormalize=False, + correction_bias=None, + *, + out=None, + topk_ids_dtype=int32, +): + if out is None: + out = _alloc_outputs(gating_output, topk, topk_ids_dtype) + topk_weights, topk_ids, token_expert_indices = out + _infinicore.moe_topk_sigmoid_vllm_( + topk_weights._underlying, + topk_ids._underlying, + token_expert_indices._underlying, + gating_output._underlying, + renormalize, + None if correction_bias is None else correction_bias._underlying, + ) + return out diff --git a/python/infinicore/ops/scaled_mm_w4a8.py b/python/infinicore/ops/scaled_mm_w4a8.py new file mode 100644 index 000000000..cbab3344f --- /dev/null +++ b/python/infinicore/ops/scaled_mm_w4a8.py @@ -0,0 +1,25 @@ +from infinicore.lib import _infinicore + + +def scaled_mm_w4a8( + a, b, a_scales, b_scales, bias=None, trans_weight=False, *, out=None +): + if out is None: + return _infinicore.scaled_mm_w4a8( + a._underlying, + b._underlying, + a_scales._underlying, + b_scales._underlying, + None if bias is None else bias._underlying, + trans_weight, + ) + _infinicore.scaled_mm_w4a8_( + out._underlying, + a._underlying, + b._underlying, + a_scales._underlying, + b_scales._underlying, + None if bias is None else bias._underlying, + trans_weight, + ) + return out diff --git a/python/infinicore/ops/w4a8_group_gemm.py b/python/infinicore/ops/w4a8_group_gemm.py new file mode 100644 index 000000000..700e77625 --- /dev/null +++ b/python/infinicore/ops/w4a8_group_gemm.py @@ -0,0 +1,28 @@ +from infinicore.lib import _infinicore + + +def w4a8_group_gemm_( + out, + input, + weight, + input_scale, + weight_scale, + tokens_per_experts, + sorted_token_ids=None, + bias=None, + trans_weight=True, + is_decode=False, +): + _infinicore.w4a8_group_gemm_( + out._underlying, + input._underlying, + weight._underlying, + input_scale._underlying, + weight_scale._underlying, + tokens_per_experts._underlying, + None if sorted_token_ids is None else sorted_token_ids._underlying, + None if bias is None else bias._underlying, + trans_weight, + is_decode, + ) + return out diff --git a/python/infinicore/ops/w8a8_group_gemm.py b/python/infinicore/ops/w8a8_group_gemm.py new file mode 100644 index 000000000..8a6e8d98f --- /dev/null +++ b/python/infinicore/ops/w8a8_group_gemm.py @@ -0,0 +1,36 @@ +from infinicore.lib import _infinicore + + +def w8a8_group_gemm_( + out, + input, + weight, + input_scale, + weight_scale, + tokens_per_experts, + sorted_token_ids=None, + bias=None, + trans_weight=True, + is_decode=False, +): + if not trans_weight: + raise RuntimeError( + "w8a8_group_gemm currently supports only trans_weight=True (TN layout)" + ) + if is_decode: + raise RuntimeError( + "w8a8_group_gemm decode path is disabled: vllm_iluvatar cuinfer decode currently reports internal error" + ) + _infinicore.w8a8_group_gemm_( + out._underlying, + input._underlying, + weight._underlying, + input_scale._underlying, + weight_scale._underlying, + tokens_per_experts._underlying, + None if sorted_token_ids is None else sorted_token_ids._underlying, + None if bias is None else bias._underlying, + trans_weight, + is_decode, + ) + return out diff --git a/src/infinicore/adaptor/aten_adaptor.cc b/src/infinicore/adaptor/aten_adaptor.cc index 73c01dfdb..1d041ebbc 100644 --- a/src/infinicore/adaptor/aten_adaptor.cc +++ b/src/infinicore/adaptor/aten_adaptor.cc @@ -1,6 +1,12 @@ #ifdef ENABLE_ATEN #include "infinicore/adaptor/aten_adaptor.hpp" +#include + +#if defined(ENABLE_ILUVATAR_API) +extern "C" int32_t torch_set_current_cuda_stream(void *stream, int32_t device_index); +#endif + namespace infinicore::adaptor { at::Tensor to_aten_tensor(const infinicore::Tensor &t) { @@ -44,6 +50,19 @@ c10::cuda::CUDAStream get_cuda_stream() { } #endif +#if defined(ENABLE_ILUVATAR_API) +void set_aten_stream_to_infinicore() { + const auto error = torch_set_current_cuda_stream( + infinicore::context::getStream(), + static_cast(infinicore::context::getDevice().getIndex())); + if (error != 0) { + throw std::runtime_error( + "torch_set_current_cuda_stream failed with error code " + + std::to_string(error)); + } +} +#endif + #if defined(ENABLE_MOORE_API) c10::musa::MUSAStream get_musa_stream() { return c10::musa::getStreamFromExternal( diff --git a/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc b/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc new file mode 100644 index 000000000..34cf54a22 --- /dev/null +++ b/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc @@ -0,0 +1,361 @@ +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#include "infinicore/adaptor/aten_adaptor.hpp" + +#if defined(ENABLE_ILUVATAR_API) +#define INFINICORE_VLLM_ILUVATAR_STREAM_GUARD() \ + infinicore::adaptor::set_aten_stream_to_infinicore() +#else +#define INFINICORE_VLLM_ILUVATAR_STREAM_GUARD() ((void)0) +#endif + +#include +#include +#include +#include +#include +#include + +namespace infinicore::adaptor::vllm_iluvatar { +namespace { + +using fused_add_rms_norm_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, float); +using dynamic_scaled_int8_quant_fn = void (*)(at::Tensor &, at::Tensor &, const at::Tensor &); +using concat_mla_q_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &); +using concat_and_cache_mla_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, const std::string &, at::Tensor &); +using concat_and_cache_mla_int8_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &); +using topk_softmax_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, const at::Tensor &, bool, std::optional); +using topk_sigmoid_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, const at::Tensor &, bool, std::optional); +using grouped_topk_fn = void (*)(at::Tensor &, at::Tensor &, const at::Tensor &, const std::optional &, int64_t, int64_t, std::string, bool); +using scaled_mm_w4a8_fn = void (*)(at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const std::optional &, bool); +using scaled_mm_w8a8_fn = scaled_mm_w4a8_fn; +using w4a8_group_gemm_fn = void (*)(at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const std::optional &, const std::optional &, bool, bool); +using w8a8_group_gemm_fn = void (*)(at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const std::optional &, const std::optional &, bool, bool); +using argsort_bincount_with_inv_pos_fn = void (*)(const at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, int64_t); +using expand_moe_input_with_inv_pos_fn = void (*)(at::Tensor &, std::optional, const at::Tensor &, const at::Tensor &, int64_t, int64_t, int64_t); +using silu_and_mul_quant_fn = void (*)(at::Tensor &, std::optional, const at::Tensor &, int64_t); +using moe_sum_vllm_fn = void (*)(at::Tensor &, const at::Tensor &, std::optional, std::optional, double, double); + +struct Symbols { + void *handle = nullptr; + fused_add_rms_norm_fn fused_add_rms_norm = nullptr; + dynamic_scaled_int8_quant_fn dynamic_scaled_int8_quant = nullptr; + concat_mla_q_fn concat_mla_q = nullptr; + concat_and_cache_mla_fn concat_and_cache_mla = nullptr; + concat_and_cache_mla_int8_fn concat_and_cache_mla_int8 = nullptr; + topk_softmax_fn topk_softmax = nullptr; + topk_sigmoid_fn topk_sigmoid = nullptr; + grouped_topk_fn grouped_topk = nullptr; + scaled_mm_w4a8_fn scaled_mm_w4a8 = nullptr; + scaled_mm_w8a8_fn scaled_mm_w8a8 = nullptr; + w4a8_group_gemm_fn w4a8_group_gemm = nullptr; + w8a8_group_gemm_fn w8a8_group_gemm = nullptr; + argsort_bincount_with_inv_pos_fn argsort_bincount_with_inv_pos = nullptr; + expand_moe_input_with_inv_pos_fn expand_moe_input_with_inv_pos = nullptr; + silu_and_mul_quant_fn silu_and_mul_quant = nullptr; + moe_sum_vllm_fn moe_sum_vllm = nullptr; + std::string error; +}; + +Symbols &symbols() { + static Symbols syms; + static std::once_flag once; + std::call_once(once, []() { +#if defined(ENABLE_ILUVATAR_API) + const char *env_path = std::getenv("INFINICORE_VLLM_ILUVATAR_SO"); + const char *paths[] = { + env_path, + "/usr/local/lib/python3.12/site-packages/vllm_iluvatar/_C.cpython-312-x86_64-linux-gnu.so", + "/usr/local/lib/python3.10/site-packages/vllm_iluvatar/_C.cpython-310-x86_64-linux-gnu.so", + }; + + for (const char *path : paths) { + if (path == nullptr || path[0] == '\0') { + continue; + } + void *handle = dlopen(path, RTLD_LAZY | RTLD_LOCAL); + if (!handle) { + syms.error = dlerror(); + continue; + } + auto fused_fn = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf18fused_add_rms_normERN2at6TensorES3_S3_f")); + auto quant_fn = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf25dynamic_scaled_int8_quantERN2at6TensorES3_RKS2_")); + auto concat_mla_q_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf12concat_mla_qERN2at6TensorES3_S3_")); + auto concat_and_cache_mla_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf20concat_and_cache_mlaERN2at6TensorES3_S3_S3_RKSsS3_")); + auto concat_and_cache_mla_int8_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf25concat_and_cache_mla_int8ERN2at6TensorES3_S3_S3_S3_S3_S3_")); + auto topk_softmax_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf12topk_softmaxERN2at6TensorES3_S3_RKS2_bSt8optionalIS2_E")); + auto topk_sigmoid_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf12topk_sigmoidERN2at6TensorES3_S3_RKS2_bSt8optionalIS2_E")); + auto grouped_topk_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf16moe_grouped_topkERN2at6TensorES3_RKS2_RKSt8optionalIS2_EllSsb")); + auto scaled_mm_w4a8_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer7cuinfer14scaled_mm_w4a8ERN2at6TensorERKS2_S5_S5_S5_RKSt8optionalIS2_Eb")); + auto scaled_mm_w8a8_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer7cuinfer9scaled_mmERN2at6TensorERKS2_S5_S5_S5_RKSt8optionalIS2_Eb")); + auto w4a8_group_gemm_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer7cuinfer15w4a8_group_gemmERN2at6TensorERKS2_S5_S5_S5_S5_RKSt8optionalIS2_ES9_bb")); + auto w8a8_group_gemm_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer7cuinfer15w8a8_group_gemmERN2at6TensorERKS2_S5_S5_S5_S5_RKSt8optionalIS2_ES9_bb")); + auto argsort_bincount_with_inv_pos_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf29argsort_bincount_with_inv_posERKN2at6TensorERS2_S5_S5_l")); + auto expand_moe_input_with_inv_pos_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf29expand_moe_input_with_inv_posERN2at6TensorESt8optionalIS2_ERKS2_S7_lll")); + auto silu_and_mul_quant_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf18silu_and_mul_quantERN2at6TensorESt8optionalIS2_ERKS2_l")); + auto moe_sum_vllm_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf7moe_sumERN2at6TensorERKS2_St8optionalIS2_ES7_dd")); + if (!fused_fn && !quant_fn && !concat_mla_q_fn_ptr && !concat_and_cache_mla_fn_ptr && !concat_and_cache_mla_int8_fn_ptr && !topk_softmax_fn_ptr && !topk_sigmoid_fn_ptr && !grouped_topk_fn_ptr && !scaled_mm_w4a8_fn_ptr && !scaled_mm_w8a8_fn_ptr && !w4a8_group_gemm_fn_ptr && !w8a8_group_gemm_fn_ptr && !argsort_bincount_with_inv_pos_fn_ptr && !expand_moe_input_with_inv_pos_fn_ptr && !silu_and_mul_quant_fn_ptr && !moe_sum_vllm_fn_ptr) { + syms.error = dlerror(); + dlclose(handle); + continue; + } + syms.handle = handle; + syms.fused_add_rms_norm = fused_fn; + syms.dynamic_scaled_int8_quant = quant_fn; + syms.concat_mla_q = concat_mla_q_fn_ptr; + syms.concat_and_cache_mla = concat_and_cache_mla_fn_ptr; + syms.concat_and_cache_mla_int8 = concat_and_cache_mla_int8_fn_ptr; + syms.topk_softmax = topk_softmax_fn_ptr; + syms.topk_sigmoid = topk_sigmoid_fn_ptr; + syms.grouped_topk = grouped_topk_fn_ptr; + syms.scaled_mm_w4a8 = scaled_mm_w4a8_fn_ptr; + syms.scaled_mm_w8a8 = scaled_mm_w8a8_fn_ptr; + syms.w4a8_group_gemm = w4a8_group_gemm_fn_ptr; + syms.w8a8_group_gemm = w8a8_group_gemm_fn_ptr; + syms.argsort_bincount_with_inv_pos = argsort_bincount_with_inv_pos_fn_ptr; + syms.expand_moe_input_with_inv_pos = expand_moe_input_with_inv_pos_fn_ptr; + syms.silu_and_mul_quant = silu_and_mul_quant_fn_ptr; + syms.moe_sum_vllm = moe_sum_vllm_fn_ptr; + syms.error.clear(); + return; + } + if (syms.error.empty()) { + syms.error = "vllm_iluvatar extension not found"; + } +#else + syms.error = "InfiniCore was not built with ENABLE_ILUVATAR_API"; +#endif + }); + return syms; +} + +} // namespace + +bool available() { + return symbols().fused_add_rms_norm != nullptr; +} + +bool dynamic_scaled_int8_quant_available() { + return symbols().dynamic_scaled_int8_quant != nullptr; +} + +bool concat_mla_q_available() { + return symbols().concat_mla_q != nullptr; +} + +bool concat_and_cache_mla_available() { + return symbols().concat_and_cache_mla != nullptr; +} + +bool concat_and_cache_mla_int8_available() { + return symbols().concat_and_cache_mla_int8 != nullptr; +} + +bool topk_softmax_available() { + return symbols().topk_softmax != nullptr; +} + +bool topk_sigmoid_available() { + return symbols().topk_sigmoid != nullptr; +} + +bool grouped_topk_available() { + return symbols().grouped_topk != nullptr; +} + +bool scaled_mm_w4a8_available() { + return symbols().scaled_mm_w4a8 != nullptr; +} + +bool scaled_mm_w8a8_available() { + return symbols().scaled_mm_w8a8 != nullptr; +} + +bool w4a8_group_gemm_available() { + return symbols().w4a8_group_gemm != nullptr; +} + +bool w8a8_group_gemm_available() { + return symbols().w8a8_group_gemm != nullptr; +} + +bool argsort_bincount_with_inv_pos_available() { + return symbols().argsort_bincount_with_inv_pos != nullptr; +} + +bool expand_moe_input_with_inv_pos_available() { + return symbols().expand_moe_input_with_inv_pos != nullptr; +} + +bool silu_and_mul_quant_available() { + return symbols().silu_and_mul_quant != nullptr; +} + +bool moe_sum_vllm_available() { + return symbols().moe_sum_vllm != nullptr; +} + +void fused_add_rms_norm(at::Tensor &input, at::Tensor &residual, at::Tensor &weight, float epsilon) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.fused_add_rms_norm) { + throw std::runtime_error("vllm_iluvatar fused_add_rms_norm unavailable: " + syms.error); + } + syms.fused_add_rms_norm(input, residual, weight, epsilon); +} + +void dynamic_scaled_int8_quant(at::Tensor &output, at::Tensor &input_scales, const at::Tensor &input) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.dynamic_scaled_int8_quant) { + throw std::runtime_error("vllm_iluvatar dynamic_scaled_int8_quant unavailable: " + syms.error); + } + syms.dynamic_scaled_int8_quant(output, input_scales, input); +} + +void concat_mla_q(at::Tensor &ql_nope, at::Tensor &q_pe, at::Tensor &q_out) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.concat_mla_q) { + throw std::runtime_error("vllm_iluvatar concat_mla_q unavailable: " + syms.error); + } + syms.concat_mla_q(ql_nope, q_pe, q_out); +} + +void concat_and_cache_mla(at::Tensor &kv_c, at::Tensor &k_pe, at::Tensor &kv_cache, at::Tensor &slot_mapping, const std::string &kv_cache_dtype, at::Tensor &scale) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.concat_and_cache_mla) { + throw std::runtime_error("vllm_iluvatar concat_and_cache_mla unavailable: " + syms.error); + } + syms.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale); +} + +void concat_and_cache_mla_int8(at::Tensor &kv_c_int8, at::Tensor &kv_c_scale, at::Tensor &k_pe_int8, at::Tensor &k_pe_scale, at::Tensor &kv_cache, at::Tensor &kv_cache_scale, at::Tensor &slot_mapping) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.concat_and_cache_mla_int8) { + throw std::runtime_error("vllm_iluvatar concat_and_cache_mla_int8 unavailable: " + syms.error); + } + syms.concat_and_cache_mla_int8(kv_c_int8, kv_c_scale, k_pe_int8, k_pe_scale, kv_cache, kv_cache_scale, slot_mapping); +} + +void topk_softmax(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional correction_bias) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.topk_softmax) { + throw std::runtime_error("vllm_iluvatar topk_softmax unavailable: " + syms.error); + } + syms.topk_softmax(topk_weights, topk_ids, token_expert_indices, gating_output, renormalize, correction_bias); +} + +void topk_sigmoid(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional correction_bias) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.topk_sigmoid) { + throw std::runtime_error("vllm_iluvatar topk_sigmoid unavailable: " + syms.error); + } + syms.topk_sigmoid(topk_weights, topk_ids, token_expert_indices, gating_output, renormalize, correction_bias); +} + +void grouped_topk(at::Tensor &topk_weights, at::Tensor &topk_ids, const at::Tensor &scores, std::optional bias, int64_t num_expert_group, int64_t topk_group, const std::string &scoring_func, bool renormalize) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.grouped_topk) { + throw std::runtime_error("vllm_iluvatar grouped_topk unavailable: " + syms.error); + } + syms.grouped_topk(topk_weights, topk_ids, scores, bias, num_expert_group, topk_group, scoring_func, renormalize); +} + +void scaled_mm_w4a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional bias, bool trans_weight) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.scaled_mm_w4a8) { + throw std::runtime_error("vllm_iluvatar scaled_mm_w4a8 unavailable: " + syms.error); + } + syms.scaled_mm_w4a8(out, a, b, a_scales, b_scales, bias, trans_weight); +} + +void scaled_mm_w8a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional bias, bool trans_weight) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.scaled_mm_w8a8) { + throw std::runtime_error("vllm_iluvatar scaled_mm W8A8 unavailable: " + syms.error); + } + syms.scaled_mm_w8a8(out, a, b, a_scales, b_scales, bias, trans_weight); +} + +void w4a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.w4a8_group_gemm) { + throw std::runtime_error("vllm_iluvatar w4a8_group_gemm unavailable: " + syms.error); + } + syms.w4a8_group_gemm(out, input, weight, input_scale, weight_scale, tokens_per_experts, sorted_token_ids, bias, trans_weight, is_decode); +} + +void w8a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.w8a8_group_gemm) { + throw std::runtime_error("vllm_iluvatar w8a8_group_gemm unavailable: " + syms.error); + } + syms.w8a8_group_gemm(out, input, weight, input_scale, weight_scale, tokens_per_experts, sorted_token_ids, bias, trans_weight, is_decode); +} + +void argsort_bincount_with_inv_pos(const at::Tensor &topk_ids, at::Tensor &tokens_per_experts, at::Tensor &sorted_indices, at::Tensor &inv_pos, int64_t num_experts) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.argsort_bincount_with_inv_pos) { + throw std::runtime_error("vllm_iluvatar argsort_bincount_with_inv_pos unavailable: " + syms.error); + } + syms.argsort_bincount_with_inv_pos(topk_ids, tokens_per_experts, sorted_indices, inv_pos, num_experts); +} + +void expand_moe_input_with_inv_pos(at::Tensor &expand_states, std::optional expand_scales, const at::Tensor &hidden_states, const at::Tensor &inv_pos, int64_t top_k, int64_t group_size, int64_t format) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.expand_moe_input_with_inv_pos) { + throw std::runtime_error("vllm_iluvatar expand_moe_input_with_inv_pos unavailable: " + syms.error); + } + syms.expand_moe_input_with_inv_pos(expand_states, expand_scales, hidden_states, inv_pos, top_k, group_size, format); +} + +void silu_and_mul_quant(at::Tensor &output, std::optional output_scale, const at::Tensor &input, int64_t format) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.silu_and_mul_quant) { + throw std::runtime_error("vllm_iluvatar silu_and_mul_quant unavailable: " + syms.error); + } + syms.silu_and_mul_quant(output, output_scale, input, format); +} + +void moe_sum_vllm(at::Tensor &output, const at::Tensor &input, std::optional topk_weights, std::optional extra_residual, double routed_scale, double residual_scale) { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &syms = symbols(); + if (!syms.moe_sum_vllm) { + throw std::runtime_error("vllm_iluvatar moe_sum unavailable: " + syms.error); + } + syms.moe_sum_vllm(output, input, topk_weights, extra_residual, routed_scale, residual_scale); +} + +#undef INFINICORE_VLLM_ILUVATAR_STREAM_GUARD + +} // namespace infinicore::adaptor::vllm_iluvatar +#endif // ENABLE_ATEN diff --git a/src/infinicore/context/allocators/pinnable_block_allocator.cc b/src/infinicore/context/allocators/pinnable_block_allocator.cc index 32e5c5e9b..3e1c5c43b 100644 --- a/src/infinicore/context/allocators/pinnable_block_allocator.cc +++ b/src/infinicore/context/allocators/pinnable_block_allocator.cc @@ -21,17 +21,14 @@ inline size_t align_up(size_t size, size_t alignment) { PinnableBlockAllocator::PinnableBlockAllocator(Device device) : device_(device) { size_classes_ = { - {32 * 1024, {}}, // 32 KB - {256 * 1024, {}}, // 256 KB - {1 * 1024 * 1024, {}}, // 1 MB - {2 * 1024 * 1024, {}}, // 2 MB - {4 * 1024 * 1024, {}}, // 4 MB - {8 * 1024 * 1024, {}}, // 8 MB - {16 * 1024 * 1024, {}}, // 16 MB - {32 * 1024 * 1024, {}}, // 32 MB - {64 * 1024 * 1024, {}}, // 64 MB - {128 * 1024 * 1024, {}}, // 128 MB - {256 * 1024 * 1024, {}}, // 256 MB + {32 * 1024, {}}, // 32 KB + {256 * 1024, {}}, // 256 KB + {1 * 1024 * 1024, {}}, // 1 MB + {2 * 1024 * 1024, {}}, // 2 MB + {4 * 1024 * 1024, {}}, // 4 MB + {8 * 1024 * 1024, {}}, // 8 MB + {16 * 1024 * 1024, {}}, // 16 MB + {32 * 1024 * 1024, {}}, // 32 MB }; } diff --git a/src/infinicore/nn/parameter.cc b/src/infinicore/nn/parameter.cc index 1b8a82762..49a8ebe4a 100644 --- a/src/infinicore/nn/parameter.cc +++ b/src/infinicore/nn/parameter.cc @@ -72,13 +72,25 @@ void Parameter::load_no_sync(const Tensor &tensor) { if (num_shards_ == 0 || num_shards_ >= tp_size_) { expected_shape[tp_dim_] *= tp_size_; + Tensor source_tensor = tensor; if (expected_shape != tensor->shape()) { - throw std::runtime_error("Shape mismatch when loading tensor into parameter. Weight: " + impl_->info() + ", Tensor: " + tensor->info() + "."); + bool degenerate_2d_transpose = false; + const auto &actual_shape = tensor->shape(); + if (expected_shape.size() == 2 && actual_shape.size() == 2 + && expected_shape[0] == actual_shape[1] + && expected_shape[1] == actual_shape[0] + && (expected_shape[0] == 1 || expected_shape[1] == 1)) { + degenerate_2d_transpose = true; + } + if (!degenerate_2d_transpose) { + throw std::runtime_error("Shape mismatch when loading tensor into parameter. Weight: " + impl_->info() + ", Tensor: " + tensor->info() + "."); + } + source_tensor = tensor->is_contiguous() ? tensor->view(expected_shape) : tensor->contiguous()->view(expected_shape); } if (tp_size_ > 1) { - impl_->copy_from(tensor->narrow({{tp_dim_, tp_rank_ * impl_->size(tp_dim_), impl_->size(tp_dim_)}})); + impl_->copy_from(source_tensor->narrow({{tp_dim_, tp_rank_ * impl_->size(tp_dim_), impl_->size(tp_dim_)}})); } else { - impl_->copy_from(tensor); + impl_->copy_from(source_tensor); } } else { if (num_shards_ == 0) { diff --git a/src/infinicore/ops/add_rms_norm/add_rms_norm.cc b/src/infinicore/ops/add_rms_norm/add_rms_norm.cc index ccba62e21..f856d7b85 100644 --- a/src/infinicore/ops/add_rms_norm/add_rms_norm.cc +++ b/src/infinicore/ops/add_rms_norm/add_rms_norm.cc @@ -1,5 +1,10 @@ #include "infinicore/ops/add_rms_norm.hpp" +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + #include "../../utils.hpp" namespace infinicore::op { @@ -27,6 +32,16 @@ void add_rms_norm_(Tensor out, Tensor residual, const Tensor &a, const Tensor &b } void add_rms_norm_inplace(Tensor input, Tensor residual, const Tensor &weight, float epsilon) { +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (input->device().getType() == Device::Type::ILUVATAR && adaptor::vllm_iluvatar::available()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(input, residual, weight); + auto input_at = adaptor::to_aten_tensor(input); + auto residual_at = adaptor::to_aten_tensor(residual); + auto weight_at = adaptor::to_aten_tensor(weight); + adaptor::vllm_iluvatar::fused_add_rms_norm(input_at, residual_at, weight_at, epsilon); + return; + } +#endif add_rms_norm_(input, residual, input, residual, weight, epsilon); } diff --git a/src/infinicore/ops/cast/cast.cc b/src/infinicore/ops/cast/cast.cc new file mode 100644 index 000000000..a1b6d4500 --- /dev/null +++ b/src/infinicore/ops/cast/cast.cc @@ -0,0 +1,21 @@ +#include "infinicore/ops/cast.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) +#include "infinicore/adaptor/aten_adaptor.hpp" +#endif +namespace infinicore::op { +void cast_(Tensor out, const Tensor &in) { + if (!out || !in || out->shape() != in->shape()) { + throw std::runtime_error("cast_ expects equal non-empty shapes"); + } + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, in); +#if defined(ENABLE_ATEN) + auto o = adaptor::to_aten_tensor(out), x = adaptor::to_aten_tensor(in); + o.copy_(x); + return; +#else + throw std::runtime_error("cast_ requires ATen"); +#endif +} +} // namespace infinicore::op diff --git a/src/infinicore/ops/concat_and_cache_mla/concat_and_cache_mla.cc b/src/infinicore/ops/concat_and_cache_mla/concat_and_cache_mla.cc new file mode 100644 index 000000000..d4a121a5c --- /dev/null +++ b/src/infinicore/ops/concat_and_cache_mla/concat_and_cache_mla.cc @@ -0,0 +1,77 @@ +#include "infinicore/ops/concat_and_cache_mla.hpp" +#include "../../utils.hpp" + +#include + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { + +namespace { + +void validate_concat_and_cache_mla(const Tensor &kv_c, + const Tensor &k_pe, + Tensor kv_cache, + const Tensor &slot_mapping, + const std::string &kv_cache_dtype, + Tensor scale) { + if (!kv_c || !k_pe || !kv_cache || !slot_mapping || !scale) { + throw std::runtime_error("concat_and_cache_mla expects non-empty kv_c, k_pe, kv_cache, slot_mapping, and scale tensors"); + } + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(kv_c, k_pe, kv_cache, slot_mapping, scale); + if (kv_cache_dtype != "auto" && kv_cache_dtype != "fp8" && kv_cache_dtype != "fp8_e4m3" && kv_cache_dtype != "fp8_e5m2") { + throw std::runtime_error("concat_and_cache_mla expects kv_cache_dtype to be auto/fp8/fp8_e4m3/fp8_e5m2"); + } + if (kv_c->ndim() != 2 || k_pe->ndim() != 2) { + throw std::runtime_error("concat_and_cache_mla expects kv_c and k_pe to be 2D [tokens, dim]"); + } + if (kv_c->size(0) != k_pe->size(0) || kv_c->size(0) != slot_mapping->numel()) { + throw std::runtime_error("concat_and_cache_mla expects kv_c/k_pe tokens to match slot_mapping numel"); + } + const auto head_dim = kv_c->size(1) + k_pe->size(1); + if (kv_cache->ndim() < 3 || kv_cache->size(kv_cache->ndim() - 1) != head_dim) { + throw std::runtime_error("concat_and_cache_mla expects kv_cache last dim == kv_c.shape[-1] + k_pe.shape[-1]"); + } + if (slot_mapping->dtype() != DataType::I64 && slot_mapping->dtype() != DataType::I32) { + throw std::runtime_error("concat_and_cache_mla expects slot_mapping dtype int64 or int32"); + } + if (scale->dtype() != DataType::F32) { + throw std::runtime_error("concat_and_cache_mla expects scale dtype float32"); + } + if (!kv_c->is_contiguous() || !k_pe->is_contiguous() || !kv_cache->is_contiguous() || !slot_mapping->is_contiguous() || !scale->is_contiguous()) { + throw std::runtime_error("concat_and_cache_mla expects contiguous tensors"); + } +} + +} // namespace + +void concat_and_cache_mla_(const Tensor &kv_c, + const Tensor &k_pe, + Tensor kv_cache, + const Tensor &slot_mapping, + const std::string &kv_cache_dtype, + Tensor scale) { + validate_concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale); + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (kv_cache->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::concat_and_cache_mla_available()) { + throw std::runtime_error("concat_and_cache_mla requires vllm_iluvatar perf extension on Iluvatar"); + } + auto kv_c_at = adaptor::to_aten_tensor(kv_c); + auto k_pe_at = adaptor::to_aten_tensor(k_pe); + auto kv_cache_at = adaptor::to_aten_tensor(kv_cache); + auto slot_mapping_at = adaptor::to_aten_tensor(slot_mapping); + auto scale_at = adaptor::to_aten_tensor(scale); + adaptor::vllm_iluvatar::concat_and_cache_mla(kv_c_at, k_pe_at, kv_cache_at, slot_mapping_at, kv_cache_dtype, scale_at); + return; + } +#endif + + throw std::runtime_error("concat_and_cache_mla currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/concat_and_cache_mla_int8/concat_and_cache_mla_int8.cc b/src/infinicore/ops/concat_and_cache_mla_int8/concat_and_cache_mla_int8.cc new file mode 100644 index 000000000..b709b7bb8 --- /dev/null +++ b/src/infinicore/ops/concat_and_cache_mla_int8/concat_and_cache_mla_int8.cc @@ -0,0 +1,89 @@ +#include "infinicore/ops/concat_and_cache_mla_int8.hpp" +#include "../../utils.hpp" + +#include + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { + +namespace { + +void validate_concat_and_cache_mla_int8(const Tensor &kv_c_int8, + const Tensor &kv_c_scale, + const Tensor &k_pe_int8, + const Tensor &k_pe_scale, + Tensor kv_cache, + Tensor kv_cache_scale, + const Tensor &slot_mapping) { + if (!kv_c_int8 || !kv_c_scale || !k_pe_int8 || !k_pe_scale || !kv_cache || !kv_cache_scale || !slot_mapping) { + throw std::runtime_error("concat_and_cache_mla_int8 expects non-empty tensors"); + } + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(kv_c_int8, kv_c_scale, k_pe_int8, k_pe_scale, kv_cache, kv_cache_scale, slot_mapping); + if (kv_c_int8->dtype() != DataType::I8 || k_pe_int8->dtype() != DataType::I8 || kv_cache->dtype() != DataType::I8) { + throw std::runtime_error("concat_and_cache_mla_int8 expects int8 kv_c_int8, k_pe_int8, and kv_cache"); + } + if (kv_c_scale->dtype() != DataType::F32 || k_pe_scale->dtype() != DataType::F32 || kv_cache_scale->dtype() != DataType::F32) { + throw std::runtime_error("concat_and_cache_mla_int8 expects float32 scales"); + } + if (slot_mapping->dtype() != DataType::I64) { + throw std::runtime_error("concat_and_cache_mla_int8 expects int64 slot_mapping, matching vLLM flatten().to(torch.long)"); + } + if (kv_c_int8->ndim() != 2 || k_pe_int8->ndim() != 2) { + throw std::runtime_error("concat_and_cache_mla_int8 expects kv_c_int8/k_pe_int8 to be 2D [tokens, dim]"); + } + const auto tokens = kv_c_int8->size(0); + if (k_pe_int8->size(0) != tokens || slot_mapping->numel() != tokens) { + throw std::runtime_error("concat_and_cache_mla_int8 expects matching token counts"); + } + if (kv_c_scale->numel() != tokens || k_pe_scale->numel() != tokens) { + throw std::runtime_error("concat_and_cache_mla_int8 expects one kv_c and k_pe scale per token"); + } + const auto head_dim = kv_c_int8->size(1) + k_pe_int8->size(1); + if (kv_cache->ndim() < 3 || kv_cache->size(kv_cache->ndim() - 1) != head_dim) { + throw std::runtime_error("concat_and_cache_mla_int8 expects kv_cache last dim == kv_c_int8.shape[-1] + k_pe_int8.shape[-1]"); + } + if (kv_cache_scale->ndim() < 3 || kv_cache_scale->size(kv_cache_scale->ndim() - 1) != 2) { + throw std::runtime_error("concat_and_cache_mla_int8 expects kv_cache_scale last dim == 2"); + } + if (!kv_c_int8->is_contiguous() || !kv_c_scale->is_contiguous() || !k_pe_int8->is_contiguous() || !k_pe_scale->is_contiguous() + || !kv_cache->is_contiguous() || !kv_cache_scale->is_contiguous() || !slot_mapping->is_contiguous()) { + throw std::runtime_error("concat_and_cache_mla_int8 expects contiguous tensors"); + } +} + +} // namespace + +void concat_and_cache_mla_int8_(const Tensor &kv_c_int8, + const Tensor &kv_c_scale, + const Tensor &k_pe_int8, + const Tensor &k_pe_scale, + Tensor kv_cache, + Tensor kv_cache_scale, + const Tensor &slot_mapping) { + validate_concat_and_cache_mla_int8(kv_c_int8, kv_c_scale, k_pe_int8, k_pe_scale, kv_cache, kv_cache_scale, slot_mapping); + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (kv_cache->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::concat_and_cache_mla_int8_available()) { + throw std::runtime_error("concat_and_cache_mla_int8 requires vllm_iluvatar perf extension on Iluvatar"); + } + auto kv_c_int8_at = adaptor::to_aten_tensor(kv_c_int8); + auto kv_c_scale_at = adaptor::to_aten_tensor(kv_c_scale); + auto k_pe_int8_at = adaptor::to_aten_tensor(k_pe_int8); + auto k_pe_scale_at = adaptor::to_aten_tensor(k_pe_scale); + auto kv_cache_at = adaptor::to_aten_tensor(kv_cache); + auto kv_cache_scale_at = adaptor::to_aten_tensor(kv_cache_scale); + auto slot_mapping_at = adaptor::to_aten_tensor(slot_mapping); + adaptor::vllm_iluvatar::concat_and_cache_mla_int8(kv_c_int8_at, kv_c_scale_at, k_pe_int8_at, k_pe_scale_at, kv_cache_at, kv_cache_scale_at, slot_mapping_at); + return; + } +#endif + + throw std::runtime_error("concat_and_cache_mla_int8 currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/concat_mla_q/concat_mla_q.cc b/src/infinicore/ops/concat_mla_q/concat_mla_q.cc new file mode 100644 index 000000000..c46c322f9 --- /dev/null +++ b/src/infinicore/ops/concat_mla_q/concat_mla_q.cc @@ -0,0 +1,83 @@ +#include "infinicore/ops/concat_mla_q.hpp" +#include "../../utils.hpp" + +#include + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { + +namespace { + +void validate_concat_mla_q(const Tensor &ql_nope, const Tensor &q_pe, Tensor q_out) { + if (!ql_nope || !q_pe || !q_out) { + throw std::runtime_error("concat_mla_q expects non-empty ql_nope, q_pe, and q_out tensors"); + } + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(ql_nope, q_pe, q_out); + if (ql_nope->dtype() != q_pe->dtype() || ql_nope->dtype() != q_out->dtype()) { + throw std::runtime_error("concat_mla_q expects ql_nope, q_pe, and q_out to have the same dtype"); + } + if (ql_nope->dtype() != DataType::F16 && ql_nope->dtype() != DataType::BF16 && ql_nope->dtype() != DataType::F32) { + throw std::runtime_error("concat_mla_q expects float16, bfloat16, or float32 tensors"); + } + if (ql_nope->ndim() != 3 || q_pe->ndim() != 3 || q_out->ndim() != 3) { + throw std::runtime_error("concat_mla_q expects 3D tensors [tokens, heads, dim]"); + } + if (ql_nope->size(0) != q_pe->size(0) || ql_nope->size(0) != q_out->size(0) + || ql_nope->size(1) != q_pe->size(1) || ql_nope->size(1) != q_out->size(1)) { + throw std::runtime_error("concat_mla_q expects matching first two dimensions"); + } + if (ql_nope->size(2) + q_pe->size(2) != q_out->size(2)) { + throw std::runtime_error("concat_mla_q expects q_out.shape[-1] == ql_nope.shape[-1] + q_pe.shape[-1]"); + } + // The current vllm_iluvatar perf kernel used by GLM-5.2 is specialized for + // MLA q concat with ql_nope_dim=512 and q_pe_dim=64. Other template stubs + // either do not cover all dtypes or have been observed to produce incorrect + // results, so keep the public wrapper constrained to the verified GLM path. + if (ql_nope->size(2) != 512 || q_pe->size(2) != 64 || q_out->size(2) != 576) { + throw std::runtime_error("concat_mla_q vllm_iluvatar bridge currently supports only GLM MLA dims 512 + 64 -> 576"); + } + if (!ql_nope->is_contiguous() || !q_pe->is_contiguous() || !q_out->is_contiguous()) { + throw std::runtime_error("concat_mla_q expects contiguous ql_nope, q_pe, and q_out tensors"); + } +} + +} // namespace + +void concat_mla_q_(const Tensor &ql_nope, const Tensor &q_pe, Tensor q_out) { + validate_concat_mla_q(ql_nope, q_pe, q_out); + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (q_out->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::concat_mla_q_available()) { + throw std::runtime_error("concat_mla_q requires vllm_iluvatar perf extension on Iluvatar"); + } + auto ql_nope_at = adaptor::to_aten_tensor(ql_nope); + auto q_pe_at = adaptor::to_aten_tensor(q_pe); + auto q_out_at = adaptor::to_aten_tensor(q_out); + adaptor::vllm_iluvatar::concat_mla_q(ql_nope_at, q_pe_at, q_out_at); + return; + } +#endif + + throw std::runtime_error("concat_mla_q currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} + +Tensor concat_mla_q(const Tensor &ql_nope, const Tensor &q_pe) { + if (!ql_nope || !q_pe) { + throw std::runtime_error("concat_mla_q expects non-empty input tensors"); + } + if (ql_nope->ndim() != 3 || q_pe->ndim() != 3) { + throw std::runtime_error("concat_mla_q expects 3D tensors [tokens, heads, dim]"); + } + Shape out_shape = ql_nope->shape(); + out_shape[2] = ql_nope->size(2) + q_pe->size(2); + auto q_out = Tensor::empty(out_shape, ql_nope->dtype(), ql_nope->device()); + concat_mla_q_(ql_nope, q_pe, q_out); + return q_out; +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/grouped_topk_vllm/grouped_topk_vllm.cc b/src/infinicore/ops/grouped_topk_vllm/grouped_topk_vllm.cc new file mode 100644 index 000000000..7c9cfb9ec --- /dev/null +++ b/src/infinicore/ops/grouped_topk_vllm/grouped_topk_vllm.cc @@ -0,0 +1,82 @@ +#include "infinicore/ops/grouped_topk_vllm.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif +namespace infinicore::op { +void grouped_topk_vllm_(Tensor topk_weights, Tensor topk_ids, const Tensor &scores, int64_t num_expert_group, int64_t topk_group, bool renormalize, float routed_scaling_factor, const Tensor &bias, const std::string &scoring_func) { + if (!topk_weights || !topk_ids || !scores) { + throw std::runtime_error("grouped_topk_vllm expects non-empty topk_weights, topk_ids, scores"); + } + if (bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_ids, scores, bias); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_ids, scores); + } + if (scores->ndim() != 2 || topk_weights->ndim() != 2 || topk_ids->ndim() != 2) { + throw std::runtime_error("grouped_topk_vllm expects 2D tensors"); + } + const auto tokens = scores->size(0), experts = scores->size(1), topk = topk_weights->size(1); + if (topk_weights->size(0) != tokens || topk_ids->size(0) != tokens || topk_ids->size(1) != topk) { + throw std::runtime_error("grouped_topk_vllm expects outputs (tokens, topk)"); + } + if (num_expert_group != 1 && num_expert_group != 8) { + throw std::runtime_error("grouped_topk_vllm currently supports num_expert_group 1 or 8"); + } + if (topk_group < 1 || topk_group > num_expert_group) { + throw std::runtime_error("grouped_topk_vllm expects 1 <= topk_group <= num_expert_group"); + } + if (!(experts == 64 || experts == 128 || experts == 160 || experts == 192 || experts == 256 || experts == 384)) { + throw std::runtime_error("grouped_topk_vllm supports expert counts 64/128/160/192/256/384"); + } + if (experts % static_cast(num_expert_group) != 0) { + throw std::runtime_error("grouped_topk_vllm expects experts divisible by num_expert_group"); + } + if (topk < 1 || topk > 32 || topk > experts) { + throw std::runtime_error("grouped_topk_vllm expects topk in [1,32] and <= experts"); + } + if (scores->dtype() != DataType::F16 && scores->dtype() != DataType::BF16) { + throw std::runtime_error("grouped_topk_vllm perf supports only fp16/bfloat16 scores"); + } + if (!bias) { + throw std::runtime_error("grouped_topk_vllm currently requires correction bias; vllm_iluvatar perf no-bias path mismatches reference"); + } + if (topk_weights->dtype() != DataType::F32) { + throw std::runtime_error("grouped_topk_vllm expects topk_weights float32"); + } + if (topk_ids->dtype() != DataType::I32 && topk_ids->dtype() != DataType::I64) { + throw std::runtime_error("grouped_topk_vllm expects topk_ids int32/int64"); + } + if (bias && (bias->numel() != experts || bias->dtype() != scores->dtype())) { + throw std::runtime_error("grouped_topk_vllm expects bias shape (experts,) and same dtype as scores"); + } + if (scoring_func != "softmax" && scoring_func != "sigmoid") { + throw std::runtime_error("grouped_topk_vllm scoring_func must be softmax or sigmoid"); + } + if (!topk_weights->is_contiguous() || !topk_ids->is_contiguous() || !scores->is_contiguous() || (bias && !bias->is_contiguous())) { + throw std::runtime_error("grouped_topk_vllm expects contiguous tensors"); + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (scores->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::grouped_topk_available()) { + throw std::runtime_error("grouped_topk requires vllm_iluvatar perf extension"); + } + auto w = adaptor::to_aten_tensor(topk_weights); + auto ids = adaptor::to_aten_tensor(topk_ids); + auto s = adaptor::to_aten_tensor(scores); + std::optional b; + if (bias) { + b = adaptor::to_aten_tensor(bias); + } + adaptor::vllm_iluvatar::grouped_topk(w, ids, s, b, num_expert_group, topk_group, scoring_func, renormalize); + if (routed_scaling_factor != 1.0f) { + w.mul_(routed_scaling_factor); + } + return; + } +#endif + throw std::runtime_error("grouped_topk_vllm currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} +} // namespace infinicore::op diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc index 0167c17df..7db38c72c 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc @@ -47,6 +47,8 @@ void run(void *planned_meta) { #ifdef ENABLE_FLASH_ATTN #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); +#elif defined(ENABLE_ILUVATAR_API) + infinicore::adaptor::set_aten_stream_to_infinicore(); #endif auto *p = reinterpret_cast(planned_meta); diff --git a/src/infinicore/ops/moe_argsort_bincount/moe_argsort_bincount.cc b/src/infinicore/ops/moe_argsort_bincount/moe_argsort_bincount.cc new file mode 100644 index 000000000..d2de5b358 --- /dev/null +++ b/src/infinicore/ops/moe_argsort_bincount/moe_argsort_bincount.cc @@ -0,0 +1,44 @@ +#include "infinicore/ops/moe_argsort_bincount.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { + +void moe_argsort_bincount_with_inv_pos_(Tensor tokens_per_experts, Tensor sorted_indices, Tensor inv_pos, const Tensor &topk_ids, int64_t num_experts) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(tokens_per_experts, sorted_indices, inv_pos, topk_ids); + if (num_experts <= 0 || num_experts > 512) { + throw std::runtime_error("moe_argsort_bincount_with_inv_pos expects 0 < num_experts <= 512"); + } + if (topk_ids->dtype() != DataType::I32 || tokens_per_experts->dtype() != DataType::I32 || sorted_indices->dtype() != DataType::I32 || inv_pos->dtype() != DataType::I32) { + throw std::runtime_error("moe_argsort_bincount_with_inv_pos expects int32 tensors"); + } + if (tokens_per_experts->ndim() != 1 || tokens_per_experts->numel() != static_cast(num_experts)) { + throw std::runtime_error("moe_argsort_bincount_with_inv_pos tokens_per_experts shape mismatch"); + } + if (sorted_indices->ndim() != 1 || inv_pos->ndim() != 1 || sorted_indices->numel() != topk_ids->numel() || inv_pos->numel() != topk_ids->numel()) { + throw std::runtime_error("moe_argsort_bincount_with_inv_pos sorted_indices/inv_pos shape mismatch"); + } + if (!topk_ids->is_contiguous() || !tokens_per_experts->is_contiguous() || !sorted_indices->is_contiguous() || !inv_pos->is_contiguous()) { + throw std::runtime_error("moe_argsort_bincount_with_inv_pos expects contiguous tensors"); + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (topk_ids->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::argsort_bincount_with_inv_pos_available()) { + throw std::runtime_error("moe_argsort_bincount_with_inv_pos requires vllm_iluvatar perf extension"); + } + auto ids = adaptor::to_aten_tensor(topk_ids); + auto tpe = adaptor::to_aten_tensor(tokens_per_experts); + auto sorted = adaptor::to_aten_tensor(sorted_indices); + auto inv = adaptor::to_aten_tensor(inv_pos); + adaptor::vllm_iluvatar::argsort_bincount_with_inv_pos(ids, tpe, sorted, inv, num_experts); + return; + } +#endif + throw std::runtime_error("moe_argsort_bincount_with_inv_pos currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_expand_input/moe_expand_input.cc b/src/infinicore/ops/moe_expand_input/moe_expand_input.cc new file mode 100644 index 000000000..49669cd53 --- /dev/null +++ b/src/infinicore/ops/moe_expand_input/moe_expand_input.cc @@ -0,0 +1,91 @@ +#include "infinicore/ops/moe_expand_input.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { + +void moe_expand_input_with_inv_pos_(Tensor expand_states, std::optional expand_scales, const Tensor &hidden_states, const Tensor &inv_pos, int64_t top_k, int64_t group_size, int64_t format) { + if (expand_scales) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(expand_states, *expand_scales, hidden_states, inv_pos); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(expand_states, hidden_states, inv_pos); + } + if (top_k <= 0 || top_k > 32) { + throw std::runtime_error("moe_expand_input_with_inv_pos expects 1 <= top_k <= 32"); + } + if (group_size != 64 && group_size != 128) { + throw std::runtime_error("moe_expand_input_with_inv_pos expects group_size 64 or 128"); + } + if (format < 0 || format > 2) { + throw std::runtime_error("moe_expand_input_with_inv_pos format must be 0 normal, 1 quant, or 2 packed"); + } + if (hidden_states->ndim() != 2 || expand_states->ndim() != 2 || inv_pos->ndim() != 1) { + throw std::runtime_error("moe_expand_input_with_inv_pos expects 2D hidden/expand and 1D inv_pos"); + } + if (hidden_states->dtype() != DataType::F16 && hidden_states->dtype() != DataType::BF16) { + throw std::runtime_error("moe_expand_input_with_inv_pos expects fp16/bfloat16 hidden_states"); + } + if (inv_pos->dtype() != DataType::I32) { + throw std::runtime_error("moe_expand_input_with_inv_pos expects int32 inv_pos"); + } + if (!hidden_states->is_contiguous() || !expand_states->is_contiguous() || !inv_pos->is_contiguous() || (expand_scales && !(*expand_scales)->is_contiguous())) { + throw std::runtime_error("moe_expand_input_with_inv_pos expects contiguous tensors"); + } + const size_t m = hidden_states->size(0); + const size_t n = hidden_states->size(1); + const size_t total = m * static_cast(top_k); + if (inv_pos->numel() != total || expand_states->size(0) != total) { + throw std::runtime_error("moe_expand_input_with_inv_pos total token shape mismatch"); + } + if (format == 0) { + if (expand_scales) { + throw std::runtime_error("moe_expand_input_with_inv_pos normal format does not take expand_scales"); + } + if (expand_states->dtype() != hidden_states->dtype() || expand_states->size(1) != n) { + throw std::runtime_error("moe_expand_input_with_inv_pos normal output shape/dtype mismatch"); + } + if (n > 16384) { + throw std::runtime_error("moe_expand_input_with_inv_pos normal format supports hidden <= 16384"); + } + } else { + const size_t n_out = ((n + static_cast(group_size) - 1) / static_cast(group_size)) * static_cast(group_size); + if (!expand_scales) { + throw std::runtime_error("moe_expand_input_with_inv_pos quant/packed format requires expand_scales"); + } + if (expand_states->dtype() != DataType::I8 || (*expand_scales)->dtype() != DataType::F32) { + throw std::runtime_error("moe_expand_input_with_inv_pos quant/packed expects int8 output and fp32 scales"); + } + if (expand_states->size(1) != n_out || (*expand_scales)->ndim() != 2 || (*expand_scales)->size(0) != total || (*expand_scales)->size(1) != 1) { + throw std::runtime_error("moe_expand_input_with_inv_pos quant/packed output shape mismatch"); + } + if (n_out > 32768) { + throw std::runtime_error("moe_expand_input_with_inv_pos quant/packed supports padded hidden <= 32768"); + } + if (format == 2 && (n_out % 64) != 0) { + throw std::runtime_error("moe_expand_input_with_inv_pos packed requires padded hidden % 64 == 0"); + } + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (hidden_states->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::expand_moe_input_with_inv_pos_available()) { + throw std::runtime_error("moe_expand_input_with_inv_pos requires vllm_iluvatar perf extension"); + } + auto out = adaptor::to_aten_tensor(expand_states); + auto hs = adaptor::to_aten_tensor(hidden_states); + auto inv = adaptor::to_aten_tensor(inv_pos); + std::optional scales; + if (expand_scales) { + scales = adaptor::to_aten_tensor(*expand_scales); + } + adaptor::vllm_iluvatar::expand_moe_input_with_inv_pos(out, scales, hs, inv, top_k, group_size, format); + return; + } +#endif + throw std::runtime_error("moe_expand_input_with_inv_pos currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_silu_and_mul_quant/moe_silu_and_mul_quant.cc b/src/infinicore/ops/moe_silu_and_mul_quant/moe_silu_and_mul_quant.cc new file mode 100644 index 000000000..c0e1bb546 --- /dev/null +++ b/src/infinicore/ops/moe_silu_and_mul_quant/moe_silu_and_mul_quant.cc @@ -0,0 +1,82 @@ +#include "infinicore/ops/moe_silu_and_mul_quant.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { + +void moe_silu_and_mul_quant_(Tensor output, std::optional output_scale, const Tensor &input, int64_t format) { + if (output_scale) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, *output_scale, input); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + } + if (format < 0 || format > 2) { + throw std::runtime_error("moe_silu_and_mul_quant format must be 0 normal, 1 quant, or 2 packed"); + } + if (input->ndim() != 2 || output->ndim() != 2) { + throw std::runtime_error("moe_silu_and_mul_quant expects 2D tensors"); + } + if (input->dtype() != DataType::F16 && input->dtype() != DataType::BF16) { + throw std::runtime_error("moe_silu_and_mul_quant expects fp16/bfloat16 input"); + } + if ((input->size(1) % 2) != 0) { + throw std::runtime_error("moe_silu_and_mul_quant input last dim must be even"); + } + if (!input->is_contiguous() || !output->is_contiguous() || (output_scale && !(*output_scale)->is_contiguous())) { + throw std::runtime_error("moe_silu_and_mul_quant expects contiguous tensors"); + } + const size_t m = input->size(0); + const size_t n = input->size(1) / 2; + if (output->size(0) != m || output->size(1) != n) { + throw std::runtime_error("moe_silu_and_mul_quant output shape mismatch"); + } + if (format == 0) { + if (output_scale) { + throw std::runtime_error("moe_silu_and_mul_quant normal format does not take output_scale"); + } + if (output->dtype() != input->dtype()) { + throw std::runtime_error("moe_silu_and_mul_quant normal output dtype mismatch"); + } + if ((n % 2) != 0 || n > 16384) { + throw std::runtime_error("moe_silu_and_mul_quant normal format requires N % 2 == 0 and N <= 16384"); + } + } else { + if (!output_scale) { + throw std::runtime_error("moe_silu_and_mul_quant quant/packed requires output_scale"); + } + if (output->dtype() != DataType::I8 || (*output_scale)->dtype() != DataType::F32) { + throw std::runtime_error("moe_silu_and_mul_quant quant/packed expects int8 output and fp32 scale"); + } + if ((*output_scale)->ndim() != 2 || (*output_scale)->size(0) != m || (*output_scale)->size(1) != 1) { + throw std::runtime_error("moe_silu_and_mul_quant scale shape mismatch"); + } + if ((n % 4) != 0 || n > 32768) { + throw std::runtime_error("moe_silu_and_mul_quant quant/packed requires N % 4 == 0 and N <= 32768"); + } + if (format == 2 && (n % 64) != 0) { + throw std::runtime_error("moe_silu_and_mul_quant packed requires N % 64 == 0"); + } + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (input->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::silu_and_mul_quant_available()) { + throw std::runtime_error("moe_silu_and_mul_quant requires vllm_iluvatar perf extension"); + } + auto out = adaptor::to_aten_tensor(output); + auto in = adaptor::to_aten_tensor(input); + std::optional scale; + if (output_scale) { + scale = adaptor::to_aten_tensor(*output_scale); + } + adaptor::vllm_iluvatar::silu_and_mul_quant(out, scale, in, format); + return; + } +#endif + throw std::runtime_error("moe_silu_and_mul_quant currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_sum_vllm/moe_sum_vllm.cc b/src/infinicore/ops/moe_sum_vllm/moe_sum_vllm.cc new file mode 100644 index 000000000..7358e2a03 --- /dev/null +++ b/src/infinicore/ops/moe_sum_vllm/moe_sum_vllm.cc @@ -0,0 +1,68 @@ +#include "infinicore/ops/moe_sum_vllm.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { + +void moe_sum_vllm_(Tensor output, const Tensor &input, std::optional topk_weights, std::optional extra_residual, double routed_scale, double residual_scale) { + if (topk_weights && extra_residual) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, *topk_weights, *extra_residual); + } else if (topk_weights) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, *topk_weights); + } else if (extra_residual) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, *extra_residual); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + } + if (input->ndim() != 3 || output->ndim() != 2) { + throw std::runtime_error("moe_sum_vllm expects input 3D and output 2D"); + } + if (input->dtype() != DataType::F16 && input->dtype() != DataType::BF16) { + throw std::runtime_error("moe_sum_vllm expects fp16/bfloat16 input"); + } + if (output->dtype() != input->dtype()) { + throw std::runtime_error("moe_sum_vllm output dtype must match input"); + } + if (!input->is_contiguous() || !output->is_contiguous() || (topk_weights && !(*topk_weights)->is_contiguous()) || (extra_residual && !(*extra_residual)->is_contiguous())) { + throw std::runtime_error("moe_sum_vllm expects contiguous tensors"); + } + const size_t n = input->size(0), t = input->size(1), h = input->size(2); + if (output->size(0) != n || output->size(1) != h) { + throw std::runtime_error("moe_sum_vllm output shape mismatch"); + } + if (h == 0 || (h % 2) != 0 || h > 16384) { + throw std::runtime_error("moe_sum_vllm requires 0 < H <= 16384 and H % 2 == 0"); + } + if (topk_weights && ((*topk_weights)->dtype() != DataType::F32 || (*topk_weights)->ndim() != 2 || (*topk_weights)->size(0) != n || (*topk_weights)->size(1) != t)) { + throw std::runtime_error("moe_sum_vllm topk_weights must be float32 [N,T]"); + } + if (extra_residual && ((*extra_residual)->dtype() != output->dtype() || (*extra_residual)->ndim() != 2 || (*extra_residual)->size(0) != n || (*extra_residual)->size(1) != h)) { + throw std::runtime_error("moe_sum_vllm extra_residual shape/dtype mismatch"); + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (output->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::moe_sum_vllm_available()) { + throw std::runtime_error("moe_sum_vllm requires vllm_iluvatar perf extension"); + } + auto out = adaptor::to_aten_tensor(output); + auto in = adaptor::to_aten_tensor(input); + std::optional tw; + if (topk_weights) { + tw = adaptor::to_aten_tensor(*topk_weights); + } + std::optional er; + if (extra_residual) { + er = adaptor::to_aten_tensor(*extra_residual); + } + adaptor::vllm_iluvatar::moe_sum_vllm(out, in, tw, er, routed_scale, residual_scale); + return; + } +#endif + throw std::runtime_error("moe_sum_vllm currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_topk_vllm/moe_topk_vllm.cc b/src/infinicore/ops/moe_topk_vllm/moe_topk_vllm.cc new file mode 100644 index 000000000..2f39a1951 --- /dev/null +++ b/src/infinicore/ops/moe_topk_vllm/moe_topk_vllm.cc @@ -0,0 +1,94 @@ +#include "infinicore/ops/moe_topk_vllm.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif +namespace infinicore::op { +namespace { +void validate_vllm_topk(Tensor topk_weights, Tensor topk_ids, Tensor token_expert_indices, const Tensor &gating_output, const Tensor &correction_bias, const char *name) { + if (!topk_weights || !topk_ids || !token_expert_indices || !gating_output) { + throw std::runtime_error(std::string(name) + " expects non-empty tensors"); + } + if (correction_bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_ids, token_expert_indices, gating_output, correction_bias); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_ids, token_expert_indices, gating_output); + } + if (gating_output->ndim() != 2 || topk_weights->ndim() != 2 || topk_ids->ndim() != 2 || token_expert_indices->ndim() != 2) { + throw std::runtime_error(std::string(name) + " expects 2D tensors"); + } + const auto tokens = gating_output->size(0), experts = gating_output->size(1), topk = topk_weights->size(1); + if (topk_weights->size(0) != tokens || topk_ids->size(0) != tokens || token_expert_indices->size(0) != tokens || topk_ids->size(1) != topk || token_expert_indices->size(1) != topk) { + throw std::runtime_error(std::string(name) + " expects output shapes (tokens, topk)"); + } + if (topk < 1 || topk > 32 || topk > experts || experts < 1 || experts > 512) { + throw std::runtime_error(std::string(name) + " supports experts in [1,512], topk in [1,32], topk<=experts"); + } + if ((gating_output->dtype() == DataType::F16 || gating_output->dtype() == DataType::BF16) && (experts % 2) != 0) { + throw std::runtime_error(std::string(name) + " requires even experts for fp16/bfloat16 gating"); + } + if (gating_output->dtype() != DataType::F16 && gating_output->dtype() != DataType::BF16 && gating_output->dtype() != DataType::F32) { + throw std::runtime_error(std::string(name) + " expects fp16/bfloat16/fp32 gating"); + } + if (topk_weights->dtype() != DataType::F32 && topk_weights->dtype() != gating_output->dtype()) { + throw std::runtime_error(std::string(name) + " expects topk_weights dtype float32 or same as gating"); + } + if (topk_ids->dtype() != DataType::I32 && topk_ids->dtype() != DataType::I64) { + throw std::runtime_error(std::string(name) + " expects topk_ids int32/int64"); + } + if (token_expert_indices->dtype() != DataType::I32) { + throw std::runtime_error(std::string(name) + " expects token_expert_indices int32"); + } + if (correction_bias && (correction_bias->numel() != experts || correction_bias->dtype() != gating_output->dtype())) { + throw std::runtime_error(std::string(name) + " expects correction_bias shape (experts,) and same dtype as gating"); + } + if (!topk_weights->is_contiguous() || !topk_ids->is_contiguous() || !token_expert_indices->is_contiguous() || !gating_output->is_contiguous() || (correction_bias && !correction_bias->is_contiguous())) { + throw std::runtime_error(std::string(name) + " expects contiguous tensors"); + } +} +} // namespace +void moe_topk_softmax_vllm_(Tensor topk_weights, Tensor topk_ids, Tensor token_expert_indices, const Tensor &gating_output, bool renormalize, const Tensor &correction_bias) { + validate_vllm_topk(topk_weights, topk_ids, token_expert_indices, gating_output, correction_bias, "moe_topk_softmax_vllm"); +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (gating_output->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::topk_softmax_available()) { + throw std::runtime_error("topk_softmax requires vllm_iluvatar perf extension"); + } + auto w = adaptor::to_aten_tensor(topk_weights); + auto ids = adaptor::to_aten_tensor(topk_ids); + auto src = adaptor::to_aten_tensor(token_expert_indices); + auto gate = adaptor::to_aten_tensor(gating_output); + std::optional bias; + if (correction_bias) { + bias = adaptor::to_aten_tensor(correction_bias); + } + adaptor::vllm_iluvatar::topk_softmax(w, ids, src, gate, renormalize, bias); + return; + } +#endif + throw std::runtime_error("moe_topk_softmax_vllm currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} +void moe_topk_sigmoid_vllm_(Tensor topk_weights, Tensor topk_ids, Tensor token_expert_indices, const Tensor &gating_output, bool renormalize, const Tensor &correction_bias) { + validate_vllm_topk(topk_weights, topk_ids, token_expert_indices, gating_output, correction_bias, "moe_topk_sigmoid_vllm"); +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (gating_output->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::topk_sigmoid_available()) { + throw std::runtime_error("topk_sigmoid requires vllm_iluvatar perf extension"); + } + auto w = adaptor::to_aten_tensor(topk_weights); + auto ids = adaptor::to_aten_tensor(topk_ids); + auto src = adaptor::to_aten_tensor(token_expert_indices); + auto gate = adaptor::to_aten_tensor(gating_output); + std::optional bias; + if (correction_bias) { + bias = adaptor::to_aten_tensor(correction_bias); + } + adaptor::vllm_iluvatar::topk_sigmoid(w, ids, src, gate, renormalize, bias); + return; + } +#endif + throw std::runtime_error("moe_topk_sigmoid_vllm currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} +} // namespace infinicore::op diff --git a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc index 4a233bf5b..760f78823 100644 --- a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc +++ b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc @@ -71,6 +71,8 @@ void run(void *planned_meta) { #else #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); +#elif defined(ENABLE_ILUVATAR_API) + infinicore::adaptor::set_aten_stream_to_infinicore(); #endif auto *p = reinterpret_cast(planned_meta); diff --git a/src/infinicore/ops/quant/dynamic_scaled_int8_quant/dynamic_scaled_int8_quant.cc b/src/infinicore/ops/quant/dynamic_scaled_int8_quant/dynamic_scaled_int8_quant.cc new file mode 100644 index 000000000..6fd83e154 --- /dev/null +++ b/src/infinicore/ops/quant/dynamic_scaled_int8_quant/dynamic_scaled_int8_quant.cc @@ -0,0 +1,73 @@ +#include "infinicore/ops/dynamic_scaled_int8_quant.hpp" +#include "../../../utils.hpp" + +#include +#include + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { + +namespace { + +void validate_dynamic_scaled_int8_quant(Tensor output, const Tensor &input, Tensor input_scales) { + if (!output || !input || !input_scales) { + throw std::runtime_error("dynamic_scaled_int8_quant expects non-empty output, input, and input_scales tensors"); + } + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, input_scales); + if (output->dtype() != DataType::I8) { + throw std::runtime_error("dynamic_scaled_int8_quant expects output dtype int8"); + } + if (input_scales->dtype() != DataType::F32) { + throw std::runtime_error("dynamic_scaled_int8_quant expects input_scales dtype float32"); + } + if (input->dtype() != DataType::F16 && input->dtype() != DataType::BF16) { + throw std::runtime_error("dynamic_scaled_int8_quant expects input dtype float16 or bfloat16"); + } + if (input->ndim() == 0 || input->size(input->ndim() - 1) == 0) { + throw std::runtime_error("dynamic_scaled_int8_quant expects input with a non-empty hidden dimension"); + } + if (output->numel() != input->numel()) { + throw std::runtime_error("dynamic_scaled_int8_quant expects output numel to equal input numel"); + } + const auto hidden_size = input->size(input->ndim() - 1); + const auto num_tokens = input->numel() / hidden_size; + if (input_scales->numel() != num_tokens) { + throw std::runtime_error("dynamic_scaled_int8_quant expects input_scales numel to equal input.numel / input.shape[-1]"); + } + if (!input->is_contiguous() || !output->is_contiguous() || !input_scales->is_contiguous()) { + throw std::runtime_error("dynamic_scaled_int8_quant expects contiguous input, output, and input_scales tensors"); + } +} + +} // namespace + +void dynamic_scaled_int8_quant_(Tensor output, const Tensor &input, Tensor input_scales) { + validate_dynamic_scaled_int8_quant(output, input, input_scales); + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (input->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::dynamic_scaled_int8_quant_available()) { + throw std::runtime_error("dynamic_scaled_int8_quant requires vllm_iluvatar perf extension on Iluvatar"); + } + auto output_at = adaptor::to_aten_tensor(output); + auto scales_at = adaptor::to_aten_tensor(input_scales); + auto input_at = adaptor::to_aten_tensor(input); + adaptor::vllm_iluvatar::dynamic_scaled_int8_quant(output_at, scales_at, input_at); + return; + } +#endif + + throw std::runtime_error("dynamic_scaled_int8_quant currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} + +Tensor dynamic_scaled_int8_quant(const Tensor &input, Tensor input_scales) { + auto output = Tensor::empty(input->shape(), DataType::I8, input->device()); + dynamic_scaled_int8_quant_(output, input, input_scales); + return output; +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/scaled_mm_w4a8/scaled_mm_w4a8.cc b/src/infinicore/ops/scaled_mm_w4a8/scaled_mm_w4a8.cc new file mode 100644 index 000000000..ff180259f --- /dev/null +++ b/src/infinicore/ops/scaled_mm_w4a8/scaled_mm_w4a8.cc @@ -0,0 +1,79 @@ +#include "infinicore/ops/scaled_mm_w4a8.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { +Tensor scaled_mm_w4a8(const Tensor &a, const Tensor &b, const Tensor &a_scales, const Tensor &b_scales, std::optional bias, bool trans_weight) { + if (a->ndim() != 2 || b->ndim() != 2) { + throw std::runtime_error("scaled_mm_w4a8 expects 2D a and b"); + } + const size_t m = a->size(0); + const size_t n = trans_weight ? b->size(0) : b->size(1) * 2; + Tensor out = Tensor::empty({m, n}, bias ? (*bias)->dtype() : DataType::F16, a->device()); + scaled_mm_w4a8_(out, a, b, a_scales, b_scales, bias, trans_weight); + return out; +} + +void scaled_mm_w4a8_(Tensor out, const Tensor &a, const Tensor &b, const Tensor &a_scales, const Tensor &b_scales, std::optional bias, bool trans_weight) { + if (bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, a, b, a_scales, b_scales, *bias); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, a, b, a_scales, b_scales); + } + if (a->ndim() != 2 || b->ndim() != 2 || out->ndim() != 2 || a_scales->ndim() != 2 || b_scales->ndim() != 2) { + throw std::runtime_error("scaled_mm_w4a8 expects 2D tensors"); + } + if (a->dtype() != DataType::I8 || b->dtype() != DataType::I8) { + throw std::runtime_error("scaled_mm_w4a8 expects int8 a and packed int8 b"); + } + if (a_scales->dtype() != DataType::F32 || b_scales->dtype() != DataType::F32) { + throw std::runtime_error("scaled_mm_w4a8 expects float32 scales"); + } + if (out->dtype() != DataType::F16 && out->dtype() != DataType::BF16) { + throw std::runtime_error("scaled_mm_w4a8 expects fp16/bfloat16 out"); + } + if (bias && ((*bias)->ndim() != 1 || (*bias)->dtype() != out->dtype() || (*bias)->numel() != out->size(1))) { + throw std::runtime_error("scaled_mm_w4a8 expects bias shape (N,) and same dtype as out"); + } + if (!out->is_contiguous() || !a->is_contiguous() || !b->is_contiguous() || !a_scales->is_contiguous() || !b_scales->is_contiguous() || (bias && !(*bias)->is_contiguous())) { + throw std::runtime_error("scaled_mm_w4a8 expects contiguous tensors"); + } + const size_t k = a->size(1); + if ((!trans_weight && b->size(0) != k) || (trans_weight && b->size(1) * 2 != k)) { + throw std::runtime_error("scaled_mm_w4a8 K dimension mismatch"); + } + const size_t n = trans_weight ? b->size(0) : b->size(1) * 2; + if (out->size(0) != a->size(0) || out->size(1) != n) { + throw std::runtime_error("scaled_mm_w4a8 out shape mismatch"); + } + if (a_scales->size(0) != a->size(0) || a_scales->size(1) != 1) { + throw std::runtime_error("scaled_mm_w4a8 expects a_scales shape (M,1)"); + } + if (b_scales->size(0) != out->size(1) || b_scales->size(1) != 1) { + throw std::runtime_error("scaled_mm_w4a8 expects b_scales shape (N,1)"); + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (out->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::scaled_mm_w4a8_available()) { + throw std::runtime_error("scaled_mm_w4a8 requires vllm_iluvatar cuinfer extension"); + } + auto o = adaptor::to_aten_tensor(out); + auto aa = adaptor::to_aten_tensor(a); + auto bb = adaptor::to_aten_tensor(b); + auto as = adaptor::to_aten_tensor(a_scales); + auto bs = adaptor::to_aten_tensor(b_scales); + std::optional bi; + if (bias) { + bi = adaptor::to_aten_tensor(*bias); + } + adaptor::vllm_iluvatar::scaled_mm_w4a8(o, aa, bb, as, bs, bi, trans_weight); + return; + } +#endif + throw std::runtime_error("scaled_mm_w4a8 currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} +} // namespace infinicore::op diff --git a/src/infinicore/ops/scaled_mm_w8a8/scaled_mm_w8a8.cc b/src/infinicore/ops/scaled_mm_w8a8/scaled_mm_w8a8.cc new file mode 100644 index 000000000..6fcb07297 --- /dev/null +++ b/src/infinicore/ops/scaled_mm_w8a8/scaled_mm_w8a8.cc @@ -0,0 +1,80 @@ +#include "infinicore/ops/scaled_mm_w8a8.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { +Tensor scaled_mm_w8a8(const Tensor &a, const Tensor &b, const Tensor &a_scales, + const Tensor &b_scales, std::optional bias, bool trans_weight) { + if (a->ndim() != 2 || b->ndim() != 2) { + throw std::runtime_error("scaled_mm_w8a8 expects 2D a and b"); + } + const size_t n = trans_weight ? b->size(0) : b->size(1); + Tensor out = Tensor::empty({a->size(0), n}, bias ? (*bias)->dtype() : DataType::F16, a->device()); + scaled_mm_w8a8_(out, a, b, a_scales, b_scales, bias, trans_weight); + return out; +} + +void scaled_mm_w8a8_(Tensor out, const Tensor &a, const Tensor &b, const Tensor &a_scales, + const Tensor &b_scales, std::optional bias, bool trans_weight) { + if (bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, a, b, a_scales, b_scales, *bias); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, a, b, a_scales, b_scales); + } + if (a->ndim() != 2 || b->ndim() != 2 || out->ndim() != 2 || a_scales->ndim() != 2 || b_scales->ndim() != 2) { + throw std::runtime_error("scaled_mm_w8a8 expects 2D tensors"); + } + if (a->dtype() != DataType::I8 || b->dtype() != DataType::I8) { + throw std::runtime_error("scaled_mm_w8a8 expects int8 a and b"); + } + if (a_scales->dtype() != DataType::F32 || b_scales->dtype() != DataType::F32) { + throw std::runtime_error("scaled_mm_w8a8 expects float32 scales"); + } + if (out->dtype() != DataType::F16 && out->dtype() != DataType::BF16) { + throw std::runtime_error("scaled_mm_w8a8 expects fp16/bfloat16 out"); + } + const size_t k = a->size(1); + const size_t n = trans_weight ? b->size(0) : b->size(1); + if ((!trans_weight && b->size(0) != k) || (trans_weight && b->size(1) != k)) { + throw std::runtime_error("scaled_mm_w8a8 K dimension mismatch"); + } + if (out->size(0) != a->size(0) || out->size(1) != n) { + throw std::runtime_error("scaled_mm_w8a8 out shape mismatch"); + } + if (a_scales->size(0) != a->size(0) || a_scales->size(1) != 1) { + throw std::runtime_error("scaled_mm_w8a8 expects a_scales (M,1)"); + } + if (b_scales->size(0) != n || b_scales->size(1) != 1) { + throw std::runtime_error("scaled_mm_w8a8 expects b_scales (N,1)"); + } + if (bias && ((*bias)->ndim() != 1 || (*bias)->dtype() != out->dtype() || (*bias)->numel() != n)) { + throw std::runtime_error("scaled_mm_w8a8 invalid bias"); + } + if (!out->is_contiguous() || !a->is_contiguous() || !b->is_contiguous() || !a_scales->is_contiguous() || !b_scales->is_contiguous() || (bias && !(*bias)->is_contiguous())) { + throw std::runtime_error("scaled_mm_w8a8 expects contiguous tensors"); + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (out->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::scaled_mm_w8a8_available()) { + throw std::runtime_error("scaled_mm_w8a8 requires vllm_iluvatar cuinfer extension"); + } + auto o = adaptor::to_aten_tensor(out); + auto aa = adaptor::to_aten_tensor(a); + auto bb = adaptor::to_aten_tensor(b); + auto as = adaptor::to_aten_tensor(a_scales); + auto bs = adaptor::to_aten_tensor(b_scales); + std::optional bi; + if (bias) { + bi = adaptor::to_aten_tensor(*bias); + } + adaptor::vllm_iluvatar::scaled_mm_w8a8(o, aa, bb, as, bs, bi, trans_weight); + return; + } +#endif + throw std::runtime_error("scaled_mm_w8a8 currently supports only Iluvatar with vllm_iluvatar"); +} +} // namespace infinicore::op diff --git a/src/infinicore/ops/vocab_parallel_embedding/vocab_parallel_embedding.cc b/src/infinicore/ops/vocab_parallel_embedding/vocab_parallel_embedding.cc new file mode 100644 index 000000000..e65152cb5 --- /dev/null +++ b/src/infinicore/ops/vocab_parallel_embedding/vocab_parallel_embedding.cc @@ -0,0 +1,26 @@ +#include "infinicore/ops/vocab_parallel_embedding.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include +#endif +namespace infinicore::op { +void vocab_parallel_embedding_(Tensor out, const Tensor &ids, const Tensor &w, int64_t start, int64_t end) { + if (!out || !ids || !w || w->ndim() != 2) { + throw std::runtime_error("vocab_parallel_embedding_: invalid tensors"); + } + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, ids, w); +#if defined(ENABLE_ATEN) + auto o = adaptor::to_aten_tensor(out), i = adaptor::to_aten_tensor(ids), wt = adaptor::to_aten_tensor(w); + auto mask = i.lt(start).logical_or(i.ge(end)); + auto local = (i - start).clamp(0, end - start - 1); + auto y = at::embedding(wt, local, -1, false, false); + y.masked_fill_(mask.unsqueeze(-1), 0); + o.copy_(y); + return; +#else + throw std::runtime_error("vocab_parallel_embedding_ requires ATen"); +#endif +} +} // namespace infinicore::op diff --git a/src/infinicore/ops/w4a8_group_gemm/w4a8_group_gemm.cc b/src/infinicore/ops/w4a8_group_gemm/w4a8_group_gemm.cc new file mode 100644 index 000000000..06376a401 --- /dev/null +++ b/src/infinicore/ops/w4a8_group_gemm/w4a8_group_gemm.cc @@ -0,0 +1,89 @@ +#include "infinicore/ops/w4a8_group_gemm.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif +namespace infinicore::op { +void w4a8_group_gemm_(Tensor out, const Tensor &input, const Tensor &weight, const Tensor &input_scale, const Tensor &weight_scale, const Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode) { + if (bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight, input_scale, weight_scale, *bias); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight, input_scale, weight_scale); + } + if (sorted_token_ids && (*sorted_token_ids)->device().getType() != Device::Type::CPU) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, *sorted_token_ids); + } + if (out->ndim() != 2 || input->ndim() != 2 || weight->ndim() != 3 || input_scale->ndim() != 2 || weight_scale->ndim() != 3 || tokens_per_experts->ndim() != 1) { + throw std::runtime_error("w4a8_group_gemm expects out/input 2D, weight/weight_scale 3D, tokens_per_experts 1D"); + } + if (out->dtype() != DataType::F16 && out->dtype() != DataType::BF16) { + throw std::runtime_error("w4a8_group_gemm expects fp16/bfloat16 output"); + } + if (input->dtype() != DataType::I8 || weight->dtype() != DataType::I8) { + throw std::runtime_error("w4a8_group_gemm expects int8 input and packed int8 weight"); + } + if (input_scale->dtype() != DataType::F32 || weight_scale->dtype() != DataType::F32) { + throw std::runtime_error("w4a8_group_gemm expects float32 scales"); + } + if (tokens_per_experts->dtype() != DataType::I32) { + throw std::runtime_error("w4a8_group_gemm expects int32 tokens_per_experts"); + } + if (sorted_token_ids && (*sorted_token_ids)->dtype() != DataType::I32) { + throw std::runtime_error("w4a8_group_gemm expects int32 sorted_token_ids"); + } + if (bias && ((*bias)->ndim() != 2 || (*bias)->dtype() != out->dtype())) { + throw std::runtime_error("w4a8_group_gemm expects bias shape (E,N) and same dtype as out"); + } + if (!out->is_contiguous() || !input->is_contiguous() || !weight->is_contiguous() || !input_scale->is_contiguous() || !weight_scale->is_contiguous() || !tokens_per_experts->is_contiguous() || (sorted_token_ids && !(*sorted_token_ids)->is_contiguous()) || (bias && !(*bias)->is_contiguous())) { + throw std::runtime_error("w4a8_group_gemm expects contiguous tensors"); + } + const size_t e = weight->size(0); + if (tokens_per_experts->numel() != e || weight_scale->size(0) != e) { + throw std::runtime_error("w4a8_group_gemm expert count mismatch"); + } + if (trans_weight) { + if (weight_scale->size(1) != weight->size(1) || weight_scale->size(2) != 1 || input->size(1) != weight->size(2) * 2 || out->size(1) != weight->size(1)) { + throw std::runtime_error("w4a8_group_gemm TN shape mismatch"); + } + } else { + if (weight_scale->size(0) != e || input->size(1) != weight->size(1) || out->size(1) != weight->size(2) * 2) { + throw std::runtime_error("w4a8_group_gemm NN shape mismatch"); + } + } + if (out->size(0) != input->size(0) || input_scale->size(0) != input->size(0) || input_scale->size(1) != 1) { + throw std::runtime_error("w4a8_group_gemm input/output scale shape mismatch"); + } + if (sorted_token_ids && (*sorted_token_ids)->numel() != out->size(0)) { + throw std::runtime_error("w4a8_group_gemm sorted_token_ids length mismatch"); + } + if (bias && ((*bias)->size(0) != e || (*bias)->size(1) != out->size(1))) { + throw std::runtime_error("w4a8_group_gemm bias shape mismatch"); + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (out->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::w4a8_group_gemm_available()) { + throw std::runtime_error("w4a8_group_gemm requires vllm_iluvatar cuinfer extension"); + } + auto o = adaptor::to_aten_tensor(out); + auto in = adaptor::to_aten_tensor(input); + auto w = adaptor::to_aten_tensor(weight); + auto is = adaptor::to_aten_tensor(input_scale); + auto ws = adaptor::to_aten_tensor(weight_scale); + auto tpe = adaptor::to_aten_tensor(tokens_per_experts); + std::optional sti; + if (sorted_token_ids) { + sti = adaptor::to_aten_tensor(*sorted_token_ids); + } + std::optional bi; + if (bias) { + bi = adaptor::to_aten_tensor(*bias); + } + adaptor::vllm_iluvatar::w4a8_group_gemm(o, in, w, is, ws, tpe, sti, bi, trans_weight, is_decode); + return; + } +#endif + throw std::runtime_error("w4a8_group_gemm currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} +} // namespace infinicore::op diff --git a/src/infinicore/ops/w8a8_group_gemm/w8a8_group_gemm.cc b/src/infinicore/ops/w8a8_group_gemm/w8a8_group_gemm.cc new file mode 100644 index 000000000..08b4a186f --- /dev/null +++ b/src/infinicore/ops/w8a8_group_gemm/w8a8_group_gemm.cc @@ -0,0 +1,89 @@ +#include "infinicore/ops/w8a8_group_gemm.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif +namespace infinicore::op { +void w8a8_group_gemm_(Tensor out, const Tensor &input, const Tensor &weight, const Tensor &input_scale, const Tensor &weight_scale, const Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode) { + if (bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight, input_scale, weight_scale, *bias); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight, input_scale, weight_scale); + } + if (!trans_weight) { + throw std::runtime_error("w8a8_group_gemm currently supports only trans_weight=True (TN layout)"); + } + if (is_decode) { + throw std::runtime_error("w8a8_group_gemm decode path is disabled: vllm_iluvatar cuinfer decode currently reports internal error"); + } + if (sorted_token_ids && (*sorted_token_ids)->device().getType() != Device::Type::CPU) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, *sorted_token_ids); + } + if (out->ndim() != 2 || input->ndim() != 2 || weight->ndim() != 3 || input_scale->ndim() != 2 || weight_scale->ndim() != 3 || tokens_per_experts->ndim() != 1) { + throw std::runtime_error("w8a8_group_gemm expects out/input 2D, weight/weight_scale 3D, tokens_per_experts 1D"); + } + if (out->dtype() != DataType::F16 && out->dtype() != DataType::BF16) { + throw std::runtime_error("w8a8_group_gemm expects fp16/bfloat16 output"); + } + if (input->dtype() != DataType::I8 || weight->dtype() != DataType::I8) { + throw std::runtime_error("w8a8_group_gemm expects int8 input and weight"); + } + if (input_scale->dtype() != DataType::F32 || weight_scale->dtype() != DataType::F32) { + throw std::runtime_error("w8a8_group_gemm expects float32 scales"); + } + if (tokens_per_experts->dtype() != DataType::I32) { + throw std::runtime_error("w8a8_group_gemm expects int32 tokens_per_experts"); + } + if (sorted_token_ids && (*sorted_token_ids)->dtype() != DataType::I32) { + throw std::runtime_error("w8a8_group_gemm expects int32 sorted_token_ids"); + } + if (bias && ((*bias)->ndim() != 2 || (*bias)->dtype() != out->dtype())) { + throw std::runtime_error("w8a8_group_gemm expects bias shape (E,N) and same dtype as out"); + } + if (!out->is_contiguous() || !input->is_contiguous() || !weight->is_contiguous() || !input_scale->is_contiguous() || !weight_scale->is_contiguous() || !tokens_per_experts->is_contiguous() || (sorted_token_ids && !(*sorted_token_ids)->is_contiguous()) || (bias && !(*bias)->is_contiguous())) { + throw std::runtime_error("w8a8_group_gemm expects contiguous tensors"); + } + const size_t e = weight->size(0); + if (tokens_per_experts->numel() != e || weight_scale->size(0) != e) { + throw std::runtime_error("w8a8_group_gemm expert count mismatch"); + } + if (weight_scale->size(1) != weight->size(1) || weight_scale->size(2) != 1 || input->size(1) != weight->size(2) || out->size(1) != weight->size(1)) { + throw std::runtime_error("w8a8_group_gemm TN shape mismatch"); + } + if (out->size(0) != input->size(0) || input_scale->size(0) != input->size(0) || input_scale->size(1) != 1) { + throw std::runtime_error("w8a8_group_gemm input/output scale shape mismatch"); + } + if (sorted_token_ids && (*sorted_token_ids)->numel() != out->size(0)) { + throw std::runtime_error("w8a8_group_gemm sorted_token_ids length mismatch"); + } + if (bias && ((*bias)->size(0) != e || (*bias)->size(1) != out->size(1))) { + throw std::runtime_error("w8a8_group_gemm bias shape mismatch"); + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (out->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::w8a8_group_gemm_available()) { + throw std::runtime_error("w8a8_group_gemm requires vllm_iluvatar cuinfer extension"); + } + auto o = adaptor::to_aten_tensor(out); + auto in = adaptor::to_aten_tensor(input); + auto w = adaptor::to_aten_tensor(weight); + auto is = adaptor::to_aten_tensor(input_scale); + auto ws = adaptor::to_aten_tensor(weight_scale); + auto tpe = adaptor::to_aten_tensor(tokens_per_experts); + std::optional sti; + if (sorted_token_ids) { + sti = adaptor::to_aten_tensor(*sorted_token_ids); + } + std::optional bi; + if (bias) { + bi = adaptor::to_aten_tensor(*bias); + } + adaptor::vllm_iluvatar::w8a8_group_gemm(o, in, w, is, ws, tpe, sti, bi, trans_weight, is_decode); + return; + } +#endif + throw std::runtime_error("w8a8_group_gemm currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} +} // namespace infinicore::op diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp index 087261382..0e9cb7f6d 100644 --- a/src/infinicore/pybind11/ops.hpp +++ b/src/infinicore/pybind11/ops.hpp @@ -36,11 +36,15 @@ #include "ops/causal_softmax.hpp" #include "ops/cdist.hpp" #include "ops/chunk_gated_delta_rule.hpp" +#include "ops/concat_and_cache_mla.hpp" +#include "ops/concat_and_cache_mla_int8.hpp" +#include "ops/concat_mla_q.hpp" #include "ops/conv2d.hpp" #include "ops/cross_entropy.hpp" #include "ops/diff.hpp" #include "ops/digamma.hpp" #include "ops/dist.hpp" +#include "ops/dynamic_scaled_int8_quant.hpp" #include "ops/embedding.hpp" #include "ops/equal.hpp" #include "ops/flash_attention.hpp" @@ -52,6 +56,7 @@ #include "ops/fmod.hpp" #include "ops/fused_gated_delta_net_gating.hpp" #include "ops/gaussian_nll_loss.hpp" +#include "ops/grouped_topk_vllm.hpp" #include "ops/hardswish.hpp" #include "ops/hardtanh.hpp" #include "ops/hinge_embedding_loss.hpp" @@ -82,7 +87,12 @@ #include "ops/mha.hpp" #include "ops/mha_kvcache.hpp" #include "ops/mha_varlen.hpp" +#include "ops/moe_argsort_bincount.hpp" +#include "ops/moe_expand_input.hpp" +#include "ops/moe_silu_and_mul_quant.hpp" +#include "ops/moe_sum_vllm.hpp" #include "ops/moe_topk_softmax.hpp" +#include "ops/moe_topk_vllm.hpp" #include "ops/mrope.hpp" #include "ops/mul.hpp" #include "ops/mul_scalar.hpp" @@ -106,6 +116,7 @@ #include "ops/rotmg.hpp" #include "ops/rwkv5_wkv.hpp" #include "ops/scal.hpp" +#include "ops/scaled_mm_w4a8.hpp" #include "ops/scatter.hpp" #include "ops/selu.hpp" #include "ops/sigmoid.hpp" @@ -131,6 +142,8 @@ #include "ops/vander.hpp" #include "ops/var.hpp" #include "ops/var_mean.hpp" +#include "ops/w4a8_group_gemm.hpp" +#include "ops/w8a8_group_gemm.hpp" namespace py = pybind11; @@ -167,6 +180,7 @@ inline void bind(py::module &m) { bind_diff(m); bind_digamma(m); bind_dist(m); + bind_dynamic_scaled_int8_quant(m); bind_flash_attention(m); bind_hinge_embedding_loss(m); bind_kv_caching(m); @@ -192,11 +206,17 @@ inline void bind(py::module &m) { bind_mha_kvcache(m); bind_mha_varlen(m); bind_mha(m); + bind_moe_argsort_bincount(m); + bind_moe_expand_input(m); + bind_moe_silu_and_mul_quant(m); + bind_moe_sum_vllm(m); bind_moe_topk_softmax(m); + bind_moe_topk_vllm(m); bind_mrope(m); bind_hardswish(m); bind_hardtanh(m); bind_gaussian_nll_loss(m); + bind_grouped_topk_vllm(m); bind_interpolate(m); bind_paged_attention(m); bind_paged_attention_prefill(m); @@ -206,6 +226,9 @@ inline void bind(py::module &m) { bind_random_sample(m); bind_cross_entropy(m); bind_conv2d(m); + bind_concat_mla_q(m); + bind_concat_and_cache_mla(m); + bind_concat_and_cache_mla_int8(m); bind_hypot(m); bind_take(m); bind_index_copy(m); @@ -224,6 +247,8 @@ inline void bind(py::module &m) { bind_logical_and(m); bind_logical_not(m); bind_vander(m); + bind_w4a8_group_gemm(m); + bind_w8a8_group_gemm(m); bind_unfold(m); bind_rope(m); bind_rot(m); @@ -236,6 +261,7 @@ inline void bind(py::module &m) { bind_flipud(m); bind_multi_margin_loss(m); bind_scatter(m); + bind_scaled_mm_w4a8(m); bind_scal(m); bind_broadcast_to(m); bind_softplus(m); diff --git a/src/infinicore/pybind11/ops/add_rms_norm.hpp b/src/infinicore/pybind11/ops/add_rms_norm.hpp index 5f9b243e5..c1c4c27c5 100644 --- a/src/infinicore/pybind11/ops/add_rms_norm.hpp +++ b/src/infinicore/pybind11/ops/add_rms_norm.hpp @@ -46,6 +46,13 @@ inline void bind_add_rms_norm(py::module &m) { weight: Scale weights epsilon: Small constant for numerical stability, default is 1e-5 )doc"); + m.def("add_rms_norm_inplace", + &op::add_rms_norm_inplace, + py::arg("input"), + py::arg("residual"), + py::arg("weight"), + py::arg("epsilon") = 1e-5f, + R"doc(In-place vLLM-style fused Add and RMSNorm: input = RMSNorm(input + residual), residual = input_old + residual.)doc"); } } // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/concat_and_cache_mla.hpp b/src/infinicore/pybind11/ops/concat_and_cache_mla.hpp new file mode 100644 index 000000000..9070091e2 --- /dev/null +++ b/src/infinicore/pybind11/ops/concat_and_cache_mla.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include + +#include "infinicore/ops/concat_and_cache_mla.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_concat_and_cache_mla(py::module &m) { + m.def("concat_and_cache_mla_", + &op::concat_and_cache_mla_, + py::arg("kv_c"), + py::arg("k_pe"), + py::arg("kv_cache"), + py::arg("slot_mapping"), + py::arg("kv_cache_dtype"), + py::arg("scale"), + R"doc(In-place MLA KV concat and paged cache write via vllm_iluvatar on Iluvatar.)doc"); +} + +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/concat_and_cache_mla_int8.hpp b/src/infinicore/pybind11/ops/concat_and_cache_mla_int8.hpp new file mode 100644 index 000000000..55ce5774c --- /dev/null +++ b/src/infinicore/pybind11/ops/concat_and_cache_mla_int8.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/concat_and_cache_mla_int8.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_concat_and_cache_mla_int8(py::module &m) { + m.def("concat_and_cache_mla_int8_", + &op::concat_and_cache_mla_int8_, + py::arg("kv_c_int8"), + py::arg("kv_c_scale"), + py::arg("k_pe_int8"), + py::arg("k_pe_scale"), + py::arg("kv_cache"), + py::arg("kv_cache_scale"), + py::arg("slot_mapping"), + R"doc(In-place int8 MLA KV cache write via vllm_iluvatar on Iluvatar.)doc"); +} + +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/concat_mla_q.hpp b/src/infinicore/pybind11/ops/concat_mla_q.hpp new file mode 100644 index 000000000..e6a164bb6 --- /dev/null +++ b/src/infinicore/pybind11/ops/concat_mla_q.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/concat_mla_q.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_concat_mla_q(py::module &m) { + m.def("concat_mla_q_", + &op::concat_mla_q_, + py::arg("ql_nope"), + py::arg("q_pe"), + py::arg("q_out"), + R"doc(In-place MLA query concat: q_out[..., :nope] = ql_nope; q_out[..., nope:] = q_pe.)doc"); + + m.def("concat_mla_q", + &op::concat_mla_q, + py::arg("ql_nope"), + py::arg("q_pe"), + R"doc(MLA query concat via vllm_iluvatar on Iluvatar.)doc"); +} + +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/dynamic_scaled_int8_quant.hpp b/src/infinicore/pybind11/ops/dynamic_scaled_int8_quant.hpp new file mode 100644 index 000000000..31667b716 --- /dev/null +++ b/src/infinicore/pybind11/ops/dynamic_scaled_int8_quant.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/dynamic_scaled_int8_quant.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_dynamic_scaled_int8_quant(py::module &m) { + m.def("dynamic_scaled_int8_quant_", + &op::dynamic_scaled_int8_quant_, + py::arg("output"), + py::arg("input"), + py::arg("input_scales"), + R"doc(In-place per-token dynamic scaled int8 quantization via vllm_iluvatar on Iluvatar.)doc"); + + m.def("dynamic_scaled_int8_quant", + &op::dynamic_scaled_int8_quant, + py::arg("input"), + py::arg("input_scales"), + R"doc(Per-token dynamic scaled int8 quantization via vllm_iluvatar on Iluvatar.)doc"); +} + +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/grouped_topk_vllm.hpp b/src/infinicore/pybind11/ops/grouped_topk_vllm.hpp new file mode 100644 index 000000000..0fa4eb6c0 --- /dev/null +++ b/src/infinicore/pybind11/ops/grouped_topk_vllm.hpp @@ -0,0 +1,17 @@ +#pragma once +#include "infinicore/ops/grouped_topk_vllm.hpp" +#include +namespace py = pybind11; +namespace infinicore::ops { +inline void bind_grouped_topk_vllm(py::module &m) { + m.def( + "grouped_topk_vllm_", [](Tensor topk_weights, Tensor topk_ids, Tensor scores, int64_t num_expert_group, int64_t topk_group, bool renormalize, float routed_scaling_factor, py::object bias, const std::string &scoring_func) { + Tensor b; + if (!bias.is_none()) { + b = bias.cast(); + } + op::grouped_topk_vllm_(topk_weights, topk_ids, scores, num_expert_group, topk_group, renormalize, routed_scaling_factor, b, scoring_func); + }, + py::arg("topk_weights"), py::arg("topk_ids"), py::arg("scores"), py::arg("num_expert_group"), py::arg("topk_group"), py::arg("renormalize"), py::arg("routed_scaling_factor"), py::arg("bias") = py::none(), py::arg("scoring_func") = "softmax", R"doc(vLLM-style grouped_topk routing.)doc"); +} +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/moe_argsort_bincount.hpp b/src/infinicore/pybind11/ops/moe_argsort_bincount.hpp new file mode 100644 index 000000000..7a9b0f7f3 --- /dev/null +++ b/src/infinicore/pybind11/ops/moe_argsort_bincount.hpp @@ -0,0 +1,10 @@ +#pragma once +#include "infinicore/ops/moe_argsort_bincount.hpp" +#include +namespace py = pybind11; + +namespace infinicore::ops { +inline void bind_moe_argsort_bincount(py::module &m) { + m.def("moe_argsort_bincount_with_inv_pos_", &op::moe_argsort_bincount_with_inv_pos_, py::arg("tokens_per_experts"), py::arg("sorted_indices"), py::arg("inv_pos"), py::arg("topk_ids"), py::arg("num_experts")); +} +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/moe_expand_input.hpp b/src/infinicore/pybind11/ops/moe_expand_input.hpp new file mode 100644 index 000000000..eb5a9679a --- /dev/null +++ b/src/infinicore/pybind11/ops/moe_expand_input.hpp @@ -0,0 +1,10 @@ +#pragma once +#include "infinicore/ops/moe_expand_input.hpp" +#include +namespace py = pybind11; + +namespace infinicore::ops { +inline void bind_moe_expand_input(py::module &m) { + m.def("moe_expand_input_with_inv_pos_", &op::moe_expand_input_with_inv_pos_, py::arg("expand_states"), py::arg("expand_scales"), py::arg("hidden_states"), py::arg("inv_pos"), py::arg("top_k"), py::arg("group_size") = 128, py::arg("format") = 0); +} +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/moe_silu_and_mul_quant.hpp b/src/infinicore/pybind11/ops/moe_silu_and_mul_quant.hpp new file mode 100644 index 000000000..f90d7d482 --- /dev/null +++ b/src/infinicore/pybind11/ops/moe_silu_and_mul_quant.hpp @@ -0,0 +1,10 @@ +#pragma once +#include "infinicore/ops/moe_silu_and_mul_quant.hpp" +#include +namespace py = pybind11; + +namespace infinicore::ops { +inline void bind_moe_silu_and_mul_quant(py::module &m) { + m.def("moe_silu_and_mul_quant_", &op::moe_silu_and_mul_quant_, py::arg("output"), py::arg("output_scale"), py::arg("input"), py::arg("format") = 0); +} +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/moe_sum_vllm.hpp b/src/infinicore/pybind11/ops/moe_sum_vllm.hpp new file mode 100644 index 000000000..3c8f6255a --- /dev/null +++ b/src/infinicore/pybind11/ops/moe_sum_vllm.hpp @@ -0,0 +1,10 @@ +#pragma once +#include "infinicore/ops/moe_sum_vllm.hpp" +#include +namespace py = pybind11; + +namespace infinicore::ops { +inline void bind_moe_sum_vllm(py::module &m) { + m.def("moe_sum_vllm_", &op::moe_sum_vllm_, py::arg("output"), py::arg("input"), py::arg("topk_weights") = std::nullopt, py::arg("extra_residual") = std::nullopt, py::arg("routed_scale") = 1.0, py::arg("residual_scale") = 1.0); +} +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/moe_topk_vllm.hpp b/src/infinicore/pybind11/ops/moe_topk_vllm.hpp new file mode 100644 index 000000000..4a9755665 --- /dev/null +++ b/src/infinicore/pybind11/ops/moe_topk_vllm.hpp @@ -0,0 +1,26 @@ +#pragma once +#include "infinicore/ops/moe_topk_vllm.hpp" +#include +namespace py = pybind11; +namespace infinicore::ops { +inline void bind_moe_topk_vllm(py::module &m) { + m.def( + "moe_topk_softmax_vllm_", [](Tensor topk_weights, Tensor topk_ids, Tensor token_expert_indices, Tensor gating_output, bool renormalize, py::object correction_bias) { + Tensor bias; + if (!correction_bias.is_none()) { + bias = correction_bias.cast(); + } + op::moe_topk_softmax_vllm_(topk_weights, topk_ids, token_expert_indices, gating_output, renormalize, bias); + }, + py::arg("topk_weights"), py::arg("topk_ids"), py::arg("token_expert_indices"), py::arg("gating_output"), py::arg("renormalize") = false, py::arg("correction_bias") = py::none(), R"doc(vLLM-style MoE topk softmax with source_rows output.)doc"); + m.def( + "moe_topk_sigmoid_vllm_", [](Tensor topk_weights, Tensor topk_ids, Tensor token_expert_indices, Tensor gating_output, bool renormalize, py::object correction_bias) { + Tensor bias; + if (!correction_bias.is_none()) { + bias = correction_bias.cast(); + } + op::moe_topk_sigmoid_vllm_(topk_weights, topk_ids, token_expert_indices, gating_output, renormalize, bias); + }, + py::arg("topk_weights"), py::arg("topk_ids"), py::arg("token_expert_indices"), py::arg("gating_output"), py::arg("renormalize") = false, py::arg("correction_bias") = py::none(), R"doc(vLLM-style MoE topk sigmoid with source_rows output.)doc"); +} +} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/scaled_mm_w4a8.hpp b/src/infinicore/pybind11/ops/scaled_mm_w4a8.hpp new file mode 100644 index 000000000..d1c7e2076 --- /dev/null +++ b/src/infinicore/pybind11/ops/scaled_mm_w4a8.hpp @@ -0,0 +1,10 @@ +#pragma once +#include "infinicore/ops/scaled_mm_w4a8.hpp" +#include +namespace py = pybind11; +namespace infinicore { +inline void bind_scaled_mm_w4a8(py::module &m) { + m.def("scaled_mm_w4a8", &op::scaled_mm_w4a8, py::arg("a"), py::arg("b"), py::arg("a_scales"), py::arg("b_scales"), py::arg("bias") = std::nullopt, py::arg("trans_weight") = false); + m.def("scaled_mm_w4a8_", &op::scaled_mm_w4a8_, py::arg("out"), py::arg("a"), py::arg("b"), py::arg("a_scales"), py::arg("b_scales"), py::arg("bias") = std::nullopt, py::arg("trans_weight") = false); +} +} // namespace infinicore diff --git a/src/infinicore/pybind11/ops/w4a8_group_gemm.hpp b/src/infinicore/pybind11/ops/w4a8_group_gemm.hpp new file mode 100644 index 000000000..bceb7b988 --- /dev/null +++ b/src/infinicore/pybind11/ops/w4a8_group_gemm.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "infinicore/ops/w4a8_group_gemm.hpp" +#include +namespace py = pybind11; +namespace infinicore { +inline void bind_w4a8_group_gemm(py::module &m) { + m.def("w4a8_group_gemm_", &op::w4a8_group_gemm_, py::arg("out"), py::arg("input"), py::arg("weight"), py::arg("input_scale"), py::arg("weight_scale"), py::arg("tokens_per_experts"), py::arg("sorted_token_ids") = std::nullopt, py::arg("bias") = std::nullopt, py::arg("trans_weight") = true, py::arg("is_decode") = false); +} +} // namespace infinicore diff --git a/src/infinicore/pybind11/ops/w8a8_group_gemm.hpp b/src/infinicore/pybind11/ops/w8a8_group_gemm.hpp new file mode 100644 index 000000000..5d3546a36 --- /dev/null +++ b/src/infinicore/pybind11/ops/w8a8_group_gemm.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "infinicore/ops/w8a8_group_gemm.hpp" +#include +namespace py = pybind11; +namespace infinicore { +inline void bind_w8a8_group_gemm(py::module &m) { + m.def("w8a8_group_gemm_", &op::w8a8_group_gemm_, py::arg("out"), py::arg("input"), py::arg("weight"), py::arg("input_scale"), py::arg("weight_scale"), py::arg("tokens_per_experts"), py::arg("sorted_token_ids") = std::nullopt, py::arg("bias") = std::nullopt, py::arg("trans_weight") = true, py::arg("is_decode") = false); +} +} // namespace infinicore diff --git a/src/infinicore/tensor/view.cc b/src/infinicore/tensor/view.cc index 051ee42c0..39db330c2 100644 --- a/src/infinicore/tensor/view.cc +++ b/src/infinicore/tensor/view.cc @@ -5,9 +5,37 @@ #include "../utils.hpp" #include +#include #include namespace infinicore { +namespace { +std::string shape_to_string(const infinicore::Shape &shape) { + std::ostringstream oss; + oss << "["; + for (size_t i = 0; i < shape.size(); ++i) { + if (i != 0) { + oss << ","; + } + oss << shape[i]; + } + oss << "]"; + return oss.str(); +} + +std::string strides_to_string(const infinicore::Strides &strides) { + std::ostringstream oss; + oss << "["; + for (size_t i = 0; i < strides.size(); ++i) { + if (i != 0) { + oss << ","; + } + oss << strides[i]; + } + oss << "]"; + return oss.str(); +} +} // namespace Tensor TensorImpl::squeeze(size_t dim) const { // Create new shape with dimension of size one removed at dim if (meta_.shape[dim] != 1) { @@ -133,14 +161,14 @@ Tensor TensorImpl::view(const Shape &new_shape) const { // Find which merged dimension contains this new dimension while (new_shape[i] > remaining_size) { if (++merged_idx >= merged_shape.size()) { - throw std::runtime_error("Incompatible shape for view operation."); + throw std::runtime_error("Incompatible shape for view operation: old_shape=" + shape_to_string(old_shape) + " old_strides=" + strides_to_string(old_strides) + " new_shape=" + shape_to_string(new_shape)); } current_stride = merged_strides[merged_idx]; remaining_size = merged_shape[merged_idx]; } if (remaining_size % new_shape[i] != 0) { - throw std::runtime_error("Incompatible shape for view operation."); + throw std::runtime_error("Incompatible shape for view operation: old_shape=" + shape_to_string(old_shape) + " old_strides=" + strides_to_string(old_strides) + " new_shape=" + shape_to_string(new_shape)); }; new_strides[i] = current_stride * (remaining_size / new_shape[i]); diff --git a/src/infiniop/ops/moe_sum/nvidia/moe_sum_nvidia.cu b/src/infiniop/ops/moe_sum/nvidia/moe_sum_nvidia.cu index 5d0859d20..2e24cb358 100644 --- a/src/infiniop/ops/moe_sum/nvidia/moe_sum_nvidia.cu +++ b/src/infiniop/ops/moe_sum/nvidia/moe_sum_nvidia.cu @@ -1,4 +1,4 @@ -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) #include "moe_sum_nvidia.cuh" @@ -117,4 +117,4 @@ infiniStatus_t Descriptor::calculate( } // namespace op::moe_sum::nvidia -#endif // ENABLE_NVIDIA_API +#endif // ENABLE_NVIDIA_API || ENABLE_ILUVATAR_API diff --git a/src/infiniop/ops/moe_sum/operator.cc b/src/infiniop/ops/moe_sum/operator.cc index 04e75434a..06764e4d3 100644 --- a/src/infiniop/ops/moe_sum/operator.cc +++ b/src/infiniop/ops/moe_sum/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/moe_sum.h" -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) #include "nvidia/moe_sum_nvidia.cuh" #endif @@ -23,6 +23,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeSumDescriptor( switch (handle->device) { #ifdef ENABLE_NVIDIA_API CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -43,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeSumWorkspaceSize( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -67,6 +73,9 @@ __INFINI_C infiniStatus_t infiniopMoeSum( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -86,6 +95,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeSumDescriptor( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu b/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu index 5d6acf37a..a2bbf298c 100644 --- a/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu +++ b/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu @@ -7,7 +7,7 @@ * Licensed under the Apache License, Version 2.0. */ -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) #include "moe_topk_sigmoid_nvidia.cuh" @@ -464,4 +464,4 @@ infiniStatus_t Descriptor::calculate( } // namespace op::moe_topk_sigmoid::nvidia -#endif // ENABLE_NVIDIA_API +#endif // ENABLE_NVIDIA_API || ENABLE_ILUVATAR_API diff --git a/src/infiniop/ops/moe_topk_sigmoid/operator.cc b/src/infiniop/ops/moe_topk_sigmoid/operator.cc index 5a7810d9f..97f255c70 100644 --- a/src/infiniop/ops/moe_topk_sigmoid/operator.cc +++ b/src/infiniop/ops/moe_topk_sigmoid/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/moe_topk_sigmoid.h" -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) #include "nvidia/moe_topk_sigmoid_nvidia.cuh" #endif @@ -22,6 +22,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeTopkSigmoidDescriptor( switch (handle->device) { #ifdef ENABLE_NVIDIA_API CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -39,6 +42,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeTopkSigmoidWorkspaceSize( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -62,6 +68,9 @@ __INFINI_C infiniStatus_t infiniopMoeTopkSigmoid( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -78,6 +87,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeTopkSigmoidDescriptor( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu b/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu index 48842e144..952b3ca5d 100644 --- a/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu +++ b/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu @@ -7,7 +7,7 @@ * Licensed under the Apache License, Version 2.0. */ -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) #include "moe_topk_softmax_nvidia.cuh" @@ -618,4 +618,4 @@ infiniStatus_t Descriptor::calculate( } // namespace op::moe_topk_softmax::nvidia -#endif // ENABLE_NVIDIA_API +#endif // ENABLE_NVIDIA_API || ENABLE_ILUVATAR_API diff --git a/src/infiniop/ops/moe_topk_softmax/operator.cc b/src/infiniop/ops/moe_topk_softmax/operator.cc index 60fcaa9dd..213ad93ce 100644 --- a/src/infiniop/ops/moe_topk_softmax/operator.cc +++ b/src/infiniop/ops/moe_topk_softmax/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/moe_topk_softmax.h" -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) #include "nvidia/moe_topk_softmax_nvidia.cuh" #endif @@ -24,6 +24,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeTopkSoftmaxDescriptor( switch (handle->device) { #ifdef ENABLE_NVIDIA_API CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -41,6 +44,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeTopkSoftmaxWorkspaceSize( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -64,6 +70,9 @@ __INFINI_C infiniStatus_t infiniopMoeTopkSoftmax( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -80,6 +89,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeTopkSoftmaxDescriptor( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; 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 diff --git a/xmake.lua b/xmake.lua index fcbcf423d..c21aa73a8 100644 --- a/xmake.lua +++ b/xmake.lua @@ -6,6 +6,11 @@ add_requires("pybind11") local GREEN = '\27[0;32m' local YELLOW = '\27[1;33m' local NC = '\27[0m' -- No Color +local PYTHON = os.getenv("PYTHON") + +if not PYTHON or PYTHON == "" then + PYTHON = is_host("windows") and "python" or "python3" +end set_encodings("utf-8") @@ -673,6 +678,12 @@ target("infinicore_cpp_api") local cuda_root = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or get_config("cuda") or "/usr/local/cuda" add_includedirs(cuda_root .. "/include") end + if has_config("iluvatar-gpu") then + local corex_root = get_config("cuda") or os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or "/usr/local/corex" + add_includedirs(corex_root .. "/include") + add_linkdirs(corex_root .. "/lib64") + add_rpathdirs(corex_root .. "/lib64") + end if has_config("infiniops") then local infiniops_root = path.absolute(get_config("infiniops-root") or "submodules/InfiniOps", os.projectdir()) if not os.isdir(infiniops_root) then @@ -770,7 +781,7 @@ target("infinicore_cpp_api") end if has_config("aten") then - local outdata = os.iorunv("python", {"-c", "import torch, os; print(os.path.dirname(torch.__file__))"}):trim() + local outdata = os.iorunv(PYTHON, {"-c", "import torch, os; print(os.path.dirname(torch.__file__))"}):trim() local TORCH_DIR = outdata target:add( @@ -785,6 +796,12 @@ target("infinicore_cpp_api") { public = true } ) + target:add( + "rpathdirs", + path.join(TORCH_DIR, "lib"), + { public = true } + ) + -- Moore mate: link torch_musa instead of torch_cuda/c10_cuda if has_config("moore-gpu") then target:add( @@ -797,7 +814,7 @@ target("infinicore_cpp_api") ) -- Detect torch_musa install path - local musa_outdata = os.iorunv("python", {"-c", "import torch_musa, os; print(os.path.dirname(torch_musa.__file__))"}):trim() + local musa_outdata = os.iorunv(PYTHON, {"-c", "import torch_musa, os; print(os.path.dirname(torch_musa.__file__))"}):trim() local TORCH_MUSA_DIR = musa_outdata local MUSA_ROOT = os.getenv("MUSA_ROOT") or os.getenv("MUSA_HOME") or os.getenv("MUSA_PATH") or "/usr/local/musa" @@ -824,11 +841,11 @@ target("infinicore_cpp_api") ) -- libpython for pybind11::scoped_interpreter / embed - local pyinc = os.iorunv("python", {"-c", + local pyinc = os.iorunv(PYTHON, {"-c", "import sysconfig; print(sysconfig.get_path('include'))"}):trim() - local pylib = os.iorunv("python", {"-c", + local pylib = os.iorunv(PYTHON, {"-c", "import sysconfig; print(sysconfig.get_config_var('LIBDIR'))"}):trim() - local pyver = os.iorunv("python", {"-c", + local pyver = os.iorunv(PYTHON, {"-c", "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')"}):trim() target:add("includedirs", pyinc, { public = true }) target:add("linkdirs", pylib, { public = true }) @@ -870,7 +887,7 @@ target("infinicore_cpp_api") -- Moore mate: force link torch_python to bypass --as-needed if has_config("moore-gpu") and has_config("aten") and has_config("flash-attn") then before_link(function (target) - local torch_dir = os.iorunv("python", {"-c", + local torch_dir = os.iorunv(PYTHON, {"-c", "import torch, os; print(os.path.dirname(torch.__file__))"}):trim() local torch_lib = path.join(torch_dir, "lib") target:add("shflags", @@ -1017,7 +1034,7 @@ target("infinicore") table.insert(pip_install_args, "--editable") end - os.execv("python", table.join({"-m", "pip", "install"}, pip_install_args, {"."})) + os.execv(PYTHON, table.join({"-m", "pip", "install"}, pip_install_args, {"."})) end) target_end() diff --git a/xmake/iluvatar.lua b/xmake/iluvatar.lua index 046e6716e..c0e2a1bc9 100644 --- a/xmake/iluvatar.lua +++ b/xmake/iluvatar.lua @@ -25,7 +25,7 @@ toolchain("iluvatar.toolchain") toolchain_end() rule("iluvatar.env") - add_deps("cuda.env", {order = true}) + add_orders("cuda.env", "iluvatar.env") after_load(function (target) local old = target:get("syslinks") local new = {} From d6d4fb1f32a9c0393109805f940d8e1c57752551 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Wed, 15 Jul 2026 11:36:24 +0000 Subject: [PATCH 2/4] pepe: deepseek v2 additional ops --- .../adaptor/vllm_iluvatar_adaptor.hpp | 12 ++ include/infinicore/ops.hpp | 2 + .../infinicore/ops/paged_attention_mla.hpp | 17 ++ include/infinicore/ops/w16a16_group_gemm.hpp | 15 ++ python/infinicore/__init__.py | 9 + python/infinicore/ops/w16a16_group_gemm.py | 28 +++ .../adaptor/vllm_iluvatar_adaptor.cc | 194 +++++++++++++++--- .../paged_attention_mla.cc | 88 ++++++++ .../w16a16_group_gemm/w16a16_group_gemm.cc | 104 ++++++++++ src/infinicore/pybind11/ops.hpp | 2 + .../pybind11/ops/w16a16_group_gemm.hpp | 21 ++ 11 files changed, 459 insertions(+), 33 deletions(-) create mode 100644 include/infinicore/ops/paged_attention_mla.hpp create mode 100644 include/infinicore/ops/w16a16_group_gemm.hpp create mode 100644 python/infinicore/ops/w16a16_group_gemm.py create mode 100644 src/infinicore/ops/paged_attention_mla/paged_attention_mla.cc create mode 100644 src/infinicore/ops/w16a16_group_gemm/w16a16_group_gemm.cc create mode 100644 src/infinicore/pybind11/ops/w16a16_group_gemm.hpp diff --git a/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp b/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp index 8f6a47c09..b84dd88a7 100644 --- a/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp +++ b/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp @@ -12,6 +12,7 @@ bool dynamic_scaled_int8_quant_available(); bool concat_mla_q_available(); bool concat_and_cache_mla_available(); bool concat_and_cache_mla_int8_available(); +bool paged_attention_mla_available(); bool topk_softmax_available(); bool topk_sigmoid_available(); bool grouped_topk_available(); @@ -19,6 +20,7 @@ bool scaled_mm_w4a8_available(); bool scaled_mm_w8a8_available(); bool w4a8_group_gemm_available(); bool w8a8_group_gemm_available(); +bool w16a16_group_gemm_available(); bool argsort_bincount_with_inv_pos_available(); bool expand_moe_input_with_inv_pos_available(); bool silu_and_mul_quant_available(); @@ -28,6 +30,15 @@ void dynamic_scaled_int8_quant(at::Tensor &output, at::Tensor &input_scales, con void concat_mla_q(at::Tensor &ql_nope, at::Tensor &q_pe, at::Tensor &q_out); void concat_and_cache_mla(at::Tensor &kv_c, at::Tensor &k_pe, at::Tensor &kv_cache, at::Tensor &slot_mapping, const std::string &kv_cache_dtype, at::Tensor &scale); void concat_and_cache_mla_int8(at::Tensor &kv_c_int8, at::Tensor &kv_c_scale, at::Tensor &k_pe_int8, at::Tensor &k_pe_scale, at::Tensor &kv_cache, at::Tensor &kv_cache_scale, at::Tensor &slot_mapping); +void paged_attention_mla(at::Tensor &output, + at::Tensor &query, + at::Tensor &kv_cache, + double scale, + at::Tensor &block_tables, + at::Tensor &context_lens, + int64_t max_context_len, + bool use_cuda_graph, + at::Tensor &softmax_lse); void topk_softmax(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional correction_bias); void topk_sigmoid(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional correction_bias); void grouped_topk(at::Tensor &topk_weights, at::Tensor &topk_ids, const at::Tensor &scores, std::optional bias, int64_t num_expert_group, int64_t topk_group, const std::string &scoring_func, bool renormalize); @@ -35,6 +46,7 @@ void scaled_mm_w4a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, c void scaled_mm_w8a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional bias, bool trans_weight); void w4a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode); void w8a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode); +void w16a16_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode); void argsort_bincount_with_inv_pos(const at::Tensor &topk_ids, at::Tensor &tokens_per_experts, at::Tensor &sorted_indices, at::Tensor &inv_pos, int64_t num_experts); void expand_moe_input_with_inv_pos(at::Tensor &expand_states, std::optional expand_scales, const at::Tensor &hidden_states, const at::Tensor &inv_pos, int64_t top_k, int64_t group_size, int64_t format); void silu_and_mul_quant(at::Tensor &output, std::optional output_scale, const at::Tensor &input, int64_t format); diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index b5c4ff18f..7058b9869 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -48,6 +48,7 @@ #include "ops/nrm2.hpp" #include "ops/ones.hpp" #include "ops/paged_attention.hpp" +#include "ops/paged_attention_mla.hpp" #include "ops/paged_attention_prefill.hpp" #include "ops/paged_caching.hpp" #include "ops/per_tensor_dequant_i8.hpp" @@ -74,3 +75,4 @@ #include "ops/swap.hpp" #include "ops/swiglu.hpp" #include "ops/topksoftmax.hpp" +#include "ops/w16a16_group_gemm.hpp" diff --git a/include/infinicore/ops/paged_attention_mla.hpp b/include/infinicore/ops/paged_attention_mla.hpp new file mode 100644 index 000000000..99034003a --- /dev/null +++ b/include/infinicore/ops/paged_attention_mla.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../tensor.hpp" + +#include + +namespace infinicore::op { + +void paged_attention_mla_(Tensor output, + const Tensor &query, + const Tensor &kv_cache, + float scale, + const Tensor &block_tables, + const Tensor &context_lens, + int64_t max_context_len); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/w16a16_group_gemm.hpp b/include/infinicore/ops/w16a16_group_gemm.hpp new file mode 100644 index 000000000..15b949f85 --- /dev/null +++ b/include/infinicore/ops/w16a16_group_gemm.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../tensor.hpp" +#include + +namespace infinicore::op { +void w16a16_group_gemm_(Tensor out, + const Tensor &input, + const Tensor &weight, + const Tensor &tokens_per_experts, + std::optional sorted_token_ids = std::nullopt, + std::optional bias = std::nullopt, + bool trans_weight = true, + bool is_decode = false); +} // namespace infinicore::op diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index b01222408..dadea41fe 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -152,6 +152,7 @@ from infinicore.ops.var_mean import var_mean from infinicore.ops.w4a8_group_gemm import w4a8_group_gemm_ from infinicore.ops.w8a8_group_gemm import w8a8_group_gemm_ +from infinicore.ops.w16a16_group_gemm import w16a16_group_gemm_ from infinicore.tensor import ( Tensor, empty, @@ -232,6 +233,10 @@ "floor", "attention", "mrope", + "moe_argsort_bincount_with_inv_pos_", + "moe_expand_input_with_inv_pos_", + "moe_silu_and_mul_quant_", + "moe_sum_vllm_", "moe_topk_sigmoid_vllm", "moe_topk_softmax_vllm", "block_diag", @@ -290,6 +295,7 @@ "rotm", "rotmg", "scal", + "scaled_mm_w4a8", "logcumsumexp", "logical_not", "logical_and", @@ -313,6 +319,9 @@ "var_mean", "moore_mate_flash_attn_prefill", "moore_mate_flash_attn_decode", + "w4a8_group_gemm_", + "w8a8_group_gemm_", + "w16a16_group_gemm_", "var", "topk", "all", diff --git a/python/infinicore/ops/w16a16_group_gemm.py b/python/infinicore/ops/w16a16_group_gemm.py new file mode 100644 index 000000000..e9597f9bd --- /dev/null +++ b/python/infinicore/ops/w16a16_group_gemm.py @@ -0,0 +1,28 @@ +from infinicore.lib import _infinicore + + +def w16a16_group_gemm_( + out, + input, + weight, + tokens_per_experts, + sorted_token_ids=None, + bias=None, + trans_weight=True, + is_decode=False, +): + if not trans_weight: + raise RuntimeError( + "w16a16_group_gemm currently supports only trans_weight=True (TN layout)" + ) + _infinicore.w16a16_group_gemm_( + out._underlying, + input._underlying, + weight._underlying, + tokens_per_experts._underlying, + None if sorted_token_ids is None else sorted_token_ids._underlying, + None if bias is None else bias._underlying, + trans_weight, + is_decode, + ) + return out diff --git a/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc b/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc index 34cf54a22..32a485627 100644 --- a/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc +++ b/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc @@ -1,6 +1,8 @@ #ifdef ENABLE_ATEN #include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" #include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/context/context.hpp" +#include "infinicore/graph/graph.hpp" #if defined(ENABLE_ILUVATAR_API) #define INFINICORE_VLLM_ILUVATAR_STREAM_GUARD() \ @@ -15,6 +17,8 @@ #include #include #include +#include +#include namespace infinicore::adaptor::vllm_iluvatar { namespace { @@ -24,6 +28,7 @@ using dynamic_scaled_int8_quant_fn = void (*)(at::Tensor &, at::Tensor &, const using concat_mla_q_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &); using concat_and_cache_mla_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, const std::string &, at::Tensor &); using concat_and_cache_mla_int8_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &); +using paged_attention_mla_fn = at::Tensor (*)(at::Tensor &, at::Tensor &, at::Tensor &, double, at::Tensor &, at::Tensor &, int64_t, bool, at::Tensor &); using topk_softmax_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, const at::Tensor &, bool, std::optional); using topk_sigmoid_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, const at::Tensor &, bool, std::optional); using grouped_topk_fn = void (*)(at::Tensor &, at::Tensor &, const at::Tensor &, const std::optional &, int64_t, int64_t, std::string, bool); @@ -31,9 +36,34 @@ using scaled_mm_w4a8_fn = void (*)(at::Tensor &, const at::Tensor &, const at::T using scaled_mm_w8a8_fn = scaled_mm_w4a8_fn; using w4a8_group_gemm_fn = void (*)(at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const std::optional &, const std::optional &, bool, bool); using w8a8_group_gemm_fn = void (*)(at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const std::optional &, const std::optional &, bool, bool); +using w16a16_group_gemm_fn = void (*)(at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const std::optional &, const std::optional &, bool, bool); using argsort_bincount_with_inv_pos_fn = void (*)(const at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, int64_t); using expand_moe_input_with_inv_pos_fn = void (*)(at::Tensor &, std::optional, const at::Tensor &, const at::Tensor &, int64_t, int64_t, int64_t); using silu_and_mul_quant_fn = void (*)(at::Tensor &, std::optional, const at::Tensor &, int64_t); + +template +class VllmIluvatarGraphOperator final : public graph::GraphOperator { +public: + explicit VllmIluvatarGraphOperator(Fn fn) : fn_(std::move(fn)) {} + + void run() const override { + fn_(); + } + +private: + mutable Fn fn_; +}; + +template +void record_or_run(Fn &&fn) { + using Op = VllmIluvatarGraphOperator>; + auto op = std::make_shared(std::forward(fn)); + if (context::isGraphRecording()) { + context::addGraphOperator(op); + } else { + op->run(); + } +} using moe_sum_vllm_fn = void (*)(at::Tensor &, const at::Tensor &, std::optional, std::optional, double, double); struct Symbols { @@ -43,6 +73,7 @@ struct Symbols { concat_mla_q_fn concat_mla_q = nullptr; concat_and_cache_mla_fn concat_and_cache_mla = nullptr; concat_and_cache_mla_int8_fn concat_and_cache_mla_int8 = nullptr; + paged_attention_mla_fn paged_attention_mla = nullptr; topk_softmax_fn topk_softmax = nullptr; topk_sigmoid_fn topk_sigmoid = nullptr; grouped_topk_fn grouped_topk = nullptr; @@ -50,6 +81,7 @@ struct Symbols { scaled_mm_w8a8_fn scaled_mm_w8a8 = nullptr; w4a8_group_gemm_fn w4a8_group_gemm = nullptr; w8a8_group_gemm_fn w8a8_group_gemm = nullptr; + w16a16_group_gemm_fn w16a16_group_gemm = nullptr; argsort_bincount_with_inv_pos_fn argsort_bincount_with_inv_pos = nullptr; expand_moe_input_with_inv_pos_fn expand_moe_input_with_inv_pos = nullptr; silu_and_mul_quant_fn silu_and_mul_quant = nullptr; @@ -88,6 +120,8 @@ Symbols &symbols() { dlsym(handle, "_ZN7pyinfer4perf20concat_and_cache_mlaERN2at6TensorES3_S3_S3_RKSsS3_")); auto concat_and_cache_mla_int8_fn_ptr = reinterpret_cast( dlsym(handle, "_ZN7pyinfer4perf25concat_and_cache_mla_int8ERN2at6TensorES3_S3_S3_S3_S3_S3_")); + auto paged_attention_mla_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer7cuinfer24vllm_paged_attention_mlaERN2at6TensorES3_S3_dS3_S3_lbS3_")); auto topk_softmax_fn_ptr = reinterpret_cast( dlsym(handle, "_ZN7pyinfer4perf12topk_softmaxERN2at6TensorES3_S3_RKS2_bSt8optionalIS2_E")); auto topk_sigmoid_fn_ptr = reinterpret_cast( @@ -102,6 +136,8 @@ Symbols &symbols() { dlsym(handle, "_ZN7pyinfer7cuinfer15w4a8_group_gemmERN2at6TensorERKS2_S5_S5_S5_S5_RKSt8optionalIS2_ES9_bb")); auto w8a8_group_gemm_fn_ptr = reinterpret_cast( dlsym(handle, "_ZN7pyinfer7cuinfer15w8a8_group_gemmERN2at6TensorERKS2_S5_S5_S5_S5_RKSt8optionalIS2_ES9_bb")); + auto w16a16_group_gemm_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer7cuinfer17w16a16_group_gemmERN2at6TensorERKS2_S5_S5_RKSt8optionalIS2_ES9_bb")); auto argsort_bincount_with_inv_pos_fn_ptr = reinterpret_cast( dlsym(handle, "_ZN7pyinfer4perf29argsort_bincount_with_inv_posERKN2at6TensorERS2_S5_S5_l")); auto expand_moe_input_with_inv_pos_fn_ptr = reinterpret_cast( @@ -110,7 +146,7 @@ Symbols &symbols() { dlsym(handle, "_ZN7pyinfer4perf18silu_and_mul_quantERN2at6TensorESt8optionalIS2_ERKS2_l")); auto moe_sum_vllm_fn_ptr = reinterpret_cast( dlsym(handle, "_ZN7pyinfer4perf7moe_sumERN2at6TensorERKS2_St8optionalIS2_ES7_dd")); - if (!fused_fn && !quant_fn && !concat_mla_q_fn_ptr && !concat_and_cache_mla_fn_ptr && !concat_and_cache_mla_int8_fn_ptr && !topk_softmax_fn_ptr && !topk_sigmoid_fn_ptr && !grouped_topk_fn_ptr && !scaled_mm_w4a8_fn_ptr && !scaled_mm_w8a8_fn_ptr && !w4a8_group_gemm_fn_ptr && !w8a8_group_gemm_fn_ptr && !argsort_bincount_with_inv_pos_fn_ptr && !expand_moe_input_with_inv_pos_fn_ptr && !silu_and_mul_quant_fn_ptr && !moe_sum_vllm_fn_ptr) { + if (!fused_fn && !quant_fn && !concat_mla_q_fn_ptr && !concat_and_cache_mla_fn_ptr && !concat_and_cache_mla_int8_fn_ptr && !paged_attention_mla_fn_ptr && !topk_softmax_fn_ptr && !topk_sigmoid_fn_ptr && !grouped_topk_fn_ptr && !scaled_mm_w4a8_fn_ptr && !scaled_mm_w8a8_fn_ptr && !w4a8_group_gemm_fn_ptr && !w8a8_group_gemm_fn_ptr && !w16a16_group_gemm_fn_ptr && !argsort_bincount_with_inv_pos_fn_ptr && !expand_moe_input_with_inv_pos_fn_ptr && !silu_and_mul_quant_fn_ptr && !moe_sum_vllm_fn_ptr) { syms.error = dlerror(); dlclose(handle); continue; @@ -121,6 +157,7 @@ Symbols &symbols() { syms.concat_mla_q = concat_mla_q_fn_ptr; syms.concat_and_cache_mla = concat_and_cache_mla_fn_ptr; syms.concat_and_cache_mla_int8 = concat_and_cache_mla_int8_fn_ptr; + syms.paged_attention_mla = paged_attention_mla_fn_ptr; syms.topk_softmax = topk_softmax_fn_ptr; syms.topk_sigmoid = topk_sigmoid_fn_ptr; syms.grouped_topk = grouped_topk_fn_ptr; @@ -128,6 +165,7 @@ Symbols &symbols() { syms.scaled_mm_w8a8 = scaled_mm_w8a8_fn_ptr; syms.w4a8_group_gemm = w4a8_group_gemm_fn_ptr; syms.w8a8_group_gemm = w8a8_group_gemm_fn_ptr; + syms.w16a16_group_gemm = w16a16_group_gemm_fn_ptr; syms.argsort_bincount_with_inv_pos = argsort_bincount_with_inv_pos_fn_ptr; syms.expand_moe_input_with_inv_pos = expand_moe_input_with_inv_pos_fn_ptr; syms.silu_and_mul_quant = silu_and_mul_quant_fn_ptr; @@ -167,6 +205,10 @@ bool concat_and_cache_mla_int8_available() { return symbols().concat_and_cache_mla_int8 != nullptr; } +bool paged_attention_mla_available() { + return symbols().paged_attention_mla != nullptr; +} + bool topk_softmax_available() { return symbols().topk_softmax != nullptr; } @@ -195,6 +237,10 @@ bool w8a8_group_gemm_available() { return symbols().w8a8_group_gemm != nullptr; } +bool w16a16_group_gemm_available() { + return symbols().w16a16_group_gemm != nullptr; +} + bool argsort_bincount_with_inv_pos_available() { return symbols().argsort_bincount_with_inv_pos != nullptr; } @@ -212,147 +258,229 @@ bool moe_sum_vllm_available() { } void fused_add_rms_norm(at::Tensor &input, at::Tensor &residual, at::Tensor &weight, float epsilon) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.fused_add_rms_norm) { throw std::runtime_error("vllm_iluvatar fused_add_rms_norm unavailable: " + syms.error); } - syms.fused_add_rms_norm(input, residual, weight, epsilon); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.fused_add_rms_norm(input, residual, weight, epsilon); + }); } void dynamic_scaled_int8_quant(at::Tensor &output, at::Tensor &input_scales, const at::Tensor &input) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.dynamic_scaled_int8_quant) { throw std::runtime_error("vllm_iluvatar dynamic_scaled_int8_quant unavailable: " + syms.error); } - syms.dynamic_scaled_int8_quant(output, input_scales, input); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.dynamic_scaled_int8_quant(output, input_scales, input); + }); } void concat_mla_q(at::Tensor &ql_nope, at::Tensor &q_pe, at::Tensor &q_out) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.concat_mla_q) { throw std::runtime_error("vllm_iluvatar concat_mla_q unavailable: " + syms.error); } - syms.concat_mla_q(ql_nope, q_pe, q_out); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.concat_mla_q(ql_nope, q_pe, q_out); + }); } void concat_and_cache_mla(at::Tensor &kv_c, at::Tensor &k_pe, at::Tensor &kv_cache, at::Tensor &slot_mapping, const std::string &kv_cache_dtype, at::Tensor &scale) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.concat_and_cache_mla) { throw std::runtime_error("vllm_iluvatar concat_and_cache_mla unavailable: " + syms.error); } - syms.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale); + }); } void concat_and_cache_mla_int8(at::Tensor &kv_c_int8, at::Tensor &kv_c_scale, at::Tensor &k_pe_int8, at::Tensor &k_pe_scale, at::Tensor &kv_cache, at::Tensor &kv_cache_scale, at::Tensor &slot_mapping) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.concat_and_cache_mla_int8) { throw std::runtime_error("vllm_iluvatar concat_and_cache_mla_int8 unavailable: " + syms.error); } - syms.concat_and_cache_mla_int8(kv_c_int8, kv_c_scale, k_pe_int8, k_pe_scale, kv_cache, kv_cache_scale, slot_mapping); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.concat_and_cache_mla_int8(kv_c_int8, kv_c_scale, k_pe_int8, k_pe_scale, kv_cache, kv_cache_scale, slot_mapping); + }); +} + +void paged_attention_mla(at::Tensor &output, + at::Tensor &query, + at::Tensor &kv_cache, + double scale, + at::Tensor &block_tables, + at::Tensor &context_lens, + int64_t max_context_len, + bool use_cuda_graph, + at::Tensor &softmax_lse) { + auto &syms = symbols(); + if (!syms.paged_attention_mla) { + throw std::runtime_error("vllm_iluvatar paged_attention_mla unavailable: " + syms.error); + } + const bool graph_recording = context::isGraphRecording(); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.paged_attention_mla( + output, query, kv_cache, scale, block_tables, context_lens, max_context_len, (use_cuda_graph || graph_recording), softmax_lse); + }); } void topk_softmax(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional correction_bias) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.topk_softmax) { throw std::runtime_error("vllm_iluvatar topk_softmax unavailable: " + syms.error); } - syms.topk_softmax(topk_weights, topk_ids, token_expert_indices, gating_output, renormalize, correction_bias); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.topk_softmax(topk_weights, topk_ids, token_expert_indices, gating_output, renormalize, correction_bias); + }); } void topk_sigmoid(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional correction_bias) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.topk_sigmoid) { throw std::runtime_error("vllm_iluvatar topk_sigmoid unavailable: " + syms.error); } - syms.topk_sigmoid(topk_weights, topk_ids, token_expert_indices, gating_output, renormalize, correction_bias); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.topk_sigmoid(topk_weights, topk_ids, token_expert_indices, gating_output, renormalize, correction_bias); + }); } void grouped_topk(at::Tensor &topk_weights, at::Tensor &topk_ids, const at::Tensor &scores, std::optional bias, int64_t num_expert_group, int64_t topk_group, const std::string &scoring_func, bool renormalize) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.grouped_topk) { throw std::runtime_error("vllm_iluvatar grouped_topk unavailable: " + syms.error); } - syms.grouped_topk(topk_weights, topk_ids, scores, bias, num_expert_group, topk_group, scoring_func, renormalize); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.grouped_topk(topk_weights, topk_ids, scores, bias, num_expert_group, topk_group, scoring_func, renormalize); + }); } void scaled_mm_w4a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional bias, bool trans_weight) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.scaled_mm_w4a8) { throw std::runtime_error("vllm_iluvatar scaled_mm_w4a8 unavailable: " + syms.error); } - syms.scaled_mm_w4a8(out, a, b, a_scales, b_scales, bias, trans_weight); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.scaled_mm_w4a8(out, a, b, a_scales, b_scales, bias, trans_weight); + }); } void scaled_mm_w8a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional bias, bool trans_weight) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.scaled_mm_w8a8) { throw std::runtime_error("vllm_iluvatar scaled_mm W8A8 unavailable: " + syms.error); } - syms.scaled_mm_w8a8(out, a, b, a_scales, b_scales, bias, trans_weight); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.scaled_mm_w8a8(out, a, b, a_scales, b_scales, bias, trans_weight); + }); } void w4a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.w4a8_group_gemm) { throw std::runtime_error("vllm_iluvatar w4a8_group_gemm unavailable: " + syms.error); } - syms.w4a8_group_gemm(out, input, weight, input_scale, weight_scale, tokens_per_experts, sorted_token_ids, bias, trans_weight, is_decode); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.w4a8_group_gemm(out, input, weight, input_scale, weight_scale, tokens_per_experts, sorted_token_ids, bias, trans_weight, is_decode); + }); } void w8a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.w8a8_group_gemm) { throw std::runtime_error("vllm_iluvatar w8a8_group_gemm unavailable: " + syms.error); } - syms.w8a8_group_gemm(out, input, weight, input_scale, weight_scale, tokens_per_experts, sorted_token_ids, bias, trans_weight, is_decode); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.w8a8_group_gemm(out, input, weight, input_scale, weight_scale, tokens_per_experts, sorted_token_ids, bias, trans_weight, is_decode); + }); +} + +void w16a16_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode) { + auto &syms = symbols(); + if (!syms.w16a16_group_gemm) { + throw std::runtime_error("vllm_iluvatar w16a16_group_gemm unavailable: " + syms.error); + } + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.w16a16_group_gemm(out, input, weight, tokens_per_experts, sorted_token_ids, bias, trans_weight, is_decode); + }); } void argsort_bincount_with_inv_pos(const at::Tensor &topk_ids, at::Tensor &tokens_per_experts, at::Tensor &sorted_indices, at::Tensor &inv_pos, int64_t num_experts) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.argsort_bincount_with_inv_pos) { throw std::runtime_error("vllm_iluvatar argsort_bincount_with_inv_pos unavailable: " + syms.error); } - syms.argsort_bincount_with_inv_pos(topk_ids, tokens_per_experts, sorted_indices, inv_pos, num_experts); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.argsort_bincount_with_inv_pos(topk_ids, tokens_per_experts, sorted_indices, inv_pos, num_experts); + }); } void expand_moe_input_with_inv_pos(at::Tensor &expand_states, std::optional expand_scales, const at::Tensor &hidden_states, const at::Tensor &inv_pos, int64_t top_k, int64_t group_size, int64_t format) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.expand_moe_input_with_inv_pos) { throw std::runtime_error("vllm_iluvatar expand_moe_input_with_inv_pos unavailable: " + syms.error); } - syms.expand_moe_input_with_inv_pos(expand_states, expand_scales, hidden_states, inv_pos, top_k, group_size, format); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.expand_moe_input_with_inv_pos(expand_states, expand_scales, hidden_states, inv_pos, top_k, group_size, format); + }); } void silu_and_mul_quant(at::Tensor &output, std::optional output_scale, const at::Tensor &input, int64_t format) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.silu_and_mul_quant) { throw std::runtime_error("vllm_iluvatar silu_and_mul_quant unavailable: " + syms.error); } - syms.silu_and_mul_quant(output, output_scale, input, format); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.silu_and_mul_quant(output, output_scale, input, format); + }); } void moe_sum_vllm(at::Tensor &output, const at::Tensor &input, std::optional topk_weights, std::optional extra_residual, double routed_scale, double residual_scale) { - INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); auto &syms = symbols(); if (!syms.moe_sum_vllm) { throw std::runtime_error("vllm_iluvatar moe_sum unavailable: " + syms.error); } - syms.moe_sum_vllm(output, input, topk_weights, extra_residual, routed_scale, residual_scale); + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.moe_sum_vllm(output, input, topk_weights, extra_residual, routed_scale, residual_scale); + }); } #undef INFINICORE_VLLM_ILUVATAR_STREAM_GUARD diff --git a/src/infinicore/ops/paged_attention_mla/paged_attention_mla.cc b/src/infinicore/ops/paged_attention_mla/paged_attention_mla.cc new file mode 100644 index 000000000..12a0b3781 --- /dev/null +++ b/src/infinicore/ops/paged_attention_mla/paged_attention_mla.cc @@ -0,0 +1,88 @@ +#include "infinicore/ops/paged_attention_mla.hpp" + +#include "../../utils.hpp" + +#include + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { +namespace { + +void validate_paged_attention_mla(const Tensor &output, + const Tensor &query, + const Tensor &kv_cache, + const Tensor &block_tables, + const Tensor &context_lens, + int64_t max_context_len) { + if (!output || !query || !kv_cache || !block_tables || !context_lens) { + throw std::runtime_error("paged_attention_mla expects non-empty tensors"); + } + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, query, kv_cache, block_tables, context_lens); + if (query->ndim() != 3 || output->ndim() != 3 || kv_cache->ndim() != 3 + || block_tables->ndim() != 2 || context_lens->ndim() != 1) { + throw std::runtime_error( + "paged_attention_mla expects query/output/cache/block_tables/context_lens ranks 3/3/3/2/1"); + } + if (output->size(0) != query->size(0) || output->size(1) != query->size(1) + || kv_cache->size(2) != query->size(2) || context_lens->size(0) != query->size(0) + || block_tables->size(0) != query->size(0)) { + throw std::runtime_error("paged_attention_mla tensor shapes are inconsistent"); + } + if (output->dtype() != query->dtype() || kv_cache->dtype() != query->dtype() + || (query->dtype() != DataType::F16 && query->dtype() != DataType::BF16)) { + throw std::runtime_error("paged_attention_mla requires matching fp16/bfloat16 data tensors"); + } + if (context_lens->dtype() != DataType::I32) { + throw std::runtime_error("paged_attention_mla expects int32 context_lens"); + } + if (block_tables->dtype() != DataType::I32) { + throw std::runtime_error("paged_attention_mla expects int32 block_tables for the current Iluvatar SO"); + } + if (!output->is_contiguous() || !query->is_contiguous() || !kv_cache->is_contiguous() + || !block_tables->is_contiguous() || !context_lens->is_contiguous()) { + throw std::runtime_error("paged_attention_mla expects contiguous tensors"); + } + if (max_context_len <= 0) { + throw std::runtime_error("paged_attention_mla expects max_context_len > 0"); + } +} + +} // namespace + +void paged_attention_mla_(Tensor output, + const Tensor &query, + const Tensor &kv_cache, + float scale, + const Tensor &block_tables, + const Tensor &context_lens, + int64_t max_context_len) { + validate_paged_attention_mla(output, query, kv_cache, block_tables, context_lens, max_context_len); + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (output->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::paged_attention_mla_available()) { + throw std::runtime_error("paged_attention_mla requires vllm_iluvatar cuinfer extension"); + } + auto output_at = adaptor::to_aten_tensor(output); + auto query_at = adaptor::to_aten_tensor(query); + auto kv_cache_at = adaptor::to_aten_tensor(kv_cache); + auto block_tables_at = adaptor::to_aten_tensor(block_tables); + auto context_lens_at = adaptor::to_aten_tensor(context_lens); + auto softmax_lse = Tensor::empty( + {query->size(0), query->size(1)}, DataType::F32, query->device()); + auto softmax_lse_at = adaptor::to_aten_tensor(softmax_lse); + adaptor::vllm_iluvatar::paged_attention_mla( + output_at, query_at, kv_cache_at, scale, block_tables_at, context_lens_at, + max_context_len, false, softmax_lse_at); + return; + } +#endif + + throw std::runtime_error("paged_attention_mla currently supports only Iluvatar builds with ATen"); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/w16a16_group_gemm/w16a16_group_gemm.cc b/src/infinicore/ops/w16a16_group_gemm/w16a16_group_gemm.cc new file mode 100644 index 000000000..8de283402 --- /dev/null +++ b/src/infinicore/ops/w16a16_group_gemm/w16a16_group_gemm.cc @@ -0,0 +1,104 @@ +#include "infinicore/ops/w16a16_group_gemm.hpp" +#include "../../utils.hpp" +#include +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { +void w16a16_group_gemm_(Tensor out, + const Tensor &input, + const Tensor &weight, + const Tensor &tokens_per_experts, + std::optional sorted_token_ids, + std::optional bias, + bool trans_weight, + bool is_decode) { + if (bias && sorted_token_ids) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight, *sorted_token_ids, *bias); + } else if (bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight, *bias); + } else if (sorted_token_ids) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight, *sorted_token_ids); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight); + } + if (!trans_weight) { + throw std::runtime_error("w16a16_group_gemm currently supports only trans_weight=True (TN layout)"); + } + if (out->ndim() != 2 || input->ndim() != 2 || weight->ndim() != 3 || tokens_per_experts->ndim() != 1) { + throw std::runtime_error("w16a16_group_gemm expects out/input 2D, weight 3D, tokens_per_experts 1D"); + } + if (out->dtype() != DataType::F16 && out->dtype() != DataType::BF16) { + throw std::runtime_error("w16a16_group_gemm expects fp16/bfloat16 tensors"); + } + if (input->dtype() != out->dtype() || weight->dtype() != out->dtype()) { + throw std::runtime_error("w16a16_group_gemm expects input, weight, and out to have the same dtype"); + } + if (tokens_per_experts->dtype() != DataType::I32) { + throw std::runtime_error("w16a16_group_gemm expects int32 tokens_per_experts"); + } + if (is_decode && tokens_per_experts->device() != out->device()) { + throw std::runtime_error("w16a16_group_gemm decode expects GPU tokens_per_experts on the output device"); + } + if (!is_decode && tokens_per_experts->device().getType() != Device::Type::CPU) { + throw std::runtime_error("w16a16_group_gemm prefill expects CPU tokens_per_experts"); + } + if (sorted_token_ids && (*sorted_token_ids)->dtype() != DataType::I32) { + throw std::runtime_error("w16a16_group_gemm expects int32 sorted_token_ids"); + } + if (bias && ((*bias)->ndim() != 2 || (*bias)->dtype() != out->dtype())) { + throw std::runtime_error("w16a16_group_gemm expects bias shape (E,N) and same dtype as out"); + } + if (!out->is_contiguous() || !input->is_contiguous() || !weight->is_contiguous() || !tokens_per_experts->is_contiguous() || (sorted_token_ids && !(*sorted_token_ids)->is_contiguous()) || (bias && !(*bias)->is_contiguous())) { + throw std::runtime_error("w16a16_group_gemm expects contiguous tensors"); + } + + const size_t experts = weight->size(0); + const size_t rows = input->size(0); + const size_t output_features = weight->size(1); + if (tokens_per_experts->numel() != experts) { + throw std::runtime_error("w16a16_group_gemm expert count mismatch"); + } + if (input->size(1) != weight->size(2) || out->size(0) != rows || out->size(1) != output_features) { + throw std::runtime_error("w16a16_group_gemm TN shape mismatch"); + } + if (sorted_token_ids && (*sorted_token_ids)->numel() != rows) { + throw std::runtime_error("w16a16_group_gemm sorted_token_ids length mismatch"); + } + if (bias && ((*bias)->size(0) != experts || (*bias)->size(1) != output_features)) { + throw std::runtime_error("w16a16_group_gemm bias shape mismatch"); + } + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (out->device().getType() == Device::Type::ILUVATAR) { + if (!adaptor::vllm_iluvatar::w16a16_group_gemm_available()) { + throw std::runtime_error("w16a16_group_gemm requires vllm_iluvatar cuinfer extension"); + } + auto output_aten = adaptor::to_aten_tensor(out); + auto input_aten = adaptor::to_aten_tensor(input); + auto weight_aten = adaptor::to_aten_tensor(weight); + auto tokens_per_experts_aten = adaptor::to_aten_tensor(tokens_per_experts); + std::optional sorted_token_ids_aten; + if (sorted_token_ids) { + sorted_token_ids_aten = adaptor::to_aten_tensor(*sorted_token_ids); + } + std::optional bias_aten; + if (bias) { + bias_aten = adaptor::to_aten_tensor(*bias); + } + adaptor::vllm_iluvatar::w16a16_group_gemm(output_aten, + input_aten, + weight_aten, + tokens_per_experts_aten, + sorted_token_ids_aten, + bias_aten, + trans_weight, + is_decode); + return; + } +#endif + throw std::runtime_error("w16a16_group_gemm currently supports only Iluvatar builds with ATen and vllm_iluvatar"); +} +} // namespace infinicore::op diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp index 0e9cb7f6d..129ed528d 100644 --- a/src/infinicore/pybind11/ops.hpp +++ b/src/infinicore/pybind11/ops.hpp @@ -142,6 +142,7 @@ #include "ops/vander.hpp" #include "ops/var.hpp" #include "ops/var_mean.hpp" +#include "ops/w16a16_group_gemm.hpp" #include "ops/w4a8_group_gemm.hpp" #include "ops/w8a8_group_gemm.hpp" @@ -249,6 +250,7 @@ inline void bind(py::module &m) { bind_vander(m); bind_w4a8_group_gemm(m); bind_w8a8_group_gemm(m); + bind_w16a16_group_gemm(m); bind_unfold(m); bind_rope(m); bind_rot(m); diff --git a/src/infinicore/pybind11/ops/w16a16_group_gemm.hpp b/src/infinicore/pybind11/ops/w16a16_group_gemm.hpp new file mode 100644 index 000000000..3a2f858d8 --- /dev/null +++ b/src/infinicore/pybind11/ops/w16a16_group_gemm.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include "infinicore/ops/w16a16_group_gemm.hpp" +#include + +namespace py = pybind11; + +namespace infinicore { +inline void bind_w16a16_group_gemm(py::module &m) { + m.def("w16a16_group_gemm_", + &op::w16a16_group_gemm_, + py::arg("out"), + py::arg("input"), + py::arg("weight"), + py::arg("tokens_per_experts"), + py::arg("sorted_token_ids") = std::nullopt, + py::arg("bias") = std::nullopt, + py::arg("trans_weight") = true, + py::arg("is_decode") = false); +} +} // namespace infinicore From f2dabd9ca485283c943a81e869a794ee3cf29dc1 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Thu, 23 Jul 2026 08:31:11 +0000 Subject: [PATCH 3/4] pepe: additional ops --- include/infiniccl.h | 16 + include/infinicore/adaptor/aten_adaptor.hpp | 2 + .../adaptor/vllm_iluvatar_adaptor.hpp | 29 ++ include/infinicore/context/context.hpp | 1 + include/infinicore/graph/graph.hpp | 9 +- include/infinicore/nn/rope.hpp | 2 + include/infinicore/ops.hpp | 6 + include/infinicore/ops/bmm_strided.hpp | 12 + include/infinicore/ops/distributed/p2p.hpp | 80 +++ include/infinicore/ops/dsa.hpp | 87 ++++ 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 + .../infinicore/ops/fused_rotary_embedding.hpp | 16 + .../ops/select_last_token_hidden.hpp | 13 + include/infiniop.h | 5 + 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 ++ .../infiniop/ops/select_last_token_hidden.h | 25 + python/infinicore/__init__.py | 22 + python/infinicore/dtype.py | 1 + python/infinicore/include/infiniccl.h | 75 +++ python/infinicore/include/infinicore.h | 77 +++ python/infinicore/include/infinicore.hpp | 8 + .../infinicore/adaptor/aten_adaptor.hpp | 86 ++++ .../adaptor/flash_attention_adaptor.hpp | 139 +++++ .../adaptor/vllm_iluvatar_adaptor.hpp | 85 ++++ .../include/infinicore/analyzer.hpp | 11 + .../infinicore/analyzer/intent_generator.hpp | 268 ++++++++++ .../analyzer/mutual_awareness_analyzer.hpp | 118 +++++ .../include/infinicore/analyzer/op_trace.hpp | 160 ++++++ .../include/infinicore/analyzer/op_type.hpp | 238 +++++++++ .../infinicore/analyzer/op_type_registry.hpp | 49 ++ .../analyzer/optimization_intent.hpp | 187 +++++++ .../infinicore/analyzer/phase_detector.hpp | 126 +++++ .../infinicore/analyzer/resource_sensor.hpp | 181 +++++++ .../include/infinicore/common/LRUCache.hpp | 137 +++++ .../include/infinicore/common/hash.hpp | 72 +++ .../include/infinicore/context/context.hpp | 58 +++ .../infinicore/include/infinicore/device.hpp | 53 ++ .../include/infinicore/device_event.hpp | 125 +++++ .../infinicore/include/infinicore/dtype.hpp | 33 ++ .../include/infinicore/graph/graph.hpp | 133 +++++ python/infinicore/include/infinicore/io.hpp | 86 ++++ .../infinicore/include/infinicore/memory.hpp | 30 ++ python/infinicore/include/infinicore/nn.hpp | 6 + .../include/infinicore/nn/embedding.hpp | 87 ++++ .../include/infinicore/nn/layer_norm.hpp | 59 +++ .../include/infinicore/nn/linear.hpp | 128 +++++ .../include/infinicore/nn/module.hpp | 174 +++++++ .../include/infinicore/nn/parameter.hpp | 39 ++ .../include/infinicore/nn/rmsnorm.hpp | 96 ++++ .../infinicore/include/infinicore/nn/rope.hpp | 114 +++++ .../infinicore/nn/rope_scaling_configs.hpp | 131 +++++ python/infinicore/include/infinicore/ops.hpp | 84 ++++ .../include/infinicore/ops/acos.hpp | 16 + .../infinicore/ops/adaptive_avg_pool1d.hpp | 18 + .../infinicore/ops/adaptive_avg_pool3d.hpp | 16 + .../infinicore/ops/adaptive_max_pool1d.hpp | 16 + .../infinicore/include/infinicore/ops/add.hpp | 14 + .../include/infinicore/ops/add_rms_norm.hpp | 18 + .../include/infinicore/ops/addbmm.hpp | 19 + .../include/infinicore/ops/addcmul.hpp | 17 + .../include/infinicore/ops/addr.hpp | 16 + .../include/infinicore/ops/affine_grid.hpp | 17 + .../infinicore/include/infinicore/ops/all.hpp | 18 + .../include/infinicore/ops/argwhere.hpp | 14 + .../include/infinicore/ops/asin.hpp | 18 + .../include/infinicore/ops/asinh.hpp | 16 + .../include/infinicore/ops/asum.hpp | 14 + .../include/infinicore/ops/atanh.hpp | 34 ++ .../include/infinicore/ops/attention.hpp | 16 + .../include/infinicore/ops/avg_pool1d.hpp | 18 + .../infinicore/ops/awq_marlin_gemm.hpp | 14 + .../include/infinicore/ops/axpy.hpp | 13 + .../include/infinicore/ops/baddbmm.hpp | 15 + .../include/infinicore/ops/bilinear.hpp | 12 + .../ops/binary_cross_entropy_with_logits.hpp | 46 ++ .../infinicore/ops/bitwise_right_shift.hpp | 16 + .../include/infinicore/ops/blas_amax.hpp | 14 + .../include/infinicore/ops/blas_amin.hpp | 14 + .../include/infinicore/ops/blas_copy.hpp | 13 + .../include/infinicore/ops/blas_dot.hpp | 14 + .../include/infinicore/ops/block_diag.hpp | 20 + .../include/infinicore/ops/bmm_strided.hpp | 12 + .../include/infinicore/ops/broadcast_to.hpp | 14 + .../include/infinicore/ops/cast.hpp | 5 + .../infinicore/include/infinicore/ops/cat.hpp | 9 + .../include/infinicore/ops/causal_conv1d.hpp | 41 ++ .../include/infinicore/ops/causal_softmax.hpp | 14 + .../include/infinicore/ops/cdist.hpp | 32 ++ .../infinicore/ops/chunk_gated_delta_rule.hpp | 53 ++ .../include/infinicore/ops/common/cache.hpp | 83 +++ .../infinicore/ops/common/dispatcher.hpp | 80 +++ .../include/infinicore/ops/common/op.hpp | 5 + .../infinicore/ops/concat_and_cache_mla.hpp | 16 + .../ops/concat_and_cache_mla_int8.hpp | 16 + .../include/infinicore/ops/concat_mla_q.hpp | 11 + .../include/infinicore/ops/conv2d.hpp | 38 ++ .../include/infinicore/ops/cross_entropy.hpp | 35 ++ .../include/infinicore/ops/deepseek_moe.hpp | 41 ++ .../include/infinicore/ops/dequantize_awq.hpp | 10 + .../include/infinicore/ops/diff.hpp | 14 + .../include/infinicore/ops/digamma.hpp | 14 + .../include/infinicore/ops/dist.hpp | 14 + .../infinicore/ops/distributed/allgather.hpp | 52 ++ .../infinicore/ops/distributed/allreduce.hpp | 24 + .../ops/distributed/reduce_scatter.hpp | 65 +++ .../infinicore/include/infinicore/ops/dsa.hpp | 87 ++++ .../ops/dynamic_scaled_int8_quant.hpp | 11 + .../include/infinicore/ops/embedding.hpp | 13 + .../include/infinicore/ops/equal.hpp | 19 + .../infinicore/ops/flash_attention.hpp | 12 + .../include/infinicore/ops/flipud.hpp | 19 + .../include/infinicore/ops/float_power.hpp | 68 +++ .../include/infinicore/ops/floor.hpp | 16 + .../include/infinicore/ops/floor_divide.hpp | 16 + .../include/infinicore/ops/fmin.hpp | 18 + .../include/infinicore/ops/fmod.hpp | 16 + .../infinicore/ops/fp8_indexer_logits.hpp | 22 + .../infinicore/ops/fp8_indexer_quant.hpp | 39 ++ .../infinicore/ops/fp8_mla_rmsnorm_cache.hpp | 22 + .../include/infinicore/ops/fp8_sparse_mla.hpp | 20 + .../ops/fused_gated_delta_net_gating.hpp | 37 ++ .../infinicore/ops/fused_rotary_embedding.hpp | 16 + .../infinicore/ops/gaussian_nll_loss.hpp | 28 ++ .../include/infinicore/ops/gelu.hpp | 16 + .../include/infinicore/ops/gelutanh.hpp | 16 + .../include/infinicore/ops/gemm.hpp | 14 + .../infinicore/ops/gptq_marlin_gemm.hpp | 18 + .../infinicore/ops/grouped_topk_vllm.hpp | 8 + .../include/infinicore/ops/hardswish.hpp | 18 + .../include/infinicore/ops/hardtanh.hpp | 18 + .../infinicore/ops/hinge_embedding_loss.hpp | 18 + .../include/infinicore/ops/huber_loss.hpp | 21 + .../include/infinicore/ops/hypot.hpp | 18 + .../include/infinicore/ops/index_add.hpp | 19 + .../include/infinicore/ops/index_copy.hpp | 18 + .../include/infinicore/ops/inner.hpp | 17 + .../include/infinicore/ops/interpolate.hpp | 30 ++ .../include/infinicore/ops/kron.hpp | 18 + .../include/infinicore/ops/kthvalue.hpp | 24 + .../include/infinicore/ops/kv_caching.hpp | 16 + .../include/infinicore/ops/layer_norm.hpp | 16 + .../include/infinicore/ops/ldexp.hpp | 24 + .../include/infinicore/ops/lerp.hpp | 27 + .../include/infinicore/ops/linear.hpp | 12 + .../infinicore/ops/linear_w4a16_awq.hpp | 12 + .../infinicore/ops/linear_w4a16_gptq_qy.hpp | 12 + .../include/infinicore/ops/linear_w8a8i8.hpp | 13 + .../include/infinicore/ops/log_softmax.hpp | 23 + .../include/infinicore/ops/logaddexp.hpp | 18 + .../include/infinicore/ops/logaddexp2.hpp | 18 + .../include/infinicore/ops/logcumsumexp.hpp | 21 + .../include/infinicore/ops/logdet.hpp | 14 + .../include/infinicore/ops/logical_and.hpp | 20 + .../include/infinicore/ops/logical_not.hpp | 23 + .../infinicore/ops/mamba_selective_scan.hpp | 43 ++ .../include/infinicore/ops/masked_select.hpp | 16 + .../include/infinicore/ops/matmul.hpp | 11 + .../infinicore/include/infinicore/ops/mha.hpp | 34 ++ .../include/infinicore/ops/mha_kvcache.hpp | 51 ++ .../include/infinicore/ops/mha_varlen.hpp | 46 ++ .../include/infinicore/ops/moe_align.hpp | 40 ++ .../infinicore/ops/moe_argsort_bincount.hpp | 7 + .../infinicore/ops/moe_expand_input.hpp | 9 + .../infinicore/ops/moe_fused_dense.hpp | 42 ++ .../include/infinicore/ops/moe_fused_gate.hpp | 44 ++ .../infinicore/ops/moe_silu_and_mul_quant.hpp | 9 + .../include/infinicore/ops/moe_sum.hpp | 15 + .../include/infinicore/ops/moe_sum_vllm.hpp | 7 + .../infinicore/ops/moe_topk_sigmoid.hpp | 32 ++ .../infinicore/ops/moe_topk_softmax.hpp | 35 ++ .../include/infinicore/ops/moe_topk_vllm.hpp | 7 + .../include/infinicore/ops/mrope.hpp | 51 ++ .../infinicore/include/infinicore/ops/mul.hpp | 14 + .../include/infinicore/ops/mul_scalar.hpp | 14 + .../infinicore/ops/multi_margin_loss.hpp | 19 + .../include/infinicore/ops/nrm2.hpp | 14 + .../ops/nsa_compress_paged_cache.hpp | 15 + .../infinicore/ops/nsa_paged_attention.hpp | 19 + .../include/infinicore/ops/ones.hpp | 16 + .../infinicore/include/infinicore/ops/pad.hpp | 25 + .../infinicore/ops/paged_attention.hpp | 20 + .../infinicore/ops/paged_attention_mla.hpp | 17 + .../ops/paged_attention_prefill.hpp | 52 ++ .../include/infinicore/ops/paged_caching.hpp | 13 + .../infinicore/ops/per_channel_quant_i8.hpp | 12 + .../infinicore/ops/per_tensor_dequant_i8.hpp | 11 + .../infinicore/ops/per_tensor_quant_i8.hpp | 13 + .../include/infinicore/ops/prelu.hpp | 16 + .../infinicore/ops/prepare_moe_input.hpp | 48 ++ .../include/infinicore/ops/quickgelu.hpp | 16 + .../include/infinicore/ops/random_sample.hpp | 22 + .../include/infinicore/ops/rearrange.hpp | 14 + .../include/infinicore/ops/reciprocal.hpp | 16 + .../ops/recurrent_gated_delta_rule.hpp | 55 ++ .../include/infinicore/ops/relu.hpp | 16 + .../include/infinicore/ops/relu6.hpp | 16 + .../include/infinicore/ops/rms_norm.hpp | 14 + .../include/infinicore/ops/rope.hpp | 28 ++ .../infinicore/include/infinicore/ops/rot.hpp | 13 + .../include/infinicore/ops/rotg.hpp | 17 + .../include/infinicore/ops/rotm.hpp | 17 + .../include/infinicore/ops/rotmg.hpp | 17 + .../include/infinicore/ops/rwkv5_wkv.hpp | 34 ++ .../include/infinicore/ops/scal.hpp | 13 + .../include/infinicore/ops/scaled_mm_i8.hpp | 13 + .../ops/scaled_mm_w4a16_gptq_qy.hpp | 13 + .../include/infinicore/ops/scaled_mm_w4a8.hpp | 9 + .../include/infinicore/ops/scaled_mm_w8a8.hpp | 12 + .../include/infinicore/ops/scatter.hpp | 21 + .../ops/select_last_token_hidden.hpp | 13 + .../include/infinicore/ops/selu.hpp | 18 + .../include/infinicore/ops/sigmoid.hpp | 16 + .../include/infinicore/ops/silu.hpp | 16 + .../include/infinicore/ops/silu_and_mul.hpp | 14 + .../include/infinicore/ops/sinh.hpp | 18 + .../include/infinicore/ops/smooth_l1_loss.hpp | 19 + .../include/infinicore/ops/softmax.hpp | 16 + .../include/infinicore/ops/softplus.hpp | 20 + .../include/infinicore/ops/softsign.hpp | 16 + .../infinicore/include/infinicore/ops/sum.hpp | 19 + .../include/infinicore/ops/swap.hpp | 13 + .../include/infinicore/ops/swiglu.hpp | 15 + .../include/infinicore/ops/take.hpp | 20 + .../infinicore/include/infinicore/ops/tan.hpp | 17 + .../include/infinicore/ops/tanhshrink.hpp | 16 + .../include/infinicore/ops/topk.hpp | 16 + .../include/infinicore/ops/topksoftmax.hpp | 14 + .../infinicore/ops/triplet_margin_loss.hpp | 24 + .../ops/triplet_margin_with_distance_loss.hpp | 24 + .../include/infinicore/ops/unfold.hpp | 35 ++ .../infinicore/ops/upsample_bilinear.hpp | 22 + .../infinicore/ops/upsample_nearest.hpp | 26 + .../include/infinicore/ops/vander.hpp | 21 + .../infinicore/include/infinicore/ops/var.hpp | 19 + .../include/infinicore/ops/var_mean.hpp | 19 + .../ops/vocab_parallel_embedding.hpp | 6 + .../infinicore/ops/w16a16_group_gemm.hpp | 15 + .../infinicore/ops/w4a8_group_gemm.hpp | 6 + .../infinicore/ops/w8a8_group_gemm.hpp | 6 + .../include/infinicore/quantization.hpp | 9 + .../include/infinicore/quantization/awq.hpp | 30 ++ .../quantization/base_quantization.hpp | 43 ++ .../quantization/compressed_tensors.hpp | 20 + .../include/infinicore/quantization/gptq.hpp | 30 ++ .../infinicore/quantization/gptq_qy.hpp | 320 ++++++++++++ .../quantization/none_quantizaiton.hpp | 20 + .../quantization/quantization_scheme.hpp | 19 + .../infinicore/include/infinicore/tensor.hpp | 336 +++++++++++++ python/infinicore/include/infiniop.h | 164 ++++++ python/infinicore/include/infiniop/handle.h | 14 + .../include/infiniop/operator_descriptor.h | 13 + python/infinicore/include/infiniop/ops/acos.h | 24 + .../infiniop/ops/adaptive_avg_pool1d.h | 34 ++ .../infiniop/ops/adaptive_avg_pool3d.h | 30 ++ .../infiniop/ops/adaptive_max_pool1d.h | 22 + python/infinicore/include/infiniop/ops/add.h | 26 + .../include/infiniop/ops/add_rms_norm.h | 32 ++ .../infinicore/include/infiniop/ops/addbmm.h | 30 ++ .../infinicore/include/infiniop/ops/addcmul.h | 57 +++ python/infinicore/include/infiniop/ops/addr.h | 30 ++ .../include/infiniop/ops/affine_grid.h | 25 + python/infinicore/include/infiniop/ops/all.h | 31 ++ .../include/infiniop/ops/argwhere.h | 29 ++ python/infinicore/include/infiniop/ops/asin.h | 24 + .../infinicore/include/infiniop/ops/asinh.h | 24 + python/infinicore/include/infiniop/ops/asum.h | 24 + .../infinicore/include/infiniop/ops/atanh.h | 24 + .../include/infiniop/ops/attention.h | 34 ++ .../include/infiniop/ops/avg_pool1d.h | 32 ++ .../include/infiniop/ops/avg_pool3d.h | 27 + .../include/infiniop/ops/awq_marlin_gemm.h | 46 ++ .../include/infiniop/ops/awq_marlin_repack.h | 27 + python/infinicore/include/infiniop/ops/axpy.h | 26 + .../ops/binary_cross_entropy_with_logits.h | 73 +++ .../infiniop/ops/bitwise_right_shift.h | 26 + .../include/infiniop/ops/blas_amax.h | 24 + .../include/infiniop/ops/blas_amin.h | 24 + .../include/infiniop/ops/blas_copy.h | 24 + .../include/infiniop/ops/blas_dot.h | 26 + .../include/infiniop/ops/block_diag.h | 25 + .../include/infiniop/ops/broadcast_to.h | 24 + .../include/infiniop/ops/causal_conv1d.h | 43 ++ .../include/infiniop/ops/causal_softmax.h | 26 + .../infinicore/include/infiniop/ops/cdist.h | 56 +++ .../infiniop/ops/chunk_gated_delta_rule.h | 49 ++ python/infinicore/include/infiniop/ops/clip.h | 28 ++ python/infinicore/include/infiniop/ops/conv.h | 25 + .../include/infiniop/ops/cross_entropy.h | 31 ++ .../include/infiniop/ops/deepseek_moe.h | 57 +++ .../ops/dequant/per_tensor_dequant_int8.h | 28 ++ .../include/infiniop/ops/dequantize_awq.h | 28 ++ .../include/infiniop/ops/dequantize_gptq.h | 30 ++ python/infinicore/include/infiniop/ops/diff.h | 26 + .../infinicore/include/infiniop/ops/digamma.h | 24 + python/infinicore/include/infiniop/ops/dist.h | 27 + python/infinicore/include/infiniop/ops/dot.h | 26 + .../include/infiniop/ops/embedding.h | 25 + .../infinicore/include/infiniop/ops/equal.h | 31 ++ python/infinicore/include/infiniop/ops/erf.h | 24 + python/infinicore/include/infiniop/ops/erfc.h | 24 + .../infinicore/include/infiniop/ops/erfinv.h | 24 + .../include/infiniop/ops/flash_attention.h | 36 ++ .../infinicore/include/infiniop/ops/flipud.h | 27 + .../include/infiniop/ops/float_power.h | 27 + .../infinicore/include/infiniop/ops/floor.h | 24 + .../include/infiniop/ops/floor_divide.h | 26 + python/infinicore/include/infiniop/ops/fmin.h | 26 + python/infinicore/include/infiniop/ops/fmod.h | 26 + .../include/infiniop/ops/fp8_indexer_logits.h | 33 ++ .../include/infiniop/ops/fp8_indexer_quant.h | 64 +++ .../infiniop/ops/fp8_mla_rmsnorm_cache.h | 30 ++ .../include/infiniop/ops/fp8_sparse_mla.h | 36 ++ .../include/infiniop/ops/fused_ffn.h | 37 ++ .../ops/fused_gated_delta_net_gating.h | 43 ++ .../include/infiniop/ops/gaussian_nll_loss.h | 31 ++ python/infinicore/include/infiniop/ops/gelu.h | 24 + .../include/infiniop/ops/gelutanh.h | 43 ++ python/infinicore/include/infiniop/ops/gemm.h | 28 ++ .../include/infiniop/ops/gptq_marlin_gemm.h | 42 ++ .../include/infiniop/ops/gptq_marlin_repack.h | 29 ++ .../include/infiniop/ops/gptq_qyblas_gemm.h | 37 ++ .../include/infiniop/ops/hardswish.h | 29 ++ .../include/infiniop/ops/hardtanh.h | 27 + .../infiniop/ops/hinge_embedding_loss.h | 28 ++ .../infinicore/include/infiniop/ops/histc.h | 28 ++ .../include/infiniop/ops/huber_loss.h | 28 ++ .../infinicore/include/infiniop/ops/hypot.h | 26 + .../include/infiniop/ops/index_add.h | 29 ++ .../include/infiniop/ops/index_copy.h | 28 ++ .../infinicore/include/infiniop/ops/inner.h | 28 ++ .../include/infiniop/ops/int8_gemm.h | 32 ++ .../include/infiniop/ops/interpolate.h | 28 ++ python/infinicore/include/infiniop/ops/kron.h | 26 + .../include/infiniop/ops/kthvalue.h | 29 ++ .../include/infiniop/ops/kv_caching.h | 31 ++ .../include/infiniop/ops/layer_norm.h | 34 ++ .../infinicore/include/infiniop/ops/ldexp.h | 24 + python/infinicore/include/infiniop/ops/lerp.h | 29 ++ .../infinicore/include/infiniop/ops/log10.h | 24 + .../infinicore/include/infiniop/ops/log1p.h | 24 + .../include/infiniop/ops/log_softmax.h | 25 + .../include/infiniop/ops/logaddexp.h | 26 + .../include/infiniop/ops/logaddexp2.h | 26 + .../include/infiniop/ops/logcumsumexp.h | 30 ++ .../infinicore/include/infiniop/ops/logdet.h | 24 + .../infinicore/include/infiniop/ops/lp_norm.h | 28 ++ .../infiniop/ops/mamba_selective_scan.h | 50 ++ .../include/infiniop/ops/masked_select.h | 28 ++ .../include/infiniop/ops/matrix_power.h | 25 + .../include/infiniop/ops/moe_align.h | 37 ++ .../include/infiniop/ops/moe_fused_dense.h | 43 ++ .../include/infiniop/ops/moe_fused_gate.h | 38 ++ .../infinicore/include/infiniop/ops/moe_sum.h | 29 ++ .../include/infiniop/ops/moe_topk_sigmoid.h | 34 ++ .../include/infiniop/ops/moe_topk_softmax.h | 35 ++ .../infinicore/include/infiniop/ops/mrope.h | 42 ++ python/infinicore/include/infiniop/ops/mul.h | 26 + .../include/infiniop/ops/mul_scalar.h | 25 + .../include/infiniop/ops/multi_margin_loss.h | 30 ++ python/infinicore/include/infiniop/ops/nrm2.h | 24 + .../infiniop/ops/nsa_compress_paged_cache.h | 39 ++ .../infiniop/ops/nsa_paged_attention.h | 47 ++ python/infinicore/include/infiniop/ops/ones.h | 24 + python/infinicore/include/infiniop/ops/pad.h | 28 ++ .../include/infiniop/ops/paged_attention.h | 93 ++++ .../infiniop/ops/paged_attention_prefill.h | 87 ++++ .../include/infiniop/ops/paged_caching.h | 77 +++ .../include/infiniop/ops/pixel_shuffle.h | 25 + .../infinicore/include/infiniop/ops/prelu.h | 26 + .../include/infiniop/ops/prepare_moe_input.h | 42 ++ .../ops/quant/per_channel_quant_int8.h | 28 ++ .../ops/quant/per_tensor_quant_int8.h | 29 ++ .../include/infiniop/ops/quickgelu.h | 42 ++ .../include/infiniop/ops/random_sample.h | 52 ++ .../include/infiniop/ops/rearrange.h | 23 + .../include/infiniop/ops/reciprocal.h | 24 + .../infiniop/ops/recurrent_gated_delta_rule.h | 46 ++ python/infinicore/include/infiniop/ops/relu.h | 24 + .../infinicore/include/infiniop/ops/relu6.h | 24 + .../include/infiniop/ops/rms_norm.h | 23 + python/infinicore/include/infiniop/ops/rope.h | 40 ++ python/infinicore/include/infiniop/ops/rot.h | 28 ++ python/infinicore/include/infiniop/ops/rotg.h | 28 ++ python/infinicore/include/infiniop/ops/rotm.h | 26 + .../infinicore/include/infiniop/ops/rotmg.h | 30 ++ .../include/infiniop/ops/rwkv5_wkv.h | 45 ++ python/infinicore/include/infiniop/ops/scal.h | 24 + .../infinicore/include/infiniop/ops/scatter.h | 30 ++ .../infiniop/ops/select_last_token_hidden.h | 25 + python/infinicore/include/infiniop/ops/selu.h | 24 + .../infinicore/include/infiniop/ops/sigmoid.h | 24 + python/infinicore/include/infiniop/ops/silu.h | 24 + .../include/infiniop/ops/silu_and_mul.h | 71 +++ python/infinicore/include/infiniop/ops/sinh.h | 24 + .../include/infiniop/ops/smooth_l1_loss.h | 27 + .../infinicore/include/infiniop/ops/softmax.h | 27 + .../include/infiniop/ops/softplus.h | 29 ++ .../include/infiniop/ops/softsign.h | 25 + python/infinicore/include/infiniop/ops/sub.h | 26 + python/infinicore/include/infiniop/ops/sum.h | 31 ++ python/infinicore/include/infiniop/ops/swap.h | 24 + .../infinicore/include/infiniop/ops/swiglu.h | 26 + python/infinicore/include/infiniop/ops/take.h | 26 + python/infinicore/include/infiniop/ops/tan.h | 26 + python/infinicore/include/infiniop/ops/tanh.h | 24 + .../include/infiniop/ops/tanhshrink.h | 24 + python/infinicore/include/infiniop/ops/topk.h | 35 ++ .../include/infiniop/ops/topkrouter.h | 28 ++ .../include/infiniop/ops/topksoftmax.h | 26 + .../infiniop/ops/triplet_margin_loss.h | 33 ++ .../ops/triplet_margin_with_distance_loss.h | 32 ++ .../infinicore/include/infiniop/ops/unfold.h | 28 ++ .../include/infiniop/ops/upsample_bilinear.h | 25 + .../include/infiniop/ops/upsample_nearest.h | 24 + .../infinicore/include/infiniop/ops/vander.h | 25 + python/infinicore/include/infiniop/ops/var.h | 33 ++ .../include/infiniop/ops/var_mean.h | 35 ++ .../infinicore/include/infiniop/ops/zeros.h | 24 + .../include/infiniop/tensor_descriptor.h | 14 + python/infinicore/include/infinirt.h | 134 +++++ python/infinicore/ops/dsa.py | 151 ++++++ python/infinicore/utils.py | 4 + src/infiniccl-test/infiniccl_test.cpp | 27 +- src/infiniccl/cuda/infiniccl_cuda.cu | 44 +- src/infiniccl/infiniccl.cc | 46 ++ src/infiniccl/infiniccl_impl.h | 16 + src/infinicore/adaptor/aten_adaptor.cc | 74 ++- .../adaptor/vllm_iluvatar_adaptor.cc | 208 +++++++- src/infinicore/context/context_impl.cc | 4 + src/infinicore/context/runtime/runtime.cc | 14 +- src/infinicore/context/runtime/runtime.hpp | 1 + src/infinicore/graph/graph.cc | 126 +++-- src/infinicore/graph/graph_manager.hpp | 1 + src/infinicore/ops/bmm_strided/bmm_strided.cc | 88 ++++ .../concat_and_cache_mla.cc | 18 +- src/infinicore/ops/distributed/p2p.cc | 202 ++++++++ src/infinicore/ops/dsa/dsa.cc | 463 +++++++++++++++++ .../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 +++ .../fused_rotary_embedding.cc | 131 +++++ .../grouped_topk_vllm/grouped_topk_vllm.cc | 3 +- .../select_last_token_hidden.cc | 54 ++ .../select_last_token_hidden_infiniop.cc | 53 ++ .../vocab_parallel_embedding.cc | 45 +- src/infinicore/pybind11/ops.hpp | 6 + src/infinicore/pybind11/ops/dsa.hpp | 43 ++ .../pybind11/ops/fp8_indexer_quant.hpp | 32 ++ .../pybind11/ops/fp8_mla_rmsnorm_cache.hpp | 28 ++ src/infinicore/tensor/tensor.cc | 89 +++- .../fp8_indexer_logits/fp8_indexer_logits.h | 67 +++ .../nvidia/fp8_indexer_logits_nvidia.cu | 177 +++++++ .../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 +++++ .../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/infinicore/ops/dsa.py | 213 ++++++++ .../ops/fp8_mla_rmsnorm_dual_cache.py | 221 ++++++++ test/infiniop/select_last_token_hidden.py | 108 ++++ 488 files changed, 20149 insertions(+), 50 deletions(-) create mode 100644 include/infinicore/ops/bmm_strided.hpp create mode 100644 include/infinicore/ops/distributed/p2p.hpp create mode 100644 include/infinicore/ops/dsa.hpp 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/fused_rotary_embedding.hpp create mode 100644 include/infinicore/ops/select_last_token_hidden.hpp 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 include/infiniop/ops/select_last_token_hidden.h create mode 100644 python/infinicore/include/infiniccl.h create mode 100644 python/infinicore/include/infinicore.h create mode 100644 python/infinicore/include/infinicore.hpp create mode 100644 python/infinicore/include/infinicore/adaptor/aten_adaptor.hpp create mode 100644 python/infinicore/include/infinicore/adaptor/flash_attention_adaptor.hpp create mode 100644 python/infinicore/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp create mode 100644 python/infinicore/include/infinicore/analyzer.hpp create mode 100644 python/infinicore/include/infinicore/analyzer/intent_generator.hpp create mode 100644 python/infinicore/include/infinicore/analyzer/mutual_awareness_analyzer.hpp create mode 100644 python/infinicore/include/infinicore/analyzer/op_trace.hpp create mode 100644 python/infinicore/include/infinicore/analyzer/op_type.hpp create mode 100644 python/infinicore/include/infinicore/analyzer/op_type_registry.hpp create mode 100644 python/infinicore/include/infinicore/analyzer/optimization_intent.hpp create mode 100644 python/infinicore/include/infinicore/analyzer/phase_detector.hpp create mode 100644 python/infinicore/include/infinicore/analyzer/resource_sensor.hpp create mode 100644 python/infinicore/include/infinicore/common/LRUCache.hpp create mode 100644 python/infinicore/include/infinicore/common/hash.hpp create mode 100644 python/infinicore/include/infinicore/context/context.hpp create mode 100644 python/infinicore/include/infinicore/device.hpp create mode 100644 python/infinicore/include/infinicore/device_event.hpp create mode 100644 python/infinicore/include/infinicore/dtype.hpp create mode 100644 python/infinicore/include/infinicore/graph/graph.hpp create mode 100644 python/infinicore/include/infinicore/io.hpp create mode 100644 python/infinicore/include/infinicore/memory.hpp create mode 100644 python/infinicore/include/infinicore/nn.hpp create mode 100644 python/infinicore/include/infinicore/nn/embedding.hpp create mode 100644 python/infinicore/include/infinicore/nn/layer_norm.hpp create mode 100644 python/infinicore/include/infinicore/nn/linear.hpp create mode 100644 python/infinicore/include/infinicore/nn/module.hpp create mode 100644 python/infinicore/include/infinicore/nn/parameter.hpp create mode 100644 python/infinicore/include/infinicore/nn/rmsnorm.hpp create mode 100644 python/infinicore/include/infinicore/nn/rope.hpp create mode 100644 python/infinicore/include/infinicore/nn/rope_scaling_configs.hpp create mode 100644 python/infinicore/include/infinicore/ops.hpp create mode 100644 python/infinicore/include/infinicore/ops/acos.hpp create mode 100644 python/infinicore/include/infinicore/ops/adaptive_avg_pool1d.hpp create mode 100644 python/infinicore/include/infinicore/ops/adaptive_avg_pool3d.hpp create mode 100644 python/infinicore/include/infinicore/ops/adaptive_max_pool1d.hpp create mode 100644 python/infinicore/include/infinicore/ops/add.hpp create mode 100644 python/infinicore/include/infinicore/ops/add_rms_norm.hpp create mode 100644 python/infinicore/include/infinicore/ops/addbmm.hpp create mode 100644 python/infinicore/include/infinicore/ops/addcmul.hpp create mode 100644 python/infinicore/include/infinicore/ops/addr.hpp create mode 100644 python/infinicore/include/infinicore/ops/affine_grid.hpp create mode 100644 python/infinicore/include/infinicore/ops/all.hpp create mode 100644 python/infinicore/include/infinicore/ops/argwhere.hpp create mode 100644 python/infinicore/include/infinicore/ops/asin.hpp create mode 100644 python/infinicore/include/infinicore/ops/asinh.hpp create mode 100644 python/infinicore/include/infinicore/ops/asum.hpp create mode 100644 python/infinicore/include/infinicore/ops/atanh.hpp create mode 100644 python/infinicore/include/infinicore/ops/attention.hpp create mode 100644 python/infinicore/include/infinicore/ops/avg_pool1d.hpp create mode 100644 python/infinicore/include/infinicore/ops/awq_marlin_gemm.hpp create mode 100644 python/infinicore/include/infinicore/ops/axpy.hpp create mode 100644 python/infinicore/include/infinicore/ops/baddbmm.hpp create mode 100644 python/infinicore/include/infinicore/ops/bilinear.hpp create mode 100644 python/infinicore/include/infinicore/ops/binary_cross_entropy_with_logits.hpp create mode 100644 python/infinicore/include/infinicore/ops/bitwise_right_shift.hpp create mode 100644 python/infinicore/include/infinicore/ops/blas_amax.hpp create mode 100644 python/infinicore/include/infinicore/ops/blas_amin.hpp create mode 100644 python/infinicore/include/infinicore/ops/blas_copy.hpp create mode 100644 python/infinicore/include/infinicore/ops/blas_dot.hpp create mode 100644 python/infinicore/include/infinicore/ops/block_diag.hpp create mode 100644 python/infinicore/include/infinicore/ops/bmm_strided.hpp create mode 100644 python/infinicore/include/infinicore/ops/broadcast_to.hpp create mode 100644 python/infinicore/include/infinicore/ops/cast.hpp create mode 100644 python/infinicore/include/infinicore/ops/cat.hpp create mode 100644 python/infinicore/include/infinicore/ops/causal_conv1d.hpp create mode 100644 python/infinicore/include/infinicore/ops/causal_softmax.hpp create mode 100644 python/infinicore/include/infinicore/ops/cdist.hpp create mode 100644 python/infinicore/include/infinicore/ops/chunk_gated_delta_rule.hpp create mode 100644 python/infinicore/include/infinicore/ops/common/cache.hpp create mode 100644 python/infinicore/include/infinicore/ops/common/dispatcher.hpp create mode 100644 python/infinicore/include/infinicore/ops/common/op.hpp create mode 100644 python/infinicore/include/infinicore/ops/concat_and_cache_mla.hpp create mode 100644 python/infinicore/include/infinicore/ops/concat_and_cache_mla_int8.hpp create mode 100644 python/infinicore/include/infinicore/ops/concat_mla_q.hpp create mode 100644 python/infinicore/include/infinicore/ops/conv2d.hpp create mode 100644 python/infinicore/include/infinicore/ops/cross_entropy.hpp create mode 100644 python/infinicore/include/infinicore/ops/deepseek_moe.hpp create mode 100644 python/infinicore/include/infinicore/ops/dequantize_awq.hpp create mode 100644 python/infinicore/include/infinicore/ops/diff.hpp create mode 100644 python/infinicore/include/infinicore/ops/digamma.hpp create mode 100644 python/infinicore/include/infinicore/ops/dist.hpp create mode 100644 python/infinicore/include/infinicore/ops/distributed/allgather.hpp create mode 100644 python/infinicore/include/infinicore/ops/distributed/allreduce.hpp create mode 100644 python/infinicore/include/infinicore/ops/distributed/reduce_scatter.hpp create mode 100644 python/infinicore/include/infinicore/ops/dsa.hpp create mode 100644 python/infinicore/include/infinicore/ops/dynamic_scaled_int8_quant.hpp create mode 100644 python/infinicore/include/infinicore/ops/embedding.hpp create mode 100644 python/infinicore/include/infinicore/ops/equal.hpp create mode 100644 python/infinicore/include/infinicore/ops/flash_attention.hpp create mode 100644 python/infinicore/include/infinicore/ops/flipud.hpp create mode 100644 python/infinicore/include/infinicore/ops/float_power.hpp create mode 100644 python/infinicore/include/infinicore/ops/floor.hpp create mode 100644 python/infinicore/include/infinicore/ops/floor_divide.hpp create mode 100644 python/infinicore/include/infinicore/ops/fmin.hpp create mode 100644 python/infinicore/include/infinicore/ops/fmod.hpp create mode 100644 python/infinicore/include/infinicore/ops/fp8_indexer_logits.hpp create mode 100644 python/infinicore/include/infinicore/ops/fp8_indexer_quant.hpp create mode 100644 python/infinicore/include/infinicore/ops/fp8_mla_rmsnorm_cache.hpp create mode 100644 python/infinicore/include/infinicore/ops/fp8_sparse_mla.hpp create mode 100644 python/infinicore/include/infinicore/ops/fused_gated_delta_net_gating.hpp create mode 100644 python/infinicore/include/infinicore/ops/fused_rotary_embedding.hpp create mode 100644 python/infinicore/include/infinicore/ops/gaussian_nll_loss.hpp create mode 100644 python/infinicore/include/infinicore/ops/gelu.hpp create mode 100644 python/infinicore/include/infinicore/ops/gelutanh.hpp create mode 100644 python/infinicore/include/infinicore/ops/gemm.hpp create mode 100644 python/infinicore/include/infinicore/ops/gptq_marlin_gemm.hpp create mode 100644 python/infinicore/include/infinicore/ops/grouped_topk_vllm.hpp create mode 100644 python/infinicore/include/infinicore/ops/hardswish.hpp create mode 100644 python/infinicore/include/infinicore/ops/hardtanh.hpp create mode 100644 python/infinicore/include/infinicore/ops/hinge_embedding_loss.hpp create mode 100644 python/infinicore/include/infinicore/ops/huber_loss.hpp create mode 100644 python/infinicore/include/infinicore/ops/hypot.hpp create mode 100644 python/infinicore/include/infinicore/ops/index_add.hpp create mode 100644 python/infinicore/include/infinicore/ops/index_copy.hpp create mode 100644 python/infinicore/include/infinicore/ops/inner.hpp create mode 100644 python/infinicore/include/infinicore/ops/interpolate.hpp create mode 100644 python/infinicore/include/infinicore/ops/kron.hpp create mode 100644 python/infinicore/include/infinicore/ops/kthvalue.hpp create mode 100644 python/infinicore/include/infinicore/ops/kv_caching.hpp create mode 100644 python/infinicore/include/infinicore/ops/layer_norm.hpp create mode 100644 python/infinicore/include/infinicore/ops/ldexp.hpp create mode 100644 python/infinicore/include/infinicore/ops/lerp.hpp create mode 100644 python/infinicore/include/infinicore/ops/linear.hpp create mode 100644 python/infinicore/include/infinicore/ops/linear_w4a16_awq.hpp create mode 100644 python/infinicore/include/infinicore/ops/linear_w4a16_gptq_qy.hpp create mode 100644 python/infinicore/include/infinicore/ops/linear_w8a8i8.hpp create mode 100644 python/infinicore/include/infinicore/ops/log_softmax.hpp create mode 100644 python/infinicore/include/infinicore/ops/logaddexp.hpp create mode 100644 python/infinicore/include/infinicore/ops/logaddexp2.hpp create mode 100644 python/infinicore/include/infinicore/ops/logcumsumexp.hpp create mode 100644 python/infinicore/include/infinicore/ops/logdet.hpp create mode 100644 python/infinicore/include/infinicore/ops/logical_and.hpp create mode 100644 python/infinicore/include/infinicore/ops/logical_not.hpp create mode 100644 python/infinicore/include/infinicore/ops/mamba_selective_scan.hpp create mode 100644 python/infinicore/include/infinicore/ops/masked_select.hpp create mode 100644 python/infinicore/include/infinicore/ops/matmul.hpp create mode 100644 python/infinicore/include/infinicore/ops/mha.hpp create mode 100644 python/infinicore/include/infinicore/ops/mha_kvcache.hpp create mode 100644 python/infinicore/include/infinicore/ops/mha_varlen.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_align.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_argsort_bincount.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_expand_input.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_fused_dense.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_fused_gate.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_silu_and_mul_quant.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_sum.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_sum_vllm.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_topk_softmax.hpp create mode 100644 python/infinicore/include/infinicore/ops/moe_topk_vllm.hpp create mode 100644 python/infinicore/include/infinicore/ops/mrope.hpp create mode 100644 python/infinicore/include/infinicore/ops/mul.hpp create mode 100644 python/infinicore/include/infinicore/ops/mul_scalar.hpp create mode 100644 python/infinicore/include/infinicore/ops/multi_margin_loss.hpp create mode 100644 python/infinicore/include/infinicore/ops/nrm2.hpp create mode 100644 python/infinicore/include/infinicore/ops/nsa_compress_paged_cache.hpp create mode 100644 python/infinicore/include/infinicore/ops/nsa_paged_attention.hpp create mode 100644 python/infinicore/include/infinicore/ops/ones.hpp create mode 100644 python/infinicore/include/infinicore/ops/pad.hpp create mode 100644 python/infinicore/include/infinicore/ops/paged_attention.hpp create mode 100644 python/infinicore/include/infinicore/ops/paged_attention_mla.hpp create mode 100644 python/infinicore/include/infinicore/ops/paged_attention_prefill.hpp create mode 100644 python/infinicore/include/infinicore/ops/paged_caching.hpp create mode 100644 python/infinicore/include/infinicore/ops/per_channel_quant_i8.hpp create mode 100644 python/infinicore/include/infinicore/ops/per_tensor_dequant_i8.hpp create mode 100644 python/infinicore/include/infinicore/ops/per_tensor_quant_i8.hpp create mode 100644 python/infinicore/include/infinicore/ops/prelu.hpp create mode 100644 python/infinicore/include/infinicore/ops/prepare_moe_input.hpp create mode 100644 python/infinicore/include/infinicore/ops/quickgelu.hpp create mode 100644 python/infinicore/include/infinicore/ops/random_sample.hpp create mode 100644 python/infinicore/include/infinicore/ops/rearrange.hpp create mode 100644 python/infinicore/include/infinicore/ops/reciprocal.hpp create mode 100644 python/infinicore/include/infinicore/ops/recurrent_gated_delta_rule.hpp create mode 100644 python/infinicore/include/infinicore/ops/relu.hpp create mode 100644 python/infinicore/include/infinicore/ops/relu6.hpp create mode 100644 python/infinicore/include/infinicore/ops/rms_norm.hpp create mode 100644 python/infinicore/include/infinicore/ops/rope.hpp create mode 100644 python/infinicore/include/infinicore/ops/rot.hpp create mode 100644 python/infinicore/include/infinicore/ops/rotg.hpp create mode 100644 python/infinicore/include/infinicore/ops/rotm.hpp create mode 100644 python/infinicore/include/infinicore/ops/rotmg.hpp create mode 100644 python/infinicore/include/infinicore/ops/rwkv5_wkv.hpp create mode 100644 python/infinicore/include/infinicore/ops/scal.hpp create mode 100644 python/infinicore/include/infinicore/ops/scaled_mm_i8.hpp create mode 100644 python/infinicore/include/infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp create mode 100644 python/infinicore/include/infinicore/ops/scaled_mm_w4a8.hpp create mode 100644 python/infinicore/include/infinicore/ops/scaled_mm_w8a8.hpp create mode 100644 python/infinicore/include/infinicore/ops/scatter.hpp create mode 100644 python/infinicore/include/infinicore/ops/select_last_token_hidden.hpp create mode 100644 python/infinicore/include/infinicore/ops/selu.hpp create mode 100644 python/infinicore/include/infinicore/ops/sigmoid.hpp create mode 100644 python/infinicore/include/infinicore/ops/silu.hpp create mode 100644 python/infinicore/include/infinicore/ops/silu_and_mul.hpp create mode 100644 python/infinicore/include/infinicore/ops/sinh.hpp create mode 100644 python/infinicore/include/infinicore/ops/smooth_l1_loss.hpp create mode 100644 python/infinicore/include/infinicore/ops/softmax.hpp create mode 100644 python/infinicore/include/infinicore/ops/softplus.hpp create mode 100644 python/infinicore/include/infinicore/ops/softsign.hpp create mode 100644 python/infinicore/include/infinicore/ops/sum.hpp create mode 100644 python/infinicore/include/infinicore/ops/swap.hpp create mode 100644 python/infinicore/include/infinicore/ops/swiglu.hpp create mode 100644 python/infinicore/include/infinicore/ops/take.hpp create mode 100644 python/infinicore/include/infinicore/ops/tan.hpp create mode 100644 python/infinicore/include/infinicore/ops/tanhshrink.hpp create mode 100644 python/infinicore/include/infinicore/ops/topk.hpp create mode 100644 python/infinicore/include/infinicore/ops/topksoftmax.hpp create mode 100644 python/infinicore/include/infinicore/ops/triplet_margin_loss.hpp create mode 100644 python/infinicore/include/infinicore/ops/triplet_margin_with_distance_loss.hpp create mode 100644 python/infinicore/include/infinicore/ops/unfold.hpp create mode 100644 python/infinicore/include/infinicore/ops/upsample_bilinear.hpp create mode 100644 python/infinicore/include/infinicore/ops/upsample_nearest.hpp create mode 100644 python/infinicore/include/infinicore/ops/vander.hpp create mode 100644 python/infinicore/include/infinicore/ops/var.hpp create mode 100644 python/infinicore/include/infinicore/ops/var_mean.hpp create mode 100644 python/infinicore/include/infinicore/ops/vocab_parallel_embedding.hpp create mode 100644 python/infinicore/include/infinicore/ops/w16a16_group_gemm.hpp create mode 100644 python/infinicore/include/infinicore/ops/w4a8_group_gemm.hpp create mode 100644 python/infinicore/include/infinicore/ops/w8a8_group_gemm.hpp create mode 100644 python/infinicore/include/infinicore/quantization.hpp create mode 100644 python/infinicore/include/infinicore/quantization/awq.hpp create mode 100644 python/infinicore/include/infinicore/quantization/base_quantization.hpp create mode 100644 python/infinicore/include/infinicore/quantization/compressed_tensors.hpp create mode 100644 python/infinicore/include/infinicore/quantization/gptq.hpp create mode 100644 python/infinicore/include/infinicore/quantization/gptq_qy.hpp create mode 100644 python/infinicore/include/infinicore/quantization/none_quantizaiton.hpp create mode 100644 python/infinicore/include/infinicore/quantization/quantization_scheme.hpp create mode 100644 python/infinicore/include/infinicore/tensor.hpp create mode 100644 python/infinicore/include/infiniop.h create mode 100644 python/infinicore/include/infiniop/handle.h create mode 100644 python/infinicore/include/infiniop/operator_descriptor.h create mode 100644 python/infinicore/include/infiniop/ops/acos.h create mode 100644 python/infinicore/include/infiniop/ops/adaptive_avg_pool1d.h create mode 100644 python/infinicore/include/infiniop/ops/adaptive_avg_pool3d.h create mode 100644 python/infinicore/include/infiniop/ops/adaptive_max_pool1d.h create mode 100644 python/infinicore/include/infiniop/ops/add.h create mode 100644 python/infinicore/include/infiniop/ops/add_rms_norm.h create mode 100644 python/infinicore/include/infiniop/ops/addbmm.h create mode 100644 python/infinicore/include/infiniop/ops/addcmul.h create mode 100644 python/infinicore/include/infiniop/ops/addr.h create mode 100644 python/infinicore/include/infiniop/ops/affine_grid.h create mode 100644 python/infinicore/include/infiniop/ops/all.h create mode 100644 python/infinicore/include/infiniop/ops/argwhere.h create mode 100644 python/infinicore/include/infiniop/ops/asin.h create mode 100644 python/infinicore/include/infiniop/ops/asinh.h create mode 100644 python/infinicore/include/infiniop/ops/asum.h create mode 100644 python/infinicore/include/infiniop/ops/atanh.h create mode 100644 python/infinicore/include/infiniop/ops/attention.h create mode 100644 python/infinicore/include/infiniop/ops/avg_pool1d.h create mode 100644 python/infinicore/include/infiniop/ops/avg_pool3d.h create mode 100644 python/infinicore/include/infiniop/ops/awq_marlin_gemm.h create mode 100644 python/infinicore/include/infiniop/ops/awq_marlin_repack.h create mode 100644 python/infinicore/include/infiniop/ops/axpy.h create mode 100644 python/infinicore/include/infiniop/ops/binary_cross_entropy_with_logits.h create mode 100644 python/infinicore/include/infiniop/ops/bitwise_right_shift.h create mode 100644 python/infinicore/include/infiniop/ops/blas_amax.h create mode 100644 python/infinicore/include/infiniop/ops/blas_amin.h create mode 100644 python/infinicore/include/infiniop/ops/blas_copy.h create mode 100644 python/infinicore/include/infiniop/ops/blas_dot.h create mode 100644 python/infinicore/include/infiniop/ops/block_diag.h create mode 100644 python/infinicore/include/infiniop/ops/broadcast_to.h create mode 100644 python/infinicore/include/infiniop/ops/causal_conv1d.h create mode 100644 python/infinicore/include/infiniop/ops/causal_softmax.h create mode 100644 python/infinicore/include/infiniop/ops/cdist.h create mode 100644 python/infinicore/include/infiniop/ops/chunk_gated_delta_rule.h create mode 100644 python/infinicore/include/infiniop/ops/clip.h create mode 100644 python/infinicore/include/infiniop/ops/conv.h create mode 100644 python/infinicore/include/infiniop/ops/cross_entropy.h create mode 100644 python/infinicore/include/infiniop/ops/deepseek_moe.h create mode 100644 python/infinicore/include/infiniop/ops/dequant/per_tensor_dequant_int8.h create mode 100644 python/infinicore/include/infiniop/ops/dequantize_awq.h create mode 100644 python/infinicore/include/infiniop/ops/dequantize_gptq.h create mode 100644 python/infinicore/include/infiniop/ops/diff.h create mode 100644 python/infinicore/include/infiniop/ops/digamma.h create mode 100644 python/infinicore/include/infiniop/ops/dist.h create mode 100644 python/infinicore/include/infiniop/ops/dot.h create mode 100644 python/infinicore/include/infiniop/ops/embedding.h create mode 100644 python/infinicore/include/infiniop/ops/equal.h create mode 100644 python/infinicore/include/infiniop/ops/erf.h create mode 100644 python/infinicore/include/infiniop/ops/erfc.h create mode 100644 python/infinicore/include/infiniop/ops/erfinv.h create mode 100644 python/infinicore/include/infiniop/ops/flash_attention.h create mode 100644 python/infinicore/include/infiniop/ops/flipud.h create mode 100644 python/infinicore/include/infiniop/ops/float_power.h create mode 100644 python/infinicore/include/infiniop/ops/floor.h create mode 100644 python/infinicore/include/infiniop/ops/floor_divide.h create mode 100644 python/infinicore/include/infiniop/ops/fmin.h create mode 100644 python/infinicore/include/infiniop/ops/fmod.h create mode 100644 python/infinicore/include/infiniop/ops/fp8_indexer_logits.h create mode 100644 python/infinicore/include/infiniop/ops/fp8_indexer_quant.h create mode 100644 python/infinicore/include/infiniop/ops/fp8_mla_rmsnorm_cache.h create mode 100644 python/infinicore/include/infiniop/ops/fp8_sparse_mla.h create mode 100644 python/infinicore/include/infiniop/ops/fused_ffn.h create mode 100644 python/infinicore/include/infiniop/ops/fused_gated_delta_net_gating.h create mode 100644 python/infinicore/include/infiniop/ops/gaussian_nll_loss.h create mode 100644 python/infinicore/include/infiniop/ops/gelu.h create mode 100644 python/infinicore/include/infiniop/ops/gelutanh.h create mode 100644 python/infinicore/include/infiniop/ops/gemm.h create mode 100644 python/infinicore/include/infiniop/ops/gptq_marlin_gemm.h create mode 100644 python/infinicore/include/infiniop/ops/gptq_marlin_repack.h create mode 100644 python/infinicore/include/infiniop/ops/gptq_qyblas_gemm.h create mode 100644 python/infinicore/include/infiniop/ops/hardswish.h create mode 100644 python/infinicore/include/infiniop/ops/hardtanh.h create mode 100644 python/infinicore/include/infiniop/ops/hinge_embedding_loss.h create mode 100644 python/infinicore/include/infiniop/ops/histc.h create mode 100644 python/infinicore/include/infiniop/ops/huber_loss.h create mode 100644 python/infinicore/include/infiniop/ops/hypot.h create mode 100644 python/infinicore/include/infiniop/ops/index_add.h create mode 100644 python/infinicore/include/infiniop/ops/index_copy.h create mode 100644 python/infinicore/include/infiniop/ops/inner.h create mode 100644 python/infinicore/include/infiniop/ops/int8_gemm.h create mode 100644 python/infinicore/include/infiniop/ops/interpolate.h create mode 100644 python/infinicore/include/infiniop/ops/kron.h create mode 100644 python/infinicore/include/infiniop/ops/kthvalue.h create mode 100644 python/infinicore/include/infiniop/ops/kv_caching.h create mode 100644 python/infinicore/include/infiniop/ops/layer_norm.h create mode 100644 python/infinicore/include/infiniop/ops/ldexp.h create mode 100644 python/infinicore/include/infiniop/ops/lerp.h create mode 100644 python/infinicore/include/infiniop/ops/log10.h create mode 100644 python/infinicore/include/infiniop/ops/log1p.h create mode 100644 python/infinicore/include/infiniop/ops/log_softmax.h create mode 100644 python/infinicore/include/infiniop/ops/logaddexp.h create mode 100644 python/infinicore/include/infiniop/ops/logaddexp2.h create mode 100644 python/infinicore/include/infiniop/ops/logcumsumexp.h create mode 100644 python/infinicore/include/infiniop/ops/logdet.h create mode 100644 python/infinicore/include/infiniop/ops/lp_norm.h create mode 100644 python/infinicore/include/infiniop/ops/mamba_selective_scan.h create mode 100644 python/infinicore/include/infiniop/ops/masked_select.h create mode 100644 python/infinicore/include/infiniop/ops/matrix_power.h create mode 100644 python/infinicore/include/infiniop/ops/moe_align.h create mode 100644 python/infinicore/include/infiniop/ops/moe_fused_dense.h create mode 100644 python/infinicore/include/infiniop/ops/moe_fused_gate.h create mode 100644 python/infinicore/include/infiniop/ops/moe_sum.h create mode 100644 python/infinicore/include/infiniop/ops/moe_topk_sigmoid.h create mode 100644 python/infinicore/include/infiniop/ops/moe_topk_softmax.h create mode 100644 python/infinicore/include/infiniop/ops/mrope.h create mode 100644 python/infinicore/include/infiniop/ops/mul.h create mode 100644 python/infinicore/include/infiniop/ops/mul_scalar.h create mode 100644 python/infinicore/include/infiniop/ops/multi_margin_loss.h create mode 100644 python/infinicore/include/infiniop/ops/nrm2.h create mode 100644 python/infinicore/include/infiniop/ops/nsa_compress_paged_cache.h create mode 100644 python/infinicore/include/infiniop/ops/nsa_paged_attention.h create mode 100644 python/infinicore/include/infiniop/ops/ones.h create mode 100644 python/infinicore/include/infiniop/ops/pad.h create mode 100644 python/infinicore/include/infiniop/ops/paged_attention.h create mode 100644 python/infinicore/include/infiniop/ops/paged_attention_prefill.h create mode 100644 python/infinicore/include/infiniop/ops/paged_caching.h create mode 100644 python/infinicore/include/infiniop/ops/pixel_shuffle.h create mode 100644 python/infinicore/include/infiniop/ops/prelu.h create mode 100644 python/infinicore/include/infiniop/ops/prepare_moe_input.h create mode 100644 python/infinicore/include/infiniop/ops/quant/per_channel_quant_int8.h create mode 100644 python/infinicore/include/infiniop/ops/quant/per_tensor_quant_int8.h create mode 100644 python/infinicore/include/infiniop/ops/quickgelu.h create mode 100644 python/infinicore/include/infiniop/ops/random_sample.h create mode 100644 python/infinicore/include/infiniop/ops/rearrange.h create mode 100644 python/infinicore/include/infiniop/ops/reciprocal.h create mode 100644 python/infinicore/include/infiniop/ops/recurrent_gated_delta_rule.h create mode 100644 python/infinicore/include/infiniop/ops/relu.h create mode 100644 python/infinicore/include/infiniop/ops/relu6.h create mode 100644 python/infinicore/include/infiniop/ops/rms_norm.h create mode 100644 python/infinicore/include/infiniop/ops/rope.h create mode 100644 python/infinicore/include/infiniop/ops/rot.h create mode 100644 python/infinicore/include/infiniop/ops/rotg.h create mode 100644 python/infinicore/include/infiniop/ops/rotm.h create mode 100644 python/infinicore/include/infiniop/ops/rotmg.h create mode 100644 python/infinicore/include/infiniop/ops/rwkv5_wkv.h create mode 100644 python/infinicore/include/infiniop/ops/scal.h create mode 100644 python/infinicore/include/infiniop/ops/scatter.h create mode 100644 python/infinicore/include/infiniop/ops/select_last_token_hidden.h create mode 100644 python/infinicore/include/infiniop/ops/selu.h create mode 100644 python/infinicore/include/infiniop/ops/sigmoid.h create mode 100644 python/infinicore/include/infiniop/ops/silu.h create mode 100644 python/infinicore/include/infiniop/ops/silu_and_mul.h create mode 100644 python/infinicore/include/infiniop/ops/sinh.h create mode 100644 python/infinicore/include/infiniop/ops/smooth_l1_loss.h create mode 100644 python/infinicore/include/infiniop/ops/softmax.h create mode 100644 python/infinicore/include/infiniop/ops/softplus.h create mode 100644 python/infinicore/include/infiniop/ops/softsign.h create mode 100644 python/infinicore/include/infiniop/ops/sub.h create mode 100644 python/infinicore/include/infiniop/ops/sum.h create mode 100644 python/infinicore/include/infiniop/ops/swap.h create mode 100644 python/infinicore/include/infiniop/ops/swiglu.h create mode 100644 python/infinicore/include/infiniop/ops/take.h create mode 100644 python/infinicore/include/infiniop/ops/tan.h create mode 100644 python/infinicore/include/infiniop/ops/tanh.h create mode 100644 python/infinicore/include/infiniop/ops/tanhshrink.h create mode 100644 python/infinicore/include/infiniop/ops/topk.h create mode 100644 python/infinicore/include/infiniop/ops/topkrouter.h create mode 100644 python/infinicore/include/infiniop/ops/topksoftmax.h create mode 100644 python/infinicore/include/infiniop/ops/triplet_margin_loss.h create mode 100644 python/infinicore/include/infiniop/ops/triplet_margin_with_distance_loss.h create mode 100644 python/infinicore/include/infiniop/ops/unfold.h create mode 100644 python/infinicore/include/infiniop/ops/upsample_bilinear.h create mode 100644 python/infinicore/include/infiniop/ops/upsample_nearest.h create mode 100644 python/infinicore/include/infiniop/ops/vander.h create mode 100644 python/infinicore/include/infiniop/ops/var.h create mode 100644 python/infinicore/include/infiniop/ops/var_mean.h create mode 100644 python/infinicore/include/infiniop/ops/zeros.h create mode 100644 python/infinicore/include/infiniop/tensor_descriptor.h create mode 100644 python/infinicore/include/infinirt.h create mode 100644 python/infinicore/ops/dsa.py create mode 100644 src/infinicore/ops/bmm_strided/bmm_strided.cc create mode 100644 src/infinicore/ops/distributed/p2p.cc create mode 100644 src/infinicore/ops/dsa/dsa.cc 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/fused_rotary_embedding/fused_rotary_embedding.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/dsa.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/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 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/infinicore/ops/dsa.py create mode 100644 test/infinicore/ops/fp8_mla_rmsnorm_dual_cache.py create mode 100644 test/infiniop/select_last_token_hidden.py diff --git a/include/infiniccl.h b/include/infiniccl.h index b338d85f5..22f97263b 100644 --- a/include/infiniccl.h +++ b/include/infiniccl.h @@ -27,6 +27,22 @@ __INFINI_C __export infiniStatus_t infinicclGroupStart(infinicclComm_t comm); __INFINI_C __export infiniStatus_t infinicclGroupEnd(infinicclComm_t comm); +__INFINI_C __export infiniStatus_t infinicclSend( + const void *sendbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream); + +__INFINI_C __export infiniStatus_t infinicclRecv( + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream); + __INFINI_C __export infiniStatus_t infinicclAllReduce( void *sendbuf, void *recvbuf, diff --git a/include/infinicore/adaptor/aten_adaptor.hpp b/include/infinicore/adaptor/aten_adaptor.hpp index ded231482..ddb8d1c58 100644 --- a/include/infinicore/adaptor/aten_adaptor.hpp +++ b/include/infinicore/adaptor/aten_adaptor.hpp @@ -32,6 +32,8 @@ inline at::ScalarType to_at_dtype(DataType dtype) { return at::kBFloat16; case DataType::I8: return at::kChar; + case DataType::U8: + return at::kByte; case DataType::I32: return at::kInt; case DataType::I64: diff --git a/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp b/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp index b84dd88a7..067c49ca3 100644 --- a/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp +++ b/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp @@ -4,10 +4,12 @@ #include #include #include +#include namespace infinicore::adaptor::vllm_iluvatar { bool available(); +bool rotary_embedding_available(); bool dynamic_scaled_int8_quant_available(); bool concat_mla_q_available(); bool concat_and_cache_mla_available(); @@ -25,7 +27,23 @@ bool argsort_bincount_with_inv_pos_available(); bool expand_moe_input_with_inv_pos_available(); bool silu_and_mul_quant_available(); bool moe_sum_vllm_available(); +bool fused_deepseek_v2_indexer_postprocess_available(); +bool indexer_k_cache_available(); +bool indexer_k_quant_and_cache_available(); +bool compute_block_sparse_mqa_logits_available(); +bool select_prefill_topk_block_indices_available(); +bool select_decode_topk_block_indices_available(); +bool map_prefill_request_block_indices_available(); +bool map_decode_request_block_indices_available(); +bool sparse_flash_mla_available(); +bool topk_indices_context_lens_available(); void fused_add_rms_norm(at::Tensor &input, at::Tensor &residual, at::Tensor &weight, float epsilon); +void rotary_embedding(at::Tensor &positions, + at::Tensor &query, + std::optional key, + int64_t head_size, + at::Tensor &cos_sin_cache, + bool is_neox); void dynamic_scaled_int8_quant(at::Tensor &output, at::Tensor &input_scales, const at::Tensor &input); void concat_mla_q(at::Tensor &ql_nope, at::Tensor &q_pe, at::Tensor &q_out); void concat_and_cache_mla(at::Tensor &kv_c, at::Tensor &k_pe, at::Tensor &kv_cache, at::Tensor &slot_mapping, const std::string &kv_cache_dtype, at::Tensor &scale); @@ -51,6 +69,17 @@ void argsort_bincount_with_inv_pos(const at::Tensor &topk_ids, at::Tensor &token void expand_moe_input_with_inv_pos(at::Tensor &expand_states, std::optional expand_scales, const at::Tensor &hidden_states, const at::Tensor &inv_pos, int64_t top_k, int64_t group_size, int64_t format); void silu_and_mul_quant(at::Tensor &output, std::optional output_scale, const at::Tensor &input, int64_t format); void moe_sum_vllm(at::Tensor &output, const at::Tensor &input, std::optional topk_weights, std::optional extra_residual, double routed_scale, double residual_scale); +void select_last_token_hidden(at::Tensor &output, at::Tensor &indices, const at::Tensor &hidden_states, const at::Tensor &input_offsets); +void fused_deepseek_v2_indexer_postprocess(at::Tensor &q_out, at::Tensor &k_out, at::Tensor &weights_out, at::Tensor &kv_cache, const at::Tensor &slot_mapping, const at::Tensor &q, const at::Tensor &kw, const at::Tensor &norm_weight, const at::Tensor &norm_bias, const at::Tensor &positions, const at::Tensor &cos_sin_cache, int64_t num_cache_tokens, bool is_neox, double eps, double weights_scale); +void indexer_k_cache(const at::Tensor &k, at::Tensor &kv_cache, const at::Tensor &slot_mapping); +void indexer_k_quant_and_cache(at::Tensor &k, at::Tensor &kv_cache, at::Tensor &slot_mapping, int64_t quant_block_size, const std::string &scale_fmt); +void compute_block_sparse_mqa_logits(const at::Tensor &q, const at::Tensor &kv_cache, const at::Tensor &cu_seqlens_q, const at::Tensor &cu_seqlens_kv, const at::Tensor &block_table, const at::Tensor &weights, at::Tensor &logits, int64_t max_q_len, int64_t max_kv_len, int64_t max_context_len); +void select_prefill_topk_block_indices(const at::Tensor &logits, const at::Tensor &cu_seqlen_ks, const at::Tensor &cu_seqlen_ke, at::Tensor &topk_indices); +void select_decode_topk_block_indices(const at::Tensor &logits, const at::Tensor &seq_lens, at::Tensor &topk_indices); +void map_prefill_request_block_indices(at::Tensor &output, const at::Tensor &req_id, const at::Tensor &block_table, const at::Tensor &token_indices, int64_t block_size, bool has_prefill_workspace, std::optional prefill_workspace_request_ids, std::optional prefill_workspace_starts); +void map_decode_request_block_indices(at::Tensor &output, const at::Tensor &req_id, const at::Tensor &block_table, const at::Tensor &token_indices, int64_t block_size); +void topk_indices_context_lens(at::Tensor &topk_lens, const at::Tensor &indices); +void sparse_flash_mla(at::Tensor &output, at::Tensor &query, at::Tensor &kv_cache, at::Tensor &indices, at::Tensor &topk_lens, float scale, std::optional attn_sink); } // namespace infinicore::adaptor::vllm_iluvatar #endif // ENABLE_ATEN diff --git a/include/infinicore/context/context.hpp b/include/infinicore/context/context.hpp index 82320ed52..9db86e97c 100644 --- a/include/infinicore/context/context.hpp +++ b/include/infinicore/context/context.hpp @@ -51,6 +51,7 @@ bool isGraphRecording(); void startGraphRecording(); void addGraphOperator(std::shared_ptr op); std::shared_ptr stopGraphRecording(); +void cancelGraphRecording(); } // namespace context diff --git a/include/infinicore/graph/graph.hpp b/include/infinicore/graph/graph.hpp index be368f92d..d1b695789 100644 --- a/include/infinicore/graph/graph.hpp +++ b/include/infinicore/graph/graph.hpp @@ -17,6 +17,12 @@ class GraphTensor : public Tensor { class GraphOperator { public: virtual void run() const = 0; + virtual bool is_device_graph_capture_safe() const { + return true; + } + virtual bool requires_stream_sync_after_run() const { + return false; + } virtual ~GraphOperator() = default; }; @@ -49,7 +55,8 @@ class Graph { private: struct DeviceGraph; - std::unique_ptr device_graph_; + struct ReplayStep; + std::vector> replay_steps_; }; } // namespace infinicore::graph diff --git a/include/infinicore/nn/rope.hpp b/include/infinicore/nn/rope.hpp index cbffb81c8..eaeba8712 100644 --- a/include/infinicore/nn/rope.hpp +++ b/include/infinicore/nn/rope.hpp @@ -87,6 +87,8 @@ class RoPE : public Module { DataType dtype() const { return dtype_; } const std::optional> &mrope_section() const { return mrope_section_; } bool mrope_interleaved() const { return mrope_interleaved_; } + const Tensor &sin_cache() const { return sin_cache_; } + const Tensor &cos_cache() const { return cos_cache_; } // String representation std::string extra_repr() const; diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index 7058b9869..3f56203d3 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -25,10 +25,15 @@ #include "ops/conv2d.hpp" #include "ops/cross_entropy.hpp" #include "ops/deepseek_moe.hpp" +#include "ops/dsa.hpp" #include "ops/embedding.hpp" #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/gelu.hpp" #include "ops/gelutanh.hpp" @@ -68,6 +73,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/bmm_strided.hpp b/include/infinicore/ops/bmm_strided.hpp new file mode 100644 index 000000000..6610ef443 --- /dev/null +++ b/include/infinicore/ops/bmm_strided.hpp @@ -0,0 +1,12 @@ +#pragma once + +#include "../tensor.hpp" + +namespace infinicore::op { + +// Batched matrix multiplication that preserves arbitrary valid tensor strides, +// including a transposed output view. The Iluvatar implementation deliberately +// mirrors vLLM's torch.bmm(..., out=transpose_view) MLA projection path. +void bmm_strided_(Tensor output, const Tensor &a, const Tensor &b); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/distributed/p2p.hpp b/include/infinicore/ops/distributed/p2p.hpp new file mode 100644 index 000000000..c4d70b859 --- /dev/null +++ b/include/infinicore/ops/distributed/p2p.hpp @@ -0,0 +1,80 @@ +#pragma once + +#include "../../graph/graph.hpp" +#include "../common/op.hpp" + +#include +#include + +namespace infinicore::op::distributed { +class Send : public graph::GraphOperator { +public: + Send(const Tensor &input, int peer, infinicclComm_t communicator); + ~Send(); + void run() const override; + bool is_device_graph_capture_safe() const override { + return false; + } + static void execute( + const Tensor &input, int peer, infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +class Recv : public graph::GraphOperator { +public: + Recv(Tensor output, int peer, infinicclComm_t communicator); + ~Recv(); + void run() const override; + bool is_device_graph_capture_safe() const override { + return false; + } + static void execute( + Tensor output, int peer, infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +class GroupedSend : public graph::GraphOperator { +public: + GroupedSend(const std::vector &inputs, int peer, + infinicclComm_t communicator); + ~GroupedSend(); + void run() const override; + bool is_device_graph_capture_safe() const override { + return false; + } + static void execute(const std::vector &inputs, int peer, + infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +class GroupedRecv : public graph::GraphOperator { +public: + GroupedRecv(const std::vector &outputs, int peer, + infinicclComm_t communicator); + ~GroupedRecv(); + void run() const override; + bool is_device_graph_capture_safe() const override { + return false; + } + static void execute(const std::vector &outputs, int peer, + infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +void send(const Tensor &input, int peer, infinicclComm_t communicator); +void recv_(Tensor output, int peer, infinicclComm_t communicator); +void send_grouped( + const std::vector &inputs, + int peer, infinicclComm_t communicator); +void recv_grouped_( + const std::vector &outputs, + int peer, infinicclComm_t communicator); +} // namespace infinicore::op::distributed diff --git a/include/infinicore/ops/dsa.hpp b/include/infinicore/ops/dsa.hpp new file mode 100644 index 000000000..ab40c9f17 --- /dev/null +++ b/include/infinicore/ops/dsa.hpp @@ -0,0 +1,87 @@ +#pragma once + +#include "../tensor.hpp" + +#include +#include + +namespace infinicore::op { + +void fused_deepseek_v2_indexer_postprocess_( + Tensor q_out, + Tensor k_out, + Tensor weights_out, + Tensor kv_cache, + const Tensor &slot_mapping, + const Tensor &q, + const Tensor &kw, + const Tensor &norm_weight, + const Tensor &norm_bias, + const Tensor &positions, + const Tensor &cos_sin_cache, + int64_t num_cache_tokens, + bool is_neox, + double eps, + double weights_scale); + +void indexer_k_cache_(const Tensor &k, Tensor kv_cache, const Tensor &slot_mapping); + +void indexer_k_quant_and_cache_( + const Tensor &k, + Tensor kv_cache, + const Tensor &slot_mapping, + int64_t quant_block_size = 128, + const std::string &scale_fmt = "ue8m0"); + +void compute_block_sparse_mqa_logits_( + Tensor logits, + const Tensor &q, + const Tensor &kv_cache, + const Tensor &cu_seqlens_q, + const Tensor &cu_seqlens_kv, + const Tensor &block_table, + const Tensor &weights, + int64_t max_q_len, + int64_t max_kv_len, + int64_t max_context_len); + +void select_prefill_topk_block_indices_( + Tensor topk_indices, + const Tensor &logits, + const Tensor &cu_seqlen_ks, + const Tensor &cu_seqlen_ke); + +void select_decode_topk_block_indices_( + Tensor topk_indices, + const Tensor &logits, + const Tensor &seq_lens); + +void map_prefill_request_block_indices_( + Tensor output, + const Tensor &req_id, + const Tensor &block_table, + const Tensor &token_indices, + int64_t block_size, + bool has_prefill_workspace = false, + std::optional prefill_workspace_request_ids = std::nullopt, + std::optional prefill_workspace_starts = std::nullopt); + +void map_decode_request_block_indices_( + Tensor output, + const Tensor &req_id, + const Tensor &block_table, + const Tensor &token_indices, + int64_t block_size); + +void topk_indices_context_lens_(Tensor topk_lens, const Tensor &indices); + +void sparse_flash_mla_( + Tensor output, + const Tensor &query, + const Tensor &kv_cache, + const Tensor &indices, + const Tensor &topk_lens, + float scale, + std::optional attn_sink = std::nullopt); + +} // namespace infinicore::op 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/fused_rotary_embedding.hpp b/include/infinicore/ops/fused_rotary_embedding.hpp new file mode 100644 index 000000000..5a5a4afba --- /dev/null +++ b/include/infinicore/ops/fused_rotary_embedding.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../tensor.hpp" + +namespace infinicore::op { + +// Apply rotary embedding to query and key in one vendor-kernel launch. This +// mirrors vLLM's rotary_embedding(positions, query, key, ...) execution path. +void fused_rotary_embedding_(Tensor query, + Tensor key, + const Tensor &positions, + int64_t head_size, + const Tensor &cos_sin_cache, + bool is_neox); + +} // 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/include/infiniop.h b/include/infiniop.h index 95bf75a0d..ddf8c47cb 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/gelu.h" #include "infiniop/ops/gelutanh.h" @@ -126,6 +130,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/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/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/python/infinicore/__init__.py b/python/infinicore/__init__.py index dadea41fe..284d087f2 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -35,6 +35,7 @@ double, dtype, float, + float8, float16, float32, float64, @@ -84,6 +85,17 @@ from infinicore.ops.diff import diff from infinicore.ops.digamma import digamma from infinicore.ops.dist import dist +from infinicore.ops.dsa import ( + compute_block_sparse_mqa_logits_, + fused_deepseek_v2_indexer_postprocess_, + indexer_k_cache_, + map_decode_request_block_indices_, + map_prefill_request_block_indices_, + select_decode_topk_block_indices_, + select_prefill_topk_block_indices_, + sparse_flash_mla_, + topk_indices_context_lens_, +) from infinicore.ops.dynamic_scaled_int8_quant import dynamic_scaled_int8_quant from infinicore.ops.equal import equal from infinicore.ops.flipud import flipud @@ -199,6 +211,7 @@ "double", "float", "float16", + "float8", "float32", "float64", "half", @@ -211,6 +224,15 @@ "short", "uint8", # Operations. + "compute_block_sparse_mqa_logits_", + "fused_deepseek_v2_indexer_postprocess_", + "indexer_k_cache_", + "map_decode_request_block_indices_", + "map_prefill_request_block_indices_", + "select_decode_topk_block_indices_", + "select_prefill_topk_block_indices_", + "sparse_flash_mla_", + "topk_indices_context_lens_", "addcmul", "atanh", "binary_cross_entropy_with_logits", 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/include/infiniccl.h b/python/infinicore/include/infiniccl.h new file mode 100644 index 000000000..b338d85f5 --- /dev/null +++ b/python/infinicore/include/infiniccl.h @@ -0,0 +1,75 @@ +#ifndef __INFINICCL_API_H__ +#define __INFINICCL_API_H__ + +#include "infinirt.h" + +typedef enum { + INFINICCL_SUM = 0, + INFINICCL_PROD = 1, + INFINICCL_MAX = 2, + INFINICCL_MIN = 3, + INFINICCL_AVG = 4, +} infinicclReduceOp_t; + +struct InfinicclComm; + +typedef struct InfinicclComm *infinicclComm_t; + +__INFINI_C __export infiniStatus_t infinicclCommInitAll( + infiniDevice_t device_type, + infinicclComm_t *comms, + int ndevice, + const int *device_ids); + +__INFINI_C __export infiniStatus_t infinicclCommDestroy(infinicclComm_t comm); + +__INFINI_C __export infiniStatus_t infinicclGroupStart(infinicclComm_t comm); + +__INFINI_C __export infiniStatus_t infinicclGroupEnd(infinicclComm_t comm); + +__INFINI_C __export infiniStatus_t infinicclAllReduce( + void *sendbuf, + void *recvbuf, + size_t count, + infiniDtype_t dataype, + infinicclReduceOp_t op, + infinicclComm_t comm, + infinirtStream_t stream); + +__INFINI_C __export infiniStatus_t infinicclAllGather( + void *sendbuf, + void *recvbuf, + size_t send_count, + infiniDtype_t dataype, + infinicclComm_t comm, + infinirtStream_t stream); + +__INFINI_C __export infiniStatus_t infinicclAllGatherV( + void *sendbuf, + void *recvbuf, + const size_t *recv_counts, + int nranks, + infiniDtype_t dataype, + infinicclComm_t comm, + infinirtStream_t stream); + +__INFINI_C __export infiniStatus_t infinicclReduceScatter( + void *sendbuf, + void *recvbuf, + size_t recv_count, + infiniDtype_t dataype, + infinicclReduceOp_t op, + infinicclComm_t comm, + infinirtStream_t stream); + +__INFINI_C __export infiniStatus_t infinicclReduceScatterV( + void *sendbuf, + void *recvbuf, + const size_t *send_counts, + int nranks, + infiniDtype_t dataype, + infinicclReduceOp_t op, + infinicclComm_t comm, + infinirtStream_t stream); + +#endif diff --git a/python/infinicore/include/infinicore.h b/python/infinicore/include/infinicore.h new file mode 100644 index 000000000..849c62c14 --- /dev/null +++ b/python/infinicore/include/infinicore.h @@ -0,0 +1,77 @@ +#ifndef __INFINICORE_API_H__ +#define __INFINICORE_API_H__ + +#if defined(_WIN32) +#define __export __declspec(dllexport) +#elif defined(__GNUC__) && ((__GNUC__ >= 4) || (__GNUC__ == 3 && __GNUC_MINOR__ >= 3)) +#define __export __attribute__((visibility("default"))) +#else +#define __export +#endif + +#ifdef __cplusplus +#define __INFINI_C extern "C" +#include +#else +#define __INFINI_C +#include +#endif + +typedef enum { + // Success + INFINI_STATUS_SUCCESS = 0, + // General Errors + INFINI_STATUS_INTERNAL_ERROR = 1, + INFINI_STATUS_NOT_IMPLEMENTED = 2, + INFINI_STATUS_BAD_PARAM = 3, + INFINI_STATUS_NULL_POINTER = 4, + INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED = 5, + INFINI_STATUS_DEVICE_NOT_FOUND = 6, + INFINI_STATUS_DEVICE_NOT_INITIALIZED = 7, + INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED = 8, + // Op Errors + INFINI_STATUS_BAD_TENSOR_DTYPE = 10, + INFINI_STATUS_BAD_TENSOR_SHAPE = 11, + INFINI_STATUS_BAD_TENSOR_STRIDES = 12, + INFINI_STATUS_INSUFFICIENT_WORKSPACE = 13, +} infiniStatus_t; + +typedef enum { + INFINI_DEVICE_CPU = 0, + INFINI_DEVICE_NVIDIA = 1, + INFINI_DEVICE_CAMBRICON = 2, + INFINI_DEVICE_ASCEND = 3, + INFINI_DEVICE_METAX = 4, + INFINI_DEVICE_MOORE = 5, + INFINI_DEVICE_ILUVATAR = 6, + INFINI_DEVICE_KUNLUN = 7, + INFINI_DEVICE_HYGON = 8, + INFINI_DEVICE_QY = 9, + INFINI_DEVICE_ALI = 10, + INFINI_DEVICE_TYPE_COUNT +} infiniDevice_t; + +typedef enum { + INFINI_DTYPE_INVALID = 0, + INFINI_DTYPE_BYTE = 1, + INFINI_DTYPE_BOOL = 2, + INFINI_DTYPE_I8 = 3, + INFINI_DTYPE_I16 = 4, + INFINI_DTYPE_I32 = 5, + INFINI_DTYPE_I64 = 6, + INFINI_DTYPE_U8 = 7, + INFINI_DTYPE_U16 = 8, + INFINI_DTYPE_U32 = 9, + INFINI_DTYPE_U64 = 10, + INFINI_DTYPE_F8 = 11, + INFINI_DTYPE_F16 = 12, + INFINI_DTYPE_F32 = 13, + INFINI_DTYPE_F64 = 14, + INFINI_DTYPE_C16 = 15, + INFINI_DTYPE_C32 = 16, + INFINI_DTYPE_C64 = 17, + INFINI_DTYPE_C128 = 18, + INFINI_DTYPE_BF16 = 19, +} infiniDtype_t; + +#endif // __INFINICORE_API_H__ diff --git a/python/infinicore/include/infinicore.hpp b/python/infinicore/include/infinicore.hpp new file mode 100644 index 000000000..ddd2f3d75 --- /dev/null +++ b/python/infinicore/include/infinicore.hpp @@ -0,0 +1,8 @@ +#pragma once + +#include "infinicore/device_event.hpp" +#include "infinicore/io.hpp" +#include "infinicore/nn.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/quantization.hpp" +#include "infinicore/tensor.hpp" diff --git a/python/infinicore/include/infinicore/adaptor/aten_adaptor.hpp b/python/infinicore/include/infinicore/adaptor/aten_adaptor.hpp new file mode 100644 index 000000000..ddb8d1c58 --- /dev/null +++ b/python/infinicore/include/infinicore/adaptor/aten_adaptor.hpp @@ -0,0 +1,86 @@ +#ifdef ENABLE_ATEN +#pragma once + +#include "../context/context.hpp" +#include "../tensor.hpp" + +#include + +#if defined(ENABLE_HYGON_API) +#include +#include +#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#include +#include +#include +#endif + +#if defined(ENABLE_MOORE_API) +#include +#include +#include +#endif + +namespace infinicore::adaptor { +inline at::ScalarType to_at_dtype(DataType dtype) { + switch (dtype) { + case DataType::F32: + return at::kFloat; + case DataType::F16: + return at::kHalf; + case DataType::BF16: + return at::kBFloat16; + case DataType::I8: + return at::kChar; + case DataType::U8: + return at::kByte; + case DataType::I32: + return at::kInt; + case DataType::I64: + return at::kLong; + default: + throw std::runtime_error("Unsupported dtype for ATen"); + } +} + +inline at::Device to_at_device(const Device &device) { + // PyTorch ATen only exposes standard device types (e.g. kCPU/kCUDA). + // Treat CUDA-compatible vendor devices as CUDA devices for ATen tensor interoperability. + if (device.getType() == Device::Type::NVIDIA + || device.getType() == Device::Type::METAX + || device.getType() == Device::Type::QY + || device.getType() == Device::Type::ILUVATAR + || device.getType() == Device::Type::HYGON) { + return at::Device(at::kCUDA, device.getIndex()); + } else if (device.getType() == Device::Type::CPU) { + return at::Device(at::kCPU); + } +#if defined(ENABLE_MOORE_API) + else if (device.getType() == Device::Type::MOORE) { + return at::Device(at::DeviceType::PrivateUse1, device.getIndex()); + } +#endif + else { + throw std::runtime_error("Unsupported device type for ATen"); + } +} + +at::Tensor to_aten_tensor(const infinicore::Tensor &t); + +#if defined(ENABLE_HYGON_API) +c10::hip::HIPStream get_hip_stream(); +#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +c10::cuda::CUDAStream get_cuda_stream(); +#endif + +#if defined(ENABLE_ILUVATAR_API) +void set_aten_stream_to_infinicore(); +#endif + +#if defined(ENABLE_MOORE_API) +c10::musa::MUSAStream get_musa_stream(); +#endif + +} // namespace infinicore::adaptor + +#endif // ENABLE_ATEN diff --git a/python/infinicore/include/infinicore/adaptor/flash_attention_adaptor.hpp b/python/infinicore/include/infinicore/adaptor/flash_attention_adaptor.hpp new file mode 100644 index 000000000..c5bf14858 --- /dev/null +++ b/python/infinicore/include/infinicore/adaptor/flash_attention_adaptor.hpp @@ -0,0 +1,139 @@ +#ifdef ENABLE_FLASH_ATTN +#pragma once +#include "aten_adaptor.hpp" + +// NVIDIA flash-attn-nvidia.so uses namespace flash. The pip/MetaX flash_attn_2_cuda extension +// exports the same entry points at global scope (no namespace), matching FLASH_NAMESPACE builds +// where the namespace is empty. +#if !defined(ENABLE_METAX_API) +namespace flash { +#endif +std::vector +mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) + const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) + const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) + std::optional &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) + std::optional &alibi_slopes_, // num_heads or batch_size x num_heads + const float p_dropout, + const float softmax_scale, + bool is_causal, + int window_size_left, + int window_size_right, + const float softcap, + const bool return_softmax, + std::optional gen_ +#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) + // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. + , + std::optional &flash_attn_mars_ext_ +#endif +); + +std::vector +mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i + const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. + const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. + std::optional &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i + const at::Tensor &cu_seqlens_q, // b+1 + const at::Tensor &cu_seqlens_k, // b+1 + std::optional &seqused_k, // b. If given, only this many elements of each batch element's keys are used. + std::optional &leftpad_k_, // batch_size + std::optional &block_table_, // batch_size x max_num_blocks_per_seq + std::optional &alibi_slopes_, // num_heads or b x num_heads + int max_seqlen_q, + const int max_seqlen_k, + const float p_dropout, + const float softmax_scale, + const bool zero_tensors, + bool is_causal, + int window_size_left, + int window_size_right, + const float softcap, + const bool return_softmax, + std::optional gen_ +#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) + // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. + , + std::optional &flash_attn_mars_ext_ +#endif +); + +std::vector +mha_bwd(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8) + const at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size + const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x head_size + const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x head_size + const at::Tensor &out, // batch_size x seqlen_q x num_heads x head_size + const at::Tensor &softmax_lse, // b x h x seqlen_q + std::optional &dq_, // batch_size x seqlen_q x num_heads x head_size + std::optional &dk_, // batch_size x seqlen_k x num_heads_k x head_size + std::optional &dv_, // batch_size x seqlen_k x num_heads_k x head_size + std::optional &alibi_slopes_, // num_heads or batch_size x num_heads + const float p_dropout, // probability to drop + const float softmax_scale, + const bool is_causal, + int window_size_left, + int window_size_right, + const float softcap, + const bool deterministic, + std::optional gen_, + std::optional &rng_state); + +std::vector +mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads, x head_size + const at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i + const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i + const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i + const at::Tensor &out, // total_q x num_heads x head_size + const at::Tensor &softmax_lse, // h x total_q, softmax logsumexp + std::optional &dq_, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i + std::optional &dk_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i + std::optional &dv_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i + const at::Tensor &cu_seqlens_q, // b+1 + const at::Tensor &cu_seqlens_k, // b+1 + std::optional &alibi_slopes_, // num_heads or b x num_heads + const int max_seqlen_q, + const int max_seqlen_k, // max sequence length to choose the kernel + const float p_dropout, // probability to drop + const float softmax_scale, + const bool zero_tensors, + const bool is_causal, + int window_size_left, + int window_size_right, + const float softcap, + const bool deterministic, + std::optional gen_, + std::optional &rng_state); + +std::vector +mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size + const at::Tensor &kcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. + const at::Tensor &vcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. + std::optional &k_, // batch_size x seqlen_knew x num_heads_k x head_size + std::optional &v_, // batch_size x seqlen_knew x num_heads_k x head_size + std::optional &seqlens_k_, // batch_size + std::optional &rotary_cos_, // seqlen_ro x (rotary_dim / 2) + std::optional &rotary_sin_, // seqlen_ro x (rotary_dim / 2) + std::optional &cache_batch_idx_, // indices to index into the KV cache + std::optional &leftpad_k_, // batch_size + std::optional &block_table_, // batch_size x max_num_blocks_per_seq + std::optional &alibi_slopes_, // num_heads or batch_size x num_heads + std::optional &out_, // batch_size x seqlen_q x num_heads x head_size + const float softmax_scale, + bool is_causal, + int window_size_left, + int window_size_right, + const float softcap, + bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2 + int num_splits +#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) + // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. + , + std::optional &flash_attn_mars_ext_ +#endif +); + +#if !defined(ENABLE_METAX_API) +} // namespace flash +#endif +#endif // ENABLE_FLASH_ATTN diff --git a/python/infinicore/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp b/python/infinicore/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp new file mode 100644 index 000000000..067c49ca3 --- /dev/null +++ b/python/infinicore/include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp @@ -0,0 +1,85 @@ +#pragma once + +#ifdef ENABLE_ATEN +#include +#include +#include +#include + +namespace infinicore::adaptor::vllm_iluvatar { + +bool available(); +bool rotary_embedding_available(); +bool dynamic_scaled_int8_quant_available(); +bool concat_mla_q_available(); +bool concat_and_cache_mla_available(); +bool concat_and_cache_mla_int8_available(); +bool paged_attention_mla_available(); +bool topk_softmax_available(); +bool topk_sigmoid_available(); +bool grouped_topk_available(); +bool scaled_mm_w4a8_available(); +bool scaled_mm_w8a8_available(); +bool w4a8_group_gemm_available(); +bool w8a8_group_gemm_available(); +bool w16a16_group_gemm_available(); +bool argsort_bincount_with_inv_pos_available(); +bool expand_moe_input_with_inv_pos_available(); +bool silu_and_mul_quant_available(); +bool moe_sum_vllm_available(); +bool fused_deepseek_v2_indexer_postprocess_available(); +bool indexer_k_cache_available(); +bool indexer_k_quant_and_cache_available(); +bool compute_block_sparse_mqa_logits_available(); +bool select_prefill_topk_block_indices_available(); +bool select_decode_topk_block_indices_available(); +bool map_prefill_request_block_indices_available(); +bool map_decode_request_block_indices_available(); +bool sparse_flash_mla_available(); +bool topk_indices_context_lens_available(); +void fused_add_rms_norm(at::Tensor &input, at::Tensor &residual, at::Tensor &weight, float epsilon); +void rotary_embedding(at::Tensor &positions, + at::Tensor &query, + std::optional key, + int64_t head_size, + at::Tensor &cos_sin_cache, + bool is_neox); +void dynamic_scaled_int8_quant(at::Tensor &output, at::Tensor &input_scales, const at::Tensor &input); +void concat_mla_q(at::Tensor &ql_nope, at::Tensor &q_pe, at::Tensor &q_out); +void concat_and_cache_mla(at::Tensor &kv_c, at::Tensor &k_pe, at::Tensor &kv_cache, at::Tensor &slot_mapping, const std::string &kv_cache_dtype, at::Tensor &scale); +void concat_and_cache_mla_int8(at::Tensor &kv_c_int8, at::Tensor &kv_c_scale, at::Tensor &k_pe_int8, at::Tensor &k_pe_scale, at::Tensor &kv_cache, at::Tensor &kv_cache_scale, at::Tensor &slot_mapping); +void paged_attention_mla(at::Tensor &output, + at::Tensor &query, + at::Tensor &kv_cache, + double scale, + at::Tensor &block_tables, + at::Tensor &context_lens, + int64_t max_context_len, + bool use_cuda_graph, + at::Tensor &softmax_lse); +void topk_softmax(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional correction_bias); +void topk_sigmoid(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional correction_bias); +void grouped_topk(at::Tensor &topk_weights, at::Tensor &topk_ids, const at::Tensor &scores, std::optional bias, int64_t num_expert_group, int64_t topk_group, const std::string &scoring_func, bool renormalize); +void scaled_mm_w4a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional bias, bool trans_weight); +void scaled_mm_w8a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional bias, bool trans_weight); +void w4a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode); +void w8a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode); +void w16a16_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &tokens_per_experts, std::optional sorted_token_ids, std::optional bias, bool trans_weight, bool is_decode); +void argsort_bincount_with_inv_pos(const at::Tensor &topk_ids, at::Tensor &tokens_per_experts, at::Tensor &sorted_indices, at::Tensor &inv_pos, int64_t num_experts); +void expand_moe_input_with_inv_pos(at::Tensor &expand_states, std::optional expand_scales, const at::Tensor &hidden_states, const at::Tensor &inv_pos, int64_t top_k, int64_t group_size, int64_t format); +void silu_and_mul_quant(at::Tensor &output, std::optional output_scale, const at::Tensor &input, int64_t format); +void moe_sum_vllm(at::Tensor &output, const at::Tensor &input, std::optional topk_weights, std::optional extra_residual, double routed_scale, double residual_scale); +void select_last_token_hidden(at::Tensor &output, at::Tensor &indices, const at::Tensor &hidden_states, const at::Tensor &input_offsets); +void fused_deepseek_v2_indexer_postprocess(at::Tensor &q_out, at::Tensor &k_out, at::Tensor &weights_out, at::Tensor &kv_cache, const at::Tensor &slot_mapping, const at::Tensor &q, const at::Tensor &kw, const at::Tensor &norm_weight, const at::Tensor &norm_bias, const at::Tensor &positions, const at::Tensor &cos_sin_cache, int64_t num_cache_tokens, bool is_neox, double eps, double weights_scale); +void indexer_k_cache(const at::Tensor &k, at::Tensor &kv_cache, const at::Tensor &slot_mapping); +void indexer_k_quant_and_cache(at::Tensor &k, at::Tensor &kv_cache, at::Tensor &slot_mapping, int64_t quant_block_size, const std::string &scale_fmt); +void compute_block_sparse_mqa_logits(const at::Tensor &q, const at::Tensor &kv_cache, const at::Tensor &cu_seqlens_q, const at::Tensor &cu_seqlens_kv, const at::Tensor &block_table, const at::Tensor &weights, at::Tensor &logits, int64_t max_q_len, int64_t max_kv_len, int64_t max_context_len); +void select_prefill_topk_block_indices(const at::Tensor &logits, const at::Tensor &cu_seqlen_ks, const at::Tensor &cu_seqlen_ke, at::Tensor &topk_indices); +void select_decode_topk_block_indices(const at::Tensor &logits, const at::Tensor &seq_lens, at::Tensor &topk_indices); +void map_prefill_request_block_indices(at::Tensor &output, const at::Tensor &req_id, const at::Tensor &block_table, const at::Tensor &token_indices, int64_t block_size, bool has_prefill_workspace, std::optional prefill_workspace_request_ids, std::optional prefill_workspace_starts); +void map_decode_request_block_indices(at::Tensor &output, const at::Tensor &req_id, const at::Tensor &block_table, const at::Tensor &token_indices, int64_t block_size); +void topk_indices_context_lens(at::Tensor &topk_lens, const at::Tensor &indices); +void sparse_flash_mla(at::Tensor &output, at::Tensor &query, at::Tensor &kv_cache, at::Tensor &indices, at::Tensor &topk_lens, float scale, std::optional attn_sink); + +} // namespace infinicore::adaptor::vllm_iluvatar +#endif // ENABLE_ATEN diff --git a/python/infinicore/include/infinicore/analyzer.hpp b/python/infinicore/include/infinicore/analyzer.hpp new file mode 100644 index 000000000..81e9ab1a9 --- /dev/null +++ b/python/infinicore/include/infinicore/analyzer.hpp @@ -0,0 +1,11 @@ +#pragma once + +// Convenience header — includes all analyzer components. + +#include "analyzer/intent_generator.hpp" +#include "analyzer/mutual_awareness_analyzer.hpp" +#include "analyzer/op_trace.hpp" +#include "analyzer/op_type.hpp" +#include "analyzer/optimization_intent.hpp" +#include "analyzer/phase_detector.hpp" +#include "analyzer/resource_sensor.hpp" diff --git a/python/infinicore/include/infinicore/analyzer/intent_generator.hpp b/python/infinicore/include/infinicore/analyzer/intent_generator.hpp new file mode 100644 index 000000000..b166e4c1e --- /dev/null +++ b/python/infinicore/include/infinicore/analyzer/intent_generator.hpp @@ -0,0 +1,268 @@ +#pragma once + +#include "op_trace.hpp" +#include "optimization_intent.hpp" + +#include +#include + +namespace infinicore::analyzer { + +/// IntentGenerator — the core "mutual awareness" logic. +/// +/// This is where task demand and resource supply are jointly +/// analyzed to produce an OptimizationIntent. It implements +/// the key insight: the same task phase has different optimization +/// needs under different resource conditions, and the same resource +/// state has different supply value under different task phases. +class IntentGenerator { +public: + IntentGenerator() = default; + + /// Generate the global semantic intent from phase detection + /// result and op trace window. + GlobalSemanticIntent generateGlobal( + PhaseType phase, + const std::vector &window, + const std::vector &device_intents) const { + + GlobalSemanticIntent intent; + intent.current_phase = phase; + intent.timestamp_ns = OpTraceEntry::now(); + + if (!window.empty()) { + intent.op_window_start = 0; + intent.op_window_end = static_cast(window.size()); + } + + // --- Compute intensity estimation --- + intent.compute_intensity = estimateComputeIntensity(phase, window); + + // --- Determine primary bottleneck (mutual awareness) --- + intent.primary_bottleneck = determineGlobalBottleneck(phase, device_intents); + + // --- Set optimization goal based on phase + bottleneck --- + intent.goal = determineGoal(phase, intent.primary_bottleneck); + + // --- Generate strategy hints --- + intent.strategy = generateStrategy(phase, intent.primary_bottleneck, device_intents); + + // --- Confidence --- + intent.confidence = computeConfidence(phase, window); + + return intent; + } + + /// Build the complete two-layer OptimizationIntent. + OptimizationIntent generate( + PhaseType phase, + const std::vector &window, + const std::vector &device_intents) const { + + OptimizationIntent result; + result.global = generateGlobal(phase, window, device_intents); + result.per_device = device_intents; + return result; + } + +private: + /// Estimate compute intensity (higher = more compute-heavy). + /// Uses a simple heuristic based on op type composition. + float estimateComputeIntensity( + PhaseType phase, + const std::vector &window) const { + + if (window.empty()) { + return 0.0f; + } + + size_t heavy_compute_ops = 0; + for (auto &e : window) { + if (isGemmMlpOp(e.op_type) || isAttentionOp(e.op_type)) { + heavy_compute_ops++; + } + } + return static_cast(heavy_compute_ops) / static_cast(window.size()); + } + + /// Determine global bottleneck by jointly considering phase and + /// per-device resource state (the core mutual awareness logic). + BottleneckType determineGlobalBottleneck( + PhaseType phase, + const std::vector &device_intents) const { + + bool any_memory_bound = false; + bool any_compute_bound = false; + bool any_bandwidth_bound = false; + bool any_communication_bound = false; + for (auto &d : device_intents) { + any_memory_bound = any_memory_bound || d.local_bottleneck == BottleneckType::MEMORY_BOUND; + any_compute_bound = any_compute_bound || d.local_bottleneck == BottleneckType::COMPUTE_BOUND; + any_bandwidth_bound = any_bandwidth_bound || d.local_bottleneck == BottleneckType::BANDWIDTH_BOUND; + any_communication_bound = any_communication_bound || d.local_bottleneck == BottleneckType::COMMUNICATION_BOUND; + } + + // --- Mutual awareness logic --- + // The same resource state has different "supply value" depending on phase: + + if (any_memory_bound) { + return BottleneckType::MEMORY_BOUND; + } + + if (phase == PhaseType::COMMUNICATION || any_communication_bound) { + return BottleneckType::COMMUNICATION_BOUND; + } + + switch (phase) { + case PhaseType::ATTENTION_DENSE: + case PhaseType::PREFILL: + // Attention/prefill is dominated by memory movement and KV access, + // so phase semantics should win unless memory/communication already + // forced an earlier return above. + if (any_bandwidth_bound) { + return BottleneckType::BANDWIDTH_BOUND; + } + return BottleneckType::BANDWIDTH_BOUND; + + case PhaseType::GEMM_MLP_DENSE: + if (any_compute_bound) { + return BottleneckType::COMPUTE_BOUND; + } + if (any_bandwidth_bound) { + return BottleneckType::BANDWIDTH_BOUND; + } + return BottleneckType::COMPUTE_BOUND; + + case PhaseType::DECODE: + if (any_bandwidth_bound) { + return BottleneckType::BANDWIDTH_BOUND; + } + if (any_compute_bound) { + return BottleneckType::COMPUTE_BOUND; + } + return BottleneckType::BANDWIDTH_BOUND; + + case PhaseType::KV_CACHE: + if (any_bandwidth_bound) { + return BottleneckType::BANDWIDTH_BOUND; + } + return BottleneckType::MEMORY_BOUND; + + default: + if (any_bandwidth_bound) { + return BottleneckType::BANDWIDTH_BOUND; + } + if (any_compute_bound) { + return BottleneckType::COMPUTE_BOUND; + } + return BottleneckType::BALANCED; + } + } + + /// Determine optimization goal based on phase and bottleneck. + OptimizationGoal determineGoal( + PhaseType phase, + BottleneckType bottleneck) const { + + // Under memory pressure, prioritize memory safety + if (bottleneck == BottleneckType::MEMORY_BOUND) { + return OptimizationGoal::MEMORY_SAFE; + } + + if (bottleneck == BottleneckType::COMMUNICATION_BOUND) { + return OptimizationGoal::STABILITY_FIRST; + } + + switch (phase) { + case PhaseType::DECODE: + // Decode latency is user-facing → latency first + return OptimizationGoal::LATENCY_FIRST; + + case PhaseType::PREFILL: + // Prefill processes a full prompt → throughput first + return OptimizationGoal::THROUGHPUT_FIRST; + + case PhaseType::ATTENTION_DENSE: + return OptimizationGoal::LATENCY_FIRST; + + case PhaseType::GEMM_MLP_DENSE: + return OptimizationGoal::THROUGHPUT_FIRST; + + default: + return OptimizationGoal::LATENCY_FIRST; + } + } + + /// Generate strategy hints from phase + bottleneck + resources. + StrategyHint generateStrategy( + PhaseType phase, + BottleneckType bottleneck, + const std::vector &device_intents) const { + + StrategyHint hint; + + // Fusion is beneficial for bandwidth-bound phases (reduce memory traffic) + hint.prefer_fused_ops = (bottleneck == BottleneckType::BANDWIDTH_BOUND) + || phase == PhaseType::DECODE; + + // In-place when memory is tight + hint.prefer_in_place = (bottleneck == BottleneckType::MEMORY_BOUND); + + // Recomputation (activation checkpointing) when memory is critical + bool extreme_memory = false; + for (auto &d : device_intents) { + if (d.memory_usage_ratio >= 0.95f) { + extreme_memory = true; + break; + } + } + hint.prefer_recomputation = extreme_memory; + + // Async comm overlap for multi-device and communication phases + hint.prefer_async_comm = (device_intents.size() > 1) + && (phase == PhaseType::GEMM_MLP_DENSE + || phase == PhaseType::COMMUNICATION); + + return hint; + } + + /// Compute confidence based on how clear the phase signal is. + float computeConfidence( + PhaseType phase, + const std::vector &window) const { + + if (window.empty() || phase == PhaseType::UNKNOWN) { + return 0.0f; + } + + // Count how many ops in the window match the detected phase + size_t matching = 0; + for (auto &e : window) { + bool match = false; + switch (phase) { + case PhaseType::ATTENTION_DENSE: + case PhaseType::PREFILL: + match = isAttentionOp(e.op_type); + break; + case PhaseType::GEMM_MLP_DENSE: + match = isGemmMlpOp(e.op_type) || isActivationOp(e.op_type); + break; + case PhaseType::KV_CACHE: + match = isKvCacheOp(e.op_type); + break; + case PhaseType::DECODE: + match = isAttentionOp(e.op_type) || isGemmMlpOp(e.op_type); + break; + default: + break; + } + if (match) { + matching++; + } + } + + return static_cast(matching) / static_cast(window.size()); + } +}; + +} // namespace infinicore::analyzer diff --git a/python/infinicore/include/infinicore/analyzer/mutual_awareness_analyzer.hpp b/python/infinicore/include/infinicore/analyzer/mutual_awareness_analyzer.hpp new file mode 100644 index 000000000..d373da694 --- /dev/null +++ b/python/infinicore/include/infinicore/analyzer/mutual_awareness_analyzer.hpp @@ -0,0 +1,118 @@ +#pragma once + +#include "intent_generator.hpp" +#include "op_trace.hpp" +#include "optimization_intent.hpp" +#include "phase_detector.hpp" +#include "resource_sensor.hpp" + +#include +#include + +namespace infinicore::analyzer { + +/// MutualAwarenessAnalyzer — the top-level facade for the +/// hardware-task mutual awareness requirements analysis module. +/// +/// This is the primary entry point exposed to external frameworks +/// (e.g., InfiniLM) via C++ function calls. It orchestrates: +/// 1. Op trace collection (via OpTraceRing) +/// 2. Phase detection (via PhaseDetector) +/// 3. Resource sensing (via ResourceSensor) +/// 4. Intent generation (via IntentGenerator) +/// +/// Usage: +/// auto& analyzer = MutualAwarenessAnalyzer::instance(); +/// // ... ops execute and get traced automatically ... +/// auto intent = analyzer.analyze(); // Produces OptimizationIntent +/// +/// Thread safety: analyze() is safe to call from any thread. +/// The analyzer reads a snapshot of the op trace ring. +class MutualAwarenessAnalyzer { +public: + /// Get the singleton instance. + static MutualAwarenessAnalyzer &instance(); + + // Non-copyable, non-movable + MutualAwarenessAnalyzer(const MutualAwarenessAnalyzer &) = delete; + MutualAwarenessAnalyzer &operator=(const MutualAwarenessAnalyzer &) = delete; + + /// Main analysis entry point. + /// Analyzes the current op trace window + resource state + /// and returns a complete OptimizationIntent. + /// + /// This is the function InfiniLM should call. + /// Latency: expected < 1ms for MVP rule-based analysis. + OptimizationIntent analyze(); + + /// Analyze with explicitly provided memory stats per device. + /// Use this when the caller can provide resource info directly. + OptimizationIntent analyze(const std::vector> &device_stats); + + /// Analyze with explicitly provided device resource snapshots. + /// This is the richer input path used by demand-analysis-oriented callers. + OptimizationIntent analyze(const std::vector &device_snapshots); + + /// Get the current phase without generating full intent. + /// Lightweight query for simple use cases. + PhaseType getCurrentPhase() const; + + /// Get the current optimization goal derived from the + /// latest analyzer result. + OptimizationGoal getCurrentOptimizationGoal() const; + + /// Get the most recent OptimizationIntent (cached from last analyze()). + const OptimizationIntent &lastIntent() const; + + /// Access the underlying components for configuration. + PhaseDetector &phaseDetector() { return phase_detector_; } + ResourceSensor &resourceSensor() { return resource_sensor_; } + OpTraceRing &opTrace() { return getGlobalOpTrace(); } + + /// Enable / disable the analyzer. + /// When disabled, analyze() returns a default intent and + /// op trace recording is skipped. + void setEnabled(bool enabled) { enabled_ = enabled; } + bool isEnabled() const { return enabled_; } + + /// Graph recording support: when graph recording stops, + /// analyze the recorded op sequence once and cache the result. + /// Subsequent calls return the cached intent without re-analysis. + void onGraphRecordingStop(); + void clearGraphCache(); + +private: + MutualAwarenessAnalyzer(); + + PhaseDetector phase_detector_; + ResourceSensor resource_sensor_; + IntentGenerator intent_generator_; + + OptimizationIntent last_intent_; + mutable std::mutex mutex_; + + bool enabled_ = true; + + // Graph recording cache + bool graph_intent_cached_ = false; + OptimizationIntent graph_cached_intent_; +}; + +// ============================================================ +// C-style API for external framework integration (e.g., InfiniLM) +// ============================================================ + +/// Analyze current state and return an OptimizationIntent. +/// This is the simplest API for external frameworks to call. +OptimizationIntent analyzeCurrentState(); + +/// Get the current detected phase. +PhaseType getCurrentPhase(); + +/// Get the current optimization goal. +OptimizationGoal getCurrentOptimizationGoal(); + +/// Enable / disable the mutual awareness analyzer. +void setAnalyzerEnabled(bool enabled); + +} // namespace infinicore::analyzer diff --git a/python/infinicore/include/infinicore/analyzer/op_trace.hpp b/python/infinicore/include/infinicore/analyzer/op_trace.hpp new file mode 100644 index 000000000..0b3018c28 --- /dev/null +++ b/python/infinicore/include/infinicore/analyzer/op_trace.hpp @@ -0,0 +1,160 @@ +#pragma once + +#include "op_type.hpp" + +#include +#include +#include +#include +#include +#include + +namespace infinicore::analyzer { + +// ============================================================ +// OpTraceEntry — a single recorded op invocation +// ============================================================ + +/// Compact record of one operator invocation for phase detection. +/// Designed to be small (~80 bytes) and cheap to fill. +struct OpTraceEntry { + OpType op_type = OpType::UNKNOWN; + + // Tensor shape summary (up to 4 dims for the primary input tensor) + static constexpr size_t MAX_DIMS = 4; + uint32_t ndim = 0; + uint32_t shape[MAX_DIMS] = {}; + + // Data type of the primary input (cast from infiniDtype_t) + uint8_t dtype = 0; + + // Device info + uint8_t device_type = 0; + int8_t device_id = -1; + + // Timestamp (nanoseconds since epoch, from steady_clock) + uint64_t timestamp_ns = 0; + + /// Fill shape from a shape vector. + void setShape(const size_t *dims, size_t n) { + ndim = static_cast(n > MAX_DIMS ? MAX_DIMS : n); + for (uint32_t i = 0; i < ndim; ++i) { + shape[i] = static_cast(dims[i]); + } + } + + /// Get current steady_clock timestamp in nanoseconds. + static uint64_t now() { + return static_cast( + std::chrono::steady_clock::now().time_since_epoch().count()); + } +}; + +// ============================================================ +// OpTraceRing — lock-free ring buffer for op trace entries +// ============================================================ + +/// A fixed-capacity ring buffer for OpTraceEntry. +/// Single-producer (op execution thread) friendly. +/// Reader can safely read a snapshot via getRecentEntries(). +/// +/// Thread safety: +/// - write() is safe to call from the single producer thread +/// (typical in InfiniCore where ops are dispatched on one thread). +/// - getRecentEntries() takes a snapshot and is safe to call from +/// any thread (may see a partially written entry at the boundary, +/// which is acceptable for heuristic phase detection). +class OpTraceRing { +public: + static constexpr size_t DEFAULT_CAPACITY = 256; + + explicit OpTraceRing(size_t capacity = DEFAULT_CAPACITY) + : capacity_(capacity), + entries_(capacity), + write_pos_(0), + total_count_(0) { + } + + /// Record a new op trace entry. + void write(const OpTraceEntry &entry) { + size_t pos = write_pos_.load(std::memory_order_relaxed); + entries_[pos % capacity_] = entry; + write_pos_.store(pos + 1, std::memory_order_release); + total_count_.fetch_add(1, std::memory_order_relaxed); + } + + /// Get the most recent N entries (ordered oldest to newest). + /// Returns fewer entries if the ring hasn't filled up yet. + std::vector getRecentEntries(size_t n) const { + size_t wp = write_pos_.load(std::memory_order_acquire); + size_t available = wp < capacity_ ? wp : capacity_; + size_t count = n < available ? n : available; + + std::vector result; + result.reserve(count); + + // Read from (wp - count) to (wp - 1) + for (size_t i = wp - count; i < wp; ++i) { + result.push_back(entries_[i % capacity_]); + } + return result; + } + + /// Get all valid entries in the ring (ordered oldest to newest). + std::vector getAllEntries() const { + return getRecentEntries(capacity_); + } + + /// Total number of ops traced since creation. + size_t totalCount() const { + return total_count_.load(std::memory_order_relaxed); + } + + /// Current number of valid entries in the ring. + size_t size() const { + size_t wp = write_pos_.load(std::memory_order_relaxed); + return wp < capacity_ ? wp : capacity_; + } + + /// Ring capacity. + size_t capacity() const { return capacity_; } + + /// Clear all entries. + void clear() { + write_pos_.store(0, std::memory_order_relaxed); + total_count_.store(0, std::memory_order_relaxed); + } + +private: + size_t capacity_; + std::vector entries_; + std::atomic write_pos_; + std::atomic total_count_; +}; + +// ============================================================ +// Global OpTrace singleton access +// ============================================================ + +/// Get the global OpTraceRing instance. +/// This is the primary entry point for recording op traces. +OpTraceRing &getGlobalOpTrace(); + +/// Record an op invocation to the global trace ring. +/// This is the function called from the INFINICORE_GRAPH_OP_RECORD_OR_RUN +/// macro hook (when ENABLE_MUTUAL_AWARENESS is defined). +inline void traceOp(OpType op_type, + const size_t *shape, size_t ndim, + uint8_t dtype, + uint8_t device_type, int8_t device_id) { + OpTraceEntry entry; + entry.op_type = op_type; + entry.setShape(shape, ndim); + entry.dtype = dtype; + entry.device_type = device_type; + entry.device_id = device_id; + entry.timestamp_ns = OpTraceEntry::now(); + getGlobalOpTrace().write(entry); +} + +} // namespace infinicore::analyzer diff --git a/python/infinicore/include/infinicore/analyzer/op_type.hpp b/python/infinicore/include/infinicore/analyzer/op_type.hpp new file mode 100644 index 000000000..cbb3ac102 --- /dev/null +++ b/python/infinicore/include/infinicore/analyzer/op_type.hpp @@ -0,0 +1,238 @@ +#pragma once + +#include +#include + +namespace infinicore::analyzer { + +/// Op types recognized by the analyzer for phase detection. +/// Not an exhaustive list of all infiniop ops — only the ones +/// relevant for phase classification in LLM inference. +enum class OpType : uint8_t { + UNKNOWN = 0, + + // --- Attention-related --- + ATTENTION, + FLASH_ATTENTION, + CAUSAL_SOFTMAX, + PAGED_ATTENTION, + PAGED_ATTENTION_PREFILL, + MHA_KVCACHE, + MHA_VARLEN, + SOFTMAX, + + // --- GEMM / MLP --- + GEMM, + LINEAR, + MATMUL, + INT8_GEMM, + SCALED_MM_I8, + + // --- Activation --- + SILU, + SILU_AND_MUL, + GELU, + SWIGLU, + RELU, + SIGMOID, + + // --- Norm --- + RMS_NORM, + ADD_RMS_NORM, + LAYER_NORM, + + // --- Embedding / Positional --- + EMBEDDING, + ROPE, + + // --- KV Cache --- + KV_CACHING, + PAGED_CACHING, + + // --- Elementwise / Reduce --- + ADD, + MUL, + SUB, + SUM, + RECIPROCAL, + + // --- Quantization --- + PER_TENSOR_QUANT_I8, + PER_TENSOR_DEQUANT_I8, + PER_CHANNEL_QUANT_I8, + DEQUANTIZE_AWQ, + DEQUANTIZE_GPTQ, + + // --- Sampling --- + RANDOM_SAMPLE, + TOPK, + TOPK_ROUTER, + TOPK_SOFTMAX, + + // --- Communication (future) --- + ALLREDUCE, + + // --- Misc --- + REARRANGE, + ONES, + ZEROS, + TAKE, + + OP_TYPE_COUNT, +}; + +/// Convert OpType to human-readable string. +inline const char *opTypeToString(OpType type) { + switch (type) { + case OpType::ATTENTION: + return "attention"; + case OpType::FLASH_ATTENTION: + return "flash_attention"; + case OpType::CAUSAL_SOFTMAX: + return "causal_softmax"; + case OpType::PAGED_ATTENTION: + return "paged_attention"; + case OpType::PAGED_ATTENTION_PREFILL: + return "paged_attention_prefill"; + case OpType::MHA_KVCACHE: + return "mha_kvcache"; + case OpType::MHA_VARLEN: + return "mha_varlen"; + case OpType::SOFTMAX: + return "softmax"; + case OpType::GEMM: + return "gemm"; + case OpType::LINEAR: + return "linear"; + case OpType::MATMUL: + return "matmul"; + case OpType::INT8_GEMM: + return "int8_gemm"; + case OpType::SCALED_MM_I8: + return "scaled_mm_i8"; + case OpType::SILU: + return "silu"; + case OpType::SILU_AND_MUL: + return "silu_and_mul"; + case OpType::GELU: + return "gelu"; + case OpType::SWIGLU: + return "swiglu"; + case OpType::RELU: + return "relu"; + case OpType::SIGMOID: + return "sigmoid"; + case OpType::RMS_NORM: + return "rms_norm"; + case OpType::ADD_RMS_NORM: + return "add_rms_norm"; + case OpType::LAYER_NORM: + return "layer_norm"; + case OpType::EMBEDDING: + return "embedding"; + case OpType::ROPE: + return "rope"; + case OpType::KV_CACHING: + return "kv_caching"; + case OpType::PAGED_CACHING: + return "paged_caching"; + case OpType::ADD: + return "add"; + case OpType::MUL: + return "mul"; + case OpType::SUB: + return "sub"; + case OpType::SUM: + return "sum"; + case OpType::RECIPROCAL: + return "reciprocal"; + case OpType::PER_TENSOR_QUANT_I8: + return "per_tensor_quant_i8"; + case OpType::PER_TENSOR_DEQUANT_I8: + return "per_tensor_dequant_i8"; + case OpType::PER_CHANNEL_QUANT_I8: + return "per_channel_quant_i8"; + case OpType::DEQUANTIZE_AWQ: + return "dequantize_awq"; + case OpType::DEQUANTIZE_GPTQ: + return "dequantize_gptq"; + case OpType::RANDOM_SAMPLE: + return "random_sample"; + case OpType::TOPK: + return "topk"; + case OpType::TOPK_ROUTER: + return "topk_router"; + case OpType::TOPK_SOFTMAX: + return "topk_softmax"; + case OpType::ALLREDUCE: + return "allreduce"; + case OpType::REARRANGE: + return "rearrange"; + case OpType::ONES: + return "ones"; + case OpType::ZEROS: + return "zeros"; + case OpType::TAKE: + return "take"; + default: + return "unknown"; + } +} + +/// Check if an op type belongs to the attention family. +inline bool isAttentionOp(OpType type) { + switch (type) { + case OpType::ATTENTION: + case OpType::FLASH_ATTENTION: + case OpType::CAUSAL_SOFTMAX: + case OpType::PAGED_ATTENTION: + case OpType::PAGED_ATTENTION_PREFILL: + case OpType::MHA_KVCACHE: + case OpType::MHA_VARLEN: + return true; + default: + return false; + } +} + +/// Check if an op type belongs to the GEMM/MLP family. +inline bool isGemmMlpOp(OpType type) { + switch (type) { + case OpType::GEMM: + case OpType::LINEAR: + case OpType::MATMUL: + case OpType::INT8_GEMM: + case OpType::SCALED_MM_I8: + return true; + default: + return false; + } +} + +/// Check if an op type is an activation function. +inline bool isActivationOp(OpType type) { + switch (type) { + case OpType::SILU: + case OpType::SILU_AND_MUL: + case OpType::GELU: + case OpType::SWIGLU: + case OpType::RELU: + case OpType::SIGMOID: + return true; + default: + return false; + } +} + +/// Check if an op type is KV cache related. +inline bool isKvCacheOp(OpType type) { + switch (type) { + case OpType::KV_CACHING: + case OpType::PAGED_CACHING: + return true; + default: + return false; + } +} + +} // namespace infinicore::analyzer diff --git a/python/infinicore/include/infinicore/analyzer/op_type_registry.hpp b/python/infinicore/include/infinicore/analyzer/op_type_registry.hpp new file mode 100644 index 000000000..4445f3a51 --- /dev/null +++ b/python/infinicore/include/infinicore/analyzer/op_type_registry.hpp @@ -0,0 +1,49 @@ +#pragma once + +#include "op_type.hpp" +#include +#include + +namespace infinicore::analyzer { + +/// Centralized class-name → OpType registry. +/// New ops only need one line added here — no changes to op headers. +inline OpType opTypeFromName(const char *name) { + static const std::unordered_map registry = { + // Attention + {"FlashAttention", OpType::FLASH_ATTENTION}, + {"CausalSoftmax", OpType::CAUSAL_SOFTMAX}, + {"PagedAttention", OpType::PAGED_ATTENTION}, + {"MhaKVCache", OpType::MHA_KVCACHE}, + {"MultiheadAttentionVarlen", OpType::MHA_VARLEN}, + // GEMM / MLP + {"Gemm", OpType::GEMM}, + {"I8Gemm", OpType::SCALED_MM_I8}, + // Activation + {"SiluAndMul", OpType::SILU_AND_MUL}, + {"SwiGLU", OpType::SWIGLU}, + // Norm + {"RMSNorm", OpType::RMS_NORM}, + {"AddRMSNorm", OpType::ADD_RMS_NORM}, + // Embedding / Positional + {"Embedding", OpType::EMBEDDING}, + {"RoPE", OpType::ROPE}, + // KV Cache + {"KVCaching", OpType::KV_CACHING}, + {"PagedCaching", OpType::PAGED_CACHING}, + // Elementwise + {"Add", OpType::ADD}, + {"Mul", OpType::MUL}, + // Quantization + {"PerTensorQuantI8", OpType::PER_TENSOR_QUANT_I8}, + {"PerTensorDequantI8", OpType::PER_TENSOR_DEQUANT_I8}, + {"PerChannelQuantI8", OpType::PER_CHANNEL_QUANT_I8}, + {"DequantizeAWQ", OpType::DEQUANTIZE_AWQ}, + // Misc + {"Rearrange", OpType::REARRANGE}, + }; + auto it = registry.find(name); + return it != registry.end() ? it->second : OpType::UNKNOWN; +} + +} // namespace infinicore::analyzer diff --git a/python/infinicore/include/infinicore/analyzer/optimization_intent.hpp b/python/infinicore/include/infinicore/analyzer/optimization_intent.hpp new file mode 100644 index 000000000..19d6c868c --- /dev/null +++ b/python/infinicore/include/infinicore/analyzer/optimization_intent.hpp @@ -0,0 +1,187 @@ +#pragma once + +#include "op_type.hpp" + +#include +#include +#include + +namespace infinicore::analyzer { + +// ============================================================ +// Phase & Bottleneck & Goal enumerations +// ============================================================ + +/// Recognized task phase types in LLM inference. +enum class PhaseType : uint8_t { + UNKNOWN = 0, + PREFILL, // Long-sequence prefill phase + DECODE, // Autoregressive decode phase + ATTENTION_DENSE, // Attention-dominated segment + GEMM_MLP_DENSE, // GEMM/MLP-dominated segment + MOE_ROUTING, // MoE routing segment (future) + KV_CACHE, // KV cache manipulation segment + COMMUNICATION, // Communication-dominated segment (future) +}; + +inline const char *phaseTypeToString(PhaseType phase) { + switch (phase) { + case PhaseType::PREFILL: + return "prefill"; + case PhaseType::DECODE: + return "decode"; + case PhaseType::ATTENTION_DENSE: + return "attention_dense"; + case PhaseType::GEMM_MLP_DENSE: + return "gemm_mlp_dense"; + case PhaseType::MOE_ROUTING: + return "moe_routing"; + case PhaseType::KV_CACHE: + return "kv_cache"; + case PhaseType::COMMUNICATION: + return "communication"; + default: + return "unknown"; + } +} + +/// Primary bottleneck type of the current phase. +enum class BottleneckType : uint8_t { + COMPUTE_BOUND = 0, + MEMORY_BOUND, + BANDWIDTH_BOUND, + COMMUNICATION_BOUND, + BALANCED, +}; + +inline const char *bottleneckTypeToString(BottleneckType bt) { + switch (bt) { + case BottleneckType::COMPUTE_BOUND: + return "compute_bound"; + case BottleneckType::MEMORY_BOUND: + return "memory_bound"; + case BottleneckType::BANDWIDTH_BOUND: + return "bandwidth_bound"; + case BottleneckType::COMMUNICATION_BOUND: + return "communication_bound"; + case BottleneckType::BALANCED: + return "balanced"; + default: + return "unknown"; + } +} + +/// Optimization goal / objective preference. +enum class OptimizationGoal : uint8_t { + LATENCY_FIRST = 0, + THROUGHPUT_FIRST, + MEMORY_SAFE, + STABILITY_FIRST, +}; + +inline const char *optimizationGoalToString(OptimizationGoal goal) { + switch (goal) { + case OptimizationGoal::LATENCY_FIRST: + return "latency_first"; + case OptimizationGoal::THROUGHPUT_FIRST: + return "throughput_first"; + case OptimizationGoal::MEMORY_SAFE: + return "memory_safe"; + case OptimizationGoal::STABILITY_FIRST: + return "stability_first"; + default: + return "unknown"; + } +} + +// ============================================================ +// Two-layer OptimizationIntent structure +// +// Design: Global Semantic Intent + Per-Device Local Refinement +// Rationale: Single-node multi-GPU scenario requires a unified +// task-level view PLUS per-device resource constraints. +// ============================================================ + +/// Strategy hints generated by the analyzer. +struct StrategyHint { + bool prefer_fused_ops = false; // Suggest op fusion? + bool prefer_in_place = false; // Suggest in-place computation? + bool prefer_recomputation = false; // Suggest recompute to save memory? + bool prefer_async_comm = false; // Suggest async comm overlap? +}; + +/// Global semantic-level intent — describes the task's overall +/// compute characteristics and optimization direction. +struct GlobalSemanticIntent { + // --- Current phase --- + PhaseType current_phase = PhaseType::UNKNOWN; + + // --- Primary bottleneck --- + BottleneckType primary_bottleneck = BottleneckType::BALANCED; + + // --- Optimization goal --- + OptimizationGoal goal = OptimizationGoal::LATENCY_FIRST; + + // --- Compute intensity estimate (ops/byte) --- + float compute_intensity = 0.0f; + + // --- Analysis confidence (0.0~1.0) --- + float confidence = 0.0f; + + // --- Strategy hints --- + StrategyHint strategy; + + // --- Production metadata --- + uint64_t timestamp_ns = 0; + uint32_t op_window_start = 0; // Start index of analyzed op window + uint32_t op_window_end = 0; // End index (exclusive) of analyzed op window +}; + +/// Per-device local intent — resource-aware refinement +/// for a specific device in multi-GPU scenario. +struct DeviceLocalIntent { + int device_id = -1; + + // Memory utilization (0.0 ~ 1.0) + float memory_usage_ratio = 0.0f; + + // Available memory in bytes + size_t memory_available_bytes = 0; + + // Device-local bottleneck (may differ from global) + BottleneckType local_bottleneck = BottleneckType::BALANCED; + + // Additional resource observations from the device snapshot + float compute_utilization = 0.0f; + float memory_bandwidth_utilization = 0.0f; + float communication_time_ratio = 0.0f; + + // Confidence of the resource-side classification (0.0 ~ 1.0) + float resource_confidence = 0.0f; +}; + +/// The complete OptimizationIntent: global task semantics +/// combined with per-device resource views. +/// +/// Produced by: MutualAwarenessAnalyzer +/// Consumed by (roadmap): +/// 1. MVP: query/record/validation (Python + C++ API) +/// 2. First: OpDispatcher / kernel selection +/// 3. Second: Allocator strategy (conservative vs aggressive) +/// 4. Third: Communication strategy (overlap, buffer, collective) +struct OptimizationIntent { + GlobalSemanticIntent global; + std::vector per_device; + + /// Get the local intent for a specific device, or nullptr if not found. + const DeviceLocalIntent *getDeviceIntent(int device_id) const { + for (auto &d : per_device) { + if (d.device_id == device_id) { + return &d; + } + } + return nullptr; + } +}; + +} // namespace infinicore::analyzer diff --git a/python/infinicore/include/infinicore/analyzer/phase_detector.hpp b/python/infinicore/include/infinicore/analyzer/phase_detector.hpp new file mode 100644 index 000000000..2f4828c6e --- /dev/null +++ b/python/infinicore/include/infinicore/analyzer/phase_detector.hpp @@ -0,0 +1,126 @@ +#pragma once + +#include "op_trace.hpp" +#include "optimization_intent.hpp" + +#include +#include + +namespace infinicore::analyzer { + +/// PhaseDetector — detects the current task phase from the +/// recent op trace window using rule-based pattern matching. +/// +/// Design choice (MVP): Fixed rule matching based on op type +/// composition in a sliding window. Will evolve to support +/// offline-generated phase templates in future iterations. +class PhaseDetector { +public: + /// Configuration for phase detection thresholds. + struct Config { + size_t window_size; + float attention_threshold; + float gemm_mlp_threshold; + float kv_cache_threshold; + uint32_t decode_seq_len_max; + uint32_t prefill_seq_len_min; + + Config() + : window_size(16), + attention_threshold(0.3f), + gemm_mlp_threshold(0.3f), + kv_cache_threshold(0.4f), + decode_seq_len_max(4), + prefill_seq_len_min(32) {} + }; + + explicit PhaseDetector(Config config = {}) : config_(config) {} + + /// Detect the current phase from a window of recent op traces. + PhaseType detect(const std::vector &window) const { + if (window.empty()) { + return PhaseType::UNKNOWN; + } + + // Count op categories in the window + size_t attention_count = 0; + size_t gemm_mlp_count = 0; + size_t kv_cache_count = 0; + size_t activation_count = 0; + size_t total = window.size(); + + // Track shape info for prefill/decode inference + uint32_t max_seq_len = 0; + uint32_t min_seq_len = UINT32_MAX; + bool has_attention_shape = false; + + for (auto &entry : window) { + if (isAttentionOp(entry.op_type)) { + attention_count++; + // For attention ops, shape[1] or shape[2] typically indicates seq_len + if (entry.ndim >= 2) { + // Heuristic: for attention-like ops, look at the sequence dimension + // Typically shape = [batch, seq_len, ...] or [batch, heads, seq_len, ...] + uint32_t seq_dim = (entry.ndim >= 3) ? entry.shape[2] : entry.shape[1]; + max_seq_len = std::max(max_seq_len, seq_dim); + min_seq_len = std::min(min_seq_len, seq_dim); + has_attention_shape = true; + } + } else if (isGemmMlpOp(entry.op_type)) { + gemm_mlp_count++; + } else if (isKvCacheOp(entry.op_type)) { + kv_cache_count++; + } else if (isActivationOp(entry.op_type)) { + activation_count++; + } + } + + float attention_ratio = static_cast(attention_count) / total; + float gemm_mlp_ratio = static_cast(gemm_mlp_count + activation_count) / total; + float kv_cache_ratio = static_cast(kv_cache_count) / total; + + // --- Phase classification --- + + // KV cache phase (high KV cache op ratio) + if (kv_cache_ratio >= config_.kv_cache_threshold) { + return PhaseType::KV_CACHE; + } + + // --- Prefill vs Decode inference from shape --- + // Self-inferred from sequence length, no external flags needed. + if (has_attention_shape) { + if (max_seq_len <= config_.decode_seq_len_max) { + return PhaseType::DECODE; + } + if (min_seq_len >= config_.prefill_seq_len_min) { + return PhaseType::PREFILL; + } + } + + // Attention-dense phase + if (attention_ratio >= config_.attention_threshold && attention_ratio >= gemm_mlp_ratio) { + return PhaseType::ATTENTION_DENSE; + } + + // GEMM/MLP-dense phase (include activation ops as co-indicators) + if (gemm_mlp_ratio >= config_.gemm_mlp_threshold && gemm_mlp_ratio >= attention_ratio) { + return PhaseType::GEMM_MLP_DENSE; + } + + return PhaseType::UNKNOWN; + } + + /// Convenience: detect from the global trace ring. + PhaseType detectFromTrace(const OpTraceRing &trace) const { + auto window = trace.getRecentEntries(config_.window_size); + return detect(window); + } + + const Config &config() const { return config_; } + void setConfig(Config config) { config_ = config; } + +private: + Config config_; +}; + +} // namespace infinicore::analyzer diff --git a/python/infinicore/include/infinicore/analyzer/resource_sensor.hpp b/python/infinicore/include/infinicore/analyzer/resource_sensor.hpp new file mode 100644 index 000000000..72c871080 --- /dev/null +++ b/python/infinicore/include/infinicore/analyzer/resource_sensor.hpp @@ -0,0 +1,181 @@ +#pragma once + +#include "infinicore/device.hpp" +#include "optimization_intent.hpp" + +#include +#include +#include + +namespace infinicore::analyzer { + +/// Memory statistics from the allocator. +struct MemoryStats { + size_t allocated_bytes = 0; // Currently allocated bytes + size_t total_capacity = 0; // Total pool capacity in bytes + size_t peak_allocated = 0; // Peak allocation since last reset + size_t allocation_count = 0; // Number of active allocations + + float usageRatio() const { + return total_capacity > 0 + ? static_cast(allocated_bytes) / static_cast(total_capacity) + : 0.0f; + } +}; + +/// A normalized resource image consumed by the analyzer. +/// +/// This is intentionally vendor-neutral. Backend-specific runtimes +/// populate these fields through infinirt, and the analyzer consumes +/// only this normalized view. +struct DeviceResourceSnapshot { + int device_id = -1; + Device::Type device_type = Device::Type::CPU; + + bool has_memory_capacity = false; + bool has_compute_utilization = false; + bool has_memory_bandwidth_utilization = false; + bool has_kernel_time_ratio = false; + bool has_communication = false; + bool kernel_time_estimated = false; + + size_t free_bytes = 0; + size_t total_bytes = 0; + size_t used_bytes = 0; + size_t reserved_bytes = 0; + + float compute_utilization = 0.0f; + float memory_bandwidth_utilization = 0.0f; + + uint64_t bytes_read = 0; + uint64_t bytes_written = 0; + + float kernel_time_ratio = 0.0f; + float idle_time_ratio = 0.0f; + + float communication_time_ratio = 0.0f; + uint64_t communication_bytes = 0; + + float load_imbalance_score = 0.0f; + + float memoryUsageRatio() const { + if (total_bytes > 0) { + return static_cast(used_bytes) / static_cast(total_bytes); + } + if (reserved_bytes > 0 && used_bytes <= reserved_bytes) { + return static_cast(used_bytes) / static_cast(reserved_bytes); + } + return 0.0f; + } + + float resourceConfidence() const { + float confidence = 0.0f; + if (has_memory_capacity) { + confidence += 0.35f; + } + if (has_compute_utilization) { + confidence += 0.25f; + } + if (has_memory_bandwidth_utilization) { + confidence += 0.25f; + } + if (has_communication) { + confidence += 0.15f; + } + return std::min(confidence, 1.0f); + } +}; + +/// ResourceSensor — gathers current resource state from +/// the runtime and allocator subsystems. +/// +/// This is the "resource supply" side of the mutual-awareness +/// equation. It aggregates device type, memory stats, and +/// potentially timing info into a resource snapshot. +class ResourceSensor { +public: + ResourceSensor() = default; + + /// Build a DeviceLocalIntent from a normalized resource image. + DeviceLocalIntent sense(const DeviceResourceSnapshot &snapshot) const { + DeviceLocalIntent intent; + intent.device_id = snapshot.device_id; + intent.memory_usage_ratio = snapshot.memoryUsageRatio(); + intent.memory_available_bytes = snapshot.free_bytes > 0 + ? snapshot.free_bytes + : ((snapshot.total_bytes >= snapshot.used_bytes) ? (snapshot.total_bytes - snapshot.used_bytes) : 0); + intent.compute_utilization = snapshot.compute_utilization; + intent.memory_bandwidth_utilization = snapshot.memory_bandwidth_utilization; + intent.communication_time_ratio = snapshot.communication_time_ratio; + intent.resource_confidence = snapshot.resourceConfidence(); + + if (snapshot.has_communication && snapshot.communication_time_ratio > high_communication_threshold_) { + intent.local_bottleneck = BottleneckType::COMMUNICATION_BOUND; + return intent; + } + + if (intent.memory_usage_ratio > high_memory_threshold_) { + intent.local_bottleneck = BottleneckType::MEMORY_BOUND; + return intent; + } + + if (snapshot.has_memory_bandwidth_utilization + && snapshot.memory_bandwidth_utilization > high_bandwidth_threshold_ + && (!snapshot.has_compute_utilization + || snapshot.memory_bandwidth_utilization >= snapshot.compute_utilization + bandwidth_margin_)) { + intent.local_bottleneck = BottleneckType::BANDWIDTH_BOUND; + return intent; + } + + if (snapshot.has_compute_utilization && snapshot.compute_utilization > high_compute_threshold_) { + intent.local_bottleneck = BottleneckType::COMPUTE_BOUND; + return intent; + } + + if (intent.memory_usage_ratio > moderate_memory_threshold_) { + intent.local_bottleneck = BottleneckType::BALANCED; + return intent; + } + + // Preserve the MVP fallback: if all we know is that memory is low, + // compute is the most likely bottleneck. + intent.local_bottleneck = BottleneckType::COMPUTE_BOUND; + return intent; + } + + /// Build a DeviceLocalIntent from current resource state. + /// In MVP, this primarily queries allocator memory stats. + /// + /// @param device_id The device ID to query + /// @param stats Memory stats from the allocator + DeviceLocalIntent sense( + int device_id, + const MemoryStats &stats, + Device::Type device_type = Device::Type::CPU) const { + DeviceResourceSnapshot snapshot; + snapshot.device_id = device_id; + snapshot.device_type = device_type; + snapshot.has_memory_capacity = stats.total_capacity > 0; + snapshot.free_bytes = stats.total_capacity >= stats.allocated_bytes + ? (stats.total_capacity - stats.allocated_bytes) + : 0; + snapshot.total_bytes = stats.total_capacity; + snapshot.used_bytes = stats.allocated_bytes; + snapshot.reserved_bytes = stats.total_capacity; + return sense(snapshot); + } + + /// Thresholds for memory-based bottleneck classification. + void setHighMemoryThreshold(float t) { high_memory_threshold_ = t; } + void setModerateMemoryThreshold(float t) { moderate_memory_threshold_ = t; } + +private: + float high_memory_threshold_ = 0.85f; + float moderate_memory_threshold_ = 0.5f; + float high_compute_threshold_ = 0.75f; + float high_bandwidth_threshold_ = 0.75f; + float high_communication_threshold_ = 0.25f; + float bandwidth_margin_ = 0.05f; +}; + +} // namespace infinicore::analyzer diff --git a/python/infinicore/include/infinicore/common/LRUCache.hpp b/python/infinicore/include/infinicore/common/LRUCache.hpp new file mode 100644 index 000000000..0c13d1d09 --- /dev/null +++ b/python/infinicore/include/infinicore/common/LRUCache.hpp @@ -0,0 +1,137 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace infinicore::common { +template +class LRUCache { +public: + using KeyValuePair = std::pair; + using ListIt = typename std::list::iterator; + using Destructor = std::function; + + explicit LRUCache(size_t capacity = 100, Destructor destructor = nullptr) + : capacity_(capacity), destructor_(destructor) { + if (capacity == 0) { + capacity_ = UINT64_MAX; // effectively unbounded + } + } + + ~LRUCache() { + cleanup(); + } + + bool contains(const Key &key) const { + return map_.find(key) != map_.end(); + } + + void put(const Key &key, const Value &value) { + auto it = map_.find(key); + if (it != map_.end()) { + if (destructor_) { + destructor_(it->second->second); + } + it->second->second = value; + touch(it); + } else { + // insert new + if (list_.size() >= capacity_) { + evictLRU(); + } + list_.emplace_front(key, value); + map_[key] = list_.begin(); + } + } + + std::optional get(const Key &key) { + auto it = map_.find(key); + if (it == map_.end()) { + return std::nullopt; + } + touch(it); + return it->second->second; + } + + std::optional get(const Key &key) const { + auto it = map_.find(key); + if (it == map_.end()) { + return std::nullopt; + } + // Note: can't touch in const context + return it->second->second; + } + + void setDestructor(Destructor destructor) { + destructor_ = destructor; + } + + void setCapacity(size_t capacity) { + capacity_ = capacity; + while (list_.size() > capacity_) { + evictLRU(); + } + } + + void clear() { + if (destructor_) { + for (auto &item : list_) { + safeDestruct(item.second); + } + } + list_.clear(); + map_.clear(); + } + + const std::list &getAllItems() const { + return list_; + } + +protected: + std::list list_; // front = most recent, back = least + +private: + void touch(typename std::unordered_map::iterator it) { + // move this key to front (most recent) + list_.splice(list_.begin(), list_, it->second); + it->second = list_.begin(); + } + + void safeDestruct(Value &value) { + if (!destructor_) { + return; + } + + try { + destructor_(value); + } catch (const std::exception &e) { + // Built-in default error handling + std::cerr << "Cache destructor error (type: " << typeid(Value).name() + << "): " << e.what() << std::endl; + } + } + + void evictLRU() { + if (!list_.empty()) { + auto &kv = list_.back(); + safeDestruct(kv.second); + map_.erase(kv.first); + list_.pop_back(); + } + } + + void cleanup() { + clear(); + } + + size_t capacity_; + std::unordered_map map_; + Destructor destructor_; +}; + +} // namespace infinicore::common diff --git a/python/infinicore/include/infinicore/common/hash.hpp b/python/infinicore/include/infinicore/common/hash.hpp new file mode 100644 index 000000000..65178e4db --- /dev/null +++ b/python/infinicore/include/infinicore/common/hash.hpp @@ -0,0 +1,72 @@ +#pragma once + +#include "../tensor.hpp" + +#include +#include + +namespace infinicore { + +// Base hash_combine for arithmetic types +template +std::enable_if_t, void> +hash_combine(size_t &seed, const T &value) { + seed ^= std::hash{}(value) + 0x9e3779b9 + (seed << 6) + (seed >> 2); +} + +// Specialization for Tensor +inline void hash_combine(size_t &seed, Tensor tensor) { + if (!tensor) { + hash_combine(seed, static_cast(0)); + return; + } + + hash_combine(seed, static_cast(tensor->dtype())); + for (Size shape : tensor->shape()) { + hash_combine(seed, shape); + } + for (Stride stride : tensor->strides()) { + hash_combine(seed, static_cast(stride)); + } +} + +// Specialization for optional +template +inline void hash_combine(size_t &seed, const std::optional &opt) { + hash_combine(seed, opt.has_value()); + if (opt) { + hash_combine(seed, *opt); + } +} + +// Specialization for std::string +inline void hash_combine(size_t &seed, const std::string &str) { + hash_combine(seed, std::hash{}(str)); +} + +// Specialization for const char* +inline void hash_combine(size_t &seed, const char *str) { + hash_combine(seed, std::string(str)); +} + +// Variadic template for multiple arguments +template +void hash_combine(size_t &seed, const First &first, const Rest &...rest) { + hash_combine(seed, first); + hash_combine(seed, rest...); +} + +// Base case for variadic template +inline void hash_combine(size_t &seed) { + // Base case - do nothing +} + +// Convenience function to hash multiple values +template +size_t hash_combine(const Types &...values) { + size_t seed = 0; + hash_combine(seed, values...); + return seed; +} + +} // namespace infinicore diff --git a/python/infinicore/include/infinicore/context/context.hpp b/python/infinicore/include/infinicore/context/context.hpp new file mode 100644 index 000000000..9db86e97c --- /dev/null +++ b/python/infinicore/include/infinicore/context/context.hpp @@ -0,0 +1,58 @@ +#pragma once + +#include "../device.hpp" +#include "../memory.hpp" + +#include "../graph/graph.hpp" + +#include +#include + +#include + +namespace infinicore { + +namespace context { +void setDevice(Device device); +Device getDevice(); +size_t getDeviceCount(Device::Type type); + +infinirtStream_t getStream(); +infiniopHandle_t getInfiniopHandle(Device device); + +void syncStream(); +void syncDevice(); +void trimMemory(); + +std::shared_ptr allocateMemory(size_t size); +std::shared_ptr allocateHostMemory(size_t size); +std::shared_ptr allocatePinnedHostMemory(size_t size); + +void memcpyH2D(void *dst, const void *src, size_t size, bool async = true); +void memcpyD2H(void *dst, const void *src, size_t size); +void memcpyD2D(void *dst, const void *src, size_t size, bool async = true); +void memcpyH2H(void *dst, const void *src, size_t size); + +void setDeviceMemory(void *ptr, int value, size_t count); +void setDeviceMemoryAsync(void *ptr, int value, size_t count, infinirtStream_t stream); + +// Timing APIs for performance measurement +infinirtEvent_t createEvent(); +infinirtEvent_t createEventWithFlags(uint32_t flags); +void recordEvent(infinirtEvent_t event, infinirtStream_t stream = nullptr); +bool queryEvent(infinirtEvent_t event); +void synchronizeEvent(infinirtEvent_t event); +void destroyEvent(infinirtEvent_t event); +float elapsedTime(infinirtEvent_t start, infinirtEvent_t end); +void streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event); + +// Graph recording APIs +bool isGraphRecording(); +void startGraphRecording(); +void addGraphOperator(std::shared_ptr op); +std::shared_ptr stopGraphRecording(); +void cancelGraphRecording(); + +} // namespace context + +} // namespace infinicore diff --git a/python/infinicore/include/infinicore/device.hpp b/python/infinicore/include/infinicore/device.hpp new file mode 100644 index 000000000..4bdab1088 --- /dev/null +++ b/python/infinicore/include/infinicore/device.hpp @@ -0,0 +1,53 @@ +#pragma once + +#include +#include + +#include "infinicore.h" + +namespace infinicore { + +class Device { +public: + using Index = std::size_t; + + enum class Type { + CPU = INFINI_DEVICE_CPU, + NVIDIA = INFINI_DEVICE_NVIDIA, + CAMBRICON = INFINI_DEVICE_CAMBRICON, + ASCEND = INFINI_DEVICE_ASCEND, + METAX = INFINI_DEVICE_METAX, + MOORE = INFINI_DEVICE_MOORE, + ILUVATAR = INFINI_DEVICE_ILUVATAR, + KUNLUN = INFINI_DEVICE_KUNLUN, + HYGON = INFINI_DEVICE_HYGON, + QY = INFINI_DEVICE_QY, + ALI = INFINI_DEVICE_ALI, + COUNT = INFINI_DEVICE_TYPE_COUNT, + }; + + Device(const Type &type = Type::CPU, const Index &index = 0); + + const Type &getType() const; + + const Index &getIndex() const; + + std::string toString() const; + + static std::string toString(const Type &type); + + bool operator==(const Device &other) const; + + bool operator!=(const Device &other) const; + + inline static Device cpu() { + return Device(Type::CPU, 0); + } + +private: + Type type_; + + Index index_; +}; + +} // namespace infinicore diff --git a/python/infinicore/include/infinicore/device_event.hpp b/python/infinicore/include/infinicore/device_event.hpp new file mode 100644 index 000000000..8ce38d645 --- /dev/null +++ b/python/infinicore/include/infinicore/device_event.hpp @@ -0,0 +1,125 @@ +#pragma once + +#include "device.hpp" +#include "infinirt.h" +#include +#include + +namespace infinicore { + +/** + * @brief A device event for timing operations and synchronization across devices. + * + * Similar to torch.cuda.Event, this class provides functionality to: + * - Record events on specific device streams + * - Synchronize with events + * - Measure elapsed time between events + * - Query event completion status + * - Make streams wait for events + */ +class DeviceEvent { +private: + infinirtEvent_t event_; // Underlying event handle + Device device_; // Device where this event was created + bool is_recorded_; // Whether the event has been recorded + +public: + /** + * @brief Construct a new DeviceEvent on the current device. + */ + DeviceEvent(); + + /** + * @brief Construct a new DeviceEvent on the current device with specific flags. + * @param flags Event creation flags (e.g., for timing, blocking sync) + */ + explicit DeviceEvent(uint32_t flags); + + /** + * @brief Construct a new DeviceEvent on a specific device. + * @param device Target device for this event + */ + explicit DeviceEvent(Device device); + + /** + * @brief Construct a new DeviceEvent on a specific device with flags. + * @param device Target device for this event + * @param flags Event creation flags + */ + DeviceEvent(Device device, uint32_t flags); + + // Disallow copying + DeviceEvent(const DeviceEvent &) = delete; + DeviceEvent &operator=(const DeviceEvent &) = delete; + + /** + * @brief Move constructor. + */ + DeviceEvent(DeviceEvent &&other) noexcept; + + /** + * @brief Move assignment operator. + */ + DeviceEvent &operator=(DeviceEvent &&other) noexcept; + + /** + * @brief Destroy the DeviceEvent and release underlying resources. + */ + ~DeviceEvent(); + + /** + * @brief Record the event on the current stream of its device. + */ + void record(); + + /** + * @brief Record the event on a specific stream. + * @param stream Stream to record the event on + */ + void record(infinirtStream_t stream); + + /** + * @brief Wait for the event to complete (blocking). + */ + void synchronize(); + + /** + * @brief Check if the event has been completed. + * @return true if completed, false otherwise + */ + bool query() const; + + /** + * @brief Calculate elapsed time between this event and another event (in milliseconds). + * @param other The other event to compare with + * @return Elapsed time in milliseconds + * @throws std::runtime_error if events are on different devices or not recorded + */ + float elapsed_time(const DeviceEvent &other) const; + + /** + * @brief Make a stream wait for this event to complete. + * @param stream Stream to make wait for this event (nullptr for current stream) + */ + void wait(infinirtStream_t stream = nullptr) const; + + /** + * @brief Get the device where this event was created. + * @return Device associated with this event + */ + Device device() const { return device_; } + + /** + * @brief Get the underlying event handle. + * @return Raw event handle + */ + infinirtEvent_t get() const { return event_; } + + /** + * @brief Check if the event has been recorded. + * @return true if recorded, false otherwise + */ + bool is_recorded() const { return is_recorded_; } +}; + +} // namespace infinicore diff --git a/python/infinicore/include/infinicore/dtype.hpp b/python/infinicore/include/infinicore/dtype.hpp new file mode 100644 index 000000000..ea3f49286 --- /dev/null +++ b/python/infinicore/include/infinicore/dtype.hpp @@ -0,0 +1,33 @@ +#pragma once + +#include +#include + +namespace infinicore { + +enum class DataType { + BYTE = INFINI_DTYPE_BYTE, + BOOL = INFINI_DTYPE_BOOL, + I8 = INFINI_DTYPE_I8, + I16 = INFINI_DTYPE_I16, + I32 = INFINI_DTYPE_I32, + I64 = INFINI_DTYPE_I64, + U8 = INFINI_DTYPE_U8, + U16 = INFINI_DTYPE_U16, + U32 = INFINI_DTYPE_U32, + U64 = INFINI_DTYPE_U64, + F8 = INFINI_DTYPE_F8, + F16 = INFINI_DTYPE_F16, + F32 = INFINI_DTYPE_F32, + F64 = INFINI_DTYPE_F64, + C16 = INFINI_DTYPE_C16, + C32 = INFINI_DTYPE_C32, + C64 = INFINI_DTYPE_C64, + C128 = INFINI_DTYPE_C128, + BF16 = INFINI_DTYPE_BF16, +}; + +std::string toString(const DataType &dtype); +size_t dsize(const DataType &dtype); + +} // namespace infinicore diff --git a/python/infinicore/include/infinicore/graph/graph.hpp b/python/infinicore/include/infinicore/graph/graph.hpp new file mode 100644 index 000000000..be368f92d --- /dev/null +++ b/python/infinicore/include/infinicore/graph/graph.hpp @@ -0,0 +1,133 @@ +#pragma once + +#include +#include + +#include "../tensor.hpp" + +namespace infinicore::graph { +// Forward declarations +class GraphManager; + +class GraphTensor : public Tensor { +public: + GraphTensor(const Tensor &); +}; + +class GraphOperator { +public: + virtual void run() const = 0; + virtual ~GraphOperator() = default; +}; + +class DispatchableGraphOperator : public GraphOperator { +public: + void run() const override; + ~DispatchableGraphOperator() override; + +protected: + using run_schema = void (*)(void *); + using cleanup_schema = void (*)(void **); + void *planned_meta_; + run_schema runner_; + cleanup_schema deleter_; +}; + +class Graph { +public: + Graph(); + ~Graph(); + + void run() const; + +protected: + void add_operator(std::shared_ptr op); + void instantiate(); + std::vector> op_list_; + + friend class GraphManager; + +private: + struct DeviceGraph; + std::unique_ptr device_graph_; +}; +} // namespace infinicore::graph + +#define INFINICORE_GRAPH_OP_CLASS(__OP_NAME__, ...) \ + class __OP_NAME__ : public graph::DispatchableGraphOperator { \ + public: \ + using schema = void (*)(__VA_ARGS__); \ + using plan_schema = void *(*)(__VA_ARGS__); \ + static common::OpDispatcher &plan_dispatcher(); \ + static common::OpDispatcher &run_dispatcher(); \ + static common::OpDispatcher &cleanup_dispatcher(); \ + __OP_NAME__(__VA_ARGS__); \ + static void execute(__VA_ARGS__); \ + }; + +#define INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(__OP_NAME__) \ + common::OpDispatcher<__OP_NAME__::plan_schema> &__OP_NAME__::plan_dispatcher() { \ + static common::OpDispatcher<__OP_NAME__::plan_schema> dispatcher_; \ + return dispatcher_; \ + } \ + common::OpDispatcher<__OP_NAME__::run_schema> &__OP_NAME__::run_dispatcher() { \ + static common::OpDispatcher<__OP_NAME__::run_schema> dispatcher_; \ + return dispatcher_; \ + } \ + common::OpDispatcher<__OP_NAME__::cleanup_schema> &__OP_NAME__::cleanup_dispatcher() { \ + static common::OpDispatcher<__OP_NAME__::cleanup_schema> dispatcher_; \ + return dispatcher_; \ + } + +#define INFINICORE_GRAPH_OP_DISPATCH(__DEVICE_TYPE__, ...) \ + planned_meta_ = plan_dispatcher().lookup(__DEVICE_TYPE__)(__VA_ARGS__); \ + runner_ = run_dispatcher().lookup(__DEVICE_TYPE__); \ + deleter_ = cleanup_dispatcher().lookup(__DEVICE_TYPE__); + +#define INFINICORE_DETAIL_FIRST_ARG(__FIRST__, ...) __FIRST__ + +#ifdef ENABLE_MUTUAL_AWARENESS +#include "../analyzer/op_trace.hpp" +#include "../analyzer/op_type_registry.hpp" + +// Trace one op invocation into the global ring. Op type is resolved by +// stringified class name through `opTypeFromName`, so new graph ops are +// automatically discoverable without modifying the op header. +#define _INFINICORE_TRACE_OP(__OP_NAME__, __TRACE_TENSOR__) \ + do { \ + auto __op_type = ::infinicore::analyzer::opTypeFromName(#__OP_NAME__); \ + auto &&__trace_tensor = (__TRACE_TENSOR__); \ + if (__trace_tensor) { \ + const auto &__trace_shape = __trace_tensor->shape(); \ + const auto __trace_device = __trace_tensor->device(); \ + ::infinicore::analyzer::traceOp( \ + __op_type, \ + __trace_shape.data(), \ + __trace_shape.size(), \ + static_cast(__trace_tensor->dtype()), \ + static_cast(__trace_device.getType()), \ + static_cast(__trace_device.getIndex())); \ + } else { \ + ::infinicore::analyzer::traceOp(__op_type, nullptr, 0, 0, 0, -1); \ + } \ + } while (0) +#else +#define _INFINICORE_TRACE_OP(__OP_NAME__, __TRACE_TENSOR__) ((void)0) +#endif + +#define INFINICORE_GRAPH_OP_RECORD_OR_RUN(__OP_NAME__, ...) \ + auto ___op = std::make_shared<__OP_NAME__>(__VA_ARGS__); \ + if (context::isGraphRecording()) { \ + context::addGraphOperator(___op); \ + } else { \ + ___op->run(); \ + } \ + _INFINICORE_TRACE_OP(__OP_NAME__, INFINICORE_DETAIL_FIRST_ARG(__VA_ARGS__)); + +#define INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(__OP_NAME__, __PLAN_F__, __RUN_F__, __CLEANUP_F__) \ + static bool registered = []() { \ + __OP_NAME__::plan_dispatcher().registerAll(__PLAN_F__, false); \ + __OP_NAME__::run_dispatcher().registerAll(__RUN_F__, false); \ + __OP_NAME__::cleanup_dispatcher().registerAll(__CLEANUP_F__, false); \ + return true; \ + }(); diff --git a/python/infinicore/include/infinicore/io.hpp b/python/infinicore/include/infinicore/io.hpp new file mode 100644 index 000000000..8d11a88b2 --- /dev/null +++ b/python/infinicore/include/infinicore/io.hpp @@ -0,0 +1,86 @@ +#pragma once + +#include "tensor.hpp" +#include + +namespace infinicore::print_options { + +/** + * @brief Sets the line width. After \a line_width chars, a new line is added. + * @param line_width The line width + */ +void set_line_width(int line_width); + +/** + * @brief Sets the threshold after which summarization is triggered (default: 1000). + * @param threshold The number of elements in the tensor that triggers summarization in the output + */ +void set_threshold(int threshold); + +/** + * @brief Sets the number of edge items. + * If the summarization is triggered, this value defines how many items of each dimension are printed. + * @param edge_items The number of edge items + */ +void set_edge_items(int edge_items); + +/** + * @brief Sets the precision for printing floating point values. + * @param precision The number of digits for floating point output + */ + +void set_precision(int precision); + +/** + * @brief Sets the sci mode of the floating point values when printing an Tensor. + * @param sci_mode The sci mode: -1 for auto decision, 0 to disable, 1 to enable + */ + +void set_sci_mode(int sci_mode); // -1: auto, 0: disable, 1: enable + +#define DEFINE_LOCAL_PRINT_OPTION(NAME) \ + class NAME { \ + public: \ + NAME(int value) : m_value(value) { id(); } \ + static int id() { \ + static int id = std::ios_base::xalloc(); \ + return id; \ + } \ + int value() const { return m_value; } \ + \ + private: \ + int m_value; \ + }; \ + \ + inline std::ostream &operator<<(std::ostream &out, const NAME &n) { \ + out.iword(NAME::id()) = n.value(); \ + return out; \ + } + +/** + * @class line_width + * io manipulator used to set the width of the lines when printing an Tensor. + * + * @code{.cpp} + * using po = infinicore::print_options; + * std::cout << po::line_width(100) << tensor << std::endl; + * @endcode + */ +DEFINE_LOCAL_PRINT_OPTION(line_width) + +/** + * io manipulator used to set the threshold after which summarization is triggered. + */ +DEFINE_LOCAL_PRINT_OPTION(threshold) + +/** + * io manipulator used to set the number of egde items if the summarization is triggered. + */ +DEFINE_LOCAL_PRINT_OPTION(edge_items) + +/** + * io manipulator used to set the precision of the floating point values when printing an Tensor. + */ +DEFINE_LOCAL_PRINT_OPTION(precision) + +} // namespace infinicore::print_options diff --git a/python/infinicore/include/infinicore/memory.hpp b/python/infinicore/include/infinicore/memory.hpp new file mode 100644 index 000000000..65d812d67 --- /dev/null +++ b/python/infinicore/include/infinicore/memory.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include "device.hpp" + +#include +#include + +namespace infinicore { + +class Memory { +public: + using Deleter = std::function; + + Memory(std::byte *data, size_t size, Device device, Deleter deleter, bool pin_memory = false); + ~Memory(); + + std::byte *data(); + Device device() const; + size_t size() const; + bool is_pinned() const; + +private: + std::byte *data_; + size_t size_; + Device device_; + Deleter deleter_; + bool is_pinned_; +}; + +} // namespace infinicore diff --git a/python/infinicore/include/infinicore/nn.hpp b/python/infinicore/include/infinicore/nn.hpp new file mode 100644 index 000000000..16f463a72 --- /dev/null +++ b/python/infinicore/include/infinicore/nn.hpp @@ -0,0 +1,6 @@ +#pragma once + +#include "nn/embedding.hpp" +#include "nn/linear.hpp" +#include "nn/rmsnorm.hpp" +#include "nn/rope.hpp" diff --git a/python/infinicore/include/infinicore/nn/embedding.hpp b/python/infinicore/include/infinicore/nn/embedding.hpp new file mode 100644 index 000000000..50a387325 --- /dev/null +++ b/python/infinicore/include/infinicore/nn/embedding.hpp @@ -0,0 +1,87 @@ +#pragma once + +#include "../ops.hpp" +#include "module.hpp" +#include + +namespace infinicore::nn { + +/** + * @brief Embedding layer that maps indices to dense vectors + * + * A simple lookup table that stores embeddings of a fixed dictionary and size. + * This module is often used to store word embeddings and retrieve them using indices. + * The input to the module is a tensor of indices, and the output is the corresponding + * embedding vectors. + * + * Similar to PyTorch's nn.Embedding: + * https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html + * + * Example: + * @code + * // Create embedding: 10000 words, 300-dimensional embeddings + * auto embedding = Embedding(10000, 300); + * + * // Input: tensor of indices [batch_size, seq_len] + * auto indices = Tensor::from_data({2, 5}, {3, 5, 12, 8, 99, 0, 1, 45, 67, 23}); + * + * // Output: [batch_size, seq_len, embedding_dim] = [2, 5, 300] + * auto embeddings = embedding.forward(indices); + * @endcode + */ +class Embedding : public Module { +public: + /** + * @brief Construct an Embedding layer + * + * @param num_embeddings Size of the dictionary of embeddings (vocabulary size) + * @param embedding_dim The size of each embedding vector + * @param padding_idx If specified, the entries at padding_idx do not contribute to gradient + * and the embedding vector at padding_idx is not updated during training + * @param dtype Data type for the embedding weights (default: DataType::F32) + * @param device Device to create the embedding weight on + */ + Embedding(size_t num_embeddings, + size_t embedding_dim, + std::optional padding_idx = std::nullopt, + const DataType &dtype = DataType::F32, + const Device &device = Device()); + + /** + * @brief Forward pass: lookup embeddings for given indices + * + * @param indices Tensor containing indices into the embedding matrix. + * Can be any shape (*), typically [batch_size] or [batch_size, seq_len] + * @return Tensor containing the embedding vectors. + * Shape: (*, embedding_dim) where * matches the input shape + * + * Example: + * Input shape: [2, 3] -> Output shape: [2, 3, embedding_dim] + * Input shape: [10] -> Output shape: [10, embedding_dim] + */ + Tensor forward(const Tensor &indices) const; + + // Module information + size_t num_embeddings() const { return num_embeddings_; } + size_t embedding_dim() const { return embedding_dim_; } + std::optional padding_idx() const { return padding_idx_; } + DataType dtype() const { return dtype_; } + + // String representation + std::string extra_repr() const; + + // Accessors for parameters + Tensor weight() const { return weight_; } + +protected: + // Parameters + INFINICORE_NN_PARAMETER(weight); + +private: + size_t num_embeddings_; // Vocabulary size + size_t embedding_dim_; // Embedding dimension + std::optional padding_idx_; // Optional padding index + DataType dtype_; // Data type for embedding weights +}; + +} // namespace infinicore::nn diff --git a/python/infinicore/include/infinicore/nn/layer_norm.hpp b/python/infinicore/include/infinicore/nn/layer_norm.hpp new file mode 100644 index 000000000..d69d0594f --- /dev/null +++ b/python/infinicore/include/infinicore/nn/layer_norm.hpp @@ -0,0 +1,59 @@ +#pragma once + +#include "../ops.hpp" +#include "module.hpp" + +namespace infinicore::nn { + +class LayerNorm : public Module { +public: + /** + * @brief Construct a LayerNorm layer + * + * @param normalized_shape Size of the feature dimension to normalize (typically hidden_size) + * @param eps Small constant for numerical stability (default: 1e-6) + * @param dtype Data type for the weight (default: DataType::F32) + * @param device Device to create the weight on + */ + LayerNorm(size_t normalized_shape, + double eps = 1e-6, + const DataType &dtype = DataType::F32, + const Device &device = Device()); + + /** + * @brief Forward pass: apply LayerNorm + * + * @param x Input tensor of shape (*, normalized_shape) where * is any number of dimensions + * @return Normalized tensor with same shape as input + * + * The normalization is applied over the last dimension. + * For example: + * Input: [batch, seq_len, hidden_size] -> normalize over hidden_size + * Input: [batch, hidden_size] -> normalize over hidden_size + */ + Tensor forward(const Tensor &x) const; + + // Module information + size_t normalized_shape() const { return normalized_shape_; } + double eps() const { return eps_; } + DataType dtype() const { return dtype_; } + + // String representation + std::string extra_repr() const; + + // Accessors for parameters + Tensor weight() const { return weight_; } + Tensor bias() const { return bias_; } + +protected: + // Parameters + INFINICORE_NN_PARAMETER(weight); + INFINICORE_NN_PARAMETER(bias); + +private: + size_t normalized_shape_; // Size of the feature dimension + double eps_; // Epsilon for numerical stability + DataType dtype_; // Data type for weight +}; + +} // namespace infinicore::nn diff --git a/python/infinicore/include/infinicore/nn/linear.hpp b/python/infinicore/include/infinicore/nn/linear.hpp new file mode 100644 index 000000000..00ea8a1a5 --- /dev/null +++ b/python/infinicore/include/infinicore/nn/linear.hpp @@ -0,0 +1,128 @@ +#pragma once + +#include "../ops.hpp" +#include "../quantization.hpp" +#include "module.hpp" +#include +#include + +namespace infinicore::nn { + +class BaseLinear : public Module { +public: + BaseLinear(size_t in_features, size_t out_features, bool bias = true, + const DataType &dtype = DataType::F32, const Device &device = Device()); + + BaseLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, + const DataType &dtype = DataType::F32, const Device &device = Device()); + + // Forward pass: output = input @ weight.T + bias + Tensor forward(Tensor &input) const; + + // Forward pass with residual connection (InfiniLM-style) + // output = input @ weight.T + bias + residual + Tensor forward(Tensor &input, Tensor &residual) const; + + // Module information + size_t in_features() const { return in_features_; } + size_t out_features() const { return out_features_; } + bool has_bias() const { return has_bias_; } + DataType dtype() const { return dtype_; } + float alpha() const { return alpha_; } + void set_alpha(float alpha) { alpha_ = alpha; } + + // Accessors for parameters + Tensor weight() const { return weight_; } + Tensor bias() const { return bias_; } + Tensor weight_scale() const { return weight_scale_; } + Tensor weight_zeros() const { return weight_zeros_; } + Tensor gidx() const { return gidx_; } + + std::shared_ptr get_quantization() const { return quantization_; } + void process_weights_after_loading() override; + +protected: + // Parameters + INFINICORE_NN_PARAMETER(weight); + INFINICORE_NN_PARAMETER(bias); + + INFINICORE_NN_PARAMETER(weight_scale); + INFINICORE_NN_PARAMETER(weight_zeros); + + INFINICORE_NN_PARAMETER(gidx); + +protected: + // Helper method for common forward computation + Tensor compute_linear(Tensor &input) const; + + size_t in_features_; + size_t out_features_; + bool has_bias_; + DataType dtype_; + float alpha_ = 1.0f; + std::shared_ptr quantization_ = std::make_shared(nullptr); +}; + +} // namespace infinicore::nn + +namespace infinicore::nn { + +class Linear : public BaseLinear { +public: + Linear(size_t in_features, size_t out_features, bool bias = true, + const DataType &dtype = DataType::F32, const Device &device = Device()); + + Linear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, + const DataType &dtype = DataType::F32, const Device &device = Device()); + + // Forward pass: output = input @ weight.T + bias + Tensor forward(Tensor &input) const; + + // String representation + std::string extra_repr() const; +}; + +class ColumnParallelLinear : public BaseLinear { +public: + ColumnParallelLinear(size_t in_features, size_t out_features, bool bias = true, + const DataType &dtype = DataType::F32, const Device &device = Device(), + Size tp_rank = 0, Size tp_size = 1); + + ColumnParallelLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, + const DataType &dtype = DataType::F32, const Device &device = Device(), + Size tp_rank = 0, Size tp_size = 1); + + // Forward pass: output = input @ weight.T + bias + Tensor forward(Tensor &input) const; + + // String representation + std::string extra_repr() const; + +protected: + Size tp_rank_ = 0; + Size tp_size_ = 1; +}; + +class RowParallelLinear : public BaseLinear { +public: + RowParallelLinear(size_t in_features, size_t out_features, bool bias = true, + const DataType &dtype = DataType::F32, const Device &device = Device(), + Size tp_rank = 0, Size tp_size = 1, infinicclComm_t communicator = nullptr); + + RowParallelLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, + const DataType &dtype = DataType::F32, const Device &device = Device(), + Size tp_rank = 0, Size tp_size = 1, infinicclComm_t communicator = nullptr); + + // Forward pass: output = input @ weight.T + bias + Tensor forward(Tensor &input) const; + + // String representation + std::string extra_repr() const; + +protected: + Size tp_rank_ = 0; + Size tp_size_ = 1; + infinicclComm_t communicator_; +}; + +} // namespace infinicore::nn diff --git a/python/infinicore/include/infinicore/nn/module.hpp b/python/infinicore/include/infinicore/nn/module.hpp new file mode 100644 index 000000000..2b5747a16 --- /dev/null +++ b/python/infinicore/include/infinicore/nn/module.hpp @@ -0,0 +1,174 @@ +#pragma once + +#include "../tensor.hpp" +#include "parameter.hpp" + +#include +#include +#include +#include +#include + +namespace infinicore::nn { +class Module { +public: + Module() = default; + + virtual ~Module() = default; + + std::unordered_map state_dict() const; + + std::vector state_dict_keys() const; + + void load_state_dict(const std::unordered_map &_state_dict); + + void load_parameter(const std::string &name, const Tensor ¶m); + + void load_parameters_no_sync(const std::unordered_map ¶ms, bool strict = true); + + void load_parameter_(const std::string &name, const Tensor ¶m); + + void load_parameter_from_blob(const std::string &name, const void *data); + + std::unordered_map modules_dict() const; + + virtual void process_weights_after_loading() {} + + virtual void reset_runtime_state() const {} + + const std::unordered_map> &children() const { + return submodules_; + } + +protected: + Tensor register_parameter(const std::string &name, Parameter param); + + Tensor register_buffer(const std::string &name, Parameter buffer); + + // Add an existing submodule to this module's hierarchy + // Template parameter M must be a type derived from Module + // Returns the submodule for convenience (allows method chaining) + template + std::shared_ptr add_module(const std::string &name, std::shared_ptr submodule) { + // Ensure M is derived from Module (compile-time check) + static_assert(std::is_base_of::value, + "Template parameter M must be derived from infinicore::nn::Module"); + + // Store in the submodules map (std::shared_ptr automatically converts to std::shared_ptr) + submodules_[name] = submodule; + + return submodule; + } + + // Create and register a new submodule by constructing it with the given arguments + // Template parameter M must be a type derived from Module + // Args are forwarded to M's constructor + template + std::shared_ptr register_module(const std::string &name, Args &&...args) { + // Ensure M is derived from Module (compile-time check) + static_assert(std::is_base_of::value, + "Template parameter M must be derived from infinicore::nn::Module"); + + // Construct the submodule + auto submodule = std::make_shared(std::forward(args)...); + + return add_module(name, submodule); + } + + // Create and register multiple submodules of the same type + // Each submodule is named as "name.0", "name.1", etc. + // Template parameter M must be a type derived from Module + template + std::vector> register_modules(size_t count, const std::string &name, Args &&...args) { + static_assert(std::is_base_of::value, + "Template parameter M must be derived from infinicore::nn::Module"); + + std::vector> modules; + modules.reserve(count); + for (size_t i = 0; i < count; i++) { + modules.push_back(register_module(name + "." + std::to_string(i), std::forward(args)...)); + } + return modules; + } + +protected: + Device device_; + std::unordered_map> submodules_; + std::unordered_map buffers_; + std::unordered_map parameters_; + +private: + void load_state_dict_recursively(const std::unordered_map &_state_dict, const std::string &prefix = ""); + void collect_all_parameters(std::unordered_map &all_params, const std::string &prefix = "") const; + void collect_all_parameter_names(std::vector &all_names, const std::string &prefix = "") const; + void collect_all_modules(std::unordered_map &out, const std::string &prefix) const; +}; + +// ============================================================================ +// PyTorch-like Macros for Convenient Module Registration +// ============================================================================ + +/** + * @brief Register submodules with automatic name inference from variable name + * + * Usage: + * @code + * class MyModel : public Module { + * protected: + * INFINICORE_NN_MODULE(Linear, layer1); + * INFINICORE_NN_MODULE(Linear, layer2); + * INFINICORE_NN_MODULE_VEC(Linear, layers); + * INFINICORE_NN_PARAMETER(scaling_factor); + * + * public: + * MyModel() { + * INFINICORE_NN_MODULE_INIT(layer1, 128, 64); + * INFINICORE_NN_MODULE_INIT(layer2, 64, 32); + * INFINICORE_NN_MODULE_VEC_INIT(layers, 3, Linear, 32, 16); + * INFINICORE_NN_PARAMETER_INIT(scaling_factor, ({1}, DataType::F32, Device())); + * } + * }; + * @endcode + */ + +// Declare a single module member variable +#define INFINICORE_NN_MODULE(ModuleType, name) \ + std::shared_ptr name##_ + +// Declare a vector of modules member variable +#define INFINICORE_NN_MODULE_VEC(ModuleType, name) \ + std::vector> name##_ + +// Initialize a module in constructor +#define INFINICORE_NN_MODULE_INIT(name, ...) \ + name##_ = this->register_module::type>(#name, ##__VA_ARGS__) + +// Initialize a vector of modules in constructor +// Usage: INFINICORE_NN_MODULE_VEC_INIT(layers, count, ModuleType, ctor_args...) +// Example: INFINICORE_NN_MODULE_VEC_INIT(layers, 3, Linear, 128, 64) +#define INFINICORE_NN_MODULE_VEC_INIT(name, count, ModuleType, ...) \ + name##_ = this->register_modules(count, #name, ##__VA_ARGS__) + +// Declare a parameter member variable +#define INFINICORE_NN_PARAMETER(name) \ + infinicore::nn::Parameter name##_ + +// Initialize a parameter in constructor +// Usage: INFINICORE_NN_PARAMETER_INIT(name, (shape, dtype, device)) +// Example: INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, DataType::F32, device)) +#define INFINICORE_NN_PARAMETER_INIT(name, args) \ + name##_ = infinicore::nn::Parameter args; \ + this->register_parameter(#name, name##_) + +// Declare a buffer member variable +#define INFINICORE_NN_BUFFER(name) \ + infinicore::nn::Parameter name##_ + +// Initialize a buffer in constructor +// Usage: INFINICORE_NN_BUFFER_INIT(name, (shape, dtype, device)) +// Example: INFINICORE_NN_BUFFER_INIT(cache, ({max_seq_len, head_dim}, DataType::F32, device)) +#define INFINICORE_NN_BUFFER_INIT(name, args) \ + name##_ = infinicore::nn::Parameter args; \ + this->register_buffer(#name, name##_) + +} // namespace infinicore::nn diff --git a/python/infinicore/include/infinicore/nn/parameter.hpp b/python/infinicore/include/infinicore/nn/parameter.hpp new file mode 100644 index 000000000..1602f58d6 --- /dev/null +++ b/python/infinicore/include/infinicore/nn/parameter.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include "../tensor.hpp" + +namespace infinicore::nn { +class Parameter : public Tensor { +public: + Parameter(); + + Parameter(const Tensor &tensor, + Size tp_dim = 0, + Size tp_rank = 0, + Size tp_size = 1, + Size num_shards = 0); + + Parameter(const Shape &shape, + const DataType &dtype, + const Device &device, + Size tp_dim = 0, + Size tp_rank = 0, + Size tp_size = 1, + Size num_shards = 0); + + Parameter(const Parameter &other); + + void load_blob(const void *data); + + void load(const Tensor &tensor); + + void load_no_sync(const Tensor &tensor); + +protected: + // Tensor parallel configs + Size tp_dim_; // dimension partitioned + Size tp_rank_; // rank of this partition among tp group + Size tp_size_; // total number of partitions + Size num_shards_ = 0; // number of logical shards, used when tp_size > num_kv_head +}; +} // namespace infinicore::nn diff --git a/python/infinicore/include/infinicore/nn/rmsnorm.hpp b/python/infinicore/include/infinicore/nn/rmsnorm.hpp new file mode 100644 index 000000000..5891819eb --- /dev/null +++ b/python/infinicore/include/infinicore/nn/rmsnorm.hpp @@ -0,0 +1,96 @@ +#pragma once + +#include "../ops.hpp" +#include "module.hpp" + +namespace infinicore::nn { + +/** + * @brief Root Mean Square Layer Normalization (RMSNorm) + * + * Applies Root Mean Square Layer Normalization over the last dimension. + * Unlike LayerNorm, RMSNorm doesn't subtract mean and doesn't use bias. + * + * Formula: y = (x / RMS(x)) * weight + * where RMS(x) = sqrt(mean(x^2) + eps) + * + * Used in LLaMA, Galactica, and other modern language models as a + * simpler and faster alternative to LayerNorm. + * + * Example: + * @code + * // Create RMSNorm for hidden size 4096 + * auto norm = RMSNorm(4096); + * + * // Input: [batch, seq_len, hidden_size] + * auto input = Tensor::randn({2, 10, 4096}); + * + * // Output: [batch, seq_len, hidden_size] + * auto output = norm.forward(input); + * @endcode + */ +class RMSNorm : public Module { +public: + /** + * @brief Construct a RMSNorm layer + * + * @param normalized_shape Size of the feature dimension to normalize (typically hidden_size) + * @param eps Small constant for numerical stability (default: 1e-6) + * @param dtype Data type for the weight (default: DataType::F32) + * @param device Device to create the weight on + */ + RMSNorm(size_t normalized_shape, + double eps = 1e-6, + const DataType &dtype = DataType::F32, + const Device &device = Device()); + + /** + * @brief Forward pass: apply RMSNorm + * + * @param x Input tensor of shape (*, normalized_shape) where * is any number of dimensions + * @return Normalized tensor with same shape as input + * + * The normalization is applied over the last dimension. + * For example: + * Input: [batch, seq_len, hidden_size] -> normalize over hidden_size + * Input: [batch, hidden_size] -> normalize over hidden_size + */ + Tensor forward(const Tensor &x) const; + + /** + * @brief Forward pass: apply RMSNorm in-place with residual + * + * @param x Input tensor of shape (*, normalized_shape) where * is any number of dimensions. + * Will be modified in-place to the normalized output. + * @param residual Residual tensor to add to input before normalization. + * Will be modified in-place to the sum of input and residual. + * + * The normalization is applied over the last dimension. + * For example: + * Input: [batch, seq_len, hidden_size] -> normalize over hidden_size + * Input: [batch, hidden_size] -> normalize over hidden_size + */ + void forward_inplace(Tensor &x, Tensor &residual) const; + + // Module information + size_t normalized_shape() const { return normalized_shape_; } + double eps() const { return eps_; } + DataType dtype() const { return dtype_; } + + // String representation + std::string extra_repr() const; + + // Accessors for parameters + Tensor weight() const { return weight_; } + +protected: + // Parameters + INFINICORE_NN_PARAMETER(weight); + +private: + size_t normalized_shape_; // Size of the feature dimension + double eps_; // Epsilon for numerical stability + DataType dtype_; // Data type for weight +}; + +} // namespace infinicore::nn diff --git a/python/infinicore/include/infinicore/nn/rope.hpp b/python/infinicore/include/infinicore/nn/rope.hpp new file mode 100644 index 000000000..eaeba8712 --- /dev/null +++ b/python/infinicore/include/infinicore/nn/rope.hpp @@ -0,0 +1,114 @@ +#pragma once + +#include "../context/context.hpp" +#include "../tensor.hpp" +#include "module.hpp" +#include "rope_scaling_configs.hpp" +#include +#include +#include +#include +#include + +namespace infinicore::nn { + +class RoPE : public Module { +public: + /** + * @brief RoPE algorithm type + */ + enum class Algo { + GPT_J = 0, // GPT-J style RoPE algorithm (Interleave even and odd dimensions) + GPT_NEOX = 1, // GPT-NeoX style RoPE algorithm (First half dimensions for sin, second half for cos) + }; + + /** + * @brief Construct a RoPE layer + * + * @param head_dim Dimension of each attention head (must be even) + * @param rotary_dim Number of dimensions to apply rotation to (must be even). + * For full rotation models, this equals head_dim; + * for partial rotation models, this equals head_dim * partial_rotary_factor. + * @param max_seq_len Maximum sequence length for pre-computed cache + * @param theta Base frequency for rotary embeddings (default: 10000.0) + * @param algo RoPE algorithm type (default: Algo::GPT_J) + * @param dtype Data type for sin/cos cache (default: DataType::F32) + * @param device Device to create the cache on + * @param scaling RoPE scaling configuration (default: nullptr) + * @param mrope_section Optional MRoPE section sizes [t, h, w], whose sum must equal rotary_dim / 2. + * When set, pair forward overloads apply MRoPE to q/k using positions [3, num_tokens]. + * @param mrope_interleaved Whether to interleave MRoPE axes/frequency sections. + */ + RoPE(size_t head_dim, + size_t rotary_dim, + size_t max_seq_len, + double theta = 10000.0, + Algo algo = Algo::GPT_J, + const DataType &dtype = DataType::F32, + const Device &device = Device(), + std::shared_ptr scaling = nullptr, + std::optional> mrope_section = std::nullopt, + bool mrope_interleaved = false); + + /** + * @brief Forward pass: apply standard RoPE to a tensor + * + * @param x Input tensor of shape (..., rotary_dim) where ... is any number of dimensions + * @param pos Position IDs tensor of shape (*,) typically [seq_len] or [batch, seq_len] + * @param in_place If true, modify input tensor in place (default: false) + * @return Rotated tensor with same shape as input + */ + Tensor forward(const Tensor &x, const Tensor &pos, bool in_place = false) const; + + /** + * @brief Apply MRoPE to q and k. + * + * Requires construction with mrope_section. q/k may be either + * [num_tokens, num_heads * head_dim] or [num_tokens, num_heads, head_dim]. + * positions is [3, num_tokens] with axes ordered as t, h, w. + */ + std::pair forward(const Tensor &q, const Tensor &k, const Tensor &positions) const; + + /** + * @brief Apply MRoPE to q and k into caller-provided outputs. + */ + std::pair forward(const Tensor &q_out, + const Tensor &k_out, + const Tensor &q, + const Tensor &k, + const Tensor &positions) const; + + // Module information + size_t rotary_dim() const { return rotary_dim_; } + size_t head_dim() const { return head_dim_; } + size_t max_seq_len() const { return max_seq_len_; } + double theta() const { return theta_; } + Algo algo() const { return algo_; } + DataType dtype() const { return dtype_; } + const std::optional> &mrope_section() const { return mrope_section_; } + bool mrope_interleaved() const { return mrope_interleaved_; } + const Tensor &sin_cache() const { return sin_cache_; } + const Tensor &cos_cache() const { return cos_cache_; } + + // String representation + std::string extra_repr() const; + +protected: + // Buffers (sin and cos cache tables) - not exposed in state_dict + INFINICORE_NN_BUFFER(sin_cache); + INFINICORE_NN_BUFFER(cos_cache); + +private: + void initialize_cache(); + size_t rotary_dim_; // Number of dimensions to apply rotation to (must be even). + size_t head_dim_; // Dimension of each attention head + size_t max_seq_len_; // Maximum sequence length + double theta_; // Base frequency for rotary embeddings + Algo algo_; // RoPE algorithm type + DataType dtype_; // Data type for cache tables + std::shared_ptr scaling_; // RoPE scaling configuration + std::optional> mrope_section_; + bool mrope_interleaved_; +}; + +} // namespace infinicore::nn diff --git a/python/infinicore/include/infinicore/nn/rope_scaling_configs.hpp b/python/infinicore/include/infinicore/nn/rope_scaling_configs.hpp new file mode 100644 index 000000000..4129e6053 --- /dev/null +++ b/python/infinicore/include/infinicore/nn/rope_scaling_configs.hpp @@ -0,0 +1,131 @@ +#pragma once +#include +#include + +namespace infinicore::nn { + +/** + * @brief Abstract base class for RoPE scaling strategies. + * Uses polymorphism to eliminate type checking (if-else) in the core RoPE loop. + */ +class RopeScalingConfig { +public: + virtual ~RopeScalingConfig() = default; + + /** + * @brief Calculate the frequency scaling factor for a specific position and dimension. + * + * @param pos Current sequence position + * @param dim_idx Current dimension index (0 to head_dim/2 - 1) + * @param base_inv_freq Pre-computed base inverse frequency for this dimension (1.0 / theta^(2j/head_dim)) + * @return Frequency scaling factor (default 1.0) + */ + virtual float get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const { + return 1.0f; + } + + /** + * @brief Calculate the magnitude scaling factor for a specific position and dimension. + * + * @param pos Current sequence position + * @param dim_idx Current dimension index (0 to head_dim/2 - 1) + * @param base_inv_freq Pre-computed base inverse frequency for this dimension + * @return Magnitude scaling factor (default 1.0) + */ + virtual float get_magnitude_scale(size_t pos, size_t dim_idx, float base_inv_freq) const { + return 1.0f; + } +}; + +/** + * @brief LongRoPE scaling configuration. + */ +class LongRopeScalingConfig : public RopeScalingConfig { +public: + LongRopeScalingConfig( + std::vector short_factor, + std::vector long_factor, + size_t original_max_position_embeddings, + float factor = 1.0f); + + float get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const override; + float get_magnitude_scale(size_t pos, size_t dim_idx, float base_inv_freq) const override; + + size_t original_max_position_embeddings() const { return original_max_position_embeddings_; } + const std::vector &short_factor() const { return short_factor_; } + const std::vector &long_factor() const { return long_factor_; } + float factor() const { return factor_; } + +private: + std::vector short_factor_; + std::vector long_factor_; + size_t original_max_position_embeddings_; + float factor_; +}; + +// TODO(rubik) implement in cpp +/** + * @brief Llama3 frequency-aware RoPE scaling configuration. + * Native support for Llama 3.1 RoPE scaling (smooth interpolation based on wavelength). + */ +class Llama3RopeScalingConfig : public RopeScalingConfig { +public: + Llama3RopeScalingConfig( + float factor, + float low_freq_factor, + float high_freq_factor, + size_t original_max_position_embeddings); + + float get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const override; + + // Llama3 does not use magnitude scaling, so it inherits the default get_magnitude_scale() returning 1.0f + +private: + float factor_; + float low_freq_factor_; + float high_freq_factor_; + size_t original_max_position_embeddings_; +}; + +/** + * @brief YaRN (Yet another RoPE extensioN) scaling configuration. + * + * rope_scaling fields: factor, original_max_position_embeddings, beta_fast, beta_slow, + * mscale, mscale_all_dim + * Model fields (must match RoPE): rotary_dim (e.g. qk_rope_head_dim), rope_theta + */ +class YarnRopeScalingConfig : public RopeScalingConfig { +public: + YarnRopeScalingConfig( + float factor, + size_t original_max_position_embeddings, + size_t rotary_dim, + float rope_theta, + int beta_fast = 32, + int beta_slow = 1, + float mscale = 1.0f, + float mscale_all_dim = 0.0f); + + float get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const override; + float get_magnitude_scale(size_t pos, size_t dim_idx, float base_inv_freq) const override; + + /** Recommended RoPE cache length: original_max_position_embeddings * factor. */ + static size_t max_seq_len(float factor, size_t original_max_position_embeddings) { + return static_cast( + static_cast(original_max_position_embeddings) * factor); + } + + float factor() const { return factor_; } + size_t original_max_position_embeddings() const { return original_max_position_embeddings_; } + +private: + float yarn_linear_ramp(size_t dim_idx) const; + + float factor_; + size_t original_max_position_embeddings_; + float magnitude_scale_; + float correction_low_; + float correction_high_; +}; + +} // namespace infinicore::nn diff --git a/python/infinicore/include/infinicore/ops.hpp b/python/infinicore/include/infinicore/ops.hpp new file mode 100644 index 000000000..3f56203d3 --- /dev/null +++ b/python/infinicore/include/infinicore/ops.hpp @@ -0,0 +1,84 @@ +#pragma once + +#include "ops/adaptive_max_pool1d.hpp" +#include "ops/add.hpp" +#include "ops/add_rms_norm.hpp" +#include "ops/addcmul.hpp" +#include "ops/asin.hpp" +#include "ops/asinh.hpp" +#include "ops/asum.hpp" +#include "ops/atanh.hpp" +#include "ops/attention.hpp" +#include "ops/avg_pool1d.hpp" +#include "ops/axpy.hpp" +#include "ops/baddbmm.hpp" +#include "ops/bilinear.hpp" +#include "ops/binary_cross_entropy_with_logits.hpp" +#include "ops/blas_amax.hpp" +#include "ops/blas_amin.hpp" +#include "ops/blas_copy.hpp" +#include "ops/blas_dot.hpp" +#include "ops/causal_conv1d.hpp" +#include "ops/causal_softmax.hpp" +#include "ops/cdist.hpp" +#include "ops/chunk_gated_delta_rule.hpp" +#include "ops/conv2d.hpp" +#include "ops/cross_entropy.hpp" +#include "ops/deepseek_moe.hpp" +#include "ops/dsa.hpp" +#include "ops/embedding.hpp" +#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/gelu.hpp" +#include "ops/gelutanh.hpp" +#include "ops/hardswish.hpp" +#include "ops/hardtanh.hpp" +#include "ops/kv_caching.hpp" +#include "ops/layer_norm.hpp" +#include "ops/linear.hpp" +#include "ops/mamba_selective_scan.hpp" +#include "ops/matmul.hpp" +#include "ops/moe_align.hpp" +#include "ops/moe_fused_dense.hpp" +#include "ops/moe_fused_gate.hpp" +#include "ops/moe_sum.hpp" +#include "ops/moe_topk_sigmoid.hpp" +#include "ops/moe_topk_softmax.hpp" +#include "ops/nrm2.hpp" +#include "ops/ones.hpp" +#include "ops/paged_attention.hpp" +#include "ops/paged_attention_mla.hpp" +#include "ops/paged_attention_prefill.hpp" +#include "ops/paged_caching.hpp" +#include "ops/per_tensor_dequant_i8.hpp" +#include "ops/per_tensor_quant_i8.hpp" +#include "ops/prepare_moe_input.hpp" +#include "ops/quickgelu.hpp" +#include "ops/random_sample.hpp" +#include "ops/rearrange.hpp" +#include "ops/reciprocal.hpp" +#include "ops/recurrent_gated_delta_rule.hpp" +#include "ops/relu.hpp" +#include "ops/rms_norm.hpp" +#include "ops/rope.hpp" +#include "ops/rot.hpp" +#include "ops/rotg.hpp" +#include "ops/rotm.hpp" +#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" +#include "ops/softmax.hpp" +#include "ops/swap.hpp" +#include "ops/swiglu.hpp" +#include "ops/topksoftmax.hpp" +#include "ops/w16a16_group_gemm.hpp" diff --git a/python/infinicore/include/infinicore/ops/acos.hpp b/python/infinicore/include/infinicore/ops/acos.hpp new file mode 100644 index 000000000..91aaee020 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/acos.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Acos { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor acos(Tensor input); +void acos_(Tensor output, Tensor input); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/adaptive_avg_pool1d.hpp b/python/infinicore/include/infinicore/ops/adaptive_avg_pool1d.hpp new file mode 100644 index 000000000..3e8a95b8d --- /dev/null +++ b/python/infinicore/include/infinicore/ops/adaptive_avg_pool1d.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class AdaptiveAvgPool1d { +public: + // Schema: execute(Output, Input) + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor adaptive_avg_pool1d(Tensor input, int64_t output_size); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/adaptive_avg_pool3d.hpp b/python/infinicore/include/infinicore/ops/adaptive_avg_pool3d.hpp new file mode 100644 index 000000000..df6a4d2c3 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/adaptive_avg_pool3d.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class AdaptiveAvgPool3D { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor y, Tensor x); + static common::OpDispatcher &dispatcher(); +}; + +Tensor adaptive_avg_pool3d(Tensor x, std::vector output_size); +void adaptive_avg_pool3d_(Tensor y, Tensor x); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/adaptive_max_pool1d.hpp b/python/infinicore/include/infinicore/ops/adaptive_max_pool1d.hpp new file mode 100644 index 000000000..51832cae0 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/adaptive_max_pool1d.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class AdaptiveMaxPool1d { +public: + using schema = void (*)(Tensor, Tensor, size_t); + static void execute(Tensor y, Tensor x, size_t output_size); + static common::OpDispatcher &dispatcher(); +}; + +Tensor adaptive_max_pool1d(Tensor x, size_t output_size); +void adaptive_max_pool1d_(Tensor y, Tensor x, size_t output_size); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/add.hpp b/python/infinicore/include/infinicore/ops/add.hpp new file mode 100644 index 000000000..528cca18a --- /dev/null +++ b/python/infinicore/include/infinicore/ops/add.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Add, Tensor, const Tensor &, const Tensor &); + +Tensor add(const Tensor &a, const Tensor &b); +void add_(Tensor c, const Tensor &a, const Tensor &b); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/add_rms_norm.hpp b/python/infinicore/include/infinicore/ops/add_rms_norm.hpp new file mode 100644 index 000000000..50064e0a4 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/add_rms_norm.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { +INFINICORE_GRAPH_OP_CLASS(AddRMSNorm, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &, float); + +// Fused Add and RMS Normalization +// Returns: (normalized_result, add_result) +// The add_result can be used as residual for subsequent layers +std::pair add_rms_norm(const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon = 1e-5f); +void add_rms_norm_(Tensor out, Tensor residual, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon = 1e-5f); +// Fused Add and RMS Normalization (inplace) +// normalized_result wil be stored in input, add_result will be stored in residual +void add_rms_norm_inplace(Tensor input, Tensor residual, const Tensor &weight, float epsilon = 1e-5f); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/addbmm.hpp b/python/infinicore/include/infinicore/ops/addbmm.hpp new file mode 100644 index 000000000..6c17a35d5 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/addbmm.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Addbmm { +public: + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float, float); + static void execute(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha); + + static common::OpDispatcher &dispatcher(); +}; +Tensor addbmm(Tensor input, Tensor batch1, Tensor batch2, float beta = 1.0f, float alpha = 1.0f); + +void addbmm_(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/addcmul.hpp b/python/infinicore/include/infinicore/ops/addcmul.hpp new file mode 100644 index 000000000..2806d1218 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/addcmul.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Addcmul { +public: + // schema: out, input, t1, t2, value + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float); + static void execute(Tensor out, Tensor input, Tensor t1, Tensor t2, float value); + static common::OpDispatcher &dispatcher(); +}; +Tensor addcmul(Tensor input, Tensor t1, Tensor t2, float value); +void addcmul_(Tensor out, Tensor input, Tensor t1, Tensor t2, float value); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/addr.hpp b/python/infinicore/include/infinicore/ops/addr.hpp new file mode 100644 index 000000000..cff78eade --- /dev/null +++ b/python/infinicore/include/infinicore/ops/addr.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Addr { +public: + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float, float); + static void execute(Tensor out, Tensor input, Tensor vec1, Tensor vec2, float beta = 1.0f, float alpha = 1.0f); + static common::OpDispatcher &dispatcher(); +}; + +Tensor addr(Tensor input, Tensor vec1, Tensor vec2, float beta = 1.0f, float alpha = 1.0f); +void addr_(Tensor out, Tensor input, Tensor vec1, Tensor vec2, float beta = 1.0f, float alpha = 1.0f); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/affine_grid.hpp b/python/infinicore/include/infinicore/ops/affine_grid.hpp new file mode 100644 index 000000000..ea025c24d --- /dev/null +++ b/python/infinicore/include/infinicore/ops/affine_grid.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +class AffineGrid { +public: + using schema = void (*)(Tensor, Tensor, bool); + static void execute(Tensor output, Tensor theta, bool align_corners); + static common::OpDispatcher &dispatcher(); +}; +Tensor affine_grid(Tensor theta, const std::vector &size, bool align_corners = false); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/all.hpp b/python/infinicore/include/infinicore/ops/all.hpp new file mode 100644 index 000000000..50d76f2d7 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/all.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include +#include +namespace infinicore::op { +class All { +public: + using schema = void (*)(Tensor, Tensor, std::vector, bool); + static void execute(Tensor output, Tensor input, std::vector dim, bool keepdim = false); + static common::OpDispatcher &dispatcher(); +}; + +Tensor all(Tensor input, std::vector dim, bool keepdim = false); +void all_(Tensor output, Tensor input, std::vector dim, bool keepdim = false); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/argwhere.hpp b/python/infinicore/include/infinicore/ops/argwhere.hpp new file mode 100644 index 000000000..24f4fa65e --- /dev/null +++ b/python/infinicore/include/infinicore/ops/argwhere.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Argwhere { +public: + using schema = void (*)(void **, size_t *, Tensor); + static void execute(void **, size_t *count, Tensor x); + static common::OpDispatcher &dispatcher(); +}; +Tensor argwhere(Tensor x); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/asin.hpp b/python/infinicore/include/infinicore/ops/asin.hpp new file mode 100644 index 000000000..bc5973978 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/asin.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Asin { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor asin(Tensor input); +void asin_(Tensor output, Tensor input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/asinh.hpp b/python/infinicore/include/infinicore/ops/asinh.hpp new file mode 100644 index 000000000..505eb97d9 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/asinh.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Asinh { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor y, Tensor x); + static common::OpDispatcher &dispatcher(); +}; + +Tensor asinh(Tensor x); +void asinh_(Tensor y, Tensor x); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/asum.hpp b/python/infinicore/include/infinicore/ops/asum.hpp new file mode 100644 index 000000000..6471e6fe4 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/asum.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Asum, const Tensor &, Tensor); + +Tensor asum(const Tensor &x); +void asum_(const Tensor &x, Tensor result); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/atanh.hpp b/python/infinicore/include/infinicore/ops/atanh.hpp new file mode 100644 index 000000000..5c3894e5c --- /dev/null +++ b/python/infinicore/include/infinicore/ops/atanh.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Atanh { +public: + // schema 定义为：void(输出 Tensor, 输入 Tensor) + using schema = void (*)(Tensor, Tensor); + + // 执行函数 + static void execute(Tensor y, Tensor a); + + // 获取算子分发器，用于多后端（CPU/CUDA 等）匹配 + static common::OpDispatcher &dispatcher(); +}; + +/** + * @brief 计算输入 Tensor 的反双曲正切值 (out-of-place) + * @param a 输入 Tensor + * @return 包含结果的新 Tensor + */ +Tensor atanh(Tensor a); + +/** + * @brief 计算输入 Tensor 的反双曲正切值 (in-place / specified output) + * @param y 输出 Tensor + * @param a 输入 Tensor + */ +void atanh_(Tensor y, Tensor a); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/attention.hpp b/python/infinicore/include/infinicore/ops/attention.hpp new file mode 100644 index 000000000..1bc447c77 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/attention.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Attention { +public: + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, size_t); + static void execute(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos); + static common::OpDispatcher &dispatcher(); +}; + +Tensor attention(Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos); +void attention_(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/avg_pool1d.hpp b/python/infinicore/include/infinicore/ops/avg_pool1d.hpp new file mode 100644 index 000000000..4bf69bc2a --- /dev/null +++ b/python/infinicore/include/infinicore/ops/avg_pool1d.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class AvgPool1d { +public: + using schema = void (*)(Tensor, Tensor, size_t, size_t, size_t); + static void execute(Tensor output, Tensor input, size_t kernel_size, size_t stride, size_t padding); + static common::OpDispatcher &dispatcher(); +}; + +Tensor avg_pool1d(Tensor input, size_t kernel_size, size_t stride = 0, size_t padding = 0); +void avg_pool1d_(Tensor output, Tensor input, size_t kernel_size, size_t stride = 0, size_t padding = 0); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/awq_marlin_gemm.hpp b/python/infinicore/include/infinicore/ops/awq_marlin_gemm.hpp new file mode 100644 index 000000000..164ee016b --- /dev/null +++ b/python/infinicore/include/infinicore/ops/awq_marlin_gemm.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(AwqMarlinGemm, Tensor, const Tensor &, const Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, int64_t, bool, bool, bool, bool); + +void awq_marlin_gemm_(Tensor c, const Tensor &a, const Tensor &b, Tensor &b_bias, Tensor &b_scales, Tensor &a_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/axpy.hpp b/python/infinicore/include/infinicore/ops/axpy.hpp new file mode 100644 index 000000000..280d5ab60 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/axpy.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(Axpy, const Tensor &, const Tensor &, Tensor); + +void axpy_(const Tensor &alpha, const Tensor &x, Tensor y); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/baddbmm.hpp b/python/infinicore/include/infinicore/ops/baddbmm.hpp new file mode 100644 index 000000000..9988000bd --- /dev/null +++ b/python/infinicore/include/infinicore/ops/baddbmm.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +Tensor baddbmm(Tensor input, Tensor batch1, Tensor batch2, + float beta = 1.0f, + float alpha = 1.0f); +void baddbmm_(Tensor out, Tensor input, Tensor batch1, Tensor batch2, + float beta = 1.0f, + float alpha = 1.0f); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/bilinear.hpp b/python/infinicore/include/infinicore/ops/bilinear.hpp new file mode 100644 index 000000000..805fa1efc --- /dev/null +++ b/python/infinicore/include/infinicore/ops/bilinear.hpp @@ -0,0 +1,12 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +Tensor bilinear(Tensor x1, Tensor x2, Tensor weight, std::optional bias); +void bilinear_(Tensor out, Tensor x1, Tensor x2, Tensor weight, std::optional bias); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/binary_cross_entropy_with_logits.hpp b/python/infinicore/include/infinicore/ops/binary_cross_entropy_with_logits.hpp new file mode 100644 index 000000000..4cdead724 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/binary_cross_entropy_with_logits.hpp @@ -0,0 +1,46 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +class BinaryCrossEntropyWithLogits { +public: + /** + * @brief BCEWithLogits 算子的函数原型 + * 参数顺序: out, logits, target, weight, pos_weight, reduction + */ + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, std::string); + + static void execute(Tensor out, + Tensor logits, + Tensor target, + Tensor weight, + Tensor pos_weight, + std::string reduction); + + static common::OpDispatcher &dispatcher(); +}; + +/** + * @brief 非原地操作接口 (Out-of-place) + */ +Tensor binary_cross_entropy_with_logits(Tensor logits, + Tensor target, + Tensor weight = {}, + Tensor pos_weight = {}, + std::string reduction = "mean"); + +/** + * @brief 显式指定输出张量的接口 + */ +void binary_cross_entropy_with_logits_(Tensor out, + Tensor logits, + Tensor target, + Tensor weight, + Tensor pos_weight, + std::string reduction); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/bitwise_right_shift.hpp b/python/infinicore/include/infinicore/ops/bitwise_right_shift.hpp new file mode 100644 index 000000000..db8d6a2e1 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/bitwise_right_shift.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "infinicore.h" + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(BitwiseRightShift, Tensor, const Tensor &, const Tensor &); + +__export Tensor bitwise_right_shift(const Tensor &input, const Tensor &other); +__export void bitwise_right_shift_(Tensor out, const Tensor &input, const Tensor &other); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/blas_amax.hpp b/python/infinicore/include/infinicore/ops/blas_amax.hpp new file mode 100644 index 000000000..a6a571f95 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/blas_amax.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(BlasAmax, const Tensor &, Tensor); + +Tensor blas_amax(const Tensor &x); +void blas_amax_(const Tensor &x, Tensor result); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/blas_amin.hpp b/python/infinicore/include/infinicore/ops/blas_amin.hpp new file mode 100644 index 000000000..a2ed21c7b --- /dev/null +++ b/python/infinicore/include/infinicore/ops/blas_amin.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(BlasAmin, const Tensor &, Tensor); + +Tensor blas_amin(const Tensor &x); +void blas_amin_(const Tensor &x, Tensor result); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/blas_copy.hpp b/python/infinicore/include/infinicore/ops/blas_copy.hpp new file mode 100644 index 000000000..dd32646a0 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/blas_copy.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(BlasCopy, const Tensor &, Tensor); + +void blas_copy_(const Tensor &x, Tensor y); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/blas_dot.hpp b/python/infinicore/include/infinicore/ops/blas_dot.hpp new file mode 100644 index 000000000..157c167f6 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/blas_dot.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(BlasDot, const Tensor &, const Tensor &, Tensor); + +Tensor blas_dot(const Tensor &x, const Tensor &y); +void blas_dot_(const Tensor &x, const Tensor &y, Tensor result); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/block_diag.hpp b/python/infinicore/include/infinicore/ops/block_diag.hpp new file mode 100644 index 000000000..89e227509 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/block_diag.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +class BlockDiag { +public: + using schema = void (*)(Tensor, const std::vector &); + static void execute(Tensor output, const std::vector &inputs); + static common::OpDispatcher &dispatcher(); +}; + +Tensor block_diag(const std::vector &inputs); +void block_diag_(Tensor output, const std::vector &inputs); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/bmm_strided.hpp b/python/infinicore/include/infinicore/ops/bmm_strided.hpp new file mode 100644 index 000000000..6610ef443 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/bmm_strided.hpp @@ -0,0 +1,12 @@ +#pragma once + +#include "../tensor.hpp" + +namespace infinicore::op { + +// Batched matrix multiplication that preserves arbitrary valid tensor strides, +// including a transposed output view. The Iluvatar implementation deliberately +// mirrors vLLM's torch.bmm(..., out=transpose_view) MLA projection path. +void bmm_strided_(Tensor output, const Tensor &a, const Tensor &b); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/broadcast_to.hpp b/python/infinicore/include/infinicore/ops/broadcast_to.hpp new file mode 100644 index 000000000..94d98e065 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/broadcast_to.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(BroadcastTo, Tensor, Tensor); + +Tensor broadcast_to(Tensor x, const std::vector &shape); +void broadcast_to_(Tensor y, Tensor x); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/cast.hpp b/python/infinicore/include/infinicore/ops/cast.hpp new file mode 100644 index 000000000..4aa907922 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/cast.hpp @@ -0,0 +1,5 @@ +#pragma once +#include "../tensor.hpp" +namespace infinicore::op { +void cast_(Tensor output, const Tensor &input); +} diff --git a/python/infinicore/include/infinicore/ops/cat.hpp b/python/infinicore/include/infinicore/ops/cat.hpp new file mode 100644 index 000000000..95be62103 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/cat.hpp @@ -0,0 +1,9 @@ +#pragma once + +#include "common/op.hpp" + +namespace infinicore::op { + +Tensor cat(std::vector tensors, int dim); +void cat_(Tensor out, std::vector tensors, int dim); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/causal_conv1d.hpp b/python/infinicore/include/infinicore/ops/causal_conv1d.hpp new file mode 100644 index 000000000..ac20b779d --- /dev/null +++ b/python/infinicore/include/infinicore/ops/causal_conv1d.hpp @@ -0,0 +1,41 @@ +#pragma once + +#include "infinicore.h" + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(CausalConv1d, + Tensor, + Tensor, + std::optional, + const Tensor &, + const Tensor &, + std::optional, + std::optional, + std::optional, + std::optional); + +__export Tensor causal_conv1d(const Tensor &qkv, + Tensor conv_state, + const Tensor &weight, + std::optional bias = std::nullopt, + std::optional cu_seqlens = std::nullopt, + std::optional initial_state_indices = std::nullopt, + std::optional final_state_indices = std::nullopt); + +__export void causal_conv1d_(Tensor out, + Tensor conv_state, + std::optional final_conv_state, + const Tensor &qkv, + const Tensor &weight, + std::optional bias, + std::optional cu_seqlens, + std::optional initial_state_indices, + std::optional final_state_indices); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/causal_softmax.hpp b/python/infinicore/include/infinicore/ops/causal_softmax.hpp new file mode 100644 index 000000000..2646852af --- /dev/null +++ b/python/infinicore/include/infinicore/ops/causal_softmax.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(CausalSoftmax, Tensor, const Tensor &); + +Tensor causal_softmax(const Tensor &input); +void causal_softmax_(Tensor output, const Tensor &input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/cdist.hpp b/python/infinicore/include/infinicore/ops/cdist.hpp new file mode 100644 index 000000000..7bc4902a0 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/cdist.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Cdist { +public: + /** + * @brief 成对距离计算算子 (Pairwise distance) + * schema: out (M, N), x1 (M, D), x2 (N, D), p (norm degree) + */ + using schema = void (*)(Tensor, Tensor, Tensor, double); + + static void execute(Tensor out, Tensor x1, Tensor x2, double p); + + static common::OpDispatcher &dispatcher(); +}; + +/** + * @brief 非原地（Out-of-place）接口 + * @return 返回形状为 (M, N) 的新 Tensor + */ +Tensor cdist(Tensor x1, Tensor x2, double p = 2.0); + +/** + * @brief 显式指定输出接口 + */ +void cdist_(Tensor out, Tensor x1, Tensor x2, double p = 2.0); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/chunk_gated_delta_rule.hpp b/python/infinicore/include/infinicore/ops/chunk_gated_delta_rule.hpp new file mode 100644 index 000000000..5102836e0 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/chunk_gated_delta_rule.hpp @@ -0,0 +1,53 @@ +#pragma once + +#include "infinicore.h" + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(ChunkGatedDeltaRule, + Tensor, + Tensor, + std::optional, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + std::optional, + std::optional, + std::optional, + bool, + size_t); + +__export Tensor chunk_gated_delta_rule(const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &g, + const Tensor &beta, + Tensor initial_state, + std::optional cu_seqlens = std::nullopt, + std::optional initial_state_indices = std::nullopt, + std::optional final_state_indices = std::nullopt, + bool use_qk_l2norm = false, + size_t chunk_size = 64); + +__export void chunk_gated_delta_rule_(Tensor out, + Tensor initial_state, + std::optional final_state, + const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &g, + const Tensor &beta, + std::optional cu_seqlens, + std::optional initial_state_indices, + std::optional final_state_indices, + bool use_qk_l2norm = false, + size_t chunk_size = 64); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/common/cache.hpp b/python/infinicore/include/infinicore/ops/common/cache.hpp new file mode 100644 index 000000000..ad3aa88b9 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/common/cache.hpp @@ -0,0 +1,83 @@ +#pragma once + +#include "../../common/LRUCache.hpp" +#include "../../context/context.hpp" +#include +#include +#include +#include + +namespace infinicore::op::common { + +template +class OpCache { +private: + using BaseCache = infinicore::common::LRUCache; + using Destructor = typename BaseCache::Destructor; + using CacheVector = std::vector; + +public: + explicit OpCache(size_t capacity = 100, Destructor destructor = nullptr) + : capacity_(capacity), destructor_(destructor) {} + + ~OpCache() { + clear(); + } + + BaseCache &getCache(Device::Type device_type, size_t device_index) { + auto &cache_vector = caches_[static_cast(device_type)]; + + if (cache_vector.size() <= device_index) { + cache_vector.resize(device_index + 1, BaseCache(capacity_, destructor_)); + } else { + cache_vector[device_index].setDestructor(destructor_); + } + + return cache_vector[device_index]; + } + + BaseCache &getCache(Device device) { + return getCache(device.getType(), device.getIndex()); + } + + void setCapacity(size_t capacity) { + capacity_ = capacity; + for (auto &vec : caches_) { + for (auto &cache : vec) { + cache.setCapacity(capacity); + } + } + } + + void clear() { + Device current_device = context::getDevice(); + + for (size_t type_idx = 0; type_idx < caches_.size(); ++type_idx) { + auto &vec = caches_[type_idx]; + for (size_t dev_idx = 0; dev_idx < vec.size(); ++dev_idx) { + Device target_device(static_cast(type_idx), dev_idx); + + if (current_device != target_device) { + context::setDevice(target_device); + } + + vec[dev_idx].clear(); + + if (current_device != target_device) { + context::setDevice(current_device); + } + } + vec.clear(); + } + + caches_ = {}; + } + +private: + size_t capacity_; + Destructor destructor_; + + std::array(Device::Type::COUNT)> caches_ = {}; +}; + +} // namespace infinicore::op::common diff --git a/python/infinicore/include/infinicore/ops/common/dispatcher.hpp b/python/infinicore/include/infinicore/ops/common/dispatcher.hpp new file mode 100644 index 000000000..9c5cac179 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/common/dispatcher.hpp @@ -0,0 +1,80 @@ +#pragma once + +#include "../../device.hpp" + +#ifdef ENABLE_MUTUAL_AWARENESS +#include "../../analyzer/optimization_intent.hpp" +#endif + +#include +#include + +namespace infinicore::op::common { +template +class OpDispatcher { +public: + void registerDevice(Device::Type device_type, Fn fn, bool override_existing = true) { + if (table_[(size_t)device_type] == nullptr || override_existing) { + table_[(size_t)device_type] = fn; + } + } + + void registerDevice(std::initializer_list device_types, Fn fn, bool override_existing = true) { + for (auto device_type : device_types) { + registerDevice(device_type, fn, override_existing); + } + } + + void registerAll(Fn fn, bool override_existing = true) { + for (size_t device_type = 0; device_type < static_cast(Device::Type::COUNT); ++device_type) { + registerDevice((Device::Type)device_type, fn, override_existing); + } + } + + Fn lookup(Device::Type device_type) const { + return table_.at((size_t)device_type); + } + +#ifdef ENABLE_MUTUAL_AWARENESS + // Goal-aware kernel registration. Backward compatible: callers that don't + // know about goals keep using the device-only overloads. Only kernels that + // want to specialize per OptimizationGoal need the goal-aware form. + static constexpr std::size_t kGoalCount = 4; + + void registerDevice(Device::Type device_type, + Fn fn, + analyzer::OptimizationGoal goal, + bool override_existing = true) { + std::size_t k = goalKey(device_type, goal); + if (goal_table_[k] == nullptr || override_existing) { + goal_table_[k] = fn; + } + } + + // Look up a kernel by (device, goal). If no goal-specific kernel is + // registered, fall back to the device-default kernel registered through + // the legacy lookup(device_type) path. + Fn lookup(Device::Type device_type, analyzer::OptimizationGoal goal) const { + std::size_t k = goalKey(device_type, goal); + Fn fn = goal_table_[k]; + if (fn != nullptr) { + return fn; + } + return lookup(device_type); + } + +private: + static std::size_t goalKey(Device::Type device_type, analyzer::OptimizationGoal goal) { + return static_cast(device_type) * kGoalCount + + static_cast(goal); + } + + std::array(Device::Type::COUNT) * kGoalCount> + goal_table_{}; +#endif + +private: + std::array(Device::Type::COUNT)> table_{}; +}; +} // namespace infinicore::op::common diff --git a/python/infinicore/include/infinicore/ops/common/op.hpp b/python/infinicore/include/infinicore/ops/common/op.hpp new file mode 100644 index 000000000..b4f0434ef --- /dev/null +++ b/python/infinicore/include/infinicore/ops/common/op.hpp @@ -0,0 +1,5 @@ +#pragma once + +#include "../../context/context.hpp" +#include "../../tensor.hpp" +#include "dispatcher.hpp" diff --git a/python/infinicore/include/infinicore/ops/concat_and_cache_mla.hpp b/python/infinicore/include/infinicore/ops/concat_and_cache_mla.hpp new file mode 100644 index 000000000..6acc70b21 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/concat_and_cache_mla.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "../tensor.hpp" +#include + +namespace infinicore::op { + +void concat_and_cache_mla_(const Tensor &kv_c, + const Tensor &k_pe, + Tensor kv_cache, + const Tensor &slot_mapping, + const std::string &kv_cache_dtype, + Tensor scale); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/concat_and_cache_mla_int8.hpp b/python/infinicore/include/infinicore/ops/concat_and_cache_mla_int8.hpp new file mode 100644 index 000000000..3de0f4d13 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/concat_and_cache_mla_int8.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "../tensor.hpp" + +namespace infinicore::op { + +void concat_and_cache_mla_int8_(const Tensor &kv_c_int8, + const Tensor &kv_c_scale, + const Tensor &k_pe_int8, + const Tensor &k_pe_scale, + Tensor kv_cache, + Tensor kv_cache_scale, + const Tensor &slot_mapping); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/concat_mla_q.hpp b/python/infinicore/include/infinicore/ops/concat_mla_q.hpp new file mode 100644 index 000000000..90d84c596 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/concat_mla_q.hpp @@ -0,0 +1,11 @@ +#pragma once + +#include "../device.hpp" +#include "../tensor.hpp" + +namespace infinicore::op { + +void concat_mla_q_(const Tensor &ql_nope, const Tensor &q_pe, Tensor q_out); +Tensor concat_mla_q(const Tensor &ql_nope, const Tensor &q_pe); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/conv2d.hpp b/python/infinicore/include/infinicore/ops/conv2d.hpp new file mode 100644 index 000000000..f1dda90ac --- /dev/null +++ b/python/infinicore/include/infinicore/ops/conv2d.hpp @@ -0,0 +1,38 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +#include +#include + +namespace infinicore::op { +class Conv2d { +public: + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, + const size_t *, const size_t *, const size_t *, size_t); + static void execute(Tensor output, + Tensor input, + Tensor weight, + Tensor bias, + const size_t *pads, + const size_t *strides, + const size_t *dilations, + size_t n); + static common::OpDispatcher &dispatcher(); +}; + +Tensor conv2d(Tensor input, + Tensor weight, + Tensor bias, + const std::vector &pads, + const std::vector &strides, + const std::vector &dilations); +void conv2d_(Tensor output, + Tensor input, + Tensor weight, + Tensor bias, + const std::vector &pads, + const std::vector &strides, + const std::vector &dilations); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/cross_entropy.hpp b/python/infinicore/include/infinicore/ops/cross_entropy.hpp new file mode 100644 index 000000000..9a6d446d2 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/cross_entropy.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class CrossEntropy { +public: + // Schema 定义：函数指针类型 + // CrossEntropy 需要接收三个 Tensor: Output (Loss), Input (Logits), Target (Labels) + using schema = void (*)(Tensor, Tensor, Tensor); + + // 执行入口 + static void execute(Tensor output, Tensor input, Tensor target); + + // 分发器访问接口 + static common::OpDispatcher &dispatcher(); +}; + +// ================================================================== +// 对外 Functional API +// ================================================================== + +// 1. Out-of-place 接口： +// 输入 Logits 和 Target，内部自动创建 Output Tensor 并返回 +Tensor cross_entropy(Tensor input, Tensor target); + +// 2. Explicit Output 接口 (类似于 In-place 风格)： +// 用户显式提供 Output Tensor 用于存储结果 +// 注意：虽然命名带有下划线 _，但通常 CrossEntropy 无法真正原地修改 input， +// 所以这里只是表示“写入指定的 output 内存” +void cross_entropy_(Tensor output, Tensor input, Tensor target); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/deepseek_moe.hpp b/python/infinicore/include/infinicore/ops/deepseek_moe.hpp new file mode 100644 index 000000000..e8e7ffef9 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/deepseek_moe.hpp @@ -0,0 +1,41 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS( + DeepseekMoe, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const std::vector &, + const std::vector &, + const std::vector &, + size_t, + size_t); + +Tensor deepseek_moe(const Tensor &hidden, + const Tensor &topk_indices, + const Tensor &topk_weights, + const std::vector &gate_weights, + const std::vector &up_weights, + const std::vector &down_weights, + size_t intermediate_size, + size_t num_experts); + +void deepseek_moe_(Tensor out, + const Tensor &hidden, + const Tensor &topk_indices, + const Tensor &topk_weights, + const std::vector &gate_weights, + const std::vector &up_weights, + const std::vector &down_weights, + size_t intermediate_size, + size_t num_experts); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/dequantize_awq.hpp b/python/infinicore/include/infinicore/ops/dequantize_awq.hpp new file mode 100644 index 000000000..50e4328f3 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/dequantize_awq.hpp @@ -0,0 +1,10 @@ +#pragma once +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { +INFINICORE_GRAPH_OP_CLASS(DequantizeAWQ, Tensor, const Tensor &, const Tensor &, const Tensor &); + +void dequantize_awq_(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zeros); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/diff.hpp b/python/infinicore/include/infinicore/ops/diff.hpp new file mode 100644 index 000000000..1def7a61a --- /dev/null +++ b/python/infinicore/include/infinicore/ops/diff.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Diff, Tensor, const Tensor &, int, int); + +Tensor diff(const Tensor &x, int n = 1, int dim = -1); +void diff_(Tensor y, const Tensor &x, int n = 1, int dim = -1); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/digamma.hpp b/python/infinicore/include/infinicore/ops/digamma.hpp new file mode 100644 index 000000000..b61fc40f6 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/digamma.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Digamma, Tensor, const Tensor &); + +Tensor digamma(const Tensor &x); +void digamma_(Tensor y, const Tensor &x); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/dist.hpp b/python/infinicore/include/infinicore/ops/dist.hpp new file mode 100644 index 000000000..459a52202 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/dist.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Dist, Tensor, const Tensor &, const Tensor &, double); + +Tensor dist(const Tensor &x1, const Tensor &x2, double p = 2.0); +void dist_(Tensor y, const Tensor &x1, const Tensor &x2, double p = 2.0); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/distributed/allgather.hpp b/python/infinicore/include/infinicore/ops/distributed/allgather.hpp new file mode 100644 index 000000000..f206cab32 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/distributed/allgather.hpp @@ -0,0 +1,52 @@ +#pragma once + +#include "../../device.hpp" +#include "../../graph/graph.hpp" +#include "../common/op.hpp" + +#include +#include + +namespace infinicore::op::distributed { + +class AllGather : public graph::GraphOperator { +public: + AllGather(Tensor output, const Tensor &input, infinicclComm_t communicator); + ~AllGather(); + void run() const override; + static void execute(Tensor output, const Tensor &input, infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +class AllGatherV : public graph::GraphOperator { +public: + AllGatherV(Tensor output, + const Tensor &input, + std::vector split_counts, + infinicclComm_t communicator); + ~AllGatherV(); + void run() const override; + static void execute(Tensor output, + const Tensor &input, + std::vector split_counts, + infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +Tensor allgather(const Tensor &input, size_t world_size, infinicclComm_t communicator); +void allgather_(Tensor output, const Tensor &input, infinicclComm_t communicator); +Tensor allgatherv(const Tensor &input, const std::vector &split_sizes, infinicclComm_t communicator); +void allgatherv_(Tensor output, const Tensor &input, const std::vector &split_sizes, infinicclComm_t communicator); +std::vector allgatherv_many(const std::vector &inputs, + const std::vector &split_sizes, + infinicclComm_t communicator); +void allgatherv_many_(const std::vector &outputs, + const std::vector &inputs, + const std::vector &split_sizes, + infinicclComm_t communicator); + +} // namespace infinicore::op::distributed diff --git a/python/infinicore/include/infinicore/ops/distributed/allreduce.hpp b/python/infinicore/include/infinicore/ops/distributed/allreduce.hpp new file mode 100644 index 000000000..39f74243a --- /dev/null +++ b/python/infinicore/include/infinicore/ops/distributed/allreduce.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "../../device.hpp" +#include "../../graph/graph.hpp" +#include "../common/op.hpp" + +#include + +namespace infinicore::op::distributed { +class AllReduce : public graph::GraphOperator { +public: + AllReduce(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); + ~AllReduce(); + void run() const override; + static void execute(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +Tensor allreduce(const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); +void allreduce_(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); + +} // namespace infinicore::op::distributed diff --git a/python/infinicore/include/infinicore/ops/distributed/reduce_scatter.hpp b/python/infinicore/include/infinicore/ops/distributed/reduce_scatter.hpp new file mode 100644 index 000000000..ee5622872 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/distributed/reduce_scatter.hpp @@ -0,0 +1,65 @@ +#pragma once + +#include "../../device.hpp" +#include "../../graph/graph.hpp" +#include "../common/op.hpp" + +#include +#include + +namespace infinicore::op::distributed { + +class ReduceScatter : public graph::GraphOperator { +public: + ReduceScatter(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); + ~ReduceScatter(); + void run() const override; + static void execute(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +class ReduceScatterV : public graph::GraphOperator { +public: + ReduceScatterV(Tensor output, + const Tensor &input, + std::vector split_counts, + infinicclReduceOp_t op, + infinicclComm_t communicator); + ~ReduceScatterV(); + void run() const override; + static void execute(Tensor output, + const Tensor &input, + std::vector split_counts, + infinicclReduceOp_t op, + infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +Tensor reduce_scatter(const Tensor &input, size_t world_size, infinicclReduceOp_t op, infinicclComm_t communicator); +void reduce_scatter_(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); +Tensor reduce_scatterv(const Tensor &input, + const std::vector &split_sizes, + size_t rank, + infinicclReduceOp_t op, + infinicclComm_t communicator); +void reduce_scatterv_(Tensor output, + const Tensor &input, + const std::vector &split_sizes, + infinicclReduceOp_t op, + infinicclComm_t communicator); +std::vector reduce_scatterv_many(const std::vector &inputs, + const std::vector &split_sizes, + size_t rank, + infinicclReduceOp_t op, + infinicclComm_t communicator); +void reduce_scatterv_many_(const std::vector &outputs, + const std::vector &inputs, + const std::vector &split_sizes, + infinicclReduceOp_t op, + infinicclComm_t communicator); + +} // namespace infinicore::op::distributed diff --git a/python/infinicore/include/infinicore/ops/dsa.hpp b/python/infinicore/include/infinicore/ops/dsa.hpp new file mode 100644 index 000000000..ab40c9f17 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/dsa.hpp @@ -0,0 +1,87 @@ +#pragma once + +#include "../tensor.hpp" + +#include +#include + +namespace infinicore::op { + +void fused_deepseek_v2_indexer_postprocess_( + Tensor q_out, + Tensor k_out, + Tensor weights_out, + Tensor kv_cache, + const Tensor &slot_mapping, + const Tensor &q, + const Tensor &kw, + const Tensor &norm_weight, + const Tensor &norm_bias, + const Tensor &positions, + const Tensor &cos_sin_cache, + int64_t num_cache_tokens, + bool is_neox, + double eps, + double weights_scale); + +void indexer_k_cache_(const Tensor &k, Tensor kv_cache, const Tensor &slot_mapping); + +void indexer_k_quant_and_cache_( + const Tensor &k, + Tensor kv_cache, + const Tensor &slot_mapping, + int64_t quant_block_size = 128, + const std::string &scale_fmt = "ue8m0"); + +void compute_block_sparse_mqa_logits_( + Tensor logits, + const Tensor &q, + const Tensor &kv_cache, + const Tensor &cu_seqlens_q, + const Tensor &cu_seqlens_kv, + const Tensor &block_table, + const Tensor &weights, + int64_t max_q_len, + int64_t max_kv_len, + int64_t max_context_len); + +void select_prefill_topk_block_indices_( + Tensor topk_indices, + const Tensor &logits, + const Tensor &cu_seqlen_ks, + const Tensor &cu_seqlen_ke); + +void select_decode_topk_block_indices_( + Tensor topk_indices, + const Tensor &logits, + const Tensor &seq_lens); + +void map_prefill_request_block_indices_( + Tensor output, + const Tensor &req_id, + const Tensor &block_table, + const Tensor &token_indices, + int64_t block_size, + bool has_prefill_workspace = false, + std::optional prefill_workspace_request_ids = std::nullopt, + std::optional prefill_workspace_starts = std::nullopt); + +void map_decode_request_block_indices_( + Tensor output, + const Tensor &req_id, + const Tensor &block_table, + const Tensor &token_indices, + int64_t block_size); + +void topk_indices_context_lens_(Tensor topk_lens, const Tensor &indices); + +void sparse_flash_mla_( + Tensor output, + const Tensor &query, + const Tensor &kv_cache, + const Tensor &indices, + const Tensor &topk_lens, + float scale, + std::optional attn_sink = std::nullopt); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/dynamic_scaled_int8_quant.hpp b/python/infinicore/include/infinicore/ops/dynamic_scaled_int8_quant.hpp new file mode 100644 index 000000000..f0a114821 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/dynamic_scaled_int8_quant.hpp @@ -0,0 +1,11 @@ +#pragma once + +#include "../device.hpp" +#include "../tensor.hpp" + +namespace infinicore::op { + +void dynamic_scaled_int8_quant_(Tensor output, const Tensor &input, Tensor input_scales); +Tensor dynamic_scaled_int8_quant(const Tensor &input, Tensor input_scales); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/embedding.hpp b/python/infinicore/include/infinicore/ops/embedding.hpp new file mode 100644 index 000000000..43f18d090 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/embedding.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(Embedding, Tensor, const Tensor &, const Tensor &); + +Tensor embedding(const Tensor &input, const Tensor &weight); +void embedding_(Tensor out, const Tensor &input, const Tensor &weight); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/equal.hpp b/python/infinicore/include/infinicore/ops/equal.hpp new file mode 100644 index 000000000..1a158bf1e --- /dev/null +++ b/python/infinicore/include/infinicore/ops/equal.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Equal { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + + static void execute(Tensor out, Tensor a, Tensor b); + static common::OpDispatcher &dispatcher(); +}; + +Tensor equal(Tensor a, Tensor b); +void equal_(Tensor out, Tensor a, Tensor b); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/flash_attention.hpp b/python/infinicore/include/infinicore/ops/flash_attention.hpp new file mode 100644 index 000000000..24e33cfb6 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/flash_attention.hpp @@ -0,0 +1,12 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(FlashAttention, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, float, bool); + +Tensor flash_attention(const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal); +void flash_attention_(Tensor out, const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/flipud.hpp b/python/infinicore/include/infinicore/ops/flipud.hpp new file mode 100644 index 000000000..7e449f9f5 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/flipud.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Flipud { +public: + // Schema signature: (Output, Input) + using schema = void (*)(Tensor, Tensor); + + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; +Tensor flipud(Tensor input); +void flipud_(Tensor output, Tensor input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/float_power.hpp b/python/infinicore/include/infinicore/ops/float_power.hpp new file mode 100644 index 000000000..c461db2eb --- /dev/null +++ b/python/infinicore/include/infinicore/ops/float_power.hpp @@ -0,0 +1,68 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class FloatPower { +public: + // ========================================================== + // Dispatcher Schemas + // ========================================================== + + // Output = Input ^ Scalar (scalar must be double!) + using schema_scalar = void (*)(Tensor output, + Tensor input, + double exponent); + + // Output = Input ^ Tensor + using schema_tensor = void (*)(Tensor output, + Tensor input, + Tensor exponent); + + // ========================================================== + // Execute Entry Points (called by functional interface) + // ========================================================== + + static void execute(Tensor output, + Tensor input, + double exponent); + + static void execute(Tensor output, + Tensor input, + Tensor exponent); + + // ========================================================== + // Dispatchers + // ========================================================== + + static common::OpDispatcher &dispatcher_scalar(); + static common::OpDispatcher &dispatcher_tensor(); +}; + +// ======================================================================= +// Functional Interface (Python-visible semantics) +// ======================================================================= + +// ------------------------------- +// 1. Scalar Exponent +// ------------------------------- + +// out-of-place: ALWAYS float64 +Tensor float_power(Tensor input, double exponent); + +// in-place +void float_power_(Tensor output, Tensor input, double exponent); + +// ------------------------------- +// 2. Tensor Exponent +// ------------------------------- + +// out-of-place: ALWAYS float64 +Tensor float_power(Tensor input, Tensor exponent); + +// in-place +void float_power_(Tensor output, Tensor input, Tensor exponent); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/floor.hpp b/python/infinicore/include/infinicore/ops/floor.hpp new file mode 100644 index 000000000..11b52571b --- /dev/null +++ b/python/infinicore/include/infinicore/ops/floor.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Floor { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor floor(Tensor input); +void floor_(Tensor output, Tensor input); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/floor_divide.hpp b/python/infinicore/include/infinicore/ops/floor_divide.hpp new file mode 100644 index 000000000..43267dce6 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/floor_divide.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class FloorDivide { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + static void execute(Tensor c, Tensor a, Tensor b); + static common::OpDispatcher &dispatcher(); +}; + +Tensor floor_divide(Tensor a, Tensor b); +void floor_divide_(Tensor c, Tensor a, Tensor b); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/fmin.hpp b/python/infinicore/include/infinicore/ops/fmin.hpp new file mode 100644 index 000000000..4ea00787f --- /dev/null +++ b/python/infinicore/include/infinicore/ops/fmin.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Fmin { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + static void execute(Tensor c, Tensor a, Tensor b); + static common::OpDispatcher &dispatcher(); +}; + +Tensor fmin(Tensor a, Tensor b); +void fmin_(Tensor c, Tensor a, Tensor b); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/fmod.hpp b/python/infinicore/include/infinicore/ops/fmod.hpp new file mode 100644 index 000000000..87b90d515 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/fmod.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Fmod { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + static void execute(Tensor c, Tensor a, Tensor b); + static common::OpDispatcher &dispatcher(); +}; + +Tensor fmod(Tensor a, Tensor b); +void fmod_(Tensor c, Tensor a, Tensor b); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/fp8_indexer_logits.hpp b/python/infinicore/include/infinicore/ops/fp8_indexer_logits.hpp new file mode 100644 index 000000000..523eb00c1 --- /dev/null +++ b/python/infinicore/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/python/infinicore/include/infinicore/ops/fp8_indexer_quant.hpp b/python/infinicore/include/infinicore/ops/fp8_indexer_quant.hpp new file mode 100644 index 000000000..8af946005 --- /dev/null +++ b/python/infinicore/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/python/infinicore/include/infinicore/ops/fp8_mla_rmsnorm_cache.hpp b/python/infinicore/include/infinicore/ops/fp8_mla_rmsnorm_cache.hpp new file mode 100644 index 000000000..e90349125 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/fp8_mla_rmsnorm_cache.hpp @@ -0,0 +1,22 @@ +#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); + +void fp8_mla_rmsnorm_cache_( + Tensor cache, + const Tensor &compressed_kv, + const Tensor &norm_weight, + const Tensor &rope, + const Tensor &slot_mapping, + double eps); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/fp8_sparse_mla.hpp b/python/infinicore/include/infinicore/ops/fp8_sparse_mla.hpp new file mode 100644 index 000000000..6ad6fa79d --- /dev/null +++ b/python/infinicore/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/python/infinicore/include/infinicore/ops/fused_gated_delta_net_gating.hpp b/python/infinicore/include/infinicore/ops/fused_gated_delta_net_gating.hpp new file mode 100644 index 000000000..a61fd9937 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/fused_gated_delta_net_gating.hpp @@ -0,0 +1,37 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(FusedGatedDeltaNetGating, + Tensor, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + float, + float); + +std::pair fused_gated_delta_net_gating(const Tensor &A_log, + const Tensor &a, + const Tensor &b, + const Tensor &dt_bias, + float beta = 1.0f, + float threshold = 20.0f); + +void fused_gated_delta_net_gating_(Tensor g, + Tensor beta_output, + const Tensor &A_log, + const Tensor &a, + const Tensor &b, + const Tensor &dt_bias, + float beta = 1.0f, + float threshold = 20.0f); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/fused_rotary_embedding.hpp b/python/infinicore/include/infinicore/ops/fused_rotary_embedding.hpp new file mode 100644 index 000000000..5a5a4afba --- /dev/null +++ b/python/infinicore/include/infinicore/ops/fused_rotary_embedding.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../tensor.hpp" + +namespace infinicore::op { + +// Apply rotary embedding to query and key in one vendor-kernel launch. This +// mirrors vLLM's rotary_embedding(positions, query, key, ...) execution path. +void fused_rotary_embedding_(Tensor query, + Tensor key, + const Tensor &positions, + int64_t head_size, + const Tensor &cos_sin_cache, + bool is_neox); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/gaussian_nll_loss.hpp b/python/infinicore/include/infinicore/ops/gaussian_nll_loss.hpp new file mode 100644 index 000000000..c40eb1e86 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/gaussian_nll_loss.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include "infinicore.h" + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(GaussianNllLoss, Tensor, const Tensor &, const Tensor &, const Tensor &, bool, double, int); + +__export Tensor gaussian_nll_loss(const Tensor &input, + const Tensor &target, + const Tensor &var, + bool full = false, + double eps = 1e-6, + int reduction = 1); + +__export void gaussian_nll_loss_(Tensor out, + const Tensor &input, + const Tensor &target, + const Tensor &var, + bool full = false, + double eps = 1e-6, + int reduction = 1); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/gelu.hpp b/python/infinicore/include/infinicore/ops/gelu.hpp new file mode 100644 index 000000000..5e8c61347 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/gelu.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Gelu { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor gelu(Tensor input); +void gelu_(Tensor output, Tensor input); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/gelutanh.hpp b/python/infinicore/include/infinicore/ops/gelutanh.hpp new file mode 100644 index 000000000..c968c81fd --- /dev/null +++ b/python/infinicore/include/infinicore/ops/gelutanh.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class GeluTanh { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor gelu_tanh(Tensor input); +void gelu_tanh_(Tensor output, Tensor input); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/gemm.hpp b/python/infinicore/include/infinicore/ops/gemm.hpp new file mode 100644 index 000000000..4f76cee26 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/gemm.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Gemm, Tensor, const Tensor &, const Tensor &, float, float); + +Tensor gemm(const Tensor &a, const Tensor &b, float alpha = 1.0f, float beta = 0.0f); +void gemm_(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/gptq_marlin_gemm.hpp b/python/infinicore/include/infinicore/ops/gptq_marlin_gemm.hpp new file mode 100644 index 000000000..b16df0ad6 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/gptq_marlin_gemm.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(GptqMarlinGemm, Tensor, const Tensor &, const Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, int64_t, bool, bool, bool, bool); +INFINICORE_GRAPH_OP_CLASS(GptqMarlinGemmWithWorkspace, Tensor, Tensor, const Tensor &, const Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, int64_t, bool, bool, bool, bool); + +void gptq_marlin_gemm_(Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float); + +size_t gptq_marlin_gemm_workspace_size(Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm); +void gptq_marlin_gemm_with_workspace_(Tensor workspace, Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float); +void gptq_marlin_gemm_with_workspace_direct_(Tensor workspace, Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/grouped_topk_vllm.hpp b/python/infinicore/include/infinicore/ops/grouped_topk_vllm.hpp new file mode 100644 index 000000000..11d30cdf4 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/grouped_topk_vllm.hpp @@ -0,0 +1,8 @@ +#pragma once +#include "../device.hpp" +#include "../tensor.hpp" +#include +#include +namespace infinicore::op { +void grouped_topk_vllm_(Tensor topk_weights, Tensor topk_ids, const Tensor &scores, int64_t num_expert_group, int64_t topk_group, bool renormalize, float routed_scaling_factor, const Tensor &bias = Tensor(), const std::string &scoring_func = "softmax"); +} diff --git a/python/infinicore/include/infinicore/ops/hardswish.hpp b/python/infinicore/include/infinicore/ops/hardswish.hpp new file mode 100644 index 000000000..15313f461 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/hardswish.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Hardswish { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor hardswish(Tensor input); +void hardswish_(Tensor output, Tensor input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/hardtanh.hpp b/python/infinicore/include/infinicore/ops/hardtanh.hpp new file mode 100644 index 000000000..511408fee --- /dev/null +++ b/python/infinicore/include/infinicore/ops/hardtanh.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class HardTanh { +public: + using schema = void (*)(Tensor, Tensor, float, float); + static void execute(Tensor output, Tensor input, float min_val, float max_val); + static common::OpDispatcher &dispatcher(); +}; + +Tensor hardtanh(Tensor input, float min_val = -1.0f, float max_val = 1.0f); +void hardtanh_(Tensor output, Tensor input, float min_val = -1.0f, float max_val = 1.0f); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/hinge_embedding_loss.hpp b/python/infinicore/include/infinicore/ops/hinge_embedding_loss.hpp new file mode 100644 index 000000000..9d0fcfbab --- /dev/null +++ b/python/infinicore/include/infinicore/ops/hinge_embedding_loss.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class HingeEmbeddingLoss { +public: + using schema = void (*)(Tensor, Tensor, Tensor, double, int); + static void execute(Tensor output, Tensor input, Tensor target, double margin, int reduction); + static common::OpDispatcher &dispatcher(); +}; + +Tensor hinge_embedding_loss(Tensor input, Tensor target, double margin = 1.0, int reduction = 1); +void hinge_embedding_loss_(Tensor output, Tensor input, Tensor target, double margin, int reduction); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/huber_loss.hpp b/python/infinicore/include/infinicore/ops/huber_loss.hpp new file mode 100644 index 000000000..c6212f0c1 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/huber_loss.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class HuberLoss { +public: + // Schema: output, input, target, delta, reduction + using schema = void (*)(Tensor, Tensor, Tensor, float, int64_t); + + static void execute(Tensor output, Tensor input, Tensor target, float delta, int64_t reduction); + static common::OpDispatcher &dispatcher(); +}; + +// delta 默认为 1.0f，reduction 默认为 1 (MEAN) +Tensor huber_loss(Tensor input, Tensor target, float delta = 1.0f, int64_t reduction = 1); +void huber_loss_(Tensor output, Tensor input, Tensor target, float delta, int64_t reduction); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/hypot.hpp b/python/infinicore/include/infinicore/ops/hypot.hpp new file mode 100644 index 000000000..24eebd44f --- /dev/null +++ b/python/infinicore/include/infinicore/ops/hypot.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Hypot { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + + static void execute(Tensor output, Tensor input_a, Tensor input_b); + static common::OpDispatcher &dispatcher(); +}; + +Tensor hypot(Tensor input_a, Tensor input_b); + +void hypot_(Tensor output, Tensor input_a, Tensor input_b); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/index_add.hpp b/python/infinicore/include/infinicore/ops/index_add.hpp new file mode 100644 index 000000000..e880114e9 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/index_add.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class IndexAdd { +public: + using schema = void (*)(Tensor, Tensor, int64_t, Tensor, Tensor, float); + static void execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha); + + static common::OpDispatcher &dispatcher(); +}; + +Tensor index_add(Tensor input, int64_t dim, Tensor index, Tensor source, float alpha = 1.0f); +void index_add_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/index_copy.hpp b/python/infinicore/include/infinicore/ops/index_copy.hpp new file mode 100644 index 000000000..b2437dead --- /dev/null +++ b/python/infinicore/include/infinicore/ops/index_copy.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class IndexCopy { +public: + using schema = void (*)(Tensor, Tensor, int64_t, Tensor, Tensor); + static void execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source); + + static common::OpDispatcher &dispatcher(); +}; +Tensor index_copy(Tensor input, int64_t dim, Tensor index, Tensor source); +void index_copy_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/inner.hpp b/python/infinicore/include/infinicore/ops/inner.hpp new file mode 100644 index 000000000..ae372c415 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/inner.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Inner { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + static void execute(Tensor out, Tensor input, Tensor other); + static common::OpDispatcher &dispatcher(); +}; + +Tensor inner(Tensor input, Tensor other); +void inner_(Tensor out, Tensor input, Tensor other); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/interpolate.hpp b/python/infinicore/include/infinicore/ops/interpolate.hpp new file mode 100644 index 000000000..2a235ef3d --- /dev/null +++ b/python/infinicore/include/infinicore/ops/interpolate.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include "infinicore.h" + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +#include +#include +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Interpolate, Tensor, const Tensor &, std::string, std::vector, std::vector, int); + +__export Tensor interpolate(const Tensor &input, + std::string mode, + std::vector size, + std::vector scale_factor, + int align_corners); + +__export void interpolate_(Tensor out, + const Tensor &input, + std::string mode, + std::vector size, + std::vector scale_factor, + int align_corners); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/kron.hpp b/python/infinicore/include/infinicore/ops/kron.hpp new file mode 100644 index 000000000..3a6834424 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/kron.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Kron { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + static void execute(Tensor output, Tensor a, Tensor b); + static common::OpDispatcher &dispatcher(); +}; + +Tensor kron(Tensor a, Tensor b); +void kron_(Tensor output, Tensor a, Tensor b); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/kthvalue.hpp b/python/infinicore/include/infinicore/ops/kthvalue.hpp new file mode 100644 index 000000000..1e9721d28 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/kthvalue.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +class Kthvalue { +public: + // Schema signature: values(out), indices(out), input, k, dim, keepdim + using schema = void (*)(Tensor, Tensor, Tensor, int64_t, int64_t, bool); + + static void execute(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim); + static common::OpDispatcher &dispatcher(); +}; + +// Functional API: Returns a tuple containing (values, indices) +std::tuple kthvalue(Tensor input, int64_t k, int64_t dim = -1, bool keepdim = false); + +// In-place/Output-provided API +void kthvalue_(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/kv_caching.hpp b/python/infinicore/include/infinicore/ops/kv_caching.hpp new file mode 100644 index 000000000..3a70c2824 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/kv_caching.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(KVCaching, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &); + +void kv_caching_(Tensor k_cache, + Tensor v_cache, + const Tensor &k, + const Tensor &v, + const Tensor &past_kv_lengths); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/layer_norm.hpp b/python/infinicore/include/infinicore/ops/layer_norm.hpp new file mode 100644 index 000000000..da6256b51 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/layer_norm.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(LayerNorm, Tensor, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &, float); + +Tensor layer_norm(const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon = 1e-5f); +void layer_norm_(Tensor y, Tensor standardization, Tensor std_deviation, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon = 1e-5f); +void layer_norm_(Tensor y, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon = 1e-5f); +void layer_norm_for_pybind(Tensor y, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon = 1e-5f); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/ldexp.hpp b/python/infinicore/include/infinicore/ops/ldexp.hpp new file mode 100644 index 000000000..bac7964c2 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/ldexp.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Ldexp { +public: + // Schema signature: output(out), input(x), other(exp) + using schema = void (*)(Tensor, Tensor, Tensor); + + static void execute(Tensor output, Tensor input, Tensor other); + static common::OpDispatcher &dispatcher(); +}; + +// Functional API: Returns a new Tensor containing input * (2^other) +Tensor ldexp(Tensor input, Tensor other); + +// In-place/Output-provided API +// Writes the result into 'output' +void ldexp_(Tensor output, Tensor input, Tensor other); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/lerp.hpp b/python/infinicore/include/infinicore/ops/lerp.hpp new file mode 100644 index 000000000..bf3e5cbec --- /dev/null +++ b/python/infinicore/include/infinicore/ops/lerp.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Lerp { +public: + using schema_t = void (*)(Tensor, Tensor, Tensor, Tensor); + using schema_s = void (*)(Tensor, Tensor, Tensor, float); + + static void execute(Tensor output, Tensor start, Tensor end, Tensor weight); + static void execute(Tensor output, Tensor start, Tensor end, float weight); + + // 【核心修改】必须声明为模板函数，才能支持 dispatcher() 和 dispatcher() + template + static common::OpDispatcher &dispatcher(); +}; + +Tensor lerp(Tensor start, Tensor end, Tensor weight); +Tensor lerp(Tensor start, Tensor end, float weight); + +void lerp_(Tensor output, Tensor start, Tensor end, Tensor weight); +void lerp_(Tensor output, Tensor start, Tensor end, float weight); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/linear.hpp b/python/infinicore/include/infinicore/ops/linear.hpp new file mode 100644 index 000000000..96ca733e4 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/linear.hpp @@ -0,0 +1,12 @@ +#pragma once + +#include "common/op.hpp" +#include + +namespace infinicore::op { + +Tensor linear(Tensor input, Tensor weight, std::optional bias, float alpha = 1.0f); + +void linear_(Tensor out, Tensor input, Tensor weight, std::optional bias, float alpha = 1.0f); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/linear_w4a16_awq.hpp b/python/infinicore/include/infinicore/ops/linear_w4a16_awq.hpp new file mode 100644 index 000000000..ebae7a685 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/linear_w4a16_awq.hpp @@ -0,0 +1,12 @@ +#pragma once + +#include "common/op.hpp" +#include + +namespace infinicore::op { + +Tensor linear_w4a16_awq(Tensor input, Tensor weight_packed, Tensor weight_scale, Tensor weight_zeros, std::optional bias); + +void linear_w4a16_awq_(Tensor out, Tensor input, Tensor weight_packed, Tensor weight_scale, Tensor weight_zeros, std::optional bias); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/linear_w4a16_gptq_qy.hpp b/python/infinicore/include/infinicore/ops/linear_w4a16_gptq_qy.hpp new file mode 100644 index 000000000..72cc9eef4 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/linear_w4a16_gptq_qy.hpp @@ -0,0 +1,12 @@ +#pragma once + +#include "common/op.hpp" +#include + +namespace infinicore::op { + +Tensor linear_w4a16_gptq_qy(Tensor in, Tensor qweight, Tensor qzeros, Tensor scales, int64_t quant_type, int64_t bit); + +void linear_w4a16_gptq_qy_(Tensor out, Tensor in, Tensor qweights, Tensor scales, Tensor qzeros, int64_t quant_type, int64_t bit); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/linear_w8a8i8.hpp b/python/infinicore/include/infinicore/ops/linear_w8a8i8.hpp new file mode 100644 index 000000000..08cadc111 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/linear_w8a8i8.hpp @@ -0,0 +1,13 @@ +#pragma once + +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +Tensor linear_w8a8i8(Tensor input, Tensor weight_packed, Tensor weight_scale, std::optional bias); + +void linear_w8a8i8_(Tensor out, Tensor input, Tensor weight_packed, Tensor weight_scale, std::optional bias); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/log_softmax.hpp b/python/infinicore/include/infinicore/ops/log_softmax.hpp new file mode 100644 index 000000000..00e1c90c1 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/log_softmax.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class LogSoftmax { +public: + // Schema signature: output(out), input, dim + using schema = void (*)(Tensor, Tensor, int64_t); + + static void execute(Tensor output, Tensor input, int64_t dim); + static common::OpDispatcher &dispatcher(); +}; + +// Functional API: Returns the result tensor +Tensor log_softmax(Tensor input, int64_t dim); + +// In-place/Output-provided API +void log_softmax_(Tensor output, Tensor input, int64_t dim); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/logaddexp.hpp b/python/infinicore/include/infinicore/ops/logaddexp.hpp new file mode 100644 index 000000000..fc84c10c2 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/logaddexp.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class LogAddExp { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + static void execute(Tensor c, Tensor a, Tensor b); + static common::OpDispatcher &dispatcher(); +}; + +Tensor logaddexp(Tensor a, Tensor b); +void logaddexp_(Tensor c, Tensor a, Tensor b); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/logaddexp2.hpp b/python/infinicore/include/infinicore/ops/logaddexp2.hpp new file mode 100644 index 000000000..7c7dbe96e --- /dev/null +++ b/python/infinicore/include/infinicore/ops/logaddexp2.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class LogAddExp2 { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + static void execute(Tensor c, Tensor a, Tensor b); + static common::OpDispatcher &dispatcher(); +}; + +Tensor logaddexp2(Tensor a, Tensor b); +void logaddexp2_(Tensor c, Tensor a, Tensor b); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/logcumsumexp.hpp b/python/infinicore/include/infinicore/ops/logcumsumexp.hpp new file mode 100644 index 000000000..212ff625d --- /dev/null +++ b/python/infinicore/include/infinicore/ops/logcumsumexp.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class LogCumSumExp { +public: + using schema = void (*)(Tensor, Tensor, int, bool, bool); + + static void execute(Tensor y, Tensor x, int axis, bool exclusive, bool reverse); + + static common::OpDispatcher &dispatcher(); +}; + +Tensor logcumsumexp(Tensor x, int axis, bool exclusive = false, bool reverse = false); + +void logcumsumexp_(Tensor y, Tensor x, int axis, bool exclusive = false, bool reverse = false); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/logdet.hpp b/python/infinicore/include/infinicore/ops/logdet.hpp new file mode 100644 index 000000000..69c1b6d0d --- /dev/null +++ b/python/infinicore/include/infinicore/ops/logdet.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Logdet, Tensor, const Tensor &); + +Tensor logdet(const Tensor &x); +void logdet_(Tensor y, const Tensor &x); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/logical_and.hpp b/python/infinicore/include/infinicore/ops/logical_and.hpp new file mode 100644 index 000000000..41f45f1c6 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/logical_and.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class LogicalAnd { +public: + // LogicalAnd 是二元操作，schema 通常定义为 (Output, Input1, Input2) + using schema = void (*)(Tensor, Tensor, Tensor); + + static void execute(Tensor output, Tensor input1, Tensor input2); + static common::OpDispatcher &dispatcher(); +}; + +Tensor logical_and(Tensor input1, Tensor input2); +void logical_and_(Tensor output, Tensor input1, Tensor input2); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/logical_not.hpp b/python/infinicore/include/infinicore/ops/logical_not.hpp new file mode 100644 index 000000000..caf6bbb76 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/logical_not.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class LogicalNot { +public: + // LogicalNot 是一元操作，schema 定义为 (Output, Input) + using schema = void (*)(Tensor, Tensor); + + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +// 构造新 Tensor 返回结果 +Tensor logical_not(Tensor input); + +// 将结果写入指定的 output Tensor +void logical_not_(Tensor output, Tensor input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/mamba_selective_scan.hpp b/python/infinicore/include/infinicore/ops/mamba_selective_scan.hpp new file mode 100644 index 000000000..3acdb1c4d --- /dev/null +++ b/python/infinicore/include/infinicore/ops/mamba_selective_scan.hpp @@ -0,0 +1,43 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS( + MambaSelectiveScan, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + Tensor); + +Tensor mamba_selective_scan(const Tensor &x, + const Tensor &dt, + const Tensor &b, + const Tensor &c, + const Tensor &a_log, + const Tensor &d, + const Tensor &gate, + const Tensor &dt_bias, + Tensor state); + +void mamba_selective_scan_(Tensor out, + const Tensor &x, + const Tensor &dt, + const Tensor &b, + const Tensor &c, + const Tensor &a_log, + const Tensor &d, + const Tensor &gate, + const Tensor &dt_bias, + Tensor state); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/masked_select.hpp b/python/infinicore/include/infinicore/ops/masked_select.hpp new file mode 100644 index 000000000..d003e3bb8 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/masked_select.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class MaskedSelect { +public: + using schema = void (*)(Tensor, Tensor, void **, size_t *); + static void execute(Tensor input, Tensor mask, void **data_ptr, size_t *dlen_ptr); + static common::OpDispatcher &dispatcher(); +}; + +Tensor masked_select(Tensor input, Tensor mask); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/matmul.hpp b/python/infinicore/include/infinicore/ops/matmul.hpp new file mode 100644 index 000000000..1fb41774e --- /dev/null +++ b/python/infinicore/include/infinicore/ops/matmul.hpp @@ -0,0 +1,11 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +Tensor matmul(Tensor a, Tensor b, float alpha = 1.0f); +void matmul_(Tensor c, Tensor a, Tensor b, float alpha = 1.0f); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/mha.hpp b/python/infinicore/include/infinicore/ops/mha.hpp new file mode 100644 index 000000000..f62598c78 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/mha.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS( + MultiheadAttention, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + std::optional, + float, + bool); + +Tensor mha(const Tensor &q, + const Tensor &k, + const Tensor &v, + std::optional alibi_slopes, + float scale, + bool is_causal); + +void mha_(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + std::optional alibi_slopes, + float scale, + bool is_causal); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/mha_kvcache.hpp b/python/infinicore/include/infinicore/ops/mha_kvcache.hpp new file mode 100644 index 000000000..69a11f61c --- /dev/null +++ b/python/infinicore/include/infinicore/ops/mha_kvcache.hpp @@ -0,0 +1,51 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +// Flash Attention KV-cache decode op. +// +// Wraps flash::mha_fwd_kvcache for single-step (decode) attention over a +// paged KV cache. +// +// Tensor shapes: +// out : [batch_size, seqlen_q, num_heads, head_size] +// q : [batch_size, seqlen_q, num_heads, head_size] +// k_cache : [num_blocks, block_size, num_heads_k, head_size] (paged layout) +// v_cache : [num_blocks, block_size, num_heads_k, head_size] (paged layout) +// seqlens_k : [batch_size] int32 — total KV length per request +// block_table : [batch_size, max_num_blocks_per_seq] int32 + +INFINICORE_GRAPH_OP_CLASS( + MhaKVCache, + Tensor, // out + const Tensor &, // q + const Tensor &, // k_cache + const Tensor &, // v_cache + const Tensor &, // seqlens_k + const Tensor &, // block_table + std::optional, // alibi_slopes + float); // scale + +Tensor mha_kvcache(const Tensor &q, + const Tensor &k_cache, + const Tensor &v_cache, + const Tensor &seqlens_k, + const Tensor &block_table, + std::optional alibi_slopes, + float scale); + +void mha_kvcache_(Tensor out, + const Tensor &q, + const Tensor &k_cache, + const Tensor &v_cache, + const Tensor &seqlens_k, + const Tensor &block_table, + std::optional alibi_slopes, + float scale); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/mha_varlen.hpp b/python/infinicore/include/infinicore/ops/mha_varlen.hpp new file mode 100644 index 000000000..4d226f383 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/mha_varlen.hpp @@ -0,0 +1,46 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS( + MultiheadAttentionVarlen, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + std::optional, + int, + int, + std::optional, + float); + +Tensor mha_varlen(const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &cum_seqlens_q, + const Tensor &cum_seqlens_k, + std::optional block_table, + int max_seqlen_q, + int max_seqlen_k, + std::optional alibi_slopes, + float scale); + +void mha_varlen_(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &cum_seqlens_q, + const Tensor &cum_seqlens_k, + std::optional block_table, + int max_seqlen_q, + int max_seqlen_k, + std::optional alibi_slopes, + float scale); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/moe_align.hpp b/python/infinicore/include/infinicore/ops/moe_align.hpp new file mode 100644 index 000000000..2ce7ca8f4 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_align.hpp @@ -0,0 +1,40 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MoeAlign, Tensor, Tensor, Tensor, const Tensor &, const size_t, const size_t, const bool); +INFINICORE_GRAPH_OP_CLASS(MoeAlignWithExpertMap, Tensor, Tensor, Tensor, const Tensor &, const Tensor &, const size_t, const size_t, const bool); + +std::tuple moe_align( + const Tensor &topk_ids, + size_t num_experts, + size_t block_size, + bool pad_sorted_token_ids = true); + +void moe_align_( + Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + size_t num_experts, + size_t block_size, + bool pad_sorted_token_ids = true); + +void moe_align_with_expert_map_( + Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + const Tensor &expert_map, + size_t num_experts, + size_t block_size, + bool pad_sorted_token_ids = true); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/moe_argsort_bincount.hpp b/python/infinicore/include/infinicore/ops/moe_argsort_bincount.hpp new file mode 100644 index 000000000..b710383b9 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_argsort_bincount.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../tensor.hpp" +#include + +namespace infinicore::op { +void moe_argsort_bincount_with_inv_pos_(Tensor tokens_per_experts, Tensor sorted_indices, Tensor inv_pos, const Tensor &topk_ids, int64_t num_experts); +} diff --git a/python/infinicore/include/infinicore/ops/moe_expand_input.hpp b/python/infinicore/include/infinicore/ops/moe_expand_input.hpp new file mode 100644 index 000000000..0f1f9cabd --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_expand_input.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "../tensor.hpp" +#include +#include + +namespace infinicore::op { +// format: 0=normal, 1=quant, 2=packed. group_size is used by quant/packed paths. +void moe_expand_input_with_inv_pos_(Tensor expand_states, std::optional expand_scales, const Tensor &hidden_states, const Tensor &inv_pos, int64_t top_k, int64_t group_size = 128, int64_t format = 0); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/moe_fused_dense.hpp b/python/infinicore/include/infinicore/ops/moe_fused_dense.hpp new file mode 100644 index 000000000..0a34c734c --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_fused_dense.hpp @@ -0,0 +1,42 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MoeFusedDense, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &); + +Tensor moe_fused_dense( + const Tensor &hidden_states, + const Tensor &w13, + const Tensor &w2, + const Tensor &topk_weights, + const Tensor &topk_ids, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded); + +void moe_fused_dense_( + Tensor output, + const Tensor &hidden_states, + const Tensor &w13, + const Tensor &w2, + const Tensor &topk_weights, + const Tensor &topk_ids, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/moe_fused_gate.hpp b/python/infinicore/include/infinicore/ops/moe_fused_gate.hpp new file mode 100644 index 000000000..e20d0a657 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_fused_gate.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MoeFusedGate, + Tensor, + Tensor, + const Tensor &, + const Tensor &, + const size_t, + const size_t, + const size_t, + const float, + const bool); + +std::tuple moe_fused_gate( + const Tensor &input, + const Tensor &bias, + size_t topk, + size_t num_expert_group, + size_t topk_group, + size_t num_fused_shared_experts, + float routed_scaling_factor, + bool apply_routed_scaling_factor_on_output); + +void moe_fused_gate_( + Tensor topk_weights, + Tensor topk_indices, + const Tensor &input, + const Tensor &bias, + size_t num_expert_group, + size_t topk_group, + size_t num_fused_shared_experts, + float routed_scaling_factor, + bool apply_routed_scaling_factor_on_output); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/moe_silu_and_mul_quant.hpp b/python/infinicore/include/infinicore/ops/moe_silu_and_mul_quant.hpp new file mode 100644 index 000000000..496c9f3e9 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_silu_and_mul_quant.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "../tensor.hpp" +#include +#include + +namespace infinicore::op { +// input shape [M, 2N]. format: 0=normal fp output, 1=quant int8+scale, 2=packed int8+scale. +void moe_silu_and_mul_quant_(Tensor output, std::optional output_scale, const Tensor &input, int64_t format = 0); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/moe_sum.hpp b/python/infinicore/include/infinicore/ops/moe_sum.hpp new file mode 100644 index 000000000..e9623e3b1 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_sum.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MoeSum, Tensor, const Tensor &); + +Tensor moe_sum(const Tensor &input); +void moe_sum_(Tensor output, const Tensor &input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/moe_sum_vllm.hpp b/python/infinicore/include/infinicore/ops/moe_sum_vllm.hpp new file mode 100644 index 000000000..93aaab35b --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_sum_vllm.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../tensor.hpp" +#include + +namespace infinicore::op { +void moe_sum_vllm_(Tensor output, const Tensor &input, std::optional topk_weights = std::nullopt, std::optional extra_residual = std::nullopt, double routed_scale = 1.0, double residual_scale = 1.0); +} diff --git a/python/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp b/python/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp new file mode 100644 index 000000000..e796c1dc8 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MoeTopkSigmoid, + Tensor, + Tensor, + const Tensor &, + const Tensor &, + const bool); + +std::tuple moe_topk_sigmoid( + const Tensor &gating_output, + size_t topk, + bool renormalize = false, + const Tensor &correction_bias = Tensor()); + +void moe_topk_sigmoid_( + Tensor topk_weights, + Tensor topk_indices, + const Tensor &gating_output, + const Tensor &correction_bias = Tensor(), + bool renormalize = false); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/moe_topk_softmax.hpp b/python/infinicore/include/infinicore/ops/moe_topk_softmax.hpp new file mode 100644 index 000000000..e6964aa88 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_topk_softmax.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MoeTopkSoftmax, + Tensor, + Tensor, + const Tensor &, + const Tensor &, + const bool, + const float); + +std::tuple moe_topk_softmax( + const Tensor &gating_output, + size_t topk, + bool renormalize = false, + float moe_softcapping = 0.0f, + const Tensor &correction_bias = Tensor()); + +void moe_topk_softmax_( + Tensor topk_weights, + Tensor topk_indices, + const Tensor &gating_output, + const Tensor &correction_bias = Tensor(), + bool renormalize = false, + float moe_softcapping = 0.0f); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/moe_topk_vllm.hpp b/python/infinicore/include/infinicore/ops/moe_topk_vllm.hpp new file mode 100644 index 000000000..cedba1746 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/moe_topk_vllm.hpp @@ -0,0 +1,7 @@ +#pragma once +#include "../device.hpp" +#include "../tensor.hpp" +namespace infinicore::op { +void moe_topk_softmax_vllm_(Tensor topk_weights, Tensor topk_ids, Tensor token_expert_indices, const Tensor &gating_output, bool renormalize = false, const Tensor &correction_bias = Tensor()); +void moe_topk_sigmoid_vllm_(Tensor topk_weights, Tensor topk_ids, Tensor token_expert_indices, const Tensor &gating_output, bool renormalize = false, const Tensor &correction_bias = Tensor()); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/mrope.hpp b/python/infinicore/include/infinicore/ops/mrope.hpp new file mode 100644 index 000000000..3ee2c1ef1 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/mrope.hpp @@ -0,0 +1,51 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MRoPE, + Tensor, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + int, + int, + int, + int, + int, + bool); + +void mrope_(Tensor q_out, + Tensor k_out, + const Tensor &q, + const Tensor &k, + const Tensor &cos, + const Tensor &sin, + const Tensor &positions, + int head_size, + int rotary_dim, + int section_t, + int section_h, + int section_w, + bool interleaved); + +std::pair mrope(const Tensor &q, + const Tensor &k, + const Tensor &cos, + const Tensor &sin, + const Tensor &positions, + int head_size, + int rotary_dim, + int section_t, + int section_h, + int section_w, + bool interleaved); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/mul.hpp b/python/infinicore/include/infinicore/ops/mul.hpp new file mode 100644 index 000000000..2eb480ddb --- /dev/null +++ b/python/infinicore/include/infinicore/ops/mul.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Mul, Tensor, const Tensor &, const Tensor &); + +Tensor mul(const Tensor &a, const Tensor &b); +void mul_(Tensor c, const Tensor &a, const Tensor &b); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/mul_scalar.hpp b/python/infinicore/include/infinicore/ops/mul_scalar.hpp new file mode 100644 index 000000000..5ee5e0d0b --- /dev/null +++ b/python/infinicore/include/infinicore/ops/mul_scalar.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MulScalar, Tensor, const Tensor &, double); + +Tensor mul_scalar(const Tensor &a, double alpha); +void mul_scalar_(Tensor c, const Tensor &a, double alpha); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/multi_margin_loss.hpp b/python/infinicore/include/infinicore/ops/multi_margin_loss.hpp new file mode 100644 index 000000000..e36cf469b --- /dev/null +++ b/python/infinicore/include/infinicore/ops/multi_margin_loss.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class MultiMarginLoss { +public: + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, int64_t, float, int64_t); + + static void execute(Tensor output, Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction); + static common::OpDispatcher &dispatcher(); +}; + +Tensor multi_margin_loss(Tensor input, Tensor target, Tensor weight = {}, int64_t p = 1, float margin = 1.0f, int64_t reduction = 1); +void multi_margin_loss_(Tensor output, Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/nrm2.hpp b/python/infinicore/include/infinicore/ops/nrm2.hpp new file mode 100644 index 000000000..c5552a4c1 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/nrm2.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Nrm2, const Tensor &, Tensor); + +Tensor nrm2(const Tensor &x); +void nrm2_(const Tensor &x, Tensor result); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/nsa_compress_paged_cache.hpp b/python/infinicore/include/infinicore/ops/nsa_compress_paged_cache.hpp new file mode 100644 index 000000000..bb5380632 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/nsa_compress_paged_cache.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(NsaCompressPagedCache, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, int, bool); + +void nsa_compress_paged_cache_(Tensor k_cmp, Tensor v_cmp, const Tensor &k_cache, const Tensor &v_cache, + const Tensor &block_tables, const Tensor &kv_lens, int nsa_block_size, + bool update_last_only = false); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/nsa_paged_attention.hpp b/python/infinicore/include/infinicore/ops/nsa_paged_attention.hpp new file mode 100644 index 000000000..00e9adc10 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/nsa_paged_attention.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(NsaPagedAttention, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, float, int, int, int); + +Tensor nsa_paged_attention(const Tensor &q, const Tensor &k_cmp, const Tensor &v_cmp, const Tensor &k_cache, const Tensor &v_cache, + const Tensor &block_tables, const Tensor &kv_lens, const Tensor &gates, + float scale, int nsa_block_size, int window_size, int select_blocks); + +void nsa_paged_attention_(Tensor out, const Tensor &q, const Tensor &k_cmp, const Tensor &v_cmp, const Tensor &k_cache, const Tensor &v_cache, + const Tensor &block_tables, const Tensor &kv_lens, const Tensor &gates, + float scale, int nsa_block_size, int window_size, int select_blocks); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/ones.hpp b/python/infinicore/include/infinicore/ops/ones.hpp new file mode 100644 index 000000000..0ea747484 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/ones.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "common/op.hpp" + +namespace infinicore::op { +class Ones { + +public: + using schema = void (*)(Tensor); + static void execute(Tensor output); + static common::OpDispatcher &dispatcher(); +}; + +Tensor ones(); +void ones_(Tensor output); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/pad.hpp b/python/infinicore/include/infinicore/ops/pad.hpp new file mode 100644 index 000000000..e598a94b4 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/pad.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +#include +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Pad, Tensor, const Tensor &, const std::vector &, const std::string &, double); + +Tensor pad(const Tensor &x, + const std::vector &pad, + const std::string &mode = "constant", + double value = 0.0); + +void pad_(Tensor y, + const Tensor &x, + const std::vector &pad, + const std::string &mode = "constant", + double value = 0.0); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/paged_attention.hpp b/python/infinicore/include/infinicore/ops/paged_attention.hpp new file mode 100644 index 000000000..8c906c95e --- /dev/null +++ b/python/infinicore/include/infinicore/ops/paged_attention.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(PagedAttention, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, std::optional, float); + +Tensor paged_attention(const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, + const Tensor &block_tables, const Tensor &kv_lens, + std::optional alibi_slopes, float scale); + +void paged_attention_(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, + const Tensor &block_tables, const Tensor &kv_lens, + std::optional alibi_slopes, float scale); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/paged_attention_mla.hpp b/python/infinicore/include/infinicore/ops/paged_attention_mla.hpp new file mode 100644 index 000000000..99034003a --- /dev/null +++ b/python/infinicore/include/infinicore/ops/paged_attention_mla.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../tensor.hpp" + +#include + +namespace infinicore::op { + +void paged_attention_mla_(Tensor output, + const Tensor &query, + const Tensor &kv_cache, + float scale, + const Tensor &block_tables, + const Tensor &context_lens, + int64_t max_context_len); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/paged_attention_prefill.hpp b/python/infinicore/include/infinicore/ops/paged_attention_prefill.hpp new file mode 100644 index 000000000..952924528 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/paged_attention_prefill.hpp @@ -0,0 +1,52 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +class PagedAttentionPrefill { +public: + /** + * @brief PagedAttentionPrefill operator signature + * * Argument order: + * 1. out: Output tensor (Packed format) + * 2. q: Current Query tensor (Packed format) + * 3. k_cache: Physical Key cache (Paged format) + * 4. v_cache: Physical Value cache (Paged format) + * 5. block_tables: Mapping table from logical blocks to physical blocks + * 6. total_kv_lens: lengths of Complete Key/Value for each request + * 7. cu_seqlens_q: Cumulative sequence lengths of Query (prefix sum for variable-length batch) + * 8. alibi_slopes: ALiBi bias slopes (optional) + * 9. scale: Scaling factor (typically 1/sqrt(head_size)) + */ + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, std::optional, float); + + static void execute(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, + Tensor block_tables, Tensor total_kv_lens, Tensor cum_seqlens_q, + std::optional alibi_slopes, float scale); + + static common::OpDispatcher &dispatcher(); +}; + +Tensor paged_attention_prefill(Tensor q, + Tensor k_cache, + Tensor v_cache, + Tensor block_tables, + Tensor total_kv_lens, + Tensor cum_seqlens_q, + std::optional alibi_slopes, + float scale); + +void paged_attention_prefill_(Tensor out, + Tensor q, + Tensor k_cache, + Tensor v_cache, + Tensor block_tables, + Tensor total_kv_lens, + Tensor cum_seqlens_q, + std::optional alibi_slopes, + float scale); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/paged_caching.hpp b/python/infinicore/include/infinicore/ops/paged_caching.hpp new file mode 100644 index 000000000..403b4b738 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/paged_caching.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(PagedCaching, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &); + +void paged_caching_(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/per_channel_quant_i8.hpp b/python/infinicore/include/infinicore/ops/per_channel_quant_i8.hpp new file mode 100644 index 000000000..0b0296248 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/per_channel_quant_i8.hpp @@ -0,0 +1,12 @@ +#pragma once +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(PerChannelQuantI8, const Tensor &, Tensor, Tensor); + +void per_channel_quant_i8_(const Tensor &x, Tensor x_packed, Tensor x_scale); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/per_tensor_dequant_i8.hpp b/python/infinicore/include/infinicore/ops/per_tensor_dequant_i8.hpp new file mode 100644 index 000000000..a6c4b1381 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/per_tensor_dequant_i8.hpp @@ -0,0 +1,11 @@ +#pragma once +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(PerTensorDequantI8, Tensor, const Tensor &, const Tensor &, const Tensor &); + +void per_tensor_dequant_i8_(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zero); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/per_tensor_quant_i8.hpp b/python/infinicore/include/infinicore/ops/per_tensor_quant_i8.hpp new file mode 100644 index 000000000..84ec99c37 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/per_tensor_quant_i8.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(PerTensorQuantI8, const Tensor &, Tensor, Tensor, Tensor, bool); + +void per_tensor_quant_i8_(const Tensor &x, Tensor x_packed, Tensor x_scale, Tensor x_zero, bool is_static); + +Tensor per_tensor_quant_i8(const Tensor &x, Tensor x_scale, Tensor x_zero, bool is_static); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/prelu.hpp b/python/infinicore/include/infinicore/ops/prelu.hpp new file mode 100644 index 000000000..e22c73d70 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/prelu.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "infinicore.h" + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Prelu, Tensor, const Tensor &, const Tensor &); + +__export Tensor prelu(const Tensor &input, const Tensor &weight); +__export void prelu_(Tensor out, const Tensor &input, const Tensor &weight); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/prepare_moe_input.hpp b/python/infinicore/include/infinicore/ops/prepare_moe_input.hpp new file mode 100644 index 000000000..c06951320 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/prepare_moe_input.hpp @@ -0,0 +1,48 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(PrepareMoeInput, + Tensor, + Tensor, + Tensor, + Tensor, + Tensor, + Tensor, + const Tensor &, + const size_t, + const size_t, + const size_t); + +struct PrepareMoeInputOutput { + Tensor expert_offsets; + Tensor problem_sizes1; + Tensor problem_sizes2; + Tensor input_permutation; + Tensor output_permutation; +}; + +PrepareMoeInputOutput prepare_moe_input( + const Tensor &topk_ids, + size_t num_experts, + size_t n, + size_t k); + +void prepare_moe_input_( + Tensor expert_offsets, + Tensor blockscale_offsets, + Tensor problem_sizes1, + Tensor problem_sizes2, + Tensor input_permutation, + Tensor output_permutation, + const Tensor &topk_ids, + size_t num_experts, + size_t n, + size_t k); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/quickgelu.hpp b/python/infinicore/include/infinicore/ops/quickgelu.hpp new file mode 100644 index 000000000..b67d7f71a --- /dev/null +++ b/python/infinicore/include/infinicore/ops/quickgelu.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class QuickGelu { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor quick_gelu(Tensor input); +void quick_gelu_(Tensor output, Tensor input); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/random_sample.hpp b/python/infinicore/include/infinicore/ops/random_sample.hpp new file mode 100644 index 000000000..a3a9d29f1 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/random_sample.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +#include "infinicore/tensor.hpp" + +namespace infinicore::op { + +class RandomSample { +public: + using schema = void (*)(Tensor, Tensor, float, float, int, float); + static void execute(Tensor indices, Tensor logits, float random_val, float topp, int topk, float temperature); + static common::OpDispatcher &dispatcher(); +}; + +// Out-of-place API +Tensor random_sample(Tensor logits, float random_val, float topp, int topk, float temperature); +// In-place API +void random_sample_(Tensor indices, Tensor logits, float random_val, float topp, int topk, float temperature); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/rearrange.hpp b/python/infinicore/include/infinicore/ops/rearrange.hpp new file mode 100644 index 000000000..5db983ef8 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/rearrange.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Rearrange, Tensor, const Tensor &); + +Tensor rearrange(const Tensor &x); +void rearrange_(Tensor y, const Tensor &x); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/reciprocal.hpp b/python/infinicore/include/infinicore/ops/reciprocal.hpp new file mode 100644 index 000000000..172fcdf36 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/reciprocal.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Reciprocal { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor y, Tensor x); + static common::OpDispatcher &dispatcher(); +}; + +Tensor reciprocal(Tensor x); +void reciprocal_(Tensor y, Tensor x); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/recurrent_gated_delta_rule.hpp b/python/infinicore/include/infinicore/ops/recurrent_gated_delta_rule.hpp new file mode 100644 index 000000000..7d837c67b --- /dev/null +++ b/python/infinicore/include/infinicore/ops/recurrent_gated_delta_rule.hpp @@ -0,0 +1,55 @@ +#pragma once + +#include "infinicore.h" + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(RecurrentGatedDeltaRule, + Tensor, + Tensor, + std::optional, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + std::optional, + std::optional, + bool); + +__export Tensor recurrent_gated_delta_rule(const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &g, + const Tensor &beta, + const Tensor &initial_state, + bool use_qk_l2norm = false); + +__export Tensor recurrent_gated_delta_rule_indexed(const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &g, + const Tensor &beta, + Tensor initial_state, + const Tensor &initial_state_indices, + const Tensor &final_state_indices, + bool use_qk_l2norm = false); + +__export void recurrent_gated_delta_rule_(Tensor out, + Tensor initial_state, + std::optional final_state, + const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &g, + const Tensor &beta, + std::optional initial_state_indices, + std::optional final_state_indices, + bool use_qk_l2norm = false); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/relu.hpp b/python/infinicore/include/infinicore/ops/relu.hpp new file mode 100644 index 000000000..59f2b26f5 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/relu.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Relu { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor relu(Tensor input); +void relu_(Tensor output, Tensor input); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/relu6.hpp b/python/infinicore/include/infinicore/ops/relu6.hpp new file mode 100644 index 000000000..7dd5635f8 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/relu6.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "infinicore.h" + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Relu6, Tensor, const Tensor &); + +__export Tensor relu6(const Tensor &input); +__export void relu6_(Tensor out, const Tensor &input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/rms_norm.hpp b/python/infinicore/include/infinicore/ops/rms_norm.hpp new file mode 100644 index 000000000..c7b2b2d72 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/rms_norm.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(RMSNorm, Tensor, const Tensor &, const Tensor &, float); + +Tensor rms_norm(const Tensor &x, const Tensor &weight, float epsilon = 1e-5f); +void rms_norm_(Tensor y, const Tensor &x, const Tensor &weight, float epsilon = 1e-5f); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/rope.hpp b/python/infinicore/include/infinicore/ops/rope.hpp new file mode 100644 index 000000000..8fd630ce1 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/rope.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../nn/rope.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(RoPE, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, infinicore::nn::RoPE::Algo); + +// Internal +void rope_(Tensor x_out, + const Tensor &x, + const Tensor &pos, + const Tensor &sin_table, + const Tensor &cos_table, + infinicore::nn::RoPE::Algo algo); + +// Public API +Tensor rope(const Tensor &x, + const Tensor &pos, + const Tensor &sin_table, + const Tensor &cos_table, + infinicore::nn::RoPE::Algo algo); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/rot.hpp b/python/infinicore/include/infinicore/ops/rot.hpp new file mode 100644 index 000000000..ff473f33e --- /dev/null +++ b/python/infinicore/include/infinicore/ops/rot.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(Rot, Tensor, Tensor, const Tensor &, const Tensor &); + +void rot_(Tensor x, Tensor y, const Tensor &c, const Tensor &s); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/rotg.hpp b/python/infinicore/include/infinicore/ops/rotg.hpp new file mode 100644 index 000000000..c65f211dd --- /dev/null +++ b/python/infinicore/include/infinicore/ops/rotg.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Rotg { +public: + using schema = void (*)(Tensor, Tensor, Tensor, Tensor); + static void execute(Tensor x, Tensor y, Tensor c, Tensor s); + static common::OpDispatcher &dispatcher(); +}; + +void rotg_(Tensor x, Tensor y, Tensor c, Tensor s); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/rotm.hpp b/python/infinicore/include/infinicore/ops/rotm.hpp new file mode 100644 index 000000000..813fc27ee --- /dev/null +++ b/python/infinicore/include/infinicore/ops/rotm.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Rotm { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + static void execute(Tensor x, Tensor y, Tensor param); + static common::OpDispatcher &dispatcher(); +}; + +void rotm_(Tensor x, Tensor y, Tensor param); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/rotmg.hpp b/python/infinicore/include/infinicore/ops/rotmg.hpp new file mode 100644 index 000000000..e245840a4 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/rotmg.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Rotmg { +public: + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor); + static void execute(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param); + static common::OpDispatcher &dispatcher(); +}; + +void rotmg_(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/rwkv5_wkv.hpp b/python/infinicore/include/infinicore/ops/rwkv5_wkv.hpp new file mode 100644 index 000000000..f1e6accd1 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/rwkv5_wkv.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS( + Rwkv5Wkv, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + Tensor); + +Tensor rwkv5_wkv(const Tensor &receptance, + const Tensor &key, + const Tensor &value, + const Tensor &time_decay, + const Tensor &time_faaaa, + Tensor state); + +void rwkv5_wkv_(Tensor out, + const Tensor &receptance, + const Tensor &key, + const Tensor &value, + const Tensor &time_decay, + const Tensor &time_faaaa, + Tensor state); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/scal.hpp b/python/infinicore/include/infinicore/ops/scal.hpp new file mode 100644 index 000000000..d6cb5ce8c --- /dev/null +++ b/python/infinicore/include/infinicore/ops/scal.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(Scal, const Tensor &, Tensor); + +void scal_(const Tensor &alpha, Tensor x); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/scaled_mm_i8.hpp b/python/infinicore/include/infinicore/ops/scaled_mm_i8.hpp new file mode 100644 index 000000000..427ed9d65 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/scaled_mm_i8.hpp @@ -0,0 +1,13 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(I8Gemm, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, std::optional); + +void scaled_mm_i8_(Tensor c, const Tensor &a_p, const Tensor &a_s, const Tensor &b_p, const Tensor &b_s, std::optional bias); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp b/python/infinicore/include/infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp new file mode 100644 index 000000000..33ee7c37b --- /dev/null +++ b/python/infinicore/include/infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp @@ -0,0 +1,13 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(GptqQyblasGemm, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, int64_t, int64_t); + +void scaled_mm_w4a16_gptq_qy_(Tensor out, const Tensor &in, const Tensor &qweight, const Tensor &scales, const Tensor &qzeros, int64_t quant_type, int64_t bit); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/scaled_mm_w4a8.hpp b/python/infinicore/include/infinicore/ops/scaled_mm_w4a8.hpp new file mode 100644 index 000000000..c7d1f2f1f --- /dev/null +++ b/python/infinicore/include/infinicore/ops/scaled_mm_w4a8.hpp @@ -0,0 +1,9 @@ +#pragma once + +#include "../tensor.hpp" +#include + +namespace infinicore::op { +void scaled_mm_w4a8_(Tensor out, const Tensor &a, const Tensor &b, const Tensor &a_scales, const Tensor &b_scales, std::optional bias = std::nullopt, bool trans_weight = false); +Tensor scaled_mm_w4a8(const Tensor &a, const Tensor &b, const Tensor &a_scales, const Tensor &b_scales, std::optional bias = std::nullopt, bool trans_weight = false); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/scaled_mm_w8a8.hpp b/python/infinicore/include/infinicore/ops/scaled_mm_w8a8.hpp new file mode 100644 index 000000000..54b7a4ee9 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/scaled_mm_w8a8.hpp @@ -0,0 +1,12 @@ +#pragma once +#include "infinicore/tensor.hpp" +#include + +namespace infinicore::op { +Tensor scaled_mm_w8a8(const Tensor &a, const Tensor &b, + const Tensor &a_scales, const Tensor &b_scales, + std::optional bias = std::nullopt, bool trans_weight = true); +void scaled_mm_w8a8_(Tensor out, const Tensor &a, const Tensor &b, + const Tensor &a_scales, const Tensor &b_scales, + std::optional bias = std::nullopt, bool trans_weight = true); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/scatter.hpp b/python/infinicore/include/infinicore/ops/scatter.hpp new file mode 100644 index 000000000..306bb1232 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/scatter.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Scatter { +public: + using schema = void (*)(Tensor, Tensor, int64_t, Tensor, Tensor, int64_t); + + static void execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction); + static common::OpDispatcher &dispatcher(); +}; + +Tensor scatter(Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction = 0); + +// In-place / 指定 Output 接口 +void scatter_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/select_last_token_hidden.hpp b/python/infinicore/include/infinicore/ops/select_last_token_hidden.hpp new file mode 100644 index 000000000..9b9907051 --- /dev/null +++ b/python/infinicore/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/include/infinicore/ops/selu.hpp b/python/infinicore/include/infinicore/ops/selu.hpp new file mode 100644 index 000000000..2aed070ed --- /dev/null +++ b/python/infinicore/include/infinicore/ops/selu.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Selu { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor selu(Tensor input); +void selu_(Tensor output, Tensor input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/sigmoid.hpp b/python/infinicore/include/infinicore/ops/sigmoid.hpp new file mode 100644 index 000000000..e5346608f --- /dev/null +++ b/python/infinicore/include/infinicore/ops/sigmoid.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "infinicore.h" + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Sigmoid, Tensor, const Tensor &); + +__export Tensor sigmoid(const Tensor &input); +__export void sigmoid_(Tensor output, const Tensor &input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/silu.hpp b/python/infinicore/include/infinicore/ops/silu.hpp new file mode 100644 index 000000000..d4cbac57d --- /dev/null +++ b/python/infinicore/include/infinicore/ops/silu.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Silu { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor silu(Tensor input); +void silu_(Tensor output, Tensor input); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/silu_and_mul.hpp b/python/infinicore/include/infinicore/ops/silu_and_mul.hpp new file mode 100644 index 000000000..1b594111d --- /dev/null +++ b/python/infinicore/include/infinicore/ops/silu_and_mul.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(SiluAndMul, Tensor, const Tensor &); + +Tensor silu_and_mul(const Tensor &x); +void silu_and_mul_(Tensor out, const Tensor &x); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/sinh.hpp b/python/infinicore/include/infinicore/ops/sinh.hpp new file mode 100644 index 000000000..31cb3c256 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/sinh.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Sinh { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor sinh(Tensor input); +void sinh_(Tensor output, Tensor input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/smooth_l1_loss.hpp b/python/infinicore/include/infinicore/ops/smooth_l1_loss.hpp new file mode 100644 index 000000000..9a67c6109 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/smooth_l1_loss.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class SmoothL1Loss { +public: + using schema = void (*)(Tensor, Tensor, Tensor, float, int64_t); + + static void execute(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction); + static common::OpDispatcher &dispatcher(); +}; + +Tensor smooth_l1_loss(Tensor input, Tensor target, float beta = 1.0f, int64_t reduction = 1); +void smooth_l1_loss_(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/softmax.hpp b/python/infinicore/include/infinicore/ops/softmax.hpp new file mode 100644 index 000000000..c96ab1810 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/softmax.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Softmax { +public: + using schema = void (*)(Tensor, Tensor, int); + static void execute(Tensor output, Tensor input, int axis); + static common::OpDispatcher &dispatcher(); +}; + +Tensor softmax(Tensor input, int axis = -1); +void softmax_(Tensor output, Tensor input, int axis = -1); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/softplus.hpp b/python/infinicore/include/infinicore/ops/softplus.hpp new file mode 100644 index 000000000..b6d665455 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/softplus.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Softplus { +public: + // 修改 1: Schema 增加 float beta, float threshold + using schema = void (*)(Tensor, Tensor, float, float); + static void execute(Tensor y, Tensor x, float beta, float threshold); + static common::OpDispatcher &dispatcher(); +}; +// default: beta = 1.0, threshold = 20.0 +Tensor softplus(Tensor x, float beta = 1.0f, float threshold = 20.0f); + +void softplus_(Tensor y, Tensor x, float beta = 1.0f, float threshold = 20.0f); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/softsign.hpp b/python/infinicore/include/infinicore/ops/softsign.hpp new file mode 100644 index 000000000..05e8e54fa --- /dev/null +++ b/python/infinicore/include/infinicore/ops/softsign.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Softsign { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor y, Tensor x); + static common::OpDispatcher &dispatcher(); +}; +// 返回新 Tensor 的函数接口 +Tensor softsign(Tensor x); +void softsign_(Tensor y, Tensor x); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/sum.hpp b/python/infinicore/include/infinicore/ops/sum.hpp new file mode 100644 index 000000000..0ead8de26 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/sum.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include +#include + +namespace infinicore::op { +class Sum { +public: + using schema = void (*)(Tensor, Tensor, std::vector, bool); + static void execute(Tensor output, Tensor input, std::vector dim, bool keepdim = false); + static common::OpDispatcher &dispatcher(); +}; + +Tensor sum(Tensor input, std::vector dim, bool keepdim = false); +void sum_(Tensor output, Tensor input, std::vector dim, bool keepdim = false); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/swap.hpp b/python/infinicore/include/infinicore/ops/swap.hpp new file mode 100644 index 000000000..aba3ad563 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/swap.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(Swap, Tensor, Tensor); + +void swap_(Tensor x, Tensor y); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/swiglu.hpp b/python/infinicore/include/infinicore/ops/swiglu.hpp new file mode 100644 index 000000000..7aa77e632 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/swiglu.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(SwiGLU, Tensor, const Tensor &, const Tensor &); + +Tensor swiglu(const Tensor &a, const Tensor &b); +void swiglu_(Tensor c, const Tensor &a, const Tensor &b); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/take.hpp b/python/infinicore/include/infinicore/ops/take.hpp new file mode 100644 index 000000000..5672005ad --- /dev/null +++ b/python/infinicore/include/infinicore/ops/take.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Take { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + + static void execute(Tensor output, Tensor input, Tensor indices); + static common::OpDispatcher &dispatcher(); +}; + +Tensor take(Tensor input, Tensor indices); + +void take_(Tensor output, Tensor input, Tensor indices); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/tan.hpp b/python/infinicore/include/infinicore/ops/tan.hpp new file mode 100644 index 000000000..72410f866 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/tan.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Tan { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor tan(Tensor input); +void tan_(Tensor output, Tensor input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/tanhshrink.hpp b/python/infinicore/include/infinicore/ops/tanhshrink.hpp new file mode 100644 index 000000000..83673b22c --- /dev/null +++ b/python/infinicore/include/infinicore/ops/tanhshrink.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { +class Tanhshrink { +public: + using schema = void (*)(Tensor, Tensor); + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +Tensor tanhshrink(Tensor input); +void tanhshrink_(Tensor output, Tensor input); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/topk.hpp b/python/infinicore/include/infinicore/ops/topk.hpp new file mode 100644 index 000000000..d8486112c --- /dev/null +++ b/python/infinicore/include/infinicore/ops/topk.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +namespace infinicore::op { +class TopK { +public: + using schema = void (*)(Tensor, Tensor, Tensor, size_t, size_t, bool, bool); + static void execute(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true); + static common::OpDispatcher &dispatcher(); +}; + +std::pair topk(Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true); +void topk_(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/topksoftmax.hpp b/python/infinicore/include/infinicore/ops/topksoftmax.hpp new file mode 100644 index 000000000..2cd3ac7c5 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/topksoftmax.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Topksoftmax, Tensor, Tensor, const Tensor &, const size_t, const int); + +void topksoftmax(Tensor values, Tensor indices, const Tensor &x, const size_t topk, const int norm = 0); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/triplet_margin_loss.hpp b/python/infinicore/include/infinicore/ops/triplet_margin_loss.hpp new file mode 100644 index 000000000..43b935eb8 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/triplet_margin_loss.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class TripletMarginLoss { +public: + // Schema signature: output, anchor, positive, negative, margin, p, eps, swap, reduction + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float, int64_t, float, bool, int64_t); + + static void execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction); + static common::OpDispatcher &dispatcher(); +}; + +// Functional API +// reduction: 0=None, 1=Mean, 2=Sum +Tensor triplet_margin_loss(Tensor anchor, Tensor positive, Tensor negative, float margin = 1.0f, int64_t p = 2, float eps = 1e-6f, bool swap = false, int64_t reduction = 1); + +// In-place / Explicit Output API +void triplet_margin_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/triplet_margin_with_distance_loss.hpp b/python/infinicore/include/infinicore/ops/triplet_margin_with_distance_loss.hpp new file mode 100644 index 000000000..d6ea449c5 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/triplet_margin_with_distance_loss.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class TripletMarginWithDistanceLoss { +public: + // Schema signature: output(out), anchor, positive, negative, margin, swap, reduction + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, double, bool, int64_t); + + static void execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction); + static common::OpDispatcher &dispatcher(); +}; + +// Functional API: Returns the result tensor +// margin default 1.0, swap default false, reduction default 1 (Mean) typically +Tensor triplet_margin_with_distance_loss(Tensor anchor, Tensor positive, Tensor negative, double margin = 1.0, bool swap = false, int64_t reduction = 1); + +// In-place/Output-provided API +void triplet_margin_with_distance_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/unfold.hpp b/python/infinicore/include/infinicore/ops/unfold.hpp new file mode 100644 index 000000000..f8c7a3b13 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/unfold.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +class Unfold { +public: + // schema: output, input, kernel_sizes, dilations, paddings, strides + using schema = void (*)(Tensor, Tensor, const std::vector &, const std::vector &, const std::vector &, const std::vector &); + + static void execute(Tensor output, Tensor input, + const std::vector &kernel_sizes, + const std::vector &dilations, + const std::vector &paddings, + const std::vector &strides); + static common::OpDispatcher &dispatcher(); +}; + +// Functional API +Tensor unfold(Tensor input, + std::vector kernel_sizes, + std::vector dilations, + std::vector paddings, + std::vector strides); + +void unfold_(Tensor output, Tensor input, + std::vector kernel_sizes, + std::vector dilations, + std::vector paddings, + std::vector strides); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/upsample_bilinear.hpp b/python/infinicore/include/infinicore/ops/upsample_bilinear.hpp new file mode 100644 index 000000000..ba9aa2379 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/upsample_bilinear.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +class UpsampleBilinear { +public: + // Schema signature: output, input, align_corners + using schema = void (*)(Tensor, Tensor, bool); + + static void execute(Tensor output, Tensor input, bool align_corners); + static common::OpDispatcher &dispatcher(); +}; + +// 需要传入 output_size (如 {H_out, W_out} 或 {N, C, H_out, W_out}) 来决定新 Tensor 的形状 +Tensor upsample_bilinear(Tensor input, std::vector output_size, bool align_corners = false); +void upsample_bilinear_(Tensor output, Tensor input, bool align_corners); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/upsample_nearest.hpp b/python/infinicore/include/infinicore/ops/upsample_nearest.hpp new file mode 100644 index 000000000..188d1b923 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/upsample_nearest.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +class UpsampleNearest { +public: + // Schema signature: output(out), input + // Note: Scales are inferred from output.shape / input.shape + using schema = void (*)(Tensor, Tensor); + + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +// Functional API: Returns the result tensor +// Requires output_size to calculate the shape of the result tensor +Tensor upsample_nearest(Tensor input, const std::vector &output_size); + +// In-place/Output-provided API +void upsample_nearest_(Tensor output, Tensor input); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/vander.hpp b/python/infinicore/include/infinicore/ops/vander.hpp new file mode 100644 index 000000000..8198cfeb8 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/vander.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class Vander { +public: + // schema: output, input, N, increasing + using schema = void (*)(Tensor, Tensor, int64_t, bool); + + static void execute(Tensor output, Tensor input, int64_t N, bool increasing); + static common::OpDispatcher &dispatcher(); +}; + +// N defaults to 0 (implying N = input.size(0), i.e., a square matrix) +Tensor vander(Tensor input, int64_t N = 0, bool increasing = false); +void vander_(Tensor output, Tensor input, int64_t N, bool increasing); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/var.hpp b/python/infinicore/include/infinicore/ops/var.hpp new file mode 100644 index 000000000..d1e01e1bf --- /dev/null +++ b/python/infinicore/include/infinicore/ops/var.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include +#include +#include +namespace infinicore::op { +class Var { +public: + using schema = void (*)(Tensor, Tensor, std::vector, bool, bool); // var_output, input, dim, unbiased, keepdim + static void execute(Tensor var_output, Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); + static common::OpDispatcher &dispatcher(); +}; + +Tensor var(Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); +void var_(Tensor var_output, Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/var_mean.hpp b/python/infinicore/include/infinicore/ops/var_mean.hpp new file mode 100644 index 000000000..a9679187c --- /dev/null +++ b/python/infinicore/include/infinicore/ops/var_mean.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include +#include +#include +namespace infinicore::op { +class Var_Mean { +public: + using schema = void (*)(Tensor, Tensor, Tensor, std::vector, bool, bool); // var_output, mean_output, input, dim, unbiased, keepdim + static void execute(Tensor var_output, Tensor mean_output, Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); + static common::OpDispatcher &dispatcher(); +}; + +std::pair var_mean(Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); +void var_mean_(Tensor var_output, Tensor mean_output, Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); + +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/vocab_parallel_embedding.hpp b/python/infinicore/include/infinicore/ops/vocab_parallel_embedding.hpp new file mode 100644 index 000000000..25390664a --- /dev/null +++ b/python/infinicore/include/infinicore/ops/vocab_parallel_embedding.hpp @@ -0,0 +1,6 @@ +#pragma once +#include "../tensor.hpp" +#include +namespace infinicore::op { +void vocab_parallel_embedding_(Tensor output, const Tensor &indices, const Tensor &weight, int64_t vocab_start, int64_t vocab_end); +} diff --git a/python/infinicore/include/infinicore/ops/w16a16_group_gemm.hpp b/python/infinicore/include/infinicore/ops/w16a16_group_gemm.hpp new file mode 100644 index 000000000..15b949f85 --- /dev/null +++ b/python/infinicore/include/infinicore/ops/w16a16_group_gemm.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../tensor.hpp" +#include + +namespace infinicore::op { +void w16a16_group_gemm_(Tensor out, + const Tensor &input, + const Tensor &weight, + const Tensor &tokens_per_experts, + std::optional sorted_token_ids = std::nullopt, + std::optional bias = std::nullopt, + bool trans_weight = true, + bool is_decode = false); +} // namespace infinicore::op diff --git a/python/infinicore/include/infinicore/ops/w4a8_group_gemm.hpp b/python/infinicore/include/infinicore/ops/w4a8_group_gemm.hpp new file mode 100644 index 000000000..60312c14c --- /dev/null +++ b/python/infinicore/include/infinicore/ops/w4a8_group_gemm.hpp @@ -0,0 +1,6 @@ +#pragma once +#include "../tensor.hpp" +#include +namespace infinicore::op { +void w4a8_group_gemm_(Tensor out, const Tensor &input, const Tensor &weight, const Tensor &input_scale, const Tensor &weight_scale, const Tensor &tokens_per_experts, std::optional sorted_token_ids = std::nullopt, std::optional bias = std::nullopt, bool trans_weight = true, bool is_decode = false); +} diff --git a/python/infinicore/include/infinicore/ops/w8a8_group_gemm.hpp b/python/infinicore/include/infinicore/ops/w8a8_group_gemm.hpp new file mode 100644 index 000000000..68239471c --- /dev/null +++ b/python/infinicore/include/infinicore/ops/w8a8_group_gemm.hpp @@ -0,0 +1,6 @@ +#pragma once +#include "../tensor.hpp" +#include +namespace infinicore::op { +void w8a8_group_gemm_(Tensor out, const Tensor &input, const Tensor &weight, const Tensor &input_scale, const Tensor &weight_scale, const Tensor &tokens_per_experts, std::optional sorted_token_ids = std::nullopt, std::optional bias = std::nullopt, bool trans_weight = true, bool is_decode = false); +} diff --git a/python/infinicore/include/infinicore/quantization.hpp b/python/infinicore/include/infinicore/quantization.hpp new file mode 100644 index 000000000..58964c412 --- /dev/null +++ b/python/infinicore/include/infinicore/quantization.hpp @@ -0,0 +1,9 @@ +#pragma once + +#include "quantization/awq.hpp" +#include "quantization/base_quantization.hpp" +#include "quantization/compressed_tensors.hpp" +#include "quantization/gptq.hpp" +#include "quantization/gptq_qy.hpp" +#include "quantization/none_quantizaiton.hpp" +#include "quantization/quantization_scheme.hpp" diff --git a/python/infinicore/include/infinicore/quantization/awq.hpp b/python/infinicore/include/infinicore/quantization/awq.hpp new file mode 100644 index 000000000..d64b1315b --- /dev/null +++ b/python/infinicore/include/infinicore/quantization/awq.hpp @@ -0,0 +1,30 @@ +#pragma once +#include "base_quantization.hpp" +namespace infinicore::quantization { + +class AWQ : public BaseQuantization { + // This is a temporary class that currently only returns AWQ_W4A16. + // Future enhancements should parse quant_config to extract detailed quantization + // information and support multiple quantization schemes. +public: + explicit AWQ(const nlohmann::json &quant_config) + : BaseQuantization(quant_config){}; + + infinicore::quantization::QuantScheme + get_quant_scheme() const override { + return infinicore::quantization::QuantScheme::AWQ_W4A16; + }; + + int get_packing_num() const { + // For AWQ, we pack 8 int4 weights into a single int32 value. + return 32 / this->get_or("bits", 4); // Default to 8 if not specified in config + } + + int get_group_size() const { + // For simplicity, we return a fixed group size here. In a more complete implementation, + // this could be extracted from quant_config_ to support different group sizes. + return this->get_or("group_size", 128); // Standard AWQ group size + } +}; + +} // namespace infinicore::quantization diff --git a/python/infinicore/include/infinicore/quantization/base_quantization.hpp b/python/infinicore/include/infinicore/quantization/base_quantization.hpp new file mode 100644 index 000000000..4ee4f7608 --- /dev/null +++ b/python/infinicore/include/infinicore/quantization/base_quantization.hpp @@ -0,0 +1,43 @@ +#pragma once +#include "nlohmann/json.hpp" +#include "quantization_scheme.hpp" + +namespace infinicore::quantization { +class BaseQuantization { + // Base class for quantization schemes. Intended to be extended to support various quantization methods. +public: + explicit BaseQuantization(const nlohmann::json &quant_config) : quant_config_(quant_config){}; + virtual ~BaseQuantization() = default; + + const nlohmann::json &get_config() const { return quant_config_; } + + virtual infinicore::quantization::QuantScheme get_quant_scheme() const = 0; + template + T get(const std::string &key) const { + if (!quant_config_.contains(key)) { + throw std::out_of_range("Key '" + key + "' not found in config."); + } + try { + return quant_config_.at(key).get(); + } catch (const nlohmann::json::type_error &e) { + throw std::runtime_error("Type conversion failed for key '" + key + "': " + std::string(e.what())); + } + } + + template + T get_or(const std::string &key, const T &default_value) const { + if (!quant_config_.contains(key) || quant_config_.at(key).is_null()) { + return default_value; + } + try { + return quant_config_.at(key).get(); + } catch (const nlohmann::json::type_error &) { + // If type conversion fails, return default value + return default_value; + } + } + +protected: + nlohmann::json quant_config_; +}; +} // namespace infinicore::quantization diff --git a/python/infinicore/include/infinicore/quantization/compressed_tensors.hpp b/python/infinicore/include/infinicore/quantization/compressed_tensors.hpp new file mode 100644 index 000000000..0549c1f00 --- /dev/null +++ b/python/infinicore/include/infinicore/quantization/compressed_tensors.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "base_quantization.hpp" +namespace infinicore::quantization { + +class CompressedTensors : public BaseQuantization { + // This is a temporary class that currently only returns COMPRESSED_TENSOR_W8A8I8. + // Future enhancements should parse quant_config to extract detailed quantization + // information and support multiple quantization schemes. +public: + explicit CompressedTensors(const nlohmann::json &quant_config) + : BaseQuantization(quant_config){}; + + infinicore::quantization::QuantScheme + get_quant_scheme() const override { + return infinicore::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8; + }; +}; + +} // namespace infinicore::quantization diff --git a/python/infinicore/include/infinicore/quantization/gptq.hpp b/python/infinicore/include/infinicore/quantization/gptq.hpp new file mode 100644 index 000000000..c03416b07 --- /dev/null +++ b/python/infinicore/include/infinicore/quantization/gptq.hpp @@ -0,0 +1,30 @@ +#pragma once +#include "base_quantization.hpp" +namespace infinicore::quantization { + +class GPTQ : public BaseQuantization { + // This is a temporary class that currently only returns GPTQ W4A16. + // Future enhancements should parse quant_config to extract detailed quantization + // information and support multiple quantization schemes. +public: + explicit GPTQ(const nlohmann::json &quant_config) + : BaseQuantization(quant_config){}; + + infinicore::quantization::QuantScheme + get_quant_scheme() const override { + return infinicore::quantization::QuantScheme::GPTQ_W4A16; + }; + + int get_packing_num() const { + // For GPTQ, we pack 8 int4 weights into a single int32 value. + return 32 / this->get_or("bits", 4); // Default to 8 if not specified in config + } + + int get_group_size() const { + // For simplicity, we return a fixed group size here. In a more complete implementation, + // this could be extracted from quant_config_ to support different group sizes. + return this->get_or("group_size", 128); // Standard GPTQ group size + } +}; + +} // namespace infinicore::quantization diff --git a/python/infinicore/include/infinicore/quantization/gptq_qy.hpp b/python/infinicore/include/infinicore/quantization/gptq_qy.hpp new file mode 100644 index 000000000..f3ca47c59 --- /dev/null +++ b/python/infinicore/include/infinicore/quantization/gptq_qy.hpp @@ -0,0 +1,320 @@ +#pragma once + +#include "../tensor.hpp" +#include "base_quantization.hpp" +#include +#include +#include +#include +#include +#include +#include + +namespace { +#ifndef INFINICORE_FLOAT16_DEFINED +#define INFINICORE_FLOAT16_DEFINED +struct float16_raw { + uint16_t data; + float16_raw() : data(0) {} + explicit float16_raw(float f) : data(fp32_to_fp16_bits(f)) {} + + static uint16_t fp32_to_fp16_bits(float value) { + union { + float f; + uint32_t u; + } f2u; + f2u.f = value; + uint32_t x = f2u.u; + + uint32_t sign = (x >> 16) & 0x8000; + int32_t exp = ((x >> 23) & 0xFF) - 127; + uint32_t mantissa = x & 0x007FFFFF; + + if (exp == 128) { + if (mantissa == 0) { + return static_cast(sign | 0x7C00); + } else { + return static_cast(sign | 0x7C00 | (mantissa >> 13)); + } + } + if (exp > 15) { + return static_cast(sign | 0x7C00); + } + if (exp < -14) { + if (exp < -24) { + return static_cast(sign); + } + mantissa |= 0x00800000; + uint32_t shift = -exp - 14; + mantissa >>= shift; + if ((mantissa & 0x1000) && ((mantissa & 0x2FFF) != 0)) { + mantissa += 0x2000; + } + return static_cast(sign | (mantissa >> 13)); + } + + uint32_t exp16 = static_cast(exp + 15) << 10; + uint32_t mantissa16 = mantissa >> 13; + if ((mantissa & 0x1000) && ((mantissa & 0x2FFF) || (mantissa16 & 1))) { + mantissa16++; + if (mantissa16 == 0x400) { + exp16 += 0x400; + mantissa16 = 0; + } + } + return static_cast(sign | exp16 | mantissa16); + } +}; +using float16_t = float16_raw; +#endif + +inline std::vector float_to_fp16_bits(const std::vector &values) { + std::vector result; + result.reserve(values.size()); + for (float f : values) { +#ifdef INFINICORE_HAS_FLOAT16 + infinicore::float16_t h(f); + result.push_back(*reinterpret_cast(&h)); +#else + result.push_back(float16_raw::fp32_to_fp16_bits(f)); +#endif + } + return result; +} +} // anonymous namespace + +namespace infinicore::quantization { + +class GPTQ_QY : public BaseQuantization { +public: + explicit GPTQ_QY(const nlohmann::json &quant_config) + : BaseQuantization(quant_config) { + int bits = weight_bits(); + if (bits != 4) { + spdlog::warn("GPTQ_QY: bits={} not fully tested, expected 4", bits); + } + } + + QuantScheme get_quant_scheme() const override { + return QuantScheme::GPTQ_W4A16_QY; + } + + int get_packing_num() const { + return 32 / weight_bits(); + } + + int get_group_size() const { + return get_or("group_size", 128); + } + + void convert_from_gptq_w4a16(const Tensor &original_qweight, + const Tensor &original_qzeros, + const Tensor &original_scales, + const Tensor &g_idx, + const Device &target_device) { + if (converted_) { + spdlog::debug("GPTQ_QY: weights already converted, skipping"); + return; + } + + const int bits = weight_bits(); + const int values_per_int32 = 32 / bits; + + { + const auto &shape = original_qweight->shape(); + assert(shape.size() == 2); + size_t M = shape[0], N = shape[1]; + + auto weight_unpacked = unpack_int32_to_nibbles_3d_(original_qweight, bits); + auto weight_packed = combine_nibbles_last_dim_(weight_unpacked, M, values_per_int32, N); + + size_t dimY = N; + size_t total_bytes = M * values_per_int32 * (N / 2); + size_t dimX = total_bytes / dimY; + + assert(dimX * dimY == total_bytes && "Weight shape calculation mismatch"); + + converted_weight_ = make_tensor_from_host_( + weight_packed.data(), + total_bytes * sizeof(uint8_t), + {dimX, dimY}, + DataType::U8, + target_device); + } + + { + const auto &shape = original_qzeros->shape(); + assert(shape.size() == 2); + size_t P = shape[0], Q = shape[1]; + + auto zeros_fp32 = unpack_zeros_to_fp32_2d_(original_qzeros, bits); + auto zeros_fp16 = ::float_to_fp16_bits(zeros_fp32); + + converted_zeros_ = make_tensor_from_host_( + zeros_fp16.data(), + zeros_fp16.size() * sizeof(uint16_t), + {P, Q * static_cast(values_per_int32)}, + DataType::F16, + target_device); + } + + { + auto scales_cpu = original_scales->to(Device::Type::CPU); + size_t num_elements = scales_cpu->numel(); + const void *raw_data = scales_cpu->data(); + + std::vector scales_fp16(num_elements); + if (scales_cpu->dtype() == DataType::F16) { + std::memcpy(scales_fp16.data(), raw_data, num_elements * sizeof(uint16_t)); + } else if (scales_cpu->dtype() == DataType::F32) { + std::vector scales_fp32(num_elements); + std::memcpy(scales_fp32.data(), raw_data, num_elements * sizeof(float)); + scales_fp16 = ::float_to_fp16_bits(scales_fp32); + } else { + spdlog::error("Unsupported scales dtype, expected F16 or F32"); + assert(false && "Unsupported scales dtype"); + } + + converted_scales_ = make_tensor_from_host_( + scales_fp16.data(), + scales_fp16.size() * sizeof(uint16_t), + original_scales->shape(), + DataType::F16, + target_device); + } + + if (g_idx->numel() > 0) { + g_idx_ = g_idx->to(target_device); + } + + converted_ = true; + } + + void release_buffers() { + converted_weight_ = Tensor(); + converted_zeros_ = Tensor(); + converted_scales_ = Tensor(); + g_idx_ = Tensor(); + } + + void convert_and_take_ownership(Tensor &weight, Tensor &zeros, Tensor &scales, + const Tensor &g_idx, const Device &target_device) { + if (converted_) { + spdlog::warn("GPTQ_QY: Already converted, skipping"); + return; + } + + convert_from_gptq_w4a16(weight, zeros, scales, g_idx, target_device); + + weight = std::move(converted_weight_); + zeros = std::move(converted_zeros_); + scales = std::move(converted_scales_); + + converted_ = false; + spdlog::debug("GPTQ_QY: Ownership transferred, internal buffers cleared."); + } + + const Tensor &get_converted_weight() const { return std::move(converted_weight_); } + const Tensor &get_converted_zeros() const { return std::move(converted_zeros_); } + const Tensor &get_converted_scales() const { return std::move(converted_scales_); } + const Tensor &get_g_idx() const { return g_idx_; } + bool is_converted() const { return converted_; } + + int weight_bits() const { return get_or("bits", 4); } + bool desc_act() const { return get_or("desc_act", false); } + +private: + static inline std::vector unpack_int32_to_nibbles_3d_(const Tensor &packed, int bits) { + assert(bits == 4 || bits == 8); + const int values_per_int32 = 32 / bits; + + auto packed_cpu = packed->to(Device::Type::CPU); + const int32_t *packed_host = reinterpret_cast(packed_cpu->data()); + + const auto &shape = packed->shape(); + assert(shape.size() == 2); + size_t M = shape[0], N = shape[1]; + + std::vector unpacked(M * values_per_int32 * N); + + for (size_t i = 0; i < M; ++i) { + for (int k = 0; k < values_per_int32; ++k) { + for (size_t j = 0; j < N; ++j) { + int32_t val = packed_host[i * N + j]; + uint8_t extracted = static_cast((val >> (k * bits)) & ((1 << bits) - 1)); + size_t idx = i * (values_per_int32 * N) + k * N + j; + unpacked[idx] = extracted; + } + } + } + return unpacked; + } + + static inline std::vector combine_nibbles_last_dim_( + const std::vector &nibbles, size_t M, size_t K, size_t N) { + assert(N % 2 == 0 && "Last dimension must be even for nibble pairing"); + + std::vector combined(M * K * (N / 2)); + size_t out_idx = 0; + + for (size_t i = 0; i < M; ++i) { + for (size_t k = 0; k < K; ++k) { + size_t row_base = i * (K * N) + k * N; + for (size_t j = 0; j < N; j += 2) { + uint8_t low = nibbles[row_base + j] & 0x0F; + uint8_t high = nibbles[row_base + j + 1] & 0x0F; + combined[out_idx++] = static_cast((high << 4) | low); + } + } + } + return combined; + } + + static inline std::vector unpack_zeros_to_fp32_2d_(const Tensor &packed_zeros, int bits) { + assert(bits == 4 || bits == 8); + const int values_per_int32 = 32 / bits; + const int mask = (1 << bits) - 1; + + auto packed_cpu = packed_zeros->to(Device::Type::CPU); + const int32_t *packed_host = reinterpret_cast(packed_cpu->data()); + + const auto &shape = packed_zeros->shape(); + assert(shape.size() == 2); + size_t P = shape[0], Q = shape[1]; + + std::vector result(P * Q * values_per_int32); + size_t out_idx = 0; + + for (size_t p = 0; p < P; ++p) { + for (size_t q = 0; q < Q; ++q) { + int32_t val = packed_host[p * Q + q]; + for (int k = 0; k < values_per_int32; ++k) { + uint8_t extracted = static_cast((val >> (k * bits)) & mask); + int dequant_val = (static_cast(extracted) + 1) & mask; + result[out_idx++] = static_cast(dequant_val); + } + } + } + return result; + } + + static inline Tensor make_tensor_from_host_(const void *data, size_t bytes, + const std::vector &shape, + DataType dtype, const Device &device) { + auto tensor = Tensor::empty(shape, dtype, Device::Type::CPU); + std::memcpy(reinterpret_cast(tensor->data()), data, bytes); + + if (device != Device::Type::CPU) { + return tensor->to(device); + } + return tensor; + } + + Tensor converted_weight_; + Tensor converted_zeros_; + Tensor converted_scales_; + Tensor g_idx_; + bool converted_ = false; +}; + +} // namespace infinicore::quantization diff --git a/python/infinicore/include/infinicore/quantization/none_quantizaiton.hpp b/python/infinicore/include/infinicore/quantization/none_quantizaiton.hpp new file mode 100644 index 000000000..5009f0adc --- /dev/null +++ b/python/infinicore/include/infinicore/quantization/none_quantizaiton.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "base_quantization.hpp" +namespace infinicore::quantization { + +class NoneQuantization : public BaseQuantization { + // This is a temporary class that currently only returns COMPRESSED_TENSOR_W8A8I8. + // Future enhancements should parse quant_config to extract detailed quantization + // information and support multiple quantization schemes. +public: + explicit NoneQuantization(const nlohmann::json &quant_config) + : BaseQuantization(quant_config){}; + + infinicore::quantization::QuantScheme + get_quant_scheme() const override { + return infinicore::quantization::QuantScheme::NONE; + }; +}; + +} // namespace infinicore::quantization diff --git a/python/infinicore/include/infinicore/quantization/quantization_scheme.hpp b/python/infinicore/include/infinicore/quantization/quantization_scheme.hpp new file mode 100644 index 000000000..3f2daff09 --- /dev/null +++ b/python/infinicore/include/infinicore/quantization/quantization_scheme.hpp @@ -0,0 +1,19 @@ +// quant.hpp +#pragma once + +namespace infinicore::quantization { + +enum class QuantScheme { + NONE, + COMPRESSED_TENSOR_W8A8I8, + AWQ_W4A16, + GPTQ_W4A16_QY, + GPTQ_W4A16, +}; + +enum class KVQuantAlgo { + NONE, + INT8, +}; + +} // namespace infinicore::quantization diff --git a/python/infinicore/include/infinicore/tensor.hpp b/python/infinicore/include/infinicore/tensor.hpp new file mode 100644 index 000000000..b1650417d --- /dev/null +++ b/python/infinicore/include/infinicore/tensor.hpp @@ -0,0 +1,336 @@ +#pragma once + +#include "device.hpp" +#include "dtype.hpp" +#include "memory.hpp" + +#include +#include +#include +#include +#include + +#include +namespace infinicore { + +using Size = std::size_t; +using Stride = std::ptrdiff_t; +using Shape = std::vector; +using Strides = std::vector; + +class TensorImpl; + +struct TensorMetaData { + Shape shape; + Strides strides; + DataType dtype; + infiniopTensorDescriptor_t desc; + + TensorMetaData(const Shape &shape, const Strides &strides, const DataType &dtype); + ~TensorMetaData(); +}; + +struct TensorData { + size_t offset; + std::shared_ptr memory; +}; + +struct TensorSliceParams { + size_t dim; + size_t start; + Size len; +}; + +class Tensor { +public: + static Tensor empty(const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory = false); + + static Tensor strided_empty(const Shape &shape, + const Strides &strides, + const DataType &dtype, + const Device &device, + bool pin_memory = false); + + static Tensor zeros(const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory = false); + + static Tensor ones(const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory = false); + + static Tensor from_blob(void *raw_ptr, + const Shape &shape, + const DataType &dtype, + const Device &device); + + static Tensor strided_from_blob(void *raw_ptr, + const Shape &shape, + const Strides &strides, + const DataType &dtype, + const Device &device); + + void reset() noexcept { + impl_.reset(); + } + + void reset(std::shared_ptr new_impl) noexcept { + impl_ = std::move(new_impl); + } + + bool empty() const noexcept { + return impl_ == nullptr; + } + + size_t use_count() const noexcept { + return impl_ ? impl_.use_count() : 0; + } + + Tensor() = default; + Tensor(const Tensor &) = default; + Tensor(Tensor &&) = default; + Tensor &operator=(const Tensor &) = default; + Tensor &operator=(Tensor &&) = default; + + TensorImpl *operator->(); + const TensorImpl *operator->() const; + + operator bool() const; + +protected: + Tensor(std::shared_ptr impl) : impl_(std::move(impl)) {} + std::shared_ptr impl_; + friend class TensorImpl; + friend std::ostream &operator<<(std::ostream &out, const Tensor &tensor); +}; + +class TensorImpl : public std::enable_shared_from_this { + +public: + TensorImpl(const Shape &shape, const DataType &dtype); + TensorImpl(const Shape &shape, const Strides &strides, const DataType &dtype); + + std::byte *data(); + const std::byte *data() const; + + const Shape &shape() const; + + const Strides &strides() const; + + bool is_contiguous() const; + + Size ndim() const; + + Size numel() const; + + Size size(size_t dim) const; + + size_t element_size() const; + + size_t nbytes() const; + + Stride stride(size_t dim) const; + + DataType dtype() const; + + Device device() const; + + infiniopTensorDescriptor_t desc() const; + + bool is_pinned() const; + + std::string info() const; + + void debug(const std::string &filename) const; + + void debug() const; + + /** + * Unsafe API that returns a new tensor with the same raw memory untracked by allocator + * This API is used for loosely tracking a piece of memory while allowing it to be reused, + * typically in a compute graph scenario. + */ + Tensor to_blob_() const; + + /** + * Unsafe API that returns a new tensor with the same memory and let allocator retracks the memory. + * Should only be used on the tensor returned by to_blob_(). + */ + Tensor resume_from_blob_() const; + + /// + /// Data Transfer APIs + /// + + /** + * Returns a new tensor with the same data on a different device. + * If the new device passed is same as the current device, the original tensor is returned. + * + * @param device The device of the new tensor + * + * @return A new tensor with the same data on the specified device + */ + Tensor to(Device device) const; + + /** + * Copy Data from another tensor to this tensor. + * + * @param src The source tensor to copy from + * + * @return A new tensor with the same data on the specified device + */ + void copy_from(Tensor src); + + /** + * Return a tensor with the same data in contiguous arrangement as current tensor. + * If this tensor is already contiguous, the original tensor is returned. + * + * @return A new tensor with the same data on the specified device + */ + Tensor contiguous() const; + + /// + /// View APIs + /// + + /** + * Returns a new tensor with a dimension of size one removed at the specified position. + * Throws runtime_error if the dimension to be removed is not of size 1. + * + * @param dim The dimension index to remove + * @return A new tensor with the removed dimension + * + * Example: + * // For a 3D tensor with shape [1, 3, 4], squeeze at dim 0 results in shape [3, 4] + * tensor->squeeze(0); + */ + Tensor squeeze(size_t dim) const; + + /** + * Returns a new tensor with a dimension of size one inserted at the specified position. + * The returned tensor shares the same underlying storage with the original tensor. + * + * @param dim The dimension index at which to insert the new dimension + * @return A new tensor with the added dimension + * + * Example: + * // For a 2D tensor with shape [3, 4], unsqueeze at dim 0 results in shape [1, 3, 4] + * // unsqueeze at dim 1 results in shape [3, 1, 4] + * // unsqueeze at dim 2 results in shape [3, 4, 1] + * tensor->unsqueeze(0); + */ + Tensor unsqueeze(size_t dim) const; + + /** + * Returns a new tensor that is a narrowed version of the current tensor. + * The returned tensor shares the same underlying storage with the original tensor. + * + * @param slices A vector of slice parameters specifying the dimension, start index, + * and length for each dimension to narrow + * @return A new tensor with narrowed dimensions + * + * Example: + * // Narrow dimension 0 from index 2 to 5 (length 3) + * // and dimension 1 from index 1 to 3 (length 2) + * tensor.narrow({{0, 2, 3}, {1, 1, 2}}); + */ + Tensor narrow(const std::vector &slices) const; + + /** + * Returns a new tensor with the dimensions permuted (reordered) according to the given order. + * The returned tensor shares the same underlying storage with the original tensor. + * + * @param order The desired ordering of dimensions + * @return A new tensor with permuted dimensions + * + * Example: + * // For a 3D tensor with shape [2, 3, 4], permute to [2, 0, 1] + * // This swaps the dimensions: dim0->dim2, dim1->dim0, dim2->dim1 + * tensor->permute({2, 0, 1}); + */ + Tensor permute(const Shape &order) const; + + /** + * Returns a new tensor with the same data but a different shape. + * The returned tensor shares the same underlying storage with the original tensor. + * The tensor is rearranged if the new shape is not compatible with the current shape. + * + * @param new_shape The desired new shape + * @return A new tensor with the specified shape + * + * Example: + * // Reshape a 2x3 tensor (6 elements) to a 3x2 tensor + * tensor->view({3, 2}); + */ + Tensor view(const Shape &new_shape) const; + + /** + * Insecurely returns a new tensor with the specified shape and strides. + * The returned tensor shares the same underlying storage with the original tensor. + * + * @param new_shape The desired new shape + * @param new_strides The desired new strides + * @return A new tensor with the specified shape and strides + * + * Example: + * // Create a non-contiguous view with custom strides + * tensor->as_strided({2, 3}, {6, 2}); // Stride of 6 for dim0, 2 for dim1 + */ + Tensor as_strided(const Shape &new_shape, const Strides &new_strides) const; + +protected: + static std::shared_ptr empty( + const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory = false); + + static std::shared_ptr strided_empty( + const Shape &shape, + const Strides &strides, + const DataType &dtype, + const Device &device, + bool pin_memory = false); + + static std::shared_ptr zeros( + const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory = false); + + static std::shared_ptr ones( + const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory = false); + + static std::shared_ptr from_blob( + void *raw_ptr, + const Shape &shape, + const DataType &dtype, + const Device &device); + + static std::shared_ptr strided_from_blob( + void *raw_ptr, + const Shape &shape, + const Strides &strides, + const DataType &dtype, + const Device &device); + + friend class Tensor; + +protected: + TensorMetaData meta_; + TensorData data_; + +private: + // Mark to indicate if the tensor is created from to_blob_() + bool to_blob_mark_ = false; +}; + +} // namespace infinicore diff --git a/python/infinicore/include/infiniop.h b/python/infinicore/include/infiniop.h new file mode 100644 index 000000000..ddf8c47cb --- /dev/null +++ b/python/infinicore/include/infiniop.h @@ -0,0 +1,164 @@ +#ifndef __INFINIOP_API_H__ +#define __INFINIOP_API_H__ + +#include "infiniop/handle.h" +#include "infiniop/ops/acos.h" +#include "infiniop/ops/adaptive_avg_pool1d.h" +#include "infiniop/ops/adaptive_avg_pool3d.h" +#include "infiniop/ops/adaptive_max_pool1d.h" +#include "infiniop/ops/add.h" +#include "infiniop/ops/add_rms_norm.h" +#include "infiniop/ops/addbmm.h" +#include "infiniop/ops/addcmul.h" +#include "infiniop/ops/addr.h" +#include "infiniop/ops/affine_grid.h" +#include "infiniop/ops/all.h" +#include "infiniop/ops/asin.h" +#include "infiniop/ops/asinh.h" +#include "infiniop/ops/asum.h" +#include "infiniop/ops/atanh.h" +#include "infiniop/ops/attention.h" +#include "infiniop/ops/avg_pool1d.h" +#include "infiniop/ops/avg_pool3d.h" +#include "infiniop/ops/awq_marlin_gemm.h" +#include "infiniop/ops/axpy.h" +#include "infiniop/ops/binary_cross_entropy_with_logits.h" +#include "infiniop/ops/blas_amax.h" +#include "infiniop/ops/blas_amin.h" +#include "infiniop/ops/blas_copy.h" +#include "infiniop/ops/blas_dot.h" +#include "infiniop/ops/block_diag.h" +#include "infiniop/ops/broadcast_to.h" +#include "infiniop/ops/causal_conv1d.h" +#include "infiniop/ops/causal_softmax.h" +#include "infiniop/ops/cdist.h" +#include "infiniop/ops/chunk_gated_delta_rule.h" +#include "infiniop/ops/clip.h" +#include "infiniop/ops/conv.h" +#include "infiniop/ops/cross_entropy.h" +#include "infiniop/ops/deepseek_moe.h" +#include "infiniop/ops/dequant/per_tensor_dequant_int8.h" +#include "infiniop/ops/dequantize_awq.h" +#include "infiniop/ops/dequantize_gptq.h" +#include "infiniop/ops/diff.h" +#include "infiniop/ops/digamma.h" +#include "infiniop/ops/dist.h" +#include "infiniop/ops/dot.h" +#include "infiniop/ops/embedding.h" +#include "infiniop/ops/equal.h" +#include "infiniop/ops/erf.h" +#include "infiniop/ops/erfc.h" +#include "infiniop/ops/erfinv.h" +#include "infiniop/ops/flash_attention.h" +#include "infiniop/ops/flipud.h" +#include "infiniop/ops/float_power.h" +#include "infiniop/ops/floor.h" +#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/gelu.h" +#include "infiniop/ops/gelutanh.h" +#include "infiniop/ops/gemm.h" +#include "infiniop/ops/gptq_marlin_gemm.h" +#include "infiniop/ops/gptq_qyblas_gemm.h" +#include "infiniop/ops/hardswish.h" +#include "infiniop/ops/hardtanh.h" +#include "infiniop/ops/hinge_embedding_loss.h" +#include "infiniop/ops/histc.h" +#include "infiniop/ops/huber_loss.h" +#include "infiniop/ops/hypot.h" +#include "infiniop/ops/index_add.h" +#include "infiniop/ops/index_copy.h" +#include "infiniop/ops/inner.h" +#include "infiniop/ops/int8_gemm.h" +#include "infiniop/ops/kron.h" +#include "infiniop/ops/kthvalue.h" +#include "infiniop/ops/kv_caching.h" +#include "infiniop/ops/layer_norm.h" +#include "infiniop/ops/ldexp.h" +#include "infiniop/ops/lerp.h" +#include "infiniop/ops/log10.h" +#include "infiniop/ops/log1p.h" +#include "infiniop/ops/log_softmax.h" +#include "infiniop/ops/logaddexp.h" +#include "infiniop/ops/logaddexp2.h" +#include "infiniop/ops/logcumsumexp.h" +#include "infiniop/ops/logdet.h" +#include "infiniop/ops/lp_norm.h" +#include "infiniop/ops/mamba_selective_scan.h" +#include "infiniop/ops/masked_select.h" +#include "infiniop/ops/matrix_power.h" +#include "infiniop/ops/moe_align.h" +#include "infiniop/ops/moe_fused_dense.h" +#include "infiniop/ops/moe_fused_gate.h" +#include "infiniop/ops/moe_sum.h" +#include "infiniop/ops/moe_topk_sigmoid.h" +#include "infiniop/ops/moe_topk_softmax.h" +#include "infiniop/ops/mrope.h" +#include "infiniop/ops/mul.h" +#include "infiniop/ops/mul_scalar.h" +#include "infiniop/ops/multi_margin_loss.h" +#include "infiniop/ops/nrm2.h" +#include "infiniop/ops/nsa_compress_paged_cache.h" +#include "infiniop/ops/nsa_paged_attention.h" +#include "infiniop/ops/ones.h" +#include "infiniop/ops/pad.h" +#include "infiniop/ops/paged_attention.h" +#include "infiniop/ops/paged_attention_prefill.h" +#include "infiniop/ops/paged_caching.h" +#include "infiniop/ops/pixel_shuffle.h" +#include "infiniop/ops/prepare_moe_input.h" +#include "infiniop/ops/quant/per_channel_quant_int8.h" +#include "infiniop/ops/quant/per_tensor_quant_int8.h" +#include "infiniop/ops/quickgelu.h" +#include "infiniop/ops/random_sample.h" +#include "infiniop/ops/rearrange.h" +#include "infiniop/ops/reciprocal.h" +#include "infiniop/ops/recurrent_gated_delta_rule.h" +#include "infiniop/ops/relu.h" +#include "infiniop/ops/rms_norm.h" +#include "infiniop/ops/rope.h" +#include "infiniop/ops/rot.h" +#include "infiniop/ops/rotg.h" +#include "infiniop/ops/rotm.h" +#include "infiniop/ops/rotmg.h" +#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" +#include "infiniop/ops/silu_and_mul.h" +#include "infiniop/ops/sinh.h" +#include "infiniop/ops/smooth_l1_loss.h" +#include "infiniop/ops/softmax.h" +#include "infiniop/ops/softplus.h" +#include "infiniop/ops/softsign.h" +#include "infiniop/ops/sub.h" +#include "infiniop/ops/sum.h" +#include "infiniop/ops/swap.h" +#include "infiniop/ops/swiglu.h" +#include "infiniop/ops/take.h" +#include "infiniop/ops/tan.h" +#include "infiniop/ops/tanh.h" +#include "infiniop/ops/tanhshrink.h" +#include "infiniop/ops/topk.h" +#include "infiniop/ops/topkrouter.h" +#include "infiniop/ops/topksoftmax.h" +#include "infiniop/ops/triplet_margin_loss.h" +#include "infiniop/ops/triplet_margin_with_distance_loss.h" +#include "infiniop/ops/unfold.h" +#include "infiniop/ops/upsample_bilinear.h" +#include "infiniop/ops/upsample_nearest.h" +#include "infiniop/ops/vander.h" +#include "infiniop/ops/var.h" +#include "infiniop/ops/var_mean.h" +#include "infiniop/ops/zeros.h" +#include "infiniop/tensor_descriptor.h" +#endif // __INFINIOP_API_H__ diff --git a/python/infinicore/include/infiniop/handle.h b/python/infinicore/include/infiniop/handle.h new file mode 100644 index 000000000..da8689542 --- /dev/null +++ b/python/infinicore/include/infiniop/handle.h @@ -0,0 +1,14 @@ +#ifndef __INFINIOP_HANDLE_API_H__ +#define __INFINIOP_HANDLE_API_H__ + +#include "../infinicore.h" + +struct InfiniopHandle; + +typedef struct InfiniopHandle *infiniopHandle_t; + +__INFINI_C __export infiniStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr); + +__INFINI_C __export infiniStatus_t infiniopDestroyHandle(infiniopHandle_t handle); + +#endif diff --git a/python/infinicore/include/infiniop/operator_descriptor.h b/python/infinicore/include/infiniop/operator_descriptor.h new file mode 100644 index 000000000..e46935bc0 --- /dev/null +++ b/python/infinicore/include/infiniop/operator_descriptor.h @@ -0,0 +1,13 @@ +#ifndef __INFINIOP_OPERATOR_DESCRIPTOR_API_H__ +#define __INFINIOP_OPERATOR_DESCRIPTOR_API_H__ + +#include "handle.h" +#include "tensor_descriptor.h" + +// Base descriptor for all operators +struct InfiniopDescriptor; + +__INFINI_C __export infiniStatus_t infiniopGetDescriptorDeviceType(const struct InfiniopDescriptor *desc_ptr, infiniDevice_t *device_type); +__INFINI_C __export infiniStatus_t infiniopGetDescriptorDeviceId(const struct InfiniopDescriptor *desc_ptr, int *device_id); + +#endif //__INFINIOP_OPERATOR_DESCRIPTOR_API_H__ diff --git a/python/infinicore/include/infiniop/ops/acos.h b/python/infinicore/include/infiniop/ops/acos.h new file mode 100644 index 000000000..011625d84 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/acos.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_ACOS_API_H__ +#define __INFINIOP_ACOS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAcosDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAcosDescriptor(infiniopHandle_t handle, + infiniopAcosDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetAcosWorkspaceSize(infiniopAcosDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAcos(infiniopAcosDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAcosDescriptor(infiniopAcosDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/adaptive_avg_pool1d.h b/python/infinicore/include/infiniop/ops/adaptive_avg_pool1d.h new file mode 100644 index 000000000..192dee03e --- /dev/null +++ b/python/infinicore/include/infiniop/ops/adaptive_avg_pool1d.h @@ -0,0 +1,34 @@ +#ifndef __INFINIOP_ADAPTIVE_AVG_POOL1D_API_H__ +#define __INFINIOP_ADAPTIVE_AVG_POOL1D_API_H__ + +#include "../operator_descriptor.h" + +// 定义算子描述符类型 +typedef struct InfiniopDescriptor *infiniopAdaptiveAvgPool1dDescriptor_t; + +// 1. 创建算子描述符 +__INFINI_C __export infiniStatus_t infiniopCreateAdaptiveAvgPool1dDescriptor( + infiniopHandle_t handle, + infiniopAdaptiveAvgPool1dDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t input_desc); + +// 2. 获取 Workspace 大小 +__INFINI_C __export infiniStatus_t infiniopGetAdaptiveAvgPool1dWorkspaceSize( + infiniopAdaptiveAvgPool1dDescriptor_t desc, + size_t *size); + +// 3. 执行计算 +__INFINI_C __export infiniStatus_t infiniopAdaptiveAvgPool1d( + infiniopAdaptiveAvgPool1dDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +// 4. 销毁描述符 +__INFINI_C __export infiniStatus_t infiniopDestroyAdaptiveAvgPool1dDescriptor( + infiniopAdaptiveAvgPool1dDescriptor_t desc); + +#endif // __INFINIOP_ADAPTIVE_AVG_POOL1D_API_H__ diff --git a/python/infinicore/include/infiniop/ops/adaptive_avg_pool3d.h b/python/infinicore/include/infiniop/ops/adaptive_avg_pool3d.h new file mode 100644 index 000000000..bd23ebce8 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/adaptive_avg_pool3d.h @@ -0,0 +1,30 @@ +#ifndef INFINIOP_ADAPTIVE_AVG_POOL3D_H_ +#define INFINIOP_ADAPTIVE_AVG_POOL3D_H_ + +#include "../operator_descriptor.h" +#include + +typedef struct InfiniopDescriptor *infiniopAdaptiveAvgPool3DDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAdaptiveAvgPool3DDescriptor( + infiniopHandle_t handle, + infiniopAdaptiveAvgPool3DDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + size_t *output_size); + +__INFINI_C __export infiniStatus_t infiniopGetAdaptiveAvgPool3DWorkspaceSize( + infiniopAdaptiveAvgPool3DDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAdaptiveAvgPool3D( + infiniopAdaptiveAvgPool3DDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAdaptiveAvgPool3DDescriptor(infiniopAdaptiveAvgPool3DDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/adaptive_max_pool1d.h b/python/infinicore/include/infiniop/ops/adaptive_max_pool1d.h new file mode 100644 index 000000000..0b42844d1 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/adaptive_max_pool1d.h @@ -0,0 +1,22 @@ +#ifndef __INFINIOP_ADAPTIVE_MAX_POOL1D_H__ +#define __INFINIOP_ADAPTIVE_MAX_POOL1D_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAdaptiveMaxPool1dDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAdaptiveMaxPool1dDescriptor( + infiniopHandle_t handle, + infiniopAdaptiveMaxPool1dDescriptor_t *desc, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + size_t output_size); + +__INFINI_C __export infiniStatus_t infiniopGetAdaptiveMaxPool1dWorkspaceSize(infiniopAdaptiveMaxPool1dDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAdaptiveMaxPool1d(infiniopAdaptiveMaxPool1dDescriptor_t desc, void *workspace, size_t workspace_size, + void *y, const void *x, void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAdaptiveMaxPool1dDescriptor(infiniopAdaptiveMaxPool1dDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/add.h b/python/infinicore/include/infiniop/ops/add.h new file mode 100644 index 000000000..8a7592c1c --- /dev/null +++ b/python/infinicore/include/infiniop/ops/add.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_ADD_API_H__ +#define __INFINIOP_ADD_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAddDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAddDescriptor(infiniopHandle_t handle, + infiniopAddDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__INFINI_C __export infiniStatus_t infiniopGetAddWorkspaceSize(infiniopAddDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAdd(infiniopAddDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAddDescriptor(infiniopAddDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/add_rms_norm.h b/python/infinicore/include/infiniop/ops/add_rms_norm.h new file mode 100644 index 000000000..2cde1db76 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/add_rms_norm.h @@ -0,0 +1,32 @@ +#ifndef __INFINIOP_ADD_RMS_NORM_API_H__ +#define __INFINIOP_ADD_RMS_NORM_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAddRMSNormDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAddRMSNormDescriptor( + infiniopHandle_t handle, + infiniopAddRMSNormDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t residual_out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t weight_desc, + float epsilon); + +__INFINI_C __export infiniStatus_t infiniopGetAddRMSNormWorkspaceSize(infiniopAddRMSNormDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAddRMSNorm(infiniopAddRMSNormDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + void *residual_out, + const void *a, + const void *b, + const void *weight, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAddRMSNormDescriptor(infiniopAddRMSNormDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/addbmm.h b/python/infinicore/include/infiniop/ops/addbmm.h new file mode 100644 index 000000000..04f918f2d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/addbmm.h @@ -0,0 +1,30 @@ +#ifndef __INFINIOP_ADDBMM_API_H__ +#define __INFINIOP_ADDBMM_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAddbmmDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAddbmmDescriptor(infiniopHandle_t handle, + infiniopAddbmmDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t batch1_desc, + infiniopTensorDescriptor_t batch2_desc, + float alpha, + float beta); + +__INFINI_C __export infiniStatus_t infiniopGetAddbmmWorkspaceSize(infiniopAddbmmDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAddbmm(infiniopAddbmmDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *batch1, + const void *batch2, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAddbmmDescriptor(infiniopAddbmmDescriptor_t desc); + +#endif // __INFINIOP_ADDBMM_API_H__ diff --git a/python/infinicore/include/infiniop/ops/addcmul.h b/python/infinicore/include/infiniop/ops/addcmul.h new file mode 100644 index 000000000..b858dccf0 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/addcmul.h @@ -0,0 +1,57 @@ +#ifndef __INFINIOP_ADDCMUL_API_H__ +#define __INFINIOP_ADDCMUL_API_H__ + +#include "../operator_descriptor.h" + +// 定义 addcmul 算子描述符类型 +typedef struct InfiniopDescriptor *infiniopAddcmulDescriptor_t; + +/** + * @brief 创建 Addcmul 算子描述符 + * @param handle 算子句柄 + * @param desc_ptr 指向返回的描述符指针 + * @param out 输出张量描述符 + * @param input 加项张量描述符 + * @param tensor1 乘项张量1描述符 + * @param tensor2 乘项张量2描述符 + * @param value 乘积的标量系数 + */ +__INFINI_C __export infiniStatus_t infiniopCreateAddcmulDescriptor(infiniopHandle_t handle, + infiniopAddcmulDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out, + infiniopTensorDescriptor_t input, + infiniopTensorDescriptor_t tensor1, + infiniopTensorDescriptor_t tensor2, + float value); + +/** + * @brief 获取 Addcmul 计算所需的临时空间大小 + */ +__INFINI_C __export infiniStatus_t infiniopGetAddcmulWorkspaceSize(infiniopAddcmulDescriptor_t desc, size_t *size); + +/** + * @brief 执行 Addcmul 计算 + * @param desc 算子描述符 + * @param workspace 临时空间指针 + * @param workspace_size 临时空间大小 + * @param out 输出数据指针 + * @param input 加项数据指针 + * @param tensor1 乘项1数据指针 + * @param tensor2 乘项2数据指针 + * @param stream 计算流 (CUDA stream 等) + */ +__INFINI_C __export infiniStatus_t infiniopAddcmul(infiniopAddcmulDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *input, + const void *tensor1, + const void *tensor2, + void *stream); + +/** + * @brief 销毁 Addcmul 算子描述符 + */ +__INFINI_C __export infiniStatus_t infiniopDestroyAddcmulDescriptor(infiniopAddcmulDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/addr.h b/python/infinicore/include/infiniop/ops/addr.h new file mode 100644 index 000000000..a2f26d330 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/addr.h @@ -0,0 +1,30 @@ +#ifndef __INFINIOP_ADDR_API_H__ +#define __INFINIOP_ADDR_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAddrDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAddrDescriptor(infiniopHandle_t handle, + infiniopAddrDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out, + infiniopTensorDescriptor_t input, + infiniopTensorDescriptor_t vec1, + infiniopTensorDescriptor_t vec2, + float beta, + float alpha); + +__INFINI_C __export infiniStatus_t infiniopGetAddrWorkspaceSize(infiniopAddrDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAddr(infiniopAddrDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *input, + const void *vec1, + const void *vec2, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAddrDescriptor(infiniopAddrDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/affine_grid.h b/python/infinicore/include/infiniop/ops/affine_grid.h new file mode 100644 index 000000000..390d4b85f --- /dev/null +++ b/python/infinicore/include/infiniop/ops/affine_grid.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_AFFINE_GRID_API_H__ +#define __INFINIOP_AFFINE_GRID_API_H__ +#include "../operator_descriptor.h" +#include + +typedef struct InfiniopDescriptor *infiniopAffineGridDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAffineGridDescriptor(infiniopHandle_t handle, + infiniopAffineGridDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + uint8_t align_corners); + +__INFINI_C __export infiniStatus_t infiniopGetAffineGridWorkspaceSize(infiniopAffineGridDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAffineGrid(infiniopAffineGridDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAffineGridDescriptor(infiniopAffineGridDescriptor_t desc); + +#endif // __INFINIOP_AFFINE_GRID_API_H__ diff --git a/python/infinicore/include/infiniop/ops/all.h b/python/infinicore/include/infiniop/ops/all.h new file mode 100644 index 000000000..41d74cf9a --- /dev/null +++ b/python/infinicore/include/infiniop/ops/all.h @@ -0,0 +1,31 @@ +#ifndef __INFINIOP_ALL_API_H__ +#define __INFINIOP_ALL_API_H__ + +#include "../operator_descriptor.h" +#include +#include +typedef struct InfiniopDescriptor *infiniopAllDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAllDescriptor(infiniopHandle_t handle, + infiniopAllDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + size_t *dim, + size_t dim_size, + bool keepdim); + +__INFINI_C __export infiniStatus_t infiniopGetAllWorkspaceSize(infiniopAllDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAll(infiniopAllDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + size_t *dim, + size_t dim_size, + bool keepdim, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAllDescriptor(infiniopAllDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/argwhere.h b/python/infinicore/include/infiniop/ops/argwhere.h new file mode 100644 index 000000000..b1d2f9f99 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/argwhere.h @@ -0,0 +1,29 @@ +#ifndef __INFINIOP_ARGWHERE_API_H__ +#define __INFINIOP_ARGWHERE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopArgwhereDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateArgwhereDescriptor( + infiniopHandle_t handle, + infiniopArgwhereDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t input_desc); + +__INFINI_C __export infiniStatus_t infiniopGetArgwhereWorkspaceSize( + infiniopArgwhereDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopArgwhere( + infiniopArgwhereDescriptor_t desc, + void *workspace, + size_t workspace_size, + void **output, + size_t *count, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyArgwhereDescriptor( + infiniopArgwhereDescriptor_t desc); + +#endif // __INFINIOP_ARGWHERE_API_H__ diff --git a/python/infinicore/include/infiniop/ops/asin.h b/python/infinicore/include/infiniop/ops/asin.h new file mode 100644 index 000000000..b4d16d481 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/asin.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_ASIN_API_H__ +#define __INFINIOP_ASIN_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAsinDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAsinDescriptor(infiniopHandle_t handle, + infiniopAsinDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input); + +__INFINI_C __export infiniStatus_t infiniopGetAsinWorkspaceSize(infiniopAsinDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAsin(infiniopAsinDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAsinDescriptor(infiniopAsinDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/asinh.h b/python/infinicore/include/infiniop/ops/asinh.h new file mode 100644 index 000000000..95c86bf79 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/asinh.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_ASINH_API_H__ +#define __INFINIOP_ASINH_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAsinhDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAsinhDescriptor(infiniopHandle_t handle, + infiniopAsinhDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetAsinhWorkspaceSize(infiniopAsinhDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAsinh(infiniopAsinhDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAsinhDescriptor(infiniopAsinhDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/asum.h b/python/infinicore/include/infiniop/ops/asum.h new file mode 100644 index 000000000..89336bce7 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/asum.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_ASUM_API_H__ +#define __INFINIOP_ASUM_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAsumDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAsumDescriptor(infiniopHandle_t handle, + infiniopAsumDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t result); + +__INFINI_C __export infiniStatus_t infiniopGetAsumWorkspaceSize(infiniopAsumDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAsum(infiniopAsumDescriptor_t desc, + void *workspace, + size_t workspace_size, + const void *x, + void *result, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAsumDescriptor(infiniopAsumDescriptor_t desc); + +#endif // __INFINIOP_ASUM_API_H__ diff --git a/python/infinicore/include/infiniop/ops/atanh.h b/python/infinicore/include/infiniop/ops/atanh.h new file mode 100644 index 000000000..0db82b164 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/atanh.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_Atanh_API_H__ +#define __INFINIOP_Atanh_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAtanhDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAtanhDescriptor(infiniopHandle_t handle, + infiniopAtanhDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t a); + +__INFINI_C __export infiniStatus_t infiniopGetAtanhWorkspaceSize(infiniopAtanhDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAtanh(infiniopAtanhDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *a, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAtanhDescriptor(infiniopAtanhDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/attention.h b/python/infinicore/include/infiniop/ops/attention.h new file mode 100644 index 000000000..21a911489 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/attention.h @@ -0,0 +1,34 @@ +#ifndef __INFINIOP_ATTENTION_API_H__ +#define __INFINIOP_ATTENTION_API_H__ + +#include "../operator_descriptor.h" +#include "gemm.h" +#include "swiglu.h" + +typedef struct InfiniopDescriptor *infiniopAttentionDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAttentionDescriptor(infiniopHandle_t handle, + infiniopAttentionDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t k_desc, + infiniopTensorDescriptor_t v_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t v_cache_desc, + size_t pos); + +__INFINI_C __export infiniStatus_t infiniopGetAttentionWorkspaceSize(infiniopAttentionDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAttention(infiniopAttentionDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *q, + const void *k, + const void *v, + void *k_cache, + void *v_cache, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAttentionDescriptor(infiniopAttentionDescriptor_t desc); +#endif diff --git a/python/infinicore/include/infiniop/ops/avg_pool1d.h b/python/infinicore/include/infiniop/ops/avg_pool1d.h new file mode 100644 index 000000000..81c489dd7 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/avg_pool1d.h @@ -0,0 +1,32 @@ +#ifndef __INFINIOP_AVG_POOL1D_API_H__ +#define __INFINIOP_AVG_POOL1D_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAvgPool1dDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAvgPool1dDescriptor( + infiniopHandle_t handle, + infiniopAvgPool1dDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + size_t kernel_size, + size_t stride, + size_t padding); + +__INFINI_C __export infiniStatus_t infiniopGetAvgPool1dWorkspaceSize( + infiniopAvgPool1dDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAvgPool1d( + infiniopAvgPool1dDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAvgPool1dDescriptor( + infiniopAvgPool1dDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/avg_pool3d.h b/python/infinicore/include/infiniop/ops/avg_pool3d.h new file mode 100644 index 000000000..a54fbb263 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/avg_pool3d.h @@ -0,0 +1,27 @@ +#ifndef __INFINIOP_AVG_POOL3D_API_H__ +#define __INFINIOP_AVG_POOL3D_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAvgPool3dDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAvgPool3dDescriptor(infiniopHandle_t handle, + infiniopAvgPool3dDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + void *kernel_size, + void *stride, + void *padding); + +__INFINI_C __export infiniStatus_t infiniopGetAvgPool3dWorkspaceSize(infiniopAvgPool3dDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAvgPool3d(infiniopAvgPool3dDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAvgPool3dDescriptor(infiniopAvgPool3dDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/awq_marlin_gemm.h b/python/infinicore/include/infiniop/ops/awq_marlin_gemm.h new file mode 100644 index 000000000..2d3fefcb5 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/awq_marlin_gemm.h @@ -0,0 +1,46 @@ +#ifndef __INFINIOP_AWQ_MARLIN_GEMM_API_H__ +#define __INFINIOP_AWQ_MARLIN_GEMM_API_H__ + +#include "../operator_descriptor.h" +#include + +typedef struct InfiniopDescriptor *infiniopAwqMarlinGemmDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAwqMarlinGemmDescriptor(infiniopHandle_t handle, + infiniopAwqMarlinGemmDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t b_bias_desc, + infiniopTensorDescriptor_t b_scales_desc, + infiniopTensorDescriptor_t a_scales_desc, + infiniopTensorDescriptor_t global_scales_desc, + infiniopTensorDescriptor_t b_zeros_desc, + infiniopTensorDescriptor_t g_idx_desc, + infiniopTensorDescriptor_t perm_desc); + +__INFINI_C __export infiniStatus_t infiniopGetAwqMarlinGemmWorkspaceSize(infiniopAwqMarlinGemmDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAwqMarlinGemm(infiniopAwqMarlinGemmDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *b_bias, + void *b_scales, + void *a_scales, + void *global_scales, + void *b_zeros, + void *g_idx, + void *perm, + int64_t b_q_type_id, + bool is_k_full, + bool use_atomic_add, + bool use_fp32_reduce, + bool is_zp_float, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAwqMarlinGemmDescriptor(infiniopAwqMarlinGemmDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/awq_marlin_repack.h b/python/infinicore/include/infiniop/ops/awq_marlin_repack.h new file mode 100644 index 000000000..017ff5568 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/awq_marlin_repack.h @@ -0,0 +1,27 @@ +#ifndef __INFINIOP_AWQ_MARLIN_REPACK_API_H__ +#define __INFINIOP_AWQ_MARLIN_REPACK_API_H__ + +#include "../operator_descriptor.h" +#include + +typedef struct InfiniopDescriptor *infiniopAwqMarlinRepackDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAwqMarlinRepackDescriptor(infiniopHandle_t handle, + infiniopAwqMarlinRepackDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + int64_t num_bits, + bool is_a_8bit); + +__INFINI_C __export infiniStatus_t infiniopGetAwqMarlinRepackWorkspaceSize(infiniopAwqMarlinRepackDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAwqMarlinRepack(infiniopAwqMarlinRepackDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAwqMarlinRepackDescriptor(infiniopAwqMarlinRepackDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/axpy.h b/python/infinicore/include/infiniop/ops/axpy.h new file mode 100644 index 000000000..ce6b2a23b --- /dev/null +++ b/python/infinicore/include/infiniop/ops/axpy.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_AXPY_API_H__ +#define __INFINIOP_AXPY_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopAxpyDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateAxpyDescriptor(infiniopHandle_t handle, + infiniopAxpyDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t alpha, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t y); + +__INFINI_C __export infiniStatus_t infiniopGetAxpyWorkspaceSize(infiniopAxpyDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopAxpy(infiniopAxpyDescriptor_t desc, + void *workspace, + size_t workspace_size, + const void *alpha, + const void *x, + void *y, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyAxpyDescriptor(infiniopAxpyDescriptor_t desc); + +#endif // __INFINIOP_AXPY_API_H__ diff --git a/python/infinicore/include/infiniop/ops/binary_cross_entropy_with_logits.h b/python/infinicore/include/infiniop/ops/binary_cross_entropy_with_logits.h new file mode 100644 index 000000000..0bec6f9d5 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/binary_cross_entropy_with_logits.h @@ -0,0 +1,73 @@ +#ifndef __INFINIOP_BINARY_CROSS_ENTROPY_WITH_LOGITS_API_H__ +#define __INFINIOP_BINARY_CROSS_ENTROPY_WITH_LOGITS_API_H__ + +#include "../operator_descriptor.h" + +// 定义归约方式枚举 +typedef enum { + INFINIOP_REDUCTION_NONE = 0, + INFINIOP_REDUCTION_MEAN = 1, + INFINIOP_REDUCTION_SUM = 2 +} infiniopReduction_t; + +// 定义 BCEWithLogits 算子描述符类型 +typedef struct InfiniopDescriptor *infiniopBCEWithLogitsDescriptor_t; + +/** + * @brief 创建 BCEWithLogits 算子描述符 + * @param handle 算子句柄 + * @param desc_ptr 指向返回的描述符指针 + * @param out 输出张量描述符 (none时与input同形状，mean/sum时为标量) + * @param logits 输入 Logits 张量描述符 + * @param target 目标标签张量描述符 + * @param weight 样本权重描述符 (可选，不需要则传 NULL) + * @param pos_weight 正样本权重描述符 (可选，不需要则传 NULL) + * @param reduction 归约方式 (none, mean, sum) + */ +__INFINI_C __export infiniStatus_t infiniopCreateBCEWithLogitsDescriptor( + infiniopHandle_t handle, + infiniopBCEWithLogitsDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out, + infiniopTensorDescriptor_t logits, + infiniopTensorDescriptor_t target, + infiniopTensorDescriptor_t weight, + infiniopTensorDescriptor_t pos_weight, + infiniopReduction_t reduction); + +/** + * @brief 获取 BCEWithLogits 计算所需的临时空间大小 + */ +__INFINI_C __export infiniStatus_t infiniopGetBCEWithLogitsWorkspaceSize( + infiniopBCEWithLogitsDescriptor_t desc, + size_t *size); + +/** + * @brief 执行 BCEWithLogits 计算 + * @param desc 算子描述符 + * @param workspace 临时空间指针 + * @param workspace_size 临时空间大小 + * @param out 输出数据指针 + * @param logits Logits 数据指针 + * @param target Target 数据指针 + * @param weight 权重数据指针 (可选，传 NULL 表示权重全为 1) + * @param pos_weight 正样本权重数据指针 (可选，传 NULL 表示权重全为 1) + * @param stream 计算流 + */ +__INFINI_C __export infiniStatus_t infiniopBCEWithLogits( + infiniopBCEWithLogitsDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *logits, + const void *target, + const void *weight, + const void *pos_weight, + void *stream); + +/** + * @brief 销毁 BCEWithLogits 算子描述符 + */ +__INFINI_C __export infiniStatus_t infiniopDestroyBCEWithLogitsDescriptor( + infiniopBCEWithLogitsDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/bitwise_right_shift.h b/python/infinicore/include/infiniop/ops/bitwise_right_shift.h new file mode 100644 index 000000000..64256fd8f --- /dev/null +++ b/python/infinicore/include/infiniop/ops/bitwise_right_shift.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_BITWISE_RIGHT_SHIFT_API_H__ +#define __INFINIOP_BITWISE_RIGHT_SHIFT_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopBitwiseRightShiftDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateBitwiseRightShiftDescriptor(infiniopHandle_t handle, + infiniopBitwiseRightShiftDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x1, + infiniopTensorDescriptor_t x2); + +__INFINI_C __export infiniStatus_t infiniopGetBitwiseRightShiftWorkspaceSize(infiniopBitwiseRightShiftDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopBitwiseRightShift(infiniopBitwiseRightShiftDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x1, + const void *x2, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyBitwiseRightShiftDescriptor(infiniopBitwiseRightShiftDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/blas_amax.h b/python/infinicore/include/infiniop/ops/blas_amax.h new file mode 100644 index 000000000..9981c7fb7 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/blas_amax.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_BLAS_AMAX_API_H__ +#define __INFINIOP_BLAS_AMAX_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopBlasAmaxDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateBlasAmaxDescriptor(infiniopHandle_t handle, + infiniopBlasAmaxDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t result); + +__INFINI_C __export infiniStatus_t infiniopGetBlasAmaxWorkspaceSize(infiniopBlasAmaxDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopBlasAmax(infiniopBlasAmaxDescriptor_t desc, + void *workspace, + size_t workspace_size, + const void *x, + void *result, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyBlasAmaxDescriptor(infiniopBlasAmaxDescriptor_t desc); + +#endif // __INFINIOP_BLAS_AMAX_API_H__ diff --git a/python/infinicore/include/infiniop/ops/blas_amin.h b/python/infinicore/include/infiniop/ops/blas_amin.h new file mode 100644 index 000000000..6bc8680ba --- /dev/null +++ b/python/infinicore/include/infiniop/ops/blas_amin.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_BLAS_AMIN_API_H__ +#define __INFINIOP_BLAS_AMIN_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopBlasAminDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateBlasAminDescriptor(infiniopHandle_t handle, + infiniopBlasAminDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t result); + +__INFINI_C __export infiniStatus_t infiniopGetBlasAminWorkspaceSize(infiniopBlasAminDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopBlasAmin(infiniopBlasAminDescriptor_t desc, + void *workspace, + size_t workspace_size, + const void *x, + void *result, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyBlasAminDescriptor(infiniopBlasAminDescriptor_t desc); + +#endif // __INFINIOP_BLAS_AMIN_API_H__ diff --git a/python/infinicore/include/infiniop/ops/blas_copy.h b/python/infinicore/include/infiniop/ops/blas_copy.h new file mode 100644 index 000000000..7c6f3611c --- /dev/null +++ b/python/infinicore/include/infiniop/ops/blas_copy.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_BLAS_COPY_API_H__ +#define __INFINIOP_BLAS_COPY_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopBlasCopyDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateBlasCopyDescriptor(infiniopHandle_t handle, + infiniopBlasCopyDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t y); + +__INFINI_C __export infiniStatus_t infiniopGetBlasCopyWorkspaceSize(infiniopBlasCopyDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopBlasCopy(infiniopBlasCopyDescriptor_t desc, + void *workspace, + size_t workspace_size, + const void *x, + void *y, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyBlasCopyDescriptor(infiniopBlasCopyDescriptor_t desc); + +#endif // __INFINIOP_BLAS_COPY_API_H__ diff --git a/python/infinicore/include/infiniop/ops/blas_dot.h b/python/infinicore/include/infiniop/ops/blas_dot.h new file mode 100644 index 000000000..9d03af2c0 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/blas_dot.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_BLAS_DOT_API_H__ +#define __INFINIOP_BLAS_DOT_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopBlasDotDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateBlasDotDescriptor(infiniopHandle_t handle, + infiniopBlasDotDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t result); + +__INFINI_C __export infiniStatus_t infiniopGetBlasDotWorkspaceSize(infiniopBlasDotDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopBlasDot(infiniopBlasDotDescriptor_t desc, + void *workspace, + size_t workspace_size, + const void *x, + const void *y, + void *result, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyBlasDotDescriptor(infiniopBlasDotDescriptor_t desc); + +#endif // __INFINIOP_BLAS_DOT_API_H__ diff --git a/python/infinicore/include/infiniop/ops/block_diag.h b/python/infinicore/include/infiniop/ops/block_diag.h new file mode 100644 index 000000000..f1829d82a --- /dev/null +++ b/python/infinicore/include/infiniop/ops/block_diag.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_BLOCK_DIAG_API_H__ +#define __INFINIOP_BLOCK_DIAG_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopBlockDiagDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateBlockDiagDescriptor(infiniopHandle_t handle, + infiniopBlockDiagDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t *x, + size_t num_inputs); + +__INFINI_C __export infiniStatus_t infiniopGetBlockDiagWorkspaceSize(infiniopBlockDiagDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopBlockDiag(infiniopBlockDiagDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void **x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyBlockDiagDescriptor(infiniopBlockDiagDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/broadcast_to.h b/python/infinicore/include/infiniop/ops/broadcast_to.h new file mode 100644 index 000000000..2830ececb --- /dev/null +++ b/python/infinicore/include/infiniop/ops/broadcast_to.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_BROADCAST_TO_API_H__ +#define __INFINIOP_BROADCAST_TO_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopBroadcastToDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateBroadcastToDescriptor(infiniopHandle_t handle, + infiniopBroadcastToDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetBroadcastToWorkspaceSize(infiniopBroadcastToDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopBroadcastTo(infiniopBroadcastToDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyBroadcastToDescriptor(infiniopBroadcastToDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/causal_conv1d.h b/python/infinicore/include/infiniop/ops/causal_conv1d.h new file mode 100644 index 000000000..431ff1858 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/causal_conv1d.h @@ -0,0 +1,43 @@ +#ifndef __INFINIOP_CAUSAL_CONV1D_API_H__ +#define __INFINIOP_CAUSAL_CONV1D_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopCausalConv1dDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateCausalConv1dDescriptor( + infiniopHandle_t handle, + infiniopCausalConv1dDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, // padded: [B, T, C]; varlen: [1, total_tokens, C] + infiniopTensorDescriptor_t conv_state_desc, // no-index: [B/num_requests, C, state_len]; pool: [pool_size, C, state_len] + infiniopTensorDescriptor_t final_conv_state_desc, // nullable when final_state_indices_desc is provided + infiniopTensorDescriptor_t qkv_desc, // padded: [B, T, C]; varlen: [1, total_tokens, C] + infiniopTensorDescriptor_t weight_desc, // [C, 1, K], depthwise; current backend supports K == 4 + infiniopTensorDescriptor_t bias_desc, // nullable; [C] + infiniopTensorDescriptor_t cu_seqlens_desc, // nullable; [num_requests + 1], int32/int64 + infiniopTensorDescriptor_t initial_state_indices_desc, // nullable; [num_requests], int32/int64 + infiniopTensorDescriptor_t final_state_indices_desc); // nullable; [num_requests], int32/int64; writes final state in-place + +__INFINI_C __export infiniStatus_t infiniopGetCausalConv1dWorkspaceSize( + infiniopCausalConv1dDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopCausalConv1d( + infiniopCausalConv1dDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + void *conv_state, + void *final_conv_state, + const void *qkv, + const void *weight, + const void *bias, + const void *cu_seqlens, + const void *initial_state_indices, + const void *final_state_indices, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyCausalConv1dDescriptor( + infiniopCausalConv1dDescriptor_t desc); + +#endif // __INFINIOP_CAUSAL_CONV1D_API_H__ diff --git a/python/infinicore/include/infiniop/ops/causal_softmax.h b/python/infinicore/include/infiniop/ops/causal_softmax.h new file mode 100644 index 000000000..13b325ffd --- /dev/null +++ b/python/infinicore/include/infiniop/ops/causal_softmax.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_CAUSAL_SOFTMAX_API_H__ +#define __INFINIOP_CAUSAL_SOFTMAX_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopCausalSoftmaxDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateCausalSoftmaxDescriptor( + infiniopHandle_t handle, + infiniopCausalSoftmaxDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc); + +__INFINI_C __export infiniStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSoftmaxDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopCausalSoftmax( + infiniopCausalSoftmaxDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyCausalSoftmaxDescriptor(infiniopCausalSoftmaxDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/cdist.h b/python/infinicore/include/infiniop/ops/cdist.h new file mode 100644 index 000000000..ab8d10877 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/cdist.h @@ -0,0 +1,56 @@ +#ifndef __INFINIOP_CDIST_API_H__ +#define __INFINIOP_CDIST_API_H__ + +#include "../operator_descriptor.h" + +// 定义 cdist 算子描述符类型 +typedef struct InfiniopDescriptor *infiniopCdistDescriptor_t; + +/** + * @brief 创建 Cdist 算子描述符 + * @param handle 算子句柄 + * @param desc_ptr 指向返回的描述符指针 + * @param y 输出张量描述符 (Shape: M x N) + * @param x1 输入张量1描述符 (Shape: M x D) + * @param x2 输入张量2描述符 (Shape: N x D) + * @param p 范数阶数 (L-p norm) + */ +__INFINI_C __export infiniStatus_t infiniopCreateCdistDescriptor( + infiniopHandle_t handle, + infiniopCdistDescriptor_t *desc_ptr, // 注意这里应该是具体类型的指针 + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x1_desc, + infiniopTensorDescriptor_t x2_desc, + double p); + +/** + * @brief 获取 Cdist 计算所需的临时空间大小 + */ +__INFINI_C __export infiniStatus_t infiniopGetCdistWorkspaceSize(infiniopCdistDescriptor_t desc, + size_t *size); + +/** + * @brief 执行 Cdist 计算 + * @param desc 算子描述符 + * @param workspace 临时空间指针 + * @param workspace_size 临时空间大小 + * @param y 输出数据指针 + * @param x1 输入1数据指针 + * @param x2 输入2数据指针 + * @param stream 计算流 (CUDA stream 等) + */ +__INFINI_C __export infiniStatus_t infiniopCdist( + infiniopCdistDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x1, + const void *x2, + void *stream); + +/** + * @brief 销毁 Cdist 算子描述符 + */ +__INFINI_C __export infiniStatus_t infiniopDestroyCdistDescriptor(infiniopCdistDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/chunk_gated_delta_rule.h b/python/infinicore/include/infiniop/ops/chunk_gated_delta_rule.h new file mode 100644 index 000000000..6caca6217 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/chunk_gated_delta_rule.h @@ -0,0 +1,49 @@ +#ifndef __INFINIOP_CHUNK_GATED_DELTA_RULE_API_H__ +#define __INFINIOP_CHUNK_GATED_DELTA_RULE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopChunkGatedDeltaRuleDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateChunkGatedDeltaRuleDescriptor( + infiniopHandle_t handle, + infiniopChunkGatedDeltaRuleDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, // padded: [B, T, Hv, Dv]; varlen: [1, total_tokens, Hv, Dv] + infiniopTensorDescriptor_t initial_state_desc, // legacy: [B, Hv, Dv, Dk]; indexed pool: [pool_size, Hv, Dv, Dk] + infiniopTensorDescriptor_t final_state_desc, // null when final_state_indices_desc is provided + infiniopTensorDescriptor_t q_desc, // padded: [B, T, Hk, Dk]; varlen: [1, total_tokens, Hk, Dk] + infiniopTensorDescriptor_t k_desc, // same shape as q + infiniopTensorDescriptor_t v_desc, // padded: [B, T, Hv, Dv]; varlen: [1, total_tokens, Hv, Dv] + infiniopTensorDescriptor_t g_desc, // padded: [B, T, Hv]; varlen: [1, total_tokens, Hv] + infiniopTensorDescriptor_t beta_desc, // same shape/dtype as g + infiniopTensorDescriptor_t cu_seqlens_desc, // nullable; [B + 1], int32/int64 + infiniopTensorDescriptor_t initial_state_indices_desc, // nullable; [B], int32/int64; enables indexed state-pool reads + infiniopTensorDescriptor_t final_state_indices_desc, // nullable; [B], int32/int64; writes final state in-place to initial_state pool + bool use_qk_l2norm, + size_t chunk_size); + +__INFINI_C __export infiniStatus_t infiniopGetChunkGatedDeltaRuleWorkspaceSize( + infiniopChunkGatedDeltaRuleDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopChunkGatedDeltaRule( + infiniopChunkGatedDeltaRuleDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + void *initial_state, + void *final_state, + const void *q, + const void *k, + const void *v, + const void *g, + const void *beta, + const void *cu_seqlens, + const void *initial_state_indices, + const void *final_state_indices, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyChunkGatedDeltaRuleDescriptor( + infiniopChunkGatedDeltaRuleDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/clip.h b/python/infinicore/include/infiniop/ops/clip.h new file mode 100644 index 000000000..965932ad8 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/clip.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_CLIP_API_H__ +#define __INFINIOP_CLIP_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopClipDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateClipDescriptor(infiniopHandle_t handle, + infiniopClipDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t min_val, + infiniopTensorDescriptor_t max_val); + +__INFINI_C __export infiniStatus_t infiniopGetClipWorkspaceSize(infiniopClipDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopClip(infiniopClipDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + const void *min_val, + const void *max_val, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyClipDescriptor(infiniopClipDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/conv.h b/python/infinicore/include/infiniop/ops/conv.h new file mode 100644 index 000000000..eea64f5e9 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/conv.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_CONV_API_H__ +#define __INFINIOP_CONV_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopConvDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateConvDescriptor(infiniopHandle_t handle, + infiniopConvDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t w_desc, + infiniopTensorDescriptor_t b_desc, + void *pads, + void *strides, + void *dilations, + size_t n); + +__INFINI_C __export infiniStatus_t infiniopGetConvWorkspaceSize(infiniopConvDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopConv(infiniopConvDescriptor_t desc, void *workspace, size_t workspace_size, void *y, const void *x, const void *w, const void *bias, void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyConvDescriptor(infiniopConvDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/cross_entropy.h b/python/infinicore/include/infiniop/ops/cross_entropy.h new file mode 100644 index 000000000..2ebd4b168 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/cross_entropy.h @@ -0,0 +1,31 @@ +#ifndef __INFINIOP_CROSS_ENTROPY_API_H__ +#define __INFINIOP_CROSS_ENTROPY_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopCrossEntropyDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateCrossEntropyDescriptor( + infiniopHandle_t handle, + infiniopCrossEntropyDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t target_desc); + +__INFINI_C __export infiniStatus_t infiniopGetCrossEntropyWorkspaceSize( + infiniopCrossEntropyDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopCrossEntropy( + infiniopCrossEntropyDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + const void *target, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyCrossEntropyDescriptor( + infiniopCrossEntropyDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/deepseek_moe.h b/python/infinicore/include/infiniop/ops/deepseek_moe.h new file mode 100644 index 000000000..ae55ae977 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/deepseek_moe.h @@ -0,0 +1,57 @@ +#ifndef __INFINIOP_DEEPSEEK_MOE_API_H__ +#define __INFINIOP_DEEPSEEK_MOE_API_H__ + +#include "../operator_descriptor.h" + +#ifdef __cplusplus +#include +#else +#include +#endif + +typedef struct InfiniopDescriptor *infiniopDeepseekMoeDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDeepseekMoeDescriptor( + infiniopHandle_t handle, + infiniopDeepseekMoeDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t hidden_desc, + infiniopTensorDescriptor_t topk_indices_desc, + infiniopTensorDescriptor_t topk_weights_desc, + size_t intermediate_size, + size_t num_experts); + +__INFINI_C __export infiniStatus_t infiniopGetDeepseekMoeWorkspaceSize( + infiniopDeepseekMoeDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopDeepseekMoe( + infiniopDeepseekMoeDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *hidden, + const void *topk_indices, + const void *topk_weights, + const void *const *gate_weights, + const void *const *up_weights, + const void *const *down_weights, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDeepseekMoeWithDevicePtrs( + infiniopDeepseekMoeDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *hidden, + const void *topk_indices, + const void *topk_weights, + const void *gate_weight_ptrs, + const void *up_weight_ptrs, + const void *down_weight_ptrs, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyDeepseekMoeDescriptor( + infiniopDeepseekMoeDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/dequant/per_tensor_dequant_int8.h b/python/infinicore/include/infiniop/ops/dequant/per_tensor_dequant_int8.h new file mode 100644 index 000000000..9614b4303 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/dequant/per_tensor_dequant_int8.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_PER_TENSOR_DEQUANT_INT8_API_H__ +#define __INFINIOP_PER_TENSOR_DEQUANT_INT8_API_H__ + +#include "../../operator_descriptor.h" + +typedef InfiniopDescriptor *infiniopPerTensorDequantI8Descriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreatePerTensorDequantI8Descriptor(infiniopHandle_t handle, + infiniopPerTensorDequantI8Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t x_packed_desc, + infiniopTensorDescriptor_t x_scale_desc, + infiniopTensorDescriptor_t x_zero_desc); + +__INFINI_C __export infiniStatus_t infiniopGetPerTensorDequantI8WorkspaceSize(infiniopPerTensorDequantI8Descriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopPerTensorDequantI8(infiniopPerTensorDequantI8Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *x, + const void *x_packed, + const void *x_scale, + const void *x_zero, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyPerTensorDequantI8Descriptor(infiniopPerTensorDequantI8Descriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/dequantize_awq.h b/python/infinicore/include/infiniop/ops/dequantize_awq.h new file mode 100644 index 000000000..f4cdc4bbe --- /dev/null +++ b/python/infinicore/include/infiniop/ops/dequantize_awq.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_DEQUANTIZE_AWQ_API_H__ +#define __INFINIOP_DEQUANTIZE_AWQ_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDequantizeAWQDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDequantizeAWQDescriptor(infiniopHandle_t handle, + infiniopDequantizeAWQDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t qweight_desc, + infiniopTensorDescriptor_t scales_desc, + infiniopTensorDescriptor_t zeros_desc); + +__INFINI_C __export infiniStatus_t infiniopGetDequantizeAWQWorkspaceSize(infiniopDequantizeAWQDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopDequantizeAWQ(infiniopDequantizeAWQDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *qweight, + const void *scales, + const void *zeros, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyDequantizeAWQDescriptor(infiniopDequantizeAWQDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/dequantize_gptq.h b/python/infinicore/include/infiniop/ops/dequantize_gptq.h new file mode 100644 index 000000000..78ff15839 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/dequantize_gptq.h @@ -0,0 +1,30 @@ +#ifndef __INFINIOP_DEQUANTIZE_GPTQ_API_H__ +#define __INFINIOP_DEQUANTIZE_GPTQ_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDequantizeGPTQDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDequantizeGPTQDescriptor(infiniopHandle_t handle, + infiniopDequantizeGPTQDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t qweight_desc, + infiniopTensorDescriptor_t scales_desc, + infiniopTensorDescriptor_t zeros_desc, + infiniopTensorDescriptor_t g_idx_desc); // add g_idx + +__INFINI_C __export infiniStatus_t infiniopGetDequantizeGPTQWorkspaceSize(infiniopDequantizeGPTQDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopDequantizeGPTQ(infiniopDequantizeGPTQDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *qweight, + const void *scales, + const void *zeros, + const void *g_idx, // add g_idx + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyDequantizeGPTQDescriptor(infiniopDequantizeGPTQDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/diff.h b/python/infinicore/include/infiniop/ops/diff.h new file mode 100644 index 000000000..f1d8dbcf9 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/diff.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_DIFF_API_H__ +#define __INFINIOP_DIFF_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDiffDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDiffDescriptor(infiniopHandle_t handle, + infiniopDiffDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + int dim, + int n); + +__INFINI_C __export infiniStatus_t infiniopGetDiffWorkspaceSize(infiniopDiffDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopDiff(infiniopDiffDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyDiffDescriptor(infiniopDiffDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/digamma.h b/python/infinicore/include/infiniop/ops/digamma.h new file mode 100644 index 000000000..2c53dc9a1 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/digamma.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_DIGAMMA_API_H__ +#define __INFINIOP_DIGAMMA_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDigammaDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDigammaDescriptor(infiniopHandle_t handle, + infiniopDigammaDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetDigammaWorkspaceSize(infiniopDigammaDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopDigamma(infiniopDigammaDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyDigammaDescriptor(infiniopDigammaDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/dist.h b/python/infinicore/include/infiniop/ops/dist.h new file mode 100644 index 000000000..45e044aa8 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/dist.h @@ -0,0 +1,27 @@ +#ifndef __INFINIOP_DIST_API_H__ +#define __INFINIOP_DIST_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDistDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDistDescriptor(infiniopHandle_t handle, + infiniopDistDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x1, + infiniopTensorDescriptor_t x2, + double p); + +__INFINI_C __export infiniStatus_t infiniopGetDistWorkspaceSize(infiniopDistDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopDist(infiniopDistDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x1, + const void *x2, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyDistDescriptor(infiniopDistDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/dot.h b/python/infinicore/include/infiniop/ops/dot.h new file mode 100644 index 000000000..1fb9197f9 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/dot.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_DOT_API_H__ +#define __INFINIOP_DOT_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDotDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDotDescriptor(infiniopHandle_t handle, + infiniopDotDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__INFINI_C __export infiniStatus_t infiniopGetDotWorkspaceSize(infiniopDotDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopDot(infiniopDotDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *a, + const void *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyDotDescriptor(infiniopDotDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/embedding.h b/python/infinicore/include/infiniop/ops/embedding.h new file mode 100644 index 000000000..5528be131 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/embedding.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_EMBEDDING_API_H__ +#define __INFINIOP_EMBEDDING_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopEmbeddingDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateEmbeddingDescriptor( + infiniopHandle_t handle, + infiniopEmbeddingDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t weight_desc); + +__INFINI_C __export infiniStatus_t infiniopEmbedding( + infiniopEmbeddingDescriptor_t desc, + void *output, + const void *input, + const void *weight, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyEmbeddingDescriptor( + infiniopEmbeddingDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/equal.h b/python/infinicore/include/infiniop/ops/equal.h new file mode 100644 index 000000000..90c4f3386 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/equal.h @@ -0,0 +1,31 @@ +#ifndef __INFINIOP_EQUAL_API_H__ +#define __INFINIOP_EQUAL_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopEqualDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateEqualDescriptor( + infiniopHandle_t handle, + infiniopEqualDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__INFINI_C __export infiniStatus_t infiniopGetEqualWorkspaceSize( + infiniopEqualDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopEqual( + infiniopEqualDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyEqualDescriptor( + infiniopEqualDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/erf.h b/python/infinicore/include/infiniop/ops/erf.h new file mode 100644 index 000000000..8786b31a3 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/erf.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_ERF_API_H__ +#define __INFINIOP_ERF_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopErfDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateErfDescriptor(infiniopHandle_t handle, + infiniopErfDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetErfWorkspaceSize(infiniopErfDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopErf(infiniopErfDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyErfDescriptor(infiniopErfDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/erfc.h b/python/infinicore/include/infiniop/ops/erfc.h new file mode 100644 index 000000000..6bceebe35 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/erfc.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_ERFC_API_H__ +#define __INFINIOP_ERFC_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopErfcDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateErfcDescriptor(infiniopHandle_t handle, + infiniopErfcDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetErfcWorkspaceSize(infiniopErfcDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopErfc(infiniopErfcDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyErfcDescriptor(infiniopErfcDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/erfinv.h b/python/infinicore/include/infiniop/ops/erfinv.h new file mode 100644 index 000000000..b14975253 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/erfinv.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_ERFINV_API_H__ +#define __INFINIOP_ERFINV_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopErfinvDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateErfinvDescriptor(infiniopHandle_t handle, + infiniopErfinvDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetErfinvWorkspaceSize(infiniopErfinvDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopErfinv(infiniopErfinvDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyErfinvDescriptor(infiniopErfinvDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/flash_attention.h b/python/infinicore/include/infiniop/ops/flash_attention.h new file mode 100644 index 000000000..06f1e14e1 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/flash_attention.h @@ -0,0 +1,36 @@ +#ifndef __INFINIOP_FLASH_ATTENTION_API_H__ +#define __INFINIOP_FLASH_ATTENTION_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFlashAttentionDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFlashAttentionDescriptor( + infiniopHandle_t handle, + infiniopFlashAttentionDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t k_desc, + infiniopTensorDescriptor_t v_desc, + infiniopTensorDescriptor_t total_kv_len, + float scale, + char is_causal); + +__INFINI_C __export infiniStatus_t infiniopGetFlashAttentionWorkspaceSize( + infiniopFlashAttentionDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopFlashAttention( + infiniopFlashAttentionDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *q, + const void *k, + const void *v, + const void *total_kv_len, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFlashAttentionDescriptor( + infiniopFlashAttentionDescriptor_t desc); +#endif diff --git a/python/infinicore/include/infiniop/ops/flipud.h b/python/infinicore/include/infiniop/ops/flipud.h new file mode 100644 index 000000000..6f5a4768d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/flipud.h @@ -0,0 +1,27 @@ +#ifndef __INFINIOP_FLIPUD_API_H__ +#define __INFINIOP_FLIPUD_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFlipudDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFlipudDescriptor(infiniopHandle_t handle, + infiniopFlipudDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input); + +// 获取工作空间大小 +__INFINI_C __export infiniStatus_t infiniopGetFlipudWorkspaceSize(infiniopFlipudDescriptor_t desc, size_t *size); + +// 执行 Flipud 算子 +__INFINI_C __export infiniStatus_t infiniopFlipud(infiniopFlipudDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +// 销毁描述符 +__INFINI_C __export infiniStatus_t infiniopDestroyFlipudDescriptor(infiniopFlipudDescriptor_t desc); + +#endif // __INFINIOP_FLIPUD_API_H__ diff --git a/python/infinicore/include/infiniop/ops/float_power.h b/python/infinicore/include/infiniop/ops/float_power.h new file mode 100644 index 000000000..8a4c9a97e --- /dev/null +++ b/python/infinicore/include/infiniop/ops/float_power.h @@ -0,0 +1,27 @@ +#ifndef __INFINIOP_FLOAT_POWER_API_H__ +#define __INFINIOP_FLOAT_POWER_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFloatPowerDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFloatPowerDescriptor(infiniopHandle_t handle, + infiniopFloatPowerDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t exponent, + float scalar_exponent); + +__INFINI_C __export infiniStatus_t infiniopGetFloatPowerWorkspaceSize(infiniopFloatPowerDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopFloatPower(infiniopFloatPowerDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + const void *exponent, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFloatPowerDescriptor(infiniopFloatPowerDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/floor.h b/python/infinicore/include/infiniop/ops/floor.h new file mode 100644 index 000000000..037237ed7 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/floor.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_FLOOR_API_H__ +#define __INFINIOP_FLOOR_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFloorDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFloorDescriptor(infiniopHandle_t handle, + infiniopFloorDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t intput); + +__INFINI_C __export infiniStatus_t infiniopGetFloorWorkspaceSize(infiniopFloorDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopFloor(infiniopFloorDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *intput, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFloorDescriptor(infiniopFloorDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/floor_divide.h b/python/infinicore/include/infiniop/ops/floor_divide.h new file mode 100644 index 000000000..300290ede --- /dev/null +++ b/python/infinicore/include/infiniop/ops/floor_divide.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_FLOOR_DIVIDE_API_H__ +#define __INFINIOP_FLOOR_DIVIDE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFloorDivideDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFloorDivideDescriptor(infiniopHandle_t handle, + infiniopFloorDivideDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__INFINI_C __export infiniStatus_t infiniopGetFloorDivideWorkspaceSize(infiniopFloorDivideDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopFloorDivide(infiniopFloorDivideDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFloorDivideDescriptor(infiniopFloorDivideDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/fmin.h b/python/infinicore/include/infiniop/ops/fmin.h new file mode 100644 index 000000000..c1f9f75bf --- /dev/null +++ b/python/infinicore/include/infiniop/ops/fmin.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_FMIN_H__ +#define __INFINIOP_FMIN_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFminDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFminDescriptor(infiniopHandle_t handle, + infiniopFminDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__INFINI_C __export infiniStatus_t infiniopGetFminWorkspaceSize(infiniopFminDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopFmin(infiniopFminDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFminDescriptor(infiniopFminDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/fmod.h b/python/infinicore/include/infiniop/ops/fmod.h new file mode 100644 index 000000000..e51f1b3c8 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/fmod.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_FMOD_API_H__ +#define __INFINIOP_FMOD_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFmodDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFmodDescriptor(infiniopHandle_t handle, + infiniopFmodDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__INFINI_C __export infiniStatus_t infiniopGetFmodWorkspaceSize(infiniopFmodDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopFmod(infiniopFmodDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFmodDescriptor(infiniopFmodDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/fp8_indexer_logits.h b/python/infinicore/include/infiniop/ops/fp8_indexer_logits.h new file mode 100644 index 000000000..3808b0292 --- /dev/null +++ b/python/infinicore/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/python/infinicore/include/infiniop/ops/fp8_indexer_quant.h b/python/infinicore/include/infiniop/ops/fp8_indexer_quant.h new file mode 100644 index 000000000..101a325a9 --- /dev/null +++ b/python/infinicore/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/python/infinicore/include/infiniop/ops/fp8_mla_rmsnorm_cache.h b/python/infinicore/include/infiniop/ops/fp8_mla_rmsnorm_cache.h new file mode 100644 index 000000000..ccbbd1f88 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/fp8_mla_rmsnorm_cache.h @@ -0,0 +1,30 @@ +#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 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, + 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/python/infinicore/include/infiniop/ops/fp8_sparse_mla.h b/python/infinicore/include/infiniop/ops/fp8_sparse_mla.h new file mode 100644 index 000000000..3d199a614 --- /dev/null +++ b/python/infinicore/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/python/infinicore/include/infiniop/ops/fused_ffn.h b/python/infinicore/include/infiniop/ops/fused_ffn.h new file mode 100644 index 000000000..6becd77d5 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/fused_ffn.h @@ -0,0 +1,37 @@ +#ifndef __INFINIOP_FUSED_FFN_API_H__ +#define __INFINIOP_FUSED_FFN_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFusedFFNDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateFusedFFNDescriptor( + infiniopHandle_t handle, + infiniopFusedFFNDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t in_desc, + infiniopTensorDescriptor_t residual_desc, + infiniopTensorDescriptor_t norm_weight_desc, + infiniopTensorDescriptor_t gate_up_weight_desc, + infiniopTensorDescriptor_t down_weight_desc, + float epsilon); + +__INFINI_C __export infiniStatus_t infiniopGetFusedFFNWorkspaceSize( + infiniopFusedFFNDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopFusedFFN( + infiniopFusedFFNDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *in, + const void *residual, + const void *norm_weight, + const void *gate_up_weight, + const void *down_weight, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyFusedFFNDescriptor( + infiniopFusedFFNDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/fused_gated_delta_net_gating.h b/python/infinicore/include/infiniop/ops/fused_gated_delta_net_gating.h new file mode 100644 index 000000000..dafea6710 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/fused_gated_delta_net_gating.h @@ -0,0 +1,43 @@ +#ifndef __INFINIOP_FUSED_GATED_DELTA_NET_GATING_API_H__ +#define __INFINIOP_FUSED_GATED_DELTA_NET_GATING_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopFusedGatedDeltaNetGatingDescriptor_t; + +__INFINI_C __export infiniStatus_t +infiniopCreateFusedGatedDeltaNetGatingDescriptor( + infiniopHandle_t handle, + infiniopFusedGatedDeltaNetGatingDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t g_desc, + infiniopTensorDescriptor_t beta_output_desc, + infiniopTensorDescriptor_t A_log_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t dt_bias_desc, + float beta, + float threshold); + +__INFINI_C __export infiniStatus_t +infiniopGetFusedGatedDeltaNetGatingWorkspaceSize( + infiniopFusedGatedDeltaNetGatingDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t +infiniopFusedGatedDeltaNetGating( + infiniopFusedGatedDeltaNetGatingDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *g, + void *beta_output, + const void *A_log, + const void *a, + const void *b, + const void *dt_bias, + void *stream); + +__INFINI_C __export infiniStatus_t +infiniopDestroyFusedGatedDeltaNetGatingDescriptor( + infiniopFusedGatedDeltaNetGatingDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/gaussian_nll_loss.h b/python/infinicore/include/infiniop/ops/gaussian_nll_loss.h new file mode 100644 index 000000000..cc09edee8 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/gaussian_nll_loss.h @@ -0,0 +1,31 @@ +#ifndef __INFINIOP_GAUSSIAN_NLL_LOSS_API_H__ +#define __INFINIOP_GAUSSIAN_NLL_LOSS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopGaussianNllLossDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateGaussianNllLossDescriptor(infiniopHandle_t handle, + infiniopGaussianNllLossDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t input, + infiniopTensorDescriptor_t target, + infiniopTensorDescriptor_t var, + int full, + double eps, + int reduction); + +__INFINI_C __export infiniStatus_t infiniopGetGaussianNllLossWorkspaceSize(infiniopGaussianNllLossDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopGaussianNllLoss(infiniopGaussianNllLossDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *input, + const void *target, + const void *var, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyGaussianNllLossDescriptor(infiniopGaussianNllLossDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/gelu.h b/python/infinicore/include/infiniop/ops/gelu.h new file mode 100644 index 000000000..3ddc21df8 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/gelu.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_GELU_API_H__ +#define __INFINIOP_GELU_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopGeluDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateGeluDescriptor(infiniopHandle_t handle, + infiniopGeluDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t intput); + +__INFINI_C __export infiniStatus_t infiniopGetGeluWorkspaceSize(infiniopGeluDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopGelu(infiniopGeluDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *intput, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyGeluDescriptor(infiniopGeluDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/gelutanh.h b/python/infinicore/include/infiniop/ops/gelutanh.h new file mode 100644 index 000000000..9e94514b8 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/gelutanh.h @@ -0,0 +1,43 @@ +#ifndef __INFINIOP_GELUTANH_API_H__ +#define __INFINIOP_GELUTANH_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopGeluTanhDescriptor_t; + +/** + * Create GELU-Tanh descriptor + * + * y = x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) + */ +__INFINI_C __export infiniStatus_t infiniopCreateGeluTanhDescriptor( + infiniopHandle_t handle, + infiniopGeluTanhDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +/** + * Query workspace size + */ +__INFINI_C __export infiniStatus_t infiniopGetGeluTanhWorkspaceSize( + infiniopGeluTanhDescriptor_t desc, + size_t *size); + +/** + * Launch GELU-Tanh operator + */ +__INFINI_C __export infiniStatus_t infiniopGeluTanh( + infiniopGeluTanhDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +/** + * Destroy descriptor + */ +__INFINI_C __export infiniStatus_t infiniopDestroyGeluTanhDescriptor( + infiniopGeluTanhDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/gemm.h b/python/infinicore/include/infiniop/ops/gemm.h new file mode 100644 index 000000000..430e37003 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/gemm.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_GEMM_API_H__ +#define __INFINIOP_GEMM_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopGemmDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateGemmDescriptor(infiniopHandle_t handle, + infiniopGemmDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc); + +__INFINI_C __export infiniStatus_t infiniopGetGemmWorkspaceSize(infiniopGemmDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopGemm(infiniopGemmDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + void const *a, + void const *b, + float alpha, + float beta, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyGemmDescriptor(infiniopGemmDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/gptq_marlin_gemm.h b/python/infinicore/include/infiniop/ops/gptq_marlin_gemm.h new file mode 100644 index 000000000..66b2e1bb5 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/gptq_marlin_gemm.h @@ -0,0 +1,42 @@ +#ifndef __INFINIOP_GPTQ_MARLIN_GEMM_API_H__ +#define __INFINIOP_GPTQ_MARLIN_GEMM_API_H__ + +#include "../operator_descriptor.h" +#include + +typedef struct InfiniopDescriptor *infiniopGptqMarlinGemmDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateGptqMarlinGemmDescriptor(infiniopHandle_t handle, + infiniopGptqMarlinGemmDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t b_scales_desc, + infiniopTensorDescriptor_t global_scales_desc, + infiniopTensorDescriptor_t b_zeros_desc, + infiniopTensorDescriptor_t g_idx_desc, + infiniopTensorDescriptor_t perm_desc); + +__INFINI_C __export infiniStatus_t infiniopGetGptqMarlinGemmWorkspaceSize(infiniopGptqMarlinGemmDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopGptqMarlinGemm(infiniopGptqMarlinGemmDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *a, + const void *b, + void *b_scales, + void *global_scales, + void *b_zeros, + void *g_idx, + void *perm, + int64_t b_q_type_id, + bool is_k_full, + bool use_atomic_add, + bool use_fp32_reduce, + bool is_zp_float, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyGptqMarlinGemmDescriptor(infiniopGptqMarlinGemmDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/gptq_marlin_repack.h b/python/infinicore/include/infiniop/ops/gptq_marlin_repack.h new file mode 100644 index 000000000..c3b588fa5 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/gptq_marlin_repack.h @@ -0,0 +1,29 @@ +#ifndef __INFINIOP_GPTQ_MARLIN_REPACK_API_H__ +#define __INFINIOP_GPTQ_MARLIN_REPACK_API_H__ + +#include "../operator_descriptor.h" +#include + +typedef struct InfiniopDescriptor *infiniopGptqMarlinRepackDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateGptqMarlinRepackDescriptor(infiniopHandle_t handle, + infiniopGptqMarlinRepackDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t perm_desc, + int64_t num_bits, + bool is_a_8bit); + +__INFINI_C __export infiniStatus_t infiniopGetGptqMarlinRepackWorkspaceSize(infiniopGptqMarlinRepackDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopGptqMarlinRepack(infiniopGptqMarlinRepackDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *perm, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyGptqMarlinRepackDescriptor(infiniopGptqMarlinRepackDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/gptq_qyblas_gemm.h b/python/infinicore/include/infiniop/ops/gptq_qyblas_gemm.h new file mode 100644 index 000000000..bb105132c --- /dev/null +++ b/python/infinicore/include/infiniop/ops/gptq_qyblas_gemm.h @@ -0,0 +1,37 @@ +#ifndef __INFINIOP_GPTQ_QYBLAS_GEMM_API_H__ +#define __INFINIOP_GPTQ_QYBLAS_GEMM_API_H__ + +#include "../operator_descriptor.h" +#include + +typedef struct InfiniopDescriptor *infiniopGptqQyblasGemmDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateGptqQyblasGemmDescriptor( + infiniopHandle_t handle, + infiniopGptqQyblasGemmDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t b_scales_desc, + infiniopTensorDescriptor_t b_zeros_desc); + +__INFINI_C __export infiniStatus_t infiniopGetGptqQyblasGemmWorkspaceSize( + infiniopGptqQyblasGemmDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopGptqQyblasGemm( + infiniopGptqQyblasGemmDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *a, + const void *b, + void *b_scale, + void *b_zero, + int64_t quant_type, + int64_t bit, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyGptqQyblasGemmDescriptor( + infiniopGptqQyblasGemmDescriptor_t desc); +#endif diff --git a/python/infinicore/include/infiniop/ops/hardswish.h b/python/infinicore/include/infiniop/ops/hardswish.h new file mode 100644 index 000000000..1cdeecf67 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/hardswish.h @@ -0,0 +1,29 @@ +#ifndef __INFINIOP_HARDSWISH_API_H__ +#define __INFINIOP_HARDSWISH_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopHardSwishDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateHardSwishDescriptor( + infiniopHandle_t handle, + infiniopHardSwishDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input); + +__INFINI_C __export infiniStatus_t infiniopGetHardSwishWorkspaceSize( + infiniopHardSwishDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopHardSwish( + infiniopHardSwishDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyHardSwishDescriptor( + infiniopHardSwishDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/hardtanh.h b/python/infinicore/include/infiniop/ops/hardtanh.h new file mode 100644 index 000000000..d2f98cedd --- /dev/null +++ b/python/infinicore/include/infiniop/ops/hardtanh.h @@ -0,0 +1,27 @@ +#ifndef __INFINIOP_HARDTANH_API_H__ +#define __INFINIOP_HARDTANH_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopHardTanhDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateHardTanhDescriptor(infiniopHandle_t handle, + infiniopHardTanhDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + float min_val, + float max_val); + +__INFINI_C __export infiniStatus_t infiniopGetHardTanhWorkspaceSize(infiniopHardTanhDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopHardTanh(infiniopHardTanhDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyHardTanhDescriptor(infiniopHardTanhDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/hinge_embedding_loss.h b/python/infinicore/include/infiniop/ops/hinge_embedding_loss.h new file mode 100644 index 000000000..6058f96cd --- /dev/null +++ b/python/infinicore/include/infiniop/ops/hinge_embedding_loss.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_HINGE_EMBEDDING_LOSS_API_H__ +#define __INFINIOP_HINGE_EMBEDDING_LOSS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopHingeEmbeddingLossDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateHingeEmbeddingLossDescriptor(infiniopHandle_t handle, + infiniopHingeEmbeddingLossDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t input, + infiniopTensorDescriptor_t target, + double margin, + int reduction); + +__INFINI_C __export infiniStatus_t infiniopGetHingeEmbeddingLossWorkspaceSize(infiniopHingeEmbeddingLossDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopHingeEmbeddingLoss(infiniopHingeEmbeddingLossDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *input, + const void *target, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyHingeEmbeddingLossDescriptor(infiniopHingeEmbeddingLossDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/histc.h b/python/infinicore/include/infiniop/ops/histc.h new file mode 100644 index 000000000..3e7232259 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/histc.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_HISTC_API_H__ +#define __INFINIOP_HISTC_API_H__ + +#include "../operator_descriptor.h" +#include + +typedef struct InfiniopDescriptor *infiniopHistcDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateHistcDescriptor(infiniopHandle_t handle, + infiniopHistcDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + int64_t bins, + double min_val, + double max_val); + +__INFINI_C __export infiniStatus_t infiniopGetHistcWorkspaceSize(infiniopHistcDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopHistc(infiniopHistcDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyHistcDescriptor(infiniopHistcDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/huber_loss.h b/python/infinicore/include/infiniop/ops/huber_loss.h new file mode 100644 index 000000000..4bf0d7716 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/huber_loss.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_HUBER_LOSS_API_H__ +#define __INFINIOP_HUBER_LOSS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopHuberLossDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateHuberLossDescriptor(infiniopHandle_t handle, + infiniopHuberLossDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + infiniopTensorDescriptor_t target, + float delta, + int reduction); + +__INFINI_C __export infiniStatus_t infiniopGetHuberLossWorkspaceSize(infiniopHuberLossDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopHuberLoss(infiniopHuberLossDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *target, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyHuberLossDescriptor(infiniopHuberLossDescriptor_t desc); + +#endif // __INFINIOP_HUBER_LOSS_API_H__ diff --git a/python/infinicore/include/infiniop/ops/hypot.h b/python/infinicore/include/infiniop/ops/hypot.h new file mode 100644 index 000000000..791c18463 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/hypot.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_HYPOT_API_H__ +#define __INFINIOP_HYPOT_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopHypotDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateHypotDescriptor(infiniopHandle_t handle, + infiniopHypotDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input_a, + infiniopTensorDescriptor_t input_b); + +__INFINI_C __export infiniStatus_t infiniopGetHypotWorkspaceSize(infiniopHypotDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopHypot(infiniopHypotDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input_a, + const void *input_b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyHypotDescriptor(infiniopHypotDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/index_add.h b/python/infinicore/include/infiniop/ops/index_add.h new file mode 100644 index 000000000..f8b23e268 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/index_add.h @@ -0,0 +1,29 @@ +#ifndef __INFINIOP_INDEX_ADD_API_H__ +#define __INFINIOP_INDEX_ADD_API_H__ +#include "../operator_descriptor.h" +#include + +typedef struct InfiniopDescriptor *infiniopIndexAddDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateIndexAddDescriptor(infiniopHandle_t handle, + infiniopIndexAddDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + int64_t dim, + infiniopTensorDescriptor_t index, + infiniopTensorDescriptor_t source, + float alpha); + +__INFINI_C __export infiniStatus_t infiniopGetIndexAddWorkspaceSize(infiniopIndexAddDescriptor_t desc, size_t *size); +__INFINI_C __export infiniStatus_t infiniopIndexAdd(infiniopIndexAddDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *index, + const void *source, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyIndexAddDescriptor(infiniopIndexAddDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/index_copy.h b/python/infinicore/include/infiniop/ops/index_copy.h new file mode 100644 index 000000000..e44032fcd --- /dev/null +++ b/python/infinicore/include/infiniop/ops/index_copy.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_INDEX_COPY_API_H__ +#define __INFINIOP_INDEX_COPY_API_H__ +#include "../operator_descriptor.h" +#include + +typedef struct InfiniopDescriptor *infiniopIndexCopyDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateIndexCopyDescriptor(infiniopHandle_t handle, + infiniopIndexCopyDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + int64_t dim, + infiniopTensorDescriptor_t index, + infiniopTensorDescriptor_t source); + +__INFINI_C __export infiniStatus_t infiniopGetIndexCopyWorkspaceSize(infiniopIndexCopyDescriptor_t desc, size_t *size); +__INFINI_C __export infiniStatus_t infiniopIndexCopy(infiniopIndexCopyDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *index, + const void *source, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyIndexCopyDescriptor(infiniopIndexCopyDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/inner.h b/python/infinicore/include/infiniop/ops/inner.h new file mode 100644 index 000000000..99fc8ab13 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/inner.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_INNER_API_H__ +#define __INFINIOP_INNER_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopInnerDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateInnerDescriptor( + infiniopHandle_t handle, + infiniopInnerDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t other_desc); + +__INFINI_C __export infiniStatus_t infiniopGetInnerWorkspaceSize(infiniopInnerDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopInner( + infiniopInnerDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *input, + const void *other, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyInnerDescriptor(infiniopInnerDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/int8_gemm.h b/python/infinicore/include/infiniop/ops/int8_gemm.h new file mode 100644 index 000000000..bc96c50a3 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/int8_gemm.h @@ -0,0 +1,32 @@ +#ifndef __INFINIOP_I8GEMM_API_H__ +#define __INFINIOP_I8GEMM_API_H__ + +#include "../operator_descriptor.h" + +typedef InfiniopDescriptor *infiniopI8GemmDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateI8GemmDescriptor(infiniopHandle_t handle, + infiniopI8GemmDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t bias_desc, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t x_scale_desc, + infiniopTensorDescriptor_t weights_desc, + infiniopTensorDescriptor_t weights_scale_desc); + +__INFINI_C __export infiniStatus_t infiniopGetI8GemmWorkspaceSize(infiniopI8GemmDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopI8Gemm(infiniopI8GemmDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *bias, + const void *x, + const void *x_scale, + const void *weights, + const void *weights_scale, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyI8GemmDescriptor(infiniopI8GemmDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/interpolate.h b/python/infinicore/include/infiniop/ops/interpolate.h new file mode 100644 index 000000000..d52615267 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/interpolate.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_INTERPOLATE_API_H__ +#define __INFINIOP_INTERPOLATE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopInterpolateDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateInterpolateDescriptor(infiniopHandle_t handle, + infiniopInterpolateDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + const char *mode, + void *size, + void *scale_factor, + int align_corners); + +__INFINI_C __export infiniStatus_t infiniopGetInterpolateWorkspaceSize(infiniopInterpolateDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopInterpolate(infiniopInterpolateDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyInterpolateDescriptor(infiniopInterpolateDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/kron.h b/python/infinicore/include/infiniop/ops/kron.h new file mode 100644 index 000000000..44608b8e2 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/kron.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_KRON_API_H__ +#define __INFINIOP_KRON_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopKronDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateKronDescriptor(infiniopHandle_t handle, + infiniopKronDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x1, + infiniopTensorDescriptor_t x2); + +__INFINI_C __export infiniStatus_t infiniopGetKronWorkspaceSize(infiniopKronDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopKron(infiniopKronDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x1, + const void *x2, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyKronDescriptor(infiniopKronDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/kthvalue.h b/python/infinicore/include/infiniop/ops/kthvalue.h new file mode 100644 index 000000000..1f16df01d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/kthvalue.h @@ -0,0 +1,29 @@ +#ifndef __INFINIOP_KTHVALUE_API_H__ +#define __INFINIOP_KTHVALUE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopKthvalueDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateKthvalueDescriptor(infiniopHandle_t handle, + infiniopKthvalueDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t values, + infiniopTensorDescriptor_t indices, + infiniopTensorDescriptor_t input, + int k, + int dim, + int keepdim); + +__INFINI_C __export infiniStatus_t infiniopGetKthvalueWorkspaceSize(infiniopKthvalueDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopKthvalue(infiniopKthvalueDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *values, + void *indices, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyKthvalueDescriptor(infiniopKthvalueDescriptor_t desc); + +#endif // __INFINIOP_KTHVALUE_API_H__ diff --git a/python/infinicore/include/infiniop/ops/kv_caching.h b/python/infinicore/include/infiniop/ops/kv_caching.h new file mode 100644 index 000000000..1e091cf37 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/kv_caching.h @@ -0,0 +1,31 @@ +#ifndef __INFINIOP_KV_CACHING_API_H__ +#define __INFINIOP_KV_CACHING_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopKVCachingDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateKVCachingDescriptor( + infiniopHandle_t handle, + infiniopKVCachingDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t k_cache, + infiniopTensorDescriptor_t v_cache, + infiniopTensorDescriptor_t k, + infiniopTensorDescriptor_t v, + infiniopTensorDescriptor_t past_kv_lengths); + +__INFINI_C __export infiniStatus_t infiniopGetKVCachingWorkspaceSize(infiniopKVCachingDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopKVCaching(infiniopKVCachingDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *past_kv_lengths, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyKVCachingDescriptor(infiniopKVCachingDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/layer_norm.h b/python/infinicore/include/infiniop/ops/layer_norm.h new file mode 100644 index 000000000..140d18f02 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/layer_norm.h @@ -0,0 +1,34 @@ +#ifndef __INFINIOP_LAYER_NORM_API_H__ +#define __INFINIOP_LAYER_NORM_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLayerNormDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLayerNormDescriptor( + infiniopHandle_t handle, + infiniopLayerNormDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_standardization_desc, + infiniopTensorDescriptor_t input_std_deviation_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t bias_desc, + float eps); + +__INFINI_C __export infiniStatus_t infiniopGetLayerNormWorkspaceSize(infiniopLayerNormDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLayerNorm(infiniopLayerNormDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + void *input_standardization, + void *input_std_deviation, + const void *input, + const void *weight, + const void *bias, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLayerNormDescriptor(infiniopLayerNormDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/ldexp.h b/python/infinicore/include/infiniop/ops/ldexp.h new file mode 100644 index 000000000..f01891212 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/ldexp.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_LDEXP_API_H__ +#define __INFINIOP_LDEXP_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLdexpDescriptor_t; +__INFINI_C __export infiniStatus_t infiniopCreateLdexpDescriptor(infiniopHandle_t handle, + infiniopLdexpDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t exp); +__INFINI_C __export infiniStatus_t infiniopGetLdexpWorkspaceSize(infiniopLdexpDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLdexp(infiniopLdexpDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + const void *exp, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLdexpDescriptor(infiniopLdexpDescriptor_t desc); + +#endif // __INFINIOP_LDEXP_API_H__ diff --git a/python/infinicore/include/infiniop/ops/lerp.h b/python/infinicore/include/infiniop/ops/lerp.h new file mode 100644 index 000000000..4b392b51c --- /dev/null +++ b/python/infinicore/include/infiniop/ops/lerp.h @@ -0,0 +1,29 @@ +#ifndef __INFINIOP_LERP_API_H__ +#define __INFINIOP_LERP_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLerpDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLerpDescriptor(infiniopHandle_t handle, + infiniopLerpDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t start, + infiniopTensorDescriptor_t end, + infiniopTensorDescriptor_t weight, + float weight_scalar); + +__INFINI_C __export infiniStatus_t infiniopGetLerpWorkspaceSize(infiniopLerpDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLerp(infiniopLerpDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *start, + const void *end, + const void *weight, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLerpDescriptor(infiniopLerpDescriptor_t desc); + +#endif // __INFINIOP_LERP_API_H__ diff --git a/python/infinicore/include/infiniop/ops/log10.h b/python/infinicore/include/infiniop/ops/log10.h new file mode 100644 index 000000000..5ca8ccd8d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/log10.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_LOG10_API_H__ +#define __INFINIOP_LOG10_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLog10Descriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLog10Descriptor(infiniopHandle_t handle, + infiniopLog10Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetLog10WorkspaceSize(infiniopLog10Descriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLog10(infiniopLog10Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLog10Descriptor(infiniopLog10Descriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/log1p.h b/python/infinicore/include/infiniop/ops/log1p.h new file mode 100644 index 000000000..93f229c7d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/log1p.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_LOG1P_API_H__ +#define __INFINIOP_LOG1P_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLog1pDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLog1pDescriptor(infiniopHandle_t handle, + infiniopLog1pDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetLog1pWorkspaceSize(infiniopLog1pDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLog1p(infiniopLog1pDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLog1pDescriptor(infiniopLog1pDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/log_softmax.h b/python/infinicore/include/infiniop/ops/log_softmax.h new file mode 100644 index 000000000..8e2bef0f3 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/log_softmax.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_LOG_SOFTMAX_API_H__ +#define __INFINIOP_LOG_SOFTMAX_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLogSoftmaxDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLogSoftmaxDescriptor(infiniopHandle_t handle, + infiniopLogSoftmaxDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + int dim); + +__INFINI_C __export infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLogSoftmax(infiniopLogSoftmaxDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLogSoftmaxDescriptor(infiniopLogSoftmaxDescriptor_t desc); + +#endif // __INFINIOP_LOG_SOFTMAX_API_H__ diff --git a/python/infinicore/include/infiniop/ops/logaddexp.h b/python/infinicore/include/infiniop/ops/logaddexp.h new file mode 100644 index 000000000..df18a130f --- /dev/null +++ b/python/infinicore/include/infiniop/ops/logaddexp.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_LOGADDEXP_API_H__ +#define __INFINIOP_LOGADDEXP_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLogAddExpDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLogAddExpDescriptor(infiniopHandle_t handle, + infiniopLogAddExpDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__INFINI_C __export infiniStatus_t infiniopGetLogAddExpWorkspaceSize(infiniopLogAddExpDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLogAddExp(infiniopLogAddExpDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLogAddExpDescriptor(infiniopLogAddExpDescriptor_t desc); + +#endif // __INFINIOP_LOGADDEXP_API_H__ diff --git a/python/infinicore/include/infiniop/ops/logaddexp2.h b/python/infinicore/include/infiniop/ops/logaddexp2.h new file mode 100644 index 000000000..8572c235e --- /dev/null +++ b/python/infinicore/include/infiniop/ops/logaddexp2.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_LOGADDEXP2_API_H__ +#define __INFINIOP_LOGADDEXP2_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLogAddExp2Descriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLogAddExp2Descriptor(infiniopHandle_t handle, + infiniopLogAddExp2Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__INFINI_C __export infiniStatus_t infiniopGetLogAddExp2WorkspaceSize(infiniopLogAddExp2Descriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLogAddExp2(infiniopLogAddExp2Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLogAddExp2Descriptor(infiniopLogAddExp2Descriptor_t desc); + +#endif // __INFINIOP_LOGADDEXP2_API_H__ diff --git a/python/infinicore/include/infiniop/ops/logcumsumexp.h b/python/infinicore/include/infiniop/ops/logcumsumexp.h new file mode 100644 index 000000000..7e00727e3 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/logcumsumexp.h @@ -0,0 +1,30 @@ +#ifndef __INFINIOP_LOGCUMSUMEXP_API_H__ +#define __INFINIOP_LOGCUMSUMEXP_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLogCumSumExpDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLogCumSumExpDescriptor(infiniopHandle_t handle, + infiniopLogCumSumExpDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + int axis, + int exclusive, + int reverse); + +/* 获取执行 LogCumSumExp 所需的临时空间大小 */ +__INFINI_C __export infiniStatus_t infiniopGetLogCumSumExpWorkspaceSize(infiniopLogCumSumExpDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLogCumSumExp(infiniopLogCumSumExpDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +/* 销毁描述符 */ +__INFINI_C __export infiniStatus_t infiniopDestroyLogCumSumExpDescriptor(infiniopLogCumSumExpDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/logdet.h b/python/infinicore/include/infiniop/ops/logdet.h new file mode 100644 index 000000000..f92fb58f3 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/logdet.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_LOGDET_API_H__ +#define __INFINIOP_LOGDET_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLogdetDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLogdetDescriptor(infiniopHandle_t handle, + infiniopLogdetDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetLogdetWorkspaceSize(infiniopLogdetDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLogdet(infiniopLogdetDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLogdetDescriptor(infiniopLogdetDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/lp_norm.h b/python/infinicore/include/infiniop/ops/lp_norm.h new file mode 100644 index 000000000..1022fdb85 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/lp_norm.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_LP_NORM_API_H__ +#define __INFINIOP_LP_NORM_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLPNormDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateLPNormDescriptor( + infiniopHandle_t handle, + infiniopLPNormDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + int axis, + int p, + float eps); + +__INFINI_C __export infiniStatus_t infiniopGetLPNormWorkspaceSize(infiniopLPNormDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopLPNorm(infiniopLPNormDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyLPNormDescriptor(infiniopLPNormDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/mamba_selective_scan.h b/python/infinicore/include/infiniop/ops/mamba_selective_scan.h new file mode 100644 index 000000000..c4eb5f574 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/mamba_selective_scan.h @@ -0,0 +1,50 @@ +#ifndef __INFINIOP_MAMBA_SELECTIVE_SCAN_API_H__ +#define __INFINIOP_MAMBA_SELECTIVE_SCAN_API_H__ + +#include "../operator_descriptor.h" +#ifdef __cplusplus +#include +#else +#include +#endif + +typedef struct InfiniopDescriptor *infiniopMambaSelectiveScanDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMambaSelectiveScanDescriptor( + infiniopHandle_t handle, + infiniopMambaSelectiveScanDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t dt_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t c_desc, + infiniopTensorDescriptor_t a_log_desc, + infiniopTensorDescriptor_t d_desc, + infiniopTensorDescriptor_t gate_desc, + infiniopTensorDescriptor_t dt_bias_desc, + infiniopTensorDescriptor_t state_desc); + +__INFINI_C __export infiniStatus_t infiniopGetMambaSelectiveScanWorkspaceSize( + infiniopMambaSelectiveScanDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMambaSelectiveScan( + infiniopMambaSelectiveScanDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *x, + const void *dt, + const void *b, + const void *c, + const void *a_log, + const void *d, + const void *gate, + const void *dt_bias, + void *state, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMambaSelectiveScanDescriptor( + infiniopMambaSelectiveScanDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/masked_select.h b/python/infinicore/include/infiniop/ops/masked_select.h new file mode 100644 index 000000000..d24ed11f7 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/masked_select.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_MASKED_SELECT_API_H__ +#define __INFINIOP_MASKED_SELECT_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMaskedSelectDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMaskedSelectDescriptor( + infiniopHandle_t handle, + infiniopMaskedSelectDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t mask_desc); + +__INFINI_C __export infiniStatus_t infiniopGetMaskedSelectWorkspaceSize(infiniopMaskedSelectDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMaskedSelect( + infiniopMaskedSelectDescriptor_t desc, + void *workspace, + size_t workspace_size, + const void *input, + const bool *mask, + void **data_ptr, + size_t *dlen, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMaskedSelectDescriptor(infiniopMaskedSelectDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/matrix_power.h b/python/infinicore/include/infiniop/ops/matrix_power.h new file mode 100644 index 000000000..639d3cf02 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/matrix_power.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_MATRIX_POWER_API_H__ +#define __INFINIOP_MATRIX_POWER_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMatrixPowerDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMatrixPowerDescriptor(infiniopHandle_t handle, + infiniopMatrixPowerDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + int n); + +__INFINI_C __export infiniStatus_t infiniopGetMatrixPowerWorkspaceSize(infiniopMatrixPowerDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMatrixPower(infiniopMatrixPowerDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMatrixPowerDescriptor(infiniopMatrixPowerDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/moe_align.h b/python/infinicore/include/infiniop/ops/moe_align.h new file mode 100644 index 000000000..582bc781f --- /dev/null +++ b/python/infinicore/include/infiniop/ops/moe_align.h @@ -0,0 +1,37 @@ +#ifndef __INFINIOP_MOE_ALIGN_API_H__ +#define __INFINIOP_MOE_ALIGN_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMoeAlignDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMoeAlignDescriptor( + infiniopHandle_t handle, + infiniopMoeAlignDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t sorted_token_ids_desc, + infiniopTensorDescriptor_t expert_ids_desc, + infiniopTensorDescriptor_t num_tokens_post_padded_desc, + infiniopTensorDescriptor_t topk_ids_desc, + size_t num_experts, + size_t block_size); + +__INFINI_C __export infiniStatus_t infiniopGetMoeAlignWorkspaceSize( + infiniopMoeAlignDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMoeAlign( + infiniopMoeAlignDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *sorted_token_ids, + void *expert_ids, + void *num_tokens_post_padded, + const void *topk_ids, + const void *expert_map, + int pad_sorted_token_ids, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMoeAlignDescriptor( + infiniopMoeAlignDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/moe_fused_dense.h b/python/infinicore/include/infiniop/ops/moe_fused_dense.h new file mode 100644 index 000000000..61812ce01 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/moe_fused_dense.h @@ -0,0 +1,43 @@ +#ifndef __INFINIOP_MOE_FUSED_DENSE_API_H__ +#define __INFINIOP_MOE_FUSED_DENSE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMoeFusedDenseDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMoeFusedDenseDescriptor( + infiniopHandle_t handle, + infiniopMoeFusedDenseDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t hidden_states_desc, + infiniopTensorDescriptor_t w13_desc, + infiniopTensorDescriptor_t w2_desc, + infiniopTensorDescriptor_t topk_weights_desc, + infiniopTensorDescriptor_t topk_ids_desc, + infiniopTensorDescriptor_t sorted_token_ids_desc, + infiniopTensorDescriptor_t expert_ids_desc, + infiniopTensorDescriptor_t num_tokens_post_padded_desc); + +__INFINI_C __export infiniStatus_t infiniopGetMoeFusedDenseWorkspaceSize( + infiniopMoeFusedDenseDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMoeFusedDense( + infiniopMoeFusedDenseDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *hidden_states, + const void *w13, + const void *w2, + const void *topk_weights, + const void *topk_ids, + const void *sorted_token_ids, + const void *expert_ids, + const void *num_tokens_post_padded, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMoeFusedDenseDescriptor( + infiniopMoeFusedDenseDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/moe_fused_gate.h b/python/infinicore/include/infiniop/ops/moe_fused_gate.h new file mode 100644 index 000000000..b64a901c1 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/moe_fused_gate.h @@ -0,0 +1,38 @@ +#ifndef __INFINIOP_MOE_FUSED_GATE_API_H__ +#define __INFINIOP_MOE_FUSED_GATE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMoeFusedGateDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMoeFusedGateDescriptor( + infiniopHandle_t handle, + infiniopMoeFusedGateDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t topk_weights_desc, + infiniopTensorDescriptor_t topk_indices_desc, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t bias_desc, + size_t num_expert_group, + size_t topk_group, + size_t num_fused_shared_experts, + float routed_scaling_factor, + bool apply_routed_scaling_factor_on_output); + +__INFINI_C __export infiniStatus_t infiniopGetMoeFusedGateWorkspaceSize( + infiniopMoeFusedGateDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMoeFusedGate( + infiniopMoeFusedGateDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *topk_weights, + void *topk_indices, + const void *input, + const void *bias, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMoeFusedGateDescriptor( + infiniopMoeFusedGateDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/moe_sum.h b/python/infinicore/include/infiniop/ops/moe_sum.h new file mode 100644 index 000000000..3b9a79f0d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/moe_sum.h @@ -0,0 +1,29 @@ +#ifndef __INFINIOP_MOE_SUM_API_H__ +#define __INFINIOP_MOE_SUM_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMoeSumDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMoeSumDescriptor( + infiniopHandle_t handle, + infiniopMoeSumDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc); + +__INFINI_C __export infiniStatus_t infiniopGetMoeSumWorkspaceSize( + infiniopMoeSumDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMoeSum( + infiniopMoeSumDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMoeSumDescriptor( + infiniopMoeSumDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/moe_topk_sigmoid.h b/python/infinicore/include/infiniop/ops/moe_topk_sigmoid.h new file mode 100644 index 000000000..0afd40e13 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/moe_topk_sigmoid.h @@ -0,0 +1,34 @@ +#ifndef __INFINIOP_MOE_TOPK_SIGMOID_API_H__ +#define __INFINIOP_MOE_TOPK_SIGMOID_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMoeTopkSigmoidDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMoeTopkSigmoidDescriptor( + infiniopHandle_t handle, + infiniopMoeTopkSigmoidDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t topk_weights_desc, + infiniopTensorDescriptor_t topk_indices_desc, + infiniopTensorDescriptor_t gating_output_desc, + infiniopTensorDescriptor_t correction_bias_desc, + bool renormalize); + +__INFINI_C __export infiniStatus_t infiniopGetMoeTopkSigmoidWorkspaceSize( + infiniopMoeTopkSigmoidDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMoeTopkSigmoid( + infiniopMoeTopkSigmoidDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *topk_weights, + void *topk_indices, + const void *gating_output, + const void *correction_bias, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMoeTopkSigmoidDescriptor( + infiniopMoeTopkSigmoidDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/moe_topk_softmax.h b/python/infinicore/include/infiniop/ops/moe_topk_softmax.h new file mode 100644 index 000000000..a79425e25 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/moe_topk_softmax.h @@ -0,0 +1,35 @@ +#ifndef __INFINIOP_MOE_TOPK_SOFTMAX_API_H__ +#define __INFINIOP_MOE_TOPK_SOFTMAX_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMoeTopkSoftmaxDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMoeTopkSoftmaxDescriptor( + infiniopHandle_t handle, + infiniopMoeTopkSoftmaxDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t topk_weights_desc, + infiniopTensorDescriptor_t topk_indices_desc, + infiniopTensorDescriptor_t gating_output_desc, + infiniopTensorDescriptor_t correction_bias_desc, + bool renormalize, + float moe_softcapping); + +__INFINI_C __export infiniStatus_t infiniopGetMoeTopkSoftmaxWorkspaceSize( + infiniopMoeTopkSoftmaxDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMoeTopkSoftmax( + infiniopMoeTopkSoftmaxDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *topk_weights, + void *topk_indices, + const void *gating_output, + const void *correction_bias, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMoeTopkSoftmaxDescriptor( + infiniopMoeTopkSoftmaxDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/mrope.h b/python/infinicore/include/infiniop/ops/mrope.h new file mode 100644 index 000000000..be33bdcce --- /dev/null +++ b/python/infinicore/include/infiniop/ops/mrope.h @@ -0,0 +1,42 @@ +#ifndef __INFINIOP_MROPE_API_H__ +#define __INFINIOP_MROPE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMRoPEDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMRoPEDescriptor( + infiniopHandle_t handle, + infiniopMRoPEDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t q_out, + infiniopTensorDescriptor_t k_out, + infiniopTensorDescriptor_t q, + infiniopTensorDescriptor_t k, + infiniopTensorDescriptor_t cos, + infiniopTensorDescriptor_t sin, + infiniopTensorDescriptor_t positions, + int head_size, + int rotary_dim, + int section_t, + int section_h, + int section_w, + bool interleaved); + +__INFINI_C __export infiniStatus_t infiniopGetMRoPEWorkspaceSize(infiniopMRoPEDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMRoPE( + infiniopMRoPEDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *q_out, + void *k_out, + const void *q, + const void *k, + void const *cos, + void const *sin, + void const *positions, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMRoPEDescriptor(infiniopMRoPEDescriptor_t desc); + +#endif // __INFINIOP_MROPE_API_H__ diff --git a/python/infinicore/include/infiniop/ops/mul.h b/python/infinicore/include/infiniop/ops/mul.h new file mode 100644 index 000000000..e3e48c39d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/mul.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_MUL_API_H__ +#define __INFINIOP_MUL_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMulDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMulDescriptor(infiniopHandle_t handle, + infiniopMulDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__INFINI_C __export infiniStatus_t infiniopGetMulWorkspaceSize(infiniopMulDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMul(infiniopMulDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMulDescriptor(infiniopMulDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/mul_scalar.h b/python/infinicore/include/infiniop/ops/mul_scalar.h new file mode 100644 index 000000000..9f90284ee --- /dev/null +++ b/python/infinicore/include/infiniop/ops/mul_scalar.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_MUL_SCALAR_API_H__ +#define __INFINIOP_MUL_SCALAR_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMulScalarDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateMulScalarDescriptor(infiniopHandle_t handle, + infiniopMulScalarDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input); + +__INFINI_C __export infiniStatus_t infiniopGetMulScalarWorkspaceSize(infiniopMulScalarDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMulScalar(infiniopMulScalarDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + double alpha, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMulScalarDescriptor(infiniopMulScalarDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/multi_margin_loss.h b/python/infinicore/include/infiniop/ops/multi_margin_loss.h new file mode 100644 index 000000000..9789e2ae4 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/multi_margin_loss.h @@ -0,0 +1,30 @@ +#ifndef __INFINIOP_MULTI_MARGIN_LOSS_API_H__ +#define __INFINIOP_MULTI_MARGIN_LOSS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopMultiMarginLossDescriptor_t; +__INFINI_C __export infiniStatus_t infiniopCreateMultiMarginLossDescriptor(infiniopHandle_t handle, + infiniopMultiMarginLossDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + infiniopTensorDescriptor_t target, + infiniopTensorDescriptor_t weight, + int p, + float margin, + int reduction); + +__INFINI_C __export infiniStatus_t infiniopGetMultiMarginLossWorkspaceSize(infiniopMultiMarginLossDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopMultiMarginLoss(infiniopMultiMarginLossDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *target, + const void *weight, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyMultiMarginLossDescriptor(infiniopMultiMarginLossDescriptor_t desc); + +#endif // __INFINIOP_MULTI_MARGIN_LOSS_API_H__ diff --git a/python/infinicore/include/infiniop/ops/nrm2.h b/python/infinicore/include/infiniop/ops/nrm2.h new file mode 100644 index 000000000..2f1eed348 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/nrm2.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_NRM2_API_H__ +#define __INFINIOP_NRM2_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopNrm2Descriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateNrm2Descriptor(infiniopHandle_t handle, + infiniopNrm2Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t result); + +__INFINI_C __export infiniStatus_t infiniopGetNrm2WorkspaceSize(infiniopNrm2Descriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopNrm2(infiniopNrm2Descriptor_t desc, + void *workspace, + size_t workspace_size, + const void *x, + void *result, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyNrm2Descriptor(infiniopNrm2Descriptor_t desc); + +#endif // __INFINIOP_NRM2_API_H__ diff --git a/python/infinicore/include/infiniop/ops/nsa_compress_paged_cache.h b/python/infinicore/include/infiniop/ops/nsa_compress_paged_cache.h new file mode 100644 index 000000000..f69d3fb5d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/nsa_compress_paged_cache.h @@ -0,0 +1,39 @@ +#ifndef __INFINIOP_NSA_COMPRESS_PAGED_CACHE_API_H__ +#define __INFINIOP_NSA_COMPRESS_PAGED_CACHE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopNsaCompressPagedCacheDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateNsaCompressPagedCacheDescriptor( + infiniopHandle_t handle, + infiniopNsaCompressPagedCacheDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t k_cmp_desc, + infiniopTensorDescriptor_t v_cmp_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t v_cache_desc, + infiniopTensorDescriptor_t block_tables_desc, + infiniopTensorDescriptor_t seq_lens_desc, + int nsa_block_size, + int update_last_only); + +__INFINI_C __export infiniStatus_t infiniopGetNsaCompressPagedCacheWorkspaceSize( + infiniopNsaCompressPagedCacheDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopNsaCompressPagedCache( + infiniopNsaCompressPagedCacheDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *k_cmp, + void *v_cmp, + const void *k_cache, + const void *v_cache, + const void *block_tables, + const void *seq_lens, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyNsaCompressPagedCacheDescriptor( + infiniopNsaCompressPagedCacheDescriptor_t desc); + +#endif // __INFINIOP_NSA_COMPRESS_PAGED_CACHE_API_H__ diff --git a/python/infinicore/include/infiniop/ops/nsa_paged_attention.h b/python/infinicore/include/infiniop/ops/nsa_paged_attention.h new file mode 100644 index 000000000..e92414f86 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/nsa_paged_attention.h @@ -0,0 +1,47 @@ +#ifndef __INFINIOP_NSA_PAGED_ATTENTION_API_H__ +#define __INFINIOP_NSA_PAGED_ATTENTION_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopNsaPagedAttentionDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateNsaPagedAttentionDescriptor( + infiniopHandle_t handle, + infiniopNsaPagedAttentionDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t k_cmp_desc, + infiniopTensorDescriptor_t v_cmp_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t v_cache_desc, + infiniopTensorDescriptor_t block_tables_desc, + infiniopTensorDescriptor_t seq_lens_desc, + infiniopTensorDescriptor_t gates_desc, + float scale, + int nsa_block_size, + int window_size, + int select_blocks); + +__INFINI_C __export infiniStatus_t infiniopGetNsaPagedAttentionWorkspaceSize( + infiniopNsaPagedAttentionDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopNsaPagedAttention( + infiniopNsaPagedAttentionDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *q, + const void *k_cmp, + const void *v_cmp, + const void *k_cache, + const void *v_cache, + const void *block_tables, + const void *seq_lens, + const void *gates, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyNsaPagedAttentionDescriptor( + infiniopNsaPagedAttentionDescriptor_t desc); + +#endif // __INFINIOP_NSA_PAGED_ATTENTION_API_H__ diff --git a/python/infinicore/include/infiniop/ops/ones.h b/python/infinicore/include/infiniop/ops/ones.h new file mode 100644 index 000000000..e66b0f884 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/ones.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_ONES_API_H__ +#define __INFINIOP_ONES_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopOnesDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateOnesDescriptor(infiniopHandle_t handle, + infiniopOnesDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetOnesWorkspaceSize(infiniopOnesDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopOnes(infiniopOnesDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyOnesDescriptor(infiniopOnesDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/pad.h b/python/infinicore/include/infiniop/ops/pad.h new file mode 100644 index 000000000..af3af576f --- /dev/null +++ b/python/infinicore/include/infiniop/ops/pad.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_PAD_API_H__ +#define __INFINIOP_PAD_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopPadDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreatePadDescriptor(infiniopHandle_t handle, + infiniopPadDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + void *pad, + size_t pad_size, + const char *mode, + double value); + +__INFINI_C __export infiniStatus_t infiniopGetPadWorkspaceSize(infiniopPadDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopPad(infiniopPadDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyPadDescriptor(infiniopPadDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/paged_attention.h b/python/infinicore/include/infiniop/ops/paged_attention.h new file mode 100644 index 000000000..7f1656ef3 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/paged_attention.h @@ -0,0 +1,93 @@ +#ifndef __INFINIOP_PAGED_ATTENTION_API_H__ +#define __INFINIOP_PAGED_ATTENTION_API_H__ + +#include "../operator_descriptor.h" + +// Define an opaque handle for the Paged Attention descriptor. +typedef struct InfiniopDescriptor *infiniopPagedAttentionDescriptor_t; + +/** + * @brief Creates a descriptor for the Paged Attention v1 operation. + * + * @param handle The library context handle. + * @param desc_ptr Pointer to the created descriptor. + * @param out_desc [Output] Shape: (num_seqs, num_heads, head_size). + * The output tensor for the attention mechanism. + * @param q_desc [Input] Shape: (num_seqs, num_heads, head_size). + * The query tensor. + * @param k_cache_desc [Input] Shape: (num_blocks, num_kv_heads, block_size, head_size). + * Paged key cache storing keys for all sequences. + * @param v_cache_desc [Input] Shape: (num_blocks, num_kv_heads, block_size, head_size). + * Paged value cache storing values for all sequences. + * @param block_tables_desc [Input] Shape: (num_seqs, max_num_blocks_per_seq). + * Maps each sequence to its physical block indices in the cache. + * Expected DType: int64_t (I64). + * @param seq_lens_desc [Input] Shape: (num_seqs,). + * The current logical length of each sequence. + * Expected DType: int64_t (I64). + * @param alibi_slopes_desc [Optional] Shape: (num_heads,). + * Slopes for ALiBi (Attention with Linear Biases). Can be NULL. + * @param scale The attention scaling factor (typically 1/sqrt(head_size)). + * @return infiniStatus_t Status code. + */ +__INFINI_C __export infiniStatus_t infiniopCreatePagedAttentionDescriptor( + infiniopHandle_t handle, + infiniopPagedAttentionDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t v_cache_desc, + infiniopTensorDescriptor_t block_tables_desc, + infiniopTensorDescriptor_t seq_lens_desc, + infiniopTensorDescriptor_t alibi_slopes_desc, + float scale); + +/** + * @brief Retrieves the workspace size required for the Paged Attention operation. + * + * @param desc The Paged Attention descriptor. + * @param size A pointer to store the required workspace size in bytes. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopGetPagedAttentionWorkspaceSize( + infiniopPagedAttentionDescriptor_t desc, size_t *size); + +/** + * @brief Executes the Paged Attention v1 operation. + * + * @param desc The Paged Attention descriptor. + * @param workspace Pointer to the workspace memory. + * @param workspace_size The size of the workspace. + * @param out Pointer to the output tensor data. + * @param q Pointer to the query tensor data. + * @param k_cache Pointer to the key cache data. + * @param v_cache Pointer to the value cache data. + * @param block_tables Pointer to the block tables data. + * @param seq_lens Pointer to the sequence lengths data. + * @param alibi_slopes Pointer to the ALiBi slopes data. Can be NULL. + * @param stream The CUDA stream for the operation. Can be NULL. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopPagedAttention( + infiniopPagedAttentionDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *q, + const void *k_cache, + const void *v_cache, + const void *block_tables, + const void *seq_lens, + const void *alibi_slopes, + void *stream); + +/** + * @brief Destroys a Paged Attention descriptor. + * + * @param desc The descriptor to be destroyed. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopDestroyPagedAttentionDescriptor( + infiniopPagedAttentionDescriptor_t desc); + +#endif // __INFINIOP_PAGED_ATTENTION_API_H__ diff --git a/python/infinicore/include/infiniop/ops/paged_attention_prefill.h b/python/infinicore/include/infiniop/ops/paged_attention_prefill.h new file mode 100644 index 000000000..e2e93076b --- /dev/null +++ b/python/infinicore/include/infiniop/ops/paged_attention_prefill.h @@ -0,0 +1,87 @@ +#ifndef __INFINIOP_PAGED_ATTENTION_PREFILL_API_H__ +#define __INFINIOP_PAGED_ATTENTION_PREFILL_API_H__ + +#include "../operator_descriptor.h" + +// Define an opaque handle for the Paged Attention Prefill descriptor. +typedef struct InfiniopDescriptor *infiniopPagedAttentionPrefillDescriptor_t; + +/** + * @brief Creates a descriptor for the Paged Attention Prefill operation. + * @param handle The handle to the InfiniOP library context. + * @param desc_ptr A pointer to store the created descriptor. + * @param out_desc Descriptor for the output tensor. + * Shape: [total_q_tokens, num_heads, head_size] + * @param q_desc Descriptor for the query tensor (packed/flattened). + * Shape: [total_q_tokens, num_heads, head_size] + * @param k_cache_desc Descriptor for the global physical key cache. + * Shape: [max_num_blocks, num_kv_heads, block_size, head_size] + * @param v_cache_desc Descriptor for the global physical value cache. + * Shape: [max_num_blocks, num_kv_heads, block_size, head_size] + * @param block_tables_desc Descriptor for the block tables mapping logic to physical blocks. + * Shape: [batch_size, max_blocks_per_seq] + * @param seq_lens_desc Descriptor for the total KV lengths of each sequence. + * Shape: [batch_size] + * @param cum_seq_lens_q_desc Descriptor for the cumulative start position (prefix sum) of each Q sequence. + * Shape: [batch_size + 1] + * @param alibi_slopes_desc Optional descriptor for the ALiBi slopes tensor. Can be NULL. + * Shape: [num_heads] + * @param scale The attention scaling factor (typically 1.0 / sqrt(head_size)). + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopCreatePagedAttentionPrefillDescriptor( + infiniopHandle_t handle, + infiniopPagedAttentionPrefillDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t v_cache_desc, + infiniopTensorDescriptor_t block_tables_desc, + infiniopTensorDescriptor_t seq_lens_desc, + infiniopTensorDescriptor_t cum_seq_lens_q_desc, + infiniopTensorDescriptor_t alibi_slopes_desc, + float scale); + +/** + * @brief Retrieves the workspace size required for the Paged Attention Prefill operation. + */ +__INFINI_C __export infiniStatus_t infiniopGetPagedAttentionPrefillWorkspaceSize( + infiniopPagedAttentionPrefillDescriptor_t desc, size_t *size); + +/** + * @brief Executes the Paged Attention Prefill operation. + * @param desc The Paged Attention Prefill descriptor. + * @param workspace Pointer to the workspace memory. + * @param workspace_size The size of the workspace. + * @param out Pointer to the output tensor data. + * @param q Pointer to the query tensor data (packed). + * @param k_cache Pointer to the global key cache data. + * @param v_cache Pointer to the global value cache data. + * @param block_tables Pointer to the block tables data. + * @param seq_lens Pointer to the KV lengths data. + * @param cum_seq_lens_q Pointer to the Q cumulative sequence lengths data (prefix sum). + * @param alibi_slopes Pointer to the ALiBi slopes data. Can be NULL. + * @param stream The device stream (e.g., cudaStream_t) for the operation. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopPagedAttentionPrefill( + infiniopPagedAttentionPrefillDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *q, + const void *k_cache, + const void *v_cache, + const void *block_tables, + const void *seq_lens, + const void *cum_seq_lens_q, + const void *alibi_slopes, + void *stream); + +/** + * @brief Destroys a Paged Attention Prefill descriptor. + */ +__INFINI_C __export infiniStatus_t infiniopDestroyPagedAttentionPrefillDescriptor( + infiniopPagedAttentionPrefillDescriptor_t desc); + +#endif // __INFINIOP_PAGED_ATTENTION_PREFILL_API_H__ diff --git a/python/infinicore/include/infiniop/ops/paged_caching.h b/python/infinicore/include/infiniop/ops/paged_caching.h new file mode 100644 index 000000000..d85125c30 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/paged_caching.h @@ -0,0 +1,77 @@ +#ifndef __INFINIOP_PAGED_CACHING_API_H__ +#define __INFINIOP_PAGED_CACHING_API_H__ + +#include "../operator_descriptor.h" + +// Define an opaque handle for the Paged Caching descriptor. +typedef struct InfiniopDescriptor *infiniopPagedCachingDescriptor_t; + +/** + * @brief Creates a descriptor for the Paged Caching operation. + * + * This function initializes a descriptor that holds all the metadata needed + * to copy key/value vectors into their respective cache pools. + * + * @param handle The handle to the InfiniOP library context. + * @param desc_ptr A pointer to store the created descriptor. + * @param k_cache_desc Descriptor for the key cache pool tensor. + * @param v_cache_desc Descriptor for the value cache pool tensor. + * @param k_desc Descriptor for the source key tensor. + * @param v_desc Descriptor for the source value tensor. + * @param slot_mapping_desc Descriptor for the slot mapping tensor. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopCreatePagedCachingDescriptor( + infiniopHandle_t handle, + infiniopPagedCachingDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t k_cache_desc, + infiniopTensorDescriptor_t v_cache_desc, + infiniopTensorDescriptor_t k_desc, + infiniopTensorDescriptor_t v_desc, + infiniopTensorDescriptor_t slot_mapping_desc); + +/** + * @brief Retrieves the workspace size required for the Paged Caching operation. + * + * @param desc The Paged Caching descriptor. + * @param size A pointer to store the required workspace size in bytes (typically 0). + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopGetPagedCachingWorkspaceSize( + infiniopPagedCachingDescriptor_t desc, size_t *size); + +/** + * @brief Executes the Paged Caching operation. + * + * @param desc The Paged Caching descriptor. + * @param workspace Pointer to the workspace memory. + * @param workspace_size The size of the workspace. + * @param k_cache Pointer to the key cache pool data. + * @param v_cache Pointer to the value cache pool data. + * @param k Pointer to the source key tensor data. + * @param v Pointer to the source value tensor data. + * @param slot_mapping Pointer to the slot mapping data. + * @param stream The CUDA stream for the operation. Can be NULL. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopPagedCaching( + infiniopPagedCachingDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *k_cache, + void *v_cache, + const void *k, + const void *v, + const void *slot_mapping, + void *stream); + +/** + * @brief Destroys a Paged Caching descriptor. + * + * @param desc The descriptor to be destroyed. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopDestroyPagedCachingDescriptor( + infiniopPagedCachingDescriptor_t desc); + +#endif // __INFINIOP_PAGED_CACHING_API_H__ diff --git a/python/infinicore/include/infiniop/ops/pixel_shuffle.h b/python/infinicore/include/infiniop/ops/pixel_shuffle.h new file mode 100644 index 000000000..fac8eff2c --- /dev/null +++ b/python/infinicore/include/infiniop/ops/pixel_shuffle.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_PIXEL_SHUFFLE_API_H__ +#define __INFINIOP_PIXEL_SHUFFLE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopPixelShuffleDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreatePixelShuffleDescriptor(infiniopHandle_t handle, + infiniopPixelShuffleDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + int upscale_factor); + +__INFINI_C __export infiniStatus_t infiniopGetPixelShuffleWorkspaceSize(infiniopPixelShuffleDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopPixelShuffle(infiniopPixelShuffleDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyPixelShuffleDescriptor(infiniopPixelShuffleDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/prelu.h b/python/infinicore/include/infiniop/ops/prelu.h new file mode 100644 index 000000000..e90a571be --- /dev/null +++ b/python/infinicore/include/infiniop/ops/prelu.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_PRELU_API_H__ +#define __INFINIOP_PRELU_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopPreluDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreatePreluDescriptor(infiniopHandle_t handle, + infiniopPreluDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t weight); + +__INFINI_C __export infiniStatus_t infiniopGetPreluWorkspaceSize(infiniopPreluDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopPrelu(infiniopPreluDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + const void *weight, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyPreluDescriptor(infiniopPreluDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/prepare_moe_input.h b/python/infinicore/include/infiniop/ops/prepare_moe_input.h new file mode 100644 index 000000000..8621d044a --- /dev/null +++ b/python/infinicore/include/infiniop/ops/prepare_moe_input.h @@ -0,0 +1,42 @@ +#ifndef __INFINIOP_PREPARE_MOE_INPUT_API_H__ +#define __INFINIOP_PREPARE_MOE_INPUT_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopPrepareMoeInputDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreatePrepareMoeInputDescriptor( + infiniopHandle_t handle, + infiniopPrepareMoeInputDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t expert_offsets_desc, + infiniopTensorDescriptor_t blockscale_offsets_desc, + infiniopTensorDescriptor_t problem_sizes1_desc, + infiniopTensorDescriptor_t problem_sizes2_desc, + infiniopTensorDescriptor_t input_permutation_desc, + infiniopTensorDescriptor_t output_permutation_desc, + infiniopTensorDescriptor_t topk_ids_desc, + size_t num_experts, + size_t n, + size_t k); + +__INFINI_C __export infiniStatus_t infiniopGetPrepareMoeInputWorkspaceSize( + infiniopPrepareMoeInputDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopPrepareMoeInput( + infiniopPrepareMoeInputDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *expert_offsets, + void *blockscale_offsets, + void *problem_sizes1, + void *problem_sizes2, + void *input_permutation, + void *output_permutation, + const void *topk_ids, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyPrepareMoeInputDescriptor( + infiniopPrepareMoeInputDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/quant/per_channel_quant_int8.h b/python/infinicore/include/infiniop/ops/quant/per_channel_quant_int8.h new file mode 100644 index 000000000..faa5efd78 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/quant/per_channel_quant_int8.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_PER_CHANNEL_QUANT_INT8_API_H__ +#define __INFINIOP_PER_CHANNEL_QUANT_INT8_API_H__ + +#include "../../operator_descriptor.h" + +typedef InfiniopDescriptor *infiniopPerChannelQuantI8Descriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreatePerChannelQuantI8Descriptor(infiniopHandle_t handle, + infiniopPerChannelQuantI8Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t x_packed_desc, + infiniopTensorDescriptor_t x_scale_desc, + infiniopTensorDescriptor_t x_zero_desc, + infiniopTensorDescriptor_t x_desc); + +__INFINI_C __export infiniStatus_t infiniopGetPerChannelQuantI8WorkspaceSize(infiniopPerChannelQuantI8Descriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopPerChannelQuantI8(infiniopPerChannelQuantI8Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *x_packed, + void *x_scale, + void *x_zero, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyPerChannelQuantI8Descriptor(infiniopPerChannelQuantI8Descriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/quant/per_tensor_quant_int8.h b/python/infinicore/include/infiniop/ops/quant/per_tensor_quant_int8.h new file mode 100644 index 000000000..16e1c2bc6 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/quant/per_tensor_quant_int8.h @@ -0,0 +1,29 @@ +#ifndef __INFINIOP_PER_TENSOR_QUANT_INT8_API_H__ +#define __INFINIOP_PER_TENSOR_QUANT_INT8_API_H__ + +#include "../../operator_descriptor.h" + +typedef InfiniopDescriptor *infiniopPerTensorQuantI8Descriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreatePerTensorQuantI8Descriptor(infiniopHandle_t handle, + infiniopPerTensorQuantI8Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t x_packed_desc, + infiniopTensorDescriptor_t x_scale_desc, + infiniopTensorDescriptor_t x_zero_desc, + infiniopTensorDescriptor_t x_desc); + +__INFINI_C __export infiniStatus_t infiniopGetPerTensorQuantI8WorkspaceSize(infiniopPerTensorQuantI8Descriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopPerTensorQuantI8(infiniopPerTensorQuantI8Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *x_packed, + void *x_scale, + void *x_zero, + const void *x, + const bool is_static, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyPerTensorQuantI8Descriptor(infiniopPerTensorQuantI8Descriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/quickgelu.h b/python/infinicore/include/infiniop/ops/quickgelu.h new file mode 100644 index 000000000..be66f9495 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/quickgelu.h @@ -0,0 +1,42 @@ +#ifndef __INFINIOP_QUICKGELU_API_H__ +#define __INFINIOP_QUICKGELU_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopQuickGeluDescriptor_t; + +/** + * Create QuickGELU descriptor + * y = x * sigmoid(1.702 * x) + */ +__INFINI_C __export infiniStatus_t infiniopCreateQuickGeluDescriptor( + infiniopHandle_t handle, + infiniopQuickGeluDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +/** + * Query workspace size + */ +__INFINI_C __export infiniStatus_t infiniopGetQuickGeluWorkspaceSize( + infiniopQuickGeluDescriptor_t desc, + size_t *size); + +/** + * Launch QuickGELU operator + */ +__INFINI_C __export infiniStatus_t infiniopQuickGelu( + infiniopQuickGeluDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +/** + * Destroy descriptor + */ +__INFINI_C __export infiniStatus_t infiniopDestroyQuickGeluDescriptor( + infiniopQuickGeluDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/random_sample.h b/python/infinicore/include/infiniop/ops/random_sample.h new file mode 100644 index 000000000..c12dc0af3 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/random_sample.h @@ -0,0 +1,52 @@ +#ifndef __INFINIOP_RANDOM_SAMPLE_API_H__ +#define __INFINIOP_RANDOM_SAMPLE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopRandomSampleDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRandomSampleDescriptor( + infiniopHandle_t handle, + infiniopRandomSampleDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t result, + infiniopTensorDescriptor_t probs); + +__INFINI_C __export infiniStatus_t infiniopGetRandomSampleWorkspaceSize( + infiniopRandomSampleDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopCreateRandomSampleBatchDescriptor( + infiniopHandle_t handle, + infiniopRandomSampleDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t result, + infiniopTensorDescriptor_t probs); + +__INFINI_C __export infiniStatus_t infiniopRandomSample( + infiniopRandomSampleDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *result, + const void *probs, + float random_val, + float topp, + int topk, + float temperature, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopRandomSampleBatch( + infiniopRandomSampleDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *result, + const void *probs, + const float *random_val, + const float *topp, + const int *topk, + const float *temperature, + int batch_size, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRandomSampleDescriptor( + infiniopRandomSampleDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/rearrange.h b/python/infinicore/include/infiniop/ops/rearrange.h new file mode 100644 index 000000000..e28aeb97d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/rearrange.h @@ -0,0 +1,23 @@ +#ifndef __INFINIOP_REARRANGE_API_H__ +#define __INFINIOP_REARRANGE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopRearrangeDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRearrangeDescriptor( + infiniopHandle_t handle, + infiniopRearrangeDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t dst, + infiniopTensorDescriptor_t src); + +__INFINI_C __export infiniStatus_t infiniopRearrange( + infiniopRearrangeDescriptor_t desc, + void *dst, + const void *src, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRearrangeDescriptor( + infiniopRearrangeDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/reciprocal.h b/python/infinicore/include/infiniop/ops/reciprocal.h new file mode 100644 index 000000000..bd1ea31c5 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/reciprocal.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_RECIPROCAL_API_H__ +#define __INFINIOP_RECIPROCAL_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopReciprocalDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateReciprocalDescriptor(infiniopHandle_t handle, + infiniopReciprocalDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetReciprocalWorkspaceSize(infiniopReciprocalDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopReciprocal(infiniopReciprocalDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyReciprocalDescriptor(infiniopReciprocalDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/recurrent_gated_delta_rule.h b/python/infinicore/include/infiniop/ops/recurrent_gated_delta_rule.h new file mode 100644 index 000000000..6a1b29c65 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/recurrent_gated_delta_rule.h @@ -0,0 +1,46 @@ +#ifndef __INFINIOP_RECURRENT_GATED_DELTA_RULE_API_H__ +#define __INFINIOP_RECURRENT_GATED_DELTA_RULE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopRecurrentGatedDeltaRuleDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRecurrentGatedDeltaRuleDescriptor( + infiniopHandle_t handle, + infiniopRecurrentGatedDeltaRuleDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, // [B, T, Hv, Dv], T must be 1; last dim contiguous + infiniopTensorDescriptor_t initial_state_desc, // legacy: [B, Hv, Dv, Dk]; indexed pool: [pool_size, Hv, Dv, Dk] + infiniopTensorDescriptor_t final_state_desc, // legacy/indexed out-of-place final state; null when final_state_indices_desc is provided + infiniopTensorDescriptor_t q_desc, // [B, T, Hk, Dk], T must be 1; last dim contiguous + infiniopTensorDescriptor_t k_desc, // [B, T, Hk, Dk], same shape as q; last dim contiguous + infiniopTensorDescriptor_t v_desc, // [B, T, Hv, Dv], Hv must be a multiple of Hk; last dim contiguous + infiniopTensorDescriptor_t g_desc, // [B, T, Hv]; may have a different fp dtype from q/k/v/out/state + infiniopTensorDescriptor_t beta_desc, // [B, T, Hv]; same dtype as g + infiniopTensorDescriptor_t initial_state_indices_desc, // nullable; [B], int32/int64; enables indexed pool mode + infiniopTensorDescriptor_t final_state_indices_desc, // nullable; [B], int32/int64; writes final state in-place to initial_state pool + bool use_qk_l2norm); + +__INFINI_C __export infiniStatus_t infiniopGetRecurrentGatedDeltaRuleWorkspaceSize( + infiniopRecurrentGatedDeltaRuleDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRecurrentGatedDeltaRule( + infiniopRecurrentGatedDeltaRuleDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + void *initial_state, + void *final_state, + const void *q, + const void *k, + const void *v, + const void *g, + const void *beta, + const void *initial_state_indices, + const void *final_state_indices, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRecurrentGatedDeltaRuleDescriptor( + infiniopRecurrentGatedDeltaRuleDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/relu.h b/python/infinicore/include/infiniop/ops/relu.h new file mode 100644 index 000000000..8b129badd --- /dev/null +++ b/python/infinicore/include/infiniop/ops/relu.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_RELU_API_H__ +#define __INFINIOP_RELU_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopReluDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateReluDescriptor(infiniopHandle_t handle, + infiniopReluDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetReluWorkspaceSize(infiniopReluDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRelu(infiniopReluDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyReluDescriptor(infiniopReluDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/relu6.h b/python/infinicore/include/infiniop/ops/relu6.h new file mode 100644 index 000000000..2246bbc55 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/relu6.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_RELU6_API_H__ +#define __INFINIOP_RELU6_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopRelu6Descriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRelu6Descriptor(infiniopHandle_t handle, + infiniopRelu6Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetRelu6WorkspaceSize(infiniopRelu6Descriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRelu6(infiniopRelu6Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRelu6Descriptor(infiniopRelu6Descriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/rms_norm.h b/python/infinicore/include/infiniop/ops/rms_norm.h new file mode 100644 index 000000000..f90da65e4 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/rms_norm.h @@ -0,0 +1,23 @@ +#ifndef __INFINIOP_RMS_NORM_API_H__ +#define __INFINIOP_RMS_NORM_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopRMSNormDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRMSNormDescriptor( + infiniopHandle_t handle, + infiniopRMSNormDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t w_desc, + float epsilon); + +__INFINI_C __export infiniStatus_t infiniopGetRMSNormWorkspaceSize(infiniopRMSNormDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRMSNorm(infiniopRMSNormDescriptor_t desc, void *workspace, size_t workspace_size, + void *y, const void *x, const void *w, void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRMSNormDescriptor(infiniopRMSNormDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/rope.h b/python/infinicore/include/infiniop/ops/rope.h new file mode 100644 index 000000000..7c91d51f9 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/rope.h @@ -0,0 +1,40 @@ +#ifndef __INFINIOP_ROPE_API_H__ +#define __INFINIOP_ROPE_API_H__ + +#include "../operator_descriptor.h" + +typedef enum { + INFINIOP_ROPE_ALGO_GPT_J = 0, // GPT-J style RoPE algorithm (Interleave even and odd dimensions) + INFINIOP_ROPE_ALGO_GPT_NEOX = 1, // GPT-NeoX style RoPE algorithm (First half dimensions for sin, second half for cos) + // Count + INFINIOP_ROPE_ALGO_COUNT = 2, +} infiniopRoPEAlgo_t; + +typedef struct InfiniopDescriptor *infiniopRoPEDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRoPEDescriptor( + infiniopHandle_t handle, + infiniopRoPEDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t pos_ids, + infiniopTensorDescriptor_t sin_table, + infiniopTensorDescriptor_t cos_table, + infiniopRoPEAlgo_t algo); + +__INFINI_C __export infiniStatus_t infiniopGetRoPEWorkspaceSize(infiniopRoPEDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRoPE( + infiniopRoPEDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void const *pos_ids, + void const *sin_table, + void const *cos_table, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRoPEDescriptor(infiniopRoPEDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/rot.h b/python/infinicore/include/infiniop/ops/rot.h new file mode 100644 index 000000000..0ecbae52d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/rot.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_ROT_API_H__ +#define __INFINIOP_ROT_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopRotDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRotDescriptor(infiniopHandle_t handle, + infiniopRotDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t s); + +__INFINI_C __export infiniStatus_t infiniopGetRotWorkspaceSize(infiniopRotDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRot(infiniopRotDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *x, + void *y, + const void *c, + const void *s, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRotDescriptor(infiniopRotDescriptor_t desc); + +#endif // __INFINIOP_ROT_API_H__ diff --git a/python/infinicore/include/infiniop/ops/rotg.h b/python/infinicore/include/infiniop/ops/rotg.h new file mode 100644 index 000000000..63c2dad46 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/rotg.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_ROTG_API_H__ +#define __INFINIOP_ROTG_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopRotgDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRotgDescriptor(infiniopHandle_t handle, + infiniopRotgDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t s); + +__INFINI_C __export infiniStatus_t infiniopGetRotgWorkspaceSize(infiniopRotgDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRotg(infiniopRotgDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *x, + void *y, + void *c, + void *s, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRotgDescriptor(infiniopRotgDescriptor_t desc); + +#endif // __INFINIOP_ROTG_API_H__ diff --git a/python/infinicore/include/infiniop/ops/rotm.h b/python/infinicore/include/infiniop/ops/rotm.h new file mode 100644 index 000000000..6cc6a636c --- /dev/null +++ b/python/infinicore/include/infiniop/ops/rotm.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_ROTM_API_H__ +#define __INFINIOP_ROTM_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopRotmDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRotmDescriptor(infiniopHandle_t handle, + infiniopRotmDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t param); + +__INFINI_C __export infiniStatus_t infiniopGetRotmWorkspaceSize(infiniopRotmDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRotm(infiniopRotmDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *x, + void *y, + const void *param, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRotmDescriptor(infiniopRotmDescriptor_t desc); + +#endif // __INFINIOP_ROTM_API_H__ diff --git a/python/infinicore/include/infiniop/ops/rotmg.h b/python/infinicore/include/infiniop/ops/rotmg.h new file mode 100644 index 000000000..0295339cb --- /dev/null +++ b/python/infinicore/include/infiniop/ops/rotmg.h @@ -0,0 +1,30 @@ +#ifndef __INFINIOP_ROTMG_API_H__ +#define __INFINIOP_ROTMG_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopRotmgDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRotmgDescriptor(infiniopHandle_t handle, + infiniopRotmgDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t d1, + infiniopTensorDescriptor_t d2, + infiniopTensorDescriptor_t x1, + infiniopTensorDescriptor_t y1, + infiniopTensorDescriptor_t param); + +__INFINI_C __export infiniStatus_t infiniopGetRotmgWorkspaceSize(infiniopRotmgDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRotmg(infiniopRotmgDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *d1, + void *d2, + void *x1, + const void *y1, + void *param, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRotmgDescriptor(infiniopRotmgDescriptor_t desc); + +#endif // __INFINIOP_ROTMG_API_H__ diff --git a/python/infinicore/include/infiniop/ops/rwkv5_wkv.h b/python/infinicore/include/infiniop/ops/rwkv5_wkv.h new file mode 100644 index 000000000..6f11c5497 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/rwkv5_wkv.h @@ -0,0 +1,45 @@ +#ifndef __INFINIOP_RWKV5_WKV_API_H__ +#define __INFINIOP_RWKV5_WKV_API_H__ + +#include "../operator_descriptor.h" + +#ifdef __cplusplus +#include +#else +#include +#endif + +typedef struct InfiniopDescriptor *infiniopRwkv5WkvDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRwkv5WkvDescriptor( + infiniopHandle_t handle, + infiniopRwkv5WkvDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t receptance_desc, + infiniopTensorDescriptor_t key_desc, + infiniopTensorDescriptor_t value_desc, + infiniopTensorDescriptor_t time_decay_desc, + infiniopTensorDescriptor_t time_faaaa_desc, + infiniopTensorDescriptor_t state_desc); + +__INFINI_C __export infiniStatus_t infiniopGetRwkv5WkvWorkspaceSize( + infiniopRwkv5WkvDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRwkv5Wkv( + infiniopRwkv5WkvDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *out, + const void *receptance, + const void *key, + const void *value, + const void *time_decay, + const void *time_faaaa, + void *state, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRwkv5WkvDescriptor( + infiniopRwkv5WkvDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/scal.h b/python/infinicore/include/infiniop/ops/scal.h new file mode 100644 index 000000000..f7903de56 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/scal.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_SCAL_API_H__ +#define __INFINIOP_SCAL_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopScalDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateScalDescriptor(infiniopHandle_t handle, + infiniopScalDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t alpha, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetScalWorkspaceSize(infiniopScalDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopScal(infiniopScalDescriptor_t desc, + void *workspace, + size_t workspace_size, + const void *alpha, + void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyScalDescriptor(infiniopScalDescriptor_t desc); + +#endif // __INFINIOP_SCAL_API_H__ diff --git a/python/infinicore/include/infiniop/ops/scatter.h b/python/infinicore/include/infiniop/ops/scatter.h new file mode 100644 index 000000000..2fa22c534 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/scatter.h @@ -0,0 +1,30 @@ +#ifndef __INFINIOP_SCATTER_API_H__ +#define __INFINIOP_SCATTER_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopScatterDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateScatterDescriptor(infiniopHandle_t handle, + infiniopScatterDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + infiniopTensorDescriptor_t indices, + infiniopTensorDescriptor_t updates, + int axis, + int reduction); + +__INFINI_C __export infiniStatus_t infiniopGetScatterWorkspaceSize(infiniopScatterDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopScatter(infiniopScatterDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *indices, + const void *updates, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyScatterDescriptor(infiniopScatterDescriptor_t desc); + +#endif // __INFINIOP_SCATTER_API_H__ diff --git a/python/infinicore/include/infiniop/ops/select_last_token_hidden.h b/python/infinicore/include/infiniop/ops/select_last_token_hidden.h new file mode 100644 index 000000000..ac86fa8b7 --- /dev/null +++ b/python/infinicore/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/python/infinicore/include/infiniop/ops/selu.h b/python/infinicore/include/infiniop/ops/selu.h new file mode 100644 index 000000000..b1bf3dfcf --- /dev/null +++ b/python/infinicore/include/infiniop/ops/selu.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_SELU_API_H__ +#define __INFINIOP_SELU_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSeluDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSeluDescriptor(infiniopHandle_t handle, + infiniopSeluDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetSeluWorkspaceSize(infiniopSeluDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSelu(infiniopSeluDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySeluDescriptor(infiniopSeluDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/sigmoid.h b/python/infinicore/include/infiniop/ops/sigmoid.h new file mode 100644 index 000000000..461fbdf7a --- /dev/null +++ b/python/infinicore/include/infiniop/ops/sigmoid.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_SIGMOID_API_H__ +#define __INFINIOP_SIGMOID_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSigmoidDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSigmoidDescriptor(infiniopHandle_t handle, + infiniopSigmoidDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetSigmoidWorkspaceSize(infiniopSigmoidDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSigmoid(infiniopSigmoidDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySigmoidDescriptor(infiniopSigmoidDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/silu.h b/python/infinicore/include/infiniop/ops/silu.h new file mode 100644 index 000000000..f864962c2 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/silu.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_SILU_API_H__ +#define __INFINIOP_SILU_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSiluDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSiluDescriptor(infiniopHandle_t handle, + infiniopSiluDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t intput); + +__INFINI_C __export infiniStatus_t infiniopGetSiluWorkspaceSize(infiniopSiluDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSilu(infiniopSiluDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *intput, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySiluDescriptor(infiniopSiluDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/silu_and_mul.h b/python/infinicore/include/infiniop/ops/silu_and_mul.h new file mode 100644 index 000000000..b1effcec8 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/silu_and_mul.h @@ -0,0 +1,71 @@ +#ifndef __INFINIOP_SILU_AND_MUL_API_H__ +#define __INFINIOP_SILU_AND_MUL_API_H__ + +#include "../operator_descriptor.h" + +/** + * @brief Opaque handle for the SiluAndMul descriptor. + */ +typedef struct InfiniopDescriptor *infiniopSiluAndMulDescriptor_t; + +/** + * @brief Creates a descriptor for the SiLU and Multiply (SiluAndMul) operation. + * + * Format: (input_shape, output_shape) + * Referencing vLLM kernel SiluAndMul interface: + * - input_shape is [..., 2*d] (last dimension is split into two halves for SiLU and multiplication) + * - output_shape is [..., d] (last dimension reduced to half) + * + * @param handle The handle to the InfiniOP library context. + * @param desc_ptr A pointer to store the created descriptor. + * @param output Descriptor for the output tensor. Shape [..., d]. + * @param input Descriptor for the input tensor. Shape [..., 2*d]. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopCreateSiluAndMulDescriptor( + infiniopHandle_t handle, + infiniopSiluAndMulDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input); + +/** + * @brief Queries the workspace size required for SiluAndMul computation. + * @param desc The SiluAndMul descriptor. + * @param size Pointer to store the required workspace size in bytes. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopGetSiluAndMulWorkspaceSize( + infiniopSiluAndMulDescriptor_t desc, + size_t *size); + +/** + * @brief Executes the SiluAndMul operation. + * + * Performs SiLU activation on the first half of the last dimension of `input`, + * multiplies element-wise with the second half, and stores the result in `output`. + * + * @param desc The SiluAndMul descriptor. + * @param workspace Pointer to workspace memory allocated according to GetWorkspaceSize(). + * @param workspace_size Size of the workspace in bytes. + * @param output Pointer to the output tensor memory. Shape [..., d]. + * @param input Pointer to the input tensor memory. Shape [..., 2*d]. + * @param stream Pointer to the execution stream (e.g., CUDA stream). Can be NULL for default stream. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopSiluAndMul( + infiniopSiluAndMulDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +/** + * @brief Destroys a previously created SiluAndMul descriptor. + * @param desc The descriptor to destroy. + * @return infiniStatus_t Status code of the operation. + */ +__INFINI_C __export infiniStatus_t infiniopDestroySiluAndMulDescriptor( + infiniopSiluAndMulDescriptor_t desc); + +#endif // __INFINIOP_SILU_AND_MUL_API_H__ diff --git a/python/infinicore/include/infiniop/ops/sinh.h b/python/infinicore/include/infiniop/ops/sinh.h new file mode 100644 index 000000000..78d36b88b --- /dev/null +++ b/python/infinicore/include/infiniop/ops/sinh.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_SINH_API_H__ +#define __INFINIOP_SINH_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSinhDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSinhDescriptor(infiniopHandle_t handle, + infiniopSinhDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetSinhWorkspaceSize(infiniopSinhDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSinh(infiniopSinhDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySinhDescriptor(infiniopSinhDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/smooth_l1_loss.h b/python/infinicore/include/infiniop/ops/smooth_l1_loss.h new file mode 100644 index 000000000..81678634f --- /dev/null +++ b/python/infinicore/include/infiniop/ops/smooth_l1_loss.h @@ -0,0 +1,27 @@ +#ifndef __INFINIOP_SMOOTH_L1_LOSS_API_H__ +#define __INFINIOP_SMOOTH_L1_LOSS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSmoothL1LossDescriptor_t; +__INFINI_C __export infiniStatus_t infiniopCreateSmoothL1LossDescriptor(infiniopHandle_t handle, + infiniopSmoothL1LossDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + infiniopTensorDescriptor_t target, + float beta, + int reduction); + +__INFINI_C __export infiniStatus_t infiniopGetSmoothL1LossWorkspaceSize(infiniopSmoothL1LossDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSmoothL1Loss(infiniopSmoothL1LossDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *target, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySmoothL1LossDescriptor(infiniopSmoothL1LossDescriptor_t desc); + +#endif // __INFINIOP_SMOOTH_L1_LOSS_API_H__ diff --git a/python/infinicore/include/infiniop/ops/softmax.h b/python/infinicore/include/infiniop/ops/softmax.h new file mode 100644 index 000000000..8ab37b123 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/softmax.h @@ -0,0 +1,27 @@ +#ifndef __INFINIOP_SOFTMAX_API_H__ +#define __INFINIOP_SOFTMAX_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSoftmaxDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSoftmaxDescriptor( + infiniopHandle_t handle, + infiniopSoftmaxDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + int axis); + +__INFINI_C __export infiniStatus_t infiniopGetSoftmaxWorkspaceSize(infiniopSoftmaxDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSoftmax( + infiniopSoftmaxDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySoftmaxDescriptor(infiniopSoftmaxDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/softplus.h b/python/infinicore/include/infiniop/ops/softplus.h new file mode 100644 index 000000000..7e0abc60b --- /dev/null +++ b/python/infinicore/include/infiniop/ops/softplus.h @@ -0,0 +1,29 @@ +#ifndef __INFINIOP_OPS_SOFTPLUS_API_H__ +#define __INFINIOP_OPS_SOFTPLUS_API_H__ +#include "../tensor_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSoftplusDescriptor_t; +__INFINI_C __export infiniStatus_t infiniopCreateSoftplusDescriptor( + infiniopHandle_t handle, + infiniopSoftplusDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + float beta, + float threshold); + +__INFINI_C __export infiniStatus_t infiniopGetSoftplusWorkspaceSize( + infiniopSoftplusDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSoftplus( + infiniopSoftplusDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySoftplusDescriptor( + infiniopSoftplusDescriptor_t desc); + +#endif // __INFINIOP_OPS_SOFTPLUS_API_H__ diff --git a/python/infinicore/include/infiniop/ops/softsign.h b/python/infinicore/include/infiniop/ops/softsign.h new file mode 100644 index 000000000..62759dc54 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/softsign.h @@ -0,0 +1,25 @@ + +#ifndef __INFINIOP_SOFTSIGN_API_H__ +#define __INFINIOP_SOFTSIGN_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSoftsignDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSoftsignDescriptor(infiniopHandle_t handle, + infiniopSoftsignDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetSoftsignWorkspaceSize(infiniopSoftsignDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSoftsign(infiniopSoftsignDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySoftsignDescriptor(infiniopSoftsignDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/sub.h b/python/infinicore/include/infiniop/ops/sub.h new file mode 100644 index 000000000..3a516c30c --- /dev/null +++ b/python/infinicore/include/infiniop/ops/sub.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_SUB_API_H__ +#define __INFINIOP_SUB_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSubDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSubDescriptor(infiniopHandle_t handle, + infiniopSubDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__INFINI_C __export infiniStatus_t infiniopGetSubWorkspaceSize(infiniopSubDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSub(infiniopSubDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySubDescriptor(infiniopSubDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/sum.h b/python/infinicore/include/infiniop/ops/sum.h new file mode 100644 index 000000000..c97104c90 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/sum.h @@ -0,0 +1,31 @@ +#ifndef __INFINIOP_SUM_API_H__ +#define __INFINIOP_SUM_API_H__ + +#include "../operator_descriptor.h" +#include +#include +typedef struct InfiniopDescriptor *infiniopSumDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSumDescriptor(infiniopHandle_t handle, + infiniopSumDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + size_t *dim, + size_t dim_size, + bool keepdim); + +__INFINI_C __export infiniStatus_t infiniopGetSumWorkspaceSize(infiniopSumDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSum(infiniopSumDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + size_t *dim, + size_t dim_size, + bool keepdim, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySumDescriptor(infiniopSumDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/swap.h b/python/infinicore/include/infiniop/ops/swap.h new file mode 100644 index 000000000..7eb14b57a --- /dev/null +++ b/python/infinicore/include/infiniop/ops/swap.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_SWAP_API_H__ +#define __INFINIOP_SWAP_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSwapDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSwapDescriptor(infiniopHandle_t handle, + infiniopSwapDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t y); + +__INFINI_C __export infiniStatus_t infiniopGetSwapWorkspaceSize(infiniopSwapDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSwap(infiniopSwapDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *x, + void *y, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySwapDescriptor(infiniopSwapDescriptor_t desc); + +#endif // __INFINIOP_SWAP_API_H__ diff --git a/python/infinicore/include/infiniop/ops/swiglu.h b/python/infinicore/include/infiniop/ops/swiglu.h new file mode 100644 index 000000000..e11749c6b --- /dev/null +++ b/python/infinicore/include/infiniop/ops/swiglu.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_SWIGLU_API_H__ +#define __INFINIOP_SWIGLU_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopSwiGLUDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateSwiGLUDescriptor(infiniopHandle_t handle, + infiniopSwiGLUDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc); + +__INFINI_C __export infiniStatus_t infiniopGetSwiGLUWorkspaceSize(infiniopSwiGLUDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopSwiGLU(infiniopSwiGLUDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + void const *a, + void const *b, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroySwiGLUDescriptor(infiniopSwiGLUDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/take.h b/python/infinicore/include/infiniop/ops/take.h new file mode 100644 index 000000000..58fb15311 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/take.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_TAKE_API_H__ +#define __INFINIOP_TAKE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopTakeDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateTakeDescriptor(infiniopHandle_t handle, + infiniopTakeDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + infiniopTensorDescriptor_t indices); + +__INFINI_C __export infiniStatus_t infiniopGetTakeWorkspaceSize(infiniopTakeDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopTake(infiniopTakeDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + const void *indices, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyTakeDescriptor(infiniopTakeDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/tan.h b/python/infinicore/include/infiniop/ops/tan.h new file mode 100644 index 000000000..1b7bef23e --- /dev/null +++ b/python/infinicore/include/infiniop/ops/tan.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_TAN_API_H__ +#define __INFINIOP_TAN_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopTanDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateTanDescriptor( + infiniopHandle_t handle, + infiniopTanDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t input_desc, + infiniopTensorDescriptor_t out_desc); + +__INFINI_C __export infiniStatus_t infiniopGetTanWorkspaceSize(infiniopTanDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopTan( + infiniopTanDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyTanDescriptor(infiniopTanDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/tanh.h b/python/infinicore/include/infiniop/ops/tanh.h new file mode 100644 index 000000000..e623d854c --- /dev/null +++ b/python/infinicore/include/infiniop/ops/tanh.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_TANH_API_H__ +#define __INFINIOP_TANH_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopTanhDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateTanhDescriptor(infiniopHandle_t handle, + infiniopTanhDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input); + +__INFINI_C __export infiniStatus_t infiniopGetTanhWorkspaceSize(infiniopTanhDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopTanh(infiniopTanhDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyTanhDescriptor(infiniopTanhDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/tanhshrink.h b/python/infinicore/include/infiniop/ops/tanhshrink.h new file mode 100644 index 000000000..b2895ccab --- /dev/null +++ b/python/infinicore/include/infiniop/ops/tanhshrink.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_TANHSHRINK_API_H__ +#define __INFINIOP_TANHSHRINK_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopTanhshrinkDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateTanhshrinkDescriptor(infiniopHandle_t handle, + infiniopTanhshrinkDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t intput); + +__INFINI_C __export infiniStatus_t infiniopGetTanhshrinkWorkspaceSize(infiniopTanhshrinkDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopTanhshrink(infiniopTanhshrinkDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *intput, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyTanhshrinkDescriptor(infiniopTanhshrinkDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/topk.h b/python/infinicore/include/infiniop/ops/topk.h new file mode 100644 index 000000000..3eaf94289 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/topk.h @@ -0,0 +1,35 @@ +#ifndef __INFINIOP_TOPK_API_H__ +#define __INFINIOP_TOPK_API_H__ + +#include "../operator_descriptor.h" +#include +#include +typedef struct InfiniopDescriptor *infiniopTopKDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateTopKDescriptor(infiniopHandle_t handle, + infiniopTopKDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t values_output_desc, + infiniopTensorDescriptor_t indices_output_desc, + infiniopTensorDescriptor_t input_desc, + size_t k, + size_t dim, + bool largest, + bool sorted); + +__INFINI_C __export infiniStatus_t infiniopGetTopKWorkspaceSize(infiniopTopKDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopTopK(infiniopTopKDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *values_output, + void *indices_output, + const void *input, + size_t k, + size_t dim, + bool largest, + bool sorted, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyTopKDescriptor(infiniopTopKDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/topkrouter.h b/python/infinicore/include/infiniop/ops/topkrouter.h new file mode 100644 index 000000000..1d56a17cb --- /dev/null +++ b/python/infinicore/include/infiniop/ops/topkrouter.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_TOPKRouter_API_H__ +#define __INFINIOP_TOPKRouter_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopTopkrouterDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateTopkrouterDescriptor(infiniopHandle_t handle, + infiniopTopkrouterDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t correction_bias_desc); + +__INFINI_C __export infiniStatus_t infiniopGetTopkrouterWorkspaceSize(infiniopTopkrouterDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopTopkrouter(infiniopTopkrouterDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *values, + void *indices, + const void *x, + const void *correction_bias, + const float routed_scaling_factor, + const size_t topk, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyTopkrouterDescriptor(infiniopTopkrouterDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/topksoftmax.h b/python/infinicore/include/infiniop/ops/topksoftmax.h new file mode 100644 index 000000000..74ea10d69 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/topksoftmax.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_TOPKSOFTMAX_API_H__ +#define __INFINIOP_TOPKSOFTMAX_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopTopksoftmaxDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateTopksoftmaxDescriptor(infiniopHandle_t handle, + infiniopTopksoftmaxDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x_desc); + +__INFINI_C __export infiniStatus_t infiniopGetTopksoftmaxWorkspaceSize(infiniopTopksoftmaxDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopTopksoftmax(infiniopTopksoftmaxDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *values, + void *indices, + const void *x, + const size_t topk, + const int norm, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyTopksoftmaxDescriptor(infiniopTopksoftmaxDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/triplet_margin_loss.h b/python/infinicore/include/infiniop/ops/triplet_margin_loss.h new file mode 100644 index 000000000..a96dca22d --- /dev/null +++ b/python/infinicore/include/infiniop/ops/triplet_margin_loss.h @@ -0,0 +1,33 @@ +#ifndef __INFINIOP_TRIPLET_MARGIN_LOSS_API_H__ +#define __INFINIOP_TRIPLET_MARGIN_LOSS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopTripletMarginLossDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateTripletMarginLossDescriptor(infiniopHandle_t handle, + infiniopTripletMarginLossDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t anchor, + infiniopTensorDescriptor_t positive, + infiniopTensorDescriptor_t negative, + float margin, + int p, + float eps, + int swap, // 0: False, 1: True + int reduction); // 0: None, 1: Mean, 2: Sum + +__INFINI_C __export infiniStatus_t infiniopGetTripletMarginLossWorkspaceSize(infiniopTripletMarginLossDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopTripletMarginLoss(infiniopTripletMarginLossDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *anchor, + const void *positive, + const void *negative, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyTripletMarginLossDescriptor(infiniopTripletMarginLossDescriptor_t desc); + +#endif // __INFINIOP_TRIPLET_MARGIN_LOSS_API_H__ diff --git a/python/infinicore/include/infiniop/ops/triplet_margin_with_distance_loss.h b/python/infinicore/include/infiniop/ops/triplet_margin_with_distance_loss.h new file mode 100644 index 000000000..c834f32f5 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/triplet_margin_with_distance_loss.h @@ -0,0 +1,32 @@ +#ifndef __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__ +#define __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopTripletMarginWithDistanceLossDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateTripletMarginWithDistanceLossDescriptor( + infiniopHandle_t handle, + infiniopTripletMarginWithDistanceLossDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t anchor, + infiniopTensorDescriptor_t positive, + infiniopTensorDescriptor_t negative, + float margin, + int swap, + int reduction); +__INFINI_C __export infiniStatus_t infiniopGetTripletMarginWithDistanceLossWorkspaceSize( + infiniopTripletMarginWithDistanceLossDescriptor_t desc, + size_t *size); +__INFINI_C __export infiniStatus_t infiniopTripletMarginWithDistanceLoss(infiniopTripletMarginWithDistanceLossDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *anchor, + const void *positive, + const void *negative, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyTripletMarginWithDistanceLossDescriptor( + infiniopTripletMarginWithDistanceLossDescriptor_t desc); +#endif // __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__ diff --git a/python/infinicore/include/infiniop/ops/unfold.h b/python/infinicore/include/infiniop/ops/unfold.h new file mode 100644 index 000000000..0d1688978 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/unfold.h @@ -0,0 +1,28 @@ +#ifndef __INFINIOP_UNFOLD_API_H__ +#define __INFINIOP_UNFOLD_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopUnfoldDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateUnfoldDescriptor(infiniopHandle_t handle, + infiniopUnfoldDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + const int *kernel_sizes, + const int *strides, + const int *paddings, + const int *dilations); + +// 获取 Unfold 工作区大小 +__INFINI_C __export infiniStatus_t infiniopGetUnfoldWorkspaceSize(infiniopUnfoldDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopUnfold(infiniopUnfoldDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); +__INFINI_C __export infiniStatus_t infiniopDestroyUnfoldDescriptor(infiniopUnfoldDescriptor_t desc); + +#endif // __INFINIOP_UNFOLD_API_H__ diff --git a/python/infinicore/include/infiniop/ops/upsample_bilinear.h b/python/infinicore/include/infiniop/ops/upsample_bilinear.h new file mode 100644 index 000000000..c53a5b24a --- /dev/null +++ b/python/infinicore/include/infiniop/ops/upsample_bilinear.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_UPSAMPLE_BILINEAR_API_H__ +#define __INFINIOP_UPSAMPLE_BILINEAR_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopUpsampleBilinearDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateUpsampleBilinearDescriptor(infiniopHandle_t handle, + infiniopUpsampleBilinearDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + int align_corners); + +__INFINI_C __export infiniStatus_t infiniopGetUpsampleBilinearWorkspaceSize(infiniopUpsampleBilinearDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopUpsampleBilinear(infiniopUpsampleBilinearDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyUpsampleBilinearDescriptor(infiniopUpsampleBilinearDescriptor_t desc); + +#endif // __INFINIOP_UPSAMPLE_BILINEAR_API_H__ diff --git a/python/infinicore/include/infiniop/ops/upsample_nearest.h b/python/infinicore/include/infiniop/ops/upsample_nearest.h new file mode 100644 index 000000000..de64badcc --- /dev/null +++ b/python/infinicore/include/infiniop/ops/upsample_nearest.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_UPSAMPLE_NEAREST_API_H__ +#define __INFINIOP_UPSAMPLE_NEAREST_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopUpsampleNearestDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateUpsampleNearestDescriptor(infiniopHandle_t handle, + infiniopUpsampleNearestDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input); + +__INFINI_C __export infiniStatus_t infiniopGetUpsampleNearestWorkspaceSize(infiniopUpsampleNearestDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopUpsampleNearest(infiniopUpsampleNearestDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyUpsampleNearestDescriptor(infiniopUpsampleNearestDescriptor_t desc); + +#endif // __INFINIOP_UPSAMPLE_NEAREST_API_H__ diff --git a/python/infinicore/include/infiniop/ops/vander.h b/python/infinicore/include/infiniop/ops/vander.h new file mode 100644 index 000000000..72773da5b --- /dev/null +++ b/python/infinicore/include/infiniop/ops/vander.h @@ -0,0 +1,25 @@ +#ifndef __INFINIOP_VANDER_API_H__ +#define __INFINIOP_VANDER_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopVanderDescriptor_t; +__INFINI_C __export infiniStatus_t infiniopCreateVanderDescriptor(infiniopHandle_t handle, + infiniopVanderDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + int N, + int increasing); + +__INFINI_C __export infiniStatus_t infiniopGetVanderWorkspaceSize(infiniopVanderDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopVander(infiniopVanderDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyVanderDescriptor(infiniopVanderDescriptor_t desc); + +#endif // __INFINIOP_VANDER_API_H__ diff --git a/python/infinicore/include/infiniop/ops/var.h b/python/infinicore/include/infiniop/ops/var.h new file mode 100644 index 000000000..7dc601a94 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/var.h @@ -0,0 +1,33 @@ +#ifndef __INFINIOP_VAR_API_H__ +#define __INFINIOP_VAR_API_H__ + +#include "../operator_descriptor.h" +#include +#include +typedef struct InfiniopDescriptor *infiniopVarDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateVarDescriptor(infiniopHandle_t handle, + infiniopVarDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t var_output_desc, + infiniopTensorDescriptor_t input_desc, + size_t *dim, + size_t dim_size, + bool unbiased, + bool keepdim); + +__INFINI_C __export infiniStatus_t infiniopGetVarWorkspaceSize(infiniopVarDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopVar(infiniopVarDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *var_output, + const void *input, + size_t *dim, + size_t dim_size, + bool unbiased, + bool keepdim, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyVarDescriptor(infiniopVarDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/var_mean.h b/python/infinicore/include/infiniop/ops/var_mean.h new file mode 100644 index 000000000..358a55636 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/var_mean.h @@ -0,0 +1,35 @@ +#ifndef __INFINIOP_VAR_MEAN_API_H__ +#define __INFINIOP_VAR_MEAN_API_H__ + +#include "../operator_descriptor.h" +#include +#include +typedef struct InfiniopDescriptor *infiniopVarMeanDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateVarMeanDescriptor(infiniopHandle_t handle, + infiniopVarMeanDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t var_output_desc, + infiniopTensorDescriptor_t mean_output_desc, + infiniopTensorDescriptor_t input_desc, + size_t *dim, + size_t dim_size, + bool unbiased, + bool keepdim); + +__INFINI_C __export infiniStatus_t infiniopGetVarMeanWorkspaceSize(infiniopVarMeanDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopVarMean(infiniopVarMeanDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *var_output, + void *mean_output, + const void *input, + size_t *dim, + size_t dim_size, + bool unbiased, + bool keepdim, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyVarMeanDescriptor(infiniopVarMeanDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/ops/zeros.h b/python/infinicore/include/infiniop/ops/zeros.h new file mode 100644 index 000000000..5d7ce03d5 --- /dev/null +++ b/python/infinicore/include/infiniop/ops/zeros.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_ZEROS_API_H__ +#define __INFINIOP_ZEROS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopZerosDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateZerosDescriptor(infiniopHandle_t handle, + infiniopZerosDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x); + +__INFINI_C __export infiniStatus_t infiniopGetZerosWorkspaceSize(infiniopZerosDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopZeros(infiniopZerosDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyZerosDescriptor(infiniopZerosDescriptor_t desc); + +#endif diff --git a/python/infinicore/include/infiniop/tensor_descriptor.h b/python/infinicore/include/infiniop/tensor_descriptor.h new file mode 100644 index 000000000..8390da81c --- /dev/null +++ b/python/infinicore/include/infiniop/tensor_descriptor.h @@ -0,0 +1,14 @@ +#ifndef __INFINIOP_TENSOR_DESCRIPTOR_API_H__ +#define __INFINIOP_TENSOR_DESCRIPTOR_API_H__ + +#include "../infinicore.h" + +struct InfiniopTensorDescriptor; + +typedef struct InfiniopTensorDescriptor *infiniopTensorDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateTensorDescriptor(infiniopTensorDescriptor_t *desc_ptr, size_t ndim, const size_t *shape, const ptrdiff_t *strides, infiniDtype_t dtype); + +__INFINI_C __export infiniStatus_t infiniopDestroyTensorDescriptor(infiniopTensorDescriptor_t desc); + +#endif // __INFINIOP_TENSOR_DESCRIPTOR__ diff --git a/python/infinicore/include/infinirt.h b/python/infinicore/include/infinirt.h new file mode 100644 index 000000000..4ada848f5 --- /dev/null +++ b/python/infinicore/include/infinirt.h @@ -0,0 +1,134 @@ +#ifndef __INFINIRT_API_H__ +#define __INFINIRT_API_H__ + +#include "infinicore.h" +#include + +typedef void *infinirtStream_t; +typedef void *infinirtEvent_t; +typedef void *infinirtGraph_t; +typedef void *infinirtGraphNode_t; +typedef void *infinirtGraphExec_t; + +// Bitmask describing which fields of an infinirtDeviceResourceSnapshot_t +// have been populated by the backend. Backends without a particular +// capability simply leave the flag clear. +typedef enum { + INFINIRT_RESOURCE_FIELD_NONE = 0, + INFINIRT_RESOURCE_FIELD_MEMORY_CAPACITY = 1ull << 0, + INFINIRT_RESOURCE_FIELD_COMPUTE_UTILIZATION = 1ull << 1, + INFINIRT_RESOURCE_FIELD_MEMORY_BANDWIDTH_UTILIZATION = 1ull << 2, + INFINIRT_RESOURCE_FIELD_KERNEL_TIME_RATIO = 1ull << 3, + INFINIRT_RESOURCE_FIELD_COMMUNICATION = 1ull << 4, +} infinirtResourceField_t; + +// Vendor-neutral resource snapshot consumed by the optional mutual +// awareness analyzer. Backends populate the fields they can observe and +// set the matching `valid_fields` bits; consumers must check the bits +// before relying on a value. +typedef struct { + infiniDevice_t device_type; + int device_id; + + uint64_t valid_fields; + uint64_t estimated_fields; + + size_t free_bytes; + size_t total_bytes; + size_t used_bytes; + size_t reserved_bytes; + + float compute_utilization; + float memory_bandwidth_utilization; + float kernel_time_ratio; + + float communication_time_ratio; + uint64_t communication_bytes; +} infinirtDeviceResourceSnapshot_t; + +__INFINI_C __export infiniStatus_t infinirtInit(); + +// Device +__INFINI_C __export infiniStatus_t infinirtGetAllDeviceCount(int *count_array); +__INFINI_C __export infiniStatus_t infinirtGetDeviceCount(infiniDevice_t device, int *count); +__INFINI_C __export infiniStatus_t infinirtSetDevice(infiniDevice_t device, int device_id); +__INFINI_C __export infiniStatus_t infinirtGetDevice(infiniDevice_t *device_ptr, int *device_id_ptr); +__INFINI_C __export infiniStatus_t infinirtGetMemInfo(infiniDevice_t device, int device_id, size_t *free_bytes, size_t *total_bytes); +__INFINI_C __export infiniStatus_t infinirtGetDeviceResourceSnapshot( + infiniDevice_t device, + int device_id, + infinirtDeviceResourceSnapshot_t *snapshot); +__INFINI_C __export infiniStatus_t infinirtDeviceSynchronize(); + +// Stream +__INFINI_C __export infiniStatus_t infinirtStreamCreate(infinirtStream_t *stream_ptr); +__INFINI_C __export infiniStatus_t infinirtStreamDestroy(infinirtStream_t stream); +__INFINI_C __export infiniStatus_t infinirtStreamSynchronize(infinirtStream_t stream); +__INFINI_C __export infiniStatus_t infinirtStreamWaitEvent(infinirtStream_t stream, infinirtEvent_t event); + +// Event +typedef enum { + INFINIRT_EVENT_COMPLETE = 0, + INFINIRT_EVENT_NOT_READY = 1, +} infinirtEventStatus_t; + +// Event flags for precise timing +typedef enum { + INFINIRT_EVENT_DEFAULT = 0x0, // Default event creation flags + INFINIRT_EVENT_DISABLE_TIMING = 0x1, // Event will not record timing data + INFINIRT_EVENT_BLOCKING_SYNC = 0x2, // Event uses blocking synchronization +} infinirtEventFlags_t; + +__INFINI_C __export infiniStatus_t infinirtEventCreate(infinirtEvent_t *event_ptr); +__INFINI_C __export infiniStatus_t infinirtEventCreateWithFlags(infinirtEvent_t *event_ptr, uint32_t flags); +__INFINI_C __export infiniStatus_t infinirtEventRecord(infinirtEvent_t event, infinirtStream_t stream); +__INFINI_C __export infiniStatus_t infinirtEventQuery(infinirtEvent_t event, infinirtEventStatus_t *status_ptr); +__INFINI_C __export infiniStatus_t infinirtEventSynchronize(infinirtEvent_t event); +__INFINI_C __export infiniStatus_t infinirtEventDestroy(infinirtEvent_t event); +__INFINI_C __export infiniStatus_t infinirtEventElapsedTime(float *ms_ptr, infinirtEvent_t start, infinirtEvent_t end); + +// Memory +typedef enum { + INFINIRT_MEMCPY_H2H = 0, + INFINIRT_MEMCPY_H2D = 1, + INFINIRT_MEMCPY_D2H = 2, + INFINIRT_MEMCPY_D2D = 3, +} infinirtMemcpyKind_t; + +__INFINI_C __export infiniStatus_t infinirtMalloc(void **p_ptr, size_t size); +__INFINI_C __export infiniStatus_t infinirtMallocHost(void **p_ptr, size_t size); +__INFINI_C __export infiniStatus_t infinirtFree(void *ptr); +__INFINI_C __export infiniStatus_t infinirtFreeHost(void *ptr); + +__INFINI_C __export infiniStatus_t infinirtMemcpy(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind); +__INFINI_C __export infiniStatus_t infinirtMemcpyAsync(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind, infinirtStream_t stream); + +// Stream-ordered memory +__INFINI_C __export infiniStatus_t infinirtMallocAsync(void **p_ptr, size_t size, infinirtStream_t stream); +__INFINI_C __export infiniStatus_t infinirtFreeAsync(void *ptr, infinirtStream_t stream); + +// Memset +__INFINI_C __export infiniStatus_t infinirtMemset(void *ptr, int value, size_t count); +__INFINI_C __export infiniStatus_t infinirtMemsetAsync(void *ptr, int value, size_t count, infinirtStream_t stream); + +// Graph +typedef enum { + INFINIRT_STREAM_CAPTURE_MODE_GLOBAL = 0, + INFINIRT_STREAM_CAPTURE_MODE_THREAD_LOCAL = 1, + INFINIRT_STREAM_CAPTURE_MODE_RELAXED = 2, + +} infinirtStreamCaptureMode_t; + +__INFINI_C __export infiniStatus_t infinirtStreamBeginCapture(infinirtStream_t stream, infinirtStreamCaptureMode_t mode); +__INFINI_C __export infiniStatus_t infinirtStreamEndCapture(infinirtStream_t stream, infinirtGraph_t *graph_ptr); +__INFINI_C __export infiniStatus_t infinirtGraphDestroy(infinirtGraph_t graph); +__INFINI_C __export infiniStatus_t infinirtGraphInstantiate( + infinirtGraphExec_t *graph_exec_ptr, + infinirtGraph_t graph, + infinirtGraphNode_t *node_ptr, + char *log_buffer, + size_t buffer_size); +__INFINI_C __export infiniStatus_t infinirtGraphExecDestroy(infinirtGraphExec_t graph_exec); +__INFINI_C __export infiniStatus_t infinirtGraphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream); + +#endif // __INFINIRT_API_H__ diff --git a/python/infinicore/ops/dsa.py b/python/infinicore/ops/dsa.py new file mode 100644 index 000000000..610e2f7f7 --- /dev/null +++ b/python/infinicore/ops/dsa.py @@ -0,0 +1,151 @@ +from infinicore.lib import _infinicore + + +def _t(tensor): + return tensor._underlying + + +def fused_deepseek_v2_indexer_postprocess_( + q_out, + k_out, + weights_out, + kv_cache, + slot_mapping, + q, + kw, + norm_weight, + norm_bias, + positions, + cos_sin_cache, + num_cache_tokens, + is_neox, + eps, + weights_scale, +): + _infinicore.fused_deepseek_v2_indexer_postprocess_( + *map( + _t, + ( + q_out, + k_out, + weights_out, + kv_cache, + slot_mapping, + q, + kw, + norm_weight, + norm_bias, + positions, + cos_sin_cache, + ), + ), + num_cache_tokens, + is_neox, + eps, + weights_scale, + ) + + +def indexer_k_cache_(k, kv_cache, slot_mapping): + _infinicore.indexer_k_cache_(_t(k), _t(kv_cache), _t(slot_mapping)) + + +def compute_block_sparse_mqa_logits_( + logits, + q, + kv_cache, + cu_seqlens_q, + cu_seqlens_kv, + block_table, + weights, + max_q_len, + max_kv_len, + max_context_len, +): + _infinicore.compute_block_sparse_mqa_logits_( + *map( + _t, + ( + logits, + q, + kv_cache, + cu_seqlens_q, + cu_seqlens_kv, + block_table, + weights, + ), + ), + max_q_len, + max_kv_len, + max_context_len, + ) + + +def select_prefill_topk_block_indices_( + topk_indices, logits, cu_seqlen_ks, cu_seqlen_ke +): + _infinicore.select_prefill_topk_block_indices_( + _t(topk_indices), _t(logits), _t(cu_seqlen_ks), _t(cu_seqlen_ke) + ) + + +def select_decode_topk_block_indices_(topk_indices, logits, seq_lens): + _infinicore.select_decode_topk_block_indices_( + _t(topk_indices), _t(logits), _t(seq_lens) + ) + + +def map_prefill_request_block_indices_( + output, + req_id, + block_table, + token_indices, + block_size, + has_prefill_workspace=False, + prefill_workspace_request_ids=None, + prefill_workspace_starts=None, +): + _infinicore.map_prefill_request_block_indices_( + _t(output), + _t(req_id), + _t(block_table), + _t(token_indices), + block_size, + has_prefill_workspace, + None + if prefill_workspace_request_ids is None + else _t(prefill_workspace_request_ids), + None if prefill_workspace_starts is None else _t(prefill_workspace_starts), + ) + + +def map_decode_request_block_indices_( + output, req_id, block_table, token_indices, block_size +): + _infinicore.map_decode_request_block_indices_( + _t(output), _t(req_id), _t(block_table), _t(token_indices), block_size + ) + + +def topk_indices_context_lens_(topk_lens, indices): + _infinicore.topk_indices_context_lens_(_t(topk_lens), _t(indices)) + + +def sparse_flash_mla_( + output, + query, + kv_cache, + indices, + topk_lens, + scale, + attn_sink=None, +): + _infinicore.sparse_flash_mla_( + _t(output), + _t(query), + _t(kv_cache), + _t(indices), + _t(topk_lens), + scale, + None if attn_sink is None else _t(attn_sink), + ) diff --git a/python/infinicore/utils.py b/python/infinicore/utils.py index fd39db0e1..b6dbc6789 100644 --- a/python/infinicore/utils.py +++ b/python/infinicore/utils.py @@ -13,6 +13,8 @@ def to_torch_dtype(infini_dtype): """Convert infinicore data type to PyTorch data type""" if infini_dtype == infinicore.float16: return torch.float16 + elif infini_dtype == infinicore.float8: + return torch.float8_e4m3fn elif infini_dtype == infinicore.float32: return torch.float32 elif infini_dtype == infinicore.float64: @@ -39,6 +41,8 @@ def to_infinicore_dtype(torch_dtype): """Convert PyTorch data type to infinicore data type""" if torch_dtype == torch.float32: return infinicore.float32 + elif torch_dtype == torch.float8_e4m3fn: + return infinicore.float8 elif torch_dtype == torch.float16: return infinicore.float16 elif torch_dtype == torch.bfloat16: diff --git a/src/infiniccl-test/infiniccl_test.cpp b/src/infiniccl-test/infiniccl_test.cpp index 0aa898484..a4b3d330d 100644 --- a/src/infiniccl-test/infiniccl_test.cpp +++ b/src/infiniccl-test/infiniccl_test.cpp @@ -1,6 +1,8 @@ #include "infiniccl_test.hpp" #include +#include +#include #include #include #include @@ -20,7 +22,12 @@ const size_t TEST_COUNTS[] = { MAX_COUNT, }; -const infiniDtype_t TEST_DTYPES[] = {INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16}; +const infiniDtype_t TEST_DTYPES[] = { + INFINI_DTYPE_F32, + INFINI_DTYPE_F16, + INFINI_DTYPE_BF16, + INFINI_DTYPE_I32, +}; const size_t WARM_UPS = 10; @@ -41,6 +48,12 @@ struct ThreadArgs { void setData(infiniDtype_t dtype, void *data, size_t count, float val) { switch (dtype) { + case INFINI_DTYPE_I32: + for (size_t i = 0; i < count; i++) { + static_cast(data)[i] = static_cast(val); + } + break; + case INFINI_DTYPE_F32: for (size_t i = 0; i < count; i++) { ((float *)data)[i] = val; @@ -90,6 +103,8 @@ int checkData(const T *actual_, const T *expected_, size_t count) { int checkData(const void *actual, const void *expected, infiniDtype_t dtype, size_t count) { switch (dtype) { + case INFINI_DTYPE_I32: + return checkData(static_cast(actual), static_cast(expected), count); case INFINI_DTYPE_F32: return checkData((const float *)actual, (const float *)expected, count); case INFINI_DTYPE_F16: @@ -160,10 +175,20 @@ int testAllReduce(infiniDevice_t device_type, int ndevice) { device_ids[i] = i; } + const char *dtype_filter = std::getenv("INFINICCL_TEST_DTYPE"); + const char *count_filter = std::getenv("INFINICCL_TEST_COUNT"); + for (infiniDtype_t dtype : TEST_DTYPES) { + if (dtype_filter != nullptr && infiniDtypeToString(dtype) != dtype_filter) { + continue; + } setData(dtype, data, MAX_COUNT, 1.0f); setData(dtype, ans, MAX_COUNT, 1.0f * ndevice); for (size_t count : TEST_COUNTS) { + if (count_filter != nullptr + && count != static_cast(std::stoull(count_filter))) { + continue; + } TEST_INFINI(infinicclCommInitAll(device_type, comms.data(), ndevice, device_ids.data())); std::cout << "Testing AllReduce with " << count << " elements of " << infiniDtypeToString(dtype) << std::endl; for (int rank = 0; rank < ndevice; rank++) { diff --git a/src/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu index 6a8442d21..1f3b40d8e 100644 --- a/src/infiniccl/cuda/infiniccl_cuda.cu +++ b/src/infiniccl/cuda/infiniccl_cuda.cu @@ -9,6 +9,13 @@ #define CHECK_NCCL(API__) CHECK_INTERNAL(API__, ncclSuccess) +// CoreX exports NCCL-compatible P2P entry points but omits their declarations +// from its nccl.h. Keep the ABI declarations local to this backend. +extern "C" ncclResult_t ncclSend( + const void *, size_t, ncclDataType_t, int, ncclComm_t, cudaStream_t); +extern "C" ncclResult_t ncclRecv( + void *, size_t, ncclDataType_t, int, ncclComm_t, cudaStream_t); + inline cudaStream_t getCudaStream(infinirtStream_t stream) { if (stream == nullptr) { return 0; @@ -93,6 +100,40 @@ infiniStatus_t groupEnd(infinicclComm_t) { return INFINI_STATUS_SUCCESS; } +infiniStatus_t send( + const void *sendbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream) { + CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, + INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, + INFINI_DTYPE_U32, INFINI_DTYPE_U64); + CHECK_OR_DO(peer >= 0 && peer < comm->world_size, + return INFINI_STATUS_BAD_PARAM); + CHECK_NCCL(ncclSend(sendbuf, count, getNcclDtype(datatype), peer, + getNcclComm(comm), getCudaStream(stream))); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t recv( + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream) { + CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, + INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, + INFINI_DTYPE_U32, INFINI_DTYPE_U64); + CHECK_OR_DO(peer >= 0 && peer < comm->world_size, + return INFINI_STATUS_BAD_PARAM); + CHECK_NCCL(ncclRecv(recvbuf, count, getNcclDtype(datatype), peer, + getNcclComm(comm), getCudaStream(stream))); + return INFINI_STATUS_SUCCESS; +} + infiniStatus_t allReduce( void *sendbuf, void *recvbuf, @@ -102,7 +143,8 @@ infiniStatus_t allReduce( infinicclComm_t comm, infinirtStream_t stream) { - CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); + CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, + INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_U32, INFINI_DTYPE_U64); CHECK_NCCL(ncclAllReduce(sendbuf, recvbuf, count, getNcclDtype(datatype), getNcclRedOp(op), getNcclComm(comm), getCudaStream(stream))); diff --git a/src/infiniccl/infiniccl.cc b/src/infiniccl/infiniccl.cc index 48dbea425..32ece43d5 100644 --- a/src/infiniccl/infiniccl.cc +++ b/src/infiniccl/infiniccl.cc @@ -115,6 +115,52 @@ __INFINI_C infiniStatus_t infinicclGroupEnd(infinicclComm_t comm) { #undef GROUP_END } +__INFINI_C infiniStatus_t infinicclSend( + const void *sendbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream) { + if (comm == nullptr || sendbuf == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + switch (comm->device_type) { + case INFINI_DEVICE_NVIDIA: + case INFINI_DEVICE_ILUVATAR: + case INFINI_DEVICE_QY: + case INFINI_DEVICE_HYGON: + case INFINI_DEVICE_ALI: + return infiniccl::cuda::send( + sendbuf, count, datatype, peer, comm, stream); + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +} + +__INFINI_C infiniStatus_t infinicclRecv( + void *recvbuf, + size_t count, + infiniDtype_t datatype, + int peer, + infinicclComm_t comm, + infinirtStream_t stream) { + if (comm == nullptr || recvbuf == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + switch (comm->device_type) { + case INFINI_DEVICE_NVIDIA: + case INFINI_DEVICE_ILUVATAR: + case INFINI_DEVICE_QY: + case INFINI_DEVICE_HYGON: + case INFINI_DEVICE_ALI: + return infiniccl::cuda::recv( + recvbuf, count, datatype, peer, comm, stream); + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +} + __INFINI_C infiniStatus_t infinicclAllReduce( void *sendbuf, void *recvbuf, diff --git a/src/infiniccl/infiniccl_impl.h b/src/infiniccl/infiniccl_impl.h index 790355bd9..480b5ab39 100644 --- a/src/infiniccl/infiniccl_impl.h +++ b/src/infiniccl/infiniccl_impl.h @@ -24,6 +24,22 @@ struct InfinicclComm { \ infiniStatus_t groupEnd(infinicclComm_t comm) IMPL; \ \ + infiniStatus_t send( \ + const void *sendbuf, \ + size_t count, \ + infiniDtype_t datatype, \ + int peer, \ + infinicclComm_t comm, \ + infinirtStream_t stream) IMPL; \ + \ + infiniStatus_t recv( \ + void *recvbuf, \ + size_t count, \ + infiniDtype_t datatype, \ + int peer, \ + infinicclComm_t comm, \ + infinirtStream_t stream) IMPL; \ + \ infiniStatus_t allReduce( \ void *sendbuf, \ void *recvbuf, \ diff --git a/src/infinicore/adaptor/aten_adaptor.cc b/src/infinicore/adaptor/aten_adaptor.cc index 1d041ebbc..f020d3af3 100644 --- a/src/infinicore/adaptor/aten_adaptor.cc +++ b/src/infinicore/adaptor/aten_adaptor.cc @@ -2,12 +2,50 @@ #include "infinicore/adaptor/aten_adaptor.hpp" #include +#include +#include #if defined(ENABLE_ILUVATAR_API) extern "C" int32_t torch_set_current_cuda_stream(void *stream, int32_t device_index); #endif namespace infinicore::adaptor { +namespace { +struct AtenTensorCacheKey { + void *data; + std::vector sizes; + std::vector strides; + int dtype; + int device_type; + int device_index; + + bool operator==(const AtenTensorCacheKey &other) const { + return data == other.data && sizes == other.sizes + && strides == other.strides && dtype == other.dtype + && device_type == other.device_type + && device_index == other.device_index; + } +}; + +struct AtenTensorCacheKeyHash { + size_t operator()(const AtenTensorCacheKey &key) const { + size_t hash = std::hash{}(key.data); + const auto combine = [&hash](size_t value) { + hash ^= value + 0x9e3779b97f4a7c15ULL + (hash << 6) + (hash >> 2); + }; + combine(std::hash{}(key.dtype)); + combine(std::hash{}(key.device_type)); + combine(std::hash{}(key.device_index)); + for (const auto value : key.sizes) { + combine(std::hash{}(value)); + } + for (const auto value : key.strides) { + combine(std::hash{}(value)); + } + return hash; + } +}; +} // namespace at::Tensor to_aten_tensor(const infinicore::Tensor &t) { void *data_ptr = (void *)(t->data()); @@ -16,7 +54,31 @@ at::Tensor to_aten_tensor(const infinicore::Tensor &t) { t->shape().begin(), t->shape().end()); - auto strides = t->strides(); + const auto tensor_strides = t->strides(); + auto strides = std::vector(tensor_strides.begin(), tensor_strides.end()); + + // vLLM keeps ATen tensors alive across calls. InfiniCore tensors are + // non-owning from the ATen point of view, so cache wrappers per worker + // thread and reuse them only when address and complete metadata match. + const bool cache_wrapper = t->numel() != 0; + static thread_local std::unordered_map + wrapper_cache; + AtenTensorCacheKey cache_key{ + data_ptr, + sizes, + strides, + static_cast(t->dtype()), + static_cast(t->device().getType()), + static_cast(t->device().getIndex()), + }; + if (cache_wrapper) { + const auto it = wrapper_cache.find(cache_key); + if (it != wrapper_cache.end()) { + return it->second; + } + } auto dtype = to_at_dtype(t->dtype()); auto device = to_at_device(t->device()); @@ -30,12 +92,20 @@ at::Tensor to_aten_tensor(const infinicore::Tensor &t) { .device(device) .requires_grad(false); - return at::from_blob( + if (t->numel() == 0) { + return at::empty_strided(sizes, strides, options); + } + + auto result = at::from_blob( data_ptr, sizes, strides, deleter_, options); + if (cache_wrapper) { + wrapper_cache.emplace(std::move(cache_key), result); + } + return result; } #if defined(ENABLE_HYGON_API) diff --git a/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc b/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc index 32a485627..87d12e7bf 100644 --- a/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc +++ b/src/infinicore/adaptor/vllm_iluvatar_adaptor.cc @@ -24,6 +24,7 @@ namespace infinicore::adaptor::vllm_iluvatar { namespace { using fused_add_rms_norm_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, float); +using rotary_embedding_fn = void (*)(at::Tensor &, at::Tensor &, std::optional, int64_t, at::Tensor &, bool); using dynamic_scaled_int8_quant_fn = void (*)(at::Tensor &, at::Tensor &, const at::Tensor &); using concat_mla_q_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &); using concat_and_cache_mla_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, const std::string &, at::Tensor &); @@ -65,10 +66,23 @@ void record_or_run(Fn &&fn) { } } using moe_sum_vllm_fn = void (*)(at::Tensor &, const at::Tensor &, std::optional, std::optional, double, double); +using fused_indexer_postprocess_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, int64_t, bool, double, double); +using indexer_k_cache_fn = void (*)(const at::Tensor &, at::Tensor &, const at::Tensor &); +using indexer_k_quant_and_cache_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, int64_t, const std::string &); +using block_sparse_logits_fn = void (*)(const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, at::Tensor &, int64_t, int64_t, int64_t); +using select_prefill_topk_fn = void (*)(const at::Tensor &, const at::Tensor &, const at::Tensor &, at::Tensor &); +using select_decode_topk_fn = void (*)(const at::Tensor &, const at::Tensor &, at::Tensor &); +using map_prefill_result = std::tuple>; +using map_prefill_indices_fn = map_prefill_result (*)(const at::Tensor &, const at::Tensor &, const at::Tensor &, int64_t, bool, std::optional, std::optional, bool, std::optional); +using map_decode_indices_fn = at::Tensor (*)(const at::Tensor &, const at::Tensor &, const at::Tensor &, int64_t, std::optional); +using sparse_flash_attn_fn = void (*)(at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, double); +using topk_indices_context_lens_fn = void (*)(at::Tensor &, const at::Tensor &); +using flash_mla_sparse_v2_fn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, at::Tensor &, float, const std::optional &); struct Symbols { void *handle = nullptr; fused_add_rms_norm_fn fused_add_rms_norm = nullptr; + rotary_embedding_fn rotary_embedding = nullptr; dynamic_scaled_int8_quant_fn dynamic_scaled_int8_quant = nullptr; concat_mla_q_fn concat_mla_q = nullptr; concat_and_cache_mla_fn concat_and_cache_mla = nullptr; @@ -86,7 +100,20 @@ struct Symbols { expand_moe_input_with_inv_pos_fn expand_moe_input_with_inv_pos = nullptr; silu_and_mul_quant_fn silu_and_mul_quant = nullptr; moe_sum_vllm_fn moe_sum_vllm = nullptr; + fused_indexer_postprocess_fn fused_indexer_postprocess = nullptr; + indexer_k_cache_fn indexer_k_cache = nullptr; + indexer_k_quant_and_cache_fn indexer_k_quant_and_cache = nullptr; + block_sparse_logits_fn block_sparse_logits = nullptr; + select_prefill_topk_fn select_prefill_topk = nullptr; + select_decode_topk_fn select_decode_topk = nullptr; + map_prefill_indices_fn map_prefill_indices = nullptr; + map_decode_indices_fn map_decode_indices = nullptr; + sparse_flash_attn_fn sparse_flash_attn = nullptr; + void *ix_handle = nullptr; + topk_indices_context_lens_fn topk_indices_context_lens = nullptr; + flash_mla_sparse_v2_fn flash_mla_sparse_v2 = nullptr; std::string error; + std::string ix_error; }; Symbols &symbols() { @@ -112,6 +139,8 @@ Symbols &symbols() { } auto fused_fn = reinterpret_cast( dlsym(handle, "_ZN7pyinfer4perf18fused_add_rms_normERN2at6TensorES3_S3_f")); + auto rotary_embedding_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf16rotary_embeddingERN2at6TensorES3_St8optionalIS2_ElS3_b")); auto quant_fn = reinterpret_cast( dlsym(handle, "_ZN7pyinfer4perf25dynamic_scaled_int8_quantERN2at6TensorES3_RKS2_")); auto concat_mla_q_fn_ptr = reinterpret_cast( @@ -146,13 +175,32 @@ Symbols &symbols() { dlsym(handle, "_ZN7pyinfer4perf18silu_and_mul_quantERN2at6TensorESt8optionalIS2_ERKS2_l")); auto moe_sum_vllm_fn_ptr = reinterpret_cast( dlsym(handle, "_ZN7pyinfer4perf7moe_sumERN2at6TensorERKS2_St8optionalIS2_ES7_dd")); - if (!fused_fn && !quant_fn && !concat_mla_q_fn_ptr && !concat_and_cache_mla_fn_ptr && !concat_and_cache_mla_int8_fn_ptr && !paged_attention_mla_fn_ptr && !topk_softmax_fn_ptr && !topk_sigmoid_fn_ptr && !grouped_topk_fn_ptr && !scaled_mm_w4a8_fn_ptr && !scaled_mm_w8a8_fn_ptr && !w4a8_group_gemm_fn_ptr && !w8a8_group_gemm_fn_ptr && !w16a16_group_gemm_fn_ptr && !argsort_bincount_with_inv_pos_fn_ptr && !expand_moe_input_with_inv_pos_fn_ptr && !silu_and_mul_quant_fn_ptr && !moe_sum_vllm_fn_ptr) { + auto fused_indexer_postprocess_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf37fused_deepseek_v2_indexer_postprocessERN2at6TensorES3_S3_S3_RKS2_S5_S5_S5_S5_S5_S5_lbdd")); + auto indexer_k_cache_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf15indexer_k_cacheERKN2at6TensorERS2_S4_")); + auto indexer_k_quant_and_cache_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf25indexer_k_quant_and_cacheERN2at6TensorES3_S3_lRKSs")); + auto block_sparse_logits_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf31compute_block_sparse_mqa_logitsERKN2at6TensorES4_S4_S4_S4_S4_RS2_lll")); + auto select_prefill_topk_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf33select_prefill_topk_block_indicesERKN2at6TensorES4_S4_RS2_")); + auto select_decode_topk_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf32select_decode_topk_block_indicesERKN2at6TensorES4_RS2_")); + auto map_prefill_indices_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf50map_prefill_request_block_indices_to_global_blocksERKN2at6TensorES4_S4_lbSt8optionalIS2_ES6_bS6_")); + auto map_decode_indices_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf49map_decode_request_block_indices_to_global_blocksERKN2at6TensorES4_S4_lSt8optionalIS2_E")); + auto sparse_flash_attn_fn_ptr = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer7cuinfer17sparse_flash_attnERN2at6TensorERKS2_S5_S5_d")); + if (!fused_fn && !rotary_embedding_fn_ptr && !quant_fn && !concat_mla_q_fn_ptr && !concat_and_cache_mla_fn_ptr && !concat_and_cache_mla_int8_fn_ptr && !paged_attention_mla_fn_ptr && !topk_softmax_fn_ptr && !topk_sigmoid_fn_ptr && !grouped_topk_fn_ptr && !scaled_mm_w4a8_fn_ptr && !scaled_mm_w8a8_fn_ptr && !w4a8_group_gemm_fn_ptr && !w8a8_group_gemm_fn_ptr && !w16a16_group_gemm_fn_ptr && !argsort_bincount_with_inv_pos_fn_ptr && !expand_moe_input_with_inv_pos_fn_ptr && !silu_and_mul_quant_fn_ptr && !moe_sum_vllm_fn_ptr) { syms.error = dlerror(); dlclose(handle); continue; } syms.handle = handle; syms.fused_add_rms_norm = fused_fn; + syms.rotary_embedding = rotary_embedding_fn_ptr; syms.dynamic_scaled_int8_quant = quant_fn; syms.concat_mla_q = concat_mla_q_fn_ptr; syms.concat_and_cache_mla = concat_and_cache_mla_fn_ptr; @@ -170,12 +218,55 @@ Symbols &symbols() { syms.expand_moe_input_with_inv_pos = expand_moe_input_with_inv_pos_fn_ptr; syms.silu_and_mul_quant = silu_and_mul_quant_fn_ptr; syms.moe_sum_vllm = moe_sum_vllm_fn_ptr; + syms.fused_indexer_postprocess = fused_indexer_postprocess_fn_ptr; + syms.indexer_k_cache = indexer_k_cache_fn_ptr; + syms.indexer_k_quant_and_cache = indexer_k_quant_and_cache_fn_ptr; + syms.block_sparse_logits = block_sparse_logits_fn_ptr; + syms.select_prefill_topk = select_prefill_topk_fn_ptr; + syms.select_decode_topk = select_decode_topk_fn_ptr; + syms.map_prefill_indices = map_prefill_indices_fn_ptr; + syms.map_decode_indices = map_decode_indices_fn_ptr; + syms.sparse_flash_attn = sparse_flash_attn_fn_ptr; syms.error.clear(); - return; + break; } - if (syms.error.empty()) { + if (!syms.handle && syms.error.empty()) { syms.error = "vllm_iluvatar extension not found"; } + + const char *ix_env_path = std::getenv("INFINICORE_IXTRITURBO_OPS_SO"); + const char *ix_paths[] = { + ix_env_path, + "/usr/local/corex-4.5.0.20260619/lib64/python3/dist-packages/ixtriturbo/_C/ops.cpython-312-x86_64-linux-gnu.so", + }; + for (const char *path : ix_paths) { + if (path == nullptr || path[0] == '\0') { + continue; + } + void *handle = dlopen(path, RTLD_LAZY | RTLD_LOCAL); + if (!handle) { + syms.ix_error = dlerror(); + continue; + } + auto lens_fn = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf25topk_indices_context_lensERN2at6TensorERKS2_")); + auto sparse_v2_fn = reinterpret_cast( + dlsym(handle, "_ZN7pyinfer4perf19flash_mla_sparse_v2ERN2at6TensorES3_S3_S3_S3_fRKSt8optionalIS2_E")); + if (!lens_fn || !sparse_v2_fn) { + const char *error = dlerror(); + syms.ix_error = error ? error : "ixtriturbo sparse MLA symbols not found"; + dlclose(handle); + continue; + } + syms.ix_handle = handle; + syms.topk_indices_context_lens = lens_fn; + syms.flash_mla_sparse_v2 = sparse_v2_fn; + syms.ix_error.clear(); + break; + } + if (!syms.ix_handle && syms.ix_error.empty()) { + syms.ix_error = "ixtriturbo ops extension not found"; + } #else syms.error = "InfiniCore was not built with ENABLE_ILUVATAR_API"; #endif @@ -189,6 +280,10 @@ bool available() { return symbols().fused_add_rms_norm != nullptr; } +bool rotary_embedding_available() { + return symbols().rotary_embedding != nullptr; +} + bool dynamic_scaled_int8_quant_available() { return symbols().dynamic_scaled_int8_quant != nullptr; } @@ -257,6 +352,17 @@ bool moe_sum_vllm_available() { return symbols().moe_sum_vllm != nullptr; } +bool fused_deepseek_v2_indexer_postprocess_available() { return symbols().fused_indexer_postprocess != nullptr; } +bool indexer_k_cache_available() { return symbols().indexer_k_cache != nullptr; } +bool indexer_k_quant_and_cache_available() { return symbols().indexer_k_quant_and_cache != nullptr; } +bool compute_block_sparse_mqa_logits_available() { return symbols().block_sparse_logits != nullptr; } +bool select_prefill_topk_block_indices_available() { return symbols().select_prefill_topk != nullptr; } +bool select_decode_topk_block_indices_available() { return symbols().select_decode_topk != nullptr; } +bool map_prefill_request_block_indices_available() { return symbols().map_prefill_indices != nullptr; } +bool map_decode_request_block_indices_available() { return symbols().map_decode_indices != nullptr; } +bool sparse_flash_mla_available() { return symbols().flash_mla_sparse_v2 != nullptr || symbols().sparse_flash_attn != nullptr; } +bool topk_indices_context_lens_available() { return symbols().topk_indices_context_lens != nullptr; } + void fused_add_rms_norm(at::Tensor &input, at::Tensor &residual, at::Tensor &weight, float epsilon) { auto &syms = symbols(); if (!syms.fused_add_rms_norm) { @@ -269,6 +375,23 @@ void fused_add_rms_norm(at::Tensor &input, at::Tensor &residual, at::Tensor &wei }); } +void rotary_embedding(at::Tensor &positions, + at::Tensor &query, + std::optional key, + int64_t head_size, + at::Tensor &cos_sin_cache, + bool is_neox) { + auto &syms = symbols(); + if (!syms.rotary_embedding) { + throw std::runtime_error("vllm_iluvatar rotary_embedding unavailable: " + syms.error); + } + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + run_syms.rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox); + }); +} + void dynamic_scaled_int8_quant(at::Tensor &output, at::Tensor &input_scales, const at::Tensor &input) { auto &syms = symbols(); if (!syms.dynamic_scaled_int8_quant) { @@ -483,6 +606,85 @@ void moe_sum_vllm(at::Tensor &output, const at::Tensor &input, std::optional prefill_workspace_request_ids, std::optional prefill_workspace_starts) { + if (!symbols().map_prefill_indices) { + throw std::runtime_error("vllm_iluvatar prefill index mapping unavailable: " + symbols().error); + } + record_or_run([=]() mutable { INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); (void)symbols().map_prefill_indices(req_id, block_table, token_indices, block_size, has_prefill_workspace, prefill_workspace_request_ids, prefill_workspace_starts, false, output); }); +} + +void map_decode_request_block_indices(at::Tensor &output, const at::Tensor &req_id, const at::Tensor &block_table, const at::Tensor &token_indices, int64_t block_size) { + if (!symbols().map_decode_indices) { + throw std::runtime_error("vllm_iluvatar decode index mapping unavailable: " + symbols().error); + } + record_or_run([=]() mutable { INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); (void)symbols().map_decode_indices(req_id, block_table, token_indices, block_size, output); }); +} + +void topk_indices_context_lens(at::Tensor &topk_lens, const at::Tensor &indices) { + if (!symbols().topk_indices_context_lens) { + throw std::runtime_error("ixtriturbo topk_indices_context_lens unavailable: " + symbols().ix_error); + } + record_or_run([=]() mutable { INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); symbols().topk_indices_context_lens(topk_lens, indices); }); +} + +void sparse_flash_mla(at::Tensor &output, at::Tensor &query, at::Tensor &kv_cache, at::Tensor &indices, at::Tensor &topk_lens, float scale, std::optional attn_sink) { + auto &syms = symbols(); + if (!syms.flash_mla_sparse_v2 && !syms.sparse_flash_attn) { + throw std::runtime_error("Iluvatar sparse FlashMLA unavailable: " + syms.error + "; " + syms.ix_error); + } + record_or_run([=]() mutable { + INFINICORE_VLLM_ILUVATAR_STREAM_GUARD(); + auto &run_syms = symbols(); + if (run_syms.flash_mla_sparse_v2) { + run_syms.flash_mla_sparse_v2(output, query, kv_cache, indices, topk_lens, scale, attn_sink); + } else { + run_syms.sparse_flash_attn(output, query, kv_cache, indices, scale); + } + }); +} + #undef INFINICORE_VLLM_ILUVATAR_STREAM_GUARD } // namespace infinicore::adaptor::vllm_iluvatar diff --git a/src/infinicore/context/context_impl.cc b/src/infinicore/context/context_impl.cc index f126beb1d..00ed9a986 100644 --- a/src/infinicore/context/context_impl.cc +++ b/src/infinicore/context/context_impl.cc @@ -208,6 +208,10 @@ std::shared_ptr stopGraphRecording() { return ContextImpl::singleton().getCurrentRuntime()->stopGraphRecording(); } +void cancelGraphRecording() { + ContextImpl::singleton().getCurrentRuntime()->cancelGraphRecording(); +} + std::shared_ptr reinstantiateBlob(std::shared_ptr blob) { setDevice(blob->device()); return ContextImpl::singleton().getCurrentRuntime()->reinstantiateBlob(blob); diff --git a/src/infinicore/context/runtime/runtime.cc b/src/infinicore/context/runtime/runtime.cc index 92ba0f737..65d244f41 100644 --- a/src/infinicore/context/runtime/runtime.cc +++ b/src/infinicore/context/runtime/runtime.cc @@ -192,9 +192,19 @@ void Runtime::addGraphOperator(std::shared_ptr op) { } std::shared_ptr Runtime::stopGraphRecording() { - auto graph = graph_manager_->stop_recording(); + try { + auto graph = graph_manager_->stop_recording(); + device_memory_allocator_->set_pin_mode(false); + return graph; + } catch (...) { + device_memory_allocator_->set_pin_mode(false); + throw; + } +} + +void Runtime::cancelGraphRecording() { + graph_manager_->cancel_recording(); device_memory_allocator_->set_pin_mode(false); - return graph; } std::string Runtime::toString() const { diff --git a/src/infinicore/context/runtime/runtime.hpp b/src/infinicore/context/runtime/runtime.hpp index 2ed382ab7..7dbb5185f 100644 --- a/src/infinicore/context/runtime/runtime.hpp +++ b/src/infinicore/context/runtime/runtime.hpp @@ -66,6 +66,7 @@ class Runtime { void startGraphRecording(); void addGraphOperator(std::shared_ptr op); std::shared_ptr stopGraphRecording(); + void cancelGraphRecording(); std::string toString() const; diff --git a/src/infinicore/graph/graph.cc b/src/infinicore/graph/graph.cc index 3b8fc57e5..9ccee8400 100644 --- a/src/infinicore/graph/graph.cc +++ b/src/infinicore/graph/graph.cc @@ -2,6 +2,7 @@ #include "../utils.hpp" #include "infinicore/context/context.hpp" +#include #include namespace infinicore::graph { @@ -38,6 +39,9 @@ struct Graph::DeviceGraph { std::vector log_buffer; DeviceGraph() { + graph = nullptr; + exec = nullptr; + node = nullptr; log_buffer.resize(4 * 1024); } @@ -55,15 +59,32 @@ struct Graph::DeviceGraph { } }; +struct Graph::ReplayStep { + std::unique_ptr device_graph; + std::shared_ptr host_op; +}; + Graph::Graph() { } void Graph::run() const { - if (device_graph_ != nullptr && device_graph_.get()->exec != nullptr) { - device_graph_.get()->launch(); - } else { + if (replay_steps_.empty()) { for (auto &op : op_list_) { op->run(); + if (op->requires_stream_sync_after_run()) { + infinicore::context::syncStream(); + } + } + return; + } + for (const auto &step : replay_steps_) { + if (step->device_graph) { + step->device_graph->launch(); + } else { + step->host_op->run(); + if (step->host_op->requires_stream_sync_after_run()) { + infinicore::context::syncStream(); + } } } } @@ -73,44 +94,86 @@ void Graph::add_operator(std::shared_ptr op) { } void Graph::instantiate() { - // Reset device graph - device_graph_ = std::make_unique(); + replay_steps_.clear(); - // warmup + // Warm the complete op list before splitting it at host-replayed P2P ops. for (size_t iter = 0; iter < 5; ++iter) { this->run(); } infinicore::context::syncStream(); - if (infinirtStreamBeginCapture( - context::getStream(), - INFINIRT_STREAM_CAPTURE_MODE_RELAXED) - != INFINI_STATUS_SUCCESS) { + // Diagnostic escape hatch: keep GraphTensor/operator replay semantics but + // bypass device-graph capture, including segmented PP capture. + if (std::getenv("INFINICORE_DISABLE_DEVICE_GRAPH_SEGMENTS") != nullptr) { + spdlog::info("device graph segments disabled; replaying recorded operators"); return; } - // Run and record - this->run(); + auto capture_segment = [&](size_t begin, size_t end) { + if (begin == end) { + return; + } + auto device_graph = std::make_unique(); + if (infinirtStreamBeginCapture( + context::getStream(), + INFINIRT_STREAM_CAPTURE_MODE_RELAXED) + != INFINI_STATUS_SUCCESS) { + throw std::runtime_error("failed to begin device graph capture"); + } - if (infinirtStreamEndCapture( - context::getStream(), - &device_graph_.get()->graph) - != INFINI_STATUS_SUCCESS) { - return; - } + for (size_t i = begin; i < end; ++i) { + op_list_[i]->run(); + } - if (infinirtGraphInstantiate( - &device_graph_.get()->exec, - device_graph_.get()->graph, - &device_graph_.get()->node, - device_graph_.get()->log_buffer.data(), - device_graph_.get()->log_buffer.size()) - != INFINI_STATUS_SUCCESS) { - static bool warned_once = false; - if (!warned_once) { - warned_once = true; - spdlog::warn("Fail to instantiate device graph: {}", std::string(device_graph_.get()->log_buffer.data())); + if (infinirtStreamEndCapture( + context::getStream(), + &device_graph->graph) + != INFINI_STATUS_SUCCESS) { + throw std::runtime_error("failed to end device graph capture"); } + + if (infinirtGraphInstantiate( + &device_graph->exec, + device_graph->graph, + &device_graph->node, + device_graph->log_buffer.data(), + device_graph->log_buffer.size()) + != INFINI_STATUS_SUCCESS) { + throw std::runtime_error( + "failed to instantiate device graph: " + + std::string(device_graph->log_buffer.data())); + } + auto step = std::make_unique(); + step->device_graph = std::move(device_graph); + replay_steps_.push_back(std::move(step)); + }; + + size_t segment_begin = 0; + for (size_t i = 0; i < op_list_.size(); ++i) { + if (op_list_[i]->is_device_graph_capture_safe()) { + continue; + } + capture_segment(segment_begin, i); + // Execute the P2P operation once between captured compute segments so + // downstream capture sees the same stream-ordered data dependency. + op_list_[i]->run(); + if (op_list_[i]->requires_stream_sync_after_run()) { + infinicore::context::syncStream(); + } + auto step = std::make_unique(); + step->host_op = op_list_[i]; + replay_steps_.push_back(std::move(step)); + segment_begin = i + 1; + } + capture_segment(segment_begin, op_list_.size()); + if (std::getenv("INFINICORE_GRAPH_DEBUG") != nullptr) { + size_t host_steps = 0; + for (const auto &step : replay_steps_) { + host_steps += step->host_op != nullptr ? 1 : 0; + } + spdlog::info( + "segmented graph: operators={}, replay_steps={}, host_steps={}", + op_list_.size(), replay_steps_.size(), host_steps); } } @@ -150,4 +213,9 @@ std::shared_ptr GraphManager::stop_recording() { return std::exchange(graph_, nullptr); } +void GraphManager::cancel_recording() { + recording_ = false; + graph_.reset(); +} + } // namespace infinicore::graph diff --git a/src/infinicore/graph/graph_manager.hpp b/src/infinicore/graph/graph_manager.hpp index 9cc040366..587d0c98f 100644 --- a/src/infinicore/graph/graph_manager.hpp +++ b/src/infinicore/graph/graph_manager.hpp @@ -16,6 +16,7 @@ class GraphManager { void start_recording(); void add_operator(std::shared_ptr op); std::shared_ptr stop_recording(); + void cancel_recording(); private: std::shared_ptr graph_; diff --git a/src/infinicore/ops/bmm_strided/bmm_strided.cc b/src/infinicore/ops/bmm_strided/bmm_strided.cc new file mode 100644 index 000000000..5419b9006 --- /dev/null +++ b/src/infinicore/ops/bmm_strided/bmm_strided.cc @@ -0,0 +1,88 @@ +#include "infinicore/ops/bmm_strided.hpp" + +#include "../../utils.hpp" +#include "infinicore/context/context.hpp" +#include "infinicore/graph/graph.hpp" + +#include +#include +#include + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#endif + +namespace infinicore::op { +namespace { + +class DeferredGraphOperator final : public graph::GraphOperator { +public: + explicit DeferredGraphOperator(std::function runner) + : runner_(std::move(runner)) {} + + void run() const override { runner_(); } + +private: + std::function runner_; +}; + +void record_or_run(std::function runner) { + auto op = std::make_shared(std::move(runner)); + if (context::isGraphRecording()) { + context::addGraphOperator(op); + } else { + op->run(); + } +} + +void validate_bmm_strided(const Tensor &output, const Tensor &a, const Tensor &b) { + if (!output || !a || !b) { + throw std::runtime_error("bmm_strided expects non-empty output, a, and b tensors"); + } + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, a, b); + if (output->dtype() != a->dtype() || output->dtype() != b->dtype()) { + throw std::runtime_error("bmm_strided expects output, a, and b to have the same dtype"); + } + if (output->dtype() != DataType::F16 + && output->dtype() != DataType::BF16 + && output->dtype() != DataType::F32) { + throw std::runtime_error("bmm_strided expects float16, bfloat16, or float32 tensors"); + } + if (output->ndim() != 3 || a->ndim() != 3 || b->ndim() != 3) { + throw std::runtime_error("bmm_strided expects three-dimensional tensors"); + } + if (a->size(0) != b->size(0) + || a->size(0) != output->size(0) + || a->size(2) != b->size(1) + || a->size(1) != output->size(1) + || b->size(2) != output->size(2)) { + throw std::runtime_error("bmm_strided tensor shapes are incompatible"); + } +} + +void run_bmm_strided(Tensor output, const Tensor &a, const Tensor &b) { +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (output->device().getType() == Device::Type::ILUVATAR) { + adaptor::set_aten_stream_to_infinicore(); + auto output_at = adaptor::to_aten_tensor(output); + const auto a_at = adaptor::to_aten_tensor(a); + const auto b_at = adaptor::to_aten_tensor(b); + at::bmm_out(output_at, a_at, b_at); + return; + } +#endif + + throw std::runtime_error( + "bmm_strided currently supports only Iluvatar builds with ATen"); +} + +} // namespace + +void bmm_strided_(Tensor output, const Tensor &a, const Tensor &b) { + validate_bmm_strided(output, a, b); + record_or_run([output, a, b] { + run_bmm_strided(output, a, b); + }); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/concat_and_cache_mla/concat_and_cache_mla.cc b/src/infinicore/ops/concat_and_cache_mla/concat_and_cache_mla.cc index d4a121a5c..b1e4edb77 100644 --- a/src/infinicore/ops/concat_and_cache_mla/concat_and_cache_mla.cc +++ b/src/infinicore/ops/concat_and_cache_mla/concat_and_cache_mla.cc @@ -22,7 +22,9 @@ void validate_concat_and_cache_mla(const Tensor &kv_c, throw std::runtime_error("concat_and_cache_mla expects non-empty kv_c, k_pe, kv_cache, slot_mapping, and scale tensors"); } INFINICORE_ASSERT_TENSORS_SAME_DEVICE(kv_c, k_pe, kv_cache, slot_mapping, scale); - if (kv_cache_dtype != "auto" && kv_cache_dtype != "fp8" && kv_cache_dtype != "fp8_e4m3" && kv_cache_dtype != "fp8_e5m2") { + if (kv_cache_dtype != "auto" && kv_cache_dtype != "fp8" + && kv_cache_dtype != "fp8_e4m3" && kv_cache_dtype != "fp8_e5m2" + && kv_cache_dtype != "fp8_ds_mla") { throw std::runtime_error("concat_and_cache_mla expects kv_cache_dtype to be auto/fp8/fp8_e4m3/fp8_e5m2"); } if (kv_c->ndim() != 2 || k_pe->ndim() != 2) { @@ -32,8 +34,18 @@ void validate_concat_and_cache_mla(const Tensor &kv_c, throw std::runtime_error("concat_and_cache_mla expects kv_c/k_pe tokens to match slot_mapping numel"); } const auto head_dim = kv_c->size(1) + k_pe->size(1); - if (kv_cache->ndim() < 3 || kv_cache->size(kv_cache->ndim() - 1) != head_dim) { - throw std::runtime_error("concat_and_cache_mla expects kv_cache last dim == kv_c.shape[-1] + k_pe.shape[-1]"); + if (kv_cache_dtype == "fp8_ds_mla") { + const auto cache_stride = kv_c->size(1) + 4 * sizeof(float) + + k_pe->size(1) * sizeof(uint16_t); + if (kv_cache->ndim() < 3 || kv_cache->dtype() != DataType::U8 + || kv_cache->size(kv_cache->ndim() - 1) != cache_stride) { + throw std::runtime_error( + "concat_and_cache_mla expects fp8_ds_mla uint8 cache with value/scales/rope layout"); + } + } else if (kv_cache->ndim() < 3 + || kv_cache->size(kv_cache->ndim() - 1) != head_dim) { + throw std::runtime_error( + "concat_and_cache_mla expects kv_cache last dim == kv_c.shape[-1] + k_pe.shape[-1]"); } if (slot_mapping->dtype() != DataType::I64 && slot_mapping->dtype() != DataType::I32) { throw std::runtime_error("concat_and_cache_mla expects slot_mapping dtype int64 or int32"); diff --git a/src/infinicore/ops/distributed/p2p.cc b/src/infinicore/ops/distributed/p2p.cc new file mode 100644 index 000000000..d4879d834 --- /dev/null +++ b/src/infinicore/ops/distributed/p2p.cc @@ -0,0 +1,202 @@ +#include "infinicore/ops/distributed/p2p.hpp" +#include "../../utils.hpp" + +namespace infinicore::op::distributed { +namespace { +struct PlannedMeta { + graph::GraphTensor tensor; + int peer; + infinicclComm_t communicator; +}; +struct GroupedPlannedMeta { + std::vector tensors; + int peer; + infinicclComm_t communicator; +}; +} // namespace + +Send::Send(const Tensor &input, int peer, infinicclComm_t communicator) { + INFINICORE_ASSERT(input->is_contiguous()); + planned_meta_ = new PlannedMeta{ + graph::GraphTensor(input), peer, communicator}; +} + +Send::~Send() { + delete reinterpret_cast(planned_meta_); +} + +void Send::run() const { + const auto *meta = reinterpret_cast(planned_meta_); + INFINICORE_CHECK_ERROR(infinicclSend( + meta->tensor->data(), meta->tensor->numel(), + static_cast(static_cast(meta->tensor->dtype())), + meta->peer, meta->communicator, infinicore::context::getStream())); +} + +void Send::execute( + const Tensor &input, int peer, infinicclComm_t communicator) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Send, input, peer, communicator); +} + +Recv::Recv(Tensor output, int peer, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->is_contiguous()); + planned_meta_ = new PlannedMeta{ + graph::GraphTensor(output), peer, communicator}; +} + +Recv::~Recv() { + delete reinterpret_cast(planned_meta_); +} + +void Recv::run() const { + const auto *meta = reinterpret_cast(planned_meta_); + INFINICORE_CHECK_ERROR(infinicclRecv( + const_cast(static_cast(meta->tensor->data())), + meta->tensor->numel(), + static_cast(static_cast(meta->tensor->dtype())), + meta->peer, meta->communicator, infinicore::context::getStream())); +} + +void Recv::execute( + Tensor output, int peer, infinicclComm_t communicator) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Recv, output, peer, communicator); +} + +GroupedSend::GroupedSend(const std::vector &inputs, int peer, + infinicclComm_t communicator) { + std::vector tensors; + tensors.reserve(inputs.size()); + for (const auto &input : inputs) { + INFINICORE_ASSERT(input->is_contiguous()); + tensors.emplace_back(input); + } + planned_meta_ = new GroupedPlannedMeta{ + std::move(tensors), peer, communicator}; +} + +GroupedSend::~GroupedSend() { + delete reinterpret_cast(planned_meta_); +} + +void GroupedSend::run() const { + const auto *meta = reinterpret_cast(planned_meta_); + INFINICORE_CHECK_ERROR(infinicclGroupStart(meta->communicator)); + try { + for (const auto &tensor : meta->tensors) { + INFINICORE_CHECK_ERROR(infinicclSend( + tensor->data(), tensor->numel(), + static_cast(static_cast(tensor->dtype())), + meta->peer, meta->communicator, + infinicore::context::getStream())); + } + } catch (...) { + (void)infinicclGroupEnd(meta->communicator); + throw; + } + INFINICORE_CHECK_ERROR(infinicclGroupEnd(meta->communicator)); +} + +void GroupedSend::execute(const std::vector &inputs, int peer, + infinicclComm_t communicator) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + GroupedSend, inputs, peer, communicator); +} + +GroupedRecv::GroupedRecv(const std::vector &outputs, int peer, + infinicclComm_t communicator) { + std::vector tensors; + tensors.reserve(outputs.size()); + for (const auto &output : outputs) { + INFINICORE_ASSERT(output->is_contiguous()); + tensors.emplace_back(output); + } + planned_meta_ = new GroupedPlannedMeta{ + std::move(tensors), peer, communicator}; +} + +GroupedRecv::~GroupedRecv() { + delete reinterpret_cast(planned_meta_); +} + +void GroupedRecv::run() const { + const auto *meta = reinterpret_cast(planned_meta_); + INFINICORE_CHECK_ERROR(infinicclGroupStart(meta->communicator)); + try { + for (const auto &tensor : meta->tensors) { + INFINICORE_CHECK_ERROR(infinicclRecv( + const_cast(static_cast(tensor->data())), + tensor->numel(), + static_cast(static_cast(tensor->dtype())), + meta->peer, meta->communicator, + infinicore::context::getStream())); + } + } catch (...) { + (void)infinicclGroupEnd(meta->communicator); + throw; + } + INFINICORE_CHECK_ERROR(infinicclGroupEnd(meta->communicator)); +} + +void GroupedRecv::execute(const std::vector &outputs, int peer, + infinicclComm_t communicator) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + GroupedRecv, outputs, peer, communicator); +} + +void send(const Tensor &input, int peer, infinicclComm_t communicator) { + Send::execute(input, peer, communicator); +} + +void recv_(Tensor output, int peer, infinicclComm_t communicator) { + Recv::execute(output, peer, communicator); +} + +namespace { +template +void run_grouped(infinicclComm_t communicator, Enqueue &&enqueue) { + // Graph recording retains each P2P operation as a host replay step. The + // short decode tensors already use that validated segmented path. + if (infinicore::context::isGraphRecording()) { + enqueue(); + return; + } + + INFINICORE_CHECK_ERROR(infinicclGroupStart(communicator)); + try { + enqueue(); + } catch (...) { + (void)infinicclGroupEnd(communicator); + throw; + } + INFINICORE_CHECK_ERROR(infinicclGroupEnd(communicator)); +} +} // namespace + +void send_grouped( + const std::vector &inputs, + int peer, infinicclComm_t communicator) { + if (infinicore::context::isGraphRecording()) { + GroupedSend::execute(inputs, peer, communicator); + return; + } + run_grouped(communicator, [&] { + for (const auto &input : inputs) { + Send::execute(input, peer, communicator); + } + }); +} + +void recv_grouped_( + const std::vector &outputs, + int peer, infinicclComm_t communicator) { + if (infinicore::context::isGraphRecording()) { + GroupedRecv::execute(outputs, peer, communicator); + return; + } + run_grouped(communicator, [&] { + for (const auto &output : outputs) { + Recv::execute(output, peer, communicator); + } + }); +} +} // namespace infinicore::op::distributed diff --git a/src/infinicore/ops/dsa/dsa.cc b/src/infinicore/ops/dsa/dsa.cc new file mode 100644 index 000000000..6a8cb652e --- /dev/null +++ b/src/infinicore/ops/dsa/dsa.cc @@ -0,0 +1,463 @@ +#include "infinicore/ops/dsa.hpp" + +#include "../../utils.hpp" +#include "infinicore/context/context.hpp" +#include "infinicore/graph/graph.hpp" +#include "infinicore/ops/fp8_sparse_mla.hpp" + +#include +#include +#include + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { +namespace { + +class DeferredGraphOperator final : public graph::GraphOperator { +public: + explicit DeferredGraphOperator(std::function runner) + : runner_(std::move(runner)) {} + + void run() const override { runner_(); } + +private: + std::function runner_; +}; + +bool defer_if_recording(std::function runner) { + if (!context::isGraphRecording()) { + return false; + } + context::addGraphOperator( + std::make_shared(std::move(runner))); + return true; +} + +void require_iluvatar(const Tensor &tensor, const char *op_name) { + if (!tensor || tensor->device().getType() != Device::Type::ILUVATAR) { + throw std::runtime_error(std::string(op_name) + " currently supports only Iluvatar tensors"); + } +} + +void require_contiguous(const Tensor &tensor, const char *op_name) { + if (!tensor || !tensor->is_contiguous()) { + throw std::runtime_error(std::string(op_name) + " expects contiguous tensors"); + } +} + +void require_i32(const Tensor &tensor, const char *name) { + if (tensor->dtype() != DataType::I32) { + throw std::runtime_error(std::string(name) + " must be int32"); + } +} + +} // namespace + +void fused_deepseek_v2_indexer_postprocess_( + Tensor q_out, + Tensor k_out, + Tensor weights_out, + Tensor kv_cache, + const Tensor &slot_mapping, + const Tensor &q, + const Tensor &kw, + const Tensor &norm_weight, + const Tensor &norm_bias, + const Tensor &positions, + const Tensor &cos_sin_cache, + int64_t num_cache_tokens, + bool is_neox, + double eps, + double weights_scale) { + require_iluvatar(q_out, "fused_deepseek_v2_indexer_postprocess"); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + q_out, k_out, weights_out, kv_cache, slot_mapping, q, kw, norm_weight, + norm_bias, positions, cos_sin_cache); + for (const auto &tensor : {q_out, k_out, weights_out, kv_cache, slot_mapping, + q, kw, norm_weight, norm_bias, positions, + cos_sin_cache}) { + require_contiguous(tensor, "fused_deepseek_v2_indexer_postprocess"); + } + if (q->ndim() != 3 || kw->ndim() != 2 || q_out->shape() != q->shape() + || weights_out->ndim() != 2 || positions->ndim() != 1 + || slot_mapping->ndim() != 1 || cos_sin_cache->ndim() != 2) { + throw std::runtime_error("fused_deepseek_v2_indexer_postprocess tensor rank/shape mismatch"); + } + if (q->dtype() != DataType::F16 && q->dtype() != DataType::BF16) { + throw std::runtime_error("fused_deepseek_v2_indexer_postprocess requires fp16/bfloat16"); + } + if (positions->dtype() != DataType::I64 || slot_mapping->dtype() != DataType::I64) { + throw std::runtime_error("fused_deepseek_v2_indexer_postprocess positions and slot_mapping must be int64"); + } + if (num_cache_tokens < 0 || static_cast(num_cache_tokens) > slot_mapping->numel()) { + throw std::runtime_error("fused_deepseek_v2_indexer_postprocess invalid num_cache_tokens"); + } + if (defer_if_recording([q_out, k_out, weights_out, kv_cache, slot_mapping, q, kw, + norm_weight, norm_bias, positions, cos_sin_cache, + num_cache_tokens, is_neox, eps, weights_scale] { + fused_deepseek_v2_indexer_postprocess_( + q_out, k_out, weights_out, kv_cache, slot_mapping, q, kw, + norm_weight, norm_bias, positions, cos_sin_cache, + num_cache_tokens, is_neox, eps, weights_scale); + })) { + return; + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (!adaptor::vllm_iluvatar::fused_deepseek_v2_indexer_postprocess_available()) { + throw std::runtime_error("vllm_iluvatar fused indexer postprocess is unavailable"); + } + auto q_out_at = adaptor::to_aten_tensor(q_out); + auto k_out_at = adaptor::to_aten_tensor(k_out); + auto weights_out_at = adaptor::to_aten_tensor(weights_out); + auto kv_cache_at = adaptor::to_aten_tensor(kv_cache); + auto slot_mapping_at = adaptor::to_aten_tensor(slot_mapping); + auto q_at = adaptor::to_aten_tensor(q); + auto kw_at = adaptor::to_aten_tensor(kw); + auto norm_weight_at = adaptor::to_aten_tensor(norm_weight); + auto norm_bias_at = adaptor::to_aten_tensor(norm_bias); + auto positions_at = adaptor::to_aten_tensor(positions); + auto cos_sin_cache_at = adaptor::to_aten_tensor(cos_sin_cache); + adaptor::vllm_iluvatar::fused_deepseek_v2_indexer_postprocess( + q_out_at, k_out_at, weights_out_at, kv_cache_at, slot_mapping_at, q_at, + kw_at, norm_weight_at, norm_bias_at, positions_at, cos_sin_cache_at, + num_cache_tokens, is_neox, eps, weights_scale); + return; +#else + throw std::runtime_error("fused_deepseek_v2_indexer_postprocess requires an Iluvatar ATen build"); +#endif +} + +void indexer_k_cache_(const Tensor &k, Tensor kv_cache, const Tensor &slot_mapping) { + require_iluvatar(k, "indexer_k_cache"); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k, kv_cache, slot_mapping); + require_contiguous(k, "indexer_k_cache"); + require_contiguous(kv_cache, "indexer_k_cache"); + require_contiguous(slot_mapping, "indexer_k_cache"); + if (k->ndim() != 2 || kv_cache->ndim() != 3 || slot_mapping->ndim() != 1 + || k->size(1) != kv_cache->size(2) || slot_mapping->dtype() != DataType::I64) { + throw std::runtime_error("indexer_k_cache tensor shape/dtype mismatch"); + } + if (defer_if_recording([k, kv_cache, slot_mapping] { + indexer_k_cache_(k, kv_cache, slot_mapping); + })) { + return; + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (!adaptor::vllm_iluvatar::indexer_k_cache_available()) { + throw std::runtime_error("vllm_iluvatar indexer_k_cache is unavailable"); + } + auto k_at = adaptor::to_aten_tensor(k); + auto cache_at = adaptor::to_aten_tensor(kv_cache); + auto slots_at = adaptor::to_aten_tensor(slot_mapping); + adaptor::vllm_iluvatar::indexer_k_cache(k_at, cache_at, slots_at); + return; +#else + throw std::runtime_error("indexer_k_cache requires an Iluvatar ATen build"); +#endif +} + +void indexer_k_quant_and_cache_( + const Tensor &k, + Tensor kv_cache, + const Tensor &slot_mapping, + int64_t quant_block_size, + const std::string &scale_fmt) { + require_iluvatar(k, "indexer_k_quant_and_cache"); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k, kv_cache, slot_mapping); + require_contiguous(k, "indexer_k_quant_and_cache"); + require_contiguous(kv_cache, "indexer_k_quant_and_cache"); + require_contiguous(slot_mapping, "indexer_k_quant_and_cache"); + if (k->ndim() != 2 || kv_cache->ndim() != 3 || slot_mapping->ndim() != 1 + || kv_cache->dtype() != DataType::U8 || slot_mapping->dtype() != DataType::I64 + || k->size(0) != slot_mapping->numel() + || kv_cache->size(2) != k->size(1) + sizeof(float) + || quant_block_size != static_cast(k->size(1))) { + throw std::runtime_error("indexer_k_quant_and_cache tensor shape/dtype mismatch"); + } + if (scale_fmt != "ue8m0" && !scale_fmt.empty()) { + throw std::runtime_error("indexer_k_quant_and_cache supports ue8m0 or empty scale_fmt"); + } + if (defer_if_recording([k, kv_cache, slot_mapping, quant_block_size, scale_fmt] { + indexer_k_quant_and_cache_( + k, kv_cache, slot_mapping, quant_block_size, scale_fmt); + })) { + return; + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (!adaptor::vllm_iluvatar::indexer_k_quant_and_cache_available()) { + throw std::runtime_error("vllm_iluvatar indexer_k_quant_and_cache is unavailable"); + } + auto k_at = adaptor::to_aten_tensor(k); + auto cache_at = adaptor::to_aten_tensor(kv_cache); + auto slots_at = adaptor::to_aten_tensor(slot_mapping); + adaptor::vllm_iluvatar::indexer_k_quant_and_cache( + k_at, cache_at, slots_at, quant_block_size, scale_fmt); + return; +#else + throw std::runtime_error( + "indexer_k_quant_and_cache requires an Iluvatar ATen build"); +#endif +} + +void compute_block_sparse_mqa_logits_( + Tensor logits, + const Tensor &q, + const Tensor &kv_cache, + const Tensor &cu_seqlens_q, + const Tensor &cu_seqlens_kv, + const Tensor &block_table, + const Tensor &weights, + int64_t max_q_len, + int64_t max_kv_len, + int64_t max_context_len) { + require_iluvatar(q, "compute_block_sparse_mqa_logits"); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + logits, q, kv_cache, cu_seqlens_q, cu_seqlens_kv, block_table, weights); + for (const auto &tensor : {logits, q, kv_cache, cu_seqlens_q, cu_seqlens_kv, + block_table, weights}) { + require_contiguous(tensor, "compute_block_sparse_mqa_logits"); + } + if (q->ndim() != 3 || kv_cache->ndim() != 3 || logits->ndim() != 2 + || weights->ndim() != 2 || block_table->ndim() != 2) { + throw std::runtime_error("compute_block_sparse_mqa_logits tensor rank mismatch"); + } + require_i32(cu_seqlens_q, "cu_seqlens_q"); + require_i32(cu_seqlens_kv, "cu_seqlens_kv"); + require_i32(block_table, "block_table"); + if (max_q_len <= 0 || max_kv_len <= 0 || max_context_len <= 0) { + throw std::runtime_error("compute_block_sparse_mqa_logits lengths must be positive"); + } + if (defer_if_recording([logits, q, kv_cache, cu_seqlens_q, cu_seqlens_kv, block_table, + weights, max_q_len, max_kv_len, max_context_len] { + compute_block_sparse_mqa_logits_( + logits, q, kv_cache, cu_seqlens_q, cu_seqlens_kv, + block_table, weights, max_q_len, max_kv_len, max_context_len); + })) { + return; + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (!adaptor::vllm_iluvatar::compute_block_sparse_mqa_logits_available()) { + throw std::runtime_error("vllm_iluvatar block sparse logits is unavailable"); + } + auto logits_at = adaptor::to_aten_tensor(logits); + auto q_at = adaptor::to_aten_tensor(q); + auto cache_at = adaptor::to_aten_tensor(kv_cache); + auto cu_q_at = adaptor::to_aten_tensor(cu_seqlens_q); + auto cu_kv_at = adaptor::to_aten_tensor(cu_seqlens_kv); + auto blocks_at = adaptor::to_aten_tensor(block_table); + auto weights_at = adaptor::to_aten_tensor(weights); + adaptor::vllm_iluvatar::compute_block_sparse_mqa_logits( + q_at, cache_at, cu_q_at, cu_kv_at, blocks_at, weights_at, logits_at, + max_q_len, max_kv_len, max_context_len); + return; +#else + throw std::runtime_error("compute_block_sparse_mqa_logits requires an Iluvatar ATen build"); +#endif +} + +void select_prefill_topk_block_indices_( + Tensor topk_indices, + const Tensor &logits, + const Tensor &cu_seqlen_ks, + const Tensor &cu_seqlen_ke) { + require_iluvatar(logits, "select_prefill_topk_block_indices"); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_indices, logits, cu_seqlen_ks, cu_seqlen_ke); + require_i32(topk_indices, "topk_indices"); + require_i32(cu_seqlen_ks, "cu_seqlen_ks"); + require_i32(cu_seqlen_ke, "cu_seqlen_ke"); + if (defer_if_recording([topk_indices, logits, cu_seqlen_ks, cu_seqlen_ke] { + select_prefill_topk_block_indices_( + topk_indices, logits, cu_seqlen_ks, cu_seqlen_ke); + })) { + return; + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + auto out_at = adaptor::to_aten_tensor(topk_indices); + auto logits_at = adaptor::to_aten_tensor(logits); + auto ks_at = adaptor::to_aten_tensor(cu_seqlen_ks); + auto ke_at = adaptor::to_aten_tensor(cu_seqlen_ke); + adaptor::vllm_iluvatar::select_prefill_topk_block_indices(logits_at, ks_at, ke_at, out_at); + return; +#else + throw std::runtime_error("select_prefill_topk_block_indices requires an Iluvatar ATen build"); +#endif +} + +void select_decode_topk_block_indices_( + Tensor topk_indices, + const Tensor &logits, + const Tensor &seq_lens) { + require_iluvatar(logits, "select_decode_topk_block_indices"); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_indices, logits, seq_lens); + require_i32(topk_indices, "topk_indices"); + require_i32(seq_lens, "seq_lens"); + if (defer_if_recording([topk_indices, logits, seq_lens] { + select_decode_topk_block_indices_(topk_indices, logits, seq_lens); + })) { + return; + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + auto out_at = adaptor::to_aten_tensor(topk_indices); + auto logits_at = adaptor::to_aten_tensor(logits); + auto lens_at = adaptor::to_aten_tensor(seq_lens); + adaptor::vllm_iluvatar::select_decode_topk_block_indices(logits_at, lens_at, out_at); + return; +#else + throw std::runtime_error("select_decode_topk_block_indices requires an Iluvatar ATen build"); +#endif +} + +void map_prefill_request_block_indices_( + Tensor output, + const Tensor &req_id, + const Tensor &block_table, + const Tensor &token_indices, + int64_t block_size, + bool has_prefill_workspace, + std::optional prefill_workspace_request_ids, + std::optional prefill_workspace_starts) { + require_iluvatar(output, "map_prefill_request_block_indices"); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, req_id, block_table, token_indices); + if (block_size <= 0 || output->shape() != token_indices->shape()) { + throw std::runtime_error("map_prefill_request_block_indices invalid output shape or block size"); + } + if (defer_if_recording([output, req_id, block_table, token_indices, block_size, + has_prefill_workspace, prefill_workspace_request_ids, + prefill_workspace_starts] { + map_prefill_request_block_indices_( + output, req_id, block_table, token_indices, block_size, + has_prefill_workspace, prefill_workspace_request_ids, + prefill_workspace_starts); + })) { + return; + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + auto output_at = adaptor::to_aten_tensor(output); + auto req_at = adaptor::to_aten_tensor(req_id); + auto blocks_at = adaptor::to_aten_tensor(block_table); + auto indices_at = adaptor::to_aten_tensor(token_indices); + std::optional request_ids_at; + std::optional starts_at; + if (prefill_workspace_request_ids) { + request_ids_at = adaptor::to_aten_tensor(*prefill_workspace_request_ids); + } + if (prefill_workspace_starts) { + starts_at = adaptor::to_aten_tensor(*prefill_workspace_starts); + } + adaptor::vllm_iluvatar::map_prefill_request_block_indices( + output_at, req_at, blocks_at, indices_at, block_size, + has_prefill_workspace, request_ids_at, starts_at); + return; +#else + throw std::runtime_error("map_prefill_request_block_indices requires an Iluvatar ATen build"); +#endif +} + +void map_decode_request_block_indices_( + Tensor output, + const Tensor &req_id, + const Tensor &block_table, + const Tensor &token_indices, + int64_t block_size) { + require_iluvatar(output, "map_decode_request_block_indices"); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, req_id, block_table, token_indices); + if (block_size <= 0 || output->shape() != token_indices->shape()) { + throw std::runtime_error("map_decode_request_block_indices invalid output shape or block size"); + } + if (defer_if_recording([output, req_id, block_table, token_indices, block_size] { + map_decode_request_block_indices_( + output, req_id, block_table, token_indices, block_size); + })) { + return; + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + auto output_at = adaptor::to_aten_tensor(output); + auto req_at = adaptor::to_aten_tensor(req_id); + auto blocks_at = adaptor::to_aten_tensor(block_table); + auto indices_at = adaptor::to_aten_tensor(token_indices); + adaptor::vllm_iluvatar::map_decode_request_block_indices( + output_at, req_at, blocks_at, indices_at, block_size); + return; +#else + throw std::runtime_error("map_decode_request_block_indices requires an Iluvatar ATen build"); +#endif +} + +void topk_indices_context_lens_(Tensor topk_lens, const Tensor &indices) { + require_iluvatar(indices, "topk_indices_context_lens"); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_lens, indices); + require_i32(topk_lens, "topk_lens"); + require_i32(indices, "indices"); + if (defer_if_recording([topk_lens, indices] { + topk_indices_context_lens_(topk_lens, indices); + })) { + return; + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + auto lens_at = adaptor::to_aten_tensor(topk_lens); + auto indices_at = adaptor::to_aten_tensor(indices); + adaptor::vllm_iluvatar::topk_indices_context_lens(lens_at, indices_at); + return; +#else + throw std::runtime_error("topk_indices_context_lens requires an Iluvatar ATen build"); +#endif +} + +void sparse_flash_mla_( + Tensor output, + const Tensor &query, + const Tensor &kv_cache, + const Tensor &indices, + const Tensor &topk_lens, + float scale, + std::optional attn_sink) { + require_iluvatar(output, "sparse_flash_mla"); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, query, kv_cache, indices, topk_lens); + if (query->ndim() != 3 || output->ndim() != 3 || kv_cache->ndim() != 3 + || indices->ndim() != 3 || topk_lens->ndim() != 1) { + throw std::runtime_error("sparse_flash_mla tensor rank mismatch"); + } + require_i32(indices, "indices"); + require_i32(topk_lens, "topk_lens"); + if (kv_cache->dtype() == DataType::U8) { + if (attn_sink.has_value()) { + throw std::runtime_error( + "fp8 sparse MLA does not support attention sinks"); + } + fp8_sparse_mla_( + output, query, kv_cache, indices, topk_lens, scale); + return; + } + if (defer_if_recording([output, query, kv_cache, indices, topk_lens, scale, attn_sink] { + sparse_flash_mla_( + output, query, kv_cache, indices, topk_lens, scale, + attn_sink); + })) { + return; + } +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (!adaptor::vllm_iluvatar::sparse_flash_mla_available()) { + throw std::runtime_error("Iluvatar sparse FlashMLA extension is unavailable"); + } + auto output_at = adaptor::to_aten_tensor(output); + auto query_at = adaptor::to_aten_tensor(query); + auto cache_at = adaptor::to_aten_tensor(kv_cache); + auto indices_at = adaptor::to_aten_tensor(indices); + auto lens_at = adaptor::to_aten_tensor(topk_lens); + std::optional sink_at; + if (attn_sink) { + sink_at = adaptor::to_aten_tensor(*attn_sink); + } + adaptor::vllm_iluvatar::sparse_flash_mla( + output_at, query_at, cache_at, indices_at, lens_at, scale, sink_at); + return; +#else + throw std::runtime_error("sparse_flash_mla requires an Iluvatar ATen build"); +#endif +} + +} // namespace infinicore::op 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/fused_rotary_embedding/fused_rotary_embedding.cc b/src/infinicore/ops/fused_rotary_embedding/fused_rotary_embedding.cc new file mode 100644 index 000000000..b0f3daa20 --- /dev/null +++ b/src/infinicore/ops/fused_rotary_embedding/fused_rotary_embedding.cc @@ -0,0 +1,131 @@ +#include "infinicore/ops/fused_rotary_embedding.hpp" + +#include "../../utils.hpp" +#include "infinicore/context/context.hpp" +#include "infinicore/graph/graph.hpp" + +#include +#include +#include + +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/vllm_iluvatar_adaptor.hpp" +#endif + +namespace infinicore::op { +namespace { + +class DeferredGraphOperator final : public graph::GraphOperator { +public: + explicit DeferredGraphOperator(std::function runner) + : runner_(std::move(runner)) {} + + void run() const override { runner_(); } + +private: + std::function runner_; +}; + +void record_or_run(std::function runner) { + auto op = std::make_shared(std::move(runner)); + if (context::isGraphRecording()) { + context::addGraphOperator(op); + } else { + op->run(); + } +} + +void validate_fused_rotary_embedding(const Tensor &query, + const Tensor &key, + const Tensor &positions, + int64_t head_size, + const Tensor &cos_sin_cache) { + if (!query || !key || !positions || !cos_sin_cache) { + throw std::runtime_error( + "fused_rotary_embedding expects non-empty query, key, positions, and cos_sin_cache tensors"); + } + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(query, key, positions, cos_sin_cache); + if (query->dtype() != key->dtype() || query->dtype() != cos_sin_cache->dtype()) { + throw std::runtime_error( + "fused_rotary_embedding expects query, key, and cos_sin_cache to have the same dtype"); + } + if (query->dtype() != DataType::F16 && query->dtype() != DataType::BF16) { + throw std::runtime_error( + "fused_rotary_embedding expects float16 or bfloat16 query and key tensors"); + } + if (positions->dtype() != DataType::I64) { + throw std::runtime_error("fused_rotary_embedding expects int64 positions"); + } + if (query->ndim() != 3 || key->ndim() != 3 || positions->ndim() != 1 + || cos_sin_cache->ndim() != 2) { + throw std::runtime_error( + "fused_rotary_embedding expects query/key [tokens,heads,dim], positions [tokens], and cache [positions,dim]"); + } + if (head_size <= 0 + || query->size(0) != key->size(0) + || query->size(0) != positions->numel() + || query->size(2) != static_cast(head_size) + || key->size(2) != static_cast(head_size) + || cos_sin_cache->size(1) != static_cast(head_size)) { + throw std::runtime_error( + "fused_rotary_embedding tensor shapes are incompatible: head_size=" + + std::to_string(head_size) + + ", query=[" + std::to_string(query->size(0)) + "," + + std::to_string(query->size(1)) + "," + std::to_string(query->size(2)) + + "], key=[" + std::to_string(key->size(0)) + "," + + std::to_string(key->size(1)) + "," + std::to_string(key->size(2)) + + "], positions_numel=" + std::to_string(positions->numel()) + + ", cache=[" + std::to_string(cos_sin_cache->size(0)) + "," + + std::to_string(cos_sin_cache->size(1)) + "]"); + } +} + +void run_fused_rotary_embedding(Tensor query, + Tensor key, + const Tensor &positions, + int64_t head_size, + const Tensor &cos_sin_cache, + bool is_neox) { +#if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) + if (query->device().getType() == Device::Type::ILUVATAR) { + auto query_at = adaptor::to_aten_tensor(query); + auto key_at = adaptor::to_aten_tensor(key); + auto positions_at = adaptor::to_aten_tensor(positions); + auto cos_sin_cache_at = adaptor::to_aten_tensor(cos_sin_cache); + adaptor::vllm_iluvatar::rotary_embedding( + positions_at, + query_at, + std::optional(key_at), + head_size, + cos_sin_cache_at, + is_neox); + return; + } +#endif + + throw std::runtime_error( + "fused_rotary_embedding currently supports only Iluvatar builds with ATen"); +} + +} // namespace + +void fused_rotary_embedding_(Tensor query, + Tensor key, + const Tensor &positions, + int64_t head_size, + const Tensor &cos_sin_cache, + bool is_neox) { + validate_fused_rotary_embedding(query, key, positions, head_size, cos_sin_cache); + record_or_run([query, key, positions, head_size, cos_sin_cache, is_neox] { + run_fused_rotary_embedding( + query, + key, + positions, + head_size, + cos_sin_cache, + is_neox); + }); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/grouped_topk_vllm/grouped_topk_vllm.cc b/src/infinicore/ops/grouped_topk_vllm/grouped_topk_vllm.cc index 7c9cfb9ec..956072ab1 100644 --- a/src/infinicore/ops/grouped_topk_vllm/grouped_topk_vllm.cc +++ b/src/infinicore/ops/grouped_topk_vllm/grouped_topk_vllm.cc @@ -1,5 +1,6 @@ #include "infinicore/ops/grouped_topk_vllm.hpp" #include "../../utils.hpp" +#include "infinicore/ops/mul_scalar.hpp" #include #if defined(ENABLE_ATEN) && defined(ENABLE_ILUVATAR_API) #include "infinicore/adaptor/aten_adaptor.hpp" @@ -72,7 +73,7 @@ void grouped_topk_vllm_(Tensor topk_weights, Tensor topk_ids, const Tensor &scor } adaptor::vllm_iluvatar::grouped_topk(w, ids, s, b, num_expert_group, topk_group, scoring_func, renormalize); if (routed_scaling_factor != 1.0f) { - w.mul_(routed_scaling_factor); + mul_scalar_(topk_weights, topk_weights, routed_scaling_factor); } return; } 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/ops/vocab_parallel_embedding/vocab_parallel_embedding.cc b/src/infinicore/ops/vocab_parallel_embedding/vocab_parallel_embedding.cc index e65152cb5..3babd454d 100644 --- a/src/infinicore/ops/vocab_parallel_embedding/vocab_parallel_embedding.cc +++ b/src/infinicore/ops/vocab_parallel_embedding/vocab_parallel_embedding.cc @@ -1,26 +1,61 @@ #include "infinicore/ops/vocab_parallel_embedding.hpp" #include "../../utils.hpp" +#include "infinicore/context/context.hpp" +#include "infinicore/graph/graph.hpp" +#include +#include #include #if defined(ENABLE_ATEN) #include "infinicore/adaptor/aten_adaptor.hpp" #include #endif namespace infinicore::op { -void vocab_parallel_embedding_(Tensor out, const Tensor &ids, const Tensor &w, int64_t start, int64_t end) { - if (!out || !ids || !w || w->ndim() != 2) { - throw std::runtime_error("vocab_parallel_embedding_: invalid tensors"); +namespace { + +class DeferredGraphOperator final : public graph::GraphOperator { +public: + explicit DeferredGraphOperator(std::function runner) + : runner_(std::move(runner)) {} + + void run() const override { runner_(); } + +private: + std::function runner_; +}; + +void record_or_run(std::function runner) { + auto op = std::make_shared(std::move(runner)); + if (context::isGraphRecording()) { + context::addGraphOperator(op); + } else { + op->run(); } - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, ids, w); +} + +void run_vocab_parallel_embedding( + Tensor out, const Tensor &ids, const Tensor &w, int64_t start, int64_t end) { #if defined(ENABLE_ATEN) + adaptor::set_aten_stream_to_infinicore(); auto o = adaptor::to_aten_tensor(out), i = adaptor::to_aten_tensor(ids), wt = adaptor::to_aten_tensor(w); auto mask = i.lt(start).logical_or(i.ge(end)); auto local = (i - start).clamp(0, end - start - 1); auto y = at::embedding(wt, local, -1, false, false); y.masked_fill_(mask.unsqueeze(-1), 0); o.copy_(y); - return; #else throw std::runtime_error("vocab_parallel_embedding_ requires ATen"); #endif } + +} // namespace + +void vocab_parallel_embedding_(Tensor out, const Tensor &ids, const Tensor &w, int64_t start, int64_t end) { + if (!out || !ids || !w || w->ndim() != 2) { + throw std::runtime_error("vocab_parallel_embedding_: invalid tensors"); + } + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, ids, w); + record_or_run([out, ids, w, start, end] { + run_vocab_parallel_embedding(out, ids, w, start, end); + }); +} } // namespace infinicore::op diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp index 129ed528d..9e70f466b 100644 --- a/src/infinicore/pybind11/ops.hpp +++ b/src/infinicore/pybind11/ops.hpp @@ -44,6 +44,7 @@ #include "ops/diff.hpp" #include "ops/digamma.hpp" #include "ops/dist.hpp" +#include "ops/dsa.hpp" #include "ops/dynamic_scaled_int8_quant.hpp" #include "ops/embedding.hpp" #include "ops/equal.hpp" @@ -54,6 +55,8 @@ #include "ops/floor_divide.hpp" #include "ops/fmin.hpp" #include "ops/fmod.hpp" +#include "ops/fp8_indexer_quant.hpp" +#include "ops/fp8_mla_rmsnorm_cache.hpp" #include "ops/fused_gated_delta_net_gating.hpp" #include "ops/gaussian_nll_loss.hpp" #include "ops/grouped_topk_vllm.hpp" @@ -182,7 +185,10 @@ inline void bind(py::module &m) { bind_digamma(m); bind_dist(m); bind_dynamic_scaled_int8_quant(m); + bind_dsa(m); bind_flash_attention(m); + bind_fp8_indexer_quant(m); + bind_fp8_mla_rmsnorm_cache(m); bind_hinge_embedding_loss(m); bind_kv_caching(m); bind_fmod(m); diff --git a/src/infinicore/pybind11/ops/dsa.hpp b/src/infinicore/pybind11/ops/dsa.hpp new file mode 100644 index 000000000..21e48a4c2 --- /dev/null +++ b/src/infinicore/pybind11/ops/dsa.hpp @@ -0,0 +1,43 @@ +#pragma once + +#include "infinicore/ops/dsa.hpp" +#include + +namespace py = pybind11; + +namespace infinicore { +inline void bind_dsa(py::module &m) { + m.def("fused_deepseek_v2_indexer_postprocess_", + &op::fused_deepseek_v2_indexer_postprocess_); + m.def("indexer_k_cache_", &op::indexer_k_cache_); + m.def("compute_block_sparse_mqa_logits_", + &op::compute_block_sparse_mqa_logits_); + m.def("select_prefill_topk_block_indices_", + &op::select_prefill_topk_block_indices_); + m.def("select_decode_topk_block_indices_", + &op::select_decode_topk_block_indices_); + m.def("map_prefill_request_block_indices_", + &op::map_prefill_request_block_indices_, + py::arg("output"), + py::arg("req_id"), + py::arg("block_table"), + py::arg("token_indices"), + py::arg("block_size"), + py::arg("has_prefill_workspace") = false, + py::arg("prefill_workspace_request_ids") = std::nullopt, + py::arg("prefill_workspace_starts") = std::nullopt); + m.def("map_decode_request_block_indices_", + &op::map_decode_request_block_indices_); + m.def("topk_indices_context_lens_", + &op::topk_indices_context_lens_); + m.def("sparse_flash_mla_", + &op::sparse_flash_mla_, + py::arg("output"), + py::arg("query"), + py::arg("kv_cache"), + py::arg("indices"), + py::arg("topk_lens"), + py::arg("scale"), + py::arg("attn_sink") = std::nullopt); +} +} // namespace infinicore 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/tensor/tensor.cc b/src/infinicore/tensor/tensor.cc index e7e54aa59..25cc63ff0 100644 --- a/src/infinicore/tensor/tensor.cc +++ b/src/infinicore/tensor/tensor.cc @@ -4,6 +4,11 @@ #include "infinicore/context/context.hpp" #include "infinicore/dtype.hpp" +#include +#include +#include +#include + #include namespace { @@ -252,8 +257,88 @@ std::shared_ptr TensorImpl::ones(const Shape &shape, const DataType &dtype, const Device &device, bool pin_memory) { - // TODO: Implement this. - return empty(shape, dtype, device, pin_memory); + auto result = empty(shape, dtype, device, pin_memory); + if (result->nbytes() == 0) { + return result; + } + + std::array one{}; + auto set_pattern = [&](const auto &value) { + if (sizeof(value) != result->element_size()) { + throw std::runtime_error("Invalid one pattern size for " + + toString(dtype)); + } + std::memcpy(one.data(), &value, sizeof(value)); + }; + + switch (dtype) { + case DataType::BYTE: + case DataType::BOOL: + case DataType::I8: + case DataType::U8: + set_pattern(uint8_t{1}); + break; + case DataType::F8: + set_pattern(uint8_t{0x38}); + break; + case DataType::I16: + case DataType::U16: + set_pattern(uint16_t{1}); + break; + case DataType::F16: + set_pattern(uint16_t{0x3c00}); + break; + case DataType::BF16: + set_pattern(uint16_t{0x3f80}); + break; + case DataType::I32: + case DataType::U32: + set_pattern(uint32_t{1}); + break; + case DataType::F32: + set_pattern(float{1.0f}); + break; + case DataType::I64: + case DataType::U64: + set_pattern(uint64_t{1}); + break; + case DataType::F64: + set_pattern(double{1.0}); + break; + case DataType::C16: { + const uint8_t value[2] = {0x38, 0}; + set_pattern(value); + break; + } + case DataType::C32: { + const uint16_t value[2] = {0x3c00, 0}; + set_pattern(value); + break; + } + case DataType::C64: { + const float value[2] = {1.0f, 0.0f}; + set_pattern(value); + break; + } + case DataType::C128: { + const double value[2] = {1.0, 0.0}; + set_pattern(value); + break; + } + } + + std::vector host(result->nbytes()); + for (size_t i = 0; i < result->numel(); ++i) { + std::memcpy(host.data() + i * result->element_size(), + one.data(), + result->element_size()); + } + if (device.getType() == Device::Type::CPU) { + std::memcpy(result->data(), host.data(), host.size()); + } else { + context::memcpyH2D(result->data(), host.data(), host.size(), false); + } + return result; } std::shared_ptr TensorImpl::from_blob( 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..4d16c92ca --- /dev/null +++ b/src/infiniop/ops/fp8_indexer_logits/nvidia/fp8_indexer_logits_nvidia.cu @@ -0,0 +1,177 @@ +#include "../../../../utils.h" +#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; +constexpr size_t LANES_PER_KEY = 4; + +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; + const size_t logical_block = blockIdx.y; + const size_t lane = threadIdx.x % LANES_PER_KEY; + const size_t key_in_block = 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; + } + + float acc = 0.0f; + if (valid) { + const uint8_t *cache_block = kv_cache + + static_cast(physical_block) * block_size * cache_stride; + const auto *key = reinterpret_cast( + cache_block + key_in_block * head_dim); + const float key_scale = *reinterpret_cast( + cache_block + block_size * head_dim + key_in_block * sizeof(float)); + for (size_t head = 0; head < num_heads; ++head) { + float dot = 0.0f; + 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 += __shfl_xor_sync(0xffffffffu, dot, 1, LANES_PER_KEY); + dot += __shfl_xor_sync(0xffffffffu, dot, 2, LANES_PER_KEY); + if (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 { + const dim3 grid( + static_cast(_num_tokens), + static_cast((_max_context_len + _block_size - 1) / _block_size)); + const size_t smem = _num_heads * _head_dim + _num_heads * sizeof(float); + fp8IndexerLogitsKernel<<(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 +} 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/infinicore/ops/dsa.py b/test/infinicore/ops/dsa.py new file mode 100644 index 000000000..329d26c83 --- /dev/null +++ b/test/infinicore/ops/dsa.py @@ -0,0 +1,213 @@ +import torch +from ixtriturbo._C import ops as ixops +from vllm_iluvatar.custom_kernels.compute_block_sparse_mqa_logits.perf_impl import ( + compute_block_sparse_mqa_logits_perf, +) +from vllm_iluvatar.custom_kernels.fused_deepseek_v2_indexer_postprocess.perf_impl import ( + fused_deepseek_v2_indexer_postprocess_perf, +) +from vllm_iluvatar.custom_kernels.indexer_k_cache.perf_impl import ( + indexer_k_cache_perf, +) +from vllm_iluvatar.custom_kernels.map_prefill_request_block_indices_to_global_blocks.perf_impl import ( + map_decode_request_block_indices_to_global_blocks_perf, + map_prefill_request_block_indices_to_global_blocks_perf, +) +from vllm_iluvatar.custom_kernels.select_prefill_topk_block_indices.perf_impl import ( + select_decode_topk_block_indices_perf, + select_prefill_topk_block_indices_perf, +) + +import infinicore + + +def ic(t): + return infinicore.from_torch(t) + + +dev = "cuda" + +# block sparse MQA logits +q = torch.randn(4, 64, 128, device=dev, dtype=torch.bfloat16) +cache = torch.randn(2, 64, 128, device=dev, dtype=torch.bfloat16) +cuq = torch.tensor([0, 4], device=dev, dtype=torch.int32) +cukv = torch.tensor([0, 4], device=dev, dtype=torch.int32) +blocks = torch.tensor([[0]], device=dev, dtype=torch.int32) +weights = torch.randn(4, 64, device=dev, dtype=torch.bfloat16) +a = torch.zeros(4, 64, device=dev, dtype=torch.float32) +b = torch.zeros_like(a) +compute_block_sparse_mqa_logits_perf(q, cuq, cukv, cache, blocks, weights, a, 4, 4, 64) +torch.cuda.synchronize() +infinicore.compute_block_sparse_mqa_logits_( + ic(b), ic(q), ic(cache), ic(cuq), ic(cukv), ic(blocks), ic(weights), 4, 4, 64 +) +infinicore.sync_device() +torch.testing.assert_close(a, b) +print("block_sparse_logits ok") + +# prefill topk +ks = torch.tensor([0, 0, 0, 0], device=dev, dtype=torch.int32) +ke = torch.tensor([1, 2, 3, 4], device=dev, dtype=torch.int32) +a = torch.empty(4, 8, device=dev, dtype=torch.int32) +b = torch.empty_like(a) +select_prefill_topk_block_indices_perf( + a.new_empty((4, 64), dtype=torch.float32).normal_(), ks, ke, a +) +# share identical logits for comparison +logits = torch.randn(4, 64, device=dev) +select_prefill_topk_block_indices_perf(logits, ks, ke, a) +torch.cuda.synchronize() +infinicore.select_prefill_topk_block_indices_(ic(b), ic(logits), ic(ks), ic(ke)) +infinicore.sync_device() +torch.testing.assert_close(a, b) +print("select_prefill_topk ok") + +# prefill mapping without workspace +req = torch.tensor([0, 0, 1, 1], device=dev, dtype=torch.int32) +bt = torch.tensor([[2, 3], [5, 6]], device=dev, dtype=torch.int32) +idx = torch.tensor( + [[0, 63, 64, -1], [1, 65, -1, -1], [0, 64, -1, -1], [3, 66, -1, -1]], + device=dev, + dtype=torch.int32, +) +a = map_prefill_request_block_indices_to_global_blocks_perf(req, bt, idx, 64) +b = torch.empty_like(idx) +torch.cuda.synchronize() +infinicore.map_prefill_request_block_indices_(ic(b), ic(req), ic(bt), ic(idx), 64) +infinicore.sync_device() +torch.testing.assert_close(a, b) +print("map_prefill ok") + +# fused indexer postprocess +T, H, D, R = 3, 64, 128, 64 +q = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) +kw = torch.randn(T, D + H, device=dev, dtype=torch.bfloat16) +nw = torch.randn(D, device=dev, dtype=torch.bfloat16) +nb = torch.randn(D, device=dev, dtype=torch.bfloat16) +pos = torch.arange(T, device=dev, dtype=torch.int64) +cs = torch.randn(16, R, device=dev, dtype=torch.bfloat16) +slots = torch.tensor([0, 65, 130], device=dev, dtype=torch.int64) +qa = torch.empty_like(q) +qb = torch.empty_like(q) +ka = torch.empty(0, D, device=dev, dtype=torch.bfloat16) +kb = torch.empty_like(ka) +wa = torch.empty(T, H, device=dev, dtype=torch.bfloat16) +wb = torch.empty_like(wa) +ca = torch.zeros(4, 64, D, device=dev, dtype=torch.bfloat16) +cb = ca.clone() +fused_deepseek_v2_indexer_postprocess_perf( + qa, ka, wa, ca, slots, q, kw, nw, nb, pos, cs, T, False, 1e-6, 0.01 +) +torch.cuda.synchronize() +infinicore.fused_deepseek_v2_indexer_postprocess_( + ic(qb), + ic(kb), + ic(wb), + ic(cb), + ic(slots), + ic(q), + ic(kw), + ic(nw), + ic(nb), + ic(pos), + ic(cs), + T, + False, + 1e-6, + 0.01, +) +infinicore.sync_device() +torch.testing.assert_close(qa, qb) +infinicore.sync_device() +torch.testing.assert_close(wa, wb) +infinicore.sync_device() +torch.testing.assert_close(ca, cb) +print("fused_indexer_postprocess ok") + +# production ixtriturbo sparse MLA v2 +T, H, D, V, K, N = 1, 64, 576, 512, 64, 64 +query = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) +kv = torch.randn(N, 1, D, device=dev, dtype=torch.bfloat16) +indices = torch.arange(K, device=dev, dtype=torch.int32).view(T, 1, K) +lens = torch.tensor([K], device=dev, dtype=torch.int32) +a = torch.empty(T, H, V, device=dev, dtype=torch.bfloat16) +b = torch.empty_like(a) +ixops.flash_mla_sparse_v2(a, query, kv, indices, lens, float(D**-0.5), None) +torch.cuda.synchronize() +infinicore.sparse_flash_mla_( + ic(b), ic(query), ic(kv), ic(indices), ic(lens), float(D**-0.5) +) +infinicore.sync_device() +torch.testing.assert_close(a, b) +print("sparse_flash_mla_v2 ok") + + +# indexer cache +keys = torch.randn(3, 128, device=dev, dtype=torch.bfloat16) +slots = torch.tensor([0, 65, 130], device=dev, dtype=torch.int64) +cache_ref = torch.zeros(4, 64, 128, device=dev, dtype=torch.bfloat16) +cache_out = cache_ref.clone() +indexer_k_cache_perf(keys, cache_ref, slots) +torch.cuda.synchronize() +infinicore.indexer_k_cache_(ic(keys), ic(cache_out), ic(slots)) +infinicore.sync_device() +torch.testing.assert_close(cache_ref, cache_out) +print("indexer_k_cache ok") + +# decode topk +logits = torch.randn(2, 128, device=dev, dtype=torch.float32) +seq_lens = torch.tensor([64, 96], device=dev, dtype=torch.int32) +topk_ref = torch.empty(2, 16, device=dev, dtype=torch.int32) +topk_out = torch.empty_like(topk_ref) +select_decode_topk_block_indices_perf(logits, seq_lens, topk_ref) +torch.cuda.synchronize() +infinicore.select_decode_topk_block_indices_(ic(topk_out), ic(logits), ic(seq_lens)) +infinicore.sync_device() +torch.testing.assert_close(topk_ref, topk_out) +print("select_decode_topk ok") + +# decode mapping +req = torch.tensor([0, 1], device=dev, dtype=torch.int32) +block_table = torch.tensor([[2, 3], [5, 6]], device=dev, dtype=torch.int32) +token_indices = torch.tensor( + [[0, 63, 64, -1], [3, 66, -1, -1]], + device=dev, + dtype=torch.int32, +) +mapped_ref = map_decode_request_block_indices_to_global_blocks_perf( + req, block_table, token_indices, 64 +) +mapped_out = torch.empty_like(token_indices) +torch.cuda.synchronize() +infinicore.map_decode_request_block_indices_( + ic(mapped_out), ic(req), ic(block_table), ic(token_indices), 64 +) +infinicore.sync_device() +torch.testing.assert_close(mapped_ref, mapped_out) +print("map_decode ok") + +# valid sparse topk lengths +indices = torch.tensor( + [[[1, 2, -1, -1]], [[3, 4, 5, -1]]], + device=dev, + dtype=torch.int32, +) +lens_ref = torch.empty(2, device=dev, dtype=torch.int32) +lens_out = torch.empty_like(lens_ref) +ixops.topk_indices_context_lens(lens_ref, indices) +torch.cuda.synchronize() +infinicore.topk_indices_context_lens_(ic(lens_out), ic(indices)) +infinicore.sync_device() +torch.testing.assert_close(lens_ref, lens_out) +print("topk_indices_context_lens ok") + +# InfiniCore graph record/replay keeps the bridge on the current InfiniCore stream. +graph_cache = torch.zeros_like(cache_ref) +torch.cuda.synchronize() +infinicore.start_graph_recording() +infinicore.indexer_k_cache_(ic(keys), ic(graph_cache), ic(slots)) +graph = infinicore.stop_graph_recording() +graph.run() +infinicore.sync_device() +torch.testing.assert_close(cache_ref, graph_cache) +print("graph_replay ok") diff --git a/test/infinicore/ops/fp8_mla_rmsnorm_dual_cache.py b/test/infinicore/ops/fp8_mla_rmsnorm_dual_cache.py new file mode 100644 index 000000000..8046bac68 --- /dev/null +++ b/test/infinicore/ops/fp8_mla_rmsnorm_dual_cache.py @@ -0,0 +1,221 @@ +import time + +import torch +from infinicore.lib import _infinicore + +import infinicore + + +def ic(tensor): + return infinicore.from_torch(tensor) + + +def raw(tensor): + return ic(tensor)._underlying + + +def sync(): + infinicore.sync_device() + + +def produce(cache, vendor_cache, compressed_kv, norm_weight, rope, slots): + _infinicore.fp8_mla_rmsnorm_dual_cache_( + raw(cache), + raw(vendor_cache), + raw(compressed_kv), + raw(norm_weight), + raw(rope), + raw(slots), + 1e-5, + ) + sync() + + +def bench(fn, warmup=5, iterations=20): + for _ in range(warmup): + fn() + sync() + begin = time.perf_counter() + for _ in range(iterations): + fn() + sync() + return (time.perf_counter() - begin) * 1000.0 / iterations + + +def validate_dual_cache(): + torch.manual_seed(7) + num_tokens = 64 + compressed_kv = torch.randn(num_tokens, 512, device="cuda", dtype=torch.bfloat16) + norm_weight = torch.randn(512, device="cuda", dtype=torch.bfloat16) + rope = torch.randn(num_tokens, 64, device="cuda", dtype=torch.bfloat16) + slots = torch.arange(num_tokens, device="cuda", dtype=torch.int64) + cache = torch.zeros(1, 64, 656, device="cuda", dtype=torch.uint8) + vendor_cache = torch.zeros(1, 64, 576, device="cuda", dtype=torch.bfloat16) + produce(cache, vendor_cache, compressed_kv, norm_weight, rope, slots) + + entries = cache.view(-1, 656)[slots] + latent_fp8 = entries[:, :512].contiguous().view(torch.float8_e4m3fn) + scales = entries[:, 512:528].contiguous().view(torch.float32) + expected_latent = ( + (latent_fp8.float().view(num_tokens, 4, 128) * scales.view(num_tokens, 4, 1)) + .reshape(num_tokens, 512) + .to(torch.bfloat16) + ) + expected_rope = entries[:, 528:].contiguous().view(torch.bfloat16) + vendor_entries = vendor_cache.view(-1, 576)[slots] + torch.testing.assert_close(vendor_entries[:, :512], expected_latent, rtol=0, atol=0) + torch.testing.assert_close(vendor_entries[:, 512:], expected_rope, rtol=0, atol=0) + print("dual cache byte semantics ok") + + num_heads = 64 + query = torch.randn(1, num_heads, 576, device="cuda", dtype=torch.bfloat16) + indices = torch.arange(num_tokens, device="cuda", dtype=torch.int32).view( + 1, 1, num_tokens + ) + lens = torch.tensor([num_tokens], device="cuda", dtype=torch.int32) + output_fp8 = torch.empty(1, num_heads, 512, device="cuda", dtype=torch.bfloat16) + output_vendor = torch.empty_like(output_fp8) + scale = float(576**-0.5) + infinicore.sparse_flash_mla_( + ic(output_fp8), + ic(query), + ic(cache.view(-1, 1, 656)), + ic(indices), + ic(lens), + scale, + ) + infinicore.sparse_flash_mla_( + ic(output_vendor), + ic(query), + ic(vendor_cache.view(-1, 1, 576)), + ic(indices), + ic(lens), + scale, + ) + sync() + diff = (output_fp8.float() - output_vendor.float()).abs() + print( + "sparse output diff:", + f"max={diff.max().item():.6f}", + f"mean={diff.mean().item():.6f}", + ) + torch.testing.assert_close(output_vendor, output_fp8, rtol=0.05, atol=0.05) + print("sparse FP8/vendor numerical comparison ok") + + +def benchmark_producer(): + torch.manual_seed(9) + for num_tokens in (1, 4, 16): + compressed_kv = torch.randn( + num_tokens, 512, device="cuda", dtype=torch.bfloat16 + ) + norm_weight = torch.randn(512, device="cuda", dtype=torch.bfloat16) + rope = torch.randn(num_tokens, 64, device="cuda", dtype=torch.bfloat16) + slots = torch.arange(num_tokens, device="cuda", dtype=torch.int64) + cache = torch.zeros(1, 64, 656, device="cuda", dtype=torch.uint8) + vendor_cache = torch.zeros(1, 64, 576, device="cuda", dtype=torch.bfloat16) + cache_raw = raw(cache) + vendor_raw = raw(vendor_cache) + compressed_raw = raw(compressed_kv) + weight_raw = raw(norm_weight) + rope_raw = raw(rope) + slots_raw = raw(slots) + + def run_fp8(): + _infinicore.fp8_mla_rmsnorm_cache_( + cache_raw, + compressed_raw, + weight_raw, + rope_raw, + slots_raw, + 1e-5, + ) + + def run_dual(): + _infinicore.fp8_mla_rmsnorm_dual_cache_( + cache_raw, + vendor_raw, + compressed_raw, + weight_raw, + rope_raw, + slots_raw, + 1e-5, + ) + + fp8_ms = bench(run_fp8, warmup=10, iterations=200) + dual_ms = bench(run_dual, warmup=10, iterations=200) + print( + f"producer T={num_tokens}: fp8={fp8_ms:.4f}ms " + f"dual={dual_ms:.4f}ms delta={dual_ms - fp8_ms:.4f}ms" + ) + + +def benchmark_sparse(): + torch.manual_seed(11) + topk = 2048 + num_heads = 64 + compressed_kv = torch.randn(topk, 512, device="cuda", dtype=torch.bfloat16) + norm_weight = torch.randn(512, device="cuda", dtype=torch.bfloat16) + rope = torch.randn(topk, 64, device="cuda", dtype=torch.bfloat16) + slots = torch.arange(topk, device="cuda", dtype=torch.int64) + cache = torch.zeros(32, 64, 656, device="cuda", dtype=torch.uint8) + vendor_cache = torch.zeros(32, 64, 576, device="cuda", dtype=torch.bfloat16) + produce(cache, vendor_cache, compressed_kv, norm_weight, rope, slots) + scale = float(576**-0.5) + + for num_tokens in (1, 4, 16): + query = torch.randn( + num_tokens, + num_heads, + 576, + device="cuda", + dtype=torch.bfloat16, + ) + indices = ( + torch.arange(topk, device="cuda", dtype=torch.int32) + .view(1, 1, topk) + .expand(num_tokens, 1, topk) + .contiguous() + ) + lens = torch.full((num_tokens,), topk, device="cuda", dtype=torch.int32) + output_fp8 = torch.empty( + num_tokens, + num_heads, + 512, + device="cuda", + dtype=torch.bfloat16, + ) + output_vendor = torch.empty_like(output_fp8) + + def run_fp8(): + infinicore.sparse_flash_mla_( + ic(output_fp8), + ic(query), + ic(cache.view(-1, 1, 656)), + ic(indices), + ic(lens), + scale, + ) + + def run_vendor(): + infinicore.sparse_flash_mla_( + ic(output_vendor), + ic(query), + ic(vendor_cache.view(-1, 1, 576)), + ic(indices), + ic(lens), + scale, + ) + + fp8_ms = bench(run_fp8) + vendor_ms = bench(run_vendor) + print( + f"T={num_tokens} K={topk}: fp8={fp8_ms:.3f}ms " + f"vendor={vendor_ms:.3f}ms speedup={fp8_ms / vendor_ms:.2f}x" + ) + + +if __name__ == "__main__": + validate_dual_cache() + benchmark_producer() + benchmark_sparse() 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 7dc838737e30203fea0c297b356eef4502e6168e Mon Sep 17 00:00:00 2001 From: wooway777 Date: Mon, 27 Jul 2026 07:37:17 +0000 Subject: [PATCH 4/4] pepe: optimize indexer --- .../nvidia/fp8_indexer_logits_nvidia.cu | 94 +++++++++++++------ 1 file changed, 64 insertions(+), 30 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 4d16c92ca..740e9d887 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 @@ -1,4 +1,3 @@ -#include "../../../../utils.h" #include "../../../devices/nvidia/nvidia_common.cuh" #include "../../../devices/nvidia/nvidia_kernel_common.cuh" #include "../../../tensor.h" @@ -9,8 +8,7 @@ namespace { constexpr size_t THREADS = 256; -constexpr size_t LANES_PER_KEY = 4; - +template INFINIOP_CUDA_KERNEL fp8IndexerLogitsKernel( float *__restrict__ logits, const cuda_fp8_e4m3 *__restrict__ q_fp8, @@ -28,9 +26,13 @@ INFINIOP_CUDA_KERNEL fp8IndexerLogitsKernel( size_t max_blocks_per_request, size_t max_context_len) { const size_t token = blockIdx.x; - const size_t logical_block = blockIdx.y; + 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 = 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); @@ -62,26 +64,36 @@ INFINIOP_CUDA_KERNEL fp8IndexerLogitsKernel( && 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; - if (valid) { - const uint8_t *cache_block = kv_cache - + static_cast(physical_block) * block_size * cache_stride; - const auto *key = reinterpret_cast( - cache_block + key_in_block * head_dim); - const float key_scale = *reinterpret_cast( - cache_block + block_size * head_dim + key_in_block * sizeof(float)); - for (size_t head = 0; head < num_heads; ++head) { - float dot = 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]); } - dot += __shfl_xor_sync(0xffffffffu, dot, 1, LANES_PER_KEY); - dot += __shfl_xor_sync(0xffffffffu, dot, 2, LANES_PER_KEY); - if (lane == 0) { - acc += fmaxf(dot * key_scale, 0.0f) * shared_weights[head]; - } + } + // 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) { @@ -155,20 +167,42 @@ infiniStatus_t Descriptor::calculate( 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)); + 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); - fp8IndexerLogitsKernel<<(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); + 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;