From 7d5212cd1a2756e2c621b7fdec5f8bc6a4cc5446 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 16 Jul 2026 11:14:40 +0800 Subject: [PATCH 1/6] refactor(runtime)!: adopt modern Infini stack BREAKING CHANGE: legacy InfiniCore runtime and inference APIs are removed in favor of InfiniRT, InfiniOps, and the migrated InfiniLM runtime. --- README.md | 246 ++--- csrc/backends/attention_backends.hpp | 8 +- csrc/config/config_factory.cpp | 13 +- csrc/config/quant_config.cpp | 14 +- csrc/config/quant_config.hpp | 24 +- csrc/debug_utils/tensor_utils.hpp | 54 +- csrc/engine/compiler/graph_compiler.hpp | 27 + csrc/engine/compiler/paged_compiler.cpp | 33 +- .../compiler/static_batching_compiler.cpp | 12 +- .../distributed/communication_group.cpp | 165 +++- .../distributed/communication_group.hpp | 13 +- csrc/engine/infer_engine.hpp | 2 +- csrc/engine/rank_worker.cpp | 6 +- csrc/infinicore/include/infinicore.hpp | 7 + .../infinicore/adaptor/aten_adaptor.hpp | 73 ++ .../adaptor/flash_attention_adaptor.hpp | 139 +++ .../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 | 232 +++++ .../infinicore/analyzer/op_type_registry.hpp | 48 + .../analyzer/optimization_intent.hpp | 187 ++++ .../infinicore/analyzer/phase_detector.hpp | 126 +++ .../infinicore/analyzer/resource_sensor.hpp | 180 ++++ .../include/infinicore/common/LRUCache.hpp | 137 +++ .../include/infinicore/common/hash.hpp | 72 ++ .../include/infinicore/context/context.hpp | 56 ++ csrc/infinicore/include/infinicore/device.hpp | 9 + .../include/infinicore/device_event.hpp | 125 +++ csrc/infinicore/include/infinicore/dtype.hpp | 15 + .../include/infinicore/graph/graph.hpp | 148 +++ csrc/infinicore/include/infinicore/io.hpp | 86 ++ csrc/infinicore/include/infinicore/memory.hpp | 30 + csrc/infinicore/include/infinicore/nn.hpp | 5 + .../include/infinicore/nn/embedding.hpp | 87 ++ .../include/infinicore/nn/layer_norm.hpp | 59 ++ .../include/infinicore/nn/module.hpp | 174 ++++ .../include/infinicore/nn/parameter.hpp | 39 + .../include/infinicore/nn/rmsnorm.hpp | 96 ++ .../infinicore/include/infinicore/nn/rope.hpp | 112 +++ .../infinicore/nn/rope_scaling_configs.hpp | 131 +++ csrc/infinicore/include/infinicore/ops.hpp | 76 ++ .../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 + .../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 + .../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 | 14 + .../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/broadcast_to.hpp | 14 + .../infinicore/include/infinicore/ops/cat.hpp | 9 + .../include/infinicore/ops/causal_conv1d.hpp | 39 + .../include/infinicore/ops/causal_softmax.hpp | 14 + .../include/infinicore/ops/cdist.hpp | 32 + .../infinicore/ops/chunk_gated_delta_rule.hpp | 51 + .../include/infinicore/ops/common/cache.hpp | 83 ++ .../infinicore/ops/common/dispatcher.hpp | 88 ++ .../include/infinicore/ops/common/op.hpp | 5 + .../include/infinicore/ops/conv2d.hpp | 38 + .../include/infinicore/ops/cross_entropy.hpp | 35 + .../include/infinicore/ops/deepseek_moe.hpp | 41 + .../include/infinicore/ops/diff.hpp | 14 + .../include/infinicore/ops/digamma.hpp | 14 + .../include/infinicore/ops/dist.hpp | 14 + .../infinicore/ops/distributed/allgather.hpp | 35 + .../infinicore/ops/distributed/allreduce.hpp | 24 + .../ops/distributed/reduce_scatter.hpp | 46 + .../include/infinicore/ops/embedding.hpp | 13 + .../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 + .../ops/fused_gated_delta_net_gating.hpp | 37 + .../infinicore/ops/gaussian_nll_loss.hpp | 26 + .../include/infinicore/ops/gelu.hpp | 16 + .../include/infinicore/ops/gelutanh.hpp | 16 + .../include/infinicore/ops/gemm.hpp | 14 + .../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 | 28 + .../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 + .../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 + .../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_fused_dense.hpp | 42 + .../include/infinicore/ops/moe_fused_gate.hpp | 44 + .../include/infinicore/ops/moe_sum.hpp | 15 + .../infinicore/ops/moe_topk_sigmoid.hpp | 32 + .../infinicore/ops/moe_topk_softmax.hpp | 35 + .../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 + .../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 | 14 + .../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 | 53 + .../include/infinicore/ops/relu.hpp | 16 + .../include/infinicore/ops/relu6.hpp | 14 + .../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 + .../include/infinicore/ops/scatter.hpp | 21 + .../include/infinicore/ops/selu.hpp | 18 + .../include/infinicore/ops/sigmoid.hpp | 14 + .../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 + csrc/infinicore/include/infinicore/tensor.hpp | 334 +++++++ csrc/infinicore/src/adaptor/aten_adaptor.cc | 56 ++ .../flash_attn/hygon/flash_attn_hygon.cc | 173 ++++ .../flash_attn/hygon/flash_attn_hygon.hpp | 60 ++ .../src/analyzer/mutual_awareness_analyzer.cc | 200 ++++ csrc/infinicore/src/analyzer/op_trace.cc | 10 + .../allocators/device_pinned_allocator.cc | 55 ++ .../allocators/device_pinned_allocator.hpp | 28 + .../src/context/allocators/host_allocator.cc | 18 + .../src/context/allocators/host_allocator.hpp | 15 + .../context/allocators/memory_allocator.hpp | 15 + .../allocators/pinnable_block_allocator.cc | 310 ++++++ .../allocators/pinnable_block_allocator.hpp | 81 ++ .../allocators/stream_ordered_allocator.cc | 36 + .../allocators/stream_ordered_allocator.hpp | 20 + csrc/infinicore/src/context/context_impl.cc | 358 +++++++ csrc/infinicore/src/context/context_impl.hpp | 49 + csrc/infinicore/src/context/internal.hpp | 11 + .../infinicore/src/context/runtime/runtime.cc | 430 ++++++++ .../src/context/runtime/runtime.hpp | 79 ++ csrc/infinicore/src/device.cc | 1 + csrc/infinicore/src/device_event.cc | 165 ++++ csrc/infinicore/src/dtype.cc | 13 + csrc/infinicore/src/graph/graph.cc | 371 +++++++ csrc/infinicore/src/graph/graph_manager.hpp | 37 + csrc/infinicore/src/io.cc | 650 ++++++++++++ csrc/infinicore/src/memory.cc | 33 + csrc/infinicore/src/nn/embedding.cc | 139 +++ csrc/infinicore/src/nn/layer_norm.cc | 27 + csrc/infinicore/src/nn/module.cc | 185 ++++ csrc/infinicore/src/nn/parameter.cc | 106 ++ csrc/infinicore/src/nn/rmsnorm.cc | 47 + csrc/infinicore/src/nn/rope.cc | 266 +++++ .../infinicore/src/nn/rope_scaling_configs.cc | 193 ++++ csrc/infinicore/src/ops/acos/acos.cc | 24 + .../adaptive_avg_pool1d.cc | 43 + .../adaptive_avg_pool3d.cc | 36 + .../adaptive_max_pool1d.cc | 30 + csrc/infinicore/src/ops/add/add.cc | 27 + csrc/infinicore/src/ops/add/add_infiniops.cc | 59 ++ .../src/ops/add_rms_norm/add_rms_norm.cc | 33 + .../add_rms_norm/add_rms_norm_infiniops.cc | 72 ++ csrc/infinicore/src/ops/addbmm/addbmm.cc | 32 + csrc/infinicore/src/ops/addcmul/addcmul.cc | 29 + csrc/infinicore/src/ops/addr/addr.cc | 32 + .../src/ops/affine_grid/affine_grid.cc | 58 ++ csrc/infinicore/src/ops/argwhere/argwhere.cc | 32 + csrc/infinicore/src/ops/asin/asin.cc | 36 + csrc/infinicore/src/ops/asinh/asinh.cc | 27 + csrc/infinicore/src/ops/asum/asum.cc | 28 + csrc/infinicore/src/ops/atanh/atanh.cc | 37 + .../infinicore/src/ops/attention/attention.cc | 40 + .../src/ops/avg_pool1d/avg_pool1d.cc | 68 ++ csrc/infinicore/src/ops/axpy/axpy.cc | 22 + csrc/infinicore/src/ops/baddbmm/baddbmm.cc | 100 ++ csrc/infinicore/src/ops/bilinear/bilinear.cc | 95 ++ .../binary_cross_entropy_with_logits.cc | 67 ++ .../bitwise_right_shift.cc | 27 + .../infinicore/src/ops/blas_amax/blas_amax.cc | 28 + .../infinicore/src/ops/blas_amin/blas_amin.cc | 28 + .../infinicore/src/ops/blas_copy/blas_copy.cc | 22 + csrc/infinicore/src/ops/blas_dot/blas_dot.cc | 28 + .../src/ops/block_diag/block_diag.cc | 63 ++ .../src/ops/broadcast_to/broadcast_to.cc | 28 + csrc/infinicore/src/ops/cat/cat.cc | 199 ++++ .../src/ops/causal_conv1d/causal_conv1d.cc | 146 +++ .../src/ops/causal_softmax/causal_softmax.cc | 27 + .../causal_softmax_infiniops.cc | 58 ++ csrc/infinicore/src/ops/cdist/cdist.cc | 56 ++ .../chunk_gated_delta_rule.cc | 167 ++++ csrc/infinicore/src/ops/conv2d/conv2d.cc | 101 ++ .../src/ops/conv2d/conv2d_infiniops.cc | 72 ++ .../src/ops/cross_entropy/cross_entropy.cc | 45 + .../src/ops/deepseek_moe/deepseek_moe.cc | 83 ++ .../per_tensor_quant_int8.cc | 20 + csrc/infinicore/src/ops/diff/diff.cc | 78 ++ csrc/infinicore/src/ops/digamma/digamma.cc | 28 + csrc/infinicore/src/ops/dist/dist.cc | 30 + .../src/ops/distributed/allgather.cc | 169 ++++ .../src/ops/distributed/allreduce.cc | 56 ++ .../src/ops/distributed/reduce_scatter.cc | 187 ++++ csrc/infinicore/src/ops/distributed/utils.hpp | 56 ++ .../infinicore/src/ops/embedding/embedding.cc | 37 + .../src/ops/embedding/embedding_infiniops.cc | 61 ++ .../ops/flash_attention/flash_attention.cc | 31 + csrc/infinicore/src/ops/flipud/flipud.cc | 27 + .../src/ops/float_power/float_power.cc | 71 ++ csrc/infinicore/src/ops/floor/floor.cc | 22 + .../src/ops/floor_divide/floor_divide.cc | 27 + csrc/infinicore/src/ops/fmin/fmin.cc | 27 + csrc/infinicore/src/ops/fmod/fmod.cc | 28 + .../fused_gated_delta_net_gating.cc | 85 ++ .../gaussian_nll_loss/gaussian_nll_loss.cc | 51 + csrc/infinicore/src/ops/gelu/gelu.cc | 37 + .../infinicore/src/ops/gelu/gelu_infiniops.cc | 41 + csrc/infinicore/src/ops/gelutanh/gelutanh.cc | 37 + .../src/ops/gelutanh/gelutanh_infiniops.cc | 38 + csrc/infinicore/src/ops/gemm/gemm.cc | 30 + .../infinicore/src/ops/gemm/gemm_infiniops.cc | 68 ++ .../infinicore/src/ops/hardswish/hardswish.cc | 38 + csrc/infinicore/src/ops/hardtanh/hardtanh.cc | 38 + .../hinge_embedding_loss.cc | 57 ++ .../src/ops/huber_loss/huber_loss.cc | 35 + csrc/infinicore/src/ops/hypot/hypot.cc | 25 + .../infinicore/src/ops/index_add/index_add.cc | 114 +++ .../src/ops/index_copy/index_copy.cc | 114 +++ csrc/infinicore/src/ops/infiniops_impl.hpp | 103 ++ csrc/infinicore/src/ops/inner/inner.cc | 46 + .../src/ops/interpolate/interpolate.cc | 145 +++ csrc/infinicore/src/ops/kron/kron.cc | 50 + csrc/infinicore/src/ops/kthvalue/kthvalue.cc | 49 + .../src/ops/kv_caching/kv_caching.cc | 42 + .../ops/kv_caching/kv_caching_infiniops.cc | 50 + .../src/ops/layer_norm/layer_norm.cc | 42 + csrc/infinicore/src/ops/ldexp/ldexp.cc | 51 + csrc/infinicore/src/ops/lerp/lerp.cc | 132 +++ csrc/infinicore/src/ops/linear/linear.cc | 60 ++ .../src/ops/linear_w8a8i8/linear_w8a8i8.cc | 66 ++ .../src/ops/log_softmax/log_softmax.cc | 34 + .../infinicore/src/ops/logaddexp/logaddexp.cc | 27 + .../src/ops/logaddexp2/logaddxep2.cc | 27 + .../src/ops/logcumsumexp/logcumsumexp.cc | 35 + csrc/infinicore/src/ops/logdet/logdet.cc | 28 + .../mamba_selective_scan.cc | 28 + .../src/ops/masked_select/masked_select.cc | 33 + csrc/infinicore/src/ops/matmul/matmul.cc | 13 + .../hygon/mha_kvcache_flashattn_hygon.cc | 167 ++++ .../src/ops/mha_kvcache/mha_kvcache.cc | 58 ++ .../ops/mha_kvcache/mha_kvcache_flashattn.cc | 136 +++ .../mha_kvcache_flashattn_moore.cc | 127 +++ .../infinicore/src/ops/moe_align/moe_align.cc | 155 +++ .../ops/moe_fused_dense/moe_fused_dense.cc | 73 ++ .../src/ops/moe_fused_gate/moe_fused_gate.cc | 101 ++ csrc/infinicore/src/ops/moe_sum/moe_sum.cc | 30 + .../ops/moe_topk_sigmoid/moe_topk_sigmoid.cc | 63 ++ .../ops/moe_topk_softmax/moe_topk_softmax.cc | 69 ++ csrc/infinicore/src/ops/mrope/mrope.cc | 100 ++ csrc/infinicore/src/ops/mul/mul.cc | 27 + .../src/ops/mul_scalar/mul_scalar.cc | 27 + .../hygon/mha_flashattn_hygon.cc | 102 ++ .../src/ops/multi_head_attention/mha.cc | 54 + .../ops/multi_head_attention/mha_flashattn.cc | 115 +++ .../hygon/mha_varlen_flashattn_hygon.cc | 266 +++++ .../multi_head_attention_varlen/mha_varlen.cc | 75 ++ .../mha_varlen_flashattn.cc | 191 ++++ .../mha_varlen_flashattn_moore.cc | 139 +++ .../multi_margin_loss/multi_margin_loss.cc | 34 + csrc/infinicore/src/ops/nrm2/nrm2.cc | 28 + .../nsa_compress_paged_cache.cc | 29 + .../nsa_paged_attention.cc | 40 + csrc/infinicore/src/ops/ones/ones.cc | 13 + csrc/infinicore/src/ops/pad/pad.cc | 56 ++ .../ops/paged_attention/paged_attention.cc | 38 + .../paged_attention_infiniops.cc | 88 ++ .../paged_attention_prefill.cc | 39 + .../paged_attention_prefill_infiniops.cc | 68 ++ .../src/ops/paged_caching/paged_caching.cc | 21 + .../paged_caching/paged_caching_infiniops.cc | 50 + .../per_channel_quant_i8.cc | 20 + csrc/infinicore/src/ops/prelu/prelu.cc | 27 + .../prepare_moe_input/prepare_moe_input.cc | 124 +++ .../per_tensor_quant_int8.cc | 26 + .../infinicore/src/ops/quickgelu/quickgelu.cc | 37 + .../src/ops/random_sample/random_sample.cc | 42 + .../random_sample/random_sample_infiniops.cc | 42 + .../infinicore/src/ops/rearrange/rearrange.cc | 33 + .../src/ops/rearrange/rearrange_infiniops.cc | 45 + .../src/ops/reciprocal/reciprocal.cc | 27 + .../recurrent_gated_delta_rule.cc | 166 ++++ csrc/infinicore/src/ops/relu/relu.cc | 37 + .../infinicore/src/ops/relu/relu_infiniops.cc | 38 + csrc/infinicore/src/ops/relu6/relu6.cc | 27 + csrc/infinicore/src/ops/rms_norm/rms_norm.cc | 26 + .../src/ops/rms_norm/rms_norm_infiniops.cc | 64 ++ csrc/infinicore/src/ops/rope/rope.cc | 46 + .../infinicore/src/ops/rope/rope_infiniops.cc | 73 ++ csrc/infinicore/src/ops/rot/rot.cc | 22 + csrc/infinicore/src/ops/rotg/rotg.cc | 22 + csrc/infinicore/src/ops/rotm/rotm.cc | 22 + csrc/infinicore/src/ops/rotmg/rotmg.cc | 22 + .../infinicore/src/ops/rwkv5_wkv/rwkv5_wkv.cc | 50 + csrc/infinicore/src/ops/scal/scal.cc | 22 + .../src/ops/scaled_mm_i8/scaled_mm_i8.cc | 21 + csrc/infinicore/src/ops/scatter/scatter.cc | 26 + csrc/infinicore/src/ops/selu/selu.cc | 40 + csrc/infinicore/src/ops/sigmoid/sigmoid.cc | 27 + .../src/ops/sigmoid/sigmoid_infiniops.cc | 58 ++ csrc/infinicore/src/ops/silu/silu.cc | 37 + .../infinicore/src/ops/silu/silu_infiniops.cc | 38 + .../src/ops/silu_and_mul/silu_and_mul.cc | 35 + .../silu_and_mul/silu_and_mul_infiniops.cc | 44 + csrc/infinicore/src/ops/sinh/sinh.cc | 40 + .../src/ops/smooth_l1_loss/smooth_l1_loss.cc | 34 + csrc/infinicore/src/ops/softmax/softmax.cc | 37 + .../src/ops/softmax/softmax_infiniops.cc | 42 + csrc/infinicore/src/ops/softplus/softplus.cc | 28 + csrc/infinicore/src/ops/softsign/softsign.cc | 27 + csrc/infinicore/src/ops/sum/sum.cc | 67 ++ csrc/infinicore/src/ops/swap/swap.cc | 22 + csrc/infinicore/src/ops/swiglu/swiglu.cc | 26 + .../src/ops/swiglu/swiglu_infiniops.cc | 61 ++ csrc/infinicore/src/ops/take/take.cc | 28 + csrc/infinicore/src/ops/tan/tan.cc | 32 + .../src/ops/tanhshrink/tanhshrink.cc | 32 + csrc/infinicore/src/ops/topk/topk.cc | 40 + .../src/ops/topksoftmax/topksoftmax.cc | 30 + .../ops/topksoftmax/topksoftmax_infiniops.cc | 52 + .../triplet_margin_loss.cc | 37 + .../triplet_margin_with_distance_loss.cc | 38 + csrc/infinicore/src/ops/unfold/unfold.cc | 72 ++ .../upsample_bilinear/upsample_bilinear.cc | 41 + .../ops/upsample_nearest/upsample_nearest.cc | 61 ++ csrc/infinicore/src/ops/vander/vander.cc | 35 + csrc/infinicore/src/ops/var/var.cc | 68 ++ csrc/infinicore/src/ops/var_mean/var_mean.cc | 69 ++ csrc/infinicore/src/pybind11/analyzer.hpp | 168 ++++ csrc/infinicore/src/pybind11/context.hpp | 39 + csrc/infinicore/src/pybind11/device.hpp | 32 + csrc/infinicore/src/pybind11/device_event.hpp | 43 + csrc/infinicore/src/pybind11/dtype.hpp | 27 + csrc/infinicore/src/pybind11/from_list.cc | 349 +++++++ csrc/infinicore/src/pybind11/graph.hpp | 17 + csrc/infinicore/src/pybind11/infinicore.cc | 35 + csrc/infinicore/src/pybind11/io.hpp | 28 + csrc/infinicore/src/pybind11/ops.hpp | 35 + csrc/infinicore/src/pybind11/ops/acos.hpp | 28 + .../src/pybind11/ops/adaptive_avg_pool1d.hpp | 28 + .../src/pybind11/ops/adaptive_avg_pool3d.hpp | 22 + .../src/pybind11/ops/adaptive_max_pool1d.hpp | 39 + csrc/infinicore/src/pybind11/ops/add.hpp | 26 + .../src/pybind11/ops/add_rms_norm.hpp | 51 + csrc/infinicore/src/pybind11/ops/addbmm.hpp | 52 + csrc/infinicore/src/pybind11/ops/addcmul.hpp | 51 + csrc/infinicore/src/pybind11/ops/addr.hpp | 29 + .../src/pybind11/ops/affine_grid.hpp | 30 + csrc/infinicore/src/pybind11/ops/all.hpp | 60 ++ csrc/infinicore/src/pybind11/ops/argwhere.hpp | 13 + csrc/infinicore/src/pybind11/ops/asin.hpp | 24 + csrc/infinicore/src/pybind11/ops/asinh.hpp | 24 + csrc/infinicore/src/pybind11/ops/asum.hpp | 24 + csrc/infinicore/src/pybind11/ops/atanh.hpp | 24 + .../infinicore/src/pybind11/ops/attention.hpp | 56 ++ .../src/pybind11/ops/avg_pool1d.hpp | 37 + csrc/infinicore/src/pybind11/ops/axpy.hpp | 20 + csrc/infinicore/src/pybind11/ops/baddbmm.hpp | 56 ++ csrc/infinicore/src/pybind11/ops/bilinear.hpp | 61 ++ .../ops/binary_cross_entropy_with_logits.hpp | 77 ++ .../src/pybind11/ops/bitwise_right_shift.hpp | 26 + .../infinicore/src/pybind11/ops/blas_amax.hpp | 24 + .../infinicore/src/pybind11/ops/blas_amin.hpp | 24 + .../infinicore/src/pybind11/ops/blas_copy.hpp | 19 + csrc/infinicore/src/pybind11/ops/blas_dot.hpp | 26 + .../src/pybind11/ops/block_diag.hpp | 20 + .../src/pybind11/ops/broadcast_to.hpp | 26 + csrc/infinicore/src/pybind11/ops/cat.hpp | 26 + .../src/pybind11/ops/causal_conv1d.hpp | 38 + .../src/pybind11/ops/causal_softmax.hpp | 24 + csrc/infinicore/src/pybind11/ops/cdist.hpp | 46 + .../pybind11/ops/chunk_gated_delta_rule.hpp | 42 + csrc/infinicore/src/pybind11/ops/conv2d.hpp | 61 ++ .../src/pybind11/ops/cross_entropy.hpp | 26 + csrc/infinicore/src/pybind11/ops/diff.hpp | 28 + csrc/infinicore/src/pybind11/ops/digamma.hpp | 24 + csrc/infinicore/src/pybind11/ops/dist.hpp | 28 + .../infinicore/src/pybind11/ops/embedding.hpp | 26 + csrc/infinicore/src/pybind11/ops/equal.hpp | 26 + .../src/pybind11/ops/flash_attention.hpp | 22 + csrc/infinicore/src/pybind11/ops/flipud.hpp | 31 + .../src/pybind11/ops/float_power.hpp | 59 ++ csrc/infinicore/src/pybind11/ops/floor.hpp | 25 + .../src/pybind11/ops/floor_divide.hpp | 26 + csrc/infinicore/src/pybind11/ops/fmin.hpp | 26 + csrc/infinicore/src/pybind11/ops/fmod.hpp | 26 + .../ops/fused_gated_delta_net_gating.hpp | 44 + .../src/pybind11/ops/gaussian_nll_loss.hpp | 34 + .../infinicore/src/pybind11/ops/hardswish.hpp | 24 + csrc/infinicore/src/pybind11/ops/hardtanh.hpp | 28 + .../src/pybind11/ops/hinge_embedding_loss.hpp | 22 + .../src/pybind11/ops/huber_loss.hpp | 46 + csrc/infinicore/src/pybind11/ops/hypot.hpp | 27 + .../infinicore/src/pybind11/ops/index_add.hpp | 31 + .../src/pybind11/ops/index_copy.hpp | 30 + csrc/infinicore/src/pybind11/ops/inner.hpp | 26 + .../src/pybind11/ops/interpolate.hpp | 33 + csrc/infinicore/src/pybind11/ops/kron.hpp | 20 + csrc/infinicore/src/pybind11/ops/kthvalue.hpp | 39 + .../src/pybind11/ops/kv_caching.hpp | 32 + .../src/pybind11/ops/layer_norm.hpp | 48 + csrc/infinicore/src/pybind11/ops/ldexp.hpp | 40 + csrc/infinicore/src/pybind11/ops/lerp.hpp | 64 ++ csrc/infinicore/src/pybind11/ops/linear.hpp | 56 ++ .../src/pybind11/ops/linear_w8a8i8.hpp | 54 + .../src/pybind11/ops/log_softmax.hpp | 32 + .../infinicore/src/pybind11/ops/logaddexp.hpp | 25 + .../src/pybind11/ops/logaddexp2.hpp | 25 + .../src/pybind11/ops/logcumsumexp.hpp | 31 + csrc/infinicore/src/pybind11/ops/logdet.hpp | 24 + .../src/pybind11/ops/logical_and.hpp | 27 + .../src/pybind11/ops/logical_not.hpp | 26 + .../src/pybind11/ops/mamba_selective_scan.hpp | 34 + .../src/pybind11/ops/masked_select.hpp | 19 + csrc/infinicore/src/pybind11/ops/matmul.hpp | 28 + csrc/infinicore/src/pybind11/ops/mha.hpp | 78 ++ .../src/pybind11/ops/mha_kvcache.hpp | 127 +++ .../src/pybind11/ops/mha_varlen.hpp | 110 +++ .../src/pybind11/ops/moe_topk_softmax.hpp | 57 ++ csrc/infinicore/src/pybind11/ops/mrope.hpp | 46 + csrc/infinicore/src/pybind11/ops/mul.hpp | 26 + .../src/pybind11/ops/mul_scalar.hpp | 26 + .../src/pybind11/ops/multi_margin_loss.hpp | 57 ++ csrc/infinicore/src/pybind11/ops/nrm2.hpp | 24 + csrc/infinicore/src/pybind11/ops/pad.hpp | 30 + .../src/pybind11/ops/paged_attention.hpp | 53 + .../pybind11/ops/paged_attention_prefill.hpp | 69 ++ .../src/pybind11/ops/paged_caching.hpp | 22 + .../src/pybind11/ops/per_channel_quant_i8.hpp | 20 + csrc/infinicore/src/pybind11/ops/prelu.hpp | 26 + .../src/pybind11/ops/random_sample.hpp | 32 + .../infinicore/src/pybind11/ops/rearrange.hpp | 24 + .../src/pybind11/ops/reciprocal.hpp | 24 + .../ops/recurrent_gated_delta_rule.hpp | 37 + csrc/infinicore/src/pybind11/ops/relu6.hpp | 24 + csrc/infinicore/src/pybind11/ops/rms_norm.hpp | 44 + csrc/infinicore/src/pybind11/ops/rope.hpp | 37 + csrc/infinicore/src/pybind11/ops/rot.hpp | 21 + csrc/infinicore/src/pybind11/ops/rotg.hpp | 21 + csrc/infinicore/src/pybind11/ops/rotm.hpp | 20 + csrc/infinicore/src/pybind11/ops/rotmg.hpp | 22 + .../infinicore/src/pybind11/ops/rwkv5_wkv.hpp | 34 + csrc/infinicore/src/pybind11/ops/scal.hpp | 19 + .../src/pybind11/ops/scaled_mm_i8.hpp | 30 + csrc/infinicore/src/pybind11/ops/scatter.hpp | 56 ++ csrc/infinicore/src/pybind11/ops/selu.hpp | 26 + csrc/infinicore/src/pybind11/ops/sigmoid.hpp | 24 + csrc/infinicore/src/pybind11/ops/silu.hpp | 24 + .../src/pybind11/ops/silu_and_mul.hpp | 29 + csrc/infinicore/src/pybind11/ops/sinh.hpp | 26 + .../src/pybind11/ops/smooth_l1_loss.hpp | 36 + csrc/infinicore/src/pybind11/ops/softplus.hpp | 30 + csrc/infinicore/src/pybind11/ops/softsign.hpp | 24 + csrc/infinicore/src/pybind11/ops/sum.hpp | 60 ++ csrc/infinicore/src/pybind11/ops/swap.hpp | 19 + csrc/infinicore/src/pybind11/ops/swiglu.hpp | 26 + csrc/infinicore/src/pybind11/ops/take.hpp | 25 + csrc/infinicore/src/pybind11/ops/tan.hpp | 24 + .../src/pybind11/ops/tanhshrink.hpp | 24 + csrc/infinicore/src/pybind11/ops/topk.hpp | 54 + .../src/pybind11/ops/topksoftmax.hpp | 32 + .../src/pybind11/ops/triplet_margin_loss.hpp | 50 + .../ops/triplet_margin_with_distance_loss.hpp | 41 + csrc/infinicore/src/pybind11/ops/unfold.hpp | 61 ++ .../src/pybind11/ops/upsample_bilinear.hpp | 35 + .../src/pybind11/ops/upsample_nearest.hpp | 32 + csrc/infinicore/src/pybind11/ops/vander.hpp | 42 + csrc/infinicore/src/pybind11/ops/var.hpp | 62 ++ csrc/infinicore/src/pybind11/ops/var_mean.hpp | 63 ++ csrc/infinicore/src/pybind11/tensor.hpp | 131 +++ csrc/infinicore/src/tensor/copy.cc | 125 +++ csrc/infinicore/src/tensor/debug.cc | 226 +++++ csrc/infinicore/src/tensor/tensor.cc | 296 ++++++ csrc/infinicore/src/tensor/view.cc | 159 +++ csrc/infinicore/src/utils.hpp | 66 ++ csrc/infinicore/utils/custom_types.cc | 85 ++ csrc/infinicore/utils/custom_types.h | 67 ++ csrc/layers/attention/attention.cpp | 4 +- .../layers/attention/backends/static_attn.cpp | 76 +- csrc/layers/linear/base_linear.hpp | 4 +- csrc/layers/linear/fused_linear.hpp | 14 +- csrc/layers/linear/linear.cpp | 2 +- csrc/layers/linear/linear.hpp | 10 +- .../moe/dispatcher/standard_dispatcher.cpp | 2 +- .../moe/dispatcher/standard_dispatcher.hpp | 2 +- .../allgather_reduce_scatter_dispatcher.cpp | 4 +- csrc/layers/moe/ep/base_ep_dispatcher.cpp | 8 +- csrc/layers/moe/ep/base_ep_dispatcher.hpp | 2 +- .../moe/ep/local_allreduce_dispatcher.cpp | 2 +- csrc/layers/moe/router/topk_router.cpp | 6 +- .../moe/runner/cuda_fused_moe_runner.cpp | 10 +- csrc/layers/quantization/awq.cpp | 79 +- csrc/layers/quantization/awq_marlin.cpp | 141 --- csrc/layers/quantization/awq_marlin.hpp | 55 -- .../layers/quantization/base_quantization.hpp | 21 +- .../quantization/compressed_tensors.cpp | 4 +- csrc/layers/quantization/gptq.cpp | 77 +- csrc/layers/quantization/gptq_marlin.cpp | 141 --- csrc/layers/quantization/gptq_marlin.hpp | 58 -- csrc/layers/quantization/gptq_qy.cpp | 259 ----- csrc/layers/quantization/gptq_qy.hpp | 139 --- csrc/layers/quantization/marlin_support.hpp | 7 - csrc/layers/quantization/marlin_utils.cpp | 476 --------- csrc/layers/quantization/marlin_utils.hpp | 26 - csrc/layers/quantization/quantization.hpp | 3 - .../quantization/quantization_scheme.hpp | 3 - .../rotary_embedding/rotary_embedding.cpp | 2 +- csrc/models/deepseek_v2/deepseek_v2_moe.cpp | 29 +- csrc/models/deepseek_v2/deepseek_v2_moe.hpp | 2 +- csrc/models/mamba/mamba_for_causal_lm.cpp | 2 +- csrc/models/minicpmv/minicpmv_model.cpp | 4 +- csrc/models/minicpmv/resampler.cpp | 10 +- csrc/models/minicpmv/siglip_vision.cpp | 4 +- .../qwen3_5/qwen3_5_fused_qkv_linear.hpp | 2 +- csrc/models/qwen3_5/qwen3_5_model.cpp | 6 +- csrc/models/qwen3_5/qwen3_5_vision.cpp | 8 +- csrc/models/qwen3_moe/qwen3_moe_experts.cpp | 4 +- .../qwen3_moe/qwen3_moe_topk_router.cpp | 4 +- csrc/models/videonsa/videonsa_attention.cpp | 26 +- .../videonsa_for_conditional_generation.cpp | 6 +- csrc/models/videonsa/videonsa_vision.cpp | 4 +- csrc/pybind11/engine/engine.hpp | 2 +- csrc/utils.hpp | 26 +- include/infinicore_infer.h | 14 - include/infinicore_infer/cache.h | 21 - include/infinicore_infer/models/deepseek.h | 163 --- include/infinicore_infer/models/jiuge.h | 102 -- include/infinicore_infer/models/jiuge_awq.h | 79 -- include/infinicore_infer/models/jiuge_gptq.h | 79 -- include/infinicore_infer/models/qwen3vl.h | 203 ---- include/infinicore_infer/weights_loader.h | 14 - pyproject.toml | 13 +- python/infinicore/__init__.py | 99 ++ python/infinicore/context.py | 46 + python/infinicore/device.py | 67 ++ python/infinicore/device_event.py | 38 + python/infinicore/dtype.py | 51 + python/infinicore/graph.py | 6 + python/infinicore/lib/__init__.py | 3 + python/infinicore/nn/__init__.py | 5 + python/infinicore/nn/functional/__init__.py | 19 + python/infinicore/nn/modules/__init__.py | 3 + python/infinicore/nn/modules/module.py | 8 + python/infinicore/nn/parameter.py | 8 + python/infinicore/ops/__init__.py | 81 ++ python/infinicore/tensor.py | 253 +++++ python/infinicore/utils.py | 98 ++ python/infinilm/__init__.py | 10 +- python/infinilm/base_config.py | 16 +- python/infinilm/exception_utils.py | 5 +- python/infinilm/generation/__init__.py | 1 + python/infinilm/generation/utils.py | 40 +- .../kv_connector/mooncake/__init__.py | 1 + python/infinilm/lib/__init__.py | 18 +- python/infinilm/llm/model_runner/__init__.py | 1 + python/infinilm/models/__init__.py | 4 +- python/infinilm/multimodal/__init__.py | 1 + python/infinilm/server/__init__.py | 1 + scripts/deepseek.py | 775 --------------- scripts/infer_task.py | 59 -- scripts/jiuge.py | 881 ----------------- scripts/jiuge_awq.py | 369 ------- scripts/jiuge_gptq.py | 369 ------- scripts/jiuge_ppl.py | 115 --- scripts/kvcache_pool.py | 90 -- scripts/launch_server.py | 329 ------ scripts/libinfinicore_infer/__init__.py | 55 -- scripts/libinfinicore_infer/base.py | 72 -- scripts/libinfinicore_infer/deepseek_v3.py | 209 ---- scripts/libinfinicore_infer/jiuge.py | 151 --- scripts/libinfinicore_infer/jiuge_awq.py | 167 ---- scripts/libinfinicore_infer/jiuge_gptq.py | 167 ---- scripts/libinfinicore_infer/qwen3vl.py | 329 ------ scripts/qwen3vl.py | 935 ------------------ scripts/test_ceval.py | 205 ---- setup.py | 94 +- src/allocator.hpp | 54 - src/allocator/memory_allocator.cpp | 135 --- src/cache.hpp | 9 - src/cache_manager/kvcache.cpp | 80 -- src/cache_manager/opcache_manager.hpp | 197 ---- src/dataloader/weights_loader.cpp | 86 -- src/dataloader/weights_loader.hpp | 61 -- src/models/deepseek_v3/deepseek_v3.cpp | 589 ----------- src/models/deepseek_v3/deepseek_v3_cache.cpp | 41 - src/models/deepseek_v3/deepseek_v3_impl.hpp | 116 --- src/models/deepseek_v3/deepseek_v3_weight.cpp | 451 --------- src/models/inference_context.cpp | 371 ------- src/models/inference_context.hpp | 182 ---- src/models/jiuge/jiuge.cpp | 474 --------- src/models/jiuge/jiuge_impl.hpp | 69 -- src/models/jiuge/jiuge_weight.hpp | 204 ---- src/models/jiuge_awq/jiuge_awq.cpp | 405 -------- src/models/jiuge_awq/jiuge_awq.hpp | 82 -- src/models/jiuge_awq/jiuge_awq_weight.cpp | 132 --- src/models/jiuge_gptq/jiuge_gptq.cpp | 405 -------- src/models/jiuge_gptq/jiuge_gptq.hpp | 82 -- src/models/jiuge_gptq/jiuge_gptq_weight.cpp | 137 --- src/models/qwen3vl/qwen3vl.cpp | 715 -------------- src/models/qwen3vl/qwen3vl_cache.cpp | 43 - src/models/qwen3vl/qwen3vl_impl.hpp | 141 --- src/models/qwen3vl/qwen3vl_weight.cpp | 636 ------------ src/tensor.hpp | 190 ---- src/tensor/strorage.cpp | 53 - src/tensor/tensor.cpp | 426 -------- src/tensor/transform.cpp | 137 --- src/utils.hpp | 127 --- test/bench/backends/infinilm.py | 11 +- .../llama/test_intermediate_validation.py | 130 ++- test/models/llama/utils.py | 63 +- .../test_infinicore_python_contracts.py | 357 +++++++ .../test_infinicore_runtime_contracts.py | 785 +++++++++++++++ xmake.lua | 107 +- 694 files changed, 30258 insertions(+), 14901 deletions(-) create mode 100644 csrc/infinicore/include/infinicore.hpp create mode 100644 csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp create mode 100644 csrc/infinicore/include/infinicore/adaptor/flash_attention_adaptor.hpp create mode 100644 csrc/infinicore/include/infinicore/analyzer.hpp create mode 100644 csrc/infinicore/include/infinicore/analyzer/intent_generator.hpp create mode 100644 csrc/infinicore/include/infinicore/analyzer/mutual_awareness_analyzer.hpp create mode 100644 csrc/infinicore/include/infinicore/analyzer/op_trace.hpp create mode 100644 csrc/infinicore/include/infinicore/analyzer/op_type.hpp create mode 100644 csrc/infinicore/include/infinicore/analyzer/op_type_registry.hpp create mode 100644 csrc/infinicore/include/infinicore/analyzer/optimization_intent.hpp create mode 100644 csrc/infinicore/include/infinicore/analyzer/phase_detector.hpp create mode 100644 csrc/infinicore/include/infinicore/analyzer/resource_sensor.hpp create mode 100644 csrc/infinicore/include/infinicore/common/LRUCache.hpp create mode 100644 csrc/infinicore/include/infinicore/common/hash.hpp create mode 100644 csrc/infinicore/include/infinicore/context/context.hpp create mode 100644 csrc/infinicore/include/infinicore/device.hpp create mode 100644 csrc/infinicore/include/infinicore/device_event.hpp create mode 100644 csrc/infinicore/include/infinicore/dtype.hpp create mode 100644 csrc/infinicore/include/infinicore/graph/graph.hpp create mode 100644 csrc/infinicore/include/infinicore/io.hpp create mode 100644 csrc/infinicore/include/infinicore/memory.hpp create mode 100644 csrc/infinicore/include/infinicore/nn.hpp create mode 100644 csrc/infinicore/include/infinicore/nn/embedding.hpp create mode 100644 csrc/infinicore/include/infinicore/nn/layer_norm.hpp create mode 100644 csrc/infinicore/include/infinicore/nn/module.hpp create mode 100644 csrc/infinicore/include/infinicore/nn/parameter.hpp create mode 100644 csrc/infinicore/include/infinicore/nn/rmsnorm.hpp create mode 100644 csrc/infinicore/include/infinicore/nn/rope.hpp create mode 100644 csrc/infinicore/include/infinicore/nn/rope_scaling_configs.hpp create mode 100644 csrc/infinicore/include/infinicore/ops.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/acos.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/adaptive_avg_pool1d.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/adaptive_avg_pool3d.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/adaptive_max_pool1d.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/add.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/add_rms_norm.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/addbmm.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/addcmul.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/addr.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/affine_grid.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/argwhere.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/asin.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/asinh.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/asum.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/atanh.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/attention.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/avg_pool1d.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/axpy.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/baddbmm.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/bilinear.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/binary_cross_entropy_with_logits.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/bitwise_right_shift.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/blas_amax.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/blas_amin.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/blas_copy.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/blas_dot.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/block_diag.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/broadcast_to.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/cat.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/causal_conv1d.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/causal_softmax.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/cdist.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/chunk_gated_delta_rule.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/common/cache.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/common/dispatcher.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/common/op.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/conv2d.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/cross_entropy.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/deepseek_moe.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/diff.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/digamma.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/dist.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/distributed/allgather.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/distributed/allreduce.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/distributed/reduce_scatter.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/embedding.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/flash_attention.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/flipud.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/float_power.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/floor.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/floor_divide.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/fmin.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/fmod.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/fused_gated_delta_net_gating.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/gaussian_nll_loss.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/gelu.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/gelutanh.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/gemm.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/hardswish.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/hardtanh.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/hinge_embedding_loss.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/huber_loss.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/hypot.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/index_add.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/index_copy.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/inner.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/interpolate.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/kron.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/kthvalue.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/kv_caching.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/layer_norm.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/ldexp.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/lerp.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/linear.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/linear_w8a8i8.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/log_softmax.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/logaddexp.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/logaddexp2.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/logcumsumexp.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/logdet.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/mamba_selective_scan.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/masked_select.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/matmul.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/mha.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/mha_kvcache.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/mha_varlen.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/moe_align.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/moe_fused_gate.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/moe_sum.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/moe_topk_softmax.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/mrope.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/mul.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/mul_scalar.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/multi_margin_loss.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/nrm2.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/nsa_compress_paged_cache.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/nsa_paged_attention.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/ones.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/pad.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/paged_attention.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/paged_attention_prefill.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/paged_caching.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/per_channel_quant_i8.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/per_tensor_dequant_i8.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/per_tensor_quant_i8.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/prelu.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/prepare_moe_input.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/quickgelu.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/random_sample.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/rearrange.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/reciprocal.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/recurrent_gated_delta_rule.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/relu.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/relu6.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/rms_norm.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/rope.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/rot.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/rotg.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/rotm.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/rotmg.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/rwkv5_wkv.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/scal.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/scaled_mm_i8.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/scatter.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/selu.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/sigmoid.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/silu.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/silu_and_mul.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/sinh.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/smooth_l1_loss.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/softmax.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/softplus.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/softsign.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/sum.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/swap.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/swiglu.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/take.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/tan.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/tanhshrink.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/topk.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/topksoftmax.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/triplet_margin_loss.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/triplet_margin_with_distance_loss.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/unfold.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/upsample_bilinear.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/upsample_nearest.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/vander.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/var.hpp create mode 100644 csrc/infinicore/include/infinicore/ops/var_mean.hpp create mode 100644 csrc/infinicore/include/infinicore/tensor.hpp create mode 100644 csrc/infinicore/src/adaptor/aten_adaptor.cc create mode 100644 csrc/infinicore/src/adaptor/flash_attn/hygon/flash_attn_hygon.cc create mode 100644 csrc/infinicore/src/adaptor/flash_attn/hygon/flash_attn_hygon.hpp create mode 100644 csrc/infinicore/src/analyzer/mutual_awareness_analyzer.cc create mode 100644 csrc/infinicore/src/analyzer/op_trace.cc create mode 100644 csrc/infinicore/src/context/allocators/device_pinned_allocator.cc create mode 100644 csrc/infinicore/src/context/allocators/device_pinned_allocator.hpp create mode 100644 csrc/infinicore/src/context/allocators/host_allocator.cc create mode 100644 csrc/infinicore/src/context/allocators/host_allocator.hpp create mode 100644 csrc/infinicore/src/context/allocators/memory_allocator.hpp create mode 100644 csrc/infinicore/src/context/allocators/pinnable_block_allocator.cc create mode 100644 csrc/infinicore/src/context/allocators/pinnable_block_allocator.hpp create mode 100644 csrc/infinicore/src/context/allocators/stream_ordered_allocator.cc create mode 100644 csrc/infinicore/src/context/allocators/stream_ordered_allocator.hpp create mode 100644 csrc/infinicore/src/context/context_impl.cc create mode 100644 csrc/infinicore/src/context/context_impl.hpp create mode 100644 csrc/infinicore/src/context/internal.hpp create mode 100644 csrc/infinicore/src/context/runtime/runtime.cc create mode 100644 csrc/infinicore/src/context/runtime/runtime.hpp create mode 100644 csrc/infinicore/src/device.cc create mode 100644 csrc/infinicore/src/device_event.cc create mode 100644 csrc/infinicore/src/dtype.cc create mode 100644 csrc/infinicore/src/graph/graph.cc create mode 100644 csrc/infinicore/src/graph/graph_manager.hpp create mode 100644 csrc/infinicore/src/io.cc create mode 100644 csrc/infinicore/src/memory.cc create mode 100644 csrc/infinicore/src/nn/embedding.cc create mode 100644 csrc/infinicore/src/nn/layer_norm.cc create mode 100644 csrc/infinicore/src/nn/module.cc create mode 100644 csrc/infinicore/src/nn/parameter.cc create mode 100644 csrc/infinicore/src/nn/rmsnorm.cc create mode 100644 csrc/infinicore/src/nn/rope.cc create mode 100644 csrc/infinicore/src/nn/rope_scaling_configs.cc create mode 100644 csrc/infinicore/src/ops/acos/acos.cc create mode 100644 csrc/infinicore/src/ops/adaptive_avg_pool1d/adaptive_avg_pool1d.cc create mode 100644 csrc/infinicore/src/ops/adaptive_avg_pool3d/adaptive_avg_pool3d.cc create mode 100644 csrc/infinicore/src/ops/adaptive_max_pool1d/adaptive_max_pool1d.cc create mode 100644 csrc/infinicore/src/ops/add/add.cc create mode 100644 csrc/infinicore/src/ops/add/add_infiniops.cc create mode 100644 csrc/infinicore/src/ops/add_rms_norm/add_rms_norm.cc create mode 100644 csrc/infinicore/src/ops/add_rms_norm/add_rms_norm_infiniops.cc create mode 100644 csrc/infinicore/src/ops/addbmm/addbmm.cc create mode 100644 csrc/infinicore/src/ops/addcmul/addcmul.cc create mode 100644 csrc/infinicore/src/ops/addr/addr.cc create mode 100644 csrc/infinicore/src/ops/affine_grid/affine_grid.cc create mode 100644 csrc/infinicore/src/ops/argwhere/argwhere.cc create mode 100644 csrc/infinicore/src/ops/asin/asin.cc create mode 100644 csrc/infinicore/src/ops/asinh/asinh.cc create mode 100644 csrc/infinicore/src/ops/asum/asum.cc create mode 100644 csrc/infinicore/src/ops/atanh/atanh.cc create mode 100644 csrc/infinicore/src/ops/attention/attention.cc create mode 100644 csrc/infinicore/src/ops/avg_pool1d/avg_pool1d.cc create mode 100644 csrc/infinicore/src/ops/axpy/axpy.cc create mode 100644 csrc/infinicore/src/ops/baddbmm/baddbmm.cc create mode 100644 csrc/infinicore/src/ops/bilinear/bilinear.cc create mode 100644 csrc/infinicore/src/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits.cc create mode 100644 csrc/infinicore/src/ops/bitwise_right_shift/bitwise_right_shift.cc create mode 100644 csrc/infinicore/src/ops/blas_amax/blas_amax.cc create mode 100644 csrc/infinicore/src/ops/blas_amin/blas_amin.cc create mode 100644 csrc/infinicore/src/ops/blas_copy/blas_copy.cc create mode 100644 csrc/infinicore/src/ops/blas_dot/blas_dot.cc create mode 100644 csrc/infinicore/src/ops/block_diag/block_diag.cc create mode 100644 csrc/infinicore/src/ops/broadcast_to/broadcast_to.cc create mode 100644 csrc/infinicore/src/ops/cat/cat.cc create mode 100644 csrc/infinicore/src/ops/causal_conv1d/causal_conv1d.cc create mode 100644 csrc/infinicore/src/ops/causal_softmax/causal_softmax.cc create mode 100644 csrc/infinicore/src/ops/causal_softmax/causal_softmax_infiniops.cc create mode 100644 csrc/infinicore/src/ops/cdist/cdist.cc create mode 100644 csrc/infinicore/src/ops/chunk_gated_delta_rule/chunk_gated_delta_rule.cc create mode 100644 csrc/infinicore/src/ops/conv2d/conv2d.cc create mode 100644 csrc/infinicore/src/ops/conv2d/conv2d_infiniops.cc create mode 100644 csrc/infinicore/src/ops/cross_entropy/cross_entropy.cc create mode 100644 csrc/infinicore/src/ops/deepseek_moe/deepseek_moe.cc create mode 100644 csrc/infinicore/src/ops/dequant/per_tensor_dequant_int8/per_tensor_quant_int8.cc create mode 100644 csrc/infinicore/src/ops/diff/diff.cc create mode 100644 csrc/infinicore/src/ops/digamma/digamma.cc create mode 100644 csrc/infinicore/src/ops/dist/dist.cc create mode 100644 csrc/infinicore/src/ops/distributed/allgather.cc create mode 100644 csrc/infinicore/src/ops/distributed/allreduce.cc create mode 100644 csrc/infinicore/src/ops/distributed/reduce_scatter.cc create mode 100644 csrc/infinicore/src/ops/distributed/utils.hpp create mode 100644 csrc/infinicore/src/ops/embedding/embedding.cc create mode 100644 csrc/infinicore/src/ops/embedding/embedding_infiniops.cc create mode 100644 csrc/infinicore/src/ops/flash_attention/flash_attention.cc create mode 100644 csrc/infinicore/src/ops/flipud/flipud.cc create mode 100644 csrc/infinicore/src/ops/float_power/float_power.cc create mode 100644 csrc/infinicore/src/ops/floor/floor.cc create mode 100644 csrc/infinicore/src/ops/floor_divide/floor_divide.cc create mode 100644 csrc/infinicore/src/ops/fmin/fmin.cc create mode 100644 csrc/infinicore/src/ops/fmod/fmod.cc create mode 100644 csrc/infinicore/src/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating.cc create mode 100644 csrc/infinicore/src/ops/gaussian_nll_loss/gaussian_nll_loss.cc create mode 100644 csrc/infinicore/src/ops/gelu/gelu.cc create mode 100644 csrc/infinicore/src/ops/gelu/gelu_infiniops.cc create mode 100644 csrc/infinicore/src/ops/gelutanh/gelutanh.cc create mode 100644 csrc/infinicore/src/ops/gelutanh/gelutanh_infiniops.cc create mode 100644 csrc/infinicore/src/ops/gemm/gemm.cc create mode 100644 csrc/infinicore/src/ops/gemm/gemm_infiniops.cc create mode 100644 csrc/infinicore/src/ops/hardswish/hardswish.cc create mode 100644 csrc/infinicore/src/ops/hardtanh/hardtanh.cc create mode 100644 csrc/infinicore/src/ops/hinge_embedding_loss/hinge_embedding_loss.cc create mode 100644 csrc/infinicore/src/ops/huber_loss/huber_loss.cc create mode 100644 csrc/infinicore/src/ops/hypot/hypot.cc create mode 100644 csrc/infinicore/src/ops/index_add/index_add.cc create mode 100644 csrc/infinicore/src/ops/index_copy/index_copy.cc create mode 100644 csrc/infinicore/src/ops/infiniops_impl.hpp create mode 100644 csrc/infinicore/src/ops/inner/inner.cc create mode 100644 csrc/infinicore/src/ops/interpolate/interpolate.cc create mode 100644 csrc/infinicore/src/ops/kron/kron.cc create mode 100644 csrc/infinicore/src/ops/kthvalue/kthvalue.cc create mode 100644 csrc/infinicore/src/ops/kv_caching/kv_caching.cc create mode 100644 csrc/infinicore/src/ops/kv_caching/kv_caching_infiniops.cc create mode 100644 csrc/infinicore/src/ops/layer_norm/layer_norm.cc create mode 100644 csrc/infinicore/src/ops/ldexp/ldexp.cc create mode 100644 csrc/infinicore/src/ops/lerp/lerp.cc create mode 100644 csrc/infinicore/src/ops/linear/linear.cc create mode 100644 csrc/infinicore/src/ops/linear_w8a8i8/linear_w8a8i8.cc create mode 100644 csrc/infinicore/src/ops/log_softmax/log_softmax.cc create mode 100644 csrc/infinicore/src/ops/logaddexp/logaddexp.cc create mode 100644 csrc/infinicore/src/ops/logaddexp2/logaddxep2.cc create mode 100644 csrc/infinicore/src/ops/logcumsumexp/logcumsumexp.cc create mode 100644 csrc/infinicore/src/ops/logdet/logdet.cc create mode 100644 csrc/infinicore/src/ops/mamba_selective_scan/mamba_selective_scan.cc create mode 100644 csrc/infinicore/src/ops/masked_select/masked_select.cc create mode 100644 csrc/infinicore/src/ops/matmul/matmul.cc create mode 100644 csrc/infinicore/src/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc create mode 100644 csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc create mode 100644 csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_flashattn.cc create mode 100644 csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_flashattn_moore.cc create mode 100644 csrc/infinicore/src/ops/moe_align/moe_align.cc create mode 100644 csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc create mode 100644 csrc/infinicore/src/ops/moe_fused_gate/moe_fused_gate.cc create mode 100644 csrc/infinicore/src/ops/moe_sum/moe_sum.cc create mode 100644 csrc/infinicore/src/ops/moe_topk_sigmoid/moe_topk_sigmoid.cc create mode 100644 csrc/infinicore/src/ops/moe_topk_softmax/moe_topk_softmax.cc create mode 100644 csrc/infinicore/src/ops/mrope/mrope.cc create mode 100644 csrc/infinicore/src/ops/mul/mul.cc create mode 100644 csrc/infinicore/src/ops/mul_scalar/mul_scalar.cc create mode 100644 csrc/infinicore/src/ops/multi_head_attention/hygon/mha_flashattn_hygon.cc create mode 100644 csrc/infinicore/src/ops/multi_head_attention/mha.cc create mode 100644 csrc/infinicore/src/ops/multi_head_attention/mha_flashattn.cc create mode 100644 csrc/infinicore/src/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc create mode 100644 csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen.cc create mode 100644 csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc create mode 100644 csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_flashattn_moore.cc create mode 100644 csrc/infinicore/src/ops/multi_margin_loss/multi_margin_loss.cc create mode 100644 csrc/infinicore/src/ops/nrm2/nrm2.cc create mode 100644 csrc/infinicore/src/ops/nsa_compress_paged_cache/nsa_compress_paged_cache.cc create mode 100644 csrc/infinicore/src/ops/nsa_paged_attention/nsa_paged_attention.cc create mode 100644 csrc/infinicore/src/ops/ones/ones.cc create mode 100644 csrc/infinicore/src/ops/pad/pad.cc create mode 100644 csrc/infinicore/src/ops/paged_attention/paged_attention.cc create mode 100644 csrc/infinicore/src/ops/paged_attention/paged_attention_infiniops.cc create mode 100644 csrc/infinicore/src/ops/paged_attention_prefill/paged_attention_prefill.cc create mode 100644 csrc/infinicore/src/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc create mode 100644 csrc/infinicore/src/ops/paged_caching/paged_caching.cc create mode 100644 csrc/infinicore/src/ops/paged_caching/paged_caching_infiniops.cc create mode 100644 csrc/infinicore/src/ops/per_channel_quant_i8/per_channel_quant_i8.cc create mode 100644 csrc/infinicore/src/ops/prelu/prelu.cc create mode 100644 csrc/infinicore/src/ops/prepare_moe_input/prepare_moe_input.cc create mode 100644 csrc/infinicore/src/ops/quant/per_tensor_quant_int8/per_tensor_quant_int8.cc create mode 100644 csrc/infinicore/src/ops/quickgelu/quickgelu.cc create mode 100644 csrc/infinicore/src/ops/random_sample/random_sample.cc create mode 100644 csrc/infinicore/src/ops/random_sample/random_sample_infiniops.cc create mode 100644 csrc/infinicore/src/ops/rearrange/rearrange.cc create mode 100644 csrc/infinicore/src/ops/rearrange/rearrange_infiniops.cc create mode 100644 csrc/infinicore/src/ops/reciprocal/reciprocal.cc create mode 100644 csrc/infinicore/src/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule.cc create mode 100644 csrc/infinicore/src/ops/relu/relu.cc create mode 100644 csrc/infinicore/src/ops/relu/relu_infiniops.cc create mode 100644 csrc/infinicore/src/ops/relu6/relu6.cc create mode 100644 csrc/infinicore/src/ops/rms_norm/rms_norm.cc create mode 100644 csrc/infinicore/src/ops/rms_norm/rms_norm_infiniops.cc create mode 100644 csrc/infinicore/src/ops/rope/rope.cc create mode 100644 csrc/infinicore/src/ops/rope/rope_infiniops.cc create mode 100644 csrc/infinicore/src/ops/rot/rot.cc create mode 100644 csrc/infinicore/src/ops/rotg/rotg.cc create mode 100644 csrc/infinicore/src/ops/rotm/rotm.cc create mode 100644 csrc/infinicore/src/ops/rotmg/rotmg.cc create mode 100644 csrc/infinicore/src/ops/rwkv5_wkv/rwkv5_wkv.cc create mode 100644 csrc/infinicore/src/ops/scal/scal.cc create mode 100644 csrc/infinicore/src/ops/scaled_mm_i8/scaled_mm_i8.cc create mode 100644 csrc/infinicore/src/ops/scatter/scatter.cc create mode 100644 csrc/infinicore/src/ops/selu/selu.cc create mode 100644 csrc/infinicore/src/ops/sigmoid/sigmoid.cc create mode 100644 csrc/infinicore/src/ops/sigmoid/sigmoid_infiniops.cc create mode 100644 csrc/infinicore/src/ops/silu/silu.cc create mode 100644 csrc/infinicore/src/ops/silu/silu_infiniops.cc create mode 100644 csrc/infinicore/src/ops/silu_and_mul/silu_and_mul.cc create mode 100644 csrc/infinicore/src/ops/silu_and_mul/silu_and_mul_infiniops.cc create mode 100644 csrc/infinicore/src/ops/sinh/sinh.cc create mode 100644 csrc/infinicore/src/ops/smooth_l1_loss/smooth_l1_loss.cc create mode 100644 csrc/infinicore/src/ops/softmax/softmax.cc create mode 100644 csrc/infinicore/src/ops/softmax/softmax_infiniops.cc create mode 100644 csrc/infinicore/src/ops/softplus/softplus.cc create mode 100644 csrc/infinicore/src/ops/softsign/softsign.cc create mode 100644 csrc/infinicore/src/ops/sum/sum.cc create mode 100644 csrc/infinicore/src/ops/swap/swap.cc create mode 100644 csrc/infinicore/src/ops/swiglu/swiglu.cc create mode 100644 csrc/infinicore/src/ops/swiglu/swiglu_infiniops.cc create mode 100644 csrc/infinicore/src/ops/take/take.cc create mode 100644 csrc/infinicore/src/ops/tan/tan.cc create mode 100644 csrc/infinicore/src/ops/tanhshrink/tanhshrink.cc create mode 100644 csrc/infinicore/src/ops/topk/topk.cc create mode 100644 csrc/infinicore/src/ops/topksoftmax/topksoftmax.cc create mode 100644 csrc/infinicore/src/ops/topksoftmax/topksoftmax_infiniops.cc create mode 100644 csrc/infinicore/src/ops/triplet_margin_loss/triplet_margin_loss.cc create mode 100644 csrc/infinicore/src/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.cc create mode 100644 csrc/infinicore/src/ops/unfold/unfold.cc create mode 100644 csrc/infinicore/src/ops/upsample_bilinear/upsample_bilinear.cc create mode 100644 csrc/infinicore/src/ops/upsample_nearest/upsample_nearest.cc create mode 100644 csrc/infinicore/src/ops/vander/vander.cc create mode 100644 csrc/infinicore/src/ops/var/var.cc create mode 100644 csrc/infinicore/src/ops/var_mean/var_mean.cc create mode 100644 csrc/infinicore/src/pybind11/analyzer.hpp create mode 100644 csrc/infinicore/src/pybind11/context.hpp create mode 100644 csrc/infinicore/src/pybind11/device.hpp create mode 100644 csrc/infinicore/src/pybind11/device_event.hpp create mode 100644 csrc/infinicore/src/pybind11/dtype.hpp create mode 100644 csrc/infinicore/src/pybind11/from_list.cc create mode 100644 csrc/infinicore/src/pybind11/graph.hpp create mode 100644 csrc/infinicore/src/pybind11/infinicore.cc create mode 100644 csrc/infinicore/src/pybind11/io.hpp create mode 100644 csrc/infinicore/src/pybind11/ops.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/acos.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/adaptive_avg_pool1d.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/adaptive_avg_pool3d.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/adaptive_max_pool1d.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/add.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/add_rms_norm.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/addbmm.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/addcmul.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/addr.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/affine_grid.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/all.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/argwhere.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/asin.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/asinh.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/asum.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/atanh.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/attention.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/avg_pool1d.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/axpy.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/baddbmm.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/bilinear.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/binary_cross_entropy_with_logits.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/bitwise_right_shift.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/blas_amax.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/blas_amin.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/blas_copy.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/blas_dot.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/block_diag.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/broadcast_to.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/cat.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/causal_conv1d.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/causal_softmax.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/cdist.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/chunk_gated_delta_rule.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/conv2d.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/cross_entropy.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/diff.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/digamma.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/dist.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/embedding.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/equal.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/flash_attention.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/flipud.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/float_power.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/floor.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/floor_divide.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/fmin.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/fmod.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/fused_gated_delta_net_gating.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/gaussian_nll_loss.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/hardswish.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/hardtanh.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/hinge_embedding_loss.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/huber_loss.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/hypot.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/index_add.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/index_copy.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/inner.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/interpolate.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/kron.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/kthvalue.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/kv_caching.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/layer_norm.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/ldexp.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/lerp.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/linear.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/linear_w8a8i8.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/log_softmax.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/logaddexp.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/logaddexp2.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/logcumsumexp.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/logdet.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/logical_and.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/logical_not.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/mamba_selective_scan.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/masked_select.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/matmul.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/mha.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/mha_kvcache.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/mha_varlen.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/moe_topk_softmax.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/mrope.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/mul.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/mul_scalar.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/multi_margin_loss.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/nrm2.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/pad.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/paged_attention.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/paged_attention_prefill.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/paged_caching.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/per_channel_quant_i8.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/prelu.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/random_sample.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/rearrange.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/reciprocal.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/recurrent_gated_delta_rule.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/relu6.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/rms_norm.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/rope.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/rot.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/rotg.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/rotm.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/rotmg.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/rwkv5_wkv.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/scal.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/scaled_mm_i8.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/scatter.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/selu.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/sigmoid.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/silu.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/silu_and_mul.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/sinh.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/smooth_l1_loss.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/softplus.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/softsign.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/sum.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/swap.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/swiglu.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/take.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/tan.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/tanhshrink.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/topk.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/topksoftmax.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/triplet_margin_loss.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/triplet_margin_with_distance_loss.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/unfold.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/upsample_bilinear.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/upsample_nearest.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/vander.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/var.hpp create mode 100644 csrc/infinicore/src/pybind11/ops/var_mean.hpp create mode 100644 csrc/infinicore/src/pybind11/tensor.hpp create mode 100644 csrc/infinicore/src/tensor/copy.cc create mode 100644 csrc/infinicore/src/tensor/debug.cc create mode 100644 csrc/infinicore/src/tensor/tensor.cc create mode 100644 csrc/infinicore/src/tensor/view.cc create mode 100644 csrc/infinicore/src/utils.hpp create mode 100644 csrc/infinicore/utils/custom_types.cc create mode 100644 csrc/infinicore/utils/custom_types.h delete mode 100644 csrc/layers/quantization/awq_marlin.cpp delete mode 100644 csrc/layers/quantization/awq_marlin.hpp delete mode 100644 csrc/layers/quantization/gptq_marlin.cpp delete mode 100644 csrc/layers/quantization/gptq_marlin.hpp delete mode 100644 csrc/layers/quantization/gptq_qy.cpp delete mode 100644 csrc/layers/quantization/gptq_qy.hpp delete mode 100644 csrc/layers/quantization/marlin_support.hpp delete mode 100644 csrc/layers/quantization/marlin_utils.cpp delete mode 100644 csrc/layers/quantization/marlin_utils.hpp delete mode 100644 include/infinicore_infer.h delete mode 100644 include/infinicore_infer/cache.h delete mode 100644 include/infinicore_infer/models/deepseek.h delete mode 100644 include/infinicore_infer/models/jiuge.h delete mode 100644 include/infinicore_infer/models/jiuge_awq.h delete mode 100644 include/infinicore_infer/models/jiuge_gptq.h delete mode 100644 include/infinicore_infer/models/qwen3vl.h delete mode 100644 include/infinicore_infer/weights_loader.h create mode 100644 python/infinicore/__init__.py create mode 100644 python/infinicore/context.py create mode 100644 python/infinicore/device.py create mode 100644 python/infinicore/device_event.py create mode 100644 python/infinicore/dtype.py create mode 100644 python/infinicore/graph.py create mode 100644 python/infinicore/lib/__init__.py create mode 100644 python/infinicore/nn/__init__.py create mode 100644 python/infinicore/nn/functional/__init__.py create mode 100644 python/infinicore/nn/modules/__init__.py create mode 100644 python/infinicore/nn/modules/module.py create mode 100644 python/infinicore/nn/parameter.py create mode 100644 python/infinicore/ops/__init__.py create mode 100644 python/infinicore/tensor.py create mode 100644 python/infinicore/utils.py create mode 100644 python/infinilm/generation/__init__.py create mode 100644 python/infinilm/kv_connector/mooncake/__init__.py create mode 100644 python/infinilm/llm/model_runner/__init__.py create mode 100644 python/infinilm/multimodal/__init__.py create mode 100644 python/infinilm/server/__init__.py delete mode 100644 scripts/deepseek.py delete mode 100644 scripts/infer_task.py delete mode 100644 scripts/jiuge.py delete mode 100644 scripts/jiuge_awq.py delete mode 100644 scripts/jiuge_gptq.py delete mode 100644 scripts/jiuge_ppl.py delete mode 100644 scripts/kvcache_pool.py delete mode 100644 scripts/launch_server.py delete mode 100644 scripts/libinfinicore_infer/__init__.py delete mode 100644 scripts/libinfinicore_infer/base.py delete mode 100644 scripts/libinfinicore_infer/deepseek_v3.py delete mode 100644 scripts/libinfinicore_infer/jiuge.py delete mode 100644 scripts/libinfinicore_infer/jiuge_awq.py delete mode 100644 scripts/libinfinicore_infer/jiuge_gptq.py delete mode 100644 scripts/libinfinicore_infer/qwen3vl.py delete mode 100644 scripts/qwen3vl.py delete mode 100644 scripts/test_ceval.py delete mode 100644 src/allocator.hpp delete mode 100644 src/allocator/memory_allocator.cpp delete mode 100644 src/cache.hpp delete mode 100644 src/cache_manager/kvcache.cpp delete mode 100644 src/cache_manager/opcache_manager.hpp delete mode 100644 src/dataloader/weights_loader.cpp delete mode 100644 src/dataloader/weights_loader.hpp delete mode 100644 src/models/deepseek_v3/deepseek_v3.cpp delete mode 100644 src/models/deepseek_v3/deepseek_v3_cache.cpp delete mode 100644 src/models/deepseek_v3/deepseek_v3_impl.hpp delete mode 100644 src/models/deepseek_v3/deepseek_v3_weight.cpp delete mode 100644 src/models/inference_context.cpp delete mode 100644 src/models/inference_context.hpp delete mode 100644 src/models/jiuge/jiuge.cpp delete mode 100644 src/models/jiuge/jiuge_impl.hpp delete mode 100644 src/models/jiuge/jiuge_weight.hpp delete mode 100644 src/models/jiuge_awq/jiuge_awq.cpp delete mode 100644 src/models/jiuge_awq/jiuge_awq.hpp delete mode 100644 src/models/jiuge_awq/jiuge_awq_weight.cpp delete mode 100644 src/models/jiuge_gptq/jiuge_gptq.cpp delete mode 100644 src/models/jiuge_gptq/jiuge_gptq.hpp delete mode 100644 src/models/jiuge_gptq/jiuge_gptq_weight.cpp delete mode 100644 src/models/qwen3vl/qwen3vl.cpp delete mode 100644 src/models/qwen3vl/qwen3vl_cache.cpp delete mode 100644 src/models/qwen3vl/qwen3vl_impl.hpp delete mode 100644 src/models/qwen3vl/qwen3vl_weight.cpp delete mode 100644 src/tensor.hpp delete mode 100644 src/tensor/strorage.cpp delete mode 100644 src/tensor/tensor.cpp delete mode 100644 src/tensor/transform.cpp delete mode 100644 src/utils.hpp create mode 100644 test/static/test_infinicore_python_contracts.py create mode 100644 test/static/test_infinicore_runtime_contracts.py diff --git a/README.md b/README.md index 4c1eaf1e1..da0c1001d 100644 --- a/README.md +++ b/README.md @@ -1,185 +1,91 @@ # InfiniLM -![star](https://atomgit.com/InfiniTensor/InfiniLM/star/badge.svg) - -本项目是基于 [`InfiniCore`](https://github.com/InfiniTensor/InfiniCore) 的推理引擎。 - -## 使用方式 -#### 一、编译并安装 `InfiniCore` -编译并安装 `InfiniCore`， 详情见 InfiniCore的 [`README`](https://github.com/InfiniTensor/InfiniCore) : - -- 注意根据提示设置好 `INFINI_ROOT` 环境变量（默认为 `$HOME/.infini`） -- 根据硬件平台，选择 xmake 构建配置 -- 编译安装InfiniCore -- 安装 C++ 库 -- 安装 Python 包 - - -#### 二、编译并安装 `InfiniLM` - - 克隆项目 - - 由于仓库中含有子模块，所以在克隆时请添加 `--recursive` 或 `--recurse-submodules`，如： - - ```shell - git clone --recursive https://github.com/InfiniTensor/InfiniLM.git - ``` - - 或者在普通克隆后进行更新： - - ```shell - git submodule update --init --recursive - ``` - - - 安装 InfiniLM Python 包 - ```bash - pip install -e . - ``` - - - 单次推理测试 - - llama示例 - ```bash - python examples/test_infer.py --device [cpu | nvidia | qy | metax | moore | iluvatar | ali | cambricon | hygon] --model= - ``` - - 例如： - ```bash - python examples/test_infer.py --device=nvidia --model=/models/TinyLlama-1.1B-Chat-v1.0 - ``` - - 分布式推理测试 - - 9g示例 - ```bash - python examples/test_infer.py [-- device nvidia] --model= --backend=cpp --tp=NDEV --batch-size=MAX_BATCH - ``` - - - 例如： 9G7B模型，cpp后端，batch_size为16，4卡分布式 - ```bash - python examples/test_infer.py --device nvidia --model=/models/9G7B_MHA/ --backend=cpp --tp=4 --batch-size=16 - ``` - - - - 推理服务测试 - - 启动推理服务 - ```bash - python python/infinilm/server/inference_server.py --device [cpu | nvidia | qy | metax | moore | iluvatar | ali | cambricon | hygon] --model= --max-new-tokens=MAX_TOKENS --max-batch-size=MAX_BATCH --tp=NDEV --temperature=TEMP --top-p=TOP_P --top-k=TOP_K --host=HOST --port=PORT - ``` - - - 单卡示例： - ```bash - CUDA_VISIBLE_DEVICES=0 python python/infinilm/server/inference_server.py --device nvidia --model=/models/9G7B_MHA/ --max-new-tokens=100 --max-batch-size=32 --tp=1 --temperature=1.0 --top-p=0.8 --top-k=1 - ``` - - - 多卡分布式示例： - ```bash - CUDA_VISIBLE_DEVICES=0,1,2,3 python python/infinilm/server/inference_server.py --device nvidia --model=/models/9G7B_MHA/ --max-new-tokens=100 --max-batch-size=32 --tp=4 --temperature=1.0 --top-p=0.8 --top-k=1 - ``` - - - 使用paged attention, flash attention后端，cuda graph等功能： - ```bash - CUDA_VISIBLE_DEVICES=0,1,2,3 python python/infinilm/server/inference_server.py --device nvidia --model=/models/9G7B_MHA/ --enable-paged-attn --attn=flash-attn --enable-graph - ``` - - - 测试推理服务性能： - ```bash - python scripts/test_perf.py --verbose - ``` - - - 单请求推理服务测试 - ```bash - python test/service/request.py --content="text:Image 1:" --content="image_url:xxx.jpg" --content="text:Image 2:" --content="image_url:xxxx.jpg" --content="text:Compare the 2 images." - ``` - - - 运行推理基准测试（C-Eval/MMLU） - - ```bash - python test/bench/test_benchmark.py --device [cpu | nvidia | qy | metax | moore | iluvatar | ali | cambricon | hygon] --model --bench {ceval|mmlu} [--backend cpp] [--tp N] [--subject SUBJECT] [--num-samples N] [--max-new-tokens N] [--output-csv PATH] [--cache-dir PATH] - ``` - - - 参数说明： - - `--subject`: 指定科目，支持单个科目、多个科目（逗号分隔）或 `all`（默认值，加载全部科目） - - `--output-csv`: 可选，指定CSV输出文件路径。如未指定则不生成CSV文件。CSV包含每个科目的结果和总体结果 - - `--cache-dir`: 可选，指定数据集缓存目录的父目录。应指向包含 `ceval___ceval-exam` 和 `cais___mmlu` 等数据集子目录的父目录（例如 `~/.cache/huggingface/datasets/`）。设置后脚本优先使用本地 CSV（`pandas.read_csv`）离线加载数据，避免 `load_dataset` 的网络请求 - - - C-Eval示例： - - 单个科目： - ```bash - python test/bench/test_benchmark.py --device nvidia /models/9G7B_MHA --bench ceval --subject middle_school_mathematics --num-samples 100 --backend cpp --tp 1 - ``` - - 多个科目（逗号分隔）： - ```bash - python test/bench/test_benchmark.py --device nvidia /models/9G7B_MHA --bench ceval --subject middle_school_mathematics,high_school_physics --backend cpp --tp 1 --output-csv results.csv - ``` - - 全部科目并输出CSV： - ```bash - python test/bench/test_benchmark.py --device nvidia /models/9G7B_MHA --bench ceval --subject all --backend cpp --tp 1 --output-csv results.csv - ``` - - 使用缓存目录加速加载： - ```bash - python test/bench/test_benchmark.py --device nvidia /models/9G7B_MHA --bench ceval --subject middle_school_mathematics --backend cpp --tp 1 --cache-dir ~/.cache/huggingface/datasets/ - ``` - > 注意：`--cache-dir` 应指向包含 `ceval___ceval-exam` 和 `cais___mmlu` 等数据集子目录的父目录，而不是直接指向这些子目录 - - - MMLU示例： - - 单个科目： - ```bash - python test/bench/test_benchmark.py --device nvidia /models/9G7B_MHA --bench mmlu --subject abstract_algebra --backend cpp --tp 1 - ``` - - 多个科目（逗号分隔）： - ```bash - python test/bench/test_benchmark.py --device nvidia /models/9G7B_MHA --bench mmlu --subject abstract_algebra,anatomy,astronomy --backend cpp --tp 1 --output-csv results.csv - ``` - - 使用缓存目录加速加载： - ```bash - python test/bench/test_benchmark.py --device nvidia /models/9G7B_MHA --bench mmlu --subject abstract_algebra --backend cpp --tp 1 --cache-dir ~/.cache/huggingface/datasets/ - ``` - > 注意：`--cache-dir` 应指向包含 `ceval___ceval-exam` 和 `cais___mmlu` 等数据集子目录的父目录，而不是直接指向这些子目录 - - - 试验中功能 - - Warm Up - ```bash - python examples/bench.py --device nvidia --model= --warmup - ``` - - Paged Attention - ```bash - python examples/bench.py --device nvidia --model= --enable-paged-attn - ``` - - CUDA Graph - ```bash - python examples/bench.py --device nvidia --model= --enable-paged-attn --enable-graph - ``` - - 选择attention后端 (使用flash attention后端需要先在InfiniCore完成相关配置和编译) - ```bash - python examples/bench.py --device nvidia --model= --enable-paged-attn [--attn=default | --attn=flash-attn] - ``` - -## 使用方式(旧版) - -- 编译并安装 `InfiniCore` 。注意根据提示设置好 `INFINI_ROOT` 环境变量（默认为 `$HOME/.infini`）。 - -- 编译并安装 `InfiniLM` - -```bash -xmake && xmake install +InfiniLM is the high-level inference engine in the InfiniTensor stack. It owns +model execution, runtime management, tensor abstractions, and the Python API. +The lower-level runtime, operator, and collective APIs are provided by +[InfiniRT](https://github.com/InfiniTensor/InfiniRT), +[InfiniOps](https://github.com/InfiniTensor/InfiniOps), and +[InfiniCCL](https://github.com/InfiniTensor/InfiniCCL), respectively. + +[InfiniCore](https://github.com/InfiniTensor/InfiniCore) pins those three +projects and provides an integration build. It no longer provides a separate +runtime API or Python package. + +## Build + +Clone both repositories with their submodules: + +```shell +git clone --recurse-submodules https://github.com/InfiniTensor/InfiniCore.git +git clone --recurse-submodules https://github.com/InfiniTensor/InfiniLM.git +``` + +Build the NVIDIA dependency stack from InfiniCore. The default operator set is +the set required by InfiniLM: + +```shell +cd InfiniCore +python3 scripts/build_integration.py --cuda-arch sm_80 --jobs 16 --test +export INFINI_ROOT="$PWD/build/integration/nvidia/prefix" +export LD_LIBRARY_PATH="$INFINI_ROOT/lib:${LD_LIBRARY_PATH:-}" +``` + +Then build and install InfiniLM: + +```shell +cd ../InfiniLM +python3 -m pip install . --no-build-isolation +``` + +The native build currently recognizes `cpu`, `nvidia`, `cambricon`, `ascend`, +`metax`, `moore`, `iluvatar`, and `hygon` platforms. The migration in this +branch is validated on NVIDIA first. + +The modern operator closure currently supports `qwen3`. Other registered model +types are rejected before worker startup until their required InfiniOps adapters +are available. + +## Inference + +Run a single-model smoke test: + +```shell +python examples/test_infer.py --device nvidia --model=/path/to/model ``` -- 运行模型推理测试 +For tensor-parallel inference: -```bash -python scripts/jiuge.py [--cpu | --nvidia | --qy | --cambricon | --ascend | --metax | --moore | --iluvatar | --kunlun | --hygon | --ali] path/to/model_dir [n_device] +```shell +python examples/test_infer.py --device nvidia --model=/path/to/model --tp=4 --batch-size=16 ``` -- 部署模型推理服务 +Start the OpenAI-compatible server: -```bash -python scripts/launch_server.py --model MODEL_PATH [-h] [--dev {cpu,nvidia,qy, cambricon,ascend,metax,moore,iluvatar,kunlun,hygon}] [--ndev NDEV] [--max-batch MAX_BATCH] [--max-new-tokens MAX_TOKENS] +```shell +python python/infinilm/server/inference_server.py --device nvidia --model=/path/to/model --tp=1 ``` -- 测试模型推理服务性能 +Paged attention and graph execution are selected by InfiniLM arguments and are +built as part of InfiniLM: -```bash -python scripts/test_perf.py +```shell +python examples/bench.py --device nvidia --model=/path/to/model --enable-paged-attn --enable-graph ``` -- 使用推理服务测试模型困惑度（Perplexity） +## Development + +Format staged files with the repository formatter: -```bash -python scripts/test_ppl.py --model MODEL_PATH [--ndev NDEV] [--max-batch MAX_BATCH] [--max-new-tokens MAX_TOKENS] +```shell +python scripts/format.py --staged ``` + +Run the static migration contracts with: + +```shell +python -m unittest discover -s test/static -p "test_*.py" +``` + +## License + +InfiniLM is licensed under the MIT License. See [LICENSE](LICENSE). diff --git a/csrc/backends/attention_backends.hpp b/csrc/backends/attention_backends.hpp index b274aacc7..cb0461e05 100644 --- a/csrc/backends/attention_backends.hpp +++ b/csrc/backends/attention_backends.hpp @@ -44,14 +44,16 @@ inline AttentionBackend parse_attention_backend(const std::string &backend) { return AttentionBackend::PAGED_ATTN; } if (backend == "flash-attn") { - return AttentionBackend::FLASH_ATTN; + throw std::invalid_argument( + "`flash-attn` is unsupported until its kernels are available in InfiniOps."); } if (backend == "flashinfer") { - return AttentionBackend::FLASHINFER; + throw std::invalid_argument( + "`flashinfer` is unsupported until its kernels are available in InfiniOps."); } throw std::invalid_argument( - "Invalid attention_backend: " + backend + ". Valid options are: static-attn, paged-attn, flash-attn, flashinfer"); + "Invalid attention_backend: " + backend + ". Valid options are: static-attn, paged-attn"); } } // namespace infinilm::backends diff --git a/csrc/config/config_factory.cpp b/csrc/config/config_factory.cpp index 0467f4536..e00522030 100644 --- a/csrc/config/config_factory.cpp +++ b/csrc/config/config_factory.cpp @@ -1,6 +1,7 @@ #include "config_factory.hpp" #include "../models/models_registry.hpp" #include +#include namespace infinilm::config { @@ -11,12 +12,18 @@ std::shared_ptr ConfigFactory::createConfig(const const std::string model_type = model_config->get("model_type"); const auto &config_map = models::get_model_config_map(); auto it = config_map.find(model_type); - if (it != config_map.end()) { - it->second(model_config); - } else { + if (it == config_map.end()) { throw std::invalid_argument("infinilm::config::ConfigFactory::createConfig: Unsupported model config type: " + model_type); } + static const std::unordered_set kModernModelTypes{"qwen3"}; + if (kModernModelTypes.find(model_type) == kModernModelTypes.end()) { + throw std::invalid_argument( + "infinilm::config::ConfigFactory::createConfig: model type `" + model_type + + "` is unavailable with the modern InfiniOps backend; supported model types: qwen3"); + } + + it->second(model_config); return model_config; } diff --git a/csrc/config/quant_config.cpp b/csrc/config/quant_config.cpp index da261ce09..95bf7815a 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -1,4 +1,5 @@ #include "quant_config.hpp" +#include namespace infinilm::config { QuantConfig::QuantConfig(const nlohmann::json &json) : quantization_config(json) { @@ -13,18 +14,17 @@ QuantConfig::get_quantization_method() const { const std::string quant_method = quantization_config.value("quant_method", ""); - // Determine the quantization scheme from the JSON config if (quant_method == "compressed-tensors") { - return std::make_shared(quantization_config); + throw std::runtime_error( + "`compressed-tensors` quantization is unsupported until its kernels are available in InfiniOps."); } else if (quant_method == "awq") { - return std::make_shared(quantization_config); + throw std::runtime_error( + "AWQ quantization is unsupported until its kernels are available in InfiniOps."); } else if (quant_method == "gptq") { - return std::make_shared(quantization_config); + throw std::runtime_error( + "GPTQ quantization is unsupported until its kernels are available in InfiniOps."); } else { return std::make_shared(quantization_config); } - // Add other schemes as needed - - return std::make_shared(quantization_config); // Default case if no matching scheme } } // namespace infinilm::config diff --git a/csrc/config/quant_config.hpp b/csrc/config/quant_config.hpp index fb0b8abf3..2675f1922 100644 --- a/csrc/config/quant_config.hpp +++ b/csrc/config/quant_config.hpp @@ -1,9 +1,9 @@ #pragma once -#include "../utils.hpp" #include "../layers/quantization/quantization.hpp" +#include "../utils.hpp" #include "nlohmann/json.hpp" #include -#include +#include namespace infinilm::config { @@ -25,23 +25,9 @@ class QuantConfig { } void set_kv_quant_scheme(infinicore::DataType kv_cache_dtype) { - try { - this->kv_cache_dtype_ = std::make_optional(kv_cache_dtype); - switch (kv_cache_dtype) { - case infinicore::DataType::I8: { - this->kv_quant_scheme = infinilm::quantization::KVQuantAlgo::INT8; - break; - } - default: { - spdlog::warn("Unsupported kv_cache_dtype: '{}', fallback to NONE", infinicore::toString(kv_cache_dtype)); - this->kv_quant_scheme = infinilm::quantization::KVQuantAlgo::NONE; - break; - } - } - } catch (const std::exception &e) { - spdlog::error("Failed to parse kv_cache_dtype '{}': {}", infinicore::toString(kv_cache_dtype), e.what()); - this->kv_quant_scheme = infinilm::quantization::KVQuantAlgo::NONE; - } + throw std::runtime_error( + "KV cache INT8 quantization is unsupported until its kernels are available in InfiniOps; requested dtype `" + + infinicore::toString(kv_cache_dtype) + "`."); } infinilm::quantization::KVQuantAlgo get_kv_quant_scheme() const { diff --git a/csrc/debug_utils/tensor_utils.hpp b/csrc/debug_utils/tensor_utils.hpp index 3cd0e2191..30644d41e 100644 --- a/csrc/debug_utils/tensor_utils.hpp +++ b/csrc/debug_utils/tensor_utils.hpp @@ -1,11 +1,11 @@ #pragma once #include "infinicore/tensor.hpp" -#include #include #include -#include +#include #include +#include namespace infinilm::models::debug_utils { @@ -21,32 +21,32 @@ inline void log_tensor_stats(const infinicore::Tensor &tensor, const std::string // Log basic info std::string shape_str = "["; for (size_t i = 0; i < shape.size(); ++i) { - if (i > 0) shape_str += ", "; + if (i > 0) { + shape_str += ", "; + } shape_str += std::to_string(shape[i]); } shape_str += "]"; - SPDLOG_INFO(" {}: shape={}, dtype={}, device={}", name, shape_str, static_cast(dtype), device.toString()); + SPDLOG_INFO(" {}: shape={}, dtype={}, device={}", name, shape_str, static_cast(dtype), device.ToString()); // For F32, F16, and BF16 tensors, compute and log statistics - if (dtype == infinicore::DataType::F32 || - dtype == infinicore::DataType::F16 || - dtype == infinicore::DataType::BF16) { + if (dtype == infinicore::DataType::kFloat32 || dtype == infinicore::DataType::kFloat16 || dtype == infinicore::DataType::kBFloat16) { // Copy to CPU if needed and compute stats - auto cpu_tensor = tensor->to(infinicore::Device(infinicore::Device::Type::CPU, 0)); + auto cpu_tensor = tensor->to(infinicore::Device(infinicore::Device::Type::kCpu, 0)); std::byte *raw_data = cpu_tensor->data(); size_t numel = cpu_tensor->numel(); if (numel > 0) { - if (dtype == infinicore::DataType::F32) { - float *data = reinterpret_cast(raw_data); + if (dtype == infinicore::DataType::kFloat32) { + float *data = reinterpret_cast(raw_data); float min_val = *std::min_element(data, data + numel); float max_val = *std::max_element(data, data + numel); float sum = std::accumulate(data, data + numel, 0.0f); float mean_val = sum / static_cast(numel); SPDLOG_INFO(" Stats: min={:.6e}, max={:.6e}, mean={:.6e}, numel={}", - min_val, max_val, mean_val, numel); + min_val, max_val, mean_val, numel); // Log sample values at specific positions if (log_samples && numel > 0) { @@ -56,9 +56,9 @@ inline void log_tensor_stats(const infinicore::Tensor &tensor, const std::string SPDLOG_INFO(" [{}] = {:.6e}", i, data[i]); } } - } else if (dtype == infinicore::DataType::F16) { + } else if (dtype == infinicore::DataType::kFloat16) { // F16 is typically uint16_t, need to convert to float for logging - uint16_t *data = reinterpret_cast(raw_data); + uint16_t *data = reinterpret_cast(raw_data); std::vector float_data(numel); for (size_t i = 0; i < numel; ++i) { // Simple F16 to F32 conversion (approximate) @@ -67,7 +67,7 @@ inline void log_tensor_stats(const infinicore::Tensor &tensor, const std::string uint32_t exp = (h >> 10) & 0x1F; uint32_t mant = h & 0x3FF; uint32_t f32 = (sign << 31) | ((exp + 112) << 23) | (mant << 13); - float_data[i] = *reinterpret_cast(&f32); + float_data[i] = *reinterpret_cast(&f32); } float min_val = *std::min_element(float_data.begin(), float_data.end()); float max_val = *std::max_element(float_data.begin(), float_data.end()); @@ -75,7 +75,7 @@ inline void log_tensor_stats(const infinicore::Tensor &tensor, const std::string float mean_val = sum / static_cast(numel); SPDLOG_INFO(" Stats (F16): min={:.6e}, max={:.6e}, mean={:.6e}, numel={}", - min_val, max_val, mean_val, numel); + min_val, max_val, mean_val, numel); if (log_samples && numel > 0) { size_t sample_count = std::min(max_samples, numel); @@ -84,15 +84,15 @@ inline void log_tensor_stats(const infinicore::Tensor &tensor, const std::string SPDLOG_INFO(" [{}] = {:.6e}", i, float_data[i]); } } - } else if (dtype == infinicore::DataType::BF16) { + } else if (dtype == infinicore::DataType::kBFloat16) { // BF16 is typically uint16_t, need to convert to float for logging - uint16_t *data = reinterpret_cast(raw_data); + uint16_t *data = reinterpret_cast(raw_data); std::vector float_data(numel); for (size_t i = 0; i < numel; ++i) { // BF16 to F32 conversion uint16_t b = data[i]; uint32_t f32 = (static_cast(b) << 16); - float_data[i] = *reinterpret_cast(&f32); + float_data[i] = *reinterpret_cast(&f32); } float min_val = *std::min_element(float_data.begin(), float_data.end()); float max_val = *std::max_element(float_data.begin(), float_data.end()); @@ -100,7 +100,7 @@ inline void log_tensor_stats(const infinicore::Tensor &tensor, const std::string float mean_val = sum / static_cast(numel); SPDLOG_INFO(" Stats (BF16): min={:.6e}, max={:.6e}, mean={:.6e}, numel={}", - min_val, max_val, mean_val, numel); + min_val, max_val, mean_val, numel); if (log_samples && numel > 0) { size_t sample_count = std::min(max_samples, numel); @@ -133,16 +133,16 @@ inline void log_tensor_positions(const infinicore::Tensor &tensor, const std::st auto dtype = tensor->dtype(); // Only log for F32 tensors (or copy to CPU) - if (dtype == infinicore::DataType::F32) { - auto cpu_tensor = tensor->to(infinicore::Device(infinicore::Device::Type::CPU, 0)); + if (dtype == infinicore::DataType::kFloat32) { + auto cpu_tensor = tensor->to(infinicore::Device(infinicore::Device::Type::kCpu, 0)); std::byte *raw_data = cpu_tensor->data(); - float *data = reinterpret_cast(raw_data); + float *data = reinterpret_cast(raw_data); SPDLOG_INFO(" {}: Logging specific positions:", name); for (const auto &pos : positions) { if (pos.size() != shape.size()) { SPDLOG_INFO(" Position {}: dimension mismatch (expected {} dims, got {})", - pos.size(), shape.size()); + pos.size(), shape.size()); continue; } @@ -162,7 +162,9 @@ inline void log_tensor_positions(const infinicore::Tensor &tensor, const std::st if (valid && idx < cpu_tensor->numel()) { std::string pos_str = "["; for (size_t i = 0; i < pos.size(); ++i) { - if (i > 0) pos_str += ", "; + if (i > 0) { + pos_str += ", "; + } pos_str += std::to_string(pos[i]); } pos_str += "]"; @@ -170,7 +172,9 @@ inline void log_tensor_positions(const infinicore::Tensor &tensor, const std::st } else { std::string pos_str = "["; for (size_t i = 0; i < pos.size(); ++i) { - if (i > 0) pos_str += ", "; + if (i > 0) { + pos_str += ", "; + } pos_str += std::to_string(pos[i]); } pos_str += "]"; diff --git a/csrc/engine/compiler/graph_compiler.hpp b/csrc/engine/compiler/graph_compiler.hpp index 5173994fd..4da6fe45e 100644 --- a/csrc/engine/compiler/graph_compiler.hpp +++ b/csrc/engine/compiler/graph_compiler.hpp @@ -5,6 +5,33 @@ namespace infinilm::engine { +class GraphRecordingGuard { +public: + GraphRecordingGuard() { + infinicore::context::startGraphRecording(); + } + + ~GraphRecordingGuard() noexcept { + if (active_) { + infinicore::context::cancelGraphRecording(); + } + } + + GraphRecordingGuard(const GraphRecordingGuard &) = delete; + GraphRecordingGuard &operator=(const GraphRecordingGuard &) = delete; + GraphRecordingGuard(GraphRecordingGuard &&) = delete; + GraphRecordingGuard &operator=(GraphRecordingGuard &&) = delete; + + std::shared_ptr finish() { + auto graph = infinicore::context::stopGraphRecording(); + active_ = false; + return graph; + } + +private: + bool active_ = true; +}; + class GraphCompiler { public: using Compiled = std::tuple< diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index f267794e2..21ed89953 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -26,30 +26,30 @@ void PagedCompiler::compile() { size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end()); compiled_map_decode_.clear(); block_tables_holder_ = infinicore::Tensor::empty( - {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice()); + {nblocks * max_batch_size}, infinicore::DataType::kInt32, infinicore::context::getDevice()); set_zeros(block_tables_holder_); auto make_decode_input = [&](size_t b) { InfinilmModel::Input input; - input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice()); + input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::kInt64, infinicore::context::getDevice()); + input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::kInt64, infinicore::context::getDevice()); + input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::kInt32, infinicore::context::getDevice()); set_zeros(input.input_ids.value()); set_zeros(input.position_ids.value()); set_zeros(input.total_sequence_lengths.value()); std::vector total_sequence_lengths_vec(b, 1); infinicore::context::memcpyH2D(input.total_sequence_lengths.value()->data(), total_sequence_lengths_vec.data(), b * sizeof(int32_t), false); - input.input_offsets = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); + input.input_offsets = infinicore::Tensor::empty({b + 1}, infinicore::DataType::kInt32, infinicore::context::getDevice()); std::vector input_offsets_vec(b + 1, 0); for (size_t i = 0; i <= b; i++) { input_offsets_vec[i] = i; } infinicore::context::memcpyH2D(input.input_offsets.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); - input.cu_seqlens = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); + input.cu_seqlens = infinicore::Tensor::empty({b + 1}, infinicore::DataType::kInt32, infinicore::context::getDevice()); infinicore::context::memcpyH2D(input.cu_seqlens.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); const size_t block_per_req = nblocks; input.block_tables = block_tables_holder_->as_strided({b, block_per_req}, {(ptrdiff_t)block_per_req, 1}); - input.slot_mapping = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); + input.slot_mapping = infinicore::Tensor::empty({b}, infinicore::DataType::kInt64, infinicore::context::getDevice()); set_zeros(input.slot_mapping.value()); // Attention reads attn_metadata from thread-local forward context. @@ -69,9 +69,7 @@ void PagedCompiler::compile() { auto input = make_decode_input(warmup_batch_size); model_->forward(input); infinicore::context::syncStream(); - // Warmup runs the eager Marlin path and may leave per-layer lock - // workspaces dirty. Reset before CUDA graph capture so capture - // starts from the same all-zero lock state as normal execution. + // Clear transient operator state before CUDA graph capture. model_->reset_runtime_state(); infinicore::context::syncStream(); } @@ -82,15 +80,12 @@ void PagedCompiler::compile() { barrier_->wait(); (void)model_->forward(input); infinicore::context::syncStream(); - // Capture must not start with stale Marlin locks from previous - // warmup/capture attempts. This reset is intentionally outside - // graph capture; the current implementation still pays a memset - // before every graph replay in get_compiled(). + // Capture must not start with stale state from previous attempts. model_->reset_runtime_state(); infinicore::context::syncStream(); - infinicore::context::startGraphRecording(); + GraphRecordingGuard recording; auto output = model_->forward(input); - auto graph = infinicore::context::stopGraphRecording(); + auto graph = recording.finish(); barrier_->wait(); auto shared_output = std::shared_ptr( @@ -135,11 +130,7 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & set_minus_one_device_async(graph_block_tables); graph_block_tables->narrow({{1, 0, block_per_req}})->copy_from(input.block_tables.value()); graph_input.slot_mapping.value()->copy_from(input.slot_mapping.value()); - // CUDA graph replay reuses the same per-layer Marlin workspaces. - // The graph itself does not contain a workspace reset, so enqueue - // one on the same stream before launch. This is correct but costs - // decode latency; the intended follow-up is a reusable global - // zero workspace/lock buffer shared by all Marlin layers. + // Reset transient state on the graph stream before replay. model_->reset_runtime_state(); auto graph = std::get<0>(result->second.compiled); diff --git a/csrc/engine/compiler/static_batching_compiler.cpp b/csrc/engine/compiler/static_batching_compiler.cpp index af2f4799f..8a2106780 100644 --- a/csrc/engine/compiler/static_batching_compiler.cpp +++ b/csrc/engine/compiler/static_batching_compiler.cpp @@ -11,10 +11,10 @@ void StaticBatchingCompiler::compile() { if (model_->get_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { size_t b = dynamic_cast(model_->get_cache_config())->max_batch_size(); InfinilmModel::Input input; - input.input_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.position_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.past_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); + input.input_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::kInt64, infinicore::context::getDevice()); + input.position_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::kInt64, infinicore::context::getDevice()); + input.past_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::kInt64, infinicore::context::getDevice()); + input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::kInt64, infinicore::context::getDevice()); std::vector total_sequence_lengths_vec(b, 1); infinicore::context::memcpyH2D(input.total_sequence_lengths.value()->data(), total_sequence_lengths_vec.data(), b * sizeof(int64_t), false); @@ -32,9 +32,9 @@ void StaticBatchingCompiler::compile() { (void)model_->forward(input); infinicore::context::syncStream(); - infinicore::context::startGraphRecording(); + GraphRecordingGuard recording; auto output = model_->forward(input); - auto graph = infinicore::context::stopGraphRecording(); + auto graph = recording.finish(); barrier_->wait(); auto shared_output = std::shared_ptr(new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)}); diff --git a/csrc/engine/distributed/communication_group.cpp b/csrc/engine/distributed/communication_group.cpp index 782faa9ec..4185e9113 100644 --- a/csrc/engine/distributed/communication_group.cpp +++ b/csrc/engine/distributed/communication_group.cpp @@ -1,27 +1,147 @@ #include "communication_group.hpp" -#include "../../utils.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include namespace infinilm::engine::distributed { +namespace { + +void checkInfiniccl(const char *operation, infinicclResult_t result) { + if (result == infinicclSuccess) { + return; + } + throw std::runtime_error("InfiniCCL operation `" + std::string(operation) + + "` failed with result " + std::to_string(static_cast(result))); +} + +void destroyCommunicators(infinicore::Device::Type device_type, + const std::vector &device_ids, + std::vector &communicators) noexcept { + infinicore::Device previous_device; + bool restore_device = false; + try { + previous_device = infinicore::context::getDevice(); + restore_device = true; + } catch (...) { + } + + for (size_t rank = 0; rank < communicators.size(); ++rank) { + if (communicators[rank] == nullptr) { + continue; + } + try { + infinicore::context::setDevice(infinicore::Device(device_type, device_ids[rank])); + (void)infinicclCommDestroy(communicators[rank]); + } catch (...) { + } + communicators[rank] = nullptr; + } + + if (restore_device) { + try { + infinicore::context::setDevice(previous_device); + } catch (...) { + } + } +} + +} // namespace CommunicationGroup::CommunicationGroup(const DistConfig &dist_config, infinicore::Device::Type device_type) : dist_config_(dist_config), device_type_(device_type), - communicators_(std::vector(dist_config.tp_device_ids.size(), nullptr)) { + communicators_(dist_config.tp_device_ids.size(), nullptr) { + const size_t world_size = dist_config_.tp_device_ids.size(); + if (world_size == 0) { + throw std::invalid_argument("tensor parallel device list must not be empty"); + } + if (world_size > static_cast(std::numeric_limits::max())) { + throw std::invalid_argument("tensor parallel world size exceeds the InfiniCCL rank limit"); + } - size_t world_size = dist_config_.tp_device_ids.size(); - size_t device_count = infinicore::context::getDeviceCount(device_type); - if (device_count < world_size) { - throw std::runtime_error("infinilm::engine::distributed::CommunicationGroup error, world size is larger than the number of available GPUs. world size: " + std::to_string(world_size) + ", device count: " + std::to_string(device_count)); + const size_t device_count = infinicore::context::getDeviceCount(device_type_); + std::unordered_set unique_device_ids; + for (int device_id : dist_config_.tp_device_ids) { + if (device_id < 0 || static_cast(device_id) >= device_count) { + throw std::invalid_argument("tensor parallel device ID " + std::to_string(device_id) + + " is outside the available range [0, " + + std::to_string(device_count) + ")"); + } + if (!unique_device_ids.insert(device_id).second) { + throw std::invalid_argument("tensor parallel device ID " + std::to_string(device_id) + + " is duplicated"); + } + } + + if (world_size == 1) { + return; + } + + infinicclUniqueId unique_id{}; + checkInfiniccl("infinicclGetUniqueId", infinicclGetUniqueId(&unique_id)); + + std::vector errors(world_size); + std::vector workers; + workers.reserve(world_size); + std::mutex start_mutex; + std::condition_variable start_cv; + bool start = false; + bool cancel = false; + try { + for (size_t rank = 0; rank < world_size; ++rank) { + workers.emplace_back([&, rank] { + { + std::unique_lock lock(start_mutex); + start_cv.wait(lock, [&] { return start || cancel; }); + if (cancel) { + return; + } + } + try { + infinicore::context::setDevice( + infinicore::Device(device_type_, dist_config_.tp_device_ids[rank])); + checkInfiniccl( + "infinicclCommInitRank", + infinicclCommInitRank(&communicators_[rank], + static_cast(world_size), + unique_id, + static_cast(rank))); + } catch (...) { + errors[rank] = std::current_exception(); + } + }); + } + } catch (...) { + { + std::lock_guard lock(start_mutex); + cancel = true; + } + start_cv.notify_all(); + for (auto &worker : workers) { + worker.join(); + } + destroyCommunicators(device_type_, dist_config_.tp_device_ids, communicators_); + throw; } - if (infinicore::context::getDevice().getType() != device_type_) { - infinicore::context::setDevice(infinicore::Device(device_type_, 0)); + { + std::lock_guard lock(start_mutex); + start = true; } - if (world_size > 1) { - RUN_INFINI(infinicclCommInitAll( - (infiniDevice_t)infinicore::context::getDevice().getType(), - communicators_.data(), - dist_config.tp_device_ids.size(), - dist_config.tp_device_ids.data())); + start_cv.notify_all(); + for (auto &worker : workers) { + worker.join(); + } + for (const auto &error : errors) { + if (error) { + destroyCommunicators(device_type_, dist_config_.tp_device_ids, communicators_); + std::rethrow_exception(error); + } } } @@ -30,24 +150,23 @@ const DistConfig &CommunicationGroup::get_dist_config() const { } RankInfo CommunicationGroup::get_rank_info(int rank) const { - RankInfo info; - info.tp_size = dist_config_.tp_device_ids.size(); + if (rank < 0 || static_cast(rank) >= dist_config_.tp_device_ids.size()) { + throw std::out_of_range("tensor parallel rank " + std::to_string(rank) + " is out of range"); + } + + RankInfo info(infinicore::Device(device_type_, dist_config_.tp_device_ids[rank])); + info.tp_size = static_cast(dist_config_.tp_device_ids.size()); info.tp_rank = rank; - info.device = infinicore::Device(device_type_, dist_config_.tp_device_ids[rank]); info.comm = communicators_[rank]; return info; } int CommunicationGroup::get_world_size() const { - return dist_config_.tp_device_ids.size(); + return static_cast(dist_config_.tp_device_ids.size()); } CommunicationGroup::~CommunicationGroup() { - if (communicators_.size() > 1) { - for (auto &comm : communicators_) { - infinicclCommDestroy(comm); - } - } + destroyCommunicators(device_type_, dist_config_.tp_device_ids, communicators_); } } // namespace infinilm::engine::distributed diff --git a/csrc/engine/distributed/communication_group.hpp b/csrc/engine/distributed/communication_group.hpp index e4f3c81a8..1ba78133e 100644 --- a/csrc/engine/distributed/communication_group.hpp +++ b/csrc/engine/distributed/communication_group.hpp @@ -2,10 +2,11 @@ #include "dist_config.hpp" -#include +#include #include #include +#include #include namespace infinilm::engine::distributed { @@ -21,12 +22,12 @@ struct RankInfo { // Communicator handle infinicclComm_t comm; - RankInfo(infinicore::Device _device = infinicore::context::getDevice()) - : tp_size(1), tp_rank(0), device(_device), comm(nullptr){}; + explicit RankInfo(infinicore::Device device = infinicore::context::getDevice()) + : device(std::move(device)), tp_size(1), tp_rank(0), comm(nullptr) {} std::string to_string() const { std::stringstream ss; - ss << "RankInfo: device=" << device.toString() << ", tp_size=" << tp_size << ", tp_rank=" << tp_rank; + ss << "RankInfo: device=" << device.ToString() << ", tp_size=" << tp_size << ", tp_rank=" << tp_rank; return ss.str(); } }; @@ -35,6 +36,10 @@ struct RankInfo { class CommunicationGroup { public: explicit CommunicationGroup(const DistConfig &dist_config, infinicore::Device::Type device_type); + CommunicationGroup(const CommunicationGroup &) = delete; + CommunicationGroup &operator=(const CommunicationGroup &) = delete; + CommunicationGroup(CommunicationGroup &&) = delete; + CommunicationGroup &operator=(CommunicationGroup &&) = delete; const DistConfig &get_dist_config() const; diff --git a/csrc/engine/infer_engine.hpp b/csrc/engine/infer_engine.hpp index 4c0c0345c..324ba0243 100644 --- a/csrc/engine/infer_engine.hpp +++ b/csrc/engine/infer_engine.hpp @@ -26,7 +26,7 @@ class InferEngine { InferEngine( const std::string &config_str, const distributed::DistConfig &distributed_config = distributed::DistConfig(), - infinicore::Device::Type device_type = infinicore::context::getDevice().getType(), + infinicore::Device::Type device_type = infinicore::context::getDevice().type(), const cache::CacheConfig *cache_config = nullptr, bool enable_graph_compiling = false, backends::AttentionBackend attention_backend = backends::AttentionBackend::Default, diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 11696e1eb..fff417d4d 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -420,7 +420,7 @@ void RankWorker::thread_loop() { // All-position speculative/MTP runs need eager mode because // hidden states are not part of compiled graph outputs. if (!local_args.sample_all_positions && compiler_ != nullptr) { - auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device::cpu())); + auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device{infinicore::Device::Type::kCpu})); if (graph != nullptr && output != nullptr) { graph->run(); logits = output->logits; @@ -450,7 +450,7 @@ void RankWorker::thread_loop() { const bool sample_all_positions = local_args.sample_all_positions; const size_t n_out = sample_all_positions ? static_cast(input_offsets[n_req]) : n_req; - auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device)}; + auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::kInt64, rank_info_.device)}; for (size_t i{0}; i < n_out; ++i) { size_t score_idx = i; @@ -464,7 +464,7 @@ void RankWorker::thread_loop() { out, score, random_val, top_p, top_k, temperature); } - output_ids = output_ids->to(infinicore::Device::cpu()); + output_ids = output_ids->to(infinicore::Device{infinicore::Device::Type::kCpu}); infinicore::context::syncStream(); diff --git a/csrc/infinicore/include/infinicore.hpp b/csrc/infinicore/include/infinicore.hpp new file mode 100644 index 000000000..4d4304e3c --- /dev/null +++ b/csrc/infinicore/include/infinicore.hpp @@ -0,0 +1,7 @@ +#pragma once + +#include "infinicore/device_event.hpp" +#include "infinicore/io.hpp" +#include "infinicore/nn.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/tensor.hpp" diff --git a/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp b/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp new file mode 100644 index 000000000..f48f640cd --- /dev/null +++ b/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp @@ -0,0 +1,73 @@ +#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) +#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::kFloat32: + return at::kFloat; + case DataType::kFloat16: + return at::kHalf; + case DataType::kBFloat16: + return at::kBFloat16; + case DataType::kInt32: + return at::kInt; + case DataType::kInt64: + 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 devices as CUDA devices for ATen interoperability. + if (device.type() == Device::Type::kNvidia || device.type() == Device::Type::kMetax || device.type() == Device::Type::kHygon) { + return at::Device(at::kCUDA, device.index()); + } else if (device.type() == Device::Type::kCpu) { + return at::Device(at::kCPU); + } +#if defined(ENABLE_MOORE_API) + else if (device.type() == Device::Type::kMoore) { + return at::Device(at::DeviceType::PrivateUse1, device.index()); + } +#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) +c10::cuda::CUDAStream get_cuda_stream(); +#endif + +#if defined(ENABLE_MOORE_API) +c10::musa::MUSAStream get_musa_stream(); +#endif + +} // namespace infinicore::adaptor + +#endif // ENABLE_ATEN diff --git a/csrc/infinicore/include/infinicore/adaptor/flash_attention_adaptor.hpp b/csrc/infinicore/include/infinicore/adaptor/flash_attention_adaptor.hpp new file mode 100644 index 000000000..c5bf14858 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/analyzer.hpp b/csrc/infinicore/include/infinicore/analyzer.hpp new file mode 100644 index 000000000..81e9ab1a9 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/analyzer/intent_generator.hpp b/csrc/infinicore/include/infinicore/analyzer/intent_generator.hpp new file mode 100644 index 000000000..b166e4c1e --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/analyzer/mutual_awareness_analyzer.hpp b/csrc/infinicore/include/infinicore/analyzer/mutual_awareness_analyzer.hpp new file mode 100644 index 000000000..d373da694 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/analyzer/op_trace.hpp b/csrc/infinicore/include/infinicore/analyzer/op_trace.hpp new file mode 100644 index 000000000..0a8296cd6 --- /dev/null +++ b/csrc/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. + 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/csrc/infinicore/include/infinicore/analyzer/op_type.hpp b/csrc/infinicore/include/infinicore/analyzer/op_type.hpp new file mode 100644 index 000000000..b92e1d887 --- /dev/null +++ b/csrc/infinicore/include/infinicore/analyzer/op_type.hpp @@ -0,0 +1,232 @@ +#pragma once + +#include +#include + +namespace infinicore::analyzer { + +/// Op types recognized by the analyzer for phase detection. +/// This is not an exhaustive list of operations, 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, + + // --- 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::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/csrc/infinicore/include/infinicore/analyzer/op_type_registry.hpp b/csrc/infinicore/include/infinicore/analyzer/op_type_registry.hpp new file mode 100644 index 000000000..8f3681623 --- /dev/null +++ b/csrc/infinicore/include/infinicore/analyzer/op_type_registry.hpp @@ -0,0 +1,48 @@ +#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}, + // Misc + {"Rearrange", OpType::REARRANGE}, + }; + auto it = registry.find(name); + return it != registry.end() ? it->second : OpType::UNKNOWN; +} + +} // namespace infinicore::analyzer diff --git a/csrc/infinicore/include/infinicore/analyzer/optimization_intent.hpp b/csrc/infinicore/include/infinicore/analyzer/optimization_intent.hpp new file mode 100644 index 000000000..19d6c868c --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/analyzer/phase_detector.hpp b/csrc/infinicore/include/infinicore/analyzer/phase_detector.hpp new file mode 100644 index 000000000..2f4828c6e --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/analyzer/resource_sensor.hpp b/csrc/infinicore/include/infinicore/analyzer/resource_sensor.hpp new file mode 100644 index 000000000..c2e8e6af8 --- /dev/null +++ b/csrc/infinicore/include/infinicore/analyzer/resource_sensor.hpp @@ -0,0 +1,180 @@ +#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 and the analyzer consumes only this normalized view. +struct DeviceResourceSnapshot { + int device_id = -1; + Device::Type device_type = Device::Type::kCpu; + + 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::kCpu) 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/csrc/infinicore/include/infinicore/common/LRUCache.hpp b/csrc/infinicore/include/infinicore/common/LRUCache.hpp new file mode 100644 index 000000000..0c13d1d09 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/common/hash.hpp b/csrc/infinicore/include/infinicore/common/hash.hpp new file mode 100644 index 000000000..65178e4db --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/context/context.hpp b/csrc/infinicore/include/infinicore/context/context.hpp new file mode 100644 index 000000000..33e92e4db --- /dev/null +++ b/csrc/infinicore/include/infinicore/context/context.hpp @@ -0,0 +1,56 @@ +#pragma once + +#include "../device.hpp" +#include "../memory.hpp" + +#include "../graph/graph.hpp" + +#include + +#include + +namespace infinicore { + +namespace context { +void setDevice(Device device); +Device getDevice(); +size_t getDeviceCount(Device::Type type); + +infini::rt::runtime::Stream getStream(); + +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, infini::rt::runtime::Stream stream); + +// Timing APIs for performance measurement +infini::rt::runtime::Event createEvent(); +infini::rt::runtime::Event createEventWithFlags(uint32_t flags); +void recordEvent(infini::rt::runtime::Event event, infini::rt::runtime::Stream stream = nullptr); +bool queryEvent(infini::rt::runtime::Event event); +void synchronizeEvent(infini::rt::runtime::Event event); +void destroyEvent(infini::rt::runtime::Event event); +float elapsedTime(infini::rt::runtime::Event start, infini::rt::runtime::Event end); +void streamWaitEvent(infini::rt::runtime::Stream stream, infini::rt::runtime::Event event); + +// Graph recording APIs +bool isGraphRecording(); +void startGraphRecording(); +void addGraphOperator(std::shared_ptr op); +std::shared_ptr stopGraphRecording(); +void cancelGraphRecording() noexcept; + +} // namespace context + +} // namespace infinicore diff --git a/csrc/infinicore/include/infinicore/device.hpp b/csrc/infinicore/include/infinicore/device.hpp new file mode 100644 index 000000000..ad981e42e --- /dev/null +++ b/csrc/infinicore/include/infinicore/device.hpp @@ -0,0 +1,9 @@ +#pragma once + +#include + +namespace infinicore { + +using Device = infini::rt::Device; + +} // namespace infinicore diff --git a/csrc/infinicore/include/infinicore/device_event.hpp b/csrc/infinicore/include/infinicore/device_event.hpp new file mode 100644 index 000000000..b09ac1553 --- /dev/null +++ b/csrc/infinicore/include/infinicore/device_event.hpp @@ -0,0 +1,125 @@ +#pragma once + +#include "device.hpp" +#include +#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: + infini::rt::runtime::Event event_ = nullptr; // 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() noexcept; + + /** + * @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(infini::rt::runtime::Stream 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(infini::rt::runtime::Stream 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 + */ + infini::rt::runtime::Event 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/csrc/infinicore/include/infinicore/dtype.hpp b/csrc/infinicore/include/infinicore/dtype.hpp new file mode 100644 index 000000000..279738118 --- /dev/null +++ b/csrc/infinicore/include/infinicore/dtype.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include + +#include +#include + +namespace infinicore { + +using DataType = infini::rt::DataType; + +std::string toString(const DataType &dtype); +std::size_t dsize(const DataType &dtype); + +} // namespace infinicore diff --git a/csrc/infinicore/include/infinicore/graph/graph.hpp b/csrc/infinicore/include/infinicore/graph/graph.hpp new file mode 100644 index 000000000..6995d7f4f --- /dev/null +++ b/csrc/infinicore/include/infinicore/graph/graph.hpp @@ -0,0 +1,148 @@ +#pragma once + +#include +#include + +#include "../tensor.hpp" + +namespace infinicore { +class Runtime; +} + +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 { +private: + // Declared first so it outlives operators and the native device graph. + std::shared_ptr<::infinicore::Runtime> runtime_lease_; + std::shared_ptr allocation_lease_; + +public: + Graph(); + ~Graph() noexcept; + + void run() const; + +protected: + void add_operator(std::shared_ptr op); + void instantiate(); + std::vector> op_list_; + + friend class GraphManager; + +private: + void retain_runtime(std::shared_ptr<::infinicore::Runtime> runtime, + std::shared_ptr allocation_lease); + + struct DeviceGraph; + std::unique_ptr device_graph_; + + friend class ::infinicore::Runtime; +}; +} // 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.type()), \ + static_cast(__trace_device.index())); \ + } 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__, ...) \ + const bool ___recording = context::isGraphRecording(); \ + auto ___op = std::make_shared<__OP_NAME__>(__VA_ARGS__); \ + if (___recording) { \ + 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/csrc/infinicore/include/infinicore/io.hpp b/csrc/infinicore/include/infinicore/io.hpp new file mode 100644 index 000000000..8d11a88b2 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/memory.hpp b/csrc/infinicore/include/infinicore/memory.hpp new file mode 100644 index 000000000..65d812d67 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/nn.hpp b/csrc/infinicore/include/infinicore/nn.hpp new file mode 100644 index 000000000..155a37083 --- /dev/null +++ b/csrc/infinicore/include/infinicore/nn.hpp @@ -0,0 +1,5 @@ +#pragma once + +#include "nn/embedding.hpp" +#include "nn/rmsnorm.hpp" +#include "nn/rope.hpp" diff --git a/csrc/infinicore/include/infinicore/nn/embedding.hpp b/csrc/infinicore/include/infinicore/nn/embedding.hpp new file mode 100644 index 000000000..754d4a3ed --- /dev/null +++ b/csrc/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::kFloat32) + * @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::kFloat32, + 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/csrc/infinicore/include/infinicore/nn/layer_norm.hpp b/csrc/infinicore/include/infinicore/nn/layer_norm.hpp new file mode 100644 index 000000000..a18b31d25 --- /dev/null +++ b/csrc/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::kFloat32) + * @param device Device to create the weight on + */ + LayerNorm(size_t normalized_shape, + double eps = 1e-6, + const DataType &dtype = DataType::kFloat32, + 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/csrc/infinicore/include/infinicore/nn/module.hpp b/csrc/infinicore/include/infinicore/nn/module.hpp new file mode 100644 index 000000000..e36767b4e --- /dev/null +++ b/csrc/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::kFloat32, 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::kFloat32, 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::kFloat32, device)) +#define INFINICORE_NN_BUFFER_INIT(name, args) \ + name##_ = infinicore::nn::Parameter args; \ + this->register_buffer(#name, name##_) + +} // namespace infinicore::nn diff --git a/csrc/infinicore/include/infinicore/nn/parameter.hpp b/csrc/infinicore/include/infinicore/nn/parameter.hpp new file mode 100644 index 000000000..1602f58d6 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/nn/rmsnorm.hpp b/csrc/infinicore/include/infinicore/nn/rmsnorm.hpp new file mode 100644 index 000000000..0c1e80a70 --- /dev/null +++ b/csrc/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::kFloat32) + * @param device Device to create the weight on + */ + RMSNorm(size_t normalized_shape, + double eps = 1e-6, + const DataType &dtype = DataType::kFloat32, + 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/csrc/infinicore/include/infinicore/nn/rope.hpp b/csrc/infinicore/include/infinicore/nn/rope.hpp new file mode 100644 index 000000000..31731d727 --- /dev/null +++ b/csrc/infinicore/include/infinicore/nn/rope.hpp @@ -0,0 +1,112 @@ +#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::kFloat32) + * @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::kFloat32, + 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_; } + + // 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/csrc/infinicore/include/infinicore/nn/rope_scaling_configs.hpp b/csrc/infinicore/include/infinicore/nn/rope_scaling_configs.hpp new file mode 100644 index 000000000..4129e6053 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops.hpp b/csrc/infinicore/include/infinicore/ops.hpp new file mode 100644 index 000000000..b5c4ff18f --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops.hpp @@ -0,0 +1,76 @@ +#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/embedding.hpp" +#include "ops/flash_attention.hpp" +#include "ops/fmin.hpp" +#include "ops/fmod.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_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/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" diff --git a/csrc/infinicore/include/infinicore/ops/acos.hpp b/csrc/infinicore/include/infinicore/ops/acos.hpp new file mode 100644 index 000000000..91aaee020 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/adaptive_avg_pool1d.hpp b/csrc/infinicore/include/infinicore/ops/adaptive_avg_pool1d.hpp new file mode 100644 index 000000000..3e8a95b8d --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/adaptive_avg_pool3d.hpp b/csrc/infinicore/include/infinicore/ops/adaptive_avg_pool3d.hpp new file mode 100644 index 000000000..df6a4d2c3 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/adaptive_max_pool1d.hpp b/csrc/infinicore/include/infinicore/ops/adaptive_max_pool1d.hpp new file mode 100644 index 000000000..51832cae0 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/add.hpp b/csrc/infinicore/include/infinicore/ops/add.hpp new file mode 100644 index 000000000..528cca18a --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/add_rms_norm.hpp b/csrc/infinicore/include/infinicore/ops/add_rms_norm.hpp new file mode 100644 index 000000000..50064e0a4 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/addbmm.hpp b/csrc/infinicore/include/infinicore/ops/addbmm.hpp new file mode 100644 index 000000000..6c17a35d5 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/addcmul.hpp b/csrc/infinicore/include/infinicore/ops/addcmul.hpp new file mode 100644 index 000000000..2806d1218 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/addr.hpp b/csrc/infinicore/include/infinicore/ops/addr.hpp new file mode 100644 index 000000000..cff78eade --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/affine_grid.hpp b/csrc/infinicore/include/infinicore/ops/affine_grid.hpp new file mode 100644 index 000000000..ea025c24d --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/argwhere.hpp b/csrc/infinicore/include/infinicore/ops/argwhere.hpp new file mode 100644 index 000000000..24f4fa65e --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/asin.hpp b/csrc/infinicore/include/infinicore/ops/asin.hpp new file mode 100644 index 000000000..bc5973978 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/asinh.hpp b/csrc/infinicore/include/infinicore/ops/asinh.hpp new file mode 100644 index 000000000..505eb97d9 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/asum.hpp b/csrc/infinicore/include/infinicore/ops/asum.hpp new file mode 100644 index 000000000..6471e6fe4 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/atanh.hpp b/csrc/infinicore/include/infinicore/ops/atanh.hpp new file mode 100644 index 000000000..5c3894e5c --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/attention.hpp b/csrc/infinicore/include/infinicore/ops/attention.hpp new file mode 100644 index 000000000..1bc447c77 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/avg_pool1d.hpp b/csrc/infinicore/include/infinicore/ops/avg_pool1d.hpp new file mode 100644 index 000000000..4bf69bc2a --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/axpy.hpp b/csrc/infinicore/include/infinicore/ops/axpy.hpp new file mode 100644 index 000000000..280d5ab60 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/baddbmm.hpp b/csrc/infinicore/include/infinicore/ops/baddbmm.hpp new file mode 100644 index 000000000..9988000bd --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/bilinear.hpp b/csrc/infinicore/include/infinicore/ops/bilinear.hpp new file mode 100644 index 000000000..805fa1efc --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/binary_cross_entropy_with_logits.hpp b/csrc/infinicore/include/infinicore/ops/binary_cross_entropy_with_logits.hpp new file mode 100644 index 000000000..4cdead724 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/bitwise_right_shift.hpp b/csrc/infinicore/include/infinicore/ops/bitwise_right_shift.hpp new file mode 100644 index 000000000..a1dc83a3d --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/bitwise_right_shift.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(BitwiseRightShift, Tensor, const Tensor &, const Tensor &); + +Tensor bitwise_right_shift(const Tensor &input, const Tensor &other); +void bitwise_right_shift_(Tensor out, const Tensor &input, const Tensor &other); + +} // namespace infinicore::op diff --git a/csrc/infinicore/include/infinicore/ops/blas_amax.hpp b/csrc/infinicore/include/infinicore/ops/blas_amax.hpp new file mode 100644 index 000000000..a6a571f95 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/blas_amin.hpp b/csrc/infinicore/include/infinicore/ops/blas_amin.hpp new file mode 100644 index 000000000..a2ed21c7b --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/blas_copy.hpp b/csrc/infinicore/include/infinicore/ops/blas_copy.hpp new file mode 100644 index 000000000..dd32646a0 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/blas_dot.hpp b/csrc/infinicore/include/infinicore/ops/blas_dot.hpp new file mode 100644 index 000000000..157c167f6 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/block_diag.hpp b/csrc/infinicore/include/infinicore/ops/block_diag.hpp new file mode 100644 index 000000000..89e227509 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/broadcast_to.hpp b/csrc/infinicore/include/infinicore/ops/broadcast_to.hpp new file mode 100644 index 000000000..94d98e065 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/cat.hpp b/csrc/infinicore/include/infinicore/ops/cat.hpp new file mode 100644 index 000000000..95be62103 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/causal_conv1d.hpp b/csrc/infinicore/include/infinicore/ops/causal_conv1d.hpp new file mode 100644 index 000000000..7123f6f05 --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/causal_conv1d.hpp @@ -0,0 +1,39 @@ +#pragma once + +#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); + +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); + +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/csrc/infinicore/include/infinicore/ops/causal_softmax.hpp b/csrc/infinicore/include/infinicore/ops/causal_softmax.hpp new file mode 100644 index 000000000..2646852af --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/cdist.hpp b/csrc/infinicore/include/infinicore/ops/cdist.hpp new file mode 100644 index 000000000..7bc4902a0 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/chunk_gated_delta_rule.hpp b/csrc/infinicore/include/infinicore/ops/chunk_gated_delta_rule.hpp new file mode 100644 index 000000000..779dff38b --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/chunk_gated_delta_rule.hpp @@ -0,0 +1,51 @@ +#pragma once + +#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); + +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); + +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/csrc/infinicore/include/infinicore/ops/common/cache.hpp b/csrc/infinicore/include/infinicore/ops/common/cache.hpp new file mode 100644 index 000000000..103f57411 --- /dev/null +++ b/csrc/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.type(), device.index()); + } + + 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::kCount)> caches_ = {}; +}; + +} // namespace infinicore::op::common diff --git a/csrc/infinicore/include/infinicore/ops/common/dispatcher.hpp b/csrc/infinicore/include/infinicore/ops/common/dispatcher.hpp new file mode 100644 index 000000000..6ab1f5ced --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/common/dispatcher.hpp @@ -0,0 +1,88 @@ +#pragma once + +#include "../../device.hpp" + +#ifdef ENABLE_MUTUAL_AWARENESS +#include "../../analyzer/optimization_intent.hpp" +#endif + +#include +#include +#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::kCount); ++device_type) { + registerDevice((Device::Type)device_type, fn, override_existing); + } + } + + Fn lookup(Device::Type device_type) const { + const auto fn = table_.at(static_cast(device_type)); + if (fn == nullptr) { + throw std::runtime_error( + "No operator implementation is registered for device `" + + std::string(Device::StringFromType(device_type)) + "`."); + } + return fn; + } + +#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::kCount) * kGoalCount> + goal_table_{}; +#endif + +private: + std::array(Device::Type::kCount)> table_{}; +}; +} // namespace infinicore::op::common diff --git a/csrc/infinicore/include/infinicore/ops/common/op.hpp b/csrc/infinicore/include/infinicore/ops/common/op.hpp new file mode 100644 index 000000000..b4f0434ef --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/conv2d.hpp b/csrc/infinicore/include/infinicore/ops/conv2d.hpp new file mode 100644 index 000000000..f1dda90ac --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/cross_entropy.hpp b/csrc/infinicore/include/infinicore/ops/cross_entropy.hpp new file mode 100644 index 000000000..9a6d446d2 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/deepseek_moe.hpp b/csrc/infinicore/include/infinicore/ops/deepseek_moe.hpp new file mode 100644 index 000000000..e8e7ffef9 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/diff.hpp b/csrc/infinicore/include/infinicore/ops/diff.hpp new file mode 100644 index 000000000..1def7a61a --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/digamma.hpp b/csrc/infinicore/include/infinicore/ops/digamma.hpp new file mode 100644 index 000000000..b61fc40f6 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/dist.hpp b/csrc/infinicore/include/infinicore/ops/dist.hpp new file mode 100644 index 000000000..459a52202 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/distributed/allgather.hpp b/csrc/infinicore/include/infinicore/ops/distributed/allgather.hpp new file mode 100644 index 000000000..eabb40ae2 --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/distributed/allgather.hpp @@ -0,0 +1,35 @@ +#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_; +}; + +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/csrc/infinicore/include/infinicore/ops/distributed/allreduce.hpp b/csrc/infinicore/include/infinicore/ops/distributed/allreduce.hpp new file mode 100644 index 000000000..7b071563a --- /dev/null +++ b/csrc/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, infinicclRedOp_t op, infinicclComm_t communicator); + ~AllReduce(); + void run() const override; + static void execute(Tensor output, const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +Tensor allreduce(const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator); +void allreduce_(Tensor output, const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator); + +} // namespace infinicore::op::distributed diff --git a/csrc/infinicore/include/infinicore/ops/distributed/reduce_scatter.hpp b/csrc/infinicore/include/infinicore/ops/distributed/reduce_scatter.hpp new file mode 100644 index 000000000..22ed6c01a --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/distributed/reduce_scatter.hpp @@ -0,0 +1,46 @@ +#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, infinicclRedOp_t op, infinicclComm_t communicator); + ~ReduceScatter(); + void run() const override; + static void execute(Tensor output, const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator); + +private: + void *planned_meta_; +}; + +Tensor reduce_scatter(const Tensor &input, size_t world_size, infinicclRedOp_t op, infinicclComm_t communicator); +void reduce_scatter_(Tensor output, const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator); +Tensor reduce_scatterv(const Tensor &input, + const std::vector &split_sizes, + size_t rank, + infinicclRedOp_t op, + infinicclComm_t communicator); +void reduce_scatterv_(Tensor output, + const Tensor &input, + const std::vector &split_sizes, + infinicclRedOp_t op, + infinicclComm_t communicator); +std::vector reduce_scatterv_many(const std::vector &inputs, + const std::vector &split_sizes, + size_t rank, + infinicclRedOp_t op, + infinicclComm_t communicator); +void reduce_scatterv_many_(const std::vector &outputs, + const std::vector &inputs, + const std::vector &split_sizes, + infinicclRedOp_t op, + infinicclComm_t communicator); + +} // namespace infinicore::op::distributed diff --git a/csrc/infinicore/include/infinicore/ops/embedding.hpp b/csrc/infinicore/include/infinicore/ops/embedding.hpp new file mode 100644 index 000000000..43f18d090 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/flash_attention.hpp b/csrc/infinicore/include/infinicore/ops/flash_attention.hpp new file mode 100644 index 000000000..24e33cfb6 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/flipud.hpp b/csrc/infinicore/include/infinicore/ops/flipud.hpp new file mode 100644 index 000000000..7e449f9f5 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/float_power.hpp b/csrc/infinicore/include/infinicore/ops/float_power.hpp new file mode 100644 index 000000000..c461db2eb --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/floor.hpp b/csrc/infinicore/include/infinicore/ops/floor.hpp new file mode 100644 index 000000000..11b52571b --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/floor_divide.hpp b/csrc/infinicore/include/infinicore/ops/floor_divide.hpp new file mode 100644 index 000000000..43267dce6 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/fmin.hpp b/csrc/infinicore/include/infinicore/ops/fmin.hpp new file mode 100644 index 000000000..4ea00787f --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/fmod.hpp b/csrc/infinicore/include/infinicore/ops/fmod.hpp new file mode 100644 index 000000000..87b90d515 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/fused_gated_delta_net_gating.hpp b/csrc/infinicore/include/infinicore/ops/fused_gated_delta_net_gating.hpp new file mode 100644 index 000000000..a61fd9937 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/gaussian_nll_loss.hpp b/csrc/infinicore/include/infinicore/ops/gaussian_nll_loss.hpp new file mode 100644 index 000000000..1dc68862f --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/gaussian_nll_loss.hpp @@ -0,0 +1,26 @@ +#pragma once + +#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); + +Tensor gaussian_nll_loss(const Tensor &input, + const Tensor &target, + const Tensor &var, + bool full = false, + double eps = 1e-6, + int reduction = 1); + +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/csrc/infinicore/include/infinicore/ops/gelu.hpp b/csrc/infinicore/include/infinicore/ops/gelu.hpp new file mode 100644 index 000000000..5e8c61347 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/gelutanh.hpp b/csrc/infinicore/include/infinicore/ops/gelutanh.hpp new file mode 100644 index 000000000..c968c81fd --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/gemm.hpp b/csrc/infinicore/include/infinicore/ops/gemm.hpp new file mode 100644 index 000000000..4f76cee26 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/hardswish.hpp b/csrc/infinicore/include/infinicore/ops/hardswish.hpp new file mode 100644 index 000000000..15313f461 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/hardtanh.hpp b/csrc/infinicore/include/infinicore/ops/hardtanh.hpp new file mode 100644 index 000000000..511408fee --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/hinge_embedding_loss.hpp b/csrc/infinicore/include/infinicore/ops/hinge_embedding_loss.hpp new file mode 100644 index 000000000..9d0fcfbab --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/huber_loss.hpp b/csrc/infinicore/include/infinicore/ops/huber_loss.hpp new file mode 100644 index 000000000..c6212f0c1 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/hypot.hpp b/csrc/infinicore/include/infinicore/ops/hypot.hpp new file mode 100644 index 000000000..24eebd44f --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/index_add.hpp b/csrc/infinicore/include/infinicore/ops/index_add.hpp new file mode 100644 index 000000000..e880114e9 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/index_copy.hpp b/csrc/infinicore/include/infinicore/ops/index_copy.hpp new file mode 100644 index 000000000..b2437dead --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/inner.hpp b/csrc/infinicore/include/infinicore/ops/inner.hpp new file mode 100644 index 000000000..ae372c415 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/interpolate.hpp b/csrc/infinicore/include/infinicore/ops/interpolate.hpp new file mode 100644 index 000000000..71bb675f7 --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/interpolate.hpp @@ -0,0 +1,28 @@ +#pragma once + +#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); + +Tensor interpolate(const Tensor &input, + std::string mode, + std::vector size, + std::vector scale_factor, + int align_corners); + +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/csrc/infinicore/include/infinicore/ops/kron.hpp b/csrc/infinicore/include/infinicore/ops/kron.hpp new file mode 100644 index 000000000..3a6834424 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/kthvalue.hpp b/csrc/infinicore/include/infinicore/ops/kthvalue.hpp new file mode 100644 index 000000000..1e9721d28 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/kv_caching.hpp b/csrc/infinicore/include/infinicore/ops/kv_caching.hpp new file mode 100644 index 000000000..3a70c2824 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/layer_norm.hpp b/csrc/infinicore/include/infinicore/ops/layer_norm.hpp new file mode 100644 index 000000000..da6256b51 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/ldexp.hpp b/csrc/infinicore/include/infinicore/ops/ldexp.hpp new file mode 100644 index 000000000..bac7964c2 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/lerp.hpp b/csrc/infinicore/include/infinicore/ops/lerp.hpp new file mode 100644 index 000000000..bf3e5cbec --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/linear.hpp b/csrc/infinicore/include/infinicore/ops/linear.hpp new file mode 100644 index 000000000..96ca733e4 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/linear_w8a8i8.hpp b/csrc/infinicore/include/infinicore/ops/linear_w8a8i8.hpp new file mode 100644 index 000000000..08cadc111 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/log_softmax.hpp b/csrc/infinicore/include/infinicore/ops/log_softmax.hpp new file mode 100644 index 000000000..00e1c90c1 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/logaddexp.hpp b/csrc/infinicore/include/infinicore/ops/logaddexp.hpp new file mode 100644 index 000000000..fc84c10c2 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/logaddexp2.hpp b/csrc/infinicore/include/infinicore/ops/logaddexp2.hpp new file mode 100644 index 000000000..7c7dbe96e --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/logcumsumexp.hpp b/csrc/infinicore/include/infinicore/ops/logcumsumexp.hpp new file mode 100644 index 000000000..212ff625d --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/logdet.hpp b/csrc/infinicore/include/infinicore/ops/logdet.hpp new file mode 100644 index 000000000..69c1b6d0d --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/mamba_selective_scan.hpp b/csrc/infinicore/include/infinicore/ops/mamba_selective_scan.hpp new file mode 100644 index 000000000..3acdb1c4d --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/masked_select.hpp b/csrc/infinicore/include/infinicore/ops/masked_select.hpp new file mode 100644 index 000000000..d003e3bb8 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/matmul.hpp b/csrc/infinicore/include/infinicore/ops/matmul.hpp new file mode 100644 index 000000000..1fb41774e --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/mha.hpp b/csrc/infinicore/include/infinicore/ops/mha.hpp new file mode 100644 index 000000000..f62598c78 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/mha_kvcache.hpp b/csrc/infinicore/include/infinicore/ops/mha_kvcache.hpp new file mode 100644 index 000000000..69a11f61c --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/mha_varlen.hpp b/csrc/infinicore/include/infinicore/ops/mha_varlen.hpp new file mode 100644 index 000000000..4d226f383 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/moe_align.hpp b/csrc/infinicore/include/infinicore/ops/moe_align.hpp new file mode 100644 index 000000000..2ce7ca8f4 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp b/csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp new file mode 100644 index 000000000..0a34c734c --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/moe_fused_gate.hpp b/csrc/infinicore/include/infinicore/ops/moe_fused_gate.hpp new file mode 100644 index 000000000..e20d0a657 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/moe_sum.hpp b/csrc/infinicore/include/infinicore/ops/moe_sum.hpp new file mode 100644 index 000000000..e9623e3b1 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp b/csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp new file mode 100644 index 000000000..e796c1dc8 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/moe_topk_softmax.hpp b/csrc/infinicore/include/infinicore/ops/moe_topk_softmax.hpp new file mode 100644 index 000000000..e6964aa88 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/mrope.hpp b/csrc/infinicore/include/infinicore/ops/mrope.hpp new file mode 100644 index 000000000..3ee2c1ef1 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/mul.hpp b/csrc/infinicore/include/infinicore/ops/mul.hpp new file mode 100644 index 000000000..2eb480ddb --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/mul_scalar.hpp b/csrc/infinicore/include/infinicore/ops/mul_scalar.hpp new file mode 100644 index 000000000..5ee5e0d0b --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/multi_margin_loss.hpp b/csrc/infinicore/include/infinicore/ops/multi_margin_loss.hpp new file mode 100644 index 000000000..e36cf469b --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/nrm2.hpp b/csrc/infinicore/include/infinicore/ops/nrm2.hpp new file mode 100644 index 000000000..c5552a4c1 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/nsa_compress_paged_cache.hpp b/csrc/infinicore/include/infinicore/ops/nsa_compress_paged_cache.hpp new file mode 100644 index 000000000..bb5380632 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/nsa_paged_attention.hpp b/csrc/infinicore/include/infinicore/ops/nsa_paged_attention.hpp new file mode 100644 index 000000000..00e9adc10 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/ones.hpp b/csrc/infinicore/include/infinicore/ops/ones.hpp new file mode 100644 index 000000000..0ea747484 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/pad.hpp b/csrc/infinicore/include/infinicore/ops/pad.hpp new file mode 100644 index 000000000..e598a94b4 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/paged_attention.hpp b/csrc/infinicore/include/infinicore/ops/paged_attention.hpp new file mode 100644 index 000000000..8c906c95e --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/paged_attention_prefill.hpp b/csrc/infinicore/include/infinicore/ops/paged_attention_prefill.hpp new file mode 100644 index 000000000..952924528 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/paged_caching.hpp b/csrc/infinicore/include/infinicore/ops/paged_caching.hpp new file mode 100644 index 000000000..403b4b738 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/per_channel_quant_i8.hpp b/csrc/infinicore/include/infinicore/ops/per_channel_quant_i8.hpp new file mode 100644 index 000000000..0b0296248 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/per_tensor_dequant_i8.hpp b/csrc/infinicore/include/infinicore/ops/per_tensor_dequant_i8.hpp new file mode 100644 index 000000000..a6c4b1381 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/per_tensor_quant_i8.hpp b/csrc/infinicore/include/infinicore/ops/per_tensor_quant_i8.hpp new file mode 100644 index 000000000..84ec99c37 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/prelu.hpp b/csrc/infinicore/include/infinicore/ops/prelu.hpp new file mode 100644 index 000000000..5e9026af8 --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/prelu.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(Prelu, Tensor, const Tensor &, const Tensor &); + +Tensor prelu(const Tensor &input, const Tensor &weight); +void prelu_(Tensor out, const Tensor &input, const Tensor &weight); + +} // namespace infinicore::op diff --git a/csrc/infinicore/include/infinicore/ops/prepare_moe_input.hpp b/csrc/infinicore/include/infinicore/ops/prepare_moe_input.hpp new file mode 100644 index 000000000..c06951320 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/quickgelu.hpp b/csrc/infinicore/include/infinicore/ops/quickgelu.hpp new file mode 100644 index 000000000..b67d7f71a --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/random_sample.hpp b/csrc/infinicore/include/infinicore/ops/random_sample.hpp new file mode 100644 index 000000000..a3a9d29f1 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/rearrange.hpp b/csrc/infinicore/include/infinicore/ops/rearrange.hpp new file mode 100644 index 000000000..5db983ef8 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/reciprocal.hpp b/csrc/infinicore/include/infinicore/ops/reciprocal.hpp new file mode 100644 index 000000000..172fcdf36 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/recurrent_gated_delta_rule.hpp b/csrc/infinicore/include/infinicore/ops/recurrent_gated_delta_rule.hpp new file mode 100644 index 000000000..66a95d251 --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/recurrent_gated_delta_rule.hpp @@ -0,0 +1,53 @@ +#pragma once + +#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); + +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); + +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); + +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/csrc/infinicore/include/infinicore/ops/relu.hpp b/csrc/infinicore/include/infinicore/ops/relu.hpp new file mode 100644 index 000000000..59f2b26f5 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/relu6.hpp b/csrc/infinicore/include/infinicore/ops/relu6.hpp new file mode 100644 index 000000000..5090a796c --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/relu6.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(Relu6, Tensor, const Tensor &); + +Tensor relu6(const Tensor &input); +void relu6_(Tensor out, const Tensor &input); + +} // namespace infinicore::op diff --git a/csrc/infinicore/include/infinicore/ops/rms_norm.hpp b/csrc/infinicore/include/infinicore/ops/rms_norm.hpp new file mode 100644 index 000000000..c7b2b2d72 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/rope.hpp b/csrc/infinicore/include/infinicore/ops/rope.hpp new file mode 100644 index 000000000..8fd630ce1 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/rot.hpp b/csrc/infinicore/include/infinicore/ops/rot.hpp new file mode 100644 index 000000000..ff473f33e --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/rotg.hpp b/csrc/infinicore/include/infinicore/ops/rotg.hpp new file mode 100644 index 000000000..c65f211dd --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/rotm.hpp b/csrc/infinicore/include/infinicore/ops/rotm.hpp new file mode 100644 index 000000000..813fc27ee --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/rotmg.hpp b/csrc/infinicore/include/infinicore/ops/rotmg.hpp new file mode 100644 index 000000000..e245840a4 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/rwkv5_wkv.hpp b/csrc/infinicore/include/infinicore/ops/rwkv5_wkv.hpp new file mode 100644 index 000000000..f1e6accd1 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/scal.hpp b/csrc/infinicore/include/infinicore/ops/scal.hpp new file mode 100644 index 000000000..d6cb5ce8c --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/scaled_mm_i8.hpp b/csrc/infinicore/include/infinicore/ops/scaled_mm_i8.hpp new file mode 100644 index 000000000..427ed9d65 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/scatter.hpp b/csrc/infinicore/include/infinicore/ops/scatter.hpp new file mode 100644 index 000000000..306bb1232 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/selu.hpp b/csrc/infinicore/include/infinicore/ops/selu.hpp new file mode 100644 index 000000000..2aed070ed --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/sigmoid.hpp b/csrc/infinicore/include/infinicore/ops/sigmoid.hpp new file mode 100644 index 000000000..429376002 --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/sigmoid.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(Sigmoid, Tensor, const Tensor &); + +Tensor sigmoid(const Tensor &input); +void sigmoid_(Tensor output, const Tensor &input); + +} // namespace infinicore::op diff --git a/csrc/infinicore/include/infinicore/ops/silu.hpp b/csrc/infinicore/include/infinicore/ops/silu.hpp new file mode 100644 index 000000000..d4cbac57d --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/silu_and_mul.hpp b/csrc/infinicore/include/infinicore/ops/silu_and_mul.hpp new file mode 100644 index 000000000..1b594111d --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/sinh.hpp b/csrc/infinicore/include/infinicore/ops/sinh.hpp new file mode 100644 index 000000000..31cb3c256 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/smooth_l1_loss.hpp b/csrc/infinicore/include/infinicore/ops/smooth_l1_loss.hpp new file mode 100644 index 000000000..9a67c6109 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/softmax.hpp b/csrc/infinicore/include/infinicore/ops/softmax.hpp new file mode 100644 index 000000000..c96ab1810 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/softplus.hpp b/csrc/infinicore/include/infinicore/ops/softplus.hpp new file mode 100644 index 000000000..b6d665455 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/softsign.hpp b/csrc/infinicore/include/infinicore/ops/softsign.hpp new file mode 100644 index 000000000..05e8e54fa --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/sum.hpp b/csrc/infinicore/include/infinicore/ops/sum.hpp new file mode 100644 index 000000000..0ead8de26 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/swap.hpp b/csrc/infinicore/include/infinicore/ops/swap.hpp new file mode 100644 index 000000000..aba3ad563 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/swiglu.hpp b/csrc/infinicore/include/infinicore/ops/swiglu.hpp new file mode 100644 index 000000000..7aa77e632 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/take.hpp b/csrc/infinicore/include/infinicore/ops/take.hpp new file mode 100644 index 000000000..5672005ad --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/tan.hpp b/csrc/infinicore/include/infinicore/ops/tan.hpp new file mode 100644 index 000000000..72410f866 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/tanhshrink.hpp b/csrc/infinicore/include/infinicore/ops/tanhshrink.hpp new file mode 100644 index 000000000..83673b22c --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/topk.hpp b/csrc/infinicore/include/infinicore/ops/topk.hpp new file mode 100644 index 000000000..d8486112c --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/topksoftmax.hpp b/csrc/infinicore/include/infinicore/ops/topksoftmax.hpp new file mode 100644 index 000000000..2cd3ac7c5 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/triplet_margin_loss.hpp b/csrc/infinicore/include/infinicore/ops/triplet_margin_loss.hpp new file mode 100644 index 000000000..43b935eb8 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/triplet_margin_with_distance_loss.hpp b/csrc/infinicore/include/infinicore/ops/triplet_margin_with_distance_loss.hpp new file mode 100644 index 000000000..d6ea449c5 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/unfold.hpp b/csrc/infinicore/include/infinicore/ops/unfold.hpp new file mode 100644 index 000000000..f8c7a3b13 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/upsample_bilinear.hpp b/csrc/infinicore/include/infinicore/ops/upsample_bilinear.hpp new file mode 100644 index 000000000..ba9aa2379 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/upsample_nearest.hpp b/csrc/infinicore/include/infinicore/ops/upsample_nearest.hpp new file mode 100644 index 000000000..188d1b923 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/vander.hpp b/csrc/infinicore/include/infinicore/ops/vander.hpp new file mode 100644 index 000000000..8198cfeb8 --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/var.hpp b/csrc/infinicore/include/infinicore/ops/var.hpp new file mode 100644 index 000000000..d1e01e1bf --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/ops/var_mean.hpp b/csrc/infinicore/include/infinicore/ops/var_mean.hpp new file mode 100644 index 000000000..a9679187c --- /dev/null +++ b/csrc/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/csrc/infinicore/include/infinicore/tensor.hpp b/csrc/infinicore/include/infinicore/tensor.hpp new file mode 100644 index 000000000..e6a5548a8 --- /dev/null +++ b/csrc/infinicore/include/infinicore/tensor.hpp @@ -0,0 +1,334 @@ +#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; + + TensorMetaData(const Shape &shape, const Strides &strides, const DataType &dtype); +}; + +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; + + infini::rt::TensorView view(); + + 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/csrc/infinicore/src/adaptor/aten_adaptor.cc b/csrc/infinicore/src/adaptor/aten_adaptor.cc new file mode 100644 index 000000000..87eb033f4 --- /dev/null +++ b/csrc/infinicore/src/adaptor/aten_adaptor.cc @@ -0,0 +1,56 @@ +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" + +namespace infinicore::adaptor { + +at::Tensor to_aten_tensor(const infinicore::Tensor &t) { + void *data_ptr = (void *)(t->data()); + + auto sizes = std::vector( + t->shape().begin(), + t->shape().end()); + + auto strides = t->strides(); + + auto dtype = to_at_dtype(t->dtype()); + auto device = to_at_device(t->device()); + + auto deleter_ = [](void * /*unused*/) mutable { + + }; + + at::TensorOptions options = at::TensorOptions() + .dtype(dtype) + .device(device) + .requires_grad(false); + + return at::from_blob( + data_ptr, + sizes, + strides, + deleter_, + options); +} + +#if defined(ENABLE_HYGON_API) +c10::hip::HIPStream get_hip_stream() { + return c10::hip::getStreamFromExternal( + hipStream_t(infinicore::context::getStream()), infinicore::context::getDevice().index()); +} +#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +c10::cuda::CUDAStream get_cuda_stream() { + return c10::cuda::getStreamFromExternal( + cudaStream_t(infinicore::context::getStream()), infinicore::context::getDevice().index()); +} +#endif + +#if defined(ENABLE_MOORE_API) +c10::musa::MUSAStream get_musa_stream() { + return c10::musa::getStreamFromExternal( + musaStream_t(infinicore::context::getStream()), infinicore::context::getDevice().index()); +} +#endif + +} // namespace infinicore::adaptor + +#endif // ENABLE_ATEN diff --git a/csrc/infinicore/src/adaptor/flash_attn/hygon/flash_attn_hygon.cc b/csrc/infinicore/src/adaptor/flash_attn/hygon/flash_attn_hygon.cc new file mode 100644 index 000000000..a26b400ba --- /dev/null +++ b/csrc/infinicore/src/adaptor/flash_attn/hygon/flash_attn_hygon.cc @@ -0,0 +1,173 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_FLASH_ATTN) + +#include "flash_attn_hygon.hpp" + +#include +#include +#include +#include +#include +#include +#include + +using mha_fwd_kvcache_fn_t = std::vector (*)( + at::Tensor &q, + const at::Tensor &kcache, + const at::Tensor &vcache, + std::optional &k_, + std::optional &v_, + std::optional &seqlens_k_, + std::optional &rotary_cos_, + std::optional &rotary_sin_, + std::optional &cache_batch_idx_, + std::optional &leftpad_k_, + std::optional &block_table_, + std::optional &alibi_slopes_, + std::optional &out_, + float softmax_scale, + bool is_causal, + int window_size_left, + int window_size_right, + float softcap, + bool is_rotary_interleaved, + int num_splits, + const std::optional &s_aux_); + +using mha_varlen_fwd_fn_t = std::vector (*)( + at::Tensor &q, + const at::Tensor &k, + const at::Tensor &v, + std::optional &out_, + const at::Tensor &cu_seqlens_q, + const at::Tensor &cu_seqlens_k, + std::optional &seqused_k, + std::optional &leftpad_k_, + std::optional &block_table_, + std::optional &alibi_slopes_, + int max_seqlen_q, + int max_seqlen_k, + float p_dropout, + float softmax_scale, + bool zero_tensors, + bool is_causal, + int window_size_left, + int window_size_right, + float softcap, + bool return_softmax, + std::optional q_descale_, + std::optional k_descale_, + std::optional v_descale_, + std::optional gen_, + const std::optional &s_aux_); + +namespace { + +void *resolve_flash_symbol(const char *name) { + void *sym = dlsym(RTLD_DEFAULT, name); + if (!sym) { + throw std::runtime_error(std::string("flash_attn symbol not found: ") + name); + } + return sym; +} + +void *resolve_flash_extension_symbol(const char *name) { + static std::string last_error; + static void *handle = [&]() -> void * { + const char *candidates[] = { + "flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so", + "/usr/local/lib/python3.10/dist-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so", + }; + for (const char *candidate : candidates) { + dlerror(); + if (void *h = dlopen(candidate, RTLD_NOW | RTLD_GLOBAL)) { + return h; + } + if (const char *err = dlerror()) { + last_error += std::string(candidate) + ": " + err + "; "; + } + } + return nullptr; + }(); + if (handle) { + dlerror(); + if (void *sym = dlsym(handle, name)) { + return sym; + } + if (const char *err = dlerror()) { + throw std::runtime_error(std::string("flash_attn extension symbol not found: ") + name + "; " + err); + } + throw std::runtime_error(std::string("flash_attn extension symbol not found: ") + name); + } + throw std::runtime_error(std::string("flash_attn extension dlopen failed for symbol: ") + name + "; " + last_error); +} + +} // namespace + +namespace flash { + +std::vector +mha_fwd_kvcache(at::Tensor &q, + const at::Tensor &kcache, + const at::Tensor &vcache, + std::optional &k_, + std::optional &v_, + std::optional &seqlens_k_, + std::optional &rotary_cos_, + std::optional &rotary_sin_, + std::optional &cache_batch_idx_, + std::optional &leftpad_k_, + std::optional &block_table_, + std::optional &alibi_slopes_, + std::optional &out_, + float softmax_scale, + bool is_causal, + int window_size_left, + int window_size_right, + float softcap, + bool is_rotary_interleaved, + int num_splits) { + static auto fn = reinterpret_cast(resolve_flash_extension_symbol("mha_fwd_kvcache")); + std::optional s_aux = std::nullopt; + return fn(q, kcache, vcache, k_, v_, seqlens_k_, rotary_cos_, rotary_sin_, + cache_batch_idx_, leftpad_k_, block_table_, alibi_slopes_, out_, + softmax_scale, is_causal, window_size_left, window_size_right, + softcap, is_rotary_interleaved, num_splits, s_aux); +} + +std::vector +vllm_mha_varlen_fwd(at::Tensor &q, + const at::Tensor &k, + const at::Tensor &v, + std::optional &out_, + const at::Tensor &cu_seqlens_q, + const at::Tensor &cu_seqlens_k, + std::optional &seqused_k, + std::optional &leftpad_k_, + std::optional &block_table_, + std::optional &alibi_slopes_, + int max_seqlen_q, + int max_seqlen_k, + float p_dropout, + float softmax_scale, + bool zero_tensors, + bool is_causal, + int window_size_left, + int window_size_right, + float softcap, + bool return_softmax, + std::optional gen_) { + static auto fn = reinterpret_cast(resolve_flash_extension_symbol("mha_varlen_fwd")); + std::optional q_descale = std::nullopt; + std::optional k_descale = std::nullopt; + std::optional v_descale = std::nullopt; + std::optional s_aux = std::nullopt; + return fn(q, k, v, out_, cu_seqlens_q, cu_seqlens_k, seqused_k, leftpad_k_, + block_table_, alibi_slopes_, max_seqlen_q, max_seqlen_k, p_dropout, + softmax_scale, zero_tensors, is_causal, window_size_left, + window_size_right, softcap, return_softmax, q_descale, k_descale, + v_descale, gen_, s_aux); +} + +} // namespace flash + +#endif diff --git a/csrc/infinicore/src/adaptor/flash_attn/hygon/flash_attn_hygon.hpp b/csrc/infinicore/src/adaptor/flash_attn/hygon/flash_attn_hygon.hpp new file mode 100644 index 000000000..8600c8711 --- /dev/null +++ b/csrc/infinicore/src/adaptor/flash_attn/hygon/flash_attn_hygon.hpp @@ -0,0 +1,60 @@ +#ifndef INFINICORE_ADAPTOR_FLASH_ATTN_HYGON_FLASH_ATTN_HYGON_HPP +#define INFINICORE_ADAPTOR_FLASH_ATTN_HYGON_FLASH_ATTN_HYGON_HPP + +#if defined(ENABLE_FLASH_ATTN) && defined(ENABLE_HYGON_API) + +#include +#include +#include + +namespace flash { + +std::vector +mha_fwd_kvcache(at::Tensor &q, + const at::Tensor &kcache, + const at::Tensor &vcache, + std::optional &k_, + std::optional &v_, + std::optional &seqlens_k_, + std::optional &rotary_cos_, + std::optional &rotary_sin_, + std::optional &cache_batch_idx_, + std::optional &leftpad_k_, + std::optional &block_table_, + std::optional &alibi_slopes_, + std::optional &out_, + float softmax_scale, + bool is_causal, + int window_size_left, + int window_size_right, + float softcap, + bool is_rotary_interleaved, + int num_splits); + +std::vector +vllm_mha_varlen_fwd(at::Tensor &q, + const at::Tensor &k, + const at::Tensor &v, + std::optional &out_, + const at::Tensor &cu_seqlens_q, + const at::Tensor &cu_seqlens_k, + std::optional &seqused_k, + std::optional &leftpad_k_, + std::optional &block_table_, + std::optional &alibi_slopes_, + int max_seqlen_q, + int max_seqlen_k, + float p_dropout, + float softmax_scale, + bool zero_tensors, + bool is_causal, + int window_size_left, + int window_size_right, + float softcap, + bool return_softmax, + std::optional gen_); + +} // namespace flash + +#endif // ENABLE_FLASH_ATTN && ENABLE_HYGON_API +#endif // INFINICORE_ADAPTOR_FLASH_ATTN_HYGON_FLASH_ATTN_HYGON_HPP diff --git a/csrc/infinicore/src/analyzer/mutual_awareness_analyzer.cc b/csrc/infinicore/src/analyzer/mutual_awareness_analyzer.cc new file mode 100644 index 000000000..1ea602f1b --- /dev/null +++ b/csrc/infinicore/src/analyzer/mutual_awareness_analyzer.cc @@ -0,0 +1,200 @@ +#include "infinicore/analyzer/mutual_awareness_analyzer.hpp" + +namespace infinicore::analyzer { + +namespace { + +std::vector collectRuntimeResourceSnapshots() { + std::vector device_snapshots; + // TODO: integrate with ContextImpl allocator stats when available. + // For now, return an empty snapshot list. + return device_snapshots; +} + +} // namespace + +// ============================================================ +// Singleton +// ============================================================ + +MutualAwarenessAnalyzer &MutualAwarenessAnalyzer::instance() { + static MutualAwarenessAnalyzer inst; + return inst; +} + +MutualAwarenessAnalyzer::MutualAwarenessAnalyzer() + : phase_detector_(), + resource_sensor_(), + intent_generator_(), + enabled_(true), + graph_intent_cached_(false) { +} + +// ============================================================ +// Main analysis entry points +// ============================================================ + +OptimizationIntent MutualAwarenessAnalyzer::analyze() { + if (!enabled_) { + return OptimizationIntent{}; + } + + // If we have a cached graph intent, return it + if (graph_intent_cached_) { + return graph_cached_intent_; + } + + // Get recent op trace window + auto &trace = getGlobalOpTrace(); + auto window = trace.getRecentEntries(phase_detector_.config().window_size); + + // Detect phase + PhaseType phase = phase_detector_.detect(window); + + auto device_snapshots = collectRuntimeResourceSnapshots(); + std::vector device_intents; + device_intents.reserve(device_snapshots.size()); + for (auto const &snapshot : device_snapshots) { + device_intents.push_back(resource_sensor_.sense(snapshot)); + } + + // Generate intent + auto intent = intent_generator_.generate(phase, window, device_intents); + + // Cache result + { + std::lock_guard lock(mutex_); + last_intent_ = intent; + } + + return intent; +} + +OptimizationIntent MutualAwarenessAnalyzer::analyze( + const std::vector> &device_stats) { + + if (!enabled_) { + return OptimizationIntent{}; + } + + // If we have a cached graph intent, return it + if (graph_intent_cached_) { + return graph_cached_intent_; + } + + // Get recent op trace window + auto &trace = getGlobalOpTrace(); + auto window = trace.getRecentEntries(phase_detector_.config().window_size); + + // Detect phase + PhaseType phase = phase_detector_.detect(window); + + // Build per-device intents from provided stats + std::vector device_intents; + device_intents.reserve(device_stats.size()); + for (auto &[dev_id, stats] : device_stats) { + device_intents.push_back(resource_sensor_.sense(dev_id, stats)); + } + + // Generate intent + auto intent = intent_generator_.generate(phase, window, device_intents); + + // Cache result + { + std::lock_guard lock(mutex_); + last_intent_ = intent; + } + + return intent; +} + +OptimizationIntent MutualAwarenessAnalyzer::analyze( + const std::vector &device_snapshots) { + + if (!enabled_) { + return OptimizationIntent{}; + } + + if (graph_intent_cached_) { + return graph_cached_intent_; + } + + auto &trace = getGlobalOpTrace(); + auto window = trace.getRecentEntries(phase_detector_.config().window_size); + + PhaseType phase = phase_detector_.detect(window); + + std::vector device_intents; + device_intents.reserve(device_snapshots.size()); + for (auto const &snapshot : device_snapshots) { + device_intents.push_back(resource_sensor_.sense(snapshot)); + } + + auto intent = intent_generator_.generate(phase, window, device_intents); + + { + std::lock_guard lock(mutex_); + last_intent_ = intent; + } + + return intent; +} + +PhaseType MutualAwarenessAnalyzer::getCurrentPhase() const { + if (!enabled_) { + return PhaseType::UNKNOWN; + } + + return phase_detector_.detectFromTrace(getGlobalOpTrace()); +} + +OptimizationGoal MutualAwarenessAnalyzer::getCurrentOptimizationGoal() const { + return const_cast(this)->analyze().global.goal; +} + +const OptimizationIntent &MutualAwarenessAnalyzer::lastIntent() const { + return last_intent_; +} + +// ============================================================ +// Graph recording support +// ============================================================ + +void MutualAwarenessAnalyzer::onGraphRecordingStop() { + if (!enabled_) { + return; + } + + // Analyze the op sequence recorded during graph capture + // and cache the result. Graph ops are static, so we only + // need to analyze once. + graph_cached_intent_ = analyze(); + graph_intent_cached_ = true; +} + +void MutualAwarenessAnalyzer::clearGraphCache() { + graph_intent_cached_ = false; + graph_cached_intent_ = OptimizationIntent{}; +} + +// ============================================================ +// C-style API for external framework integration +// ============================================================ + +OptimizationIntent analyzeCurrentState() { + return MutualAwarenessAnalyzer::instance().analyze(); +} + +PhaseType getCurrentPhase() { + return MutualAwarenessAnalyzer::instance().getCurrentPhase(); +} + +OptimizationGoal getCurrentOptimizationGoal() { + return MutualAwarenessAnalyzer::instance().getCurrentOptimizationGoal(); +} + +void setAnalyzerEnabled(bool enabled) { + MutualAwarenessAnalyzer::instance().setEnabled(enabled); +} + +} // namespace infinicore::analyzer diff --git a/csrc/infinicore/src/analyzer/op_trace.cc b/csrc/infinicore/src/analyzer/op_trace.cc new file mode 100644 index 000000000..566a83d09 --- /dev/null +++ b/csrc/infinicore/src/analyzer/op_trace.cc @@ -0,0 +1,10 @@ +#include "infinicore/analyzer/op_trace.hpp" + +namespace infinicore::analyzer { + +OpTraceRing &getGlobalOpTrace() { + static OpTraceRing instance(OpTraceRing::DEFAULT_CAPACITY); + return instance; +} + +} // namespace infinicore::analyzer diff --git a/csrc/infinicore/src/context/allocators/device_pinned_allocator.cc b/csrc/infinicore/src/context/allocators/device_pinned_allocator.cc new file mode 100644 index 000000000..53827bf9f --- /dev/null +++ b/csrc/infinicore/src/context/allocators/device_pinned_allocator.cc @@ -0,0 +1,55 @@ +#include "device_pinned_allocator.hpp" + +#include "../../utils.hpp" + +namespace infinicore { +DevicePinnedHostAllocator::DevicePinnedHostAllocator(Device device) : MemoryAllocator(), owner_(device) {} + +DevicePinnedHostAllocator::~DevicePinnedHostAllocator() noexcept { + try { + gc(); + } catch (const std::exception &error) { + try { + spdlog::warn("freeing pinned host memory failed during allocator cleanup: {}", error.what()); + } catch (...) { + } + } catch (...) { + try { + spdlog::warn("freeing pinned host memory failed during allocator cleanup"); + } catch (...) { + } + } +} + +std::byte *DevicePinnedHostAllocator::allocate(size_t size) { + if (size == 0) { + return nullptr; + } + void *ptr; + INFINICORE_CHECK_ERROR(infini::rt::runtime::MallocHost(&ptr, size)); + return (std::byte *)ptr; +} + +void DevicePinnedHostAllocator::deallocate(std::byte *ptr) { + if (ptr == nullptr) { + return; + } + if (owner_ == context::getDevice()) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::FreeHost(ptr)); + gc(); + } else { + std::lock_guard lock{gc_mutex_}; + gc_queue_.push(ptr); + } +} + +void DevicePinnedHostAllocator::gc() { + std::lock_guard lock{gc_mutex_}; + while (gc_queue_.empty() == false) { + std::byte *p = gc_queue_.front(); + INFINICORE_CHECK_ERROR(infini::rt::runtime::FreeHost(p)); + gc_queue_.pop(); + } +} + +} // namespace infinicore diff --git a/csrc/infinicore/src/context/allocators/device_pinned_allocator.hpp b/csrc/infinicore/src/context/allocators/device_pinned_allocator.hpp new file mode 100644 index 000000000..1458d8758 --- /dev/null +++ b/csrc/infinicore/src/context/allocators/device_pinned_allocator.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include "memory_allocator.hpp" + +#include "../context_impl.hpp" + +#include +#include + +namespace infinicore { +class DevicePinnedHostAllocator : public MemoryAllocator { +public: + explicit DevicePinnedHostAllocator(Device device); + ~DevicePinnedHostAllocator() noexcept; + + std::byte *allocate(size_t size) override; + void deallocate(std::byte *ptr) override; + + void gc(); + +private: + Device owner_; + + std::mutex gc_mutex_; + std::queue gc_queue_; +}; + +} // namespace infinicore diff --git a/csrc/infinicore/src/context/allocators/host_allocator.cc b/csrc/infinicore/src/context/allocators/host_allocator.cc new file mode 100644 index 000000000..0f1cc7072 --- /dev/null +++ b/csrc/infinicore/src/context/allocators/host_allocator.cc @@ -0,0 +1,18 @@ +#include "host_allocator.hpp" + +namespace infinicore { +std::byte *HostAllocator::allocate(size_t size) { + if (size == 0) { + return nullptr; + } + return (std::byte *)std::malloc(size); +} + +void HostAllocator::deallocate(std::byte *ptr) { + if (ptr == nullptr) { + return; + } + std::free(ptr); +} + +} // namespace infinicore diff --git a/csrc/infinicore/src/context/allocators/host_allocator.hpp b/csrc/infinicore/src/context/allocators/host_allocator.hpp new file mode 100644 index 000000000..57cc00898 --- /dev/null +++ b/csrc/infinicore/src/context/allocators/host_allocator.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "memory_allocator.hpp" + +namespace infinicore { +class HostAllocator : public MemoryAllocator { +public: + HostAllocator() = default; + ~HostAllocator() = default; + + std::byte *allocate(size_t size) override; + void deallocate(std::byte *ptr) override; +}; + +} // namespace infinicore diff --git a/csrc/infinicore/src/context/allocators/memory_allocator.hpp b/csrc/infinicore/src/context/allocators/memory_allocator.hpp new file mode 100644 index 000000000..7fce31460 --- /dev/null +++ b/csrc/infinicore/src/context/allocators/memory_allocator.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "infinicore/memory.hpp" + +#include + +namespace infinicore { +class MemoryAllocator { +public: + virtual ~MemoryAllocator() = default; + + virtual std::byte *allocate(size_t size) = 0; + virtual void deallocate(std::byte *ptr) = 0; +}; +} // namespace infinicore diff --git a/csrc/infinicore/src/context/allocators/pinnable_block_allocator.cc b/csrc/infinicore/src/context/allocators/pinnable_block_allocator.cc new file mode 100644 index 000000000..f036fd910 --- /dev/null +++ b/csrc/infinicore/src/context/allocators/pinnable_block_allocator.cc @@ -0,0 +1,310 @@ +#include "pinnable_block_allocator.hpp" + +#include "../context_impl.hpp" + +#include "../../utils.hpp" + +#include +#include +#include + +namespace infinicore { + +PinnableBlockAllocator::PinLease::PinLease( + PinnableBlockAllocator *allocator, + std::vector> blocks) noexcept + : allocator_(allocator), blocks_(std::move(blocks)) {} + +PinnableBlockAllocator::PinLease::~PinLease() noexcept { + allocator_->release_frozen_blocks_(blocks_); +} + +// ------------------- Helper functions ------------------- + +// Round up size to nearest multiple of alignment +inline size_t align_up(size_t size, size_t alignment) { + return (size + alignment - 1) / alignment * alignment; +} + +// ------------------- Constructor ------------------- +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 + }; +} + +void PinnableBlockAllocator::begin_pin_mode() { + std::lock_guard lock(mutex_); + if (pinned_mode_) { + throw std::runtime_error("allocator graph capture is already active"); + } + INFINICORE_ASSERT(capture_frozen_blocks_.empty()); + INFINICORE_ASSERT(capture_frozen_block_set_.empty()); + pinned_mode_ = true; + pin_owner_ = std::this_thread::get_id(); +} + +std::shared_ptr +PinnableBlockAllocator::commit_pin_mode() { + std::vector> blocks; + { + std::lock_guard lock(mutex_); + if (!pinned_mode_) { + return nullptr; + } + if (pin_owner_ != std::this_thread::get_id()) { + throw std::runtime_error("cannot commit allocator pin mode: another thread owns the graph capture"); + } + pinned_mode_ = false; + pin_owner_ = {}; + blocks = std::move(capture_frozen_blocks_); + capture_frozen_block_set_.clear(); + } + + try { + return std::make_shared(this, std::move(blocks)); + } catch (...) { + release_frozen_blocks_(blocks); + throw; + } +} + +void PinnableBlockAllocator::cancel_pin_mode() { + std::lock_guard lock(mutex_); + const auto caller = std::this_thread::get_id(); + if (pinned_mode_ && pin_owner_ != caller) { + throw std::runtime_error("cannot cancel allocator pin mode: another thread owns the graph capture"); + } + for (const auto &block : capture_frozen_blocks_) { + if (block->pin_count > 0) { + --block->pin_count; + } + } + capture_frozen_blocks_.clear(); + capture_frozen_block_set_.clear(); + pinned_mode_ = false; + pin_owner_ = {}; +} + +void PinnableBlockAllocator::freeze_for_capture_(const std::shared_ptr &block) { + if (pinned_mode_ && pin_owner_ == std::this_thread::get_id() + && capture_frozen_block_set_.insert(block.get()).second) { + ++block->pin_count; + capture_frozen_blocks_.push_back(block); + } +} + +void PinnableBlockAllocator::retain_for_capture(void *ptr) { + std::lock_guard lock(mutex_); + if (!pinned_mode_ || pin_owner_ != std::this_thread::get_id()) { + throw std::runtime_error("allocator graph capture is not active on this thread"); + } + + const auto it = all_blocks_.find(ptr); + if (it == all_blocks_.end()) { + throw std::runtime_error( + "graph capture requires tensors backed by InfiniCore-owned memory"); + } + freeze_for_capture_(it->second); +} + +void PinnableBlockAllocator::release_frozen_blocks_( + const std::vector> &blocks) noexcept { + try { + std::lock_guard lock(mutex_); + for (const auto &block : blocks) { + if (block->pin_count > 0) { + --block->pin_count; + } + } + } catch (...) { + } +} + +// ------------------- allocate ------------------- +std::byte *PinnableBlockAllocator::allocate(size_t size) { + if (size == 0) { + return nullptr; + } + std::lock_guard lock(mutex_); + // Align size to 256 bytes for GPU + size = align_up(size, 256); + + std::shared_ptr block; + + // 1. Try size-class allocation for small/medium + for (auto &cls : size_classes_) { + if (size <= cls.block_size) { + const auto free_block = std::find_if( + cls.free_blocks.begin(), cls.free_blocks.end(), + [](const auto &block) { + return !block->in_use && block->pin_count == 0; + }); + if (free_block != cls.free_blocks.end()) { + block = *free_block; + cls.free_blocks.erase(free_block); + block->in_use = true; + block->use_count = 1; + freeze_for_capture_(block); + return reinterpret_cast(block->ptr); + } + // Allocate a new block for this class + block = std::make_shared(); + block->size = cls.block_size; + block->in_use = true; + block->use_count = 1; + freeze_for_capture_(block); + + INFINICORE_CHECK_ERROR(infini::rt::runtime::Malloc(&block->ptr, block->size)); + + all_blocks_[block->ptr] = block; + return reinterpret_cast(block->ptr); + } + } + + // 2. Large block allocation + // Try to reuse an unpinned free large block. + auto it = std::find_if(large_blocks_.begin(), large_blocks_.end(), + [size](const std::shared_ptr &b) { + return b->size >= size && !b->in_use && b->pin_count == 0; + }); + + if (it != large_blocks_.end()) { + block = *it; + block->in_use = true; + block->use_count = 1; + freeze_for_capture_(block); + return reinterpret_cast(block->ptr); + } + + // Allocate new large block + block = std::make_shared(); + block->size = size; + block->in_use = true; + block->use_count = 1; + freeze_for_capture_(block); + + INFINICORE_CHECK_ERROR(infini::rt::runtime::Malloc(&block->ptr, block->size)); + + large_blocks_.push_back(block); + all_blocks_[block->ptr] = block; + + return reinterpret_cast(block->ptr); +} + +// ------------------- deallocate ------------------- +void PinnableBlockAllocator::deallocate(std::byte *ptr) { + if (ptr == nullptr) { + return; + } + + std::lock_guard lock(mutex_); + + auto it = all_blocks_.find(reinterpret_cast(ptr)); + if (it == all_blocks_.end()) { + throw std::runtime_error("Pointer not allocated by this allocator"); + } + + auto block = it->second; + if (!block->in_use || block->use_count == 0) { + throw std::runtime_error("Double free detected in PinnableBlockAllocator"); + } + + --block->use_count; + if (block->use_count > 0) { + return; + } + + block->in_use = false; + for (auto &cls : size_classes_) { + if (block->size == cls.block_size) { + cls.free_blocks.push_back(block); + break; + } + } +} + +size_t PinnableBlockAllocator::mark_in_use_(void *ptr, bool in_use) { + std::lock_guard lock(mutex_); + + auto it = all_blocks_.find(reinterpret_cast(ptr)); + if (it == all_blocks_.end()) { + throw std::runtime_error("Pointer not allocated by this allocator"); + } + + auto block = it->second; + if (in_use) { + for (auto &cls : size_classes_) { + if (block->size == cls.block_size) { + cls.free_blocks.erase( + std::remove( + cls.free_blocks.begin(), + cls.free_blocks.end(), + block), + cls.free_blocks.end()); + break; + } + } + block->in_use = true; + ++block->use_count; + } else if (block->use_count > 0) { + --block->use_count; + block->in_use = block->use_count > 0; + } + return it->second->size; +} + +// ------------------- trim ------------------- +void PinnableBlockAllocator::trim() { + std::lock_guard lock(mutex_); + // Free non-frozen size-class blocks + for (auto &cls : size_classes_) { + for (auto it = cls.free_blocks.begin(); it != cls.free_blocks.end();) { + if ((*it)->pin_count == 0) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::Free((*it)->ptr)); + all_blocks_.erase((*it)->ptr); + it = cls.free_blocks.erase(it); + } else { + ++it; + } + } + } + // Free non-frozen large blocks + for (auto it = large_blocks_.begin(); it != large_blocks_.end();) { + if ((*it)->pin_count == 0 && !(*it)->in_use) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::Free((*it)->ptr)); + all_blocks_.erase((*it)->ptr); + it = large_blocks_.erase(it); + } else { + ++it; + } + } +} + +// ------------------- Destructor ------------------- +PinnableBlockAllocator::~PinnableBlockAllocator() { + std::lock_guard lock(mutex_); + for (auto &p : all_blocks_) { + if (p.second->ptr) { + (void)infini::rt::runtime::Free(p.second->ptr); + } + } + all_blocks_.clear(); + large_blocks_.clear(); + for (auto &cls : size_classes_) { + cls.free_blocks.clear(); + } +} + +} // namespace infinicore diff --git a/csrc/infinicore/src/context/allocators/pinnable_block_allocator.hpp b/csrc/infinicore/src/context/allocators/pinnable_block_allocator.hpp new file mode 100644 index 000000000..c70b56d55 --- /dev/null +++ b/csrc/infinicore/src/context/allocators/pinnable_block_allocator.hpp @@ -0,0 +1,81 @@ +#pragma once + +#include "memory_allocator.hpp" + +#include +#include +#include +#include +#include + +namespace infinicore { +class PinnableBlockAllocator : public MemoryAllocator { + // Represents a single memory block + struct Block { + void *ptr = nullptr; // Device pointer + size_t size = 0; // Block size in bytes + size_t pin_count = 0; // Number of graphs retaining this block + bool in_use = false; // Wether the block is currently in use + size_t use_count = 0; // Number of Memory owners for this block + }; + + // A simple size-class allocator for small/medium blocks + struct SizeClass { + size_t block_size; // Fixed size for this class + std::vector> free_blocks; + }; + +public: + class PinLease { + public: + PinLease(PinnableBlockAllocator *allocator, + std::vector> blocks) noexcept; + ~PinLease() noexcept; + + PinLease(const PinLease &) = delete; + PinLease &operator=(const PinLease &) = delete; + + private: + PinnableBlockAllocator *allocator_; + std::vector> blocks_; + }; + + PinnableBlockAllocator(Device device); + ~PinnableBlockAllocator(); + + std::byte *allocate(size_t size) override; + void deallocate(std::byte *ptr) override; + + // Manage one graph-capture pin transaction. + void begin_pin_mode(); + std::shared_ptr commit_pin_mode(); + void cancel_pin_mode(); + void retain_for_capture(void *ptr); + + // internal use only, force set in_use flag for a mem block + // return the size of the block + size_t mark_in_use_(void *ptr, bool in_use); + + // trim cached blocks back to GPU (not pinned) + void trim(); + +private: + Device device_; + + bool pinned_mode_ = false; + std::thread::id pin_owner_; + + std::vector size_classes_; + std::vector> large_blocks_; + std::unordered_map> all_blocks_; + std::vector> capture_frozen_blocks_; + std::unordered_set capture_frozen_block_set_; + + std::mutex mutex_; // Thread safety + + void freeze_for_capture_(const std::shared_ptr &block); + void release_frozen_blocks_( + const std::vector> &blocks) noexcept; +}; + +} // namespace infinicore diff --git a/csrc/infinicore/src/context/allocators/stream_ordered_allocator.cc b/csrc/infinicore/src/context/allocators/stream_ordered_allocator.cc new file mode 100644 index 000000000..799807644 --- /dev/null +++ b/csrc/infinicore/src/context/allocators/stream_ordered_allocator.cc @@ -0,0 +1,36 @@ +#include "stream_ordered_allocator.hpp" + +#include "../../utils.hpp" + +namespace infinicore { +StreamOrderedAllocator::StreamOrderedAllocator(Device device) : MemoryAllocator(), device_(device) {} + +std::byte *StreamOrderedAllocator::allocate(size_t size) { + if (size == 0) { + return nullptr; + } + void *ptr = nullptr; + if (device_.type() != Device::Type::kCpu) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::MallocAsync( + &ptr, + size, + context::getStream())); + } else { + INFINICORE_CHECK_ERROR(infini::rt::runtime::Malloc(&ptr, size)); + } + return (std::byte *)ptr; +} + +void StreamOrderedAllocator::deallocate(std::byte *ptr) { + if (ptr == nullptr) { + return; + } + if (device_.type() != Device::Type::kCpu) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::FreeAsync( + ptr, + context::getStream())); + } else { + INFINICORE_CHECK_ERROR(infini::rt::runtime::Free(ptr)); + } +} +} // namespace infinicore diff --git a/csrc/infinicore/src/context/allocators/stream_ordered_allocator.hpp b/csrc/infinicore/src/context/allocators/stream_ordered_allocator.hpp new file mode 100644 index 000000000..ed2fdd6e6 --- /dev/null +++ b/csrc/infinicore/src/context/allocators/stream_ordered_allocator.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "memory_allocator.hpp" + +#include "../context_impl.hpp" + +namespace infinicore { +class StreamOrderedAllocator : public MemoryAllocator { +public: + explicit StreamOrderedAllocator(Device device); + ~StreamOrderedAllocator() = default; + + std::byte *allocate(size_t size) override; + void deallocate(std::byte *ptr) override; + +private: + Device device_; +}; + +} // namespace infinicore diff --git a/csrc/infinicore/src/context/context_impl.cc b/csrc/infinicore/src/context/context_impl.cc new file mode 100644 index 000000000..919ff07b8 --- /dev/null +++ b/csrc/infinicore/src/context/context_impl.cc @@ -0,0 +1,358 @@ +#include "context_impl.hpp" +#include "internal.hpp" + +#include "../utils.hpp" + +#include +#include +#include +#include + +namespace infinicore { +namespace { + +constexpr std::array(Device::Type::kCount)> kDefaultDevicePriority{ + Device::Type::kNvidia, + Device::Type::kCambricon, + Device::Type::kAscend, + Device::Type::kMetax, + Device::Type::kMoore, + Device::Type::kIluvatar, + Device::Type::kHygon, + Device::Type::kCpu, +}; + +void warn_graph_cleanup_failure(const char *operation, const char *detail) noexcept { + try { + spdlog::warn("{} failed during graph cleanup: {}", operation, detail); + } catch (...) { + } +} + +} // namespace + +thread_local std::shared_ptr ContextImpl::current_runtime_; +thread_local std::shared_ptr ContextImpl::graph_runtime_; + +std::shared_ptr ContextImpl::getOrCreateRuntimeLocked(Device device, const std::thread::id &thread_id) { + const auto type_index = static_cast(device.type()); + const auto device_index = static_cast(device.index()); + INFINICORE_ASSERT(type_index < runtime_table_.size()); + INFINICORE_ASSERT(device_index < runtime_table_[type_index].size()); + + auto &thread_runtimes = runtime_table_[type_index][device_index]; + if (const auto found = thread_runtimes.find(thread_id); found != thread_runtimes.end()) { + if (auto runtime = found->second.lock()) { + return runtime; + } + thread_runtimes.erase(found); + } + + auto runtime = std::shared_ptr(new Runtime(device)); + thread_runtimes.emplace(thread_id, runtime); + return runtime; +} + +template +void ContextImpl::initializeDeviceType() { + if constexpr (infini::rt::DeviceEnabled::value) { + int device_count = 0; + infini::rt::set_runtime_device_type(device_type); + INFINICORE_CHECK_ERROR(infini::rt::runtime::GetDeviceCount(&device_count)); + INFINICORE_ASSERT(device_count >= 0); + runtime_table_[static_cast(device_type)].resize(static_cast(device_count)); + } +} + +Runtime *ContextImpl::getCurrentRuntime() { + if (current_runtime_ != nullptr) { + spdlog::debug("getCurrentRuntime() returning {} (ptr={})", current_runtime_->device().ToString(), static_cast(current_runtime_.get())); + return current_runtime_.get(); + } + + spdlog::debug("current_runtime_ is null, performing lazy initialization"); + const auto thread_id = std::this_thread::get_id(); + { + std::lock_guard lock{runtime_table_mutex_}; + for (const auto device_type : kDefaultDevicePriority) { + const auto type_index = static_cast(device_type); + if (!runtime_table_[type_index].empty()) { + current_runtime_ = getOrCreateRuntimeLocked( + Device{device_type, 0}, thread_id); + break; + } + } + } + + INFINICORE_ASSERT(current_runtime_ != nullptr); + current_runtime_->activate(); + spdlog::debug("Lazy init: Set current_runtime_ to {} (ptr={})", current_runtime_->device().ToString(), static_cast(current_runtime_.get())); + return current_runtime_.get(); +} + +void ContextImpl::setDevice(Device device) { + if (device == getCurrentRuntime()->device()) { + current_runtime_->activate(); + return; + } + + std::shared_ptr runtime; + { + std::lock_guard lock{runtime_table_mutex_}; + runtime = getOrCreateRuntimeLocked(device, std::this_thread::get_id()); + } + current_runtime_ = std::move(runtime); + current_runtime_->activate(); +} + +size_t ContextImpl::getDeviceCount(Device::Type type) { + const auto type_index = static_cast(type); + if (type_index >= runtime_table_.size()) { + throw std::invalid_argument("invalid device type"); + } + std::lock_guard lock{runtime_table_mutex_}; + return runtime_table_[type_index].size(); +} + +bool ContextImpl::isGraphRecording() { + getCurrentRuntime(); + return graph_runtime_ != nullptr + && graph_runtime_ == current_runtime_ + && graph_runtime_->isGraphRecording(); +} + +void ContextImpl::startGraphRecording() { + getCurrentRuntime(); + if (graph_runtime_ != nullptr) { + throw std::runtime_error("graph recording is already active on this thread"); + } + current_runtime_->startGraphRecording(); + graph_runtime_ = current_runtime_; +} + +void ContextImpl::addGraphOperator(std::shared_ptr op) { + getCurrentRuntime(); + if (graph_runtime_ == nullptr || graph_runtime_ != current_runtime_) { + throw std::runtime_error("cannot record a graph operator on a non-capture device"); + } + graph_runtime_->addGraphOperator(std::move(op)); +} + +std::shared_ptr ContextImpl::stopGraphRecording() { + getCurrentRuntime(); + if (graph_runtime_ == nullptr) { + return current_runtime_->stopGraphRecording(); + } + + auto owner = graph_runtime_; + auto previous = current_runtime_; + current_runtime_ = owner; + bool stop_attempted = false; + try { + current_runtime_->activate(); + stop_attempted = true; + auto graph = owner->stopGraphRecording(); + graph_runtime_.reset(); + current_runtime_ = previous; + current_runtime_->activate(); + return graph; + } catch (...) { + const auto original_error = std::current_exception(); + if (stop_attempted) { + graph_runtime_.reset(); + } + current_runtime_ = previous; + try { + current_runtime_->activate(); + } catch (const std::exception &error) { + warn_graph_cleanup_failure("restoring the previous runtime", error.what()); + } catch (...) { + warn_graph_cleanup_failure("restoring the previous runtime", "unknown error"); + } + std::rethrow_exception(original_error); + } +} + +void ContextImpl::cancelGraphRecording() noexcept { + auto owner = std::exchange(graph_runtime_, nullptr); + if (owner == nullptr) { + return; + } + + auto previous = current_runtime_; + current_runtime_ = owner; + try { + owner->activate(); + } catch (const std::exception &error) { + warn_graph_cleanup_failure("activating the graph runtime", error.what()); + } catch (...) { + warn_graph_cleanup_failure("activating the graph runtime", "unknown error"); + } + owner->cancelGraphRecording(); + + current_runtime_ = previous; + if (current_runtime_ != nullptr) { + try { + current_runtime_->activate(); + } catch (const std::exception &error) { + warn_graph_cleanup_failure("restoring the previous runtime", error.what()); + } catch (...) { + warn_graph_cleanup_failure("restoring the previous runtime", "unknown error"); + } + } +} + +ContextImpl &ContextImpl::singleton() { + static ContextImpl instance; + return instance; +} + +ContextImpl::ContextImpl() { + initializeDeviceType(); + initializeDeviceType(); + initializeDeviceType(); + initializeDeviceType(); + initializeDeviceType(); + initializeDeviceType(); + initializeDeviceType(); + initializeDeviceType(); + + const bool has_runtime = std::any_of( + runtime_table_.begin(), runtime_table_.end(), + [](const auto &devices) { return !devices.empty(); }); + INFINICORE_ASSERT(has_runtime); +} + +namespace context { + +void setDevice(Device device) { + ContextImpl::singleton().setDevice(device); +} + +Device getDevice() { + return ContextImpl::singleton().getCurrentRuntime()->device(); +} + +size_t getDeviceCount(Device::Type type) { + return ContextImpl::singleton().getDeviceCount(type); +} + +infini::rt::runtime::Stream getStream() { + return ContextImpl::singleton().getCurrentRuntime()->stream(); +} + +void syncStream() { + return ContextImpl::singleton().getCurrentRuntime()->syncStream(); +} + +void syncDevice() { + return ContextImpl::singleton().getCurrentRuntime()->syncDevice(); +} + +void trimMemory() { + return ContextImpl::singleton().getCurrentRuntime()->trimMemory(); +} + +std::shared_ptr allocateMemory(size_t size) { + return ContextImpl::singleton().getCurrentRuntime()->allocateMemory(size); +} + +std::shared_ptr allocateHostMemory(size_t size) { + setDevice(Device{Device::Type::kCpu}); + return allocateMemory(size); +} + +std::shared_ptr allocatePinnedHostMemory(size_t size) { + return ContextImpl::singleton().getCurrentRuntime()->allocatePinnedHostMemory(size); +} + +void memcpyH2D(void *dst, const void *src, size_t size, bool async) { + return ContextImpl::singleton().getCurrentRuntime()->memcpyH2D(dst, src, size, async); +} + +void memcpyD2H(void *dst, const void *src, size_t size) { + return ContextImpl::singleton().getCurrentRuntime()->memcpyD2H(dst, src, size); +} + +void memcpyD2D(void *dst, const void *src, size_t size, bool async) { + return ContextImpl::singleton().getCurrentRuntime()->memcpyD2D(dst, src, size, async); +} + +void memcpyH2H(void *dst, const void *src, size_t size) { + setDevice(Device{Device::Type::kCpu}); + return ContextImpl::singleton().getCurrentRuntime()->memcpyD2D(dst, src, size); +} + +void setDeviceMemory(void *ptr, int value, size_t count) { + return ContextImpl::singleton().getCurrentRuntime()->setDeviceMemory(ptr, value, count); +} + +void setDeviceMemoryAsync(void *ptr, int value, size_t count, infini::rt::runtime::Stream stream) { + return ContextImpl::singleton().getCurrentRuntime()->setDeviceMemoryAsync(ptr, value, count, stream); +} + +// Timing API implementations +infini::rt::runtime::Event createEvent() { + return ContextImpl::singleton().getCurrentRuntime()->createEvent(); +} + +infini::rt::runtime::Event createEventWithFlags(uint32_t flags) { + return ContextImpl::singleton().getCurrentRuntime()->createEventWithFlags(flags); +} + +void recordEvent(infini::rt::runtime::Event event, infini::rt::runtime::Stream stream) { + ContextImpl::singleton().getCurrentRuntime()->recordEvent(event, stream); +} + +bool queryEvent(infini::rt::runtime::Event event) { + return ContextImpl::singleton().getCurrentRuntime()->queryEvent(event); +} + +void synchronizeEvent(infini::rt::runtime::Event event) { + ContextImpl::singleton().getCurrentRuntime()->synchronizeEvent(event); +} + +void destroyEvent(infini::rt::runtime::Event event) { + ContextImpl::singleton().getCurrentRuntime()->destroyEvent(event); +} + +float elapsedTime(infini::rt::runtime::Event start, infini::rt::runtime::Event end) { + return ContextImpl::singleton().getCurrentRuntime()->elapsedTime(start, end); +} + +void streamWaitEvent(infini::rt::runtime::Stream stream, infini::rt::runtime::Event event) { + ContextImpl::singleton().getCurrentRuntime()->streamWaitEvent(stream, event); +} + +bool isGraphRecording() { + return ContextImpl::singleton().isGraphRecording(); +} + +void startGraphRecording() { + ContextImpl::singleton().startGraphRecording(); +} + +void addGraphOperator(std::shared_ptr op) { + ContextImpl::singleton().addGraphOperator(std::move(op)); +} + +std::shared_ptr stopGraphRecording() { + return ContextImpl::singleton().stopGraphRecording(); +} + +void cancelGraphRecording() noexcept { + ContextImpl::singleton().cancelGraphRecording(); +} + +std::shared_ptr reinstantiateBlob(std::shared_ptr blob) { + setDevice(blob->device()); + return ContextImpl::singleton().getCurrentRuntime()->reinstantiateBlob(blob); +} + +void retainGraphMemory(const std::shared_ptr &memory) { + ContextImpl::singleton().getCurrentRuntime()->retainGraphMemory(memory); +} + +} // namespace context + +} // namespace infinicore diff --git a/csrc/infinicore/src/context/context_impl.hpp b/csrc/infinicore/src/context/context_impl.hpp new file mode 100644 index 000000000..2cc852225 --- /dev/null +++ b/csrc/infinicore/src/context/context_impl.hpp @@ -0,0 +1,49 @@ +#pragma once +#include "infinicore/context/context.hpp" +#include "runtime/runtime.hpp" + +#include +#include +#include +#include +#include + +namespace infinicore { +class ContextImpl { +private: + using ThreadRuntimes = std::unordered_map>; + + mutable std::mutex runtime_table_mutex_; + // Runtimes are isolated by device and thread so their streams, allocators, + // and graph state cannot race with another thread. + std::array, static_cast(Device::Type::kCount)> runtime_table_; + // Active runtime for current thread. Can use "static thread local" because context is a process singleton. + static thread_local std::shared_ptr current_runtime_; + static thread_local std::shared_ptr graph_runtime_; + + std::shared_ptr getOrCreateRuntimeLocked(Device device, const std::thread::id &thread_id); + + template + void initializeDeviceType(); + +protected: + ContextImpl(); + +public: + Runtime *getCurrentRuntime(); + + void setDevice(Device); + + size_t getDeviceCount(Device::Type type); + + bool isGraphRecording(); + void startGraphRecording(); + void addGraphOperator(std::shared_ptr op); + std::shared_ptr stopGraphRecording(); + void cancelGraphRecording() noexcept; + + static ContextImpl &singleton(); + + friend class Runtime; +}; +} // namespace infinicore diff --git a/csrc/infinicore/src/context/internal.hpp b/csrc/infinicore/src/context/internal.hpp new file mode 100644 index 000000000..11de00b05 --- /dev/null +++ b/csrc/infinicore/src/context/internal.hpp @@ -0,0 +1,11 @@ +#pragma once + +#include "infinicore/device.hpp" +#include "infinicore/memory.hpp" + +#include "infinicore/graph/graph.hpp" + +namespace infinicore::context { +std::shared_ptr reinstantiateBlob(std::shared_ptr blob); +void retainGraphMemory(const std::shared_ptr &memory); +}; // namespace infinicore::context diff --git a/csrc/infinicore/src/context/runtime/runtime.cc b/csrc/infinicore/src/context/runtime/runtime.cc new file mode 100644 index 000000000..e12983764 --- /dev/null +++ b/csrc/infinicore/src/context/runtime/runtime.cc @@ -0,0 +1,430 @@ +#include "runtime.hpp" + +#include "../../utils.hpp" + +#include "../allocators/device_pinned_allocator.hpp" +#include "../allocators/host_allocator.hpp" +#include "../allocators/pinnable_block_allocator.hpp" +#include "../allocators/stream_ordered_allocator.hpp" +#include "../context_impl.hpp" + +#include + +namespace infinicore { +namespace { + +void warn_runtime_cleanup_failure(const char *operation, infini::rt::runtime::Error status) noexcept { + if (status == infini::rt::runtime::kSuccess) { + return; + } + try { + spdlog::warn("{} failed during Runtime cleanup with error code {}", + operation, + static_cast(status)); + } catch (...) { + } +} + +void warn_runtime_cleanup_failure(const char *operation, const char *detail) noexcept { + try { + spdlog::warn("{} failed during Runtime cleanup: {}", operation, detail); + } catch (...) { + } +} + +} // namespace + +Runtime::Runtime(Device device) : device_(device), graph_manager_(std::make_unique()) { + activate(); + if (device_.type() == Device::Type::kCpu) { + device_memory_allocator_ = std::make_unique(device); + } else { + device_memory_allocator_ = std::make_unique(device); + pinned_host_memory_allocator_ = std::make_unique(device); + } +} +Runtime::~Runtime() noexcept { + Runtime *restore_runtime = ContextImpl::current_runtime_.get(); + infini::rt::set_runtime_device_type(device_.type()); + const auto set_device_status = infini::rt::runtime::SetDevice(device_.index()); + warn_runtime_cleanup_failure("selecting the runtime device", set_device_status); + + try { + std::lock_guard lock{stream_mutex_}; + if (stream_ != nullptr) { + const auto synchronize_status = infini::rt::runtime::StreamSynchronize(stream_); + warn_runtime_cleanup_failure("synchronizing the runtime stream", synchronize_status); + } + } catch (const std::exception &error) { + warn_runtime_cleanup_failure("synchronizing the runtime stream", error.what()); + } catch (...) { + warn_runtime_cleanup_failure("synchronizing the runtime stream", "unknown error"); + } + + graph_manager_.reset(); + pinned_host_memory_allocator_.reset(); + device_memory_allocator_.reset(); + try { + std::lock_guard lock{stream_mutex_}; + if (stream_ != nullptr) { + const auto destroy_status = infini::rt::runtime::StreamDestroy(stream_); + warn_runtime_cleanup_failure("destroying the runtime stream", destroy_status); + stream_ = nullptr; + } + } catch (const std::exception &error) { + warn_runtime_cleanup_failure("destroying the runtime stream", error.what()); + } catch (...) { + warn_runtime_cleanup_failure("destroying the runtime stream", "unknown error"); + } + + if (restore_runtime != nullptr && restore_runtime != this) { + infini::rt::set_runtime_device_type(restore_runtime->device().type()); + warn_runtime_cleanup_failure( + "restoring the active runtime device", + infini::rt::runtime::SetDevice(restore_runtime->device().index())); + } +} + +Runtime *Runtime::activate() { + INFINICORE_ASSERT(device_.type() != infini::rt::Device::Type::kCount); + infini::rt::set_runtime_device_type(device_.type()); + INFINICORE_CHECK_ERROR(infini::rt::runtime::SetDevice(device_.index())); + return this; +} + +Device Runtime::device() const { + return device_; +} + +infini::rt::runtime::Stream Runtime::stream() const { + infini::rt::set_runtime_device_type(device_.type()); + INFINICORE_CHECK_ERROR(infini::rt::runtime::SetDevice(device_.index())); + + std::lock_guard lock{stream_mutex_}; + if (stream_ == nullptr) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::StreamCreate(&stream_)); + } + return stream_; +} + +void Runtime::syncStream() { + INFINICORE_CHECK_ERROR(infini::rt::runtime::StreamSynchronize(stream())); +} + +void Runtime::syncDevice() { + INFINICORE_CHECK_ERROR(infini::rt::runtime::DeviceSynchronize()); +} + +void Runtime::syncStreamForCleanup() noexcept { + Runtime *restore_runtime = ContextImpl::current_runtime_.get(); + infini::rt::set_runtime_device_type(device_.type()); + const auto set_device_status = infini::rt::runtime::SetDevice(device_.index()); + warn_runtime_cleanup_failure("selecting the graph runtime device", set_device_status); + + if (set_device_status == infini::rt::runtime::kSuccess) { + try { + std::lock_guard lock{stream_mutex_}; + if (stream_ != nullptr) { + const auto synchronize_status = infini::rt::runtime::StreamSynchronize(stream_); + warn_runtime_cleanup_failure("synchronizing the graph runtime stream", synchronize_status); + } + } catch (const std::exception &error) { + warn_runtime_cleanup_failure("synchronizing the graph runtime stream", error.what()); + } catch (...) { + warn_runtime_cleanup_failure("synchronizing the graph runtime stream", "unknown error"); + } + } + + if (restore_runtime != nullptr && restore_runtime != this) { + infini::rt::set_runtime_device_type(restore_runtime->device().type()); + warn_runtime_cleanup_failure( + "restoring the active runtime device", + infini::rt::runtime::SetDevice(restore_runtime->device().index())); + } +} + +void Runtime::trimMemory() { + device_memory_allocator_->trim(); +} + +void Runtime::releaseDeviceMemory(std::byte *ptr) noexcept { + Runtime *restore_runtime = ContextImpl::current_runtime_.get(); + const bool cross_runtime_release = restore_runtime != this + && device_.type() != Device::Type::kCpu; + + bool release_is_safe = true; + if (cross_runtime_release) { + infini::rt::set_runtime_device_type(device_.type()); + const auto set_device_status = infini::rt::runtime::SetDevice(device_.index()); + warn_runtime_cleanup_failure("selecting the allocation device", set_device_status); + if (set_device_status == infini::rt::runtime::kSuccess) { + const auto synchronize_status = infini::rt::runtime::DeviceSynchronize(); + warn_runtime_cleanup_failure("synchronizing a cross-runtime allocation", synchronize_status); + release_is_safe = synchronize_status == infini::rt::runtime::kSuccess; + } else { + release_is_safe = false; + } + } + + if (release_is_safe) { + try { + device_memory_allocator_->deallocate(ptr); + } catch (const std::exception &error) { + warn_runtime_cleanup_failure("releasing device memory", error.what()); + } catch (...) { + warn_runtime_cleanup_failure("releasing device memory", "unknown error"); + } + } + + if (cross_runtime_release && restore_runtime != nullptr) { + infini::rt::set_runtime_device_type(restore_runtime->device().type()); + warn_runtime_cleanup_failure( + "restoring the active runtime device", + infini::rt::runtime::SetDevice(restore_runtime->device().index())); + } +} + +std::shared_ptr Runtime::allocateMemory(size_t size) { + auto runtime = shared_from_this(); + std::byte *data_ptr = runtime->device_memory_allocator_->allocate(size); + return std::make_shared( + data_ptr, size, device_, + [runtime](std::byte *p) { + runtime->releaseDeviceMemory(p); + }); +} + +std::shared_ptr Runtime::allocatePinnedHostMemory(size_t size) { + if (!pinned_host_memory_allocator_) { + spdlog::warn("For CPU devices, pinned memory is not supported, falling back to regular host memory"); + return allocateMemory(size); + } + auto runtime = shared_from_this(); + std::byte *data_ptr = runtime->pinned_host_memory_allocator_->allocate(size); + return std::make_shared( + data_ptr, size, Device{Device::Type::kCpu}, + [runtime](std::byte *p) { + runtime->pinned_host_memory_allocator_->deallocate(p); + }, + true); +} + +std::shared_ptr Runtime::reinstantiateBlob(std::shared_ptr blob) { + std::lock_guard lock(reinstantiated_blob_mutex_); + + auto ptr = blob->data(); + auto it = reinstantiated_blobs_.find(ptr); + if (it != reinstantiated_blobs_.end()) { + if (auto memory = it->second.lock()) { + return memory; + } + } + + auto runtime = shared_from_this(); + runtime->device_memory_allocator_->mark_in_use_(ptr, true); + auto memory = std::make_shared( + blob->data(), blob->size(), device_, + [runtime](std::byte *p) { + runtime->releaseDeviceMemory(p); + }); + reinstantiated_blobs_[ptr] = memory; + return memory; +} + +void Runtime::retainGraphMemory(const std::shared_ptr &memory) { + if (memory->device() != device_) { + throw std::runtime_error( + "graph capture cannot retain memory from a different device"); + } + device_memory_allocator_->retain_for_capture(memory->data()); +} + +void Runtime::memcpyH2D(void *dst, const void *src, size_t size, bool async) { + if (device_.type() == Device::Type::kCpu) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::Memcpy(dst, src, size, infini::rt::runtime::kMemcpyHostToDevice)); + return; + } + + const auto current_stream = stream(); + INFINICORE_CHECK_ERROR( + infini::rt::runtime::MemcpyAsync(dst, src, size, infini::rt::runtime::kMemcpyHostToDevice, current_stream)); + if (!async) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::StreamSynchronize(current_stream)); + } +} + +void Runtime::memcpyD2H(void *dst, const void *src, size_t size) { + if (device_.type() == Device::Type::kCpu) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::Memcpy(dst, src, size, infini::rt::runtime::kMemcpyDeviceToHost)); + return; + } + + const auto current_stream = stream(); + INFINICORE_CHECK_ERROR( + infini::rt::runtime::MemcpyAsync(dst, src, size, infini::rt::runtime::kMemcpyDeviceToHost, current_stream)); + INFINICORE_CHECK_ERROR(infini::rt::runtime::StreamSynchronize(current_stream)); +} + +void Runtime::memcpyD2D(void *dst, const void *src, size_t size, bool async) { + if (device_.type() == Device::Type::kCpu) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::Memcpy(dst, src, size, infini::rt::runtime::kMemcpyDeviceToDevice)); + return; + } + + const auto current_stream = stream(); + INFINICORE_CHECK_ERROR( + infini::rt::runtime::MemcpyAsync(dst, src, size, infini::rt::runtime::kMemcpyDeviceToDevice, current_stream)); + if (!async) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::StreamSynchronize(current_stream)); + } +} + +void Runtime::setDeviceMemory(void *ptr, int value, size_t count) { + if (device_.type() == Device::Type::kCpu) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::Memset(ptr, value, count)); + return; + } + + const auto current_stream = stream(); + INFINICORE_CHECK_ERROR(infini::rt::runtime::MemsetAsync(ptr, value, count, current_stream)); + INFINICORE_CHECK_ERROR(infini::rt::runtime::StreamSynchronize(current_stream)); +} + +void Runtime::setDeviceMemoryAsync(void *ptr, int value, size_t count, infini::rt::runtime::Stream stream) { + if (device_.type() != Device::Type::kCpu) { + if (stream == nullptr) { + stream = this->stream(); + } + INFINICORE_CHECK_ERROR(infini::rt::runtime::MemsetAsync(ptr, value, count, stream)); + } else { + INFINICORE_CHECK_ERROR(infini::rt::runtime::Memset(ptr, value, count)); + } +} + +// Timing method implementations +infini::rt::runtime::Event Runtime::createEvent() { + infini::rt::runtime::Event event = nullptr; + INFINICORE_CHECK_ERROR(infini::rt::runtime::EventCreate(&event)); + return event; +} + +infini::rt::runtime::Event Runtime::createEventWithFlags(uint32_t flags) { + infini::rt::runtime::Event event = nullptr; + INFINICORE_CHECK_ERROR(infini::rt::runtime::EventCreateWithFlags(&event, flags)); + return event; +} + +void Runtime::recordEvent(infini::rt::runtime::Event event, infini::rt::runtime::Stream stream) { + if (stream == nullptr) { + stream = this->stream(); + } + INFINICORE_CHECK_ERROR(infini::rt::runtime::EventRecord(event, stream)); +} + +bool Runtime::queryEvent(infini::rt::runtime::Event event) { + // InfiniRT does not expose a portable not-ready value, so every + // non-success query result is conservatively reported as incomplete. + return infini::rt::runtime::EventQuery(event) == infini::rt::runtime::kSuccess; +} + +void Runtime::synchronizeEvent(infini::rt::runtime::Event event) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::EventSynchronize(event)); +} + +void Runtime::destroyEvent(infini::rt::runtime::Event event) { + INFINICORE_CHECK_ERROR(infini::rt::runtime::EventDestroy(event)); +} + +float Runtime::elapsedTime(infini::rt::runtime::Event start, infini::rt::runtime::Event end) { + float ms; + INFINICORE_CHECK_ERROR(infini::rt::runtime::EventElapsedTime(&ms, start, end)); + return ms; +} + +void Runtime::streamWaitEvent(infini::rt::runtime::Stream stream, infini::rt::runtime::Event event) { + // Use current stream if no specific stream is provided + if (stream == nullptr) { + stream = this->stream(); + } + INFINICORE_CHECK_ERROR(infini::rt::runtime::StreamWaitEvent(stream, event, 0)); +} + +graph::GraphManager::CaptureState Runtime::graphCaptureState() const { + return graph_manager_->capture_state(); +} + +bool Runtime::isGraphRecording() const { + return graph_manager_->is_recording(); +} + +void Runtime::startGraphRecording() { + graph_manager_->start_recording(); + try { + device_memory_allocator_->begin_pin_mode(); + } catch (...) { + graph_manager_->cancel_recording(); + throw; + } +} + +void Runtime::addGraphOperator(std::shared_ptr op) { + return graph_manager_->add_operator(op); +} + +std::shared_ptr Runtime::stopGraphRecording() { + if (!graph_manager_->is_recording()) { + return graph_manager_->stop_recording(); + } + + std::shared_ptr graph; + try { + graph = graph_manager_->stop_recording(); + } catch (...) { + const auto original_error = std::current_exception(); + try { + device_memory_allocator_->cancel_pin_mode(); + } catch (...) { + } + graph_manager_->finish_recording(); + std::rethrow_exception(original_error); + } + + std::shared_ptr allocation_lease; + try { + allocation_lease = device_memory_allocator_->commit_pin_mode(); + } catch (...) { + const auto original_error = std::current_exception(); + try { + device_memory_allocator_->cancel_pin_mode(); + } catch (...) { + } + graph_manager_->finish_recording(); + std::rethrow_exception(original_error); + } + graph->retain_runtime(shared_from_this(), std::move(allocation_lease)); + graph_manager_->finish_recording(); + return graph; +} + +void Runtime::cancelGraphRecording() noexcept { + try { + graph_manager_->cancel_recording(); + } catch (const std::exception &error) { + warn_runtime_cleanup_failure("canceling graph recording", error.what()); + } catch (...) { + warn_runtime_cleanup_failure("canceling graph recording", "unknown error"); + } + try { + device_memory_allocator_->cancel_pin_mode(); + } catch (const std::exception &error) { + warn_runtime_cleanup_failure("rolling back graph allocations", error.what()); + } catch (...) { + warn_runtime_cleanup_failure("rolling back graph allocations", "unknown error"); + } +} + +std::string Runtime::toString() const { + return fmt::format("Runtime({})", device_.ToString()); +} + +} // namespace infinicore diff --git a/csrc/infinicore/src/context/runtime/runtime.hpp b/csrc/infinicore/src/context/runtime/runtime.hpp new file mode 100644 index 000000000..5df5eb0e2 --- /dev/null +++ b/csrc/infinicore/src/context/runtime/runtime.hpp @@ -0,0 +1,79 @@ +#pragma once + +#include "../allocators/pinnable_block_allocator.hpp" + +#include "infinicore/context/context.hpp" + +#include "../../graph/graph_manager.hpp" + +#include +#include +#include + +namespace infinicore { +class ContextImpl; +class Runtime : public std::enable_shared_from_this { +private: + Device device_; + mutable std::mutex stream_mutex_; + mutable infini::rt::runtime::Stream stream_ = nullptr; + std::unique_ptr device_memory_allocator_; + std::unique_ptr pinned_host_memory_allocator_; + std::unique_ptr graph_manager_; + std::mutex reinstantiated_blob_mutex_; + std::unordered_map> reinstantiated_blobs_; + + void releaseDeviceMemory(std::byte *ptr) noexcept; + void syncStreamForCleanup() noexcept; + +protected: + Runtime(Device device); + +public: + ~Runtime() noexcept; + + Runtime *activate(); + + Device device() const; + infini::rt::runtime::Stream stream() const; + + void syncStream(); + void syncDevice(); + void trimMemory(); + + std::shared_ptr allocateMemory(size_t size); + std::shared_ptr allocatePinnedHostMemory(size_t size); + std::shared_ptr reinstantiateBlob(std::shared_ptr blob); + void retainGraphMemory(const std::shared_ptr &memory); + + 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 setDeviceMemory(void *ptr, int value, size_t count); + void setDeviceMemoryAsync(void *ptr, int value, size_t count, infini::rt::runtime::Stream stream); + + // Timing methods + infini::rt::runtime::Event createEvent(); + infini::rt::runtime::Event createEventWithFlags(uint32_t flags); + void recordEvent(infini::rt::runtime::Event event, infini::rt::runtime::Stream stream = nullptr); + bool queryEvent(infini::rt::runtime::Event event); + void synchronizeEvent(infini::rt::runtime::Event event); + void destroyEvent(infini::rt::runtime::Event event); + float elapsedTime(infini::rt::runtime::Event start, infini::rt::runtime::Event end); + void streamWaitEvent(infini::rt::runtime::Stream stream, infini::rt::runtime::Event event); + + // Graph + graph::GraphManager::CaptureState graphCaptureState() const; + bool isGraphRecording() const; + void startGraphRecording(); + void addGraphOperator(std::shared_ptr op); + std::shared_ptr stopGraphRecording(); + void cancelGraphRecording() noexcept; + + std::string toString() const; + + friend class ContextImpl; + friend class graph::Graph; +}; +} // namespace infinicore diff --git a/csrc/infinicore/src/device.cc b/csrc/infinicore/src/device.cc new file mode 100644 index 000000000..7d4576a1f --- /dev/null +++ b/csrc/infinicore/src/device.cc @@ -0,0 +1 @@ +#include "infinicore/device.hpp" diff --git a/csrc/infinicore/src/device_event.cc b/csrc/infinicore/src/device_event.cc new file mode 100644 index 000000000..0102fe493 --- /dev/null +++ b/csrc/infinicore/src/device_event.cc @@ -0,0 +1,165 @@ +#include "infinicore.hpp" + +namespace infinicore { +namespace { + +void warn_cleanup_failure(const char *operation, const char *detail) noexcept { + try { + spdlog::warn("{} failed during DeviceEvent cleanup: {}", operation, detail); + } catch (...) { + } +} + +void warn_cleanup_failure(const char *operation, infini::rt::runtime::Error status) noexcept { + try { + spdlog::warn("{} failed during DeviceEvent cleanup with error code {}", + operation, + static_cast(status)); + } catch (...) { + } +} + +class ScopedDevice { +public: + explicit ScopedDevice(const Device &target) + : original_(context::getDevice()), changed_(original_ != target) { + if (changed_) { + try { + context::setDevice(target); + } catch (...) { + restore(); + throw; + } + } + } + + ScopedDevice(const ScopedDevice &) = delete; + ScopedDevice &operator=(const ScopedDevice &) = delete; + + ~ScopedDevice() noexcept { restore(); } + +private: + void restore() noexcept { + if (!changed_) { + return; + } + changed_ = false; + try { + context::setDevice(original_); + } catch (const std::exception &error) { + warn_cleanup_failure("restoring the previous device", error.what()); + } catch (...) { + warn_cleanup_failure("restoring the previous device", "unknown error"); + } + } + + Device original_; + bool changed_; +}; + +void destroy_event_noexcept(infini::rt::runtime::Event event, const Device &device) noexcept { + if (event == nullptr) { + return; + } + try { + ScopedDevice guard{device}; + const auto status = infini::rt::runtime::EventDestroy(event); + if (status != infini::rt::runtime::kSuccess) { + warn_cleanup_failure("destroying the event", status); + } + } catch (const std::exception &error) { + warn_cleanup_failure("destroying the event", error.what()); + } catch (...) { + warn_cleanup_failure("destroying the event", "unknown error"); + } +} + +} // namespace + +DeviceEvent::DeviceEvent() + : device_(context::getDevice()), is_recorded_(false) { + event_ = context::createEvent(); +} + +DeviceEvent::DeviceEvent(uint32_t flags) + : device_(context::getDevice()), is_recorded_(false) { + event_ = context::createEventWithFlags(flags); +} + +DeviceEvent::DeviceEvent(Device device) + : device_(device), is_recorded_(false) { + ScopedDevice guard{device_}; + event_ = context::createEvent(); +} + +DeviceEvent::DeviceEvent(Device device, uint32_t flags) + : device_(device), is_recorded_(false) { + ScopedDevice guard{device_}; + event_ = context::createEventWithFlags(flags); +} + +DeviceEvent::DeviceEvent(DeviceEvent &&other) noexcept + : event_(other.event_), device_(other.device_), is_recorded_(other.is_recorded_) { + other.event_ = nullptr; + other.is_recorded_ = false; +} + +DeviceEvent &DeviceEvent::operator=(DeviceEvent &&other) noexcept { + if (this != &other) { + destroy_event_noexcept(event_, device_); + event_ = other.event_; + device_ = other.device_; + is_recorded_ = other.is_recorded_; + other.event_ = nullptr; + other.is_recorded_ = false; + } + return *this; +} + +DeviceEvent::~DeviceEvent() noexcept { + destroy_event_noexcept(event_, device_); +} + +void DeviceEvent::record() { + ScopedDevice guard{device_}; + context::recordEvent(event_); + is_recorded_ = true; +} + +void DeviceEvent::record(infini::rt::runtime::Stream stream) { + ScopedDevice guard{device_}; + context::recordEvent(event_, stream); + is_recorded_ = true; +} + +void DeviceEvent::synchronize() { + ScopedDevice guard{device_}; + context::synchronizeEvent(event_); +} + +bool DeviceEvent::query() const { + ScopedDevice guard{device_}; + return context::queryEvent(event_); +} + +float DeviceEvent::elapsed_time(const DeviceEvent &other) const { + // Both events must be on the same device + if (device_ != other.device_) { + throw std::runtime_error("Cannot measure elapsed time between events on different devices"); + } + + // Both events must be recorded + if (!is_recorded_ || !other.is_recorded_) { + throw std::runtime_error("Both events must be recorded before measuring elapsed time"); + } + + ScopedDevice guard{device_}; + return context::elapsedTime(event_, other.event_); +} + +void DeviceEvent::wait(infini::rt::runtime::Stream stream) const { + ScopedDevice guard{device_}; + context::streamWaitEvent(stream, event_); +} + +} // namespace infinicore diff --git a/csrc/infinicore/src/dtype.cc b/csrc/infinicore/src/dtype.cc new file mode 100644 index 000000000..c3ffd9b38 --- /dev/null +++ b/csrc/infinicore/src/dtype.cc @@ -0,0 +1,13 @@ +#include + +namespace infinicore { + +std::string toString(const DataType &dtype) { + return std::string{infini::rt::kDataTypeToDesc.at(dtype)}; +} + +std::size_t dsize(const DataType &dtype) { + return infini::rt::kDataTypeToSize.at(dtype); +} + +} // namespace infinicore diff --git a/csrc/infinicore/src/graph/graph.cc b/csrc/infinicore/src/graph/graph.cc new file mode 100644 index 000000000..60384a516 --- /dev/null +++ b/csrc/infinicore/src/graph/graph.cc @@ -0,0 +1,371 @@ +#include "graph_manager.hpp" + +#include "../context/runtime/runtime.hpp" +#include "../utils.hpp" +#include "infinicore/context/context.hpp" + +#ifdef USE_INFINIRT_GRAPH +#include +#endif + +#include + +namespace infinicore::graph { + +#ifdef USE_INFINIRT_GRAPH +namespace rt_runtime = ::infini::rt::runtime; + +namespace { + +void warn_runtime_failure(const char *operation, rt_runtime::Error status) noexcept { + if (status == rt_runtime::kSuccess) { + return; + } + try { + spdlog::warn("{} failed during graph cleanup with error code {}", + operation, + static_cast(status)); + } catch (...) { + } +} + +void warn_runtime_failure(const char *operation, const char *detail) noexcept { + try { + spdlog::warn("{} failed during graph cleanup: {}", operation, detail); + } catch (...) { + } +} + +class RuntimeDeviceGuard { +public: + RuntimeDeviceGuard(::infini::rt::Device::Type target_type, int target_index) noexcept + : previous_type_(::infini::rt::runtime_device_type()) { + ::infini::rt::set_runtime_device_type(previous_type_); + const auto get_status = rt_runtime::GetDevice(&previous_index_); + restore_index_ = get_status == rt_runtime::kSuccess; + warn_runtime_failure("reading the previous device", get_status); + + if (target_type == ::infini::rt::Device::Type::kCount) { + warn_runtime_failure("selecting the graph device", "invalid device type"); + restore(); + return; + } + ::infini::rt::set_runtime_device_type(target_type); + const auto set_status = rt_runtime::SetDevice(target_index); + active_ = set_status == rt_runtime::kSuccess; + warn_runtime_failure("selecting the graph device", set_status); + if (!active_) { + restore(); + } + } + + RuntimeDeviceGuard(const RuntimeDeviceGuard &) = delete; + RuntimeDeviceGuard &operator=(const RuntimeDeviceGuard &) = delete; + + ~RuntimeDeviceGuard() noexcept { restore(); } + + bool active() const { return active_; } + +private: + void restore() noexcept { + if (restored_) { + return; + } + restored_ = true; + ::infini::rt::set_runtime_device_type(previous_type_); + if (restore_index_) { + warn_runtime_failure("restoring the previous device", rt_runtime::SetDevice(previous_index_)); + } + } + + ::infini::rt::Device::Type previous_type_; + int previous_index_ = 0; + bool restore_index_ = false; + bool active_ = false; + bool restored_ = false; +}; + +class StreamCaptureGuard { +public: + explicit StreamCaptureGuard(rt_runtime::Stream stream) : stream_(stream) {} + + StreamCaptureGuard(const StreamCaptureGuard &) = delete; + StreamCaptureGuard &operator=(const StreamCaptureGuard &) = delete; + + ~StreamCaptureGuard() noexcept { abort(); } + + rt_runtime::Error begin() { + const auto status = rt_runtime::StreamBeginCapture( + stream_, rt_runtime::StreamCaptureMode::kStreamCaptureModeRelaxed); + active_ = status == rt_runtime::kSuccess; + return status; + } + + rt_runtime::Error end(rt_runtime::Graph *graph) { + active_ = false; + return rt_runtime::StreamEndCapture(stream_, graph); + } + +private: + void abort() noexcept { + if (!active_) { + return; + } + active_ = false; + rt_runtime::Graph abandoned_graph = nullptr; + const auto end_status = rt_runtime::StreamEndCapture(stream_, &abandoned_graph); + warn_runtime_failure("ending an abandoned stream capture", end_status); + if (end_status == rt_runtime::kSuccess && abandoned_graph != nullptr) { + warn_runtime_failure("destroying an abandoned graph", rt_runtime::GraphDestroy(abandoned_graph)); + } + } + + rt_runtime::Stream stream_; + bool active_ = false; +}; + +} // namespace +#endif + +/* ========================= + * GraphTensor + * ========================= */ + +GraphTensor::GraphTensor(const Tensor &tensor) : Tensor(tensor->to_blob_()) { +} + +/* ========================= + * GraphOperator + * ========================= */ + +void DispatchableGraphOperator::run() const { + runner_(planned_meta_); +} + +DispatchableGraphOperator::~DispatchableGraphOperator() { + if (deleter_) { + deleter_(&planned_meta_); + } +} + +/* ========================= + * Graph + * ========================= */ + +#ifdef USE_INFINIRT_GRAPH +struct Graph::DeviceGraph { + rt_runtime::Graph graph = nullptr; + rt_runtime::GraphExec exec = nullptr; + rt_runtime::Stream stream = nullptr; + ::infini::rt::Device::Type device_type = ::infini::rt::Device::Type::kCount; + int device_index = 0; + + ~DeviceGraph() noexcept { + if (exec == nullptr && graph == nullptr) { + return; + } + RuntimeDeviceGuard guard{device_type, device_index}; + if (!guard.active()) { + warn_runtime_failure("activating the graph device for cleanup", "device selection failed"); + return; + } + if (exec) { + warn_runtime_failure("destroying the graph executable", rt_runtime::GraphExecDestroy(exec)); + exec = nullptr; + } + if (graph) { + warn_runtime_failure("destroying the graph", rt_runtime::GraphDestroy(graph)); + graph = nullptr; + } + } + + void launch() { + RuntimeDeviceGuard guard{device_type, device_index}; + INFINICORE_ASSERT(guard.active()); + INFINICORE_CHECK_ERROR(rt_runtime::GraphLaunch(exec, stream)); + } +}; +#else +struct Graph::DeviceGraph {}; +#endif + +Graph::Graph() { +} + +void Graph::retain_runtime(std::shared_ptr<::infinicore::Runtime> runtime, + std::shared_ptr allocation_lease) { + runtime_lease_ = std::move(runtime); + allocation_lease_ = std::move(allocation_lease); +} + +void Graph::run() const { + (void)context::isGraphRecording(); +#ifdef USE_INFINIRT_GRAPH + if (device_graph_ != nullptr && device_graph_.get()->exec != nullptr) { + device_graph_.get()->launch(); + return; + } +#endif + for (auto &op : op_list_) { + op->run(); + } +} + +void Graph::add_operator(std::shared_ptr op) { + op_list_.push_back(op); +} + +void Graph::instantiate() { +#ifdef USE_INFINIRT_GRAPH + // Reset device graph + device_graph_ = std::make_unique(); + auto current_device = context::getDevice(); + device_graph_->device_type = current_device.type(); + device_graph_->device_index = current_device.index(); + device_graph_->stream = context::getStream(); + RuntimeDeviceGuard device_guard{device_graph_->device_type, device_graph_->device_index}; + if (!device_guard.active()) { + spdlog::warn("InfiniRT graph runtime failed to select the current device. Falling back to eager execution."); + device_graph_.reset(); + return; + } + + // warmup + for (size_t iter = 0; iter < 5; ++iter) { + this->run(); + } + infinicore::context::syncStream(); + + StreamCaptureGuard capture_guard{device_graph_->stream}; + auto begin_status = capture_guard.begin(); + if (begin_status != rt_runtime::kSuccess) { + spdlog::warn("Fail to begin device graph capture."); + device_graph_.reset(); + return; + } + + // Run and record + this->run(); + + auto end_status = capture_guard.end(&device_graph_->graph); + if (end_status != rt_runtime::kSuccess) { + spdlog::warn("Fail to end device graph capture."); + device_graph_.reset(); + return; + } + + auto instantiate_status = rt_runtime::GraphInstantiate( + &device_graph_->exec, + device_graph_->graph); + if (instantiate_status != rt_runtime::kSuccess) { + static bool warned_once = false; + if (!warned_once) { + warned_once = true; + spdlog::warn("Fail to instantiate device graph."); + } + device_graph_.reset(); + return; + } + static bool logged_once = false; + if (!logged_once) { + logged_once = true; + spdlog::info("Using InfiniRT C++ graph runtime API for graph capture and replay."); + } +#endif +} + +Graph::~Graph() noexcept { + if (runtime_lease_ != nullptr) { + runtime_lease_->syncStreamForCleanup(); + } +} + +/* ========================= + * GraphManager + * ========================= */ + +bool GraphManager::is_recording() const { + std::lock_guard lock{mutex_}; + if (!recording_) { + return false; + } + const auto state = capture_owner_ == std::this_thread::get_id() + ? CaptureState::kActiveOwner + : CaptureState::kActiveNonOwner; + if (state == CaptureState::kActiveNonOwner) { + throw std::runtime_error("cannot access the shared runtime stream: another thread owns the graph capture"); + } + return true; +} + +GraphManager::CaptureState GraphManager::capture_state() const { + std::lock_guard lock{mutex_}; + if (!recording_) { + return CaptureState::kInactive; + } + return capture_owner_ == std::this_thread::get_id() + ? CaptureState::kActiveOwner + : CaptureState::kActiveNonOwner; +} + +void GraphManager::start_recording() { + auto graph = std::make_shared(); + const auto caller = std::this_thread::get_id(); + std::lock_guard lock{mutex_}; + if (recording_ && capture_owner_ != caller) { + throw std::runtime_error("cannot start graph recording: another thread owns the capture"); + } + if (recording_) { + spdlog::warn("Graph is already recording. Previous recording will be dropped."); + } + recording_ = true; + capture_owner_ = caller; + graph_ = std::move(graph); +} + +void GraphManager::add_operator(std::shared_ptr op) { + std::lock_guard lock{mutex_}; + INFINICORE_ASSERT(recording_ && capture_owner_ == std::this_thread::get_id()); + graph_->add_operator(op); +} + +std::shared_ptr GraphManager::stop_recording() { + std::shared_ptr graph; + { + std::lock_guard lock{mutex_}; + if (!recording_) { + spdlog::warn("Graph is not recording. Please start recording first."); + return nullptr; + } + if (capture_owner_ != std::this_thread::get_id()) { + throw std::runtime_error("cannot stop graph recording: another thread owns the capture"); + } + graph = std::exchange(graph_, nullptr); + } +#ifdef USE_INFINIRT_GRAPH + graph->instantiate(); +#endif + return graph; +} + +void GraphManager::finish_recording() { + std::lock_guard lock{mutex_}; + INFINICORE_ASSERT(recording_ && capture_owner_ == std::this_thread::get_id()); + recording_ = false; + capture_owner_ = {}; +} + +void GraphManager::cancel_recording() { + std::shared_ptr graph; + { + std::lock_guard lock{mutex_}; + if (recording_ && capture_owner_ == std::this_thread::get_id()) { + recording_ = false; + capture_owner_ = {}; + graph = std::exchange(graph_, nullptr); + } + } + graph.reset(); +} + +} // namespace infinicore::graph diff --git a/csrc/infinicore/src/graph/graph_manager.hpp b/csrc/infinicore/src/graph/graph_manager.hpp new file mode 100644 index 000000000..68eff29c5 --- /dev/null +++ b/csrc/infinicore/src/graph/graph_manager.hpp @@ -0,0 +1,37 @@ +#pragma once + +#include "infinicore/graph/graph.hpp" + +#include +#include +#include + +namespace infinicore::graph { + +class GraphManager { +public: + enum class CaptureState { + kInactive, + kActiveOwner, + kActiveNonOwner, + }; + + GraphManager() = default; + ~GraphManager() = default; + + CaptureState capture_state() const; + bool is_recording() const; + void start_recording(); + void add_operator(std::shared_ptr op); + std::shared_ptr stop_recording(); + void finish_recording(); + void cancel_recording(); + +private: + mutable std::mutex mutex_; + std::shared_ptr graph_; + std::thread::id capture_owner_; + bool recording_ = false; +}; + +} // namespace infinicore::graph diff --git a/csrc/infinicore/src/io.cc b/csrc/infinicore/src/io.cc new file mode 100644 index 000000000..29961104f --- /dev/null +++ b/csrc/infinicore/src/io.cc @@ -0,0 +1,650 @@ +/* + * This file contains some code derived from xtensor (https://github.com/xtensor-stack/xtensor), + * which is licensed under the BSD-3-Clause license. + */ + +#include "infinicore/io.hpp" +#include "../utils/custom_types.h" +#include "infinicore/context/context.hpp" +#include "infinicore/dtype.hpp" +#include "infinicore/tensor.hpp" +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using Tensor = infinicore::Tensor; +using TensorSliceParams = infinicore::TensorSliceParams; +using DataType = infinicore::DataType; +using Device = infinicore::Device; +using TensorImpl = infinicore::TensorImpl; +using Size = infinicore::Size; + +/** + * @brief This function extracts a scalar or sub-tensor from a tensor using a vector of indexes. + */ +inline Tensor at_impl(const Tensor &tensor, const std::vector &indexes) { + if (indexes.size() > tensor->ndim()) { + throw std::runtime_error("at_impl:: Number of indexes (" + std::to_string(indexes.size()) + ") exceeds tensor dimensions (" + std::to_string(tensor->ndim()) + ")"); + } + + for (size_t i = 0; i < indexes.size(); i++) { + if (indexes[i] >= tensor->shape()[i]) { + throw std::runtime_error("at_impl :: Index " + std::to_string(indexes[i]) + " is out of bounds for dimension " + std::to_string(i)); + } + } + + std::vector slices; + slices.reserve(indexes.size()); + for (size_t i = 0; i < indexes.size(); i++) { + slices.push_back({i, indexes[i], 1}); + } + + Tensor result = tensor->narrow(slices); + for (size_t i = 0; i < indexes.size(); i++) { + result = result->squeeze(0); + } + + return result; +} + +template +Tensor at(const Tensor &tensor, Args... args) { + std::vector indexes = {static_cast(args)...}; + return at_impl(tensor, indexes); +} + +[[maybe_unused]] Tensor at(const Tensor &tensor, std::initializer_list indexes) { + std::vector indexes_vec(indexes.begin(), indexes.end()); + return at_impl(tensor, indexes_vec); +} + +Tensor at(const Tensor &tensor, const std::vector &indexes) { + return at_impl(tensor, indexes); +} + +/** + * @brief read a value from raw data pointer based on DataType. + */ +template +T item_impl(const std::byte *data, DataType dtype) { + switch (dtype) { + case DataType::kFloat16: { + const fp16_t *ptr = reinterpret_cast(data); + float f = _f16_to_f32(ptr[0]); + return static_cast(f); + } + case DataType::kBFloat16: { + const bf16_t *ptr = reinterpret_cast(data); + float f = _bf16_to_f32(ptr[0]); + return static_cast(f); + } + default: + break; + } + + const T *ptr = reinterpret_cast(data); + return ptr[0]; +} + +/** + * @brief Extracts a scalar value from a single-element tensor. + * The tensor must have exactly one element and must be located on CPU device. + * + * @code{.cpp} + * float value = item(tensor, dtype); // Extract as float + * @endcode + */ +template +T item(const Tensor &tensor) { + if (tensor->numel() != 1) { + throw std::runtime_error("item() can only be called on a tensor with exactly one element, but got " + std::to_string(tensor->numel()) + " elements"); + } + + if (tensor->device().type() != Device::Type::kCpu) { + throw std::runtime_error("item() can only be called on a CPU tensor, but got device: " + tensor->device().ToString()); + } + + const std::byte *data = tensor->data(); + DataType dtype = tensor->dtype(); + return item_impl(data, dtype); +} +} // namespace + +namespace infinicore { +namespace print_options { + +template +class fmtflags_guard { +public: + explicit fmtflags_guard(S &stream) + : m_stream(stream), m_flags(stream.flags()) {} + ~fmtflags_guard() { m_stream.flags(m_flags); } + +private: + S &m_stream; + std::ios_base::fmtflags m_flags; +}; + +struct PrintOptionsImpl { + int edge_items = 3; // default edge items: 3 means print 3 items of each dimension. + int line_width = 80; // default line width: 75 means print 75 chars per line. + int threshold = 1000; // default threshold: 1000 means print 1000 elements of the tensor. + int precision = 4; // default precision: -1 means no precision limit. + int sci_mode = -1; // default sci_mode: -1 means auto decision. +}; + +inline PrintOptionsImpl &print_options() { + static PrintOptionsImpl po; + return po; +} + +void set_line_width(int line_width) { + print_options().line_width = line_width; +} + +void set_threshold(int threshold) { + print_options().threshold = threshold; +} + +void set_edge_items(int edge_items) { + print_options().edge_items = edge_items; +} + +void set_precision(int precision) { + print_options().precision = precision; +} + +void set_sci_mode(int sci_mode) { + print_options().sci_mode = sci_mode; +} + +/** + * @brief read print options from the out stream and global settings. + */ +inline print_options::PrintOptionsImpl get_print_options(std::ostream &out) { + print_options::PrintOptionsImpl res; + +// Macro to read option from stream, apply default if not set, or reset stream value +#define PROCESS_PRINT_OPTION(OPTION) \ + res.OPTION = static_cast(out.iword(print_options::OPTION::id())); \ + if (res.OPTION > 0) { \ + out.iword(print_options::OPTION::id()) = long(-1); \ + } else { \ + res.OPTION = print_options::print_options().OPTION; \ + } + + // Process all print options + PROCESS_PRINT_OPTION(edge_items); + PROCESS_PRINT_OPTION(line_width); + PROCESS_PRINT_OPTION(threshold); + PROCESS_PRINT_OPTION(precision); + + res.sci_mode = print_options::print_options().sci_mode; + +#undef PROCESS_PRINT_OPTION + return res; +} + +template +struct Printer; + +/** + * @brief Printer specialization for floating-point types (float, double, long double). + */ +template +struct Printer::value>> { + using value_type = T; + using cache_type = std::vector; + using cache_iterator = typename cache_type::const_iterator; + + explicit Printer(std::streamsize precision, int sci_mode = 0) : m_precision(precision), m_sci_mode(sci_mode) {} + + void calculate() { + m_precision = m_precision > m_required_precision ? m_required_precision : m_precision; + m_it = m_cache.cbegin(); + + if (m_scientific) { + // 3 = sign, number and dot and 4 = "e+00" + m_width = m_precision + 7; + if (m_large_exponent) { + // = e+000 (additional number) + m_width += 1; + } + } else { + std::streamsize decimals = 1; // print a leading 0 + if (std::floor(m_max) != 0) { + decimals += std::streamsize(std::log10(std::floor(m_max))); + } + // 2 => sign and dot + m_width = 2 + decimals + m_precision; + } + if (!m_required_precision) { + --m_width; + } + } + + std::ostream &print_next(std::ostream &out) { + if ((1 == m_sci_mode) || ((-1 == m_sci_mode) && m_scientific)) { + if (!m_large_exponent) { + out << std::scientific; + out.width(m_width); + out << (*m_it); + } else { + std::stringstream buf; + buf.width(m_width); + buf << std::scientific; + buf.precision(m_precision); + buf << (*m_it); + std::string res = buf.str(); + + if (res[res.size() - 4] == 'e') { + res.erase(0, 1); + res.insert(res.size() - 2, "0"); + } + out << res; + } + } else { + std::stringstream buf; + buf.width(m_width); + buf << std::fixed; + buf.precision(m_precision); + buf << (*m_it); + if (!m_required_precision && !std::isinf(*m_it) && !std::isnan(*m_it)) { + buf << '.'; + } + std::string res = buf.str(); + auto sit = res.rbegin(); + while (*sit == '0') { + *sit = ' '; + ++sit; + } + out << res; + } + ++m_it; + return out; + } + + void update(const value_type &val) { + if (val != 0 && !std::isinf(val) && !std::isnan(val)) { + if (!m_scientific || !m_large_exponent) { + int exponent = 1 + int(std::log10(std::abs(val))); + if (exponent <= -5 || exponent > 7) { + m_scientific = true; + m_required_precision = m_precision; + if (exponent <= -100 || exponent >= 100) { + m_large_exponent = true; + } + } + } + + if (std::abs(val) > m_max) { + m_max = std::abs(val); + } + if (m_required_precision < m_precision) { + while (std::floor(val * std::pow(10, m_required_precision)) != val * std::pow(10, m_required_precision)) { + m_required_precision++; + } + } + } + m_cache.push_back(val); + } + + std::streamsize width() const { return m_width; } + +private: + bool m_large_exponent = false; + bool m_scientific = false; + + std::streamsize m_width = 9; + std::streamsize m_precision; + std::streamsize m_required_precision = 0; + value_type m_max = 0; + int m_sci_mode = -1; + cache_type m_cache; + cache_iterator m_it; +}; + +/** + * @brief Printer specialization for integer types (signed and unsigned integers). + */ +template +struct Printer< + T, std::enable_if_t::value && !std::is_same::value>> { + using value_type = T; + using cache_type = std::vector; + using cache_iterator = typename cache_type::const_iterator; + + explicit Printer(std::streamsize, int sci_mode = 0) {} + + void calculate() { + m_it = m_cache.cbegin(); + m_width = 1 + std::streamsize((m_max > 0) ? std::log10(m_max) : 0) + m_sign; + } + + std::ostream &print_next(std::ostream &out) { + // + enables printing of chars etc. as numbers + // TODO should chars be printed as numbers? + out.width(m_width); + out << +(*m_it); + ++m_it; + return out; + } + + void update(const value_type &val) { + // For unsigned types, abs is not needed (always non-negative) + // For signed types, we need to take absolute value + value_type abs_val; + if constexpr (std::is_signed::value) { + abs_val = (val < 0) ? -val : val; + } else { + abs_val = val; + } + + if (abs_val > m_max) { + m_max = abs_val; + } + + if (std::is_signed::value && val < 0) { + m_sign = true; + } + m_cache.push_back(val); + } + + std::streamsize width() { return m_width; } + +private: + std::streamsize m_width; + bool m_sign = false; + value_type m_max = 0; + + cache_type m_cache; + cache_iterator m_it; +}; + +/** + * @brief Printer specialization for bool type. + */ +template +struct Printer< + T, std::enable_if_t::value>> { + using value_type = bool; + using cache_type = std::vector; + using cache_iterator = typename cache_type::const_iterator; + + explicit Printer(std::streamsize, int sci_mode = 0) {} + + void calculate() { + m_it = m_cache.cbegin(); + } + + std::ostream &print_next(std::ostream &out) { + if (*m_it) { + out << " true"; + } else { + out << "false"; + } + // TODO: the following std::setw(5) isn't working correctly on OSX. + // out << std::boolalpha << std::setw(m_width) << (*m_it); + ++m_it; + return out; + } + + void update(const value_type &val) { m_cache.push_back(val); } + + std::streamsize width() { return m_width; } + +private: + std::streamsize m_width = 5; + cache_type m_cache; + cache_iterator m_it; +}; + +} // namespace print_options +} // namespace infinicore + +namespace infinicore { +namespace print_options { + +/** + * @brief Recursively traverses tensor dimensions to collect values for printing. + */ +template +void recurser_run(Printer &printer, + const Tensor &tensor, + std::vector indexes, + std::size_t lim = 0) { + + using size_type = Size; + const auto view = at(tensor, indexes); + if (view->ndim() == 0) { + T value = item(view); + printer.update(value); + } else { + size_type i = 0; + for (; i != static_cast(view->shape()[0] - 1); ++i) { + if (lim && size_type(view->shape()[0]) > (lim * 2) && i == lim) { + i = static_cast(view->shape()[0]) - lim; + if (lim <= 1) { + break; + } + } + indexes.push_back(static_cast(i)); + recurser_run(printer, tensor, indexes, lim); + indexes.pop_back(); + } + indexes.push_back(static_cast(i)); + recurser_run(printer, tensor, indexes, lim); + indexes.pop_back(); + } +} + +/** + * @brief Recursively prints tensor elements with proper formatting. + */ +template +std::ostream &xoutput(std::ostream &out, + const Tensor &tensor, + std::vector &indexes, + Printer &printer, + std::size_t blanks, + std::streamsize element_width, + std::size_t edge_items, + std::size_t line_width) { + + using size_type = Size; + const auto view = at(tensor, indexes); + if (view->ndim() == 0) { + printer.print_next(out); + } else { + std::string indents(blanks, ' '); + + size_type i = 0; + size_type elems_on_line = 0; + const size_type ewp2 = static_cast(element_width) + size_type(2); + const size_type line_lim = static_cast(std::floor(line_width / ewp2)); + + out << '['; + for (; i != size_type(view->shape()[0] - 1); ++i) { + + if (edge_items && size_type(view->shape()[0]) > (edge_items * 2) && i == edge_items) { + if (view->ndim() == 1 && line_lim != 0 && elems_on_line >= line_lim) { + out << " ...,"; + } else if (view->ndim() > 1) { + elems_on_line = 0; + out << "...," << std::endl + << indents; + } else { + out << "..., "; + } + i = size_type(view->shape()[0]) - edge_items; + if (edge_items <= 1) { + break; + } + } + if (view->ndim() == 1 && line_lim != 0 && elems_on_line >= line_lim) { + out << std::endl + << indents; + elems_on_line = 0; + } + + indexes.push_back(static_cast(i)); + xoutput(out, tensor, indexes, printer, blanks + 1, element_width, edge_items, + line_width) + << ','; + indexes.pop_back(); + elems_on_line++; + + if ((view->ndim() == 1) && !(line_lim != 0 && elems_on_line >= line_lim)) { + ; // out << ' '; + } else if (view->ndim() > 1) { + out << std::endl + << indents; + } + } + if (view->ndim() == 1 && line_lim != 0 && elems_on_line >= line_lim) { + out << std::endl + << indents; + } + + indexes.push_back(static_cast(i)); + xoutput(out, tensor, indexes, printer, blanks + 1, element_width, edge_items, + line_width) + << ']'; + indexes.pop_back(); + } + return out; +} + +template +std::ostream &pretty_print(const Tensor &original_tensor, + std::ostream &out = std::cout) { + Tensor tensor = original_tensor->to(Device::Type::kCpu); + bool on_cpu = original_tensor->device() == Device::Type::kCpu; + std::string device_str = original_tensor->device().ToString(); + infinicore::context::syncDevice(); + + fmtflags_guard guard(out); + + std::size_t edge_items = 0; + Size sz = tensor->numel(); + auto po = get_print_options(out); + + if (sz > static_cast(po.threshold)) { + edge_items = static_cast(po.edge_items); + } + if (sz == 0) { + out << "[]"; + return out; + } + + auto temp_precision = out.precision(); + auto precision = temp_precision; + + if (po.precision != -1) { + out.precision(static_cast(po.precision)); + precision = static_cast(po.precision); + } + + Printer printer(precision, po.sci_mode); + std::vector indexes = {}; + + recurser_run(printer, tensor, indexes, edge_items); + + printer.calculate(); + indexes.clear(); + + auto element_width = printer.width(); + + out << "tensor("; + xoutput(out, + tensor, + indexes, + printer, + 1 + 7, + element_width, + edge_items, + static_cast(po.line_width)); + + if (!on_cpu) { + out << ", device=" << '\'' << device_str << '\''; + } + + out << ", dtype=infinicore." << toString(tensor->dtype()) << ")\n"; + out.precision(temp_precision); // restore precision + return out; +} + +} // namespace print_options +} // namespace infinicore + +namespace infinicore { +std::ostream &operator<<(std::ostream &out, const Tensor &tensor) { + if (!tensor) { + out << "tensor([])\n"; + return out; + } + + switch (tensor->dtype()) { + case DataType::kInt8: // 3 + { + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kInt16: // 4 + { + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kInt32: // 5 + { + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kInt64: // 6 + { + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kUInt8: // 7 + { + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kUInt16: // 8 + { + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kUInt32: // 9 + { + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kUInt64: // 10 + { + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kFloat16: // 12 + { + // Convert F16 to F32 for printing + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kFloat32: // 13 + { + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kFloat64: // 14 + { + return infinicore::print_options::pretty_print(tensor, out); + } + case DataType::kBFloat16: // 19 + { + // Convert BF16 to F32 for printing + return infinicore::print_options::pretty_print(tensor, out); + } + default: + throw std::runtime_error("cant not print unknown dtype tensor : " + toString(tensor->dtype())); + } + + return out; +} +} // namespace infinicore diff --git a/csrc/infinicore/src/memory.cc b/csrc/infinicore/src/memory.cc new file mode 100644 index 000000000..0a7cc2df3 --- /dev/null +++ b/csrc/infinicore/src/memory.cc @@ -0,0 +1,33 @@ +#include "infinicore/memory.hpp" + +namespace infinicore { + +Memory::Memory(std::byte *data, + size_t size, + Device device, + Memory::Deleter deleter, + bool pin_memory) + : data_{data}, size_{size}, device_{device}, deleter_{deleter}, is_pinned_(pin_memory) {} + +Memory::~Memory() { + if (deleter_) { + deleter_(data_); + } +} + +std::byte *Memory::data() { + return data_; +} + +Device Memory::device() const { + return device_; +} + +size_t Memory::size() const { + return size_; +} + +bool Memory::is_pinned() const { + return is_pinned_; +} +} // namespace infinicore diff --git a/csrc/infinicore/src/nn/embedding.cc b/csrc/infinicore/src/nn/embedding.cc new file mode 100644 index 000000000..cc36b5287 --- /dev/null +++ b/csrc/infinicore/src/nn/embedding.cc @@ -0,0 +1,139 @@ +#include "infinicore/nn/embedding.hpp" +#include "infinicore/context/context.hpp" +#include "infinicore/ops.hpp" +#include +#include +#include + +namespace infinicore::nn { + +Embedding::Embedding(size_t num_embeddings, + size_t embedding_dim, + std::optional padding_idx, + const DataType &dtype, + const Device &device) + : num_embeddings_(num_embeddings), + embedding_dim_(embedding_dim), + padding_idx_(padding_idx), + dtype_(dtype) { + + device_ = device; + + // Validate padding_idx + if (padding_idx_.has_value()) { + int64_t idx = padding_idx_.value(); + if (idx < 0 || idx >= static_cast(num_embeddings)) { + throw std::invalid_argument( + "padding_idx must be within num_embeddings range, got " + std::to_string(idx) + " for num_embeddings=" + std::to_string(num_embeddings)); + } + } + + // Initialize parameter using macro + INFINICORE_NN_PARAMETER_INIT(weight, ({num_embeddings, embedding_dim}, dtype_, device)); + + // If padding_idx is specified, initialize that row to zeros + if (padding_idx_.has_value()) { + // TODO: Set weight[padding_idx] to zeros + // This would require a slice operation + } + + SPDLOG_DEBUG("Created Embedding module: num_embeddings={}, embedding_dim={}, dtype={}, padding_idx={}", + num_embeddings, embedding_dim, static_cast(dtype_), + padding_idx_.has_value() ? std::to_string(padding_idx_.value()) : "None"); +} + +Tensor Embedding::forward(const Tensor &indices) const { + // TODO: Implement on-device embedding for all devices, then remove the condition and the classic approach + auto device_type = device_.type(); + if (device_type != Device::Type::kCpu) { + // Use op::embedding which supports device-side input and batch dimension + return op::embedding(indices->contiguous()->to(device_), weight_); + } + + // Get the shape of indices + auto indices_shape = indices->shape(); + + // Output shape: indices_shape + [embedding_dim] + std::vector output_shape = indices_shape; + output_shape.push_back(embedding_dim_); + + // Create output tensor on the same device as weight + auto out = Tensor::empty(output_shape, weight_->dtype(), weight_->device()); + + // Flatten indices for sequential row copies + auto cpu_device = Device(Device::Type::kCpu, 0); + auto indices_cpu = indices->to(cpu_device)->contiguous(); + + // Calculate total number of lookups + size_t num_lookups = 1; + for (auto dim : indices_shape) { + num_lookups *= dim; + } + + const size_t row_bytes = embedding_dim_ * dsize(weight_->dtype()); + + // Helper lambda to read index based on dtype with bounds checking + auto read_index = [&](size_t i) -> int64_t { + auto dtype = indices_cpu->dtype(); + if (dtype == DataType::kInt32) { + const auto *data = reinterpret_cast(indices_cpu->data()); + return static_cast(data[i]); + } else if (dtype == DataType::kInt64) { + const auto *data = reinterpret_cast(indices_cpu->data()); + return data[i]; + } else if (dtype == DataType::kUInt32) { + const auto *data = reinterpret_cast(indices_cpu->data()); + return static_cast(data[i]); + } else if (dtype == DataType::kUInt64) { + const auto *data = reinterpret_cast(indices_cpu->data()); + uint64_t val = data[i]; + // Check if value can fit in int64_t + if (val > static_cast(std::numeric_limits::max())) { + throw std::out_of_range("Index value out of range for int64_t: " + std::to_string(val)); + } + return static_cast(val); + } else { + throw std::runtime_error("Embedding indices must be integer type, got dtype=" + std::to_string(static_cast(dtype))); + } + }; + + if (weight_->device().type() == Device::Type::kCpu) { + // CPU path: memcpy row by row + const auto *weight_base = reinterpret_cast(weight_->data()); + auto *out_base = reinterpret_cast(out->data()); + for (size_t i = 0; i < num_lookups; ++i) { + int64_t idx = read_index(i); + if (idx < 0 || idx >= static_cast(num_embeddings_)) { + throw std::out_of_range( + "Index out of range: " + std::to_string(idx) + " (num_embeddings=" + std::to_string(num_embeddings_) + ")"); + } + std::memcpy(out_base + i * row_bytes, weight_base + idx * row_bytes, row_bytes); + } + } else { + // Device fallback: copy rows through Tensor slices so device runtimes own stride/stream handling. + auto flat_out = out->view({num_lookups, embedding_dim_}); + for (size_t i = 0; i < num_lookups; ++i) { + int64_t idx = read_index(i); + if (idx < 0 || idx >= static_cast(num_embeddings_)) { + throw std::out_of_range( + "Index out of range: " + std::to_string(idx) + " (num_embeddings=" + std::to_string(num_embeddings_) + ")"); + } + auto dst = flat_out->narrow({{0, i, 1}}); + auto src = weight_->narrow({{0, static_cast(idx), 1}}); + dst->copy_from(src); + } + } + + return out; +} + +std::string Embedding::extra_repr() const { + std::string repr = "Embedding(num_embeddings=" + std::to_string(num_embeddings_) + ", embedding_dim=" + std::to_string(embedding_dim_) + ", dtype=" + std::to_string(static_cast(dtype_)); + if (padding_idx_.has_value()) { + repr += ", padding_idx=" + std::to_string(padding_idx_.value()); + } + repr += ")"; + return repr; +} + +} // namespace infinicore::nn diff --git a/csrc/infinicore/src/nn/layer_norm.cc b/csrc/infinicore/src/nn/layer_norm.cc new file mode 100644 index 000000000..45d3b9cb2 --- /dev/null +++ b/csrc/infinicore/src/nn/layer_norm.cc @@ -0,0 +1,27 @@ +#include "infinicore/nn/layer_norm.hpp" +#include "infinicore/ops.hpp" +#include +#include + +namespace infinicore::nn { + +LayerNorm::LayerNorm(size_t normalized_shape, double eps, const DataType &dtype, const Device &device) + : normalized_shape_(normalized_shape), + eps_(eps), + dtype_(dtype) { + + device_ = device; + + INFINICORE_NN_PARAMETER_INIT(weight, ({normalized_shape}, dtype_, device)); + INFINICORE_NN_PARAMETER_INIT(bias, ({normalized_shape}, dtype_, device)); +} + +Tensor LayerNorm::forward(const Tensor &x) const { + return op::layer_norm(x, weight_, bias_, static_cast(eps_)); +} + +std::string LayerNorm::extra_repr() const { + return "LayerNorm(normalized_shape=" + std::to_string(normalized_shape_) + ", eps=" + std::to_string(eps_) + ", dtype=" + std::to_string(static_cast(dtype_)) + ")"; +} + +} // namespace infinicore::nn diff --git a/csrc/infinicore/src/nn/module.cc b/csrc/infinicore/src/nn/module.cc new file mode 100644 index 000000000..5ca255322 --- /dev/null +++ b/csrc/infinicore/src/nn/module.cc @@ -0,0 +1,185 @@ +#include "infinicore/nn/module.hpp" +#include +#include + +namespace infinicore::nn { +namespace { +Parameter get_state_dict_parameter( + const std::unordered_map &state_dict, + const std::string &name) { + auto it = state_dict.find(name); + if (it == state_dict.end()) { + throw std::runtime_error("Parameter '" + name + "' not found in module."); + } + return it->second; +} +} // namespace + +std::unordered_map Module::state_dict() const { + std::unordered_map result; + collect_all_parameters(result, ""); + return result; +} + +std::vector Module::state_dict_keys() const { + std::vector result; + collect_all_parameter_names(result, ""); + return result; +} + +void Module::load_state_dict(const std::unordered_map &_state_dict) { + load_state_dict_recursively(_state_dict, ""); +} + +void Module::load_parameter(const std::string &name, const Tensor ¶m) { + auto param_it = parameters_.find(name); + if (param_it != parameters_.end()) { + try { + param_it->second.load(param); + } catch (const std::exception &e) { + throw std::runtime_error("Error loading parameter '" + name + "'. \n" + e.what()); + } + return; + } + + std::shared_ptr matched_submodule; + std::string matched_prefix; + for (const auto &[sub_name, submodule] : submodules_) { + if (name.size() <= sub_name.size() || name.compare(0, sub_name.size(), sub_name) != 0 || name[sub_name.size()] != '.') { + continue; + } + if (sub_name.size() > matched_prefix.size()) { + matched_prefix = sub_name; + matched_submodule = submodule; + } + } + + if (matched_submodule) { + try { + matched_submodule->load_parameter(name.substr(matched_prefix.size() + 1), param); + } catch (const std::exception &e) { + throw std::runtime_error("Error loading parameter '" + name + "'. \n" + e.what()); + } + return; + } + + spdlog::debug("load_parameter: Parameter '{}' not found. Available direct params={}, submodules={}", + name, parameters_.size(), submodules_.size()); + throw std::runtime_error("Parameter '" + name + "' not found in module."); +} + +void Module::load_parameters_no_sync(const std::unordered_map ¶ms, bool strict) { + auto all_params = state_dict(); + for (const auto &[name, param] : params) { + auto it = all_params.find(name); + if (it == all_params.end()) { + if (strict) { + throw std::runtime_error("Parameter '" + name + "' not found in module."); + } + continue; + } + auto existing_param = it->second; + try { + existing_param.load_no_sync(param); + } catch (const std::exception &e) { + throw std::runtime_error("Error loading parameter '" + name + "'. \n" + e.what()); + } + } +} + +void Module::load_parameter_(const std::string &name, const Tensor ¶m) { + // This function only handles direct parameters (no hierarchical traversal) + auto it = parameters_.find(name); + if (it != parameters_.end()) { + auto existing_param = it->second; + try { + existing_param.load(param); + } catch (const std::exception &e) { + throw std::runtime_error("Error loading parameter '" + name + "'. \n" + e.what()); + } + return; + } + + // Parameter not found + spdlog::debug("load_parameter_: Parameter '{}' not found. Available: {} params", + name, parameters_.size()); + throw std::runtime_error("Parameter '" + name + "' not found in module."); +} + +void Module::load_parameter_from_blob(const std::string &name, const void *data) { + auto param = parameters_[name]; + param.load_blob(data); +} + +Tensor Module::register_parameter(const std::string &name, Parameter param) { + parameters_[name] = param; + return param; +} + +Tensor Module::register_buffer(const std::string &name, Parameter buffer) { + buffers_[name] = buffer; + return buffer; +} + +void Module::load_state_dict_recursively(const std::unordered_map &_state_dict, const std::string &prefix) { + // Load direct parameters with the given prefix + for (const auto &[param_name, param] : parameters_) { + std::string full_name = prefix.empty() ? param_name : prefix + "." + param_name; + auto it = _state_dict.find(full_name); + if (it != _state_dict.end()) { + load_parameter_(param_name, it->second); + } + } + + // Recursively load parameters from submodules with extended prefix + for (const auto &[sub_name, submodule] : submodules_) { + std::string sub_prefix = prefix.empty() ? sub_name : prefix + "." + sub_name; + submodule->load_state_dict_recursively(_state_dict, sub_prefix); + } +} + +void Module::collect_all_parameters(std::unordered_map &all_params, const std::string &prefix) const { + // Add direct parameters with the given prefix + for (const auto &[param_name, param] : parameters_) { + std::string full_name = prefix.empty() ? param_name : prefix + "." + param_name; + all_params[full_name] = param; + } + + // Recursively collect parameters from submodules with extended prefix + for (const auto &[sub_name, submodule] : submodules_) { + std::string sub_prefix = prefix.empty() ? sub_name : prefix + "." + sub_name; + submodule->collect_all_parameters(all_params, sub_prefix); + } +} + +void Module::collect_all_parameter_names(std::vector &all_names, const std::string &prefix) const { + for (const auto &[param_name, _] : parameters_) { + all_names.push_back(prefix.empty() ? param_name : prefix + "." + param_name); + } + + for (const auto &[sub_name, submodule] : submodules_) { + std::string sub_prefix = prefix.empty() ? sub_name : prefix + "." + sub_name; + submodule->collect_all_parameter_names(all_names, sub_prefix); + } +} + +std::unordered_map Module::modules_dict() const { + std::unordered_map result; + collect_all_modules(result, ""); + return result; +} + +void Module::collect_all_modules(std::unordered_map &out, const std::string &prefix) const { + // 记录当前模块（跳过根节点的空前缀，可按需改为 "root"） + if (!prefix.empty()) { + out[prefix] = const_cast(this); + } + + // 递归遍历子模块 + for (const auto &[name, sub] : submodules_) { + std::string sub_prefix = prefix.empty() ? name : prefix + "." + name; + sub->collect_all_modules(out, sub_prefix); + } +} + +} // namespace infinicore::nn diff --git a/csrc/infinicore/src/nn/parameter.cc b/csrc/infinicore/src/nn/parameter.cc new file mode 100644 index 000000000..476764b7a --- /dev/null +++ b/csrc/infinicore/src/nn/parameter.cc @@ -0,0 +1,106 @@ +#include "infinicore/nn/parameter.hpp" + +#include "infinicore/context/context.hpp" + +#include +#include + +namespace infinicore::nn { +Parameter::Parameter() + : Tensor() { +} + +inline Shape get_partition_shape_(const Shape &shape, Size tp_dim, Size tp_size, Size num_shards) { + if (tp_size <= 1) { + return shape; + } + Shape part_shape = shape; + if (tp_dim < shape.size()) { + Size partition_factor = (num_shards > 0) ? num_shards : tp_size; + if (shape[tp_dim] % partition_factor != 0) { + throw std::runtime_error("Tensor dimension " + std::to_string(tp_dim) + " with size " + std::to_string(shape[tp_dim]) + " is not divisible by " + (num_shards > 0 ? "num_shards " : "tp_size ") + std::to_string(partition_factor) + "."); + } + part_shape[tp_dim] = shape[tp_dim] / partition_factor; + } + return part_shape; +} + +Parameter::Parameter(const Tensor &tensor, Size tp_dim, Size tp_rank, Size tp_size, Size num_shards) : Tensor(tensor), tp_dim_(tp_dim), tp_rank_(tp_rank), tp_size_(tp_size), num_shards_(num_shards) { + if (tp_rank_ >= tp_size_) { + throw std::runtime_error("Tensor parallel rank " + std::to_string(tp_rank_) + " must be less than tensor parallel size " + std::to_string(tp_size_) + "."); + } +} + +Parameter::Parameter( + const Shape &shape, + const DataType &dtype, + const Device &device, + Size tp_dim, + Size tp_rank, + Size tp_size, + Size num_shards) + : Parameter(Tensor::empty(get_partition_shape_(shape, tp_dim, tp_size, num_shards), dtype, device, false), tp_dim, tp_rank, tp_size, num_shards) { +} + +Parameter::Parameter(const Parameter &other) + : Tensor(other), + tp_dim_(other.tp_dim_), + tp_rank_(other.tp_rank_), + tp_size_(other.tp_size_), + num_shards_(other.num_shards_) {} + +void Parameter::load_blob(const void *data) { + Shape expected_shape = Shape(impl_->shape()); + expected_shape[tp_dim_] *= tp_size_; + auto buffer = Tensor::empty(expected_shape, impl_->dtype(), Device(Device::Type::kCpu, 0), true); + std::memcpy(buffer->data(), data, buffer->nbytes()); + this->load(buffer); +} + +void Parameter::load(const Tensor &tensor) { + load_no_sync(tensor); + infinicore::context::syncStream(); +} + +void Parameter::load_no_sync(const Tensor &tensor) { + if (impl_->dtype() != tensor->dtype()) { + throw std::runtime_error("Dtype mismatch when loading tensor into parameter. Weight: " + impl_->info() + ", Tensor: " + tensor->info() + "."); + } + + Shape expected_shape = Shape(impl_->shape()); + + if (num_shards_ == 0 || num_shards_ >= tp_size_) { + expected_shape[tp_dim_] *= tp_size_; + + if (expected_shape != tensor->shape()) { + throw std::runtime_error("Shape mismatch when loading tensor into parameter. Weight: " + impl_->info() + ", Tensor: " + tensor->info() + "."); + } + if (tp_size_ > 1) { + impl_->copy_from(tensor->narrow({{tp_dim_, tp_rank_ * impl_->size(tp_dim_), impl_->size(tp_dim_)}})); + } else { + impl_->copy_from(tensor); + } + } else { + if (num_shards_ == 0) { + throw std::runtime_error("num_shards_ is 0 but entered new logic branch!"); + } + + Size replica_size = tp_size_ / num_shards_; + if (replica_size == 0) { + throw std::runtime_error("replica_size is 0! tp_size_=" + std::to_string(tp_size_) + ", num_shards_=" + std::to_string(num_shards_)); + } + + Size shard_id = tp_rank_ / replica_size; + Size shard_size = impl_->size(tp_dim_); + Size offset = shard_id * shard_size; + + expected_shape[tp_dim_] *= num_shards_; + + if (offset + shard_size > tensor->shape()[tp_dim_]) { + throw std::runtime_error("Slice out of bounds! offset=" + std::to_string(offset) + ", shard_size=" + std::to_string(shard_size) + ", tensor_dim=" + std::to_string(tensor->shape()[tp_dim_])); + } + + impl_->copy_from(tensor->narrow({{tp_dim_, offset, shard_size}})); + } +} +} // namespace infinicore::nn diff --git a/csrc/infinicore/src/nn/rmsnorm.cc b/csrc/infinicore/src/nn/rmsnorm.cc new file mode 100644 index 000000000..a18e183f2 --- /dev/null +++ b/csrc/infinicore/src/nn/rmsnorm.cc @@ -0,0 +1,47 @@ +#include "infinicore/nn/rmsnorm.hpp" +#include "infinicore/ops.hpp" +#include +#include + +namespace infinicore::nn { + +RMSNorm::RMSNorm(size_t normalized_shape, double eps, const DataType &dtype, const Device &device) + : normalized_shape_(normalized_shape), + eps_(eps), + dtype_(dtype) { + + device_ = device; + + INFINICORE_NN_PARAMETER_INIT(weight, ({normalized_shape}, dtype_, device)); +} + +Tensor RMSNorm::forward(const Tensor &x) const { + // Delegate to the InfiniOps-backed operation. + return op::rms_norm(x, weight_, static_cast(eps_)); +} + +void RMSNorm::forward_inplace(Tensor &x, Tensor &residual) const { + if (!residual) { + residual = x; + x = op::rms_norm(x, weight_, static_cast(eps_)); + } else { + if (device_.type() == Device::Type::kCpu + || device_.type() == Device::Type::kNvidia + || device_.type() == Device::Type::kIluvatar + || device_.type() == Device::Type::kMetax + || device_.type() == Device::Type::kMoore + || device_.type() == Device::Type::kCambricon + || device_.type() == Device::Type::kHygon) { + op::add_rms_norm_inplace(x, residual, weight_, static_cast(eps_)); + } else { + op::add_(residual, x, residual); + op::rms_norm_(x, residual, weight_, static_cast(eps_)); + } + } +} + +std::string RMSNorm::extra_repr() const { + return "RMSNorm(normalized_shape=" + std::to_string(normalized_shape_) + ", eps=" + std::to_string(eps_) + ", dtype=" + std::to_string(static_cast(dtype_)) + ")"; +} + +} // namespace infinicore::nn diff --git a/csrc/infinicore/src/nn/rope.cc b/csrc/infinicore/src/nn/rope.cc new file mode 100644 index 000000000..e88e4e732 --- /dev/null +++ b/csrc/infinicore/src/nn/rope.cc @@ -0,0 +1,266 @@ +#include "infinicore/nn/rope.hpp" +#include "../../utils/custom_types.h" +#include "../utils.hpp" +#include "infinicore/ops/mrope.hpp" +#include "infinicore/ops/rope.hpp" +#include +#include +#include +#include +#include +#include +#include + +namespace infinicore::nn { + +RoPE::RoPE(size_t head_dim, + size_t rotary_dim, + size_t max_seq_len, + double theta, + Algo algo, + const DataType &dtype, + const Device &device, + std::shared_ptr scaling, + std::optional> mrope_section, + bool mrope_interleaved) + : rotary_dim_(rotary_dim), + head_dim_(head_dim), + max_seq_len_(max_seq_len), + theta_(theta), + algo_(algo), + dtype_(dtype), + scaling_(scaling), + mrope_section_(mrope_section), + mrope_interleaved_(mrope_interleaved) { + if (rotary_dim % 2 != 0) { + throw std::invalid_argument("rotary_dim must be even for RoPE, got " + std::to_string(rotary_dim)); + } + assert((rotary_dim > 0) && (rotary_dim <= head_dim_)); + if (mrope_section_.has_value()) { + const auto §ion = mrope_section_.value(); + if (section.size() != 3 || section[0] <= 0 || section[1] <= 0 || section[2] <= 0) { + throw std::invalid_argument("mrope_section must contain 3 positive values"); + } + if (2 * static_cast(section[0] + section[1] + section[2]) != rotary_dim_) { + throw std::invalid_argument("MRoPE section sum must equal rotary_dim / 2"); + } + } + device_ = device; + + // Initialize cache tables + initialize_cache(); +} + +void RoPE::initialize_cache() { + size_t cache_dim = rotary_dim_ / 2; + + // Create sin and cos cache tables: [max_seq_len, cache_dim] + INFINICORE_NN_BUFFER_INIT(sin_cache, ({max_seq_len_, cache_dim}, dtype_, device_)); + INFINICORE_NN_BUFFER_INIT(cos_cache, ({max_seq_len_, cache_dim}, dtype_, device_)); + + // Pre-compute sin and cos values + // Frequency generation always uses GPT-J style (theta^(-2j/rotary_dim)). + // The rotation algorithm (algo_) controls how dimensions are paired in the kernel. + + // Compute on CPU first, then copy to device + auto cpu_device = Device(Device::Type::kCpu, 0); + + // Allocate CPU buffers + std::vector sin_data(max_seq_len_ * cache_dim); + std::vector cos_data(max_seq_len_ * cache_dim); + + for (size_t pos = 0; pos < max_seq_len_; pos++) { + for (size_t dim_idx = 0; dim_idx < cache_dim; dim_idx++) { + // 1. Base inverse frequency (shared across all RoPE types) + float base_inv_freq = 1.0f / std::pow(static_cast(theta_), 2.0f * static_cast(dim_idx) / static_cast(rotary_dim_)); + + // 2. Polymorphic scaling resolution + // Passing pre-computed base_inv_freq avoids redundant pow() calculations in subclasses + float freq_scale = scaling_ ? scaling_->get_freq_scale(pos, dim_idx, base_inv_freq) : 1.0f; + float mag_scale = scaling_ ? scaling_->get_magnitude_scale(pos, dim_idx, base_inv_freq) : 1.0f; + + // 3. Compute final angle and sin/cos values + float angle = static_cast(pos) * base_inv_freq * freq_scale; + + sin_data[pos * cache_dim + dim_idx] = std::sin(angle) * mag_scale; + cos_data[pos * cache_dim + dim_idx] = std::cos(angle) * mag_scale; + } + } + + // Convert to target dtype on CPU (matching Python's numpy astype conversion pattern) + // Python: np_array.astype(ml_dtypes.bfloat16, copy=True) converts F32 -> BF16 + if (dtype_ == DataType::kFloat32) { + // Direct use of F32 data + auto sin_f32_cpu = Tensor::from_blob(sin_data.data(), {max_seq_len_, cache_dim}, DataType::kFloat32, cpu_device); + auto cos_f32_cpu = Tensor::from_blob(cos_data.data(), {max_seq_len_, cache_dim}, DataType::kFloat32, cpu_device); + sin_cache_->copy_from(sin_f32_cpu); + cos_cache_->copy_from(cos_f32_cpu); + } else if (dtype_ == DataType::kBFloat16) { + // Convert F32 to BF16 using the same conversion as Python's ml_dtypes.bfloat16 + // This uses round-to-nearest-even (matching _f32_to_bf16 implementation) + std::vector sin_bf16_data(max_seq_len_ * cache_dim); + std::vector cos_bf16_data(max_seq_len_ * cache_dim); + + for (size_t i = 0; i < sin_data.size(); i++) { + sin_bf16_data[i] = utils::cast(sin_data[i]); + cos_bf16_data[i] = utils::cast(cos_data[i]); + } + + auto sin_bf16_cpu = Tensor::from_blob(sin_bf16_data.data(), {max_seq_len_, cache_dim}, DataType::kBFloat16, cpu_device); + auto cos_bf16_cpu = Tensor::from_blob(cos_bf16_data.data(), {max_seq_len_, cache_dim}, DataType::kBFloat16, cpu_device); + + // copy_from handles cross-device copying to target device + sin_cache_->copy_from(sin_bf16_cpu); + cos_cache_->copy_from(cos_bf16_cpu); + } else if (dtype_ == DataType::kFloat16) { + // Convert F32 to F16 + std::vector sin_f16_data(max_seq_len_ * cache_dim); + std::vector cos_f16_data(max_seq_len_ * cache_dim); + + for (size_t i = 0; i < sin_data.size(); i++) { + sin_f16_data[i] = utils::cast(sin_data[i]); + cos_f16_data[i] = utils::cast(cos_data[i]); + } + + auto sin_f16_cpu = Tensor::from_blob(sin_f16_data.data(), {max_seq_len_, cache_dim}, DataType::kFloat16, cpu_device); + auto cos_f16_cpu = Tensor::from_blob(cos_f16_data.data(), {max_seq_len_, cache_dim}, DataType::kFloat16, cpu_device); + + sin_cache_->copy_from(sin_f16_cpu); + cos_cache_->copy_from(cos_f16_cpu); + } else { + throw std::runtime_error( + "RoPE cache dtype conversion not yet supported for dtype: " + + std::to_string(static_cast(dtype_))); + } +} + +Tensor RoPE::forward(const Tensor &x, const Tensor &pos, bool in_place) const { + if (mrope_section_.has_value()) { + throw std::runtime_error("MRoPE single-tensor forward is not implemented; use fused forward(q, k, positions) instead"); + } + Tensor y; + if (in_place) { + y = Tensor(x); + } else { + y = Tensor::empty(x->shape(), x->dtype(), x->device()); + if (rotary_dim_ < head_dim_) { + y->copy_from(x); + } + } + + size_t ndim = x->ndim(); + op::rope_(y->narrow({{ndim - 1, 0, rotary_dim_}}), + x->narrow({{ndim - 1, 0, rotary_dim_}}), + pos, sin_cache_, cos_cache_, algo_); + return y; +} + +static Tensor mrope_flatten_input(const Tensor &x, size_t head_dim, const char *name) { + if (x->ndim() == 2) { + if (x->size(1) % head_dim != 0) { + throw std::runtime_error(std::string("MRoPE expects ") + name + " hidden size to be a multiple of head_dim"); + } + return x; + } + if (x->ndim() == 3 && x->size(2) == head_dim) { + return x->view({x->size(0), x->size(1) * head_dim}); + } + throw std::runtime_error(std::string("MRoPE expects ") + name + " with shape [num_tokens, num_heads * head_dim] or [num_tokens, num_heads, head_dim]"); +} + +static Tensor mrope_flatten_output(const Tensor &x, size_t head_dim, const char *name) { + if (x->ndim() == 2) { + if (x->size(1) % head_dim != 0) { + throw std::runtime_error(std::string("MRoPE expects ") + name + " hidden size to be a multiple of head_dim"); + } + return x->view({x->size(0), x->size(1)}); + } + if (x->ndim() == 3 && x->size(2) == head_dim) { + return x->view({x->size(0), x->size(1) * head_dim}); + } + throw std::runtime_error(std::string("MRoPE expects ") + name + " with shape [num_tokens, num_heads * head_dim] or [num_tokens, num_heads, head_dim]"); +} + +std::pair RoPE::forward(const Tensor &q, const Tensor &k, const Tensor &positions) const { + if (!mrope_section_.has_value()) { + auto q_out = Tensor::empty(q->shape(), q->dtype(), q->device()); + auto k_out = Tensor::empty(k->shape(), k->dtype(), k->device()); + return forward(q_out, k_out, q, k, positions); + } + auto q_flat = mrope_flatten_input(q, head_dim_, "q"); + auto k_flat = mrope_flatten_input(k, head_dim_, "k"); + auto q_out = Tensor::empty(q_flat->shape(), q_flat->dtype(), q_flat->device()); + auto k_out = Tensor::empty(k_flat->shape(), k_flat->dtype(), k_flat->device()); + const auto §ion = mrope_section_.value(); + op::mrope_(q_out, + k_out, + q_flat, + k_flat, + cos_cache_, + sin_cache_, + positions, + static_cast(head_dim_), + static_cast(rotary_dim_), + section[0], + section[1], + section[2], + mrope_interleaved_); + return {q_out->view(q->shape()), k_out->view(k->shape())}; +} + +std::pair RoPE::forward(const Tensor &q_out, + const Tensor &k_out, + const Tensor &q, + const Tensor &k, + const Tensor &positions) const { + if (!mrope_section_.has_value()) { + auto apply_standard = [this, &positions](Tensor out, const Tensor &in) { + if (rotary_dim_ < head_dim_) { + out->copy_from(in); + } + size_t ndim = in->ndim(); + op::rope_(out->narrow({{ndim - 1, 0, rotary_dim_}}), + in->narrow({{ndim - 1, 0, rotary_dim_}}), + positions, + sin_cache_, + cos_cache_, + algo_); + }; + apply_standard(q_out, q); + apply_standard(k_out, k); + return {q_out, k_out}; + } + auto q_flat = mrope_flatten_input(q, head_dim_, "q"); + auto k_flat = mrope_flatten_input(k, head_dim_, "k"); + auto q_out_flat = mrope_flatten_output(q_out, head_dim_, "q_out"); + auto k_out_flat = mrope_flatten_output(k_out, head_dim_, "k_out"); + const auto §ion = mrope_section_.value(); + op::mrope_(q_out_flat, + k_out_flat, + q_flat, + k_flat, + cos_cache_, + sin_cache_, + positions, + static_cast(head_dim_), + static_cast(rotary_dim_), + section[0], + section[1], + section[2], + mrope_interleaved_); + return {q_out, k_out}; +} + +std::string RoPE::extra_repr() const { + std::string algo_str = (algo_ == Algo::GPT_J) ? "GPT_J" : "GPT_NEOX"; + std::string repr = "RoPE(head_dim=" + std::to_string(head_dim_) + ", rotary_dim=" + std::to_string(rotary_dim_) + ", max_seq_len=" + std::to_string(max_seq_len_) + ", theta=" + std::to_string(theta_) + ", algo=" + algo_str + ", dtype=" + std::to_string(static_cast(dtype_)); + if (mrope_section_.has_value()) { + const auto §ion = mrope_section_.value(); + repr += ", mrope_section=[" + std::to_string(section[0]) + "," + std::to_string(section[1]) + "," + std::to_string(section[2]) + "]"; + repr += ", mrope_interleaved=" + std::string(mrope_interleaved_ ? "true" : "false"); + } + repr += ")"; + return repr; +} + +} // namespace infinicore::nn diff --git a/csrc/infinicore/src/nn/rope_scaling_configs.cc b/csrc/infinicore/src/nn/rope_scaling_configs.cc new file mode 100644 index 000000000..d0901e7d3 --- /dev/null +++ b/csrc/infinicore/src/nn/rope_scaling_configs.cc @@ -0,0 +1,193 @@ +#include "infinicore/nn/rope_scaling_configs.hpp" +#include +#include +#include +#include +#include + +namespace infinicore::nn { + +namespace { +// Define a portable PI constant to avoid relying on the non-standard M_PI macro +// which is missing on MSVC (Windows) by default. +constexpr float kPi = 3.14159265358979323846f; +} // anonymous namespace + +// LongRopeScalingConfig Implementation +LongRopeScalingConfig::LongRopeScalingConfig( + std::vector short_factor, + std::vector long_factor, + size_t original_max_position_embeddings, + float factor) + : short_factor_(std::move(short_factor)), + long_factor_(std::move(long_factor)), + original_max_position_embeddings_(original_max_position_embeddings), + factor_(factor == 1.0f ? 1.0f : std::sqrt(1 + std::log(factor) / std::log(original_max_position_embeddings))) {} + +float LongRopeScalingConfig::get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const { + float _ext = (pos < original_max_position_embeddings_) ? short_factor_[dim_idx] : long_factor_[dim_idx]; + // The base inv_freq is multiplied by this scale. + // Original: inv_freq = 1.0f / (_ext * pow(theta, 2j/head_dim)) + // New: inv_freq = base_inv_freq * (1.0f / _ext) + return 1.0f / _ext; +} + +float LongRopeScalingConfig::get_magnitude_scale(size_t pos, size_t dim_idx, float base_inv_freq) const { + return factor_; +} + +// Llama3RopeScalingConfig Implementation +Llama3RopeScalingConfig::Llama3RopeScalingConfig( + float factor, + float low_freq_factor, + float high_freq_factor, + size_t original_max_position_embeddings) + : factor_(factor), + low_freq_factor_(low_freq_factor), + high_freq_factor_(high_freq_factor), + original_max_position_embeddings_(original_max_position_embeddings) {} + +float Llama3RopeScalingConfig::get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const { + // Calculate the wavelength corresponding to the current inverse frequency + float wavelen = 2.0f * static_cast(kPi) / base_inv_freq; + + // Compute the wavelength thresholds that separate high, mid, and low frequencies + float low_freq_wavelen = static_cast(original_max_position_embeddings_) / low_freq_factor_; + float high_freq_wavelen = static_cast(original_max_position_embeddings_) / high_freq_factor_; + + float scale = 1.0f; + + if (wavelen < low_freq_wavelen) { + // High-frequency band: short wavelengths retain the original scale + scale = 1.0f; + } else if (wavelen > high_freq_wavelen) { + // Low-frequency band: long wavelengths are directly scaled by the factor + scale = factor_; + } else { + // Mid-frequency band: apply smooth linear interpolation between 1.0 and factor_ + float smooth = (static_cast(original_max_position_embeddings_) / wavelen - low_freq_factor_) / (high_freq_factor_ - low_freq_factor_); + scale = 1.0f - smooth + smooth * factor_; + } + + // The framework applies the scale multiplicatively (inv_freq = base_inv_freq * return_value). + // Since the Llama3 logic divides the frequency (inv_freq = base_inv_freq / scale), + // we return the inverse of the computed scale. + return 1.0f / scale; +} + +namespace { + +float yarn_find_correction_dim( + int num_rotations, + size_t rotary_dim, + float base, + size_t original_max_position_embeddings) { + return (static_cast(rotary_dim) + * std::log(static_cast(original_max_position_embeddings) + / (static_cast(num_rotations) * 2.0f * kPi))) + / (2.0f * std::log(base)); +} + +std::pair yarn_find_correction_range( + int low_rot, + int high_rot, + size_t rotary_dim, + float base, + size_t original_max_position_embeddings, + bool truncate) { + float low = yarn_find_correction_dim( + low_rot, rotary_dim, base, original_max_position_embeddings); + float high = yarn_find_correction_dim( + high_rot, rotary_dim, base, original_max_position_embeddings); + if (truncate) { + low = std::floor(low); + high = std::ceil(high); + } + low = std::max(low, 0.0f); + high = std::min(high, static_cast(rotary_dim) - 1.0f); + return {low, high}; +} + +float yarn_get_mscale(float scale, float mscale_coeff = 1.0f) { + if (scale <= 1.0f) { + return 1.0f; + } + return 0.1f * mscale_coeff * std::log(scale) + 1.0f; +} + +} // anonymous namespace + +// YarnRopeScalingConfig Implementation +YarnRopeScalingConfig::YarnRopeScalingConfig( + float factor, + size_t original_max_position_embeddings, + size_t rotary_dim, + float rope_theta, + int beta_fast, + int beta_slow, + float mscale, + float mscale_all_dim) + : factor_(factor), + original_max_position_embeddings_(original_max_position_embeddings) { + if (factor <= 0.0f) { + throw std::invalid_argument( + "YarnRopeScalingConfig factor must be positive, got " + + std::to_string(factor)); + } + if (original_max_position_embeddings == 0) { + throw std::invalid_argument( + "YarnRopeScalingConfig original_max_position_embeddings must be positive"); + } + if (rope_theta <= 0.0f) { + throw std::invalid_argument( + "YarnRopeScalingConfig rope_theta must be positive, got " + + std::to_string(rope_theta)); + } + if (rotary_dim < 2 || rotary_dim % 2 != 0) { + throw std::invalid_argument( + "YarnRopeScalingConfig rotary_dim must be a positive even number, got " + + std::to_string(rotary_dim)); + } + + // vLLM: yarn_find_correction_range(beta_fast, beta_slow, ...) + auto [low, high] = yarn_find_correction_range( + beta_fast, + beta_slow, + rotary_dim, + rope_theta, + original_max_position_embeddings, + true); + correction_low_ = low; + correction_high_ = high; + + magnitude_scale_ = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim); +} + +float YarnRopeScalingConfig::yarn_linear_ramp(size_t dim_idx) const { + float low = correction_low_; + float high = correction_high_; + if (low == high) { + high += 0.001f; // Prevent singularity (matches vLLM) + } + float linear = (static_cast(dim_idx) - low) / (high - low); + return std::clamp(linear, 0.0f, 1.0f); +} + +float YarnRopeScalingConfig::get_freq_scale( + size_t /*pos*/, + size_t dim_idx, + float /*base_inv_freq*/) const { + constexpr float kExtrapolationFactor = 1.0f; + float ramp = yarn_linear_ramp(dim_idx); + float inv_freq_mask = (1.0f - ramp) * kExtrapolationFactor; + return (1.0f - inv_freq_mask) / factor_ + inv_freq_mask; +} + +float YarnRopeScalingConfig::get_magnitude_scale( + size_t /*pos*/, + size_t /*dim_idx*/, + float /*base_inv_freq*/) const { + return magnitude_scale_; +} + +} // namespace infinicore::nn diff --git a/csrc/infinicore/src/ops/acos/acos.cc b/csrc/infinicore/src/ops/acos/acos.cc new file mode 100644 index 000000000..049ae35e9 --- /dev/null +++ b/csrc/infinicore/src/ops/acos/acos.cc @@ -0,0 +1,24 @@ +#include "infinicore/ops/acos.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Acos::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Acos::execute(Tensor output, Tensor input) { + dispatcher().lookup(context::getDevice().type())(output, input); +} + +Tensor acos(Tensor input) { + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + acos_(output, input); + return output; +} + +void acos_(Tensor output, Tensor input) { + Acos::execute(output, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/adaptive_avg_pool1d/adaptive_avg_pool1d.cc b/csrc/infinicore/src/ops/adaptive_avg_pool1d/adaptive_avg_pool1d.cc new file mode 100644 index 000000000..e7833a195 --- /dev/null +++ b/csrc/infinicore/src/ops/adaptive_avg_pool1d/adaptive_avg_pool1d.cc @@ -0,0 +1,43 @@ +#include "infinicore/ops/adaptive_avg_pool1d.hpp" +#include +#include + +namespace infinicore::op { + +common::OpDispatcher &AdaptiveAvgPool1d::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void AdaptiveAvgPool1d::execute(Tensor output, Tensor input) { + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No AdaptiveAvgPool1d implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor adaptive_avg_pool1d(Tensor input, int64_t output_size) { + size_t ndim = input->ndim(); + if (ndim != 2 && ndim != 3) { + throw std::runtime_error("AdaptiveAvgPool1d: Input tensor must be 2D or 3D."); + } + + if (output_size <= 0) { + throw std::runtime_error("AdaptiveAvgPool1d: output_size must be positive."); + } + + auto out_shape = input->shape(); + out_shape[ndim - 1] = output_size; + + auto output = Tensor::empty(out_shape, input->dtype(), input->device()); + + AdaptiveAvgPool1d::execute(output, input); + + return output; +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/adaptive_avg_pool3d/adaptive_avg_pool3d.cc b/csrc/infinicore/src/ops/adaptive_avg_pool3d/adaptive_avg_pool3d.cc new file mode 100644 index 000000000..337522f91 --- /dev/null +++ b/csrc/infinicore/src/ops/adaptive_avg_pool3d/adaptive_avg_pool3d.cc @@ -0,0 +1,36 @@ +#include "infinicore/ops/adaptive_avg_pool3d.hpp" +#include +#include +namespace infinicore::op { + +common::OpDispatcher &AdaptiveAvgPool3D::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void AdaptiveAvgPool3D::execute(Tensor y, Tensor x) { + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + if (func == nullptr) { + throw std::runtime_error("No AdaptiveAvgPool3D implementation found for device type: " + std::to_string(static_cast(device_type))); + } + func(y, x); +} + +Tensor adaptive_avg_pool3d(Tensor x, std::vector output_size) { + + // Create output tensor shap + Shape y_shape = x->shape(); + y_shape[2] = output_size[0]; // D dimension + y_shape[3] = output_size[1]; // H dimension + y_shape[4] = output_size[2]; // W dimension + + auto y = Tensor::empty(y_shape, x->dtype(), x->device()); + adaptive_avg_pool3d_(y, x); + return y; +} + +void adaptive_avg_pool3d_(Tensor y, Tensor x) { + AdaptiveAvgPool3D::execute(y, x); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/adaptive_max_pool1d/adaptive_max_pool1d.cc b/csrc/infinicore/src/ops/adaptive_max_pool1d/adaptive_max_pool1d.cc new file mode 100644 index 000000000..3612496ae --- /dev/null +++ b/csrc/infinicore/src/ops/adaptive_max_pool1d/adaptive_max_pool1d.cc @@ -0,0 +1,30 @@ +#include "infinicore/ops/adaptive_max_pool1d.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &AdaptiveMaxPool1d::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void AdaptiveMaxPool1d::execute(Tensor y, Tensor x, size_t output_size) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + infinicore::context::setDevice(y->device()); + dispatcher().lookup(y->device().type())(y, x, output_size); +} + +Tensor adaptive_max_pool1d(Tensor x, size_t output_size) { + infinicore::Shape y_shape = x->shape(); + y_shape.back() = output_size; + auto y = Tensor::empty(y_shape, x->dtype(), x->device()); + adaptive_max_pool1d_(y, x, output_size); + return y; +} + +void adaptive_max_pool1d_(Tensor y, Tensor x, size_t output_size) { + AdaptiveMaxPool1d::execute(y, x, output_size); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/add/add.cc b/csrc/infinicore/src/ops/add/add.cc new file mode 100644 index 000000000..9aa7be15d --- /dev/null +++ b/csrc/infinicore/src/ops/add/add.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/add.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Add); + +Add::Add(Tensor c, const Tensor &a, const Tensor &b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + INFINICORE_GRAPH_OP_DISPATCH(c->device().type(), c, a, b); +} + +void Add::execute(Tensor c, const Tensor &a, const Tensor &b) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Add, c, a, b); +} + +Tensor add(const Tensor &a, const Tensor &b) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + add_(c, a, b); + return c; +} + +void add_(Tensor c, const Tensor &a, const Tensor &b) { + Add::execute(c, a, b); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/add/add_infiniops.cc b/csrc/infinicore/src/ops/add/add_infiniops.cc new file mode 100644 index 000000000..7c45f1ce6 --- /dev/null +++ b/csrc/infinicore/src/ops/add/add_infiniops.cc @@ -0,0 +1,59 @@ +#include "infinicore/ops/add.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +namespace infinicore::op::add_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +struct PlannedMeta { + TensorMeta c, a, b; + graph::GraphTensor c_tensor, a_tensor, b_tensor; +}; + +} // namespace + +void *plan(Tensor c, const Tensor &a, const Tensor &b) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(c->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + + return new PlannedMeta{ + TensorMeta(c), + TensorMeta(a), + TensorMeta(b), + graph::GraphTensor(c), + graph::GraphTensor(a), + graph::GraphTensor(b)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + infini::ops::Add::Call( + handle, + config, + planned->a.tensor(planned->a_tensor), + planned->b.tensor(planned->b_tensor), + planned->c.tensor(planned->c_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Add::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(Add::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(Add::cleanup_dispatcher(), &cleanup); + return true; +}(); + +} // namespace infinicore::op::add_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/add_rms_norm/add_rms_norm.cc b/csrc/infinicore/src/ops/add_rms_norm/add_rms_norm.cc new file mode 100644 index 000000000..64f45b893 --- /dev/null +++ b/csrc/infinicore/src/ops/add_rms_norm/add_rms_norm.cc @@ -0,0 +1,33 @@ +#include "infinicore/ops/add_rms_norm.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(AddRMSNorm); + +AddRMSNorm::AddRMSNorm(Tensor y, Tensor residual_out, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, residual_out, a, b, weight); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), y, residual_out, a, b, weight, epsilon); +} + +void AddRMSNorm::execute(Tensor y, Tensor residual_out, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(AddRMSNorm, y, residual_out, a, b, weight, epsilon); +} + +std::pair add_rms_norm(const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { + auto y = Tensor::empty(a->shape(), a->dtype(), a->device()); + auto residual_out = Tensor::empty(a->shape(), a->dtype(), a->device()); + add_rms_norm_(y, residual_out, a, b, weight, epsilon); + return std::make_pair(y, residual_out); +} + +void add_rms_norm_(Tensor out, Tensor residual, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { + AddRMSNorm::execute(out, residual, a, b, weight, epsilon); +} + +void add_rms_norm_inplace(Tensor input, Tensor residual, const Tensor &weight, float epsilon) { + add_rms_norm_(input, residual, input, residual, weight, epsilon); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/add_rms_norm/add_rms_norm_infiniops.cc b/csrc/infinicore/src/ops/add_rms_norm/add_rms_norm_infiniops.cc new file mode 100644 index 000000000..b06964e20 --- /dev/null +++ b/csrc/infinicore/src/ops/add_rms_norm/add_rms_norm_infiniops.cc @@ -0,0 +1,72 @@ +#include "infinicore/ops/add_rms_norm.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/add_rms_norm.h" + +#include + +namespace infinicore::op::add_rms_norm_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +struct PlannedMeta { + TensorMeta out, residual, a, b, weight; + graph::GraphTensor out_tensor, residual_tensor, a_tensor, b_tensor, weight_tensor; + float epsilon; +}; + +} // namespace + +void *plan(Tensor out, Tensor residual, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, residual, a, b, weight); + + return new PlannedMeta{ + TensorMeta(out), + TensorMeta(residual), + TensorMeta(a), + TensorMeta(b), + TensorMeta(weight), + graph::GraphTensor(out), + graph::GraphTensor(residual), + graph::GraphTensor(a), + graph::GraphTensor(b), + graph::GraphTensor(weight), + epsilon}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + infini::ops::AddRmsNorm::Call( + handle, + config, + planned->a.tensor(planned->a_tensor), + planned->b.tensor(planned->b_tensor), + planned->weight.tensor(planned->weight_tensor), + std::optional{planned->epsilon}, + planned->out.tensor(planned->out_tensor), + planned->residual.tensor(planned->residual_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(AddRMSNorm::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(AddRMSNorm::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(AddRMSNorm::cleanup_dispatcher(), &cleanup); + return true; +}(); + +} // namespace infinicore::op::add_rms_norm_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/addbmm/addbmm.cc b/csrc/infinicore/src/ops/addbmm/addbmm.cc new file mode 100644 index 000000000..999a4331f --- /dev/null +++ b/csrc/infinicore/src/ops/addbmm/addbmm.cc @@ -0,0 +1,32 @@ +#include "infinicore/ops/addbmm.hpp" +#include "../../utils.hpp" +#include "infinicore/ops/addbmm.hpp" + +namespace infinicore::op { + +// 1. 初始化 Dispatcher +common::OpDispatcher &Addbmm::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Addbmm::execute(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha) { + + // 切换上下文 + infinicore::context::setDevice(output->device()); + + // 分发计算 + dispatcher().lookup(output->device().type())(output, input, batch1, batch2, beta, alpha); +} + +Tensor addbmm(Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha) { + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + Addbmm::execute(output, input, batch1, batch2, beta, alpha); + return output; +} + +void addbmm_(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha) { + Addbmm::execute(output, input, batch1, batch2, beta, alpha); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/addcmul/addcmul.cc b/csrc/infinicore/src/ops/addcmul/addcmul.cc new file mode 100644 index 000000000..a5a805e09 --- /dev/null +++ b/csrc/infinicore/src/ops/addcmul/addcmul.cc @@ -0,0 +1,29 @@ +#include "infinicore/ops/addcmul.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Addcmul::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +// 执行核心逻辑：设备校验与后端分发 +void Addcmul::execute(Tensor out, Tensor input, Tensor t1, Tensor t2, float value) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, t1, t2); + infinicore::context::setDevice(out->device()); + dispatcher().lookup(out->device().type())(out, input, t1, t2, value); +} + +// Out-of-place 接口：自动创建输出 Tensor +Tensor addcmul(Tensor input, Tensor t1, Tensor t2, float value) { + auto out = Tensor::empty(input->shape(), input->dtype(), input->device()); + addcmul_(out, input, t1, t2, value); + return out; +} + +void addcmul_(Tensor out, Tensor input, Tensor t1, Tensor t2, float value) { + Addcmul::execute(out, input, t1, t2, value); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/addr/addr.cc b/csrc/infinicore/src/ops/addr/addr.cc new file mode 100644 index 000000000..51198655f --- /dev/null +++ b/csrc/infinicore/src/ops/addr/addr.cc @@ -0,0 +1,32 @@ +#include "infinicore/ops/addr.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Addr::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Addr::execute(Tensor out, Tensor input, Tensor vec1, Tensor vec2, float beta, float alpha) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, vec1, vec2); + infinicore::context::setDevice(out->device()); + dispatcher().lookup(out->device().type())(out, input, vec1, vec2, beta, alpha); +} + +Tensor addr(Tensor input, Tensor vec1, Tensor vec2, float beta, float alpha) { + + size_t n = vec1->shape()[0]; + size_t m = vec2->shape()[0]; + + // Create output tensor + Tensor out = Tensor::empty({n, m}, input->dtype(), input->device()); + addr_(out, input, vec1, vec2, beta, alpha); + return out; +} + +void addr_(Tensor out, Tensor input, Tensor vec1, Tensor vec2, float beta, float alpha) { + Addr::execute(out, input, vec1, vec2, beta, alpha); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/affine_grid/affine_grid.cc b/csrc/infinicore/src/ops/affine_grid/affine_grid.cc new file mode 100644 index 000000000..a33924f85 --- /dev/null +++ b/csrc/infinicore/src/ops/affine_grid/affine_grid.cc @@ -0,0 +1,58 @@ +#include "infinicore/ops/affine_grid.hpp" +#include +#include +#include + +namespace infinicore::op { + +common::OpDispatcher &AffineGrid::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void AffineGrid::execute(Tensor output, Tensor theta, bool align_corners) { + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No AffineGrid implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, theta, align_corners); +} + +Tensor affine_grid(Tensor theta, const std::vector &size, bool align_corners) { + if (theta->ndim() != 3) { + throw std::runtime_error("AffineGrid: Theta tensor must be 3D (N, 2, 3)."); + } + if (theta->shape()[1] != 2 || theta->shape()[2] != 3) { + throw std::runtime_error("AffineGrid: Theta tensor shape must be (N, 2, 3)."); + } + + if (size.size() != 4) { + throw std::runtime_error("AffineGrid: target size length must be 4 (N, C, H, W)."); + } + + if (static_cast(theta->shape()[0]) != size[0]) { + throw std::runtime_error("AffineGrid: Theta batch size does not match target size batch."); + } + + if (!theta->is_contiguous()) { + theta = theta->contiguous(); + } + + std::vector out_shape; + out_shape.reserve(4); + out_shape.push_back(static_cast(size[0])); + out_shape.push_back(static_cast(size[2])); + out_shape.push_back(static_cast(size[3])); + out_shape.push_back(2); + + auto output = Tensor::empty(out_shape, theta->dtype(), theta->device()); + + AffineGrid::execute(output, theta, align_corners); + + return output; +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/argwhere/argwhere.cc b/csrc/infinicore/src/ops/argwhere/argwhere.cc new file mode 100644 index 000000000..934a886e5 --- /dev/null +++ b/csrc/infinicore/src/ops/argwhere/argwhere.cc @@ -0,0 +1,32 @@ +#include "infinicore/ops/argwhere.hpp" +#include "../../utils.hpp" +#include "infinicore/dtype.hpp" +#include "infinicore/tensor.hpp" +#include + +namespace infinicore::op { + +common::OpDispatcher &Argwhere::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void Argwhere::execute(void **y, size_t *count, Tensor x) { + auto device_type = context::getDevice().type(); + + auto func = dispatcher().lookup(device_type); + if (func == nullptr) { + throw std::runtime_error("Argwhere op not implemented for device type " + std::to_string(static_cast(device_type))); + } + func(y, count, x); +} +Tensor argwhere(Tensor x) { + void *y = nullptr; + size_t count = 0; + Argwhere::execute(&y, &count, x); + auto result = Tensor::from_blob(y, Shape{count, x->ndim()}, DataType::kInt64, Device(Device::Type::kCpu)); + result = result->to(x->device()); + return result; +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/asin/asin.cc b/csrc/infinicore/src/ops/asin/asin.cc new file mode 100644 index 000000000..b4df3daaa --- /dev/null +++ b/csrc/infinicore/src/ops/asin/asin.cc @@ -0,0 +1,36 @@ +#include "infinicore/ops/asin.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &Asin::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; +void Asin::execute(Tensor output, Tensor input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Asin implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor asin(Tensor input) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + asin_(output, input); + return output; +} + +void asin_(Tensor output, Tensor input) { + Asin::execute(output, input); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/asinh/asinh.cc b/csrc/infinicore/src/ops/asinh/asinh.cc new file mode 100644 index 000000000..1a0730b83 --- /dev/null +++ b/csrc/infinicore/src/ops/asinh/asinh.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/asinh.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Asinh::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Asinh::execute(Tensor y, Tensor x) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + infinicore::context::setDevice(y->device()); + dispatcher().lookup(y->device().type())(y, x); +} + +Tensor asinh(Tensor x) { + auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); + asinh_(y, x); + return y; +} + +void asinh_(Tensor y, Tensor x) { + Asinh::execute(y, x); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/asum/asum.cc b/csrc/infinicore/src/ops/asum/asum.cc new file mode 100644 index 000000000..11b24f485 --- /dev/null +++ b/csrc/infinicore/src/ops/asum/asum.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/asum.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Asum); + +Asum::Asum(const Tensor &x, Tensor result) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, result); + INFINICORE_GRAPH_OP_DISPATCH(result->device().type(), x, result); +} + +void Asum::execute(const Tensor &x, Tensor result) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Asum, x, result); +} + +Tensor asum(const Tensor &x) { + auto result = Tensor::empty({}, x->dtype(), x->device()); + asum_(x, result); + return result; +} + +void asum_(const Tensor &x, Tensor result) { + Asum::execute(x, result); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/atanh/atanh.cc b/csrc/infinicore/src/ops/atanh/atanh.cc new file mode 100644 index 000000000..d7e1daade --- /dev/null +++ b/csrc/infinicore/src/ops/atanh/atanh.cc @@ -0,0 +1,37 @@ +#include "infinicore/ops/atanh.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +// 获取单例分发器 +common::OpDispatcher &Atanh::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +// 执行入口：负责设备切换和后端查找 +void Atanh::execute(Tensor y, Tensor a) { + // 确保输入和输出在同一个设备上 + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, a); + + // 切换当前上下文到目标设备 + infinicore::context::setDevice(y->device()); + + // 根据设备类型（CPU/CUDA等）查找对应的实现并执行 + dispatcher().lookup(y->device().type())(y, a); +} + +// Out-of-place 接口：自动创建结果 Tensor +Tensor atanh(Tensor a) { + // 创建一个与输入形状、类型、设备完全相同的空 Tensor + auto y = Tensor::empty(a->shape(), a->dtype(), a->device()); + atanh_(y, a); + return y; +} + +// In-place 或指定输出接口 +void atanh_(Tensor y, Tensor a) { + Atanh::execute(y, a); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/attention/attention.cc b/csrc/infinicore/src/ops/attention/attention.cc new file mode 100644 index 000000000..eafc48f95 --- /dev/null +++ b/csrc/infinicore/src/ops/attention/attention.cc @@ -0,0 +1,40 @@ +#include "infinicore/ops/attention.hpp" +#include "../../utils.hpp" + +#ifdef ENABLE_MUTUAL_AWARENESS +#include "infinicore/analyzer/mutual_awareness_analyzer.hpp" +#endif + +namespace infinicore::op { + +common::OpDispatcher &Attention::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Attention::execute(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k, v, k_cache, v_cache); + infinicore::context::setDevice(out->device()); +#ifdef ENABLE_MUTUAL_AWARENESS + auto goal = analyzer::MutualAwarenessAnalyzer::instance().getCurrentOptimizationGoal(); + dispatcher().lookup(out->device().type(), goal)(out, q, k, v, k_cache, v_cache, pos); +#else + dispatcher().lookup(out->device().type())(out, q, k, v, k_cache, v_cache, pos); +#endif +} + +Tensor attention(Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos) { + size_t n_q_head = q->shape()[0]; + size_t seq_len = q->shape()[1]; + size_t head_dim = q->shape()[2]; + Shape shape = {seq_len, n_q_head, head_dim}; + auto out = Tensor::empty(shape, q->dtype(), q->device()); + attention_(out, q, k, v, k_cache, v_cache, pos); + return out; +} + +void attention_(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos) { + Attention::execute(out, q, k, v, k_cache, v_cache, pos); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/avg_pool1d/avg_pool1d.cc b/csrc/infinicore/src/ops/avg_pool1d/avg_pool1d.cc new file mode 100644 index 000000000..0e4ded14f --- /dev/null +++ b/csrc/infinicore/src/ops/avg_pool1d/avg_pool1d.cc @@ -0,0 +1,68 @@ +#include "infinicore/ops/avg_pool1d.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &AvgPool1d::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void AvgPool1d::execute( + Tensor output, + Tensor input, + size_t kernel_size, + size_t stride, + size_t padding) { + + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + if (stride == 0) { + stride = kernel_size; + } + + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error( + "No AvgPool1d implementation for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input, kernel_size, stride, padding); +} + +Tensor avg_pool1d(Tensor input, size_t kernel_size, size_t stride, size_t padding) { + if (stride == 0) { + stride = kernel_size; + } + + const auto &shape = input->shape(); + if (shape.size() != 3) { + throw std::runtime_error("AvgPool1d expects tensors with shape [N, C, L]"); + } + + const size_t n = shape[0]; + const size_t c = shape[1]; + const size_t l_in = shape[2]; + + if (l_in + 2 * padding < kernel_size) { + throw std::runtime_error("AvgPool1d kernel_size is larger than padded length"); + } + + const size_t out_width = (l_in + 2 * padding - kernel_size) / stride + 1; + + Shape out_shape = {n, c, out_width}; + auto output = Tensor::empty(out_shape, input->dtype(), input->device()); + avg_pool1d_(output, input, kernel_size, stride, padding); + return output; +} + +void avg_pool1d_(Tensor output, Tensor input, size_t kernel_size, size_t stride, size_t padding) { + AvgPool1d::execute(output, input, kernel_size, stride, padding); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/axpy/axpy.cc b/csrc/infinicore/src/ops/axpy/axpy.cc new file mode 100644 index 000000000..63d228c2c --- /dev/null +++ b/csrc/infinicore/src/ops/axpy/axpy.cc @@ -0,0 +1,22 @@ +#include "infinicore/ops/axpy.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Axpy); + +Axpy::Axpy(const Tensor &alpha, const Tensor &x, Tensor y) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(alpha, x, y); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), alpha, x, y); +} + +void Axpy::execute(const Tensor &alpha, const Tensor &x, Tensor y) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Axpy, alpha, x, y); +} + +void axpy_(const Tensor &alpha, const Tensor &x, Tensor y) { + Axpy::execute(alpha, x, y); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/baddbmm/baddbmm.cc b/csrc/infinicore/src/ops/baddbmm/baddbmm.cc new file mode 100644 index 000000000..c4f8efaf5 --- /dev/null +++ b/csrc/infinicore/src/ops/baddbmm/baddbmm.cc @@ -0,0 +1,100 @@ +#include "infinicore/ops/baddbmm.hpp" +#include "infinicore/ops/gemm.hpp" +#include "infinicore/ops/rearrange.hpp" + +namespace infinicore::op { + +// 内联的 BLAS 兼容性检查，减少函数调用开销 +inline bool is_blas_compatible(const Tensor &t) { + const auto ndim = t->ndim(); + if (ndim == 2) { + const auto rs = t->stride(0); + const auto cs = t->stride(1); + if (rs != 1 && cs != 1) { + return false; + } + if (rs == 1 && cs == 1) { + return t->shape()[0] == 1 || t->shape()[1] == 1; + } + return true; + } else if (ndim == 3) { + const auto rs = t->stride(1); + const auto cs = t->stride(2); + if (t->shape()[0] > 1 && t->stride(0) == 0) { + return false; + } + if (rs != 1 && cs != 1) { + return false; + } + if (rs == 1 && cs == 1) { + return t->shape()[1] == 1 || t->shape()[2] == 1; + } + return true; + } + return false; +} + +inline void prepare_gemm_input(Tensor &output, Tensor &input, const size_t batch_size, const size_t m, const size_t n) { + const auto input_ndim = input->ndim(); + if (input_ndim == 2) { + rearrange_(output, input->as_strided( + {batch_size, m, n}, + {0, input->stride(0), input->stride(1)})); + } else if (input_ndim == 3 && input->shape()[0] == 1 && batch_size > 1) { + rearrange_(output, input->as_strided( + {batch_size, m, n}, + {0, input->stride(1), input->stride(2)})); + } else { + rearrange_(output, input); + } +} + +Tensor baddbmm(Tensor input, Tensor batch1, Tensor batch2, + float beta, + float alpha) { + const size_t batch_size = batch1->shape()[0]; + const size_t m = batch1->shape()[1]; + const size_t n = batch2->shape()[2]; + + const Tensor &a = is_blas_compatible(batch1) ? batch1 : rearrange(batch1); + const Tensor &b = is_blas_compatible(batch2) ? batch2 : rearrange(batch2); + + if (beta == 0.0f) { + return gemm(a, b, alpha, 0.0f); + } + + Tensor result = Tensor::empty({batch_size, m, n}, a->dtype(), a->device()); + + prepare_gemm_input(result, input, batch_size, m, n); + + gemm_(result, a, b, alpha, beta); + return result; +} + +void baddbmm_(Tensor out, Tensor input, Tensor batch1, Tensor batch2, + float beta, + float alpha) { + const size_t batch_size = batch1->shape()[0]; + const size_t m = batch1->shape()[1]; + const size_t n = batch2->shape()[2]; + + const Tensor &a = is_blas_compatible(batch1) ? batch1 : rearrange(batch1); + const Tensor &b = is_blas_compatible(batch2) ? batch2 : rearrange(batch2); + + const bool out_is_usable = out->is_contiguous() && out->ndim() == 3 && out->shape()[0] == batch_size && out->shape()[1] == m && out->shape()[2] == n; + + if (out_is_usable) { + if (beta != 0.0f && input->data() != out->data()) { + prepare_gemm_input(out, input, batch_size, m, n); + } + gemm_(out, a, b, alpha, beta); + } else { + Tensor result = Tensor::empty({batch_size, m, n}, a->dtype(), a->device()); + if (beta != 0.0f) { + prepare_gemm_input(result, input, batch_size, m, n); + } + gemm_(result, a, b, alpha, beta); + rearrange_(out, result); + } +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/bilinear/bilinear.cc b/csrc/infinicore/src/ops/bilinear/bilinear.cc new file mode 100644 index 000000000..930f9a31a --- /dev/null +++ b/csrc/infinicore/src/ops/bilinear/bilinear.cc @@ -0,0 +1,95 @@ +#include "infinicore/ops/bilinear.hpp" +#include "infinicore/ops/add.hpp" +#include "infinicore/ops/matmul.hpp" +#include "infinicore/ops/rearrange.hpp" + +namespace infinicore::op { + +namespace { +inline bool is_gemm_compatible_3d(const Tensor &t) { + if (t->ndim() != 3) { + return false; + } + + const auto batch = t->shape()[0]; + const auto rows = t->shape()[1]; + const auto cols = t->shape()[2]; + const auto bs = t->stride(0); + const auto rs = t->stride(1); + const auto cs = t->stride(2); + + if (rs != 1 && cs != 1) { + return false; + } + + if (cs == 1) { + if (rs < static_cast(cols)) { + return false; + } + } else { + if (cs < static_cast(rows)) { + return false; + } + } + + if (batch > 1 && bs == 0) { + return false; + } + + return true; +} + +inline Tensor ensure_gemm_compatible(const Tensor &t) { + if (t->ndim() == 2) { + return t->is_contiguous() ? t : rearrange(t); + } else if (t->ndim() == 3) { + return is_gemm_compatible_3d(t) ? t : rearrange(t); + } + return t->is_contiguous() ? t : rearrange(t); +} + +} // anonymous namespace + +Tensor bilinear(Tensor x1, Tensor x2, Tensor weight, std::optional bias) { + const size_t batch_size = x1->shape()[0]; + const size_t in1_features = x1->shape()[1]; + const size_t in2_features = x2->shape()[1]; + const size_t out_features = weight->shape()[0]; + + Tensor x1_compat = ensure_gemm_compatible(x1); + Tensor x2_compat = ensure_gemm_compatible(x2); + Tensor weight_cont = weight->is_contiguous() ? weight : weight->contiguous(); + + Tensor weight_permuted = weight_cont->permute({1, 0, 2}); + Tensor weight_permuted_cont = weight_permuted->is_contiguous() + ? weight_permuted + : weight_permuted->contiguous(); + Tensor weight_matrix = weight_permuted_cont->view({in1_features, out_features * in2_features}); + + Tensor intermediate = matmul(x1_compat, weight_matrix, 1.0f); + + Tensor intermediate_3d = intermediate->view({batch_size, out_features, in2_features}); + Tensor intermediate_transposed = intermediate_3d->permute({0, 2, 1}); + Tensor intermediate_compat = ensure_gemm_compatible(intermediate_transposed); + + Tensor x2_row = x2_compat->view({batch_size, 1, in2_features}); + Tensor x2_row_compat = ensure_gemm_compatible(x2_row); + + Tensor out_3d = matmul(x2_row_compat, intermediate_compat, 1.0f); + Tensor out = out_3d->view({batch_size, out_features}); + + if (bias) { + Tensor bias_broadcast = (*bias)->as_strided( + {batch_size, out_features}, + {0, (*bias)->strides()[0]}); + out = add(out, bias_broadcast); + } + return out; +} + +void bilinear_(Tensor out, Tensor x1, Tensor x2, Tensor weight, std::optional bias) { + Tensor result = bilinear(x1, x2, weight, bias); + rearrange_(out, result); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits.cc b/csrc/infinicore/src/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits.cc new file mode 100644 index 000000000..36285e938 --- /dev/null +++ b/csrc/infinicore/src/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits.cc @@ -0,0 +1,67 @@ +#include "infinicore/ops/binary_cross_entropy_with_logits.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +// 静态调度器实例化 +common::OpDispatcher &BinaryCrossEntropyWithLogits::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +/** + * 执行核心逻辑：设备校验、上下文设置与后端分发 + */ +void BinaryCrossEntropyWithLogits::execute(Tensor out, Tensor logits, Tensor target, Tensor weight, Tensor pos_weight, std::string reduction) { + // 1. 校验所有已定义的 Tensor 是否在同一设备上 + // 使用宏或循环校验 logits, target, out 以及可选的 weight/pos_weight + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, logits, target); + if (weight) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, weight); + } + if (pos_weight) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, pos_weight); + } + + // 2. 设置当前设备上下文 + infinicore::context::setDevice(out->device()); + + // 3. 根据设备类型查找并执行具体的后端实现（如 CUDA 或 CPU 实现） + dispatcher().lookup(out->device().type())(out, logits, target, weight, pos_weight, reduction); +} + +/** + * Out-of-place 接口：根据 reduction 自动创建输出 Tensor + */ +Tensor binary_cross_entropy_with_logits(Tensor logits, Tensor target, Tensor weight, Tensor pos_weight, std::string reduction) { + std::vector out_shape; + + // 1. 根据归约方式确定输出形状 + if (reduction == "none") { + // 不归约，形状与输入 logits 一致 + auto in_shape = logits->shape(); + for (auto dim : in_shape) { + out_shape.push_back(static_cast(dim)); + } + } else { + // mean 或 sum 归约，输出为标量 (空 shape 向量表示 0-dim tensor) + out_shape = {}; + } + + // 2. 创建输出 Tensor + auto out = Tensor::empty(out_shape, logits->dtype(), logits->device()); + + // 3. 调用显式接口执行计算 + binary_cross_entropy_with_logits_(out, logits, target, weight, pos_weight, reduction); + + return out; +} + +/** + * 显式指定输出接口 + */ +void binary_cross_entropy_with_logits_(Tensor out, Tensor logits, Tensor target, Tensor weight, Tensor pos_weight, std::string reduction) { + BinaryCrossEntropyWithLogits::execute(out, logits, target, weight, pos_weight, reduction); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/bitwise_right_shift/bitwise_right_shift.cc b/csrc/infinicore/src/ops/bitwise_right_shift/bitwise_right_shift.cc new file mode 100644 index 000000000..ba48d1f18 --- /dev/null +++ b/csrc/infinicore/src/ops/bitwise_right_shift/bitwise_right_shift.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/bitwise_right_shift.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BitwiseRightShift); + +BitwiseRightShift::BitwiseRightShift(Tensor out, const Tensor &input, const Tensor &other) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, other); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), out, input, other); +} + +void BitwiseRightShift::execute(Tensor out, const Tensor &input, const Tensor &other) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(BitwiseRightShift, out, input, other); +} + +Tensor bitwise_right_shift(const Tensor &input, const Tensor &other) { + auto out = Tensor::empty(input->shape(), input->dtype(), input->device()); + bitwise_right_shift_(out, input, other); + return out; +} + +void bitwise_right_shift_(Tensor out, const Tensor &input, const Tensor &other) { + BitwiseRightShift::execute(out, input, other); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/blas_amax/blas_amax.cc b/csrc/infinicore/src/ops/blas_amax/blas_amax.cc new file mode 100644 index 000000000..deffd37f6 --- /dev/null +++ b/csrc/infinicore/src/ops/blas_amax/blas_amax.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/blas_amax.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BlasAmax); + +BlasAmax::BlasAmax(const Tensor &x, Tensor result) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, result); + INFINICORE_GRAPH_OP_DISPATCH(result->device().type(), x, result); +} + +void BlasAmax::execute(const Tensor &x, Tensor result) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(BlasAmax, x, result); +} + +Tensor blas_amax(const Tensor &x) { + auto result = Tensor::empty({}, DataType::kInt32, x->device()); + blas_amax_(x, result); + return result; +} + +void blas_amax_(const Tensor &x, Tensor result) { + BlasAmax::execute(x, result); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/blas_amin/blas_amin.cc b/csrc/infinicore/src/ops/blas_amin/blas_amin.cc new file mode 100644 index 000000000..2fb3696ab --- /dev/null +++ b/csrc/infinicore/src/ops/blas_amin/blas_amin.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/blas_amin.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BlasAmin); + +BlasAmin::BlasAmin(const Tensor &x, Tensor result) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, result); + INFINICORE_GRAPH_OP_DISPATCH(result->device().type(), x, result); +} + +void BlasAmin::execute(const Tensor &x, Tensor result) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(BlasAmin, x, result); +} + +Tensor blas_amin(const Tensor &x) { + auto result = Tensor::empty({}, DataType::kInt32, x->device()); + blas_amin_(x, result); + return result; +} + +void blas_amin_(const Tensor &x, Tensor result) { + BlasAmin::execute(x, result); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/blas_copy/blas_copy.cc b/csrc/infinicore/src/ops/blas_copy/blas_copy.cc new file mode 100644 index 000000000..211220921 --- /dev/null +++ b/csrc/infinicore/src/ops/blas_copy/blas_copy.cc @@ -0,0 +1,22 @@ +#include "infinicore/ops/blas_copy.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BlasCopy); + +BlasCopy::BlasCopy(const Tensor &x, Tensor y) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), x, y); +} + +void BlasCopy::execute(const Tensor &x, Tensor y) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(BlasCopy, x, y); +} + +void blas_copy_(const Tensor &x, Tensor y) { + BlasCopy::execute(x, y); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/blas_dot/blas_dot.cc b/csrc/infinicore/src/ops/blas_dot/blas_dot.cc new file mode 100644 index 000000000..772154a90 --- /dev/null +++ b/csrc/infinicore/src/ops/blas_dot/blas_dot.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/blas_dot.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BlasDot); + +BlasDot::BlasDot(const Tensor &x, const Tensor &y, Tensor result) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y, result); + INFINICORE_GRAPH_OP_DISPATCH(result->device().type(), x, y, result); +} + +void BlasDot::execute(const Tensor &x, const Tensor &y, Tensor result) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(BlasDot, x, y, result); +} + +Tensor blas_dot(const Tensor &x, const Tensor &y) { + auto result = Tensor::empty({}, x->dtype(), x->device()); + blas_dot_(x, y, result); + return result; +} + +void blas_dot_(const Tensor &x, const Tensor &y, Tensor result) { + BlasDot::execute(x, y, result); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/block_diag/block_diag.cc b/csrc/infinicore/src/ops/block_diag/block_diag.cc new file mode 100644 index 000000000..5afcc9124 --- /dev/null +++ b/csrc/infinicore/src/ops/block_diag/block_diag.cc @@ -0,0 +1,63 @@ +#include "infinicore/ops/block_diag.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &BlockDiag::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void BlockDiag::execute(Tensor output, const std::vector &inputs) { + if (inputs.empty()) { + throw std::runtime_error("block_diag expects at least one input tensor"); + } + + for (const auto &x : inputs) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, x); + } + + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error( + "No BlockDiag implementation found for device type: " + + std::to_string(static_cast(device_type))); + } + + func(output, inputs); +} + +Tensor block_diag(const std::vector &inputs) { + if (inputs.empty()) { + throw std::runtime_error("block_diag expects at least one input tensor"); + } + + const auto &device = inputs.front()->device(); + const auto dtype = inputs.front()->dtype(); + + Size total_rows = 0; + Size total_cols = 0; + for (const auto &x : inputs) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(inputs.front(), x); + INFINICORE_ASSERT(x->dtype() == dtype); + INFINICORE_ASSERT(x->ndim() == 2); + total_rows += x->size(0); + total_cols += x->size(1); + } + + auto output = Tensor::empty({total_rows, total_cols}, dtype, device); + block_diag_(output, inputs); + return output; +} + +void block_diag_(Tensor output, const std::vector &inputs) { + BlockDiag::execute(output, inputs); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/broadcast_to/broadcast_to.cc b/csrc/infinicore/src/ops/broadcast_to/broadcast_to.cc new file mode 100644 index 000000000..1949ae1b2 --- /dev/null +++ b/csrc/infinicore/src/ops/broadcast_to/broadcast_to.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/broadcast_to.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BroadcastTo); + +BroadcastTo::BroadcastTo(Tensor y, Tensor x) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), y, x); +} + +void BroadcastTo::execute(Tensor y, Tensor x) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(BroadcastTo, y, x); +} + +Tensor broadcast_to(Tensor x, const std::vector &shape) { + Shape target_shape(shape.begin(), shape.end()); + auto y = Tensor::empty(target_shape, x->dtype(), x->device()); + broadcast_to_(y, x); + return y; +} + +void broadcast_to_(Tensor y, Tensor x) { + BroadcastTo::execute(y, x); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/cat/cat.cc b/csrc/infinicore/src/ops/cat/cat.cc new file mode 100644 index 000000000..350dcd04c --- /dev/null +++ b/csrc/infinicore/src/ops/cat/cat.cc @@ -0,0 +1,199 @@ +#include "infinicore/ops/cat.hpp" +#include "infinicore/context/context.hpp" +#include +#include +namespace infinicore::op { + +namespace { + +bool use_slice_copy_cat(Device::Type device_type, int dim, int ndim) { + // Keep this path limited to CUDA-like backends we have validated or can + // trace to the nvidia rearrange implementation. Other devices may be + // correct through copy_from, but their performance impact is unverified. + return dim == ndim - 1 + && (device_type == Device::Type::kNvidia + || device_type == Device::Type::kHygon + || device_type == Device::Type::kIluvatar); +} + +class CatInfo { + + CatInfo() = default; + +public: + int dim; + int ndim; + size_t tensors_size; + + std::vector contiguous_dim; + std::vector copy_size; + + static CatInfo create(Tensor &out, std::vector &tensors, int dim) { + + int ndim = out->ndim(); + size_t tensors_size = tensors.size(); + + std::vector contiguous_dim(tensors_size, ndim); + std::vector copy_size(tensors_size, dsize(out->dtype())); + + for (int i = 0; i < tensors_size; i++) { + if (tensors[i]->ndim() == 1) { + continue; + } + if (tensors[i]->stride(ndim - 1) == 1) { + contiguous_dim[i] = ndim - 1; + for (int j = ndim - 2; j >= dim; j--) { + if (tensors[i]->stride(j) == tensors[i]->stride(j + 1) * tensors[i]->shape()[j + 1]) { + contiguous_dim[i] = j; + } + } + } + for (int j = contiguous_dim[i]; j < ndim; j++) { + copy_size[i] *= tensors[i]->shape()[j]; + } + } + + return CatInfo{dim, ndim, tensors_size, contiguous_dim, copy_size}; + } +}; + +void low_dim_copy( + CatInfo &info, Tensor &tensor, Tensor &out, + std::byte *tensor_ptr, std::byte *out_ptr, + int depth, int tensor_pos) { + + if (depth != info.contiguous_dim[tensor_pos]) { + std::byte *now_tensor_ptr = tensor_ptr; + std::byte *now_out_ptr = out_ptr; + + for (int i = 0; i < tensor->shape()[depth]; i++) { + + low_dim_copy(info, tensor, out, now_tensor_ptr, now_out_ptr, depth + 1, tensor_pos); + + now_tensor_ptr += tensor->stride(depth) * dsize(tensor->dtype()); + now_out_ptr += out->stride(depth) * dsize(out->dtype()); + } + } else { + if (out->device().type() == Device::Type::kCpu) { + + std::memcpy(out_ptr, tensor_ptr, info.copy_size[tensor_pos]); + } else { + + context::memcpyD2D(out_ptr, tensor_ptr, info.copy_size[tensor_pos]); + } + } +} + +void high_dim_split( + CatInfo &info, std::vector &tensors, Tensor &out, + std::vector tensors_ptr, std::byte *out_ptr, + int depth) { + + if (depth != info.dim) { + std::vector now_tensors_ptr = tensors_ptr; + std::byte *now_out_ptr = out_ptr; + + for (int i = 0; i < out->shape()[depth]; i++) { + + high_dim_split(info, tensors, out, now_tensors_ptr, now_out_ptr, depth + 1); + + for (int i = 0; i < info.tensors_size; i++) { + if (tensors[i]->ndim() == 1) { + continue; + } + now_tensors_ptr[i] += tensors[i]->stride(depth) * dsize(tensors[i]->dtype()); + } + now_out_ptr += out->stride(depth) * dsize(out->dtype()); + } + } else { + std::byte *now_out_ptr = out_ptr; + + for (int i = 0; i < info.tensors_size; i++) { + if (tensors[i]->ndim() == 1) { + continue; + } + + low_dim_copy(info, tensors[i], out, tensors_ptr[i], now_out_ptr, depth, i); + + now_out_ptr += tensors[i]->shape()[depth] * out->stride(depth) * dsize(out->dtype()); + } + } +} + +} // namespace + +Tensor cat(std::vector tensors, int dim) { + assert(tensors.size() >= 2); + int ndim = tensors[0]->ndim(); + assert(-ndim <= dim && dim < ndim); + dim = (dim + ndim) % ndim; + + Shape shape = tensors[0]->shape(); + for (int i = 1; i < tensors.size(); i++) { + assert(tensors[i]->ndim() == dim || tensors[i]->ndim() == 1); + if (tensors[i]->ndim() != ndim) { + continue; + } + shape[dim] += tensors[i]->shape()[dim]; + } + + auto out = Tensor::empty(shape, tensors[0]->dtype(), tensors[0]->device()); + cat_(out, tensors, dim); + return out; +} + +void cat_(Tensor out, std::vector tensors, int dim) { + // assert the parameter properties are correct. + assert(tensors.size() >= 2); + int ndim = out->ndim(); + assert(-ndim <= dim && dim < ndim); + dim = (dim + ndim) % ndim; + + size_t dim_shape = 0; + for (auto &tensor : tensors) { + assert(tensor->ndim() == ndim || tensor->ndim() == 1); + if (tensor->ndim() == 1) { + assert(tensor->shape()[0] == 0); + continue; + } + for (int i = 0; i < ndim; i++) { + if (i != dim) { + assert(tensor->shape()[i] == out->shape()[i]); + } else { + dim_shape += tensor->shape()[i]; + } + } + } + assert(dim_shape == out->shape()[dim]); + + if (use_slice_copy_cat(out->device().type(), dim, ndim)) { + // The generic path recursively issues one memcpy for every outer + // index. Concatenating MLA tensors on the last dimension can therefore + // enqueue hundreds of tiny D2D copy calls per layer. A strided output + // slice is semantically identical and copy_from lowers it to a single + // rearrange kernel per input tensor on CUDA-like backends. + size_t offset = 0; + for (auto &tensor : tensors) { + if (tensor->ndim() == 1) { + continue; + } + const size_t length = tensor->shape()[dim]; + auto output_slice = out->narrow({{static_cast(dim), offset, length}}); + output_slice->copy_from(tensor); + offset += length; + } + return; + } + + // Get info + CatInfo info = CatInfo::create(out, tensors, dim); + std::vector tensors_ptr(tensors.size()); + for (int i = 0; i < tensors.size(); i++) { + tensors_ptr[i] = tensors[i]->data(); + } + std::byte *out_ptr = out->data(); + + high_dim_split(info, tensors, out, tensors_ptr, out_ptr, 0); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/causal_conv1d/causal_conv1d.cc b/csrc/infinicore/src/ops/causal_conv1d/causal_conv1d.cc new file mode 100644 index 000000000..067dcec7c --- /dev/null +++ b/csrc/infinicore/src/ops/causal_conv1d/causal_conv1d.cc @@ -0,0 +1,146 @@ +#include "infinicore/ops/causal_conv1d.hpp" +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(CausalConv1d); + +CausalConv1d::CausalConv1d(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) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, conv_state, qkv, weight); + if (final_conv_state.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_conv_state.value()); + } + if (bias.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, bias.value()); + } + if (cu_seqlens.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, cu_seqlens.value()); + } + if (initial_state_indices.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state_indices.value()); + } + if (final_state_indices.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state_indices.value()); + } + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), + out, + conv_state, + final_conv_state, + qkv, + weight, + bias, + cu_seqlens, + initial_state_indices, + final_state_indices); +} + +void CausalConv1d::execute(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) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(CausalConv1d, + out, + conv_state, + final_conv_state, + qkv, + weight, + bias, + cu_seqlens, + initial_state_indices, + final_state_indices); +} + +static void check_3d_tensor(const Tensor &x, const char *name, const char *shape_hint) { + if (x->shape().size() != 3) { + throw std::runtime_error(std::string("causal_conv1d expects ") + name + " with shape " + shape_hint); + } +} + +static size_t request_count_from_args(const Tensor &qkv, std::optional cu_seqlens) { + if (cu_seqlens.has_value()) { + const auto &cu_shape = cu_seqlens.value()->shape(); + if (cu_shape.size() != 1 || cu_shape[0] < 2) { + throw std::runtime_error("causal_conv1d expects cu_seqlens with shape [num_requests + 1]"); + } + return cu_shape[0] - 1; + } + return qkv->shape()[0]; +} + +static Shape final_state_shape(const Tensor &qkv, + const Tensor &conv_state, + const Tensor &weight, + std::optional cu_seqlens) { + const auto request_count = request_count_from_args(qkv, cu_seqlens); + return {request_count, conv_state->shape()[1], weight->shape()[2] - 1}; +} + +Tensor causal_conv1d(const Tensor &qkv, + Tensor conv_state, + const Tensor &weight, + std::optional bias, + std::optional cu_seqlens, + std::optional initial_state_indices, + std::optional final_state_indices) { + check_3d_tensor(qkv, "qkv", "[B, T, C] or [1, total_tokens, C]"); + check_3d_tensor(conv_state, "conv_state", "[B/num_requests, C, state_len] or [pool_size, C, state_len]"); + check_3d_tensor(weight, "weight", "[C, 1, state_len + 1]"); + auto out = Tensor::empty(qkv->shape(), qkv->dtype(), qkv->device()); + std::optional final_conv_state = std::nullopt; + if (!final_state_indices.has_value()) { + final_conv_state = Tensor::empty(final_state_shape(qkv, conv_state, weight, cu_seqlens), + conv_state->dtype(), + conv_state->device()); + } + causal_conv1d_(out, + conv_state, + final_conv_state, + qkv, + weight, + bias, + cu_seqlens, + initial_state_indices, + final_state_indices); + return out; +} + +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) { + check_3d_tensor(out, "out", "[B, T, C] or [1, total_tokens, C]"); + check_3d_tensor(qkv, "qkv", "[B, T, C] or [1, total_tokens, C]"); + check_3d_tensor(conv_state, "conv_state", "[B/num_requests, C, state_len] or [pool_size, C, state_len]"); + check_3d_tensor(weight, "weight", "[C, 1, state_len + 1]"); + CausalConv1d::execute(out, + conv_state, + final_conv_state, + qkv, + weight, + bias, + cu_seqlens, + initial_state_indices, + final_state_indices); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/causal_softmax/causal_softmax.cc b/csrc/infinicore/src/ops/causal_softmax/causal_softmax.cc new file mode 100644 index 000000000..b7ca95566 --- /dev/null +++ b/csrc/infinicore/src/ops/causal_softmax/causal_softmax.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/causal_softmax.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(CausalSoftmax); + +CausalSoftmax::CausalSoftmax(Tensor output, const Tensor &input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + INFINICORE_GRAPH_OP_DISPATCH(output->device().type(), output, input); +} + +void CausalSoftmax::execute(Tensor output, const Tensor &input) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(CausalSoftmax, output, input); +} + +Tensor causal_softmax(const Tensor &input) { + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + causal_softmax_(output, input); + return output; +} + +void causal_softmax_(Tensor output, const Tensor &input) { + CausalSoftmax::execute(output, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/causal_softmax/causal_softmax_infiniops.cc b/csrc/infinicore/src/ops/causal_softmax/causal_softmax_infiniops.cc new file mode 100644 index 000000000..61e5379db --- /dev/null +++ b/csrc/infinicore/src/ops/causal_softmax/causal_softmax_infiniops.cc @@ -0,0 +1,58 @@ +#include "infinicore/ops/causal_softmax.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/causal_softmax.h" + +namespace infinicore::op::causal_softmax_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +struct PlannedMeta { + TensorMeta output, input; + graph::GraphTensor output_tensor, input_tensor; +}; + +} // namespace + +void *plan(Tensor output, const Tensor &input) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + + return new PlannedMeta{ + TensorMeta(output), + TensorMeta(input), + graph::GraphTensor(output), + graph::GraphTensor(input)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + infini::ops::CausalSoftmax::Call( + handle, + config, + planned->input.tensor(planned->input_tensor), + planned->output.tensor(planned->output_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(CausalSoftmax::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(CausalSoftmax::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(CausalSoftmax::cleanup_dispatcher(), &cleanup); + return true; +}(); + +} // namespace infinicore::op::causal_softmax_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/cdist/cdist.cc b/csrc/infinicore/src/ops/cdist/cdist.cc new file mode 100644 index 000000000..184defd77 --- /dev/null +++ b/csrc/infinicore/src/ops/cdist/cdist.cc @@ -0,0 +1,56 @@ +#include "infinicore/ops/cdist.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +// 静态调度器实例化 +common::OpDispatcher &Cdist::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +/** + * 执行核心逻辑：设备校验与后端分发 + */ +void Cdist::execute(Tensor out, Tensor x1, Tensor x2, double p) { + // 校验三个 Tensor 是否在同一设备上 + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, x1, x2); + + // 设置当前设备上下文 + infinicore::context::setDevice(out->device()); + + // 根据设备类型（CUDA/CPU/etc.）查找并执行注册的算子实现 + dispatcher().lookup(out->device().type())(out, x1, x2, p); +} + +/** + * Out-of-place 接口：自动创建输出 Tensor + * x1: (M, D), x2: (N, D) -> out: (M, N) + */ +Tensor cdist(Tensor x1, Tensor x2, double p) { + // 1. 获取输入维度 + auto shape1 = x1->shape(); // 假设为 {M, D} + auto shape2 = x2->shape(); // 假设为 {N, D} + + // 将原来的 std::vector 修改为 std::vector + std::vector out_shape = { + static_cast(shape1[0]), + static_cast(shape2[0])}; + + // 或者使用更简洁的初始化列表方式，强制转换类型 + auto out = Tensor::empty({(uint64_t)shape1[0], (uint64_t)shape2[0]}, x1->dtype(), x1->device()); + + // 5. 调用执行接口 + cdist_(out, x1, x2, p); + + return out; +} + +/** + * 显式指定输出接口 + */ +void cdist_(Tensor out, Tensor x1, Tensor x2, double p) { + Cdist::execute(out, x1, x2, p); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/chunk_gated_delta_rule/chunk_gated_delta_rule.cc b/csrc/infinicore/src/ops/chunk_gated_delta_rule/chunk_gated_delta_rule.cc new file mode 100644 index 000000000..2afbaf00b --- /dev/null +++ b/csrc/infinicore/src/ops/chunk_gated_delta_rule/chunk_gated_delta_rule.cc @@ -0,0 +1,167 @@ +#include "infinicore/ops/chunk_gated_delta_rule.hpp" +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(ChunkGatedDeltaRule); + +ChunkGatedDeltaRule::ChunkGatedDeltaRule(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, + size_t chunk_size) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state, q, k, v, g, beta); + if (final_state.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state.value()); + } + if (cu_seqlens.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, cu_seqlens.value()); + } + if (initial_state_indices.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state_indices.value()); + } + if (final_state_indices.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state_indices.value()); + } + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), + out, + initial_state, + final_state, + q, + k, + v, + g, + beta, + cu_seqlens, + initial_state_indices, + final_state_indices, + use_qk_l2norm, + chunk_size); +} + +void ChunkGatedDeltaRule::execute(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, + size_t chunk_size) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(ChunkGatedDeltaRule, + out, + initial_state, + final_state, + q, + k, + v, + g, + beta, + cu_seqlens, + initial_state_indices, + final_state_indices, + use_qk_l2norm, + chunk_size); +} + +static void check_4d_sequence_tensor(const Tensor &x, const char *name) { + if (x->shape().size() != 4) { + throw std::runtime_error(std::string("chunk_gated_delta_rule expects ") + name + " with shape [B, T, H, D] or [1, total_tokens, H, D]"); + } +} + +static Shape chunk_final_state_shape(const Tensor &q, + const Tensor &v, + std::optional cu_seqlens) { + const auto &q_shape = q->shape(); + const auto &v_shape = v->shape(); + size_t B = cu_seqlens.has_value() ? cu_seqlens.value()->shape()[0] - 1 : v_shape[0]; + size_t Hv = v_shape[2]; + size_t Dk = q_shape[3]; + size_t Dv = v_shape[3]; + return {B, Hv, Dv, Dk}; +} + +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::optional initial_state_indices, + std::optional final_state_indices, + bool use_qk_l2norm, + size_t chunk_size) { + check_4d_sequence_tensor(q, "q"); + check_4d_sequence_tensor(k, "k"); + check_4d_sequence_tensor(v, "v"); + auto out = Tensor::empty(v->shape(), v->dtype(), v->device()); + std::optional final_state = std::nullopt; + if (!final_state_indices.has_value()) { + final_state = Tensor::empty(chunk_final_state_shape(q, v, cu_seqlens), + initial_state->dtype(), + initial_state->device()); + } + chunk_gated_delta_rule_(out, + initial_state, + final_state, + q, + k, + v, + g, + beta, + cu_seqlens, + initial_state_indices, + final_state_indices, + use_qk_l2norm, + chunk_size); + return out; +} + +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, + size_t chunk_size) { + check_4d_sequence_tensor(q, "q"); + check_4d_sequence_tensor(k, "k"); + check_4d_sequence_tensor(v, "v"); + ChunkGatedDeltaRule::execute(out, + initial_state, + final_state, + q, + k, + v, + g, + beta, + cu_seqlens, + initial_state_indices, + final_state_indices, + use_qk_l2norm, + chunk_size); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/conv2d/conv2d.cc b/csrc/infinicore/src/ops/conv2d/conv2d.cc new file mode 100644 index 000000000..8f84d923c --- /dev/null +++ b/csrc/infinicore/src/ops/conv2d/conv2d.cc @@ -0,0 +1,101 @@ +#include "infinicore/ops/conv2d.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &Conv2d::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Conv2d::execute(Tensor output, + Tensor input, + Tensor weight, + Tensor bias, + const size_t *pads, + const size_t *strides, + const size_t *dilations, + size_t n) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, weight); + if (bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, bias); + } + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Conv2d implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input, weight, bias, pads, strides, dilations, n); +} + +Tensor conv2d(Tensor input, + Tensor weight, + Tensor bias, + const std::vector &pads, + const std::vector &strides, + const std::vector &dilations) { + const auto &in_shape = input->shape(); // [N, C_in, H_in, W_in] + const auto &w_shape = weight->shape(); // [C_out, C_in, kH, kW] + + // ------------------------------- + // Extract dimensions + // ------------------------------- + size_t N = in_shape[0]; + size_t C_in = in_shape[1]; + size_t H_in = in_shape[2]; + size_t W_in = in_shape[3]; + + size_t C_out = w_shape[0]; + size_t kH = w_shape[2]; + size_t kW = w_shape[3]; + + size_t pad_h = pads[0]; + size_t pad_w = pads[1]; + + size_t stride_h = strides[0]; + size_t stride_w = strides[1]; + + size_t dil_h = dilations[0]; + size_t dil_w = dilations[1]; + + auto calc_out = [](size_t in, size_t pad, size_t dilation, size_t kernel, size_t stride) { + return (in + 2 * pad - dilation * (kernel - 1) - 1) / stride + 1; + }; + size_t H_out = calc_out(H_in, pad_h, dil_h, kH, stride_h); + size_t W_out = calc_out(W_in, pad_w, dil_w, kW, stride_w); + if ((int64_t)H_out <= 0 || (int64_t)W_out <= 0) { + throw std::runtime_error("Invalid conv2d output shape (negative or zero)"); + } + Shape out_shape = {N, C_out, H_out, W_out}; + + auto output = Tensor::empty(out_shape, input->dtype(), input->device()); + conv2d_(output, input, weight, bias, pads, strides, dilations); + return output; +} + +void conv2d_(Tensor output, + Tensor input, + Tensor weight, + Tensor bias, + const std::vector &pads, + const std::vector &strides, + const std::vector &dilations) { + if (pads.size() != strides.size() || pads.size() != dilations.size()) { + throw std::runtime_error("conv2d_: pads/strides/dilations must have the same size"); + } + Conv2d::execute(output, + input, + weight, + bias, + pads.data(), + strides.data(), + dilations.data(), + pads.size()); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/conv2d/conv2d_infiniops.cc b/csrc/infinicore/src/ops/conv2d/conv2d_infiniops.cc new file mode 100644 index 000000000..2b4816c94 --- /dev/null +++ b/csrc/infinicore/src/ops/conv2d/conv2d_infiniops.cc @@ -0,0 +1,72 @@ +#include "infinicore/ops/conv2d.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/conv_infinilm.h" + +#include +#include + +namespace infinicore::op::conv2d_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +std::vector toInt64Vector(const size_t *values, size_t n) { + std::vector result; + result.reserve(n); + for (size_t i = 0; i < n; ++i) { + result.push_back(static_cast(values[i])); + } + return result; +} + +void calculate(Tensor output, + Tensor input, + Tensor weight, + Tensor bias, + const size_t *pads, + const size_t *strides, + const size_t *dilations, + size_t n) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, weight); + if (bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, bias); + } + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + TensorMeta output_meta(output); + TensorMeta input_meta(input); + TensorMeta weight_meta(weight); + std::optional bias_meta; + if (bias) { + bias_meta.emplace(bias); + } + + infini::ops::ConvInfinilm::Call( + handle, + config, + input_meta.tensor(input), + weight_meta.tensor(weight), + bias_meta ? std::optional{bias_meta->tensor(bias)} : std::nullopt, + toInt64Vector(pads, n), + toInt64Vector(strides, n), + toInt64Vector(dilations, n), + int64_t{1}, + output_meta.tensor(output)); +} + +} // namespace + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Conv2d::dispatcher(), &calculate); + return true; +}(); + +} // namespace infinicore::op::conv2d_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/cross_entropy/cross_entropy.cc b/csrc/infinicore/src/ops/cross_entropy/cross_entropy.cc new file mode 100644 index 000000000..121c4bdf9 --- /dev/null +++ b/csrc/infinicore/src/ops/cross_entropy/cross_entropy.cc @@ -0,0 +1,45 @@ +#include "infinicore/ops/cross_entropy.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &CrossEntropy::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void CrossEntropy::execute(Tensor output, Tensor input, Tensor target) { + + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(input, target); + + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No CrossEntropy implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input, target); +} + +Tensor cross_entropy(Tensor input, Tensor target) { + + Shape shape = target->shape(); + + auto output = Tensor::empty(shape, input->dtype(), input->device()); + + cross_entropy_(output, input, target); + return output; +} + +void cross_entropy_(Tensor output, Tensor input, Tensor target) { + CrossEntropy::execute(output, input, target); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/deepseek_moe/deepseek_moe.cc b/csrc/infinicore/src/ops/deepseek_moe/deepseek_moe.cc new file mode 100644 index 000000000..baadf2a41 --- /dev/null +++ b/csrc/infinicore/src/ops/deepseek_moe/deepseek_moe.cc @@ -0,0 +1,83 @@ +#include "infinicore/ops/deepseek_moe.hpp" +#include "../../utils.hpp" +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(DeepseekMoe); + +namespace { + +void check_weights(const std::vector &gate_weights, + const std::vector &up_weights, + const std::vector &down_weights, + size_t num_experts) { + if (gate_weights.size() != num_experts || up_weights.size() != num_experts || down_weights.size() != num_experts) { + throw std::runtime_error("DeepseekMoe: expert weight vector size mismatch"); + } +} + +} // namespace + +DeepseekMoe::DeepseekMoe(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) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, hidden, topk_indices, topk_weights); + check_weights(gate_weights, up_weights, down_weights, num_experts); + for (size_t i = 0; i < num_experts; ++i) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, gate_weights[i], up_weights[i], down_weights[i]); + } + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), + out, hidden, topk_indices, topk_weights, + gate_weights, up_weights, down_weights, + intermediate_size, num_experts); +} + +void DeepseekMoe::execute(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) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + DeepseekMoe, + out, hidden, topk_indices, topk_weights, + gate_weights, up_weights, down_weights, + intermediate_size, 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) { + DeepseekMoe::execute(out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, intermediate_size, num_experts); +} + +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) { + auto out = Tensor::empty(hidden->shape(), hidden->dtype(), hidden->device()); + deepseek_moe_(out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, intermediate_size, num_experts); + return out; +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/dequant/per_tensor_dequant_int8/per_tensor_quant_int8.cc b/csrc/infinicore/src/ops/dequant/per_tensor_dequant_int8/per_tensor_quant_int8.cc new file mode 100644 index 000000000..43a898eda --- /dev/null +++ b/csrc/infinicore/src/ops/dequant/per_tensor_dequant_int8/per_tensor_quant_int8.cc @@ -0,0 +1,20 @@ +#include "../../../utils.hpp" +#include "infinicore/ops/per_tensor_dequant_i8.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PerTensorDequantI8); + +PerTensorDequantI8::PerTensorDequantI8(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zero) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, x_packed, x_scale, x_zero); + INFINICORE_GRAPH_OP_DISPATCH(x->device().type(), x, x_packed, x_scale, x_zero); +} + +void PerTensorDequantI8::execute(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zero) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(PerTensorDequantI8, x, x_packed, x_scale, x_zero); +} + +void per_tensor_dequant_i8_(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zero) { + PerTensorDequantI8::execute(x, x_packed, x_scale, x_zero); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/diff/diff.cc b/csrc/infinicore/src/ops/diff/diff.cc new file mode 100644 index 000000000..7601b763a --- /dev/null +++ b/csrc/infinicore/src/ops/diff/diff.cc @@ -0,0 +1,78 @@ +#include "infinicore/ops/diff.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Diff); + +Diff::Diff(Tensor y, const Tensor &x, int dim, int n) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), y, x, dim, n); +} + +void Diff::execute(Tensor y, const Tensor &x, int dim, int n) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Diff, y, x, dim, n); +} + +static int normalize_dim(int dim, size_t ndim) { + if (ndim == 0) { + throw std::runtime_error("diff: input tensor must have at least one dimension."); + } + if (dim < 0) { + dim += static_cast(ndim); + } + if (dim < 0 || static_cast(dim) >= ndim) { + throw std::runtime_error("diff: dim out of range."); + } + return dim; +} + +Tensor diff(const Tensor &x, int n, int dim) { + if (n < 0) { + throw std::runtime_error("diff: n must be non-negative."); + } + Shape y_shape = x->shape(); + const int d = normalize_dim(dim, y_shape.size()); + const auto dim_size = y_shape[static_cast(d)]; + y_shape[static_cast(d)] = (dim_size >= static_cast(n)) ? (dim_size - static_cast(n)) : 0; + + auto y = Tensor::empty(y_shape, x->dtype(), x->device()); + if (n == 0) { + y->copy_from(x); + return y; + } + if (dim_size <= static_cast(n)) { + // Empty output by definition; nothing to compute. + return y; + } + + diff_(y, x, n, dim); + return y; +} + +void diff_(Tensor y, const Tensor &x, int n, int dim) { + if (n < 0) { + throw std::runtime_error("diff_: n must be non-negative."); + } + const int d = normalize_dim(dim, x->shape().size()); + Shape expected = x->shape(); + const auto dim_size = expected[static_cast(d)]; + expected[static_cast(d)] = (dim_size >= static_cast(n)) ? (dim_size - static_cast(n)) : 0; + if (y->shape() != expected) { + throw std::runtime_error("diff_: output tensor has incorrect shape."); + } + if (n == 0) { + y->copy_from(x); + return; + } + if (x->shape()[static_cast(d)] <= static_cast(n)) { + // Empty output by definition; nothing to compute. + return; + } + Diff::execute(y, x, d, n); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/digamma/digamma.cc b/csrc/infinicore/src/ops/digamma/digamma.cc new file mode 100644 index 000000000..6997b4a4a --- /dev/null +++ b/csrc/infinicore/src/ops/digamma/digamma.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/digamma.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Digamma); + +Digamma::Digamma(Tensor y, const Tensor &x) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), y, x); +} + +void Digamma::execute(Tensor y, const Tensor &x) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Digamma, y, x); +} + +Tensor digamma(const Tensor &x) { + auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); + digamma_(y, x); + return y; +} + +void digamma_(Tensor y, const Tensor &x) { + Digamma::execute(y, x); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/dist/dist.cc b/csrc/infinicore/src/ops/dist/dist.cc new file mode 100644 index 000000000..523f669c3 --- /dev/null +++ b/csrc/infinicore/src/ops/dist/dist.cc @@ -0,0 +1,30 @@ +#include "infinicore/ops/dist.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dist); + +Dist::Dist(Tensor y, const Tensor &x1, const Tensor &x2, double p) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x1, x2); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), y, x1, x2, p); +} + +void Dist::execute(Tensor y, const Tensor &x1, const Tensor &x2, double p) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dist, y, x1, x2, p); +} + +Tensor dist(const Tensor &x1, const Tensor &x2, double p) { + auto y = Tensor::empty({}, x1->dtype(), x1->device()); + dist_(y, x1, x2, p); + return y; +} + +void dist_(Tensor y, const Tensor &x1, const Tensor &x2, double p) { + Dist::execute(y, x1, x2, p); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/distributed/allgather.cc b/csrc/infinicore/src/ops/distributed/allgather.cc new file mode 100644 index 000000000..f2aadad20 --- /dev/null +++ b/csrc/infinicore/src/ops/distributed/allgather.cc @@ -0,0 +1,169 @@ +#include "infinicore/ops/distributed/allgather.hpp" +#include "../../utils.hpp" +#include "utils.hpp" + +#include "infinicore/context/context.hpp" + +#include +#include +#include + +namespace infinicore::op::distributed { +namespace { + +bool all_equal_to(const std::vector &values, size_t expected) { + return !values.empty() && std::all_of(values.begin(), values.end(), [&](size_t value) { + return value == expected; + }); +} + +void validate_gather_input(const Tensor &input, const std::vector &split_sizes) { + INFINICORE_ASSERT(input); + INFINICORE_ASSERT(input->is_contiguous()); + INFINICORE_ASSERT(input->ndim() > 0); + INFINICORE_ASSERT(input->shape()[0] > 0); + INFINICORE_ASSERT(!split_sizes.empty()); + (void)detail::toInfinicclDataType(input->dtype()); + if (!all_equal_to(split_sizes, input->shape()[0])) { + throw std::runtime_error("InfiniCCL does not support variable-count all-gather"); + } + INFINICORE_ASSERT(input->shape()[0] <= std::numeric_limits::max() / split_sizes.size()); +} + +void validate_allgatherv(Tensor output, + const Tensor &input, + const std::vector &split_sizes) { + INFINICORE_ASSERT(output && input); + INFINICORE_ASSERT(output->dtype() == input->dtype()); + (void)detail::toInfinicclDataType(input->dtype()); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); + INFINICORE_ASSERT(input->ndim() > 0); + INFINICORE_ASSERT(output->ndim() == input->ndim()); + INFINICORE_ASSERT(input->shape()[0] > 0); + INFINICORE_ASSERT(!split_sizes.empty()); + if (!all_equal_to(split_sizes, input->shape()[0])) { + throw std::runtime_error("InfiniCCL does not support variable-count all-gather"); + } + INFINICORE_ASSERT(input->shape()[0] <= std::numeric_limits::max() / split_sizes.size()); + INFINICORE_ASSERT(output->shape()[0] == input->shape()[0] * split_sizes.size()); + for (size_t dim = 1; dim < input->ndim(); ++dim) { + INFINICORE_ASSERT(output->shape()[dim] == input->shape()[dim]); + } + INFINICORE_ASSERT(input->numel() <= std::numeric_limits::max() / split_sizes.size()); + INFINICORE_ASSERT(output->numel() == input->numel() * split_sizes.size()); +} + +} // namespace + +struct AllGatherPlannedMeta { + graph::GraphTensor output, input; + infinicclComm_t communicator; +}; + +AllGather::AllGather(Tensor output, const Tensor &input, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->dtype() == input->dtype()); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); + INFINICORE_ASSERT(input->numel() > 0); + INFINICORE_ASSERT(output->numel() > 0); + INFINICORE_ASSERT(output->numel() % input->numel() == 0); + planned_meta_ = new AllGatherPlannedMeta{graph::GraphTensor(output), graph::GraphTensor(input), communicator}; +} + +AllGather::~AllGather() { + if (planned_meta_) { + auto *meta = reinterpret_cast(planned_meta_); + delete meta; + } +} + +void AllGather::run() const { + auto *meta = reinterpret_cast(planned_meta_); + detail::checkInfiniccl( + "infinicclAllGather", + infinicclAllGather(meta->input->data(), + meta->output->data(), + meta->input->numel(), + detail::toInfinicclDataType(meta->input->dtype()), + meta->communicator, + reinterpret_cast(infinicore::context::getStream()))); +} + +void AllGather::execute(Tensor output, const Tensor &input, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->dtype() == input->dtype()); + INFINICORE_ASSERT(input->numel() > 0); + INFINICORE_ASSERT(output->numel() > 0); + INFINICORE_GRAPH_OP_RECORD_OR_RUN(AllGather, output, input, communicator); +} + +Tensor allgather(const Tensor &input, size_t world_size, infinicclComm_t communicator) { + INFINICORE_ASSERT(input->ndim() > 0); + INFINICORE_ASSERT(world_size > 0); + auto shape = input->shape(); + shape[0] *= world_size; + auto output = Tensor::empty(shape, input->dtype(), input->device()); + allgather_(output, input, communicator); + return output; +} + +void allgather_(Tensor output, const Tensor &input, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->dtype() == input->dtype()); + AllGather::execute(output, input, communicator); +} + +Tensor allgatherv(const Tensor &input, const std::vector &split_sizes, infinicclComm_t communicator) { + validate_gather_input(input, split_sizes); + auto shape = input->shape(); + shape[0] *= split_sizes.size(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + allgatherv_(output, input, split_sizes, communicator); + return output; +} + +void allgatherv_(Tensor output, const Tensor &input, const std::vector &split_sizes, infinicclComm_t communicator) { + validate_allgatherv(output, input, split_sizes); + AllGather::execute(output, input, communicator); +} + +std::vector allgatherv_many(const std::vector &inputs, + const std::vector &split_sizes, + infinicclComm_t communicator) { + INFINICORE_ASSERT(!split_sizes.empty()); + std::vector outputs; + outputs.reserve(inputs.size()); + if (inputs.empty()) { + return outputs; + } + for (const auto &input : inputs) { + validate_gather_input(input, split_sizes); + } + const size_t total_dim0 = inputs.front()->shape()[0] * split_sizes.size(); + for (const auto &input : inputs) { + auto shape = input->shape(); + shape[0] = total_dim0; + outputs.push_back(Tensor::empty(shape, input->dtype(), input->device())); + } + allgatherv_many_(outputs, inputs, split_sizes, communicator); + return outputs; +} + +void allgatherv_many_(const std::vector &outputs, + const std::vector &inputs, + const std::vector &split_sizes, + infinicclComm_t communicator) { + INFINICORE_ASSERT(outputs.size() == inputs.size()); + INFINICORE_ASSERT(!split_sizes.empty()); + if (inputs.empty()) { + return; + } + + for (size_t i = 0; i < inputs.size(); ++i) { + validate_allgatherv(outputs[i], inputs[i], split_sizes); + } + for (size_t i = 0; i < inputs.size(); ++i) { + AllGather::execute(outputs[i], inputs[i], communicator); + } +} + +} // namespace infinicore::op::distributed diff --git a/csrc/infinicore/src/ops/distributed/allreduce.cc b/csrc/infinicore/src/ops/distributed/allreduce.cc new file mode 100644 index 000000000..208a8e028 --- /dev/null +++ b/csrc/infinicore/src/ops/distributed/allreduce.cc @@ -0,0 +1,56 @@ +#include "infinicore/ops/distributed/allreduce.hpp" +#include "../../utils.hpp" +#include "utils.hpp" + +namespace infinicore::op::distributed { + +struct PlannedMeta { + graph::GraphTensor output, input; + infinicclRedOp_t op; + infinicclComm_t communicator; +}; + +AllReduce::AllReduce(Tensor output, const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->dtype() == input->dtype()); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); + INFINICORE_ASSERT(output->numel() == input->numel()); + planned_meta_ = new PlannedMeta{graph::GraphTensor(output), graph::GraphTensor(input), op, communicator}; +} +AllReduce::~AllReduce() { + if (planned_meta_) { + PlannedMeta *meta = reinterpret_cast(planned_meta_); + delete meta; + } +} + +void AllReduce::run() const { + PlannedMeta *meta = reinterpret_cast(planned_meta_); + + detail::checkInfiniccl( + "infinicclAllReduce", + infinicclAllReduce(meta->input->data(), + meta->output->data(), + meta->input->numel(), + detail::toInfinicclDataType(meta->input->dtype()), + meta->op, + meta->communicator, + reinterpret_cast(infinicore::context::getStream()))); +} + +void AllReduce::execute(Tensor output, const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->dtype() == input->dtype()); + INFINICORE_GRAPH_OP_RECORD_OR_RUN(AllReduce, output, input, op, communicator); +} + +Tensor allreduce(const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator) { + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + allreduce_(output, input, op, communicator); + return output; +} + +void allreduce_(Tensor output, const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->dtype() == input->dtype()); + AllReduce::execute(output, input, op, communicator); +} +} // namespace infinicore::op::distributed diff --git a/csrc/infinicore/src/ops/distributed/reduce_scatter.cc b/csrc/infinicore/src/ops/distributed/reduce_scatter.cc new file mode 100644 index 000000000..d0363b302 --- /dev/null +++ b/csrc/infinicore/src/ops/distributed/reduce_scatter.cc @@ -0,0 +1,187 @@ +#include "infinicore/ops/distributed/reduce_scatter.hpp" +#include "../../utils.hpp" +#include "utils.hpp" + +#include "infinicore/context/context.hpp" + +#include +#include +#include + +namespace infinicore::op::distributed { +namespace { + +bool all_equal_to(const std::vector &values, size_t expected) { + return !values.empty() && std::all_of(values.begin(), values.end(), [&](size_t value) { + return value == expected; + }); +} + +void validate_scatter_input(const Tensor &input, + const std::vector &split_sizes) { + INFINICORE_ASSERT(input); + INFINICORE_ASSERT(input->is_contiguous()); + INFINICORE_ASSERT(input->ndim() > 0); + INFINICORE_ASSERT(input->shape()[0] > 0); + INFINICORE_ASSERT(!split_sizes.empty()); + (void)detail::toInfinicclDataType(input->dtype()); + const size_t local_dim0 = split_sizes.front(); + if (!all_equal_to(split_sizes, local_dim0)) { + throw std::runtime_error("InfiniCCL does not support variable-count reduce-scatter"); + } + INFINICORE_ASSERT(local_dim0 > 0); + INFINICORE_ASSERT(local_dim0 <= std::numeric_limits::max() / split_sizes.size()); + INFINICORE_ASSERT(input->shape()[0] == local_dim0 * split_sizes.size()); +} + +void validate_reduce_scatterv(Tensor output, + const Tensor &input, + const std::vector &split_sizes) { + INFINICORE_ASSERT(output && input); + INFINICORE_ASSERT(output->dtype() == input->dtype()); + (void)detail::toInfinicclDataType(input->dtype()); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); + INFINICORE_ASSERT(input->ndim() > 0); + INFINICORE_ASSERT(output->ndim() == input->ndim()); + INFINICORE_ASSERT(input->shape()[0] > 0); + INFINICORE_ASSERT(!split_sizes.empty()); + if (!all_equal_to(split_sizes, output->shape()[0])) { + throw std::runtime_error("InfiniCCL does not support variable-count reduce-scatter"); + } + INFINICORE_ASSERT(output->shape()[0] > 0); + INFINICORE_ASSERT(output->shape()[0] <= std::numeric_limits::max() / split_sizes.size()); + INFINICORE_ASSERT(input->shape()[0] == output->shape()[0] * split_sizes.size()); + for (size_t dim = 1; dim < input->ndim(); ++dim) { + INFINICORE_ASSERT(output->shape()[dim] == input->shape()[dim]); + } + INFINICORE_ASSERT(output->numel() <= std::numeric_limits::max() / split_sizes.size()); + INFINICORE_ASSERT(input->numel() == output->numel() * split_sizes.size()); +} + +} // namespace + +struct ReduceScatterPlannedMeta { + graph::GraphTensor output, input; + infinicclRedOp_t op; + infinicclComm_t communicator; +}; + +ReduceScatter::ReduceScatter(Tensor output, const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->dtype() == input->dtype()); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); + INFINICORE_ASSERT(input->numel() > 0); + INFINICORE_ASSERT(output->numel() > 0); + INFINICORE_ASSERT(input->numel() % output->numel() == 0); + planned_meta_ = new ReduceScatterPlannedMeta{graph::GraphTensor(output), graph::GraphTensor(input), op, communicator}; +} + +ReduceScatter::~ReduceScatter() { + if (planned_meta_) { + auto *meta = reinterpret_cast(planned_meta_); + delete meta; + } +} + +void ReduceScatter::run() const { + auto *meta = reinterpret_cast(planned_meta_); + detail::checkInfiniccl( + "infinicclReduceScatter", + infinicclReduceScatter(meta->input->data(), + meta->output->data(), + meta->output->numel(), + detail::toInfinicclDataType(meta->input->dtype()), + meta->op, + meta->communicator, + reinterpret_cast(infinicore::context::getStream()))); +} + +void ReduceScatter::execute(Tensor output, const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->dtype() == input->dtype()); + INFINICORE_ASSERT(input->numel() > 0); + INFINICORE_ASSERT(output->numel() > 0); + INFINICORE_GRAPH_OP_RECORD_OR_RUN(ReduceScatter, output, input, op, communicator); +} + +Tensor reduce_scatter(const Tensor &input, size_t world_size, infinicclRedOp_t op, infinicclComm_t communicator) { + INFINICORE_ASSERT(input->ndim() > 0); + INFINICORE_ASSERT(world_size > 0); + INFINICORE_ASSERT(input->shape()[0] % world_size == 0); + auto shape = input->shape(); + shape[0] /= world_size; + auto output = Tensor::empty(shape, input->dtype(), input->device()); + reduce_scatter_(output, input, op, communicator); + return output; +} + +void reduce_scatter_(Tensor output, const Tensor &input, infinicclRedOp_t op, infinicclComm_t communicator) { + INFINICORE_ASSERT(output->dtype() == input->dtype()); + ReduceScatter::execute(output, input, op, communicator); +} + +Tensor reduce_scatterv(const Tensor &input, + const std::vector &split_sizes, + size_t rank, + infinicclRedOp_t op, + infinicclComm_t communicator) { + INFINICORE_ASSERT(!split_sizes.empty()); + validate_scatter_input(input, split_sizes); + INFINICORE_ASSERT(rank < split_sizes.size()); + auto shape = input->shape(); + shape[0] = split_sizes[rank]; + auto output = Tensor::empty(shape, input->dtype(), input->device()); + reduce_scatterv_(output, input, split_sizes, op, communicator); + return output; +} + +void reduce_scatterv_(Tensor output, + const Tensor &input, + const std::vector &split_sizes, + infinicclRedOp_t op, + infinicclComm_t communicator) { + validate_reduce_scatterv(output, input, split_sizes); + ReduceScatter::execute(output, input, op, communicator); +} + +std::vector reduce_scatterv_many(const std::vector &inputs, + const std::vector &split_sizes, + size_t rank, + infinicclRedOp_t op, + infinicclComm_t communicator) { + INFINICORE_ASSERT(!split_sizes.empty()); + INFINICORE_ASSERT(rank < split_sizes.size()); + std::vector outputs; + outputs.reserve(inputs.size()); + for (const auto &input : inputs) { + validate_scatter_input(input, split_sizes); + } + for (const auto &input : inputs) { + auto shape = input->shape(); + shape[0] = split_sizes[rank]; + outputs.push_back(Tensor::empty(shape, input->dtype(), input->device())); + } + reduce_scatterv_many_(outputs, inputs, split_sizes, op, communicator); + return outputs; +} + +void reduce_scatterv_many_(const std::vector &outputs, + const std::vector &inputs, + const std::vector &split_sizes, + infinicclRedOp_t op, + infinicclComm_t communicator) { + INFINICORE_ASSERT(outputs.size() == inputs.size()); + INFINICORE_ASSERT(!split_sizes.empty()); + if (inputs.empty()) { + return; + } + + for (size_t i = 0; i < inputs.size(); ++i) { + validate_reduce_scatterv(outputs[i], inputs[i], split_sizes); + } + for (size_t i = 0; i < inputs.size(); ++i) { + ReduceScatter::execute(outputs[i], inputs[i], op, communicator); + } +} + +} // namespace infinicore::op::distributed diff --git a/csrc/infinicore/src/ops/distributed/utils.hpp b/csrc/infinicore/src/ops/distributed/utils.hpp new file mode 100644 index 000000000..87001749b --- /dev/null +++ b/csrc/infinicore/src/ops/distributed/utils.hpp @@ -0,0 +1,56 @@ +#pragma once + +#include "infinicore/dtype.hpp" + +#include + +#include +#include + +namespace infinicore::op::distributed::detail { + +inline infinicclDataType_t toInfinicclDataType(DataType dtype) { + switch (dtype) { + case DataType::kInt8: + return infinicclInt8; + case DataType::kInt16: + return infinicclInt16; + case DataType::kInt32: + return infinicclInt32; + case DataType::kInt64: + return infinicclInt64; + case DataType::kUInt8: + return infinicclUInt8; + case DataType::kUInt16: + return infinicclUInt16; + case DataType::kUInt32: + return infinicclUInt32; + case DataType::kUInt64: + return infinicclUInt64; + case DataType::kFloat16: + return infinicclFloat16; + case DataType::kBFloat16: + return infinicclBFloat16; + case DataType::kFloat32: + return infinicclFloat32; + case DataType::kFloat64: + return infinicclFloat64; + } + throw std::invalid_argument("unsupported data type for InfiniCCL"); +} + +inline void checkInfiniccl(const char *operation, infinicclResult_t result) { + if (result == infinicclSuccess) { + return; + } + + const auto result_string = std::to_string(static_cast(result)); + if (result == infinicclNotSupported) { + throw std::runtime_error("InfiniCCL operation `" + std::string(operation) + + "` is not supported (result " + result_string + ")"); + } + throw std::runtime_error("InfiniCCL operation `" + std::string(operation) + + "` failed with result " + result_string); +} + +} // namespace infinicore::op::distributed::detail diff --git a/csrc/infinicore/src/ops/embedding/embedding.cc b/csrc/infinicore/src/ops/embedding/embedding.cc new file mode 100644 index 000000000..27bca28d3 --- /dev/null +++ b/csrc/infinicore/src/ops/embedding/embedding.cc @@ -0,0 +1,37 @@ +#include "infinicore/ops/embedding.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Embedding); + +Embedding::Embedding(Tensor out, const Tensor &input, const Tensor &weight) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), out, input, weight); +} + +void Embedding::execute(Tensor out, const Tensor &input, const Tensor &weight) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Embedding, out, input, weight); +} + +Tensor embedding(const Tensor &input, // LongTensor of arbitrary shape containing the indices to extract + const Tensor &weight // Weight: Embedding matrix of floating point type with shape (V, embedding_dim), where V = maximum index + 1 +) { + auto input_shape = input->shape(); + auto weight_shape = weight->shape(); + auto embedding_dim = weight_shape[1]; + + // Assign memory to out variables + auto output_shape = input_shape; + output_shape.push_back(embedding_dim); + Tensor inputs_embeds = Tensor::empty(output_shape, weight->dtype(), weight->device()); + + embedding_(inputs_embeds, input, weight); + return inputs_embeds; +} + +void embedding_(Tensor out, const Tensor &input, const Tensor &weight) { + Embedding::execute(out, input, weight); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/embedding/embedding_infiniops.cc b/csrc/infinicore/src/ops/embedding/embedding_infiniops.cc new file mode 100644 index 000000000..fafdba31a --- /dev/null +++ b/csrc/infinicore/src/ops/embedding/embedding_infiniops.cc @@ -0,0 +1,61 @@ +#include "infinicore/ops/embedding.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/embedding.h" + +namespace infinicore::op::embedding_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +struct PlannedMeta { + TensorMeta out, input, weight; + graph::GraphTensor out_tensor, input_tensor, weight_tensor; +}; + +} // namespace + +void *plan(Tensor out, const Tensor &input, const Tensor &weight) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight); + + return new PlannedMeta{ + TensorMeta(out), + TensorMeta(input), + TensorMeta(weight), + graph::GraphTensor(out), + graph::GraphTensor(input), + graph::GraphTensor(weight)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + infini::ops::Embedding::Call( + handle, + config, + planned->input.tensor(planned->input_tensor), + planned->weight.tensor(planned->weight_tensor), + planned->out.tensor(planned->out_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Embedding::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(Embedding::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(Embedding::cleanup_dispatcher(), &cleanup); + return true; +}(); + +} // namespace infinicore::op::embedding_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/flash_attention/flash_attention.cc b/csrc/infinicore/src/ops/flash_attention/flash_attention.cc new file mode 100644 index 000000000..e392d8793 --- /dev/null +++ b/csrc/infinicore/src/ops/flash_attention/flash_attention.cc @@ -0,0 +1,31 @@ +#include "infinicore/ops/flash_attention.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(FlashAttention); + +FlashAttention::FlashAttention(Tensor out, const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k, v); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), + out, q, k, v, total_kv_len, scale, is_causal); +} + +void FlashAttention::execute(Tensor out, const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(FlashAttention, out, q, k, v, total_kv_len, scale, is_causal); +} + +Tensor flash_attention(const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal) { + Shape shape = q->shape(); + int idx = shape.size() - 1; + shape[idx] = v->shape()[idx]; + auto out = Tensor::empty(shape, q->dtype(), q->device()); + flash_attention_(out, q, k, v, total_kv_len, scale, is_causal); + return out; +} + +void flash_attention_(Tensor out, const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal) { + FlashAttention::execute(out, q, k, v, total_kv_len, scale, is_causal); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/flipud/flipud.cc b/csrc/infinicore/src/ops/flipud/flipud.cc new file mode 100644 index 000000000..5b8a807ef --- /dev/null +++ b/csrc/infinicore/src/ops/flipud/flipud.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/flipud.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &Flipud::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +// 2. 静态执行函数 +void Flipud::execute(Tensor output, Tensor input) { + dispatcher().lookup(context::getDevice().type())(output, input); +} +Tensor flipud(Tensor input) { + // Flipud 操作不改变张量的形状和数据类型 + // Output shape == Input shape + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + + flipud_(output, input); + return output; +} +void flipud_(Tensor output, Tensor input) { + Flipud::execute(output, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/float_power/float_power.cc b/csrc/infinicore/src/ops/float_power/float_power.cc new file mode 100644 index 000000000..ac6f65296 --- /dev/null +++ b/csrc/infinicore/src/ops/float_power/float_power.cc @@ -0,0 +1,71 @@ +#include "infinicore/ops/float_power.hpp" +#include "infinicore/tensor.hpp" + +namespace infinicore::op { + +// ======================================================================= +// 1. Dispatcher 单例 +// ======================================================================= + +common::OpDispatcher &FloatPower::dispatcher_scalar() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +common::OpDispatcher &FloatPower::dispatcher_tensor() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +// ======================================================================= +// 2. Execute (执行入口) +// ======================================================================= + +void FloatPower::execute(Tensor output, Tensor input, double exponent) { + dispatcher_scalar() + .lookup(context::getDevice().type())(output, input, exponent); +} + +void FloatPower::execute(Tensor output, Tensor input, Tensor exponent) { + dispatcher_tensor() + .lookup(context::getDevice().type())(output, input, exponent); +} + +// ======================================================================= +// 3. Functional interface (out-of-place) -> 强制提升为 F64 +// ======================================================================= + +Tensor float_power(Tensor input, double exponent) { + auto output = Tensor::empty( + input->shape(), + infinicore::DataType::kFloat64, + input->device()); + + float_power_(output, input, exponent); + return output; +} + +Tensor float_power(Tensor input, Tensor exponent) { + Shape output_shape = input->shape(); + auto output = Tensor::empty( + output_shape, + infinicore::DataType::kFloat64, + input->device()); + + float_power_(output, input, exponent); + return output; +} + +// ======================================================================= +// 4. Explicit / in-place +// ======================================================================= + +void float_power_(Tensor output, Tensor input, double exponent) { + FloatPower::execute(output, input, exponent); +} + +void float_power_(Tensor output, Tensor input, Tensor exponent) { + FloatPower::execute(output, input, exponent); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/floor/floor.cc b/csrc/infinicore/src/ops/floor/floor.cc new file mode 100644 index 000000000..49063040d --- /dev/null +++ b/csrc/infinicore/src/ops/floor/floor.cc @@ -0,0 +1,22 @@ +#include "infinicore/ops/floor.hpp" + +namespace infinicore::op { +common::OpDispatcher &Floor::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; +void Floor::execute(Tensor output, Tensor input) { + dispatcher().lookup(context::getDevice().type())(output, input); +} + +Tensor floor(Tensor input) { + + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + floor_(output, input); + return output; +} +void floor_(Tensor output, Tensor input) { + Floor::execute(output, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/floor_divide/floor_divide.cc b/csrc/infinicore/src/ops/floor_divide/floor_divide.cc new file mode 100644 index 000000000..d7ffeb6d0 --- /dev/null +++ b/csrc/infinicore/src/ops/floor_divide/floor_divide.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/floor_divide.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &FloorDivide::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void FloorDivide::execute(Tensor c, Tensor a, Tensor b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + infinicore::context::setDevice(c->device()); + dispatcher().lookup(c->device().type())(c, a, b); +} + +Tensor floor_divide(Tensor a, Tensor b) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + floor_divide_(c, a, b); + return c; +} + +void floor_divide_(Tensor c, Tensor a, Tensor b) { + FloorDivide::execute(c, a, b); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/fmin/fmin.cc b/csrc/infinicore/src/ops/fmin/fmin.cc new file mode 100644 index 000000000..87f1df3ab --- /dev/null +++ b/csrc/infinicore/src/ops/fmin/fmin.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/fmin.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Fmin::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void Fmin::execute(Tensor c, Tensor a, Tensor b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + infinicore::context::setDevice(c->device()); + dispatcher().lookup(c->device().type())(c, a, b); +} + +Tensor fmin(Tensor a, Tensor b) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + fmin_(c, a, b); + return c; +} + +void fmin_(Tensor c, Tensor a, Tensor b) { + Fmin::execute(c, a, b); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/fmod/fmod.cc b/csrc/infinicore/src/ops/fmod/fmod.cc new file mode 100644 index 000000000..6572e6c94 --- /dev/null +++ b/csrc/infinicore/src/ops/fmod/fmod.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/fmod.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Fmod::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Fmod::execute(Tensor c, Tensor a, Tensor b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + infinicore::context::setDevice(c->device()); + dispatcher().lookup(c->device().type())(c, a, b); +} + +Tensor fmod(Tensor a, Tensor b) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + fmod_(c, a, b); + return c; +} + +void fmod_(Tensor c, Tensor a, Tensor b) { + Fmod::execute(c, a, b); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating.cc b/csrc/infinicore/src/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating.cc new file mode 100644 index 000000000..580b407b3 --- /dev/null +++ b/csrc/infinicore/src/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating.cc @@ -0,0 +1,85 @@ +#include "infinicore/ops/fused_gated_delta_net_gating.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(FusedGatedDeltaNetGating); + +FusedGatedDeltaNetGating::FusedGatedDeltaNetGating(Tensor g, + Tensor beta_output, + const Tensor &A_log, + const Tensor &a, + const Tensor &b, + const Tensor &dt_bias, + float beta, + float threshold) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(g, beta_output, A_log, a, b, dt_bias); + INFINICORE_GRAPH_OP_DISPATCH(g->device().type(), g, beta_output, A_log, a, b, dt_bias, beta, threshold); +} + +void FusedGatedDeltaNetGating::execute(Tensor g, + Tensor beta_output, + const Tensor &A_log, + const Tensor &a, + const Tensor &b, + const Tensor &dt_bias, + float beta, + float threshold) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(FusedGatedDeltaNetGating, g, beta_output, A_log, a, b, dt_bias, beta, threshold); +} + +static void validate_inputs(const Tensor &A_log, + const Tensor &a, + const Tensor &b, + const Tensor &dt_bias) { + if (a->shape().size() != 3 || b->shape().size() != 3) { + throw std::runtime_error("fused_gated_delta_net_gating expects a and b with shape [batch_size, seq_len, hidden]"); + } + if (a->shape() != b->shape()) { + throw std::runtime_error("fused_gated_delta_net_gating expects a and b to have the same shape"); + } + if (A_log->shape().size() != 1 || dt_bias->shape().size() != 1) { + throw std::runtime_error("fused_gated_delta_net_gating expects A_log and dt_bias with shape [hidden]"); + } + if (A_log->shape()[0] != a->shape()[2] || dt_bias->shape()[0] != a->shape()[2]) { + throw std::runtime_error("fused_gated_delta_net_gating hidden dimension mismatch"); + } +} + +std::pair fused_gated_delta_net_gating(const Tensor &A_log, + const Tensor &a, + const Tensor &b, + const Tensor &dt_bias, + float beta, + float threshold) { + validate_inputs(A_log, a, b, dt_bias); + + Tensor g = Tensor::empty(a->shape(), DataType::kFloat32, a->device()); + Tensor beta_output = Tensor::empty(a->shape(), DataType::kFloat32, a->device()); + fused_gated_delta_net_gating_(g, beta_output, A_log, a, b, dt_bias, beta, threshold); + return {g, beta_output}; +} + +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, + float threshold) { + validate_inputs(A_log, a, b, dt_bias); + if (g->shape() != a->shape() || beta_output->shape() != a->shape()) { + throw std::runtime_error("fused_gated_delta_net_gating_ expects outputs with shape [batch_size, seq_len, hidden]"); + } + if (g->dtype() != DataType::kFloat32 || beta_output->dtype() != DataType::kFloat32) { + throw std::runtime_error("fused_gated_delta_net_gating_ expects float32 outputs"); + } + + FusedGatedDeltaNetGating::execute(g, beta_output, A_log, a, b, dt_bias, beta, threshold); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/gaussian_nll_loss/gaussian_nll_loss.cc b/csrc/infinicore/src/ops/gaussian_nll_loss/gaussian_nll_loss.cc new file mode 100644 index 000000000..83ecaafbf --- /dev/null +++ b/csrc/infinicore/src/ops/gaussian_nll_loss/gaussian_nll_loss.cc @@ -0,0 +1,51 @@ +#include "infinicore/ops/gaussian_nll_loss.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(GaussianNllLoss); + +GaussianNllLoss::GaussianNllLoss(Tensor out, + const Tensor &input, + const Tensor &target, + const Tensor &var, + bool full, + double eps, + int reduction) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, target, var); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), out, input, target, var, full, eps, reduction); +} + +void GaussianNllLoss::execute(Tensor out, + const Tensor &input, + const Tensor &target, + const Tensor &var, + bool full, + double eps, + int reduction) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(GaussianNllLoss, out, input, target, var, full, eps, reduction); +} + +Tensor gaussian_nll_loss(const Tensor &input, + const Tensor &target, + const Tensor &var, + bool full, + double eps, + int reduction) { + std::vector out_shape = (reduction == 0) ? input->shape() : std::vector{}; + auto out = Tensor::empty(out_shape, input->dtype(), input->device()); + gaussian_nll_loss_(out, input, target, var, full, eps, reduction); + return out; +} + +void gaussian_nll_loss_(Tensor out, + const Tensor &input, + const Tensor &target, + const Tensor &var, + bool full, + double eps, + int reduction) { + GaussianNllLoss::execute(out, input, target, var, full, eps, reduction); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/gelu/gelu.cc b/csrc/infinicore/src/ops/gelu/gelu.cc new file mode 100644 index 000000000..5921f8f6b --- /dev/null +++ b/csrc/infinicore/src/ops/gelu/gelu.cc @@ -0,0 +1,37 @@ +#include "infinicore/ops/gelu.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &Gelu::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Gelu::execute(Tensor output, Tensor input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Gelu implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor gelu(Tensor input) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + gelu_(output, input); + return output; +} + +void gelu_(Tensor output, Tensor input) { + Gelu::execute(output, input); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/gelu/gelu_infiniops.cc b/csrc/infinicore/src/ops/gelu/gelu_infiniops.cc new file mode 100644 index 000000000..0c9002496 --- /dev/null +++ b/csrc/infinicore/src/ops/gelu/gelu_infiniops.cc @@ -0,0 +1,41 @@ +#include "infinicore/ops/gelu.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/gelu_infinilm.h" + +#include + +namespace infinicore::op::gelu_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +void calculate(Tensor output, Tensor input) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + TensorMeta output_meta(output); + TensorMeta input_meta(input); + infini::ops::GeluInfinilm::Call( + handle, + config, + input_meta.tensor(input), + std::string{"none"}, + output_meta.tensor(output)); +} + +} // namespace + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Gelu::dispatcher(), &calculate); + return true; +}(); + +} // namespace infinicore::op::gelu_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/gelutanh/gelutanh.cc b/csrc/infinicore/src/ops/gelutanh/gelutanh.cc new file mode 100644 index 000000000..95b98c202 --- /dev/null +++ b/csrc/infinicore/src/ops/gelutanh/gelutanh.cc @@ -0,0 +1,37 @@ +#include "infinicore/ops/gelutanh.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &GeluTanh::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void GeluTanh::execute(Tensor output, Tensor input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No GeluTanh implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor gelu_tanh(Tensor input) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + gelu_tanh_(output, input); + return output; +} + +void gelu_tanh_(Tensor output, Tensor input) { + GeluTanh::execute(output, input); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/gelutanh/gelutanh_infiniops.cc b/csrc/infinicore/src/ops/gelutanh/gelutanh_infiniops.cc new file mode 100644 index 000000000..2ee51e6a5 --- /dev/null +++ b/csrc/infinicore/src/ops/gelutanh/gelutanh_infiniops.cc @@ -0,0 +1,38 @@ +#include "infinicore/ops/gelutanh.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/gelutanh_infinilm.h" + +namespace infinicore::op::gelutanh_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +void calculate(Tensor output, Tensor input) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + TensorMeta output_meta(output); + TensorMeta input_meta(input); + infini::ops::GelutanhInfinilm::Call( + handle, + config, + input_meta.tensor(input), + output_meta.tensor(output)); +} + +} // namespace + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(GeluTanh::dispatcher(), &calculate); + return true; +}(); + +} // namespace infinicore::op::gelutanh_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/gemm/gemm.cc b/csrc/infinicore/src/ops/gemm/gemm.cc new file mode 100644 index 000000000..017840c1f --- /dev/null +++ b/csrc/infinicore/src/ops/gemm/gemm.cc @@ -0,0 +1,30 @@ +#include "infinicore/ops/gemm.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Gemm); + +Gemm::Gemm(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + INFINICORE_GRAPH_OP_DISPATCH(c->device().type(), c, a, b, alpha, beta); +} + +void Gemm::execute(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Gemm, c, a, b, alpha, beta); +} + +Tensor gemm(const Tensor &a, const Tensor &b, float alpha, float beta) { + Shape shape = a->shape(); + Size size = a->ndim(); + shape[size - 1] = b->size(size - 1); + auto c = Tensor::empty(shape, a->dtype(), a->device()); + gemm_(c, a, b, alpha, beta); + return c; +} + +void gemm_(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta) { + Gemm::execute(c, a, b, alpha, beta); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/gemm/gemm_infiniops.cc b/csrc/infinicore/src/ops/gemm/gemm_infiniops.cc new file mode 100644 index 000000000..735d2ed0d --- /dev/null +++ b/csrc/infinicore/src/ops/gemm/gemm_infiniops.cc @@ -0,0 +1,68 @@ +#include "infinicore/ops/gemm.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include + +namespace infinicore::op::gemm_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +struct PlannedMeta { + TensorMeta c, a, b; + graph::GraphTensor c_tensor, a_tensor, b_tensor; + float alpha, beta; +}; + +} // namespace + +void *plan(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(c->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + + return new PlannedMeta{ + TensorMeta(c), + TensorMeta(a), + TensorMeta(b), + graph::GraphTensor(c), + graph::GraphTensor(a), + graph::GraphTensor(b), + alpha, + beta}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + infini::ops::Gemm::Call( + handle, + config, + planned->a.tensor(planned->a_tensor), + planned->b.tensor(planned->b_tensor), + std::optional{planned->alpha}, + std::optional{planned->beta}, + std::optional{}, + std::optional{}, + planned->c.tensor(planned->c_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Gemm::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(Gemm::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(Gemm::cleanup_dispatcher(), &cleanup); + return true; +}(); + +} // namespace infinicore::op::gemm_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/hardswish/hardswish.cc b/csrc/infinicore/src/ops/hardswish/hardswish.cc new file mode 100644 index 000000000..29fffa9b3 --- /dev/null +++ b/csrc/infinicore/src/ops/hardswish/hardswish.cc @@ -0,0 +1,38 @@ +#include "infinicore/ops/hardswish.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &Hardswish::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void Hardswish::execute(Tensor output, Tensor input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error( + "No Hardswish implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor hardswish(Tensor input) { + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + hardswish_(output, input); + return output; +} + +void hardswish_(Tensor output, Tensor input) { + Hardswish::execute(output, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/hardtanh/hardtanh.cc b/csrc/infinicore/src/ops/hardtanh/hardtanh.cc new file mode 100644 index 000000000..90649311e --- /dev/null +++ b/csrc/infinicore/src/ops/hardtanh/hardtanh.cc @@ -0,0 +1,38 @@ +#include "infinicore/ops/hardtanh.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &HardTanh::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void HardTanh::execute(Tensor output, Tensor input, float min_val, float max_val) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + if (func == nullptr) { + throw std::runtime_error( + "No HardTanh implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input, min_val, max_val); +} + +Tensor hardtanh(Tensor input, float min_val, float max_val) { + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + hardtanh_(output, input, min_val, max_val); + return output; +} + +void hardtanh_(Tensor output, Tensor input, float min_val, float max_val) { + HardTanh::execute(output, input, min_val, max_val); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/hinge_embedding_loss/hinge_embedding_loss.cc b/csrc/infinicore/src/ops/hinge_embedding_loss/hinge_embedding_loss.cc new file mode 100644 index 000000000..9b76bac85 --- /dev/null +++ b/csrc/infinicore/src/ops/hinge_embedding_loss/hinge_embedding_loss.cc @@ -0,0 +1,57 @@ +#include "infinicore/ops/hinge_embedding_loss.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &HingeEmbeddingLoss::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void HingeEmbeddingLoss::execute( + Tensor output, + Tensor input, + Tensor target, + double margin, + int reduction) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, target); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error( + "No HingeEmbeddingLoss implementation found for device type: " + + std::to_string(static_cast(device_type))); + } + + func(output, input, target, margin, reduction); +} + +Tensor hinge_embedding_loss(Tensor input, Tensor target, double margin, int reduction) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(input, target); + INFINICORE_ASSERT(input->dtype() == target->dtype()); + + Shape output_shape; + if (reduction == 0) { + output_shape = input->shape(); + } + + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + hinge_embedding_loss_(output, input, target, margin, reduction); + return output; +} + +void hinge_embedding_loss_( + Tensor output, + Tensor input, + Tensor target, + double margin, + int reduction) { + HingeEmbeddingLoss::execute(output, input, target, margin, reduction); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/huber_loss/huber_loss.cc b/csrc/infinicore/src/ops/huber_loss/huber_loss.cc new file mode 100644 index 000000000..06cda28c7 --- /dev/null +++ b/csrc/infinicore/src/ops/huber_loss/huber_loss.cc @@ -0,0 +1,35 @@ +#include "infinicore/ops/huber_loss.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &HuberLoss::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void HuberLoss::execute(Tensor output, Tensor input, Tensor target, float delta, int64_t reduction) { + dispatcher().lookup(context::getDevice().type())(output, input, target, delta, reduction); +} + +// 3. 函数式接口 +Tensor huber_loss(Tensor input, Tensor target, float delta, int64_t reduction) { + Shape output_shape; + if (reduction == 0) { // None + // HuberLoss 是 Element-wise 的，reduction='none' 时输出形状通常与输入一致 + output_shape = input->shape(); + } else { + output_shape = {}; // Scalar + } + + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + + huber_loss_(output, input, target, delta, reduction); + return output; +} + +void huber_loss_(Tensor output, Tensor input, Tensor target, float delta, int64_t reduction) { + HuberLoss::execute(output, input, target, delta, reduction); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/hypot/hypot.cc b/csrc/infinicore/src/ops/hypot/hypot.cc new file mode 100644 index 000000000..96a804d30 --- /dev/null +++ b/csrc/infinicore/src/ops/hypot/hypot.cc @@ -0,0 +1,25 @@ +#include "infinicore/ops/hypot.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &Hypot::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Hypot::execute(Tensor output, Tensor input_a, Tensor input_b) { + // lookup 需要传入设备类型，然后调用返回的函数指针 + dispatcher().lookup(context::getDevice().type())(output, input_a, input_b); +} +Tensor hypot(Tensor input_a, Tensor input_b) { + auto output = Tensor::empty(input_a->shape(), input_a->dtype(), input_a->device()); + + hypot_(output, input_a, input_b); + return output; +} +void hypot_(Tensor output, Tensor input_a, Tensor input_b) { + Hypot::execute(output, input_a, input_b); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/index_add/index_add.cc b/csrc/infinicore/src/ops/index_add/index_add.cc new file mode 100644 index 000000000..807cabb5a --- /dev/null +++ b/csrc/infinicore/src/ops/index_add/index_add.cc @@ -0,0 +1,114 @@ +#include "infinicore/ops/index_add.hpp" +#include "infinicore/tensor.hpp" +#include +#include +#include + +namespace infinicore::op { + +common::OpDispatcher &IndexAdd::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void IndexAdd::execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha) { + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No IndexAdd implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input, dim, index, source, alpha); +} + +static void check_index_add_args(const Tensor &input, int64_t &dim, const Tensor &index, const Tensor &source) { + int64_t ndim = static_cast(input->ndim()); + + if (dim < 0) { + dim += ndim; + } + if (dim < 0 || dim >= ndim) { + throw std::runtime_error("IndexAdd: Dimension out of range."); + } + + if (index->ndim() != 1) { + throw std::runtime_error("IndexAdd: Index tensor must be 1D."); + } + + // 使用 DataType::kInt64 和 I32 + if (index->dtype() != DataType::kInt64 && index->dtype() != DataType::kInt32) { + throw std::runtime_error("IndexAdd: Index tensor must be I32 or I64."); + } + + if (source->ndim() != input->ndim()) { + throw std::runtime_error("IndexAdd: Source tensor must have same number of dimensions as input tensor."); + } + + auto in_shape = input->shape(); + auto src_shape = source->shape(); + auto idx_len = index->shape()[0]; + + for (int64_t i = 0; i < ndim; ++i) { + if (i == dim) { + if (src_shape[i] != idx_len) { + throw std::runtime_error("IndexAdd: Source dimension mismatch."); + } + } else { + if (src_shape[i] != in_shape[i]) { + throw std::runtime_error("IndexAdd: Source non-index dimension mismatch."); + } + } + } +} + +// 1. Out-of-place 接口 +Tensor index_add(Tensor input, int64_t dim, Tensor index, Tensor source, float alpha) { + check_index_add_args(input, dim, index, source); + + Tensor output = Tensor::empty(input->shape(), input->dtype(), input->device()); + output->copy_from(input); + if (!index->is_contiguous()) { + index = index->contiguous(); + } + if (!source->is_contiguous()) { + source = source->contiguous(); + } + IndexAdd::execute(output, output, dim, index, source, alpha); + + return output; +} + +// 2. In-place 接口 +void index_add_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha) { + check_index_add_args(input, dim, index, source); + + if (output->shape() != input->shape()) { + throw std::runtime_error("IndexAdd (In-place): Output shape must match Input shape."); + } + + if (output.operator->() != input.operator->()) { + output->copy_from(input); + } + + if (!index->is_contiguous()) { + index = index->contiguous(); + } + if (!source->is_contiguous()) { + source = source->contiguous(); + } + + if (!output->is_contiguous()) { + + Tensor contiguous_out = output->contiguous(); + + IndexAdd::execute(contiguous_out, contiguous_out, dim, index, source, alpha); + + output->copy_from(contiguous_out); + } else { + // 正常路径: Output 已经是连续的，直接原地执行 + IndexAdd::execute(output, input, dim, index, source, alpha); + } +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/index_copy/index_copy.cc b/csrc/infinicore/src/ops/index_copy/index_copy.cc new file mode 100644 index 000000000..52931be74 --- /dev/null +++ b/csrc/infinicore/src/ops/index_copy/index_copy.cc @@ -0,0 +1,114 @@ +#include "infinicore/ops/index_copy.hpp" +#include "infinicore/tensor.hpp" +#include +#include +#include + +namespace infinicore::op { + +// ========================================================= +// Dispatcher & Execute +// ========================================================= + +common::OpDispatcher &IndexCopy::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; +void IndexCopy::execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source) { + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No IndexCopy implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input, dim, index, source); +} +static void check_index_copy_args(const Tensor &input, int64_t &dim, const Tensor &index, const Tensor &source) { + int64_t ndim = static_cast(input->ndim()); + + if (dim < 0) { + dim += ndim; + } + if (dim < 0 || dim >= ndim) { + throw std::runtime_error("IndexCopy: Dimension out of range."); + } + + if (index->ndim() != 1) { + throw std::runtime_error("IndexCopy: Index tensor must be 1D."); + } + + // 使用 DataType::kInt64 和 I32 + if (index->dtype() != DataType::kInt64 && index->dtype() != DataType::kInt32) { + throw std::runtime_error("IndexCopy: Index tensor must be I32 or I64."); + } + + if (source->ndim() != input->ndim()) { + throw std::runtime_error("IndexCopy: Source tensor must have same number of dimensions as input tensor."); + } + + auto in_shape = input->shape(); + auto src_shape = source->shape(); + auto idx_len = index->shape()[0]; + + for (int64_t i = 0; i < ndim; ++i) { + if (i == dim) { + if (src_shape[i] != idx_len) { + throw std::runtime_error("IndexCopy: Source dimension mismatch."); + } + } else { + if (src_shape[i] != in_shape[i]) { + throw std::runtime_error("IndexCopy: Source non-index dimension mismatch."); + } + } + } +} + +Tensor index_copy(Tensor input, int64_t dim, Tensor index, Tensor source) { + check_index_copy_args(input, dim, index, source); + Tensor output = Tensor::empty(input->shape(), input->dtype(), input->device()); + output->copy_from(input); + if (!index->is_contiguous()) { + index = index->contiguous(); + } + if (!source->is_contiguous()) { + source = source->contiguous(); + } + IndexCopy::execute(output, output, dim, index, source); + + return output; +} + +// 2. In-place 接口 +void index_copy_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source) { + check_index_copy_args(input, dim, index, source); + + if (output->shape() != input->shape()) { + throw std::runtime_error("IndexCopy (In-place): Output shape must match Input shape."); + } + + if (output.operator->() != input.operator->()) { + output->copy_from(input); + } + + if (!index->is_contiguous()) { + index = index->contiguous(); + } + if (!source->is_contiguous()) { + source = source->contiguous(); + } + + if (!output->is_contiguous()) { + // 策略: Copy -> Compute -> CopyBack + Tensor contiguous_out = output->contiguous(); + + IndexCopy::execute(contiguous_out, contiguous_out, dim, index, source); + + // 写回结果 + output->copy_from(contiguous_out); + } else { + IndexCopy::execute(output, input, dim, index, source); + } +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/infiniops_impl.hpp b/csrc/infinicore/src/ops/infiniops_impl.hpp new file mode 100644 index 000000000..39bfb5164 --- /dev/null +++ b/csrc/infinicore/src/ops/infiniops_impl.hpp @@ -0,0 +1,103 @@ +#pragma once + +#include "../utils.hpp" +#include "infinicore/tensor.hpp" + +#include + +#include "config.h" +#include "data_type.h" +#include "handle.h" +#include "infini/ops.h" +#include "tensor.h" + +namespace infinicore::op::infiniops { + +inline infini::ops::DataType toInfiniOpsDtype(DataType dtype) { + switch (dtype) { + case DataType::kInt8: + return infini::ops::DataType::kInt8; + case DataType::kInt16: + return infini::ops::DataType::kInt16; + case DataType::kInt32: + return infini::ops::DataType::kInt32; + case DataType::kInt64: + return infini::ops::DataType::kInt64; + case DataType::kUInt8: + return infini::ops::DataType::kUInt8; + case DataType::kUInt16: + return infini::ops::DataType::kUInt16; + case DataType::kUInt32: + return infini::ops::DataType::kUInt32; + case DataType::kUInt64: + return infini::ops::DataType::kUInt64; + case DataType::kFloat16: + return infini::ops::DataType::kFloat16; + case DataType::kBFloat16: + return infini::ops::DataType::kBFloat16; + case DataType::kFloat32: + return infini::ops::DataType::kFloat32; + case DataType::kFloat64: + return infini::ops::DataType::kFloat64; + default: + throw std::runtime_error("InfiniOps backend does not support this tensor dtype."); + } +} + +inline infini::ops::Device toInfiniOpsDevice(const Device &device) { + switch (device.type()) { + case Device::Type::kNvidia: + return infini::ops::Device{infini::ops::Device::Type::kNvidia, static_cast(device.index())}; + case Device::Type::kMetax: + return infini::ops::Device{infini::ops::Device::Type::kMetax, static_cast(device.index())}; + case Device::Type::kMoore: + return infini::ops::Device{infini::ops::Device::Type::kMoore, static_cast(device.index())}; + case Device::Type::kIluvatar: + return infini::ops::Device{infini::ops::Device::Type::kIluvatar, static_cast(device.index())}; + default: + throw std::runtime_error("InfiniOps backend does not support this device type."); + } +} + +inline bool isSupportedDevice(Device::Type device_type) { + switch (device_type) { + case Device::Type::kNvidia: + case Device::Type::kMetax: + case Device::Type::kMoore: + case Device::Type::kIluvatar: + return true; + default: + return false; + } +} + +template +void registerSupportedDevices(Dispatcher &dispatcher, Function function) { + dispatcher.registerDevice(Device::Type::kNvidia, function); + dispatcher.registerDevice(Device::Type::kMetax, function); + dispatcher.registerDevice(Device::Type::kMoore, function); + dispatcher.registerDevice(Device::Type::kIluvatar, function); +} + +struct TensorMeta { + Shape shape; + Strides strides; + infini::ops::DataType dtype; + infini::ops::Device device; + + explicit TensorMeta(const Tensor &tensor) + : shape(tensor->shape()), + strides(tensor->strides()), + dtype(toInfiniOpsDtype(tensor->dtype())), + device(toInfiniOpsDevice(tensor->device())) {} + + infini::ops::Tensor tensor(const void *data) const { + return infini::ops::Tensor(const_cast(data), shape, dtype, device, strides); + } + + infini::ops::Tensor tensor(const Tensor &tensor) const { + return this->tensor(tensor->data()); + } +}; + +} // namespace infinicore::op::infiniops diff --git a/csrc/infinicore/src/ops/inner/inner.cc b/csrc/infinicore/src/ops/inner/inner.cc new file mode 100644 index 000000000..b1881a05f --- /dev/null +++ b/csrc/infinicore/src/ops/inner/inner.cc @@ -0,0 +1,46 @@ +#include +#include + +namespace infinicore::op { + +common::OpDispatcher &Inner::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Inner::execute(Tensor out, Tensor input, Tensor other) { + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Inner implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(out, input, other); +} + +Tensor inner(Tensor input, Tensor other) { + size_t input_ndim = input->ndim(); + size_t other_ndim = other->ndim(); + + assert(input->shape()[input_ndim - 1] == other->shape()[other_ndim - 1]); + + Shape out_shape; + for (int i = 0; i < input_ndim - 1; i++) { + out_shape.push_back(input->shape()[i]); + } + for (int i = 0; i < other_ndim - 1; i++) { + out_shape.push_back(other->shape()[i]); + } + auto out = Tensor::zeros(out_shape, input->dtype(), input->device()); + + inner_(out, input, other); + return out; +} + +void inner_(Tensor out, Tensor input, Tensor other) { + + Inner::execute(out, input, other); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/interpolate/interpolate.cc b/csrc/infinicore/src/ops/interpolate/interpolate.cc new file mode 100644 index 000000000..64fc82540 --- /dev/null +++ b/csrc/infinicore/src/ops/interpolate/interpolate.cc @@ -0,0 +1,145 @@ +#include "infinicore/ops/interpolate.hpp" +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Interpolate); + +Interpolate::Interpolate(Tensor out, + const Tensor &input, + std::string mode, + std::vector size, + std::vector scale_factor, + int align_corners) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), out, input, std::move(mode), std::move(size), std::move(scale_factor), align_corners); +} + +void Interpolate::execute(Tensor out, + const Tensor &input, + std::string mode, + std::vector size, + std::vector scale_factor, + int align_corners) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Interpolate, out, input, std::move(mode), std::move(size), std::move(scale_factor), align_corners); +} + +static std::vector infer_interpolate_shape( + const std::vector &input_shape, + const std::vector &size, + const std::vector &scale_factor) { + if (input_shape.size() < 3) { + throw std::runtime_error("interpolate expects input with at least 3 dimensions"); + } + + const size_t spatial_ndim = input_shape.size() - 2; + std::vector out_shape = input_shape; + + const bool has_size = !size.empty(); + const bool has_scale = !scale_factor.empty(); + if (has_size == has_scale) { + throw std::runtime_error("interpolate expects exactly one of size or scale_factor"); + } + + if (has_size) { + if (size.size() != spatial_ndim) { + throw std::runtime_error("interpolate size dimensionality mismatch"); + } + for (size_t i = 0; i < spatial_ndim; ++i) { + if (size[i] < 0) { + throw std::runtime_error("interpolate size values must be non-negative"); + } + out_shape[i + 2] = static_cast(size[i]); + } + return out_shape; + } + + if (scale_factor.size() != spatial_ndim) { + throw std::runtime_error("interpolate scale_factor dimensionality mismatch"); + } + for (size_t i = 1; i < spatial_ndim; ++i) { + if (scale_factor[i] != scale_factor[0]) { + throw std::runtime_error("interpolate only supports scalar/uniform scale_factor"); + } + } + const double scale = scale_factor[0]; + if (!std::isfinite(scale) || scale < 0.0) { + throw std::runtime_error("interpolate scale_factor must be finite and non-negative"); + } + for (size_t i = 0; i < spatial_ndim; ++i) { + out_shape[i + 2] = static_cast(static_cast(input_shape[i + 2]) * scale); + } + return out_shape; +} + +static void normalize_interpolate_params( + const std::vector &input_shape, + std::vector &size, + std::vector &scale_factor) { + if (input_shape.size() < 3) { + throw std::runtime_error("interpolate expects input with at least 3 dimensions"); + } + + const size_t spatial_ndim = input_shape.size() - 2; + const bool has_size = !size.empty(); + const bool has_scale = !scale_factor.empty(); + if (has_size == has_scale) { + throw std::runtime_error("interpolate expects exactly one of size or scale_factor"); + } + + if (has_size) { + if (size.size() == 1 && spatial_ndim > 1) { + size.assign(spatial_ndim, size[0]); + } + if (size.size() != spatial_ndim) { + throw std::runtime_error("interpolate size dimensionality mismatch"); + } + for (size_t i = 0; i < spatial_ndim; ++i) { + if (size[i] < 0) { + throw std::runtime_error("interpolate size values must be non-negative"); + } + } + return; + } + + if (scale_factor.size() == 1 && spatial_ndim > 1) { + scale_factor.assign(spatial_ndim, scale_factor[0]); + } + if (scale_factor.size() != spatial_ndim) { + throw std::runtime_error("interpolate scale_factor dimensionality mismatch"); + } + for (size_t i = 1; i < spatial_ndim; ++i) { + if (scale_factor[i] != scale_factor[0]) { + throw std::runtime_error("interpolate only supports scalar/uniform scale_factor"); + } + } + if (!std::isfinite(scale_factor[0]) || scale_factor[0] < 0.0) { + throw std::runtime_error("interpolate scale_factor must be finite and non-negative"); + } +} + +Tensor interpolate(const Tensor &input, + std::string mode, + std::vector size, + std::vector scale_factor, + int align_corners) { + normalize_interpolate_params(input->shape(), size, scale_factor); + auto out_shape = infer_interpolate_shape(input->shape(), size, scale_factor); + auto out = Tensor::empty(out_shape, input->dtype(), input->device()); + interpolate_(out, input, std::move(mode), std::move(size), std::move(scale_factor), align_corners); + return out; +} + +void interpolate_(Tensor out, + const Tensor &input, + std::string mode, + std::vector size, + std::vector scale_factor, + int align_corners) { + normalize_interpolate_params(input->shape(), size, scale_factor); + Interpolate::execute(out, input, std::move(mode), std::move(size), std::move(scale_factor), align_corners); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/kron/kron.cc b/csrc/infinicore/src/ops/kron/kron.cc new file mode 100644 index 000000000..06a11c56c --- /dev/null +++ b/csrc/infinicore/src/ops/kron/kron.cc @@ -0,0 +1,50 @@ +#include "infinicore/ops/kron.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &Kron::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void Kron::execute(Tensor output, Tensor a, Tensor b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, a, b); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error( + "No Kron implementation found for device type: " + + std::to_string(static_cast(device_type))); + } + + func(output, a, b); +} + +Tensor kron(Tensor a, Tensor b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(a, b); + INFINICORE_ASSERT(a->dtype() == b->dtype()); + INFINICORE_ASSERT(a->shape().size() == b->shape().size()); + + const auto &a_shape = a->shape(); + const auto &b_shape = b->shape(); + Shape y_shape(a_shape.size()); + for (size_t i = 0; i < a_shape.size(); ++i) { + y_shape[i] = a_shape[i] * b_shape[i]; + } + + auto output = Tensor::empty(y_shape, a->dtype(), a->device()); + kron_(output, a, b); + return output; +} + +void kron_(Tensor output, Tensor a, Tensor b) { + Kron::execute(output, a, b); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/kthvalue/kthvalue.cc b/csrc/infinicore/src/ops/kthvalue/kthvalue.cc new file mode 100644 index 000000000..b912aa9be --- /dev/null +++ b/csrc/infinicore/src/ops/kthvalue/kthvalue.cc @@ -0,0 +1,49 @@ +#include "infinicore/ops/kthvalue.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &Kthvalue::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Kthvalue::execute(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim) { + dispatcher().lookup(context::getDevice().type())(values, indices, input, k, dim, keepdim); +} + +// 3. 函数式接口 +std::tuple kthvalue(Tensor input, int64_t k, int64_t dim, bool keepdim) { + auto input_shape = input->shape(); + int64_t ndim = input_shape.size(); + + // 处理负数维度 + if (dim < 0) { + dim += ndim; + } + + Shape output_shape; + if (keepdim) { + output_shape = input_shape; + output_shape[dim] = 1; + } else { + output_shape.reserve(ndim - 1); + for (int64_t i = 0; i < ndim; ++i) { + if (i != dim) { + output_shape.push_back(input_shape[i]); + } + } + } + + // values 与 input 类型一致 + auto values = Tensor::empty(output_shape, input->dtype(), input->device()); + auto indices = Tensor::empty(output_shape, DataType::kInt64, input->device()); + kthvalue_(values, indices, input, k, dim, keepdim); + return std::make_tuple(values, indices); +} + +void kthvalue_(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim) { + Kthvalue::execute(values, indices, input, k, dim, keepdim); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/kv_caching/kv_caching.cc b/csrc/infinicore/src/ops/kv_caching/kv_caching.cc new file mode 100644 index 000000000..b934b0b0a --- /dev/null +++ b/csrc/infinicore/src/ops/kv_caching/kv_caching.cc @@ -0,0 +1,42 @@ +#include "infinicore/ops/kv_caching.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(KVCaching); + +KVCaching::KVCaching(Tensor k_cache, + Tensor v_cache, + const Tensor &k, + const Tensor &v, + const Tensor &past_kv_lengths) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, past_kv_lengths); + INFINICORE_GRAPH_OP_DISPATCH(k_cache->device().type(), + k_cache, + v_cache, + k, + v, + past_kv_lengths); +} + +void KVCaching::execute(Tensor k_cache, + Tensor v_cache, + const Tensor &k, + const Tensor &v, + const Tensor &past_kv_lengths) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(KVCaching, + k_cache, + v_cache, + k, + v, + past_kv_lengths); +} + +void kv_caching_(Tensor k_cache, + Tensor v_cache, + const Tensor &k, + const Tensor &v, + const Tensor &past_kv_lengths) { + KVCaching::execute(k_cache, v_cache, k, v, past_kv_lengths); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/kv_caching/kv_caching_infiniops.cc b/csrc/infinicore/src/ops/kv_caching/kv_caching_infiniops.cc new file mode 100644 index 000000000..8d1587574 --- /dev/null +++ b/csrc/infinicore/src/ops/kv_caching/kv_caching_infiniops.cc @@ -0,0 +1,50 @@ +#include "infinicore/ops/kv_caching.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/kv_caching_infinilm.h" + +namespace infinicore::op::kv_caching_impl::infiniops { +namespace { +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; +struct PlannedMeta { + TensorMeta k_cache, v_cache, k, v, past_kv_lengths; + graph::GraphTensor k_cache_tensor, v_cache_tensor, k_tensor, v_tensor, past_kv_lengths_tensor; +}; +} // namespace + +void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &past_kv_lengths) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(k_cache->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, past_kv_lengths); + return new PlannedMeta{TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(k), TensorMeta(v), TensorMeta(past_kv_lengths), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(past_kv_lengths)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + infini::ops::KvCachingInfinilm::Call( + handle, + config, + planned->k.tensor(planned->k_tensor), + planned->v.tensor(planned->v_tensor), + planned->past_kv_lengths.tensor(planned->past_kv_lengths_tensor), + planned->k_cache.tensor(planned->k_cache_tensor), + planned->v_cache.tensor(planned->v_cache_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(KVCaching::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(KVCaching::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(KVCaching::cleanup_dispatcher(), &cleanup); + return true; +}(); +} // namespace infinicore::op::kv_caching_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/layer_norm/layer_norm.cc b/csrc/infinicore/src/ops/layer_norm/layer_norm.cc new file mode 100644 index 000000000..faa73a912 --- /dev/null +++ b/csrc/infinicore/src/ops/layer_norm/layer_norm.cc @@ -0,0 +1,42 @@ +#include "infinicore/ops/layer_norm.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(LayerNorm); + +LayerNorm::LayerNorm(Tensor y, Tensor standardization, Tensor std_deviation, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, standardization, std_deviation, x, weight); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), y, standardization, std_deviation, x, weight, bias, epsilon); +} + +void LayerNorm::execute(Tensor y, Tensor standardization, Tensor std_deviation, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(LayerNorm, y, standardization, std_deviation, x, weight, bias, epsilon); +} + +Tensor layer_norm(const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { + auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); + auto reduced_shape = x->shape(); + reduced_shape.pop_back(); + auto standardization = Tensor::empty(x->shape(), x->dtype(), x->device()); + auto std_deviation = Tensor::empty(reduced_shape, x->dtype(), x->device()); + layer_norm_(y, standardization, std_deviation, x, weight, bias, epsilon); + return y; +} + +void layer_norm_(Tensor y, Tensor standardization, Tensor std_deviation, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { + LayerNorm::execute(y, standardization, std_deviation, x, weight, bias, epsilon); +} + +void layer_norm_(Tensor y, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { + auto reduced_shape = x->shape(); + reduced_shape.pop_back(); + auto standardization = Tensor::empty(x->shape(), x->dtype(), x->device()); + auto std_deviation = Tensor::empty(reduced_shape, x->dtype(), x->device()); + LayerNorm::execute(y, standardization, std_deviation, x, weight, bias, epsilon); +} + +void layer_norm_for_pybind(Tensor y, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { + layer_norm_(y, x, weight, bias, epsilon); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/ldexp/ldexp.cc b/csrc/infinicore/src/ops/ldexp/ldexp.cc new file mode 100644 index 000000000..69549f466 --- /dev/null +++ b/csrc/infinicore/src/ops/ldexp/ldexp.cc @@ -0,0 +1,51 @@ +#include "infinicore/ops/ldexp.hpp" +#include // for std::max + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &Ldexp::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Ldexp::execute(Tensor output, Tensor input, Tensor other) { + dispatcher().lookup(context::getDevice().type())(output, input, other); +} + +// 2. 函数式接口 +Tensor ldexp(Tensor input, Tensor other) { + // 计算广播后的输出形状 (Broadcasting Logic) + const auto &shape_a = input->shape(); + const auto &shape_b = other->shape(); + + size_t ndim_a = shape_a.size(); + size_t ndim_b = shape_b.size(); + size_t ndim_out = std::max(ndim_a, ndim_b); + + Shape output_shape(ndim_out); + + // 从后往前对齐维度进行广播检查 + for (size_t i = 0; i < ndim_out; ++i) { + // 获取对应的维度大小，若越界则视为 1 (右对齐) + int64_t dim_a = (i >= ndim_out - ndim_a) ? shape_a[i - (ndim_out - ndim_a)] : 1; + int64_t dim_b = (i >= ndim_out - ndim_b) ? shape_b[i - (ndim_out - ndim_b)] : 1; + output_shape[i] = std::max(dim_a, dim_b); + } + + // 分配输出 Tensor + // ldexp 的输出类型通常跟随 input (尾数)，设备跟随 input + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + + // 调用 Explicit output 接口 + ldexp_(output, input, other); + + return output; +} + +// 3. Explicit Output 接口 +void ldexp_(Tensor output, Tensor input, Tensor other) { + Ldexp::execute(output, input, other); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/lerp/lerp.cc b/csrc/infinicore/src/ops/lerp/lerp.cc new file mode 100644 index 000000000..a88de4023 --- /dev/null +++ b/csrc/infinicore/src/ops/lerp/lerp.cc @@ -0,0 +1,132 @@ +#include "infinicore/ops/lerp.hpp" +#include // for std::max +#include // for std::runtime_error +#include + +namespace infinicore::op { + +// ======================================================================== +// 0. 内部辅助函数：手动实现形状广播推导 +// ======================================================================== +namespace { + +Shape compute_broadcast_shape(const std::vector &shapes) { + if (shapes.empty()) { + return {}; + } + + // 1. 找出最大的维度数 (Max Rank) + size_t max_ndim = 0; + for (const auto &shape : shapes) { + max_ndim = std::max(max_ndim, shape.size()); + } + + Shape out_shape(max_ndim); + + // 2. 从右向左遍历每一个维度 (Standard Broadcasting Rule) + for (size_t i = 0; i < max_ndim; ++i) { + size_t current_dim_val = 1; + bool set = false; + + for (const auto &shape : shapes) { + // 计算当前 shape 对应的维度索引 (从右对齐) + // 比如 max_ndim=4, 当前 shape_ndim=2, i=0 (最右边) + // shape index = 2 - 1 - 0 = 1 + if (i < shape.size()) { + size_t dim = shape[shape.size() - 1 - i]; + + if (dim == 1) { + continue; // 1 可以被广播，忽略 + } + + if (!set) { + current_dim_val = dim; + set = true; + } else if (current_dim_val != dim) { + // 维度不相等，且都不为 1，无法广播 + throw std::runtime_error( + "Lerp: Shapes are not broadcastable. Mismatch at dimension offset " + std::to_string(i)); + } + } + } + // 填充输出形状 (从右向左填，或者填好后由 vector 自动管理) + out_shape[max_ndim - 1 - i] = current_dim_val; + } + + return out_shape; +} + +} // namespace + +// ======================================================================== +// 1. 定义 Dispatcher 单例 +// ======================================================================== + +template <> +common::OpDispatcher &Lerp::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +template <> +common::OpDispatcher &Lerp::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +// ======================================================================== +// 2. Execute 静态方法实现 +// ======================================================================== + +void Lerp::execute(Tensor output, Tensor start, Tensor end, Tensor weight) { + dispatcher().lookup(context::getDevice().type())(output, start, end, weight); +} + +void Lerp::execute(Tensor output, Tensor start, Tensor end, float weight) { + dispatcher().lookup(context::getDevice().type())(output, start, end, weight); +} + +// ======================================================================== +// 3. 函数式接口 (Functional API) - 集成形状推导 +// ======================================================================== + +Tensor lerp(Tensor start, Tensor end, Tensor weight) { + // 1. 调用本地实现的推导函数，计算 start, end, weight 三者的广播形状 + Shape output_shape = compute_broadcast_shape({start->shape(), + end->shape(), + weight->shape()}); + + // 2. 分配输出内存 + auto output = Tensor::empty(output_shape, start->dtype(), start->device()); + + // 3. 执行计算 + lerp_(output, start, end, weight); + return output; +} + +Tensor lerp(Tensor start, Tensor end, float weight) { + // 1. 计算 start, end 两者的广播形状 (标量 weight 不参与形状计算) + Shape output_shape = compute_broadcast_shape({start->shape(), + end->shape()}); + + // 2. 分配输出内存 + auto output = Tensor::empty(output_shape, start->dtype(), start->device()); + + // 3. 执行计算 + lerp_(output, start, end, weight); + return output; +} + +// ======================================================================== +// 4. In-place / Output-buffer 接口 +// ======================================================================== + +void lerp_(Tensor output, Tensor start, Tensor end, Tensor weight) { + Lerp::execute(output, start, end, weight); +} + +void lerp_(Tensor output, Tensor start, Tensor end, float weight) { + Lerp::execute(output, start, end, weight); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/linear/linear.cc b/csrc/infinicore/src/ops/linear/linear.cc new file mode 100644 index 000000000..976bf0bf9 --- /dev/null +++ b/csrc/infinicore/src/ops/linear/linear.cc @@ -0,0 +1,60 @@ +#include "infinicore/ops/linear.hpp" +#include "infinicore/ops/gemm.hpp" +#include "infinicore/ops/rearrange.hpp" + +namespace infinicore::op { + +Tensor linear(Tensor input, + Tensor weight, + std::optional bias, + float alpha) { + + Size ndim = input->ndim(); + Size out_features = weight->shape()[0]; + + // Assign memory to out variables + auto output_shape = input->shape(); + output_shape[ndim - 1] = out_features; + auto out = Tensor::empty(output_shape, input->dtype(), input->device()); + + // Inplace Calculate + linear_(out, input, weight, bias, alpha); + return out; +} + +void linear_(Tensor out, + Tensor input, + Tensor weight, + std::optional bias, + float alpha) { + + auto weight_shape = weight->shape(); + Size out_features = weight_shape[0]; + Size in_features = weight_shape[1]; + + Size ndim = input->ndim(); + assert(out->ndim() == ndim); + + // Calculate the number of features + Size N = 1; + auto input_shape = input->shape(); + for (size_t i = 0; i < ndim - 1; ++i) { + N *= input_shape[i]; + } + + // linear transformation + Tensor out_view = out->view({N, out_features}); + // Add bias + float beta = 0.0f; + if (bias.has_value()) { + rearrange_(out_view, + bias.value()->as_strided({N, out_features}, {0, 1})); + beta = 1.0f; + } + + gemm_(out_view, + input->view({N, in_features}), + weight->permute({1, 0}), alpha, beta); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/linear_w8a8i8/linear_w8a8i8.cc b/csrc/infinicore/src/ops/linear_w8a8i8/linear_w8a8i8.cc new file mode 100644 index 000000000..f3774382e --- /dev/null +++ b/csrc/infinicore/src/ops/linear_w8a8i8/linear_w8a8i8.cc @@ -0,0 +1,66 @@ +#include "infinicore/ops/linear_w8a8i8.hpp" +#include "infinicore/ops/per_channel_quant_i8.hpp" +#include "infinicore/ops/scaled_mm_i8.hpp" + +namespace infinicore::op { + +Tensor linear_w8a8i8(Tensor input, + Tensor weight_packed, + Tensor weight_scale, + std::optional bias) { + + // Input is of shape [M, K], Weight_packed is of shape [N, K],stirdes is [N, 1] + Size ndim = input->ndim(); + Size out_features = weight_packed->shape()[0]; + + // Assign memory to out variables + auto output_shape = input->shape(); + output_shape[ndim - 1] = out_features; + auto out = Tensor::empty(output_shape, input->dtype(), input->device()); + + // Inplace Calculate + linear_w8a8i8_(out, input, weight_packed, weight_scale, bias); + return out; +} + +void linear_w8a8i8_(Tensor out, + Tensor input, + Tensor weight_packed, + Tensor weight_scale, + std::optional bias) { + + auto weight_packed_shape = weight_packed->shape(); + Size out_features = weight_packed_shape[0]; + Size in_features = weight_packed_shape[1]; + + Size ndim = input->ndim(); + assert(out->ndim() == ndim); + + Size N = 1; + auto input_shape = input->shape(); + for (size_t i = 0; i < ndim - 1; ++i) { + N *= input_shape[i]; + } + + auto input_packed = Tensor::empty( + {N, input_shape[ndim - 1]}, + DataType::kInt8, + input->device()); + auto input_scale = Tensor::empty( + {N, 1}, + DataType::kFloat32, + input->device()); + op::per_channel_quant_i8_(input->view({N, in_features}), input_packed, input_scale); + if (bias.has_value()) { + bias = std::make_optional(bias.value()->as_strided({N, out_features}, {0, 1})); + } + op::scaled_mm_i8_( + out->view({N, out_features}), + input_packed, + input_scale, + weight_packed->permute({1, 0}), + weight_scale, + bias); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/log_softmax/log_softmax.cc b/csrc/infinicore/src/ops/log_softmax/log_softmax.cc new file mode 100644 index 000000000..6c23eb42b --- /dev/null +++ b/csrc/infinicore/src/ops/log_softmax/log_softmax.cc @@ -0,0 +1,34 @@ +#include "infinicore/ops/log_softmax.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &LogSoftmax::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void LogSoftmax::execute(Tensor output, Tensor input, int64_t dim) { + dispatcher().lookup(context::getDevice().type())(output, input, dim); +} + +// 3. 函数式接口 +Tensor log_softmax(Tensor input, int64_t dim) { + int64_t ndim = input->shape().size(); + + // 处理负数维度 + if (dim < 0) { + dim += ndim; + } + + // LogSoftmax 输出形状与输入一致，dtype 与 input 一致 + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + log_softmax_(output, input, dim); + return output; +} + +void log_softmax_(Tensor output, Tensor input, int64_t dim) { + LogSoftmax::execute(output, input, dim); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/logaddexp/logaddexp.cc b/csrc/infinicore/src/ops/logaddexp/logaddexp.cc new file mode 100644 index 000000000..c9f1a9a99 --- /dev/null +++ b/csrc/infinicore/src/ops/logaddexp/logaddexp.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/logaddexp.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &LogAddExp::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void LogAddExp::execute(Tensor c, Tensor a, Tensor b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + infinicore::context::setDevice(c->device()); + dispatcher().lookup(c->device().type())(c, a, b); +} + +Tensor logaddexp(Tensor a, Tensor b) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + logaddexp_(c, a, b); + return c; +} + +void logaddexp_(Tensor c, Tensor a, Tensor b) { + LogAddExp::execute(c, a, b); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/logaddexp2/logaddxep2.cc b/csrc/infinicore/src/ops/logaddexp2/logaddxep2.cc new file mode 100644 index 000000000..07333c575 --- /dev/null +++ b/csrc/infinicore/src/ops/logaddexp2/logaddxep2.cc @@ -0,0 +1,27 @@ +#include "../../utils.hpp" +#include "infinicore/ops/logaddexp2.hpp" + +namespace infinicore::op { + +common::OpDispatcher &LogAddExp2::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void LogAddExp2::execute(Tensor c, Tensor a, Tensor b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + infinicore::context::setDevice(c->device()); + dispatcher().lookup(c->device().type())(c, a, b); +} + +Tensor logaddexp2(Tensor a, Tensor b) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + logaddexp2_(c, a, b); + return c; +} + +void logaddexp2_(Tensor c, Tensor a, Tensor b) { + LogAddExp2::execute(c, a, b); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/logcumsumexp/logcumsumexp.cc b/csrc/infinicore/src/ops/logcumsumexp/logcumsumexp.cc new file mode 100644 index 000000000..80b007b8f --- /dev/null +++ b/csrc/infinicore/src/ops/logcumsumexp/logcumsumexp.cc @@ -0,0 +1,35 @@ +#include "infinicore/ops/logcumsumexp.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +// 初始化 Dispatcher 单例 +common::OpDispatcher &LogCumSumExp::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +// 算子执行逻辑：校验设备并分发任务 +void LogCumSumExp::execute(Tensor y, Tensor x, int axis, bool exclusive, bool reverse) { + // 确保输入输出张量在同一设备上 + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + + // 切换到目标设备的上下文 + infinicore::context::setDevice(y->device()); + dispatcher().lookup(y->device().type())(y, x, axis, exclusive, reverse); +} + +// 函数式接口：自动创建输出张量并返回 +Tensor logcumsumexp(Tensor x, int axis, bool exclusive, bool reverse) { + // 创建一个与输入 x 形状、类型和设备相同的空张量作为输出 + auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); + logcumsumexp_(y, x, axis, exclusive, reverse); + return y; +} + +// 原地/指定输出接口 +void logcumsumexp_(Tensor y, Tensor x, int axis, bool exclusive, bool reverse) { + LogCumSumExp::execute(y, x, axis, exclusive, reverse); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/logdet/logdet.cc b/csrc/infinicore/src/ops/logdet/logdet.cc new file mode 100644 index 000000000..bc6a76d04 --- /dev/null +++ b/csrc/infinicore/src/ops/logdet/logdet.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/logdet.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Logdet); + +Logdet::Logdet(Tensor y, const Tensor &x) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), y, x); +} + +void Logdet::execute(Tensor y, const Tensor &x) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Logdet, y, x); +} + +Tensor logdet(const Tensor &x) { + auto y = Tensor::empty({}, x->dtype(), x->device()); + logdet_(y, x); + return y; +} + +void logdet_(Tensor y, const Tensor &x) { + Logdet::execute(y, x); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/mamba_selective_scan/mamba_selective_scan.cc b/csrc/infinicore/src/ops/mamba_selective_scan/mamba_selective_scan.cc new file mode 100644 index 000000000..666831868 --- /dev/null +++ b/csrc/infinicore/src/ops/mamba_selective_scan/mamba_selective_scan.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/mamba_selective_scan.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MambaSelectiveScan); + +MambaSelectiveScan::MambaSelectiveScan(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) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, x, dt, b, c, a_log, d, gate, dt_bias, state); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), out, x, dt, b, c, a_log, d, gate, dt_bias, state); +} +void MambaSelectiveScan::execute(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) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(MambaSelectiveScan, out, x, dt, b, c, a_log, d, gate, dt_bias, state); +} +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) { + auto output = Tensor::empty(x->shape(), x->dtype(), x->device()); + mamba_selective_scan_(output, x, dt, b, c, a_log, d, gate, dt_bias, state); + return output; +} +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) { + MambaSelectiveScan::execute(out, x, dt, b, c, a_log, d, gate, dt_bias, state); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/masked_select/masked_select.cc b/csrc/infinicore/src/ops/masked_select/masked_select.cc new file mode 100644 index 000000000..2aed7560a --- /dev/null +++ b/csrc/infinicore/src/ops/masked_select/masked_select.cc @@ -0,0 +1,33 @@ +#include +#include + +namespace infinicore::op { + +common::OpDispatcher &MaskedSelect::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void MaskedSelect::execute(Tensor input, Tensor mask, void **data_ptr, size_t *dlen_ptr) { + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No MaskedSelect implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(input, mask, data_ptr, dlen_ptr); +} + +Tensor masked_select(Tensor input, Tensor mask) { + + std::byte *data; + size_t dlen; + MaskedSelect::execute(input, mask, (void **)&data, &dlen); + + auto out = Tensor::from_blob(data, {dlen}, input->dtype(), input->device()); + + return out; +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/matmul/matmul.cc b/csrc/infinicore/src/ops/matmul/matmul.cc new file mode 100644 index 000000000..58119bc10 --- /dev/null +++ b/csrc/infinicore/src/ops/matmul/matmul.cc @@ -0,0 +1,13 @@ +#include "infinicore/ops/matmul.hpp" +#include "infinicore/ops/gemm.hpp" + +namespace infinicore::op { + +Tensor matmul(Tensor a, Tensor b, float alpha) { + return gemm(a, b, alpha, 0.0f); +} + +void matmul_(Tensor c, Tensor a, Tensor b, float alpha) { + Gemm::execute(c, a, b, alpha, 0.0f); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc b/csrc/infinicore/src/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc new file mode 100644 index 000000000..d4d27a2f9 --- /dev/null +++ b/csrc/infinicore/src/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc @@ -0,0 +1,167 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_FLASH_ATTN) +#include "infinicore/ops/mha_kvcache.hpp" + +#include "../../../adaptor/flash_attn/hygon/flash_attn_hygon.hpp" +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include +#include +#include +#include + +#include +#include + +namespace infinicore::op::mha_kvcache_impl::flashattn { + +struct PlannedMeta { + graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; + std::optional alibi_slopes; + float scale; +}; + +void *plan(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) { + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k_cache), + graph::GraphTensor(v_cache), + graph::GraphTensor(seqlens_k), + graph::GraphTensor(block_table), + alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, + scale}; +} + +void run(void *planned_meta) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + // Paged KV caches must be contiguous for flash-attn; avoid extra copies for q/metadata when already dense. + const bool out_need_copy_back = !p->out->is_contiguous(); + Tensor out_work = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); + auto out_tensor = infinicore::adaptor::to_aten_tensor(out_work); + auto q = infinicore::adaptor::to_aten_tensor(p->q); + auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache); + auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache); + auto seqlens_k = std::optional(infinicore::adaptor::to_aten_tensor(p->seqlens_k)); + auto block_table = std::optional(infinicore::adaptor::to_aten_tensor(p->block_table)); + auto alibi_slopes = p->alibi_slopes + ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) + : std::nullopt; + + if (std::getenv("INFINICORE_HYGON_ATEN_FALLBACK")) { + namespace idx = at::indexing; + auto seqlens_t = infinicore::adaptor::to_aten_tensor(p->seqlens_k); + auto block_table_t = infinicore::adaptor::to_aten_tensor(p->block_table); + auto seqlens_cpu = seqlens_t.to(at::kCPU); + auto block_table_cpu = block_table_t.to(at::kCPU); + + auto result = at::empty_like(out_tensor); + const int64_t batch_size = q.size(0); + const int64_t seqlen_q = q.size(1); + const int64_t num_heads = q.size(2); + const int64_t block_size = k_cache.size(1); + const int64_t num_kv_heads = k_cache.size(2); + const int64_t group_size = num_heads / num_kv_heads; + + for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + const int64_t seq_len = seqlens_cpu.index({batch_idx}).item(); + std::vector keys; + std::vector values; + keys.reserve(seq_len); + values.reserve(seq_len); + for (int64_t logical_pos = 0; logical_pos < seq_len; ++logical_pos) { + const int64_t block_id = block_table_cpu.index({batch_idx, logical_pos / block_size}).item(); + const int64_t off = logical_pos % block_size; + keys.push_back(k_cache.index({block_id, off, idx::Slice(), idx::Slice()})); + values.push_back(v_cache.index({block_id, off, idx::Slice(), idx::Slice()})); + } + auto K = at::stack(keys, 0); + auto V = at::stack(values, 0); + if (group_size > 1) { + K = K.repeat_interleave(group_size, 1); + V = V.repeat_interleave(group_size, 1); + } + auto cur_q = q.index({batch_idx}); + auto scores = at::matmul(cur_q.permute({1, 0, 2}).to(at::kFloat), K.permute({1, 2, 0}).to(at::kFloat)) * p->scale; + auto mask = at::full({seqlen_q, seq_len}, -std::numeric_limits::infinity(), q.options().dtype(at::kFloat)); + const int64_t prefix_len = seq_len - seqlen_q; + for (int64_t query_pos = 0; query_pos < seqlen_q; ++query_pos) { + mask.index_put_({query_pos, idx::Slice(0, prefix_len + query_pos + 1)}, 0.0); + } + auto attn = at::softmax(scores + mask.unsqueeze(0), -1).to(q.dtype()); + auto cur_out = at::matmul(attn, V.permute({1, 0, 2})).permute({1, 0, 2}); + result.index_put_({batch_idx}, cur_out); + } + + out_tensor.copy_(result); + if (out_need_copy_back) { + p->out->copy_from(out_work); + } + return; + } + + std::optional k_new = std::nullopt; + std::optional v_new = std::nullopt; + std::optional rotary_cos = std::nullopt; + std::optional rotary_sin = std::nullopt; + std::optional cache_batch_idx = std::nullopt; + std::optional leftpad_k = std::nullopt; + const bool use_dynamic_out = q.dim() == 4 && k_cache.dim() == 4 + && q.size(1) == 1 && q.size(2) > k_cache.size(2) + && q.size(3) % 8 == 0 && !alibi_slopes.has_value(); + + auto out = use_dynamic_out ? std::optional(std::nullopt) + : std::optional(out_tensor); + + auto result = flash::mha_fwd_kvcache( + q, + k_cache, + v_cache, + k_new, + v_new, + seqlens_k, + rotary_cos, + rotary_sin, + cache_batch_idx, + leftpad_k, + block_table, + alibi_slopes, + out, + p->scale, + true, + -1, + -1, + 0.0f, + false, + 0); + + if (!result.empty() && result[0].defined()) { + out_tensor.copy_(result[0]); + } + if (out_need_copy_back) { + p->out->copy_from(out_work); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MhaKVCache::plan_dispatcher().registerDevice(Device::Type::kHygon, &plan); + MhaKVCache::run_dispatcher().registerDevice(Device::Type::kHygon, &run); + MhaKVCache::cleanup_dispatcher().registerDevice(Device::Type::kHygon, &cleanup); + return true; +}(); + +} // namespace infinicore::op::mha_kvcache_impl::flashattn +#endif // ENABLE_HYGON_API && ENABLE_FLASH_ATTN diff --git a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc new file mode 100644 index 000000000..5d24d8e09 --- /dev/null +++ b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc @@ -0,0 +1,58 @@ +#include "infinicore/ops/mha_kvcache.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MhaKVCache); + +MhaKVCache::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) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, seqlens_k, block_table); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), + out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, scale); +} + +void MhaKVCache::execute(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) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MhaKVCache, + out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, 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) { + MhaKVCache::execute(out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, 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) { + // Output shape matches q: [batch_size, seqlen_q, num_heads, head_size] + auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); + mha_kvcache_(out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, scale); + return out; +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_flashattn.cc b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_flashattn.cc new file mode 100644 index 000000000..0167c17df --- /dev/null +++ b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_flashattn.cc @@ -0,0 +1,136 @@ +#include "infinicore/ops/mha_kvcache.hpp" + +#include "infinicore/adaptor/flash_attention_adaptor.hpp" + +#include + +#ifdef ENABLE_FLASH_ATTN +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#include +#endif +#endif + +#if defined(ENABLE_METAX_API) +#define INFINICORE_FLASH_OP(name) ::name +#else +#define INFINICORE_FLASH_OP(name) flash::name +#endif + +namespace infinicore::op::mha_kvcache_impl::flashattn { + +struct PlannedMeta { + graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; + std::optional alibi_slopes; + float scale; +}; + +void *plan(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) { + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k_cache), + graph::GraphTensor(v_cache), + graph::GraphTensor(seqlens_k), + graph::GraphTensor(block_table), + alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, + scale}; +} + +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()); +#endif + auto *p = reinterpret_cast(planned_meta); + + // Paged KV caches must be contiguous for flash-attn; avoid extra copies for q/metadata when already dense. + const bool out_need_copy_back = !p->out->is_contiguous(); + Tensor out_work = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); + auto out_tensor = infinicore::adaptor::to_aten_tensor(out_work); + auto q = infinicore::adaptor::to_aten_tensor(p->q); +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) + auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache); + auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache); +#elif defined(ENABLE_QY_API) + Tensor k_cache_work = p->k_cache->contiguous(); + Tensor v_cache_work = p->v_cache->contiguous(); + auto k_cache = infinicore::adaptor::to_aten_tensor(k_cache_work); + auto v_cache = infinicore::adaptor::to_aten_tensor(v_cache_work); +#endif + auto seqlens_k = std::optional(infinicore::adaptor::to_aten_tensor(p->seqlens_k)); + auto block_table = std::optional(infinicore::adaptor::to_aten_tensor(p->block_table)); + auto alibi_slopes = p->alibi_slopes + ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) + : std::nullopt; + + std::optional k_new = std::nullopt; + std::optional v_new = std::nullopt; + std::optional rotary_cos = std::nullopt; + std::optional rotary_sin = std::nullopt; + std::optional cache_batch_idx = std::nullopt; + std::optional leftpad_k = std::nullopt; + + const bool use_dynamic_out = q.dim() == 4 && k_cache.dim() == 4 + && q.size(1) == 1 && q.size(2) > k_cache.size(2) + && q.size(3) % 8 == 0 && !alibi_slopes.has_value(); + + auto out = use_dynamic_out ? std::optional(std::nullopt) + : std::optional(out_tensor); + +#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) + std::optional flash_attn_mars_ext = std::nullopt; +#endif + + auto result = INFINICORE_FLASH_OP(mha_fwd_kvcache)( + q, + k_cache, + v_cache, + k_new, + v_new, + seqlens_k, + rotary_cos, + rotary_sin, + cache_batch_idx, + leftpad_k, + block_table, + alibi_slopes, + out, + p->scale, + true, + -1, + -1, + 0.0f, + false, + 0 +#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) + , + flash_attn_mars_ext +#endif + ); + + if (use_dynamic_out) { + out_tensor.copy_(result[0]); + } + if (out_need_copy_back) { + p->out->copy_from(out_work); + } +#else + throw std::runtime_error("FlashAttention is not enabled in this build"); +#endif +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MhaKVCache, &plan, &run, &cleanup); + +} // namespace infinicore::op::mha_kvcache_impl::flashattn diff --git a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_flashattn_moore.cc b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_flashattn_moore.cc new file mode 100644 index 000000000..8de343096 --- /dev/null +++ b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_flashattn_moore.cc @@ -0,0 +1,127 @@ +#if defined(ENABLE_MOORE_MATE_FLASH_ATTN) + +#include "infinicore/ops/mha_kvcache.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include +#include +#include +#include + +namespace infinicore::op::mha_kvcache_impl::flashattn_moore { + +namespace py = pybind11; + +// Lightweight RAII: Binds MUSA streams, +// avoiding the need to include +namespace { +class LocalMUSAStreamGuard { +public: + explicit LocalMUSAStreamGuard(const c10::musa::MUSAStream &s) + : prev_(c10::musa::getCurrentMUSAStream(s.device_index())) { + c10::musa::setCurrentMUSAStream(s); + } + ~LocalMUSAStreamGuard() { + c10::musa::setCurrentMUSAStream(prev_); + } + LocalMUSAStreamGuard(const LocalMUSAStreamGuard &) = delete; + LocalMUSAStreamGuard &operator=(const LocalMUSAStreamGuard &) = delete; + +private: + c10::musa::MUSAStream prev_; +}; +} // namespace + +struct PlannedMeta { + graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; + std::optional alibi_slopes; + float scale; +}; + +void *plan(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) { + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k_cache), + graph::GraphTensor(v_cache), + graph::GraphTensor(seqlens_k), + graph::GraphTensor(block_table), + alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, + scale}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + if (p->alibi_slopes.has_value()) { + throw std::runtime_error( + "[mha_kvcache/moore] ALiBi not supported by mate flash_attn_with_kvcache"); + } + + LocalMUSAStreamGuard guard(infinicore::adaptor::get_musa_stream()); + + auto out_tensor = infinicore::adaptor::to_aten_tensor(p->out); + auto q_4d = infinicore::adaptor::to_aten_tensor(p->q); + auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache); + auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache); + auto seqlens_k = infinicore::adaptor::to_aten_tensor(p->seqlens_k); + auto block_table = infinicore::adaptor::to_aten_tensor(p->block_table); + + auto q_3d = q_4d.squeeze(1); + + const int64_t block_size = k_cache.size(1); + const int64_t max_seq_len = block_table.size(1) * block_size; + + try { + py::gil_scoped_acquire gil; + py::module_ wrapper = py::module_::import("infinicore.ops.moore_mate_flash_attn"); + + py::object py_q = py::cast(q_3d); + py::object py_k_cache = py::cast(k_cache); + py::object py_v_cache = py::cast(v_cache); + py::object py_seqlens_k = py::cast(seqlens_k); + py::object py_blk_tbl = py::cast(block_table); + + py::object result = wrapper.attr("moore_mate_flash_attn_decode")( + py_q, + py_k_cache, + py_v_cache, + py_blk_tbl, + py_seqlens_k, + p->scale, + block_size, + max_seq_len); + + at::Tensor result_t = result.cast(); + out_tensor.copy_(result_t.unsqueeze(1)); + + result = py::none(); + py_q = py_k_cache = py_v_cache = py_seqlens_k = py_blk_tbl = py::none(); + } catch (const py::error_already_set &e) { + throw std::runtime_error( + std::string("[mha_kvcache/moore] Python error: ") + e.what()); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MhaKVCache::plan_dispatcher().registerDevice(Device::Type::kMoore, &plan); + MhaKVCache::run_dispatcher().registerDevice(Device::Type::kMoore, &run); + MhaKVCache::cleanup_dispatcher().registerDevice(Device::Type::kMoore, &cleanup); + return true; +}(); + +} // namespace infinicore::op::mha_kvcache_impl::flashattn_moore + +#endif // ENABLE_MOORE_MATE_FLASH_ATTN diff --git a/csrc/infinicore/src/ops/moe_align/moe_align.cc b/csrc/infinicore/src/ops/moe_align/moe_align.cc new file mode 100644 index 000000000..27635932c --- /dev/null +++ b/csrc/infinicore/src/ops/moe_align/moe_align.cc @@ -0,0 +1,155 @@ +#include "infinicore/ops/moe_align.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeAlign); +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeAlignWithExpertMap); + +MoeAlign::MoeAlign(Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + const size_t num_experts, + const size_t block_size, + const bool pad_sorted_token_ids) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(sorted_token_ids, expert_ids, num_tokens_post_padded, topk_ids); + INFINICORE_GRAPH_OP_DISPATCH( + sorted_token_ids->device().type(), + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + num_experts, + block_size, + pad_sorted_token_ids); +} + +void MoeAlign::execute(Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + const size_t num_experts, + const size_t block_size, + const bool pad_sorted_token_ids) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MoeAlign, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + num_experts, + block_size, + pad_sorted_token_ids); +} + +MoeAlignWithExpertMap::MoeAlignWithExpertMap(Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + const Tensor &expert_map, + const size_t num_experts, + const size_t block_size, + const bool pad_sorted_token_ids) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(sorted_token_ids, expert_ids, num_tokens_post_padded, topk_ids, expert_map); + INFINICORE_GRAPH_OP_DISPATCH( + sorted_token_ids->device().type(), + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + expert_map, + num_experts, + block_size, + pad_sorted_token_ids); +} + +void MoeAlignWithExpertMap::execute(Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + const Tensor &expert_map, + const size_t num_experts, + const size_t block_size, + const bool pad_sorted_token_ids) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MoeAlignWithExpertMap, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + expert_map, + num_experts, + block_size, + pad_sorted_token_ids); +} + +std::tuple moe_align( + const Tensor &topk_ids, + size_t num_experts, + size_t block_size, + bool pad_sorted_token_ids) { + auto shape = topk_ids->shape(); + INFINICORE_ASSERT(shape.size() == 2); + const size_t numel = shape[0] * shape[1]; + const size_t align_num_experts = num_experts + 1; + const size_t max_num_tokens_padded = numel < align_num_experts + ? numel * block_size + : numel + align_num_experts * (block_size - 1); + const size_t sorted_token_ids_capacity = ((max_num_tokens_padded + 3) / 4) * 4; + const size_t max_num_blocks = (max_num_tokens_padded + block_size - 1) / block_size; + + auto sorted_token_ids = Tensor::empty({sorted_token_ids_capacity}, DataType::kInt32, topk_ids->device()); + auto expert_ids = Tensor::empty({max_num_blocks}, DataType::kInt32, topk_ids->device()); + auto num_tokens_post_padded = Tensor::empty({1}, DataType::kInt32, topk_ids->device()); + + moe_align_( + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + num_experts, + block_size, + pad_sorted_token_ids); + + return {sorted_token_ids, expert_ids, num_tokens_post_padded}; +} + +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) { + MoeAlign::execute( + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + num_experts, + block_size, + pad_sorted_token_ids); +} + +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) { + MoeAlignWithExpertMap::execute( + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + expert_map, + num_experts, + block_size, + pad_sorted_token_ids); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc b/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc new file mode 100644 index 000000000..b0f4f9111 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc @@ -0,0 +1,73 @@ +#include "infinicore/ops/moe_fused_dense.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeFusedDense); + +MoeFusedDense::MoeFusedDense(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) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + output, hidden_states, w13, w2, topk_weights, topk_ids, + sorted_token_ids, expert_ids, num_tokens_post_padded); + INFINICORE_GRAPH_OP_DISPATCH( + output->device().type(), output, hidden_states, w13, w2, topk_weights, topk_ids, + sorted_token_ids, expert_ids, num_tokens_post_padded); +} + +void MoeFusedDense::execute(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) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MoeFusedDense, output, hidden_states, w13, w2, topk_weights, topk_ids, + sorted_token_ids, expert_ids, num_tokens_post_padded); +} + +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) { + auto shape = hidden_states->shape(); + INFINICORE_ASSERT(shape.size() == 2); + auto output = Tensor::empty(shape, hidden_states->dtype(), hidden_states->device()); + moe_fused_dense_( + output, hidden_states, w13, w2, topk_weights, topk_ids, + sorted_token_ids, expert_ids, num_tokens_post_padded); + return output; +} + +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) { + MoeFusedDense::execute( + output, hidden_states, w13, w2, topk_weights, topk_ids, + sorted_token_ids, expert_ids, num_tokens_post_padded); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/moe_fused_gate/moe_fused_gate.cc b/csrc/infinicore/src/ops/moe_fused_gate/moe_fused_gate.cc new file mode 100644 index 000000000..a32c4fef0 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_fused_gate/moe_fused_gate.cc @@ -0,0 +1,101 @@ +#include "infinicore/ops/moe_fused_gate.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeFusedGate); + +MoeFusedGate::MoeFusedGate(Tensor topk_weights, + Tensor topk_indices, + const Tensor &input, + const Tensor &bias, + const size_t num_expert_group, + const size_t topk_group, + const size_t num_fused_shared_experts, + const float routed_scaling_factor, + const bool apply_routed_scaling_factor_on_output) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_indices, input, bias); + INFINICORE_GRAPH_OP_DISPATCH( + topk_weights->device().type(), + topk_weights, + topk_indices, + input, + bias, + num_expert_group, + topk_group, + num_fused_shared_experts, + routed_scaling_factor, + apply_routed_scaling_factor_on_output); +} + +void MoeFusedGate::execute(Tensor topk_weights, + Tensor topk_indices, + const Tensor &input, + const Tensor &bias, + const size_t num_expert_group, + const size_t topk_group, + const size_t num_fused_shared_experts, + const float routed_scaling_factor, + const bool apply_routed_scaling_factor_on_output) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MoeFusedGate, + topk_weights, + topk_indices, + input, + bias, + num_expert_group, + topk_group, + num_fused_shared_experts, + routed_scaling_factor, + apply_routed_scaling_factor_on_output); +} + +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) { + auto shape = input->shape(); + INFINICORE_ASSERT(shape.size() == 2); + auto topk_weights = Tensor::empty({shape[0], topk}, DataType::kFloat32, input->device()); + auto topk_indices = Tensor::empty({shape[0], topk}, DataType::kInt32, input->device()); + moe_fused_gate_( + topk_weights, + topk_indices, + input, + bias, + num_expert_group, + topk_group, + num_fused_shared_experts, + routed_scaling_factor, + apply_routed_scaling_factor_on_output); + return {topk_weights, topk_indices}; +} + +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) { + MoeFusedGate::execute( + topk_weights, + topk_indices, + input, + bias, + num_expert_group, + topk_group, + num_fused_shared_experts, + routed_scaling_factor, + apply_routed_scaling_factor_on_output); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/moe_sum/moe_sum.cc b/csrc/infinicore/src/ops/moe_sum/moe_sum.cc new file mode 100644 index 000000000..b676aa074 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_sum/moe_sum.cc @@ -0,0 +1,30 @@ +#include "infinicore/ops/moe_sum.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeSum); + +MoeSum::MoeSum(Tensor output, const Tensor &input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + INFINICORE_GRAPH_OP_DISPATCH(output->device().type(), output, input); +} + +void MoeSum::execute(Tensor output, const Tensor &input) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(MoeSum, output, input); +} + +Tensor moe_sum(const Tensor &input) { + auto shape = input->shape(); + INFINICORE_ASSERT(shape.size() == 3); + auto output = Tensor::empty({shape[0], shape[2]}, input->dtype(), input->device()); + moe_sum_(output, input); + return output; +} + +void moe_sum_(Tensor output, const Tensor &input) { + MoeSum::execute(output, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/moe_topk_sigmoid/moe_topk_sigmoid.cc b/csrc/infinicore/src/ops/moe_topk_sigmoid/moe_topk_sigmoid.cc new file mode 100644 index 000000000..f9428d732 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_topk_sigmoid/moe_topk_sigmoid.cc @@ -0,0 +1,63 @@ +#include "infinicore/ops/moe_topk_sigmoid.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeTopkSigmoid); + +MoeTopkSigmoid::MoeTopkSigmoid(Tensor topk_weights, + Tensor topk_indices, + const Tensor &gating_output, + const Tensor &correction_bias, + const bool renormalize) { + if (correction_bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_indices, gating_output, correction_bias); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_indices, gating_output); + } + INFINICORE_GRAPH_OP_DISPATCH( + topk_weights->device().type(), + topk_weights, + topk_indices, + gating_output, + correction_bias, + renormalize); +} + +void MoeTopkSigmoid::execute(Tensor topk_weights, + Tensor topk_indices, + const Tensor &gating_output, + const Tensor &correction_bias, + const bool renormalize) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MoeTopkSigmoid, + topk_weights, + topk_indices, + gating_output, + correction_bias, + renormalize); +} + +std::tuple moe_topk_sigmoid( + const Tensor &gating_output, + size_t topk, + bool renormalize, + const Tensor &correction_bias) { + auto shape = gating_output->shape(); + INFINICORE_ASSERT(shape.size() == 2); + auto topk_weights = Tensor::empty({shape[0], topk}, DataType::kFloat32, gating_output->device()); + auto topk_indices = Tensor::empty({shape[0], topk}, DataType::kInt32, gating_output->device()); + moe_topk_sigmoid_(topk_weights, topk_indices, gating_output, correction_bias, renormalize); + return {topk_weights, topk_indices}; +} + +void moe_topk_sigmoid_(Tensor topk_weights, + Tensor topk_indices, + const Tensor &gating_output, + const Tensor &correction_bias, + bool renormalize) { + MoeTopkSigmoid::execute(topk_weights, topk_indices, gating_output, correction_bias, renormalize); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/moe_topk_softmax/moe_topk_softmax.cc b/csrc/infinicore/src/ops/moe_topk_softmax/moe_topk_softmax.cc new file mode 100644 index 000000000..ecb3ba4a1 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_topk_softmax/moe_topk_softmax.cc @@ -0,0 +1,69 @@ +#include "infinicore/ops/moe_topk_softmax.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeTopkSoftmax); + +MoeTopkSoftmax::MoeTopkSoftmax(Tensor topk_weights, + Tensor topk_indices, + const Tensor &gating_output, + const Tensor &correction_bias, + const bool renormalize, + const float moe_softcapping) { + if (correction_bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_indices, gating_output, correction_bias); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_indices, gating_output); + } + INFINICORE_GRAPH_OP_DISPATCH( + topk_weights->device().type(), + topk_weights, + topk_indices, + gating_output, + correction_bias, + renormalize, + moe_softcapping); +} + +void MoeTopkSoftmax::execute(Tensor topk_weights, + Tensor topk_indices, + const Tensor &gating_output, + const Tensor &correction_bias, + const bool renormalize, + const float moe_softcapping) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MoeTopkSoftmax, + topk_weights, + topk_indices, + gating_output, + correction_bias, + renormalize, + moe_softcapping); +} + +std::tuple moe_topk_softmax( + const Tensor &gating_output, + size_t topk, + bool renormalize, + float moe_softcapping, + const Tensor &correction_bias) { + auto shape = gating_output->shape(); + INFINICORE_ASSERT(shape.size() == 2); + auto topk_weights = Tensor::empty({shape[0], topk}, DataType::kFloat32, gating_output->device()); + auto topk_indices = Tensor::empty({shape[0], topk}, DataType::kInt32, gating_output->device()); + moe_topk_softmax_(topk_weights, topk_indices, gating_output, correction_bias, renormalize, moe_softcapping); + return {topk_weights, topk_indices}; +} + +void moe_topk_softmax_(Tensor topk_weights, + Tensor topk_indices, + const Tensor &gating_output, + const Tensor &correction_bias, + bool renormalize, + float moe_softcapping) { + MoeTopkSoftmax::execute(topk_weights, topk_indices, gating_output, correction_bias, renormalize, moe_softcapping); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/mrope/mrope.cc b/csrc/infinicore/src/ops/mrope/mrope.cc new file mode 100644 index 000000000..a4694acf8 --- /dev/null +++ b/csrc/infinicore/src/ops/mrope/mrope.cc @@ -0,0 +1,100 @@ +#include "infinicore/ops/mrope.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MRoPE); + +MRoPE::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) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q_out, k_out, q, k, cos, sin, positions); + INFINICORE_GRAPH_OP_DISPATCH(q->device().type(), + q_out, + k_out, + q, + k, + cos, + sin, + positions, + head_size, + rotary_dim, + section_t, + section_h, + section_w, + interleaved); +} + +void MRoPE::execute(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) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(MRoPE, + q_out, + k_out, + q, + k, + cos, + sin, + positions, + head_size, + rotary_dim, + section_t, + section_h, + section_w, + interleaved); +} + +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) { + MRoPE::execute(q_out, k_out, q, k, cos, sin, positions, head_size, rotary_dim, section_t, section_h, section_w, 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) { + auto q_out = Tensor::empty(q->shape(), q->dtype(), q->device()); + auto k_out = Tensor::empty(k->shape(), k->dtype(), k->device()); + mrope_(q_out, k_out, q, k, cos, sin, positions, head_size, rotary_dim, section_t, section_h, section_w, interleaved); + return {q_out, k_out}; +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/mul/mul.cc b/csrc/infinicore/src/ops/mul/mul.cc new file mode 100644 index 000000000..0aa457222 --- /dev/null +++ b/csrc/infinicore/src/ops/mul/mul.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/mul.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Mul); + +Mul::Mul(Tensor c, const Tensor &a, const Tensor &b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + INFINICORE_GRAPH_OP_DISPATCH(c->device().type(), c, a, b); +} + +void Mul::execute(Tensor c, const Tensor &a, const Tensor &b) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Mul, c, a, b); +} + +Tensor mul(const Tensor &a, const Tensor &b) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + mul_(c, a, b); + return c; +} + +void mul_(Tensor c, const Tensor &a, const Tensor &b) { + Mul::execute(c, a, b); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/mul_scalar/mul_scalar.cc b/csrc/infinicore/src/ops/mul_scalar/mul_scalar.cc new file mode 100644 index 000000000..283b31d4a --- /dev/null +++ b/csrc/infinicore/src/ops/mul_scalar/mul_scalar.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/mul_scalar.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MulScalar); + +MulScalar::MulScalar(Tensor c, const Tensor &a, double alpha) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a); + INFINICORE_GRAPH_OP_DISPATCH(c->device().type(), c, a, alpha); +} + +void MulScalar::execute(Tensor c, const Tensor &a, double alpha) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(MulScalar, c, a, alpha); +} + +Tensor mul_scalar(const Tensor &a, double alpha) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + mul_scalar_(c, a, alpha); + return c; +} + +void mul_scalar_(Tensor c, const Tensor &a, double alpha) { + MulScalar::execute(c, a, alpha); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/multi_head_attention/hygon/mha_flashattn_hygon.cc b/csrc/infinicore/src/ops/multi_head_attention/hygon/mha_flashattn_hygon.cc new file mode 100644 index 000000000..1b3734289 --- /dev/null +++ b/csrc/infinicore/src/ops/multi_head_attention/hygon/mha_flashattn_hygon.cc @@ -0,0 +1,102 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_FLASH_ATTN) +#include "infinicore/ops/mha.hpp" + +#include "../../../adaptor/flash_attn/hygon/flash_attn_hygon.hpp" +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include +#include + +#include + +namespace infinicore::op::mha_impl::flashattn { +struct PlannedMeta { + graph::GraphTensor out, q, k, v; + std::optional alibi_slopes; + float scale; + bool is_causal; +}; + +void *plan(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + std::optional alibi_slopes, + float scale, + bool is_causal) { + + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k), + graph::GraphTensor(v), + alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, + scale, + is_causal}; +} + +void run(void *planned_meta) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + auto q = infinicore::adaptor::to_aten_tensor(p->q); + auto k = infinicore::adaptor::to_aten_tensor(p->k); + auto v = infinicore::adaptor::to_aten_tensor(p->v); + + const bool out_need_copy_back = !p->out->is_contiguous(); + Tensor out_work_ic = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); + auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic); + auto scale = p->scale; + auto is_causal = p->is_causal; + + const int64_t num_q_heads = q.size(2); + const int64_t num_kv_heads = k.size(2); + if (q.dim() == 4 && k.dim() == 4 && v.dim() == 4 && q.size(1) == 1 && k.size(1) == 1 && v.size(1) == 1) { + at::Tensor value = v; + if (num_kv_heads != num_q_heads) { + const int64_t repeat = num_q_heads / num_kv_heads; + value = v.unsqueeze(3).repeat({1, 1, 1, repeat, 1}).reshape(q.sizes()); + } + out_work.copy_(value); + if (out_need_copy_back) { + p->out->copy_from(out_work_ic); + } + return; + } + + at::Tensor k_attn = k; + at::Tensor v_attn = v; + if (num_kv_heads != num_q_heads) { + const int64_t repeat = num_q_heads / num_kv_heads; + k_attn = k.unsqueeze(3).repeat({1, 1, 1, repeat, 1}).reshape(q.sizes()); + v_attn = v.unsqueeze(3).repeat({1, 1, 1, repeat, 1}).reshape(q.sizes()); + } + + auto attn_weight = at::matmul(q.permute({0, 2, 1, 3}), k_attn.permute({0, 2, 3, 1})) * scale; + if (is_causal) { + auto mask = at::tril(at::ones_like(attn_weight), -1).flip({-2, -1}); + auto neg_inf = at::full_like(attn_weight, -std::numeric_limits::infinity()); + attn_weight = at::where(mask == 1, neg_inf, attn_weight); + } + auto attn = at::softmax(attn_weight, -1); + auto result = at::matmul(attn, v_attn.permute({0, 2, 1, 3})).permute({0, 2, 1, 3}); + out_work.copy_(result); + if (out_need_copy_back) { + p->out->copy_from(out_work_ic); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MultiheadAttention::plan_dispatcher().registerDevice(Device::Type::kHygon, &plan); + MultiheadAttention::run_dispatcher().registerDevice(Device::Type::kHygon, &run); + MultiheadAttention::cleanup_dispatcher().registerDevice(Device::Type::kHygon, &cleanup); + return true; +}(); + +} // namespace infinicore::op::mha_impl::flashattn +#endif // ENABLE_HYGON_API && ENABLE_FLASH_ATTN diff --git a/csrc/infinicore/src/ops/multi_head_attention/mha.cc b/csrc/infinicore/src/ops/multi_head_attention/mha.cc new file mode 100644 index 000000000..895037b3b --- /dev/null +++ b/csrc/infinicore/src/ops/multi_head_attention/mha.cc @@ -0,0 +1,54 @@ +#include "infinicore/ops/mha.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MultiheadAttention); + +MultiheadAttention::MultiheadAttention(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + std::optional alibi_slopes, + float scale, + bool is_causal) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k, v); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), + out, q, k, v, alibi_slopes, scale, is_causal); +} + +void MultiheadAttention::execute(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + std::optional alibi_slopes, + float scale, + bool is_causal) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MultiheadAttention, + out, q, k, v, alibi_slopes, scale, is_causal); +} + +Tensor mha( + const Tensor &q, + const Tensor &k, + const Tensor &v, + std::optional alibi_slopes, + float scale, + bool is_causal) { + auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); + mha_(out, q, k, v, alibi_slopes, scale, is_causal); + return out; +} + +void mha_(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + std::optional alibi_slopes, + float scale, + bool is_causal) { + MultiheadAttention::execute(out, q, k, v, alibi_slopes, scale, is_causal); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/multi_head_attention/mha_flashattn.cc b/csrc/infinicore/src/ops/multi_head_attention/mha_flashattn.cc new file mode 100644 index 000000000..13c96b94d --- /dev/null +++ b/csrc/infinicore/src/ops/multi_head_attention/mha_flashattn.cc @@ -0,0 +1,115 @@ +#include "infinicore/ops/mha.hpp" + +#include "infinicore/adaptor/flash_attention_adaptor.hpp" + +#include + +#ifdef ENABLE_FLASH_ATTN +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#include +#endif +#endif + +namespace infinicore::op::mha_impl::flashattn { + +struct PlannedMeta { + graph::GraphTensor out, q, k, v; + std::optional alibi_slopes; + float scale; + bool is_causal; +}; + +void *plan(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + std::optional alibi_slopes, + float scale, + bool is_causal) { + + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k), + graph::GraphTensor(v), + alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, + scale, + is_causal}; +} + +namespace { + +// Only support nv for now +#if defined(ENABLE_FLASH_ATTN) && defined(ENABLE_NVIDIA_API) +// MetaX/hpcc pip `flash_attn_2_cuda` exports `mha_fwd` at global scope (no namespace), +// while NVIDIA `flash-attn-nvidia.so` uses `flash::mha_fwd`. +#if defined(ENABLE_METAX_API) +#define INFINICORE_FLASH_OP(name) ::name +#else +#define INFINICORE_FLASH_OP(name) flash::name +#endif + +#endif // ENABLE_FLASH_ATTN +} // namespace + +void run(void *planned_meta) { +// Only support nv for now +#if defined(ENABLE_FLASH_ATTN) && defined(ENABLE_NVIDIA_API) + c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); + auto *p = reinterpret_cast(planned_meta); + + auto q = infinicore::adaptor::to_aten_tensor(p->q); + auto k = infinicore::adaptor::to_aten_tensor(p->k); + auto v = infinicore::adaptor::to_aten_tensor(p->v); + + const bool out_need_copy_back = !p->out->is_contiguous(); + Tensor out_work_ic = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); + auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic); + auto out = std::optional(out_work); + + auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; + auto scale = p->scale; + auto is_causal = p->is_causal; + +#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) + std::optional flash_attn_mars_ext = std::nullopt; +#endif + + INFINICORE_FLASH_OP(mha_fwd) + ( + q, + k, + v, + out, + alibi_slopes, + 0.0, + scale, + is_causal, + -1, + -1, + 0.0, + false, + std::nullopt +#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) + , + flash_attn_mars_ext +#endif + ); + + if (out_need_copy_back) { + p->out->copy_from(out_work_ic); + } + +#else + throw std::runtime_error("FlashAttention is not enabled in this build"); +#endif +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MultiheadAttention, &plan, &run, &cleanup); + +} // namespace infinicore::op::mha_impl::flashattn diff --git a/csrc/infinicore/src/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc b/csrc/infinicore/src/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc new file mode 100644 index 000000000..278df60e9 --- /dev/null +++ b/csrc/infinicore/src/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc @@ -0,0 +1,266 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_FLASH_ATTN) +#include "infinicore/ops/mha_varlen.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" +#include +#include +#include +#endif + +#include "../../../adaptor/flash_attn/hygon/flash_attn_hygon.hpp" + +#include +#include +#include +#include + +namespace infinicore::op::mha_varlen_impl::flashattn { + +struct PlannedMeta { + graph::GraphTensor out, q, k, v, cum_seqlens_q, cum_seqlens_k; + std::optional block_table; + int max_seqlen_q, max_seqlen_k; + std::optional alibi_slopes; + float scale; +}; + +void *plan(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) { + + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k), + graph::GraphTensor(v), + graph::GraphTensor(cum_seqlens_q), + graph::GraphTensor(cum_seqlens_k), + block_table ? std::optional(graph::GraphTensor(*block_table)) : std::nullopt, + max_seqlen_q, + max_seqlen_k, + alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, + scale}; +} + +void run(void *planned_meta) { +#ifndef ENABLE_ATEN + (void)planned_meta; + throw std::runtime_error("ATen is not enabled in this build"); +#else + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + auto q = infinicore::adaptor::to_aten_tensor(p->q); + auto k = infinicore::adaptor::to_aten_tensor(p->k); + auto v = infinicore::adaptor::to_aten_tensor(p->v); + + const bool out_need_copy_back = !p->out->is_contiguous(); + Tensor out_work_ic = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); + auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic); + + auto cu_seqlens_q = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_q); + auto cu_seqlens_kv = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_k); + + const bool dense_sdpa = !p->block_table.has_value() + && !p->alibi_slopes.has_value() + && q.dim() == 3 && k.dim() == 3 && v.dim() == 3 + && p->max_seqlen_q > 0 && p->max_seqlen_k > 0 + && p->max_seqlen_q == p->max_seqlen_k + && cu_seqlens_q.dim() == 1 + && cu_seqlens_q.size(0) == cu_seqlens_kv.size(0) + && q.size(0) == (cu_seqlens_q.size(0) - 1) * p->max_seqlen_q + && k.size(0) == (cu_seqlens_kv.size(0) - 1) * p->max_seqlen_k + && ((q.size(2) > 256) || (v.size(2) != q.size(2))); + if (dense_sdpa) { + const int64_t batch_size = cu_seqlens_q.size(0) - 1; + const int64_t seqlen = p->max_seqlen_q; + const int64_t num_heads = q.size(1); + const int64_t num_kv_heads = k.size(1); + const int64_t head_dim = q.size(2); + const int64_t value_dim = v.size(2); + auto q_4d = q.reshape({batch_size, seqlen, num_heads, head_dim}).permute({0, 2, 1, 3}); + auto k_4d = k.reshape({batch_size, seqlen, num_kv_heads, head_dim}).permute({0, 2, 1, 3}); + auto v_4d = v.reshape({batch_size, seqlen, num_kv_heads, value_dim}).permute({0, 2, 1, 3}); + if (num_heads != num_kv_heads) { + if (num_heads % num_kv_heads != 0) { + throw std::runtime_error("mha_varlen dense SDPA fallback requires num_heads to be divisible by num_kv_heads"); + } + const int64_t groups = num_heads / num_kv_heads; + k_4d = k_4d.unsqueeze(2).expand({batch_size, num_kv_heads, groups, seqlen, head_dim}).reshape({batch_size, num_heads, seqlen, head_dim}); + v_4d = v_4d.unsqueeze(2).expand({batch_size, num_kv_heads, groups, seqlen, value_dim}).reshape({batch_size, num_heads, seqlen, value_dim}); + } + auto result = at::scaled_dot_product_attention( + q_4d, + k_4d, + v_4d, + std::nullopt, + 0.0, + true, + std::optional(static_cast(p->scale))); + out_work.copy_(result.permute({0, 2, 1, 3}).reshape({q.size(0), num_heads, value_dim})); + if (out_need_copy_back) { + p->out->copy_from(out_work_ic); + } + return; + } + + if (std::getenv("INFINICORE_HYGON_ATEN_FALLBACK")) { + namespace idx = at::indexing; + auto cu_q_cpu = cu_seqlens_q.to(at::kCPU); + auto cu_k_cpu = cu_seqlens_kv.to(at::kCPU); + const int64_t num_seqs = cu_q_cpu.size(0) - 1; + auto result = at::zeros_like(out_work); + + if (!p->block_table.has_value()) { + for (int64_t i = 0; i < num_seqs; ++i) { + const int64_t q_start = cu_q_cpu.index({i}).item(); + const int64_t q_end = cu_q_cpu.index({i + 1}).item(); + const int64_t k_start = cu_k_cpu.index({i}).item(); + const int64_t k_end = cu_k_cpu.index({i + 1}).item(); + auto cur_q = q.index({idx::Slice(q_start, q_end)}).unsqueeze(0).transpose(1, 2); + auto cur_k = k.index({idx::Slice(k_start, k_end)}).unsqueeze(0).transpose(1, 2); + auto cur_v = v.index({idx::Slice(k_start, k_end)}).unsqueeze(0).transpose(1, 2); + auto cur_out = at::scaled_dot_product_attention( + cur_q, cur_k, cur_v, std::nullopt, 0.0, true, std::optional(static_cast(p->scale))); + result.index_put_({idx::Slice(q_start, q_end)}, cur_out.transpose(1, 2).squeeze(0)); + } + } else { + auto block_table_t = infinicore::adaptor::to_aten_tensor(*p->block_table); + auto block_table_cpu = block_table_t.to(at::kCPU); + const int64_t block_size = k.size(1); + for (int64_t i = 0; i < num_seqs; ++i) { + const int64_t q_start = cu_q_cpu.index({i}).item(); + const int64_t q_end = cu_q_cpu.index({i + 1}).item(); + const int64_t q_len = q_end - q_start; + const int64_t h_len = (cu_k_cpu.index({i + 1}).item() - cu_k_cpu.index({i}).item()) - q_len; + const int64_t total_len = h_len + q_len; + auto cur_q = q.index({idx::Slice(q_start, q_end)}); + std::vector keys; + std::vector values; + keys.reserve(total_len); + values.reserve(total_len); + for (int64_t j = 0; j < total_len; ++j) { + const int64_t b_id = block_table_cpu.index({i, j / block_size}).item(); + const int64_t off = j % block_size; + keys.push_back(k.index({b_id, off, idx::Slice(), idx::Slice()})); + values.push_back(v.index({b_id, off, idx::Slice(), idx::Slice()})); + } + auto K = at::stack(keys, 0); + auto V = at::stack(values, 0); + const int64_t q_heads = cur_q.size(1); + const int64_t kv_heads = K.size(1); + if (q_heads != kv_heads) { + const int64_t repeat = q_heads / kv_heads; + K = K.repeat_interleave(repeat, 1); + V = V.repeat_interleave(repeat, 1); + } + auto scores = at::matmul(cur_q.permute({1, 0, 2}).to(at::kFloat), K.permute({1, 2, 0}).to(at::kFloat)) * p->scale; + auto mask = at::full({q_len, total_len}, -std::numeric_limits::infinity(), q.options().dtype(at::kFloat)); + for (int64_t t = 0; t < q_len; ++t) { + mask.index_put_({t, idx::Slice(0, h_len + t + 1)}, 0.0); + } + auto attn = at::softmax(scores + mask.unsqueeze(0), -1).to(q.dtype()); + auto cur_out = at::matmul(attn, V.permute({1, 0, 2})).permute({1, 0, 2}); + result.index_put_({idx::Slice(q_start, q_end)}, cur_out); + } + } + + out_work.copy_(result); + if (out_need_copy_back) { + p->out->copy_from(out_work_ic); + } + return; + } + + auto out = std::optional(out_work); + std::optional seqused_k = std::nullopt; + std::optional leftpad_k = std::nullopt; + auto block_table = p->block_table ? std::optional(infinicore::adaptor::to_aten_tensor(*p->block_table)) : std::nullopt; + auto max_seqlen_q = p->max_seqlen_q; + auto max_seqlen_k = p->max_seqlen_k; + auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; + auto scale = p->scale; + + if (alibi_slopes.has_value()) { + throw std::runtime_error("[mha_varlen/hygon] ALiBi is not supported by the direct libflash_attention varlen_fwd ABI"); + } + auto q_work = q.contiguous(); + auto k_work = k.contiguous(); + auto v_work = v.contiguous(); + if (block_table.has_value() && k.dim() == 4 && v.dim() == 4) { + const int64_t num_blocks = k.size(0); + const int64_t block_size = k.size(1); + const int64_t num_kv_heads = k.size(2); + const int64_t head_dim = k.size(3); + if (block_size % 64 != 0) { + throw std::runtime_error("[mha_varlen/hygon] flash-attn requires paged KV block size to be divisible by 64"); + } + const int64_t pages_per_block = block_size / 64; + k_work = k_work.reshape({num_blocks, pages_per_block, 64, num_kv_heads, head_dim}) + .reshape({num_blocks * pages_per_block, 64, num_kv_heads, head_dim}) + .contiguous(); + v_work = v_work.reshape({num_blocks, pages_per_block, 64, num_kv_heads, head_dim}) + .reshape({num_blocks * pages_per_block, 64, num_kv_heads, head_dim}) + .contiguous(); + if (pages_per_block != 1) { + auto offsets = at::arange(pages_per_block, block_table->options()).view({1, 1, pages_per_block}); + block_table = ((*block_table).unsqueeze(-1) * pages_per_block + offsets) + .reshape({block_table->size(0), block_table->size(1) * pages_per_block}) + .contiguous(); + } + } + auto result = flash::vllm_mha_varlen_fwd( + q_work, + k_work, + v_work, + out, + cu_seqlens_q, + cu_seqlens_kv, + seqused_k, + leftpad_k, + block_table, + alibi_slopes, + max_seqlen_q, + max_seqlen_k, + 0.0, + scale, + false, + true, + -1, + -1, + 0.0, + false, + std::nullopt); + if (!result.empty() && result[0].defined()) { + out_work.copy_(result[0]); + } + if (out_need_copy_back) { + p->out->copy_from(out_work_ic); + } + +#endif +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MultiheadAttentionVarlen::plan_dispatcher().registerDevice(Device::Type::kHygon, &plan); + MultiheadAttentionVarlen::run_dispatcher().registerDevice(Device::Type::kHygon, &run); + MultiheadAttentionVarlen::cleanup_dispatcher().registerDevice(Device::Type::kHygon, &cleanup); + return true; +}(); + +} // namespace infinicore::op::mha_varlen_impl::flashattn +#endif // ENABLE_HYGON_API && ENABLE_FLASH_ATTN diff --git a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen.cc b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen.cc new file mode 100644 index 000000000..f6d4612bd --- /dev/null +++ b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen.cc @@ -0,0 +1,75 @@ +#include "infinicore/ops/mha_varlen.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MultiheadAttentionVarlen); + +MultiheadAttentionVarlen::MultiheadAttentionVarlen(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &cum_seqlens_q, + const Tensor &cum_seqlens_kv, + std::optional block_table, + int max_seqlen_q, + int max_seqlen_k, + std::optional alibi_slopes, + float scale) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k, v, cum_seqlens_q, cum_seqlens_kv); + if (block_table.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, block_table.value()); + } + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), + out, q, k, v, cum_seqlens_q, cum_seqlens_kv, block_table, max_seqlen_q, max_seqlen_k, alibi_slopes, scale); +} + +void MultiheadAttentionVarlen::execute(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &cum_seqlens_q, + const Tensor &cum_seqlens_kv, + std::optional block_table, + int max_seqlen_q, + int max_seqlen_k, + std::optional alibi_slopes, + float scale) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MultiheadAttentionVarlen, + out, q, k, v, cum_seqlens_q, cum_seqlens_kv, block_table, max_seqlen_q, max_seqlen_k, alibi_slopes, scale); +} + +Tensor mha_varlen( + const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &cum_seqlens_q, + const Tensor &cum_seqlens_kv, + std::optional block_table, + int max_seqlen_q, + int max_seqlen_k, + std::optional alibi_slopes, + float scale) { + auto out_shape = q->shape(); + out_shape.back() = v->shape().back(); + auto out = Tensor::empty(out_shape, q->dtype(), q->device()); + mha_varlen_(out, q, k, v, cum_seqlens_q, cum_seqlens_kv, block_table, max_seqlen_q, max_seqlen_k, alibi_slopes, scale); + return out; +} + +void mha_varlen_(Tensor out, + const Tensor &q, + const Tensor &k, + const Tensor &v, + const Tensor &cum_seqlens_q, + const Tensor &cum_seqlens_kv, + std::optional block_table, + int max_seqlen_q, + int max_seqlen_k, + std::optional alibi_slopes, + float scale) { + MultiheadAttentionVarlen::execute(out, q, k, v, cum_seqlens_q, cum_seqlens_kv, block_table, max_seqlen_q, max_seqlen_k, alibi_slopes, scale); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc new file mode 100644 index 000000000..4a233bf5b --- /dev/null +++ b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc @@ -0,0 +1,191 @@ +#include "infinicore/ops/mha_varlen.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" +#include +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) +#include +#endif +#endif + +#ifdef ENABLE_FLASH_ATTN +#include "infinicore/adaptor/flash_attention_adaptor.hpp" +#endif + +#include + +namespace infinicore::op::mha_varlen_impl::flashattn { + +struct PlannedMeta { + graph::GraphTensor out, q, k, v, cum_seqlens_q, cum_seqlens_k; + std::optional block_table; + int max_seqlen_q, max_seqlen_k; + std::optional alibi_slopes; + float scale; +}; + +void *plan(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) { + + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k), + graph::GraphTensor(v), + graph::GraphTensor(cum_seqlens_q), + graph::GraphTensor(cum_seqlens_k), + block_table ? std::optional(graph::GraphTensor(*block_table)) : std::nullopt, + max_seqlen_q, + max_seqlen_k, + alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, + scale}; +} + +namespace { + +#ifdef ENABLE_FLASH_ATTN +// MetaX/hpcc pip `flash_attn_2_cuda` exports `mha_varlen_fwd` at global scope (no namespace), +// while NVIDIA `flash-attn-nvidia.so` uses `flash::mha_varlen_fwd`. +#if defined(ENABLE_METAX_API) +#define INFINICORE_FLASH_OP(name) ::name +#else +#define INFINICORE_FLASH_OP(name) flash::name +#endif + +#endif // ENABLE_FLASH_ATTN +} // namespace + +void run(void *planned_meta) { +#ifndef ENABLE_ATEN + (void)planned_meta; + throw std::runtime_error("ATen is not enabled in this build"); +#else +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) + c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); +#endif + auto *p = reinterpret_cast(planned_meta); + + auto q = infinicore::adaptor::to_aten_tensor(p->q); + auto k = infinicore::adaptor::to_aten_tensor(p->k); + auto v = infinicore::adaptor::to_aten_tensor(p->v); + + const bool out_need_copy_back = !p->out->is_contiguous(); + Tensor out_work_ic = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); + auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic); + + auto cu_seqlens_q = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_q); + auto cu_seqlens_kv = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_k); + + const bool dense_sdpa = !p->block_table.has_value() + && !p->alibi_slopes.has_value() + && q.dim() == 3 && k.dim() == 3 && v.dim() == 3 + && p->max_seqlen_q > 0 && p->max_seqlen_k > 0 + && p->max_seqlen_q == p->max_seqlen_k + && cu_seqlens_q.dim() == 1 + && cu_seqlens_q.size(0) == cu_seqlens_kv.size(0) + && q.size(0) == (cu_seqlens_q.size(0) - 1) * p->max_seqlen_q + && k.size(0) == (cu_seqlens_kv.size(0) - 1) * p->max_seqlen_k + && ((q.size(2) > 256) || (v.size(2) != q.size(2))); + if (dense_sdpa) { + const int64_t batch_size = cu_seqlens_q.size(0) - 1; + const int64_t seqlen = p->max_seqlen_q; + const int64_t num_heads = q.size(1); + const int64_t num_kv_heads = k.size(1); + const int64_t head_dim = q.size(2); + const int64_t value_dim = v.size(2); + auto q_4d = q.reshape({batch_size, seqlen, num_heads, head_dim}).permute({0, 2, 1, 3}); + auto k_4d = k.reshape({batch_size, seqlen, num_kv_heads, head_dim}).permute({0, 2, 1, 3}); + auto v_4d = v.reshape({batch_size, seqlen, num_kv_heads, value_dim}).permute({0, 2, 1, 3}); + if (num_heads != num_kv_heads) { + if (num_heads % num_kv_heads != 0) { + throw std::runtime_error("mha_varlen dense SDPA fallback requires num_heads to be divisible by num_kv_heads"); + } + const int64_t groups = num_heads / num_kv_heads; + k_4d = k_4d.unsqueeze(2).expand({batch_size, num_kv_heads, groups, seqlen, head_dim}).reshape({batch_size, num_heads, seqlen, head_dim}); + v_4d = v_4d.unsqueeze(2).expand({batch_size, num_kv_heads, groups, seqlen, value_dim}).reshape({batch_size, num_heads, seqlen, value_dim}); + } + auto result = at::scaled_dot_product_attention( + q_4d, + k_4d, + v_4d, + std::nullopt, + 0.0, + true, + std::optional(static_cast(p->scale))); + out_work.copy_(result.permute({0, 2, 1, 3}).reshape({q.size(0), num_heads, value_dim})); + if (out_need_copy_back) { + p->out->copy_from(out_work_ic); + } + return; + } + +#ifdef ENABLE_FLASH_ATTN + auto out = std::optional(out_work); + std::optional seqused_k = std::nullopt; + std::optional leftpad_k = std::nullopt; + auto block_table = p->block_table ? std::optional(infinicore::adaptor::to_aten_tensor(*p->block_table)) : std::nullopt; + auto max_seqlen_q = p->max_seqlen_q; + auto max_seqlen_k = p->max_seqlen_k; + auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; + auto scale = p->scale; + +#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) + std::optional flash_attn_mars_ext = std::nullopt; +#endif + + INFINICORE_FLASH_OP(mha_varlen_fwd) + ( + q, + k, + v, + out, + cu_seqlens_q, + cu_seqlens_kv, + seqused_k, + leftpad_k, + block_table, + alibi_slopes, + max_seqlen_q, + max_seqlen_k, + 0.0, + scale, + false, + true, + -1, + -1, + 0.0, + false, + std::nullopt +#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) + , + flash_attn_mars_ext +#endif + ); + + if (out_need_copy_back) { + p->out->copy_from(out_work_ic); + } + +#else + throw std::runtime_error("FlashAttention is not enabled in this build and dense SDPA fallback is not applicable"); +#endif +#endif +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MultiheadAttentionVarlen, &plan, &run, &cleanup); + +} // namespace infinicore::op::mha_varlen_impl::flashattn diff --git a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_flashattn_moore.cc b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_flashattn_moore.cc new file mode 100644 index 000000000..e33459f10 --- /dev/null +++ b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_flashattn_moore.cc @@ -0,0 +1,139 @@ +#if defined(ENABLE_MOORE_MATE_FLASH_ATTN) + +#include "infinicore/ops/mha_varlen.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include +#include +#include +#include + +namespace infinicore::op::mha_varlen_impl::flashattn_moore { + +namespace py = pybind11; + +namespace { +class LocalMUSAStreamGuard { +public: + explicit LocalMUSAStreamGuard(const c10::musa::MUSAStream &s) + : prev_(c10::musa::getCurrentMUSAStream(s.device_index())) { + c10::musa::setCurrentMUSAStream(s); + } + ~LocalMUSAStreamGuard() { + c10::musa::setCurrentMUSAStream(prev_); + } + LocalMUSAStreamGuard(const LocalMUSAStreamGuard &) = delete; + LocalMUSAStreamGuard &operator=(const LocalMUSAStreamGuard &) = delete; + +private: + c10::musa::MUSAStream prev_; +}; +} // namespace + +struct PlannedMeta { + graph::GraphTensor out, q, k, v, cum_seqlens_q, cum_seqlens_k; + std::optional block_table; + int max_seqlen_q, max_seqlen_k; + std::optional alibi_slopes; + float scale; +}; + +void *plan(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) { + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(q), + graph::GraphTensor(k), + graph::GraphTensor(v), + graph::GraphTensor(cum_seqlens_q), + graph::GraphTensor(cum_seqlens_k), + block_table ? std::optional(graph::GraphTensor(*block_table)) : std::nullopt, + max_seqlen_q, + max_seqlen_k, + alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, + scale}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + if (p->alibi_slopes.has_value()) { + throw std::runtime_error( + "[mha_varlen/moore] ALiBi not supported by mate flash_attn_varlen"); + } + if (!p->block_table.has_value()) { + throw std::runtime_error( + "[mha_varlen/moore] dense KV is not supported by mate flash_attn_varlen"); + } + + LocalMUSAStreamGuard guard(infinicore::adaptor::get_musa_stream()); + + auto out_tensor = infinicore::adaptor::to_aten_tensor(p->out); + auto q_tensor = infinicore::adaptor::to_aten_tensor(p->q); + auto k_cache = infinicore::adaptor::to_aten_tensor(p->k); + auto v_cache = infinicore::adaptor::to_aten_tensor(p->v); + auto cu_seqlens_q = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_q); + auto cu_seqlens_k = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_k); + auto block_table = infinicore::adaptor::to_aten_tensor(*p->block_table); + + const int64_t block_size = k_cache.size(1); + + int max_seqlen_q_bound = static_cast(q_tensor.size(0)); + int max_seqlen_k_bound = static_cast(q_tensor.size(0)); + + try { + py::gil_scoped_acquire gil; + py::module_ wrapper = py::module_::import("infinicore.ops.moore_mate_flash_attn"); + + py::object py_q = py::cast(q_tensor); + py::object py_k = py::cast(k_cache); + py::object py_v = py::cast(v_cache); + py::object py_cuq = py::cast(cu_seqlens_q); + py::object py_cuk = py::cast(cu_seqlens_k); + py::object py_blk = py::cast(block_table); + + py::object result = wrapper.attr("moore_mate_flash_attn_prefill")( + py_q, + py_k, + py_v, + py_cuq, + py_cuk, + py_blk, + p->scale, + max_seqlen_q_bound, + max_seqlen_k_bound, + block_size, + true); + + at::Tensor result_t = result.cast(); + out_tensor.copy_(result_t); + } catch (const py::error_already_set &e) { + throw std::runtime_error( + std::string("[mha_varlen/moore] Python error: ") + e.what()); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MultiheadAttentionVarlen::plan_dispatcher().registerDevice(Device::Type::kMoore, &plan); + MultiheadAttentionVarlen::run_dispatcher().registerDevice(Device::Type::kMoore, &run); + MultiheadAttentionVarlen::cleanup_dispatcher().registerDevice(Device::Type::kMoore, &cleanup); + return true; +}(); + +} // namespace infinicore::op::mha_varlen_impl::flashattn_moore + +#endif // ENABLE_MOORE_MATE_FLASH_ATTN diff --git a/csrc/infinicore/src/ops/multi_margin_loss/multi_margin_loss.cc b/csrc/infinicore/src/ops/multi_margin_loss/multi_margin_loss.cc new file mode 100644 index 000000000..5a71ca282 --- /dev/null +++ b/csrc/infinicore/src/ops/multi_margin_loss/multi_margin_loss.cc @@ -0,0 +1,34 @@ +#include "infinicore/ops/multi_margin_loss.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &MultiMarginLoss::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void MultiMarginLoss::execute(Tensor output, Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction) { + dispatcher().lookup(context::getDevice().type())(output, input, target, weight, p, margin, reduction); +} + +// 3. 函数式接口 +Tensor multi_margin_loss(Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction) { + Shape output_shape; + if (reduction == 0) { // None + output_shape = {input->shape()[0]}; + } else { + output_shape = {}; // Scalar + } + + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + + multi_margin_loss_(output, input, target, weight, p, margin, reduction); + return output; +} + +void multi_margin_loss_(Tensor output, Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction) { + MultiMarginLoss::execute(output, input, target, weight, p, margin, reduction); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/nrm2/nrm2.cc b/csrc/infinicore/src/ops/nrm2/nrm2.cc new file mode 100644 index 000000000..6bb20930a --- /dev/null +++ b/csrc/infinicore/src/ops/nrm2/nrm2.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/nrm2.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Nrm2); + +Nrm2::Nrm2(const Tensor &x, Tensor result) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, result); + INFINICORE_GRAPH_OP_DISPATCH(result->device().type(), x, result); +} + +void Nrm2::execute(const Tensor &x, Tensor result) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Nrm2, x, result); +} + +Tensor nrm2(const Tensor &x) { + auto result = Tensor::empty({}, x->dtype(), x->device()); + nrm2_(x, result); + return result; +} + +void nrm2_(const Tensor &x, Tensor result) { + Nrm2::execute(x, result); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/nsa_compress_paged_cache/nsa_compress_paged_cache.cc b/csrc/infinicore/src/ops/nsa_compress_paged_cache/nsa_compress_paged_cache.cc new file mode 100644 index 000000000..08bee736b --- /dev/null +++ b/csrc/infinicore/src/ops/nsa_compress_paged_cache/nsa_compress_paged_cache.cc @@ -0,0 +1,29 @@ +#include "infinicore/ops/nsa_compress_paged_cache.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(NsaCompressPagedCache); + +NsaCompressPagedCache::NsaCompressPagedCache(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) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens); + INFINICORE_GRAPH_OP_DISPATCH(k_cmp->device().type(), k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, nsa_block_size, update_last_only); +} + +void NsaCompressPagedCache::execute(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) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + NsaCompressPagedCache, + k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, nsa_block_size, update_last_only); +} + +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) { + NsaCompressPagedCache::execute(k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, nsa_block_size, update_last_only); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/nsa_paged_attention/nsa_paged_attention.cc b/csrc/infinicore/src/ops/nsa_paged_attention/nsa_paged_attention.cc new file mode 100644 index 000000000..c3fc868d8 --- /dev/null +++ b/csrc/infinicore/src/ops/nsa_paged_attention/nsa_paged_attention.cc @@ -0,0 +1,40 @@ +#include "infinicore/ops/nsa_paged_attention.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(NsaPagedAttention); + +NsaPagedAttention::NsaPagedAttention(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) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, gates); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), + out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, gates, scale, nsa_block_size, window_size, select_blocks); +} + +void NsaPagedAttention::execute(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) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + NsaPagedAttention, + out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, gates, scale, nsa_block_size, window_size, select_blocks); +} + +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) { + auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); + nsa_paged_attention_(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, gates, scale, nsa_block_size, window_size, select_blocks); + return out; +} + +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) { + NsaPagedAttention::execute(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, gates, scale, nsa_block_size, window_size, select_blocks); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/ones/ones.cc b/csrc/infinicore/src/ops/ones/ones.cc new file mode 100644 index 000000000..c28403eaf --- /dev/null +++ b/csrc/infinicore/src/ops/ones/ones.cc @@ -0,0 +1,13 @@ +#include "infinicore/ops/ones.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Ones::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Ones::execute(Tensor output) { +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/pad/pad.cc b/csrc/infinicore/src/ops/pad/pad.cc new file mode 100644 index 000000000..59ba5c6a7 --- /dev/null +++ b/csrc/infinicore/src/ops/pad/pad.cc @@ -0,0 +1,56 @@ +#include "infinicore/ops/pad.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Pad); + +Pad::Pad(Tensor y, const Tensor &x, const std::vector &pad, const std::string &mode, double value) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), y, x, pad, mode, value); +} + +void Pad::execute(Tensor y, const Tensor &x, const std::vector &pad, const std::string &mode, double value) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Pad, y, x, pad, mode, value); +} + +static Shape infer_padded_shape(const Shape &in_shape, const std::vector &pad) { + if (pad.empty() || (pad.size() % 2) != 0) { + throw std::runtime_error("pad: pad must have even length."); + } + + Shape out_shape = in_shape; + const size_t ndim = out_shape.size(); + const size_t dims_padded = pad.size() / 2; + if (dims_padded > ndim) { + throw std::runtime_error("pad: pad has more dimensions than input."); + } + + for (size_t j = 0; j < dims_padded; ++j) { + const int left = pad[2 * j]; + const int right = pad[2 * j + 1]; + if (left < 0 || right < 0) { + throw std::runtime_error("pad: negative pad is not supported."); + } + const size_t dim = ndim - 1 - j; + out_shape[dim] += static_cast(left + right); + } + + return out_shape; +} + +Tensor pad(const Tensor &x, const std::vector &pad, const std::string &mode, double value) { + auto y_shape = infer_padded_shape(x->shape(), pad); + auto y = Tensor::empty(y_shape, x->dtype(), x->device()); + pad_(y, x, pad, mode, value); + return y; +} + +void pad_(Tensor y, const Tensor &x, const std::vector &pad, const std::string &mode, double value) { + Pad::execute(y, x, pad, mode, value); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/paged_attention/paged_attention.cc b/csrc/infinicore/src/ops/paged_attention/paged_attention.cc new file mode 100644 index 000000000..9dbffaaa8 --- /dev/null +++ b/csrc/infinicore/src/ops/paged_attention/paged_attention.cc @@ -0,0 +1,38 @@ +#include "infinicore/ops/paged_attention.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PagedAttention); + +PagedAttention::PagedAttention(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, + const Tensor &block_tables, const Tensor &kv_lens, + std::optional alibi_slopes, float scale) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, kv_lens); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), + out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); +} + +void PagedAttention::execute(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, + const Tensor &block_tables, const Tensor &kv_lens, + std::optional alibi_slopes, float scale) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + PagedAttention, + out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); +} + +Tensor paged_attention(const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, + const Tensor &block_tables, const Tensor &kv_lens, + std::optional alibi_slopes, float scale) { + auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); + paged_attention_(out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); + return out; +} + +void paged_attention_(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, + const Tensor &block_tables, const Tensor &kv_lens, + std::optional alibi_slopes, float scale) { + PagedAttention::execute(out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/paged_attention/paged_attention_infiniops.cc b/csrc/infinicore/src/ops/paged_attention/paged_attention_infiniops.cc new file mode 100644 index 000000000..76beca8fc --- /dev/null +++ b/csrc/infinicore/src/ops/paged_attention/paged_attention_infiniops.cc @@ -0,0 +1,88 @@ +#include "infinicore/ops/paged_attention.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/paged_attention_infinilm.h" + +#include +#include + +namespace infinicore::op::paged_attention_impl::infiniops { +namespace { +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +constexpr std::size_t kMaxPagedAttentionSplits = 8; + +std::size_t WorkspaceSizeInBytes(const Tensor &q) { + return kMaxPagedAttentionSplits + * static_cast(q->size(0)) + * static_cast(q->size(1)) + * static_cast(q->size(2) + 2) + * sizeof(float); +} + +struct PlannedMeta { + TensorMeta out, q, k_cache, v_cache, block_tables, cache_lens; + std::optional alibi_slopes; + graph::GraphTensor workspace, out_tensor, q_tensor, k_cache_tensor, v_cache_tensor, block_tables_tensor, cache_lens_tensor; + std::optional alibi_slopes_tensor; + float scale; +}; +} // namespace + +void *plan(Tensor out, + const Tensor &q, + const Tensor &k_cache, + const Tensor &v_cache, + const Tensor &block_tables, + const Tensor &cache_lens, + std::optional alibi_slopes, + float scale) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, cache_lens); + if (alibi_slopes) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, *alibi_slopes); + } + return new PlannedMeta{ + TensorMeta(out), TensorMeta(q), TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(block_tables), TensorMeta(cache_lens), + alibi_slopes ? std::optional{TensorMeta(*alibi_slopes)} : std::nullopt, + graph::GraphTensor(Tensor::empty({WorkspaceSizeInBytes(q)}, DataType::kUInt8, out->device())), + graph::GraphTensor(out), graph::GraphTensor(q), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(block_tables), graph::GraphTensor(cache_lens), + alibi_slopes ? std::optional{graph::GraphTensor(*alibi_slopes)} : std::nullopt, + scale}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + handle.set_workspace(planned->workspace->data()); + handle.set_workspace_size_in_bytes(planned->workspace->numel()); + infini::ops::Config config; + infini::ops::PagedAttentionInfinilm::Call( + handle, + config, + planned->q.tensor(planned->q_tensor), + planned->k_cache.tensor(planned->k_cache_tensor), + planned->v_cache.tensor(planned->v_cache_tensor), + planned->block_tables.tensor(planned->block_tables_tensor), + planned->cache_lens.tensor(planned->cache_lens_tensor), + planned->alibi_slopes ? std::optional{planned->alibi_slopes->tensor(planned->alibi_slopes_tensor.value()->data())} : std::nullopt, + planned->scale, + planned->out.tensor(planned->out_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(PagedAttention::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(PagedAttention::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(PagedAttention::cleanup_dispatcher(), &cleanup); + return true; +}(); +} // namespace infinicore::op::paged_attention_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/paged_attention_prefill/paged_attention_prefill.cc b/csrc/infinicore/src/ops/paged_attention_prefill/paged_attention_prefill.cc new file mode 100644 index 000000000..bf276c715 --- /dev/null +++ b/csrc/infinicore/src/ops/paged_attention_prefill/paged_attention_prefill.cc @@ -0,0 +1,39 @@ +#include "infinicore/ops/paged_attention_prefill.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &PagedAttentionPrefill::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void PagedAttentionPrefill::execute(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, + Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, + std::optional alibi_slopes, float scale) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q); + + infinicore::context::setDevice(out->device()); + + dispatcher().lookup(out->device().type())(out, q, k_cache, v_cache, block_tables, + kv_lens, cum_seqlens_q, alibi_slopes, scale); +} + +Tensor paged_attention_prefill(Tensor q, Tensor k_cache, Tensor v_cache, + Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, + std::optional alibi_slopes, float scale) { + + auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); + paged_attention_prefill_(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale); + return out; +} + +void paged_attention_prefill_(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, + Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, + std::optional alibi_slopes, float scale) { + + PagedAttentionPrefill::execute(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc b/csrc/infinicore/src/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc new file mode 100644 index 000000000..a7235887d --- /dev/null +++ b/csrc/infinicore/src/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc @@ -0,0 +1,68 @@ +#include "infinicore/ops/paged_attention_prefill.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/paged_attention_prefill_infinilm.h" + +#include + +namespace infinicore::op::paged_attention_prefill_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +void calculate(Tensor out, + Tensor q, + Tensor k_cache, + Tensor v_cache, + Tensor block_tables, + Tensor kv_lens, + Tensor cum_seqlens_q, + std::optional alibi_slopes, + float scale) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q); + if (alibi_slopes) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, *alibi_slopes); + } + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + TensorMeta out_meta(out); + TensorMeta q_meta(q); + TensorMeta k_cache_meta(k_cache); + TensorMeta v_cache_meta(v_cache); + TensorMeta block_tables_meta(block_tables); + TensorMeta kv_lens_meta(kv_lens); + TensorMeta cum_seqlens_q_meta(cum_seqlens_q); + std::optional alibi_slopes_meta; + if (alibi_slopes) { + alibi_slopes_meta.emplace(*alibi_slopes); + } + + infini::ops::PagedAttentionPrefillInfinilm::Call( + handle, + config, + q_meta.tensor(q), + k_cache_meta.tensor(k_cache), + v_cache_meta.tensor(v_cache), + block_tables_meta.tensor(block_tables), + kv_lens_meta.tensor(kv_lens), + cum_seqlens_q_meta.tensor(cum_seqlens_q), + alibi_slopes_meta ? std::optional{alibi_slopes_meta->tensor(*alibi_slopes)} : std::nullopt, + scale, + out_meta.tensor(out)); +} + +} // namespace + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(PagedAttentionPrefill::dispatcher(), &calculate); + return true; +}(); + +} // namespace infinicore::op::paged_attention_prefill_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/paged_caching/paged_caching.cc b/csrc/infinicore/src/ops/paged_caching/paged_caching.cc new file mode 100644 index 000000000..48d36d83d --- /dev/null +++ b/csrc/infinicore/src/ops/paged_caching/paged_caching.cc @@ -0,0 +1,21 @@ +#include "infinicore/ops/paged_caching.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PagedCaching); + +PagedCaching::PagedCaching(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, slot_mapping); + INFINICORE_GRAPH_OP_DISPATCH(k->device().type(), k_cache, v_cache, k, v, slot_mapping); +} + +void PagedCaching::execute(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(PagedCaching, k_cache, v_cache, k, v, slot_mapping); +} + +void paged_caching_(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { + PagedCaching::execute(k_cache, v_cache, k, v, slot_mapping); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/paged_caching/paged_caching_infiniops.cc b/csrc/infinicore/src/ops/paged_caching/paged_caching_infiniops.cc new file mode 100644 index 000000000..eea166908 --- /dev/null +++ b/csrc/infinicore/src/ops/paged_caching/paged_caching_infiniops.cc @@ -0,0 +1,50 @@ +#include "infinicore/ops/paged_caching.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/paged_caching_infinilm.h" + +namespace infinicore::op::paged_caching_impl::infiniops { +namespace { +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; +struct PlannedMeta { + TensorMeta k_cache, v_cache, k, v, slot_mapping; + graph::GraphTensor k_cache_tensor, v_cache_tensor, k_tensor, v_tensor, slot_mapping_tensor; +}; +} // namespace + +void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(k_cache->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, slot_mapping); + return new PlannedMeta{TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(k), TensorMeta(v), TensorMeta(slot_mapping), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(slot_mapping)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + infini::ops::PagedCachingInfinilm::Call( + handle, + config, + planned->k.tensor(planned->k_tensor), + planned->v.tensor(planned->v_tensor), + planned->slot_mapping.tensor(planned->slot_mapping_tensor), + planned->k_cache.tensor(planned->k_cache_tensor), + planned->v_cache.tensor(planned->v_cache_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(PagedCaching::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(PagedCaching::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(PagedCaching::cleanup_dispatcher(), &cleanup); + return true; +}(); +} // namespace infinicore::op::paged_caching_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/per_channel_quant_i8/per_channel_quant_i8.cc b/csrc/infinicore/src/ops/per_channel_quant_i8/per_channel_quant_i8.cc new file mode 100644 index 000000000..741212473 --- /dev/null +++ b/csrc/infinicore/src/ops/per_channel_quant_i8/per_channel_quant_i8.cc @@ -0,0 +1,20 @@ +#include "infinicore/ops/per_channel_quant_i8.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PerChannelQuantI8); + +PerChannelQuantI8::PerChannelQuantI8(const Tensor &x, Tensor x_packed, Tensor x_scale) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, x_packed, x_scale); + INFINICORE_GRAPH_OP_DISPATCH(x->device().type(), x, x_packed, x_scale); +} + +void PerChannelQuantI8::execute(const Tensor &x, Tensor x_packed, Tensor x_scale) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(PerChannelQuantI8, x, x_packed, x_scale); +} + +void per_channel_quant_i8_(const Tensor &x, Tensor x_packed, Tensor x_scale) { + PerChannelQuantI8::execute(x, x_packed, x_scale); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/prelu/prelu.cc b/csrc/infinicore/src/ops/prelu/prelu.cc new file mode 100644 index 000000000..c6f41cb95 --- /dev/null +++ b/csrc/infinicore/src/ops/prelu/prelu.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/prelu.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Prelu); + +Prelu::Prelu(Tensor out, const Tensor &input, const Tensor &weight) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), out, input, weight); +} + +void Prelu::execute(Tensor out, const Tensor &input, const Tensor &weight) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Prelu, out, input, weight); +} + +Tensor prelu(const Tensor &input, const Tensor &weight) { + auto out = Tensor::empty(input->shape(), input->dtype(), input->device()); + prelu_(out, input, weight); + return out; +} + +void prelu_(Tensor out, const Tensor &input, const Tensor &weight) { + Prelu::execute(out, input, weight); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/prepare_moe_input/prepare_moe_input.cc b/csrc/infinicore/src/ops/prepare_moe_input/prepare_moe_input.cc new file mode 100644 index 000000000..7d7047ade --- /dev/null +++ b/csrc/infinicore/src/ops/prepare_moe_input/prepare_moe_input.cc @@ -0,0 +1,124 @@ +#include "infinicore/ops/prepare_moe_input.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PrepareMoeInput); + +PrepareMoeInput::PrepareMoeInput(Tensor expert_offsets, + Tensor blockscale_offsets, + Tensor problem_sizes1, + Tensor problem_sizes2, + Tensor input_permutation, + Tensor output_permutation, + const Tensor &topk_ids, + const size_t num_experts, + const size_t n, + const size_t k) { + if (blockscale_offsets) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + expert_offsets, blockscale_offsets, problem_sizes1, problem_sizes2, input_permutation, output_permutation, topk_ids); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + expert_offsets, problem_sizes1, problem_sizes2, input_permutation, output_permutation, topk_ids); + } + INFINICORE_GRAPH_OP_DISPATCH( + expert_offsets->device().type(), + expert_offsets, + blockscale_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + topk_ids, + num_experts, + n, + k); +} + +void PrepareMoeInput::execute(Tensor expert_offsets, + Tensor blockscale_offsets, + Tensor problem_sizes1, + Tensor problem_sizes2, + Tensor input_permutation, + Tensor output_permutation, + const Tensor &topk_ids, + const size_t num_experts, + const size_t n, + const size_t k) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + PrepareMoeInput, + expert_offsets, + blockscale_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + topk_ids, + num_experts, + n, + k); +} + +PrepareMoeInputOutput prepare_moe_input( + const Tensor &topk_ids, + size_t num_experts, + size_t n, + size_t k) { + auto shape = topk_ids->shape(); + INFINICORE_ASSERT(shape.size() == 2); + const size_t topk_length = shape[0] * shape[1]; + + auto expert_offsets = Tensor::empty({num_experts + 1}, DataType::kInt32, topk_ids->device()); + auto problem_sizes1 = Tensor::empty({num_experts, 3}, DataType::kInt32, topk_ids->device()); + auto problem_sizes2 = Tensor::empty({num_experts, 3}, DataType::kInt32, topk_ids->device()); + auto input_permutation = Tensor::empty({topk_length}, DataType::kInt32, topk_ids->device()); + auto output_permutation = Tensor::empty({topk_length}, DataType::kInt32, topk_ids->device()); + + prepare_moe_input_( + expert_offsets, + Tensor(), + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + topk_ids, + num_experts, + n, + k); + + return PrepareMoeInputOutput{ + expert_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + }; +} + +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) { + PrepareMoeInput::execute( + expert_offsets, + blockscale_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + topk_ids, + num_experts, + n, + k); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/quant/per_tensor_quant_int8/per_tensor_quant_int8.cc b/csrc/infinicore/src/ops/quant/per_tensor_quant_int8/per_tensor_quant_int8.cc new file mode 100644 index 000000000..df4c2a6af --- /dev/null +++ b/csrc/infinicore/src/ops/quant/per_tensor_quant_int8/per_tensor_quant_int8.cc @@ -0,0 +1,26 @@ +#include "../../../utils.hpp" +#include "infinicore/ops/per_tensor_quant_i8.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PerTensorQuantI8); + +PerTensorQuantI8::PerTensorQuantI8(const Tensor &x, Tensor x_packed, Tensor x_scale, Tensor x_zero, bool is_static) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, x_packed, x_scale, x_zero); + INFINICORE_GRAPH_OP_DISPATCH(x->device().type(), x, x_packed, x_scale, x_zero, is_static); +} + +void PerTensorQuantI8::execute(const Tensor &x, Tensor x_packed, Tensor x_scale, Tensor x_zero, bool is_static) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(PerTensorQuantI8, x, x_packed, x_scale, x_zero, is_static); +} + +void per_tensor_quant_i8_(const Tensor &x, Tensor x_packed, Tensor x_scale, Tensor x_zero, bool is_static) { + PerTensorQuantI8::execute(x, x_packed, x_scale, x_zero, is_static); +} + +Tensor per_tensor_quant_i8(const Tensor &x, Tensor x_scale, Tensor x_zero, bool is_static) { + auto x_packed = Tensor::strided_empty(x->shape(), x->strides(), infinicore::DataType::kInt8, x->device()); + PerTensorQuantI8::execute(x, x_packed, x_scale, x_zero, is_static); + return x_packed; +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/quickgelu/quickgelu.cc b/csrc/infinicore/src/ops/quickgelu/quickgelu.cc new file mode 100644 index 000000000..edf6f958c --- /dev/null +++ b/csrc/infinicore/src/ops/quickgelu/quickgelu.cc @@ -0,0 +1,37 @@ +#include "infinicore/ops/quickgelu.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &QuickGelu::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void QuickGelu::execute(Tensor output, Tensor input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No QuickGelu implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor quick_gelu(Tensor input) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + quick_gelu_(output, input); + return output; +} + +void quick_gelu_(Tensor output, Tensor input) { + QuickGelu::execute(output, input); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/random_sample/random_sample.cc b/csrc/infinicore/src/ops/random_sample/random_sample.cc new file mode 100644 index 000000000..175971293 --- /dev/null +++ b/csrc/infinicore/src/ops/random_sample/random_sample.cc @@ -0,0 +1,42 @@ +#include "infinicore/ops/random_sample.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &RandomSample::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void RandomSample::execute( + Tensor indices, Tensor logits, + float random_val, float topp, int topk, float temperature) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(indices, logits); + infinicore::context::setDevice(logits->device()); + dispatcher().lookup(logits->device().type())( + indices, logits, random_val, topp, topk, temperature); +} + +Tensor random_sample( + Tensor logits, + float random_val, + float topp, + int topk, + float temperature) { + auto indices = Tensor::empty({}, DataType::kInt32, logits->device()); + random_sample_(indices, logits, random_val, topp, topk, temperature); + return indices; +} + +void random_sample_( + Tensor indices, + Tensor logits, + float random_val, + float topp, + int topk, + float temperature) { + RandomSample::execute(indices, logits, random_val, topp, topk, temperature); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/random_sample/random_sample_infiniops.cc b/csrc/infinicore/src/ops/random_sample/random_sample_infiniops.cc new file mode 100644 index 000000000..d7d045097 --- /dev/null +++ b/csrc/infinicore/src/ops/random_sample/random_sample_infiniops.cc @@ -0,0 +1,42 @@ +#include "infinicore/ops/random_sample.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/random_sample_infinilm.h" + +namespace infinicore::op::random_sample_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +void calculate(Tensor indices, Tensor logits, float random_val, float topp, int topk, float temperature) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(indices->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(indices, logits); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + TensorMeta indices_meta(indices); + TensorMeta logits_meta(logits); + infini::ops::RandomSampleInfinilm::Call( + handle, + config, + logits_meta.tensor(logits), + random_val, + topp, + static_cast(topk), + temperature, + indices_meta.tensor(indices)); +} + +} // namespace + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(RandomSample::dispatcher(), &calculate); + return true; +}(); + +} // namespace infinicore::op::random_sample_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/rearrange/rearrange.cc b/csrc/infinicore/src/ops/rearrange/rearrange.cc new file mode 100644 index 000000000..11beb559e --- /dev/null +++ b/csrc/infinicore/src/ops/rearrange/rearrange.cc @@ -0,0 +1,33 @@ +#include "infinicore/ops/rearrange.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Rearrange); + +Rearrange::Rearrange(Tensor y, const Tensor &x) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), y, x); +} + +void Rearrange::execute(Tensor y, const Tensor &x) { + const bool recording = context::isGraphRecording(); + auto op = std::make_shared(y, x); + if (recording) { + context::addGraphOperator(op); + } else { + op->run(); + } +} + +Tensor rearrange(const Tensor &x) { + auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); + rearrange_(y, x); + return y; +} + +void rearrange_(Tensor y, const Tensor &x) { + Rearrange::execute(y, x); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/rearrange/rearrange_infiniops.cc b/csrc/infinicore/src/ops/rearrange/rearrange_infiniops.cc new file mode 100644 index 000000000..b1b7e9213 --- /dev/null +++ b/csrc/infinicore/src/ops/rearrange/rearrange_infiniops.cc @@ -0,0 +1,45 @@ +#include "infinicore/ops/rearrange.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/rearrange_infinilm.h" + +namespace infinicore::op::rearrange_impl::infiniops { +namespace { +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; +struct PlannedMeta { + TensorMeta y, x; + graph::GraphTensor y_tensor, x_tensor; +}; +} // namespace + +void *plan(Tensor y, const Tensor &x) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(y->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + return new PlannedMeta{TensorMeta(y), TensorMeta(x), graph::GraphTensor(y), graph::GraphTensor(x)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + infini::ops::RearrangeInfinilm::Call( + handle, config, planned->x.tensor(planned->x_tensor), planned->y.tensor(planned->y_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Rearrange::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(Rearrange::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(Rearrange::cleanup_dispatcher(), &cleanup); + return true; +}(); + +} // namespace infinicore::op::rearrange_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/reciprocal/reciprocal.cc b/csrc/infinicore/src/ops/reciprocal/reciprocal.cc new file mode 100644 index 000000000..a646c375c --- /dev/null +++ b/csrc/infinicore/src/ops/reciprocal/reciprocal.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/reciprocal.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Reciprocal::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Reciprocal::execute(Tensor y, Tensor x) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + infinicore::context::setDevice(y->device()); + dispatcher().lookup(y->device().type())(y, x); +} + +Tensor reciprocal(Tensor x) { + auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); + reciprocal_(y, x); + return y; +} + +void reciprocal_(Tensor y, Tensor x) { + Reciprocal::execute(y, x); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule.cc b/csrc/infinicore/src/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule.cc new file mode 100644 index 000000000..da97e5a14 --- /dev/null +++ b/csrc/infinicore/src/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule.cc @@ -0,0 +1,166 @@ +#include "infinicore/ops/recurrent_gated_delta_rule.hpp" +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RecurrentGatedDeltaRule); + +RecurrentGatedDeltaRule::RecurrentGatedDeltaRule(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) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state, q, k, v, g, beta); + if (final_state.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state.value()); + } + if (initial_state_indices.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state_indices.value()); + } + if (final_state_indices.has_value()) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state_indices.value()); + } + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), + out, + initial_state, + final_state, + q, + k, + v, + g, + beta, + initial_state_indices, + final_state_indices, + use_qk_l2norm); +} + +void RecurrentGatedDeltaRule::execute(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) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(RecurrentGatedDeltaRule, + out, + initial_state, + final_state, + q, + k, + v, + g, + beta, + initial_state_indices, + final_state_indices, + use_qk_l2norm); +} + +static Tensor ensure_4d_sequence_tensor(const Tensor &x, const char *name) { + if (x->shape().size() == 4) { + return x; + } + if (x->shape().size() == 3) { + return x->unsqueeze(1); + } + throw std::runtime_error(std::string("recurrent_gated_delta_rule expects ") + name + " with shape [B, T, H, D] or [B, H, D]"); +} + +static Shape recurrent_output_shape(const Tensor &v) { + const auto &shape = v->shape(); + return {shape[0], shape[1], shape[2], shape[3]}; +} + +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) { + Tensor q4 = ensure_4d_sequence_tensor(q, "q"); + Tensor k4 = ensure_4d_sequence_tensor(k, "k"); + Tensor v4 = ensure_4d_sequence_tensor(v, "v"); + auto out = Tensor::empty(recurrent_output_shape(v4), v4->dtype(), v4->device()); + Shape final_state_shape = {v4->shape()[0], v4->shape()[2], v4->shape()[3], q4->shape()[3]}; + auto final_state = Tensor::empty(final_state_shape, initial_state->dtype(), initial_state->device()); + recurrent_gated_delta_rule_(out, + initial_state, + final_state, + q4, + k4, + v4, + g, + beta, + std::nullopt, + std::nullopt, + use_qk_l2norm); + return out; +} + +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) { + Tensor q4 = ensure_4d_sequence_tensor(q, "q"); + Tensor k4 = ensure_4d_sequence_tensor(k, "k"); + Tensor v4 = ensure_4d_sequence_tensor(v, "v"); + auto out = Tensor::empty(recurrent_output_shape(v4), v4->dtype(), v4->device()); + recurrent_gated_delta_rule_(out, + initial_state, + std::nullopt, + q4, + k4, + v4, + g, + beta, + initial_state_indices, + final_state_indices, + use_qk_l2norm); + return out; +} + +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) { + Tensor q4 = ensure_4d_sequence_tensor(q, "q"); + Tensor k4 = ensure_4d_sequence_tensor(k, "k"); + Tensor v4 = ensure_4d_sequence_tensor(v, "v"); + RecurrentGatedDeltaRule::execute(out, + initial_state, + final_state, + q4, + k4, + v4, + g, + beta, + initial_state_indices, + final_state_indices, + use_qk_l2norm); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/relu/relu.cc b/csrc/infinicore/src/ops/relu/relu.cc new file mode 100644 index 000000000..2084134a6 --- /dev/null +++ b/csrc/infinicore/src/ops/relu/relu.cc @@ -0,0 +1,37 @@ +#include "infinicore/ops/relu.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &Relu::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Relu::execute(Tensor output, Tensor input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Relu implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor relu(Tensor input) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + relu_(output, input); + return output; +} + +void relu_(Tensor output, Tensor input) { + Relu::execute(output, input); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/relu/relu_infiniops.cc b/csrc/infinicore/src/ops/relu/relu_infiniops.cc new file mode 100644 index 000000000..07f453cb1 --- /dev/null +++ b/csrc/infinicore/src/ops/relu/relu_infiniops.cc @@ -0,0 +1,38 @@ +#include "infinicore/ops/relu.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/relu_infinilm.h" + +namespace infinicore::op::relu_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +void calculate(Tensor output, Tensor input) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + TensorMeta output_meta(output); + TensorMeta input_meta(input); + infini::ops::ReluInfinilm::Call( + handle, + config, + input_meta.tensor(input), + output_meta.tensor(output)); +} + +} // namespace + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Relu::dispatcher(), &calculate); + return true; +}(); + +} // namespace infinicore::op::relu_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/relu6/relu6.cc b/csrc/infinicore/src/ops/relu6/relu6.cc new file mode 100644 index 000000000..0ae79f553 --- /dev/null +++ b/csrc/infinicore/src/ops/relu6/relu6.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/relu6.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Relu6); + +Relu6::Relu6(Tensor out, const Tensor &input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), out, input); +} + +void Relu6::execute(Tensor out, const Tensor &input) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Relu6, out, input); +} + +Tensor relu6(const Tensor &input) { + auto out = Tensor::empty(input->shape(), input->dtype(), input->device()); + relu6_(out, input); + return out; +} + +void relu6_(Tensor out, const Tensor &input) { + Relu6::execute(out, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/rms_norm/rms_norm.cc b/csrc/infinicore/src/ops/rms_norm/rms_norm.cc new file mode 100644 index 000000000..28aa23ad7 --- /dev/null +++ b/csrc/infinicore/src/ops/rms_norm/rms_norm.cc @@ -0,0 +1,26 @@ +#include "infinicore/ops/rms_norm.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RMSNorm); + +RMSNorm::RMSNorm(Tensor y, const Tensor &x, const Tensor &weight, float epsilon) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x, weight); + INFINICORE_GRAPH_OP_DISPATCH(y->device().type(), y, x, weight, epsilon); +} + +void RMSNorm::execute(Tensor y, const Tensor &x, const Tensor &weight, float epsilon) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(RMSNorm, y, x, weight, epsilon); +} + +Tensor rms_norm(const Tensor &x, const Tensor &weight, float epsilon) { + auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); + rms_norm_(y, x, weight, epsilon); + return y; +} + +void rms_norm_(Tensor y, const Tensor &x, const Tensor &weight, float epsilon) { + RMSNorm::execute(y, x, weight, epsilon); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/rms_norm/rms_norm_infiniops.cc b/csrc/infinicore/src/ops/rms_norm/rms_norm_infiniops.cc new file mode 100644 index 000000000..27f5f0b87 --- /dev/null +++ b/csrc/infinicore/src/ops/rms_norm/rms_norm_infiniops.cc @@ -0,0 +1,64 @@ +#include "infinicore/ops/rms_norm.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/rms_norm.h" + +namespace infinicore::op::rms_norm_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +struct PlannedMeta { + TensorMeta y, x, weight; + graph::GraphTensor y_tensor, x_tensor, weight_tensor; + float epsilon; +}; + +} // namespace + +void *plan(Tensor y, const Tensor &x, const Tensor &weight, float epsilon) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(y->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x, weight); + + return new PlannedMeta{ + TensorMeta(y), + TensorMeta(x), + TensorMeta(weight), + graph::GraphTensor(y), + graph::GraphTensor(x), + graph::GraphTensor(weight), + epsilon}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + infini::ops::RmsNorm::Call( + handle, + config, + planned->x.tensor(planned->x_tensor), + planned->weight.tensor(planned->weight_tensor), + planned->epsilon, + planned->y.tensor(planned->y_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(RMSNorm::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(RMSNorm::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(RMSNorm::cleanup_dispatcher(), &cleanup); + return true; +}(); + +} // namespace infinicore::op::rms_norm_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/rope/rope.cc b/csrc/infinicore/src/ops/rope/rope.cc new file mode 100644 index 000000000..a5d9ebbb3 --- /dev/null +++ b/csrc/infinicore/src/ops/rope/rope.cc @@ -0,0 +1,46 @@ +#include "infinicore/ops/rope.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RoPE); + +RoPE::RoPE(Tensor x_out, + const Tensor &x, + const Tensor &pos, + const Tensor &sin_table, + const Tensor &cos_table, + infinicore::nn::RoPE::Algo algo) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x_out, x, pos, sin_table, cos_table); + INFINICORE_GRAPH_OP_DISPATCH(x_out->device().type(), x_out, x, pos, sin_table, cos_table, algo); +} + +void RoPE::execute(Tensor x_out, + const Tensor &x, + const Tensor &pos, + const Tensor &sin_table, + const Tensor &cos_table, + infinicore::nn::RoPE::Algo algo) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(RoPE, x_out, x, pos, sin_table, cos_table, algo); +} + +void rope_(Tensor x_out, + const Tensor &x, + const Tensor &pos, + const Tensor &sin_table, + const Tensor &cos_table, + infinicore::nn::RoPE::Algo algo) { + RoPE::execute(x_out, x, pos, sin_table, cos_table, algo); +} + +Tensor rope(const Tensor &x, + const Tensor &pos, + const Tensor &sin_table, + const Tensor &cos_table, + infinicore::nn::RoPE::Algo algo) { + auto x_out = Tensor::empty(x->shape(), x->dtype(), x->device()); + rope_(x_out, x, pos, sin_table, cos_table, algo); + return x_out; +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/rope/rope_infiniops.cc b/csrc/infinicore/src/ops/rope/rope_infiniops.cc new file mode 100644 index 000000000..f22f3fe81 --- /dev/null +++ b/csrc/infinicore/src/ops/rope/rope_infiniops.cc @@ -0,0 +1,73 @@ +#include "infinicore/ops/rope.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/rotary_embedding_infinilm.h" + +#include + +namespace infinicore::op::rope_impl::infiniops { +namespace { +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; +struct PlannedMeta { + TensorMeta x_out, x, pos, sin, cos; + graph::GraphTensor x_out_tensor, x_tensor, pos_tensor, sin_tensor, cos_tensor; + bool is_neox; +}; + +bool toInfiniOpsIsNeox(infinicore::nn::RoPE::Algo algo) { + switch (algo) { + case infinicore::nn::RoPE::Algo::GPT_J: + return true; + case infinicore::nn::RoPE::Algo::GPT_NEOX: + return false; + default: + throw std::runtime_error("Unsupported RoPE algorithm"); + } +} +} // namespace + +void *plan(Tensor x_out, + const Tensor &x, + const Tensor &pos, + const Tensor &sin, + const Tensor &cos, + infinicore::nn::RoPE::Algo algo) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(x_out->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x_out, x, pos, sin, cos); + return new PlannedMeta{ + TensorMeta(x_out), TensorMeta(x), TensorMeta(pos), TensorMeta(sin), TensorMeta(cos), + graph::GraphTensor(x_out), graph::GraphTensor(x), graph::GraphTensor(pos), graph::GraphTensor(sin), graph::GraphTensor(cos), + toInfiniOpsIsNeox(algo)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + infini::ops::RotaryEmbeddingInfinilm::Call( + handle, + config, + planned->x.tensor(planned->x_tensor), + planned->pos.tensor(planned->pos_tensor), + planned->sin.tensor(planned->sin_tensor), + planned->cos.tensor(planned->cos_tensor), + planned->is_neox, + planned->x_out.tensor(planned->x_out_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(RoPE::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(RoPE::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(RoPE::cleanup_dispatcher(), &cleanup); + return true; +}(); +} // namespace infinicore::op::rope_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/rot/rot.cc b/csrc/infinicore/src/ops/rot/rot.cc new file mode 100644 index 000000000..7df86354e --- /dev/null +++ b/csrc/infinicore/src/ops/rot/rot.cc @@ -0,0 +1,22 @@ +#include "infinicore/ops/rot.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Rot); + +Rot::Rot(Tensor x, Tensor y, const Tensor &c, const Tensor &s) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y, c, s); + INFINICORE_GRAPH_OP_DISPATCH(x->device().type(), x, y, c, s); +} + +void Rot::execute(Tensor x, Tensor y, const Tensor &c, const Tensor &s) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Rot, x, y, c, s); +} + +void rot_(Tensor x, Tensor y, const Tensor &c, const Tensor &s) { + Rot::execute(x, y, c, s); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/rotg/rotg.cc b/csrc/infinicore/src/ops/rotg/rotg.cc new file mode 100644 index 000000000..b80308197 --- /dev/null +++ b/csrc/infinicore/src/ops/rotg/rotg.cc @@ -0,0 +1,22 @@ +#include "infinicore/ops/rotg.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Rotg::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Rotg::execute(Tensor x, Tensor y, Tensor c, Tensor s) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y, c, s); + infinicore::context::setDevice(x->device()); + dispatcher().lookup(x->device().type())(x, y, c, s); +} + +void rotg_(Tensor x, Tensor y, Tensor c, Tensor s) { + Rotg::execute(x, y, c, s); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/rotm/rotm.cc b/csrc/infinicore/src/ops/rotm/rotm.cc new file mode 100644 index 000000000..957bc679b --- /dev/null +++ b/csrc/infinicore/src/ops/rotm/rotm.cc @@ -0,0 +1,22 @@ +#include "infinicore/ops/rotm.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Rotm::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Rotm::execute(Tensor x, Tensor y, Tensor param) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y, param); + infinicore::context::setDevice(x->device()); + dispatcher().lookup(x->device().type())(x, y, param); +} + +void rotm_(Tensor x, Tensor y, Tensor param) { + Rotm::execute(x, y, param); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/rotmg/rotmg.cc b/csrc/infinicore/src/ops/rotmg/rotmg.cc new file mode 100644 index 000000000..afc5caa67 --- /dev/null +++ b/csrc/infinicore/src/ops/rotmg/rotmg.cc @@ -0,0 +1,22 @@ +#include "infinicore/ops/rotmg.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Rotmg::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Rotmg::execute(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(d1, d2, x1, y1, param); + infinicore::context::setDevice(d1->device()); + dispatcher().lookup(d1->device().type())(d1, d2, x1, y1, param); +} + +void rotmg_(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param) { + Rotmg::execute(d1, d2, x1, y1, param); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/rwkv5_wkv/rwkv5_wkv.cc b/csrc/infinicore/src/ops/rwkv5_wkv/rwkv5_wkv.cc new file mode 100644 index 000000000..67eac1f4a --- /dev/null +++ b/csrc/infinicore/src/ops/rwkv5_wkv/rwkv5_wkv.cc @@ -0,0 +1,50 @@ +#include "infinicore/ops/rwkv5_wkv.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Rwkv5Wkv); + +Rwkv5Wkv::Rwkv5Wkv(Tensor out, + const Tensor &receptance, + const Tensor &key, + const Tensor &value, + const Tensor &time_decay, + const Tensor &time_faaaa, + Tensor state) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, receptance, key, value, time_decay, time_faaaa, state); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), out, receptance, key, value, time_decay, time_faaaa, state); +} + +void Rwkv5Wkv::execute(Tensor out, + const Tensor &receptance, + const Tensor &key, + const Tensor &value, + const Tensor &time_decay, + const Tensor &time_faaaa, + Tensor state) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Rwkv5Wkv, out, receptance, key, value, time_decay, time_faaaa, state); +} + +Tensor rwkv5_wkv(const Tensor &receptance, + const Tensor &key, + const Tensor &value, + const Tensor &time_decay, + const Tensor &time_faaaa, + Tensor state) { + auto output = Tensor::empty(receptance->shape(), receptance->dtype(), receptance->device()); + rwkv5_wkv_(output, receptance, key, value, time_decay, time_faaaa, state); + return output; +} + +void rwkv5_wkv_(Tensor out, + const Tensor &receptance, + const Tensor &key, + const Tensor &value, + const Tensor &time_decay, + const Tensor &time_faaaa, + Tensor state) { + Rwkv5Wkv::execute(out, receptance, key, value, time_decay, time_faaaa, state); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/scal/scal.cc b/csrc/infinicore/src/ops/scal/scal.cc new file mode 100644 index 000000000..5458ce9d4 --- /dev/null +++ b/csrc/infinicore/src/ops/scal/scal.cc @@ -0,0 +1,22 @@ +#include "infinicore/ops/scal.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Scal); + +Scal::Scal(const Tensor &alpha, Tensor x) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(alpha, x); + INFINICORE_GRAPH_OP_DISPATCH(x->device().type(), alpha, x); +} + +void Scal::execute(const Tensor &alpha, Tensor x) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Scal, alpha, x); +} + +void scal_(const Tensor &alpha, Tensor x) { + Scal::execute(alpha, x); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/scaled_mm_i8/scaled_mm_i8.cc b/csrc/infinicore/src/ops/scaled_mm_i8/scaled_mm_i8.cc new file mode 100644 index 000000000..5d28ee3ff --- /dev/null +++ b/csrc/infinicore/src/ops/scaled_mm_i8/scaled_mm_i8.cc @@ -0,0 +1,21 @@ +#include "infinicore/ops/scaled_mm_i8.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(I8Gemm); + +I8Gemm::I8Gemm(Tensor c, const Tensor &a_p, const Tensor &a_s, const Tensor &b_p, const Tensor &b_s, std::optional bias) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a_p, a_s, b_p, b_s); + INFINICORE_GRAPH_OP_DISPATCH(c->device().type(), c, a_p, a_s, b_p, b_s, bias); +} +void I8Gemm::execute(Tensor c, const Tensor &a_p, const Tensor &a_s, const Tensor &b_p, const Tensor &b_s, std::optional bias) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(I8Gemm, c, a_p, a_s, b_p, b_s, bias); +} + +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) { + I8Gemm::execute(c, a_p, a_s, b_p, b_s, bias); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/scatter/scatter.cc b/csrc/infinicore/src/ops/scatter/scatter.cc new file mode 100644 index 000000000..6e9d53701 --- /dev/null +++ b/csrc/infinicore/src/ops/scatter/scatter.cc @@ -0,0 +1,26 @@ +#include "infinicore/ops/scatter.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Scatter::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Scatter::execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction) { + dispatcher().lookup(context::getDevice().type())(output, input, dim, index, src, reduction); +} + +Tensor scatter(Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction) { + // 创建与 input 形状、数据类型、设备一致的 Output Tensor + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + scatter_(output, input, dim, index, src, reduction); + + return output; +} + +void scatter_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction) { + Scatter::execute(output, input, dim, index, src, reduction); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/selu/selu.cc b/csrc/infinicore/src/ops/selu/selu.cc new file mode 100644 index 000000000..f4b3afe29 --- /dev/null +++ b/csrc/infinicore/src/ops/selu/selu.cc @@ -0,0 +1,40 @@ +#include "infinicore/ops/selu.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &Selu::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void Selu::execute(Tensor output, Tensor input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error( + "No Selu implementation found for device type: " + + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor selu(Tensor input) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + selu_(output, input); + return output; +} + +void selu_(Tensor output, Tensor input) { + Selu::execute(output, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/sigmoid/sigmoid.cc b/csrc/infinicore/src/ops/sigmoid/sigmoid.cc new file mode 100644 index 000000000..3f426aca8 --- /dev/null +++ b/csrc/infinicore/src/ops/sigmoid/sigmoid.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/sigmoid.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Sigmoid); + +Sigmoid::Sigmoid(Tensor output, const Tensor &input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + INFINICORE_GRAPH_OP_DISPATCH(output->device().type(), output, input); +} + +void Sigmoid::execute(Tensor output, const Tensor &input) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Sigmoid, output, input); +} + +Tensor sigmoid(const Tensor &input) { + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + sigmoid_(output, input); + return output; +} + +void sigmoid_(Tensor output, const Tensor &input) { + Sigmoid::execute(output, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/sigmoid/sigmoid_infiniops.cc b/csrc/infinicore/src/ops/sigmoid/sigmoid_infiniops.cc new file mode 100644 index 000000000..941d39495 --- /dev/null +++ b/csrc/infinicore/src/ops/sigmoid/sigmoid_infiniops.cc @@ -0,0 +1,58 @@ +#include "infinicore/ops/sigmoid.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/sigmoid_infinilm.h" + +namespace infinicore::op::sigmoid_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +struct PlannedMeta { + TensorMeta output, input; + graph::GraphTensor output_tensor, input_tensor; +}; + +} // namespace + +void *plan(Tensor output, const Tensor &input) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + + return new PlannedMeta{ + TensorMeta(output), + TensorMeta(input), + graph::GraphTensor(output), + graph::GraphTensor(input)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + infini::ops::SigmoidInfinilm::Call( + handle, + config, + planned->input.tensor(planned->input_tensor), + planned->output.tensor(planned->output_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Sigmoid::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(Sigmoid::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(Sigmoid::cleanup_dispatcher(), &cleanup); + return true; +}(); + +} // namespace infinicore::op::sigmoid_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/silu/silu.cc b/csrc/infinicore/src/ops/silu/silu.cc new file mode 100644 index 000000000..68e402d15 --- /dev/null +++ b/csrc/infinicore/src/ops/silu/silu.cc @@ -0,0 +1,37 @@ +#include "infinicore/ops/silu.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &Silu::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Silu::execute(Tensor output, Tensor input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Silu implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor silu(Tensor input) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + silu_(output, input); + return output; +} + +void silu_(Tensor output, Tensor input) { + Silu::execute(output, input); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/silu/silu_infiniops.cc b/csrc/infinicore/src/ops/silu/silu_infiniops.cc new file mode 100644 index 000000000..63ea5be94 --- /dev/null +++ b/csrc/infinicore/src/ops/silu/silu_infiniops.cc @@ -0,0 +1,38 @@ +#include "infinicore/ops/silu.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/silu.h" + +namespace infinicore::op::silu_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +void calculate(Tensor output, Tensor input) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + TensorMeta output_meta(output); + TensorMeta input_meta(input); + infini::ops::Silu::Call( + handle, + config, + input_meta.tensor(input), + output_meta.tensor(output)); +} + +} // namespace + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Silu::dispatcher(), &calculate); + return true; +}(); + +} // namespace infinicore::op::silu_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/silu_and_mul/silu_and_mul.cc b/csrc/infinicore/src/ops/silu_and_mul/silu_and_mul.cc new file mode 100644 index 000000000..04bc975cd --- /dev/null +++ b/csrc/infinicore/src/ops/silu_and_mul/silu_and_mul.cc @@ -0,0 +1,35 @@ +#include "infinicore/ops/silu_and_mul.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(SiluAndMul); + +SiluAndMul::SiluAndMul(Tensor out, const Tensor &x) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, x); + INFINICORE_GRAPH_OP_DISPATCH(out->device().type(), out, x); +} + +void SiluAndMul::execute(Tensor out, const Tensor &x) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(SiluAndMul, out, x); +} + +Tensor silu_and_mul(const Tensor &x) { + Shape shape = x->shape(); + size_t ndim = x->ndim(); + + if (shape[ndim - 1] % 2 != 0) { + throw std::runtime_error("SiluAndMul input last dim must be even."); + } + shape[ndim - 1] /= 2; + + auto out = Tensor::empty(shape, x->dtype(), x->device()); + silu_and_mul_(out, x); + return out; +} + +void silu_and_mul_(Tensor out, const Tensor &x) { + SiluAndMul::execute(out, x); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/silu_and_mul/silu_and_mul_infiniops.cc b/csrc/infinicore/src/ops/silu_and_mul/silu_and_mul_infiniops.cc new file mode 100644 index 000000000..92b482962 --- /dev/null +++ b/csrc/infinicore/src/ops/silu_and_mul/silu_and_mul_infiniops.cc @@ -0,0 +1,44 @@ +#include "infinicore/ops/silu_and_mul.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/silu_and_mul_infinilm.h" + +namespace infinicore::op::silu_and_mul_impl::infiniops { +namespace { +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; +struct PlannedMeta { + TensorMeta output, input; + graph::GraphTensor output_tensor, input_tensor; +}; +} // namespace + +void *plan(Tensor output, const Tensor &input) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + return new PlannedMeta{TensorMeta(output), TensorMeta(input), graph::GraphTensor(output), graph::GraphTensor(input)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + infini::ops::SiluAndMulInfinilm::Call( + handle, config, planned->input.tensor(planned->input_tensor), planned->output.tensor(planned->output_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(SiluAndMul::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(SiluAndMul::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(SiluAndMul::cleanup_dispatcher(), &cleanup); + return true; +}(); +} // namespace infinicore::op::silu_and_mul_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/sinh/sinh.cc b/csrc/infinicore/src/ops/sinh/sinh.cc new file mode 100644 index 000000000..380d3b139 --- /dev/null +++ b/csrc/infinicore/src/ops/sinh/sinh.cc @@ -0,0 +1,40 @@ +#include "infinicore/ops/sinh.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &Sinh::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} + +void Sinh::execute(Tensor output, Tensor input) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error( + "No Sinh implementation found for device type: " + + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor sinh(Tensor input) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + sinh_(output, input); + return output; +} + +void sinh_(Tensor output, Tensor input) { + Sinh::execute(output, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/smooth_l1_loss/smooth_l1_loss.cc b/csrc/infinicore/src/ops/smooth_l1_loss/smooth_l1_loss.cc new file mode 100644 index 000000000..a3ea10b17 --- /dev/null +++ b/csrc/infinicore/src/ops/smooth_l1_loss/smooth_l1_loss.cc @@ -0,0 +1,34 @@ +#include "infinicore/ops/smooth_l1_loss.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &SmoothL1Loss::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void SmoothL1Loss::execute(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction) { + dispatcher().lookup(context::getDevice().type())(output, input, target, beta, reduction); +} + +Tensor smooth_l1_loss(Tensor input, Tensor target, float beta, int64_t reduction) { + Shape output_shape; + if (reduction == 0) { + // Reduction::None -> 输出形状与输入一致 + output_shape = input->shape(); + } else { + output_shape = {}; + } + + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + + smooth_l1_loss_(output, input, target, beta, reduction); + return output; +} + +void smooth_l1_loss_(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction) { + SmoothL1Loss::execute(output, input, target, beta, reduction); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/softmax/softmax.cc b/csrc/infinicore/src/ops/softmax/softmax.cc new file mode 100644 index 000000000..e5ef657cb --- /dev/null +++ b/csrc/infinicore/src/ops/softmax/softmax.cc @@ -0,0 +1,37 @@ +#include "infinicore/ops/softmax.hpp" + +#include "../../utils.hpp" + +#include + +namespace infinicore::op { + +common::OpDispatcher &Softmax::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Softmax::execute(Tensor output, Tensor input, int axis) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(output->device()); + auto device_type = output->device().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Softmax implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input, axis); +} + +Tensor softmax(Tensor input, int axis) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + softmax_(output, input, axis); + return output; +} + +void softmax_(Tensor output, Tensor input, int axis) { + Softmax::execute(output, input, axis); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/softmax/softmax_infiniops.cc b/csrc/infinicore/src/ops/softmax/softmax_infiniops.cc new file mode 100644 index 000000000..833b41afd --- /dev/null +++ b/csrc/infinicore/src/ops/softmax/softmax_infiniops.cc @@ -0,0 +1,42 @@ +#include "infinicore/ops/softmax.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/softmax_infinilm.h" + +#include + +namespace infinicore::op::softmax_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +void calculate(Tensor output, Tensor input, int axis) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + TensorMeta output_meta(output); + TensorMeta input_meta(input); + infini::ops::SoftmaxInfinilm::Call( + handle, + config, + input_meta.tensor(input), + static_cast(axis), + std::optional{}, + output_meta.tensor(output)); +} + +} // namespace + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Softmax::dispatcher(), &calculate); + return true; +}(); + +} // namespace infinicore::op::softmax_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/softplus/softplus.cc b/csrc/infinicore/src/ops/softplus/softplus.cc new file mode 100644 index 000000000..620d9fb05 --- /dev/null +++ b/csrc/infinicore/src/ops/softplus/softplus.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/softplus.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Softplus::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Softplus::execute(Tensor y, Tensor x, float beta, float threshold) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + infinicore::context::setDevice(y->device()); + + dispatcher().lookup(y->device().type())(y, x, beta, threshold); +} + +Tensor softplus(Tensor x, float beta, float threshold) { + auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); + softplus_(y, x, beta, threshold); + return y; +} + +void softplus_(Tensor y, Tensor x, float beta, float threshold) { + Softplus::execute(y, x, beta, threshold); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/softsign/softsign.cc b/csrc/infinicore/src/ops/softsign/softsign.cc new file mode 100644 index 000000000..84827b30f --- /dev/null +++ b/csrc/infinicore/src/ops/softsign/softsign.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/softsign.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &Softsign::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Softsign::execute(Tensor y, Tensor x) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + infinicore::context::setDevice(y->device()); + dispatcher().lookup(y->device().type())(y, x); +} + +Tensor softsign(Tensor x) { + auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); + softsign_(y, x); + return y; +} + +void softsign_(Tensor y, Tensor x) { + Softsign::execute(y, x); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/sum/sum.cc b/csrc/infinicore/src/ops/sum/sum.cc new file mode 100644 index 000000000..7b0e8794c --- /dev/null +++ b/csrc/infinicore/src/ops/sum/sum.cc @@ -0,0 +1,67 @@ +#include "infinicore/ops/sum.hpp" + +#include "../../utils.hpp" +#include +#include + +namespace infinicore::op { + +common::OpDispatcher &Sum::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; +void Sum::execute(Tensor output, Tensor input, std::vector dim, bool keepdim) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); + infinicore::context::setDevice(input->device()); + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Sum implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input, dim, keepdim); +} + +Tensor sum(Tensor input, std::vector dim, bool keepdim) { + auto in_shape = input->shape(); + std::vector out_shape; + if (dim.empty()) { + for (size_t i = 0; i < in_shape.size(); i++) { + dim.push_back(i); + } + } + std::sort(dim.begin(), dim.end()); + if (dim.size() == in_shape.size() && !keepdim) { + out_shape = {}; + } else { + if (keepdim) { + size_t j = 0; + for (size_t i = 0; i < in_shape.size(); i++) { + if (j < dim.size() && dim[j] == i) { + out_shape.push_back(1); + j++; + } else { + out_shape.push_back(in_shape[i]); + } + } + } else { + size_t j = 0; + for (size_t i = 0; i < in_shape.size(); i++) { + if (j < dim.size() && dim[j] == i) { + j++; + } else { + out_shape.push_back(in_shape[i]); + } + } + } + } + auto output = Tensor::empty(out_shape, input->dtype(), input->device()); + sum_(output, input, dim, keepdim); + return output; +} + +void sum_(Tensor output, Tensor input, std::vector dim, bool keepdim) { + Sum::execute(output, input, dim, keepdim); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/swap/swap.cc b/csrc/infinicore/src/ops/swap/swap.cc new file mode 100644 index 000000000..f5625b1e4 --- /dev/null +++ b/csrc/infinicore/src/ops/swap/swap.cc @@ -0,0 +1,22 @@ +#include "infinicore/ops/swap.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Swap); + +Swap::Swap(Tensor x, Tensor y) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y); + INFINICORE_GRAPH_OP_DISPATCH(x->device().type(), x, y); +} + +void Swap::execute(Tensor x, Tensor y) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Swap, x, y); +} + +void swap_(Tensor x, Tensor y) { + Swap::execute(x, y); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/swiglu/swiglu.cc b/csrc/infinicore/src/ops/swiglu/swiglu.cc new file mode 100644 index 000000000..39deacd4e --- /dev/null +++ b/csrc/infinicore/src/ops/swiglu/swiglu.cc @@ -0,0 +1,26 @@ +#include "infinicore/ops/swiglu.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(SwiGLU); + +SwiGLU::SwiGLU(Tensor c, const Tensor &a, const Tensor &b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + INFINICORE_GRAPH_OP_DISPATCH(c->device().type(), c, a, b); +} + +void SwiGLU::execute(Tensor c, const Tensor &a, const Tensor &b) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(SwiGLU, c, a, b); +} + +Tensor swiglu(const Tensor &a, const Tensor &b) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + swiglu_(c, a, b); + return c; +} + +void swiglu_(Tensor c, const Tensor &a, const Tensor &b) { + SwiGLU::execute(c, a, b); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/swiglu/swiglu_infiniops.cc b/csrc/infinicore/src/ops/swiglu/swiglu_infiniops.cc new file mode 100644 index 000000000..9860e2a56 --- /dev/null +++ b/csrc/infinicore/src/ops/swiglu/swiglu_infiniops.cc @@ -0,0 +1,61 @@ +#include "infinicore/ops/swiglu.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/swiglu.h" + +namespace infinicore::op::swiglu_impl::infiniops { +namespace { + +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; + +struct PlannedMeta { + TensorMeta c, a, b; + graph::GraphTensor c_tensor, a_tensor, b_tensor; +}; + +} // namespace + +void *plan(Tensor c, const Tensor &a, const Tensor &b) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(c->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + + return new PlannedMeta{ + TensorMeta(c), + TensorMeta(a), + TensorMeta(b), + graph::GraphTensor(c), + graph::GraphTensor(a), + graph::GraphTensor(b)}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + + infini::ops::Swiglu::Call( + handle, + config, + planned->a.tensor(planned->a_tensor), + planned->b.tensor(planned->b_tensor), + planned->c.tensor(planned->c_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::cleanup_dispatcher(), &cleanup); + return true; +}(); + +} // namespace infinicore::op::swiglu_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/take/take.cc b/csrc/infinicore/src/ops/take/take.cc new file mode 100644 index 000000000..08ff8d337 --- /dev/null +++ b/csrc/infinicore/src/ops/take/take.cc @@ -0,0 +1,28 @@ +#include "infinicore/ops/take.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &Take::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +// 2. Execute 实现：查找对应设备的核函数并执行 +void Take::execute(Tensor output, Tensor input, Tensor indices) { + dispatcher().lookup(context::getDevice().type())(output, input, indices); +} + +Tensor take(Tensor input, Tensor indices) { + // 【关键区别】Take 的输出形状取决于 indices 的形状，但数据类型取决于 input + auto output = Tensor::empty(indices->shape(), input->dtype(), input->device()); + + take_(output, input, indices); + return output; +} + +void take_(Tensor output, Tensor input, Tensor indices) { + Take::execute(output, input, indices); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/tan/tan.cc b/csrc/infinicore/src/ops/tan/tan.cc new file mode 100644 index 000000000..d2283b5cc --- /dev/null +++ b/csrc/infinicore/src/ops/tan/tan.cc @@ -0,0 +1,32 @@ +#include "infinicore/ops/tan.hpp" +#include + +namespace infinicore::op { + +common::OpDispatcher &Tan::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Tan::execute(Tensor output, Tensor input) { + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Tan implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor tan(Tensor input) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + tan_(output, input); + return output; +} + +void tan_(Tensor output, Tensor input) { + Tan::execute(output, input); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/tanhshrink/tanhshrink.cc b/csrc/infinicore/src/ops/tanhshrink/tanhshrink.cc new file mode 100644 index 000000000..8bf9eb3cc --- /dev/null +++ b/csrc/infinicore/src/ops/tanhshrink/tanhshrink.cc @@ -0,0 +1,32 @@ +#include "infinicore/ops/tanhshrink.hpp" +#include + +namespace infinicore::op { + +common::OpDispatcher &Tanhshrink::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Tanhshrink::execute(Tensor output, Tensor input) { + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Tanhshrink implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(output, input); +} + +Tensor tanhshrink(Tensor input) { + Shape shape = input->shape(); + auto output = Tensor::empty(shape, input->dtype(), input->device()); + tanhshrink_(output, input); + return output; +} + +void tanhshrink_(Tensor output, Tensor input) { + Tanhshrink::execute(output, input); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/topk/topk.cc b/csrc/infinicore/src/ops/topk/topk.cc new file mode 100644 index 000000000..219709e0e --- /dev/null +++ b/csrc/infinicore/src/ops/topk/topk.cc @@ -0,0 +1,40 @@ +#include "infinicore/ops/topk.hpp" + +#include "../../utils.hpp" +#include +#include + +namespace infinicore::op { + +common::OpDispatcher &TopK::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; +void TopK::execute(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest, bool sorted) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(values_output, input); + infinicore::context::setDevice(input->device()); + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Topk implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(values_output, indices_output, input, k, dim, largest, sorted); +} + +std::pair topk(Tensor input, size_t k, size_t dim, bool largest, bool sorted) { + auto in_shape = input->shape(); + std::vector out_shape = in_shape; + out_shape[dim] = k; + + auto values_output = Tensor::empty(out_shape, input->dtype(), input->device()); + auto indices_output = Tensor::empty(out_shape, DataType::kInt32, input->device()); + topk_(values_output, indices_output, input, k, dim, largest, sorted); + return {values_output, indices_output}; +} + +void topk_(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest, bool sorted) { + TopK::execute(values_output, indices_output, input, k, dim, largest, sorted); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/topksoftmax/topksoftmax.cc b/csrc/infinicore/src/ops/topksoftmax/topksoftmax.cc new file mode 100644 index 000000000..48f1bc602 --- /dev/null +++ b/csrc/infinicore/src/ops/topksoftmax/topksoftmax.cc @@ -0,0 +1,30 @@ +#include "infinicore/ops/topksoftmax.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Topksoftmax); + +Topksoftmax::Topksoftmax(Tensor values, + Tensor indices, + const Tensor &x, + const size_t topk, + const int norm) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(values, indices, x); + INFINICORE_GRAPH_OP_DISPATCH(values->device().type(), values, indices, x, topk, norm); +} + +void Topksoftmax::execute(Tensor values, + Tensor indices, + const Tensor &x, + const size_t topk, + const int norm) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Topksoftmax, values, indices, x, topk, norm); +} + +void topksoftmax(Tensor values, Tensor indices, const Tensor &x, const size_t topk, const int norm) { + Topksoftmax::execute(values, indices, x, topk, norm); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/topksoftmax/topksoftmax_infiniops.cc b/csrc/infinicore/src/ops/topksoftmax/topksoftmax_infiniops.cc new file mode 100644 index 000000000..8b20b1aff --- /dev/null +++ b/csrc/infinicore/src/ops/topksoftmax/topksoftmax_infiniops.cc @@ -0,0 +1,52 @@ +#include "infinicore/ops/topksoftmax.hpp" + +#ifdef ENABLE_INFINIOPS_API +#include "../infiniops_impl.hpp" + +#include "base/topksoftmax_infinilm.h" + +namespace infinicore::op::topksoftmax_impl::infiniops { +namespace { +using TensorMeta = ::infinicore::op::infiniops::TensorMeta; +struct PlannedMeta { + TensorMeta values, indices, x; + graph::GraphTensor values_tensor, indices_tensor, x_tensor; + size_t topk; + int norm; +}; +} // namespace + +void *plan(Tensor values, Tensor indices, const Tensor &x, const size_t topk, const int norm) { + INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(values->device().type())); + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(values, indices, x); + return new PlannedMeta{TensorMeta(values), TensorMeta(indices), TensorMeta(x), graph::GraphTensor(values), graph::GraphTensor(indices), graph::GraphTensor(x), topk, norm}; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + infini::ops::Handle handle; + handle.set_stream(context::getStream()); + infini::ops::Config config; + infini::ops::TopksoftmaxInfinilm::Call( + handle, + config, + planned->x.tensor(planned->x_tensor), + static_cast(planned->topk), + planned->norm != 0, + planned->values.tensor(planned->values_tensor), + planned->indices.tensor(planned->indices_tensor)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + ::infinicore::op::infiniops::registerSupportedDevices(Topksoftmax::plan_dispatcher(), &plan); + ::infinicore::op::infiniops::registerSupportedDevices(Topksoftmax::run_dispatcher(), &run); + ::infinicore::op::infiniops::registerSupportedDevices(Topksoftmax::cleanup_dispatcher(), &cleanup); + return true; +}(); +} // namespace infinicore::op::topksoftmax_impl::infiniops +#endif diff --git a/csrc/infinicore/src/ops/triplet_margin_loss/triplet_margin_loss.cc b/csrc/infinicore/src/ops/triplet_margin_loss/triplet_margin_loss.cc new file mode 100644 index 000000000..78e6099f5 --- /dev/null +++ b/csrc/infinicore/src/ops/triplet_margin_loss/triplet_margin_loss.cc @@ -0,0 +1,37 @@ +#include "infinicore/ops/triplet_margin_loss.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &TripletMarginLoss::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void TripletMarginLoss::execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction) { + dispatcher().lookup(context::getDevice().type())(output, anchor, positive, negative, margin, p, eps, swap, reduction); +} + +// 3. 函数式接口 +Tensor triplet_margin_loss(Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction) { + Shape output_shape; + if (reduction == 0) { // None + // TripletMarginLoss 输入通常为 (N, D)，reduction='none' 时输出为 (N) + // 取第 0 维作为 Batch Size + output_shape = {anchor->shape()[0]}; + } else { + output_shape = {}; // Scalar + } + + // 使用 anchor 的属性创建输出 Tensor + auto output = Tensor::empty(output_shape, anchor->dtype(), anchor->device()); + + triplet_margin_loss_(output, anchor, positive, negative, margin, p, eps, swap, reduction); + return output; +} + +void triplet_margin_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction) { + TripletMarginLoss::execute(output, anchor, positive, negative, margin, p, eps, swap, reduction); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.cc b/csrc/infinicore/src/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.cc new file mode 100644 index 000000000..dfeddd6c6 --- /dev/null +++ b/csrc/infinicore/src/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.cc @@ -0,0 +1,38 @@ +#include "infinicore/ops/triplet_margin_with_distance_loss.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &TripletMarginWithDistanceLoss::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void TripletMarginWithDistanceLoss::execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { + dispatcher().lookup(context::getDevice().type())(output, anchor, positive, negative, margin, swap, reduction); +} + +// 3. 函数式接口 +Tensor triplet_margin_with_distance_loss(Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { + Shape out_shape; + + // 推断输出形状 + if (reduction == 0) { + // Reduction::None -> 输出形状取决于输入的广播结果 + out_shape = anchor->shape(); + } else { + // Reduction::Mean 或 Reduction::Sum -> 输出为标量 + out_shape = {}; + } + + auto output = Tensor::empty(out_shape, anchor->dtype(), anchor->device()); + + triplet_margin_with_distance_loss_(output, anchor, positive, negative, margin, swap, reduction); + return output; +} + +void triplet_margin_with_distance_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { + TripletMarginWithDistanceLoss::execute(output, anchor, positive, negative, margin, swap, reduction); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/unfold/unfold.cc b/csrc/infinicore/src/ops/unfold/unfold.cc new file mode 100644 index 000000000..56afcd21d --- /dev/null +++ b/csrc/infinicore/src/ops/unfold/unfold.cc @@ -0,0 +1,72 @@ +#include "infinicore/ops/unfold.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &Unfold::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +// 2. Execute 方法实现 +void Unfold::execute(Tensor output, Tensor input, + const std::vector &kernel_sizes, + const std::vector &dilations, + const std::vector &paddings, + const std::vector &strides) { + dispatcher().lookup(context::getDevice().type())(output, input, kernel_sizes, dilations, paddings, strides); +} + +// 3. 函数式接口 +Tensor unfold(Tensor input, + std::vector kernel_sizes, + std::vector dilations, + std::vector paddings, + std::vector strides) { + + // 基础维度校验与获取 + const auto &input_shape = input->shape(); + int64_t n_dim = input->ndim(); + int64_t spatial_dims = n_dim - 2; // N, C, D1, D2... -> spatial starts at 2 + int64_t N = input_shape[0]; + int64_t C = input_shape[1]; + + // 计算 dim 1: C * kernel_sizes[0] * kernel_sizes[1] ... + int64_t output_dim1 = C; + for (auto k : kernel_sizes) { + output_dim1 *= k; + } + + int64_t L = 1; + for (int i = 0; i < spatial_dims; ++i) { + int64_t input_dim = input_shape[i + 2]; + int64_t k = kernel_sizes[i]; + int64_t p = paddings[i]; + int64_t d = dilations[i]; + int64_t s = strides[i]; + + // 公式: out = floor((in + 2*p - d*(k-1) - 1) / s + 1) + int64_t output_spatial = (input_dim + 2 * p - d * (k - 1) - 1) / s + 1; + L *= output_spatial; + } + Shape output_shape = { + static_cast(N), + static_cast(output_dim1), + static_cast(L)}; + + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + + unfold_(output, input, kernel_sizes, dilations, paddings, strides); + return output; +} + +// 4. In-place / 显式输出接口 +void unfold_(Tensor output, Tensor input, + std::vector kernel_sizes, + std::vector dilations, + std::vector paddings, + std::vector strides) { + Unfold::execute(output, input, kernel_sizes, dilations, paddings, strides); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/upsample_bilinear/upsample_bilinear.cc b/csrc/infinicore/src/ops/upsample_bilinear/upsample_bilinear.cc new file mode 100644 index 000000000..c4fcbd2fb --- /dev/null +++ b/csrc/infinicore/src/ops/upsample_bilinear/upsample_bilinear.cc @@ -0,0 +1,41 @@ +#include "infinicore/ops/upsample_bilinear.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &UpsampleBilinear::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void UpsampleBilinear::execute(Tensor output, Tensor input, bool align_corners) { + dispatcher().lookup(context::getDevice().type())(output, input, align_corners); +} + +// 3. 函数式接口 +Tensor upsample_bilinear(Tensor input, std::vector output_size, bool align_corners) { + // 构造输出 Shape + // 假设 input 是 (N, C, H_in, W_in) 或 (C, H_in, W_in) + // output_size 通常只包含 (H_out, W_out) + Shape input_shape = input->shape(); + size_t ndim = input_shape.size(); + + Shape output_shape = input_shape; + + // 更新最后两个维度为 output_size 指定的大小 + if (output_size.size() == 2 && ndim >= 2) { + output_shape[ndim - 2] = output_size[0]; + output_shape[ndim - 1] = output_size[1]; + } + + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + + upsample_bilinear_(output, input, align_corners); + return output; +} + +void upsample_bilinear_(Tensor output, Tensor input, bool align_corners) { + UpsampleBilinear::execute(output, input, align_corners); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/upsample_nearest/upsample_nearest.cc b/csrc/infinicore/src/ops/upsample_nearest/upsample_nearest.cc new file mode 100644 index 000000000..adfac98fd --- /dev/null +++ b/csrc/infinicore/src/ops/upsample_nearest/upsample_nearest.cc @@ -0,0 +1,61 @@ +#include "infinicore/ops/upsample_nearest.hpp" +#include +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &UpsampleNearest::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void UpsampleNearest::execute(Tensor output, Tensor input) { + dispatcher().lookup(context::getDevice().type())(output, input); +} + +// 3. 函数式接口 +Tensor upsample_nearest(Tensor input, const std::vector &output_size) { + Shape input_shape = input->shape(); + size_t ndim = input_shape.size(); + + // 校验 + if (ndim < 3 || ndim > 4) { + if (ndim != 3 && ndim != 4) { + throw std::runtime_error("upsample_nearest: Only supports 3D (N,C,W) or 4D (N,C,H,W) input"); + } + } + + Shape output_shape = input_shape; + + if (ndim == 3) { + // [N, C, W] + // output_size 可能是 [W_out] (size=1) 或者 [1, W_out] (size=2) + int64_t target_w = 0; + if (output_size.size() == 1) { + target_w = output_size[0]; + } else if (output_size.size() == 2) { + target_w = output_size[1]; + } else { + throw std::runtime_error("upsample_nearest: output_size for 3D input must be [w] or [1, w]"); + } + output_shape[2] = target_w; + + } else if (ndim == 4) { + // [N, C, H, W] + if (output_size.size() != 2) { + throw std::runtime_error("upsample_nearest: output_size for 4D input must be [h, w]"); + } + output_shape[2] = output_size[0]; + output_shape[3] = output_size[1]; + } + + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + + upsample_nearest_(output, input); + return output; +} + +void upsample_nearest_(Tensor output, Tensor input) { + UpsampleNearest::execute(output, input); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/vander/vander.cc b/csrc/infinicore/src/ops/vander/vander.cc new file mode 100644 index 000000000..fe29ff800 --- /dev/null +++ b/csrc/infinicore/src/ops/vander/vander.cc @@ -0,0 +1,35 @@ +#include "infinicore/ops/vander.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &Vander::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +// 2. Execute 方法实现 +void Vander::execute(Tensor output, Tensor input, int64_t N, bool increasing) { + dispatcher().lookup(context::getDevice().type())(output, input, N, increasing); +} + +// 3. 函数式接口 +Tensor vander(Tensor input, int64_t N, bool increasing) { + int64_t input_size = input->shape()[0]; + int64_t cols = (N > 0) ? N : input_size; + Shape output_shape = { + static_cast(input_size), + static_cast(cols)}; + + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + + vander_(output, input, N, increasing); + return output; +} + +// 4. In-place / 显式输出接口 +void vander_(Tensor output, Tensor input, int64_t N, bool increasing) { + Vander::execute(output, input, N, increasing); +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/var/var.cc b/csrc/infinicore/src/ops/var/var.cc new file mode 100644 index 000000000..1d941cfcc --- /dev/null +++ b/csrc/infinicore/src/ops/var/var.cc @@ -0,0 +1,68 @@ +#include "infinicore/ops/var.hpp" + +#include "../../utils.hpp" +#include +#include + +namespace infinicore::op { + +common::OpDispatcher &Var::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Var::execute(Tensor var_output, Tensor input, std::vector dim, bool unbiased, bool keepdim) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(var_output, input); + infinicore::context::setDevice(input->device()); + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Var implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(var_output, input, dim, unbiased, keepdim); +} + +Tensor var(Tensor input, std::vector dim, bool unbiased, bool keepdim) { + auto in_shape = input->shape(); + std::vector out_shape; + if (dim.empty()) { + for (size_t i = 0; i < in_shape.size(); i++) { + dim.push_back(i); + } + } + std::sort(dim.begin(), dim.end()); + if (dim.size() == in_shape.size() && !keepdim) { + out_shape = {}; + } else { + if (keepdim) { + size_t j = 0; + for (size_t i = 0; i < in_shape.size(); i++) { + if (j < dim.size() && dim[j] == i) { + out_shape.push_back(1); + j++; + } else { + out_shape.push_back(in_shape[i]); + } + } + } else { + size_t j = 0; + for (size_t i = 0; i < in_shape.size(); i++) { + if (j < dim.size() && dim[j] == i) { + j++; + } else { + out_shape.push_back(in_shape[i]); + } + } + } + } + auto var_output = Tensor::empty(out_shape, input->dtype(), input->device()); + var_(var_output, input, dim, unbiased, keepdim); + return var_output; +} + +void var_(Tensor var_output, Tensor input, std::vector dim, bool unbiased, bool keepdim) { + Var::execute(var_output, input, dim, unbiased, keepdim); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/ops/var_mean/var_mean.cc b/csrc/infinicore/src/ops/var_mean/var_mean.cc new file mode 100644 index 000000000..b0a1c8fc7 --- /dev/null +++ b/csrc/infinicore/src/ops/var_mean/var_mean.cc @@ -0,0 +1,69 @@ +#include "infinicore/ops/var_mean.hpp" + +#include "../../utils.hpp" +#include +#include + +namespace infinicore::op { + +common::OpDispatcher &Var_Mean::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void Var_Mean::execute(Tensor var_output, Tensor mean_output, Tensor input, std::vector dim, bool unbiased, bool keepdim) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(var_output, mean_output, input); + infinicore::context::setDevice(input->device()); + auto device_type = context::getDevice().type(); + auto func = dispatcher().lookup(device_type); + + if (func == nullptr) { + throw std::runtime_error("No Var_Mean implementation found for device type: " + std::to_string(static_cast(device_type))); + } + + func(var_output, mean_output, input, dim, unbiased, keepdim); +} + +std::pair var_mean(Tensor input, std::vector dim, bool unbiased, bool keepdim) { + auto in_shape = input->shape(); + std::vector out_shape; + if (dim.empty()) { + for (size_t i = 0; i < in_shape.size(); i++) { + dim.push_back(i); + } + } + std::sort(dim.begin(), dim.end()); + if (dim.size() == in_shape.size() && !keepdim) { + out_shape = {}; + } else { + if (keepdim) { + size_t j = 0; + for (size_t i = 0; i < in_shape.size(); i++) { + if (j < dim.size() && dim[j] == i) { + out_shape.push_back(1); + j++; + } else { + out_shape.push_back(in_shape[i]); + } + } + } else { + size_t j = 0; + for (size_t i = 0; i < in_shape.size(); i++) { + if (j < dim.size() && dim[j] == i) { + j++; + } else { + out_shape.push_back(in_shape[i]); + } + } + } + } + auto var_output = Tensor::empty(out_shape, input->dtype(), input->device()); + auto mean_output = Tensor::empty(out_shape, input->dtype(), input->device()); + var_mean_(var_output, mean_output, input, dim, unbiased, keepdim); + return {var_output, mean_output}; +} + +void var_mean_(Tensor var_output, Tensor mean_output, Tensor input, std::vector dim, bool unbiased, bool keepdim) { + Var_Mean::execute(var_output, mean_output, input, dim, unbiased, keepdim); +} +} // namespace infinicore::op diff --git a/csrc/infinicore/src/pybind11/analyzer.hpp b/csrc/infinicore/src/pybind11/analyzer.hpp new file mode 100644 index 000000000..1d799991a --- /dev/null +++ b/csrc/infinicore/src/pybind11/analyzer.hpp @@ -0,0 +1,168 @@ +#pragma once + +#include +#include + +#include "infinicore/analyzer/mutual_awareness_analyzer.hpp" + +namespace py = pybind11; + +namespace infinicore::analyzer::pybind { + +inline void bind(py::module &m) { + auto analyzer_mod = m.def_submodule("analyzer", + "Hardware-Task Mutual Awareness Analysis Module"); + + // --- Enums --- + py::enum_(analyzer_mod, "PhaseType") + .value("UNKNOWN", PhaseType::UNKNOWN) + .value("PREFILL", PhaseType::PREFILL) + .value("DECODE", PhaseType::DECODE) + .value("ATTENTION_DENSE", PhaseType::ATTENTION_DENSE) + .value("GEMM_MLP_DENSE", PhaseType::GEMM_MLP_DENSE) + .value("MOE_ROUTING", PhaseType::MOE_ROUTING) + .value("KV_CACHE", PhaseType::KV_CACHE) + .value("COMMUNICATION", PhaseType::COMMUNICATION) + .export_values(); + + py::enum_(analyzer_mod, "BottleneckType") + .value("COMPUTE_BOUND", BottleneckType::COMPUTE_BOUND) + .value("MEMORY_BOUND", BottleneckType::MEMORY_BOUND) + .value("BANDWIDTH_BOUND", BottleneckType::BANDWIDTH_BOUND) + .value("COMMUNICATION_BOUND", BottleneckType::COMMUNICATION_BOUND) + .value("BALANCED", BottleneckType::BALANCED) + .export_values(); + + py::enum_(analyzer_mod, "OptimizationGoal") + .value("LATENCY_FIRST", OptimizationGoal::LATENCY_FIRST) + .value("THROUGHPUT_FIRST", OptimizationGoal::THROUGHPUT_FIRST) + .value("MEMORY_SAFE", OptimizationGoal::MEMORY_SAFE) + .value("STABILITY_FIRST", OptimizationGoal::STABILITY_FIRST) + .export_values(); + + py::enum_(analyzer_mod, "OpType") + .value("UNKNOWN", OpType::UNKNOWN) + .value("ATTENTION", OpType::ATTENTION) + .value("FLASH_ATTENTION", OpType::FLASH_ATTENTION) + .value("GEMM", OpType::GEMM) + .value("LINEAR", OpType::LINEAR) + .value("MATMUL", OpType::MATMUL) + .value("SILU", OpType::SILU) + .value("GELU", OpType::GELU) + .value("RMS_NORM", OpType::RMS_NORM) + .value("KV_CACHING", OpType::KV_CACHING) + .value("PAGED_CACHING", OpType::PAGED_CACHING) + .value("EMBEDDING", OpType::EMBEDDING) + .value("ROPE", OpType::ROPE) + .export_values(); + + // --- StrategyHint --- + py::class_(analyzer_mod, "StrategyHint") + .def(py::init<>()) + .def_readwrite("prefer_fused_ops", &StrategyHint::prefer_fused_ops) + .def_readwrite("prefer_in_place", &StrategyHint::prefer_in_place) + .def_readwrite("prefer_recomputation", &StrategyHint::prefer_recomputation) + .def_readwrite("prefer_async_comm", &StrategyHint::prefer_async_comm) + .def("__repr__", [](const StrategyHint &s) { + return ""; + }); + + // --- GlobalSemanticIntent --- + py::class_(analyzer_mod, "GlobalSemanticIntent") + .def(py::init<>()) + .def_readwrite("current_phase", &GlobalSemanticIntent::current_phase) + .def_readwrite("primary_bottleneck", &GlobalSemanticIntent::primary_bottleneck) + .def_readwrite("goal", &GlobalSemanticIntent::goal) + .def_readwrite("compute_intensity", &GlobalSemanticIntent::compute_intensity) + .def_readwrite("confidence", &GlobalSemanticIntent::confidence) + .def_readwrite("strategy", &GlobalSemanticIntent::strategy) + .def_readwrite("timestamp_ns", &GlobalSemanticIntent::timestamp_ns) + .def("__repr__", [](const GlobalSemanticIntent &i) { + return ""; + }); + + // --- DeviceLocalIntent --- + py::class_(analyzer_mod, "DeviceLocalIntent") + .def(py::init<>()) + .def_readwrite("device_id", &DeviceLocalIntent::device_id) + .def_readwrite("memory_usage_ratio", &DeviceLocalIntent::memory_usage_ratio) + .def_readwrite("memory_available_bytes", &DeviceLocalIntent::memory_available_bytes) + .def_readwrite("local_bottleneck", &DeviceLocalIntent::local_bottleneck) + .def_readwrite("compute_utilization", &DeviceLocalIntent::compute_utilization) + .def_readwrite("memory_bandwidth_utilization", &DeviceLocalIntent::memory_bandwidth_utilization) + .def_readwrite("communication_time_ratio", &DeviceLocalIntent::communication_time_ratio) + .def_readwrite("resource_confidence", &DeviceLocalIntent::resource_confidence) + .def("__repr__", [](const DeviceLocalIntent &d) { + return ""; + }); + + // --- OptimizationIntent --- + py::class_(analyzer_mod, "OptimizationIntent") + .def(py::init<>()) + .def_readwrite("global_intent", &OptimizationIntent::global) + .def_readwrite("per_device", &OptimizationIntent::per_device) + .def("get_device_intent", &OptimizationIntent::getDeviceIntent, + py::return_value_policy::reference, + py::arg("device_id")) + .def("__repr__", [](const OptimizationIntent &i) { + return ""; + }); + + // --- Top-level functions --- + analyzer_mod.def("analyze", &analyzeCurrentState, + "Analyze current state and return an OptimizationIntent"); + analyzer_mod.def("get_current_phase", &getCurrentPhase, + "Get the current detected task phase"); + analyzer_mod.def("set_enabled", &setAnalyzerEnabled, + "Enable/disable the mutual awareness analyzer", + py::arg("enabled")); + analyzer_mod.def( + "trace_op_for_test", + [](OpType op_type, + const std::vector &shape, + uint8_t dtype, + uint8_t device_type, + int device_id) { + traceOp(op_type, shape.data(), shape.size(), dtype, device_type, static_cast(device_id)); + }, + "Inject an OpTrace entry for testing", + py::arg("op_type"), + py::arg("shape"), + py::arg("dtype") = 0, + py::arg("device_type") = 0, + py::arg("device_id") = 0); + analyzer_mod.def( + "clear_trace", []() { + getGlobalOpTrace().clear(); + MutualAwarenessAnalyzer::instance().clearGraphCache(); + }, + "Clear the global OpTrace ring and analyzer graph cache"); + + // --- Access to analyzer instance for advanced usage --- + analyzer_mod.def("get_analyzer", &MutualAwarenessAnalyzer::instance, + py::return_value_policy::reference, + "Get the MutualAwarenessAnalyzer singleton instance"); + + py::class_(analyzer_mod, "MutualAwarenessAnalyzer") + .def("analyze", py::overload_cast<>(&MutualAwarenessAnalyzer::analyze)) + .def("get_current_phase", &MutualAwarenessAnalyzer::getCurrentPhase) + .def("last_intent", &MutualAwarenessAnalyzer::lastIntent, + py::return_value_policy::reference) + .def("set_enabled", &MutualAwarenessAnalyzer::setEnabled) + .def("is_enabled", &MutualAwarenessAnalyzer::isEnabled) + .def("on_graph_recording_stop", &MutualAwarenessAnalyzer::onGraphRecordingStop) + .def("clear_graph_cache", &MutualAwarenessAnalyzer::clearGraphCache); +} + +} // namespace infinicore::analyzer::pybind diff --git a/csrc/infinicore/src/pybind11/context.hpp b/csrc/infinicore/src/pybind11/context.hpp new file mode 100644 index 000000000..a59b8fde3 --- /dev/null +++ b/csrc/infinicore/src/pybind11/context.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include + +#include + +#include "infinicore.hpp" + +namespace py = pybind11; + +namespace infinicore::context { + +inline void bind(py::module &m) { + // Device management + m.def("get_device", &getDevice, "Get the current active device"); + m.def("get_device_count", &getDeviceCount, + "Get the number of available devices of a specific type", + py::arg("device_type")); + m.def("set_device", &setDevice, + "Set the current active device", + py::arg("device")); + + // Stream and handle management + m.def( + "get_stream", []() { return reinterpret_cast(getStream()); }, + "Get the current stream as an integer handle"); + + // Synchronization + m.def("sync_stream", &syncStream, "Synchronize the current stream"); + m.def("sync_device", &syncDevice, "Synchronize the current device"); + + // Graph + m.def("is_graph_recording", &isGraphRecording, "Check if graph recording is turned on"); + m.def("start_graph_recording", &startGraphRecording, "Start graph recording"); + m.def("stop_graph_recording", &stopGraphRecording, "Stop graph recording and return the graph"); + m.def("cancel_graph_recording", &cancelGraphRecording, "Cancel the active graph recording"); +} + +} // namespace infinicore::context diff --git a/csrc/infinicore/src/pybind11/device.hpp b/csrc/infinicore/src/pybind11/device.hpp new file mode 100644 index 000000000..ec1d7fd6f --- /dev/null +++ b/csrc/infinicore/src/pybind11/device.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include + +#include "infinicore.hpp" + +namespace py = pybind11; + +namespace infinicore::device { + +inline void bind(py::module &m) { + py::class_ device(m, "Device"); + + py::enum_(device, "Type") + .value("CPU", Device::Type::kCpu) + .value("NVIDIA", Device::Type::kNvidia) + .value("CAMBRICON", Device::Type::kCambricon) + .value("ASCEND", Device::Type::kAscend) + .value("METAX", Device::Type::kMetax) + .value("MOORE", Device::Type::kMoore) + .value("ILUVATAR", Device::Type::kIluvatar) + .value("HYGON", Device::Type::kHygon); + + device + .def(py::init(), + py::arg("type") = Device::Type::kCpu, py::arg("index") = 0) + .def_property_readonly("type", &Device::type) + .def_property_readonly("index", &Device::index) + .def("__str__", &Device::ToString); +} + +} // namespace infinicore::device diff --git a/csrc/infinicore/src/pybind11/device_event.hpp b/csrc/infinicore/src/pybind11/device_event.hpp new file mode 100644 index 000000000..1c2845c6b --- /dev/null +++ b/csrc/infinicore/src/pybind11/device_event.hpp @@ -0,0 +1,43 @@ +#pragma once + +#include "infinicore.hpp" +#include + +#include + +namespace py = pybind11; + +namespace infinicore::device_event { + +inline void bind(py::module &m) { + py::class_(m, "DeviceEvent") + .def(py::init<>(), "Construct a DeviceEvent on the current device") + .def(py::init(), "Construct a DeviceEvent with specific flags", py::arg("flags")) + .def(py::init(), "Construct a DeviceEvent on a specific device", py::arg("device")) + .def(py::init(), "Construct a DeviceEvent on a specific device with flags", + py::arg("device"), py::arg("flags")) + + .def("record", py::overload_cast<>(&DeviceEvent::record), + "Record the event on the current stream of its device") + .def( + "record", [](DeviceEvent &event, std::uintptr_t stream) { + event.record(reinterpret_cast(stream)); + }, + "Record the event on a specific stream", py::arg("stream")) + + .def("synchronize", &DeviceEvent::synchronize, "Wait for the event to complete (blocking)") + .def("query", &DeviceEvent::query, "Check if the event has been completed") + + .def("elapsed_time", &DeviceEvent::elapsed_time, "Calculate elapsed time between this event and another event (in milliseconds)", py::arg("other")) + + .def( + "wait", [](const DeviceEvent &event, std::uintptr_t stream) { + event.wait(reinterpret_cast(stream)); + }, + "Make a stream wait for this event to complete", py::arg("stream") = 0) + + .def_property_readonly("device", &DeviceEvent::device, "Get the device where this event was created") + .def_property_readonly("is_recorded", &DeviceEvent::is_recorded, "Check if the event has been recorded"); +} + +} // namespace infinicore::device_event diff --git a/csrc/infinicore/src/pybind11/dtype.hpp b/csrc/infinicore/src/pybind11/dtype.hpp new file mode 100644 index 000000000..c743f1667 --- /dev/null +++ b/csrc/infinicore/src/pybind11/dtype.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include + +#include "infinicore.hpp" + +namespace py = pybind11; + +namespace infinicore::dtype { + +inline void bind(py::module &m) { + py::enum_(m, "DataType") + .value("INT8", DataType::kInt8) + .value("INT16", DataType::kInt16) + .value("INT32", DataType::kInt32) + .value("INT64", DataType::kInt64) + .value("UINT8", DataType::kUInt8) + .value("UINT16", DataType::kUInt16) + .value("UINT32", DataType::kUInt32) + .value("UINT64", DataType::kUInt64) + .value("FLOAT16", DataType::kFloat16) + .value("BFLOAT16", DataType::kBFloat16) + .value("FLOAT32", DataType::kFloat32) + .value("FLOAT64", DataType::kFloat64); +} + +} // namespace infinicore::dtype diff --git a/csrc/infinicore/src/pybind11/from_list.cc b/csrc/infinicore/src/pybind11/from_list.cc new file mode 100644 index 000000000..d54458501 --- /dev/null +++ b/csrc/infinicore/src/pybind11/from_list.cc @@ -0,0 +1,349 @@ +#include "infinicore.hpp" +#include +#include +#include +#include +#include +#include + +namespace py = pybind11; + +namespace infinicore::tensor { + +namespace { + +constexpr const char *kElementTypeError = "List elements must be bool, int, or float"; + +// Zero-copy view over list / tuple / generic sequence; items() exposes the +// underlying PyObject** for hot loop indexing. + +class SequenceView { +public: + enum class Kind : uint8_t { List, + Tuple, + Fast }; + + explicit SequenceView(py::handle obj, const char *err = "Could not convert object to sequence") + : SequenceView(obj.ptr(), err) {} + + explicit SequenceView(PyObject *obj, const char *err = "Could not convert object to sequence") { + if (PyList_Check(obj)) { + kind_ = Kind::List; + borrowed_ = obj; + fast_ = nullptr; + return; + } + if (PyTuple_Check(obj)) { + kind_ = Kind::Tuple; + borrowed_ = obj; + fast_ = nullptr; + return; + } + kind_ = Kind::Fast; + borrowed_ = nullptr; + fast_ = PySequence_Fast(obj, err); + if (!fast_) { + throw py::error_already_set(); + } + } + + ~SequenceView() { + if (fast_ != nullptr) { + Py_DECREF(fast_); + } + } + + SequenceView(const SequenceView &) = delete; + SequenceView &operator=(const SequenceView &) = delete; + + bool is_list() const { + return kind_ == Kind::List; + } + + Py_ssize_t size() const { + return size_of(ptr()); + } + + PyObject *item(Py_ssize_t index) const { + return item_at(ptr(), index); + } + + // Direct pointer to ob_item; valid for list / tuple / PySequence_Fast result. + PyObject **items() const { + return PySequence_Fast_ITEMS(ptr()); + } + + static Py_ssize_t size_of(PyObject *obj) { + if (PyList_Check(obj)) { + return PyList_GET_SIZE(obj); + } + if (PyTuple_Check(obj)) { + return PyTuple_GET_SIZE(obj); + } + return PySequence_Size(obj); + } + + static PyObject *item_at(PyObject *obj, Py_ssize_t index) { + if (PyList_Check(obj)) { + return PyList_GET_ITEM(obj, index); + } + if (PyTuple_Check(obj)) { + return PyTuple_GET_ITEM(obj, index); + } + return PySequence_GetItem(obj, index); + } + +private: + PyObject *ptr() const { + return borrowed_ != nullptr ? borrowed_ : fast_; + } + + Kind kind_ = Kind::Fast; + PyObject *borrowed_ = nullptr; + PyObject *fast_ = nullptr; +}; + +// Shape discovery + uniform per-scalar traversal. + +inline bool is_python_scalar(PyObject *obj) { + return PyBool_Check(obj) || PyFloat_Check(obj) || PyLong_Check(obj); +} + +inline bool is_python_sequence(PyObject *obj) { + return PyList_Check(obj) || PyTuple_Check(obj); +} + +struct ListLayout { + SequenceView seq; + int ndim = 0; + Shape shape; + bool list_of_lists = false; + + explicit ListLayout(py::handle obj) + : seq(obj) { + if (!is_python_sequence(obj.ptr())) { + throw py::type_error("Input data must be a list or tuple"); + } + if (seq.size() == 0) { + throw py::value_error("Input data cannot be empty"); + } + + PyObject *first = seq.item(0); + if (is_python_scalar(first)) { + ndim = 1; + shape = Shape{static_cast(seq.size())}; + return; + } + + if (!is_python_sequence(first)) { + throw py::type_error("List elements must be scalars or nested lists"); + } + if (SequenceView::size_of(first) == 0) { + throw py::value_error("Nested list cannot be empty"); + } + if (!is_python_scalar(SequenceView::item_at(first, 0))) { + throw py::value_error("Only 1D and 2D lists are supported"); + } + + ndim = 2; + const Size rows = static_cast(seq.size()); + const Size cols = static_cast(SequenceView::size_of(first)); + list_of_lists = seq.is_list() && PyList_Check(first); + + for (Py_ssize_t i = 1; i < seq.size(); ++i) { + PyObject *row_obj = seq.item(i); + if (!is_python_sequence(row_obj)) { + throw py::value_error("Input must be a regular 2D list with equal row lengths"); + } + if (static_cast(SequenceView::size_of(row_obj)) != cols) { + throw py::value_error("Input must be a regular 2D list with equal row lengths"); + } + if (list_of_lists && !PyList_Check(row_obj)) { + list_of_lists = false; + } + } + + shape = Shape{rows, cols}; + } +}; + +template +void for_each_scalar(const ListLayout &layout, Fn &&fn) { + const Py_ssize_t rows = layout.seq.size(); + PyObject **rows_items = layout.seq.items(); + + if (layout.ndim == 1) { + for (Py_ssize_t i = 0; i < rows; ++i) { + fn(rows_items[i]); + } + return; + } + + const Py_ssize_t cols = static_cast(layout.shape[1]); + if (layout.list_of_lists) { + for (Py_ssize_t i = 0; i < rows; ++i) { + PyObject **row_items = PySequence_Fast_ITEMS(rows_items[i]); + for (Py_ssize_t j = 0; j < cols; ++j) { + fn(row_items[j]); + } + } + return; + } + + for (Py_ssize_t i = 0; i < rows; ++i) { + PyObject *row = rows_items[i]; + for (Py_ssize_t j = 0; j < cols; ++j) { + fn(SequenceView::item_at(row, j)); + } + } +} + +// read_pylong_fast inlines compact int reads on Python >= 3.12 via +// PyUnstable_Long_*; falls back to PyLong_AsLongLong otherwise. +inline int64_t read_pylong_fast(PyObject *obj) { +#if PY_VERSION_HEX >= 0x030C0000 + PyLongObject *lo = reinterpret_cast(obj); + if (PyUnstable_Long_IsCompact(lo)) { + return static_cast(PyUnstable_Long_CompactValue(lo)); + } +#endif + const int64_t value = PyLong_AsLongLong(obj); + if (value == -1 && PyErr_Occurred()) { + throw py::error_already_set(); + } + return value; +} + +inline uint64_t read_pyuint64(PyObject *obj) { + const auto value = PyLong_AsUnsignedLongLong(obj); + if (value == std::numeric_limits::max() && PyErr_Occurred()) { + throw py::error_already_set(); + } + return static_cast(value); +} + +// Convert int64 -> Dst with range check for narrow ints (i8/i16/i32, +// u8/u16/u32) and uint64 (forbids negatives); bool/i64/float/double truncate. +template +inline Dst narrow_int_to(int64_t v) { + if constexpr (std::is_floating_point_v) { + return static_cast(v); + } else if constexpr (std::is_same_v) { + return v != 0; + } else if constexpr (std::is_same_v) { + return v; + } else if constexpr (std::is_same_v) { + if (v < 0) { + throw std::overflow_error("Integer value out of range for target dtype"); + } + return static_cast(v); + } else if constexpr (std::is_unsigned_v) { + if (v < 0 || static_cast(v) > std::numeric_limits::max()) { + throw std::overflow_error("Integer value out of range for target dtype"); + } + return static_cast(v); + } else { + if (v < std::numeric_limits::min() || v > std::numeric_limits::max()) { + throw std::overflow_error("Integer value out of range for target dtype"); + } + return static_cast(v); + } +} + +template +inline Dst read_pyint(PyObject *obj) { + if constexpr (std::is_same_v) { + return read_pyuint64(obj); + } else { + return narrow_int_to(read_pylong_fast(obj)); + } +} + +// Write path: branch order is PyLong/PyFloat exact first (most inputs), then +// bool (int subclass), then long/float subclass fallback. F16/BF16 share +// write_dtype_half via float pivot; other dtypes share write_dtype_native. +template +void write_dtype_native(const ListLayout &layout, std::byte *dest) { + Dst *out = reinterpret_cast(dest); + for_each_scalar(layout, [&](PyObject *obj) { + if (PyLong_CheckExact(obj)) { + *out++ = read_pyint(obj); + } else if (PyFloat_CheckExact(obj)) { + *out++ = static_cast(PyFloat_AS_DOUBLE(obj)); + } else if (PyBool_Check(obj)) { + *out++ = static_cast(obj == Py_True); + } else if (PyLong_Check(obj)) { + *out++ = read_pyint(obj); + } else if (PyFloat_Check(obj)) { + *out++ = static_cast(PyFloat_AS_DOUBLE(obj)); + } else { + throw py::type_error(kElementTypeError); + } + }); +} + +template +void write_dtype_half(const ListLayout &layout, std::byte *dest) { + Dst *out = reinterpret_cast(dest); + for_each_scalar(layout, [&](PyObject *obj) { + float f; + if (PyLong_CheckExact(obj)) { + f = static_cast(read_pylong_fast(obj)); + } else if (PyFloat_CheckExact(obj)) { + f = static_cast(PyFloat_AS_DOUBLE(obj)); + } else if (PyBool_Check(obj)) { + f = (obj == Py_True) ? 1.0f : 0.0f; + } else if (PyLong_Check(obj)) { + f = static_cast(read_pylong_fast(obj)); + } else if (PyFloat_Check(obj)) { + f = static_cast(PyFloat_AS_DOUBLE(obj)); + } else { + throw py::type_error(kElementTypeError); + } + *out++ = Dst::FromFloat(f); + }); +} + +void write_with_dtype(const ListLayout &layout, DataType dtype, std::byte *dest) { + switch (dtype) { + case DataType::kInt8: + return write_dtype_native(layout, dest); + case DataType::kInt16: + return write_dtype_native(layout, dest); + case DataType::kInt32: + return write_dtype_native(layout, dest); + case DataType::kInt64: + return write_dtype_native(layout, dest); + case DataType::kUInt8: + return write_dtype_native(layout, dest); + case DataType::kUInt16: + return write_dtype_native(layout, dest); + case DataType::kUInt32: + return write_dtype_native(layout, dest); + case DataType::kUInt64: + return write_dtype_native(layout, dest); + case DataType::kFloat16: + return write_dtype_half(layout, dest); + case DataType::kBFloat16: + return write_dtype_half(layout, dest); + case DataType::kFloat32: + return write_dtype_native(layout, dest); + case DataType::kFloat64: + return write_dtype_native(layout, dest); + default: + throw py::type_error( + std::string("Unsupported dtype for from_list: ") + toString(dtype)); + } +} + +} // namespace + +// Entry (exported for pybind registration in tensor.hpp). +Tensor from_list_py(py::handle data, DataType dtype) { + const ListLayout layout(data); + auto tensor = Tensor::empty(layout.shape, dtype, Device(Device::Type::kCpu, 0)); + write_with_dtype(layout, dtype, tensor->data()); + return tensor; +} + +} // namespace infinicore::tensor diff --git a/csrc/infinicore/src/pybind11/graph.hpp b/csrc/infinicore/src/pybind11/graph.hpp new file mode 100644 index 000000000..d45c9b32c --- /dev/null +++ b/csrc/infinicore/src/pybind11/graph.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include +#include + +#include "infinicore.hpp" + +namespace py = pybind11; + +namespace infinicore::graph { +inline void bind(py::module_ &m) { + py::class_>(m, "Graph") + .def(py::init<>()) // allow construction + .def("run", &infinicore::graph::Graph::run); +} +} // namespace infinicore::graph diff --git a/csrc/infinicore/src/pybind11/infinicore.cc b/csrc/infinicore/src/pybind11/infinicore.cc new file mode 100644 index 000000000..9820eab1b --- /dev/null +++ b/csrc/infinicore/src/pybind11/infinicore.cc @@ -0,0 +1,35 @@ +#include +#include + +#include "../utils.hpp" +#include "context.hpp" +#include "device.hpp" +#include "device_event.hpp" +#include "dtype.hpp" +#include "graph.hpp" +#include "io.hpp" +#include "ops.hpp" +#include "tensor.hpp" + +#ifdef ENABLE_MUTUAL_AWARENESS +#include "analyzer.hpp" +#endif + +namespace infinicore { + +PYBIND11_MODULE(_infinicore, m) { + context::bind(m); + device::bind(m); + device_event::bind(m); + dtype::bind(m); + ops::bind(m); + tensor::bind(m); + io::bind(m); + graph::bind(m); + +#ifdef ENABLE_MUTUAL_AWARENESS + analyzer::pybind::bind(m); +#endif +} + +} // namespace infinicore diff --git a/csrc/infinicore/src/pybind11/io.hpp b/csrc/infinicore/src/pybind11/io.hpp new file mode 100644 index 000000000..4f1f6408e --- /dev/null +++ b/csrc/infinicore/src/pybind11/io.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include "infinicore.hpp" +#include +#include + +namespace py = pybind11; + +namespace infinicore::io { + +inline void bind(py::module &m) { + m.def( + "set_printoptions", [](int precision, int threshold, int edge_items, int line_width, py::object sci_mode) { + infinicore::print_options::set_precision(precision); + infinicore::print_options::set_threshold(threshold); + infinicore::print_options::set_edge_items(edge_items); + infinicore::print_options::set_line_width(line_width); + + // Handle sci_mode: None -> -1 (auto), True -> 1 (enable), False -> 0 (disable) + int sci_mode_value = -1; // default: auto + if (!sci_mode.is_none()) { + sci_mode_value = static_cast(py::cast(sci_mode)); // True -> 1, False -> 0 + } + + infinicore::print_options::set_sci_mode(sci_mode_value); }, pybind11::arg("precision"), pybind11::arg("threshold"), pybind11::arg("edge_items"), pybind11::arg("line_width"), pybind11::arg("sci_mode")); +} + +} // namespace infinicore::io diff --git a/csrc/infinicore/src/pybind11/ops.hpp b/csrc/infinicore/src/pybind11/ops.hpp new file mode 100644 index 000000000..4fbba73eb --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include + +#include "ops/add.hpp" +#include "ops/cat.hpp" +#include "ops/causal_softmax.hpp" +#include "ops/embedding.hpp" +#include "ops/linear.hpp" +#include "ops/matmul.hpp" +#include "ops/random_sample.hpp" +#include "ops/rms_norm.hpp" +#include "ops/rope.hpp" +#include "ops/silu.hpp" +#include "ops/silu_and_mul.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind(py::module &m) { + bind_add(m); + bind_cat(m); + bind_causal_softmax(m); + bind_embedding(m); + bind_linear(m); + bind_matmul(m); + bind_random_sample(m); + bind_rms_norm(m); + bind_rope(m); + bind_silu(m); + bind_silu_and_mul(m); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/acos.hpp b/csrc/infinicore/src/pybind11/ops/acos.hpp new file mode 100644 index 000000000..6f3c52dd0 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/acos.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include "infinicore/ops/acos.hpp" // 引用核心算子头文件 +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_acos(py::module &m) { + // 绑定 out-of-place 接口: output = acos(input) + m.def("acos", + &op::acos, + py::arg("input"), + R"doc(Computes the inverse cosine (arccosine) of each element of input. + +Returns a new tensor with the arccosine of the elements of input. +The range of the result is [0, pi].)doc"); + + // 绑定 in-place 接口: acos_(output, input) + m.def("acos_", + &op::acos_, + py::arg("output"), + py::arg("input"), + R"doc(In-place acos operation. Writes result into output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/adaptive_avg_pool1d.hpp b/csrc/infinicore/src/pybind11/ops/adaptive_avg_pool1d.hpp new file mode 100644 index 000000000..8d58232f6 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/adaptive_avg_pool1d.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include + +#include "infinicore/ops/adaptive_avg_pool1d.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_adaptive_avg_pool1d(py::module &m) { + // 绑定函数接口: output = adaptive_avg_pool1d(input, output_size) + m.def("adaptive_avg_pool1d", + &op::adaptive_avg_pool1d, + py::arg("input"), + py::arg("output_size"), + R"doc(Applies a 1D adaptive average pooling over an input signal composed of several input planes. + +Args: + input (Tensor): Input tensor of shape (C, L) or (N, C, L). + output_size (int): The target output size. + +Returns: + Tensor: Output tensor of shape (C, output_size) or (N, C, output_size). +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/adaptive_avg_pool3d.hpp b/csrc/infinicore/src/pybind11/ops/adaptive_avg_pool3d.hpp new file mode 100644 index 000000000..758baf232 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/adaptive_avg_pool3d.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include + +#include "infinicore/ops/adaptive_avg_pool3d.hpp" + +namespace py = pybind11; +namespace infinicore::ops { +inline void bind_adaptive_avg_pool3d(py::module &m) { + m.def("adaptive_avg_pool3d", + &op::adaptive_avg_pool3d, + py::arg("x"), + py::arg("output_size"), + R"doc( Adaptive Average Pooling 3D.)doc"); + + m.def("adaptive_avg_pool3d_", + &op::adaptive_avg_pool3d_, + py::arg("y"), + py::arg("x"), + R"doc(In-place, Adaptive Average Pooling 3D.)doc"); +} +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/adaptive_max_pool1d.hpp b/csrc/infinicore/src/pybind11/ops/adaptive_max_pool1d.hpp new file mode 100644 index 000000000..99dc3bbc5 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/adaptive_max_pool1d.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include + +#include "infinicore/ops/adaptive_max_pool1d.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_adaptive_max_pool1d(py::module &m) { + m.def("adaptive_max_pool1d", + &op::adaptive_max_pool1d, + py::arg("x"), + py::arg("output_size"), + R"doc(1D Adaptive Max Pooling. + +Args: + x: Input tensor of shape (N, C, L_in) or (N, L_in) + output_size: Target output size L_out +Returns: + Output tensor of shape (N, C, L_out) or (N, L_out) +)doc"); + + m.def("adaptive_max_pool1d_", + &op::adaptive_max_pool1d_, + py::arg("y"), + py::arg("x"), + py::arg("output_size"), + R"doc(In-place 1D Adaptive Max Pooling. + +Args: + y: Output tensor of shape (N, C, L_out) or (N, L_out) + x: Input tensor of shape (N, C, L_in) or (N, L_in) + output_size: Target output size L_out +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/add.hpp b/csrc/infinicore/src/pybind11/ops/add.hpp new file mode 100644 index 000000000..fe7ac852f --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/add.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/add.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_add(py::module &m) { + m.def("add", + &op::add, + py::arg("a"), + py::arg("b"), + R"doc(Addition of two tensors.)doc"); + + m.def("add_", + &op::add_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + R"doc(In-place tensor addition.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/add_rms_norm.hpp b/csrc/infinicore/src/pybind11/ops/add_rms_norm.hpp new file mode 100644 index 000000000..5f9b243e5 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/add_rms_norm.hpp @@ -0,0 +1,51 @@ +#pragma once + +#include + +#include "infinicore/ops/add_rms_norm.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_add_rms_norm(py::module &m) { + m.def("add_rms_norm", + &op::add_rms_norm, + py::arg("a"), + py::arg("b"), + py::arg("weight"), + py::arg("epsilon") = 1e-5f, + R"doc(Fused Add and RMS Normalization. + +Args: + a: First input tensor + b: Second input tensor + weight: Scale weights + epsilon: Small constant for numerical stability, default is 1e-5 + +Returns: + Tuple of (normalized_result, add_result): (RMSNorm(a + b) * weight, a + b) + The add_result can be used as residual for subsequent layers. +)doc"); + + m.def("add_rms_norm_", + &op::add_rms_norm_, + py::arg("y"), + py::arg("residual_out"), + py::arg("a"), + py::arg("b"), + py::arg("weight"), + py::arg("epsilon") = 1e-5f, + R"doc(In-place Fused Add and RMS Normalization. + +Args: + y: Output tensor for normalized result + residual_out: Output tensor for add result (a + b) before normalization + a: First input tensor + b: Second input tensor + weight: Scale weights + epsilon: Small constant for numerical stability, default is 1e-5 +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/addbmm.hpp b/csrc/infinicore/src/pybind11/ops/addbmm.hpp new file mode 100644 index 000000000..3756d9e59 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/addbmm.hpp @@ -0,0 +1,52 @@ +#pragma once + +#include "infinicore/ops/addbmm.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_addbmm(py::module &m) { + // ----------------------------------------------------------- + // 1. Out-of-place 接口: output = addbmm(...) + // ----------------------------------------------------------- + m.def("addbmm", + &op::addbmm, + py::arg("input"), + py::arg("batch1"), + py::arg("batch2"), + py::arg("beta") = 1.0f, + py::arg("alpha") = 1.0f, + R"doc(Performs a batch matrix-matrix product of matrices stored in batch1 and batch2, +with a reduced add step (summing over all matrices in the batch). + +.. math:: + \text{out} = \beta \times \text{input} + \alpha \times \sum_{i=0}^{b-1} (\text{batch1}_i \mathbin{@} \text{batch2}_i) + +Args: + input (Tensor): Matrix to be added. Shape (n, p). + batch1 (Tensor): The first batch of matrices to be multiplied. Shape (b, n, m). + batch2 (Tensor): The second batch of matrices to be multiplied. Shape (b, m, p). + beta (float, optional): Multiplier for input. Default: 1.0. + alpha (float, optional): Multiplier for batch1 @ batch2. Default: 1.0. + +Returns: + Tensor: Output tensor of shape (n, p). +)doc"); + + // ----------------------------------------------------------- + // 2. [新增] In-place 接口: addbmm_(out, ...) + // ----------------------------------------------------------- + m.def("addbmm_", + &op::addbmm_, // 绑定到 C++ 的 void addbmm_(...) + py::arg("out"), // 第一个参数通常是输出 Tensor + py::arg("input"), + py::arg("batch1"), + py::arg("batch2"), + py::arg("beta") = 1.0f, + py::arg("alpha") = 1.0f, + "In-place version of addbmm"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/addcmul.hpp b/csrc/infinicore/src/pybind11/ops/addcmul.hpp new file mode 100644 index 000000000..74c019add --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/addcmul.hpp @@ -0,0 +1,51 @@ +#pragma once + +#include + +#include "infinicore/ops/addcmul.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_addcmul(py::module &m) { + // 绑定 out-of-place 接口: out = addcmul(input, t1, t2, value) + m.def("addcmul", + &op::addcmul, + py::arg("input"), + py::arg("tensor1"), + py::arg("tensor2"), + py::arg("value") = 1.0f, + R"doc(Performs the element-wise multiplication of tensor1 by tensor2, +multiplies the result by value and adds it to input. + +Args: + input: Tensor to be added + tensor1: First tensor for multiplication + tensor2: Second tensor for multiplication + value: Scalar multiplier for tensor1 * tensor2 (default: 1.0) + +Returns: + The output tensor +)doc"); + + // 绑定 in-place / specified output 接口: addcmul_(out, input, t1, t2, value) + m.def("addcmul_", + &op::addcmul_, + py::arg("out"), + py::arg("input"), + py::arg("tensor1"), + py::arg("tensor2"), + py::arg("value") = 1.0f, + R"doc(In-place version of addcmul. + +Args: + out: The destination tensor to store the result + input: Tensor to be added + tensor1: First tensor for multiplication + tensor2: Second tensor for multiplication + value: Scalar multiplier for tensor1 * tensor2 (default: 1.0) +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/addr.hpp b/csrc/infinicore/src/pybind11/ops/addr.hpp new file mode 100644 index 000000000..912b2b5da --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/addr.hpp @@ -0,0 +1,29 @@ +#pragma once + +#include "infinicore/ops/addr.hpp" +#include + +namespace py = pybind11; +namespace infinicore::ops { +inline void bind_addr(py::module_ &m) { + m.def( + "addr", + &op::addr, + py::arg("input"), + py::arg("vec1"), + py::arg("vec2"), + py::arg("alpha"), + py::arg("beta"), + R"doc(Addr.)doc"); + m.def( + "addr_", + &op::addr_, + py::arg("out"), + py::arg("input"), + py::arg("vec1"), + py::arg("vec2"), + py::arg("beta"), + py::arg("alpha"), + R"doc(Addr.)doc"); +} +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/affine_grid.hpp b/csrc/infinicore/src/pybind11/ops/affine_grid.hpp new file mode 100644 index 000000000..c12279bcc --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/affine_grid.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include "infinicore/ops/affine_grid.hpp" +#include +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_affine_grid(py::module &m) { + // 绑定函数接口: grid = affine_grid(theta, size, align_corners) + m.def("affine_grid", + &op::affine_grid, + py::arg("theta"), + py::arg("size"), + py::arg("align_corners") = false, // 设置默认值 + R"doc(Generates a 2D or 3D flow field (sampling grid), given a batch of affine matrices theta. + +Args: + theta (Tensor): Input affine matrices of shape (N, 2, 3) for 2D or (N, 3, 4) for 3D. + size (List[int]): The target output image size. Usually (N, C, H, W) for 2D or (N, C, D, H, W) for 3D. + align_corners (bool, optional): Geometrically, we consider the pixels of the input as squares rather than points. If set to True, the extrema (-1 and 1) are considered as referring to the center points of the input's corner pixels. If set to False, they are instead considered as referring to the corner points of the input's corner pixels, making the sampling more resolution agnostic. Defaults to False. + +Returns: + Tensor: Output tensor of shape (N, H, W, 2). +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/all.hpp b/csrc/infinicore/src/pybind11/ops/all.hpp new file mode 100644 index 000000000..4ccac685b --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/all.hpp @@ -0,0 +1,60 @@ +#pragma once + +#include + +#include "infinicore/ops/all.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_all(Tensor input, py::object dim, bool keepdim) { + if (dim.is_none()) { + std::vector dim_vec; + for (int i = 0; i < input->shape().size(); i++) { + dim_vec.push_back(i); + } + return op::all(input, dim_vec, keepdim); + } else if (py::isinstance(dim) || py::isinstance(dim)) { + return op::all(input, dim.cast>(), keepdim); + } else if (py::isinstance(dim)) { + return op::all(input, std::vector(1, dim.cast()), keepdim); + } else { + throw std::invalid_argument("dim must be a tuple or an integer"); + } +} + +void py_all_(Tensor output, Tensor input, py::object dim, bool keepdim) { + if (dim.is_none()) { + std::vector dim_vec; + for (int i = 0; i < input->shape().size(); i++) { + dim_vec.push_back(i); + } + op::all_(output, input, dim_vec, keepdim); + } else if (py::isinstance(dim) || py::isinstance(dim)) { + op::all_(output, input, dim.cast>(), keepdim); + } else if (py::isinstance(dim)) { + op::all_(output, input, std::vector(1, dim.cast()), keepdim); + } else { + throw std::invalid_argument("dim must be a tuple or an integer"); + } +} + +inline void bind_all(py::module &m) { + m.def("all", + &py_all, + py::arg("input"), + py::arg("dim"), + py::arg("keepdim"), + R"doc(All of input tensor along the given dimensions.)doc"); + + m.def("all_", + &py_all_, + py::arg("output"), + py::arg("input"), + py::arg("dim"), + py::arg("keepdim"), + R"doc(In-place tensor all.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/argwhere.hpp b/csrc/infinicore/src/pybind11/ops/argwhere.hpp new file mode 100644 index 000000000..cb91142c8 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/argwhere.hpp @@ -0,0 +1,13 @@ +#pragma once + +#include "infinicore/ops/argwhere.hpp" +#include +namespace py = pybind11; +namespace infinicore::ops { +inline void bind_argwhere(py::module &m) { + m.def("argwhere", + &op::argwhere, + py::arg("x"), + R"doc(Argwhere.)doc"); +} +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/asin.hpp b/csrc/infinicore/src/pybind11/ops/asin.hpp new file mode 100644 index 000000000..ebd14adbc --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/asin.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/asin.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_asin(py::module &m) { + m.def("asin", + &op::asin, + py::arg("input"), + R"doc(Arcsin activation function.)doc"); + + m.def("asin_", + &op::asin_, + py::arg("output"), + py::arg("input"), + R"doc(In-place arcsin activation function.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/asinh.hpp b/csrc/infinicore/src/pybind11/ops/asinh.hpp new file mode 100644 index 000000000..bf1fcca23 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/asinh.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/asinh.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_asinh(py::module &m) { + m.def("asinh", + &op::asinh, + py::arg("x"), + R"doc(Element-wise inverse hyperbolic sine function.)doc"); + + m.def("asinh_", + &op::asinh_, + py::arg("y"), + py::arg("x"), + R"doc(In-place element-wise inverse hyperbolic sine function.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/asum.hpp b/csrc/infinicore/src/pybind11/ops/asum.hpp new file mode 100644 index 000000000..b094d12f5 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/asum.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/asum.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_asum(py::module &m) { + m.def("asum", + &op::asum, + py::arg("x"), + R"doc(BLAS level-1 asum.)doc"); + + m.def("asum_", + &op::asum_, + py::arg("x"), + py::arg("result"), + R"doc(In-place BLAS level-1 asum.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/atanh.hpp b/csrc/infinicore/src/pybind11/ops/atanh.hpp new file mode 100644 index 000000000..44fd74169 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/atanh.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/atanh.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_atanh(py::module &m) { + m.def("atanh", + &op::atanh, + py::arg("a"), + R"doc(Inverse hyperbolic tangent of a tensor.)doc"); + + m.def("atanh_", + &op::atanh_, + py::arg("y"), + py::arg("a"), + R"doc(Compute inverse hyperbolic tangent and store in the provided output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/attention.hpp b/csrc/infinicore/src/pybind11/ops/attention.hpp new file mode 100644 index 000000000..dc13bcac6 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/attention.hpp @@ -0,0 +1,56 @@ +#pragma once + +#include + +#include "infinicore/ops/attention.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_attention(py::module &m) { + m.def("attention", + &op::attention, + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("pos"), + R"doc(Attention mechanism with KV caching. + +Args: + q: Query tensor + k: Key tensor + v: Value tensor + k_cache: Key cache tensor + v_cache: Value cache tensor + pos: Current position in the sequence + +Returns: + Output tensor from attention computation +)doc"); + + m.def("attention_", + &op::Attention::execute, + py::arg("out"), + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("pos"), + R"doc(In-place attention mechanism with KV caching. + +Args: + out: Output tensor + q: Query tensor + k: Key tensor + v: Value tensor + k_cache: Key cache tensor + v_cache: Value cache tensor + pos: Current position in the sequence +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/avg_pool1d.hpp b/csrc/infinicore/src/pybind11/ops/avg_pool1d.hpp new file mode 100644 index 000000000..32394552a --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/avg_pool1d.hpp @@ -0,0 +1,37 @@ +#pragma once + +#include +#include + +#include "infinicore/ops/avg_pool1d.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_avg_pool1d(py::module &m) { + m.def( + "avg_pool1d", + [](::infinicore::Tensor input, size_t kernel_size, std::optional stride, size_t padding) { + return op::avg_pool1d(input, kernel_size, stride.value_or(0), padding); + }, + py::arg("input"), + py::arg("kernel_size"), + py::arg("stride") = py::none(), + py::arg("padding") = 0, + R"doc(AvgPool1d out-of-place.)doc"); + + m.def( + "avg_pool1d_", + [](::infinicore::Tensor output, ::infinicore::Tensor input, size_t kernel_size, std::optional stride, size_t padding) { + op::avg_pool1d_(output, input, kernel_size, stride.value_or(0), padding); + }, + py::arg("output"), + py::arg("input"), + py::arg("kernel_size"), + py::arg("stride") = py::none(), + py::arg("padding") = 0, + R"doc(AvgPool1d in-place variant writing to provided output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/axpy.hpp b/csrc/infinicore/src/pybind11/ops/axpy.hpp new file mode 100644 index 000000000..fd0aff633 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/axpy.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include + +#include "infinicore/ops/axpy.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_axpy(py::module &m) { + m.def("axpy_", + &op::axpy_, + py::arg("alpha"), + py::arg("x"), + py::arg("y"), + R"doc(In-place BLAS level-1 axpy, updating y.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/baddbmm.hpp b/csrc/infinicore/src/pybind11/ops/baddbmm.hpp new file mode 100644 index 000000000..bb5709afc --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/baddbmm.hpp @@ -0,0 +1,56 @@ +#pragma once + +#include + +#include "infinicore/ops/baddbmm.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_baddbmm(Tensor input, Tensor batch1, Tensor batch2, float beta = 1.0f, float alpha = 1.0f) { + return op::baddbmm(input, batch1, batch2, beta, alpha); +} + +void py_baddbmm_(Tensor out, Tensor input, Tensor batch1, Tensor batch2, float beta = 1.0f, float alpha = 1.0f) { + op::baddbmm_(out, input, batch1, batch2, beta, alpha); +} + +inline void bind_baddbmm(py::module &m) { + m.def("baddbmm", + &py_baddbmm, + py::arg("input"), + py::arg("batch1"), + py::arg("batch2"), + py::arg("beta") = 1.0f, + py::arg("alpha") = 1.0f, + R"doc(Batched matrix-matrix product with addition. +Args: + input: Input tensor + batch1: First batch of matrices + batch2: Second batch of matrices + beta: Scaling factor for input tensor + alpha: Scaling factor for the product of batch1 and batch2 +Returns: + Output tensor after baddbmm operation +)doc"); + m.def("baddbmm_", + &py_baddbmm_, + py::arg("out"), + py::arg("input"), + py::arg("batch1"), + py::arg("batch2"), + py::arg("beta") = 1.0f, + py::arg("alpha") = 1.0f, + R"doc(In-place batched matrix-matrix product with addition. +Args: + out: Output tensor + input: Input tensor + batch1: First batch of matrices + batch2: Second batch of matrices + beta: Scaling factor for input tensor + alpha: Scaling factor for the product of batch1 and batch2 +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/bilinear.hpp b/csrc/infinicore/src/pybind11/ops/bilinear.hpp new file mode 100644 index 000000000..2288f3602 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/bilinear.hpp @@ -0,0 +1,61 @@ +#pragma once + +#include + +#include "infinicore/ops/bilinear.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_bilinear(Tensor x1, Tensor x2, Tensor weight, pybind11::object bias) { + std::optional bias_tensor = std::nullopt; + if (!bias.is_none()) { + bias_tensor = bias.cast(); + } + return op::bilinear(x1, x2, weight, bias_tensor); +} + +void py_bilinear_(Tensor out, Tensor x1, Tensor x2, Tensor weight, pybind11::object bias) { + std::optional bias_tensor = std::nullopt; + if (!bias.is_none()) { + bias_tensor = bias.cast(); + } + op::bilinear_(out, x1, x2, weight, bias_tensor); +} + +inline void bind_bilinear(py::module &m) { + m.def("bilinear", + &py_bilinear, + py::arg("x1"), + py::arg("x2"), + py::arg("weight"), + py::arg("bias"), + R"doc(Bilinear transformation of two input tensors. +Args: + x1: First input tensor + x2: Second input tensor + weight: Weight tensor + bias: Bias tensor (optional) +Returns: + Output tensor after bilinear transformation +)doc"); + + m.def("bilinear_", + &py_bilinear_, + py::arg("out"), + py::arg("x1"), + py::arg("x2"), + py::arg("weight"), + py::arg("bias"), + R"doc(In-place bilinear transformation of two input tensors. +Args: + out: Output tensor + x1: First input tensor + x2: Second input tensor + weight: Weight tensor + bias: Bias tensor (optional) +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/binary_cross_entropy_with_logits.hpp b/csrc/infinicore/src/pybind11/ops/binary_cross_entropy_with_logits.hpp new file mode 100644 index 000000000..2e9527ffb --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/binary_cross_entropy_with_logits.hpp @@ -0,0 +1,77 @@ +#pragma once + +#include "infinicore/ops/binary_cross_entropy_with_logits.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_binary_cross_entropy_with_logits(py::module &m) { + // 1. 绑定 out-of-place 接口: out = binary_cross_entropy_with_logits(...) + m.def( + "binary_cross_entropy_with_logits", + [](Tensor logits, + Tensor target, + py::object weight, + py::object pos_weight, + std::string reduction) { + Tensor w = weight.is_none() ? Tensor() : weight.cast(); + Tensor pw = pos_weight.is_none() ? Tensor() : pos_weight.cast(); + + return op::binary_cross_entropy_with_logits( + logits, target, w, pw, reduction); + }, + py::arg("input"), + py::arg("target"), + py::arg("weight") = py::none(), + py::arg("pos_weight") = py::none(), + py::arg("reduction") = "mean", + R"doc(Measures Binary Cross Entropy between target and output logits. + + +Args: + input: Tensor of arbitrary shape as unnormalized scores (logits). + target: Tensor of the same shape as input with values between 0 and 1. + weight: Optional rescaling weight for each loss component. + pos_weight: Optional weight for positive examples (must be broadcastable). + reduction: Specfies the reduction to apply: 'none' | 'mean' | 'sum'. + +Returns: + A tensor representing the loss. +)doc"); + + // 2. 绑定指定输出接口: binary_cross_entropy_with_logits_(out, ...) + m.def( + "binary_cross_entropy_with_logits_", + [](Tensor output, + Tensor logits, + Tensor target, + py::object weight, + py::object pos_weight, + std::string reduction) { + Tensor w = weight.is_none() ? Tensor() : weight.cast(); + Tensor pw = pos_weight.is_none() ? Tensor() : pos_weight.cast(); + + return op::binary_cross_entropy_with_logits_( + output, logits, target, w, pw, reduction); + }, + py::arg("out"), + py::arg("input"), + py::arg("target"), + py::arg("weight") = py::none(), + py::arg("pos_weight") = py::none(), + py::arg("reduction") = "mean", + R"doc(Specified output version of binary_cross_entropy_with_logits. + +Args: + out: The destination tensor to store the loss. + input: Logits tensor. + target: Target tensor. + weight: Optional sample weight. + pos_weight: Optional positive class weight. + reduction: Specfies the reduction to apply. +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/bitwise_right_shift.hpp b/csrc/infinicore/src/pybind11/ops/bitwise_right_shift.hpp new file mode 100644 index 000000000..a540a33bc --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/bitwise_right_shift.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/bitwise_right_shift.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_bitwise_right_shift(py::module &m) { + m.def("bitwise_right_shift", + &op::bitwise_right_shift, + py::arg("input"), + py::arg("other"), + R"doc(Element-wise bitwise right shift.)doc"); + + m.def("bitwise_right_shift_", + &op::bitwise_right_shift_, + py::arg("out"), + py::arg("input"), + py::arg("other"), + R"doc(In-place element-wise bitwise right shift.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/blas_amax.hpp b/csrc/infinicore/src/pybind11/ops/blas_amax.hpp new file mode 100644 index 000000000..51e8cfe4d --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/blas_amax.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/blas_amax.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_blas_amax(py::module &m) { + m.def("blas_amax", + &op::blas_amax, + py::arg("x"), + R"doc(BLAS level-1 amax.)doc"); + + m.def("blas_amax_", + &op::blas_amax_, + py::arg("x"), + py::arg("result"), + R"doc(In-place BLAS level-1 amax.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/blas_amin.hpp b/csrc/infinicore/src/pybind11/ops/blas_amin.hpp new file mode 100644 index 000000000..8961a9363 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/blas_amin.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/blas_amin.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_blas_amin(py::module &m) { + m.def("blas_amin", + &op::blas_amin, + py::arg("x"), + R"doc(BLAS level-1 amin.)doc"); + + m.def("blas_amin_", + &op::blas_amin_, + py::arg("x"), + py::arg("result"), + R"doc(In-place BLAS level-1 amin.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/blas_copy.hpp b/csrc/infinicore/src/pybind11/ops/blas_copy.hpp new file mode 100644 index 000000000..c348ac38b --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/blas_copy.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include + +#include "infinicore/ops/blas_copy.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_blas_copy(py::module &m) { + m.def("blas_copy_", + &op::blas_copy_, + py::arg("x"), + py::arg("y"), + R"doc(In-place BLAS level-1 copy from x to y.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/blas_dot.hpp b/csrc/infinicore/src/pybind11/ops/blas_dot.hpp new file mode 100644 index 000000000..73b4f0bc9 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/blas_dot.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/blas_dot.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_blas_dot(py::module &m) { + m.def("blas_dot", + &op::blas_dot, + py::arg("x"), + py::arg("y"), + R"doc(BLAS level-1 dot.)doc"); + + m.def("blas_dot_", + &op::blas_dot_, + py::arg("x"), + py::arg("y"), + py::arg("result"), + R"doc(In-place BLAS level-1 dot.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/block_diag.hpp b/csrc/infinicore/src/pybind11/ops/block_diag.hpp new file mode 100644 index 000000000..b8f9100fe --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/block_diag.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include +#include + +#include "infinicore/ops/block_diag.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_block_diag(py::module &m) { + m.def( + "block_diag", + &op::block_diag, + py::arg("tensors"), + R"doc(Construct a block diagonal matrix from a list of 2D tensors.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/broadcast_to.hpp b/csrc/infinicore/src/pybind11/ops/broadcast_to.hpp new file mode 100644 index 000000000..6ae1e6482 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/broadcast_to.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include +#include + +#include "infinicore/ops/broadcast_to.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_broadcast_to(py::module &m) { + m.def("broadcast_to", + &op::broadcast_to, + py::arg("x"), + py::arg("shape"), + R"doc(Broadcast tensor to target shape.)doc"); + + m.def("broadcast_to_", + &op::broadcast_to_, + py::arg("y"), + py::arg("x"), + R"doc(In-place/Out broadcast tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/cat.hpp b/csrc/infinicore/src/pybind11/ops/cat.hpp new file mode 100644 index 000000000..093b15ff2 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/cat.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "infinicore/ops/cat.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_cat(py::module &m) { + + m.def("cat", + &op::cat, + py::arg("tensors"), + py::arg("dim") = 0, + R"doc(opertor: torch.cat, out-of-place mode)doc"); + + m.def("cat_", + &op::cat_, + py::arg("out"), + py::arg("tensors"), + py::arg("dim") = 0, + R"doc(opertor: torch.cat, in-place mode)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/causal_conv1d.hpp b/csrc/infinicore/src/pybind11/ops/causal_conv1d.hpp new file mode 100644 index 000000000..4fbeb3c69 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/causal_conv1d.hpp @@ -0,0 +1,38 @@ +#pragma once + +#include +#include + +#include "infinicore/ops/causal_conv1d.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_causal_conv1d(py::module &m) { + m.def("causal_conv1d", + &op::causal_conv1d, + py::arg("qkv"), + py::arg("conv_state"), + py::arg("weight"), + py::arg("bias") = std::nullopt, + py::arg("cu_seqlens") = std::nullopt, + py::arg("initial_state_indices") = std::nullopt, + py::arg("final_state_indices") = std::nullopt, + R"doc(Causal depthwise Conv1d. Returns out only. + +Padded mode: + qkv/out: [B, T, C], conv_state: [B, C, state_len]. + +Continuous-batch mode: + pass cu_seqlens [num_requests + 1]; qkv/out: [1, total_tokens, C]. + +Indexed pool mode: + conv_state is [pool_size, C, state_len]. Provide initial_state_indices [num_requests] + to read states. Provide final_state_indices [num_requests] to write final states + in-place to conv_state. The current backend supports K == 4, where + weight is [C, 1, K] and conv_state is [*, C, K - 1]. +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/causal_softmax.hpp b/csrc/infinicore/src/pybind11/ops/causal_softmax.hpp new file mode 100644 index 000000000..926a96d90 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/causal_softmax.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/causal_softmax.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_causal_softmax(py::module &m) { + m.def("causal_softmax", + &op::causal_softmax, + py::arg("input"), + R"doc(Causal softmax activation function.)doc"); + + m.def("causal_softmax_", + &op::causal_softmax_, + py::arg("output"), + py::arg("input"), + R"doc(In-place causal softmax activation function.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/cdist.hpp b/csrc/infinicore/src/pybind11/ops/cdist.hpp new file mode 100644 index 000000000..a52c24af3 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/cdist.hpp @@ -0,0 +1,46 @@ +#pragma once + +#include + +#include "infinicore/ops/cdist.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_cdist(py::module &m) { + // 1. 绑定 out-of-place 接口: out = cdist(x1, x2, p) + m.def("cdist", + &op::cdist, + py::arg("x1"), + py::arg("x2"), + py::arg("p") = 2.0, + R"doc(Computes batched pairwise distance between vectors in x1 and x2 using p-norm. + +Args: + x1: First set of vectors, shape (M, D) + x2: Second set of vectors, shape (N, D) + p: The p-norm to apply (default: 2.0) + +Returns: + A matrix containing pairwise distances, shape (M, N) +)doc"); + + // 2. 绑定 in-place / specified output 接口: cdist_(out, x1, x2, p) + m.def("cdist_", + &op::cdist_, + py::arg("out"), + py::arg("x1"), + py::arg("x2"), + py::arg("p") = 2.0, + R"doc(In-place version of cdist. Stores the results in the 'out' tensor. + +Args: + out: The destination tensor, shape (M, N) + x1: First set of vectors, shape (M, D) + x2: Second set of vectors, shape (N, D) + p: The p-norm to apply (default: 2.0) +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/chunk_gated_delta_rule.hpp b/csrc/infinicore/src/pybind11/ops/chunk_gated_delta_rule.hpp new file mode 100644 index 000000000..c39174f86 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/chunk_gated_delta_rule.hpp @@ -0,0 +1,42 @@ +#pragma once + +#include +#include + +#include "infinicore/ops/chunk_gated_delta_rule.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_chunk_gated_delta_rule(py::module &m) { + m.def("chunk_gated_delta_rule", + &op::chunk_gated_delta_rule, + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("g"), + py::arg("beta"), + py::arg("initial_state"), + py::arg("cu_seqlens") = std::nullopt, + py::arg("initial_state_indices") = std::nullopt, + py::arg("final_state_indices") = std::nullopt, + py::arg("use_qk_l2norm") = false, + py::arg("chunk_size") = 64, + R"doc(Chunk gated delta rule. Returns out only. + +Padded mode: + q/k: [B, T, Hk, Dk], v/out: [B, T, Hv, Dv], g/beta: [B, T, Hv], + initial_state: [B, Hv, Dv, Dk]. + +Continuous-batch mode: + pass cu_seqlens [B + 1]; q/k: [1, total_tokens, Hk, Dk], + v/out: [1, total_tokens, Hv, Dv], g/beta: [1, total_tokens, Hv]. + +Indexed pool mode: + initial_state is [pool_size, Hv, Dv, Dk]. Provide both initial_state_indices + and final_state_indices [B]; final state is written in-place to initial_state. +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/conv2d.hpp b/csrc/infinicore/src/pybind11/ops/conv2d.hpp new file mode 100644 index 000000000..f282840e8 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/conv2d.hpp @@ -0,0 +1,61 @@ +#pragma once + +#include "infinicore/ops/conv2d.hpp" + +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_conv2d(Tensor input, + Tensor weight, + pybind11::object bias, + const std::vector &padding, + const std::vector &stride, + const std::vector &dilation) { + Tensor bias_tensor; + if (!bias.is_none()) { + bias_tensor = bias.cast(); + } + return op::conv2d(input, weight, bias_tensor, padding, stride, dilation); +} + +void py_conv2d_(Tensor out, + Tensor input, + Tensor weight, + pybind11::object bias, + const std::vector &padding, + const std::vector &stride, + const std::vector &dilation) { + Tensor bias_tensor; + if (!bias.is_none()) { + bias_tensor = bias.cast(); + } + op::conv2d_(out, input, weight, bias_tensor, padding, stride, dilation); +} + +inline void bind_conv2d(py::module &m) { + m.def("conv2d", + &ops::py_conv2d, + py::arg("input"), + py::arg("weight"), + py::arg("bias") = py::none(), + py::arg("padding") = std::vector{0, 0}, + py::arg("stride") = std::vector{1, 1}, + py::arg("dilation") = std::vector{1, 1}, + R"doc(Applies a 2D convolution over an input tensor.)doc"); + + m.def("conv2d_", + &ops::py_conv2d_, + py::arg("out"), + py::arg("input"), + py::arg("weight"), + py::arg("bias") = py::none(), + py::arg("padding") = std::vector{0, 0}, + py::arg("stride") = std::vector{1, 1}, + py::arg("dilation") = std::vector{1, 1}, + R"doc(In-place 2D convolution.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/cross_entropy.hpp b/csrc/infinicore/src/pybind11/ops/cross_entropy.hpp new file mode 100644 index 000000000..8105642a6 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/cross_entropy.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/cross_entropy.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_cross_entropy(py::module &m) { + m.def("cross_entropy", + &op::cross_entropy, + py::arg("logits"), + py::arg("target"), + R"doc(Token-wise cross entropy loss without reduction.)doc"); + + m.def("cross_entropy_", + &op::cross_entropy_, + py::arg("loss"), + py::arg("logits"), + py::arg("target"), + R"doc(Write cross entropy loss into a provided tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/diff.hpp b/csrc/infinicore/src/pybind11/ops/diff.hpp new file mode 100644 index 000000000..4a83d9607 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/diff.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include + +#include "infinicore/ops/diff.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_diff(py::module &m) { + m.def("diff", + &op::diff, + py::arg("x"), + py::arg("n") = 1, + py::arg("dim") = -1, + R"doc(Difference of adjacent elements along a dimension.)doc"); + + m.def("diff_", + &op::diff_, + py::arg("y"), + py::arg("x"), + py::arg("n") = 1, + py::arg("dim") = -1, + R"doc(Out variant of diff.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/digamma.hpp b/csrc/infinicore/src/pybind11/ops/digamma.hpp new file mode 100644 index 000000000..3110c7a07 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/digamma.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/digamma.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_digamma(py::module &m) { + m.def("digamma", + &op::digamma, + py::arg("x"), + R"doc(Digamma function.)doc"); + + m.def("digamma_", + &op::digamma_, + py::arg("y"), + py::arg("x"), + R"doc(Out variant of digamma.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/dist.hpp b/csrc/infinicore/src/pybind11/ops/dist.hpp new file mode 100644 index 000000000..9aec0cca6 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/dist.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include + +#include "infinicore/ops/dist.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_dist(py::module &m) { + m.def("dist", + &op::dist, + py::arg("x1"), + py::arg("x2"), + py::arg("p") = 2.0, + R"doc(p-norm distance between two tensors.)doc"); + + m.def("dist_", + &op::dist_, + py::arg("y"), + py::arg("x1"), + py::arg("x2"), + py::arg("p") = 2.0, + R"doc(Out variant of dist.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/embedding.hpp b/csrc/infinicore/src/pybind11/ops/embedding.hpp new file mode 100644 index 000000000..44e14b61c --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/embedding.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "infinicore/ops/embedding.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_embedding(py::module &m) { + + m.def("embedding", + &op::embedding, + py::arg("input"), + py::arg("weight"), + R"doc(Generate a simple lookup table that looks up embeddings in a fixed dictionary and size..)doc"); + + m.def("embedding_", + &op::embedding_, + py::arg("out"), + py::arg("input"), + py::arg("weight"), + R"doc(In-place, Generate a simple lookup table that looks up embeddings in a fixed dictionary and size..)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/equal.hpp b/csrc/infinicore/src/pybind11/ops/equal.hpp new file mode 100644 index 000000000..d14a6b61d --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/equal.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/equal.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_equal(py::module &m) { + m.def("equal", + &op::equal, + py::arg("a"), + py::arg("b"), + R"doc(Elementwise equality returning a bool tensor.)doc"); + + m.def("equal_", + &op::equal_, + py::arg("out"), + py::arg("a"), + py::arg("b"), + R"doc(In-place elementwise equality writing into `out`.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/flash_attention.hpp b/csrc/infinicore/src/pybind11/ops/flash_attention.hpp new file mode 100644 index 000000000..6e3766796 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/flash_attention.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include + +#include "infinicore/ops/flash_attention.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_flash_attention(py::module &m) { + m.def("flash_attention", + &op::flash_attention, + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("total_kv_len"), + py::arg("scale"), + py::arg("is_causal")); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/flipud.hpp b/csrc/infinicore/src/pybind11/ops/flipud.hpp new file mode 100644 index 000000000..97c5641d6 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/flipud.hpp @@ -0,0 +1,31 @@ +#pragma once + +#include "infinicore/ops/flipud.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_flipud(py::module &m) { + // 1. 绑定 out-of-place 接口: output = flipud(input) + m.def("flipud", + &op::flipud, + py::arg("input"), + R"doc(Flip array in the up/down direction. + + Flips the entries in axis 0 (preserving the shape). + + Args: + input (Tensor): The input tensor. + )doc"); + + // 2. 绑定 explicit output 接口: flipud_(output, input) + m.def("flipud_", + &op::flipud_, + py::arg("output"), + py::arg("input"), + R"doc(Explicit output FlipUD operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/float_power.hpp b/csrc/infinicore/src/pybind11/ops/float_power.hpp new file mode 100644 index 000000000..8ee3a5d77 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/float_power.hpp @@ -0,0 +1,59 @@ +#include "../tensor.hpp" +#include "infinicore/ops/float_power.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +using infinicore::Tensor; +using infinicore::op::float_power; +using infinicore::op::float_power_; + +inline Tensor unwrap(py::handle obj) { + try { + return obj.cast(); + } catch (...) {} + + if (py::hasattr(obj, "_underlying")) { + return obj.attr("_underlying").cast(); + } + + throw py::type_error("Expected infinicore.Tensor, but got " + py::repr(obj.get_type()).cast()); +} + +void bind_float_power(py::module &m) { + + // --- Out-of-place: float_power(input, exponent) --- + m.def( + "float_power", [](py::object input_obj, py::object exp_obj) -> Tensor { + Tensor input = unwrap(input_obj); + + // 处理标量指数的情况 (float 或 int) + if (py::isinstance(exp_obj) || py::isinstance(exp_obj)) { + return float_power(input, exp_obj.cast()); + } + + // 处理张量指数的情况 + Tensor exponent = unwrap(exp_obj); + return float_power(input, exponent); + }, + py::arg("input"), py::arg("exponent")); + + // --- In-place: float_power_(out, input, exponent) --- + m.def( + "float_power_", [](py::object out_obj, py::object input_obj, py::object exp_obj) { + Tensor out = unwrap(out_obj); + Tensor input = unwrap(input_obj); + + if (py::isinstance(exp_obj) || py::isinstance(exp_obj)) { + float_power_(out, input, exp_obj.cast()); + } else { + Tensor exponent = unwrap(exp_obj); + float_power_(out, input, exponent); + } + }, + py::arg("out"), py::arg("input"), py::arg("exponent")); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/floor.hpp b/csrc/infinicore/src/pybind11/ops/floor.hpp new file mode 100644 index 000000000..a209ea5bc --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/floor.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include "infinicore/ops/floor.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_floor(py::module &m) { + // 绑定 out-of-place 接口: output = floor(input) + m.def("floor", + &op::floor, + py::arg("input"), + R"doc(Computes the floor of each element of input.)doc"); + + // 绑定 in-place 接口: floor_(output, input) + m.def("floor_", + &op::floor_, + py::arg("output"), + py::arg("input"), + R"doc(In-place floor operation. Writes result into output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/floor_divide.hpp b/csrc/infinicore/src/pybind11/ops/floor_divide.hpp new file mode 100644 index 000000000..626767218 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/floor_divide.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/floor_divide.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_floor_divide(py::module &m) { + m.def("floor_divide", + &op::floor_divide, + py::arg("a"), + py::arg("b"), + R"doc(Floor division of two tensors.)doc"); + + m.def("floor_divide_", + &op::floor_divide_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + R"doc(In-place tensor floor division.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/fmin.hpp b/csrc/infinicore/src/pybind11/ops/fmin.hpp new file mode 100644 index 000000000..afdbe4b33 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/fmin.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/fmin.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_fmin(py::module &m) { + m.def("fmin", + &op::fmin, + py::arg("a"), + py::arg("b"), + R"doc(fmin of two tensors.)doc"); + + m.def("fmin_", + &op::fmin_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + R"doc(In-place tensor fmin.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/fmod.hpp b/csrc/infinicore/src/pybind11/ops/fmod.hpp new file mode 100644 index 000000000..97af57da2 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/fmod.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/fmod.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_fmod(py::module &m) { + m.def("fmod", + &op::fmod, + py::arg("a"), + py::arg("b"), + R"doc(Element-wise floating point remainder of division of two tensors.)doc"); + + m.def("fmod_", + &op::fmod_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + R"doc(In-place element-wise floating point remainder of division of two tensors.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/fused_gated_delta_net_gating.hpp b/csrc/infinicore/src/pybind11/ops/fused_gated_delta_net_gating.hpp new file mode 100644 index 000000000..e43566977 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/fused_gated_delta_net_gating.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include + +#include "infinicore/ops/fused_gated_delta_net_gating.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_fused_gated_delta_net_gating(py::module &m) { + m.def( + "fused_gated_delta_net_gating", + [](const Tensor &A_log, + const Tensor &a, + const Tensor &b, + const Tensor &dt_bias, + float beta, + float threshold) { + auto result = op::fused_gated_delta_net_gating(A_log, a, b, dt_bias, beta, threshold); + return py::make_tuple(result.first, result.second); + }, + py::arg("A_log"), + py::arg("a"), + py::arg("b"), + py::arg("dt_bias"), + py::arg("beta") = 1.0f, + py::arg("threshold") = 20.0f, + R"doc(Fused GatedDeltaNet gating out-of-place.)doc"); + + m.def("fused_gated_delta_net_gating_", + &op::fused_gated_delta_net_gating_, + py::arg("g"), + py::arg("beta_output"), + py::arg("A_log"), + py::arg("a"), + py::arg("b"), + py::arg("dt_bias"), + py::arg("beta") = 1.0f, + py::arg("threshold") = 20.0f, + R"doc(Fused GatedDeltaNet gating writing to provided outputs.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/gaussian_nll_loss.hpp b/csrc/infinicore/src/pybind11/ops/gaussian_nll_loss.hpp new file mode 100644 index 000000000..974f8e29b --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/gaussian_nll_loss.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include + +#include "infinicore/ops/gaussian_nll_loss.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_gaussian_nll_loss(py::module &m) { + m.def("gaussian_nll_loss", + &op::gaussian_nll_loss, + py::arg("input"), + py::arg("target"), + py::arg("var"), + py::arg("full") = false, + py::arg("eps") = 1e-6, + py::arg("reduction") = 1, + R"doc(Gaussian negative log-likelihood loss.)doc"); + + m.def("gaussian_nll_loss_", + &op::gaussian_nll_loss_, + py::arg("out"), + py::arg("input"), + py::arg("target"), + py::arg("var"), + py::arg("full") = false, + py::arg("eps") = 1e-6, + py::arg("reduction") = 1, + R"doc(In-place Gaussian negative log-likelihood loss.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/hardswish.hpp b/csrc/infinicore/src/pybind11/ops/hardswish.hpp new file mode 100644 index 000000000..daaccec62 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/hardswish.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/hardswish.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_hardswish(py::module &m) { + m.def("hardswish", + &op::hardswish, + py::arg("input"), + R"doc(Out-of-place Hardswish activation.)doc"); + + m.def("hardswish_", + &op::hardswish_, + py::arg("output"), + py::arg("input"), + R"doc(In-place Hardswish activation.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/hardtanh.hpp b/csrc/infinicore/src/pybind11/ops/hardtanh.hpp new file mode 100644 index 000000000..ff9abb872 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/hardtanh.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include + +#include "infinicore/ops/hardtanh.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_hardtanh(py::module &m) { + m.def("hardtanh", + &op::hardtanh, + py::arg("input"), + py::arg("min_val") = -1.0f, + py::arg("max_val") = 1.0f, + R"doc(Apply the HardTanh activation.)doc"); + + m.def("hardtanh_", + &op::hardtanh_, + py::arg("output"), + py::arg("input"), + py::arg("min_val") = -1.0f, + py::arg("max_val") = 1.0f, + R"doc(In-place HardTanh activation.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/hinge_embedding_loss.hpp b/csrc/infinicore/src/pybind11/ops/hinge_embedding_loss.hpp new file mode 100644 index 000000000..fe481d853 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/hinge_embedding_loss.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include + +#include "infinicore/ops/hinge_embedding_loss.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_hinge_embedding_loss(py::module &m) { + m.def( + "hinge_embedding_loss", + &op::hinge_embedding_loss, + py::arg("input"), + py::arg("target"), + py::arg("margin") = 1.0, + py::arg("reduction") = 1, + R"doc(Hinge embedding loss.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/huber_loss.hpp b/csrc/infinicore/src/pybind11/ops/huber_loss.hpp new file mode 100644 index 000000000..94f93636c --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/huber_loss.hpp @@ -0,0 +1,46 @@ +#pragma once + +#include "infinicore/ops/huber_loss.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_huber_loss(py::module &m) { + m.def( + "huber_loss", + [](const Tensor &input, const Tensor &target, float delta, int reduction) { + return op::huber_loss(input, target, delta, reduction); + }, + py::arg("input"), + py::arg("target"), + py::arg("delta") = 1.0f, + py::arg("reduction") = 1, + R"doc(Computes the Huber Loss between input and target. + + Args: + input (Tensor): Input tensor of arbitrary shape. + target (Tensor): Ground truth labels, same shape as input. + delta (float, optional): The threshold at which to change between delta-scaled L1 and L2 loss. Default: 1.0. + reduction (int, optional): Specifies the reduction to apply to the output: 0=None, 1=Mean, 2=Sum. Default: 1. + )doc"); + + // ------------------------------------------------------------------------- + // 2. 绑定 in-place 接口 (huber_loss_) + // ------------------------------------------------------------------------- + m.def( + "huber_loss_", + [](Tensor &output, const Tensor &input, const Tensor &target, float delta, int reduction) { + // 调用底层 + op::huber_loss_(output, input, target, delta, reduction); + }, + py::arg("output"), + py::arg("input"), + py::arg("target"), + py::arg("delta") = 1.0f, + py::arg("reduction") = 1, + R"doc(Explicit output Huber Loss operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/hypot.hpp b/csrc/infinicore/src/pybind11/ops/hypot.hpp new file mode 100644 index 000000000..70d08ef19 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/hypot.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include "infinicore/ops/hypot.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_hypot(py::module &m) { + // 绑定 out-of-place 接口: output = hypot(input, other) + m.def("hypot", + &op::hypot, + py::arg("input"), + py::arg("other"), + R"doc(Computes the hypotenuse of input and other arguments, i.e. sqrt(input^2 + other^2).)doc"); + + // 绑定 in-place / 指定输出接口: hypot_(output, input, other) + m.def("hypot_", + &op::hypot_, + py::arg("output"), + py::arg("input"), + py::arg("other"), + R"doc(In-place hypot operation. Writes result into output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/index_add.hpp b/csrc/infinicore/src/pybind11/ops/index_add.hpp new file mode 100644 index 000000000..4fadef4d1 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/index_add.hpp @@ -0,0 +1,31 @@ +#pragma once + +#include "infinicore/ops/index_add.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_index_add(py::module &m) { + m.def("index_add", + &op::index_add, + py::arg("input"), + py::arg("dim"), + py::arg("index"), + py::arg("source"), + py::arg("alpha") = 1.0f, + R"doc(Accumulate elements of source into input by adding to the indices in the order given in index. + Formula: output[index[i]] = input[index[i]] + alpha * source[i])doc"); + m.def("index_add_", + &op::index_add_, + py::arg("output"), + py::arg("input"), + py::arg("dim"), + py::arg("index"), + py::arg("source"), + py::arg("alpha") = 1.0f, + R"doc(In-place version of index_add. Writes result into output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/index_copy.hpp b/csrc/infinicore/src/pybind11/ops/index_copy.hpp new file mode 100644 index 000000000..c631be95c --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/index_copy.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include "infinicore/ops/index_copy.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_index_copy(py::module &m) { + // 1. Out-of-place version (returns new tensor) + m.def("index_copy", + &op::index_copy, + py::arg("input"), + py::arg("dim"), + py::arg("index"), + py::arg("source"), + R"doc(Copies elements of source into input at the indices given in index. + Formula: output[index[i]] = source[i])doc"); + m.def("index_copy_", + &op::index_copy_, + py::arg("output"), + py::arg("input"), + py::arg("dim"), + py::arg("index"), + py::arg("source"), + R"doc(In-place version of index_copy. Writes result into output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/inner.hpp b/csrc/infinicore/src/pybind11/ops/inner.hpp new file mode 100644 index 000000000..b54d6117b --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/inner.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "infinicore/ops/inner.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_inner(py::module &m) { + + m.def("inner", + &op::inner, + py::arg("input"), + py::arg("other"), + R"doc(opertor: torch.inner, out-of-place mode)doc"); + + m.def("inner_", + &op::inner_, + py::arg("out"), + py::arg("input"), + py::arg("other"), + R"doc(opertor: torch.inner, in-place mode)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/interpolate.hpp b/csrc/infinicore/src/pybind11/ops/interpolate.hpp new file mode 100644 index 000000000..26fb3b812 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/interpolate.hpp @@ -0,0 +1,33 @@ +#pragma once + +#include +#include + +#include "infinicore/ops/interpolate.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_interpolate(py::module &m) { + m.def("interpolate", + &op::interpolate, + py::arg("input"), + py::arg("mode"), + py::arg("size"), + py::arg("scale_factor"), + py::arg("align_corners"), + R"doc(Interpolate (upsample/downsample) a tensor.)doc"); + + m.def("interpolate_", + &op::interpolate_, + py::arg("out"), + py::arg("input"), + py::arg("mode"), + py::arg("size"), + py::arg("scale_factor"), + py::arg("align_corners"), + R"doc(In-place interpolate (writes to out).)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/kron.hpp b/csrc/infinicore/src/pybind11/ops/kron.hpp new file mode 100644 index 000000000..751ef879f --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/kron.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include + +#include "infinicore/ops/kron.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_kron(py::module &m) { + m.def( + "kron", + &op::kron, + py::arg("a"), + py::arg("b"), + R"doc(Kronecker product.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/kthvalue.hpp b/csrc/infinicore/src/pybind11/ops/kthvalue.hpp new file mode 100644 index 000000000..a0f67ae2a --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/kthvalue.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include "infinicore/ops/kthvalue.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_kthvalue(py::module &m) { + // 1. 绑定 functional 接口: (values, indices) = kthvalue(input, k, dim, keepdim) + m.def("kthvalue", + &op::kthvalue, + py::arg("input"), + py::arg("k"), + py::arg("dim") = -1, + py::arg("keepdim") = false, + R"doc(Returns the k-th smallest element of each row of the input tensor in the given dimension. + + Args: + input (Tensor): The input tensor. + k (int): The k value. + dim (int): The dimension to find the k-th value along. + keepdim (bool): Whether to keep the output dimension. + )doc"); + + // 2. 绑定 explicit output 接口: kthvalue_(values, indices, input, k, dim, keepdim) + m.def("kthvalue_", + &op::kthvalue_, + py::arg("values"), + py::arg("indices"), + py::arg("input"), + py::arg("k"), + py::arg("dim") = -1, + py::arg("keepdim") = false, + R"doc(Explicit output Kthvalue operation. Writes results into values and indices tensors.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/kv_caching.hpp b/csrc/infinicore/src/pybind11/ops/kv_caching.hpp new file mode 100644 index 000000000..2864312b2 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/kv_caching.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include + +#include "infinicore/ops/kv_caching.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_kv_caching(py::module &m) { + m.def("kv_caching_", + &op::kv_caching_, + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("k"), + py::arg("v"), + py::arg("past_kv_lengths"), + R"doc(In-place Key-Value Caching. + +Updates the KV cache in-place with new key and value tensors. + +Args: + k_cache: Key cache tensor to update in-place + v_cache: Value cache tensor to update in-place + k: New key tensor to append + v: New value tensor to append + past_kv_lengths: Tensor containing current sequence lengths for each batch +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/layer_norm.hpp b/csrc/infinicore/src/pybind11/ops/layer_norm.hpp new file mode 100644 index 000000000..5ca6a87d3 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/layer_norm.hpp @@ -0,0 +1,48 @@ +#pragma once + +#include + +#include "infinicore/ops/layer_norm.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_layer_norm(py::module &m) { + m.def("layer_norm", + &op::layer_norm, + py::arg("x"), + py::arg("weight"), + py::arg("bias"), + py::arg("epsilon") = 1e-5f, + R"doc(Layer Normalization. + +Args: + x: Input tensor + weight: Scale weights + bias: Bias weights + epsilon: Small constant for numerical stability, default is 1e-5 + +Returns: + Normalized tensor with same shape as input +)doc"); + + m.def("layer_norm_", + &op::layer_norm_for_pybind, + py::arg("y"), + py::arg("x"), + py::arg("weight"), + py::arg("bias"), + py::arg("epsilon") = 1e-5f, + R"doc(In-place Layer Normalization. + +Args: + y: Output tensor + x: Input tensor + weight: Scale weights + bias: Bias weights + epsilon: Small constant for numerical stability, default is 1e-5 +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/ldexp.hpp b/csrc/infinicore/src/pybind11/ops/ldexp.hpp new file mode 100644 index 000000000..948c03d68 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/ldexp.hpp @@ -0,0 +1,40 @@ +#pragma once + +#include + +#include "infinicore/ops/ldexp.hpp" + +namespace py = pybind11; + +#pragma once + +#include "infinicore/ops/ldexp.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_ldexp(py::module &m) { + // 1. 绑定 functional 接口: output = ldexp(input, other) + m.def("ldexp", + &op::ldexp, + py::arg("input"), + py::arg("other"), + R"doc(Multiplies input by 2 raised to the power of other. + + Args: + input (Tensor): The input tensor (mantissa). + other (Tensor): The exponent tensor. + )doc"); + + // 2. 绑定 explicit output 接口: ldexp_(output, input, other) + m.def("ldexp_", + &op::ldexp_, + py::arg("output"), + py::arg("input"), + py::arg("other"), + R"doc(Explicit output Ldexp operation. Writes result into output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/lerp.hpp b/csrc/infinicore/src/pybind11/ops/lerp.hpp new file mode 100644 index 000000000..5e220c95f --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/lerp.hpp @@ -0,0 +1,64 @@ +#pragma once + +#include "infinicore/ops/lerp.hpp" +#include +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_lerp(py::module &m) { + // 定义函数指针别名，用于区分重载 + using LerpTensorFunc = Tensor (*)(Tensor, Tensor, Tensor); + using LerpScalarFunc = Tensor (*)(Tensor, Tensor, float); + using LerpTensorInplaceFunc = void (*)(Tensor, Tensor, Tensor, Tensor); + using LerpScalarInplaceFunc = void (*)(Tensor, Tensor, Tensor, float); + + // ======================================================================== + // 1. 绑定 functional 接口 + // ======================================================================== + + // 重载 1: weight 为 Tensor + m.def("lerp", + static_cast(&op::lerp), + py::arg("start"), + py::arg("end"), + py::arg("weight"), + R"doc(Does a linear interpolation of two tensors start and end based on a tensor weight. + + output = start + weight * (end - start) + )doc"); + + // 重载 2: weight 为 float + m.def("lerp", + static_cast(&op::lerp), + py::arg("start"), + py::arg("end"), + py::arg("weight"), + R"doc(Does a linear interpolation of two tensors start and end based on a scalar weight.)doc"); + + // ======================================================================== + // 2. 绑定 explicit output 接口 (In-place) + // ======================================================================== + + // 重载 1: weight 为 Tensor + m.def("lerp_", + static_cast(&op::lerp_), + py::arg("output"), + py::arg("start"), + py::arg("end"), + py::arg("weight"), + R"doc(Explicit output Lerp operation with tensor weight. Writes the result into the output tensor.)doc"); + + // 重载 2: weight 为 float + m.def("lerp_", + static_cast(&op::lerp_), + py::arg("output"), + py::arg("start"), + py::arg("end"), + py::arg("weight"), + R"doc(Explicit output Lerp operation with scalar weight. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/linear.hpp b/csrc/infinicore/src/pybind11/ops/linear.hpp new file mode 100644 index 000000000..e8f9a90a2 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/linear.hpp @@ -0,0 +1,56 @@ +#pragma once + +#include "infinicore/ops/linear.hpp" + +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_linear(Tensor input, + Tensor weight, + pybind11::object bias, + float alpha = 1.0f) { + std::optional bias_tensor = std::nullopt; + if (!bias.is_none()) { + bias_tensor = bias.cast(); + } + return op::linear(input, weight, bias_tensor, alpha); +} + +void py_linear_(Tensor out, + Tensor input, + Tensor weight, + pybind11::object bias, + float alpha = 1.0f) { + + std::optional bias_tensor = std::nullopt; + if (!bias.is_none()) { + bias_tensor = bias.cast(); + } + + op::linear_(out, input, weight, bias_tensor, alpha); +} + +inline void bind_linear(py::module &m) { + + m.def("linear", + &ops::py_linear, + py::arg("input"), + py::arg("weight"), + py::arg("bias") = py::none(), + py::arg("alpha") = 1.0f, + R"doc(Applies a linear transformation to the incoming data: y=alpha*xA^T+b.)doc"); + + m.def("linear_", + &ops::py_linear_, + py::arg("out"), + py::arg("input"), + py::arg("weight"), + py::arg("bias") = py::none(), + py::arg("alpha") = 1.0f, + R"doc(In-place, applies a linear transformation to the incoming data: y=alpha*xA^T+b.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/linear_w8a8i8.hpp b/csrc/infinicore/src/pybind11/ops/linear_w8a8i8.hpp new file mode 100644 index 000000000..926d554b1 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/linear_w8a8i8.hpp @@ -0,0 +1,54 @@ +#pragma once + +#include + +#include "infinicore/ops/linear_w8a8i8.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_linear_w8a8i8(Tensor input, + Tensor weight_packed, + Tensor weight_scale, + pybind11::object bias) { + std::optional bias_tensor = std::nullopt; + if (!bias.is_none()) { + bias_tensor = bias.cast(); + } + return op::linear_w8a8i8(input, weight_packed, weight_scale, bias_tensor); +} + +void py_linear_w8a8i8_(Tensor out, + Tensor input, + Tensor weight_packed, + Tensor weight_scale, + pybind11::object bias) { + + std::optional bias_tensor = std::nullopt; + if (!bias.is_none()) { + bias_tensor = bias.cast(); + } + + op::linear_w8a8i8_(out, input, weight_packed, weight_scale, bias_tensor); +} + +inline void bind_linear_w8a8i8(py::module &m) { + m.def("linear_w8a8i8", + &ops::py_linear_w8a8i8, + py::arg("input"), + py::arg("weight_packed"), + py::arg("weight_scale"), + py::arg("bias") = py::none(), + R"doc(linear_w8a8i8.)doc"); + m.def("linear_w8a8i8_", + &ops::py_linear_w8a8i8_, + py::arg("out"), + py::arg("input"), + py::arg("weight_packed"), + py::arg("weight_scale"), + py::arg("bias") = py::none(), + R"doc(linear_w8a8i8_.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/log_softmax.hpp b/csrc/infinicore/src/pybind11/ops/log_softmax.hpp new file mode 100644 index 000000000..b2eb8f47d --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/log_softmax.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include "infinicore/ops/log_softmax.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_log_softmax(py::module &m) { + // 1. 绑定 functional 接口: output = log_softmax(input, dim) + m.def("log_softmax", + &op::log_softmax, + py::arg("input"), + py::arg("dim"), + R"doc(Applies a softmax followed by a logarithm. + + Args: + input (Tensor): The input tensor. + dim (int): A dimension along which log_softmax will be computed. + )doc"); + + // 2. 绑定 explicit output 接口: log_softmax_(output, input, dim) + m.def("log_softmax_", + &op::log_softmax_, + py::arg("output"), + py::arg("input"), + py::arg("dim"), + R"doc(Explicit output LogSoftmax operation. Writes results into output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/logaddexp.hpp b/csrc/infinicore/src/pybind11/ops/logaddexp.hpp new file mode 100644 index 000000000..8f1bc6b18 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/logaddexp.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include + +#include "infinicore/ops/logaddexp.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_logaddexp(py::module &m) { + m.def("logaddexp", + &op::logaddexp, + py::arg("a"), + py::arg("b"), + R"doc(Logarithm of the sum of exponentiations of the inputs.)doc"); + m.def("logaddexp_", + &op::logaddexp_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + R"doc(In-place logaddexp operation. Writes results into c tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/logaddexp2.hpp b/csrc/infinicore/src/pybind11/ops/logaddexp2.hpp new file mode 100644 index 000000000..77ee154d7 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/logaddexp2.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include + +#include "infinicore/ops/logaddexp2.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_logaddexp2(py::module &m) { + m.def("logaddexp2", + &op::logaddexp2, + py::arg("a"), + py::arg("b"), + R"doc(Logarithm of the sum of exponentiations of the inputs in base-2.)doc"); + m.def("logaddexp2_", + &op::logaddexp2_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + R"doc(In-place logaddexp2 operation. Writes results into c tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/logcumsumexp.hpp b/csrc/infinicore/src/pybind11/ops/logcumsumexp.hpp new file mode 100644 index 000000000..aace6cc07 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/logcumsumexp.hpp @@ -0,0 +1,31 @@ +#pragma once + +#include "infinicore/ops/logcumsumexp.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_logcumsumexp(py::module &m) { + // 绑定非原地操作接口 (返回新 Tensor) + m.def("logcumsumexp", + &op::logcumsumexp, + py::arg("input"), + py::arg("dim"), + py::arg("exclusive") = false, + py::arg("reverse") = false, + R"doc(Computes the logarithm of the cumulative summation of the exponentiation of elements.)doc"); + + // 绑定原地/指定输出接口 + m.def("logcumsumexp_", + &op::logcumsumexp_, + py::arg("out"), + py::arg("input"), + py::arg("dim"), + py::arg("exclusive") = false, + py::arg("reverse") = false, + R"doc(In-place version of logcumsumexp.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/logdet.hpp b/csrc/infinicore/src/pybind11/ops/logdet.hpp new file mode 100644 index 000000000..e2a8f2d28 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/logdet.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/logdet.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_logdet(py::module &m) { + m.def("logdet", + &op::logdet, + py::arg("x"), + R"doc(Log determinant of a square matrix (NaN if determinant is negative).)doc"); + + m.def("logdet_", + &op::logdet_, + py::arg("y"), + py::arg("x"), + R"doc(Out variant of logdet.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/logical_and.hpp b/csrc/infinicore/src/pybind11/ops/logical_and.hpp new file mode 100644 index 000000000..1ca34fdc0 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/logical_and.hpp @@ -0,0 +1,27 @@ +#pragma once +#include + +#include "infinicore/ops/logical_and.hpp" + +namespace py = pybind11; +namespace infinicore::ops { + +inline void bind_logical_and(py::module &m) { + // 绑定常规函数: logical_and(input, other) -> Tensor + m.def("logical_and", + &op::logical_and, + py::arg("input"), + py::arg("other"), + R"doc(Computes the element-wise logical AND of the given input tensors.)doc"); + + // 绑定底层输出指定函数: logical_and_(output, input, other) + // 对应 Python 调用: _infinicore.logical_and_(out, input, other) + m.def("logical_and_", + &op::logical_and_, + py::arg("output"), + py::arg("input"), + py::arg("other"), + R"doc(Explicit output logical AND computation.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/logical_not.hpp b/csrc/infinicore/src/pybind11/ops/logical_not.hpp new file mode 100644 index 000000000..94edd2388 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/logical_not.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "infinicore/ops/logical_not.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_logical_not(py::module &m) { + // Out-of-place: output = logical_not(input) + m.def("logical_not", + &op::logical_not, + py::arg("input"), + R"doc(Logical NOT of the tensor.)doc"); + + // In-place / Explicit Output: logical_not_(output, input) + // 对应 C++: void logical_not_(Tensor output, Tensor input) + m.def("logical_not_", + &op::logical_not_, + py::arg("output"), + py::arg("input"), + R"doc(In-place logical NOT computation.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/mamba_selective_scan.hpp b/csrc/infinicore/src/pybind11/ops/mamba_selective_scan.hpp new file mode 100644 index 000000000..7bf6e9508 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/mamba_selective_scan.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include + +#include "infinicore/ops/mamba_selective_scan.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_mamba_selective_scan(py::module &m) { + m.def("mamba_selective_scan", + &op::mamba_selective_scan, + py::arg("x"), + py::arg("dt"), + py::arg("b"), + py::arg("c"), + py::arg("a_log"), + py::arg("d"), + py::arg("gate"), + py::arg("dt_bias"), + py::arg("state"), + R"doc(Mamba selective scan. Returns out and updates state in-place. + +Shapes: + x, dt, gate, out: [batch, seq_len, intermediate] + b, c: [batch, seq_len, state_size] + a_log: [intermediate, state_size] + d, dt_bias: [intermediate] + state: [batch, intermediate, state_size], float32 +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/masked_select.hpp b/csrc/infinicore/src/pybind11/ops/masked_select.hpp new file mode 100644 index 000000000..949566b6b --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/masked_select.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include "infinicore/ops/masked_select.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_masked_select(py::module &m) { + + m.def("masked_select", + &op::masked_select, + py::arg("input"), + py::arg("mask"), + R"doc(opertor: torch.masked_select, out-of-place mode)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/matmul.hpp b/csrc/infinicore/src/pybind11/ops/matmul.hpp new file mode 100644 index 000000000..c00bd2566 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/matmul.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include + +#include "infinicore/ops/matmul.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_matmul(py::module &m) { + m.def("matmul", + &op::matmul, + py::arg("a"), + py::arg("b"), + py::arg("alpha") = 1.0f, + R"doc(Matrix multiplication of two tensors.)doc"); + + m.def("matmul_", + &op::matmul_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + py::arg("alpha") = 1.0f, + R"doc(In-place matrix multiplication.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/mha.hpp b/csrc/infinicore/src/pybind11/ops/mha.hpp new file mode 100644 index 000000000..58dd21734 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/mha.hpp @@ -0,0 +1,78 @@ +#pragma once + +#include + +#include "infinicore/ops/mha.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_mha(Tensor q, + Tensor k, + Tensor v, + pybind11::object alibi_slopes, + float scale, + bool is_causal) { + std::optional alibi_slopes_tensor = std::nullopt; + if (!alibi_slopes.is_none()) { + alibi_slopes_tensor = alibi_slopes.cast(); + } + + return op::mha( + q, + k, + v, + alibi_slopes_tensor, + scale, + is_causal); +} + +void py_mha_(Tensor out, + Tensor q, + Tensor k, + Tensor v, + pybind11::object alibi_slopes, + float scale, + bool is_causal) { + std::optional alibi_slopes_tensor = std::nullopt; + if (!alibi_slopes.is_none()) { + alibi_slopes_tensor = alibi_slopes.cast(); + } + + op::mha_( + out, + q, + k, + v, + alibi_slopes_tensor, + scale, + is_causal); +} + +inline void bind_mha(py::module &m) { + m.def( + "mha", + &ops::py_mha, + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("alibi_slopes"), + py::arg("scale"), + py::arg("is_causal"), + R"doc(Variable-length multi-head attention.)doc"); + + m.def( + "mha_", + &ops::py_mha_, + py::arg("out"), + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("alibi_slopes"), + py::arg("scale"), + py::arg("is_causal"), + R"doc(In-place variable-length multi-head attention.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/mha_kvcache.hpp b/csrc/infinicore/src/pybind11/ops/mha_kvcache.hpp new file mode 100644 index 000000000..38934233e --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/mha_kvcache.hpp @@ -0,0 +1,127 @@ +#pragma once + +#include + +#include "infinicore/ops/mha_kvcache.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_mha_kvcache(Tensor q, + Tensor k_cache, + Tensor v_cache, + Tensor seqlens_k, + Tensor block_table, + pybind11::object alibi_slopes, + float scale) { + std::optional alibi_slopes_tensor = std::nullopt; + if (!alibi_slopes.is_none()) { + alibi_slopes_tensor = alibi_slopes.cast(); + } + + return op::mha_kvcache( + q, + k_cache, + v_cache, + seqlens_k, + block_table, + alibi_slopes_tensor, + scale); +} + +void py_mha_kvcache_(Tensor out, + Tensor q, + Tensor k_cache, + Tensor v_cache, + Tensor seqlens_k, + Tensor block_table, + pybind11::object alibi_slopes, + float scale) { + std::optional alibi_slopes_tensor = std::nullopt; + if (!alibi_slopes.is_none()) { + alibi_slopes_tensor = alibi_slopes.cast(); + } + + op::mha_kvcache_( + out, + q, + k_cache, + v_cache, + seqlens_k, + block_table, + alibi_slopes_tensor, + scale); +} + +inline void bind_mha_kvcache(py::module &m) { + m.def( + "mha_kvcache", + &ops::py_mha_kvcache, + py::arg("q"), + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("seqlens_k"), + py::arg("block_table"), + py::arg("alibi_slopes"), + py::arg("scale"), + R"doc(Flash attention KV-cache decode for single-step attention over a paged KV cache. + +Parameters +---------- +q : Tensor + Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] +k_cache : Tensor + Key cache tensor of shape [num_blocks, block_size, num_heads_k, head_size] (paged layout) +v_cache : Tensor + Value cache tensor of shape [num_blocks, block_size, num_heads_k, head_size] (paged layout) +seqlens_k : Tensor + Total KV length per request of shape [batch_size] (int32) +block_table : Tensor + Block mapping table of shape [batch_size, max_num_blocks_per_seq] (int32) +alibi_slopes : Optional[Tensor] + ALiBi slopes tensor, if None then ALiBi is disabled +scale : float + Scaling factor for attention scores (typically 1.0/sqrt(head_size)) + +Returns +------- +Tensor + Output tensor of shape [batch_size, seqlen_q, num_heads, head_size] +)doc"); + + m.def( + "mha_kvcache_", + &ops::py_mha_kvcache_, + py::arg("out"), + py::arg("q"), + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("seqlens_k"), + py::arg("block_table"), + py::arg("alibi_slopes"), + py::arg("scale"), + R"doc(In-place flash attention KV-cache decode. + +Parameters +---------- +out : Tensor + Output tensor of shape [batch_size, seqlen_q, num_heads, head_size] +q : Tensor + Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] +k_cache : Tensor + Key cache tensor of shape [num_blocks, block_size, num_heads_k, head_size] (paged layout) +v_cache : Tensor + Value cache tensor of shape [num_blocks, block_size, num_heads_k, head_size] (paged layout) +seqlens_k : Tensor + Total KV length per request of shape [batch_size] (int32) +block_table : Tensor + Block mapping table of shape [batch_size, max_num_blocks_per_seq] (int32) +alibi_slopes : Optional[Tensor] + ALiBi slopes tensor, if None then ALiBi is disabled +scale : float + Scaling factor for attention scores (typically 1.0/sqrt(head_size)) +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/mha_varlen.hpp b/csrc/infinicore/src/pybind11/ops/mha_varlen.hpp new file mode 100644 index 000000000..b1b463a93 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/mha_varlen.hpp @@ -0,0 +1,110 @@ +#pragma once + +#include + +#include "infinicore/ops/mha_varlen.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_mha_varlen(Tensor q, + Tensor k, + Tensor v, + Tensor cum_seqlens_q, + Tensor cum_seqlens_k, + pybind11::object block_table, + int max_seqlen_q, + int max_seqlen_k, + pybind11::object alibi_slopes, + float scale) { + std::optional block_table_tensor = std::nullopt; + if (!block_table.is_none()) { + block_table_tensor = block_table.cast(); + } + std::optional alibi_slopes_tensor = std::nullopt; + if (!alibi_slopes.is_none()) { + alibi_slopes_tensor = alibi_slopes.cast(); + } + + return op::mha_varlen( + q, + k, + v, + cum_seqlens_q, + cum_seqlens_k, + block_table_tensor, + max_seqlen_q, + max_seqlen_k, + alibi_slopes_tensor, + scale); +} + +void py_mha_varlen_(Tensor out, + Tensor q, + Tensor k, + Tensor v, + Tensor cum_seqlens_q, + Tensor cum_seqlens_k, + pybind11::object block_table, + int max_seqlen_q, + int max_seqlen_k, + pybind11::object alibi_slopes, + float scale) { + std::optional block_table_tensor = std::nullopt; + if (!block_table.is_none()) { + block_table_tensor = block_table.cast(); + } + std::optional alibi_slopes_tensor = std::nullopt; + if (!alibi_slopes.is_none()) { + alibi_slopes_tensor = alibi_slopes.cast(); + } + + op::mha_varlen_( + out, + q, + k, + v, + cum_seqlens_q, + cum_seqlens_k, + block_table_tensor, + max_seqlen_q, + max_seqlen_k, + alibi_slopes_tensor, + scale); +} + +inline void bind_mha_varlen(py::module &m) { + m.def( + "mha_varlen", + &ops::py_mha_varlen, + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("cum_seqlens_q"), + py::arg("cum_seqlens_k"), + py::arg("block_table"), + py::arg("max_seqlen_q"), + py::arg("max_seqlen_k"), + py::arg("alibi_slopes"), + py::arg("scale"), + R"doc(Variable-length multi-head attention.)doc"); + + m.def( + "mha_varlen_", + &ops::py_mha_varlen_, + py::arg("out"), + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("cum_seqlens_q"), + py::arg("cum_seqlens_k"), + py::arg("block_table"), + py::arg("max_seqlen_q"), + py::arg("max_seqlen_k"), + py::arg("alibi_slopes"), + py::arg("scale"), + R"doc(In-place variable-length multi-head attention.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/moe_topk_softmax.hpp b/csrc/infinicore/src/pybind11/ops/moe_topk_softmax.hpp new file mode 100644 index 000000000..5466fcd79 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/moe_topk_softmax.hpp @@ -0,0 +1,57 @@ +#pragma once + +#include + +#include "infinicore/ops/moe_topk_softmax.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +std::tuple py_moe_topk_softmax(Tensor gating_output, + size_t topk, + bool renormalize, + float moe_softcapping, + py::object correction_bias) { + Tensor bias; + if (!correction_bias.is_none()) { + bias = correction_bias.cast(); + } + return op::moe_topk_softmax(gating_output, topk, renormalize, moe_softcapping, bias); +} + +void py_moe_topk_softmax_(Tensor topk_weights, + Tensor topk_indices, + Tensor gating_output, + py::object correction_bias, + bool renormalize, + float moe_softcapping) { + Tensor bias; + if (!correction_bias.is_none()) { + bias = correction_bias.cast(); + } + op::moe_topk_softmax_(topk_weights, topk_indices, gating_output, bias, renormalize, moe_softcapping); +} + +inline void bind_moe_topk_softmax(py::module &m) { + m.def("moe_topk_softmax", + &py_moe_topk_softmax, + py::arg("gating_output"), + py::arg("topk"), + py::arg("renormalize") = false, + py::arg("moe_softcapping") = 0.0f, + py::arg("correction_bias") = py::none(), + R"doc(MoE top-k softmax.)doc"); + + m.def("moe_topk_softmax_", + &py_moe_topk_softmax_, + py::arg("topk_weights"), + py::arg("topk_indices"), + py::arg("gating_output"), + py::arg("correction_bias") = py::none(), + py::arg("renormalize") = false, + py::arg("moe_softcapping") = 0.0f, + R"doc(In-place MoE top-k softmax.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/mrope.hpp b/csrc/infinicore/src/pybind11/ops/mrope.hpp new file mode 100644 index 000000000..765584880 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/mrope.hpp @@ -0,0 +1,46 @@ +#pragma once + +#include +#include + +#include "infinicore/ops/mrope.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_mrope(py::module &m) { + m.def("mrope", + &op::mrope, + py::arg("q"), + py::arg("k"), + py::arg("cos"), + py::arg("sin"), + py::arg("positions"), + py::arg("head_size"), + py::arg("rotary_dim"), + py::arg("section_t"), + py::arg("section_h"), + py::arg("section_w"), + py::arg("interleaved"), + R"doc(Multimodal rotary position embedding for q and k.)doc"); + + m.def("mrope_", + &op::mrope_, + py::arg("q_out"), + py::arg("k_out"), + py::arg("q"), + py::arg("k"), + py::arg("cos"), + py::arg("sin"), + py::arg("positions"), + py::arg("head_size"), + py::arg("rotary_dim"), + py::arg("section_t"), + py::arg("section_h"), + py::arg("section_w"), + py::arg("interleaved"), + R"doc(In-place multimodal rotary position embedding for q and k.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/mul.hpp b/csrc/infinicore/src/pybind11/ops/mul.hpp new file mode 100644 index 000000000..fb8e4144b --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/mul.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/mul.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_mul(py::module &m) { + m.def("mul", + &op::mul, + py::arg("a"), + py::arg("b"), + R"doc(Element-wise multiplication of two tensors.)doc"); + + m.def("mul_", + &op::mul_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + R"doc(In-place element-wise tensor multiplication.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/mul_scalar.hpp b/csrc/infinicore/src/pybind11/ops/mul_scalar.hpp new file mode 100644 index 000000000..2bce862bc --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/mul_scalar.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/mul_scalar.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_mul_scalar(py::module &m) { + m.def("mul_scalar", + &op::mul_scalar, + py::arg("a"), + py::arg("alpha"), + R"doc(Multiply a tensor by a host scalar.)doc"); + + m.def("mul_scalar_", + &op::mul_scalar_, + py::arg("c"), + py::arg("a"), + py::arg("alpha"), + R"doc(Out-of-place tensor-scalar multiplication into c.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/multi_margin_loss.hpp b/csrc/infinicore/src/pybind11/ops/multi_margin_loss.hpp new file mode 100644 index 000000000..6593c2e97 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/multi_margin_loss.hpp @@ -0,0 +1,57 @@ +#pragma once + +#include "infinicore/ops/multi_margin_loss.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_multi_margin_loss(py::module &m) { + m.def( + "multi_margin_loss", + [](const Tensor &input, const Tensor &target, py::object weight, int p, float margin, int reduction) { + Tensor weight_tensor; + if (!weight.is_none()) { + weight_tensor = weight.cast(); + } + return op::multi_margin_loss(input, target, weight_tensor, p, margin, reduction); + }, + py::arg("input"), + py::arg("target"), + py::arg("weight") = py::none(), // Python 端看到默认值是 None + py::arg("p") = 1, + py::arg("margin") = 1.0f, + py::arg("reduction") = 1, + R"doc(Computes the Multi Margin Loss between input and target. + + Args: + input (Tensor): Input tensor of shape (N, C). + target (Tensor): Ground truth labels of shape (N,). + weight (Tensor, optional): Manual rescaling weight given to each class. If given, has to be a Tensor of size C. + p (int, optional): The norm degree for pairwise distance. p=1 or p=2. Default: 1. + margin (float, optional): Margin value. Default: 1.0. + reduction (int, optional): Specifies the reduction to apply to the output: 0=None, 1=Mean, 2=Sum. Default: 1. + )doc"); + + m.def( + "multi_margin_loss_", + [](Tensor &output, const Tensor &input, const Tensor &target, py::object weight, int p, float margin, int reduction) { + Tensor weight_tensor; + if (!weight.is_none()) { + weight_tensor = weight.cast(); + } + // 调用底层 + op::multi_margin_loss_(output, input, target, weight_tensor, p, margin, reduction); + }, + py::arg("output"), + py::arg("input"), + py::arg("target"), + py::arg("weight") = py::none(), + py::arg("p") = 1, + py::arg("margin") = 1.0f, + py::arg("reduction") = 1, + R"doc(Explicit output Multi Margin Loss operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/nrm2.hpp b/csrc/infinicore/src/pybind11/ops/nrm2.hpp new file mode 100644 index 000000000..02b21f53b --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/nrm2.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/nrm2.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_nrm2(py::module &m) { + m.def("nrm2", + &op::nrm2, + py::arg("x"), + R"doc(BLAS level-1 nrm2.)doc"); + + m.def("nrm2_", + &op::nrm2_, + py::arg("x"), + py::arg("result"), + R"doc(In-place BLAS level-1 nrm2.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/pad.hpp b/csrc/infinicore/src/pybind11/ops/pad.hpp new file mode 100644 index 000000000..19c090308 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/pad.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include + +#include "infinicore/ops/pad.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_pad(py::module &m) { + m.def("pad", + &op::pad, + py::arg("x"), + py::arg("pad"), + py::arg("mode") = std::string("constant"), + py::arg("value") = 0.0, + R"doc(Pad a tensor (PyTorch padding order).)doc"); + + m.def("pad_", + &op::pad_, + py::arg("y"), + py::arg("x"), + py::arg("pad"), + py::arg("mode") = std::string("constant"), + py::arg("value") = 0.0, + R"doc(Out variant of pad.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/paged_attention.hpp b/csrc/infinicore/src/pybind11/ops/paged_attention.hpp new file mode 100644 index 000000000..ab77c87a4 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/paged_attention.hpp @@ -0,0 +1,53 @@ +#pragma once + +#include + +#include "infinicore/ops/paged_attention.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_paged_attention(Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, pybind11::object alibi_slopes, float scale) { + std::optional alibi_slopes_tensor = std::nullopt; + if (!alibi_slopes.is_none()) { + alibi_slopes_tensor = alibi_slopes.cast(); + } + return op::paged_attention(q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes_tensor, scale); +} + +void py_paged_attention_(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, pybind11::object alibi_slopes, float scale) { + std::optional alibi_slopes_tensor = std::nullopt; + if (!alibi_slopes.is_none()) { + alibi_slopes_tensor = alibi_slopes.cast(); + } + + op::paged_attention_(out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes_tensor, scale); +} + +inline void bind_paged_attention(py::module &m) { + m.def("paged_attention", + &ops::py_paged_attention, + py::arg("q"), + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("block_tables"), + py::arg("cache_lens"), + py::arg("alibi_slopes"), + py::arg("scale"), + R"doc(Paged attention of query and key cache tensors.)doc"); + + m.def("paged_attention_", + &ops::py_paged_attention_, + py::arg("out"), + py::arg("q"), + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("block_tables"), + py::arg("cache_lens"), + py::arg("alibi_slopes"), + py::arg("scale"), + R"doc(In-place paged attention of query and key cache tensors.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/paged_attention_prefill.hpp b/csrc/infinicore/src/pybind11/ops/paged_attention_prefill.hpp new file mode 100644 index 000000000..13b2c2683 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/paged_attention_prefill.hpp @@ -0,0 +1,69 @@ +#pragma once + +#include "infinicore/ops/paged_attention_prefill.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_paged_attention_prefill(Tensor q, + Tensor k_cache, + Tensor v_cache, + Tensor block_tables, + Tensor history_lens, + Tensor cu_seqlens_q, + py::object alibi_slopes, + float scale) { + std::optional alibi_slopes_tensor = std::nullopt; + if (!alibi_slopes.is_none()) { + alibi_slopes_tensor = alibi_slopes.cast(); + } + return op::paged_attention_prefill( + q, k_cache, v_cache, block_tables, history_lens, cu_seqlens_q, alibi_slopes_tensor, scale); +} + +void py_paged_attention_prefill_(Tensor out, + Tensor q, + Tensor k_cache, + Tensor v_cache, + Tensor block_tables, + Tensor history_lens, + Tensor cu_seqlens_q, + py::object alibi_slopes, + float scale) { + std::optional alibi_slopes_tensor = std::nullopt; + if (!alibi_slopes.is_none()) { + alibi_slopes_tensor = alibi_slopes.cast(); + } + op::paged_attention_prefill_(out, q, k_cache, v_cache, block_tables, history_lens, cu_seqlens_q, alibi_slopes_tensor, scale); +} + +inline void bind_paged_attention_prefill(py::module &m) { + m.def("paged_attention_prefill", + &ops::py_paged_attention_prefill, + py::arg("q"), + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("block_tables"), + py::arg("history_lens"), + py::arg("cu_seqlens_q"), + py::arg("alibi_slopes") = py::none(), + py::arg("scale") = 1.0, + R"doc(Paged attention prefill for packed variable-length queries.)doc"); + + m.def("paged_attention_prefill_", + &ops::py_paged_attention_prefill_, + py::arg("out"), + py::arg("q"), + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("block_tables"), + py::arg("history_lens"), + py::arg("cu_seqlens_q"), + py::arg("alibi_slopes") = py::none(), + py::arg("scale") = 1.0, + R"doc(In-place paged attention prefill for packed variable-length queries.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/paged_caching.hpp b/csrc/infinicore/src/pybind11/ops/paged_caching.hpp new file mode 100644 index 000000000..4320b4eef --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/paged_caching.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include + +#include "infinicore/ops/paged_caching.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_paged_caching(py::module &m) { + m.def("paged_caching_", + &op::paged_caching_, + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("k"), + py::arg("v"), + py::arg("slot_mapping"), + R"doc(Paged caching of key and value tensors.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/per_channel_quant_i8.hpp b/csrc/infinicore/src/pybind11/ops/per_channel_quant_i8.hpp new file mode 100644 index 000000000..da6f9f592 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/per_channel_quant_i8.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include + +#include "infinicore/ops/per_channel_quant_i8.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_per_channel_quant_i8(py::module &m) { + m.def("per_channel_quant_i8_", + &op::per_channel_quant_i8_, + py::arg("x"), + py::arg("x_packed"), + py::arg("x_scale"), + R"doc(Per-channel quantization of a tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/prelu.hpp b/csrc/infinicore/src/pybind11/ops/prelu.hpp new file mode 100644 index 000000000..ee195c52f --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/prelu.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/prelu.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_prelu(py::module &m) { + m.def("prelu", + &op::prelu, + py::arg("input"), + py::arg("weight"), + R"doc(Parametric ReLU.)doc"); + + m.def("prelu_", + &op::prelu_, + py::arg("out"), + py::arg("input"), + py::arg("weight"), + R"doc(In-place Parametric ReLU (writes to out).)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/random_sample.hpp b/csrc/infinicore/src/pybind11/ops/random_sample.hpp new file mode 100644 index 000000000..e5beb9e4e --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/random_sample.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include + +#include "infinicore/ops/random_sample.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_random_sample(py::module &m) { + m.def("random_sample", + &op::random_sample, + py::arg("logits"), + py::arg("random_val"), + py::arg("topp"), + py::arg("topk"), + py::arg("temperature"), + R"doc(Random sampling: returns an int32 scalar index.)doc"); + + m.def("random_sample_", + &op::random_sample_, + py::arg("indices"), + py::arg("logits"), + py::arg("random_val"), + py::arg("topp"), + py::arg("topk"), + py::arg("temperature"), + R"doc(In-place random sampling into provided int32 scalar tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/rearrange.hpp b/csrc/infinicore/src/pybind11/ops/rearrange.hpp new file mode 100644 index 000000000..816b00079 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/rearrange.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/rearrange.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_rearrange(py::module &m) { + m.def("rearrange", + &op::rearrange, + py::arg("x"), + R"doc(Matrix rearrangement of a tensor.)doc"); + + m.def("rearrange_", + &op::rearrange_, + py::arg("y"), + py::arg("x"), + R"doc(In-place tensor rearrangement.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/reciprocal.hpp b/csrc/infinicore/src/pybind11/ops/reciprocal.hpp new file mode 100644 index 000000000..3e49205f6 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/reciprocal.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/reciprocal.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_reciprocal(py::module &m) { + m.def("reciprocal", + &op::reciprocal, + py::arg("x"), + R"doc(Computes the reciprocal of the input tensor.)doc"); + + m.def("reciprocal_", + &op::reciprocal_, + py::arg("y"), + py::arg("x"), + R"doc(Computes the reciprocal of the input tensor and stores in the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/recurrent_gated_delta_rule.hpp b/csrc/infinicore/src/pybind11/ops/recurrent_gated_delta_rule.hpp new file mode 100644 index 000000000..127b62042 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/recurrent_gated_delta_rule.hpp @@ -0,0 +1,37 @@ +#pragma once + +#include + +#include "infinicore/ops/recurrent_gated_delta_rule.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_recurrent_gated_delta_rule(py::module &m) { + m.def("recurrent_gated_delta_rule", + &op::recurrent_gated_delta_rule, + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("g"), + py::arg("beta"), + py::arg("initial_state"), + py::arg("use_qk_l2norm") = false, + R"doc(Recurrent gated delta rule. Returns out only.)doc"); + + m.def("recurrent_gated_delta_rule_indexed", + &op::recurrent_gated_delta_rule_indexed, + py::arg("q"), + py::arg("k"), + py::arg("v"), + py::arg("g"), + py::arg("beta"), + py::arg("initial_state"), + py::arg("initial_state_indices"), + py::arg("final_state_indices"), + py::arg("use_qk_l2norm") = false, + R"doc(Recurrent gated delta rule with indexed in-place state pool. Returns out only.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/relu6.hpp b/csrc/infinicore/src/pybind11/ops/relu6.hpp new file mode 100644 index 000000000..e69e3568a --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/relu6.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/relu6.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_relu6(py::module &m) { + m.def("relu6", + &op::relu6, + py::arg("input"), + R"doc(ReLU6 activation.)doc"); + + m.def("relu6_", + &op::relu6_, + py::arg("out"), + py::arg("input"), + R"doc(In-place ReLU6 activation (writes to out).)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/rms_norm.hpp b/csrc/infinicore/src/pybind11/ops/rms_norm.hpp new file mode 100644 index 000000000..1fd899c44 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/rms_norm.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include + +#include "infinicore/ops/rms_norm.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_rms_norm(py::module &m) { + m.def("rms_norm", + &op::rms_norm, + py::arg("x"), + py::arg("weight"), + py::arg("epsilon") = 1e-5f, + R"doc(Root Mean Square Normalization. + +Args: + x: Input tensor + weight: Scale weights + epsilon: Small constant for numerical stability, default is 1e-5 + +Returns: + Normalized tensor with same shape as input +)doc"); + + m.def("rms_norm_", + &op::rms_norm_, + py::arg("y"), + py::arg("x"), + py::arg("weight"), + py::arg("epsilon") = 1e-5f, + R"doc(In-place Root Mean Square Normalization. + +Args: + y: Output tensor + x: Input tensor + weight: Scale weights + epsilon: Small constant for numerical stability, default is 1e-5 +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/rope.hpp b/csrc/infinicore/src/pybind11/ops/rope.hpp new file mode 100644 index 000000000..a1836955a --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/rope.hpp @@ -0,0 +1,37 @@ +#pragma once + +#include + +#include "infinicore/ops/rope.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_rope(py::module &m) { + + py::enum_(m, "RoPEAlgo") + .value("GPT_J", infinicore::nn::RoPE::Algo::GPT_J) + .value("GPT_NEOX", infinicore::nn::RoPE::Algo::GPT_NEOX); + + m.def("rope", + &op::rope, + py::arg("x"), + py::arg("pos"), + py::arg("sin_table"), + py::arg("cos_table"), + py::arg("algo"), + R"doc( Rotary Position Embedding(RoPE).)doc"); + + m.def("rope_", + &op::rope_, + py::arg("x_out"), + py::arg("x"), + py::arg("pos"), + py::arg("sin_table"), + py::arg("cos_table"), + py::arg("algo"), + R"doc(In-place, Rotary Position Embedding(RoPE).)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/rot.hpp b/csrc/infinicore/src/pybind11/ops/rot.hpp new file mode 100644 index 000000000..359cb9745 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/rot.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include + +#include "infinicore/ops/rot.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_rot(py::module &m) { + m.def("rot_", + &op::rot_, + py::arg("x"), + py::arg("y"), + py::arg("c"), + py::arg("s"), + R"doc(In-place BLAS level-1 rot, updating x and y.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/rotg.hpp b/csrc/infinicore/src/pybind11/ops/rotg.hpp new file mode 100644 index 000000000..a37e79336 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/rotg.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include + +#include "infinicore/ops/rotg.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_rotg(py::module &m) { + m.def("rotg_", + py::overload_cast(&op::rotg_), + py::arg("x"), + py::arg("y"), + py::arg("c"), + py::arg("s"), + R"doc(In-place BLAS level-1 rotg, updating x, y, c, and s.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/rotm.hpp b/csrc/infinicore/src/pybind11/ops/rotm.hpp new file mode 100644 index 000000000..a88db38ba --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/rotm.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include + +#include "infinicore/ops/rotm.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_rotm(py::module &m) { + m.def("rotm_", + py::overload_cast(&op::rotm_), + py::arg("x"), + py::arg("y"), + py::arg("param"), + R"doc(In-place BLAS level-1 rotm, updating x and y.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/rotmg.hpp b/csrc/infinicore/src/pybind11/ops/rotmg.hpp new file mode 100644 index 000000000..72b816751 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/rotmg.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include + +#include "infinicore/ops/rotmg.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_rotmg(py::module &m) { + m.def("rotmg_", + py::overload_cast(&op::rotmg_), + py::arg("d1"), + py::arg("d2"), + py::arg("x1"), + py::arg("y1"), + py::arg("param"), + R"doc(In-place BLAS level-1 rotmg, updating d1, d2, x1, and param.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/rwkv5_wkv.hpp b/csrc/infinicore/src/pybind11/ops/rwkv5_wkv.hpp new file mode 100644 index 000000000..ec59316d0 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/rwkv5_wkv.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include + +#include "infinicore/ops/rwkv5_wkv.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_rwkv5_wkv(py::module &m) { + m.def("rwkv5_wkv", + &op::rwkv5_wkv, + py::arg("receptance"), + py::arg("key"), + py::arg("value"), + py::arg("time_decay"), + py::arg("time_faaaa"), + py::arg("state"), + R"doc(RWKV5 weighted key-value recurrence. Updates state in-place and returns output.)doc"); + + m.def("rwkv5_wkv_", + &op::rwkv5_wkv_, + py::arg("out"), + py::arg("receptance"), + py::arg("key"), + py::arg("value"), + py::arg("time_decay"), + py::arg("time_faaaa"), + py::arg("state"), + R"doc(Explicit-output RWKV5 weighted key-value recurrence. Updates out and state in-place.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/scal.hpp b/csrc/infinicore/src/pybind11/ops/scal.hpp new file mode 100644 index 000000000..75a2c5ca8 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/scal.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include + +#include "infinicore/ops/scal.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_scal(py::module &m) { + m.def("scal_", + &op::scal_, + py::arg("alpha"), + py::arg("x"), + R"doc(In-place BLAS level-1 scal, updating x.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/scaled_mm_i8.hpp b/csrc/infinicore/src/pybind11/ops/scaled_mm_i8.hpp new file mode 100644 index 000000000..c3d46d9df --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/scaled_mm_i8.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include + +#include "infinicore/ops/scaled_mm_i8.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_scaled_mm_i8(py::module &m) { + m.def("scaled_mm_i8", + &op::scaled_mm_i8, + py::arg("a_p"), + py::arg("a_s"), + py::arg("b_p"), + py::arg("b_s"), + py::arg("bias"), + R"doc(Scaled matrix multiplication of two tensors.)doc"); + + m.def("scaled_mm_i8_", + &op::scaled_mm_i8_, + py::arg("a"), + py::arg("b"), + py::arg("a_scale"), + py::arg("b_scale"), + R"doc(In-place Scaled matrix multiplication of two tensors.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/scatter.hpp b/csrc/infinicore/src/pybind11/ops/scatter.hpp new file mode 100644 index 000000000..950ed49e8 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/scatter.hpp @@ -0,0 +1,56 @@ +#pragma once + +#include "infinicore/ops/scatter.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_scatter(py::module &m) { + // ========================================================================= + // 1. 绑定 out-of-place 接口: scatter + // ========================================================================= + // 为了匹配测试脚本的行为（将所有 Tensor 作为位置参数传入，属性作为 kwargs 传入）， + // 我们将参数顺序调整为: input, index, src, dim, reduction + // ========================================================================= + m.def( + "scatter", + [](const Tensor &input, const Tensor &index, const Tensor &src, int64_t dim, int64_t reduction) { + // 调用底层 C++ 实现时，必须恢复正确的参数顺序: (input, dim, index, src, reduction) + return op::scatter(input, dim, index, src, reduction); + }, + py::arg("input"), + py::arg("index"), + py::arg("src"), + py::arg("dim"), // 关键修改：将 dim 移到 Tensor 参数之后 + py::arg("reduction") = 0, + R"doc( + Scatter operator. + Note: Parameter order in this binding is adapted for the test runner: (input, index, src, dim, reduction). + )doc"); + + // ========================================================================= + // 2. 绑定 in-place 接口: scatter_ + // ========================================================================= + // 参数顺序调整为: output, input, index, src, dim, reduction + // ========================================================================= + m.def( + "scatter_", + [](Tensor &output, const Tensor &input, const Tensor &index, const Tensor &src, int64_t dim, int64_t reduction) { + // 调用底层 C++ 实现 + op::scatter_(output, input, dim, index, src, reduction); + }, + py::arg("output"), + py::arg("input"), + py::arg("index"), + py::arg("src"), + py::arg("dim"), // 关键修改：将 dim 移到 Tensor 参数之后 + py::arg("reduction") = 0, + R"doc( + In-place Scatter operator. + Writes result into output. + )doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/selu.hpp b/csrc/infinicore/src/pybind11/ops/selu.hpp new file mode 100644 index 000000000..ea7ffe553 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/selu.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/selu.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_selu(py::module &m) { + m.def( + "selu", + &op::selu, + py::arg("input"), + R"doc(SELU activation function.)doc"); + + m.def( + "selu_", + &op::selu_, + py::arg("output"), + py::arg("input"), + R"doc(In-place SELU activation function.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/sigmoid.hpp b/csrc/infinicore/src/pybind11/ops/sigmoid.hpp new file mode 100644 index 000000000..8f24c0c8e --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/sigmoid.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/sigmoid.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_sigmoid(py::module &m) { + m.def("sigmoid", + &op::sigmoid, + py::arg("input"), + R"doc(Out-of-place Sigmoid activation.)doc"); + + m.def("sigmoid_", + &op::sigmoid_, + py::arg("output"), + py::arg("input"), + R"doc(In-place Sigmoid activation.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/silu.hpp b/csrc/infinicore/src/pybind11/ops/silu.hpp new file mode 100644 index 000000000..52aa7c3a8 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/silu.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/silu.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_silu(py::module &m) { + m.def("silu", + &op::silu, + py::arg("input"), + R"doc(SiLU (Swish) activation function.)doc"); + + m.def("silu_", + &op::silu_, + py::arg("output"), + py::arg("input"), + R"doc(In-place SiLU (Swish) activation function.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/silu_and_mul.hpp b/csrc/infinicore/src/pybind11/ops/silu_and_mul.hpp new file mode 100644 index 000000000..009c3b533 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/silu_and_mul.hpp @@ -0,0 +1,29 @@ +#pragma once + +#include + +#include "infinicore/ops/silu_and_mul.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_silu_and_mul(py::module &m) { + m.def("silu_and_mul", + &op::silu_and_mul, + py::arg("input"), + R"doc( + SiLU and Mul (SwiGLU) activation function. + Input should be [..., 2*d], output will be [..., d]. + )doc"); + + m.def("silu_and_mul_", + &op::silu_and_mul_, + py::arg("output"), + py::arg("input"), + R"doc( + In-place or destination-specified SiLU and Mul (SwiGLU) activation function. + )doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/sinh.hpp b/csrc/infinicore/src/pybind11/ops/sinh.hpp new file mode 100644 index 000000000..69ff34e99 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/sinh.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/sinh.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_sinh(py::module &m) { + m.def( + "sinh", + &op::sinh, + py::arg("input"), + R"doc(Sinh activation function.)doc"); + + m.def( + "sinh_", + &op::sinh_, + py::arg("output"), + py::arg("input"), + R"doc(In-place Sinh activation function.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/smooth_l1_loss.hpp b/csrc/infinicore/src/pybind11/ops/smooth_l1_loss.hpp new file mode 100644 index 000000000..2408fc7b1 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/smooth_l1_loss.hpp @@ -0,0 +1,36 @@ +#pragma once + +#include "infinicore/ops/smooth_l1_loss.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_smooth_l1_loss(py::module &m) { + // 1. 绑定 out-of-place 接口: output = smooth_l1_loss(input, target, beta, reduction) + m.def("smooth_l1_loss", + &op::smooth_l1_loss, + py::arg("input"), + py::arg("target"), + py::arg("beta") = 1.0f, + py::arg("reduction") = 1, + R"doc(Computes the Smooth L1 Loss between input and target. + + Args: + input (Tensor): Predicted values. + target (Tensor): Ground truth values. + beta (float, optional): The threshold at which to change between L1 and L2 loss. Default: 1.0. + reduction (int, optional): Specifies the reduction to apply to the output: 0=None, 1=Mean, 2=Sum. Default: 1. + )doc"); + m.def("smooth_l1_loss_", + &op::smooth_l1_loss_, + py::arg("output"), + py::arg("input"), + py::arg("target"), + py::arg("beta") = 1.0f, + py::arg("reduction") = 1, + R"doc(Explicit output Smooth L1 Loss operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/softplus.hpp b/csrc/infinicore/src/pybind11/ops/softplus.hpp new file mode 100644 index 000000000..768f99f38 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/softplus.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include + +#include "infinicore/ops/softplus.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_softplus(py::module &m) { + // Functional interface: returns a new Tensor + m.def("softplus", + &op::softplus, + py::arg("x"), + py::arg("beta") = 1.0f, + py::arg("threshold") = 20.0f, + R"doc(Computes the softplus function element-wise: y = 1/beta * log(1 + exp(beta * x)).)doc"); + + // In-place/Out-variant interface: writes to provided output Tensor + m.def("softplus_", + &op::softplus_, + py::arg("y"), + py::arg("x"), + py::arg("beta") = 1.0f, + py::arg("threshold") = 20.0f, + R"doc(In-place softplus activation. Writes result into y.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/softsign.hpp b/csrc/infinicore/src/pybind11/ops/softsign.hpp new file mode 100644 index 000000000..1191abe02 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/softsign.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/softsign.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_softsign(py::module &m) { + m.def("softsign", + &op::softsign, + py::arg("x"), + R"doc(Softsign activation function.)doc"); + + m.def("softsign_", + &op::softsign_, + py::arg("y"), + py::arg("x"), + R"doc(In-place softsign activation.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/sum.hpp b/csrc/infinicore/src/pybind11/ops/sum.hpp new file mode 100644 index 000000000..50fef7539 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/sum.hpp @@ -0,0 +1,60 @@ +#pragma once + +#include + +#include "infinicore/ops/sum.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_sum(Tensor input, py::object dim, bool keepdim) { + if (dim.is_none()) { + std::vector dim_vec; + for (int i = 0; i < input->shape().size(); i++) { + dim_vec.push_back(i); + } + return op::sum(input, dim_vec, keepdim); + } else if (py::isinstance(dim) || py::isinstance(dim)) { + return op::sum(input, dim.cast>(), keepdim); + } else if (py::isinstance(dim)) { + return op::sum(input, std::vector(1, dim.cast()), keepdim); + } else { + throw std::invalid_argument("dim must be a tuple or an integer"); + } +} + +void py_sum_(Tensor output, Tensor input, py::object dim, bool keepdim) { + if (dim.is_none()) { + std::vector dim_vec; + for (int i = 0; i < input->shape().size(); i++) { + dim_vec.push_back(i); + } + op::sum_(output, input, dim_vec, keepdim); + } else if (py::isinstance(dim) || py::isinstance(dim)) { + op::sum_(output, input, dim.cast>(), keepdim); + } else if (py::isinstance(dim)) { + op::sum_(output, input, std::vector(1, dim.cast()), keepdim); + } else { + throw std::invalid_argument("dim must be a tuple or an integer"); + } +} + +inline void bind_sum(py::module &m) { + m.def("sum", + &py_sum, + py::arg("input"), + py::arg("dim"), + py::arg("keepdim"), + R"doc(Sum of input tensor along the given dimensions.)doc"); + + m.def("sum_", + &py_sum_, + py::arg("output"), + py::arg("input"), + py::arg("dim"), + py::arg("keepdim"), + R"doc(In-place tensor sum.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/swap.hpp b/csrc/infinicore/src/pybind11/ops/swap.hpp new file mode 100644 index 000000000..0b8f2ae3e --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/swap.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include + +#include "infinicore/ops/swap.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_swap(py::module &m) { + m.def("swap_", + &op::swap_, + py::arg("x"), + py::arg("y"), + R"doc(In-place BLAS level-1 swap.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/swiglu.hpp b/csrc/infinicore/src/pybind11/ops/swiglu.hpp new file mode 100644 index 000000000..a5c98609f --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/swiglu.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include + +#include "infinicore/ops/swiglu.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_swiglu(py::module &m) { + m.def("swiglu", + &op::swiglu, + py::arg("a"), + py::arg("b"), + R"doc(SwiGLU activation function.)doc"); + + m.def("swiglu_", + &op::swiglu_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + R"doc(In-place SwiGLU activation function.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/take.hpp b/csrc/infinicore/src/pybind11/ops/take.hpp new file mode 100644 index 000000000..cd040f4ed --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/take.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include "infinicore/ops/take.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_take(py::module &m) { + m.def("take", + &op::take, + py::arg("input"), + py::arg("indices"), + R"doc(Extracts elements from the input tensor along the given indices. +The input tensor is treated as a flattened 1D array.)doc"); + m.def("take_", + &op::take_, + py::arg("output"), + py::arg("input"), + py::arg("indices"), + R"doc(Explicit output take operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/tan.hpp b/csrc/infinicore/src/pybind11/ops/tan.hpp new file mode 100644 index 000000000..acd62e048 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/tan.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/tan.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_tan(py::module &m) { + m.def("tan", + &op::tan, + py::arg("input"), + R"doc(opertor: torch.tan, out-of-place mode)doc"); + + m.def("tan_", + &op::tan_, + py::arg("output"), + py::arg("input"), + R"doc(opertor: torch.tan, in-place mode)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/tanhshrink.hpp b/csrc/infinicore/src/pybind11/ops/tanhshrink.hpp new file mode 100644 index 000000000..3780555f4 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/tanhshrink.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include "infinicore/ops/tanhshrink.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_tanhshrink(py::module &m) { + m.def("tanhshrink", + &op::tanhshrink, + py::arg("input"), + R"doc(opertor: torch.tanhshrink, out-of-place mode)doc"); + + m.def("tanhshrink_", + &op::tanhshrink_, + py::arg("output"), + py::arg("input"), + R"doc(opertor: torch.tanhshrink, in-place mode)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/topk.hpp b/csrc/infinicore/src/pybind11/ops/topk.hpp new file mode 100644 index 000000000..1341f39fa --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/topk.hpp @@ -0,0 +1,54 @@ +#pragma once + +#include +#include // 添加这行 + +#include "infinicore/ops/topk.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +std::pair py_topk(Tensor input, size_t k, int dim, bool largest, bool sorted) { + if (dim == -1) { + return op::topk(input, k, input->ndim() - 1, largest, sorted); + } else if (dim >= 0) { + return op::topk(input, k, static_cast(dim), largest, sorted); + } else { + throw std::invalid_argument("invalid argument: dim"); + } +} + +void py_topk_(Tensor values_output, Tensor indices_output, Tensor input, size_t k, int dim, bool largest, bool sorted) { + if (dim == -1) { + op::topk_(values_output, indices_output, input, k, input->ndim() - 1, largest, sorted); + } else if (dim >= 0) { + op::topk_(values_output, indices_output, input, k, static_cast(dim), largest, sorted); + } else { + throw std::invalid_argument("invalid argument: dim"); + } +} + +inline void bind_topk(py::module &m) { + m.def("topk", + &py_topk, + py::arg("input"), + py::arg("k"), + py::arg("dim"), + py::arg("largest"), + py::arg("sorted"), + R"doc(topk of input tensor along the given dimensions.)doc"); + + m.def("topk_", + &py_topk_, + py::arg("values_output"), + py::arg("indices_output"), + py::arg("input"), + py::arg("k"), + py::arg("dim"), + py::arg("largest"), + py::arg("sorted"), + R"doc(In-place tensor topk_.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/topksoftmax.hpp b/csrc/infinicore/src/pybind11/ops/topksoftmax.hpp new file mode 100644 index 000000000..ca93edff0 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/topksoftmax.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include + +#include "infinicore/ops/topksoftmax.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_topksoftmax(py::module &m) { + m.def("topksoftmax", + &op::topksoftmax, + py::arg("values"), + py::arg("indices"), + py::arg("x"), + py::arg("topk"), + py::arg("norm") = 0, + R"doc(In-place Top-k Softmax. + +Writes results to pre-allocated values and indices tensors. + +Args: + values: Output tensor for softmax weights [N, topk] + indices: Output tensor for selected indices [N, topk], int32 + x: Input tensor [N, width], router logits + topk: Number of top values to select + norm: Whether to re-normalize top-k probabilities (1=yes, 0=no), default 0 +)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/triplet_margin_loss.hpp b/csrc/infinicore/src/pybind11/ops/triplet_margin_loss.hpp new file mode 100644 index 000000000..f003bac94 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/triplet_margin_loss.hpp @@ -0,0 +1,50 @@ +#pragma once + +#include "infinicore/ops/triplet_margin_loss.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_triplet_margin_loss(py::module &m) { + // 1. 绑定 functional 接口: output = triplet_margin_loss(anchor, positive, negative, ...) + m.def("triplet_margin_loss", + &op::triplet_margin_loss, + py::arg("anchor"), + py::arg("positive"), + py::arg("negative"), + py::arg("margin") = 1.0f, + py::arg("p") = 2, + py::arg("eps") = 1e-6f, + py::arg("swap") = false, + py::arg("reduction") = 1, + R"doc(Computes the triplet margin loss. + + Args: + anchor (Tensor): The anchor tensor. + positive (Tensor): The positive tensor. + negative (Tensor): The negative tensor. + margin (float): Default: 1.0. + p (int): The norm degree for pairwise distance. Default: 2. + eps (float): Small constant for numerical stability. Default: 1e-6. + swap (bool): The distance swap is described in the paper Learning shallow convolutional feature descriptors with triplet losses. Default: False. + reduction (int): Specifies the reduction to apply to the output: 0 (none), 1 (mean), 2 (sum). Default: 1. + )doc"); + + // 2. 绑定 explicit output 接口: triplet_margin_loss_(output, anchor, positive, negative, ...) + m.def("triplet_margin_loss_", + &op::triplet_margin_loss_, + py::arg("output"), + py::arg("anchor"), + py::arg("positive"), + py::arg("negative"), + py::arg("margin") = 1.0f, + py::arg("p") = 2, + py::arg("eps") = 1e-6f, + py::arg("swap") = false, + py::arg("reduction") = 1, + R"doc(Explicit output TripletMarginLoss operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/triplet_margin_with_distance_loss.hpp b/csrc/infinicore/src/pybind11/ops/triplet_margin_with_distance_loss.hpp new file mode 100644 index 000000000..e4b20208a --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/triplet_margin_with_distance_loss.hpp @@ -0,0 +1,41 @@ +#pragma once + +#include "infinicore/ops/triplet_margin_with_distance_loss.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_triplet_margin_with_distance_loss(py::module &m) { + m.def("triplet_margin_with_distance_loss", + &op::triplet_margin_with_distance_loss, + py::arg("anchor"), + py::arg("positive"), + py::arg("negative"), + py::arg("margin") = 1.0, + py::arg("swap") = false, + py::arg("reduction") = 1, + R"doc(Computes the triplet margin loss with distance. + + Args: + anchor (Tensor): The anchor input tensor. + positive (Tensor): The positive input tensor. + negative (Tensor): The negative input tensor. + margin (float, optional): Default: 1.0. + swap (bool, optional): The distance swap is described in the paper Learning shallow convolutional feature descriptors with triplet losses. Default: False. + reduction (int, optional): Specifies the reduction to apply to the output: 0 (None), 1 (Mean), 2 (Sum). Default: 1. + )doc"); + m.def("triplet_margin_with_distance_loss_", + &op::triplet_margin_with_distance_loss_, + py::arg("output"), + py::arg("anchor"), + py::arg("positive"), + py::arg("negative"), + py::arg("margin"), + py::arg("swap"), + py::arg("reduction"), + R"doc(Explicit output TripletMarginWithDistanceLoss operation. Writes results into output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/unfold.hpp b/csrc/infinicore/src/pybind11/ops/unfold.hpp new file mode 100644 index 000000000..b09d2ff70 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/unfold.hpp @@ -0,0 +1,61 @@ +#pragma once + +#include "infinicore/ops/unfold.hpp" +#include +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_unfold(py::module &m) { + // ------------------------------------------------------------------------- + // 1. 绑定函数式接口 (unfold) + // ------------------------------------------------------------------------- + m.def( + "unfold", + [](const Tensor &input, + std::vector kernel_sizes, + std::vector dilations, + std::vector paddings, + std::vector strides) { + return op::unfold(input, kernel_sizes, dilations, paddings, strides); + }, + py::arg("input"), + py::arg("kernel_sizes"), + py::arg("dilations"), + py::arg("paddings"), + py::arg("strides"), + R"doc(Extracts sliding local blocks from a batched input tensor. + + Args: + input (Tensor): The input tensor. + kernel_sizes (List[int]): The size of the sliding blocks. + dilations (List[int]): The parameter that controls the stride of elements within the neighborhood. + paddings (List[int]): Implicit zero padding to be added on both sides of input. + strides (List[int]): The stride of the sliding blocks. + )doc"); + + // ------------------------------------------------------------------------- + // 2. 绑定 in-place 接口 (unfold_) + // ------------------------------------------------------------------------- + m.def( + "unfold_", + [](Tensor &output, + const Tensor &input, + std::vector kernel_sizes, + std::vector dilations, + std::vector paddings, + std::vector strides) { + op::unfold_(output, input, kernel_sizes, dilations, paddings, strides); + }, + py::arg("output"), + py::arg("input"), + py::arg("kernel_sizes"), + py::arg("dilations"), + py::arg("paddings"), + py::arg("strides"), + R"doc(Explicit output Unfold operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/upsample_bilinear.hpp b/csrc/infinicore/src/pybind11/ops/upsample_bilinear.hpp new file mode 100644 index 000000000..4a7619ffa --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/upsample_bilinear.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include "infinicore/ops/upsample_bilinear.hpp" +#include +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_upsample_bilinear(py::module &m) { + // 1. 绑定 functional 接口: output = upsample_bilinear(input, output_size, align_corners) + m.def("upsample_bilinear", + &op::upsample_bilinear, + py::arg("input"), + py::arg("output_size"), + py::arg("align_corners") = false, + R"doc(Upsample the input using bilinear interpolation. + + Args: + input (Tensor): The input tensor. + output_size (List[int]): The output spatial size (e.g. [H_out, W_out]). + align_corners (bool): If True, the corner pixels of the input and output tensors are aligned. + )doc"); + + // 2. 绑定 explicit output 接口: upsample_bilinear_(output, input, align_corners) + m.def("upsample_bilinear_", + &op::upsample_bilinear_, + py::arg("output"), + py::arg("input"), + py::arg("align_corners") = false, + R"doc(Explicit output UpsampleBilinear operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/upsample_nearest.hpp b/csrc/infinicore/src/pybind11/ops/upsample_nearest.hpp new file mode 100644 index 000000000..7e4154e9b --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/upsample_nearest.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include "infinicore/ops/upsample_nearest.hpp" +#include +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_upsample_nearest(py::module &m) { + // 1. 绑定 functional 接口: output = upsample_nearest(input, output_size) + m.def("upsample_nearest", + &op::upsample_nearest, + py::arg("input"), + py::arg("output_size"), + R"doc(Upsample the input using nearest neighbor interpolation. + + Args: + input (Tensor): The input tensor. + output_size (List[int]): The output spatial size (e.g. [H_out, W_out]). + )doc"); + + // 2. 绑定 explicit output 接口: upsample_nearest_(output, input) + m.def("upsample_nearest_", + &op::upsample_nearest_, + py::arg("output"), + py::arg("input"), + R"doc(Explicit output UpsampleNearest operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/vander.hpp b/csrc/infinicore/src/pybind11/ops/vander.hpp new file mode 100644 index 000000000..89a851f08 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/vander.hpp @@ -0,0 +1,42 @@ +#pragma once + +#include "infinicore/ops/vander.hpp" +#include + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_vander(py::module &m) { + m.def( + "vander", + [](const Tensor &input, int64_t N, bool increasing) { + return op::vander(input, N, increasing); + }, + py::arg("input"), + py::arg("N") = 0, + py::arg("increasing") = false, + R"doc(Generates a Vandermonde matrix. + + Args: + input (Tensor): 1-D input tensor. + N (int, optional): Number of columns in the output. If 0, defaults to input size (square matrix). Default: 0. + increasing (bool, optional): Order of the powers. If True, powers increase (x^0, x^1...). Default: False. + )doc"); + + // ------------------------------------------------------------------------- + // 2. 绑定 in-place 接口 (vander_) + // ------------------------------------------------------------------------- + m.def( + "vander_", + [](Tensor &output, const Tensor &input, int64_t N, bool increasing) { + op::vander_(output, input, N, increasing); + }, + py::arg("output"), + py::arg("input"), + py::arg("N") = 0, + py::arg("increasing") = false, + R"doc(Explicit output Vander operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/var.hpp b/csrc/infinicore/src/pybind11/ops/var.hpp new file mode 100644 index 000000000..9668fef5f --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/var.hpp @@ -0,0 +1,62 @@ +#pragma once + +#include + +#include "infinicore/ops/var.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +Tensor py_var(Tensor input, py::object dim, bool unbiased, bool keepdim) { + if (dim.is_none()) { + std::vector dim_vec; + for (int i = 0; i < input->shape().size(); i++) { + dim_vec.push_back(i); + } + return op::var(input, dim_vec, unbiased, keepdim); + } else if (py::isinstance(dim) || py::isinstance(dim)) { + return op::var(input, dim.cast>(), unbiased, keepdim); + } else if (py::isinstance(dim)) { + return op::var(input, std::vector(1, dim.cast()), unbiased, keepdim); + } else { + throw std::invalid_argument("dim must be a tuple or an integer"); + } +} + +void py_var_(Tensor var_output, Tensor input, py::object dim, bool unbiased, bool keepdim) { + if (dim.is_none()) { + std::vector dim_vec; + for (int i = 0; i < input->shape().size(); i++) { + dim_vec.push_back(i); + } + op::var_(var_output, input, dim_vec, unbiased, keepdim); + } else if (py::isinstance(dim) || py::isinstance(dim)) { + op::var_(var_output, input, dim.cast>(), unbiased, keepdim); + } else if (py::isinstance(dim)) { + op::var_(var_output, input, std::vector(1, dim.cast()), unbiased, keepdim); + } else { + throw std::invalid_argument("dim must be a list/tuple or an integer"); + } +} + +inline void bind_var(py::module &m) { + m.def("var", + &py_var, + py::arg("input"), + py::arg("dim"), + py::arg("unbiased"), + py::arg("keepdim"), + R"doc(Var of input tensor along the given dimensions.)doc"); + + m.def("var_", + &py_var_, + py::arg("var_output"), + py::arg("input"), + py::arg("dim"), + py::arg("unbiased"), + py::arg("keepdim"), + R"doc(In-place tensor Var .)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/ops/var_mean.hpp b/csrc/infinicore/src/pybind11/ops/var_mean.hpp new file mode 100644 index 000000000..986ec49f7 --- /dev/null +++ b/csrc/infinicore/src/pybind11/ops/var_mean.hpp @@ -0,0 +1,63 @@ +#pragma once + +#include + +#include "infinicore/ops/var_mean.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +std::pair py_var_mean(Tensor input, py::object dim, bool unbiased, bool keepdim) { + if (dim.is_none()) { + std::vector dim_vec; + for (int i = 0; i < input->shape().size(); i++) { + dim_vec.push_back(i); + } + return op::var_mean(input, dim_vec, unbiased, keepdim); + } else if (py::isinstance(dim) || py::isinstance(dim)) { + return op::var_mean(input, dim.cast>(), unbiased, keepdim); + } else if (py::isinstance(dim)) { + return op::var_mean(input, std::vector(1, dim.cast()), unbiased, keepdim); + } else { + throw std::invalid_argument("dim must be a tuple or an integer"); + } +} + +void py_var_mean_(Tensor var_output, Tensor mean_output, Tensor input, py::object dim, bool unbiased, bool keepdim) { + if (dim.is_none()) { + std::vector dim_vec; + for (int i = 0; i < input->shape().size(); i++) { + dim_vec.push_back(i); + } + op::var_mean_(var_output, mean_output, input, dim_vec, unbiased, keepdim); + } else if (py::isinstance(dim) || py::isinstance(dim)) { + op::var_mean_(var_output, mean_output, input, dim.cast>(), unbiased, keepdim); + } else if (py::isinstance(dim)) { + op::var_mean_(var_output, mean_output, input, std::vector(1, dim.cast()), unbiased, keepdim); + } else { + throw std::invalid_argument("dim must be a list/tuple or an integer"); + } +} + +inline void bind_var_mean(py::module &m) { + m.def("var_mean", + &py_var_mean, + py::arg("input"), + py::arg("dim"), + py::arg("unbiased"), + py::arg("keepdim"), + R"doc(Var & Mean of input tensor along the given dimensions.)doc"); + + m.def("var_mean_", + &py_var_mean_, + py::arg("var_output"), + py::arg("mean_output"), + py::arg("input"), + py::arg("dim"), + py::arg("unbiased"), + py::arg("keepdim"), + R"doc(In-place tensor Var & Mean .)doc"); +} + +} // namespace infinicore::ops diff --git a/csrc/infinicore/src/pybind11/tensor.hpp b/csrc/infinicore/src/pybind11/tensor.hpp new file mode 100644 index 000000000..d9573dd1a --- /dev/null +++ b/csrc/infinicore/src/pybind11/tensor.hpp @@ -0,0 +1,131 @@ +#pragma once + +#include "infinicore.hpp" +#include +#include +#include +#include + +namespace py = pybind11; + +namespace infinicore::tensor { + +Tensor from_list_py(py::handle data, DataType dtype); + +inline void bind(py::module &m) { + py::class_(m, "Tensor") + .def_property_readonly("shape", [](const Tensor &tensor) { return tensor->shape(); }) + .def_property_readonly("strides", [](const Tensor &tensor) { return tensor->strides(); }) + .def_property_readonly("ndim", [](const Tensor &tensor) { return tensor->ndim(); }) + .def_property_readonly("dtype", [](const Tensor &tensor) { return tensor->dtype(); }) + .def_property_readonly("device", [](const Tensor &tensor) { return tensor->device(); }) + .def("data_ptr", [](const Tensor &tensor) { return reinterpret_cast(tensor->data()); }) + .def("size", [](const Tensor &tensor, std::size_t dim) { return tensor->size(dim); }) + .def("stride", [](const Tensor &tensor, std::size_t dim) { return tensor->stride(dim); }) + .def("numel", [](const Tensor &tensor) { return tensor->numel(); }) + .def("is_contiguous", [](const Tensor &tensor) { return tensor->is_contiguous(); }) + .def("is_pinned", [](const Tensor &tensor) { return tensor->is_pinned(); }) + .def("info", [](const Tensor &tensor) { return tensor->info(); }) + + .def("debug", [](const Tensor &tensor) { return tensor->debug(); }) + .def("debug", [](const Tensor &tensor, const std::string &filename) { return tensor->debug(filename); }) + + .def("copy_", [](Tensor &tensor, const Tensor &other) { tensor->copy_from(other); }) + .def("to", [](const Tensor &tensor, const Device &device) { return tensor->to(device); }) + .def("contiguous", [](const Tensor &tensor) { return tensor->contiguous(); }) + + .def("as_strided", [](const Tensor &tensor, const Shape &shape, const Strides &strides) { return tensor->as_strided(shape, strides); }) + .def("narrow", [](const Tensor &tensor, std::size_t dim, std::size_t start, std::size_t length) { return tensor->narrow({{dim, start, length}}); }) + .def("permute", [](const Tensor &tensor, const Shape &dims) { return tensor->permute(dims); }) + .def("view", [](const Tensor &tensor, const Shape &shape) { return tensor->view(shape); }) + .def("unsqueeze", [](const Tensor &tensor, std::size_t dim) { return tensor->unsqueeze(dim); }) + .def("squeeze", [](const Tensor &tensor, std::size_t dim) { return tensor->squeeze(dim); }) + .def("reset", static_cast(&Tensor::reset)) + .def("use_count", &Tensor::use_count) + .def("__str__", [](const Tensor &tensor) { + std::ostringstream oss; + oss << tensor; + return oss.str(); + }) + .def("__repr__", [](const Tensor &tensor) { + std::ostringstream oss; + oss << tensor; + return oss.str(); + }) + .def("__bool__", [](const Tensor &tensor) { + return bool(tensor); + }); + + using EmptyFuncType = Tensor (*)(const Shape &, const DataType &, const Device &, bool); + using StridedEmptyFuncType = Tensor (*)(const Shape &, const Strides &, const DataType &, const Device &, bool); + + m.def("empty", static_cast(&Tensor::empty), + py::arg("shape"), + py::arg("dtype"), + py::arg("device"), + py::arg("pin_memory") = false); + + m.def("strided_empty", static_cast(&Tensor::strided_empty), + py::arg("shape"), + py::arg("strides"), + py::arg("dtype"), + py::arg("device"), + py::arg("pin_memory") = false); + + m.def("zeros", static_cast(&Tensor::zeros), + py::arg("shape"), + py::arg("dtype"), + py::arg("device"), + py::arg("pin_memory") = false); + + m.def( + "_from_numpy_copy", [](py::buffer buffer, const DataType &dtype) { + const py::buffer_info info = buffer.request(); + if (info.itemsize < 0 || static_cast(info.itemsize) != dsize(dtype)) { + throw py::value_error("NumPy item size does not match the target dtype"); + } + + Shape shape; + shape.reserve(info.shape.size()); + for (const auto dim : info.shape) { + if (dim < 0) { + throw py::value_error("NumPy shape must not contain negative dimensions"); + } + shape.push_back(static_cast(dim)); + } + + py::ssize_t expected_stride = info.itemsize; + for (py::ssize_t dim = info.ndim; dim-- > 0;) { + if (info.shape[dim] > 1 && info.strides[dim] != expected_stride) { + throw py::value_error("NumPy array must be C-contiguous"); + } + expected_stride *= info.shape[dim]; + } + + auto result = Tensor::empty(shape, dtype, Device(Device::Type::kCpu, 0)); + if (result->nbytes() != static_cast(info.size) * dsize(dtype)) { + throw py::value_error("NumPy buffer size does not match its shape"); + } + if (result->nbytes() != 0) { + std::memcpy(result->data(), info.ptr, result->nbytes()); + } + return result; + }, + py::arg("array"), py::arg("dtype")); + + m.def( + "from_blob", [](uintptr_t raw_ptr, Shape &shape, const DataType &dtype, const Device &device) { + return Tensor{infinicore::Tensor::from_blob(reinterpret_cast(raw_ptr), shape, dtype, device)}; + }, + pybind11::arg("raw_ptr"), pybind11::arg("shape"), pybind11::arg("dtype"), pybind11::arg("device")); + + m.def( + "strided_from_blob", [](uintptr_t raw_ptr, Shape &shape, Strides &strides, const DataType &dtype, const Device &device) { + return Tensor{infinicore::Tensor::strided_from_blob(reinterpret_cast(raw_ptr), shape, strides, dtype, device)}; + }, + pybind11::arg("raw_ptr"), pybind11::arg("shape"), pybind11::arg("strides"), pybind11::arg("dtype"), pybind11::arg("device")); + + m.def("from_list", &from_list_py, py::arg("data"), py::arg("dtype")); +} + +} // namespace infinicore::tensor diff --git a/csrc/infinicore/src/tensor/copy.cc b/csrc/infinicore/src/tensor/copy.cc new file mode 100644 index 000000000..77e145c11 --- /dev/null +++ b/csrc/infinicore/src/tensor/copy.cc @@ -0,0 +1,125 @@ +#include "infinicore/context/context.hpp" +#include "infinicore/dtype.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/tensor.hpp" + +#include +#include + +namespace infinicore { +namespace { + +void copyCpuStrided(std::byte *dst, + const std::byte *src, + const Shape &shape, + const Strides &dst_strides, + const Strides &src_strides, + size_t element_size, + size_t dim = 0) { + if (dim == shape.size()) { + std::memmove(dst, src, element_size); + return; + } + + const auto byte_size = static_cast(element_size); + for (size_t index = 0; index < shape[dim]; ++index) { + copyCpuStrided( + dst + static_cast(index) * dst_strides[dim] * byte_size, + src + static_cast(index) * src_strides[dim] * byte_size, + shape, + dst_strides, + src_strides, + element_size, + dim + 1); + } +} + +} // namespace + +Tensor TensorImpl::to(Device device) const { + if (device == data_.memory->device()) { + return Tensor(const_cast(this)->shared_from_this()); + } else { + std::shared_ptr _t = empty(meta_.shape, meta_.dtype, device); + _t->copy_from(Tensor(const_cast(this)->shared_from_this())); + return Tensor(_t); + } +} + +void TensorImpl::copy_from(Tensor src) { + if (src->shape() != this->shape()) { + throw std::runtime_error( + "Cannot copy from tensor with different shape. Src: " + src->info() + " Dst: " + this->info()); + } + if (src->dtype() != this->dtype()) { + throw std::runtime_error( + "Cannot copy from tensor with different dtype. Src: " + src->info() + " Dst: " + this->info()); + } + if (this->device() == src->device()) { + if (this->device().type() == Device::Type::kCpu) { + if (this->is_contiguous() && src->is_contiguous()) { + if (this->nbytes() != 0) { + std::memmove(this->data(), src->data(), this->nbytes()); + } + } else { + auto host_staging = Tensor::empty( + this->shape(), this->dtype(), Device{Device::Type::kCpu}); + copyCpuStrided( + host_staging->data(), src->data(), this->shape(), + host_staging->strides(), src->strides(), this->element_size()); + copyCpuStrided( + this->data(), host_staging->data(), this->shape(), + this->strides(), host_staging->strides(), this->element_size()); + } + } else { + op::rearrange_(Tensor(const_cast(this)->shared_from_this()), src); + } + } else { + if (!src->is_contiguous()) { + src = src->contiguous(); + } + + // Use nbytes() to get the actual tensor size, not the full memory size + size_t copy_size = std::min(this->nbytes(), src->nbytes()); + if (this->device().type() == Device::Type::kCpu) { + if (this->is_contiguous()) { + context::setDevice(src->device()); + context::memcpyD2H(this->data(), src->data(), copy_size); + } else { + auto local_src = Tensor::empty(this->shape(), this->dtype(), this->device()); + context::setDevice(src->device()); + context::memcpyD2H(local_src->data(), src->data(), copy_size); + this->copy_from(local_src); + } + } else if (src->device().type() == Device::Type::kCpu) { + context::setDevice(this->device()); + // copy_from does not retain the host source after it returns. + if (this->is_contiguous()) { + context::memcpyH2D(this->data(), src->data(), copy_size, false); + } else { + auto local_src = Tensor::empty(this->shape(), this->dtype(), this->device()); + context::memcpyH2D(local_src->data(), src->data(), copy_size, false); + op::rearrange_(Tensor(const_cast(this)->shared_from_this()), local_src); + } + } else { + auto host_staging = Tensor::empty( + this->shape(), this->dtype(), Device{Device::Type::kCpu}); + host_staging->copy_from(src); + this->copy_from(host_staging); + } + } +} + +Tensor TensorImpl::contiguous() const { + if (is_contiguous()) { + return Tensor(const_cast(this)->shared_from_this()); + } else if (device().type() == Device::Type::kCpu) { + auto result = Tensor::empty(shape(), dtype(), device()); + result->copy_from(Tensor(const_cast(this)->shared_from_this())); + return result; + } else { + return op::rearrange(Tensor(const_cast(this)->shared_from_this())); + } +} + +} // namespace infinicore diff --git a/csrc/infinicore/src/tensor/debug.cc b/csrc/infinicore/src/tensor/debug.cc new file mode 100644 index 000000000..4c3d88f16 --- /dev/null +++ b/csrc/infinicore/src/tensor/debug.cc @@ -0,0 +1,226 @@ +#include "infinicore/context/context.hpp" +#include "infinicore/dtype.hpp" +#include "infinicore/tensor.hpp" + +#include +#include +#include +#include +#include +#include + +namespace infinicore { + +inline float f16_to_f32(uint16_t h) { + uint32_t sign = (h & 0x8000) << 16; + int32_t exponent = (h >> 10) & 0x1F; + uint32_t mantissa = h & 0x3FF; + + uint32_t f32; + if (exponent == 31) { + if (mantissa != 0) { + f32 = sign | 0x7F800000 | (mantissa << 13); + } else { + f32 = sign | 0x7F800000; + } + } else if (exponent == 0) { + if (mantissa == 0) { + f32 = sign; + } else { + exponent = -14; + while ((mantissa & 0x400) == 0) { + mantissa <<= 1; + exponent--; + } + mantissa &= 0x3FF; + f32 = sign | ((exponent + 127) << 23) | (mantissa << 13); + } + } else { + f32 = sign | ((exponent + 127 - 15) << 23) | (mantissa << 13); + } + + float result; + std::memcpy(&result, &f32, sizeof(result)); + return result; +} + +inline float bf16_to_f32(uint16_t val) { + uint32_t bits32 = static_cast(val) << 16; + float out; + std::memcpy(&out, &bits32, sizeof(out)); + return out; +} + +// Template function for printing data recursively +template +void print_data(const T *data, const Shape &shape, const Strides &strides, size_t dim) { + if (dim == shape.size() - 1) { + for (size_t i = 0; i < shape[dim]; i++) { + std::cout << data[i * strides[dim]] << " "; + } + std::cout << std::endl; + } else if (dim < shape.size() - 1) { + for (size_t i = 0; i < shape[dim]; i++) { + print_data(data + i * strides[dim], shape, strides, dim + 1); + } + } +} + +// Specialization for F16 (uint16_t) +template <> +void print_data(const uint16_t *data, const Shape &shape, const Strides &strides, size_t dim) { + if (dim == shape.size() - 1) { + for (size_t i = 0; i < shape[dim]; i++) { + std::cout << f16_to_f32(data[i * strides[dim]]) << " "; + } + std::cout << std::endl; + } else if (dim < shape.size() - 1) { + for (size_t i = 0; i < shape[dim]; i++) { + print_data(data + i * strides[dim], shape, strides, dim + 1); + } + } +} + +// Function for printing BF16 data +void print_data_bf16(const uint16_t *data, const Shape &shape, const Strides &strides, size_t dim) { + if (dim == shape.size() - 1) { + for (size_t i = 0; i < shape[dim]; i++) { + std::cout << bf16_to_f32(data[i * strides[dim]]) << " "; + } + std::cout << std::endl; + } else if (dim < shape.size() - 1) { + for (size_t i = 0; i < shape[dim]; i++) { + print_data_bf16(data + i * strides[dim], shape, strides, dim + 1); + } + } +} + +// Function for printing I8 data +void print_data_i8(const int8_t *data, const Shape &shape, const Strides &strides, size_t dim) { + if (dim == shape.size() - 1) { + for (size_t i = 0; i < shape[dim]; i++) { + std::cout << static_cast(data[i * strides[dim]]) << " "; + } + std::cout << std::endl; + } else if (dim < shape.size() - 1) { + for (size_t i = 0; i < shape[dim]; i++) { + print_data_i8(data + i * strides[dim], shape, strides, dim + 1); + } + } +} + +// Function for printing U8 data +void print_data_u8(const uint8_t *data, const Shape &shape, const Strides &strides, size_t dim) { + if (dim == shape.size() - 1) { + for (size_t i = 0; i < shape[dim]; i++) { + std::cout << static_cast(data[i * strides[dim]]) << " "; + } + std::cout << std::endl; + } else if (dim < shape.size() - 1) { + for (size_t i = 0; i < shape[dim]; i++) { + print_data_u8(data + i * strides[dim], shape, strides, dim + 1); + } + } +} + +// Template function for writing data recursively to binary file (handles non-contiguous tensors) +template +void write_binary_data(std::ofstream &out, const T *data, const Shape &shape, const Strides &strides, size_t dim) { + if (dim == shape.size() - 1) { + // Write the innermost dimension + for (size_t i = 0; i < shape[dim]; i++) { + out.write(reinterpret_cast(&data[i * strides[dim]]), sizeof(T)); + } + } else { + // Recursively process higher dimensions + for (size_t i = 0; i < shape[dim]; i++) { + write_binary_data(out, data + i * strides[dim], shape, strides, dim + 1); + } + } +} + +void TensorImpl::debug(const std::string &filename) const { + // Synchronize device if needed + context::syncDevice(); + std::cout << info() << std::endl; + std::unique_ptr allocated_memory; // RAII: 自动管理内存 + auto cpu_tensor = this->contiguous()->to(Device{Device::Type::kCpu}); + const std::byte *cpu_data = cpu_tensor->data(); + // If filename is provided, save to binary file + if (!filename.empty()) { + std::ofstream outFile(filename, std::ios::binary); + if (!outFile) { + std::cerr << "Error opening file for writing: " << filename << "\n"; + return; // allocated_memory 会自动释放（RAII） + } + // Fast path: contiguous tensor, write in one go + size_t mem_size = cpu_tensor->numel() * dsize(cpu_tensor->dtype()); + outFile.write(reinterpret_cast(cpu_data), mem_size); + // 显式关闭文件并检查是否成功 + outFile.close(); + if (!outFile) { + std::cerr << "Error: Failed to write data to file: " << filename << "\n"; + return; + } + std::cout << "Data written to binary file: " << filename; + std::cout << "\n"; + return; + } + // Print data based on dtype + switch (cpu_tensor->dtype()) { + case DataType::kFloat16: + print_data(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kFloat32: + print_data(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kFloat64: + print_data(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kUInt64: + print_data(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kInt64: + print_data(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kUInt32: + print_data(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kInt32: + print_data(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kUInt16: + print_data(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kInt16: + print_data(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kUInt8: + print_data_u8(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kInt8: + print_data_i8(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + case DataType::kBFloat16: + print_data_bf16(reinterpret_cast(cpu_data), + cpu_tensor->shape(), cpu_tensor->strides(), 0); + break; + } +} + +void TensorImpl::debug() const { + this->debug(""); +} + +} // namespace infinicore diff --git a/csrc/infinicore/src/tensor/tensor.cc b/csrc/infinicore/src/tensor/tensor.cc new file mode 100644 index 000000000..c56669f82 --- /dev/null +++ b/csrc/infinicore/src/tensor/tensor.cc @@ -0,0 +1,296 @@ +#include "infinicore/tensor.hpp" +#include "../context/internal.hpp" +#include "../utils.hpp" +#include "infinicore/context/context.hpp" +#include "infinicore/dtype.hpp" + +#include + +namespace { +// Helper function to calculate contiguous strides +inline infinicore::Strides calculate_contiguous_strides(const infinicore::Shape &shape) { + infinicore::Strides strides(shape.size()); + infinicore::Stride stride = 1; + for (int i = shape.size() - 1; i >= 0; --i) { + strides[i] = stride; + stride *= shape[i]; + } + return strides; +} +} // namespace + +namespace infinicore { +TensorImpl *Tensor::operator->() { return impl_.get(); } + +const TensorImpl *Tensor::operator->() const { return impl_.get(); } + +Tensor Tensor::empty(const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory) { + return Tensor{TensorImpl::empty(shape, dtype, device, pin_memory)}; +} + +Tensor Tensor::strided_empty(const Shape &shape, + const Strides &strides, + const DataType &dtype, + const Device &device, + bool pin_memory) { + return Tensor{TensorImpl::strided_empty(shape, strides, dtype, device, pin_memory)}; +} + +Tensor Tensor::zeros(const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory) { + return Tensor{TensorImpl::zeros(shape, dtype, device, pin_memory)}; +} + +Tensor Tensor::ones(const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory) { + return Tensor{TensorImpl::ones(shape, dtype, device, pin_memory)}; +} + +Tensor Tensor::from_blob(void *raw_ptr, const Shape &shape, const DataType &dtype, const Device &device) { + return Tensor{TensorImpl::from_blob(raw_ptr, shape, dtype, device)}; +} + +Tensor Tensor::strided_from_blob(void *raw_ptr, const Shape &shape, const Strides &strides, const DataType &dtype, const Device &device) { + return Tensor{TensorImpl::strided_from_blob(raw_ptr, shape, strides, dtype, device)}; +} + +Tensor::operator bool() const { + return impl_ != nullptr; +} + +TensorMetaData::TensorMetaData(const Shape &_shape, const Strides &_strides, const DataType &_dtype) + : shape(_shape), strides(_strides), dtype(_dtype) {} + +TensorImpl::TensorImpl(const Shape &shape, const DataType &dtype) + : meta_(TensorMetaData(shape, calculate_contiguous_strides(shape), dtype)) {} + +TensorImpl::TensorImpl(const Shape &shape, const Strides &strides, const DataType &dtype) + : meta_(TensorMetaData(shape, strides, dtype)) {} + +std::byte *TensorImpl::data() { + return data_.memory->data() + data_.offset; +} + +const std::byte *TensorImpl::data() const { + return data_.memory->data() + data_.offset; +} + +const Shape &TensorImpl::shape() const { + return meta_.shape; +} + +const Strides &TensorImpl::strides() const { + return meta_.strides; +} + +Size TensorImpl::ndim() const { + return meta_.shape.size(); +} + +bool TensorImpl::is_contiguous() const { + Stride expected_stride = 1; + for (int i = meta_.shape.size() - 1; i >= 0; --i) { + if (meta_.strides[i] != expected_stride) { + return false; + } + expected_stride *= meta_.shape[i]; + } + return true; +} + +Size TensorImpl::numel() const { + Size total = 1; + for (const auto &dim : meta_.shape) { + total *= dim; + } + return total; +} + +size_t TensorImpl::element_size() const { + return dsize(dtype()); +} + +size_t TensorImpl::nbytes() const { + return numel() * element_size(); +} + +Size TensorImpl::size(size_t dim) const { + return meta_.shape[dim]; +} + +Stride TensorImpl::stride(size_t dim) const { + return meta_.strides[dim]; +} + +DataType TensorImpl::dtype() const { + return meta_.dtype; +} + +Device TensorImpl::device() const { + return data_.memory->device(); +} + +infini::rt::TensorView TensorImpl::view() { + return infini::rt::TensorView{data(), shape(), dtype(), device(), strides()}; +} + +bool TensorImpl::is_pinned() const { + return data_.memory->is_pinned(); +} + +std::string TensorImpl::info() const { + std::stringstream ss; + + ss << "Tensor: " + << "shape[ "; + for (auto s : this->shape()) { + ss << s << " "; + } + ss << "] strides[ "; + for (auto s : this->strides()) { + ss << s << " "; + } + ss << "] dtype=" << toString(this->dtype()); + ss << " device=" << this->device().ToString(); + + return ss.str(); +} + +std::shared_ptr TensorImpl::empty(const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory) { + auto t = std::shared_ptr(new TensorImpl(shape, dtype)); + t->data_.offset = 0; + + if (device.type() == Device::Type::kCpu) { + if (pin_memory) { + if (context::getDevice().type() == Device::Type::kCpu) { + spdlog::warn("Tensor memory is not pinned by any device with CPU runtime."); + t->data_.memory = context::allocateHostMemory(t->numel() * dsize(dtype)); + } else { + t->data_.memory = context::allocatePinnedHostMemory(t->numel() * dsize(dtype)); + } + } else { + t->data_.memory = context::allocateHostMemory(t->numel() * dsize(dtype)); + } + } else { + context::setDevice(device); + t->data_.memory = context::allocateMemory(t->numel() * dsize(dtype)); + } + + return t; +} + +std::shared_ptr TensorImpl::strided_empty( + const Shape &shape, + const Strides &strides, + const DataType &dtype, + const Device &device, + bool pin_memory) { + + auto impl = std::shared_ptr(new TensorImpl(shape, strides, dtype)); + impl->data_.offset = 0; + + size_t max_offset = 0; + + for (size_t i = 0; i < shape.size(); ++i) { + if (shape[i] > 0) { + max_offset += (shape[i] - 1) * strides[i]; + } + } + + size_t required_elements = max_offset + 1; + size_t required_bytes = required_elements * dsize(dtype); + + if (device.type() == Device::Type::kCpu) { + if (pin_memory) { + if (context::getDevice().type() == Device::Type::kCpu) { + spdlog::warn("Tensor memory is not pinned by any device with CPU runtime."); + impl->data_.memory = context::allocateHostMemory(required_bytes); + } else { + impl->data_.memory = context::allocatePinnedHostMemory(required_bytes); + } + } else { + impl->data_.memory = context::allocateHostMemory(required_bytes); + } + } else { + context::setDevice(device); + impl->data_.memory = context::allocateMemory(required_bytes); + } + + return impl; +} + +std::shared_ptr TensorImpl::zeros(const Shape &shape, + const DataType &dtype, + const Device &device, + bool pin_memory) { + + auto result = empty(shape, dtype, device, pin_memory); + context::setDeviceMemoryAsync(result->data(), 0, result->nbytes(), context::getStream()); + return result; +} + +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); +} + +std::shared_ptr TensorImpl::from_blob( + void *raw_ptr, + const Shape &shape, + const DataType &dtype, + const Device &device) { + auto t = std::shared_ptr(new TensorImpl(shape, dtype)); + t->data_.offset = 0; + t->data_.memory = std::make_shared((std::byte *)raw_ptr, t->numel() * dsize(dtype), device, nullptr); + return t; +} + +std::shared_ptr TensorImpl::strided_from_blob( + void *raw_ptr, + const Shape &shape, + const Strides &strides, + const DataType &dtype, + const Device &device) { + auto t = std::shared_ptr(new TensorImpl(shape, strides, dtype)); + t->data_.offset = 0; + t->data_.memory = std::make_shared((std::byte *)raw_ptr, t->numel() * dsize(dtype), device, nullptr); + return t; +} + +Tensor TensorImpl::to_blob_() const { + if (context::isGraphRecording()) { + context::retainGraphMemory(data_.memory); + } + auto t = std::shared_ptr(new TensorImpl(shape(), strides(), dtype())); + t->data_.offset = this->data_.offset; + t->data_.memory = std::make_shared(this->data_.memory->data(), this->data_.memory->size(), this->data_.memory->device(), nullptr); + t->to_blob_mark_ = true; + return Tensor{t}; +} + +Tensor TensorImpl::resume_from_blob_() const { + auto t = std::shared_ptr(new TensorImpl(shape(), strides(), dtype())); + t->data_.offset = this->data_.offset; + if (to_blob_mark_) { + t->data_.memory = context::reinstantiateBlob(this->data_.memory); + } else { + t->data_.memory = this->data_.memory; + } + + return Tensor{t}; +} + +} // namespace infinicore diff --git a/csrc/infinicore/src/tensor/view.cc b/csrc/infinicore/src/tensor/view.cc new file mode 100644 index 000000000..051ee42c0 --- /dev/null +++ b/csrc/infinicore/src/tensor/view.cc @@ -0,0 +1,159 @@ +#include "infinicore/context/context.hpp" +#include "infinicore/dtype.hpp" +#include "infinicore/tensor.hpp" + +#include "../utils.hpp" + +#include +#include + +namespace infinicore { +Tensor TensorImpl::squeeze(size_t dim) const { + // Create new shape with dimension of size one removed at dim + if (meta_.shape[dim] != 1) { + spdlog::error("Dimension {} is not of size 1 for squeeze operation on {}.", dim, this->info()); + throw std::runtime_error("Invalid squeeze operation on tensor."); + } + Shape new_shape = meta_.shape; + new_shape.erase(new_shape.begin() + dim); + Strides new_strides = meta_.strides; + new_strides.erase(new_strides.begin() + dim); + + auto tensor_impl = std::make_shared(new_shape, new_strides, meta_.dtype); + tensor_impl->data_ = data_; + + return Tensor(tensor_impl); +} + +Tensor TensorImpl::unsqueeze(size_t dim) const { + // Create new shape with dimension of size one inserted at dim + Shape new_shape = meta_.shape; + new_shape.insert(new_shape.begin() + dim, 1); + + // Create new strides with stride of zero for the new dimension + Strides new_strides = meta_.strides; + new_strides.insert(new_strides.begin() + dim, 0); + + auto tensor_impl = std::make_shared(new_shape, new_strides, meta_.dtype); + tensor_impl->data_ = data_; + + return Tensor(tensor_impl); +} + +Tensor TensorImpl::narrow(const std::vector &slices) const { + // Create new shape and calculate offset + Shape new_shape = meta_.shape; + size_t offset = data_.offset; + + for (const auto &slice : slices) { + if (meta_.shape[slice.dim] < slice.start + slice.len) { + spdlog::error("Invalid slice [dim={}, start={}, len={}] on {}.", slice.dim, slice.start, slice.len, this->info()); + throw std::runtime_error("Invalid slice on tensor."); + } + new_shape[slice.dim] = slice.len; + offset += slice.start * meta_.strides[slice.dim] * dsize(meta_.dtype); + } + + // Create new tensor with the same strides but narrowed shape + auto tensor_impl = std::make_shared(new_shape, meta_.strides, meta_.dtype); + tensor_impl->data_.offset = offset; + tensor_impl->data_.memory = data_.memory; + + return Tensor(tensor_impl); +} + +Tensor TensorImpl::permute(const Shape &order) const { + // Validate input + INFINICORE_ASSERT(meta_.shape.size() == order.size()); + + // Check that order contains all indices from 0 to n-1 exactly once + for (size_t i = 0; i < order.size(); i++) { + INFINICORE_ASSERT(std::find(order.begin(), order.end(), i) != order.end()); + } + + // Permute shape and strides + Shape new_shape(order.size()); + Strides new_strides(order.size()); + + for (size_t i = 0; i < order.size(); i++) { + new_shape[i] = meta_.shape[order[i]]; + new_strides[i] = meta_.strides[order[i]]; + } + + auto tensor_impl = std::make_shared(new_shape, new_strides, meta_.dtype); + tensor_impl->data_ = data_; + + return Tensor(tensor_impl); +} + +Tensor TensorImpl::view(const Shape &new_shape) const { + // Step 1: Validate total size + Size numel = 1; + for (Size dim : meta_.shape) { + numel *= dim; + } + + Size new_numel = 1; + for (Size dim : new_shape) { + new_numel *= dim; + } + + assert(numel == new_numel); + + // Step 2: Get current shape and strides + const Shape &old_shape = meta_.shape; + const Strides &old_strides = meta_.strides; + + // Step 3: Create merged shape and strides + Shape merged_shape; + Strides merged_strides; + + if (!old_shape.empty()) { + merged_shape.push_back(old_shape[0]); + merged_strides.push_back(old_strides[0]); + + for (size_t i = 1; i < old_shape.size(); ++i) { + if (old_strides[i] * static_cast(old_shape[i]) == merged_strides.back()) { + merged_shape.back() *= old_shape[i]; + merged_strides.back() = old_strides[i]; + } else { + merged_shape.push_back(old_shape[i]); + merged_strides.push_back(old_strides[i]); + } + } + } + + // Step 4: Compute new strides by splitting merged dimensions + Strides new_strides(new_shape.size()); + size_t merged_idx = 0; + Stride current_stride = merged_strides[0]; + Size remaining_size = merged_shape[0]; + + for (size_t i = 0; i < new_shape.size(); ++i) { + // 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."); + } + 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."); + }; + + new_strides[i] = current_stride * (remaining_size / new_shape[i]); + remaining_size /= new_shape[i]; + } + + return this->as_strided(new_shape, new_strides); +} + +Tensor TensorImpl::as_strided(const Shape &new_shape, const Strides &new_strides) const { + auto tensor_impl = std::make_shared(new_shape, new_strides, meta_.dtype); + tensor_impl->data_ = data_; + + return Tensor(tensor_impl); +} +} // namespace infinicore diff --git a/csrc/infinicore/src/utils.hpp b/csrc/infinicore/src/utils.hpp new file mode 100644 index 000000000..c76176600 --- /dev/null +++ b/csrc/infinicore/src/utils.hpp @@ -0,0 +1,66 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +inline struct SpdlogInitializer { + SpdlogInitializer() { + if (!std::getenv("INFINICORE_LOG_LEVEL")) { + spdlog::set_level(spdlog::level::info); + } else { + spdlog::cfg::load_env_levels("INFINICORE_LOG_LEVEL"); + } + // Set pattern for logging + // Using SPDLOG_* macros enables source location support (%s and %#) + // Format: [timestamp] [level] [file:line] message + spdlog::set_pattern("[%Y-%m-%d %H:%M:%S.%e] [%^%l%$] [%s:%#] %v"); + } +} spdlog_initializer; + +#define STRINGIZE_(x) #x +#define STRINGIZE(x) STRINGIZE_(x) + +#define INFINICORE_CHECK_ERROR(call) \ + do { \ + SPDLOG_DEBUG("Entering `" #call "` at `" __FILE__ ":" STRINGIZE(__LINE__) "`."); \ + const auto ret = (call); \ + SPDLOG_DEBUG("Exiting `" #call "` at `" __FILE__ ":" STRINGIZE(__LINE__) "`."); \ + if (ret != infini::rt::runtime::kSuccess) { \ + throw std::runtime_error("`" #call "` failed with error code " \ + + std::to_string(static_cast(ret)) \ + + " from " + std::string(__func__) \ + + " at " + std::string(__FILE__) \ + + ":" + std::to_string(__LINE__) + "."); \ + } \ + } while (false) + +#define INFINICORE_ASSERT_TENSORS_SAME_DEVICE(FIRST___, ...) \ + do { \ + const auto &first_device___ = (FIRST___)->device(); \ + for (const auto &tensor___ : {__VA_ARGS__}) { \ + if (first_device___ != (tensor___)->device()) { \ + throw std::runtime_error("Tensor devices mismatch " \ + + first_device___.ToString() + " vs " \ + + (tensor___)->device().ToString() \ + + " from " + std::string(__func__) \ + + " at " + std::string(__FILE__) \ + + ":" + std::to_string(__LINE__) + "."); \ + } \ + } \ + } while (0) + +#define INFINICORE_ASSERT(CONDITION__) \ + do { \ + if (!(CONDITION__)) { \ + SPDLOG_ERROR( \ + "Assertion `{}` failed from {} at {}:{}", \ + #CONDITION__, __func__, __FILE__, __LINE__); \ + throw std::runtime_error( \ + std::string("Assertion `") + #CONDITION__ + "` failed from " + __func__ + " at " + __FILE__ + ":" + std::to_string(__LINE__)); \ + } \ + } while (0) diff --git a/csrc/infinicore/utils/custom_types.cc b/csrc/infinicore/utils/custom_types.cc new file mode 100644 index 000000000..1a6da3c70 --- /dev/null +++ b/csrc/infinicore/utils/custom_types.cc @@ -0,0 +1,85 @@ +#include "custom_types.h" +#include +#include + +float _f16_to_f32(fp16_t val) { + uint16_t h = val._v; + uint32_t sign = (h & 0x8000) << 16; + int32_t exponent = (h >> 10) & 0x1F; + uint32_t mantissa = h & 0x3FF; + + uint32_t f32; + if (exponent == 31) { + if (mantissa != 0) { + f32 = sign | 0x7F800000 | (mantissa << 13); + } else { + f32 = sign | 0x7F800000; + } + } else if (exponent == 0) { + if (mantissa == 0) { + f32 = sign; + } else { + exponent = -14; + while ((mantissa & 0x400) == 0) { + mantissa <<= 1; + exponent--; + } + mantissa &= 0x3FF; + f32 = sign | ((exponent + 127) << 23) | (mantissa << 13); + } + } else { + f32 = sign | ((exponent + 127 - 15) << 23) | (mantissa << 13); + } + + float result; + memcpy(&result, &f32, sizeof(result)); + return result; +} + +fp16_t _f32_to_f16(float val) { + uint32_t f32; + memcpy(&f32, &val, sizeof(f32)); // Read the bits of the float32 + uint16_t sign = (f32 >> 16) & 0x8000; // Extract the sign bit + int32_t exponent = ((f32 >> 23) & 0xFF) - 127; // Extract and de-bias the exponent + uint32_t mantissa = f32 & 0x7FFFFF; // Extract the mantissa (fraction part) + + if (exponent >= 16) { // Special cases for Inf and NaN + // NaN + if (exponent == 128 && mantissa != 0) { + return fp16_t{static_cast(sign | 0x7E00)}; + } + // Infinity + return fp16_t{static_cast(sign | 0x7C00)}; + } else if (exponent >= -14) { // Normalized case + return fp16_t{(uint16_t)(sign | ((exponent + 15) << 10) | (mantissa >> 13))}; + } else if (exponent >= -24) { + mantissa |= 0x800000; // Add implicit leading 1 + mantissa >>= (-14 - exponent); + return fp16_t{(uint16_t)(sign | (mantissa >> 13))}; + } else { + // Too small for subnormal: return signed zero + return fp16_t{(uint16_t)sign}; + } +} + +float _bf16_to_f32(bf16_t val) { + // 只需把 bf16 放到 float32 高 16 bit，其余 16 位置 0。 + uint32_t bits32 = static_cast(val._v) << 16; + + float out; + std::memcpy(&out, &bits32, sizeof(out)); + return out; +} + +bf16_t _f32_to_bf16(float val) { + uint32_t bits32; + std::memcpy(&bits32, &val, sizeof(bits32)); + + // 截断前先加 0x7FFF，再根据第 16 位（有效位的最低位）的奇偶做 round-to-nearest-even + const uint32_t rounding_bias = 0x00007FFF + // 0111 1111 1111 1111 + ((bits32 >> 16) & 1); // 尾数的有效位的最低位奇数时 +1，即实现舍入偶数 + + uint16_t bf16_bits = static_cast((bits32 + rounding_bias) >> 16); + + return bf16_t{bf16_bits}; +} diff --git a/csrc/infinicore/utils/custom_types.h b/csrc/infinicore/utils/custom_types.h new file mode 100644 index 000000000..23be702ff --- /dev/null +++ b/csrc/infinicore/utils/custom_types.h @@ -0,0 +1,67 @@ +#ifndef __INFINIUTILS_CUSTOM_TYPES_H__ +#define __INFINIUTILS_CUSTOM_TYPES_H__ +#include +#include + +struct CustomFloat16 { + uint16_t _v; +}; +typedef struct CustomFloat16 fp16_t; + +struct CustomBFloat16 { + uint16_t _v; +}; +typedef struct CustomBFloat16 bf16_t; + +inline bool operator==(const CustomFloat16 &lhs, const CustomFloat16 &rhs) { + return lhs._v == rhs._v; +} + +inline bool operator!=(const CustomFloat16 &lhs, const CustomFloat16 &rhs) { + return !(lhs == rhs); +} + +inline bool operator==(const CustomBFloat16 &lhs, const CustomBFloat16 &rhs) { + return lhs._v == rhs._v; +} + +inline bool operator!=(const CustomBFloat16 &lhs, const CustomBFloat16 &rhs) { + return !(lhs == rhs); +} + +float _f16_to_f32(fp16_t val); +fp16_t _f32_to_f16(float val); + +float _bf16_to_f32(bf16_t val); +bf16_t _f32_to_bf16(float val); + +namespace utils { +// General template for non-fp16_t conversions +template +TypeTo cast(TypeFrom val) { + if constexpr (std::is_same::value) { + return val; + } else if constexpr (std::is_same::value && std::is_same::value) { + return _f32_to_f16(val); + } else if constexpr (std::is_same::value && !std::is_same::value) { + return _f32_to_f16(static_cast(val)); + } else if constexpr (std::is_same::value && std::is_same::value) { + return _f16_to_f32(val); + } else if constexpr (std::is_same::value && !std::is_same::value) { + return static_cast(_f16_to_f32(val)); + } else if constexpr (std::is_same::value && std::is_same::value) { + return _f32_to_bf16(val); + } else if constexpr (std::is_same::value && !std::is_same::value) { + return _f32_to_bf16(static_cast(val)); + } else if constexpr (std::is_same::value && std::is_same::value) { + return _bf16_to_f32(val); + } else if constexpr (std::is_same::value && !std::is_same::value) { + return static_cast(_bf16_to_f32(val)); + } else { + return static_cast(val); + } +} + +} // namespace utils + +#endif diff --git a/csrc/layers/attention/attention.cpp b/csrc/layers/attention/attention.cpp index 16506ef02..53d1ab663 100644 --- a/csrc/layers/attention/attention.cpp +++ b/csrc/layers/attention/attention.cpp @@ -145,9 +145,9 @@ void init_kv_cache_quant_params(std::functionpermute({0, 2, 1, 3}) - ->contiguous() - ->view({batch_size, seq_len, num_heads_ * head_dim_}); // [bs, seq_len, n_q_head * head_dim] - } else { - size_t total_seq_len = reinterpret_cast(total_sequence_lengths.value()->to(infinicore::Device::cpu())->data())[0]; - - if (infinilm::quantization::KVQuantAlgo::NONE != this->kv_quant_scheme_) { - infinilm::KVQuantUtils::dequantize( - k_total, v_total, - this->kv_quant_scheme_, - k_scale, - v_scale, - q_reshaped); - } - - k_total = k_total->narrow({{2, 0, total_seq_len}}); // [bs, n_kv_head, total_seq_len, head_dim] - v_total = v_total->narrow({{2, 0, total_seq_len}}); // [bs, n_kv_head, total_seq_len, head_dim] - - // Compute attention - size_t ngroup = num_heads_ / num_kv_heads_; - auto Q = q_reshaped->contiguous()->view({batch_size * num_kv_heads_, ngroup * seq_len, head_dim_}); - auto K = k_total->view({batch_size * num_kv_heads_, total_seq_len, head_dim_}); - auto V = v_total->view({batch_size * num_kv_heads_, total_seq_len, head_dim_}); - - auto K_transposed = K->permute({0, 2, 1}); // [bs * n_kv_head, head_dim, total_seq_len] - - auto attn_weight = infinicore::op::matmul(Q, K_transposed, scale_); // [bs * n_kv_head, ng * seq_len, total_seq_len] - - auto attn_weight_softmax = attn_weight->view({batch_size * num_heads_, seq_len, total_seq_len}); - infinicore::op::causal_softmax_(attn_weight_softmax, attn_weight_softmax); - - auto out = infinicore::op::matmul(attn_weight, V); // [bs * n_kv_head, ng * seq_len, head_dim] - - attn_output = out->view({batch_size, num_heads_, seq_len, head_dim_}) - ->permute({0, 2, 1, 3}) - ->contiguous() - ->view({batch_size, seq_len, num_heads_ * head_dim_}); // [bs, seq_len, n_q_head * head_dim] + size_t total_seq_len = reinterpret_cast(total_sequence_lengths.value()->to(infinicore::Device{infinicore::Device::Type::kCpu})->data())[0]; + + if (infinilm::quantization::KVQuantAlgo::NONE != this->kv_quant_scheme_) { + infinilm::KVQuantUtils::dequantize( + k_total, v_total, + this->kv_quant_scheme_, + k_scale, + v_scale, + q_reshaped); } - return attn_output; + + k_total = k_total->narrow({{2, 0, total_seq_len}}); // [bs, n_kv_head, total_seq_len, head_dim] + v_total = v_total->narrow({{2, 0, total_seq_len}}); // [bs, n_kv_head, total_seq_len, head_dim] + + // Compute attention. + size_t ngroup = num_heads_ / num_kv_heads_; + auto Q = q_reshaped->contiguous()->view({batch_size * num_kv_heads_, ngroup * seq_len, head_dim_}); + auto K = k_total->view({batch_size * num_kv_heads_, total_seq_len, head_dim_}); + auto V = v_total->view({batch_size * num_kv_heads_, total_seq_len, head_dim_}); + + auto K_transposed = K->permute({0, 2, 1}); // [bs * n_kv_head, head_dim, total_seq_len] + + auto attn_weight = infinicore::op::matmul(Q, K_transposed, scale_); // [bs * n_kv_head, ng * seq_len, total_seq_len] + + auto attn_weight_softmax = attn_weight->view({batch_size * num_heads_, seq_len, total_seq_len}); + infinicore::op::causal_softmax_(attn_weight_softmax, attn_weight_softmax); + + auto out = infinicore::op::matmul(attn_weight, V); // [bs * n_kv_head, ng * seq_len, head_dim] + + return out->view({batch_size, num_heads_, seq_len, head_dim_}) + ->permute({0, 2, 1, 3}) + ->contiguous() + ->view({batch_size, seq_len, num_heads_ * head_dim_}); // [bs, seq_len, n_q_head * head_dim] } std::tuple StaticAttentionImpl::do_kv_cache_update(const AttentionLayer &layer, @@ -117,7 +107,7 @@ std::tuple StaticAttentionImpl::do_kv_ca ASSERT_EQ(batch_size, max_batch_size); - size_t cache_pos = reinterpret_cast(past_sequence_lengths->to(infinicore::Device::cpu())->data())[0]; + size_t cache_pos = reinterpret_cast(past_sequence_lengths->to(infinicore::Device{infinicore::Device::Type::kCpu})->data())[0]; auto result_len = cache_pos + update_len; ASSERT(result_len <= max_seq_len); diff --git a/csrc/layers/linear/base_linear.hpp b/csrc/layers/linear/base_linear.hpp index a36954836..ae71b703f 100644 --- a/csrc/layers/linear/base_linear.hpp +++ b/csrc/layers/linear/base_linear.hpp @@ -3,7 +3,7 @@ #include "../quantization/quantization.hpp" #include "infinicore/nn/module.hpp" #include "infinicore/ops.hpp" -#include +#include #include namespace infinilm::nn { @@ -15,7 +15,7 @@ class BaseLinear : public infinicore::nn::Module { BaseLinear(size_t in_features, size_t out_features, std::shared_ptr quantization = std::make_shared(nullptr), bool bias = true, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), int split_dim = -1, int tp_rank = 0, int tp_size = 1, int tp_num_heads = -1); diff --git a/csrc/layers/linear/fused_linear.hpp b/csrc/layers/linear/fused_linear.hpp index 8773a081c..88a17f12a 100644 --- a/csrc/layers/linear/fused_linear.hpp +++ b/csrc/layers/linear/fused_linear.hpp @@ -14,7 +14,7 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear { size_t num_q_head, size_t num_k_head, size_t num_v_head, bool q_bias, bool k_bias, bool v_bias, std::shared_ptr quantization = nullptr, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); @@ -23,7 +23,7 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear { size_t num_q_head, size_t num_kv_head, std::shared_ptr quantization = nullptr, bool bias = false, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); @@ -34,7 +34,7 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear { const std::string &q_name, const std::string &k_name, const std::string &v_name, RegisterParamFn register_fn, std::shared_ptr quantization = nullptr, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); @@ -45,7 +45,7 @@ class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear { RegisterParamFn register_fn, std::shared_ptr quantization = nullptr, bool bias = false, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); @@ -98,13 +98,13 @@ class GateUpParallelLinear : public infinilm::nn::ColumnParallelLinear { GateUpParallelLinear(size_t hidden_size, size_t intermediate_size, std::shared_ptr quantization = nullptr, bool bias = false, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); GateUpParallelLinear(size_t hidden_size, size_t intermediate_size, bool gate_bias, bool up_bias, std::shared_ptr quantization = nullptr, - const infinicore::DataType &dtype = infinicore::DataType::F32, const infinicore::Device &device = infinicore::Device(), + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); GateUpParallelLinear(size_t hidden_size, size_t intermediate_size, @@ -112,7 +112,7 @@ class GateUpParallelLinear : public infinilm::nn::ColumnParallelLinear { RegisterParamFn register_fn, std::shared_ptr quantization = nullptr, bool bias = false, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); diff --git a/csrc/layers/linear/linear.cpp b/csrc/layers/linear/linear.cpp index 84982409f..1a99a7bf3 100644 --- a/csrc/layers/linear/linear.cpp +++ b/csrc/layers/linear/linear.cpp @@ -88,7 +88,7 @@ infinicore::Tensor RowParallelLinear::forward(infinicore::Tensor &input) const { auto output = BaseLinear::forward(input); if ((tp_size_ > 1) && (communicator_ != nullptr)) { - infinicore::op::distributed::allreduce_(output, output, INFINICCL_SUM, communicator_); + infinicore::op::distributed::allreduce_(output, output, infinicclSum, communicator_); } return output; } diff --git a/csrc/layers/linear/linear.hpp b/csrc/layers/linear/linear.hpp index 566cee77c..27b233c42 100644 --- a/csrc/layers/linear/linear.hpp +++ b/csrc/layers/linear/linear.hpp @@ -4,7 +4,7 @@ #include "base_linear.hpp" #include "infinicore/nn/module.hpp" #include "infinicore/ops.hpp" -#include +#include #include namespace infinilm::nn { @@ -13,14 +13,14 @@ class Linear : public BaseLinear { public: // Without quantization (backward compat) Linear(size_t in_features, size_t out_features, bool bias, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device()); // With quantization Linear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device()); infinicore::Tensor forward(infinicore::Tensor &input) const; @@ -39,7 +39,7 @@ class ColumnParallelLinear : public BaseLinear { ColumnParallelLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), infinicore::Size tp_rank = 0, infinicore::Size tp_size = 1, int tp_num_heads = -1); @@ -64,7 +64,7 @@ class RowParallelLinear : public BaseLinear { RowParallelLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), infinicore::Size tp_rank = 0, infinicore::Size tp_size = 1, infinicclComm_t communicator = nullptr); diff --git a/csrc/layers/moe/dispatcher/standard_dispatcher.cpp b/csrc/layers/moe/dispatcher/standard_dispatcher.cpp index 27e5c0b23..55ab4231d 100644 --- a/csrc/layers/moe/dispatcher/standard_dispatcher.cpp +++ b/csrc/layers/moe/dispatcher/standard_dispatcher.cpp @@ -33,7 +33,7 @@ infinicore::Tensor StandardDispatcher::combine(const CombineInput &combine_input infinicore::op::distributed::allreduce_( combine_input.hidden_states, combine_input.hidden_states, - INFINICCL_SUM, + infinicclSum, communicator_); } return combine_input.hidden_states; diff --git a/csrc/layers/moe/dispatcher/standard_dispatcher.hpp b/csrc/layers/moe/dispatcher/standard_dispatcher.hpp index d6103c9e8..e6631ec09 100644 --- a/csrc/layers/moe/dispatcher/standard_dispatcher.hpp +++ b/csrc/layers/moe/dispatcher/standard_dispatcher.hpp @@ -3,7 +3,7 @@ #include "../common/moe_types.hpp" #include "base_dispatcher.hpp" -#include +#include namespace infinilm::layers::moe { diff --git a/csrc/layers/moe/ep/allgather_reduce_scatter_dispatcher.cpp b/csrc/layers/moe/ep/allgather_reduce_scatter_dispatcher.cpp index 0b62c87eb..a9cfbd8bd 100644 --- a/csrc/layers/moe/ep/allgather_reduce_scatter_dispatcher.cpp +++ b/csrc/layers/moe/ep/allgather_reduce_scatter_dispatcher.cpp @@ -11,7 +11,7 @@ namespace infinilm::layers::moe { namespace { bool same_device(const infinicore::Tensor &tensor, const infinicore::Device &device) { - return tensor && tensor->device().getType() == device.getType() && tensor->device().getIndex() == device.getIndex(); + return tensor && tensor->device().type() == device.type() && tensor->device().index() == device.index(); } void ensure_tensor(infinicore::Tensor &tensor, @@ -71,7 +71,7 @@ void AllGatherReduceScatterDispatcher::reduce_scatter_dim0_( output, input, equal_split_sizes(local_dim0), - INFINICCL_SUM, + infinicclSum, communicator_); } diff --git a/csrc/layers/moe/ep/base_ep_dispatcher.cpp b/csrc/layers/moe/ep/base_ep_dispatcher.cpp index 613293e8b..6ca9631c2 100644 --- a/csrc/layers/moe/ep/base_ep_dispatcher.cpp +++ b/csrc/layers/moe/ep/base_ep_dispatcher.cpp @@ -41,7 +41,7 @@ infinicore::Tensor BaseEPDispatcher::expert_map(const infinicore::Device &device if (config_.ep_size == 1) { return infinicore::Tensor(); } - if (expert_map_ && expert_map_->device().getType() == device.getType() && expert_map_->device().getIndex() == device.getIndex()) { + if (expert_map_ && expert_map_->device().type() == device.type() && expert_map_->device().index() == device.index()) { return expert_map_; } @@ -59,9 +59,9 @@ infinicore::Tensor BaseEPDispatcher::expert_map(const infinicore::Device &device auto cpu = infinicore::Tensor::from_blob( map.data(), {num_experts_}, - infinicore::DataType::I32, - infinicore::Device(infinicore::Device::Type::CPU, 0)); - expert_map_ = infinicore::Tensor::empty({num_experts_}, infinicore::DataType::I32, device); + infinicore::DataType::kInt32, + infinicore::Device(infinicore::Device::Type::kCpu, 0)); + expert_map_ = infinicore::Tensor::empty({num_experts_}, infinicore::DataType::kInt32, device); expert_map_->copy_from(cpu); return expert_map_; } diff --git a/csrc/layers/moe/ep/base_ep_dispatcher.hpp b/csrc/layers/moe/ep/base_ep_dispatcher.hpp index d5d9b75e5..a773e3f71 100644 --- a/csrc/layers/moe/ep/base_ep_dispatcher.hpp +++ b/csrc/layers/moe/ep/base_ep_dispatcher.hpp @@ -3,7 +3,7 @@ #include "../dispatcher/base_dispatcher.hpp" #include "ep_config.hpp" -#include +#include #include namespace infinilm::layers::moe { diff --git a/csrc/layers/moe/ep/local_allreduce_dispatcher.cpp b/csrc/layers/moe/ep/local_allreduce_dispatcher.cpp index e9bf28761..f24ade038 100644 --- a/csrc/layers/moe/ep/local_allreduce_dispatcher.cpp +++ b/csrc/layers/moe/ep/local_allreduce_dispatcher.cpp @@ -14,7 +14,7 @@ void LocalAllReduceDispatcher::allreduce_(infinicore::Tensor tensor) const { if (!tensor) { return; } - infinicore::op::distributed::allreduce_(tensor, tensor, INFINICCL_SUM, communicator_); + infinicore::op::distributed::allreduce_(tensor, tensor, infinicclSum, communicator_); } DispatchOutput LocalAllReduceDispatcher::dispatch( diff --git a/csrc/layers/moe/router/topk_router.cpp b/csrc/layers/moe/router/topk_router.cpp index e70e55cf8..5fad57f23 100644 --- a/csrc/layers/moe/router/topk_router.cpp +++ b/csrc/layers/moe/router/topk_router.cpp @@ -58,7 +58,7 @@ TopKRouter::TopKRouter(std::shared_ptr model_conf ({num_experts_, model_config->get("hidden_size")}, model_config->get_dtype(), device)); if (use_correction_bias_) { - INFINICORE_NN_PARAMETER_INIT(e_score_correction_bias, ({num_experts_}, infinicore::DataType::F32, device)); + INFINICORE_NN_PARAMETER_INIT(e_score_correction_bias, ({num_experts_}, infinicore::DataType::kFloat32, device)); } if (router_backend_ == TopKRouterBackend::FusedGate) { @@ -83,8 +83,8 @@ std::tuple TopKRouter::forward(const inf size_t ntoken = hidden_states->shape()[0]; auto router_logits = infinicore::op::linear(hidden_states, weight_, std::nullopt, 1.0f); - auto router_scores = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::F32, hidden_states->device()); - auto router_indices = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::I32, hidden_states->device()); + auto router_scores = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::kFloat32, hidden_states->device()); + auto router_indices = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::kInt32, hidden_states->device()); const infinicore::Tensor correction_bias = use_correction_bias_ ? static_cast(e_score_correction_bias_) : infinicore::Tensor(); diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index c37abcc88..e33e16495 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -23,7 +23,7 @@ CudaFusedMoeRunner::CudaFusedMoeRunner(size_t num_local_experts, namespace { bool same_device(const infinicore::Tensor &tensor, const infinicore::Device &device) { - return tensor && tensor->device().getType() == device.getType() && tensor->device().getIndex() == device.getIndex(); + return tensor && tensor->device().type() == device.type() && tensor->device().index() == device.index(); } void ensure_tensor(infinicore::Tensor &tensor, @@ -59,7 +59,7 @@ void check_packed_weight_tensor(const infinicore::Tensor &tensor, if (!tensor) { throw std::runtime_error("MoE fused dense core requires " + name); } - if (tensor->device().getType() != device.getType() || tensor->device().getIndex() != device.getIndex()) { + if (tensor->device().type() != device.type() || tensor->device().index() != device.index()) { throw std::runtime_error("MoE fused dense core requires packed weights on the hidden_states device"); } if (tensor->dtype() != dtype) { @@ -112,7 +112,7 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchO throw std::runtime_error("MoE sorted_token_ids workspace was not initialized before graph capture"); } workspace.sorted_token_ids = infinicore::Tensor::empty( - {sorted_token_ids_capacity}, infinicore::DataType::I32, device); + {sorted_token_ids_capacity}, infinicore::DataType::kInt32, device); workspace.sorted_token_ids_capacity = sorted_token_ids_capacity; } if (!workspace.expert_ids || workspace.expert_ids_capacity < max_num_blocks) { @@ -120,7 +120,7 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchO throw std::runtime_error("MoE expert_ids workspace was not initialized before graph capture"); } workspace.expert_ids = infinicore::Tensor::empty( - {max_num_blocks}, infinicore::DataType::I32, device); + {max_num_blocks}, infinicore::DataType::kInt32, device); workspace.expert_ids_capacity = max_num_blocks; } if (!workspace.num_tokens_post_padded) { @@ -128,7 +128,7 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchO throw std::runtime_error("MoE num_tokens_post_padded workspace was not initialized before graph capture"); } workspace.num_tokens_post_padded = infinicore::Tensor::empty( - {1}, infinicore::DataType::I32, device); + {1}, infinicore::DataType::kInt32, device); } if (dispatch_output.expert_map) { diff --git a/csrc/layers/quantization/awq.cpp b/csrc/layers/quantization/awq.cpp index 87cdb1058..9aeb472e2 100644 --- a/csrc/layers/quantization/awq.cpp +++ b/csrc/layers/quantization/awq.cpp @@ -1,9 +1,5 @@ #include "awq.hpp" -#include "awq_marlin.hpp" -#include "infinicore/ops/linear_w4a16_awq.hpp" -#include "marlin_support.hpp" -#include "marlin_utils.hpp" -#include +#include namespace infinilm::quantization { @@ -24,9 +20,9 @@ std::vector AWQ::get_param_layout( int packing_num = get_packing_num(); std::vector descs; - descs.push_back({"qweight", {in_features, out_features / packing_num}, infinicore::DataType::I32, awq_tp_dim, tp_rank, tp_size}); + descs.push_back({"qweight", {in_features, out_features / packing_num}, infinicore::DataType::kInt32, awq_tp_dim, tp_rank, tp_size}); descs.push_back({"scales", {in_features / group_size, out_features}, dtype, awq_tp_dim, tp_rank, tp_size}); - descs.push_back({"qzeros", {in_features / group_size, out_features / packing_num}, infinicore::DataType::I32, awq_tp_dim, tp_rank, tp_size}); + descs.push_back({"qzeros", {in_features / group_size, out_features / packing_num}, infinicore::DataType::kInt32, awq_tp_dim, tp_rank, tp_size}); if (bias) { descs.push_back({"bias", {out_features}, dtype, -1, 0, 1}); } @@ -34,71 +30,20 @@ std::vector AWQ::get_param_layout( } infinicore::Tensor AWQ::forward( - const ParamsMap ¶ms, - const infinicore::Tensor &input, - bool has_bias, + const ParamsMap &, + const infinicore::Tensor &, + bool, float /*alpha*/) const { - - auto input_contiguous = input->is_contiguous() ? input : input->contiguous(); - auto qweight = params.at("qweight"); - auto scales = params.at("scales"); - auto qzeros = params.at("qzeros"); - - std::optional bias_opt; - if (has_bias) { - bias_opt = params.at("bias"); - } - - return infinicore::op::linear_w4a16_awq(input_contiguous->contiguous(), qweight, scales, qzeros, bias_opt); + throw std::runtime_error( + "AWQ quantization is unsupported until its kernels are available in InfiniOps."); } std::shared_ptr AWQ::process_weights_after_loading( - ParamsMap ¶ms, - const infinicore::Device &device, + ParamsMap &, + const infinicore::Device &, int /*split_dim*/) const { - if (device.getType() != infinicore::Device::Type::NVIDIA) { - return nullptr; - } - -#if INFINILM_ENABLE_MARLIN - const int bits = get_or("bits", get_or("w_bit", 4)); - if (bits != 4) { - return nullptr; - } - - auto qweight = params.at("qweight"); - const size_t input_size_per_partition = qweight->size(0); - const size_t output_size_per_partition = qweight->size(1) * get_packing_num(); - const int group_size = get_group_size(); - if (!marlin::supports_shape(input_size_per_partition, output_size_per_partition, group_size)) { - return nullptr; - } - - params["qweight"] = marlin::awq_marlin_repack( - qweight, - input_size_per_partition, - output_size_per_partition, - bits); - params["scales"] = marlin::permute_scales( - params.at("scales"), - input_size_per_partition, - output_size_per_partition, - group_size); - params["qzeros"] = marlin::awq_to_marlin_zero_points( - params.at("qzeros"), - input_size_per_partition / static_cast(group_size == -1 ? input_size_per_partition : group_size), - output_size_per_partition, - bits); - params["g_idx"] = marlin::make_empty_i32(device); - params["perm"] = marlin::make_empty_i32(device); - params["a_scales"] = marlin::make_empty_i32(device); - params["global_scales"] = marlin::make_empty_i32(device); - - return std::make_shared(get_config(), input_size_per_partition, output_size_per_partition); -#else - (void)params; - return nullptr; -#endif + throw std::runtime_error( + "AWQ quantization is unsupported until its kernels are available in InfiniOps."); } std::vector AWQ::split_params( diff --git a/csrc/layers/quantization/awq_marlin.cpp b/csrc/layers/quantization/awq_marlin.cpp deleted file mode 100644 index 47a4c269e..000000000 --- a/csrc/layers/quantization/awq_marlin.cpp +++ /dev/null @@ -1,141 +0,0 @@ -#include "awq_marlin.hpp" -#include "marlin_support.hpp" - -#if INFINILM_ENABLE_MARLIN -#include "marlin_utils.hpp" - -#include "../../utils.hpp" -#include "infinicore/ops/add.hpp" -#include "infinicore/ops/gptq_marlin_gemm.hpp" - -#include - -namespace infinilm::quantization { - -infinicore::Tensor AWQMarlin::get_workspace( - infinicore::Tensor out, - const infinicore::Tensor &a, - const infinicore::Tensor &b, - infinicore::Tensor &b_scales, - infinicore::Tensor &global_scales, - infinicore::Tensor &b_zeros, - infinicore::Tensor &g_idx, - infinicore::Tensor &perm) const { - const auto required = infinicore::op::gptq_marlin_gemm_workspace_size( - out, a, b, b_scales, global_scales, b_zeros, g_idx, perm); - if (!workspace_ || workspace_->numel() < required || workspace_->device() != out->device()) { - workspace_ = infinicore::Tensor::empty({required}, infinicore::DataType::U8, out->device()); - set_zeros(workspace_); - } - return workspace_; -} - -std::vector AWQMarlin::get_param_layout( - size_t, size_t, int, int, int, int, const infinicore::DataType &, bool) const { - return {}; -} - -void AWQMarlin::reset_runtime_state() const { - if (workspace_) { - set_zeros_device_async(workspace_); - } -} - -infinicore::Tensor AWQMarlin::forward( - const ParamsMap ¶ms, - const infinicore::Tensor &input, - bool has_bias, - float /*alpha*/) const { - auto input_contiguous = input->is_contiguous() ? input : input->contiguous(); - const auto &shape = input_contiguous->shape(); - const size_t k = shape.back(); - const size_t m = input_contiguous->numel() / k; - auto flat_input = input_contiguous->view({m, k}); - auto output = infinicore::Tensor::empty({m, output_size_per_partition_}, input->dtype(), input->device()); - - auto qweight = params.at("qweight"); - auto scales = params.at("scales"); - auto qzeros = params.at("qzeros"); - auto g_idx = params.at("g_idx"); - auto perm = params.at("perm"); - auto global_scales = params.at("global_scales"); - - auto workspace = get_workspace(output, flat_input, qweight, scales, global_scales, qzeros, g_idx, perm); - - infinicore::op::gptq_marlin_gemm_with_workspace_( - workspace, - output, - flat_input, - qweight, - scales, - global_scales, - qzeros, - g_idx, - perm, - marlin::UINT4_ID, - true, - false, - true, - false); - - if (has_bias) { - auto bias = params.at("bias"); - infinicore::op::add_(output, output, bias->as_strided(output->shape(), {0, 1})); - } - - auto out_shape = shape; - out_shape.back() = output_size_per_partition_; - return output->view(out_shape); -} - -std::vector AWQMarlin::split_params( - const std::unordered_map &, - const std::vector &, - int, - int, int, int) const { - return {}; -} - -} // namespace infinilm::quantization -#else -#include - -namespace infinilm::quantization { - -std::vector AWQMarlin::get_param_layout( - size_t, size_t, int, int, int, int, const infinicore::DataType &, bool) const { - return {}; -} - -void AWQMarlin::reset_runtime_state() const {} - -infinicore::Tensor AWQMarlin::forward( - const ParamsMap &, - const infinicore::Tensor &, - bool, - float) const { - throw std::runtime_error("AWQ Marlin is not available because InfiniCore was built without Marlin GEMM headers."); -} - -std::vector AWQMarlin::split_params( - const std::unordered_map &, - const std::vector &, - int, - int, int, int) const { - return {}; -} - -infinicore::Tensor AWQMarlin::get_workspace( - infinicore::Tensor, - const infinicore::Tensor &, - const infinicore::Tensor &, - infinicore::Tensor &, - infinicore::Tensor &, - infinicore::Tensor &, - infinicore::Tensor &, - infinicore::Tensor &) const { - throw std::runtime_error("AWQ Marlin is not available because InfiniCore was built without Marlin GEMM headers."); -} - -} // namespace infinilm::quantization -#endif diff --git a/csrc/layers/quantization/awq_marlin.hpp b/csrc/layers/quantization/awq_marlin.hpp deleted file mode 100644 index a22fb43b3..000000000 --- a/csrc/layers/quantization/awq_marlin.hpp +++ /dev/null @@ -1,55 +0,0 @@ -#pragma once - -#include "base_quantization.hpp" - -namespace infinilm::quantization { - -class AWQMarlin : public BaseQuantization { -public: - AWQMarlin(const nlohmann::json &quant_config, size_t input_size_per_partition, size_t output_size_per_partition) - : BaseQuantization(quant_config), - input_size_per_partition_(input_size_per_partition), - output_size_per_partition_(output_size_per_partition) {} - - QuantScheme get_quant_scheme() const override { return QuantScheme::AWQ_MARLIN_W4A16; } - - std::vector get_param_layout( - size_t, size_t, int, int, int, int, const infinicore::DataType &, bool) const override; - - infinicore::Tensor forward( - const ParamsMap ¶ms, - const infinicore::Tensor &input, - bool has_bias, - float alpha = 1.0f) const override; - - std::vector split_params( - const std::unordered_map ¶ms, - const std::vector &splits, - int narrow_dim, - int tp_rank, int tp_size, int tp_num_heads) const override; - - void reset_runtime_state() const override; - -private: - infinicore::Tensor get_workspace( - infinicore::Tensor out, - const infinicore::Tensor &a, - const infinicore::Tensor &b, - infinicore::Tensor &b_scales, - infinicore::Tensor &global_scales, - infinicore::Tensor &b_zeros, - infinicore::Tensor &g_idx, - infinicore::Tensor &perm) const; - - size_t input_size_per_partition_; - size_t output_size_per_partition_; - // Per-layer Marlin workspace. It must be all-zero before each launch - // because the current InfiniCore Marlin kernels use it as lock state. - // TODO: replace per-layer memset with a shared global zero workspace, or - // update the kernels so the lock region is self-reset at completion. The - // remaining gap to vLLM is mainly from two sources: TP communication cost - // and this workspace memset/reset path. - mutable infinicore::Tensor workspace_; -}; - -} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/base_quantization.hpp b/csrc/layers/quantization/base_quantization.hpp index 6c37f4121..9a81c2588 100644 --- a/csrc/layers/quantization/base_quantization.hpp +++ b/csrc/layers/quantization/base_quantization.hpp @@ -37,7 +37,7 @@ struct SplitParam { class BaseQuantization : public std::enable_shared_from_this { public: - explicit BaseQuantization(const nlohmann::json &quant_config) : quant_config_(quant_config) {}; + explicit BaseQuantization(const nlohmann::json &quant_config) : quant_config_(quant_config){}; virtual ~BaseQuantization() = default; const nlohmann::json &get_config() const { return quant_config_; } @@ -50,14 +50,16 @@ class BaseQuantization : public std::enable_shared_from_this { int split_dim, int tp_rank, int tp_size, int tp_num_heads, const infinicore::DataType &dtype, - bool bias) const = 0; + bool bias) const + = 0; // Forward pass using the registered parameters virtual infinicore::Tensor forward( const ParamsMap ¶ms, const infinicore::Tensor &input, bool has_bias, - float alpha = 1.0f) const = 0; + float alpha = 1.0f) const + = 0; // Dimension for fused-split (gate/up, q/k/v) of a column-parallel weight. // For NoneQuantization weight [out, in], split is on dim0. @@ -78,11 +80,11 @@ class BaseQuantization : public std::enable_shared_from_this { const std::unordered_map ¶ms, const std::vector &splits, int narrow_dim, - int tp_rank, int tp_size, int tp_num_heads) const = 0; + int tp_rank, int tp_size, int tp_num_heads) const + = 0; - // Post-loading weight processing (e.g., GPTQ->GPTQ_QY conversion). - // Returns a replacement quantization object if the scheme changed (e.g. GPTQ -> GPTQ_QY), - // or nullptr if no replacement is needed. + // Post-loading weight processing. Returns a replacement quantization object if the scheme + // changed, or nullptr if no replacement is needed. virtual std::shared_ptr process_weights_after_loading( ParamsMap ¶ms, const infinicore::Device &device, @@ -94,10 +96,7 @@ class BaseQuantization : public std::enable_shared_from_this { } // Reset transient buffers whose contents affect the next kernel launch. - // Marlin currently uses a zero-initialized lock/workspace region; stale - // lock values from warmup, graph capture, or a previous graph replay can - // make the next launch wait forever. Quantization schemes without such - // runtime state can keep the default no-op implementation. + // Quantization schemes without runtime state keep the default no-op. virtual void reset_runtime_state() const {} template diff --git a/csrc/layers/quantization/compressed_tensors.cpp b/csrc/layers/quantization/compressed_tensors.cpp index ff5617a1e..d2e0d136d 100644 --- a/csrc/layers/quantization/compressed_tensors.cpp +++ b/csrc/layers/quantization/compressed_tensors.cpp @@ -15,13 +15,13 @@ std::vector CompressedTensors::get_param_layout( bool bias) const { std::vector descs; - descs.push_back({"weight", {out_features, in_features}, infinicore::DataType::I8, split_dim, tp_rank, tp_size}); + descs.push_back({"weight", {out_features, in_features}, infinicore::DataType::kInt8, split_dim, tp_rank, tp_size}); // weight_scale is per-output-channel [out_features, 1]; always split on // dim0 (output dimension) for ColumnParallel, and don't split for RowParallel. int scale_split_dim = (split_dim == 0) ? 0 : -1; int scale_tp_size = (split_dim == 0) ? tp_size : 1; int scale_tp_rank = (split_dim == 0) ? tp_rank : 0; - descs.push_back({"weight_scale", {out_features, 1}, infinicore::DataType::F32, scale_split_dim, scale_tp_rank, scale_tp_size}); + descs.push_back({"weight_scale", {out_features, 1}, infinicore::DataType::kFloat32, scale_split_dim, scale_tp_rank, scale_tp_size}); if (bias) { descs.push_back({"bias", {out_features}, dtype, -1, 0, 1}); } diff --git a/csrc/layers/quantization/gptq.cpp b/csrc/layers/quantization/gptq.cpp index ed72fd1be..933c4889e 100644 --- a/csrc/layers/quantization/gptq.cpp +++ b/csrc/layers/quantization/gptq.cpp @@ -1,8 +1,5 @@ #include "gptq.hpp" -#include "gptq_marlin.hpp" -#include "gptq_qy.hpp" -#include "marlin_support.hpp" -#include "marlin_utils.hpp" +#include namespace infinilm::quantization { @@ -20,11 +17,11 @@ std::vector GPTQ::get_param_layout( int group_size = get_group_size(); std::vector descs; - descs.push_back({"qweight", {in_features / 8, out_features}, infinicore::DataType::I32, gptq_tp_dim, tp_rank, tp_size}); - descs.push_back({"qzeros", {in_features / group_size, out_features / 8}, infinicore::DataType::I32, gptq_tp_dim, tp_rank, tp_size}); + descs.push_back({"qweight", {in_features / 8, out_features}, infinicore::DataType::kInt32, gptq_tp_dim, tp_rank, tp_size}); + descs.push_back({"qzeros", {in_features / group_size, out_features / 8}, infinicore::DataType::kInt32, gptq_tp_dim, tp_rank, tp_size}); descs.push_back({"scales", {in_features / group_size, out_features}, dtype, gptq_tp_dim, tp_rank, tp_size}); const bool row_parallel = split_dim == 1; - descs.push_back({"g_idx", {in_features}, infinicore::DataType::I32, row_parallel ? 0 : -1, row_parallel ? tp_rank : 0, row_parallel ? tp_size : 1}); + descs.push_back({"g_idx", {in_features}, infinicore::DataType::kInt32, row_parallel ? 0 : -1, row_parallel ? tp_rank : 0, row_parallel ? tp_size : 1}); if (bias) { descs.push_back({"bias", {out_features}, dtype, -1, 0, 1}); } @@ -36,68 +33,16 @@ infinicore::Tensor GPTQ::forward( const infinicore::Tensor & /*input*/, bool /*has_bias*/, float /*alpha*/) const { - throw std::runtime_error("GPTQ_W4A16 must be converted to GPTQ_QY before forward pass. " - "Call process_weights_after_loading() first."); + throw std::runtime_error( + "GPTQ quantization is unsupported until its kernels are available in InfiniOps."); } std::shared_ptr GPTQ::process_weights_after_loading( - ParamsMap ¶ms, - const infinicore::Device &device, - int split_dim) const { - - if (device.getType() == infinicore::Device::Type::QY) { - return GPTQ_QY::convert_from_gptq(params, device, get_config()); - } - - if (device.getType() == infinicore::Device::Type::NVIDIA) { -#if INFINILM_ENABLE_MARLIN - const int bits = get_or("bits", get_or("w_bit", 4)); - const bool is_sym = get_or("sym", true); - if (bits == 4 && is_sym) { - auto qweight = params.at("qweight"); - const size_t input_size_per_partition = qweight->size(0) * get_packing_num(); - const size_t output_size_per_partition = qweight->size(1); - const int group_size = get_group_size(); - if (marlin::supports_shape(input_size_per_partition, output_size_per_partition, group_size)) { - const bool desc_act = get_or("desc_act", false); - const bool row_parallel = split_dim == 1; - const bool is_k_full = (!desc_act) || (desc_act && !row_parallel); - - if (desc_act) { - infinicore::Tensor perm; - params["g_idx"] = marlin::sort_g_idx(params.at("g_idx"), perm); - params["perm"] = perm; - } else { - params["g_idx"] = marlin::make_empty_i32(device); - params["perm"] = marlin::make_empty_i32(device); - } - - params["qweight"] = marlin::gptq_marlin_repack( - qweight, - params.at("perm"), - input_size_per_partition, - output_size_per_partition, - bits); - params["scales"] = marlin::permute_scales( - params.at("scales"), - input_size_per_partition, - output_size_per_partition, - group_size); - params["qzeros"] = marlin::make_empty_i32(device); - params["global_scales"] = marlin::make_empty_i32(device); - - return std::make_shared( - get_config(), - input_size_per_partition, - output_size_per_partition, - is_k_full); - } - } -#else - (void)split_dim; -#endif - } - return std::const_pointer_cast(shared_from_this()); + ParamsMap &, + const infinicore::Device &, + int) const { + throw std::runtime_error( + "GPTQ quantization is unsupported until its kernels are available in InfiniOps."); } std::vector GPTQ::split_params( diff --git a/csrc/layers/quantization/gptq_marlin.cpp b/csrc/layers/quantization/gptq_marlin.cpp deleted file mode 100644 index e95d85b5a..000000000 --- a/csrc/layers/quantization/gptq_marlin.cpp +++ /dev/null @@ -1,141 +0,0 @@ -#include "gptq_marlin.hpp" -#include "marlin_support.hpp" - -#if INFINILM_ENABLE_MARLIN -#include "marlin_utils.hpp" - -#include "../../utils.hpp" -#include "infinicore/ops/add.hpp" -#include "infinicore/ops/gptq_marlin_gemm.hpp" - -#include - -namespace infinilm::quantization { - -infinicore::Tensor GPTQMarlin::get_workspace( - infinicore::Tensor out, - const infinicore::Tensor &a, - const infinicore::Tensor &b, - infinicore::Tensor &b_scales, - infinicore::Tensor &global_scales, - infinicore::Tensor &b_zeros, - infinicore::Tensor &g_idx, - infinicore::Tensor &perm) const { - const auto required = infinicore::op::gptq_marlin_gemm_workspace_size( - out, a, b, b_scales, global_scales, b_zeros, g_idx, perm); - if (!workspace_ || workspace_->numel() < required || workspace_->device() != out->device()) { - workspace_ = infinicore::Tensor::empty({required}, infinicore::DataType::U8, out->device()); - set_zeros(workspace_); - } - return workspace_; -} - -std::vector GPTQMarlin::get_param_layout( - size_t, size_t, int, int, int, int, const infinicore::DataType &, bool) const { - return {}; -} - -void GPTQMarlin::reset_runtime_state() const { - if (workspace_) { - set_zeros_device_async(workspace_); - } -} - -infinicore::Tensor GPTQMarlin::forward( - const ParamsMap ¶ms, - const infinicore::Tensor &input, - bool has_bias, - float /*alpha*/) const { - auto input_contiguous = input->is_contiguous() ? input : input->contiguous(); - const auto &shape = input_contiguous->shape(); - const size_t k = shape.back(); - const size_t m = input_contiguous->numel() / k; - auto flat_input = input_contiguous->view({m, k}); - auto output = infinicore::Tensor::empty({m, output_size_per_partition_}, input->dtype(), input->device()); - - auto qweight = params.at("qweight"); - auto scales = params.at("scales"); - auto qzeros = params.at("qzeros"); - auto g_idx = params.at("g_idx"); - auto perm = params.at("perm"); - auto global_scales = params.at("global_scales"); - - auto workspace = get_workspace(output, flat_input, qweight, scales, global_scales, qzeros, g_idx, perm); - - infinicore::op::gptq_marlin_gemm_with_workspace_( - workspace, - output, - flat_input, - qweight, - scales, - global_scales, - qzeros, - g_idx, - perm, - marlin::UINT4B8_ID, - is_k_full_, - false, - true, - false); - - if (has_bias) { - auto bias = params.at("bias"); - infinicore::op::add_(output, output, bias->as_strided(output->shape(), {0, 1})); - } - - auto out_shape = shape; - out_shape.back() = output_size_per_partition_; - return output->view(out_shape); -} - -std::vector GPTQMarlin::split_params( - const std::unordered_map &, - const std::vector &, - int, - int, int, int) const { - return {}; -} - -} // namespace infinilm::quantization -#else -#include - -namespace infinilm::quantization { - -std::vector GPTQMarlin::get_param_layout( - size_t, size_t, int, int, int, int, const infinicore::DataType &, bool) const { - return {}; -} - -void GPTQMarlin::reset_runtime_state() const {} - -infinicore::Tensor GPTQMarlin::forward( - const ParamsMap &, - const infinicore::Tensor &, - bool, - float) const { - throw std::runtime_error("GPTQ Marlin is not available because InfiniCore was built without Marlin GEMM headers."); -} - -std::vector GPTQMarlin::split_params( - const std::unordered_map &, - const std::vector &, - int, - int, int, int) const { - return {}; -} - -infinicore::Tensor GPTQMarlin::get_workspace( - infinicore::Tensor, - const infinicore::Tensor &, - const infinicore::Tensor &, - infinicore::Tensor &, - infinicore::Tensor &, - infinicore::Tensor &, - infinicore::Tensor &, - infinicore::Tensor &) const { - throw std::runtime_error("GPTQ Marlin is not available because InfiniCore was built without Marlin GEMM headers."); -} - -} // namespace infinilm::quantization -#endif diff --git a/csrc/layers/quantization/gptq_marlin.hpp b/csrc/layers/quantization/gptq_marlin.hpp deleted file mode 100644 index c16c79438..000000000 --- a/csrc/layers/quantization/gptq_marlin.hpp +++ /dev/null @@ -1,58 +0,0 @@ -#pragma once - -#include "base_quantization.hpp" - -namespace infinilm::quantization { - -class GPTQMarlin : public BaseQuantization { -public: - GPTQMarlin(const nlohmann::json &quant_config, size_t input_size_per_partition, - size_t output_size_per_partition, bool is_k_full) - : BaseQuantization(quant_config), - input_size_per_partition_(input_size_per_partition), - output_size_per_partition_(output_size_per_partition), - is_k_full_(is_k_full) {} - - QuantScheme get_quant_scheme() const override { return QuantScheme::GPTQ_MARLIN_W4A16; } - - std::vector get_param_layout( - size_t, size_t, int, int, int, int, const infinicore::DataType &, bool) const override; - - infinicore::Tensor forward( - const ParamsMap ¶ms, - const infinicore::Tensor &input, - bool has_bias, - float alpha = 1.0f) const override; - - std::vector split_params( - const std::unordered_map ¶ms, - const std::vector &splits, - int narrow_dim, - int tp_rank, int tp_size, int tp_num_heads) const override; - - void reset_runtime_state() const override; - -private: - infinicore::Tensor get_workspace( - infinicore::Tensor out, - const infinicore::Tensor &a, - const infinicore::Tensor &b, - infinicore::Tensor &b_scales, - infinicore::Tensor &global_scales, - infinicore::Tensor &b_zeros, - infinicore::Tensor &g_idx, - infinicore::Tensor &perm) const; - - size_t input_size_per_partition_; - size_t output_size_per_partition_; - bool is_k_full_; - // Per-layer Marlin workspace. It must be all-zero before each launch - // because the current InfiniCore Marlin kernels use it as lock state. - // TODO: replace per-layer memset with a shared global zero workspace, or - // update the kernels so the lock region is self-reset at completion. The - // remaining gap to vLLM is mainly from two sources: TP communication cost - // and this workspace memset/reset path. - mutable infinicore::Tensor workspace_; -}; - -} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/gptq_qy.cpp b/csrc/layers/quantization/gptq_qy.cpp deleted file mode 100644 index 4098e452d..000000000 --- a/csrc/layers/quantization/gptq_qy.cpp +++ /dev/null @@ -1,259 +0,0 @@ -#include "gptq_qy.hpp" -#include "infinicore/ops.hpp" -#include "infinicore/ops/linear_w4a16_gptq_qy.hpp" -#include - -namespace infinilm::quantization { - -std::vector GPTQ_QY::get_param_layout( - size_t in_features, size_t out_features, - int split_dim, int tp_rank, int tp_size, - int /*tp_num_heads*/, - const infinicore::DataType &dtype, - bool bias) const { - - // GPTQ_QY weight layout is transposed: qweight [in_features/2, out_features] - // ColumnParallel (split_dim=0, split output) → tp_dim=1 - // RowParallel (split_dim=1, split input) → tp_dim=0 - int gptq_tp_dim = (split_dim >= 0) ? (1 - split_dim) : -1; - int group_size = get_group_size(); - - std::vector descs; - descs.push_back({"qweight", {in_features / 2, out_features}, infinicore::DataType::U8, gptq_tp_dim, tp_rank, tp_size}); - descs.push_back({"qzeros", {in_features / group_size, out_features}, dtype, gptq_tp_dim, tp_rank, tp_size}); - descs.push_back({"scales", {in_features / group_size, out_features}, dtype, gptq_tp_dim, tp_rank, tp_size}); - descs.push_back({"g_idx", {in_features}, infinicore::DataType::I32, -1, 0, 1}); - if (bias) { - descs.push_back({"bias", {out_features}, dtype, -1, 0, 1}); - } - return descs; -} - -infinicore::Tensor GPTQ_QY::forward( - const ParamsMap ¶ms, - const infinicore::Tensor &input, - bool has_bias, - float /*alpha*/) const { - auto input_contiguous = input->is_contiguous() ? input : input->contiguous(); - auto qweight = params.at("qweight"); - auto qzeros = params.at("qzeros"); - auto scales = params.at("scales"); - - auto output = infinicore::op::linear_w4a16_gptq_qy(input_contiguous->contiguous(), qweight, qzeros, scales, 0, 4); - - if (has_bias) { - auto bias = params.at("bias"); - infinicore::op::add_(output, output, bias->as_strided(output->shape(), {0, 0, 1})); - } - return output; -} - -std::vector GPTQ_QY::split_params( - const std::unordered_map ¶ms, - const std::vector &splits, - int /*narrow_dim*/, - int tp_rank, int tp_size, int tp_num_heads) const { - - // GPTQ_QY parameters have output dimension on dim1. - int fused_dim = get_fused_split_dim(); - std::vector result; - auto qw_it = params.find("qweight"); - auto qz_it = params.find("qzeros"); - auto sc_it = params.find("scales"); - auto gidx_it = params.find("g_idx"); - auto bias_it = params.find("bias"); - - for (const auto &s : splits) { - result.push_back({s.prefix + ".qweight", - infinicore::nn::Parameter( - qw_it->second->narrow({{static_cast(fused_dim), s.start, s.size}}), - fused_dim, tp_rank, tp_size, s.num_shards)}); - result.push_back({s.prefix + ".qzeros", - infinicore::nn::Parameter( - qz_it->second->narrow({{static_cast(fused_dim), s.start, s.size}}), - fused_dim, tp_rank, tp_size, s.num_shards)}); - result.push_back({s.prefix + ".scales", - infinicore::nn::Parameter( - sc_it->second->narrow({{static_cast(fused_dim), s.start, s.size}}), - fused_dim, tp_rank, tp_size, s.num_shards)}); - result.push_back({s.prefix + ".g_idx", - infinicore::nn::Parameter( - gidx_it->second->narrow({{0, 0, gidx_it->second->size(0)}}), - 0, 0, 1, 0)}); - if (bias_it != params.end()) { - result.push_back({s.prefix + ".bias", - infinicore::nn::Parameter( - bias_it->second->narrow({{0, s.start, s.size}}), - 0, tp_rank, tp_size, s.num_shards)}); - } - } - return result; -} - -// ---- Conversion from GPTQ_W4A16 ---- - -std::shared_ptr GPTQ_QY::convert_from_gptq( - ParamsMap ¶ms, - const infinicore::Device &device, - const nlohmann::json &quant_config) { - - auto gptq_qy = std::make_shared(quant_config); - const int bits = gptq_qy->weight_bits(); - const int values_per_int32 = 32 / bits; - - const auto &original_qweight = params.at("qweight"); - const auto &original_qzeros = params.at("qzeros"); - const auto &original_scales = params.at("scales"); - const auto &g_idx = params.at("g_idx"); - - { - 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"); - - params["qweight"] = make_tensor_from_host_( - weight_packed.data(), total_bytes * sizeof(uint8_t), - {dimX, dimY}, infinicore::DataType::U8, 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 = infinilm::detail::float_to_fp16_bits(zeros_fp32); - - params["qzeros"] = make_tensor_from_host_( - zeros_fp16.data(), zeros_fp16.size() * sizeof(uint16_t), - {P, Q * static_cast(values_per_int32)}, - infinicore::DataType::F16, device); - } - - { - auto scales_cpu = original_scales->to(infinicore::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() == infinicore::DataType::F16) { - std::memcpy(scales_fp16.data(), raw_data, num_elements * sizeof(uint16_t)); - } else if (scales_cpu->dtype() == infinicore::DataType::F32) { - std::vector scales_fp32(num_elements); - std::memcpy(scales_fp32.data(), raw_data, num_elements * sizeof(float)); - scales_fp16 = infinilm::detail::float_to_fp16_bits(scales_fp32); - } else { - spdlog::error("Unsupported scales dtype, expected F16 or F32"); - assert(false && "Unsupported scales dtype"); - } - - params["scales"] = make_tensor_from_host_( - scales_fp16.data(), scales_fp16.size() * sizeof(uint16_t), - original_scales->shape(), infinicore::DataType::F16, device); - } - - if (g_idx->numel() > 0) { - params["g_idx"] = g_idx->to(device); - } - - return gptq_qy; -} - -// ---- Private helpers ---- - -std::vector GPTQ_QY::unpack_int32_to_nibbles_3d_(const infinicore::Tensor &packed, int bits) { - assert(bits == 4 || bits == 8); - const int values_per_int32 = 32 / bits; - - auto packed_cpu = packed->to(infinicore::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; -} - -std::vector GPTQ_QY::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; -} - -std::vector GPTQ_QY::unpack_zeros_to_fp32_2d_(const infinicore::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(infinicore::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; -} - -infinicore::Tensor GPTQ_QY::make_tensor_from_host_(const void *data, size_t bytes, - const std::vector &shape, - infinicore::DataType dtype, const infinicore::Device &device) { - auto tensor = infinicore::Tensor::empty(shape, dtype, infinicore::Device::Type::CPU); - std::memcpy(reinterpret_cast(tensor->data()), data, bytes); - - if (device != infinicore::Device::Type::CPU) { - return tensor->to(device); - } - return tensor; -} - -} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/gptq_qy.hpp b/csrc/layers/quantization/gptq_qy.hpp deleted file mode 100644 index 634b4aaf7..000000000 --- a/csrc/layers/quantization/gptq_qy.hpp +++ /dev/null @@ -1,139 +0,0 @@ -#pragma once - -#include "base_quantization.hpp" -#include "infinicore/tensor.hpp" -#include -#include -#include -#include -#include -#include -#include - -namespace infinilm::detail { - -inline 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); - } - 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); -} - -inline std::vector float_to_fp16_bits(const std::vector &values) { - std::vector result; - result.reserve(values.size()); - for (float f : values) { - result.push_back(fp32_to_fp16_bits(f)); - } - return result; -} - -} // namespace infinilm::detail - -namespace infinilm::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); - } - - int weight_bits() const { return get_or("bits", 4); } - bool desc_act() const { return get_or("desc_act", false); } - - // Parameter layout for GPTQ_QY (already converted format) - std::vector get_param_layout( - size_t in_features, size_t out_features, - int split_dim, int tp_rank, int tp_size, - int tp_num_heads, - const infinicore::DataType &dtype, - bool bias) const override; - - int get_fused_split_dim() const override { return 1; } - - infinicore::Tensor forward( - const ParamsMap ¶ms, - const infinicore::Tensor &input, - bool has_bias, - float alpha = 1.0f) const override; - - // Split fused linear parameters into named sub-parameters - std::vector split_params( - const std::unordered_map ¶ms, - const std::vector &splits, - int narrow_dim, - int tp_rank, int tp_size, int tp_num_heads) const override; - - // Convert from GPTQ_W4A16 format and update params in-place. - // Returns a new GPTQ_QY quantization instance. Returns nullptr if - // the device is not QY. - static std::shared_ptr convert_from_gptq( - ParamsMap ¶ms, - const infinicore::Device &device, - const nlohmann::json &quant_config); - -private: - static std::vector unpack_int32_to_nibbles_3d_(const infinicore::Tensor &packed, int bits); - static std::vector combine_nibbles_last_dim_(const std::vector &nibbles, size_t M, size_t K, size_t N); - static std::vector unpack_zeros_to_fp32_2d_(const infinicore::Tensor &packed_zeros, int bits); - static infinicore::Tensor make_tensor_from_host_(const void *data, size_t bytes, - const std::vector &shape, - infinicore::DataType dtype, const infinicore::Device &device); -}; - -} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/marlin_support.hpp b/csrc/layers/quantization/marlin_support.hpp deleted file mode 100644 index 3a230c44f..000000000 --- a/csrc/layers/quantization/marlin_support.hpp +++ /dev/null @@ -1,7 +0,0 @@ -#pragma once - -#if __has_include("infinicore/ops/gptq_marlin_gemm.hpp") && __has_include("infiniop/ops/awq_marlin_repack.h") && __has_include("infiniop/ops/gptq_marlin_repack.h") -#define INFINILM_ENABLE_MARLIN 1 -#else -#define INFINILM_ENABLE_MARLIN 0 -#endif diff --git a/csrc/layers/quantization/marlin_utils.cpp b/csrc/layers/quantization/marlin_utils.cpp deleted file mode 100644 index 47546e244..000000000 --- a/csrc/layers/quantization/marlin_utils.cpp +++ /dev/null @@ -1,476 +0,0 @@ -#include "marlin_utils.hpp" -#include "marlin_support.hpp" - -#include "infinicore/context/context.hpp" -#include -#include -#include -#include - -#if INFINILM_ENABLE_MARLIN -#include -#include -#endif - -namespace infinilm::quantization::marlin { - -namespace { - -#if INFINILM_ENABLE_MARLIN -void check_infiniop_status(infiniStatus_t status, const char *expr) { - if (status != INFINI_STATUS_SUCCESS) { - throw std::runtime_error(std::string(expr) + " failed with status " + std::to_string(static_cast(status))); - } -} - -template -class DescriptorGuard { -public: - explicit DescriptorGuard(Desc desc) : desc_(desc) {} - DescriptorGuard(const DescriptorGuard &) = delete; - DescriptorGuard &operator=(const DescriptorGuard &) = delete; - ~DescriptorGuard() { - if (desc_ != nullptr) { - Destroy(desc_); - } - } - Desc get() const { - return desc_; - } - -private: - Desc desc_; -}; - -infinicore::Tensor make_workspace(size_t workspace_size, const infinicore::Device &device) { - if (workspace_size == 0) { - return infinicore::Tensor(); - } - return infinicore::Tensor::empty({workspace_size}, infinicore::DataType::U8, device); -} - -void *workspace_data(infinicore::Tensor &workspace) { - return workspace ? workspace->data() : nullptr; -} - -infinicore::Tensor awq_marlin_repack_gpu(const infinicore::Tensor &qweight, size_t size_k, size_t size_n, int num_bits) { - const size_t pack_factor = 32 / num_bits; - auto qweight_contiguous = qweight->is_contiguous() ? qweight : qweight->contiguous(); - auto output = infinicore::Tensor::empty({size_k / 16, size_n * 16 / pack_factor}, infinicore::DataType::I32, qweight_contiguous->device()); - - infiniopAwqMarlinRepackDescriptor_t raw_desc = nullptr; - check_infiniop_status( - infiniopCreateAwqMarlinRepackDescriptor( - infinicore::context::getInfiniopHandle(qweight_contiguous->device()), - &raw_desc, - output->desc(), - qweight_contiguous->desc(), - num_bits, - false), - "infiniopCreateAwqMarlinRepackDescriptor"); - DescriptorGuard desc(raw_desc); - - size_t workspace_size = 0; - check_infiniop_status( - infiniopGetAwqMarlinRepackWorkspaceSize(desc.get(), &workspace_size), - "infiniopGetAwqMarlinRepackWorkspaceSize"); - auto workspace = make_workspace(workspace_size, qweight_contiguous->device()); - - check_infiniop_status( - infiniopAwqMarlinRepack( - desc.get(), - workspace_data(workspace), - workspace_size, - output->data(), - qweight_contiguous->data(), - infinicore::context::getStream()), - "infiniopAwqMarlinRepack"); - infinicore::context::syncStream(); - return output; -} - -infinicore::Tensor gptq_marlin_repack_gpu( - const infinicore::Tensor &qweight, - const infinicore::Tensor &perm, - size_t size_k, - size_t size_n, - int num_bits) { - const size_t pack_factor = 32 / num_bits; - auto qweight_contiguous = qweight->is_contiguous() ? qweight : qweight->contiguous(); - auto output = infinicore::Tensor::empty({size_k / 16, size_n * 16 / pack_factor}, infinicore::DataType::I32, qweight_contiguous->device()); - auto perm_desc = (perm && perm->numel() != 0) ? perm->desc() : nullptr; - const void *perm_data = (perm && perm->numel() != 0) ? perm->data() : nullptr; - - infiniopGptqMarlinRepackDescriptor_t raw_desc = nullptr; - check_infiniop_status( - infiniopCreateGptqMarlinRepackDescriptor( - infinicore::context::getInfiniopHandle(qweight_contiguous->device()), - &raw_desc, - output->desc(), - qweight_contiguous->desc(), - perm_desc, - num_bits, - false), - "infiniopCreateGptqMarlinRepackDescriptor"); - DescriptorGuard desc(raw_desc); - - size_t workspace_size = 0; - check_infiniop_status( - infiniopGetGptqMarlinRepackWorkspaceSize(desc.get(), &workspace_size), - "infiniopGetGptqMarlinRepackWorkspaceSize"); - auto workspace = make_workspace(workspace_size, qweight_contiguous->device()); - - check_infiniop_status( - infiniopGptqMarlinRepack( - desc.get(), - workspace_data(workspace), - workspace_size, - output->data(), - qweight_contiguous->data(), - perm_data, - infinicore::context::getStream()), - "infiniopGptqMarlinRepack"); - infinicore::context::syncStream(); - return output; -} -#endif - -std::vector scale_perm() { - std::vector perm; - perm.reserve(64); - for (int i = 0; i < 8; ++i) { - for (int j = 0; j < 8; ++j) { - perm.push_back(i + 8 * j); - } - } - return perm; -} - -std::vector scale_perm_single() { - std::vector perm; - perm.reserve(32); - for (int i = 0; i < 4; ++i) { - for (int j : {0, 1, 8, 9, 16, 17, 24, 25}) { - perm.push_back(2 * i + j); - } - } - return perm; -} - -infinicore::Tensor to_cpu_contiguous(const infinicore::Tensor &tensor) { - return tensor->contiguous()->to(infinicore::Device::cpu()); -} - -infinicore::Tensor copy_to_device(const void *data, size_t bytes, const std::vector &shape, - infinicore::DataType dtype, const infinicore::Device &device) { - auto cpu = infinicore::Tensor::empty(shape, dtype, infinicore::Device::cpu()); - if (bytes != 0) { - std::memcpy(cpu->data(), data, bytes); - } - return device == infinicore::Device::cpu() ? cpu : cpu->to(device); -} - -std::vector unpack_cols(const int32_t *packed, size_t size_k, size_t size_n, int num_bits) { - const size_t pack_factor = 32 / num_bits; - std::vector out(size_k * size_n, 0); - const uint32_t mask = (1u << num_bits) - 1u; - for (size_t r = 0; r < size_k; ++r) { - for (size_t c_pack = 0; c_pack < size_n / pack_factor; ++c_pack) { - uint32_t word = static_cast(packed[r * (size_n / pack_factor) + c_pack]); - for (size_t i = 0; i < pack_factor; ++i) { - out[r * size_n + c_pack * pack_factor + i] = static_cast(word & mask); - word >>= num_bits; - } - } - } - return out; -} - -std::vector pack_cols(const std::vector &unpacked, size_t size_k, size_t size_n, int num_bits) { - const size_t pack_factor = 32 / num_bits; - std::vector out(size_k * (size_n / pack_factor), 0); - for (size_t r = 0; r < size_k; ++r) { - for (size_t c_pack = 0; c_pack < size_n / pack_factor; ++c_pack) { - uint32_t word = 0; - for (size_t i = 0; i < pack_factor; ++i) { - uint32_t value = static_cast(unpacked[r * size_n + c_pack * pack_factor + i]); - word |= value << (num_bits * i); - } - out[r * (size_n / pack_factor) + c_pack] = static_cast(word); - } - } - return out; -} - -void check_repack_shape(size_t size_k, size_t size_n, int num_bits) { - if (num_bits != 4 && num_bits != 8) { - throw std::runtime_error("marlin repack: num_bits must be 4 or 8"); - } - if (size_k % 16 != 0 || size_n % 64 != 0) { - throw std::runtime_error("marlin repack: size_k must be divisible by 16 and size_n by 64"); - } -} - -uint32_t pack_repack_values(const uint32_t *vals, int num_bits, bool upper_half) { - if (num_bits == 4) { - constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; - uint32_t word = 0; - for (int i = 0; i < 8; ++i) { - word |= vals[pack_idx[i]] << (i * 4); - } - return word; - } - - constexpr int pack_idx[4] = {0, 2, 1, 3}; - uint32_t word = 0; - const int offset = upper_half ? 4 : 0; - for (int i = 0; i < 4; ++i) { - word |= vals[offset + pack_idx[i]] << (i * 8); - } - return word; -} - -uint32_t awq_value(const int32_t *packed, size_t size_n, size_t k, size_t n, int num_bits) { - const size_t pack_factor = 32 / num_bits; - const uint32_t mask = (1u << num_bits) - 1u; - const size_t n_pack = n / pack_factor; - const size_t n_pos = n % pack_factor; - constexpr int undo4[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - constexpr int undo8[4] = {0, 2, 1, 3}; - const int pos = num_bits == 4 ? undo4[n_pos] : undo8[n_pos]; - const auto word = static_cast(packed[k * (size_n / pack_factor) + n_pack]); - return (word >> (pos * num_bits)) & mask; -} - -uint32_t gptq_value(const int32_t *packed, size_t size_n, size_t k, size_t n, int num_bits) { - const size_t pack_factor = 32 / num_bits; - const uint32_t mask = (1u << num_bits) - 1u; - const size_t k_pack = k / pack_factor; - const size_t k_pos = k % pack_factor; - const auto word = static_cast(packed[k_pack * size_n + n]); - return (word >> (k_pos * num_bits)) & mask; -} - -template -std::vector repack_to_marlin_tiles(size_t size_k, size_t size_n, int num_bits, Getter get_value) { - constexpr size_t tile_k = 16; - constexpr size_t tile_n = 64; - const size_t pack_factor = 32 / num_bits; - const size_t k_tiles = size_k / tile_k; - const size_t n_tiles = size_n / tile_n; - const size_t tile_words = tile_k * tile_n / pack_factor; - std::vector out(k_tiles * n_tiles * tile_words, 0); - constexpr int tc_offsets[4] = {0, 1, 8, 9}; - - for (size_t kt = 0; kt < k_tiles; ++kt) { - for (size_t nt = 0; nt < n_tiles; ++nt) { - const size_t out_offset = (kt * n_tiles + nt) * tile_words; - for (int warp = 0; warp < 4; ++warp) { - for (int th = 0; th < 32; ++th) { - const int tc_col = th / 4; - const int tc_row = (th % 4) * 2; - const size_t n0 = nt * tile_n + warp * 16 + tc_col; - - uint32_t vals[8]; - for (int i = 0; i < 4; ++i) { - const size_t k = kt * tile_k + static_cast(tc_row + tc_offsets[i]); - vals[i] = get_value(k, n0); - vals[4 + i] = get_value(k, n0 + 8); - } - - if (num_bits == 4) { - out[out_offset + static_cast(th * 4 + warp)] = static_cast(pack_repack_values(vals, num_bits, false)); - } else { - out[out_offset + static_cast(th * 8 + warp * 2)] = static_cast(pack_repack_values(vals, num_bits, false)); - out[out_offset + static_cast(th * 8 + warp * 2 + 1)] = static_cast(pack_repack_values(vals, num_bits, true)); - } - } - } - } - } - return out; -} - -} // namespace - -bool supports_shape(size_t input_size_per_partition, size_t output_size_per_partition, int group_size) { - if (output_size_per_partition % 64 != 0 || input_size_per_partition % 128 != 0) { - return false; - } - if (!(group_size == -1 || group_size == 32 || group_size == 64 || group_size == 128)) { - return false; - } - return group_size == -1 || input_size_per_partition % static_cast(group_size) == 0; -} - -infinicore::Tensor make_empty_i32(const infinicore::Device &device) { - return infinicore::Tensor::empty({0, 0}, infinicore::DataType::I32, device); -} - -infinicore::Tensor make_i32_tensor(const std::vector &data, const std::vector &shape, const infinicore::Device &device) { - return copy_to_device(data.data(), data.size() * sizeof(int32_t), shape, infinicore::DataType::I32, device); -} - -infinicore::Tensor awq_marlin_repack(const infinicore::Tensor &qweight, size_t size_k, size_t size_n, int num_bits) { - check_repack_shape(size_k, size_n, num_bits); - const size_t pack_factor = 32 / num_bits; -#if INFINILM_ENABLE_MARLIN - if (qweight->dtype() != infinicore::DataType::I32 || qweight->shape() != std::vector{size_k, size_n / pack_factor}) { - throw std::runtime_error("awq_marlin_repack: unexpected qweight shape or dtype"); - } - if (qweight->device().getType() == infinicore::Device::Type::NVIDIA) { - return awq_marlin_repack_gpu(qweight, size_k, size_n, num_bits); - } -#endif - auto cpu = to_cpu_contiguous(qweight); - if (cpu->dtype() != infinicore::DataType::I32 || cpu->shape() != std::vector{size_k, size_n / pack_factor}) { - throw std::runtime_error("awq_marlin_repack: unexpected qweight shape or dtype"); - } - auto *packed = reinterpret_cast(cpu->data()); - auto out = repack_to_marlin_tiles(size_k, size_n, num_bits, [&](size_t k, size_t n) { - return awq_value(packed, size_n, k, n, num_bits); - }); - return make_i32_tensor(out, {size_k / 16, size_n * 16 / pack_factor}, qweight->device()); -} - -infinicore::Tensor gptq_marlin_repack(const infinicore::Tensor &qweight, const infinicore::Tensor &perm, size_t size_k, size_t size_n, int num_bits) { - check_repack_shape(size_k, size_n, num_bits); - const size_t pack_factor = 32 / num_bits; -#if INFINILM_ENABLE_MARLIN - if (qweight->dtype() != infinicore::DataType::I32 || qweight->shape() != std::vector{size_k / pack_factor, size_n}) { - throw std::runtime_error("gptq_marlin_repack: unexpected qweight shape or dtype"); - } - if (perm && perm->numel() != 0 && (perm->dtype() != infinicore::DataType::I32 || perm->numel() != size_k)) { - throw std::runtime_error("gptq_marlin_repack: unexpected perm shape or dtype"); - } - if (qweight->device().getType() == infinicore::Device::Type::NVIDIA) { - return gptq_marlin_repack_gpu(qweight, perm, size_k, size_n, num_bits); - } -#endif - auto cpu = to_cpu_contiguous(qweight); - if (cpu->dtype() != infinicore::DataType::I32 || cpu->shape() != std::vector{size_k / pack_factor, size_n}) { - throw std::runtime_error("gptq_marlin_repack: unexpected qweight shape or dtype"); - } - auto *packed = reinterpret_cast(cpu->data()); - - std::vector perm_data; - if (perm && perm->numel() != 0) { - auto perm_cpu = to_cpu_contiguous(perm); - if (perm_cpu->dtype() != infinicore::DataType::I32 || perm_cpu->numel() != size_k) { - throw std::runtime_error("gptq_marlin_repack: unexpected perm shape or dtype"); - } - auto *src = reinterpret_cast(perm_cpu->data()); - perm_data.assign(src, src + perm_cpu->numel()); - } - - auto out = repack_to_marlin_tiles(size_k, size_n, num_bits, [&](size_t k, size_t n) { - const size_t src_k = perm_data.empty() ? k : static_cast(perm_data[k]); - return gptq_value(packed, size_n, src_k, n, num_bits); - }); - return make_i32_tensor(out, {size_k / 16, size_n * 16 / pack_factor}, qweight->device()); -} - -infinicore::Tensor sort_g_idx(const infinicore::Tensor &g_idx, infinicore::Tensor &sort_indices) { - auto cpu = to_cpu_contiguous(g_idx); - const auto size = cpu->numel(); - auto *gidx_data = reinterpret_cast(cpu->data()); - - std::vector indices(size); - std::iota(indices.begin(), indices.end(), 0); - std::stable_sort(indices.begin(), indices.end(), [gidx_data](int32_t lhs, int32_t rhs) { - return gidx_data[lhs] < gidx_data[rhs]; - }); - - std::vector sorted(size); - for (size_t i = 0; i < size; ++i) { - sorted[i] = gidx_data[indices[i]]; - } - - sort_indices = make_i32_tensor(indices, {size}, g_idx->device()); - return make_i32_tensor(sorted, {size}, g_idx->device()); -} - -infinicore::Tensor permute_scales(const infinicore::Tensor &scales, size_t size_k, size_t size_n, int group_size) { - auto cpu = to_cpu_contiguous(scales); - const auto elem_size = cpu->element_size(); - const auto *src = reinterpret_cast(cpu->data()); - std::vector dst(cpu->nbytes()); - - const auto perm = (group_size < static_cast(size_k) && group_size != -1) ? scale_perm() : scale_perm_single(); - const size_t block = perm.size(); - if (cpu->numel() % block != 0) { - throw std::runtime_error("marlin permute_scales: scale tensor size is not compatible with Marlin permutation"); - } - for (size_t row = 0; row < cpu->numel() / block; ++row) { - for (size_t i = 0; i < block; ++i) { - std::memcpy( - dst.data() + (row * block + i) * elem_size, - src + (row * block + static_cast(perm[i])) * elem_size, - elem_size); - } - } - return copy_to_device(dst.data(), dst.size(), {cpu->size(0), size_n}, cpu->dtype(), scales->device()); -} - -infinicore::Tensor awq_to_marlin_zero_points(const infinicore::Tensor &qzeros, size_t size_k, size_t size_n, int num_bits) { - auto cpu = to_cpu_contiguous(qzeros); - const size_t pack_factor = 32 / num_bits; - if (cpu->shape() != std::vector{size_k, size_n / pack_factor}) { - throw std::runtime_error("awq_to_marlin_zero_points: unexpected qzeros shape"); - } - - auto unpacked = unpack_cols(reinterpret_cast(cpu->data()), size_k, size_n, num_bits); - const std::vector undo_interleave = num_bits == 4 ? std::vector{0, 4, 1, 5, 2, 6, 3, 7} - : std::vector{0, 2, 1, 3}; - std::vector unpermuted(unpacked.size()); - for (size_t row = 0; row < unpacked.size() / undo_interleave.size(); ++row) { - for (size_t i = 0; i < undo_interleave.size(); ++i) { - unpermuted[row * undo_interleave.size() + i] = unpacked[row * undo_interleave.size() + static_cast(undo_interleave[i])]; - } - } - - auto perm = scale_perm(); - if (unpermuted.size() % perm.size() != 0) { - throw std::runtime_error("awq_to_marlin_zero_points: zero-point tensor size is not compatible with Marlin permutation"); - } - std::vector permuted(unpermuted.size()); - for (size_t row = 0; row < unpermuted.size() / perm.size(); ++row) { - for (size_t i = 0; i < perm.size(); ++i) { - permuted[row * perm.size() + i] = unpermuted[row * perm.size() + static_cast(perm[i])]; - } - } - - const std::vector interleave = num_bits == 4 ? std::vector{0, 2, 4, 6, 1, 3, 5, 7} - : std::vector{0, 2, 1, 3}; - std::vector interleaved(permuted.size()); - for (size_t row = 0; row < permuted.size() / interleave.size(); ++row) { - for (size_t i = 0; i < interleave.size(); ++i) { - interleaved[row * interleave.size() + i] = permuted[row * interleave.size() + static_cast(interleave[i])]; - } - } - - auto packed = pack_cols(interleaved, size_k, size_n, num_bits); - return make_i32_tensor(packed, {size_k, size_n / pack_factor}, qzeros->device()); -} - -infinicore::Tensor permute_bias(const infinicore::Tensor &bias) { - auto cpu = to_cpu_contiguous(bias); - const auto elem_size = cpu->element_size(); - const auto *src = reinterpret_cast(cpu->data()); - std::vector dst(cpu->nbytes()); - const auto perm = scale_perm_single(); - if (cpu->numel() % perm.size() != 0) { - throw std::runtime_error("marlin permute_bias: bias tensor size is not compatible with Marlin permutation"); - } - for (size_t row = 0; row < cpu->numel() / perm.size(); ++row) { - for (size_t i = 0; i < perm.size(); ++i) { - std::memcpy( - dst.data() + (row * perm.size() + i) * elem_size, - src + (row * perm.size() + static_cast(perm[i])) * elem_size, - elem_size); - } - } - return copy_to_device(dst.data(), dst.size(), cpu->shape(), cpu->dtype(), bias->device()); -} - -} // namespace infinilm::quantization::marlin diff --git a/csrc/layers/quantization/marlin_utils.hpp b/csrc/layers/quantization/marlin_utils.hpp deleted file mode 100644 index edeca9aa5..000000000 --- a/csrc/layers/quantization/marlin_utils.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include "infinicore/tensor.hpp" - -#include -#include - -namespace infinilm::quantization::marlin { - -constexpr int64_t UINT4_ID = 1125899906843648LL; -constexpr int64_t UINT4B8_ID = 1125899907892224LL; - -bool supports_shape(size_t input_size_per_partition, size_t output_size_per_partition, int group_size); - -infinicore::Tensor make_empty_i32(const infinicore::Device &device); -infinicore::Tensor make_i32_tensor(const std::vector &data, const std::vector &shape, const infinicore::Device &device); - -infinicore::Tensor awq_marlin_repack(const infinicore::Tensor &qweight, size_t size_k, size_t size_n, int num_bits); -infinicore::Tensor gptq_marlin_repack(const infinicore::Tensor &qweight, const infinicore::Tensor &perm, size_t size_k, size_t size_n, int num_bits); - -infinicore::Tensor sort_g_idx(const infinicore::Tensor &g_idx, infinicore::Tensor &sort_indices); -infinicore::Tensor permute_scales(const infinicore::Tensor &scales, size_t size_k, size_t size_n, int group_size); -infinicore::Tensor awq_to_marlin_zero_points(const infinicore::Tensor &qzeros, size_t size_k, size_t size_n, int num_bits); -infinicore::Tensor permute_bias(const infinicore::Tensor &bias); - -} // namespace infinilm::quantization::marlin diff --git a/csrc/layers/quantization/quantization.hpp b/csrc/layers/quantization/quantization.hpp index 4fcd9e61f..a86316c19 100644 --- a/csrc/layers/quantization/quantization.hpp +++ b/csrc/layers/quantization/quantization.hpp @@ -1,11 +1,8 @@ #pragma once #include "awq.hpp" -#include "awq_marlin.hpp" #include "base_quantization.hpp" #include "compressed_tensors.hpp" #include "gptq.hpp" -#include "gptq_marlin.hpp" -#include "gptq_qy.hpp" #include "none_quantization.hpp" #include "quantization_scheme.hpp" diff --git a/csrc/layers/quantization/quantization_scheme.hpp b/csrc/layers/quantization/quantization_scheme.hpp index 4ab57c71a..8cb595ef3 100644 --- a/csrc/layers/quantization/quantization_scheme.hpp +++ b/csrc/layers/quantization/quantization_scheme.hpp @@ -6,10 +6,7 @@ enum class QuantScheme { NONE, COMPRESSED_TENSOR_W8A8I8, AWQ_W4A16, - AWQ_MARLIN_W4A16, - GPTQ_W4A16_QY, GPTQ_W4A16, - GPTQ_MARLIN_W4A16, }; enum class KVQuantAlgo { diff --git a/csrc/layers/rotary_embedding/rotary_embedding.cpp b/csrc/layers/rotary_embedding/rotary_embedding.cpp index 13f924dac..2294c5228 100644 --- a/csrc/layers/rotary_embedding/rotary_embedding.cpp +++ b/csrc/layers/rotary_embedding/rotary_embedding.cpp @@ -33,7 +33,7 @@ std::string make_cache_key(size_t head_dim, << "_theta_" << std::setprecision(17) << rope_theta << "_algo_" << static_cast(algo) << "_dtype_" << static_cast(dtype) - << "_dev" << device.toString(); + << "_dev" << device.ToString(); if (mrope_section.has_value()) { oss << "_mrope"; diff --git a/csrc/models/deepseek_v2/deepseek_v2_moe.cpp b/csrc/models/deepseek_v2/deepseek_v2_moe.cpp index a18351ac8..7cc8463ce 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_moe.cpp +++ b/csrc/models/deepseek_v2/deepseek_v2_moe.cpp @@ -13,12 +13,11 @@ namespace { bool supports_fused_deepseek_moe(infinicore::Device::Type device_type) { switch (device_type) { - case infinicore::Device::Type::NVIDIA: - case infinicore::Device::Type::ALI: - case infinicore::Device::Type::HYGON: - case infinicore::Device::Type::ILUVATAR: - case infinicore::Device::Type::METAX: - case infinicore::Device::Type::MOORE: + case infinicore::Device::Type::kNvidia: + case infinicore::Device::Type::kHygon: + case infinicore::Device::Type::kIluvatar: + case infinicore::Device::Type::kMetax: + case infinicore::Device::Type::kMoore: return true; default: return false; @@ -44,8 +43,8 @@ DeepseekV2TopKRouter::forward(const infinicore::Tensor &hidden_states) const { ASSERT(hidden_states->ndim() == 2); const size_t ntoken = hidden_states->shape()[0]; auto router_logits = infinicore::op::linear(hidden_states, weight_, std::nullopt, 1.0f); - auto router_scores = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::F32, hidden_states->device()); - auto router_indices = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::I32, hidden_states->device()); + auto router_scores = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::kFloat32, hidden_states->device()); + auto router_indices = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::kInt32, hidden_states->device()); infinicore::op::topksoftmax(router_scores, router_indices, router_logits, num_experts_per_tok_, norm_topk_prob_); return {router_scores, router_indices}; } @@ -78,8 +77,8 @@ DeepseekV2Experts::DeepseekV2Experts(std::shared_ptrto(infinicore::Device::Type::CPU); - auto top_k_index_cpu = top_k_index->to(infinicore::Device::Type::CPU); + auto top_k_weights_cpu = top_k_weights->to(infinicore::Device::Type::kCpu); + auto top_k_index_cpu = top_k_index->to(infinicore::Device::Type::kCpu); auto *top_k_index_ptr = reinterpret_cast(top_k_index_cpu->data()); auto *top_k_weights_ptr = reinterpret_cast(top_k_weights_cpu->data()); @@ -111,20 +110,20 @@ infinicore::Tensor DeepseekV2Experts::forward(const infinicore::Tensor &hidden_s const infinicore::Tensor &top_k_index, const infinicore::Tensor &top_k_weights) const { ASSERT(hidden_states->ndim() == 2); - if (supports_fused_deepseek_moe(hidden_states->device().getType()) - && (hidden_states->dtype() == infinicore::DataType::BF16 - || hidden_states->dtype() == infinicore::DataType::F16)) { + if (supports_fused_deepseek_moe(hidden_states->device().type()) + && (hidden_states->dtype() == infinicore::DataType::kBFloat16 + || hidden_states->dtype() == infinicore::DataType::kFloat16)) { try { auto output = infinicore::op::deepseek_moe(hidden_states, top_k_index, top_k_weights, gate_weights_, up_weights_, down_weights_, local_moe_intermediate_size_, num_experts_); if (tp_size_ > 1 && communicator_ != nullptr) { - infinicore::op::distributed::allreduce_(output, output, INFINICCL_SUM, communicator_); + infinicore::op::distributed::allreduce_(output, output, infinicclSum, communicator_); } return output; } catch (const std::exception &e) { spdlog::warn("DeepseekV2Experts: deepseek_moe unavailable on {}, falling back to CPU-routed experts: {}", - static_cast(hidden_states->device().getType()), e.what()); + static_cast(hidden_states->device().type()), e.what()); } } return forward_cpu_routed_(hidden_states, top_k_index, top_k_weights); diff --git a/csrc/models/deepseek_v2/deepseek_v2_moe.hpp b/csrc/models/deepseek_v2/deepseek_v2_moe.hpp index 11af3669c..0fb608c4c 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_moe.hpp +++ b/csrc/models/deepseek_v2/deepseek_v2_moe.hpp @@ -8,7 +8,7 @@ #include "infinicore/device.hpp" #include "infinicore/nn/module.hpp" #include "infinicore/tensor.hpp" -#include +#include #include #include diff --git a/csrc/models/mamba/mamba_for_causal_lm.cpp b/csrc/models/mamba/mamba_for_causal_lm.cpp index 8a08103e0..0e7dee4e4 100644 --- a/csrc/models/mamba/mamba_for_causal_lm.cpp +++ b/csrc/models/mamba/mamba_for_causal_lm.cpp @@ -174,7 +174,7 @@ void MambaForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { cache_vec.reserve(num_layers * 2); for (size_t i = 0; i < num_layers; ++i) { cache_vec.push_back(infinicore::Tensor::zeros({max_batch_size, intermediate, conv_kernel - 1}, dtype, device)); - cache_vec.push_back(infinicore::Tensor::zeros({max_batch_size, intermediate, state_size}, infinicore::DataType::F32, device)); + cache_vec.push_back(infinicore::Tensor::zeros({max_batch_size, intermediate, state_size}, infinicore::DataType::kFloat32, device)); } } diff --git a/csrc/models/minicpmv/minicpmv_model.cpp b/csrc/models/minicpmv/minicpmv_model.cpp index bbe1d227a..3a6942514 100644 --- a/csrc/models/minicpmv/minicpmv_model.cpp +++ b/csrc/models/minicpmv/minicpmv_model.cpp @@ -42,7 +42,7 @@ MiniCPMVModel::MiniCPMVModel(std::shared_ptr mode void MiniCPMVModel::replace_embeddings(infinicore::Tensor inputs_embeds, const infinicore::Tensor &vision_hidden, const infinicore::Tensor &image_bound) const { - auto bounds_cpu = image_bound->to(infinicore::Device::cpu()); + auto bounds_cpu = image_bound->to(infinicore::Device{infinicore::Device::Type::kCpu}); auto batch_size = inputs_embeds->size(0); ASSERT_EQ(batch_size, 1); @@ -79,7 +79,7 @@ InfinilmModel::Output MiniCPMVModel::forward(const InfinilmModel::Input &input) // inputs_embeds concat tokens from all requests, while images are processed per request // slice inputs_embeds using request offsets to get the embedding of each request - infinicore::Tensor input_offsets_cpu = input.input_offsets.value()->to(infinicore::Device::cpu()); + infinicore::Tensor input_offsets_cpu = input.input_offsets.value()->to(infinicore::Device{infinicore::Device::Type::kCpu}); int32_t *offsets = (int32_t *)(input_offsets_cpu->data()); for (size_t i : global_state::get_forward_context().mm_metadata.image_req_ids.value()) { auto pixel_values = input.pixel_values.value().at(i); diff --git a/csrc/models/minicpmv/resampler.cpp b/csrc/models/minicpmv/resampler.cpp index 45a858d4b..a236306c8 100644 --- a/csrc/models/minicpmv/resampler.cpp +++ b/csrc/models/minicpmv/resampler.cpp @@ -30,18 +30,18 @@ void compute_2d_sincos_pos_embed(float *out, size_t embed_dim, size_t h, size_t } void write_pos_embed(void *dst, infinicore::DataType dtype, const float *src, size_t n) { - if (dtype == infinicore::DataType::F32) { + if (dtype == infinicore::DataType::kFloat32) { std::memcpy(dst, src, n * sizeof(float)); return; } - if (dtype == infinicore::DataType::F16) { + if (dtype == infinicore::DataType::kFloat16) { auto *out = reinterpret_cast(dst); for (size_t i = 0; i < n; ++i) { out[i] = f32_to_f16(src[i]); } return; } - if (dtype == infinicore::DataType::BF16) { + if (dtype == infinicore::DataType::kBFloat16) { auto *out = reinterpret_cast(dst); for (size_t i = 0; i < n; ++i) { out[i] = f32_to_bf16(src[i]); @@ -149,7 +149,7 @@ Resampler::Resampler(size_t num_queries, INFINICORE_NN_BUFFER_INIT(embedding_table, ({num_patches, num_patches, embed_dim_}, dtype, device_)); std::vector buf(num_patches * num_patches * embed_dim_); compute_2d_sincos_pos_embed(buf.data(), embed_dim_, num_patches, num_patches); - auto embedding_table_cpu = infinicore::Tensor::zeros({num_patches, num_patches, embed_dim_}, dtype, infinicore::Device::cpu()); + auto embedding_table_cpu = infinicore::Tensor::zeros({num_patches, num_patches, embed_dim_}, dtype, infinicore::Device{infinicore::Device::Type::kCpu}); write_pos_embed(embedding_table_cpu->data(), embedding_table_cpu->dtype(), buf.data(), num_patches * num_patches * embed_dim_); embedding_table_->copy_from(embedding_table_cpu); } @@ -166,7 +166,7 @@ infinicore::Tensor Resampler::forward(const infinicore::Tensor &x, kv = ln_kv_->forward(kv); // Build positional embeddings on CPU - auto tgt_cpu = tgt_sizes->to(infinicore::Device::cpu()); + auto tgt_cpu = tgt_sizes->to(infinicore::Device{infinicore::Device::Type::kCpu}); int64_t *tgt_sizes_ptr = (int64_t *)(tgt_cpu->data()); auto pos_embeddings = infinicore::Tensor::zeros(kv->shape(), kv->dtype(), kv->device()); diff --git a/csrc/models/minicpmv/siglip_vision.cpp b/csrc/models/minicpmv/siglip_vision.cpp index 23a23f2f1..0ebf3f997 100644 --- a/csrc/models/minicpmv/siglip_vision.cpp +++ b/csrc/models/minicpmv/siglip_vision.cpp @@ -46,14 +46,14 @@ infinicore::Tensor SiglipVisionEmbeddings::forward(const infinicore::Tensor &pix auto embeddings = patch_embeds->view({batch_size, hidden_size_, seq_len})->permute({0, 2, 1}); // Build position ids on CPU - auto pos_ids_cpu = infinicore::Tensor::zeros({batch_size, seq_len}, infinicore::DataType::I64, infinicore::Device::cpu()); + auto pos_ids_cpu = infinicore::Tensor::zeros({batch_size, seq_len}, infinicore::DataType::kInt64, infinicore::Device{infinicore::Device::Type::kCpu}); auto *pos_ptr = reinterpret_cast(pos_ids_cpu->data()); const size_t num_patches_per_side = static_cast(std::sqrt(static_cast(num_positions_))); std::vector tgt_sizes_host; - auto tgt_cpu = tgt_sizes->to(infinicore::Device::cpu()); + auto tgt_cpu = tgt_sizes->to(infinicore::Device{infinicore::Device::Type::kCpu}); auto n = tgt_cpu->numel(); tgt_sizes_host.resize(n); std::memcpy(tgt_sizes_host.data(), tgt_cpu->data(), n * sizeof(int64_t)); diff --git a/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.hpp b/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.hpp index c155030a2..6b9d571f3 100644 --- a/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.hpp +++ b/csrc/models/qwen3_5/qwen3_5_fused_qkv_linear.hpp @@ -16,7 +16,7 @@ class Qwen35FusedQKVLinear : public infinilm::nn::ColumnParallelLinear { infinilm::layers::linear::RegisterParamFn register_fn, std::shared_ptr quantization = nullptr, bool bias = false, - const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::DataType &dtype = infinicore::DataType::kFloat32, const infinicore::Device &device = infinicore::Device(), engine::distributed::RankInfo rank_info = engine::distributed::RankInfo()); diff --git a/csrc/models/qwen3_5/qwen3_5_model.cpp b/csrc/models/qwen3_5/qwen3_5_model.cpp index 82c8b0d4d..7d5fce53e 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.cpp +++ b/csrc/models/qwen3_5/qwen3_5_model.cpp @@ -12,14 +12,14 @@ namespace infinilm::models::qwen3_5 { namespace { std::vector tensor_to_i32_vector(const infinicore::Tensor &tensor) { - auto cpu_tensor = tensor->to(infinicore::Device::cpu()); + auto cpu_tensor = tensor->to(infinicore::Device{infinicore::Device::Type::kCpu}); std::vector values(cpu_tensor->numel()); - if (cpu_tensor->dtype() == infinicore::DataType::I32) { + if (cpu_tensor->dtype() == infinicore::DataType::kInt32) { const auto *ptr = reinterpret_cast(cpu_tensor->data()); values.assign(ptr, ptr + cpu_tensor->numel()); return values; } - if (cpu_tensor->dtype() == infinicore::DataType::I64) { + if (cpu_tensor->dtype() == infinicore::DataType::kInt64) { const auto *ptr = reinterpret_cast(cpu_tensor->data()); for (size_t i = 0; i < cpu_tensor->numel(); ++i) { values[i] = static_cast(ptr[i]); diff --git a/csrc/models/qwen3_5/qwen3_5_vision.cpp b/csrc/models/qwen3_5/qwen3_5_vision.cpp index d079c429d..588210711 100644 --- a/csrc/models/qwen3_5/qwen3_5_vision.cpp +++ b/csrc/models/qwen3_5/qwen3_5_vision.cpp @@ -30,14 +30,14 @@ size_t get_size_or_first(const nlohmann::json &config, const char *key, size_t d } std::vector tensor_to_i64_vector(const infinicore::Tensor &tensor) { - auto cpu_tensor = tensor->to(infinicore::Device::cpu()); + auto cpu_tensor = tensor->to(infinicore::Device{infinicore::Device::Type::kCpu}); std::vector values(cpu_tensor->numel()); - if (cpu_tensor->dtype() == infinicore::DataType::I64) { + if (cpu_tensor->dtype() == infinicore::DataType::kInt64) { const auto *ptr = reinterpret_cast(cpu_tensor->data()); values.assign(ptr, ptr + cpu_tensor->numel()); return values; } - if (cpu_tensor->dtype() == infinicore::DataType::I32) { + if (cpu_tensor->dtype() == infinicore::DataType::kInt32) { const auto *ptr = reinterpret_cast(cpu_tensor->data()); for (size_t i = 0; i < cpu_tensor->numel(); ++i) { values[i] = static_cast(ptr[i]); @@ -271,7 +271,7 @@ infinicore::Tensor Qwen35VisionModel::build_rotary_position_ids(const infinicore total_tokens += static_cast(grid[i]) * static_cast(grid[i + 1]) * static_cast(grid[i + 2]); } - auto position_ids_cpu = infinicore::Tensor::empty({2, total_tokens}, infinicore::DataType::I64, infinicore::Device::cpu()); + auto position_ids_cpu = infinicore::Tensor::empty({2, total_tokens}, infinicore::DataType::kInt64, infinicore::Device{infinicore::Device::Type::kCpu}); auto *position_ids = reinterpret_cast(position_ids_cpu->data()); size_t out_token = 0; diff --git a/csrc/models/qwen3_moe/qwen3_moe_experts.cpp b/csrc/models/qwen3_moe/qwen3_moe_experts.cpp index d08b3bf77..edc85ff76 100644 --- a/csrc/models/qwen3_moe/qwen3_moe_experts.cpp +++ b/csrc/models/qwen3_moe/qwen3_moe_experts.cpp @@ -24,8 +24,8 @@ infinicore::Tensor Qwen3MoeExperts::forward(const infinicore::Tensor &hidden_sta const infinicore::Tensor &top_k_weights) const { ASSERT(hidden_states->ndim() == 2); - auto top_k_weights_cpu = top_k_weights->to(infinicore::Device::Type::CPU); - auto top_k_index_cpu = top_k_index->to(infinicore::Device::Type::CPU); + auto top_k_weights_cpu = top_k_weights->to(infinicore::Device::Type::kCpu); + auto top_k_index_cpu = top_k_index->to(infinicore::Device::Type::kCpu); int *top_k_index_ptr = reinterpret_cast(top_k_index_cpu->data()); float *top_k_weights_ptr = reinterpret_cast(top_k_weights_cpu->data()); diff --git a/csrc/models/qwen3_moe/qwen3_moe_topk_router.cpp b/csrc/models/qwen3_moe/qwen3_moe_topk_router.cpp index d155f279a..9669afbbd 100644 --- a/csrc/models/qwen3_moe/qwen3_moe_topk_router.cpp +++ b/csrc/models/qwen3_moe/qwen3_moe_topk_router.cpp @@ -25,8 +25,8 @@ std::tuple Qwen3MoeTopKRouter::forward(c size_t ntoken = hidden_states->shape()[0]; auto router_logits = infinicore::op::linear(hidden_states, weight_, std::nullopt); - auto router_scores = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::F32, hidden_states->device()); - auto router_indices = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::I32, hidden_states->device()); + auto router_scores = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::kFloat32, hidden_states->device()); + auto router_indices = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::kInt32, hidden_states->device()); infinicore::op::moe_topk_softmax_(router_scores, router_indices, router_logits, infinicore::Tensor(), norm_topk_prob_, 0.0f); diff --git a/csrc/models/videonsa/videonsa_attention.cpp b/csrc/models/videonsa/videonsa_attention.cpp index 828d10860..ea68f1f00 100644 --- a/csrc/models/videonsa/videonsa_attention.cpp +++ b/csrc/models/videonsa/videonsa_attention.cpp @@ -32,7 +32,7 @@ constexpr size_t kNsaBlockSize = 64; constexpr int kNsaSelectBlocks = 4; infinicore::Tensor scalar_tensor(float value, const infinicore::Device &device) { - auto cpu = infinicore::Tensor::from_blob(&value, {1}, infinicore::DataType::F32, infinicore::Device::cpu()); + auto cpu = infinicore::Tensor::from_blob(&value, {1}, infinicore::DataType::kFloat32, infinicore::Device{infinicore::Device::Type::kCpu}); return cpu->to(device); } @@ -59,38 +59,38 @@ std::pair build_mrope_cache(size_t max_s } } - const auto cpu = infinicore::Device::cpu(); + const auto cpu = infinicore::Device{infinicore::Device::Type::kCpu}; auto sin_cache = infinicore::Tensor::empty({max_seq_len, cache_dim}, dtype, device); auto cos_cache = infinicore::Tensor::empty({max_seq_len, cache_dim}, dtype, device); - if (dtype == infinicore::DataType::F32) { - auto sin_cpu = infinicore::Tensor::from_blob(sin_data.data(), {max_seq_len, cache_dim}, infinicore::DataType::F32, cpu); - auto cos_cpu = infinicore::Tensor::from_blob(cos_data.data(), {max_seq_len, cache_dim}, infinicore::DataType::F32, cpu); + if (dtype == infinicore::DataType::kFloat32) { + auto sin_cpu = infinicore::Tensor::from_blob(sin_data.data(), {max_seq_len, cache_dim}, infinicore::DataType::kFloat32, cpu); + auto cos_cpu = infinicore::Tensor::from_blob(cos_data.data(), {max_seq_len, cache_dim}, infinicore::DataType::kFloat32, cpu); sin_cache->copy_from(sin_cpu); cos_cache->copy_from(cos_cpu); return {sin_cache, cos_cache}; } - if (dtype == infinicore::DataType::BF16) { + if (dtype == infinicore::DataType::kBFloat16) { std::vector sin_bf16(numel); std::vector cos_bf16(numel); for (size_t i = 0; i < numel; ++i) { sin_bf16[i] = f32_to_bf16(sin_data[i]); cos_bf16[i] = f32_to_bf16(cos_data[i]); } - auto sin_cpu = infinicore::Tensor::from_blob(sin_bf16.data(), {max_seq_len, cache_dim}, infinicore::DataType::BF16, cpu); - auto cos_cpu = infinicore::Tensor::from_blob(cos_bf16.data(), {max_seq_len, cache_dim}, infinicore::DataType::BF16, cpu); + auto sin_cpu = infinicore::Tensor::from_blob(sin_bf16.data(), {max_seq_len, cache_dim}, infinicore::DataType::kBFloat16, cpu); + auto cos_cpu = infinicore::Tensor::from_blob(cos_bf16.data(), {max_seq_len, cache_dim}, infinicore::DataType::kBFloat16, cpu); sin_cache->copy_from(sin_cpu); cos_cache->copy_from(cos_cpu); return {sin_cache, cos_cache}; } - if (dtype == infinicore::DataType::F16) { + if (dtype == infinicore::DataType::kFloat16) { std::vector sin_f16(numel); std::vector cos_f16(numel); for (size_t i = 0; i < numel; ++i) { sin_f16[i] = f32_to_f16(sin_data[i]); cos_f16[i] = f32_to_f16(cos_data[i]); } - auto sin_cpu = infinicore::Tensor::from_blob(sin_f16.data(), {max_seq_len, cache_dim}, infinicore::DataType::F16, cpu); - auto cos_cpu = infinicore::Tensor::from_blob(cos_f16.data(), {max_seq_len, cache_dim}, infinicore::DataType::F16, cpu); + auto sin_cpu = infinicore::Tensor::from_blob(sin_f16.data(), {max_seq_len, cache_dim}, infinicore::DataType::kFloat16, cpu); + auto cos_cpu = infinicore::Tensor::from_blob(cos_f16.data(), {max_seq_len, cache_dim}, infinicore::DataType::kFloat16, cpu); sin_cache->copy_from(sin_cpu); cos_cache->copy_from(cos_cpu); return {sin_cache, cos_cache}; @@ -187,7 +187,7 @@ infinicore::Tensor expand_head_gate(const infinicore::Tensor &gate, size_t head_ expand_indices[row * head_dim + d] = static_cast(row); } } - auto indices = infinicore::Tensor::empty({seq_len, num_heads, head_dim}, infinicore::DataType::I64, gate->device()); + auto indices = infinicore::Tensor::empty({seq_len, num_heads, head_dim}, infinicore::DataType::kInt64, gate->device()); infinicore::context::memcpyH2D(indices->data(), expand_indices.data(), expand_indices.size() * sizeof(int64_t), false); return infinicore::op::take(flat_gate, indices); } @@ -312,7 +312,7 @@ infinicore::Tensor VideoNSAAttention::forward(const infinicore::Tensor &position auto &kv_cache = forward_context.kv_cache_vec[layer_idx_]; auto k_cache_layer = kv_cache->narrow({{0, 0, 1}})->squeeze(0); auto v_cache_layer = kv_cache->narrow({{0, 1, 1}})->squeeze(0); - const size_t cache_pos = reinterpret_cast(attn_metadata.past_sequence_lengths.value()->to(infinicore::Device::cpu())->data())[0]; + const size_t cache_pos = reinterpret_cast(attn_metadata.past_sequence_lengths.value()->to(infinicore::Device{infinicore::Device::Type::kCpu})->data())[0]; const size_t total_seq_len = cache_pos + seq_len; k_cache_layer->narrow({{2, cache_pos, seq_len}})->copy_from(k_static->permute({0, 2, 1, 3})); v_cache_layer->narrow({{2, cache_pos, seq_len}})->copy_from(v_static->permute({0, 2, 1, 3})); diff --git a/csrc/models/videonsa/videonsa_for_conditional_generation.cpp b/csrc/models/videonsa/videonsa_for_conditional_generation.cpp index 54fbad26b..85ec32643 100644 --- a/csrc/models/videonsa/videonsa_for_conditional_generation.cpp +++ b/csrc/models/videonsa/videonsa_for_conditional_generation.cpp @@ -43,7 +43,7 @@ VideoNSAForConditionalGeneration::VideoNSAForConditionalGeneration(std::shared_p void VideoNSAForConditionalGeneration::replace_embeddings(infinicore::Tensor inputs_embeds, const infinicore::Tensor &vision_hidden, const infinicore::Tensor &image_bound) const { - auto bounds_cpu = image_bound->to(infinicore::Device::cpu()); + auto bounds_cpu = image_bound->to(infinicore::Device{infinicore::Device::Type::kCpu}); auto out_slice = inputs_embeds->squeeze(0); auto bound_slice = bounds_cpu->squeeze(0); auto bound_count = bound_slice->size(0); @@ -70,7 +70,7 @@ infinilm::InfinilmModel::Output VideoNSAForConditionalGeneration::forward(const } auto input_ids = input.input_ids.value(); auto inputs_embeds = model_->embed_tokens(input_ids); - auto input_offsets_cpu = input.input_offsets.value()->to(infinicore::Device::cpu()); + auto input_offsets_cpu = input.input_offsets.value()->to(infinicore::Device{infinicore::Device::Type::kCpu}); int32_t *offsets = reinterpret_cast(input_offsets_cpu->data()); const auto &image_req_ids = global_state::get_forward_context().mm_metadata.image_req_ids.value(); @@ -93,7 +93,7 @@ infinilm::InfinilmModel::Output VideoNSAForConditionalGeneration::forward(const size_t vision_offset = 0; for (size_t media_idx = 0; media_idx < image_req_ids.size(); ++media_idx) { const size_t req_id = image_req_ids[media_idx]; - auto bounds_cpu = input.image_bound.value().at(media_idx)->to(infinicore::Device::cpu())->squeeze(0); + auto bounds_cpu = input.image_bound.value().at(media_idx)->to(infinicore::Device{infinicore::Device::Type::kCpu})->squeeze(0); auto bound_count = bounds_cpu->size(0); auto bounds = reinterpret_cast(bounds_cpu->data()); size_t vision_len = 0; diff --git a/csrc/models/videonsa/videonsa_vision.cpp b/csrc/models/videonsa/videonsa_vision.cpp index 0b7a88faa..96b776618 100644 --- a/csrc/models/videonsa/videonsa_vision.cpp +++ b/csrc/models/videonsa/videonsa_vision.cpp @@ -188,7 +188,7 @@ VideoNSAVisionModel::VideoNSAVisionModel(const nlohmann::json &config, } VideoNSAVisionModel::WindowMetadata VideoNSAVisionModel::build_window_metadata_(const infinicore::Tensor &grid_thw) const { - auto grid_cpu = grid_thw->to(infinicore::Device::cpu()); + auto grid_cpu = grid_thw->to(infinicore::Device{infinicore::Device::Type::kCpu}); const int64_t *grid = reinterpret_cast(grid_cpu->data()); const size_t n = grid_cpu->size(0); const int64_t vit_window = static_cast(window_size_ / spatial_merge_size_ / patch_size_); @@ -264,7 +264,7 @@ infinicore::Tensor VideoNSAVisionModel::gather_rows_(const infinicore::Tensor &h flat_indices[r * width + c] = src_row * static_cast(width) + static_cast(c); } } - auto indices = infinicore::Tensor::empty({rows, width}, infinicore::DataType::I64, hidden_states->device()); + auto indices = infinicore::Tensor::empty({rows, width}, infinicore::DataType::kInt64, hidden_states->device()); infinicore::context::memcpyH2D(indices->data(), flat_indices.data(), flat_indices.size() * sizeof(int64_t), false); return infinicore::op::take(hidden_states->contiguous(), indices); } diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index b73ce06ae..fa4c3f289 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -85,7 +85,7 @@ inline void bind_infer_engine(py::module &m) { }), py::arg("config_str") = "", py::arg("distributed_config") = distributed::DistConfig(), - py::arg("device_type") = infinicore::context::getDevice().getType(), + py::arg("device_type") = infinicore::context::getDevice().type(), py::arg("cache_config") = py::none(), py::arg("enable_graph_compiling") = false, py::arg("attention_backend") = "default", diff --git a/csrc/utils.hpp b/csrc/utils.hpp index de1bf263f..a5b81da41 100644 --- a/csrc/utils.hpp +++ b/csrc/utils.hpp @@ -1,7 +1,6 @@ #pragma once -#include #include -#include +#include #include #include @@ -25,17 +24,6 @@ inline void assertTrue(int expr, const char *msg, const char *function, const ch spdlog::error("Error at {} in function {} at file {}, line {}", #EXPR, function, file, line); \ exit(EXIT_FAILURE) -#define RUN_INFINI(API) \ - do { \ - auto api_result_ = (API); \ - if (api_result_ != INFINI_STATUS_SUCCESS) { \ - std::cerr << "Error Code " << api_result_ << " in `" << #API << "`" \ - << " from " << __func__ \ - << " at " << __FILE__ << ":" << __LINE__ << std::endl; \ - throw std::runtime_error("InfiniCore C API Error"); \ - } \ - } while (0) - inline float f16_to_f32(uint16_t h) { uint32_t sign = (h & 0x8000) << 16; // Extract the sign bit int32_t exponent = (h >> 10) & 0x1F; // Extract the exponent @@ -149,13 +137,13 @@ inline void hash_combine(size_t &seed, size_t value) { inline infinicore::DataType parse_dtype(const std::string &dtype_str) { static const std::unordered_map dtype_map = { - {"float32", infinicore::DataType::F32}, - {"float16", infinicore::DataType::F16}, - {"bfloat16", infinicore::DataType::BF16}, - {"int8", infinicore::DataType::I8}, + {"float32", infinicore::DataType::kFloat32}, + {"float16", infinicore::DataType::kFloat16}, + {"bfloat16", infinicore::DataType::kBFloat16}, + {"int8", infinicore::DataType::kInt8}, // 可根据需要扩展 - {"int32", infinicore::DataType::I32}, - {"int64", infinicore::DataType::I64}, + {"int32", infinicore::DataType::kInt32}, + {"int64", infinicore::DataType::kInt64}, }; auto it = dtype_map.find(dtype_str); diff --git a/include/infinicore_infer.h b/include/infinicore_infer.h deleted file mode 100644 index 5b2ceb99e..000000000 --- a/include/infinicore_infer.h +++ /dev/null @@ -1,14 +0,0 @@ -#ifndef INFINICORE_INFER_H -#define INFINICORE_INFER_H - -#include "infinicore_infer/cache.h" -#include "infinicore_infer/weights_loader.h" - - -#include "infinicore_infer/models/deepseek.h" -#include "infinicore_infer/models/jiuge.h" -#include "infinicore_infer/models/jiuge_awq.h" -#include "infinicore_infer/models/qwen3vl.h" - - -#endif /* INFINICORE_INFER_H */ diff --git a/include/infinicore_infer/cache.h b/include/infinicore_infer/cache.h deleted file mode 100644 index 522f22352..000000000 --- a/include/infinicore_infer/cache.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef CACHE_H -#define CACHE_H - -#include - -__INFINI_C __export struct KVCache *createKVCache( - size_t nlayers, - size_t max_len, - size_t nkvh_, - size_t dk, - size_t dv, - infiniDtype_t dtype, - infiniDevice_t device, - int *dev_ids, - size_t ndev); - -__INFINI_C __export struct KVCache *duplicateKVCache(const KVCache *kv_cache, size_t seq_len); - -__INFINI_C __export void dropKVCache(KVCache *kv_cache); - -#endif /* CACHE_H */ diff --git a/include/infinicore_infer/models/deepseek.h b/include/infinicore_infer/models/deepseek.h deleted file mode 100644 index d7d2e6863..000000000 --- a/include/infinicore_infer/models/deepseek.h +++ /dev/null @@ -1,163 +0,0 @@ -#ifndef DEEPSEEK_V3_WEIGHTS_H -#define DEEPSEEK_V3_WEIGHTS_H - -#include -#include -#include - -#include -#include - -struct DeepSeekV3Weights; - -// Function pointer signatures -typedef void (*load_global_fn)(DeepSeekV3Weights *, void *cpu_ptr); -typedef void (*load_layer_fn)(DeepSeekV3Weights *, void *cpu_ptr, size_t layer_id); -typedef void (*load_layer_linear_fn)(DeepSeekV3Weights *, void *weight_ptr, void *scale_ptr, void *zero_ptr, size_t layer_id); -typedef void (*load_layer_mlp_fn)( - DeepSeekV3Weights *, - void *gate_weight_ptr, void *gate_scale_ptr, void *gate_zero_ptr, - void *up_weight_ptr, void *up_scale_ptr, void *up_zero_ptr, - void *down_weight_ptr, void *down_scale_ptr, void *down_zero_ptr, - size_t layer_id); -typedef void (*load_layer_expert_mlp_fn)( - DeepSeekV3Weights *, - void *gate_weight_ptr, void *gate_scale_ptr, void *gate_zero_ptr, - void *up_weight_ptr, void *up_scale_ptr, void *up_zero_ptr, - void *down_weight_ptr, void *down_scale_ptr, void *down_zero_ptr, - size_t layer_id, size_t expert_id); - -// Struct containing all weight loading functions -typedef struct { - // Global - load_global_fn load_input_embd; - load_global_fn load_output_norm; - load_global_fn load_output_embd; - - // Attention - load_layer_fn load_attn_norm; - load_layer_linear_fn load_attn_q_a_proj; - load_layer_fn load_attn_q_a_layernorm; - load_layer_linear_fn load_attn_q_b_proj; - load_layer_linear_fn load_attn_kv_a_proj_with_mqa; - load_layer_fn load_attn_kv_a_layernorm; - load_layer_linear_fn load_attn_kv_b_proj; - load_layer_linear_fn load_attn_o_proj; - - // MLP - load_layer_fn load_mlp_norm; - // MLP dense part - load_layer_mlp_fn load_mlp_dense; - - // MLP sparse gating - load_layer_fn load_mlp_gate_weight; - load_layer_fn load_mlp_gate_bias; - - // Shared experts - load_layer_mlp_fn load_mlp_shared_experts; - - // Per-expert functions - load_layer_expert_mlp_fn load_mlp_experts; - -} DeepSeekV3WeightLoader; - -struct DeepSeekV3Model; - -typedef struct { - infiniDtype_t dt_logits; - infiniDtype_t dt_norm; - infiniDtype_t dt_quant_weight; - infiniDtype_t dt_quant_scale; - infiniDtype_t dt_quant_zero; - infiniDtype_t dt_gate_weight; - infiniDtype_t dt_gate_bias; - - size_t n_sparse_layer; - size_t n_dense_layer; - size_t d; - size_t nh; - size_t nkvh; - size_t d_rope; - size_t d_nope; - size_t r_q; - size_t r_kv; - size_t d_qk; - size_t d_v; - - float routed_scale; - size_t nexperts; - size_t kexperts; - size_t di; - size_t di_moe; - size_t dctx; - size_t dvoc; - - float epsilon; - float rope_theta; - uint32_t end_token; - -} DeepSeekV3Meta; - -//////////////////// APIs /////////////////////// -/// @brief 创建模型 -/// @param device 协处理器种类 -/// @param ndev 协处理器数量 -/// @param dev_ids 协处理器编号，长度为 ndev -__INFINI_C __export struct DeepSeekV3Model * -createDeepSeekV3Model(const DeepSeekV3Meta *, - const DeepSeekV3Weights *); - -__INFINI_C DeepSeekV3Weights * -createDeepSeekV3Weights(const DeepSeekV3Meta *meta, - infiniDevice_t device, - int ndev, - const int *dev_ids); - -__INFINI_C __export DeepSeekV3WeightLoader * -createDeepSeekV3WeightLoader(); - -/// @brief 销毁模型 -__INFINI_C __export void destroyDeepSeekV3Model(struct DeepSeekV3Model *); - -__INFINI_C __export struct DeepSeekV3Cache * -createDeepSeekV3Cache(const struct DeepSeekV3Model *); - -__INFINI_C __export void -dropDeepSeekV3Cache(const struct DeepSeekV3Model *, - struct DeepSeekV3Cache *); - -/// @brief 批次推理一轮，并采样出新的 token -/// @param tokens 输入 token 地址 -/// @param ntok 输入 token 数量 -/// @param nreq 请求数量 -/// @param req_lens 每个请求的 token 数量 -/// @param req_pos 每个请求的起始位置 -/// @param kv_caches 每个请求的 KV Cache -/// @param temperature 采样温度（0. 表示贪心采样） -/// @param topk 采样 topk（1 表示贪心采样） -/// @param topp 采样 topp -/// @param output 输出 token 数组，每个请求一个输出，长度至少为nreq -__INFINI_C __export void -inferBatchDeepSeekV3(struct DeepSeekV3Model *, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct DeepSeekV3Cache **caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output); - -/// @brief 批次推理一轮，输出 output embedding 后的 logits -/// @param tokens 输入 token 地址 -/// @param ntok 输入 token 数量 -/// @param nreq 请求数量 -/// @param req_lens 每个请求的 token 数量 -/// @param req_pos 每个请求的起始位置 -/// @param kv_caches 每个请求的 KV Cache -/// @param logits 输出 token 数组，每个请求一个输出，长度至少为nreq -__INFINI_C __export void -forwardBatchDeepSeekV3(struct DeepSeekV3Model *, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct DeepSeekV3Cache **caches, - void *logits); - -#endif // DEEPSEEK_V3_WEIGHTS_H \ No newline at end of file diff --git a/include/infinicore_infer/models/jiuge.h b/include/infinicore_infer/models/jiuge.h deleted file mode 100644 index 824ea8a88..000000000 --- a/include/infinicore_infer/models/jiuge.h +++ /dev/null @@ -1,102 +0,0 @@ -#ifndef MODEL_JIUGE_H -#define MODEL_JIUGE_H - -#include -#include -#include - -#include - -struct JiugeModel; - -typedef struct -{ - infiniDtype_t dt_logits; - size_t nlayer, d, nh, nkvh, dh, di, dctx, dvoc; - float epsilon, theta; - uint32_t end_token; -} JiugeMeta; - -typedef struct -{ - size_t nlayer; - infiniDtype_t dt_norm, dt_mat; - // 0 if linear weights are passed as W, any other value if passed as W^T (default format in pytorch) - int transpose_linear_weights; - // [dvoc, d] - const void *input_embd; - // [d] - const void *output_norm; - // [dvoc, d] - const void *output_embd; - // nlayer * [d] - const void *const *attn_norm; - // nlayer * [ndev, (nh + 2 * nkvh) / ndev * dh, d] - const void *const *attn_qkv; - // nlayer * [ndev, (nh + 2 * nkvh) / ndev * dh] - const void *const *attn_qkv_b; - // nlayer * [dh] - const void *const *attn_q_norm; - // nlayer * [dh] - const void *const *attn_k_norm; - // nlayer * [ndev, d, nkvh / ndev * dh] - const void *const *attn_o; - // nlayer * [d] - const void *const *ffn_norm; - // nlayer * [ndev, 2 * di / ndev, d] - const void *const *ffn_gate_up; - // nlayer * [ndev, d, di / ndev] - const void *const *ffn_down; -} JiugeWeights; - -//////////////////// APIs /////////////////////// -/// @brief 创建模型 -/// @param device 协处理器种类 -/// @param ndev 协处理器数量 -/// @param dev_ids 协处理器编号，长度为 ndev -__INFINI_C __export struct JiugeModel * -createJiugeModel(const JiugeMeta *, - const JiugeWeights *, - infiniDevice_t device, - int ndev, - const int *dev_ids); - -/// @brief 销毁模型 -__INFINI_C __export void -destroyJiugeModel(struct JiugeModel *); - -/// @brief 批次推理一轮，并采样出新的 token -/// @param tokens 输入 token 地址 -/// @param ntok 输入 token 数量 -/// @param nreq 请求数量 -/// @param req_lens 每个请求的 token 数量 -/// @param req_pos 每个请求的起始位置 -/// @param kv_caches 每个请求的 KV Cache -/// @param temperature 采样温度（0. 表示贪心采样） -/// @param topk 采样 topk（1 表示贪心采样） -/// @param topp 采样 topp -/// @param output 输出 token 数组，每个请求一个输出，长度至少为nreq -__INFINI_C __export void -inferBatchJiuge(struct JiugeModel *, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output); - -/// @brief 批次推理一轮，输出 output embedding 后的 logits -/// @param tokens 输入 token 地址 -/// @param ntok 输入 token 数量 -/// @param nreq 请求数量 -/// @param req_lens 每个请求的 token 数量 -/// @param req_pos 每个请求的起始位置 -/// @param kv_caches 每个请求的 KV Cache -/// @param logits 输出 token 数组，每个请求一个输出，长度至少为nreq -__INFINI_C __export void -forwardBatchJiuge(struct JiugeModel *, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - void *logits); - -#endif diff --git a/include/infinicore_infer/models/jiuge_awq.h b/include/infinicore_infer/models/jiuge_awq.h deleted file mode 100644 index 631eb3cee..000000000 --- a/include/infinicore_infer/models/jiuge_awq.h +++ /dev/null @@ -1,79 +0,0 @@ -#ifndef MODEL_JIUGE_AWQ_H -#define MODEL_JIUGE_AWQ_H - -#include -#include -#include - -#include - -#include "../weights_loader.h" - -struct JiugeAWQModel; - -typedef struct -{ - infiniDtype_t dt_logits; - infiniDtype_t dt_linear_w; - infiniDtype_t dt_norm_w; - size_t nlayer, d, nh, nkvh, dh, di, dctx, dvoc; - float epsilon, theta; - uint32_t end_token; - size_t nbit; - size_t quant_group_size; - char has_qkv_bias; -} JiugeAWQMeta; - -//////////////////// APIs /////////////////////// -__INFINI_C __export struct ModelWeights * -createJiugeAWQWeights(const JiugeAWQMeta *, - infiniDevice_t device, - int ndev, - const int *dev_ids); -/// @brief 创建模型 -/// @param device 协处理器种类 -/// @param ndev 协处理器数量 -/// @param dev_ids 协处理器编号，长度为 ndev -__INFINI_C __export struct JiugeAWQModel * -createJiugeAWQModel(const JiugeAWQMeta *, - const ModelWeights *); - -/// @brief 销毁模型 -__INFINI_C __export void -destroyJiugeAWQModel(struct JiugeAWQModel *); - -/// @brief 批次推理一轮，并采样出新的 token -/// @param tokens 输入 token 地址 -/// @param ntok 输入 token 数量 -/// @param nreq 请求数量 -/// @param req_lens 每个请求的 token 数量 -/// @param req_pos 每个请求的起始位置 -/// @param kv_caches 每个请求的 KV Cache -/// @param temperature 采样温度（0. 表示贪心采样） -/// @param topk 采样 topk（1 表示贪心采样） -/// @param topp 采样 topp -/// @param output 输出 token 数组，每个请求一个输出，长度至少为nreq -__INFINI_C __export void -inferBatchJiugeAWQ(struct JiugeAWQModel *, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output); - -/// @brief 批次推理一轮，输出 output embedding 后的 logits -/// @param tokens 输入 token 地址 -/// @param ntok 输入 token 数量 -/// @param nreq 请求数量 -/// @param req_lens 每个请求的 token 数量 -/// @param req_pos 每个请求的起始位置 -/// @param kv_caches 每个请求的 KV Cache -/// @param logits 输出 token 数组，每个请求一个输出，长度至少为nreq -__INFINI_C __export void -forwardBatchJiugeAWQ(struct JiugeAWQModel *, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - void *logits); - -#endif diff --git a/include/infinicore_infer/models/jiuge_gptq.h b/include/infinicore_infer/models/jiuge_gptq.h deleted file mode 100644 index eaf06a02c..000000000 --- a/include/infinicore_infer/models/jiuge_gptq.h +++ /dev/null @@ -1,79 +0,0 @@ -#ifndef MODEL_JIUGE_GPTQ_H -#define MODEL_JIUGE_GPTQ_H - -#include -#include -#include - -#include - -#include "../weights_loader.h" - -struct JiugeGPTQModel; - -typedef struct -{ - infiniDtype_t dt_logits; - infiniDtype_t dt_linear_w; - infiniDtype_t dt_norm_w; - size_t nlayer, d, nh, nkvh, dh, di, dctx, dvoc; - float epsilon, theta; - uint32_t end_token; - size_t nbit; - size_t quant_group_size; - char has_qkv_bias; -} JiugeGPTQMeta; - -//////////////////// APIs /////////////////////// -__INFINI_C __export struct ModelWeights * -createJiugeGPTQWeights(const JiugeGPTQMeta *, - infiniDevice_t device, - int ndev, - const int *dev_ids); -/// @brief 创建模型 -/// @param device 协处理器种类 -/// @param ndev 协处理器数量 -/// @param dev_ids 协处理器编号，长度为 ndev -__INFINI_C __export struct JiugeGPTQModel * -createJiugeGPTQModel(const JiugeGPTQMeta *, - const ModelWeights *); - -/// @brief 销毁模型 -__INFINI_C __export void -destroyJiugeGPTQModel(struct JiugeGPTQModel *); - -/// @brief 批次推理一轮，并采样出新的 token -/// @param tokens 输入 token 地址 -/// @param ntok 输入 token 数量 -/// @param nreq 请求数量 -/// @param req_lens 每个请求的 token 数量 -/// @param req_pos 每个请求的起始位置 -/// @param kv_caches 每个请求的 KV Cache -/// @param temperature 采样温度（0. 表示贪心采样） -/// @param topk 采样 topk（1 表示贪心采样） -/// @param topp 采样 topp -/// @param output 输出 token 数组，每个请求一个输出，长度至少为nreq -__INFINI_C __export void -inferBatchJiugeGPTQ(struct JiugeGPTQModel *, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output); - -/// @brief 批次推理一轮，输出 output embedding 后的 logits -/// @param tokens 输入 token 地址 -/// @param ntok 输入 token 数量 -/// @param nreq 请求数量 -/// @param req_lens 每个请求的 token 数量 -/// @param req_pos 每个请求的起始位置 -/// @param kv_caches 每个请求的 KV Cache -/// @param logits 输出 token 数组，每个请求一个输出，长度至少为nreq -__INFINI_C __export void -forwardBatchJiugeGPTQ(struct JiugeGPTQModel *, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - void *logits); - -#endif diff --git a/include/infinicore_infer/models/qwen3vl.h b/include/infinicore_infer/models/qwen3vl.h deleted file mode 100644 index 3f28f94e8..000000000 --- a/include/infinicore_infer/models/qwen3vl.h +++ /dev/null @@ -1,203 +0,0 @@ -#ifndef QWEN3VL_WEIGHTS_H -#define QWEN3VL_WEIGHTS_H - -#include -#include -#include - -#include -#include - -struct Qwen3vlWeights; - -// Function pointer signatures -typedef void (*qwen3vl_load_global_fn)(Qwen3vlWeights *, void *cpu_ptr); -typedef void (*qwen3vl_load_layer_fn)(Qwen3vlWeights *, void *cpu_ptr, size_t layer_id); -// Struct containing all weight loading functions -typedef struct { - // Global - qwen3vl_load_global_fn load_input_embd; - qwen3vl_load_global_fn load_output_norm; - qwen3vl_load_global_fn load_output_embd; - - // Attention - qwen3vl_load_layer_fn load_attn_norm; - qwen3vl_load_layer_fn load_attn_q_norm; - qwen3vl_load_layer_fn load_attn_k_norm; - qwen3vl_load_layer_fn load_attn_qkv_proj; - qwen3vl_load_layer_fn load_attn_o_proj; - - // MLP - qwen3vl_load_layer_fn load_mlp_norm; - qwen3vl_load_layer_fn load_mlp_gate_up; - qwen3vl_load_layer_fn load_mlp_down; - -} Qwen3vlLangWeightLoader; - -typedef struct { - // Patch_embed - qwen3vl_load_global_fn load_patch_embed_weight; - qwen3vl_load_global_fn load_patch_embed_bias; - qwen3vl_load_global_fn load_pos_embed_weight; - - // blocks attn - qwen3vl_load_layer_fn load_attn_proj_weight; - qwen3vl_load_layer_fn load_attn_proj_bias; - qwen3vl_load_layer_fn load_attn_qkv_weight; - qwen3vl_load_layer_fn load_attn_qkv_bias; - - // block mlp - qwen3vl_load_layer_fn load_mlp_linear_fc1_weight; - qwen3vl_load_layer_fn load_mlp_linear_fc1_bias; - qwen3vl_load_layer_fn load_mlp_linear_fc2_weight; - qwen3vl_load_layer_fn load_mlp_linear_fc2_bias; - - // block norm - qwen3vl_load_layer_fn load_norm1_weight; - qwen3vl_load_layer_fn load_norm1_bias; - qwen3vl_load_layer_fn load_norm2_weight; - qwen3vl_load_layer_fn load_norm2_bias; - - // deepstack_merger - qwen3vl_load_layer_fn load_deepstack_merger_linear_fc1_weight; - qwen3vl_load_layer_fn load_deepstack_merger_linear_fc1_bias; - qwen3vl_load_layer_fn load_deepstack_merger_linear_fc2_weight; - qwen3vl_load_layer_fn load_deepstack_merger_linear_fc2_bias; - qwen3vl_load_layer_fn load_deepstack_merger_norm_weight; - qwen3vl_load_layer_fn load_deepstack_merger_norm_bias; - - // merger - qwen3vl_load_global_fn load_merger_linear_fc1_weight; - qwen3vl_load_global_fn load_merger_linear_fc1_bias; - qwen3vl_load_global_fn load_merger_linear_fc2_weight; - qwen3vl_load_global_fn load_merger_linear_fc2_bias; - qwen3vl_load_global_fn load_merger_norm_weight; - qwen3vl_load_global_fn load_merger_norm_bias; - -} Qwen3vlVisWeightLoader; - -typedef struct { - Qwen3vlLangWeightLoader lang_loader; - Qwen3vlVisWeightLoader vis_loader; -} Qwen3vlWeightLoader; - -struct Qwen3vlModel; - -typedef struct { - size_t bos_token_id; - size_t eos_token_id; - size_t head_dim; - size_t hidden_size; - float initializer_range; - size_t intermediate_size; - size_t max_tokens; - size_t num_attention_heads; - size_t num_hidden_layers; - size_t num_key_value_heads; - float rms_norm_eps; - size_t mrope_section[3]; - size_t rope_theta; - size_t vocab_size; -} Qwen3vlTextMeta; - -typedef struct { - size_t depth; - size_t deepstack_visual_indexes[3]; - size_t hidden_size; - size_t in_channels; - float initializer_range; - size_t intermediate_size; - size_t num_heads; - size_t num_position_embeddings; - size_t out_hidden_size; - size_t patch_size; - size_t spatial_merge_size; - size_t temporal_patch_size; -} Qwen3vlVisMeta; - -typedef struct { - infiniDtype_t dtype; // INFINI_DTYPE_BF16 - - Qwen3vlTextMeta text_meta; - Qwen3vlVisMeta vis_meta; - - size_t image_token_id; - size_t video_token_id; - size_t vision_end_token_id; - size_t vision_start_token_id; -} Qwen3vlMeta; - -//////////////////// APIs /////////////////////// -/// @brief 创建模型 -/// @param device 协处理器种类 -/// @param ndev 协处理器数量 -/// @param dev_ids 协处理器编号，长度为 ndev -__INFINI_C __export struct Qwen3vlModel * -createQwen3vlModel(const Qwen3vlMeta *, - const Qwen3vlWeights *); - -__INFINI_C Qwen3vlWeights * -createQwen3vlWeights(const Qwen3vlMeta *meta, - infiniDevice_t device, - int ndev, - const int *dev_ids, - bool transpose_weight); - -__INFINI_C __export Qwen3vlWeightLoader * -createQwen3vlWeightLoader(); - -/// @brief 销毁模型 -__INFINI_C __export void destroyQwen3vlModel(struct Qwen3vlModel *); - -__INFINI_C __export struct Qwen3vlCache * -createQwen3vlCache(const struct Qwen3vlModel *); - -__INFINI_C __export void -dropQwen3vlCache(const struct Qwen3vlModel *, - struct Qwen3vlCache *); - -/// @brief 批次推理一轮，并采样出新的 token -/// @param tokens 输入 token 地址 -/// @param ntok 输入 token 数量 -/// @param nreq 请求数量 -/// @param req_lens 每个请求的 token 数量 -/// @param req_pos 每个请求的起始位置 -/// @param kv_caches 每个请求的 KV Cache -/// @param temperature 采样温度（0. 表示贪心采样） -/// @param topk 采样 topk（1 表示贪心采样） -/// @param topp 采样 topp -/// @param output 输出 token 数组，每个请求一个输出，长度至少为nreq -__INFINI_C __export void -inferBatchQwen3vl(struct Qwen3vlModel *, - const uint32_t *tokens, uint32_t ntok, - void *pixel_values, uint32_t total_patches, - uint32_t *image_grid_thw, uint32_t num_images, - void *pixel_values_videos, uint32_t total_patches_videos, - uint32_t *video_grid_thw, uint32_t num_videos, - uint32_t patch_features, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct Qwen3vlCache **caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output); - -/// @brief 批次推理一轮，输出 output embedding 后的 logits -/// @param tokens 输入 token 地址 -/// @param ntok 输入 token 数量 -/// @param nreq 请求数量 -/// @param req_lens 每个请求的 token 数量 -/// @param req_pos 每个请求的起始位置 -/// @param kv_caches 每个请求的 KV Cache -/// @param logits 输出 token 数组，每个请求一个输出，长度至少为nreq -__INFINI_C __export void -forwardBatchQwen3vl(struct Qwen3vlModel *, - const uint32_t *tokens, uint32_t ntok, - void *pixel_values, uint32_t total_patches, - uint32_t *image_grid_thw, uint32_t num_images, - void *pixel_values_videos, uint32_t total_patches_videos, - uint32_t *video_grid_thw, uint32_t num_videos, - uint32_t patch_features, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct Qwen3vlCache **caches, - void *logits); - -#endif // QWEN3VL_WEIGHTS_H diff --git a/include/infinicore_infer/weights_loader.h b/include/infinicore_infer/weights_loader.h deleted file mode 100644 index 82eafe595..000000000 --- a/include/infinicore_infer/weights_loader.h +++ /dev/null @@ -1,14 +0,0 @@ -#ifndef WEIGHTS_LOADER_H -#define WEIGHTS_LOADER_H - -#include - -struct ModelWeights; - -__INFINI_C __export void -loadModelWeight(struct ModelWeights *weights, const char *name, void *data); - -__INFINI_C __export void -loadModelWeightDistributed(struct ModelWeights *weights, const char *name, void *data, int *ranks, int nrank); - -#endif // WEIGHTS_LOADER_H diff --git a/pyproject.toml b/pyproject.toml index 8afedbf99..32c447f6b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,18 @@ name = "InfiniLM" version = "0.1.0" description = "InfiniLM model implementations" readme = "README.md" -dependencies = [] +dependencies = [ + "janus", + "numpy", + "Pillow", + "safetensors", + "tokenizers", + "torch", + "tqdm", + "transformers", + "typing-extensions", + "xxhash", +] requires-python = ">=3.10" classifiers = [ "Programming Language :: Python :: 3", diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py new file mode 100644 index 000000000..c7e17f6f2 --- /dev/null +++ b/python/infinicore/__init__.py @@ -0,0 +1,99 @@ +from . import context, nn, utils +from .context import ( + cancel_graph_recording, + get_device, + get_device_count, + get_stream, + is_graph_recording, + set_device, + start_graph_recording, + stop_graph_recording, + sync_device, + sync_stream, +) +from .device import device +from .device_event import DeviceEvent +from .dtype import ( + bfloat16, + double, + dtype, + float, + float16, + float32, + float64, + half, + int, + int8, + int16, + int32, + int64, + long, + short, + uint8, + uint16, + uint32, + uint64, +) +from .ops import add, cat, matmul +from .tensor import ( + Tensor, + empty, + empty_like, + from_blob, + from_list, + from_numpy, + from_torch, + strided_empty, + zeros, +) + +__all__ = [ + "context", + "nn", + "utils", + "Tensor", + "device", + "DeviceEvent", + "dtype", + "empty", + "empty_like", + "strided_empty", + "zeros", + "from_blob", + "from_list", + "from_numpy", + "from_torch", + "add", + "cat", + "matmul", + "get_device", + "get_device_count", + "get_stream", + "set_device", + "sync_device", + "sync_stream", + "is_graph_recording", + "start_graph_recording", + "stop_graph_recording", + "cancel_graph_recording", + "bfloat16", + "half", + "float16", + "float", + "float32", + "double", + "float64", + "short", + "int8", + "int16", + "int", + "int32", + "long", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", +] + +use_ntops = False diff --git a/python/infinicore/context.py b/python/infinicore/context.py new file mode 100644 index 000000000..c73425749 --- /dev/null +++ b/python/infinicore/context.py @@ -0,0 +1,46 @@ +from infinicore.device import device +from infinicore.lib import _infinicore + + +def get_device(): + return device._from_underlying(_infinicore.get_device()) + + +def get_device_count(device_type): + return _infinicore.get_device_count(device(device_type)._underlying.type) + + +def set_device(target): + _infinicore.set_device(device(target)._underlying) + + +def get_stream(): + return _infinicore.get_stream() + + +def sync_stream(): + _infinicore.sync_stream() + + +def sync_device(): + _infinicore.sync_device() + + +def is_graph_recording(): + return _infinicore.is_graph_recording() + + +def start_graph_recording(target=None): + if target is not None: + set_device(target) + _infinicore.start_graph_recording() + + +def stop_graph_recording(): + from infinicore.graph import Graph + + return Graph(_infinicore.stop_graph_recording()) + + +def cancel_graph_recording(): + _infinicore.cancel_graph_recording() diff --git a/python/infinicore/device.py b/python/infinicore/device.py new file mode 100644 index 000000000..58f8069cf --- /dev/null +++ b/python/infinicore/device.py @@ -0,0 +1,67 @@ +import builtins + +from infinicore.lib import _infinicore + +_TYPE_ALIASES = { + "cpu": (_infinicore.Device.Type.CPU, "cpu"), + "cuda": (_infinicore.Device.Type.NVIDIA, "cuda"), + "nvidia": (_infinicore.Device.Type.NVIDIA, "cuda"), + "cambricon": (_infinicore.Device.Type.CAMBRICON, "mlu"), + "mlu": (_infinicore.Device.Type.CAMBRICON, "mlu"), + "ascend": (_infinicore.Device.Type.ASCEND, "npu"), + "npu": (_infinicore.Device.Type.ASCEND, "npu"), + "metax": (_infinicore.Device.Type.METAX, "metax"), + "moore": (_infinicore.Device.Type.MOORE, "musa"), + "musa": (_infinicore.Device.Type.MOORE, "musa"), + "iluvatar": (_infinicore.Device.Type.ILUVATAR, "iluvatar"), + "hygon": (_infinicore.Device.Type.HYGON, "hygon"), +} + +_NATIVE_TO_NAME = {native: name for native, name in set(_TYPE_ALIASES.values())} + + +class device: + def __init__(self, type="cpu", index=None): + if isinstance(type, device): + self.type = type.type + self.index = type.index + self._underlying = type._underlying + return + + if not isinstance(type, str): + raise TypeError("device type must be a string or infinicore.device") + if ":" in type: + if index is not None: + raise ValueError("index must not be provided twice") + type, raw_index = type.rsplit(":", 1) + index = builtins.int(raw_index) + + try: + native_type, canonical_name = _TYPE_ALIASES[type.lower()] + except KeyError as error: + raise ValueError(f"unsupported device type: {type}") from error + + self.type = canonical_name + self.index = 0 if index is None else builtins.int(index) + self._underlying = _infinicore.Device(native_type, self.index) + + @classmethod + def _from_underlying(cls, underlying): + instance = cls.__new__(cls) + instance.type = _NATIVE_TO_NAME[underlying.type] + instance.index = underlying.index + instance._underlying = underlying + return instance + + def __repr__(self): + return f"device(type='{self.type}', index={self.index})" + + def __str__(self): + return f"{self.type}:{self.index}" + + def __eq__(self, other): + return ( + isinstance(other, device) + and self._underlying.type == other._underlying.type + and self.index == other.index + ) diff --git a/python/infinicore/device_event.py b/python/infinicore/device_event.py new file mode 100644 index 000000000..edc0ed0c0 --- /dev/null +++ b/python/infinicore/device_event.py @@ -0,0 +1,38 @@ +from infinicore.device import device +from infinicore.lib import _infinicore + + +class DeviceEvent: + def __init__(self, target=None, flags=None): + args = [] + if target is not None: + args.append(device(target)._underlying) + if flags is not None: + args.append(flags) + self._underlying = _infinicore.DeviceEvent(*args) + + @property + def device(self): + return device._from_underlying(self._underlying.device) + + @property + def is_recorded(self): + return self._underlying.is_recorded + + def record(self, stream=None): + if stream is None: + self._underlying.record() + else: + self._underlying.record(stream) + + def synchronize(self): + self._underlying.synchronize() + + def query(self): + return self._underlying.query() + + def elapsed_time(self, other): + return self._underlying.elapsed_time(other._underlying) + + def wait(self, stream=0): + self._underlying.wait(stream) diff --git a/python/infinicore/dtype.py b/python/infinicore/dtype.py new file mode 100644 index 000000000..a3ade4319 --- /dev/null +++ b/python/infinicore/dtype.py @@ -0,0 +1,51 @@ +from infinicore.lib import _infinicore + + +class dtype: + def __init__(self, underlying): + self._underlying = underlying + + def __repr__(self): + return _DTYPE_NAMES[self._underlying] + + def __eq__(self, other): + return isinstance(other, dtype) and self._underlying == other._underlying + + def __hash__(self): + return hash(self._underlying) + + +_DTYPE_NAMES = { + _infinicore.DataType.INT8: "int8", + _infinicore.DataType.INT16: "int16", + _infinicore.DataType.INT32: "int32", + _infinicore.DataType.INT64: "int64", + _infinicore.DataType.UINT8: "uint8", + _infinicore.DataType.UINT16: "uint16", + _infinicore.DataType.UINT32: "uint32", + _infinicore.DataType.UINT64: "uint64", + _infinicore.DataType.FLOAT16: "float16", + _infinicore.DataType.BFLOAT16: "bfloat16", + _infinicore.DataType.FLOAT32: "float32", + _infinicore.DataType.FLOAT64: "float64", +} + +int8 = dtype(_infinicore.DataType.INT8) +int16 = dtype(_infinicore.DataType.INT16) +int32 = dtype(_infinicore.DataType.INT32) +int64 = dtype(_infinicore.DataType.INT64) +uint8 = dtype(_infinicore.DataType.UINT8) +uint16 = dtype(_infinicore.DataType.UINT16) +uint32 = dtype(_infinicore.DataType.UINT32) +uint64 = dtype(_infinicore.DataType.UINT64) +float16 = dtype(_infinicore.DataType.FLOAT16) +bfloat16 = dtype(_infinicore.DataType.BFLOAT16) +float32 = dtype(_infinicore.DataType.FLOAT32) +float64 = dtype(_infinicore.DataType.FLOAT64) + +half = float16 +float = float32 +double = float64 +short = int16 +int = int32 +long = int64 diff --git a/python/infinicore/graph.py b/python/infinicore/graph.py new file mode 100644 index 000000000..ecb2cddea --- /dev/null +++ b/python/infinicore/graph.py @@ -0,0 +1,6 @@ +class Graph: + def __init__(self, underlying): + self._underlying = underlying + + def run(self): + self._underlying.run() diff --git a/python/infinicore/lib/__init__.py b/python/infinicore/lib/__init__.py new file mode 100644 index 000000000..c00c79e94 --- /dev/null +++ b/python/infinicore/lib/__init__.py @@ -0,0 +1,3 @@ +from . import _infinicore + +__all__ = ["_infinicore"] diff --git a/python/infinicore/nn/__init__.py b/python/infinicore/nn/__init__.py new file mode 100644 index 000000000..de5d39589 --- /dev/null +++ b/python/infinicore/nn/__init__.py @@ -0,0 +1,5 @@ +from . import functional +from .modules import Module +from .parameter import Parameter + +__all__ = ["functional", "Module", "Parameter"] diff --git a/python/infinicore/nn/functional/__init__.py b/python/infinicore/nn/functional/__init__.py new file mode 100644 index 000000000..6d229738f --- /dev/null +++ b/python/infinicore/nn/functional/__init__.py @@ -0,0 +1,19 @@ +from infinicore.ops import ( + causal_softmax, + embedding, + linear, + random_sample, + rms_norm, + silu, + silu_and_mul, +) + +__all__ = [ + "causal_softmax", + "embedding", + "linear", + "random_sample", + "rms_norm", + "silu", + "silu_and_mul", +] diff --git a/python/infinicore/nn/modules/__init__.py b/python/infinicore/nn/modules/__init__.py new file mode 100644 index 000000000..697d34c93 --- /dev/null +++ b/python/infinicore/nn/modules/__init__.py @@ -0,0 +1,3 @@ +from .module import Module + +__all__ = ["Module"] diff --git a/python/infinicore/nn/modules/module.py b/python/infinicore/nn/modules/module.py new file mode 100644 index 000000000..0843ce440 --- /dev/null +++ b/python/infinicore/nn/modules/module.py @@ -0,0 +1,8 @@ +class Module: + """Compatibility base for InfiniLM Python-side type annotations.""" + + def __call__(self, *args, **kwargs): + return self.forward(*args, **kwargs) + + def forward(self, *args, **kwargs): + raise NotImplementedError diff --git a/python/infinicore/nn/parameter.py b/python/infinicore/nn/parameter.py new file mode 100644 index 000000000..8ec15abd5 --- /dev/null +++ b/python/infinicore/nn/parameter.py @@ -0,0 +1,8 @@ +from infinicore.tensor import Tensor + + +class Parameter(Tensor): + def __init__(self, data): + if not isinstance(data, Tensor): + raise TypeError("Parameter data must be an infinicore.Tensor") + super().__init__(data._underlying, owner=data) diff --git a/python/infinicore/ops/__init__.py b/python/infinicore/ops/__init__.py new file mode 100644 index 000000000..8df50d3a5 --- /dev/null +++ b/python/infinicore/ops/__init__.py @@ -0,0 +1,81 @@ +from infinicore.lib import _infinicore +from infinicore.tensor import Tensor + + +def _unwrap(value): + return value._underlying + + +def add(input, other, *, out=None): + if out is None: + return Tensor(_infinicore.add(_unwrap(input), _unwrap(other))) + _infinicore.add_(_unwrap(out), _unwrap(input), _unwrap(other)) + return out + + +def cat(tensors, dim=0, *, out=None): + underlying = [_unwrap(tensor) for tensor in tensors] + if out is None: + return Tensor(_infinicore.cat(underlying, dim)) + _infinicore.cat_(_unwrap(out), underlying, dim) + return out + + +def matmul(input, other, *, alpha=1.0, out=None): + if out is None: + return Tensor(_infinicore.matmul(_unwrap(input), _unwrap(other), alpha)) + _infinicore.matmul_(_unwrap(out), _unwrap(input), _unwrap(other), alpha) + return out + + +def linear(input, weight, bias=None, *, alpha=1.0, out=None): + raw_bias = None if bias is None else _unwrap(bias) + if out is None: + return Tensor( + _infinicore.linear(_unwrap(input), _unwrap(weight), raw_bias, alpha) + ) + _infinicore.linear_(_unwrap(out), _unwrap(input), _unwrap(weight), raw_bias, alpha) + return out + + +def embedding(input, weight, *, out=None): + if out is None: + return Tensor(_infinicore.embedding(_unwrap(input), _unwrap(weight))) + _infinicore.embedding_(_unwrap(out), _unwrap(input), _unwrap(weight)) + return out + + +def rms_norm(input, weight, epsilon=1e-5, *, out=None): + if out is None: + return Tensor(_infinicore.rms_norm(_unwrap(input), _unwrap(weight), epsilon)) + _infinicore.rms_norm_(_unwrap(out), _unwrap(input), _unwrap(weight), epsilon) + return out + + +def causal_softmax(input, *, out=None): + if out is None: + return Tensor(_infinicore.causal_softmax(_unwrap(input))) + _infinicore.causal_softmax_(_unwrap(out), _unwrap(input)) + return out + + +def silu(input, *, out=None): + if out is None: + return Tensor(_infinicore.silu(_unwrap(input))) + _infinicore.silu_(_unwrap(out), _unwrap(input)) + return out + + +def silu_and_mul(input, *, out=None): + if out is None: + return Tensor(_infinicore.silu_and_mul(_unwrap(input))) + _infinicore.silu_and_mul_(_unwrap(out), _unwrap(input)) + return out + + +def random_sample(logits, random_val, topp, topk, temperature, *, out=None): + arguments = (_unwrap(logits), random_val, topp, topk, temperature) + if out is None: + return Tensor(_infinicore.random_sample(*arguments)) + _infinicore.random_sample_(_unwrap(out), *arguments) + return out diff --git a/python/infinicore/tensor.py b/python/infinicore/tensor.py new file mode 100644 index 000000000..4f1c81cea --- /dev/null +++ b/python/infinicore/tensor.py @@ -0,0 +1,253 @@ +from infinicore.device import device as Device +from infinicore.dtype import dtype as DType +from infinicore.lib import _infinicore + + +class Tensor: + def __init__(self, underlying, *, owner=None): + self._underlying = underlying + self._owner = owner + + @property + def shape(self): + return self._underlying.shape + + @property + def strides(self): + return self._underlying.strides + + @property + def ndim(self): + return self._underlying.ndim + + @property + def dtype(self): + return DType(self._underlying.dtype) + + @property + def device(self): + return Device._from_underlying(self._underlying.device) + + def data_ptr(self): + return self._underlying.data_ptr() + + def size(self, dim=None): + return self.shape if dim is None else self._underlying.size(dim) + + def stride(self, dim=None): + return self.strides if dim is None else self._underlying.stride(dim) + + def numel(self): + return self._underlying.numel() + + def is_contiguous(self): + return self._underlying.is_contiguous() + + def is_pinned(self): + return self._underlying.is_pinned() + + def copy_(self, source): + self._underlying.copy_(source._underlying) + return self + + def to(self, target=None, *, device=None): + if target is not None and device is not None: + raise TypeError("device must be provided once") + target = device if device is not None else target + if target is None: + return self + return Tensor(self._underlying.to(Device(target)._underlying), owner=self) + + def to_numpy(self): + """Return an owning, C-contiguous NumPy copy of this tensor.""" + import ctypes + + import numpy as np + + from infinicore.utils import infinicore_to_numpy_dtype + + source = self if self.device.type == "cpu" else self.to(Device("cpu")) + if not source.is_contiguous(): + source = source.contiguous() + + result = np.empty(source.shape, dtype=infinicore_to_numpy_dtype(source.dtype)) + if result.nbytes: + ctypes.memmove(result.ctypes.data, source.data_ptr(), result.nbytes) + return result + + def contiguous(self): + return Tensor(self._underlying.contiguous(), owner=self._owner) + + def as_strided(self, size, stride): + return Tensor(self._underlying.as_strided(size, stride), owner=self._owner) + + def narrow(self, dim, start, length): + return Tensor(self._underlying.narrow(dim, start, length), owner=self._owner) + + def permute(self, dims): + return Tensor(self._underlying.permute(dims), owner=self._owner) + + def view(self, shape): + return Tensor(self._underlying.view(shape), owner=self._owner) + + def squeeze(self, dim): + return Tensor(self._underlying.squeeze(dim), owner=self._owner) + + def unsqueeze(self, dim): + return Tensor(self._underlying.unsqueeze(dim), owner=self._owner) + + def debug(self, filename=None): + if filename is None: + return self._underlying.debug() + return self._underlying.debug(filename) + + def __add__(self, other): + from infinicore.ops import add + + return add(self, other) + + def __iadd__(self, other): + from infinicore.ops import add + + return add(self, other, out=self) + + def __matmul__(self, other): + from infinicore.ops import matmul + + return matmul(self, other) + + def __repr__(self): + return repr(self._underlying) + + +def _normalize_factory_args(dtype, device): + from infinicore import float32 + + return ( + float32 if dtype is None else dtype, + Device("cpu") if device is None else Device(device), + ) + + +def empty(size, *, dtype=None, device=None, pin_memory=False): + dtype, device = _normalize_factory_args(dtype, device) + return Tensor( + _infinicore.empty(size, dtype._underlying, device._underlying, pin_memory) + ) + + +def empty_like(input, *, dtype=None, device=None): + return empty( + input.shape, + dtype=input.dtype if dtype is None else dtype, + device=input.device if device is None else device, + ) + + +def strided_empty(size, strides, *, dtype=None, device=None, pin_memory=False): + dtype, device = _normalize_factory_args(dtype, device) + return Tensor( + _infinicore.strided_empty( + size, strides, dtype._underlying, device._underlying, pin_memory + ) + ) + + +def zeros(size, *, dtype=None, device=None, pin_memory=False): + dtype, device = _normalize_factory_args(dtype, device) + return Tensor( + _infinicore.zeros(size, dtype._underlying, device._underlying, pin_memory) + ) + + +def from_blob(data_ptr, size, *, dtype, device): + return Tensor( + _infinicore.from_blob( + data_ptr, size, dtype._underlying, Device(device)._underlying + ) + ) + + +def from_torch(torch_tensor, *, device=None): + import torch + + from infinicore.context import get_device, sync_stream + from infinicore.utils import to_infinicore_dtype + + owner = torch_tensor.detach().contiguous() + if owner.device.type not in ("cpu", "cuda"): + raise ValueError( + "from_torch() currently supports CPU and CUDA-compatible tensors" + ) + + index = 0 if owner.device.index is None else owner.device.index + + if device is None: + if owner.device.type == "cuda": + current = get_device() + if current.type not in ("cpu", "cuda"): + raise ValueError( + "Torch reports a CUDA-compatible device without its vendor; " + "pass device= explicitly" + ) + target = Device("nvidia", index) + else: + target = Device(owner.device.type, index) + elif isinstance(device, str) and ":" not in device: + target = Device(device, index) + else: + target = Device(device) + + if target.index != index: + raise ValueError("device index must match the source Torch tensor") + if owner.device.type == "cuda": + if target.type not in ("cuda", "metax", "iluvatar", "hygon"): + raise ValueError("device is not compatible with a Torch CUDA tensor") + elif target.type != owner.device.type: + raise ValueError("device type must match the source Torch tensor") + + dtype = to_infinicore_dtype(owner.dtype) + borrowed = Tensor( + _infinicore.from_blob( + owner.data_ptr(), + list(owner.shape), + dtype._underlying, + target._underlying, + ), + owner=owner, + ) + result = empty(list(owner.shape), dtype=dtype, device=target) + + if owner.device.type == "cuda": + torch.cuda.synchronize(owner.device) + + result.copy_(borrowed) + if owner.device.type == "cuda": + sync_stream() + return result + + +def from_numpy(array, *, dtype=None, device=None): + import numpy as np + + from infinicore.utils import ( + infinicore_to_numpy_dtype, + numpy_to_infinicore_dtype, + ) + + source = np.asarray(array) + if dtype is None: + dtype = numpy_to_infinicore_dtype(source.dtype) + numpy_dtype = infinicore_to_numpy_dtype(dtype) + owner = np.ascontiguousarray(array, dtype=numpy_dtype) + + result = Tensor(_infinicore._from_numpy_copy(owner, dtype._underlying)) + cpu = Device("cpu") + target = cpu if device is None else Device(device) + return result if target == cpu else result.to(target) + + +def from_list(data, *, dtype): + if dtype is None: + raise TypeError("from_list() requires dtype") + return Tensor(_infinicore.from_list(data, dtype._underlying)) diff --git a/python/infinicore/utils.py b/python/infinicore/utils.py new file mode 100644 index 000000000..f04870409 --- /dev/null +++ b/python/infinicore/utils.py @@ -0,0 +1,98 @@ +def to_torch_dtype(value): + import torch + + import infinicore + + mapping = { + infinicore.int8: torch.int8, + infinicore.int16: torch.int16, + infinicore.int32: torch.int32, + infinicore.int64: torch.int64, + infinicore.uint8: torch.uint8, + infinicore.float16: torch.float16, + infinicore.bfloat16: torch.bfloat16, + infinicore.float32: torch.float32, + infinicore.float64: torch.float64, + } + for name in ("uint16", "uint32", "uint64"): + torch_dtype = getattr(torch, name, None) + if torch_dtype is not None: + mapping[getattr(infinicore, name)] = torch_dtype + try: + return mapping[value] + except KeyError as error: + raise ValueError(f"unsupported infinicore dtype: {value}") from error + + +def to_infinicore_dtype(value): + import torch + + import infinicore + + mapping = { + torch.int8: infinicore.int8, + torch.int16: infinicore.int16, + torch.int32: infinicore.int32, + torch.int64: infinicore.int64, + torch.uint8: infinicore.uint8, + torch.float16: infinicore.float16, + torch.bfloat16: infinicore.bfloat16, + torch.float32: infinicore.float32, + torch.float64: infinicore.float64, + } + for name in ("uint16", "uint32", "uint64"): + torch_dtype = getattr(torch, name, None) + if torch_dtype is not None: + mapping[torch_dtype] = getattr(infinicore, name) + try: + return mapping[value] + except KeyError as error: + raise ValueError(f"unsupported torch dtype: {value}") from error + + +def numpy_to_infinicore_dtype(value): + import numpy as np + + import infinicore + + mapping = { + np.dtype("int8"): infinicore.int8, + np.dtype("int16"): infinicore.int16, + np.dtype("int32"): infinicore.int32, + np.dtype("int64"): infinicore.int64, + np.dtype("uint8"): infinicore.uint8, + np.dtype("uint16"): infinicore.uint16, + np.dtype("uint32"): infinicore.uint32, + np.dtype("uint64"): infinicore.uint64, + np.dtype("float16"): infinicore.float16, + np.dtype("float32"): infinicore.float32, + np.dtype("float64"): infinicore.float64, + } + try: + return mapping[np.dtype(value)] + except KeyError as error: + raise ValueError(f"unsupported numpy dtype: {value}") from error + + +def infinicore_to_numpy_dtype(value): + import numpy as np + + import infinicore + + mapping = { + infinicore.int8: np.dtype("int8"), + infinicore.int16: np.dtype("int16"), + infinicore.int32: np.dtype("int32"), + infinicore.int64: np.dtype("int64"), + infinicore.uint8: np.dtype("uint8"), + infinicore.uint16: np.dtype("uint16"), + infinicore.uint32: np.dtype("uint32"), + infinicore.uint64: np.dtype("uint64"), + infinicore.float16: np.dtype("float16"), + infinicore.float32: np.dtype("float32"), + infinicore.float64: np.dtype("float64"), + } + try: + return mapping[value] + except KeyError as error: + raise ValueError(f"unsupported infinicore dtype: {value}") from error diff --git a/python/infinilm/__init__.py b/python/infinilm/__init__.py index f552a2cc9..687606313 100644 --- a/python/infinilm/__init__.py +++ b/python/infinilm/__init__.py @@ -1,19 +1,13 @@ -from .models import AutoLlamaModel -from . import distributed -from . import cache -from . import llm -from . import base_config - +from . import base_config, cache, distributed, llm from .llm import ( LLM, AsyncLLMEngine, - SamplingParams, RequestOutput, + SamplingParams, TokenOutput, ) __all__ = [ - "AutoLlamaModel", "distributed", "cache", "llm", diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 3b8f6533d..05bcd76ad 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -193,8 +193,8 @@ def _add_common_args(self): type=str, default="auto", help=( - "device platform: auto, cpu, nvidia, qy, metax, moore, iluvatar, " - "ali, cambricon, ascend, kunlun, hygon, or backend name " + "device platform: auto, cpu, nvidia, metax, moore, iluvatar, " + "cambricon, ascend, hygon, or backend name " "(cuda/mlu/musa/npu)" ), ) @@ -491,7 +491,6 @@ def detect_device(self): ("moore", ["mthreads-gmi"]), ("metax", ["mx-smi", "ht-smi"]), ("hygon", ["hy-smi"]), - ("ali", ["ppu-smi"]), ("iluvatar", ["ixsmi"]), ("nvidia", ["nvidia-smi"]), ] @@ -513,21 +512,24 @@ def get_device_str(self, device): "musa": "musa", "npu": "npu", "nvidia": "cuda", - "qy": "cuda", "cambricon": "mlu", "ascend": "npu", "metax": "cuda", "moore": "musa", "iluvatar": "cuda", - "kunlun": "cuda", "hygon": "cuda", - "ali": "cuda", } device = device.lower() if device == "auto": device = self.detect_device() print(f"Auto-detected device platform: {device}") - return DEVICE_STR_MAP.get(device, "cpu") + try: + return DEVICE_STR_MAP[device] + except KeyError: + supported = ", ".join(sorted(DEVICE_STR_MAP)) + raise ValueError( + f"unsupported device platform {device!r}; expected one of: {supported}" + ) from None def __repr__(self): """String representation of configuration""" diff --git a/python/infinilm/exception_utils.py b/python/infinilm/exception_utils.py index fef44d248..b66ac54b4 100644 --- a/python/infinilm/exception_utils.py +++ b/python/infinilm/exception_utils.py @@ -1,5 +1,5 @@ -import os import logging +import os from typing import Iterator logger = logging.getLogger(__name__) @@ -42,9 +42,8 @@ def is_oom_exception(e: BaseException) -> bool: # Common patterns observed for allocator failures. # Keep this allowlist small to avoid hard-exiting on unrelated errors. patterns = ( - # MetaX / infinirt allocator + # MetaX allocator "hcmalloc", - "infinirtmalloc", "out of memory", # CUDA / driver / runtime alloc failures "cuda out of memory", diff --git a/python/infinilm/generation/__init__.py b/python/infinilm/generation/__init__.py new file mode 100644 index 000000000..5878aa28c --- /dev/null +++ b/python/infinilm/generation/__init__.py @@ -0,0 +1 @@ +"""Generation helpers.""" diff --git a/python/infinilm/generation/utils.py b/python/infinilm/generation/utils.py index bad9c2613..750d4d418 100644 --- a/python/infinilm/generation/utils.py +++ b/python/infinilm/generation/utils.py @@ -1,47 +1,13 @@ import time from typing import Optional -import infinicore -from ..cache_utils import Cache, DynamicCache -import numpy as np - -def infini_to_ctype_dtype(infini_dtype): - """Convert PyTorch data type to infinicore data type""" - import ctypes +import infinicore - if infini_dtype == infinicore.int32: - return ctypes.c_int32 - elif infini_dtype == infinicore.float32: - return ctypes.c_float - elif infini_dtype == infinicore.int64: - return ctypes.c_int64 - else: - raise ValueError(f"Unsupported py_dtype: {infini_dtype}") +from ..cache_utils import Cache, DynamicCache def infini_to_numpy(infini_tensor: infinicore.Tensor): - if infini_tensor.device.type != "cpu": - infini_tensor_cpu = infini_tensor.to(infinicore.device("cpu", 0)) - else: - infini_tensor_cpu = infini_tensor - - # 获取数据指针和形状信息 - data_ptr = infini_tensor_cpu.data_ptr() - num_elements = infini_tensor_cpu.numel() - original_shape = infini_tensor_cpu.shape - - # 创建1D NumPy数组（共享内存） - ArrayType = infini_to_ctype_dtype(infini_tensor_cpu.dtype) * num_elements - array = ArrayType.from_address(data_ptr) - np_flat = np.ctypeslib.as_array(array) - - # 重塑为原始形状 - np_array = np_flat.reshape(original_shape) - - return np.copy(np_array) - - -infinicore.Tensor.to_numpy = infini_to_numpy + return infini_tensor.to_numpy() class GenerationMixin: diff --git a/python/infinilm/kv_connector/mooncake/__init__.py b/python/infinilm/kv_connector/mooncake/__init__.py new file mode 100644 index 000000000..202d934b5 --- /dev/null +++ b/python/infinilm/kv_connector/mooncake/__init__.py @@ -0,0 +1 @@ +"""Mooncake KV connector package.""" diff --git a/python/infinilm/lib/__init__.py b/python/infinilm/lib/__init__.py index 67c9ce400..29694ac97 100644 --- a/python/infinilm/lib/__init__.py +++ b/python/infinilm/lib/__init__.py @@ -1,19 +1,5 @@ -""" -InfiniLM C++ extension module -""" +import infinicore # noqa: F401 - register shared pybind types first -import sys -import os -from pathlib import Path - -# Ensure the directory containing this __init__.py is on sys.path -# This allows importing the .so file from the same directory -_lib_dir = Path(__file__).parent -if str(_lib_dir) not in sys.path: - sys.path.insert(0, str(_lib_dir)) - -# Import the compiled C++ module -# The .so file should be installed in this directory by xmake -import _infinilm +from . import _infinilm __all__ = ["_infinilm"] diff --git a/python/infinilm/llm/model_runner/__init__.py b/python/infinilm/llm/model_runner/__init__.py new file mode 100644 index 000000000..a57d3af40 --- /dev/null +++ b/python/infinilm/llm/model_runner/__init__.py @@ -0,0 +1 @@ +"""Model runner package.""" diff --git a/python/infinilm/models/__init__.py b/python/infinilm/models/__init__.py index 9e3f2adf5..a9a2c5b3b 100644 --- a/python/infinilm/models/__init__.py +++ b/python/infinilm/models/__init__.py @@ -1,3 +1 @@ -from .llama import AutoLlamaModel - -__all__ = ["AutoLlamaModel"] +__all__ = [] diff --git a/python/infinilm/multimodal/__init__.py b/python/infinilm/multimodal/__init__.py new file mode 100644 index 000000000..505ce9efb --- /dev/null +++ b/python/infinilm/multimodal/__init__.py @@ -0,0 +1 @@ +"""Multimodal input helpers.""" diff --git a/python/infinilm/server/__init__.py b/python/infinilm/server/__init__.py new file mode 100644 index 000000000..f4f5bc6d5 --- /dev/null +++ b/python/infinilm/server/__init__.py @@ -0,0 +1 @@ +"""Inference server package.""" diff --git a/scripts/deepseek.py b/scripts/deepseek.py deleted file mode 100644 index 29d81ce15..000000000 --- a/scripts/deepseek.py +++ /dev/null @@ -1,775 +0,0 @@ -import ctypes -from typing import List, Sequence - -from tqdm import tqdm - -from libinfinicore_infer import ( - DeepSeekV3Model, - DeepSeekV3MetaCStruct, - DeepSeekV3CacheCStruct, - DataType, - DeviceType, -) -from infer_task import InferTask, KVCache - -from ctypes import POINTER, c_float, c_int, c_uint, c_void_p, byref -import os -from pathlib import Path -import safetensors -import sys -import time -import json -import math -import torch -import transformers - -torch.set_default_device("cpu") - - -class DeepseekR1WeightsNaming: - def __init__(self, dense_replace=3): - self.dense_replace = dense_replace - - def input_embd(self): - return "model.embed_tokens.weight" - - def output_norm(self): - return "model.norm.weight" - - def output_embd(self): - return "lm_head.weight" - - # MLA - def attn_norm(self, i): - return f"model.layers.{i}.input_layernorm.weight" - - def attn_kv_a_layernorm(self, i): - return f"model.layers.{i}.self_attn.kv_a_layernorm.weight" - - def attn_kv_a_proj_with_mqa_weight(self, i): - return f"model.layers.{i}.self_attn.kv_a_proj_with_mqa.qweight" - - def attn_kv_a_proj_with_mqa_scale(self, i): - return f"model.layers.{i}.self_attn.kv_a_proj_with_mqa.scales" - - def attn_kv_a_proj_with_mqa_zero(self, i): - return f"model.layers.{i}.self_attn.kv_a_proj_with_mqa.qzeros" - - def attn_kv_b_proj_weight(self, i): - return f"model.layers.{i}.self_attn.kv_b_proj.qweight" - - def attn_kv_b_proj_scale(self, i): - return f"model.layers.{i}.self_attn.kv_b_proj.scales" - - def attn_kv_b_proj_zero(self, i): - return f"model.layers.{i}.self_attn.kv_b_proj.qzeros" - - def attn_o_proj_weight(self, i): - return f"model.layers.{i}.self_attn.o_proj.qweight" - - def attn_o_proj_scale(self, i): - return f"model.layers.{i}.self_attn.o_proj.scales" - - def attn_o_proj_zero(self, i): - return f"model.layers.{i}.self_attn.o_proj.qzeros" - - def attn_q_a_layernorm(self, i): - return f"model.layers.{i}.self_attn.q_a_layernorm.weight" - - def attn_q_a_proj_weight(self, i): - return f"model.layers.{i}.self_attn.q_a_proj.qweight" - - def attn_q_a_proj_scale(self, i): - return f"model.layers.{i}.self_attn.q_a_proj.scales" - - def attn_q_a_proj_zero(self, i): - return f"model.layers.{i}.self_attn.q_a_proj.qzeros" - - def attn_q_b_proj_weight(self, i): - return f"model.layers.{i}.self_attn.q_b_proj.qweight" - - def attn_q_b_proj_scale(self, i): - return f"model.layers.{i}.self_attn.q_b_proj.scales" - - def attn_q_b_proj_zero(self, i): - return f"model.layers.{i}.self_attn.q_b_proj.qzeros" - - # MLP - - def mlp_norm(self, i): - return f"model.layers.{i}.post_attention_layernorm.weight" - - # First self.dense_replace layers are dense - def mlp_down_proj_weight(self, i): - assert i < self.dense_replace - return f"model.layers.{i}.mlp.down_proj.qweight" - - def mlp_down_proj_scale(self, i): - assert i < self.dense_replace - return f"model.layers.{i}.mlp.down_proj.scales" - - def mlp_down_proj_zero(self, i): - assert i < self.dense_replace - return f"model.layers.{i}.mlp.down_proj.qzeros" - - def mlp_up_proj_weight(self, i): - assert i < self.dense_replace - return f"model.layers.{i}.mlp.up_proj.qweight" - - def mlp_up_proj_scale(self, i): - assert i < self.dense_replace - return f"model.layers.{i}.mlp.up_proj.scales" - - def mlp_up_proj_zero(self, i): - assert i < self.dense_replace - return f"model.layers.{i}.mlp.up_proj.qzeros" - - def mlp_gate_proj_weight(self, i): - assert i < self.dense_replace - return f"model.layers.{i}.mlp.gate_proj.qweight" - - def mlp_gate_proj_scale(self, i): - assert i < self.dense_replace - return f"model.layers.{i}.mlp.gate_proj.scales" - - def mlp_gate_proj_zero(self, i): - assert i < self.dense_replace - return f"model.layers.{i}.mlp.gate_proj.qzeros" - - # Latter layers are sparse - # Gating - def mlp_gate_weight(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.gate.weight" - - def mlp_gate_bias(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.gate.e_score_correction_bias" - - # Experts - def mlp_shared_experts_down_proj_weight(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.shared_experts.down_proj.qweight" - - def mlp_shared_experts_down_proj_scale(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.shared_experts.down_proj.scales" - - def mlp_shared_experts_down_proj_zero(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.shared_experts.down_proj.qzeros" - - def mlp_shared_experts_gate_proj_weight(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.shared_experts.gate_proj.qweight" - - def mlp_shared_experts_gate_proj_scale(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.shared_experts.gate_proj.scales" - - def mlp_shared_experts_gate_proj_zero(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.shared_experts.gate_proj.qzeros" - - def mlp_shared_experts_up_proj_weight(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.shared_experts.up_proj.qweight" - - def mlp_shared_experts_up_proj_scale(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.shared_experts.up_proj.scales" - - def mlp_shared_experts_up_proj_zero(self, i): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.shared_experts.up_proj.qzeros" - - # Experts - def mlp_experts_down_proj_weight(self, i, e): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.experts.{e}.down_proj.qweight" - - def mlp_experts_down_proj_scale(self, i, e): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.experts.{e}.down_proj.scales" - - def mlp_experts_down_proj_zero(self, i, e): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.experts.{e}.down_proj.qzeros" - - def mlp_experts_gate_proj_weight(self, i, e): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.experts.{e}.gate_proj.qweight" - - def mlp_experts_gate_proj_scale(self, i, e): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.experts.{e}.gate_proj.scales" - - def mlp_experts_gate_proj_zero(self, i, e): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.experts.{e}.gate_proj.qzeros" - - def mlp_experts_up_proj_weight(self, i, e): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.experts.{e}.up_proj.qweight" - - def mlp_experts_up_proj_scale(self, i, e): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.experts.{e}.up_proj.scales" - - def mlp_experts_up_proj_zero(self, i, e): - assert i >= self.dense_replace - return f"model.layers.{i}.mlp.experts.{e}.up_proj.qzeros" - - -class DeepSeekV3Meta(DeepSeekV3MetaCStruct): - def __init__(self, config, dtype=torch.float16, max_tokens=None): - if dtype == torch.float16: - dt_ = DataType.INFINI_DTYPE_F16 - elif dtype == torch.bfloat16: - dt_ = DataType.INFINI_DTYPE_BF16 - else: - dt_ = DataType.INFINI_DTYPE_F16 - - super().__init__( - # dtypes - dt_logits=DataType.INFINI_DTYPE_F16, - dt_norm=DataType.INFINI_DTYPE_BF16, - dt_quant_weight=DataType.INFINI_DTYPE_I32, - dt_quant_scale=DataType.INFINI_DTYPE_F16, - dt_quant_zero=DataType.INFINI_DTYPE_I32, - dt_gate_weight=DataType.INFINI_DTYPE_BF16, - dt_gate_bias=DataType.INFINI_DTYPE_BF16, - # sizes - n_sparse_layer=config["num_hidden_layers"], - n_dense_layer=config.get("first_k_dense_replace", 0), - d=config["hidden_size"], - nh=config["num_attention_heads"], - nkvh=config.get("num_key_value_heads", config["num_attention_heads"]), - d_rope=config["qk_rope_head_dim"], - d_nope=config["qk_nope_head_dim"], - r_q=config["q_lora_rank"], - r_kv=config["kv_lora_rank"], - d_qk=config["qk_nope_head_dim"] + config["qk_rope_head_dim"], - d_v=config["v_head_dim"], - # routing / experts / vocab / ctx - routed_scale=config.get("routed_scaling_factor", 1.0), - nexperts=config["n_routed_experts"], - kexperts=config["num_experts_per_tok"], - di=config["intermediate_size"], - di_moe=config["moe_intermediate_size"], - dctx=( - config["max_position_embeddings"] if max_tokens is None else max_tokens - ), - dvoc=config["vocab_size"], - # misc - epsilon=config.get("rms_norm_eps", 1e-6), - rope_theta=config.get("rope_theta", 10000.0), - end_token=config.get("eos_token_id", 2), - ) - self.torch_dtype_logits = dtype - - -def load_specific_tensor(model_dir, tensor_name): - """ - Load a specific tensor from a sharded safetensors model using its index JSON. - """ - index_file = os.path.join(model_dir, "model.safetensors.index.json") - if not os.path.exists(index_file): - raise FileNotFoundError(f"Index file not found: {index_file}") - - with open(index_file, "r") as f: - index = json.load(f) - - # Get mapping: tensor name -> file name - weight_map = index["weight_map"] - if tensor_name not in weight_map: - raise KeyError(f"{tensor_name} not found in index") - - filename = weight_map[tensor_name] - tensor_file = os.path.join(model_dir, filename) - - # Open only the relevant file and tensor - with safetensors.safe_open(tensor_file, framework="pt", device="cpu") as f: - tensor = f.get_tensor(tensor_name) - return tensor - - -def load_deepseek_weights( - meta: DeepSeekV3Meta, - weights, - model_path: str, - ndev: int, -): - model_instance = DeepSeekV3Model() - weight_loader = model_instance.create_weight_loader() - names = DeepseekR1WeightsNaming() - input_embd = load_specific_tensor(model_path, names.input_embd()).to( - meta.torch_dtype_logits - ) - weight_loader.contents.load_input_embd(weights, input_embd.data_ptr()) - del input_embd - - output_norm = load_specific_tensor(model_path, names.output_norm()) - weight_loader.contents.load_output_norm(weights, output_norm.data_ptr()) - del output_norm - - output_embd = load_specific_tensor(model_path, names.output_embd()) - weight_loader.contents.load_output_embd(weights, output_embd.data_ptr()) - del output_embd - - # ------------------------------- - # Per-layer weights - # ------------------------------- - - def load_quant(w_name, s_name, zero_name, split_dim=0): - weight = load_specific_tensor(model_path, w_name) - scale = load_specific_tensor(model_path, s_name) - zero = load_specific_tensor(model_path, zero_name) - if split_dim == 0 or ndev == 1: - return weight, scale, zero - elif split_dim == 1: - weight = ( - weight.reshape(weight.shape[0], ndev, -1).permute(1, 0, 2).contiguous() - ) - scale = ( - scale.reshape(scale.shape[0], ndev, -1).permute(1, 0, 2).contiguous() - ) - zero = zero.reshape(zero.shape[0], ndev, -1).permute(1, 0, 2).contiguous() - return weight, scale, zero - else: - raise ValueError("split_dim must be 0 or 1") - - for i in tqdm( - range(meta.n_sparse_layer + meta.n_dense_layer), desc="Loading layers" - ): - - # Attention norms + projections - attn_norm = load_specific_tensor(model_path, names.attn_norm(i)) - weight_loader.contents.load_attn_norm(weights, attn_norm.data_ptr(), i) - del attn_norm - - load_attn_q_a_layernorm = load_specific_tensor( - model_path, names.attn_q_a_layernorm(i) - ) - weight_loader.contents.load_attn_q_a_layernorm( - weights, load_attn_q_a_layernorm.data_ptr(), i - ) - del load_attn_q_a_layernorm - - attn_kv_a_layernorm = load_specific_tensor( - model_path, names.attn_kv_a_layernorm(i) - ) - weight_loader.contents.load_attn_kv_a_layernorm( - weights, attn_kv_a_layernorm.data_ptr(), i - ) - del attn_kv_a_layernorm - - w, s, z = load_quant( - names.attn_q_a_proj_weight(i), - names.attn_q_a_proj_scale(i), - names.attn_q_a_proj_zero(i), - ) - weight_loader.contents.load_attn_q_a_proj( - weights, w.data_ptr(), s.data_ptr(), z.data_ptr(), i - ) - - w, s, z = load_quant( - names.attn_q_b_proj_weight(i), - names.attn_q_b_proj_scale(i), - names.attn_q_b_proj_zero(i), - ) - weight_loader.contents.load_attn_q_b_proj( - weights, w.data_ptr(), s.data_ptr(), z.data_ptr(), i - ) - - w, s, z = load_quant( - names.attn_kv_a_proj_with_mqa_weight(i), - names.attn_kv_a_proj_with_mqa_scale(i), - names.attn_kv_a_proj_with_mqa_zero(i), - ) - weight_loader.contents.load_attn_kv_a_proj_with_mqa( - weights, w.data_ptr(), s.data_ptr(), z.data_ptr(), i - ) - - w, s, z = load_quant( - names.attn_kv_b_proj_weight(i), - names.attn_kv_b_proj_scale(i), - names.attn_kv_b_proj_zero(i), - ) - - weight_loader.contents.load_attn_kv_b_proj( - weights, w.data_ptr(), s.data_ptr(), z.data_ptr(), i - ) - - w, s, z = load_quant( - names.attn_o_proj_weight(i), - names.attn_o_proj_scale(i), - names.attn_o_proj_zero(i), - 1, - ) - - weight_loader.contents.load_attn_o_proj( - weights, w.data_ptr(), s.data_ptr(), z.data_ptr(), i - ) - - # ------------------------------- - # MLP: dense or sparse - # ------------------------------- - mlp_norm = load_specific_tensor(model_path, names.mlp_norm(i)) - weight_loader.contents.load_mlp_norm(weights, mlp_norm.data_ptr(), i) - - if i < meta.n_dense_layer: - # Dense MLP is grouped into one call - w_gate, s_gate, z_gate = load_quant( - names.mlp_gate_proj_weight(i), - names.mlp_gate_proj_scale(i), - names.mlp_gate_proj_zero(i), - ) - w_up, s_up, z_up = load_quant( - names.mlp_up_proj_weight(i), - names.mlp_up_proj_scale(i), - names.mlp_up_proj_zero(i), - ) - w_down, s_down, z_down = load_quant( - names.mlp_down_proj_weight(i), - names.mlp_down_proj_scale(i), - names.mlp_down_proj_zero(i), - 1, - ) - weight_loader.contents.load_mlp_dense( - weights, - w_gate.data_ptr(), - s_gate.data_ptr(), - z_gate.data_ptr(), - w_up.data_ptr(), - s_up.data_ptr(), - z_up.data_ptr(), - w_down.data_ptr(), - s_down.data_ptr(), - z_down.data_ptr(), - i, - ) - - else: - # Sparse MLP gating - mlp_gate_weight = load_specific_tensor(model_path, names.mlp_gate_weight(i)) - weight_loader.contents.load_mlp_gate_weight( - weights, mlp_gate_weight.data_ptr(), i - ) - del mlp_gate_weight - - mlp_gate_bias = load_specific_tensor(model_path, names.mlp_gate_bias(i)) - weight_loader.contents.load_mlp_gate_bias( - weights, mlp_gate_bias.data_ptr(), i - ) - del mlp_gate_bias - - # Shared experts - w_gate, s_gate, z_gate = load_quant( - names.mlp_shared_experts_gate_proj_weight(i), - names.mlp_shared_experts_gate_proj_scale(i), - names.mlp_shared_experts_gate_proj_zero(i), - ) - w_up, s_up, z_up = load_quant( - names.mlp_shared_experts_up_proj_weight(i), - names.mlp_shared_experts_up_proj_scale(i), - names.mlp_shared_experts_up_proj_zero(i), - ) - w_down, s_down, z_down = load_quant( - names.mlp_shared_experts_down_proj_weight(i), - names.mlp_shared_experts_down_proj_scale(i), - names.mlp_shared_experts_down_proj_zero(i), - 1, - ) - weight_loader.contents.load_mlp_shared_experts( - weights, - w_gate.data_ptr(), - s_gate.data_ptr(), - z_gate.data_ptr(), - w_up.data_ptr(), - s_up.data_ptr(), - z_up.data_ptr(), - w_down.data_ptr(), - s_down.data_ptr(), - z_down.data_ptr(), - i, - ) - - # Per-expert MLP - for e in range(meta.nexperts): - w_gate, s_gate, z_gate = load_quant( - names.mlp_experts_gate_proj_weight(i, e), - names.mlp_experts_gate_proj_scale(i, e), - names.mlp_experts_gate_proj_zero(i, e), - ) - w_up, s_up, z_up = load_quant( - names.mlp_experts_up_proj_weight(i, e), - names.mlp_experts_up_proj_scale(i, e), - names.mlp_experts_up_proj_zero(i, e), - ) - w_down, s_down, z_down = load_quant( - names.mlp_experts_down_proj_weight(i, e), - names.mlp_experts_down_proj_scale(i, e), - names.mlp_experts_down_proj_zero(i, e), - 1, - ) - weight_loader.contents.load_mlp_experts( - weights, - w_gate.data_ptr(), - s_gate.data_ptr(), - z_gate.data_ptr(), - w_up.data_ptr(), - s_up.data_ptr(), - z_up.data_ptr(), - w_down.data_ptr(), - s_down.data_ptr(), - z_down.data_ptr(), - i, - e, - ) - - -class DeepSeekV3BatchedTask: - def __init__(self, tasks: List[InferTask]): - self.tasks = tasks - self.nreq = len(tasks) - - # Precompute fields - token_lists = [t.tokens for t in tasks] - self.req_lens_list = [len(toks) for toks in token_lists] - self.req_pos_list = [t.pos for t in tasks] - self.kv_cache_ptrs = [t.kvcache().data() for t in tasks] - self.temperaturas_list = [t.temperature for t in tasks] - self.topks_list = [t.topk for t in tasks] - self.topps_list = [t.topp for t in tasks] - - # Flatten token lists - flat_tokens = [tok for toks in token_lists for tok in toks] - self.ntok = len(flat_tokens) - - # Convert to ctypes arrays in one pass - self.tokens = (c_uint * self.ntok)(*flat_tokens) - self.req_lens = (c_uint * self.nreq)(*self.req_lens_list) - self.req_pos = (c_uint * self.nreq)(*self.req_pos_list) - self.kv_caches = (POINTER(DeepSeekV3CacheCStruct) * self.nreq)( - *self.kv_cache_ptrs - ) - self.temperaturas = (c_float * self.nreq)(*self.temperaturas_list) - self.topks = (c_uint * self.nreq)(*self.topks_list) - self.topps = (c_float * self.nreq)(*self.topps_list) - - def input_args(self): - return ( - self.tokens, - self.ntok, - self.req_lens, - self.nreq, - self.req_pos, - self.kv_caches, - self.temperaturas, - self.topks, - self.topps, - ) - - -class DeepSeekV3ForCauslLM: - def __init__( - self, model_dir_path, device=DeviceType.DEVICE_TYPE_CPU, ndev=1, max_tokens=None - ): - with open(os.path.join(model_dir_path, "config.json"), "r") as f: - config = json.load(f) - self.config = config - eos_token_id = self.config["eos_token_id"] - self.eos_token_id = ( - [eos_token_id] if type(eos_token_id) == int else eos_token_id - ) - - print(model_dir_path) - - if "deepseek_v3" == config["model_type"]: - self.meta = DeepSeekV3Meta( - config, max_tokens=max_tokens, dtype=torch.float16 - ) - self.tokenizer = transformers.AutoTokenizer.from_pretrained(model_dir_path) - else: - raise ValueError("Unsupported model architecture") - - print(f"Creating model on {ndev} devices...") - load_start_time = time.time() - dev_ids = (c_int * ndev)(*[i for i in range(ndev)]) - - self.model_instance = DeepSeekV3Model() - weights = self.model_instance.create_weights( - byref(self.meta), - device, - ndev, - dev_ids, - ) - # Load weights from host - load_deepseek_weights(self.meta, weights, model_dir_path, ndev) - # Create model instance - self.model_ptr = self.model_instance.create_model( - byref(self.meta), - weights, - ) - load_end_time = time.time() - print(f"Time used: {load_end_time - load_start_time:.3f}s") - - def max_context_len(self): - return self.meta.dctx - - def create_kv_cache(self): - return self.model_instance.create_cache(self.model_ptr) - - def drop_kv_cache(self, kv_cache): - self.model_instance.drop_cache(self.model_ptr, kv_cache) - - def batch_infer_one_round(self, tasks: List[InferTask]): - output = (c_uint * len(tasks))() - batch_inputs = DeepSeekV3BatchedTask(tasks) - self.model_instance.infer_batch( - self.model_ptr, - *(batch_inputs.input_args()), - output, - ) - return list(output) - - def generate(self, input_content, max_steps, topp_=1.0, topk_=1, temperature_=1.0): - input_content = self.tokenizer.apply_chat_template( - conversation=[{"role": "user", "content": input_content}], - add_generation_prompt=True, - tokenize=False, - ) - - tokens = self.tokenizer.encode(input_content) - infer_task = InferTask( - 0, - tokens, - self.max_context_len(), - temperature_, - topk_, - topp_, - self.eos_token_id, - ) - infer_task.bind_kvcache(KVCache(self)) - print(input_content, end="", flush=True) - steps = 0 - total_time = 0 - output_content = "" - - for step_i in range(max_steps): - start_time = time.time() - output_tokens = self.batch_infer_one_round([infer_task]) - end_time = time.time() - steps += 1 - output_str = self.tokenizer.decode(output_tokens[0]) - output_content += output_str - print(output_str, end="", flush=True) - if output_tokens[0] in self.eos_token_id: - break - infer_task.next(output_tokens[0]) - - if step_i > 0: - total_time += end_time - start_time - - print("\n") - avg_time = total_time * 1000 / (steps - 1) - print(f"Time per step: {avg_time:.3f}ms") - - infer_task._kv_cache.drop(self) - return output_content, avg_time - - # def perplexity(self, test_sequences: List[Sequence[int]], batch_size=10): - # tasks = [ - # InferTask(i, [], self.max_context_len(), 1.0, 1, 1.0, self.eos_token_id) - # for i in range(batch_size) - # ] - # kv_caches = [KVCache(self) for _ in range(batch_size)] - - # nll = 0.0 - # total_len = 0 - - # for i in range(0, len(test_sequences), batch_size): - # batch_id = 0 - # true_tokens = [] - # while batch_id < batch_size and batch_id + i < len(test_sequences): - # input_tokens = test_sequences[i + batch_id][:-1] - # true_tokens.extend(test_sequences[i + batch_id][1:]) - # tasks[batch_id].tokens = input_tokens - # tasks[batch_id].bind_kvcache(kv_caches[batch_id]) - # batch_id += 1 - - # batch_inputs = DeepSeekV3BatchedTask(tasks[:batch_id]) - # logits = torch.zeros( - # (batch_inputs.ntok, self.meta.dvoc), dtype=self.meta.torch_dtype_logits - # ) - # forward_batch_deepseek_v3( - # self.model_instance, - # batch_inputs.tokens, - # batch_inputs.ntok, - # batch_inputs.req_lens, - # batch_inputs.nreq, - # batch_inputs.req_pos, - # batch_inputs.kv_caches, - # logits.data_ptr(), - # ) - - # logits = logits.float() - # token_ids = torch.tensor(true_tokens, dtype=torch.int64) # [ntok,] - # log_probs = torch.nn.functional.log_softmax(logits, dim=-1) # (ntok, vocab) - # token_logprobs = log_probs[ - # torch.arange(batch_inputs.ntok), token_ids - # ] # (ntok,) - - # start = 0 - # for l in batch_inputs.req_lens_list: - # nll += -token_logprobs[start : start + l].sum().item() - # start += l - # total_len += token_logprobs.numel() - - # for task in tasks: - # task.release_kvcache() - - # return math.exp(nll / total_len) - - def destroy_model_instance(self): - self.model_instance.destroy_model(self.model_ptr) - print("Model destroyed") - - -def test(): - if len(sys.argv) < 3: - print( - "Usage: python deepseek.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] [n_device]" - ) - sys.exit(1) - model_path = sys.argv[2] - device_type = DeviceType.DEVICE_TYPE_CPU - if sys.argv[1] == "--cpu": - device_type = DeviceType.DEVICE_TYPE_CPU - elif sys.argv[1] == "--nvidia": - device_type = DeviceType.DEVICE_TYPE_NVIDIA - elif sys.argv[1] == "--cambricon": - device_type = DeviceType.DEVICE_TYPE_CAMBRICON - elif sys.argv[1] == "--ascend": - device_type = DeviceType.DEVICE_TYPE_ASCEND - elif sys.argv[1] == "--metax": - device_type = DeviceType.DEVICE_TYPE_METAX - elif sys.argv[1] == "--moore": - device_type = DeviceType.DEVICE_TYPE_MOORE - elif sys.argv[1] == "--iluvatar": - device_type = DeviceType.DEVICE_TYPE_ILUVATAR - else: - print( - "Usage: python deepseek.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] [n_device]" - ) - sys.exit(1) - - ndev = int(sys.argv[3]) if len(sys.argv) > 3 else 1 - model = DeepSeekV3ForCauslLM(model_path, device_type, ndev, max_tokens=1024) - model.generate("山东最高的山是？", 50) - model.destroy_model_instance() - - -if __name__ == "__main__": - test() diff --git a/scripts/infer_task.py b/scripts/infer_task.py deleted file mode 100644 index 0d1231b77..000000000 --- a/scripts/infer_task.py +++ /dev/null @@ -1,59 +0,0 @@ -class InferTask: - def __init__(self, id, tokens, max_tokens, temperature, topk, topp, end_tokens): - self.id = id - self.finish_reason = None - self.tokens = tokens - self.max_tokens = max_tokens - self.temperature = temperature - self.topk = topk - self.topp = topp - self.end_tokens = end_tokens - self._kv_cache = None - self.pos = 0 - - def bind_kvcache(self, kv_cache, pos=0): - self._kv_cache = kv_cache - self.pos = pos - self.tokens = self.tokens[pos:] - - def release_kvcache(self): - cache = self._kv_cache - self._kv_cache = None - return cache - - def kvcache(self): - return self._kv_cache - - def next(self, out_token): - self._kv_cache.update_tokens(self.tokens, self.pos) - - self.pos += len(self.tokens) - if out_token == None or out_token in self.end_tokens: - self.finish_reason = "stop" - elif self.pos >= self.max_tokens: - self.finish_reason = "length" - else: - self.tokens = [out_token] - - -class KVCache: - def __init__(self, model): - self._kvcache = model.create_kv_cache() - self.tokens = [0 for _ in range(model.max_context_len())] - - def data(self): - return self._kvcache - - def drop(self, model): - model.drop_kv_cache(self._kvcache) - - def update_tokens(self, tokens, pos): - end = pos + len(tokens) - max_len = len(self.tokens) - - # If overflow, truncate tokens to fit - if end > max_len: - tokens = tokens[: max_len - pos] - end = max_len - - self.tokens[pos:end] = tokens diff --git a/scripts/jiuge.py b/scripts/jiuge.py deleted file mode 100644 index 35b2c8ca3..000000000 --- a/scripts/jiuge.py +++ /dev/null @@ -1,881 +0,0 @@ -from typing import List, Sequence -import math -import os -from pathlib import Path -import safetensors -import sys -import time -import json -import torch -import transformers - -from libinfinicore_infer import ( - JiugeModel, - JiugeMetaCStruct, - JiugeWeightsCStruct, - DataType, - DeviceType, - KVCacheCStruct, -) -from infer_task import InferTask, KVCache - -from ctypes import POINTER, c_float, c_int, c_uint, c_void_p, byref - -torch.set_default_device("cpu") - - -class LlamaWeightsNaming: - def input_embd(self): - return "model.embed_tokens.weight" - - def output_norm(self): - return "model.norm.weight" - - def output_embd(self): - return "lm_head.weight" - - def attn_norm(self, i): - return f"model.layers.{i}.input_layernorm.weight" - - def attn_q(self, i): - return f"model.layers.{i}.self_attn.q_proj.weight" - - def attn_k(self, i): - return f"model.layers.{i}.self_attn.k_proj.weight" - - def attn_v(self, i): - return f"model.layers.{i}.self_attn.v_proj.weight" - - def attn_o(self, i): - return f"model.layers.{i}.self_attn.o_proj.weight" - - def attn_q_b(self, i): - return f"model.layers.{i}.self_attn.q_proj.bias" - - def attn_k_b(self, i): - return f"model.layers.{i}.self_attn.k_proj.bias" - - def attn_v_b(self, i): - return f"model.layers.{i}.self_attn.v_proj.bias" - - def attn_q_norm(self, i): - return f"model.layers.{i}.self_attn.q_norm.weight" - - def attn_k_norm(self, i): - return f"model.layers.{i}.self_attn.k_norm.weight" - - def ffn_norm(self, i): - return f"model.layers.{i}.post_attention_layernorm.weight" - - def gate(self, i): - return f"model.layers.{i}.mlp.gate_proj.weight" - - def up(self, i): - return f"model.layers.{i}.mlp.up_proj.weight" - - def down(self, i): - return f"model.layers.{i}.mlp.down_proj.weight" - - def match(state_dict): - return ( - "model.norm.weight" in state_dict - and "model.layers.0.self_attn.q_proj.weight" in state_dict - ) - - -class JiugeMetaFromLlama(JiugeMetaCStruct): - def __init__(self, config, dtype=torch.float16, max_tokens=None): - if dtype == torch.float16: - dt_ = DataType.INFINI_DTYPE_F16 - elif dtype == torch.float32: - dt_ = DataType.INFINI_DTYPE_F32 - elif dtype == torch.bfloat16: - dt_ = DataType.INFINI_DTYPE_BF16 - else: - dt_ = DataType.INFINI_DTYPE_F16 - - self.scale_input = 1.0 - self.scale_output = 1.0 - self.scale_o = 1.0 - self.scale_down = 1.0 - if ( - config["model_type"] in ["fm9g", "minicpm"] - and "scale_emb" in config - and "scale_depth" in config - and "dim_model_base" in config - ): - self.scale_input = config["scale_emb"] - self.scale_output = config["hidden_size"] // config["dim_model_base"] - self.scale_o = config["scale_depth"] / math.sqrt( - config["num_hidden_layers"] - ) - self.scale_down = config["scale_depth"] / math.sqrt( - config["num_hidden_layers"] - ) - - super().__init__( - dt_logits=dt_, - nlayer=config["num_hidden_layers"], - d=config["hidden_size"], - nh=config["num_attention_heads"], - nkvh=( - config["num_key_value_heads"] - if "num_key_value_heads" in config - else config["num_attention_heads"] - ), - dh=( - config["head_dim"] - if "head_dim" in config - else config["hidden_size"] // config["num_attention_heads"] - ), - di=config["intermediate_size"], - dctx=( - config["max_position_embeddings"] if max_tokens is None else max_tokens - ), - dvoc=config["vocab_size"], - epsilon=config["rms_norm_eps"], - theta=(config["rope_theta"] if "rope_theta" in config else 100000.0), - end_token=2, - ) - self.torch_dtype_logits = dtype - - -class JiugeWeightsImpl(JiugeWeightsCStruct): - def __init__( - self, - meta, - naming, - state_dict, - torch_dt_mat=torch.float16, - torch_dt_norm=torch.float32, - ndev=1, - transpose_weight=True, - ): - nlayer = meta.nlayer - nh = meta.nh - nkvh = meta.nkvh - dh = meta.dh - d = meta.d - di = meta.di - scale_input = meta.scale_input - scale_output = meta.scale_output - scale_o = meta.scale_o - scale_down = meta.scale_down - assert nh % nkvh == 0 - assert nh % ndev == 0 - assert nkvh % ndev == 0 - assert di % ndev == 0 - torch_dt_logits = meta.torch_dtype_logits - if torch_dt_mat == torch.float16: - self.dt_mat = DataType.INFINI_DTYPE_F16 - elif torch_dt_mat == torch.float32: - self.dt_mat = DataType.INFINI_DTYPE_F32 - elif torch_dt_mat == torch.bfloat16: - self.dt_mat = DataType.INFINI_DTYPE_BF16 - else: - raise ValueError("Unsupported proj weight data type") - if torch_dt_norm == torch.float16: - self.dt_norm = DataType.INFINI_DTYPE_F16 - elif torch_dt_norm == torch.float32: - self.dt_norm = DataType.INFINI_DTYPE_F32 - elif torch_dt_norm == torch.bfloat16: - self.dt_norm = DataType.INFINI_DTYPE_BF16 - else: - raise ValueError("Unsupported norm weight data type") - - input_embd_naming = ( - naming.input_embd() - if naming.input_embd() in state_dict - else naming.output_embd() - ) - output_embd_naming = ( - naming.output_embd() - if naming.output_embd() in state_dict - else naming.input_embd() - ) - self.transpose_linear_weights = 1 if transpose_weight else 0 - self.nlayer = nlayer - self.input_embd_tensor = ( - state_dict[input_embd_naming].to(torch_dt_logits) * scale_input - ) - self.input_embd = self.input_embd_tensor.data_ptr() - self.output_norm_tensor = ( - state_dict[naming.output_norm()].to(torch_dt_norm) * scale_output - ) - self.output_norm = self.output_norm_tensor.data_ptr() - self.output_embd_tensor = state_dict[output_embd_naming].to(torch_dt_mat) - if not transpose_weight: - self.output_embd_tensor = self.output_embd_tensor.transpose( - 0, 1 - ).contiguous() - self.output_embd = self.output_embd_tensor.data_ptr() - - self.attn_norm_tensors = [ - state_dict[naming.attn_norm(i)].to(torch_dt_norm) for i in range(nlayer) - ] - self.attn_norm_ptrs = [ - self.attn_norm_tensors[i].data_ptr() for i in range(nlayer) - ] - self.attn_norm = (c_void_p * nlayer)(*self.attn_norm_ptrs) - - def qkv_slices(_i): - _Q = ( - state_dict[naming.attn_q(_i)] - .reshape([nh, 2, dh // 2, d]) - .transpose(1, 2) - ) - _K = ( - state_dict[naming.attn_k(_i)] - .reshape([nkvh, 2, dh // 2, d]) - .transpose(1, 2) - ) - _V = state_dict[naming.attn_v(_i)].reshape([nkvh, dh // 2, 2, d]) - _result = [] - _nh = nh // ndev - _nkvh = nkvh // ndev - for _idev in range(ndev): - _result.append(_Q[_idev * _nh : (_idev + 1) * _nh, :, :, :]) - _result.append(_K[_idev * _nkvh : (_idev + 1) * _nkvh, :, :, :]) - _result.append(_V[_idev * _nkvh : (_idev + 1) * _nkvh, :, :]) - return _result - - self.qkv_tensor = [ - torch.concat(qkv_slices(i)).to(torch_dt_mat) for i in range(nlayer) - ] - if not transpose_weight: - for i in range(nlayer): - self.qkv_tensor[i] = ( - self.qkv_tensor[i] - .reshape(ndev, (nh + 2 * nkvh) // ndev * dh, d) - .transpose(1, 2) - .contiguous() - ) - self.qkv_tensor_ptrs = [self.qkv_tensor[i].data_ptr() for i in range(nlayer)] - self.attn_qkv = (c_void_p * nlayer)(*self.qkv_tensor_ptrs) - - def qkv_b_slices(_i): - _QB = ( - state_dict[naming.attn_q_b(_i)] - .reshape([nh, 2, dh // 2]) - .transpose(1, 2) - ) - _KB = ( - state_dict[naming.attn_k_b(_i)] - .reshape([nkvh, 2, dh // 2]) - .transpose(1, 2) - ) - _VB = state_dict[naming.attn_v_b(_i)].reshape([nkvh, dh // 2, 2]) - _result = [] - _nh = nh // ndev - _nkvh = nkvh // ndev - for _idev in range(ndev): - _result.append(_QB[_idev * _nh : (_idev + 1) * _nh, :, :].flatten()) - _result.append(_KB[_idev * _nkvh : (_idev + 1) * _nkvh, :, :].flatten()) - _result.append(_VB[_idev * _nkvh : (_idev + 1) * _nkvh, :, :].flatten()) - return _result - - if naming.attn_q_b(0) in state_dict: - self.qkv_b_tensors = [ - torch.concat(qkv_b_slices(i)).to(torch_dt_logits) for i in range(nlayer) - ] - self.qkv_b_tensor_ptrs = [ - self.qkv_b_tensors[i].data_ptr() for i in range(nlayer) - ] - self.attn_qkv_b = (c_void_p * nlayer)(*self.qkv_b_tensor_ptrs) - else: - self.attn_qkv_b = None - - if naming.attn_q_norm(0) in state_dict: - self.attn_q_norm_tensors = [ - state_dict[naming.attn_q_norm(i)] - .reshape([2, dh // 2]) - .transpose(0, 1) - .contiguous() - .to(torch_dt_norm) - for i in range(nlayer) - ] - self.attn_q_norm_ptrs = [ - self.attn_q_norm_tensors[i].data_ptr() for i in range(nlayer) - ] - self.attn_q_norm = (c_void_p * nlayer)(*self.attn_q_norm_ptrs) - self.attn_k_norm_tensors = [ - state_dict[naming.attn_k_norm(i)] - .reshape([2, dh // 2]) - .transpose(0, 1) - .contiguous() - .to(torch_dt_norm) - for i in range(nlayer) - ] - self.attn_k_norm_ptrs = [ - self.attn_k_norm_tensors[i].data_ptr() for i in range(nlayer) - ] - self.attn_k_norm = (c_void_p * nlayer)(*self.attn_k_norm_ptrs) - else: - self.attn_q_norm = None - self.attn_k_norm = None - - self.attn_o_tensor = [ - ( - state_dict[naming.attn_o(i)] - .to(torch_dt_mat) - .reshape([d, ndev, nh // ndev * dh]) - .transpose(0, 1) - .contiguous() - if transpose_weight - else state_dict[naming.attn_o(i)] - .transpose(0, 1) - .to(torch_dt_mat) - .contiguous() - ) - * scale_o - for i in range(nlayer) - ] - self.attn_o_ptrs = [self.attn_o_tensor[i].data_ptr() for i in range(nlayer)] - self.attn_o = (c_void_p * nlayer)(*self.attn_o_ptrs) - - self.ffn_norm_tensors = [ - state_dict[naming.ffn_norm(i)].to(torch_dt_norm) for i in range(nlayer) - ] - self.ffn_norm_ptrs = [ - self.ffn_norm_tensors[i].data_ptr() for i in range(nlayer) - ] - self.ffn_norm = (c_void_p * nlayer)(*self.ffn_norm_ptrs) - - def gate_up_slices(_i): - _result = [] - _di = di // ndev - for _idev in range(ndev): - _start = _idev * _di - _end = (_idev + 1) * _di - _result.append(state_dict[naming.gate(_i)][_start:_end, :]) - _result.append(state_dict[naming.up(_i)][_start:_end, :]) - return _result - - self.gate_up_tensors = [ - torch.concat(gate_up_slices(i)).to(torch_dt_mat) for i in range(nlayer) - ] - if not transpose_weight: - for i in range(nlayer): - self.gate_up_tensors[i] = ( - self.gate_up_tensors[i] - .reshape(ndev, 2 * di // ndev, d) - .transpose(1, 2) - .contiguous() - ) - self.gate_up_ptrs = [self.gate_up_tensors[i].data_ptr() for i in range(nlayer)] - self.ffn_gate_up = (c_void_p * nlayer)(*self.gate_up_ptrs) - - self.ffn_down_tensor = [ - ( - state_dict[naming.down(i)] - .to(torch_dt_mat) - .reshape([d, ndev, di // ndev]) - .transpose(0, 1) - .contiguous() - if transpose_weight - else state_dict[naming.down(i)] - .transpose(0, 1) - .to(torch_dt_mat) - .contiguous() - ) - * scale_down - for i in range(nlayer) - ] - self.ffn_down_ptrs = [self.ffn_down_tensor[i].data_ptr() for i in range(nlayer)] - self.ffn_down = (c_void_p * nlayer)(*self.ffn_down_ptrs) - - -class JiugeBatchedTask: - def __init__(self, tasks: List[InferTask]): - self.tasks = tasks - self.nreq = len(tasks) - - # Precompute fields - token_lists = [t.tokens for t in tasks] - self.req_lens_list = [len(toks) for toks in token_lists] - self.req_pos_list = [t.pos for t in tasks] - self.kv_cache_ptrs = [t.kvcache().data() for t in tasks] - self.temperaturas_list = [t.temperature for t in tasks] - self.topks_list = [t.topk for t in tasks] - self.topps_list = [t.topp for t in tasks] - - # Flatten token lists - flat_tokens = [tok for toks in token_lists for tok in toks] - self.ntok = len(flat_tokens) - - # Convert to ctypes arrays in one pass - self.tokens = (c_uint * self.ntok)(*flat_tokens) - self.req_lens = (c_uint * self.nreq)(*self.req_lens_list) - self.req_pos = (c_uint * self.nreq)(*self.req_pos_list) - self.kv_caches = (POINTER(KVCacheCStruct) * self.nreq)(*self.kv_cache_ptrs) - self.temperaturas = (c_float * self.nreq)(*self.temperaturas_list) - self.topks = (c_uint * self.nreq)(*self.topks_list) - self.topps = (c_float * self.nreq)(*self.topps_list) - - def input_args(self): - return ( - self.tokens, - self.ntok, - self.req_lens, - self.nreq, - self.req_pos, - self.kv_caches, - self.temperaturas, - self.topks, - self.topps, - ) - - -class JiugeForCauslLM: - def __init__( - self, model_dir_path, device=DeviceType.DEVICE_TYPE_CPU, ndev=1, max_tokens=None - ): - def load_all_safetensors_from_dir(dir_path_: str): - tensors_ = {} - dir_path_ = Path(dir_path_) - for file in sorted(dir_path_.glob("*.safetensors")): - data_ = safetensors.safe_open(file, "pt") - for name_ in data_.keys(): - tensors_[name_] = data_.get_tensor(name_) - return tensors_ - - print("Loading model weights to host...") - load_start_time = time.time() - - with open(os.path.join(model_dir_path, "config.json"), "r") as f: - config = json.load(f) - self.config = config - eos_token_id = self.config["eos_token_id"] - self.eos_token_id = ( - [eos_token_id] if type(eos_token_id) == int else eos_token_id - ) - transpose_weight = ( - device != DeviceType.DEVICE_TYPE_ASCEND - ) # y = xW is faster than y=xW^T on Ascend - - self.jiuge_model = JiugeModel() - - if "llama" == config["model_type"]: - model = ( - transformers.LlamaForCausalLM.from_pretrained(model_dir_path) - .cpu() - .half() - ) - self.meta = JiugeMetaFromLlama(config, max_tokens=max_tokens) - self.tokenizer = transformers.AutoTokenizer.from_pretrained(model_dir_path) - self.weights = JiugeWeightsImpl( - self.meta, - LlamaWeightsNaming(), - model.state_dict(), - ndev=ndev, - transpose_weight=transpose_weight, - ) - elif "fm9g" == config["model_type"] or "minicpm" == config["model_type"]: - if any( - file.suffix == ".safetensors" for file in Path(model_dir_path).iterdir() - ): - state_dict = load_all_safetensors_from_dir(model_dir_path) - else: - state_dict = torch.load( - os.path.join(model_dir_path, "pytorch_model.bin"), - weights_only=True, - map_location="cpu", - ) - if LlamaWeightsNaming.match(state_dict): - self.meta = JiugeMetaFromLlama(config, max_tokens=max_tokens) - self.weights = JiugeWeightsImpl( - self.meta, - LlamaWeightsNaming(), - state_dict, - ndev=ndev, - transpose_weight=transpose_weight, - ) - self.tokenizer = transformers.AutoTokenizer.from_pretrained( - model_dir_path, trust_remote_code=True - ) - else: - raise ValueError("Unsupported weight naming") - elif "fm9g7b" == config["model_type"]: - if any( - file.suffix == ".safetensors" for file in Path(model_dir_path).iterdir() - ): - state_dict = load_all_safetensors_from_dir(model_dir_path) - else: - state_dict = torch.load( - os.path.join(model_dir_path, "pytorch_model.bin"), - weights_only=True, - map_location="cpu", - ) - if LlamaWeightsNaming.match(state_dict): - self.meta = JiugeMetaFromLlama(config, max_tokens=max_tokens) - self.weights = JiugeWeightsImpl( - self.meta, - LlamaWeightsNaming(), - state_dict, - ndev=ndev, - transpose_weight=transpose_weight, - ) - self.tokenizer = transformers.AutoTokenizer.from_pretrained( - model_dir_path, trust_remote_code=True - ) - else: - raise ValueError("Unsupported weight naming") - elif "qwen2" == config["model_type"] or "qwen3" == config["model_type"]: - state_dict = load_all_safetensors_from_dir(model_dir_path) - if LlamaWeightsNaming.match(state_dict): - self.meta = JiugeMetaFromLlama(config, max_tokens=max_tokens) - self.weights = JiugeWeightsImpl( - self.meta, - LlamaWeightsNaming(), - state_dict, - ndev=ndev, - transpose_weight=transpose_weight, - ) - self.tokenizer = transformers.AutoTokenizer.from_pretrained( - model_dir_path - ) - else: - raise ValueError("Unsupported model architecture") - - if "llama" == config["model_type"]: - from tokenizers import decoders as _dec - - backend = getattr(self.tokenizer, "backend_tokenizer", None) - target = getattr(backend, "_tokenizer", backend) - norm = getattr(target, "normalizer", None) - dec = getattr(target, "decoder", None) - sn = repr(norm)[:800] if norm is not None else "" - sd = repr(dec)[:800] if dec is not None else "" - has_prepend = "Prepend" in sn - has_strip = "Strip" in sd - if has_prepend and has_strip: - target.decoder = _dec.Sequence( - [ - _dec.Replace("▁", " "), - _dec.ByteFallback(), - _dec.Fuse(), - ] - ) - - load_end_time = time.time() - print(f"Time used: {load_end_time - load_start_time:.3f}s") - - print(f"Creating model on {ndev} devices...") - load_start_time = time.time() - self.dev_ids = (c_int * ndev)(*[i for i in range(ndev)]) - self.ndev = ndev - self.device = device - - self.model_instance = self.jiuge_model.create_model( - byref(self.meta), - byref(self.weights), - device, - ndev, - self.dev_ids, - ) - load_end_time = time.time() - print(f"Time used: {load_end_time - load_start_time:.3f}s") - - def max_context_len(self): - return self.meta.dctx - - def create_kv_cache(self): - return self.jiuge_model.create_kv_cache( - self.meta.nlayer, - self.meta.dctx, - self.meta.nkvh, - self.meta.dh, - self.meta.dh, - self.meta.dt_logits, - self.device, - self.dev_ids, - self.ndev, - ) - - def drop_kv_cache(self, kv_cache): - self.jiuge_model.drop_kv_cache(kv_cache) - - def batch_infer_one_round(self, tasks: List[InferTask]): - output = (c_uint * len(tasks))() - batch_inputs = JiugeBatchedTask(tasks) - self.jiuge_model.infer_batch( - self.model_instance, - *(batch_inputs.input_args()), - output, - ) - return list(output) - - def generate( - self, - input_content, - max_steps, - topp_=1.0, - topk_=1, - temperature_=1.0, - verbose=False, - ): - input_content = self.tokenizer.apply_chat_template( - conversation=[{"role": "user", "content": input_content}], - add_generation_prompt=True, - tokenize=False, - ) - print(input_content, end="", flush=True) - tokens = self.tokenizer.encode(input_content) - infer_task = InferTask( - 0, - tokens, - self.max_context_len(), - temperature_, - topk_, - topp_, - self.eos_token_id, - ) - infer_task.bind_kvcache(KVCache(self)) - - steps = 0 - total_time = 0 - prefill_time = 0 - decode_time = 0 - output_content = "" - - # Prefill phase - process initial prompt - prefill_start_time = time.time() - output_tokens = self.batch_infer_one_round([infer_task]) - prefill_end_time = time.time() - prefill_time = prefill_end_time - prefill_start_time - steps += 1 - - output_str = self.tokenizer.decode(output_tokens[0]) - output_content += output_str - print(output_str, end="", flush=True) - if output_tokens[0] in self.eos_token_id: - # If generation ends after prefill, calculate metrics - total_time = prefill_time - total_tokens = len(tokens) + 1 # input tokens + first output token - - print("\n") - print(f"Time per step: {total_time * 1000:.3f}ms") - - if verbose: - overall_throughput = total_tokens / total_time - prefill_throughput = len(tokens) / prefill_time - decode_throughput = 1 / 0.001 # Avoid division by zero, use small value - - print("=" * 50) - print("PERFORMANCE METRICS") - print("=" * 50) - print(f"Input tokens: {len(tokens)}") - print(f"Generated tokens: 1") - print(f"Total tokens: {total_tokens}") - print(f"Total time: {total_time * 1000:.3f}ms") - print(f"Prefill time: {prefill_time * 1000:.3f}ms") - print(f"Decode time: 0.000ms") - print("-" * 50) - print(f"Time per step: {total_time * 1000:.3f}ms") - print( - f"Avg prefill time per token: {prefill_time * 1000 / len(tokens):.3f}ms" - ) - print(f"Avg decode time per token: N/A") - print("-" * 50) - print(f"Overall throughput: {overall_throughput:.2f} tokens/s") - print(f"Prefill throughput: {prefill_throughput:.2f} tokens/s") - print(f"Decode throughput: N/A") - print("=" * 50) - - return output_content, total_time * 1000 - - infer_task.next(output_tokens[0]) - - # Decode phase - generate subsequent tokens - decode_start_time = time.time() - for step_i in range(1, max_steps): - start_time = time.time() - output_tokens = self.batch_infer_one_round([infer_task]) - end_time = time.time() - steps += 1 - output_str = self.tokenizer.decode(output_tokens[0]) - - output_content += output_str - print(output_str, end="", flush=True) - if output_tokens[0] in self.eos_token_id: - break - infer_task.next(output_tokens[0]) - - if step_i > 0: - total_time += end_time - start_time - - decode_end_time = time.time() - decode_time = decode_end_time - decode_start_time - - print("\n") - - # Calculate performance metrics - total_time = prefill_time + decode_time - input_tokens = len(tokens) - generated_tokens = steps # including first token from prefill - - # Time per token calculations - avg_time_per_step = ( - total_time * 1000 / (steps - 1) if steps > 1 else total_time * 1000 - ) - - print(f"Time per step: {avg_time_per_step:.3f}ms") - - # Only print detailed metrics if verbose flag is set - if verbose: - total_tokens = input_tokens + generated_tokens - - # Throughput calculations - overall_throughput = total_tokens / total_time # tokens per second - prefill_throughput = input_tokens / prefill_time if prefill_time > 0 else 0 - decode_throughput = ( - (generated_tokens - 1) / decode_time if decode_time > 0 else 0 - ) # exclude first token from prefill - - # Time per token calculations - avg_prefill_time_per_token = ( - prefill_time * 1000 / input_tokens if input_tokens > 0 else 0 - ) - avg_decode_time_per_token = ( - decode_time * 1000 / (generated_tokens - 1) - if generated_tokens > 1 - else 0 - ) - - print("=" * 50) - print("PERFORMANCE METRICS") - print("=" * 50) - print(f"Input tokens: {input_tokens}") - print(f"Generated tokens: {generated_tokens}") - print(f"Total tokens: {total_tokens}") - print(f"Total time: {total_time * 1000:.3f}ms") - print(f"Prefill time: {prefill_time * 1000:.3f}ms") - print(f"Decode time: {decode_time * 1000:.3f}ms") - print("-" * 50) - print(f"Time per step: {avg_time_per_step:.3f}ms") - print(f"Avg prefill time per token: {avg_prefill_time_per_token:.3f}ms") - print(f"Avg decode time per token: {avg_decode_time_per_token:.3f}ms") - print("-" * 50) - print(f"Overall throughput: {overall_throughput:.2f} tokens/s") - print(f"Prefill throughput: {prefill_throughput:.2f} tokens/s") - print(f"Decode throughput: {decode_throughput:.2f} tokens/s") - print("=" * 50) - - infer_task._kv_cache.drop(self) - return output_content, avg_time_per_step - - def perplexity(self, test_sequences: List[Sequence[int]], batch_size=10): - tasks = [ - InferTask(i, [], self.max_context_len(), 1.0, 1, 1.0, self.eos_token_id) - for i in range(batch_size) - ] - kv_caches = [KVCache(self) for _ in range(batch_size)] - - nll = 0.0 - total_len = 0 - - for i in range(0, len(test_sequences), batch_size): - batch_id = 0 - true_tokens = [] - while batch_id < batch_size and batch_id + i < len(test_sequences): - input_tokens = test_sequences[i + batch_id][:-1] - true_tokens.extend(test_sequences[i + batch_id][1:]) - tasks[batch_id].tokens = input_tokens - tasks[batch_id].bind_kvcache(kv_caches[batch_id]) - batch_id += 1 - - batch_inputs = JiugeBatchedTask(tasks[:batch_id]) - logits = torch.zeros( - (batch_inputs.ntok, self.meta.dvoc), dtype=self.meta.torch_dtype_logits - ) - self.jiuge_model.forward_batch( - self.model_instance, - batch_inputs.tokens, - batch_inputs.ntok, - batch_inputs.req_lens, - batch_inputs.nreq, - batch_inputs.req_pos, - batch_inputs.kv_caches, - logits.data_ptr(), - ) - - logits = logits.float() - token_ids = torch.tensor(true_tokens, dtype=torch.int64) # [ntok,] - log_probs = torch.nn.functional.log_softmax(logits, dim=-1) # (ntok, vocab) - token_logprobs = log_probs[ - torch.arange(batch_inputs.ntok), token_ids - ] # (ntok,) - - start = 0 - for l in batch_inputs.req_lens_list: - nll += -token_logprobs[start : start + l].sum().item() - start += l - total_len += token_logprobs.numel() - - for task in tasks: - task.release_kvcache() - - return math.exp(nll / total_len) - - def destroy_model_instance(self): - self.jiuge_model.destroy_model(self.model_instance) - print("Model destroyed") - - -def test(): - if len(sys.argv) < 3: - print( - "Usage: python jiuge.py [--cpu | --nvidia| --qy| --cambricon | --ascend | --metax | --moore | --iluvatar | --kunlun | --hygon] [n_device] [--verbose]" - ) - sys.exit(1) - - # Parse command line arguments - model_path = sys.argv[2] - device_type = DeviceType.DEVICE_TYPE_CPU - verbose = False - - # Check for verbose flag - for arg in sys.argv: - if arg == "--verbose": - verbose = True - break - - if sys.argv[1] == "--cpu": - device_type = DeviceType.DEVICE_TYPE_CPU - elif sys.argv[1] == "--nvidia": - device_type = DeviceType.DEVICE_TYPE_NVIDIA - elif sys.argv[1] == "--qy": - device_type = DeviceType.DEVICE_TYPE_QY - elif sys.argv[1] == "--cambricon": - device_type = DeviceType.DEVICE_TYPE_CAMBRICON - elif sys.argv[1] == "--ascend": - device_type = DeviceType.DEVICE_TYPE_ASCEND - elif sys.argv[1] == "--metax": - device_type = DeviceType.DEVICE_TYPE_METAX - elif sys.argv[1] == "--moore": - device_type = DeviceType.DEVICE_TYPE_MOORE - elif sys.argv[1] == "--iluvatar": - device_type = DeviceType.DEVICE_TYPE_ILUVATAR - elif sys.argv[1] == "--kunlun": - device_type = DeviceType.DEVICE_TYPE_KUNLUN - elif sys.argv[1] == "--hygon": - device_type = DeviceType.DEVICE_TYPE_HYGON - elif sys.argv[1] == "--ali": - device_type = DeviceType.DEVICE_TYPE_ALI - else: - print( - "Usage: python jiuge.py [--cpu | --nvidia| --qy| --cambricon | --ascend | --metax | --moore | --iluvatar | --kunlun | --hygon | --ali] [n_device] [--verbose]" - ) - sys.exit(1) - - # Find n_device argument (skip --verbose) - ndev_args = [arg for arg in sys.argv[3:] if arg != "--verbose"] - ndev = int(ndev_args[0]) if ndev_args else 1 - - model = JiugeForCauslLM(model_path, device_type, ndev) - model.generate("山东最高的山是？", 500, verbose=verbose) - model.destroy_model_instance() - - -if __name__ == "__main__": - test() diff --git a/scripts/jiuge_awq.py b/scripts/jiuge_awq.py deleted file mode 100644 index 5191efeba..000000000 --- a/scripts/jiuge_awq.py +++ /dev/null @@ -1,369 +0,0 @@ -from typing import List, Sequence -import math -import os -from pathlib import Path -import safetensors -import sys -import time -import json -import torch -import transformers - -from libinfinicore_infer import ( - JiugeAWQModel, - JiugeAWQMetaCStruct, - DataType, - DeviceType, - KVCacheCStruct, -) -from infer_task import InferTask, KVCache - -from ctypes import POINTER, c_float, c_int, c_uint, c_void_p, byref - -torch.set_default_device("cpu") - - -class JiugeAWQMetaFromConfig(JiugeAWQMetaCStruct): - def __init__(self, config, dtype=torch.float16, max_tokens=None): - if config["torch_dtype"] == "float16": - dt_ = DataType.INFINI_DTYPE_F16 - elif config["torch_dtype"] == "float32": - dt_ = DataType.INFINI_DTYPE_F32 - elif config["torch_dtype"] == "bfloat16": - dt_ = DataType.INFINI_DTYPE_BF16 - else: - dt_ = DataType.INFINI_DTYPE_F16 - - self.scale_input = 1.0 - self.scale_output = 1.0 - self.scale_o = 1.0 - self.scale_down = 1.0 - if ( - config["model_type"] in ["fm9g", "minicpm"] - and "scale_emb" in config - and "scale_depth" in config - and "dim_model_base" in config - ): - self.scale_input = config["scale_emb"] - self.scale_output = config["hidden_size"] // config["dim_model_base"] - self.scale_o = config["scale_depth"] / math.sqrt( - config["num_hidden_layers"] - ) - self.scale_down = config["scale_depth"] / math.sqrt( - config["num_hidden_layers"] - ) - - has_qkv_bias = ( - 1 if "attention_bias" in config and config["attention_bias"] else 0 - ) - if config["model_type"] in ["qwen2", "qwen3"]: - has_qkv_bias = 1 - - eos_token_id = ( - config["eos_token_id"][0] - if type(config["eos_token_id"]) == list - else config["eos_token_id"] - ) - - super().__init__( - dt_logits=dt_, - dt_linear_w=DataType.INFINI_DTYPE_I32, - dt_norm_w=dt_, - nlayer=config["num_hidden_layers"], - d=config["hidden_size"], - nh=config["num_attention_heads"], - nkvh=( - config["num_key_value_heads"] - if "num_key_value_heads" in config - else config["num_attention_heads"] - ), - dh=config["hidden_size"] // config["num_attention_heads"], - di=config["intermediate_size"], - dctx=( - config["max_position_embeddings"] if max_tokens is None else max_tokens - ), - dvoc=config["vocab_size"], - epsilon=config["rms_norm_eps"], - theta=(config["rope_theta"] if "rope_theta" in config else 100000.0), - end_token=eos_token_id, - nbit=config["quantization_config"]["bits"], - quant_group_size=config["quantization_config"]["group_size"], - has_qkv_bias=has_qkv_bias, - ) - self.torch_dtype_logits = dtype - - -class JiugeAWQBatchedTask: - def __init__(self, tasks: List[InferTask]): - self.tasks = tasks - self.nreq = len(tasks) - - # Precompute fields - token_lists = [t.tokens for t in tasks] - self.req_lens_list = [len(toks) for toks in token_lists] - self.req_pos_list = [t.pos for t in tasks] - self.kv_cache_ptrs = [t.kvcache().data() for t in tasks] - self.temperaturas_list = [t.temperature for t in tasks] - self.topks_list = [t.topk for t in tasks] - self.topps_list = [t.topp for t in tasks] - - # Flatten token lists - flat_tokens = [tok for toks in token_lists for tok in toks] - self.ntok = len(flat_tokens) - - # Convert to ctypes arrays in one pass - self.tokens = (c_uint * self.ntok)(*flat_tokens) - self.req_lens = (c_uint * self.nreq)(*self.req_lens_list) - self.req_pos = (c_uint * self.nreq)(*self.req_pos_list) - self.kv_caches = (POINTER(KVCacheCStruct) * self.nreq)(*self.kv_cache_ptrs) - self.temperaturas = (c_float * self.nreq)(*self.temperaturas_list) - self.topks = (c_uint * self.nreq)(*self.topks_list) - self.topps = (c_float * self.nreq)(*self.topps_list) - - def input_args(self): - return ( - self.tokens, - self.ntok, - self.req_lens, - self.nreq, - self.req_pos, - self.kv_caches, - self.temperaturas, - self.topks, - self.topps, - ) - - -class JiugeAWQForCausalLM: - def __init__( - self, model_dir_path, device=DeviceType.DEVICE_TYPE_CPU, ndev=1, max_tokens=None - ): - - load_start_time = time.time() - print(f"Creating model on {ndev} devices...") - with open(os.path.join(model_dir_path, "config.json"), "r") as f: - config = json.load(f) - self.config = config - eos_token_id = self.config["eos_token_id"] - self.eos_token_id = ( - [eos_token_id] if type(eos_token_id) == int else eos_token_id - ) - self.dev_ids = (c_int * ndev)(*[i for i in range(ndev)]) - self.ndev = ndev - self.device = device - self.meta = JiugeAWQMetaFromConfig(config, max_tokens=max_tokens) - - self.jiuge_awq_model = JiugeAWQModel() - - self.weights = self.jiuge_awq_model.create_weights( - byref(self.meta), - self.device, - ndev, - self.dev_ids, - ) - self.tokenizer = transformers.AutoTokenizer.from_pretrained( - model_dir_path, trust_remote_code=True - ) - - load_end_time = time.time() - print(f"Time used: {load_end_time - load_start_time:.3f}s") - - load_start_time = time.time() - print("Loading model weights to host...") - - self.load_all_safetensors_from_dir(os.path.join(model_dir_path)) - - self.model_instance = self.jiuge_awq_model.create_model( - byref(self.meta), - self.weights, - ) - load_end_time = time.time() - print(f"Time used: {load_end_time - load_start_time:.3f}s") - - def load_all_safetensors_from_dir(self, dir_path_: str): - dir_path_ = Path(dir_path_) - for file in sorted(dir_path_.glob("*.safetensors")): - with safetensors.safe_open(file, framework="pt", device="cpu") as f: - for key in f.keys(): - # print(key) - tensor = f.get_tensor(key) - if "o_proj.scales" in key: - tensor = tensor * self.meta.scale_o - elif "down_proj.scales" in key: - tensor = tensor * self.meta.scale_down - elif "embed_tokens.weight" in key: - tensor = tensor * self.meta.scale_input - elif "lm_head.weight" in key: - tensor = tensor * self.meta.scale_output - self.jiuge_awq_model.load_weight( - self.weights, key, tensor.data_ptr() - ) - - def max_context_len(self): - return self.meta.dctx - - def create_kv_cache(self): - return self.jiuge_awq_model.create_kv_cache( - self.meta.nlayer, - self.meta.dctx, - self.meta.nkvh, - self.meta.dh, - self.meta.dh, - self.meta.dt_logits, - self.device, - self.dev_ids, - self.ndev, - ) - - def drop_kv_cache(self, kv_cache): - self.jiuge_awq_model.drop_kv_cache(kv_cache) - - def batch_infer_one_round(self, tasks: List[InferTask]): - output = (c_uint * len(tasks))() - batch_inputs = JiugeAWQBatchedTask(tasks) - self.jiuge_awq_model.infer_batch( - self.model_instance, - *(batch_inputs.input_args()), - output, - ) - return list(output) - - def generate(self, input_content, max_steps, topp_=1.0, topk_=1, temperature_=1.0): - input_content = self.tokenizer.apply_chat_template( - conversation=[{"role": "user", "content": input_content}], - add_generation_prompt=True, - tokenize=False, - ) - print(input_content, end="", flush=True) - tokens = self.tokenizer.encode(input_content) - infer_task = InferTask( - 0, - tokens, - self.max_context_len(), - temperature_, - topk_, - topp_, - self.eos_token_id, - ) - infer_task.bind_kvcache(KVCache(self)) - - steps = 0 - total_time = 0 - output_content = "" - - for step_i in range(max_steps): - start_time = time.time() - output_tokens = self.batch_infer_one_round([infer_task]) - end_time = time.time() - steps += 1 - output_str = self.tokenizer.decode(output_tokens[0]) - output_content += output_str - print(output_str, end="", flush=True) - if output_tokens[0] in self.eos_token_id: - break - infer_task.next(output_tokens[0]) - - if step_i > 0: - total_time += end_time - start_time - - print("\n") - avg_time = total_time * 1000 / (steps - 1) - print(f"Time per step: {avg_time:.3f}ms") - - infer_task._kv_cache.drop(self) - return output_content, avg_time - - def perplexity(self, test_sequences: List[Sequence[int]], batch_size=10): - tasks = [ - InferTask(i, [], self.max_context_len(), 1.0, 1, 1.0, self.eos_token_id) - for i in range(batch_size) - ] - kv_caches = [KVCache(self) for _ in range(batch_size)] - - nll = 0.0 - total_len = 0 - - for i in range(0, len(test_sequences), batch_size): - batch_id = 0 - true_tokens = [] - while batch_id < batch_size and batch_id + i < len(test_sequences): - input_tokens = test_sequences[i + batch_id][:-1] - true_tokens.extend(test_sequences[i + batch_id][1:]) - tasks[batch_id].tokens = input_tokens - tasks[batch_id].bind_kvcache(kv_caches[batch_id]) - batch_id += 1 - - batch_inputs = JiugeAWQBatchedTask(tasks[:batch_id]) - logits = torch.zeros( - (batch_inputs.ntok, self.meta.dvoc), dtype=self.meta.torch_dtype_logits - ) - self.jiuge_awq_model.forward_batch( - self.model_instance, - batch_inputs.tokens, - batch_inputs.ntok, - batch_inputs.req_lens, - batch_inputs.nreq, - batch_inputs.req_pos, - batch_inputs.kv_caches, - logits.data_ptr(), - ) - - logits = logits.float() - token_ids = torch.tensor(true_tokens, dtype=torch.int64) # [ntok,] - log_probs = torch.nn.functional.log_softmax(logits, dim=-1) # (ntok, vocab) - token_logprobs = log_probs[ - torch.arange(batch_inputs.ntok), token_ids - ] # (ntok,) - - start = 0 - for l in batch_inputs.req_lens_list: - nll += -token_logprobs[start : start + l].sum().item() - start += l - total_len += token_logprobs.numel() - - for task in tasks: - task.release_kvcache() - - return math.exp(nll / total_len) - - def destroy_model_instance(self): - self.jiuge_awq_model.destroy_model(self.model_instance) - print("Model destroyed") - - -def test(): - if len(sys.argv) < 3: - print( - "Usage: python jiuge_awq.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] [n_device]" - ) - sys.exit(1) - model_path = sys.argv[2] - device_type = DeviceType.DEVICE_TYPE_CPU - if sys.argv[1] == "--cpu": - device_type = DeviceType.DEVICE_TYPE_CPU - elif sys.argv[1] == "--nvidia": - device_type = DeviceType.DEVICE_TYPE_NVIDIA - elif sys.argv[1] == "--cambricon": - device_type = DeviceType.DEVICE_TYPE_CAMBRICON - elif sys.argv[1] == "--ascend": - device_type = DeviceType.DEVICE_TYPE_ASCEND - elif sys.argv[1] == "--metax": - device_type = DeviceType.DEVICE_TYPE_METAX - elif sys.argv[1] == "--moore": - device_type = DeviceType.DEVICE_TYPE_MOORE - elif sys.argv[1] == "--iluvatar": - device_type = DeviceType.DEVICE_TYPE_ILUVATAR - else: - print( - "Usage: python main_jiuge_awq.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] [n_device]" - ) - sys.exit(1) - - ndev = int(sys.argv[3]) if len(sys.argv) > 3 else 1 - model = JiugeAWQForCausalLM(model_path, device_type, ndev) - model.generate("山东最高的山是？", 500) - model.destroy_model_instance() - - -if __name__ == "__main__": - test() diff --git a/scripts/jiuge_gptq.py b/scripts/jiuge_gptq.py deleted file mode 100644 index e510847cc..000000000 --- a/scripts/jiuge_gptq.py +++ /dev/null @@ -1,369 +0,0 @@ -from typing import List, Sequence -import math -import os -from pathlib import Path -import safetensors -import sys -import time -import json -import torch -import transformers - -from libinfinicore_infer import ( - JiugeGPTQModel, - JiugeGPTQMetaCStruct, - DataType, - DeviceType, - KVCacheCStruct, -) -from infer_task import InferTask, KVCache - -from ctypes import POINTER, c_float, c_int, c_uint, c_void_p, byref - -torch.set_default_device("cpu") - - -class JiugeGPTQMetaFromConfig(JiugeGPTQMetaCStruct): - def __init__(self, config, dtype=torch.float16, max_tokens=None): - if config["torch_dtype"] == "float16": - dt_ = DataType.INFINI_DTYPE_F16 - elif config["torch_dtype"] == "float32": - dt_ = DataType.INFINI_DTYPE_F32 - elif config["torch_dtype"] == "bfloat16": - dt_ = DataType.INFINI_DTYPE_BF16 - else: - dt_ = DataType.INFINI_DTYPE_F16 - - # GPTQ 没有 scale 字段 - - has_qkv_bias = 1 if ("attention_bias" in config and config["attention_bias"]) else 0 - if config["model_type"] in ["qwen2", "qwen3"]: - has_qkv_bias = 1 - eos_token_id = ( - config["eos_token_id"][0] - if type(config["eos_token_id"]) == list - else config["eos_token_id"] - ) - - super().__init__( - dt_logits=dt_, - dt_linear_w=DataType.INFINI_DTYPE_I32, - dt_norm_w=dt_, - nlayer=config["num_hidden_layers"], - d=config["hidden_size"], - nh=config["num_attention_heads"], - nkvh=( - config["num_key_value_heads"] - if "num_key_value_heads" in config - else config["num_attention_heads"] - ), - dh=config["hidden_size"] // config["num_attention_heads"], - di=config["intermediate_size"], - dctx=( - config["max_position_embeddings"] if max_tokens is None else max_tokens - ), - dvoc=config["vocab_size"], - epsilon=config["rms_norm_eps"], - theta=(config["rope_theta"] if "rope_theta" in config else 100000.0), - end_token=eos_token_id, - nbit=config["quantization_config"]["bits"], - quant_group_size=config["quantization_config"]["group_size"], - has_qkv_bias=has_qkv_bias, - ) - self.torch_dtype_logits = dtype - - -class JiugeGPTQBatchedTask: - def __init__(self, tasks: List[InferTask]): - self.tasks = tasks - self.nreq = len(tasks) - - # Precompute fields - token_lists = [t.tokens for t in tasks] - self.req_lens_list = [len(toks) for toks in token_lists] - self.req_pos_list = [t.pos for t in tasks] - self.kv_cache_ptrs = [t.kvcache().data() for t in tasks] - self.temperaturas_list = [t.temperature for t in tasks] - self.topks_list = [t.topk for t in tasks] - self.topps_list = [t.topp for t in tasks] - - # Flatten token lists - flat_tokens = [tok for toks in token_lists for tok in toks] - self.ntok = len(flat_tokens) - - # Convert to ctypes arrays in one pass - self.tokens = (c_uint * self.ntok)(*flat_tokens) - self.req_lens = (c_uint * self.nreq)(*self.req_lens_list) - self.req_pos = (c_uint * self.nreq)(*self.req_pos_list) - self.kv_caches = (POINTER(KVCacheCStruct) * self.nreq)(*self.kv_cache_ptrs) - self.temperaturas = (c_float * self.nreq)(*self.temperaturas_list) - self.topks = (c_uint * self.nreq)(*self.topks_list) - self.topps = (c_float * self.nreq)(*self.topps_list) - - def input_args(self): - return ( - self.tokens, - self.ntok, - self.req_lens, - self.nreq, - self.req_pos, - self.kv_caches, - self.temperaturas, - self.topks, - self.topps, - ) - - -class JiugeGPTQForCausalLM: - def __init__( - self, model_dir_path, device=DeviceType.DEVICE_TYPE_CPU, ndev=1, max_tokens=None - ): - - load_start_time = time.time() - print(f"Creating model on {ndev} devices...") - with open(os.path.join(model_dir_path, "config.json"), "r") as f: - config = json.load(f) - self.config = config - eos_token_id = self.config["eos_token_id"] - self.eos_token_id = ( - [eos_token_id] if type(eos_token_id) == int else eos_token_id - ) - self.dev_ids = (c_int * ndev)(*[i for i in range(ndev)]) - self.ndev = ndev - self.device = device - self.meta = JiugeGPTQMetaFromConfig(config, max_tokens=max_tokens) - - self.jiuge_gptq_model = JiugeGPTQModel() - - self.weights = self.jiuge_gptq_model.create_weights( - byref(self.meta), - self.device, - ndev, - self.dev_ids, - ) - self.tokenizer = transformers.AutoTokenizer.from_pretrained( - model_dir_path, trust_remote_code=True - ) - - load_end_time = time.time() - print(f"Time used: {load_end_time - load_start_time:.3f}s") - - load_start_time = time.time() - print("Loading model weights to host...") - - self.load_all_safetensors_from_dir(os.path.join(model_dir_path)) - - self.model_instance = self.jiuge_gptq_model.create_model( - byref(self.meta), - self.weights, - ) - load_end_time = time.time() - print(f"Time used: {load_end_time - load_start_time:.3f}s") - - def load_all_safetensors_from_dir(self, dir_path_: str): - dir_path_ = Path(dir_path_) - lm_head_loaded = False - embed_tokens_tensor = None - for file in sorted(dir_path_.glob("*.safetensors")): - with safetensors.safe_open(file, framework="pt", device="cpu") as f: - for key in f.keys(): - tensor = f.get_tensor(key) - if key == "lm_head.weight": - lm_head_loaded = True - elif key == "model.embed_tokens.weight": - embed_tokens_tensor = tensor - # Some GPTQ exports (e.g. Qwen2.5 GPTQ with desc_act=False) - # store a dummy g_idx tensor (often all zeros). Our backend - # still consumes g_idx, so we convert the dummy to identity - # to keep dequant correct. - tensor_ptr = tensor.data_ptr() - if key.endswith(".g_idx") and tensor.dtype == torch.int32 and tensor.numel() > 1: - try: - if int(tensor.max().item()) == 0: - if not hasattr(self, "_gidx_fix_buffers"): - self._gidx_fix_buffers = [] - fixed = torch.arange(tensor.numel(), dtype=torch.int32) - self._gidx_fix_buffers.append(fixed) - tensor_ptr = fixed.data_ptr() - except Exception: - pass - - self.jiuge_gptq_model.load_weight( - self.weights, key, tensor_ptr - ) - if not lm_head_loaded and embed_tokens_tensor is not None: - print("lm_head.weight missing, tying to embed_tokens.weight") - self.jiuge_gptq_model.load_weight( - self.weights, "lm_head.weight", embed_tokens_tensor.data_ptr() - ) - elif not lm_head_loaded: - raise RuntimeError("lm_head.weight missing and embed_tokens.weight not found") - - def max_context_len(self): - return self.meta.dctx - - def create_kv_cache(self): - return self.jiuge_gptq_model.create_kv_cache( - self.meta.nlayer, - self.meta.dctx, - self.meta.nkvh, - self.meta.dh, - self.meta.dh, - self.meta.dt_logits, - self.device, - self.dev_ids, - self.ndev, - ) - - def drop_kv_cache(self, kv_cache): - self.jiuge_gptq_model.drop_kv_cache(kv_cache) - - def batch_infer_one_round(self, tasks: List[InferTask]): - output = (c_uint * len(tasks))() - batch_inputs = JiugeGPTQBatchedTask(tasks) - self.jiuge_gptq_model.infer_batch( - self.model_instance, - *(batch_inputs.input_args()), - output, - ) - return list(output) - - def generate(self, input_content, max_steps, topp_=1.0, topk_=1, temperature_=1.0): - input_content = self.tokenizer.apply_chat_template( - conversation=[{"role": "user", "content": input_content}], - add_generation_prompt=True, - tokenize=False, - ) - print(input_content, end="", flush=True) - tokens = self.tokenizer.encode(input_content) - infer_task = InferTask( - 0, - tokens, - self.max_context_len(), - temperature_, - topk_, - topp_, - self.eos_token_id, - ) - infer_task.bind_kvcache(KVCache(self)) - - steps = 0 - total_time = 0 - output_content = "" - - for step_i in range(max_steps): - start_time = time.time() - output_tokens = self.batch_infer_one_round([infer_task]) - end_time = time.time() - steps += 1 - output_str = self.tokenizer.decode(output_tokens[0]) - output_content += output_str - print(output_str, end="", flush=True) - if output_tokens[0] in self.eos_token_id: - break - infer_task.next(output_tokens[0]) - - if step_i > 0: - total_time += end_time - start_time - - print("\n") - avg_time = total_time * 1000 / (steps - 1) - print(f"Time per step: {avg_time:.3f}ms") - - infer_task._kv_cache.drop(self) - return output_content, avg_time - - def perplexity(self, test_sequences: List[Sequence[int]], batch_size=10): - tasks = [ - InferTask(i, [], self.max_context_len(), 1.0, 1, 1.0, self.eos_token_id) - for i in range(batch_size) - ] - kv_caches = [KVCache(self) for _ in range(batch_size)] - - nll = 0.0 - total_len = 0 - - for i in range(0, len(test_sequences), batch_size): - batch_id = 0 - true_tokens = [] - while batch_id < batch_size and batch_id + i < len(test_sequences): - input_tokens = test_sequences[i + batch_id][:-1] - true_tokens.extend(test_sequences[i + batch_id][1:]) - tasks[batch_id].tokens = input_tokens - tasks[batch_id].bind_kvcache(kv_caches[batch_id]) - batch_id += 1 - - batch_inputs = JiugeGPTQBatchedTask(tasks[:batch_id]) - logits = torch.zeros( - (batch_inputs.ntok, self.meta.dvoc), dtype=self.meta.torch_dtype_logits - ) - self.jiuge_gptq_model.forward_batch( - self.model_instance, - batch_inputs.tokens, - batch_inputs.ntok, - batch_inputs.req_lens, - batch_inputs.nreq, - batch_inputs.req_pos, - batch_inputs.kv_caches, - logits.data_ptr(), - ) - - logits = logits.float() - token_ids = torch.tensor(true_tokens, dtype=torch.int64) # [ntok,] - log_probs = torch.nn.functional.log_softmax(logits, dim=-1) # (ntok, vocab) - token_logprobs = log_probs[ - torch.arange(batch_inputs.ntok), token_ids - ] # (ntok,) - - start = 0 - for l in batch_inputs.req_lens_list: - nll += -token_logprobs[start : start + l].sum().item() - start += l - total_len += token_logprobs.numel() - - for task in tasks: - task.release_kvcache() - - return math.exp(nll / total_len) - - def destroy_model_instance(self): - self.jiuge_gptq_model.destroy_model(self.model_instance) - print("Model destroyed") - - -def test(): - if len(sys.argv) < 3: - print( - "Usage: python jiuge_gptq.py [--cpu | --nvidia | --cambricon | --ascend | --metax | --moore | --iluvatar] [n_device]" - ) - sys.exit(1) - model_path = sys.argv[2] - device_type = DeviceType.DEVICE_TYPE_CPU - if sys.argv[1] == "--cpu": - device_type = DeviceType.DEVICE_TYPE_CPU - elif sys.argv[1] == "--nvidia": - device_type = DeviceType.DEVICE_TYPE_NVIDIA - elif sys.argv[1] == "--cambricon": - device_type = DeviceType.DEVICE_TYPE_CAMBRICON - elif sys.argv[1] == "--ascend": - device_type = DeviceType.DEVICE_TYPE_ASCEND - elif sys.argv[1] == "--metax": - device_type = DeviceType.DEVICE_TYPE_METAX - elif sys.argv[1] == "--moore": - device_type = DeviceType.DEVICE_TYPE_MOORE - elif sys.argv[1] == "--iluvatar": - device_type = DeviceType.DEVICE_TYPE_ILUVATAR - else: - print( - "Usage: python jiuge_gptq.py [--cpu | --nvidia | --cambricon | --ascend | --metax | --moore | --iluvatar] [n_device]" - ) - sys.exit(1) - - ndev = int(sys.argv[3]) if len(sys.argv) > 3 else 1 - model = JiugeGPTQForCausalLM(model_path, device_type, ndev) - model.generate("山东最高的山是？", 500) - model.destroy_model_instance() - - -if __name__ == "__main__": - test() diff --git a/scripts/jiuge_ppl.py b/scripts/jiuge_ppl.py deleted file mode 100644 index 923d209c6..000000000 --- a/scripts/jiuge_ppl.py +++ /dev/null @@ -1,115 +0,0 @@ -import torch -from transformers import AutoTokenizer, AutoModelForCausalLM -from datasets import load_dataset -from jiuge import JiugeForCauslLM -from libinfinicore_infer import DeviceType - -DEVICE_TYPE_MAP = { - "cpu": DeviceType.DEVICE_TYPE_CPU, - "nvidia": DeviceType.DEVICE_TYPE_NVIDIA, - "qy": DeviceType.DEVICE_TYPE_QY, - "cambricon": DeviceType.DEVICE_TYPE_CAMBRICON, - "ascend": DeviceType.DEVICE_TYPE_ASCEND, - "metax": DeviceType.DEVICE_TYPE_METAX, - "moore": DeviceType.DEVICE_TYPE_MOORE, - "iluvatar": DeviceType.DEVICE_TYPE_ILUVATAR, - "kunlun": DeviceType.DEVICE_TYPE_KUNLUN, - "hygon": DeviceType.DEVICE_TYPE_HYGON, -} - -TORCH_DEVICE_TYPE_MAP = { - "cpu": "cpu", - "nvidia": "cuda", - "qy": "cuda", - "cambricon": "mlu", - "ascend": "npu", - "metax": "cuda", - "moore": "cuda", - "iluvatar": "cuda", - "kunlun": "cuda", - "hygon": "cuda", -} - - -def test_torch(input_ids_list, device_): - device = TORCH_DEVICE_TYPE_MAP[device_] - model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).to( - device - ) - model.eval() - - total_neg_log_likelihood = 0 - total_tokens = 0 - - with torch.no_grad(): - for input_ids in input_ids_list: - input_ids = torch.tensor(input_ids, device=device) - # shift inputs and labels - inputs = input_ids[:-1].unsqueeze(0) # [1, seq_len-1] - labels = input_ids[1:].unsqueeze(0) # [1, seq_len-1] - - outputs = model(inputs, use_cache=False) - logits = outputs.logits # [1, seq_len-1, vocab_size] - - log_probs = torch.nn.functional.log_softmax(logits, dim=-1) - # gather log probs of true tokens - true_token_log_probs = log_probs.gather( - dim=-1, index=labels.unsqueeze(-1) - ).squeeze(-1) - - total_neg_log_likelihood += -true_token_log_probs.sum().item() - total_tokens += labels.numel() - - perplexity = torch.exp(torch.tensor(total_neg_log_likelihood / total_tokens)) - return perplexity - - -def test_infinicore(input_ids_list, device_, ndev_): - device = DEVICE_TYPE_MAP[device_] - - model = JiugeForCauslLM( - model_path, device, max_tokens=len(input_ids_list[0]), ndev=ndev_ - ) - perplexity = model.perplexity(input_ids_list) - model.destroy_model_instance() - return perplexity - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("--model-path", type=str, required=True) - parser.add_argument( - "--dev", type=str, default="cpu", choices=DEVICE_TYPE_MAP.keys() - ) - parser.add_argument( - "--ndev", - type=int, - default=1, - help="Number of devices to use (default: 1)", - ) - args = parser.parse_args() - - seq_len = 512 - - model_path = args.model_path - tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) - dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test") - - texts = dataset["text"] - texts = [t.strip() for t in texts if len(t.strip()) > 0] - - input_ids_list = [] - for text in texts: - ids = tokenizer.encode(text) - # split long sequences into chunks - for i in range(0, len(ids) - seq_len + 1, seq_len): - input_ids_list.append(ids[i : i + seq_len]) - - perplexity = test_infinicore(input_ids_list, args.dev, args.ndev) - print(f"InfiniCore Perplexity: {perplexity:.2f}") - - if args.ndev == 1: # Todo: support multi-device testing with torch - perplexity = test_torch(input_ids_list, args.dev) - print(f"Torch Perplexity: {perplexity.item():.2f}") diff --git a/scripts/kvcache_pool.py b/scripts/kvcache_pool.py deleted file mode 100644 index 81914535d..000000000 --- a/scripts/kvcache_pool.py +++ /dev/null @@ -1,90 +0,0 @@ -from infer_task import KVCache - -import asyncio -from typing import List -import threading - - -class KVCachePool: - def __init__(self, model, max_caches: int = 32): - self.max_caches = max_caches - self.model = model - self._available: List[KVCache] = [] - self.num_caches = len(self._available) - self._lock = threading.Lock() - self._not_empty = threading.Condition(self._lock) - self._shutdown = False - - def acquire_sync(self, infer_task): - with self._not_empty: - while True: - if self._shutdown: - raise RuntimeError( - "KVCachePool is shutting down; cannot acquire new cache." - ) - if len(self._available) == 0: - if self.num_caches < self.max_caches: - self.num_caches += 1 - print( - f"[INFO] Task {infer_task.id} created new KVCachePoolItem" - ) - return infer_task.bind_kvcache(KVCache(self.model), 0) - else: - self._not_empty.wait() - else: - max_match, max_match_index = self.find_most_matching_cache( - infer_task.tokens - ) - kvcache = self._available.pop(max_match_index) - print( - f"[INFO] Task {infer_task.id} reused KVCachePoolItem {max_match_index} with {max_match} matches" - ) - return infer_task.bind_kvcache(kvcache, max_match) - - def release_sync(self, infer_task): - with self._not_empty: - print(f"[INFO] Task {infer_task.id} returned KVCachePoolItem to pool") - self._available.append(infer_task.release_kvcache()) - self._not_empty.notify() - - async def acquire(self, infer_task): - loop = asyncio.get_event_loop() - return await loop.run_in_executor(None, self.acquire_sync, infer_task) - - async def release(self, infer_task): - loop = asyncio.get_event_loop() - return await loop.run_in_executor(None, self.release_sync, infer_task) - - def find_most_matching_cache(self, tokens: List[int]): - max_match = 0 - max_match_index = 0 - - def first_different_index(a_, b_): - for i_, (x_, y_) in enumerate(zip(a_, b_)): - if x_ != y_: - return i_ - return min(len(a_), len(b_)) - - for i, kvcache in enumerate(self._available): - common_elements = first_different_index(tokens, kvcache.tokens) - # print(f"{tokens}") - # print(f"{kvcache.tokens[:len(tokens)]}") - if common_elements > max_match: - max_match = common_elements - max_match_index = i - - return (min(max_match, len(tokens) - 1), max_match_index) - - def finalize(self): - with self._not_empty: - self._shutdown = True - while len(self._available) < self.num_caches: - self._not_empty.wait() - - for kvcache in self._available: - if kvcache is not None: - kvcache.drop(self.model) - - self._available.clear() - self.max_caches = 0 - self.num_caches = 0 diff --git a/scripts/launch_server.py b/scripts/launch_server.py deleted file mode 100644 index d04d4f69d..000000000 --- a/scripts/launch_server.py +++ /dev/null @@ -1,329 +0,0 @@ -from jiuge import JiugeForCauslLM -from jiuge_awq import JiugeAWQForCausalLM -from libinfinicore_infer import DeviceType -from infer_task import InferTask -from kvcache_pool import KVCachePool - -import argparse -import queue -from fastapi import FastAPI, Request -from fastapi.responses import StreamingResponse, JSONResponse -import contextlib -import uvicorn -import time -import uuid -import json -import threading -import janus - - -DEVICE_TYPE_MAP = { - "cpu": DeviceType.DEVICE_TYPE_CPU, - "nvidia": DeviceType.DEVICE_TYPE_NVIDIA, - "qy": DeviceType.DEVICE_TYPE_QY, - "cambricon": DeviceType.DEVICE_TYPE_CAMBRICON, - "ascend": DeviceType.DEVICE_TYPE_ASCEND, - "metax": DeviceType.DEVICE_TYPE_METAX, - "moore": DeviceType.DEVICE_TYPE_MOORE, - "iluvatar": DeviceType.DEVICE_TYPE_ILUVATAR, - "kunlun": DeviceType.DEVICE_TYPE_KUNLUN, - "hygon": DeviceType.DEVICE_TYPE_HYGON, -} - - -def parse_args(): - parser = argparse.ArgumentParser(description="Launch the LLM inference server.") - parser.add_argument( - "--model-path", - type=str, - help="Path to the model directory", - ) - parser.add_argument( - "--dev", - type=str, - choices=DEVICE_TYPE_MAP.keys(), - default="cpu", - help="Device type to run the model on (default: cpu)", - ) - parser.add_argument( - "--ndev", - type=int, - default=1, - help="Number of devices to use (default: 1)", - ) - parser.add_argument( - "--max-batch", - type=int, - default=3, - help="Maximum number of requests that can be batched together (default: 3)", - ) - parser.add_argument( - "--max-tokens", - type=int, - required=False, - default=None, - help="Max token sequence length that model will handle (follows model config if not provided)", - ) - parser.add_argument( - "--awq", - action="store_true", - help="Whether to use AWQ quantized model (default: False)", - ) - # Add support for GPTQ - parser.add_argument( - "--gptq", - action="store_true", - help="Whether to use GPTQ quantized model (default: False)", - ) - return parser.parse_args() - - -args = parse_args() -device_type = DEVICE_TYPE_MAP[args.dev] -model_path = args.model_path -ndev = args.ndev -max_tokens = args.max_tokens -USE_AWQ = args.awq -USE_GPTQ = args.gptq -MAX_BATCH = args.max_batch -print( - f"Using MAX_BATCH={MAX_BATCH}. Try reduce this value if out of memory error occurs." -) - - -def chunk_json(id_, content=None, role=None, finish_reason=None): - delta = {} - if content: - delta["content"] = content - if role: - delta["role"] = role - return { - "id": id_, - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": "jiuge", - "system_fingerprint": None, - "choices": [ - { - "index": 0, - "text": content, - "delta": delta, - "logprobs": None, - "finish_reason": finish_reason, - } - ], - } - - -# A wrapper for InferTask that supports async output queue -class AsyncInferTask(InferTask): - def __init__(self, id, tokens, max_tokens, temperature, topk, topp, end_tokens): - super().__init__(id, tokens, max_tokens, temperature, topk, topp, end_tokens) - self.output_queue = janus.Queue() - print(f"[INFO] Create InferTask {self.id}") - - def output(self, out_token): - self.next(out_token) - self.output_queue.sync_q.put(out_token) - - -@contextlib.asynccontextmanager -async def lifespan(app: FastAPI): - # Startup - if USE_AWQ: - app.state.model = JiugeAWQForCausalLM( - model_path, device_type, ndev, max_tokens=max_tokens - ) - elif USE_GPTQ: - app.state.model = JiugeGPTQForCausalLM( - model_path, device_type, ndev, max_tokens=max_tokens - ) - else: - app.state.model = JiugeForCauslLM( - model_path, device_type, ndev, max_tokens=max_tokens - ) - app.state.kv_cache_pool = KVCachePool(app.state.model, MAX_BATCH) - app.state.request_queue = janus.Queue() - worker_thread = threading.Thread(target=worker_loop, args=(app,), daemon=True) - worker_thread.start() - - try: - yield # The app runs here - finally: - # Shutdown - app.state.request_queue.sync_q.put(None) - worker_thread.join() - app.state.request_queue.shutdown() - - app.state.kv_cache_pool.finalize() - app.state.model.destroy_model_instance() - - -App = FastAPI(lifespan=lifespan) - - -# App loop: take requests from the queue, do inference, and put unfinished requests back into the queue. -def worker_loop(app): - while True: - try: - task = app.state.request_queue.sync_q.get(timeout=0.01) - except queue.Empty: - continue - - if task is None: - return - - batch = [task] - while len(batch) < MAX_BATCH: - try: - req = app.state.request_queue.sync_q.get_nowait() - if req is not None: - batch.append(req) - except queue.Empty: - break - output_tokens = app.state.model.batch_infer_one_round(batch) - for task, token in zip(batch, output_tokens): - task.output(token) - if task.finish_reason is None: - app.state.request_queue.sync_q.put(task) - else: - print(f"[INFO] Task {task.id} finished infer.") - app.state.kv_cache_pool.release_sync(task) - - -def build_task(id_, request_data, request: Request): - messages = request_data.get("messages", []) - input_content = request.app.state.model.tokenizer.apply_chat_template( - conversation=messages, - add_generation_prompt=True, - tokenize=False, - ) - tokens = request.app.state.model.tokenizer.encode(input_content) - return AsyncInferTask( - id_, - tokens, - request_data.get("max_tokens", request.app.state.model.max_context_len()), - request_data.get("temperature", 1.0), - request_data.get("top_k", 1), - request_data.get("top_p", 1.0), - request.app.state.model.eos_token_id, - ) - - -async def chat_stream(id_, request_data, request: Request): - try: - infer_task = build_task(id_, request_data, request) - await request.app.state.kv_cache_pool.acquire(infer_task) - - # Initial empty content - chunk = json.dumps( - chunk_json(id_, content="", role="assistant"), ensure_ascii=False - ) - yield f"data: {chunk}\n\n" - - request.app.state.request_queue.sync_q.put(infer_task) - - while True: - if await request.is_disconnected(): - print("Client disconnected. Aborting stream.") - break - if ( - infer_task.finish_reason is not None - and infer_task.output_queue.async_q.empty() - ): - chunk = json.dumps( - chunk_json(id_, finish_reason=infer_task.finish_reason), - ensure_ascii=False, - ) - yield f"data: {chunk}\n\n" - break - - token = await infer_task.output_queue.async_q.get() - content = request.app.state.model.tokenizer.decode(token) - - chunk = json.dumps(chunk_json(id_, content=content), ensure_ascii=False) - yield f"data: {chunk}\n\n" - - except Exception as e: - print(f"[Error] ID : {id_} Exception: {e}") - finally: - if infer_task.finish_reason is None: - infer_task.finish_reason = "cancel" - - -async def chat(id_, request_data, request: Request): - try: - infer_task = build_task(id_, request_data, request) - await request.app.state.kv_cache_pool.acquire(infer_task) - request.app.state.request_queue.sync_q.put(infer_task) - output = [] - while True: - if ( - infer_task.finish_reason is not None - and infer_task.output_queue.async_q.empty() - ): - break - - token = await infer_task.output_queue.async_q.get() - content = request.app.state.model.tokenizer.decode(token) - output.append(content) - - output_text = "".join(output).strip() - response = chunk_json( - id_, - content=output_text, - role="assistant", - finish_reason=infer_task.finish_reason or "stop", - ) - return response - - except Exception as e: - print(f"[Error] ID: {id_} Exception: {e}") - return JSONResponse(content={"error": str(e)}, status_code=500) - finally: - if infer_task.finish_reason is None: - infer_task.finish_reason = "cancel" - - -@App.post("/chat/completions") -async def chat_completions(request: Request): - data = await request.json() - print('-----------------------------------------') - print(data) - print('-----------------------------------------') - - if not data.get("messages"): - if not data.get("prompt"): - return JSONResponse(content={"error": "No message provided"}, status_code=400) - else: - data['messages'] = [{"role": "user", "content": data.get("prompt")}] - - stream = data.get("stream", False) - id_ = f"cmpl-{uuid.uuid4().hex}" - if stream: - return StreamingResponse( - chat_stream(id_, data, request), media_type="text/event-stream" - ) - else: - response = await chat(id_, data, request) - return JSONResponse(content=response) - - -if __name__ == "__main__": - uvicorn.run(App, host="0.0.0.0", port=8000) - -""" -curl -N -H "Content-Type: application/json" \ - -X POST http://127.0.0.1:8000/chat/completions \ - -d '{ - "model": "jiuge", - "messages": [ - {"role": "user", "content": "山东最高的山是？"} - ], - "temperature": 1.0, - "top_k": 50, - "top_p": 0.8, - "max_tokens": 512, - "stream": true - }' -""" diff --git a/scripts/libinfinicore_infer/__init__.py b/scripts/libinfinicore_infer/__init__.py deleted file mode 100644 index 8a6fae747..000000000 --- a/scripts/libinfinicore_infer/__init__.py +++ /dev/null @@ -1,55 +0,0 @@ -from .base import DataType, DeviceType, KVCacheCStruct -from .jiuge import JiugeModel, JiugeMetaCStruct, JiugeWeightsCStruct -# 为了区分, 给 ModelWeightsCStruct 别名 -from .jiuge_awq import JiugeAWQModel, JiugeAWQMetaCStruct, ModelWeightsCStruct as AWQModelWeightsCStruct -# 添加 GPTQ 模块 -from .jiuge_gptq import JiugeGPTQModel, JiugeGPTQMetaCStruct, ModelWeightsCStruct as GPTQModelWeightsCStruct -from .deepseek_v3 import ( - DeepSeekV3Model, - DeepSeekV3MetaCStruct, - DeepSeekV3WeightsCStruct, - DeepSeekV3WeightLoaderCStruct, - DeepSeekV3CacheCStruct, -) -from .qwen3vl import ( - Qwen3vlModel, - Qwen3vlMetaCStruct, - TextMetaCStruct, - VisMetaCStruct, - Qwen3vlWeightsCStruct, - Qwen3vlWeightLoaderCStruct, - Qwen3vlVisWeightLoaderCStruct, - Qwen3vlLangWeightLoaderCStruct, - Qwen3vlCacheCStruct, -) - -__all__ = [ - "DataType", - "DeviceType", - "KVCacheCStruct", - "JiugeModel", - "JiugeMetaCStruct", - "JiugeWeightsCStruct", - "JiugeAWQModel", - "JiugeAWQMetaCStruct", - "AWQModelWeightsCStruct", - # Add GPTQ module - "JiugeGPTQModel", - "JiugeGPTQMetaCStruct", - "GPTQModelWeightsCStruct", - "DeepSeekV3Model", - "DeepSeekV3MetaCStruct", - "DeepSeekV3WeightsCStruct", - "DeepSeekV3WeightLoaderCStruct", - "DeepSeekV3CacheCStruct", - "Qwen3vlModel", - "Qwen3vlMetaCStruct", - "TextMetaCStruct", - "VisMetaCStruct", - "Qwen3vlWeightsCStruct", - "Qwen3vlWeightLoaderCStruct", - "Qwen3vlVisWeightLoaderCStruct", - "Qwen3vlLangWeightLoaderCStruct", - "Qwen3vlCacheCStruct", - "ModelRegister", -] diff --git a/scripts/libinfinicore_infer/base.py b/scripts/libinfinicore_infer/base.py deleted file mode 100644 index 2af0c238a..000000000 --- a/scripts/libinfinicore_infer/base.py +++ /dev/null @@ -1,72 +0,0 @@ -import ctypes -from ctypes import c_char, c_char_p, c_size_t, c_uint, c_int, c_float, c_void_p, POINTER -import os - - -class DataType(ctypes.c_int): - 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 - - -class DeviceType(ctypes.c_int): - DEVICE_TYPE_CPU = 0 - DEVICE_TYPE_NVIDIA = 1 - DEVICE_TYPE_CAMBRICON = 2 - DEVICE_TYPE_ASCEND = 3 - DEVICE_TYPE_METAX = 4 - DEVICE_TYPE_MOORE = 5 - DEVICE_TYPE_ILUVATAR = 6 - DEVICE_TYPE_KUNLUN = 7 - DEVICE_TYPE_HYGON = 8 - DEVICE_TYPE_QY = 9 - DEVICE_TYPE_ALI = 10 - - -class KVCacheCStruct(ctypes.Structure): - pass - - -# Model registration system -_model_registry = [] - - -def register_model(model_class): - """Decorator to register a model class""" - _model_registry.append(model_class) - return model_class - - -def register_lib_functions(lib): - """Register all model functions with the library""" - for model_class in _model_registry: - model_class.register_lib(lib) - - -class BaseModel: - def __init__(self): - self.lib = self._load_library() - register_lib_functions(self.lib) - - def _load_library(self): - lib_path = os.path.join( - os.environ.get("INFINI_ROOT"), "lib", "libinfinicore_infer.so" - ) - return ctypes.CDLL(lib_path) diff --git a/scripts/libinfinicore_infer/deepseek_v3.py b/scripts/libinfinicore_infer/deepseek_v3.py deleted file mode 100644 index b2c380b72..000000000 --- a/scripts/libinfinicore_infer/deepseek_v3.py +++ /dev/null @@ -1,209 +0,0 @@ -from .base import BaseModel, DataType, DeviceType, KVCacheCStruct, register_model -from ctypes import ( - c_size_t, - c_uint, - c_int, - c_float, - c_void_p, - POINTER, - Structure, - CFUNCTYPE, -) - - -class DeepSeekV3MetaCStruct(Structure): - _fields_ = [ - ("dt_logits", DataType), - ("dt_norm", DataType), - ("dt_quant_weight", DataType), - ("dt_quant_scale", DataType), - ("dt_quant_zero", DataType), - ("dt_gate_weight", DataType), - ("dt_gate_bias", DataType), - ("n_sparse_layer", c_size_t), - ("n_dense_layer", c_size_t), - ("d", c_size_t), - ("nh", c_size_t), - ("nkvh", c_size_t), - ("d_rope", c_size_t), - ("d_nope", c_size_t), - ("r_q", c_size_t), - ("r_kv", c_size_t), - ("d_qk", c_size_t), - ("d_v", c_size_t), - ("routed_scale", c_float), - ("nexperts", c_size_t), - ("kexperts", c_size_t), - ("di", c_size_t), - ("di_moe", c_size_t), - ("dctx", c_size_t), - ("dvoc", c_size_t), - ("epsilon", c_float), - ("rope_theta", c_float), - ("end_token", c_uint), - ] - - -class DeepSeekV3WeightsCStruct(Structure): - pass - - -class DeepSeekV3ModelCStruct(Structure): - pass - - -class DeepSeekV3CacheCStruct(Structure): - pass - - -load_global_fn = CFUNCTYPE(None, POINTER(DeepSeekV3WeightsCStruct), c_void_p) -load_layer_fn = CFUNCTYPE(None, POINTER(DeepSeekV3WeightsCStruct), c_void_p, c_size_t) -load_layer_linear_fn = CFUNCTYPE( - None, POINTER(DeepSeekV3WeightsCStruct), c_void_p, c_void_p, c_void_p, c_size_t -) -load_layer_mlp_fn = CFUNCTYPE( - None, - POINTER(DeepSeekV3WeightsCStruct), - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_size_t, -) -load_layer_expert_mlp_fn = CFUNCTYPE( - None, - POINTER(DeepSeekV3WeightsCStruct), - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_void_p, - c_size_t, - c_size_t, -) - - -class DeepSeekV3WeightLoaderCStruct(Structure): - _fields_ = [ - ("load_input_embd", load_global_fn), - ("load_output_norm", load_global_fn), - ("load_output_embd", load_global_fn), - ("load_attn_norm", load_layer_fn), - ("load_attn_q_a_proj", load_layer_linear_fn), - ("load_attn_q_a_layernorm", load_layer_fn), - ("load_attn_q_b_proj", load_layer_linear_fn), - ("load_attn_kv_a_proj_with_mqa", load_layer_linear_fn), - ("load_attn_kv_a_layernorm", load_layer_fn), - ("load_attn_kv_b_proj", load_layer_linear_fn), - ("load_attn_o_proj", load_layer_linear_fn), - ("load_mlp_norm", load_layer_fn), - ("load_mlp_dense", load_layer_mlp_fn), - ("load_mlp_gate_weight", load_layer_fn), - ("load_mlp_gate_bias", load_layer_fn), - ("load_mlp_shared_experts", load_layer_mlp_fn), - ("load_mlp_experts", load_layer_expert_mlp_fn), - ] - - -@register_model -class DeepSeekV3Model(BaseModel): - @classmethod - def register_lib(cls, lib): - """Register DeepSeekV3 model functions with the library""" - lib.createDeepSeekV3WeightLoader.argtypes = [] - lib.createDeepSeekV3WeightLoader.restype = POINTER( - DeepSeekV3WeightLoaderCStruct - ) - - lib.createDeepSeekV3Weights.argtypes = [ - POINTER(DeepSeekV3MetaCStruct), - DeviceType, - c_int, - POINTER(c_int), - ] - lib.createDeepSeekV3Weights.restype = POINTER(DeepSeekV3WeightsCStruct) - - lib.createDeepSeekV3Model.argtypes = [ - POINTER(DeepSeekV3MetaCStruct), - POINTER(DeepSeekV3WeightsCStruct), - ] - lib.createDeepSeekV3Model.restype = POINTER(DeepSeekV3ModelCStruct) - - lib.destroyDeepSeekV3Model.argtypes = [POINTER(DeepSeekV3ModelCStruct)] - - lib.createDeepSeekV3Cache.argtypes = [POINTER(DeepSeekV3ModelCStruct)] - lib.createDeepSeekV3Cache.restype = POINTER(DeepSeekV3CacheCStruct) - - lib.dropDeepSeekV3Cache.argtypes = [ - POINTER(DeepSeekV3ModelCStruct), - POINTER(DeepSeekV3CacheCStruct), - ] - - lib.inferBatchDeepSeekV3.argtypes = [ - POINTER(DeepSeekV3ModelCStruct), - POINTER(c_uint), - c_uint, - POINTER(c_uint), - c_uint, - POINTER(c_uint), - POINTER(POINTER(DeepSeekV3CacheCStruct)), - POINTER(c_float), - POINTER(c_uint), - POINTER(c_float), - POINTER(c_uint), - ] - - def create_weight_loader(self): - return self.lib.createDeepSeekV3WeightLoader() - - def create_weights(self, meta, device_type, ndev, dev_ids): - return self.lib.createDeepSeekV3Weights(meta, device_type, ndev, dev_ids) - - def create_model(self, meta, weights): - return self.lib.createDeepSeekV3Model(meta, weights) - - def destroy_model(self, model): - self.lib.destroyDeepSeekV3Model(model) - - def create_cache(self, model): - return self.lib.createDeepSeekV3Cache(model) - - def drop_cache(self, model, cache): - self.lib.dropDeepSeekV3Cache(model, cache) - - def infer_batch( - self, - model, - tokens, - ntok, - req_lens, - nreq, - req_pos, - caches, - temperature, - topk, - topp, - output, - ): - self.lib.inferBatchDeepSeekV3( - model, - tokens, - ntok, - req_lens, - nreq, - req_pos, - caches, - temperature, - topk, - topp, - output, - ) diff --git a/scripts/libinfinicore_infer/jiuge.py b/scripts/libinfinicore_infer/jiuge.py deleted file mode 100644 index fe2abf108..000000000 --- a/scripts/libinfinicore_infer/jiuge.py +++ /dev/null @@ -1,151 +0,0 @@ -from .base import BaseModel, DataType, DeviceType, KVCacheCStruct, register_model -from ctypes import c_size_t, c_uint, c_int, c_float, c_void_p, POINTER, Structure, byref - - -class JiugeMetaCStruct(Structure): - _fields_ = [ - ("dt_logits", DataType), - ("nlayer", c_size_t), - ("d", c_size_t), - ("nh", c_size_t), - ("nkvh", c_size_t), - ("dh", c_size_t), - ("di", c_size_t), - ("dctx", c_size_t), - ("dvoc", c_size_t), - ("epsilon", c_float), - ("theta", c_float), - ("end_token", c_uint), - ] - - -class JiugeWeightsCStruct(Structure): - _fields_ = [ - ("nlayer", c_size_t), - ("dt_norm", DataType), - ("dt_mat", DataType), - ("transpose_linear_weights", c_int), - ("input_embd", c_void_p), - ("output_norm", c_void_p), - ("output_embd", c_void_p), - ("attn_norm", POINTER(c_void_p)), - ("attn_qkv", POINTER(c_void_p)), - ("attn_qkv_b", POINTER(c_void_p)), - ("attn_q_norm", POINTER(c_void_p)), - ("attn_k_norm", POINTER(c_void_p)), - ("attn_o", POINTER(c_void_p)), - ("ffn_norm", POINTER(c_void_p)), - ("ffn_gate_up", POINTER(c_void_p)), - ("ffn_down", POINTER(c_void_p)), - ] - - -class JiugeModelCStruct(Structure): - pass - - -@register_model -class JiugeModel(BaseModel): - @classmethod - def register_lib(cls, lib): - lib.createJiugeModel.restype = POINTER(JiugeModelCStruct) - lib.createJiugeModel.argtypes = [ - POINTER(JiugeMetaCStruct), - POINTER(JiugeWeightsCStruct), - DeviceType, - c_int, - POINTER(c_int), - ] - - lib.destroyJiugeModel.argtypes = [POINTER(JiugeModelCStruct)] - - lib.createKVCache.argtypes = [ - c_size_t, - c_size_t, - c_size_t, - c_size_t, - c_size_t, - DataType, - DeviceType, - POINTER(c_int), - c_size_t, - ] - lib.createKVCache.restype = POINTER(KVCacheCStruct) - - lib.dropKVCache.argtypes = [POINTER(KVCacheCStruct)] - - lib.inferBatchJiuge.argtypes = [ - POINTER(JiugeModelCStruct), - POINTER(c_uint), - c_uint, - POINTER(c_uint), - c_uint, - POINTER(c_uint), - POINTER(POINTER(KVCacheCStruct)), - POINTER(c_float), - POINTER(c_uint), - POINTER(c_float), - POINTER(c_uint), - ] - - lib.forwardBatchJiuge.argtypes = [ - POINTER(JiugeModelCStruct), - POINTER(c_uint), - c_uint, - POINTER(c_uint), - c_uint, - POINTER(c_uint), - POINTER(POINTER(KVCacheCStruct)), - c_void_p, - ] - - def create_model(self, meta, weights, device_type, ndev, dev_ids): - return self.lib.createJiugeModel(meta, weights, device_type, ndev, dev_ids) - - def destroy_model(self, model): - self.lib.destroyJiugeModel(model) - - def create_kv_cache( - self, nlayer, max_len, nkvh, dk, dv, dtype, device, dev_ids, ndev - ): - return self.lib.createKVCache( - nlayer, max_len, nkvh, dk, dv, dtype, device, dev_ids, ndev - ) - - def drop_kv_cache(self, kv_cache): - self.lib.dropKVCache(kv_cache) - - def infer_batch( - self, - model, - tokens, - ntok, - req_lens, - nreq, - req_pos, - kv_caches, - temperature, - topk, - topp, - output, - ): - self.lib.inferBatchJiuge( - model, - tokens, - ntok, - req_lens, - nreq, - req_pos, - kv_caches, - temperature, - topk, - topp, - output, - ) - - def forward_batch( - self, model, tokens, ntok, req_lens, nreq, req_pos, kv_caches, logits - ): - self.lib.forwardBatchJiuge( - model, tokens, ntok, req_lens, nreq, req_pos, kv_caches, logits - ) diff --git a/scripts/libinfinicore_infer/jiuge_awq.py b/scripts/libinfinicore_infer/jiuge_awq.py deleted file mode 100644 index 9d5873953..000000000 --- a/scripts/libinfinicore_infer/jiuge_awq.py +++ /dev/null @@ -1,167 +0,0 @@ -from .base import BaseModel, DataType, DeviceType, KVCacheCStruct, register_model -from ctypes import ( - c_size_t, - c_uint, - c_int, - c_float, - c_void_p, - POINTER, - Structure, - c_char, - c_char_p, -) - - -class JiugeAWQMetaCStruct(Structure): - _fields_ = [ - ("dt_logits", DataType), - ("dt_linear_w", DataType), - ("dt_norm_w", DataType), - ("nlayer", c_size_t), - ("d", c_size_t), - ("nh", c_size_t), - ("nkvh", c_size_t), - ("dh", c_size_t), - ("di", c_size_t), - ("dctx", c_size_t), - ("dvoc", c_size_t), - ("epsilon", c_float), - ("theta", c_float), - ("end_token", c_uint), - ("nbit", c_size_t), - ("quant_group_size", c_size_t), - ("has_qkv_bias", c_char), - ] - - -class ModelWeightsCStruct(Structure): - pass - - -class JiugeAWQModelCStruct(Structure): - pass - - -@register_model -class JiugeAWQModel(BaseModel): - @classmethod - def register_lib(cls, lib): - """Register JiugeAWQ model functions with the library""" - lib.createJiugeAWQWeights.restype = POINTER(ModelWeightsCStruct) - lib.createJiugeAWQWeights.argtypes = [ - POINTER(JiugeAWQMetaCStruct), - DeviceType, - c_int, - POINTER(c_int), - ] - - lib.createJiugeAWQModel.restype = POINTER(JiugeAWQModelCStruct) - lib.createJiugeAWQModel.argtypes = [ - POINTER(JiugeAWQMetaCStruct), - POINTER(ModelWeightsCStruct), - ] - - lib.destroyJiugeAWQModel.argtypes = [POINTER(JiugeAWQModelCStruct)] - - lib.createKVCache.argtypes = [ - c_size_t, - c_size_t, - c_size_t, - c_size_t, - c_size_t, - DataType, - DeviceType, - POINTER(c_int), - c_size_t, - ] - lib.createKVCache.restype = POINTER(KVCacheCStruct) - - lib.dropKVCache.argtypes = [POINTER(KVCacheCStruct)] - - lib.inferBatchJiugeAWQ.argtypes = [ - POINTER(JiugeAWQModelCStruct), - POINTER(c_uint), - c_uint, - POINTER(c_uint), - c_uint, - POINTER(c_uint), - POINTER(POINTER(KVCacheCStruct)), - POINTER(c_float), - POINTER(c_uint), - POINTER(c_float), - POINTER(c_uint), - ] - - lib.forwardBatchJiugeAWQ.argtypes = [ - POINTER(JiugeAWQModelCStruct), - POINTER(c_uint), - c_uint, - POINTER(c_uint), - c_uint, - POINTER(c_uint), - POINTER(POINTER(KVCacheCStruct)), - c_void_p, - ] - - lib.JiugeAWQLoadWeight.argtypes = [ - POINTER(ModelWeightsCStruct), - c_char_p, - c_void_p, - ] - - def create_weights(self, meta, device_type, ndev, dev_ids): - return self.lib.createJiugeAWQWeights(meta, device_type, ndev, dev_ids) - - def create_model(self, meta, weights): - return self.lib.createJiugeAWQModel(meta, weights) - - def destroy_model(self, model): - self.lib.destroyJiugeAWQModel(model) - - def create_kv_cache( - self, nlayer, max_len, nkvh, dk, dv, dtype, device, dev_ids, ndev - ): - return self.lib.createKVCache( - nlayer, max_len, nkvh, dk, dv, dtype, device, dev_ids, ndev - ) - - def drop_kv_cache(self, kv_cache): - self.lib.dropKVCache(kv_cache) - - def load_weight(self, weights, name, data): - self.lib.JiugeAWQLoadWeight(weights, name.encode("utf-8"), data) - - def infer_batch( - self, - model, - tokens, - ntok, - req_lens, - nreq, - req_pos, - kv_caches, - temperature, - topk, - topp, - output, - ): - self.lib.inferBatchJiugeAWQ( - model, - tokens, - ntok, - req_lens, - nreq, - req_pos, - kv_caches, - temperature, - topk, - topp, - output, - ) - - def forward_batch( - self, model, tokens, ntok, req_lens, nreq, req_pos, kv_caches, logits - ): - self.lib.forwardBatchJiugeAWQ( - model, tokens, ntok, req_lens, nreq, req_pos, kv_caches, logits - ) diff --git a/scripts/libinfinicore_infer/jiuge_gptq.py b/scripts/libinfinicore_infer/jiuge_gptq.py deleted file mode 100644 index 59d9af451..000000000 --- a/scripts/libinfinicore_infer/jiuge_gptq.py +++ /dev/null @@ -1,167 +0,0 @@ -from .base import BaseModel, DataType, DeviceType, KVCacheCStruct, register_model -from ctypes import ( - c_size_t, - c_uint, - c_int, - c_float, - c_void_p, - POINTER, - Structure, - c_char, - c_char_p, -) - - -class JiugeGPTQMetaCStruct(Structure): - _fields_ = [ - ("dt_logits", DataType), - ("dt_linear_w", DataType), - ("dt_norm_w", DataType), - ("nlayer", c_size_t), - ("d", c_size_t), - ("nh", c_size_t), - ("nkvh", c_size_t), - ("dh", c_size_t), - ("di", c_size_t), - ("dctx", c_size_t), - ("dvoc", c_size_t), - ("epsilon", c_float), - ("theta", c_float), - ("end_token", c_uint), - ("nbit", c_size_t), - ("quant_group_size", c_size_t), - ("has_qkv_bias", c_char), - ] - - -class ModelWeightsCStruct(Structure): - pass - - -class JiugeGPTQModelCStruct(Structure): - pass - - -@register_model -class JiugeGPTQModel(BaseModel): - @classmethod - def register_lib(cls, lib): - """Register JiugeGPTQ model functions with the library""" - lib.createJiugeGPTQWeights.restype = POINTER(ModelWeightsCStruct) - lib.createJiugeGPTQWeights.argtypes = [ - POINTER(JiugeGPTQMetaCStruct), - DeviceType, - c_int, - POINTER(c_int), - ] - - lib.createJiugeGPTQModel.restype = POINTER(JiugeGPTQModelCStruct) - lib.createJiugeGPTQModel.argtypes = [ - POINTER(JiugeGPTQMetaCStruct), - POINTER(ModelWeightsCStruct), - ] - - lib.destroyJiugeGPTQModel.argtypes = [POINTER(JiugeGPTQModelCStruct)] - - lib.createKVCache.argtypes = [ - c_size_t, - c_size_t, - c_size_t, - c_size_t, - c_size_t, - DataType, - DeviceType, - POINTER(c_int), - c_size_t, - ] - lib.createKVCache.restype = POINTER(KVCacheCStruct) - - lib.dropKVCache.argtypes = [POINTER(KVCacheCStruct)] - - lib.inferBatchJiugeGPTQ.argtypes = [ - POINTER(JiugeGPTQModelCStruct), - POINTER(c_uint), - c_uint, - POINTER(c_uint), - c_uint, - POINTER(c_uint), - POINTER(POINTER(KVCacheCStruct)), - POINTER(c_float), - POINTER(c_uint), - POINTER(c_float), - POINTER(c_uint), - ] - - lib.forwardBatchJiugeGPTQ.argtypes = [ - POINTER(JiugeGPTQModelCStruct), - POINTER(c_uint), - c_uint, - POINTER(c_uint), - c_uint, - POINTER(c_uint), - POINTER(POINTER(KVCacheCStruct)), - c_void_p, - ] - - lib.JiugeGPTQLoadWeight.argtypes = [ - POINTER(ModelWeightsCStruct), - c_char_p, - c_void_p, - ] - - def create_weights(self, meta, device_type, ndev, dev_ids): - return self.lib.createJiugeGPTQWeights(meta, device_type, ndev, dev_ids) - - def create_model(self, meta, weights): - return self.lib.createJiugeGPTQModel(meta, weights) - - def destroy_model(self, model): - self.lib.destroyJiugeGPTQModel(model) - - def create_kv_cache( - self, nlayer, max_len, nkvh, dk, dv, dtype, device, dev_ids, ndev - ): - return self.lib.createKVCache( - nlayer, max_len, nkvh, dk, dv, dtype, device, dev_ids, ndev - ) - - def drop_kv_cache(self, kv_cache): - self.lib.dropKVCache(kv_cache) - - def load_weight(self, weights, name, data): - self.lib.JiugeGPTQLoadWeight(weights, name.encode("utf-8"), data) - - def infer_batch( - self, - model, - tokens, - ntok, - req_lens, - nreq, - req_pos, - kv_caches, - temperature, - topk, - topp, - output, - ): - self.lib.inferBatchJiugeGPTQ( - model, - tokens, - ntok, - req_lens, - nreq, - req_pos, - kv_caches, - temperature, - topk, - topp, - output, - ) - - def forward_batch( - self, model, tokens, ntok, req_lens, nreq, req_pos, kv_caches, logits - ): - self.lib.forwardBatchJiugeGPTQ( - model, tokens, ntok, req_lens, nreq, req_pos, kv_caches, logits - ) diff --git a/scripts/libinfinicore_infer/qwen3vl.py b/scripts/libinfinicore_infer/qwen3vl.py deleted file mode 100644 index bb3517d87..000000000 --- a/scripts/libinfinicore_infer/qwen3vl.py +++ /dev/null @@ -1,329 +0,0 @@ -from .base import BaseModel, DataType, DeviceType, KVCacheCStruct, register_model -from ctypes import ( - c_size_t, - c_uint, - c_uint16, - c_int, - c_float, - c_void_p, - c_bool, - POINTER, - Structure, - CFUNCTYPE, -) - - -class TextMetaCStruct(Structure): - _fields_ = [ - ("bos_token_id", c_size_t), - ("eos_token_id", c_size_t), - ("head_dim", c_size_t), - ("hidden_size", c_size_t), - ("initializer_range", c_float), - ("_pad1", c_float), - ("intermediate_size", c_size_t), - ("max_tokens", c_size_t), - ("num_attention_heads", c_size_t), - ("num_hidden_layers", c_size_t), - ("num_key_value_heads", c_size_t), - ("rms_norm_eps", c_float), - ("_pad2", c_float), - ("mrope_section", c_size_t * 3), - ("rope_theta", c_size_t), - ("vocab_size", c_size_t), - ] - - -class VisMetaCStruct(Structure): - _fields_ = [ - ("depth", c_size_t), - ("deepstack_visual_indexes", c_size_t * 3), - ("hidden_size", c_size_t), - ("in_channels", c_size_t), - ("initializer_range", c_float), - ("_pad1", c_float), - ("intermediate_size", c_size_t), - ("num_heads", c_size_t), - ("num_position_embeddings", c_size_t), - ("out_hidden_size", c_size_t), - ("patch_size", c_size_t), - ("spatial_merge_size", c_size_t), - ("temporal_patch_size", c_size_t), - ] - - -class Qwen3vlMetaCStruct(Structure): - _fields_ = [ - ("dtype", DataType), - ("_pad_dtype", c_uint), - ("text_meta", TextMetaCStruct), - ("vis_meta", VisMetaCStruct), - # Token ids - ("image_token_id", c_size_t), - ("video_token_id", c_size_t), - ("vision_end_token_id", c_size_t), - ("vision_start_token_id", c_size_t), - ] - - -class Qwen3vlWeightsCStruct(Structure): - pass - - -class Qwen3vlModelCStruct(Structure): - pass - - -class Qwen3vlCacheCStruct(Structure): - pass - - -load_global_fn = CFUNCTYPE(None, POINTER(Qwen3vlWeightsCStruct), c_void_p) -load_layer_fn = CFUNCTYPE(None, POINTER(Qwen3vlWeightsCStruct), c_void_p, c_size_t) - - -class Qwen3vlLangWeightLoaderCStruct(Structure): - _fields_ = [ - # Global - ("load_input_embd", load_global_fn), - ("load_output_norm", load_global_fn), - ("load_output_embd", load_global_fn), - # Attention - ("load_attn_norm", load_layer_fn), - ("load_attn_q_norm", load_layer_fn), - ("load_attn_k_norm", load_layer_fn), - ("load_attn_qkv_proj", load_layer_fn), - ("load_attn_o_proj", load_layer_fn), - # MLP - ("load_mlp_norm", load_layer_fn), - ("load_mlp_gate_up", load_layer_fn), - ("load_mlp_down", load_layer_fn), - ] - - -class Qwen3vlVisWeightLoaderCStruct(Structure): - _fields_ = [ - # Patch embed - ("load_patch_embed_weight", load_global_fn), - ("load_patch_embed_bias", load_global_fn), - ("load_pos_embed_weight", load_global_fn), - # Blocks attention - ("load_attn_proj_weight", load_layer_fn), - ("load_attn_proj_bias", load_layer_fn), - ("load_attn_qkv_weight", load_layer_fn), - ("load_attn_qkv_bias", load_layer_fn), - # Blocks MLP - ("load_mlp_linear_fc1_weight", load_layer_fn), - ("load_mlp_linear_fc1_bias", load_layer_fn), - ("load_mlp_linear_fc2_weight", load_layer_fn), - ("load_mlp_linear_fc2_bias", load_layer_fn), - # Blocks norm - ("load_norm1_weight", load_layer_fn), - ("load_norm1_bias", load_layer_fn), - ("load_norm2_weight", load_layer_fn), - ("load_norm2_bias", load_layer_fn), - # Deepstack merger - ("load_deepstack_merger_linear_fc1_weight", load_layer_fn), - ("load_deepstack_merger_linear_fc1_bias", load_layer_fn), - ("load_deepstack_merger_linear_fc2_weight", load_layer_fn), - ("load_deepstack_merger_linear_fc2_bias", load_layer_fn), - ("load_deepstack_merger_norm_weight", load_layer_fn), - ("load_deepstack_merger_norm_bias", load_layer_fn), - # Merger - ("load_merger_linear_fc1_weight", load_global_fn), - ("load_merger_linear_fc1_bias", load_global_fn), - ("load_merger_linear_fc2_weight", load_global_fn), - ("load_merger_linear_fc2_bias", load_global_fn), - ("load_merger_norm_weight", load_global_fn), - ("load_merger_norm_bias", load_global_fn), - ] - - -class Qwen3vlWeightLoaderCStruct(Structure): - _fields_ = [ - ("lang_loader", Qwen3vlLangWeightLoaderCStruct), - ("vis_loader", Qwen3vlVisWeightLoaderCStruct), - ] - - -@register_model -class Qwen3vlModel(BaseModel): - @classmethod - def register_lib(cls, lib): - """Register Qwen3vl model functions with the library""" - lib.createQwen3vlWeightLoader.argtypes = [] - lib.createQwen3vlWeightLoader.restype = POINTER(Qwen3vlWeightLoaderCStruct) - - lib.createQwen3vlWeights.argtypes = [ - POINTER(Qwen3vlMetaCStruct), - DeviceType, - c_int, - POINTER(c_int), - c_bool, - ] - lib.createQwen3vlWeights.restype = POINTER(Qwen3vlWeightsCStruct) - - lib.createQwen3vlModel.argtypes = [ - POINTER(Qwen3vlMetaCStruct), - POINTER(Qwen3vlWeightsCStruct), - ] - lib.createQwen3vlModel.restype = POINTER(Qwen3vlModelCStruct) - - lib.destroyQwen3vlModel.argtypes = [POINTER(Qwen3vlModelCStruct)] - - lib.createQwen3vlCache.argtypes = [POINTER(Qwen3vlModelCStruct)] - lib.createQwen3vlCache.restype = POINTER(Qwen3vlCacheCStruct) - - lib.dropQwen3vlCache.argtypes = [ - POINTER(Qwen3vlModelCStruct), - POINTER(Qwen3vlCacheCStruct), - ] - - lib.inferBatchQwen3vl.argtypes = [ - POINTER(Qwen3vlModelCStruct), - POINTER(c_uint), - c_uint, - c_void_p, # pixel_values, - c_uint, # total_patches, - POINTER(c_uint), # image_grid_thw, - c_uint, # num_images, - c_void_p, # pixel_values_videos, - c_uint, # total_patches_videos, - POINTER(c_uint), # video_grid_thw, - c_uint, # num_videos, - c_uint, # patch_features, - POINTER(c_uint), - c_uint, - POINTER(c_uint), - POINTER(POINTER(Qwen3vlCacheCStruct)), - POINTER(c_float), - POINTER(c_uint), - POINTER(c_float), - POINTER(c_uint), - ] - - lib.forwardBatchQwen3vl.argtypes = [ - POINTER(Qwen3vlModelCStruct), - POINTER(c_uint), - c_uint, - c_void_p, # pixel_values, - c_uint, # total_patches, - POINTER(c_uint), # image_grid_thw, - c_uint, # num_images, - c_void_p, # pixel_values_videos, - c_uint, # total_patches_videos, - POINTER(c_uint), # video_grid_thw, - c_uint, # num_videos, - c_uint, # patch_features, - POINTER(c_uint), - c_uint, - POINTER(c_uint), - POINTER(POINTER(Qwen3vlCacheCStruct)), - c_void_p, - ] - - def create_weight_loader(self): - return self.lib.createQwen3vlWeightLoader() - - def create_weights(self, meta, device_type, ndev, dev_ids, transpose_weight): - return self.lib.createQwen3vlWeights( - meta, device_type, ndev, dev_ids, transpose_weight - ) - - def create_model(self, meta, weights): - return self.lib.createQwen3vlModel(meta, weights) - - def destroy_model(self, model): - self.lib.destroyQwen3vlModel(model) - - def create_cache(self, model): - return self.lib.createQwen3vlCache(model) - - def drop_cache(self, model, cache): - self.lib.dropQwen3vlCache(model, cache) - - def infer_batch( - self, - model, - tokens, - ntok, - pixel_values, - total_patches, - image_grid_thw, - num_images, - pixel_values_videos, - total_patches_videos, - video_grid_thw, - num_videos, - patch_features, - req_lens, - nreq, - req_pos, - caches, - temperature, - topk, - topp, - output, - ): - self.lib.inferBatchQwen3vl( - model, - tokens, - ntok, - pixel_values, - total_patches, - image_grid_thw, - num_images, - pixel_values_videos, - total_patches_videos, - video_grid_thw, - num_videos, - patch_features, - req_lens, - nreq, - req_pos, - caches, - temperature, - topk, - topp, - output, - ) - - def forward_batch( - self, - model, - tokens, - ntok, - pixel_values, - total_patches, - image_grid_thw, - num_images, - pixel_values_videos, - total_patches_videos, - video_grid_thw, - num_videos, - patch_features, - req_lens, - nreq, - req_pos, - caches, - logits, - ): - self.lib.forwardBatchQwen3vl( - model, - tokens, - ntok, - pixel_values, - total_patches, - image_grid_thw, - num_images, - pixel_values_videos, - total_patches_videos, - video_grid_thw, - num_videos, - patch_features, - req_lens, - nreq, - req_pos, - caches, - logits, - ) diff --git a/scripts/qwen3vl.py b/scripts/qwen3vl.py deleted file mode 100644 index 917bce24e..000000000 --- a/scripts/qwen3vl.py +++ /dev/null @@ -1,935 +0,0 @@ -import ctypes -from typing import List, Sequence - -from tqdm import tqdm - -from libinfinicore_infer import ( - Qwen3vlModel, - Qwen3vlMetaCStruct, - TextMetaCStruct, - VisMetaCStruct, - Qwen3vlWeightsCStruct, - Qwen3vlCacheCStruct, - DataType, - DeviceType, -) -from infer_task import InferTask, KVCache - -from ctypes import POINTER, c_float, c_int, c_uint, c_uint16, c_void_p, byref, c_bool -import os -from pathlib import Path -import safetensors -import sys -import time -import json -import math -import torch -import transformers - -torch.set_default_device("cpu") - - -class Qwen3vlLangWeightsNaming: - def input_embd(self): - return "model.language_model.embed_tokens.weight" - - def output_embd(self): - return "model.language_model.embed_tokens.weight" - - def output_norm(self): - return "model.language_model.norm.weight" - - def attn_norm(self, i): - return f"model.language_model.layers.{i}.input_layernorm.weight" - - def attn_q_proj(self, i): - return f"model.language_model.layers.{i}.self_attn.q_proj.weight" - - def attn_q_norm(self, i): - return f"model.language_model.layers.{i}.self_attn.q_norm.weight" - - def attn_k_proj(self, i): - return f"model.language_model.layers.{i}.self_attn.k_proj.weight" - - def attn_k_norm(self, i): - return f"model.language_model.layers.{i}.self_attn.k_norm.weight" - - def attn_o_proj(self, i): - return f"model.language_model.layers.{i}.self_attn.o_proj.weight" - - def attn_v_proj(self, i): - return f"model.language_model.layers.{i}.self_attn.v_proj.weight" - - def mlp_norm(self, i): - return f"model.language_model.layers.{i}.post_attention_layernorm.weight" - - def mlp_gate(self, i): - return f"model.language_model.layers.{i}.mlp.gate_proj.weight" - - def mlp_down(self, i): - return f"model.language_model.layers.{i}.mlp.down_proj.weight" - - def mlp_up(self, i): - return f"model.language_model.layers.{i}.mlp.up_proj.weight" - - -class Qwen3vlVisWeightsNaming: - def patch_embed_weight(self): - return "model.visual.patch_embed.proj.weight" - - def patch_embed_bias(self): - return "model.visual.patch_embed.proj.bias" - - def pos_embed_weight(self): - return "model.visual.pos_embed.weight" - - def attn_proj_weight(self, i): - return f"model.visual.blocks.{i}.attn.proj.weight" - - def attn_proj_bias(self, i): - return f"model.visual.blocks.{i}.attn.proj.bias" - - def attn_qkv_weight(self, i): - return f"model.visual.blocks.{i}.attn.qkv.weight" - - def attn_qkv_bias(self, i): - return f"model.visual.blocks.{i}.attn.qkv.bias" - - def mlp_linear_fc1_weight(self, i): - return f"model.visual.blocks.{i}.mlp.linear_fc1.weight" - - def mlp_linear_fc1_bias(self, i): - return f"model.visual.blocks.{i}.mlp.linear_fc1.bias" - - def mlp_linear_fc2_weight(self, i): - return f"model.visual.blocks.{i}.mlp.linear_fc2.weight" - - def mlp_linear_fc2_bias(self, i): - return f"model.visual.blocks.{i}.mlp.linear_fc2.bias" - - def norm1_weight(self, i): - return f"model.visual.blocks.{i}.norm1.weight" - - def norm1_bias(self, i): - return f"model.visual.blocks.{i}.norm1.bias" - - def norm2_weight(self, i): - return f"model.visual.blocks.{i}.norm2.weight" - - def norm2_bias(self, i): - return f"model.visual.blocks.{i}.norm2.bias" - - def deepstack_merger_linear_fc1_weight(self, i): - return f"model.visual.deepstack_merger_list.{i}.linear_fc1.weight" - - def deepstack_merger_linear_fc1_bias(self, i): - return f"model.visual.deepstack_merger_list.{i}.linear_fc1.bias" - - def deepstack_merger_linear_fc2_weight(self, i): - return f"model.visual.deepstack_merger_list.{i}.linear_fc2.weight" - - def deepstack_merger_linear_fc2_bias(self, i): - return f"model.visual.deepstack_merger_list.{i}.linear_fc2.bias" - - def deepstack_merger_norm_weight(self, i): - return f"model.visual.deepstack_merger_list.{i}.norm.weight" - - def deepstack_merger_norm_bias(self, i): - return f"model.visual.deepstack_merger_list.{i}.norm.bias" - - def merger_linear_fc1_weight(self): - return "model.visual.merger.linear_fc1.weight" - - def merger_linear_fc1_bias(self): - return "model.visual.merger.linear_fc1.bias" - - def merger_linear_fc2_weight(self): - return "model.visual.merger.linear_fc2.weight" - - def merger_linear_fc2_bias(self): - return "model.visual.merger.linear_fc2.bias" - - def merger_norm_weight(self): - return "model.visual.merger.norm.weight" - - def merger_norm_bias(self): - return "model.visual.merger.norm.bias" - - -class Qwen3vlMeta(Qwen3vlMetaCStruct): - def __init__(self, config, max_tokens=None): - if config["text_config"]["dtype"] == "float16": - dt_ = DataType.INFINI_DTYPE_F16 - self.torch_dtype = torch.float16 - elif config["text_config"]["dtype"] == "float32": - dt_ = DataType.INFINI_DTYPE_F32 - self.torch_dtype = torch.float32 - elif config["text_config"]["dtype"] == "bfloat16": - dt_ = DataType.INFINI_DTYPE_BF16 - self.torch_dtype = torch.bfloat16 - else: - raise ValueError( - f"Unsupported text dtype: {config['text_config']['dtype']}" - ) - - super().__init__( - dtype=dt_, - image_token_id=config["image_token_id"], - video_token_id=config["video_token_id"], - vision_end_token_id=config["vision_end_token_id"], - vision_start_token_id=config["vision_start_token_id"], - text_meta=TextMetaCStruct( - bos_token_id=config["text_config"]["bos_token_id"], - eos_token_id=config["text_config"]["eos_token_id"], - head_dim=config["text_config"]["head_dim"], - hidden_size=config["text_config"]["hidden_size"], - initializer_range=config["text_config"]["initializer_range"], - intermediate_size=config["text_config"]["intermediate_size"], - max_tokens=( - config["text_config"]["max_position_embeddings"] - if max_tokens is None - else max_tokens - ), - num_attention_heads=config["text_config"]["num_attention_heads"], - num_hidden_layers=config["text_config"]["num_hidden_layers"], - num_key_value_heads=config["text_config"]["num_key_value_heads"], - rms_norm_eps=config["text_config"]["rms_norm_eps"], - mrope_section=(ctypes.c_ulong * 3)( - *config["text_config"]["rope_scaling"]["mrope_section"] - ), - rope_theta=config["text_config"]["rope_theta"], - vocab_size=config["text_config"]["vocab_size"], - ), - vis_meta=VisMetaCStruct( - depth=config["vision_config"]["depth"], - deepstack_visual_indexes=(ctypes.c_ulong * 3)( - *config["vision_config"]["deepstack_visual_indexes"] - ), - hidden_size=config["vision_config"]["hidden_size"], - in_channels=config["vision_config"]["in_channels"], - initializer_range=config["vision_config"]["initializer_range"], - intermediate_size=config["vision_config"]["intermediate_size"], - num_heads=config["vision_config"]["num_heads"], - num_position_embeddings=config["vision_config"][ - "num_position_embeddings" - ], - out_hidden_size=config["vision_config"]["out_hidden_size"], - patch_size=config["vision_config"]["patch_size"], - spatial_merge_size=config["vision_config"]["spatial_merge_size"], - temporal_patch_size=config["vision_config"]["temporal_patch_size"], - ), - ) - - -def load_specific_tensor(model_dir, tensor_name): - """ - Load a specific tensor from a safetensors model. - Supports both sharded models (with index.json) and single file models. - """ - - # Try to load from individual .safetensors files - safetensors_files = [f for f in os.listdir(model_dir) if f.endswith(".safetensors")] - if not safetensors_files: - raise FileNotFoundError(f"No .safetensors files found in {model_dir}") - - # Try to find the tensor in each file - for filename in safetensors_files: - tensor_file = os.path.join(model_dir, filename) - try: - with safetensors.safe_open(tensor_file, framework="pt", device="cpu") as f: - if tensor_name in f.keys(): - tensor = f.get_tensor(tensor_name) - return tensor - except Exception: - continue - - # If we reach here, tensor was not found in any file - raise KeyError(f"{tensor_name} not found in any .safetensors files") - - -def load_Qwen3vl_weights( - meta: Qwen3vlMeta, - weights, - model_path: str, - ndev: int, -): - # torch load weights, and reshape for qkv_proj / mlp_gate_up stack, attn / mlp parallel - # weight loader function load from specific offset according to idev, and transpose - model_instance = Qwen3vlModel() - weight_loader = model_instance.create_weight_loader() - vis_names = Qwen3vlVisWeightsNaming() - lang_names = Qwen3vlLangWeightsNaming() - - nkvh = meta.text_meta.num_key_value_heads - nh = meta.text_meta.num_attention_heads - dh = meta.text_meta.head_dim - d = meta.text_meta.hidden_size - di = meta.text_meta.intermediate_size - - assert nh % nkvh == 0 - assert nh % ndev == 0 - assert nkvh % ndev == 0 - assert di % ndev == 0 - - # ------------------------------- - # Language_model weights - # ------------------------------- - input_embd = load_specific_tensor(model_path, lang_names.input_embd()).to( - meta.torch_dtype - ) - weight_loader.contents.lang_loader.load_input_embd(weights, input_embd.data_ptr()) - del input_embd - - output_norm = load_specific_tensor(model_path, lang_names.output_norm()).to( - meta.torch_dtype - ) - weight_loader.contents.lang_loader.load_output_norm(weights, output_norm.data_ptr()) - del output_norm - - output_embd = load_specific_tensor(model_path, lang_names.output_embd()).to( - meta.torch_dtype - ) - weight_loader.contents.lang_loader.load_output_embd(weights, output_embd.data_ptr()) - del output_embd - - for i in range(meta.text_meta.num_hidden_layers): - attn_norm = load_specific_tensor(model_path, lang_names.attn_norm(i)).to( - meta.torch_dtype - ) - weight_loader.contents.lang_loader.load_attn_norm( - weights, attn_norm.data_ptr(), i - ) - del attn_norm - - attn_q_proj = load_specific_tensor(model_path, lang_names.attn_q_proj(i)) - attn_k_proj = load_specific_tensor(model_path, lang_names.attn_k_proj(i)) - attn_v_proj = load_specific_tensor(model_path, lang_names.attn_v_proj(i)) - - _Q = attn_q_proj.reshape(nh, dh, d) - _K = attn_k_proj.reshape(nkvh, dh, d) - _V = attn_v_proj.reshape(nkvh, dh, d) - - qkv_proj = [] - _nh = nh // ndev - _nkvh = nkvh // ndev - for _idev in range(ndev): - qkv_proj.append(_Q[_idev * _nh : (_idev + 1) * _nh, :, :]) - qkv_proj.append(_K[_idev * _nkvh : (_idev + 1) * _nkvh, :, :]) - qkv_proj.append(_V[_idev * _nkvh : (_idev + 1) * _nkvh, :, :]) - attn_qkv_proj = torch.cat(qkv_proj, dim=0).to(meta.torch_dtype).contiguous() - - weight_loader.contents.lang_loader.load_attn_qkv_proj( - weights, attn_qkv_proj.data_ptr(), i - ) - del attn_qkv_proj - - attn_q_norm = load_specific_tensor(model_path, lang_names.attn_q_norm(i)).to( - meta.torch_dtype - ) - weight_loader.contents.lang_loader.load_attn_q_norm( - weights, attn_q_norm.data_ptr(), i - ) - del attn_q_norm - - attn_k_norm = load_specific_tensor(model_path, lang_names.attn_k_norm(i)).to( - meta.torch_dtype - ) - weight_loader.contents.lang_loader.load_attn_k_norm( - weights, attn_k_norm.data_ptr(), i - ) - del attn_k_norm - - attn_o_proj = load_specific_tensor(model_path, lang_names.attn_o_proj(i)) - attn_o_proj = ( - attn_o_proj.to(meta.torch_dtype) - .reshape([d, ndev, nh // ndev * dh]) - .transpose(0, 1) - .contiguous() - ) - weight_loader.contents.lang_loader.load_attn_o_proj( - weights, attn_o_proj.data_ptr(), i - ) - del attn_o_proj - - mlp_norm = load_specific_tensor(model_path, lang_names.mlp_norm(i)).to( - meta.torch_dtype - ) - weight_loader.contents.lang_loader.load_mlp_norm( - weights, mlp_norm.data_ptr(), i - ) - del mlp_norm - - mlp_gate = load_specific_tensor(model_path, lang_names.mlp_gate(i)) - mlp_up = load_specific_tensor(model_path, lang_names.mlp_up(i)) - - gate_up = [] - _di = di // ndev - for _idev in range(ndev): - _start = _idev * _di - _end = (_idev + 1) * _di - gate_up.append(mlp_gate[_start:_end, :]) - gate_up.append(mlp_up[_start:_end, :]) - mlp_gate_up = torch.cat(gate_up, dim=0).to(meta.torch_dtype).contiguous() - - weight_loader.contents.lang_loader.load_mlp_gate_up( - weights, mlp_gate_up.data_ptr(), i - ) - del mlp_gate_up - - mlp_down = load_specific_tensor(model_path, lang_names.mlp_down(i)) - mlp_down = ( - mlp_down.to(meta.torch_dtype) - .reshape([d, ndev, di // ndev]) - .transpose(0, 1) - .contiguous() - ) - weight_loader.contents.lang_loader.load_mlp_down( - weights, mlp_down.data_ptr(), i - ) - del mlp_down - - # ------------------------------- - # Vision head weights - # ------------------------------- - patch_embed_weight = load_specific_tensor( - model_path, vis_names.patch_embed_weight() - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_patch_embed_weight( - weights, patch_embed_weight.data_ptr() - ) - del patch_embed_weight - - patch_embed_bias = load_specific_tensor( - model_path, vis_names.patch_embed_bias() - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_patch_embed_bias( - weights, patch_embed_bias.data_ptr() - ) - del patch_embed_bias - - pos_embed_weight = load_specific_tensor( - model_path, vis_names.pos_embed_weight() - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_pos_embed_weight( - weights, pos_embed_weight.data_ptr() - ) - del pos_embed_weight - - for i in range(meta.vis_meta.depth): - attn_proj_weight = load_specific_tensor( - model_path, vis_names.attn_proj_weight(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_attn_proj_weight( - weights, attn_proj_weight.data_ptr(), i - ) - del attn_proj_weight - - attn_proj_bias = load_specific_tensor( - model_path, vis_names.attn_proj_bias(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_attn_proj_bias( - weights, attn_proj_bias.data_ptr(), i - ) - del attn_proj_bias - - attn_qkv_weight = load_specific_tensor( - model_path, vis_names.attn_qkv_weight(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_attn_qkv_weight( - weights, attn_qkv_weight.data_ptr(), i - ) - del attn_qkv_weight - - attn_qkv_bias = load_specific_tensor(model_path, vis_names.attn_qkv_bias(i)).to( - meta.torch_dtype - ) - weight_loader.contents.vis_loader.load_attn_qkv_bias( - weights, attn_qkv_bias.data_ptr(), i - ) - del attn_qkv_bias - - mlp_linear_fc1_weight = load_specific_tensor( - model_path, vis_names.mlp_linear_fc1_weight(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_mlp_linear_fc1_weight( - weights, mlp_linear_fc1_weight.data_ptr(), i - ) - del mlp_linear_fc1_weight - - mlp_linear_fc1_bias = load_specific_tensor( - model_path, vis_names.mlp_linear_fc1_bias(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_mlp_linear_fc1_bias( - weights, mlp_linear_fc1_bias.data_ptr(), i - ) - del mlp_linear_fc1_bias - - mlp_linear_fc2_weight = load_specific_tensor( - model_path, vis_names.mlp_linear_fc2_weight(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_mlp_linear_fc2_weight( - weights, mlp_linear_fc2_weight.data_ptr(), i - ) - del mlp_linear_fc2_weight - - mlp_linear_fc2_bias = load_specific_tensor( - model_path, vis_names.mlp_linear_fc2_bias(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_mlp_linear_fc2_bias( - weights, mlp_linear_fc2_bias.data_ptr(), i - ) - del mlp_linear_fc2_bias - - norm1_weight = load_specific_tensor(model_path, vis_names.norm1_weight(i)).to( - meta.torch_dtype - ) - weight_loader.contents.vis_loader.load_norm1_weight( - weights, norm1_weight.data_ptr(), i - ) - del norm1_weight - - norm1_bias = load_specific_tensor(model_path, vis_names.norm1_bias(i)).to( - meta.torch_dtype - ) - weight_loader.contents.vis_loader.load_norm1_bias( - weights, norm1_bias.data_ptr(), i - ) - del norm1_bias - - norm2_weight = load_specific_tensor(model_path, vis_names.norm2_weight(i)).to( - meta.torch_dtype - ) - weight_loader.contents.vis_loader.load_norm2_weight( - weights, norm2_weight.data_ptr(), i - ) - del norm2_weight - - norm2_bias = load_specific_tensor(model_path, vis_names.norm2_bias(i)).to( - meta.torch_dtype - ) - weight_loader.contents.vis_loader.load_norm2_bias( - weights, norm2_bias.data_ptr(), i - ) - del norm2_bias - - for i in range(len(meta.vis_meta.deepstack_visual_indexes)): - deepstack_merger_linear_fc1_weight = load_specific_tensor( - model_path, vis_names.deepstack_merger_linear_fc1_weight(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_deepstack_merger_linear_fc1_weight( - weights, deepstack_merger_linear_fc1_weight.data_ptr(), i - ) - del deepstack_merger_linear_fc1_weight - - deepstack_merger_linear_fc1_bias = load_specific_tensor( - model_path, vis_names.deepstack_merger_linear_fc1_bias(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_deepstack_merger_linear_fc1_bias( - weights, deepstack_merger_linear_fc1_bias.data_ptr(), i - ) - del deepstack_merger_linear_fc1_bias - - deepstack_merger_linear_fc2_weight = load_specific_tensor( - model_path, vis_names.deepstack_merger_linear_fc2_weight(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_deepstack_merger_linear_fc2_weight( - weights, deepstack_merger_linear_fc2_weight.data_ptr(), i - ) - del deepstack_merger_linear_fc2_weight - - deepstack_merger_linear_fc2_bias = load_specific_tensor( - model_path, vis_names.deepstack_merger_linear_fc2_bias(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_deepstack_merger_linear_fc2_bias( - weights, deepstack_merger_linear_fc2_bias.data_ptr(), i - ) - del deepstack_merger_linear_fc2_bias - - deepstack_merger_norm_weight = load_specific_tensor( - model_path, vis_names.deepstack_merger_norm_weight(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_deepstack_merger_norm_weight( - weights, deepstack_merger_norm_weight.data_ptr(), i - ) - del deepstack_merger_norm_weight - - deepstack_merger_norm_bias = load_specific_tensor( - model_path, vis_names.deepstack_merger_norm_bias(i) - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_deepstack_merger_norm_bias( - weights, deepstack_merger_norm_bias.data_ptr(), i - ) - del deepstack_merger_norm_bias - - merger_linear_fc1_weight = load_specific_tensor( - model_path, vis_names.merger_linear_fc1_weight() - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_merger_linear_fc1_weight( - weights, merger_linear_fc1_weight.data_ptr() - ) - del merger_linear_fc1_weight - - merger_linear_fc1_bias = load_specific_tensor( - model_path, vis_names.merger_linear_fc1_bias() - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_merger_linear_fc1_bias( - weights, merger_linear_fc1_bias.data_ptr() - ) - del merger_linear_fc1_bias - - merger_linear_fc2_weight = load_specific_tensor( - model_path, vis_names.merger_linear_fc2_weight() - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_merger_linear_fc2_weight( - weights, merger_linear_fc2_weight.data_ptr() - ) - del merger_linear_fc2_weight - - merger_linear_fc2_bias = load_specific_tensor( - model_path, vis_names.merger_linear_fc2_bias() - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_merger_linear_fc2_bias( - weights, merger_linear_fc2_bias.data_ptr() - ) - del merger_linear_fc2_bias - - merger_norm_weight = load_specific_tensor( - model_path, vis_names.merger_norm_weight() - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_merger_norm_weight( - weights, merger_norm_weight.data_ptr() - ) - del merger_norm_weight - - merger_norm_bias = load_specific_tensor( - model_path, vis_names.merger_norm_bias() - ).to(meta.torch_dtype) - weight_loader.contents.vis_loader.load_merger_norm_bias( - weights, merger_norm_bias.data_ptr() - ) - del merger_norm_bias - - -class Qwen3vlBatchedTask: - def __init__( - self, - tasks: List[InferTask], - all_pixel_values=None, - all_image_grid_thw=None, - all_pixel_values_videos=None, - all_video_grid_thw=None, - ): - self.tasks = tasks - self.nreq = len(tasks) - - # Precompute fields - token_lists = [t.tokens for t in tasks] - self.req_lens_list = [len(toks) for toks in token_lists] - self.req_pos_list = [t.pos for t in tasks] - self.kv_cache_ptrs = [t.kvcache().data() for t in tasks] - self.temperaturas_list = [t.temperature for t in tasks] - self.topks_list = [t.topk for t in tasks] - self.topps_list = [t.topp for t in tasks] - - # Flatten token lists - flat_tokens = [tok for toks in token_lists for tok in toks] - self.ntok = len(flat_tokens) - - # Convert to ctypes arrays in one pass - self.tokens = (c_uint * self.ntok)(*flat_tokens) - self.req_lens = (c_uint * self.nreq)(*self.req_lens_list) - self.req_pos = (c_uint * self.nreq)(*self.req_pos_list) - self.kv_caches = (POINTER(Qwen3vlCacheCStruct) * self.nreq)(*self.kv_cache_ptrs) - self.temperaturas = (c_float * self.nreq)(*self.temperaturas_list) - self.topks = (c_uint * self.nreq)(*self.topks_list) - self.topps = (c_float * self.nreq)(*self.topps_list) - - # initialize visual encoder inputs - self.pixel_values = None - self.total_patches = 0 - self.image_grid_thw = None - self.num_images = 0 - self.pixel_values_videos = None - self.total_patches_videos = 0 - self.video_grid_thw = None - self.num_videos = 0 - self.patch_features = 0 - - # Prepare visual encoder inputs - # all_pixel_values = [t.inputs['pixel_values'] for t in tasks if 'pixel_values' in t.inputs] - # all_image_grid_thw = [t.inputs['image_grid_thw'] for t in tasks if 'image_grid_thw' in t.inputs] - # all_pixel_values_videos = [t.inputs['pixel_values_videos'] for t in tasks if 'pixel_values_videos' in t.inputs] - # all_video_grid_thw = [t.inputs['video_grid_thw'] for t in tasks if 'video_grid_thw' in t.inputs] - - if all_pixel_values is not None: - print(all_pixel_values.shape) - concat_pixel_values = ( - torch.cat(all_pixel_values, dim=0) - if isinstance(all_pixel_values, list) - else all_pixel_values - ) # (total_patches, features) - self.total_patches = concat_pixel_values.shape[0] - self.patch_features = concat_pixel_values.shape[1] - self.flat_pixels = ( - concat_pixel_values.flatten().to(torch.bfloat16).contiguous() - ) - self.pixel_values = self.flat_pixels.data_ptr() - - if all_image_grid_thw is not None: - concat_grid_thw = ( - torch.cat(all_image_grid_thw, dim=0) - if isinstance(all_image_grid_thw, list) - else all_image_grid_thw - ) # (total_images, 3) - self.num_images = concat_grid_thw.shape[0] - self.flat_grid = ( - concat_grid_thw.flatten().to(torch.int32).contiguous().tolist() - ) - self.image_grid_thw = (c_uint * len(self.flat_grid))(*self.flat_grid) - - if all_pixel_values_videos is not None: - concat_pixel_values_videos = torch.cat( - all_pixel_values_videos, dim=0 - ) # (total_patches_videos, features) - self.total_patches_videos = concat_pixel_values_videos.shape[0] - self.patch_features_videos = concat_pixel_values_videos.shape[1] - print(self.patch_features_videos, flush=True) - self.flat_pixels_videos = ( - concat_pixel_values_videos.flatten().to(torch.bfloat16).contiguous() - ) - self.pixel_values_videos = self.flat_pixels_videos.ctypes.data_as(c_void_p) - - if all_video_grid_thw is not None: - concat_grid_thw_videos = torch.cat( - all_video_grid_thw, dim=0 - ) # (total_videos, 3) - self.num_videos = concat_grid_thw_videos.shape[0] - flat_grid_videos = ( - concat_grid_thw_videos.flatten().to(torch.int32).contiguous() - ) - self.video_grid_thw = (c_uint * len(flat_grid_videos))( - *flat_grid_videos.tolist() - ) - - def input_args(self): - return ( - self.tokens, - self.ntok, - self.pixel_values, - self.total_patches, - self.image_grid_thw, - self.num_images, - self.pixel_values_videos, - self.total_patches_videos, - self.video_grid_thw, - self.num_videos, - self.patch_features, - self.req_lens, - self.nreq, - self.req_pos, - self.kv_caches, - self.temperaturas, - self.topks, - self.topps, - ) - - -# 需要处理 visual encoder的cache 和 image video输入 -class Qwen3vlForCauslLM: - def __init__( - self, model_dir_path, device=DeviceType.DEVICE_TYPE_CPU, ndev=1, max_tokens=None - ): - with open(os.path.join(model_dir_path, "config.json"), "r") as f: - config = json.load(f) - self.config = config - eos_token_id = self.config["text_config"]["eos_token_id"] - self.eos_token_id = ( - [eos_token_id] if type(eos_token_id) == int else eos_token_id - ) - - print(model_dir_path) - - if "qwen3_vl" == config["model_type"]: - self.meta = Qwen3vlMeta(config, max_tokens=max_tokens) - self.processor = transformers.AutoProcessor.from_pretrained(model_dir_path) - self.tokenizer = transformers.AutoTokenizer.from_pretrained(model_dir_path) - else: - raise ValueError("Unsupported model architecture") - - print(f"Creating model on {ndev} devices...") - load_start_time = time.time() - dev_ids = (c_int * ndev)(*[i for i in range(ndev)]) - - self.model_instance = Qwen3vlModel() - weights = self.model_instance.create_weights( - byref(self.meta), device, ndev, dev_ids, c_bool(True) - ) - print("Loading weights...") - # Load weights from host - load_Qwen3vl_weights(self.meta, weights, model_dir_path, ndev) - # Create model instance - self.model_ptr = self.model_instance.create_model( - byref(self.meta), - weights, - ) - load_end_time = time.time() - print(f"Time used: {load_end_time - load_start_time:.3f}s") - - def max_context_len(self): - return self.meta.text_meta.max_tokens - - def create_kv_cache(self): - return self.model_instance.create_cache(self.model_ptr) - - def drop_kv_cache(self, kv_cache): - self.model_instance.drop_cache(self.model_ptr, kv_cache) - - def batch_infer_one_round( - self, - tasks: List[InferTask], - all_pixel_values=None, - all_image_grid_thw=None, - all_pixel_values_videos=None, - all_video_grid_thw=None, - ): - output = (c_uint * len(tasks))() - batch_inputs = Qwen3vlBatchedTask( - tasks, - all_pixel_values, - all_image_grid_thw, - all_pixel_values_videos, - all_video_grid_thw, - ) - self.model_instance.infer_batch( - self.model_ptr, - *(batch_inputs.input_args()), - output, - ) - return list(output) - - def generate( - self, input_content, max_steps=0, topp_=1.0, topk_=1, temperature_=1.0 - ): - inputs = self.processor.apply_chat_template( - conversation=[{"role": "user", "content": input_content}], - tokenize=True, - add_generation_prompt=True, - return_dict=True, - return_tensors="pt", - ) - tokens = inputs["input_ids"][0].tolist() - pixel_values = inputs["pixel_values"] if "pixel_values" in inputs else None - image_grid_thw = ( - inputs["image_grid_thw"] if "image_grid_thw" in inputs else None - ) - pixel_values_videos = ( - inputs["pixel_values_videos"] if "pixel_values_videos" in inputs else None - ) - video_grid_thw = ( - inputs["video_grid_thw"] if "video_grid_thw" in inputs else None - ) - - infer_task = InferTask( - 0, - tokens, - self.max_context_len(), - temperature_, - topk_, - topp_, - self.eos_token_id, - ) - infer_task.bind_kvcache(KVCache(self)) - print(input_content) - steps = 0 - total_time = 0 - output_content = "" - - # print(inputs['input_ids'][0].tolist(), flush=True) - - for step_i in range(max_steps if max_steps > 0 else self.max_context_len()): - start_time = time.time() - output_tokens = self.batch_infer_one_round( - [infer_task], - pixel_values, - image_grid_thw, - pixel_values_videos, - video_grid_thw, - ) - # print(output_tokens) - end_time = time.time() - steps += 1 - output_str = self.tokenizer.decode(output_tokens[0]) - output_content += output_str - print(output_str, end="", flush=True) - pixel_values = None - image_grid_thw = None - pixel_values_videos = None - video_grid_thw = None - if output_tokens[0] in self.eos_token_id: - break - infer_task.next(output_tokens[0]) - - if step_i > 0: - total_time += end_time - start_time - - print("\n") - avg_time = total_time * 1000 / steps if steps > 0 else -1 - # print(output_content, flush=True) - print(f"Time per step: {avg_time:.3f}ms") - - infer_task._kv_cache.drop(self) - return output_content, avg_time - - def destroy_model_instance(self): - self.model_instance.destroy_model(self.model_ptr) - print("Model destroyed") - - -def test(): - if len(sys.argv) < 3: - print( - "Usage: python qwen3vl.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] [n_device]" - ) - sys.exit(1) - model_path = sys.argv[2] - device_type = DeviceType.DEVICE_TYPE_CPU - if sys.argv[1] == "--cpu": - device_type = DeviceType.DEVICE_TYPE_CPU - elif sys.argv[1] == "--nvidia": - device_type = DeviceType.DEVICE_TYPE_NVIDIA - elif sys.argv[1] == "--cambricon": - device_type = DeviceType.DEVICE_TYPE_CAMBRICON - elif sys.argv[1] == "--ascend": - device_type = DeviceType.DEVICE_TYPE_ASCEND - elif sys.argv[1] == "--metax": - device_type = DeviceType.DEVICE_TYPE_METAX - elif sys.argv[1] == "--moore": - device_type = DeviceType.DEVICE_TYPE_MOORE - elif sys.argv[1] == "--iluvatar": - device_type = DeviceType.DEVICE_TYPE_ILUVATAR - else: - print( - "Usage: python qwen3vl.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore] [n_device]" - ) - sys.exit(1) - ndev = int(sys.argv[3]) if len(sys.argv) > 3 else 1 - - img_url = None - if len(sys.argv) > 4: - img_url = sys.argv[4] - - model = Qwen3vlForCauslLM(model_path, device_type, ndev, max_tokens=1024) - input_content = ( - [ - {"type": "text", "text": "Describe this image."}, - {"type": "image", "url": img_url}, - ] - if img_url is not None - else [{"type": "text", "text": "山东最高的山是？"}] - ) - model.generate(input_content) - model.destroy_model_instance() - - -if __name__ == "__main__": - test() diff --git a/scripts/test_ceval.py b/scripts/test_ceval.py deleted file mode 100644 index 0b05dacf9..000000000 --- a/scripts/test_ceval.py +++ /dev/null @@ -1,205 +0,0 @@ -import sys -from jiuge import * -from datasets import load_dataset - -# Support AWQ and GPTQ -from libinfinicore_infer import ( - JiugeAWQModel, - JiugeAWQMetaCStruct, - JiugeGPTQModel, - JiugeGPTQMetaCStruct, - DeviceType, - KVCacheCStruct, -) - -# Import missing classes -from jiuge import KVCache, InferTask - -class JiugeForCeval(JiugeForCauslLM): - def __init__( - self, model_dir_path, device=DeviceType.DEVICE_TYPE_CPU, ndev=1, max_tokens=None - ): - super().__init__(model_dir_path, device, ndev, max_tokens) - pass - - def generate(self, conversation, max_steps, topp_=1.0, topk_=1, temperature_=1.0): - input_content = ( - self.tokenizer.apply_chat_template( - conversation=conversation, - add_generation_prompt=True, - tokenize=False, - ) - + "正确答案是" - ) - - print(input_content, end="", flush=True) - - tokens = self.tokenizer.encode(input_content) - infer_task = InferTask( - 0, - tokens, - self.max_context_len(), - temperature_, - topk_, - topp_, - self.eos_token_id, - ) - infer_task.bind_kvcache(KVCache(self)) - - steps = 0 - total_time = 0 - output_content = "" - - for step_i in range(max_steps): - start_time = time.time() - output_tokens = self.batch_infer_one_round([infer_task]) - end_time = time.time() - steps += 1 - output_str = self.tokenizer.decode(output_tokens[0]) - output_content += output_str - print(output_str, end="", flush=True) - if output_tokens[0] in self.eos_token_id: - break - infer_task.next(output_tokens[0]) - - if step_i > 0: - total_time += end_time - start_time - - print("\n") - avg_time = total_time * 1000 / (steps - 1 + 1e-9) - print(f"Time per step: {avg_time:.3f}ms") - - infer_task._kv_cache.drop(self) - return output_content, avg_time - - -def test(): - if len(sys.argv) < 3: - print( - "Usage: python test_ceval.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore | --iluvatar | --kunlun | --hygon] [n_device]" - ) - sys.exit(1) - - # ----------------------------- - # 支持 AWQ/GPTQ 参数 - # ----------------------------- - use_awq = False - use_gptq = False - device_arg = None - model_path = None - ndev = 1 - - # 解析 sys.argv - for arg in sys.argv[1:]: - if arg.lower() == "--awq": - use_awq = True - elif arg.lower() == "--gptq": - use_gptq = True - elif arg.startswith("--"): - device_arg = arg.lower() - else: - if model_path is None: - model_path = arg - else: - ndev = int(arg) - - # ----------------------------- - # 设置设备类型 - # ----------------------------- - device_type = DeviceType.DEVICE_TYPE_CPU - if device_arg == "--cpu": - device_type = DeviceType.DEVICE_TYPE_CPU - elif device_arg == "--nvidia": - device_type = DeviceType.DEVICE_TYPE_NVIDIA - elif device_arg == "--cambricon": - device_type = DeviceType.DEVICE_TYPE_CAMBRICON - elif device_arg == "--ascend": - device_type = DeviceType.DEVICE_TYPE_ASCEND - elif device_arg == "--metax": - device_type = DeviceType.DEVICE_TYPE_METAX - elif device_arg == "--moore": - device_type = DeviceType.DEVICE_TYPE_MOORE - elif device_arg == "--iluvatar": - device_type = DeviceType.DEVICE_TYPE_ILUVATAR - elif device_arg == "--kunlun": - device_type = DeviceType.DEVICE_TYPE_KUNLUN - elif device_arg == "--hygon": - device_type = DeviceType.DEVICE_TYPE_HYGON - else: - print( - "Usage: python test_ceval.py [--cpu | --nvidia| --cambricon | --ascend | --metax | --moore | --iluvatar | --kunlun | --hygon] [n_device]" - ) - sys.exit(1) - - # ----------------------------- - # 加载 CEval 数据集 - # ----------------------------- - dataset = load_dataset(r"ceval/ceval-exam", name="middle_school_mathematics") - # dataset = load_dataset(r"ceval/ceval-exam", name="high_school_history") - # dataset = load_dataset(r"ceval/ceval-exam", name="high_school_chinese") - # dataset = load_dataset(r"ceval/ceval-exam", name="high_school_physics") - # dataset = load_dataset(r"ceval/ceval-exam", name="middle_school_geography") - # dataset = load_dataset(r"ceval/ceval-exam", name="middle_school_physics") - - samples = dataset["val"] - - # ----------------------------- - # 初始化模型 - # ----------------------------- - if use_awq: - model = JiugeAWQModel(model_path, device_type, ndev) - elif use_gptq: - model = JiugeGPTQModel(model_path, device_type, ndev) - else: - model = JiugeForCeval(model_path, device_type, ndev) - - answers_list = [] - for sample in samples: - input_content = f"'question':{sample['question']},'A': {sample['A']}, 'B':{sample['B']}, 'C': {sample['C']},'D': {sample['D']}。" - conversation = [ - { - "role": "system", - "content": "请从question的A，B，C，D四个选项中选择正确的选项。例如，标准答案：A。", - }, - {"role": "user", "content": input_content}, - ] - - answer = sample["answer"] - output_content, avg_time = model.generate( - conversation, 500, topp_=1.0, topk_=1, temperature_=1.0 - ) - print("标准答案：", answer) - answers_list.append( - {"id": sample["id"], "output_content": output_content, "answer": answer} - ) - - model.destroy_model_instance() - - print("-------------------------------------------------------------") - - # ----------------------------- - # 计算正确率 - # ----------------------------- - import re - true_num = 0 - all_num = 0 - for cont in answers_list: - id = cont["id"] - output = cont["output_content"] - answer = cont["answer"] - - all_num = all_num + 1 - # 提取模型输出的第一个选项 A/B/C/D - match = re.search(r"[A-D]", output) - model_answer = match.group(0) if match else "" - if model_answer == answer: - true_num = true_num + 1 - print(f"id {id} : ", "正确") - else: - print(f"id {id}: ", "错误") - - print(f"成绩: {true_num}/{all_num}", true_num / all_num) - - -if __name__ == "__main__": - test() diff --git a/setup.py b/setup.py index ac73b956a..7274e314e 100644 --- a/setup.py +++ b/setup.py @@ -1,16 +1,93 @@ +import os +import shutil import subprocess from pathlib import Path -from setuptools import setup +from setuptools import Distribution, find_packages, setup from setuptools.command.build import build from setuptools.command.develop import develop from setuptools.command.egg_info import egg_info +PROJECT_ROOT = Path(__file__).resolve().parent +INFINI_LIBRARY_NAMES = ("infiniops", "infiniccl", "infinirt") +INFINICORE_NATIVE_ARTIFACTS = [ + "_infinicore*.so", + "_infinicore*.dylib", + "_infinicore*.dll", + "_infinicore*.pyd", + "libinfinicore_runtime.so", + "libinfinicore_runtime.so.*", + "libinfinicore_runtime.dylib", + "infinicore_runtime.dll", + "libinfinicore_runtime.dll", + "libinfiniops.so", + "libinfiniops.so.*", + "libinfiniops.dylib", + "infiniops.dll", + "libinfiniops.dll", + "libinfiniccl.so", + "libinfiniccl.so.*", + "libinfiniccl.dylib", + "infiniccl.dll", + "libinfiniccl.dll", + "libinfinirt.so", + "libinfinirt.so.*", + "libinfinirt.dylib", + "infinirt.dll", + "libinfinirt.dll", +] +INFINILM_EXTENSION_ARTIFACTS = [ + "_infinilm*.so", + "_infinilm*.dylib", + "_infinilm*.dll", + "_infinilm*.pyd", +] +SHARED_LIBRARY_PATTERNS = ( + "lib{name}.so", + "lib{name}.so.*", + "lib{name}.dylib", + "{name}.dll", + "lib{name}.dll", +) + + +class BinaryDistribution(Distribution): + def has_ext_modules(self): + return True + + +def stage_runtime_dependencies(): + """Stage the modern Infini libraries next to the shared runtime.""" + infini_root = Path(os.environ.get("INFINI_ROOT", Path.home() / ".infini")) + search_directories = (infini_root / "lib", infini_root / "lib64") + destination = PROJECT_ROOT / "python/infinicore/lib" + destination.mkdir(parents=True, exist_ok=True) + + for library in INFINI_LIBRARY_NAMES: + sources = {} + for directory in search_directories: + for pattern in SHARED_LIBRARY_PATTERNS: + for source in sorted(directory.glob(pattern.format(name=library))): + if source.is_file(): + sources.setdefault(source.name, source) + + if not sources: + searched = ", ".join(str(directory) for directory in search_directories) + raise FileNotFoundError( + f"Could not find the `{library}` shared library in {searched}. " + "Set `INFINI_ROOT` to an installed Infini stack." + ) + + for source in sources.values(): + shutil.copy2(source, destination / source.name) + def build_cpp_module(): - """Build and install the C++ extension module""" - subprocess.run(["xmake", "build", "_infinilm"], check=True) - subprocess.run(["xmake", "install", "_infinilm"], check=True) + """Build and install the C++ extension modules.""" + stage_runtime_dependencies() + for target in ("infinicore_runtime", "_infinicore", "_infinilm"): + subprocess.run(["xmake", "build", target], check=True) + subprocess.run(["xmake", "install", target], check=True) class Build(build): @@ -37,7 +114,14 @@ def run(self): version="0.1.0", description="InfiniLM model implementations", package_dir={"": "python"}, - packages=["infinilm", "infinilm.models", "infinilm.lib", "infinilm.distributed"], + packages=find_packages(where="python"), + package_data={ + "infinicore.lib": INFINICORE_NATIVE_ARTIFACTS, + "infinilm.lib": INFINILM_EXTENSION_ARTIFACTS, + }, + include_package_data=True, + distclass=BinaryDistribution, + zip_safe=False, cmdclass={ "build": Build, "develop": Develop, diff --git a/src/allocator.hpp b/src/allocator.hpp deleted file mode 100644 index dc043fe97..000000000 --- a/src/allocator.hpp +++ /dev/null @@ -1,54 +0,0 @@ -#ifndef ALLOCATOR_HPP -#define ALLOCATOR_HPP - -#include "infinicore_infer.h" -#include -#include -#include -#include - -class AllocatorBase { -public: - virtual void *alloc(size_t size) = 0; - virtual void release(void *ptr) = 0; - virtual ~AllocatorBase() = default; -}; - -class MemoryPool : public AllocatorBase { -public: - static constexpr size_t DEFAULT_ALIGNMENT = 512; - - explicit MemoryPool(size_t initialSize = 0, size_t alignment = DEFAULT_ALIGNMENT); - ~MemoryPool(); - - void *alloc(size_t size) override; - void release(void *ptr) override; - - size_t getAlignment() const { return _alignment; } - -private: - struct Block { - void *base; - void *ptr; - size_t size; - bool is_free; - - Block(void *b, void *p, size_t s, bool f) - : base(b), ptr(p), size(s), is_free(f) {} - - bool operator<(const Block &other) const { - return ptr < other.ptr; - } - }; - - void *allocateNewRegion(size_t size); - void tryCoalesce(const Block &block); - - size_t _alignment; - std::vector _base_regions; - std::set _all_blocks; - std::multimap::iterator> _free_blocks; - std::unordered_map::iterator> _ptr_to_block; -}; - -#endif diff --git a/src/allocator/memory_allocator.cpp b/src/allocator/memory_allocator.cpp deleted file mode 100644 index 003c01d4d..000000000 --- a/src/allocator/memory_allocator.cpp +++ /dev/null @@ -1,135 +0,0 @@ -#include "../allocator.hpp" -#include "../utils.hpp" - -MemoryPool::MemoryPool(size_t initialSize, size_t alignment) - : _alignment(alignment) { - // Validate alignment is power of two - if ((alignment & (alignment - 1)) != 0 || alignment == 0) { - throw std::invalid_argument("Alignment must be a power of two"); - } - - if (initialSize > 0) { - allocateNewRegion(initialSize); - } -} - -MemoryPool::~MemoryPool() { - for (void *region : _base_regions) { - RUN_INFINI(infinirtFree(region)); - } -} - -void *MemoryPool::alloc(size_t size) { - if (size == 0) { - return nullptr; - } - - // Calculate aligned size - const size_t aligned_size = (size + _alignment - 1) & ~(_alignment - 1); - - // Find the first block with enough space (after alignment) - auto it = _free_blocks.lower_bound(aligned_size); - if (it == _free_blocks.end()) { - allocateNewRegion(aligned_size); - it = _free_blocks.lower_bound(aligned_size); - if (it == _free_blocks.end()) { - throw std::bad_alloc(); - } - } - - auto block_it = it->second; - Block block = *block_it; - _free_blocks.erase(it); - _all_blocks.erase(block_it); - - // Align the pointer within the block - size_t alignment_padding = reinterpret_cast(block.ptr) - reinterpret_cast(block.ptr); - - // Calculate remaining space after allocation - const size_t remaining = block.size - aligned_size - alignment_padding; - - // Create allocated block - Block alloc_block(block.base, block.ptr, aligned_size, false); - auto alloc_it = _all_blocks.insert(alloc_block).first; - _ptr_to_block[block.ptr] = alloc_it; - - // Split remaining space if it's large enough - if (remaining >= _alignment) { - void *rem_ptr = static_cast(block.ptr) + aligned_size; - Block rem_block(block.base, rem_ptr, remaining, true); - auto rem_it = _all_blocks.insert(rem_block).first; - _free_blocks.emplace(remaining, rem_it); - } - - return block.ptr; -} - -void MemoryPool::release(void *ptr) { - if (ptr == nullptr) { - return; - } - - auto it = _ptr_to_block.find(ptr); - if (it == _ptr_to_block.end()) { - throw std::runtime_error("Invalid pointer to free"); - } - - auto block_it = it->second; - Block block = *block_it; - _all_blocks.erase(block_it); - block.is_free = true; - auto new_it = _all_blocks.insert(block).first; - _ptr_to_block.erase(ptr); - tryCoalesce(*new_it); -} - -void *MemoryPool::allocateNewRegion(size_t size) { - // Allocate exactly the requested size - void *ptr = nullptr; - RUN_INFINI(infinirtMalloc(&ptr, size)); - _base_regions.push_back(ptr); - - // Align the pointer within the allocated region - size_t alignment_padding = reinterpret_cast(ptr) - reinterpret_cast(ptr); - size_t usable_size = size - alignment_padding; - - Block new_block(ptr, ptr, usable_size, true); - auto it = _all_blocks.insert(new_block).first; - _free_blocks.emplace(usable_size, it); - - return ptr; -} - -void MemoryPool::tryCoalesce(const Block &block) { - auto it = _all_blocks.find(block); - if (it == _all_blocks.end()) { - return; - } - - Block merged = *it; - auto next = std::next(it); - auto prev = (it == _all_blocks.begin()) ? _all_blocks.end() : std::prev(it); - - _all_blocks.erase(it); - _free_blocks.erase(merged.size); - - // Coalesce with next - if (next != _all_blocks.end() && next->is_free && static_cast(merged.ptr) + merged.size == next->ptr) { - _free_blocks.erase(next->size); - merged.size += next->size; - _all_blocks.erase(next); - } - - // Coalesce with prev - if (prev != _all_blocks.end() && prev->is_free && static_cast(prev->ptr) + prev->size == merged.ptr) { - _free_blocks.erase(prev->size); - merged.ptr = prev->ptr; - merged.size += prev->size; - merged.base = prev->base; - _all_blocks.erase(prev); - } - - merged.is_free = true; - auto new_it = _all_blocks.insert(merged).first; - _free_blocks.emplace(merged.size, new_it); -} diff --git a/src/cache.hpp b/src/cache.hpp deleted file mode 100644 index 4ec036ba3..000000000 --- a/src/cache.hpp +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "tensor.hpp" -#include -#include - -struct KVCache { - std::vector>> k, v; -}; diff --git a/src/cache_manager/kvcache.cpp b/src/cache_manager/kvcache.cpp deleted file mode 100644 index 99d07dfa9..000000000 --- a/src/cache_manager/kvcache.cpp +++ /dev/null @@ -1,80 +0,0 @@ -#include "../cache.hpp" - -__INFINI_C struct KVCache *createKVCache( - size_t nlayers, - size_t max_len, - size_t nkvh_, - size_t dk, - size_t dv, - infiniDtype_t dtype, - infiniDevice_t device, - int *dev_ids, - size_t ndev) { - - KVCache *cache = new KVCache(); - auto nkvh = nkvh_ / ndev; - - auto shape_k = std::vector{max_len, nkvh, dk}; - auto shape_v = std::vector{max_len, nkvh, dv}; - for (unsigned int idev = 0; idev < ndev; idev++) { - RUN_INFINI(infinirtSetDevice(device, dev_ids[idev])); - auto kcache = std::vector>(); - auto vcache = std::vector>(); - for (unsigned int layer = 0; layer < nlayers; layer++) { - kcache.push_back(Tensor::buffer(dtype, shape_k)); - vcache.push_back(Tensor::buffer(dtype, shape_v)); - } - cache->k.push_back(kcache); - cache->v.push_back(vcache); - } - - return cache; -} - -__INFINI_C struct KVCache *duplicateKVCache(const KVCache *kv_cache, size_t seq_len) { - auto ndev = kv_cache->k.size(); - auto nlayers = kv_cache->k[0].size(); - auto device = kv_cache->k[0][0]->deviceType(); - auto dtype = kv_cache->k[0][0]->dtype(); - auto shape_k = kv_cache->k[0][0]->shape(); - auto shape_v = kv_cache->v[0][0]->shape(); - auto size_k = seq_len * shape_k[1] * shape_k[2] * dsize(dtype); - auto size_v = seq_len * shape_v[1] * shape_v[2] * dsize(dtype); - KVCache *new_kv_cache = new KVCache(); - for (unsigned int idev = 0; idev < ndev; idev++) { - RUN_INFINI(infinirtSetDevice(device, kv_cache->k[idev][0]->deviceId())); - for (unsigned int layer = 0; layer < nlayers; layer++) { - auto kcache = std::vector>(); - auto vcache = std::vector>(); - for (unsigned int layer = 0; layer < nlayers; layer++) { - kcache.push_back(Tensor::buffer(dtype, shape_k)); - vcache.push_back(Tensor::buffer(dtype, shape_v)); - } - new_kv_cache->k.push_back(kcache); - new_kv_cache->v.push_back(vcache); - RUN_INFINI(infinirtMemcpy(new_kv_cache->k[idev][layer]->data(), - kv_cache->k[idev][layer]->data(), - size_k, - INFINIRT_MEMCPY_D2D)); - RUN_INFINI(infinirtMemcpy(new_kv_cache->v[idev][layer]->data(), - kv_cache->v[idev][layer]->data(), - size_v, - INFINIRT_MEMCPY_D2D)); - } - } - return new_kv_cache; -} - -__INFINI_C void dropKVCache(KVCache *kv_cache) { - auto ndev = kv_cache->k.size(); - auto nlayers = kv_cache->k[0].size(); - auto device = kv_cache->k[0][0]->deviceType(); - for (unsigned int idev = 0; idev < ndev; idev++) { - RUN_INFINI(infinirtSetDevice(device, kv_cache->k[idev][0]->deviceId())); - for (unsigned int layer = 0; layer < nlayers; layer++) { - kv_cache->k[idev][layer].reset(); - kv_cache->v[idev][layer].reset(); - } - } - delete kv_cache; -} diff --git a/src/cache_manager/opcache_manager.hpp b/src/cache_manager/opcache_manager.hpp deleted file mode 100644 index 7f8d38330..000000000 --- a/src/cache_manager/opcache_manager.hpp +++ /dev/null @@ -1,197 +0,0 @@ -#ifndef CACHE_MANAGER_HPP -#define CACHE_MANAGER_HPP - -#include -#include -#include -#include - -#include "../tensor.hpp" -#include "../utils.hpp" -#include "infinicore_infer.h" - -class IDescriptorDestroyer { -public: - virtual ~IDescriptorDestroyer() = default; - virtual void destroy(void *descriptor) = 0; -}; - -template -class DescriptorDestroyer : public IDescriptorDestroyer { - using DestroyFunc = infiniStatus_t (*)(DescriptorType); - DestroyFunc destroyFunc; - -public: - DescriptorDestroyer(DestroyFunc func) : destroyFunc(func) {} - - void destroy(void *descriptor) override { - destroyFunc(*static_cast(descriptor)); - } -}; - -template -class LRUDescriptorCache { -private: - struct CacheNode { - size_t key; - DescriptorType desc; - CacheNode *prev; - CacheNode *next; - - CacheNode() : key(0), desc(), prev(nullptr), next(nullptr) {} - CacheNode(size_t k, const DescriptorType &d) : key(k), desc(d), prev(nullptr), next(nullptr) {} - }; - - std::unordered_map cache; - CacheNode *head; - CacheNode *tail; - const size_t capacity; - size_t size; - std::unique_ptr destroyer; - - void removeNode(CacheNode *node) { - node->prev->next = node->next; - node->next->prev = node->prev; - if (destroyer) { - destroyer->destroy(&node->desc); - } - cache.erase(node->key); - delete node; - --size; - } - - void addToTop(CacheNode *node) { - node->next = head->next; - node->next->prev = node; - node->prev = head; - head->next = node; - cache[node->key] = node; - if (++size > capacity) { - removeNode(tail->prev); - } - } - - void moveToTop(CacheNode *node) { - node->prev->next = node->next; - node->next->prev = node->prev; - node->next = head->next; - node->next->prev = node; - node->prev = head; - head->next = node; - } - -public: - template - LRUDescriptorCache(size_t c, DestroyFunc destroyFunc) - : capacity(c), size(0), destroyer(std::make_unique>(destroyFunc)) { - head = new CacheNode(); - tail = new CacheNode(); - head->next = tail; - tail->prev = head; - } - - ~LRUDescriptorCache() { - while (head->next != tail) { - removeNode(head->next); - } - delete head; - delete tail; - } - - bool get(size_t key, DescriptorType &out_desc) { - auto it = cache.find(key); - if (it == cache.end()) { - return false; - } - - CacheNode *node = it->second; - moveToTop(node); - out_desc = node->desc; - return true; - } - - void put(size_t key, const DescriptorType &descriptor) { - auto it = cache.find(key); - if (it != cache.end()) { - // Key already exists, update the descriptor - CacheNode *node = it->second; - if (destroyer) { - destroyer->destroy(&node->desc); - } - node->desc = descriptor; - moveToTop(node); - return; - } - - // Check if we need to evict - if (size >= capacity) { - removeNode(tail->prev); - } - - // Create new node and add to top - CacheNode *node = new CacheNode(key, descriptor); - addToTop(node); - } - - LRUDescriptorCache(const LRUDescriptorCache &) = delete; - LRUDescriptorCache &operator=(const LRUDescriptorCache &) = delete; -}; - -// Helper macro to generate the destroy function name -#define DESTROY_FUNC(OpType) infiniopDestroy##OpType##Descriptor - -// Declare cache and access functions -#define DECLARE_OP_CACHE(OpType) \ - LRUDescriptorCache OpType##_cache; \ - bool get##OpType##Descriptor(size_t key, infiniop##OpType##Descriptor_t &desc) { \ - return OpType##_cache.get(key, desc); \ - } \ - void put##OpType##Descriptor(size_t key, const infiniop##OpType##Descriptor_t &desc) { \ - OpType##_cache.put(key, desc); \ - } - -class CacheManager { -public: - DECLARE_OP_CACHE(Add) - DECLARE_OP_CACHE(Conv) - DECLARE_OP_CACHE(Mul) - DECLARE_OP_CACHE(RMSNorm) - DECLARE_OP_CACHE(Gemm) - DECLARE_OP_CACHE(RoPE) - DECLARE_OP_CACHE(Rearrange) - DECLARE_OP_CACHE(CausalSoftmax) - DECLARE_OP_CACHE(Topkrouter) - DECLARE_OP_CACHE(SwiGLU) - DECLARE_OP_CACHE(Silu) - DECLARE_OP_CACHE(RandomSample) - DECLARE_OP_CACHE(DequantizeAWQ) - DECLARE_OP_CACHE(DequantizeGPTQ) - - CacheManager(size_t capacity = 100) - : Add_cache(capacity, DESTROY_FUNC(Add)), - Conv_cache(capacity, DESTROY_FUNC(Conv)), - Mul_cache(capacity, DESTROY_FUNC(Mul)), - RMSNorm_cache(capacity, DESTROY_FUNC(RMSNorm)), - Gemm_cache(capacity, DESTROY_FUNC(Gemm)), - RoPE_cache(capacity, DESTROY_FUNC(RoPE)), - Rearrange_cache(capacity, DESTROY_FUNC(Rearrange)), - CausalSoftmax_cache(capacity, DESTROY_FUNC(CausalSoftmax)), - Topkrouter_cache(capacity, DESTROY_FUNC(Topkrouter)), - SwiGLU_cache(capacity, DESTROY_FUNC(SwiGLU)), - Silu_cache(capacity, DESTROY_FUNC(Silu)), - RandomSample_cache(capacity, DESTROY_FUNC(RandomSample)), - DequantizeAWQ_cache(capacity, DESTROY_FUNC(DequantizeAWQ)), - DequantizeGPTQ_cache(capacity, DESTROY_FUNC(DequantizeGPTQ)) {} - - template - static size_t createDescriptorKey(Tensors... tensors) { - size_t seed = 0; - (..., (tensors ? hash_combine(seed, tensors->seed()) : (void)0)); - return seed; - } -}; - -#undef DESTROY_FUNC -#undef DECLARE_OP_CACHE - -#endif // CACHE_MANAGER_HPP diff --git a/src/dataloader/weights_loader.cpp b/src/dataloader/weights_loader.cpp deleted file mode 100644 index 71e634860..000000000 --- a/src/dataloader/weights_loader.cpp +++ /dev/null @@ -1,86 +0,0 @@ -#include "weights_loader.hpp" -#include "infinicore_infer/weights_loader.h" - -#include "../utils.hpp" - -#include -#include - -namespace infinicore::weights { - -void Weight::load(const void *host_data, infinirtStream_t stream) { - if (_dist_type == DistributionType::FULL) { - _tensor->load(host_data, stream); - } else if (_dist_type == DistributionType::ROW || _tensor->ndim() == 1) { // 1D column-distributed is same as row-distributed - _tensor->load((const char *)host_data + _rank * _tensor->numel() * dsize(_tensor->dtype()), stream); - } else if (_dist_type == DistributionType::COLUMN && _tensor->ndim() > 1) { // _dist_type == DistributionType::COLUMN - void *rearranged_ptr; - RUN_INFINI(infinirtMallocHost(&rearranged_ptr, _tensor->numel() * dsize(_tensor->dtype()))); - size_t row_size = _tensor->shape()[_tensor->ndim() - 1] * dsize(_tensor->dtype()); - size_t host_offset = _rank * row_size; - size_t host_row_size = _nrank * row_size; - size_t rows = std::accumulate(_tensor->shape().begin(), _tensor->shape().end() - 1, size_t(1), std::multiplies()); - for (size_t row = 0; row < rows; row++) { - memcpy((char *)rearranged_ptr + row * row_size, - (char *)host_data + host_offset + row * host_row_size, - row_size); - } - _tensor->load(rearranged_ptr, stream); - RUN_INFINI(infinirtFreeHost(rearranged_ptr)); - } else { - std::cerr << "Unsupported distribution type: " << _dist_type << std::endl; - std::abort(); - } -}; - -Loader::Loader(infiniDevice_t dev, const std::vector &dev_ids) : _device(dev), _dev_ids(dev_ids) { - _streams.resize(_dev_ids.size()); - _weights_maps.resize(_dev_ids.size()); - for (int rank = 0; rank < int(_dev_ids.size()); rank++) { - RUN_INFINI(infinirtSetDevice(_device, _dev_ids[rank])); - _weights_maps[rank] = std::unordered_map>(); - RUN_INFINI(infinirtStreamCreate(&_streams[rank])); - } -} -void Loader::register_weight(const std::string &name, std::shared_ptr tensor, int rank, DistributionType dist_type) { - _weights_maps[rank][name] = std::make_shared(tensor, rank, _dev_ids.size(), dist_type); -} -void Loader::load(const std::string &name, const void *host_data) { - for (int rank = 0; rank < int(_dev_ids.size()); rank++) { - RUN_INFINI(infinirtSetDevice(_device, _dev_ids[rank])); - auto it = _weights_maps[rank].find(name); - if (it == _weights_maps[rank].end()) { - std::cerr << "Weight " << name << " not found in rank " << rank << std::endl; - std::abort(); - } - - _weights_maps[rank][name]->load(host_data, _streams[rank]); - } - for (int rank = int(_dev_ids.size() - 1); rank >= 0; rank--) { - RUN_INFINI(infinirtSetDevice(_device, _dev_ids[rank])); - RUN_INFINI(infinirtStreamSynchronize(_streams[rank])); - } -} - -void Loader::finalize() { - int dev_id; - RUN_INFINI(infinirtGetDevice(nullptr, &dev_id)); - for (int rank = 0; rank < int(_dev_ids.size()); rank++) { - RUN_INFINI(infinirtSetDevice(_device, _dev_ids[rank])); - RUN_INFINI(infinirtStreamSynchronize(_streams[rank])); - RUN_INFINI(infinirtStreamDestroy(_streams[rank])); - } - RUN_INFINI(infinirtSetDevice(_device, dev_id)); -} -std::shared_ptr Loader::get(const std::string &name, int rank) { - return _weights_maps[rank][name]->tensor(); -} - -} // namespace infinicore::weights - -__INFINI_C void -loadModelWeight(struct ModelWeights *weights_, const char *name, void *data) { - std::string name_str(name); - auto weights = reinterpret_cast(weights_); - weights->load(name_str, data); -} diff --git a/src/dataloader/weights_loader.hpp b/src/dataloader/weights_loader.hpp deleted file mode 100644 index 931d72722..000000000 --- a/src/dataloader/weights_loader.hpp +++ /dev/null @@ -1,61 +0,0 @@ -#ifndef WEIGHTS_LOADER_HPP -#define WEIGHTS_LOADER_HPP - -#include "../tensor.hpp" - -#include -#include - -namespace infinicore { - -namespace weights { -enum DistributionType { - FULL, - ROW, - COLUMN -}; -class Weight { -private: - std::shared_ptr _tensor; - int _rank; - int _nrank; - DistributionType _dist_type; - -public: - Weight(std::shared_ptr tensor, - int rank = 0, - int nrank = 1, - DistributionType dist_type = DistributionType::FULL) - : _tensor(tensor), _rank(rank), _nrank(nrank), _dist_type(dist_type) {} - std::shared_ptr tensor() const { return _tensor; } - int rank() const { return _rank; } - int nrank() const { return _nrank; } - void load(const void *host_data, infinirtStream_t stream = nullptr); -}; - -class Loader { -protected: - std::vector>> _weights_maps; - infiniDevice_t _device; - std::vector _dev_ids; - std::vector _streams; - -public: - Loader(infiniDevice_t, const std::vector &dev_ids); - - /// @brief register a tensor to the loader - /// @param name name (aka key) of the tensor - /// @param tensor - /// @param rank the rank of the weight tensor (default 0) - /// @param dist_type either FULL, or distributed by ROW or COLUMN (default FULL) - void register_weight(const std::string &name, std::shared_ptr tensor, int rank = 0, DistributionType dist_type = DistributionType::FULL); - void load(const std::string &name, const void *host_data); - void finalize(); - std::shared_ptr get(const std::string &name, int rank = 0); - const std::vector &devIds() const { return _dev_ids; } - infiniDevice_t device() const { return _device; } -}; -} // namespace weights -} // namespace infinicore - -#endif // WEIGHTS_LOADER_HPP diff --git a/src/models/deepseek_v3/deepseek_v3.cpp b/src/models/deepseek_v3/deepseek_v3.cpp deleted file mode 100644 index db22d87df..000000000 --- a/src/models/deepseek_v3/deepseek_v3.cpp +++ /dev/null @@ -1,589 +0,0 @@ -#include "deepseek_v3_impl.hpp" - -#include "../../tensor.hpp" -#include "../../utils.hpp" -#include "../inference_context.hpp" -#include "infinicore_infer.h" - -#include -#include -#include - -void createDeviceResource(DeepSeekV3DeviceResource *rsrc, const DeepSeekV3Meta *meta, - std::shared_ptr weights, - infiniDevice_t device, int idev, - int ndev, int dev_id, - infinicclComm_t comm) { - RUN_INFINI(infinirtSetDevice(device, dev_id)); - RUN_INFINI(infinirtStreamSynchronize(weights->load_stream)); - infiniopHandle_t handle; - infiniopCreateHandle(&handle); - infinirtStream_t stream; - infinirtStreamCreate(&stream); - - auto memory_pool = std::make_shared(); - - *rsrc = DeepSeekV3DeviceResource{ - device, - dev_id, - handle, - weights, - stream, - comm, - memory_pool, - }; - RUN_INFINI(infinirtDeviceSynchronize()); -} - -void releaseDeviceResource(DeepSeekV3DeviceResource &res) { - infinirtDeviceSynchronize(); - - res.weights.reset(); - - infiniopDestroyHandle(res.handle); - res.handle = nullptr; - infinirtStreamDestroy(res.stream); - res.stream = nullptr; - infinicclCommDestroy(res.comm); - res.comm = nullptr; -} - -void inferDeviceBatch(const DeepSeekV3Meta &meta, DeepSeekV3DeviceResource &rsrc, - uint32_t idev, uint32_t ndev, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct DeepSeekV3Cache **caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output, void *last_logits) { - - auto dt_logits = meta.dt_logits; - // auto dt_norm = meta.dt_norm; - // auto dt_quant_weight = meta.dt_quant_weight; - // auto dt_quant_scale = meta.dt_quant_scale; - // auto dt_quant_zero = meta.dt_quant_zero; - // auto dt_gate_weight = meta.dt_gate_weight; - // auto dt_gate_bias = meta.dt_gate_bias; - auto n_dense_layer = meta.n_dense_layer; - auto n_sparse_layer = meta.n_sparse_layer; - auto nlayer = n_dense_layer + n_sparse_layer; - size_t nh = meta.nh / size_t(ndev); - - auto d = meta.d; - auto d_rope = meta.d_rope; - auto d_nope = meta.d_nope; - auto r_q = meta.r_q; - auto r_kv = meta.r_kv; - auto d_qk = meta.d_qk; - auto d_v = meta.d_v; - // auto routed_scale = meta.routed_scale; - // auto nexperts = meta.nexperts; - // auto kexperts = meta.kexperts; - - auto di = meta.di / size_t(ndev); - auto dvoc = meta.dvoc; - - auto stream = rsrc.stream; - - auto weights = rsrc.weights; - - // Allocate buffers - auto logits_in = Tensor::buffer(dt_logits, {ntok, d}, rsrc.memory_pool); - auto logits_out = Tensor::buffer(dt_logits, {ntok, d}, rsrc.memory_pool); - - auto q_a_buf = Tensor::buffer(dt_logits, {ntok, r_q}, rsrc.memory_pool); - auto q_buf = Tensor::buffer(dt_logits, {ntok, nh * d_qk}, rsrc.memory_pool); - auto kv_a_buf = Tensor::buffer(dt_logits, {ntok, r_kv + d_rope}, rsrc.memory_pool); - auto o_buf = Tensor::buffer(dt_logits, {ntok, nh * d_v}, rsrc.memory_pool); - - auto prob_buf = Tensor::buffer(dt_logits, {nreq, dvoc}, rsrc.memory_pool); - auto result_buf = Tensor::buffer(INFINI_DTYPE_I64, {nreq}, rsrc.memory_pool); - auto result_cpu = std::vector(nreq); - - // Prepare inputs - auto batch_pos_ids = std::vector(ntok); - size_t req_start = 0; - for (uint32_t req = 0; req < nreq; req++) { - for (uint32_t i = 0; i < req_lens[req]; i++) { - batch_pos_ids[req_start + i] = req_pos[req] + i; - } - req_start += req_lens[req]; - } - - std::shared_ptr pos_ids_buf; - if (rsrc.device == INFINI_DEVICE_CPU) { - pos_ids_buf = Tensor::weight(batch_pos_ids.data(), INFINI_DTYPE_U32, {ntok}); - } else { - pos_ids_buf = Tensor::buffer(INFINI_DTYPE_U32, {ntok}, rsrc.memory_pool); - RUN_INFINI(infinirtMemcpyAsync(pos_ids_buf->data(), batch_pos_ids.data(), sizeof(uint32_t) * ntok, - INFINIRT_MEMCPY_H2D, stream)); - } - for (uint32_t i = 0; i < ntok; i++) { - RUN_INFINI(infinirtMemcpyAsync(logits_in->data(i * d), - weights->w_in_embd->data(tokens[i] * d), - dsize(dt_logits) * d, INFINIRT_MEMCPY_D2D, stream)); - } - - // Attention - // attention inner - size_t max_qk_size = 0; - size_t max_seq_len = 0; - size_t max_total_len = 0; - - for (uint32_t req = 0; req < nreq; req++) { - auto past_len = req_pos[req]; - auto seq_len = req_lens[req]; - auto total_len = past_len + seq_len; - - max_qk_size = std::max(max_qk_size, size_t(seq_len * total_len)); - max_seq_len = std::max(max_seq_len, size_t(seq_len)); - max_total_len = std::max(max_total_len, size_t(total_len)); - } - auto full_k_buf = Tensor::buffer(dt_logits, {max_total_len, nh * d_qk}, rsrc.memory_pool); - auto kv_b_buf = Tensor::buffer(dt_logits, {max_total_len, nh * (d_nope + d_v)}, rsrc.memory_pool); - auto attn_score_buf = Tensor::buffer(dt_logits, {nh, max_qk_size}, rsrc.memory_pool); - auto attn_val_buf = Tensor::buffer(dt_logits, {nh, max_seq_len, d_v}, rsrc.memory_pool); - - // Compute - for (uint32_t layer = 0; layer < nlayer; layer++) { - // 1. Attention - // rms norm - rmsnorm(logits_out, logits_in, weights->w_layers[layer].mla_norm, meta.epsilon); - // q_proj - dequant_linear(q_a_buf, logits_out, - weights->w_layers[layer].mla->q_a_proj->w, - weights->w_layers[layer].mla->q_a_proj->s, - weights->w_layers[layer].mla->q_a_proj->z, - 1.0, 0.0, nullptr, nullptr); - rmsnorm(q_a_buf, q_a_buf, weights->w_layers[layer].mla->q_a_norm, meta.epsilon); - dequant_linear(q_buf, q_a_buf, - weights->w_layers[layer].mla->q_b_proj->w, - weights->w_layers[layer].mla->q_b_proj->s, - weights->w_layers[layer].mla->q_b_proj->z, - 1.0, 0.0, nullptr, nullptr); - auto q_rot = q_buf->view({ntok, nh, d_qk})->slice(2, d_nope, d_rope); - rope_v2(q_rot, q_rot, pos_ids_buf, weights->sin_table, weights->cos_table); - // kv_proj - dequant_linear(kv_a_buf, logits_out, - weights->w_layers[layer].mla->kv_a_proj->w, - weights->w_layers[layer].mla->kv_a_proj->s, - weights->w_layers[layer].mla->kv_a_proj->z, - 1.0, 0.0, nullptr, nullptr); - auto kv_pass = kv_a_buf->slice(1, 0, r_kv); - rmsnorm(kv_pass, kv_pass, weights->w_layers[layer].mla->kv_a_norm, meta.epsilon); - auto k_rot = kv_a_buf->slice(1, r_kv, d_rope)->view({ntok, 1, d_rope}); - rope_v2(k_rot, k_rot, pos_ids_buf, weights->sin_table, weights->cos_table); - - size_t token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto past_len = req_pos[req]; - auto seq_len = req_lens[req]; - auto total_len = past_len + seq_len; - auto o_req = o_buf->slice({{0, token_offset, seq_len}})->view({seq_len, nh, d_v}); - auto q_req = q_buf->slice({{0, token_offset, seq_len}}); - auto kv_a_req = kv_a_buf->slice({{0, token_offset, seq_len}}); - auto kv_pass_req = kv_a_req->slice(1, 0, r_kv); - auto k_rot_req = kv_a_req->slice(1, r_kv, d_rope); - - // concat cache - rearrange(caches[req]->kv_pass[idev][layer]->slice(0, past_len, seq_len), kv_pass_req); - rearrange(caches[req]->k_rot[idev][layer]->slice(0, past_len, seq_len), k_rot_req); - // kv_b_proj - auto kv_b_req = kv_b_buf->slice(0, 0, total_len); - dequant_linear(kv_b_req, caches[req]->kv_pass[idev][layer]->slice(0, 0, total_len), - weights->w_layers[layer].mla->kv_b_proj->w, - weights->w_layers[layer].mla->kv_b_proj->s, - weights->w_layers[layer].mla->kv_b_proj->z, - 1.0, 0.0, nullptr, nullptr); - auto full_v_req = kv_b_req->slice(1, nh * d_nope, nh * d_v); - // concat k - auto full_k_req = full_k_buf->slice(0, 0, total_len); - auto full_k_pass_req = full_k_req->slice(1, 0, nh * d_nope); - auto full_k_rot_req = full_k_req->slice(1, nh * d_nope, nh * d_rope); - rearrange(full_k_pass_req, kv_b_req->slice(1, 0, nh * d_nope)); - rearrange(full_k_rot_req->view({total_len, nh, d_rope}), k_rot_req->view_as({total_len, nh, d_rope}, {ptrdiff_t(d_rope), 0, 1})); // expand k_rot - - // self attention - auto attn_score_req = attn_score_buf->slice(1, 0, seq_len * total_len)->view({nh, seq_len, total_len}); - linear(attn_score_req, - q_req->view({seq_len, nh, d_qk})->permute({1, 0, 2}), - full_k_req->view({total_len, nh, d_qk})->permute({1, 2, 0}), - 1.f / float(sqrt(d_qk)), 0.f, nullptr, nullptr); - // softmax - causalSoftmax(attn_score_req, attn_score_req); - // attn val - auto attn_val_req = attn_val_buf->slice(1, 0, seq_len)->view({nh, seq_len, d_v}); - linear(attn_val_req, attn_score_req, full_v_req->view({total_len, nh, d_v})->permute({1, 0, 2}), 1.f, 0.f, nullptr, nullptr); - // rearrange attn val - rearrange(o_req, attn_val_req->permute({1, 0, 2})); - - token_offset += seq_len; - } - - // o_proj - dequant_linear(logits_in, o_buf, - weights->w_layers[layer].mla->o_proj->w, - weights->w_layers[layer].mla->o_proj->s, - weights->w_layers[layer].mla->o_proj->z, - 1.0, 0.0, idev == 0 ? logits_in : nullptr, nullptr); // only rank 0 adds residual - - // All_reduce if distributed - if (rsrc.comm != nullptr) { - RUN_INFINI(infinicclAllReduce( - logits_in->data(), logits_in->data(), ntok * d, dt_logits, - INFINICCL_SUM, rsrc.comm, stream)); - RUN_INFINI(infinirtStreamSynchronize(stream)); - } - // 2. MLP - rmsnorm(logits_out, logits_in, weights->w_layers[layer].mlp_norm, meta.epsilon); - - if (layer < n_dense_layer) { - auto gate_dense = Tensor::buffer(dt_logits, {ntok, di}, rsrc.memory_pool); - auto up_dense = Tensor::buffer(dt_logits, {ntok, di}, rsrc.memory_pool); - dequant_linear(gate_dense, logits_out, - weights->w_layers[layer].dense_mlp->gate->w, - weights->w_layers[layer].dense_mlp->gate->s, - weights->w_layers[layer].dense_mlp->gate->z, 1.0, 0.0, nullptr, nullptr); - dequant_linear(up_dense, logits_out, - weights->w_layers[layer].dense_mlp->up->w, - weights->w_layers[layer].dense_mlp->up->s, - weights->w_layers[layer].dense_mlp->up->z, 1.0, 0.0, nullptr, nullptr); - swiglu(gate_dense, up_dense, gate_dense); - dequant_linear(logits_in, gate_dense, - weights->w_layers[layer].dense_mlp->down->w, - weights->w_layers[layer].dense_mlp->down->s, - weights->w_layers[layer].dense_mlp->down->z, - 1.0, 0.0, idev == 0 ? logits_in : nullptr, nullptr); // only rank 0 adds residual - } else { - - // ------------------------------------------------------------------------ // - // 后面几层，用的 稀疏MLP // - // ------------------------------------------------------------------------ // - // 需要提前申请的缓存，给每个MLP使用 - auto moe_gate_buf = Tensor::buffer(dt_logits, {ntok, meta.di_moe}, rsrc.memory_pool); - auto moe_up_buf = Tensor::buffer(dt_logits, {ntok, meta.di_moe}, rsrc.memory_pool); - - // 需要提前申请的缓存 - std::shared_ptr shared_states = Tensor::buffer(dt_logits, {ntok, d}, rsrc.memory_pool); // 用于存储共享专家的输出 - std::shared_ptr router_states_sum = Tensor::buffer(dt_logits, {ntok, d}, rsrc.memory_pool); // 用于存储路由专家的加权输出 - - // 需要提前申请的缓存 - std::shared_ptr router_logits = Tensor::buffer(dt_logits, {ntok, meta.nexperts}, rsrc.memory_pool); // nx256，路由专家的权重 - - std::shared_ptr values_gpu = Tensor::buffer(infiniDtype_t::INFINI_DTYPE_F32, {ntok * 8}, rsrc.memory_pool); // 用于存储topkrouter的输出，每个expert对应的加权权重。 - std::shared_ptr indices_gpu = Tensor::buffer(infiniDtype_t::INFINI_DTYPE_I32, {ntok * 8}, rsrc.memory_pool); // 用于存储topkrouter的输出，要经过哪些专家id（从256个中选8个） - std::vector values_cpu(ntok * 8, 0.f); // 用于存储topkrouter的输出，每个expert对应的加权权重。（从256个中选8个） - std::vector indices_cpu(ntok * 8, 0); // 用于存储topkrouter的输出，要经过哪些专家的索引。 - - // config 参数 - float routed_scaling_factor = meta.routed_scale; // config.json的超参"routed_scaling_factor"，是固定值 2.5 - size_t topk = 8; // config.json的超参"num_experts_per_tok", 是固定值 8 - - // 明确输入输出变量 - std::shared_ptr hidden_states = logits_out; // logits_out 是整个 MoE的输入，重新起名字为 hidden_states - - // ------------------------------------------------------------------------ // - // 开始计算 // - // ------------------------------------------------------------------------ // - // (1) 共享专家： hidden_states 经过一个共享专家 - { - // 输入: hidden_states - // 输出: shared_states - dequant_linear(moe_gate_buf, hidden_states, - weights->w_layers[layer].share_expert->gate->w, - weights->w_layers[layer].share_expert->gate->s, - weights->w_layers[layer].share_expert->gate->z, 1.0, 0.0, nullptr, nullptr); - dequant_linear(moe_up_buf, hidden_states, - weights->w_layers[layer].share_expert->up->w, - weights->w_layers[layer].share_expert->up->s, - weights->w_layers[layer].share_expert->up->z, 1.0, 0.0, nullptr, nullptr); - swiglu(moe_gate_buf, moe_up_buf, moe_gate_buf); - dequant_linear(shared_states, moe_gate_buf, - weights->w_layers[layer].share_expert->down->w, - weights->w_layers[layer].share_expert->down->s, - weights->w_layers[layer].share_expert->down->z, 1.0, 0.0, nullptr, nullptr); // only rank 0 adds residual - } - - // (2) topk操作： hidden_states 经过 topkrouter - { - // 输入: hidden_states - // 输出: values_cpu，indices_cpu - auto gate_weight = weights->w_layers[layer].route->w; - gemm(router_logits, hidden_states, gate_weight, 1.0, 0.0); // 非量化的版本 - - auto gate_correction_bias = weights->w_layers[layer].route->b; - topkrouter(values_gpu, indices_gpu, router_logits, gate_correction_bias, routed_scaling_factor, topk); - RUN_INFINI(infinirtMemcpy((void *)values_cpu.data(), values_gpu->data(), values_cpu.size() * sizeof(float), INFINIRT_MEMCPY_D2H)); - RUN_INFINI(infinirtMemcpy((void *)indices_cpu.data(), indices_gpu->data(), indices_cpu.size() * sizeof(int), INFINIRT_MEMCPY_D2H)); - } - - // (3) MoE操作： hidden_states经过一个8个路由专家 - // 输入: hidden_states, values_cpu，indices_cpu - // 输出: router_states_sum - for (size_t itok = 0; itok < ntok; ++itok) { // 先遍历每一个token，再遍历该toekn经过对应的专家 - - std::shared_ptr hidden_states_i = hidden_states->slice(0, itok, 1); - std::shared_ptr router_states_sum_i = router_states_sum->slice(0, itok, 1); - std::shared_ptr moe_gate_buf_i = moe_gate_buf->slice(0, itok, 1); - std::shared_ptr moe_up_buf_i = moe_up_buf->slice(0, itok, 1); - - // 经过第一个专家 : C = alpha * AB - { - // 输入: hidden_states - // 输出: router_states_sum_i - int index = indices_cpu[itok * topk]; - float alpha = values_cpu[itok * topk]; - - dequant_linear(moe_gate_buf_i, hidden_states_i, - weights->w_layers[layer].experts[index]->gate->w, - weights->w_layers[layer].experts[index]->gate->s, - weights->w_layers[layer].experts[index]->gate->z, 1.0, 0.0, nullptr, nullptr); - dequant_linear(moe_up_buf_i, hidden_states_i, - weights->w_layers[layer].experts[index]->up->w, - weights->w_layers[layer].experts[index]->up->s, - weights->w_layers[layer].experts[index]->up->z, 1.0, 0.0, nullptr, nullptr); - - swiglu(moe_gate_buf_i, moe_up_buf_i, moe_gate_buf_i); - - dequant_linear(router_states_sum_i, moe_gate_buf_i, - weights->w_layers[layer].experts[index]->down->w, - weights->w_layers[layer].experts[index]->down->s, - weights->w_layers[layer].experts[index]->down->z, alpha, 0.0, nullptr, nullptr); // only rank 0 adds residual - } - - // 经过后续的专家 : C = alpha * AB + C_last - for (size_t k = 1; k < topk; ++k) { - int index = indices_cpu[itok * topk + k]; - float alpha = values_cpu[itok * topk + k]; - - dequant_linear(moe_gate_buf_i, hidden_states_i, - weights->w_layers[layer].experts[index]->gate->w, - weights->w_layers[layer].experts[index]->gate->s, - weights->w_layers[layer].experts[index]->gate->z, 1.0, 0.0, nullptr, nullptr); - dequant_linear(moe_up_buf_i, hidden_states_i, - weights->w_layers[layer].experts[index]->up->w, - weights->w_layers[layer].experts[index]->up->s, - weights->w_layers[layer].experts[index]->up->z, 1.0, 0.0, nullptr, nullptr); - - swiglu(moe_gate_buf_i, moe_up_buf_i, moe_gate_buf_i); - - dequant_linear(router_states_sum_i, moe_gate_buf_i, - weights->w_layers[layer].experts[index]->down->w, - weights->w_layers[layer].experts[index]->down->s, - weights->w_layers[layer].experts[index]->down->z, alpha, 0.0, router_states_sum_i, nullptr); // only rank 0 adds residual - } - } - - // (4) 最后两个类型的专家求和 - // 输入: 共享专家结果shared_states， 路由专家结果router_states_sum - // 输出: logits_out - add(shared_states, shared_states, router_states_sum); - - // (5) 最后的残差连接 - add(logits_in, shared_states, logits_in); - } - - // All_reduce if distributed - if (rsrc.comm != nullptr) { - RUN_INFINI(infinicclAllReduce( - logits_in->data(), logits_in->data(), ntok * d, dt_logits, - INFINICCL_SUM, rsrc.comm, stream)); - RUN_INFINI(infinirtStreamSynchronize(stream)); - } - } - // Sample and Output - if (idev == 0) { - if (last_logits != nullptr) { - rmsnorm(logits_out, logits_in, weights->w_out_norm, meta.epsilon); - auto last_logits_buf = Tensor::buffer(dt_logits, {ntok, dvoc}, rsrc.memory_pool); - linear(last_logits_buf, logits_out, weights->w_out_embd, 1.0, 0.0, nullptr, nullptr); - RUN_INFINI(infinirtStreamSynchronize(stream)); - RUN_INFINI(infinirtMemcpy(last_logits, last_logits_buf->data(), dsize(dt_logits) * ntok * dvoc, INFINIRT_MEMCPY_D2H)); - } - if (output != nullptr) { - size_t token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto seq_len = req_lens[req]; - token_offset += seq_len; - rmsnorm(logits_out->slice(0, req, 1), - logits_in->slice(0, token_offset - 1, 1), - weights->w_out_norm, - meta.epsilon); - } - linear(prob_buf, logits_out->slice(0, 0, nreq), weights->w_out_embd, 1.0, 0.0, nullptr, nullptr); - std::random_device _rd; - std::mt19937 gen(_rd()); - token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto seq_len = req_lens[req]; - float random_val = std::uniform_real_distribution(0, 1)(gen); - randomSample(result_buf->slice(0, req, 1)->view_as({}, {}), - prob_buf->slice(0, req, 1)->view_as({dvoc}, {1}), - random_val, topp[req], topk[req], temperature[req]); - token_offset += seq_len; - } - RUN_INFINI(infinirtStreamSynchronize(stream)); - RUN_INFINI(infinirtMemcpy(result_cpu.data(), result_buf->data(), - sizeof(int64_t) * nreq, INFINIRT_MEMCPY_D2H)); - for (uint32_t req = 0; req < nreq; req++) { - output[req] = uint32_t(result_cpu[req]); - } - } - } -} - -__INFINI_C void -inferBatchDeepSeekV3(struct DeepSeekV3Model *model, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct DeepSeekV3Cache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output) { - model->req.tokens = tokens; - model->req.ntok = ntok; - model->req.req_lens = req_lens; - model->req.nreq = nreq; - model->req.req_pos = req_pos; - model->req.kv_caches = kv_caches; - model->req.output = output; - model->req.logits = nullptr; - model->req.temperature = temperature; - model->req.topk = topk; - model->req.topp = topp; - - for (size_t idev = 0; idev < model->dev_ids.size(); idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].proceed = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - for (size_t i = model->dev_ids.size(); i > 0; i--) { - auto idev = i - 1; - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].cv_done.wait(lock, [&] { return !(model->states[idev].proceed); }); - lock.unlock(); - } -} - -__INFINI_C void -forwardBatchDeepSeekV3(struct DeepSeekV3Model *model, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct DeepSeekV3Cache **kv_caches, - void *logits) { - model->req.tokens = tokens; - model->req.ntok = ntok; - model->req.req_lens = req_lens; - model->req.nreq = nreq; - model->req.req_pos = req_pos; - model->req.kv_caches = kv_caches; - model->req.output = nullptr; - model->req.logits = logits; - model->req.temperature = nullptr; - model->req.topk = nullptr; - model->req.topp = nullptr; - - for (size_t idev = 0; idev < model->dev_ids.size(); idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].proceed = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - for (size_t i = model->dev_ids.size(); i > 0; i--) { - auto idev = i - 1; - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].cv_done.wait(lock, [&] { return !(model->states[idev].proceed); }); - lock.unlock(); - } -} - -void launchDevice(const DeepSeekV3Meta &meta, std::shared_ptr weights, DeepSeekV3DeviceResource *rsrc, InferState &state, InferRequest &req, - infiniDevice_t device, int idev, int ndev, int dev_id, infinicclComm_t comm) { - // Create Device Resource - createDeviceResource(rsrc, &meta, weights, device, idev, ndev, dev_id, comm); - - CacheManager cache_manager(100); - InferenceContext ctx(rsrc->handle, rsrc->memory_pool, &cache_manager, rsrc->stream); - - // Set the inference context for this thread - setInferenceContext(&ctx); - - { - std::unique_lock lock(state.mtx); - state.loaded = true; - lock.unlock(); - state.cv_load.notify_one(); - } - - // Infer Loop - while (true) { - std::unique_lock lock(state.mtx); - state.cv_start.wait(lock, [&] { return state.proceed || state.exit_flag; }); - // quit if exit_flag is set - if (state.exit_flag) { - break; - } - - inferDeviceBatch(meta, *rsrc, idev, ndev, req.tokens, req.ntok, - req.req_lens, req.nreq, req.req_pos, req.kv_caches, - req.temperature, req.topk, req.topp, req.output, req.logits); - - state.proceed = false; - lock.unlock(); - state.cv_done.notify_one(); - } - - // Clean-Up - releaseDeviceResource(*rsrc); - setInferenceContext(nullptr); // Clear the context when done -} - -DeepSeekV3Model::DeepSeekV3Model(const DeepSeekV3Meta *_meta, const DeepSeekV3Weights *weights) : meta(*_meta) { - auto device_weights = weights->device_weights; - int ndev = device_weights.size(); - device = device_weights[0]->device; - dev_ids.resize(ndev); - for (int i = 0; i < ndev; i++) { - dev_ids[i] = device_weights[i]->dev_id; - } - dev_resources = std::vector(ndev); - states = std::vector(ndev); - threads.resize(ndev); - RUN_INFINI(infinirtInit()); - auto comms = std::vector(ndev, nullptr); - if (ndev > 1) { - RUN_INFINI(infinicclCommInitAll(device, comms.data(), ndev, dev_ids.data())); - } - for (int i = 0; i < ndev; i++) { - threads[i] = std::thread(launchDevice, std::cref(meta), device_weights[i], &dev_resources[i], std::ref(states[i]), std::ref(req), device, i, ndev, dev_ids[i], comms[i]); - } - for (int i = 0; i < ndev; i++) { - std::unique_lock lock(states[i].mtx); - states[i].cv_load.wait(lock, [&] { return states[i].loaded; }); - lock.unlock(); - } -} - -__INFINI_C struct DeepSeekV3Model * -createDeepSeekV3Model(const DeepSeekV3Meta *_meta, - const DeepSeekV3Weights *weights) { - DeepSeekV3Model *model = new DeepSeekV3Model(_meta, weights); - return model; -} - -__INFINI_C void -destroyDeepSeekV3Model(struct DeepSeekV3Model *model) { - auto ndev = model->dev_resources.size(); - - for (size_t idev = 0; idev < ndev; idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].exit_flag = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - - for (size_t idev = 0; idev < ndev; idev++) { - model->threads[idev].join(); - } - - delete model; -} diff --git a/src/models/deepseek_v3/deepseek_v3_cache.cpp b/src/models/deepseek_v3/deepseek_v3_cache.cpp deleted file mode 100644 index d6448a618..000000000 --- a/src/models/deepseek_v3/deepseek_v3_cache.cpp +++ /dev/null @@ -1,41 +0,0 @@ -#include "deepseek_v3_impl.hpp" - -__INFINI_C struct DeepSeekV3Cache * -createDeepSeekV3Cache(const struct DeepSeekV3Model *model) { - DeepSeekV3Cache *cache = new DeepSeekV3Cache(); - auto ndev = model->dev_resources.size(); - auto nlayer = model->meta.n_dense_layer + model->meta.n_sparse_layer; - auto max_len = model->meta.dctx; - auto d_rope = model->meta.d_rope; - auto r_kv = model->meta.r_kv; - auto kv_pass_shape = std::vector{max_len, r_kv}; - auto k_rot_shape = std::vector{max_len, d_rope}; - for (size_t idev = 0; idev < ndev; idev++) { - RUN_INFINI(infinirtSetDevice(model->device, model->dev_ids[idev])); - auto kv_pass_cache = std::vector>(); - auto k_rot_cache = std::vector>(); - for (size_t layer = 0; layer < nlayer; layer++) { - kv_pass_cache.push_back(Tensor::buffer(model->meta.dt_logits, kv_pass_shape)); - k_rot_cache.push_back(Tensor::buffer(model->meta.dt_logits, k_rot_shape)); - } - cache->kv_pass.push_back(kv_pass_cache); - cache->k_rot.push_back(k_rot_cache); - } - - return cache; -} - -__INFINI_C void -dropDeepSeekV3Cache(const struct DeepSeekV3Model *model, - struct DeepSeekV3Cache *cache) { - auto ndev = model->dev_resources.size(); - auto nlayer = model->meta.n_dense_layer + model->meta.n_sparse_layer; - for (size_t idev = 0; idev < ndev; idev++) { - RUN_INFINI(infinirtSetDevice(model->device, model->dev_ids[idev])); - for (size_t layer = 0; layer < nlayer; layer++) { - cache->kv_pass[idev][layer].reset(); - cache->k_rot[idev][layer].reset(); - } - } - delete cache; -} \ No newline at end of file diff --git a/src/models/deepseek_v3/deepseek_v3_impl.hpp b/src/models/deepseek_v3/deepseek_v3_impl.hpp deleted file mode 100644 index d47510746..000000000 --- a/src/models/deepseek_v3/deepseek_v3_impl.hpp +++ /dev/null @@ -1,116 +0,0 @@ -#ifndef DEEPSEEK_V3_IMPL_H -#define DEEPSEEK_V3_IMPL_H - -#include "infinicore_infer.h" - -#include "../../allocator.hpp" -#include "../../tensor.hpp" - -#include -#include -#include -#include -#include - -struct QuantLinearWeight { - std::shared_ptr w; - std::shared_ptr s; - std::shared_ptr z; -}; - -struct MLAWeight { - std::shared_ptr kv_a_norm, q_a_norm; - std::shared_ptr kv_a_proj, kv_b_proj, o_proj, q_a_proj, q_b_proj; -}; - -struct GateWeight { - std::shared_ptr w; - std::shared_ptr b; -}; - -struct MLPWeight { - std::shared_ptr gate, up, down; -}; - -struct LayerWeight { - std::shared_ptr mla_norm; - std::shared_ptr mla; - std::shared_ptr mlp_norm; - std::shared_ptr dense_mlp; - std::shared_ptr route; - std::shared_ptr share_expert; - std::vector> experts; -}; - -struct DeepSeekV3DeviceWeights { - std::shared_ptr w_in_embd, w_out_norm, w_out_embd, sin_table, - cos_table; - std::vector w_layers; - infiniDevice_t device; - int dev_id; - infinirtStream_t load_stream; -}; - -struct DeepSeekV3Weights { - std::vector> device_weights; - - DeepSeekV3Weights(const DeepSeekV3Meta *meta, - infiniDevice_t device, - int ndev, - const int *dev_ids); -}; - -struct DeepSeekV3DeviceResource { - // Device - infiniDevice_t device; - int device_id; - infiniopHandle_t handle; - // Weights - std::shared_ptr weights; - // Streams - infinirtStream_t stream; - // Communicator - infinicclComm_t comm; - - std::shared_ptr memory_pool; -}; - -struct InferState { - std::mutex mtx; - std::condition_variable cv_load, cv_start, cv_done; - bool loaded = false; - bool proceed = false; - bool exit_flag = false; -}; - -struct InferRequest { - const uint32_t *tokens; - uint32_t ntok; - const uint32_t *req_lens; - uint32_t nreq; - const uint32_t *req_pos; - struct DeepSeekV3Cache **kv_caches; - const float *temperature; - const uint32_t *topk; - const float *topp; - uint32_t *output; - void *logits; -}; - -struct DeepSeekV3Model { - DeepSeekV3Meta meta; - infiniDevice_t device; - std::vector dev_ids; - std::vector dev_resources; - std::vector states; - std::vector threads; - InferRequest req; - - DeepSeekV3Model(const DeepSeekV3Meta *, const DeepSeekV3Weights *weights); -}; - -struct DeepSeekV3Cache { - std::vector>> kv_pass, k_rot; -}; - -#endif diff --git a/src/models/deepseek_v3/deepseek_v3_weight.cpp b/src/models/deepseek_v3/deepseek_v3_weight.cpp deleted file mode 100644 index 20a8851de..000000000 --- a/src/models/deepseek_v3/deepseek_v3_weight.cpp +++ /dev/null @@ -1,451 +0,0 @@ -#include "deepseek_v3_impl.hpp" - -#include - -inline std::shared_ptr getInEmbd( - const DeepSeekV3Meta *meta) { - auto shape = std::vector({meta->dvoc, meta->d}); - return Tensor::weight(nullptr, meta->dt_logits, shape); -} - -inline std::shared_ptr getOutNorm( - const DeepSeekV3Meta *meta) { - auto shape = std::vector({meta->d}); - return Tensor::weight(nullptr, meta->dt_norm, shape); -} - -inline std::shared_ptr getOutEmbd( - const DeepSeekV3Meta *meta) { - - auto shape = std::vector({meta->dvoc, meta->d}); - return Tensor::weight(nullptr, meta->dt_logits, shape) - ->permute({1, 0}); -} - -inline std::shared_ptr getMLANorm( - const DeepSeekV3Meta *meta) { - auto shape = std::vector({meta->d}); - return Tensor::weight(nullptr, meta->dt_norm, shape); -} - -inline std::shared_ptr getQuantLinear( - const DeepSeekV3Meta *meta, size_t in_dim, size_t out_dim) { - auto qw = std::make_shared(); - auto shape_w = std::vector({in_dim, out_dim / 8}); - qw->w = Tensor::weight(nullptr, INFINI_DTYPE_I32, shape_w); - qw->s = Tensor::weight(nullptr, meta->dt_quant_scale, {in_dim / 64, out_dim}); - qw->z = Tensor::weight(nullptr, INFINI_DTYPE_I32, {in_dim / 64, out_dim / 8}); - return qw; -} - -// ------------------- MLA Weights ------------------- -inline std::shared_ptr getMLPNorm( - const DeepSeekV3Meta *meta) { - auto shape = std::vector({meta->d}); - return Tensor::weight(nullptr, meta->dt_norm, shape); -} - -inline std::shared_ptr getMLA(const DeepSeekV3Meta *meta, int ndev) { - auto mla = std::make_shared(); - - mla->q_a_proj = getQuantLinear(meta, meta->d, meta->r_q); - mla->q_a_norm = Tensor::weight(nullptr, meta->dt_norm, {meta->r_q}); - mla->q_b_proj = getQuantLinear(meta, meta->r_q, meta->nh / ndev * meta->d_qk); - - mla->kv_a_proj = getQuantLinear(meta, meta->d, meta->r_kv + meta->d_rope); - mla->kv_a_norm = Tensor::weight(nullptr, meta->dt_norm, {meta->r_kv}); - mla->kv_b_proj = getQuantLinear(meta, meta->r_kv, meta->nh / ndev * (meta->d_nope + meta->d_v)); - - mla->o_proj = getQuantLinear(meta, meta->nh / ndev * meta->d_v, meta->d); - return mla; -} - -// ------------------- Dense MLP ------------------- - -inline std::shared_ptr getMLP(const DeepSeekV3Meta *meta, size_t d, size_t di) { - auto mlp = std::make_shared(); - mlp->gate = getQuantLinear(meta, d, di); - mlp->up = getQuantLinear(meta, d, di); - mlp->down = getQuantLinear(meta, di, d); - return mlp; -} - -inline std::shared_ptr getDenseMLP(const DeepSeekV3Meta *meta, int ndev) { - return getMLP(meta, meta->d, meta->di / ndev); -} - -// ------------------- Sparse Route + Experts ------------------- - -inline std::shared_ptr getRouteWeight( - const DeepSeekV3Meta *meta) { - auto gw = std::make_shared(); - gw->w = Tensor::weight(nullptr, meta->dt_gate_weight, {meta->nexperts, meta->d})->permute({1, 0}); - gw->b = Tensor::weight(nullptr, meta->dt_gate_bias, {meta->nexperts}); - return gw; -} - -inline std::shared_ptr getShareExpert(const DeepSeekV3Meta *meta, int ndev) { - return getMLP(meta, meta->d, meta->di_moe / ndev); -} - -inline std::vector> getExperts(const DeepSeekV3Meta *meta, int ndev) { - std::vector> experts(meta->nexperts); - for (size_t i = 0; i < meta->nexperts; i++) { - experts[i] = getMLP(meta, meta->d, meta->di_moe / ndev); - } - return experts; -} - -inline std::shared_ptr getSinTable(const DeepSeekV3Meta *meta) { - auto half_dh = meta->d_rope / 2; - auto unit = dsize(meta->dt_logits); - void *table = std::malloc(meta->dctx * half_dh * unit); - - for (size_t i = 0; i < meta->dctx; i++) { - for (size_t j = 0; j < half_dh; j++) { - float _sin = std::sin( - static_cast(i) / std::pow(meta->rope_theta, static_cast(j) / half_dh)); - if (meta->dt_logits == INFINI_DTYPE_F16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_f16(_sin); - } else if (meta->dt_logits == INFINI_DTYPE_BF16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_bf16(_sin); - } else if (meta->dt_logits == INFINI_DTYPE_F32) { - ((float *)table)[i * half_dh + j] = _sin; - } else { - std::cout << "unsupported data type" << std::endl; - exit(1); - } - } - } - auto shape = std::vector({meta->dctx, half_dh}); - auto tensor = Tensor::weight(table, meta->dt_logits, shape); - std::free(table); - return tensor; -} - -inline std::shared_ptr getCosTable(const DeepSeekV3Meta *meta) { - auto half_dh = meta->d_rope / 2; - auto unit = dsize(meta->dt_logits); - void *table = std::malloc(meta->dctx * half_dh * unit); - - for (size_t i = 0; i < meta->dctx; i++) { - for (size_t j = 0; j < half_dh; j++) { - float _cos = std::cos( - static_cast(i) / std::pow(meta->rope_theta, static_cast(j) / half_dh)); - if (meta->dt_logits == INFINI_DTYPE_F16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_f16(_cos); - } else if (meta->dt_logits == INFINI_DTYPE_BF16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_bf16(_cos); - } else if (meta->dt_logits == INFINI_DTYPE_F32) { - ((float *)table)[i * half_dh + j] = _cos; - } else { - std::cout << "unsupported data type" << std::endl; - exit(1); - } - } - } - auto shape = std::vector({meta->dctx, half_dh}); - auto tensor = Tensor::weight(table, meta->dt_logits, shape); - std::free(table); - return tensor; -} - -DeepSeekV3Weights::DeepSeekV3Weights( - const DeepSeekV3Meta *meta, infiniDevice_t device, int ndev, const int *dev_ids) { - device_weights = std::vector>(ndev); - for (int dev = 0; dev < ndev; dev++) { - int dev_id = dev_ids[dev]; - RUN_INFINI(infinirtSetDevice(device, dev_id)); - device_weights[dev] = std::make_shared(); - device_weights[dev]->device = device; - device_weights[dev]->dev_id = dev_id; - RUN_INFINI(infinirtStreamCreate(&device_weights[dev]->load_stream)); - - device_weights[dev]->w_in_embd = getInEmbd(meta); - device_weights[dev]->w_out_norm = getOutNorm(meta); - device_weights[dev]->w_out_embd = getOutEmbd(meta); - device_weights[dev]->sin_table = getSinTable(meta); - device_weights[dev]->cos_table = getCosTable(meta); - - device_weights[dev]->w_layers = std::vector(meta->n_dense_layer + meta->n_sparse_layer); - - for (size_t layer = 0; layer < meta->n_dense_layer + meta->n_sparse_layer; layer++) { - device_weights[dev]->w_layers[layer].mla_norm = getMLANorm(meta); - device_weights[dev]->w_layers[layer].mla = getMLA(meta, ndev); - device_weights[dev]->w_layers[layer].mlp_norm = getMLPNorm(meta); - if (layer < meta->n_dense_layer) { - device_weights[dev]->w_layers[layer].dense_mlp = getDenseMLP(meta, ndev); - } else { - device_weights[dev]->w_layers[layer].route = getRouteWeight(meta); - device_weights[dev]->w_layers[layer].share_expert = getShareExpert(meta, ndev); - device_weights[dev]->w_layers[layer].experts = getExperts(meta, ndev); - } - } - } -} - -// --- Global -void load_input_embd(DeepSeekV3Weights *weights, void *cpu_ptr) { - std::cout << "Loading input embedding from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_in_embd->load(cpu_ptr, weight->load_stream); - } -} - -void load_output_norm(DeepSeekV3Weights *weights, void *cpu_ptr) { - std::cout << "Loading output norm from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_out_norm->load(cpu_ptr, weight->load_stream); - } -} - -void load_output_embd(DeepSeekV3Weights *weights, void *cpu_ptr) { - std::cout << "Loading output embedding from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_out_embd->load(cpu_ptr, weight->load_stream); - } -} - -// --- Attention -void load_attn_norm(DeepSeekV3Weights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading attention norm " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_layers[layer].mla_norm->load(cpu_ptr, weight->load_stream); - } -} - -void load_attn_q_a_proj(DeepSeekV3Weights *weights, - void *weight_ptr, void *scale_ptr, void *zero_ptr, size_t layer) { - std::cout << "Loading attention q_a_proj " << layer << " from " << weight_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_layers[layer].mla->q_a_proj->w->load(weight_ptr, weight->load_stream); - weight->w_layers[layer].mla->q_a_proj->s->load(scale_ptr, weight->load_stream); - weight->w_layers[layer].mla->q_a_proj->z->load(zero_ptr, weight->load_stream); - } -} - -void load_attn_q_a_layernorm(DeepSeekV3Weights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading attention q_a_layernorm " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_layers[layer].mla->q_a_norm->load(cpu_ptr, weight->load_stream); - } -} - -inline void load_dist_linear(void *w_ptr, void *s_ptr, void *z_ptr, std::shared_ptr w, std::shared_ptr s, std::shared_ptr z, size_t ndev, size_t dev, infinirtStream_t stream) { - auto w_offset = w->shape()[0] * w->shape()[1] / ndev * dev * dsize(w->dtype()); - auto s_offset = s->shape()[0] * s->shape()[1] / ndev * dev * dsize(s->dtype()); - auto z_offset = z->shape()[0] * z->shape()[1] / ndev * dev * dsize(z->dtype()); - w->load(reinterpret_cast(w_ptr) + w_offset, stream); - s->load(reinterpret_cast(s_ptr) + s_offset, stream); - z->load(reinterpret_cast(z_ptr) + z_offset, stream); -} - -void load_attn_q_b_proj(DeepSeekV3Weights *weights, - void *weight_ptr, void *scale_ptr, void *zero_ptr, size_t layer) { - std::cout << "Loading attention q_b_proj " << layer << " from " << weight_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - auto w = weight->w_layers[layer].mla->q_b_proj->w; - auto s = weight->w_layers[layer].mla->q_b_proj->s; - auto z = weight->w_layers[layer].mla->q_b_proj->z; - load_dist_linear(weight_ptr, scale_ptr, zero_ptr, w, s, z, weights->device_weights.size(), dev, weight->load_stream); - } -} - -void load_attn_kv_a_proj_with_mqa(DeepSeekV3Weights *weights, - void *weight_ptr, void *scale_ptr, void *zero_ptr, size_t layer) { - std::cout << "Loading attention kv_a_proj_with_mqa " << layer << " from " << weight_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_layers[layer].mla->kv_a_proj->w->load(weight_ptr, weight->load_stream); - weight->w_layers[layer].mla->kv_a_proj->s->load(scale_ptr, weight->load_stream); - weight->w_layers[layer].mla->kv_a_proj->z->load(zero_ptr, weight->load_stream); - } -} - -void load_attn_kv_a_layernorm(DeepSeekV3Weights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading attention kv_a_layernorm " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_layers[layer].mla->kv_a_norm->load(cpu_ptr, weight->load_stream); - } -} - -void load_attn_kv_b_proj(DeepSeekV3Weights *weights, - void *weight_ptr, void *scale_ptr, void *zero_ptr, size_t layer) { - std::cout << "Loading attention kv_b_proj " << layer << " from " << weight_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - auto w = weight->w_layers[layer].mla->kv_b_proj->w; - auto s = weight->w_layers[layer].mla->kv_b_proj->s; - auto z = weight->w_layers[layer].mla->kv_b_proj->z; - load_dist_linear(weight_ptr, scale_ptr, zero_ptr, w, s, z, weights->device_weights.size(), dev, weight->load_stream); - } -} - -void load_attn_o_proj(DeepSeekV3Weights *weights, - void *weight_ptr, void *scale_ptr, void *zero_ptr, size_t layer) { - std::cout << "Loading attention o_proj " << layer << " from " << weight_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - auto w = weight->w_layers[layer].mla->o_proj->w; - auto s = weight->w_layers[layer].mla->o_proj->s; - auto z = weight->w_layers[layer].mla->o_proj->z; - load_dist_linear(weight_ptr, scale_ptr, zero_ptr, w, s, z, weights->device_weights.size(), dev, weight->load_stream); - } -} - -// --- MLP -void load_mlp_norm(DeepSeekV3Weights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading mlp norm " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_layers[layer].mlp_norm->load(cpu_ptr, weight->load_stream); - } -} - -void load_mlp_dense(DeepSeekV3Weights *weights, - void *gate_weight_ptr, void *gate_scale_ptr, void *gate_zero_ptr, - void *up_weight_ptr, void *up_scale_ptr, void *up_zero_ptr, - void *down_weight_ptr, void *down_scale_ptr, void *down_zero_ptr, - size_t layer_id) { - std::cout << "Loading mlp dense " << layer_id << " from " << gate_weight_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - auto gate_w = weight->w_layers[layer_id].dense_mlp->gate->w; - auto gate_s = weight->w_layers[layer_id].dense_mlp->gate->s; - auto gate_z = weight->w_layers[layer_id].dense_mlp->gate->z; - auto up_w = weight->w_layers[layer_id].dense_mlp->up->w; - auto up_s = weight->w_layers[layer_id].dense_mlp->up->s; - auto up_z = weight->w_layers[layer_id].dense_mlp->up->z; - auto down_w = weight->w_layers[layer_id].dense_mlp->down->w; - auto down_s = weight->w_layers[layer_id].dense_mlp->down->s; - auto down_z = weight->w_layers[layer_id].dense_mlp->down->z; - load_dist_linear(gate_weight_ptr, gate_scale_ptr, gate_zero_ptr, gate_w, gate_s, gate_z, weights->device_weights.size(), dev, weight->load_stream); - load_dist_linear(up_weight_ptr, up_scale_ptr, up_zero_ptr, up_w, up_s, up_z, weights->device_weights.size(), dev, weight->load_stream); - load_dist_linear(down_weight_ptr, down_scale_ptr, down_zero_ptr, down_w, down_s, down_z, weights->device_weights.size(), dev, weight->load_stream); - } -} - -void load_mlp_gate_weight(DeepSeekV3Weights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading mlp gate weight " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_layers[layer].route->w->load(cpu_ptr, weight->load_stream); - } -} - -void load_mlp_gate_bias(DeepSeekV3Weights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading mlp gate bias " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_layers[layer].route->b->load(cpu_ptr, weight->load_stream); - } -} - -void load_mlp_shared_experts(DeepSeekV3Weights *weights, - void *gate_weight_ptr, void *gate_scale_ptr, void *gate_zero_ptr, - void *up_weight_ptr, void *up_scale_ptr, void *up_zero_ptr, - void *down_weight_ptr, void *down_scale_ptr, void *down_zero_ptr, - size_t layer_id) { - std::cout << "Loading mlp shared experts " << layer_id << " from " << gate_weight_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - auto gate_w = weight->w_layers[layer_id].share_expert->gate->w; - auto gate_s = weight->w_layers[layer_id].share_expert->gate->s; - auto gate_z = weight->w_layers[layer_id].share_expert->gate->z; - auto up_w = weight->w_layers[layer_id].share_expert->up->w; - auto up_s = weight->w_layers[layer_id].share_expert->up->s; - auto up_z = weight->w_layers[layer_id].share_expert->up->z; - auto down_w = weight->w_layers[layer_id].share_expert->down->w; - auto down_s = weight->w_layers[layer_id].share_expert->down->s; - auto down_z = weight->w_layers[layer_id].share_expert->down->z; - load_dist_linear(gate_weight_ptr, gate_scale_ptr, gate_zero_ptr, gate_w, gate_s, gate_z, weights->device_weights.size(), dev, weight->load_stream); - load_dist_linear(up_weight_ptr, up_scale_ptr, up_zero_ptr, up_w, up_s, up_z, weights->device_weights.size(), dev, weight->load_stream); - load_dist_linear(down_weight_ptr, down_scale_ptr, down_zero_ptr, down_w, down_s, down_z, weights->device_weights.size(), dev, weight->load_stream); - } -} - -void load_mlp_experts(DeepSeekV3Weights *weights, - void *gate_weight_ptr, void *gate_scale_ptr, void *gate_zero_ptr, - void *up_weight_ptr, void *up_scale_ptr, void *up_zero_ptr, - void *down_weight_ptr, void *down_scale_ptr, void *down_zero_ptr, - size_t layer_id, size_t expert_id) { - std::cout << "Loading mlp expert " << layer_id << " expert " << expert_id - << " from " << gate_weight_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - auto gate_w = weight->w_layers[layer_id].experts[expert_id]->gate->w; - auto gate_s = weight->w_layers[layer_id].experts[expert_id]->gate->s; - auto gate_z = weight->w_layers[layer_id].experts[expert_id]->gate->z; - auto up_w = weight->w_layers[layer_id].experts[expert_id]->up->w; - auto up_s = weight->w_layers[layer_id].experts[expert_id]->up->s; - auto up_z = weight->w_layers[layer_id].experts[expert_id]->up->z; - auto down_w = weight->w_layers[layer_id].experts[expert_id]->down->w; - auto down_s = weight->w_layers[layer_id].experts[expert_id]->down->s; - auto down_z = weight->w_layers[layer_id].experts[expert_id]->down->z; - load_dist_linear(gate_weight_ptr, gate_scale_ptr, gate_zero_ptr, gate_w, gate_s, gate_z, weights->device_weights.size(), dev, weight->load_stream); - load_dist_linear(up_weight_ptr, up_scale_ptr, up_zero_ptr, up_w, up_s, up_z, weights->device_weights.size(), dev, weight->load_stream); - load_dist_linear(down_weight_ptr, down_scale_ptr, down_zero_ptr, down_w, down_s, down_z, weights->device_weights.size(), dev, weight->load_stream); - } -} - -static DeepSeekV3WeightLoader weight_loader = { - // Global - .load_input_embd = load_input_embd, - .load_output_norm = load_output_norm, - .load_output_embd = load_output_embd, - // Attention - .load_attn_norm = load_attn_norm, - .load_attn_q_a_proj = load_attn_q_a_proj, - .load_attn_q_a_layernorm = load_attn_q_a_layernorm, - .load_attn_q_b_proj = load_attn_q_b_proj, - .load_attn_kv_a_proj_with_mqa = load_attn_kv_a_proj_with_mqa, - .load_attn_kv_a_layernorm = load_attn_kv_a_layernorm, - .load_attn_kv_b_proj = load_attn_kv_b_proj, - .load_attn_o_proj = load_attn_o_proj, - // MLP - .load_mlp_norm = load_mlp_norm, - .load_mlp_dense = load_mlp_dense, - .load_mlp_gate_weight = load_mlp_gate_weight, - .load_mlp_gate_bias = load_mlp_gate_bias, - .load_mlp_shared_experts = load_mlp_shared_experts, - .load_mlp_experts = load_mlp_experts, -}; - -__INFINI_C DeepSeekV3Weights * -createDeepSeekV3Weights(const DeepSeekV3Meta *meta, - infiniDevice_t device, - int ndev, - const int *dev_ids) { - auto weights = new DeepSeekV3Weights(meta, device, ndev, dev_ids); - return weights; -}; - -__INFINI_C DeepSeekV3WeightLoader * -createDeepSeekV3WeightLoader() { - return &weight_loader; -} diff --git a/src/models/inference_context.cpp b/src/models/inference_context.cpp deleted file mode 100644 index 0a0e6f49a..000000000 --- a/src/models/inference_context.cpp +++ /dev/null @@ -1,371 +0,0 @@ -#include "inference_context.hpp" -#include "../tensor.hpp" -#include "../utils.hpp" - -InferenceContext::InferenceContext(infiniopHandle_t op_handle_, std::shared_ptr memory_pool_, CacheManager *cache_manager, infinirtStream_t stream) - : op_handle(op_handle_), memory_pool(memory_pool_), cache_manager(cache_manager), stream(stream) {} - -void InferenceContext::ensure_workspace(size_t required_size) { - if (required_size > current_workspace_size || !workspace_storage) { - workspace_storage = Storage::createFromPool(required_size, memory_pool); - current_workspace_size = required_size; - } -} - -void InferenceContext::add(std::shared_ptr c, - std::shared_ptr a, - std::shared_ptr b) { - size_t key = CacheManager::createDescriptorKey(c, a, b); - - infiniopAddDescriptor_t desc; - if (!cache_manager->getAddDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateAddDescriptor(op_handle, &desc, c->desc(), a->desc(), b->desc())); - cache_manager->putAddDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetAddWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopAdd( - desc, workspace, workspace_size, - c->data(), a->data(), b->data(), stream)); -} - -void InferenceContext::conv(std::shared_ptr y, - std::shared_ptr x, - std::shared_ptr w, - std::shared_ptr bias, - void *pads, - void *strides, - void *dilations, - size_t n) { - size_t key = CacheManager::createDescriptorKey(y, x, w, bias); - // Combine additional parameters into the key for unique identification - hash_combine(key, std::hash()(pads)); - hash_combine(key, std::hash()(strides)); - hash_combine(key, std::hash()(dilations)); - hash_combine(key, std::hash()(n)); - - infiniopConvDescriptor_t desc; - if (!cache_manager->getConvDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateConvDescriptor( - op_handle, &desc, y->desc(), x->desc(), w->desc(), - bias ? bias->desc() : nullptr, pads, strides, dilations, n)); - cache_manager->putConvDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetConvWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopConv( - desc, workspace, workspace_size, - y->data(), x->data(), w->data(), - bias ? bias->data() : nullptr, stream)); -} - -void InferenceContext::mul(std::shared_ptr c, - std::shared_ptr a, - std::shared_ptr b) { - size_t key = CacheManager::createDescriptorKey(c, a, b); - - infiniopMulDescriptor_t desc; - if (!cache_manager->getMulDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateMulDescriptor(op_handle, &desc, c->desc(), a->desc(), b->desc())); - cache_manager->putMulDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetMulWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopMul( - desc, workspace, workspace_size, - c->data(), a->data(), b->data(), stream)); -} - -void InferenceContext::rmsnorm(std::shared_ptr y, - std::shared_ptr x, - std::shared_ptr w, - float epsilon) { - size_t key = CacheManager::createDescriptorKey(y, x, w); - - infiniopRMSNormDescriptor_t desc; - if (!cache_manager->getRMSNormDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateRMSNormDescriptor( - op_handle, &desc, y->desc(), x->desc(), w->desc(), epsilon)); - cache_manager->putRMSNormDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetRMSNormWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopRMSNorm( - desc, workspace, workspace_size, - y->data(), x->data(), w->data(), stream)); -} - -void InferenceContext::gemm(std::shared_ptr c, - std::shared_ptr a, - std::shared_ptr b, - float alpha, float beta) { - size_t key = CacheManager::createDescriptorKey(c, a, b); - - infiniopGemmDescriptor_t desc; - if (!cache_manager->getGemmDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateGemmDescriptor(op_handle, &desc, c->desc(), a->desc(), b->desc())); - cache_manager->putGemmDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetGemmWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopGemm( - desc, workspace, workspace_size, - c->data(), a->data(), b->data(), alpha, beta, stream)); -} - -void InferenceContext::rearrange(std::shared_ptr dst, - std::shared_ptr src) { - size_t key = CacheManager::createDescriptorKey(dst, src); - - infiniopRearrangeDescriptor_t desc; - if (!cache_manager->getRearrangeDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateRearrangeDescriptor(op_handle, &desc, dst->desc(), src->desc())); - cache_manager->putRearrangeDescriptor(key, desc); - } - - RUN_INFINI(infiniopRearrange( - desc, - dst->data(), - src->data(), - stream)); -} - -void InferenceContext::rope(std::shared_ptr q, - std::shared_ptr k, - std::shared_ptr pos, - std::shared_ptr sin, - std::shared_ptr cos, - infiniopRoPEAlgo_t algo) { - size_t key = CacheManager::createDescriptorKey(q, k, pos, sin, cos); - hash_combine(key, std::hash()(algo)); - - infiniopRoPEDescriptor_t desc; - if (!cache_manager->getRoPEDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateRoPEDescriptor( - op_handle, &desc, q->desc(), k->desc(), - pos->desc(), sin->desc(), cos->desc(), algo)); - cache_manager->putRoPEDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetRoPEWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopRoPE( - desc, workspace, workspace_size, - q->data(), k->data(), pos->data(), - sin->data(), cos->data(), stream)); -} - -void InferenceContext::causalSoftmax(std::shared_ptr y, - std::shared_ptr x) { - size_t key = CacheManager::createDescriptorKey(y, x); - - infiniopCausalSoftmaxDescriptor_t desc; - if (!cache_manager->getCausalSoftmaxDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateCausalSoftmaxDescriptor( - op_handle, &desc, y->desc(), x->desc())); - cache_manager->putCausalSoftmaxDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetCausalSoftmaxWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopCausalSoftmax(desc, workspace, workspace_size, - y->data(), x->data(), stream)); -} - -void InferenceContext::topkrouter(std::shared_ptr values, // F32 - std::shared_ptr indices, // I32 - std::shared_ptr x, - std::shared_ptr correction_bias, // F32 - float routed_scaling_factor, - size_t topk) { - size_t key = CacheManager::createDescriptorKey(values, indices, x, correction_bias); - - infiniopTopkrouterDescriptor_t desc; - if (!cache_manager->getTopkrouterDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateTopkrouterDescriptor( - op_handle, &desc, x->desc(), correction_bias->desc())); - cache_manager->putTopkrouterDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetTopkrouterWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopTopkrouter(desc, workspace, workspace_size, - values->data(), indices->data(), x->data(), correction_bias->data(), - routed_scaling_factor, topk, stream)); -} - -void InferenceContext::swiglu(std::shared_ptr out, - std::shared_ptr up, - std::shared_ptr gate) { - size_t key = CacheManager::createDescriptorKey(out, up, gate); - - infiniopSwiGLUDescriptor_t desc; - if (!cache_manager->getSwiGLUDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateSwiGLUDescriptor( - op_handle, &desc, out->desc(), up->desc(), gate->desc())); - cache_manager->putSwiGLUDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetSwiGLUWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopSwiGLU(desc, workspace, workspace_size, - out->data(), up->data(), gate->data(), stream)); -} - -void InferenceContext::silu(std::shared_ptr out, - std::shared_ptr input) { - size_t key = CacheManager::createDescriptorKey(out, input); - - infiniopSiluDescriptor_t desc; - if (!cache_manager->getSiluDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateSiluDescriptor( - op_handle, &desc, out->desc(), input->desc())); - cache_manager->putSiluDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetSiluWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopSilu(desc, workspace, workspace_size, - out->data(), input->data(), stream)); -} - -void InferenceContext::randomSample(std::shared_ptr out, - std::shared_ptr prob, - float random_val, float top_p, uint32_t top_k, float temperature) { - size_t key = CacheManager::createDescriptorKey(out, prob); - - infiniopRandomSampleDescriptor_t desc; - if (!cache_manager->getRandomSampleDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateRandomSampleDescriptor( - op_handle, &desc, out->desc(), prob->desc())); - cache_manager->putRandomSampleDescriptor(key, desc); - } - - size_t workspace_size = 0; - RUN_INFINI(infiniopGetRandomSampleWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - void *workspace = workspace_storage->memory(); - - RUN_INFINI(infiniopRandomSample( - desc, workspace, workspace_size, - out->data(), prob->data(), - random_val, top_p, top_k, temperature, - stream)); -} - -void InferenceContext::linear(std::shared_ptr c, - std::shared_ptr a, - std::shared_ptr b, - float alpha, float beta, - std::shared_ptr residual, - std::shared_ptr bias) { - bool residual_flag = residual != nullptr; - - if (bias && !residual) { - int ndim_diff = c->ndim() - 1; - ASSERT_EQ(bias->ndim(), 1); - ASSERT_EQ(bias->shape()[0], c->shape()[ndim_diff]); - std::vector strides(ndim_diff, 0); - strides.push_back(bias->strides()[0]); - rearrange(c, bias->view_as(c->shape(), strides)); - residual = c; - } - - if (residual) { - if (residual->data() == c->data()) { - if (beta == 0.0) { - gemm(c, a, b, alpha, 1.0); - } else { - auto c_copy = Tensor::buffer(c->dtype(), c->shape(), memory_pool); - c_copy->copyFrom(c, op_handle, stream); - gemm(c, a, b, alpha, beta); - add(c, c, c_copy); - } - } else { - gemm(c, a, b, alpha, beta); - add(c, c, residual); - } - } else { - gemm(c, a, b, alpha, beta); - } - - if (bias && residual_flag) { - int ndim_diff = c->ndim() - 1; - ASSERT_EQ(bias->ndim(), 1); - ASSERT_EQ(bias->shape()[0], c->shape()[ndim_diff]); - std::vector strides(ndim_diff, 0); - strides.push_back(bias->strides()[0]); - add(c, c, bias->view_as(c->shape(), strides)); - } -} - -void InferenceContext::dequant(std::shared_ptr weight, - std::shared_ptr in_w, - std::shared_ptr in_s, - std::shared_ptr in_z, - QuantType type, - std::shared_ptr in_g_idx) { - size_t key = CacheManager::createDescriptorKey(weight, in_w, in_s, in_z, in_g_idx); - if (type == QuantType::AWQ) { - // unchanged - infiniopDequantizeAWQDescriptor_t desc; - if (!cache_manager->getDequantizeAWQDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateDequantizeAWQDescriptor(op_handle, &desc, - weight->desc(), in_w->desc(), in_s->desc(), in_z->desc())); - cache_manager->putDequantizeAWQDescriptor(key, desc); - } - size_t workspace_size = 0; - RUN_INFINI(infiniopGetDequantizeAWQWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - RUN_INFINI(infiniopDequantizeAWQ(desc, workspace_storage->memory(), workspace_size, - weight->data(), in_w->data(), in_s->data(), in_z->data(), stream)); - } else if (type == QuantType::GPTQ) { - ASSERT(in_g_idx && "GPTQ dequant requires g_idx"); - infiniopDequantizeGPTQDescriptor_t desc; - if (!cache_manager->getDequantizeGPTQDescriptor(key, desc)) { - RUN_INFINI(infiniopCreateDequantizeGPTQDescriptor(op_handle, &desc, - weight->desc(), in_w->desc(), in_s->desc(), in_z->desc(), in_g_idx->desc())); - cache_manager->putDequantizeGPTQDescriptor(key, desc); - } - size_t workspace_size = 0; - RUN_INFINI(infiniopGetDequantizeGPTQWorkspaceSize(desc, &workspace_size)); - ensure_workspace(workspace_size); - RUN_INFINI(infiniopDequantizeGPTQ(desc, workspace_storage->memory(), workspace_size, - weight->data(), in_w->data(), in_s->data(), in_z->data(), in_g_idx->data(), stream)); - } -} diff --git a/src/models/inference_context.hpp b/src/models/inference_context.hpp deleted file mode 100644 index 07b2ed20b..000000000 --- a/src/models/inference_context.hpp +++ /dev/null @@ -1,182 +0,0 @@ -#pragma once - -#include "../cache_manager/opcache_manager.hpp" - -#include - -enum class QuantType { - AWQ, - GPTQ -}; - -struct InferenceContext { - infiniopHandle_t op_handle; - std::shared_ptr memory_pool; - CacheManager *cache_manager; - infinirtStream_t stream; - std::shared_ptr workspace_storage; - size_t current_workspace_size = 0; - - InferenceContext(infiniopHandle_t op_handle, std::shared_ptr memory_pool, CacheManager *cache_manager, infinirtStream_t stream); - - void ensure_workspace(size_t required_size); - - void add(std::shared_ptr c, - std::shared_ptr a, - std::shared_ptr b); - void conv(std::shared_ptr y, - std::shared_ptr x, - std::shared_ptr w, - std::shared_ptr bias, - void *pads, void *strides, void *dilations, size_t n); - void mul(std::shared_ptr c, - std::shared_ptr a, - std::shared_ptr b); - void rmsnorm(std::shared_ptr y, - std::shared_ptr x, - std::shared_ptr w, - float epsilon); - void gemm(std::shared_ptr c, - std::shared_ptr a, - std::shared_ptr b, - float alpha, float beta); - void rearrange(std::shared_ptr dst, - std::shared_ptr src); - void rope(std::shared_ptr q, - std::shared_ptr k, - std::shared_ptr pos, - std::shared_ptr sin, - std::shared_ptr cos, - infiniopRoPEAlgo_t algo); - void causalSoftmax(std::shared_ptr y, - std::shared_ptr x); - - void topkrouter(std::shared_ptr values, // F32 - std::shared_ptr indices, // I32 - std::shared_ptr x, - std::shared_ptr correction_bias, // F32 - float routed_scaling_factor, - size_t topk); - - void swiglu(std::shared_ptr out, - std::shared_ptr up, - std::shared_ptr gate); - void silu(std::shared_ptr out, - std::shared_ptr input); - void randomSample(std::shared_ptr out, - std::shared_ptr prob, - float random_val, float top_p, uint32_t top_k, float temperature); - - void linear(std::shared_ptr c, - std::shared_ptr a, - std::shared_ptr b, - float alpha, float beta, - std::shared_ptr residual, - std::shared_ptr bias); - void dequant(std::shared_ptr weight, - std::shared_ptr in_w, - std::shared_ptr in_s, - std::shared_ptr in_z, - QuantType type, - std::shared_ptr in_g_idx = nullptr); -}; - -namespace { -thread_local InferenceContext *tls_inference_context = nullptr; -} - -inline InferenceContext &getInferenceContext() { - assert(tls_inference_context != nullptr && "InferenceContext not set for this thread"); - return *tls_inference_context; -} - -inline void setInferenceContext(InferenceContext *ctx) { - tls_inference_context = ctx; -} - -inline void add(std::shared_ptr c, std::shared_ptr a, std::shared_ptr b) { - getInferenceContext().add(c, a, b); -} - -inline void conv(std::shared_ptr y, std::shared_ptr x, std::shared_ptr w, std::shared_ptr bias, - void *pads, void *strides, void *dilations, size_t n) { - getInferenceContext().conv(y, x, w, bias, pads, strides, dilations, n); -} - -inline void mul(std::shared_ptr c, std::shared_ptr a, std::shared_ptr b) { - getInferenceContext().mul(c, a, b); -} - -inline void rmsnorm(std::shared_ptr y, std::shared_ptr x, - std::shared_ptr w, float epsilon) { - getInferenceContext().rmsnorm(y, x, w, epsilon); -} - -inline void gemm(std::shared_ptr c, std::shared_ptr a, - std::shared_ptr b, float alpha, float beta) { - getInferenceContext().gemm(c, a, b, alpha, beta); -} - -inline void rearrange(std::shared_ptr dst, std::shared_ptr src) { - getInferenceContext().rearrange(dst, src); -} - -inline void rope(std::shared_ptr q, std::shared_ptr k, - std::shared_ptr pos, std::shared_ptr sin, - std::shared_ptr cos) { - getInferenceContext().rope(q, k, pos, sin, cos, INFINIOP_ROPE_ALGO_GPT_J); -} - -inline void rope_v2(std::shared_ptr q, std::shared_ptr k, - std::shared_ptr pos, std::shared_ptr sin, - std::shared_ptr cos) { - getInferenceContext().rope(q, k, pos, sin, cos, INFINIOP_ROPE_ALGO_GPT_NEOX); -} - -inline void causalSoftmax(std::shared_ptr y, std::shared_ptr x) { - getInferenceContext().causalSoftmax(y, x); -} - -inline void topkrouter(std::shared_ptr values, // F32 - std::shared_ptr indices, // I32 - std::shared_ptr x, - std::shared_ptr correction_bias, // F32 - float routed_scaling_factor, - size_t topk) { - - getInferenceContext().topkrouter(values, // F32 - indices, // I32 - x, - correction_bias, // F32 - routed_scaling_factor, - topk); -} - -inline void swiglu(std::shared_ptr out, std::shared_ptr up, - std::shared_ptr gate) { - getInferenceContext().swiglu(out, up, gate); -} - -inline void silu(std::shared_ptr out, std::shared_ptr input) { - getInferenceContext().silu(out, input); -} - -inline void randomSample(std::shared_ptr out, std::shared_ptr prob, - float random_val, float top_p, uint32_t top_k, float temperature) { - getInferenceContext().randomSample(out, prob, random_val, top_p, top_k, temperature); -} - -inline void linear(std::shared_ptr c, std::shared_ptr a, - std::shared_ptr b, float alpha, float beta, - std::shared_ptr residual, std::shared_ptr bias) { - getInferenceContext().linear(c, a, b, alpha, beta, residual, bias); -} - -inline void dequant_linear(std::shared_ptr out, std::shared_ptr x, - std::shared_ptr w_w, std::shared_ptr w_s, std::shared_ptr w_z, - float alpha, float beta, std::shared_ptr residual, std::shared_ptr bias, - QuantType type = QuantType::AWQ, std::shared_ptr w_g_idx = nullptr) { - auto w = Tensor::buffer(x->dtype(), {x->shape()[1], out->shape()[1]}, getInferenceContext().memory_pool); - getInferenceContext().dequant(w, w_w, w_s, w_z, type, w_g_idx); - getInferenceContext().linear(out, x, w, alpha, beta, residual, bias); -} diff --git a/src/models/jiuge/jiuge.cpp b/src/models/jiuge/jiuge.cpp deleted file mode 100644 index 8b65d8f4f..000000000 --- a/src/models/jiuge/jiuge.cpp +++ /dev/null @@ -1,474 +0,0 @@ -#include "jiuge_impl.hpp" -#include "jiuge_weight.hpp" - -#include "../../tensor.hpp" -#include "../../utils.hpp" -#include "../inference_context.hpp" -#include "infinicore_infer.h" - -#include -#include -#include - -void createDeviceResource(JiugeDeviceResource *rsrc, const JiugeMeta *meta, - const JiugeWeights *weights, - infiniDevice_t device, int idev, - int ndev, int dev_id, - infinicclComm_t comm) { - RUN_INFINI(infinirtSetDevice(device, dev_id)); - infiniopHandle_t handle; - infiniopCreateHandle(&handle); - infinirtStream_t stream; - infinirtStreamCreate(&stream); - - std::vector> w_attn_norm, w_attn_qkv, b_attn_qkv, w_attn_q_norm, w_attn_k_norm, w_attn_out, - w_ffn_norm, w_ffn_gate_up, w_ffn_down; - for (size_t layer = 0; layer < meta->nlayer; layer++) { - w_attn_norm.push_back( - getAttnNorm(meta, weights, layer)); - w_attn_qkv.push_back( - getAttnQKV(meta, weights, layer, idev, ndev)); - if (weights->attn_qkv_b != nullptr) { - b_attn_qkv.push_back( - getAttnQKVBias(meta, weights, layer, idev, ndev)); - } - - if (weights->attn_q_norm != nullptr) { - w_attn_q_norm.push_back( - getAttnQNorm(meta, weights, layer)); - w_attn_k_norm.push_back( - getAttnKNorm(meta, weights, layer)); - } - w_attn_out.push_back( - getAttnO(meta, weights, layer, idev, ndev)); - w_ffn_norm.push_back( - getFFNNorm(meta, weights, layer)); - w_ffn_gate_up.push_back( - getFFNGateUp(meta, weights, layer, idev, ndev)); - w_ffn_down.push_back( - getFFNDown(meta, weights, layer, idev, ndev)); - } - - auto memory_pool = std::make_shared(128 * 1024 * 1024); - - *rsrc = JiugeDeviceResource{ - device, - dev_id, - handle, - getInEmbd(meta, weights), - getOutNorm(meta, weights), - getOutEmbd(meta, weights), - getSinTable(meta), - getCosTable(meta), - w_attn_norm, - w_attn_qkv, - b_attn_qkv, - w_attn_q_norm, - w_attn_k_norm, - w_attn_out, - w_ffn_norm, - w_ffn_gate_up, - w_ffn_down, - stream, - comm, - memory_pool, - }; - RUN_INFINI(infinirtDeviceSynchronize()); -} - -void releaseDeviceResource(JiugeDeviceResource &res) { - infinirtDeviceSynchronize(); - // Release individual Tensors - res.w_in_embd.reset(); - res.w_out_norm.reset(); - res.w_out_embd.reset(); - res.sin_table.reset(); - res.cos_table.reset(); - for (auto &t : res.w_attn_norm) { - t.reset(); - } - res.w_attn_norm.clear(); - for (auto &t : res.w_attn_qkv) { - t.reset(); - } - res.w_attn_qkv.clear(); - for (auto &t : res.b_attn_qkv) { - t.reset(); - } - res.b_attn_qkv.clear(); - for (auto &t : res.w_attn_out) { - t.reset(); - } - res.w_attn_out.clear(); - for (auto &t : res.w_ffn_norm) { - t.reset(); - } - res.w_ffn_norm.clear(); - for (auto &t : res.w_ffn_gate_up) { - t.reset(); - } - res.w_ffn_gate_up.clear(); - for (auto &t : res.w_ffn_down) { - t.reset(); - } - res.w_ffn_down.clear(); - infiniopDestroyHandle(res.handle); - res.handle = nullptr; - infinirtStreamDestroy(res.stream); - res.stream = nullptr; - infinicclCommDestroy(res.comm); - res.comm = nullptr; -} - -void inferDeviceBatch(const JiugeMeta &meta, JiugeDeviceResource &rsrc, - uint32_t idev, uint32_t ndev, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output, void *last_logits) { - auto nlayer = meta.nlayer; - auto nkvh = meta.nkvh / ndev; - auto nh = meta.nh / ndev; - auto ngroup = nh / nkvh; - // auto dctx = meta.dctx; - auto dh = meta.dh; - auto d = meta.d; - auto dt_logits = meta.dt_logits; - auto di = meta.di / ndev; - auto dvoc = meta.dvoc; - auto stream = rsrc.stream; - bool has_qkv_bias = rsrc.b_attn_qkv.size() > 0; - bool has_qk_norm = rsrc.w_attn_q_norm.size() > 0 && rsrc.w_attn_k_norm.size() > 0; - - // Allocate buffers - auto logits_in = Tensor::buffer(dt_logits, {ntok, d}, rsrc.memory_pool); - auto logits_out = Tensor::buffer(dt_logits, {ntok, d}, rsrc.memory_pool); - auto qkv_buf = Tensor::buffer(dt_logits, {ntok, (nh + nkvh * 2) * dh}, rsrc.memory_pool); - auto gate_up_buf = Tensor::buffer(dt_logits, {ntok, 2 * di}, rsrc.memory_pool); - auto o_buf = Tensor::buffer(dt_logits, {ntok, nh * dh}, rsrc.memory_pool); - auto prob_buf = Tensor::buffer(dt_logits, {nreq, dvoc}, rsrc.memory_pool); - auto result_buf = Tensor::buffer(INFINI_DTYPE_I64, {nreq}, rsrc.memory_pool); - auto result_cpu = std::vector(nreq); - - auto qkv_rope = qkv_buf->view({ntok, nh + nkvh * 2, dh}); - auto q_buf = qkv_rope->slice(1, 0, nh); - auto k_buf = qkv_rope->slice(1, nh, nkvh); - - // Prepare inputs - auto batch_pos_ids = std::vector(ntok); - size_t req_start = 0; - for (uint32_t req = 0; req < nreq; req++) { - for (uint32_t i = 0; i < req_lens[req]; i++) { - batch_pos_ids[req_start + i] = req_pos[req] + i; - } - req_start += req_lens[req]; - } - - std::shared_ptr pos_ids_buf; - if (rsrc.device == INFINI_DEVICE_CPU) { - pos_ids_buf = Tensor::weight(batch_pos_ids.data(), INFINI_DTYPE_U32, {ntok}); - } else { - pos_ids_buf = Tensor::buffer(INFINI_DTYPE_U32, {ntok}, rsrc.memory_pool); - RUN_INFINI(infinirtMemcpyAsync(pos_ids_buf->data(), batch_pos_ids.data(), sizeof(uint32_t) * ntok, - INFINIRT_MEMCPY_H2D, stream)); - } - for (uint32_t i = 0; i < ntok; i++) { - RUN_INFINI(infinirtMemcpyAsync(logits_in->data(i * d), - rsrc.w_in_embd->data(tokens[i] * d), - dsize(dt_logits) * d, INFINIRT_MEMCPY_D2D, stream)); - } - - // Attention - // attention inner - size_t max_qk_size = 0; - size_t max_seq_len = 0; - - for (uint32_t req = 0; req < nreq; req++) { - auto past_len = req_pos[req]; - auto seq_len = req_lens[req]; - auto total_len = past_len + seq_len; - - max_qk_size = std::max(max_qk_size, size_t(seq_len * total_len)); - max_seq_len = std::max(max_seq_len, size_t(seq_len)); - } - - auto qk_buf = Tensor::buffer(dt_logits, {nh * max_qk_size}, rsrc.memory_pool); - auto rearrange_q_buf = Tensor::buffer(dt_logits, {nkvh, ngroup * max_seq_len, dh}, rsrc.memory_pool); - auto q_rearrange = rearrange_q_buf->view({nkvh, ngroup, max_seq_len, dh}); - auto attn_val_buf = Tensor::buffer(dt_logits, {nkvh, ngroup * max_seq_len, dh}, rsrc.memory_pool); - auto attn_val_gemm = attn_val_buf->view({nkvh, ngroup, max_seq_len, dh}); - - // MLP buffers - auto gate_buf = gate_up_buf->slice(1, 0, di); - auto up_buf = gate_up_buf->slice(1, di, di); - - // Compute - for (uint32_t layer = 0; layer < nlayer; layer++) { - // 1. Attention - // rms norm - rmsnorm(logits_out, logits_in, rsrc.w_attn_norm[layer], meta.epsilon); - // qkv_proj - linear(qkv_buf, logits_out, rsrc.w_attn_qkv[layer], 1.0, 0.0, nullptr, has_qkv_bias ? rsrc.b_attn_qkv[layer] : nullptr); - - if (has_qk_norm) { - rmsnorm(q_buf, q_buf, rsrc.w_attn_q_norm[layer], meta.epsilon); - rmsnorm(k_buf, k_buf, rsrc.w_attn_k_norm[layer], meta.epsilon); - } - - // rope - rope(q_buf, q_buf, pos_ids_buf, rsrc.sin_table, rsrc.cos_table); - rope(k_buf, k_buf, pos_ids_buf, rsrc.sin_table, rsrc.cos_table); - - size_t token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto past_len = req_pos[req]; - auto seq_len = req_lens[req]; - auto total_len = past_len + seq_len; - auto o = o_buf->slice({{0, token_offset, seq_len}})->view({seq_len, nkvh, ngroup, dh})->permute({1, 2, 0, 3}); - auto q = qkv_rope->slice({{0, token_offset, seq_len}, {1, 0, nh}})->view({seq_len, nkvh, ngroup, dh})->permute({1, 2, 0, 3}); - auto k = qkv_rope->slice({{0, token_offset, seq_len}, {1, nh, nkvh}}); - auto v = qkv_rope->slice({{0, token_offset, seq_len}, {1, nh + nkvh, nkvh}}); - - // self attention - // concat - rearrange(kv_caches[req]->k[idev][layer]->slice(0, past_len, seq_len), k); - rearrange(kv_caches[req]->v[idev][layer]->slice(0, past_len, seq_len), v); - // qk - rearrange(q_rearrange->slice(2, 0, seq_len), q); - auto qk_gemm = qk_buf->slice(0, 0, nh * seq_len * total_len)->view({nkvh, ngroup * seq_len, total_len}); - auto k_gemm = kv_caches[req]->k[idev][layer]->slice(0, 0, total_len)->permute({1, 2, 0}); - linear(qk_gemm, rearrange_q_buf->slice(1, 0, ngroup * seq_len), k_gemm, 1.f / float(sqrt(dh)), 0.f, nullptr, nullptr); - // softmax - auto qk_softmax = qk_gemm->view({nh, seq_len, total_len}); - causalSoftmax(qk_softmax, qk_softmax); - auto v_gemm = kv_caches[req]->v[idev][layer]->slice(0, 0, total_len)->permute({1, 0, 2}); - linear(attn_val_buf->slice(1, 0, ngroup * seq_len), qk_gemm, v_gemm, 1.f, 0.f, nullptr, nullptr); - // rearrange attn val - rearrange(o, attn_val_gemm->slice(2, 0, seq_len)); - - token_offset += seq_len; - } - - // o_proj - linear(logits_in, o_buf, rsrc.w_attn_out[layer], 1.0, 0.0, idev == 0 ? logits_in : nullptr, nullptr); // only rank 0 adds residual - - // All_reduce if distributed - if (rsrc.comm != nullptr) { - RUN_INFINI(infinicclAllReduce( - logits_in->data(), logits_in->data(), ntok * d, dt_logits, - INFINICCL_SUM, rsrc.comm, stream)); - RUN_INFINI(infinirtStreamSynchronize(stream)); - } - // 2. FFN - rmsnorm(logits_out, logits_in, rsrc.w_ffn_norm[layer], meta.epsilon); - linear(gate_up_buf, logits_out, rsrc.w_ffn_gate_up[layer], 1.0, 0.0, nullptr, nullptr); - swiglu(gate_buf, up_buf, gate_buf); - linear(logits_in, gate_buf, rsrc.w_ffn_down[layer], 1.0, 0.0, idev == 0 ? logits_in : nullptr, nullptr); // only rank 0 adds residual - - // All_reduce if distributed - if (rsrc.comm != nullptr) { - RUN_INFINI(infinicclAllReduce( - logits_in->data(), logits_in->data(), ntok * d, dt_logits, - INFINICCL_SUM, rsrc.comm, stream)); - RUN_INFINI(infinirtStreamSynchronize(stream)); - } - } - // Sample and Output - if (idev == 0) { - if (last_logits != nullptr) { - rmsnorm(logits_out, logits_in, rsrc.w_out_norm, meta.epsilon); - auto last_logits_buf = Tensor::buffer(dt_logits, {ntok, dvoc}, rsrc.memory_pool); - linear(last_logits_buf, logits_out, rsrc.w_out_embd, 1.0, 0.0, nullptr, nullptr); - RUN_INFINI(infinirtStreamSynchronize(stream)); - RUN_INFINI(infinirtMemcpy(last_logits, last_logits_buf->data(), dsize(dt_logits) * ntok * dvoc, INFINIRT_MEMCPY_D2H)); - } - if (output != nullptr) { - size_t token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto seq_len = req_lens[req]; - token_offset += seq_len; - rmsnorm(logits_out->slice(0, req, 1), - logits_in->slice(0, token_offset - 1, 1), - rsrc.w_out_norm, - meta.epsilon); - } - linear(prob_buf, logits_out->slice(0, 0, nreq), rsrc.w_out_embd, 1.0, 0.0, nullptr, nullptr); - std::random_device _rd; - std::mt19937 gen(_rd()); - token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto seq_len = req_lens[req]; - float random_val = std::uniform_real_distribution(0, 1)(gen); - randomSample(result_buf->slice(0, req, 1)->view_as({}, {}), - prob_buf->slice(0, req, 1)->view_as({dvoc}, {1}), - random_val, topp[req], topk[req], temperature[req]); - token_offset += seq_len; - } - RUN_INFINI(infinirtStreamSynchronize(stream)); - RUN_INFINI(infinirtMemcpy(result_cpu.data(), result_buf->data(), - sizeof(int64_t) * nreq, INFINIRT_MEMCPY_D2H)); - for (uint32_t req = 0; req < nreq; req++) { - output[req] = uint32_t(result_cpu[req]); - } - } - } -} - -__INFINI_C void -inferBatchJiuge(struct JiugeModel *model, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output) { - model->req.tokens = tokens; - model->req.ntok = ntok; - model->req.req_lens = req_lens; - model->req.nreq = nreq; - model->req.req_pos = req_pos; - model->req.kv_caches = kv_caches; - model->req.output = output; - model->req.logits = nullptr; - model->req.temperature = temperature; - model->req.topk = topk; - model->req.topp = topp; - - for (size_t idev = 0; idev < model->dev_ids.size(); idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].proceed = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - for (size_t i = model->dev_ids.size(); i > 0; i--) { - auto idev = i - 1; - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].cv_done.wait(lock, [&] { return !(model->states[idev].proceed); }); - lock.unlock(); - } -} - -__INFINI_C void -forwardBatchJiuge(struct JiugeModel *model, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - void *logits) { - model->req.tokens = tokens; - model->req.ntok = ntok; - model->req.req_lens = req_lens; - model->req.nreq = nreq; - model->req.req_pos = req_pos; - model->req.kv_caches = kv_caches; - model->req.output = nullptr; - model->req.logits = logits; - model->req.temperature = nullptr; - model->req.topk = nullptr; - model->req.topp = nullptr; - - for (size_t idev = 0; idev < model->dev_ids.size(); idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].proceed = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - for (size_t i = model->dev_ids.size(); i > 0; i--) { - auto idev = i - 1; - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].cv_done.wait(lock, [&] { return !(model->states[idev].proceed); }); - lock.unlock(); - } -} - -void launchDevice(const JiugeMeta &meta, const JiugeWeights *weights, JiugeDeviceResource *rsrc, InferState &state, InferRequest &req, - infiniDevice_t device, int idev, int ndev, int dev_id, infinicclComm_t comm) { - // Create Device Resource - createDeviceResource(rsrc, &meta, weights, device, idev, ndev, dev_id, comm); - - CacheManager cache_manager(100); - InferenceContext ctx(rsrc->handle, rsrc->memory_pool, &cache_manager, rsrc->stream); - - // Set the inference context for this thread - setInferenceContext(&ctx); - - { - std::unique_lock lock(state.mtx); - state.loaded = true; - lock.unlock(); - state.cv_load.notify_one(); - } - - // Infer Loop - while (true) { - std::unique_lock lock(state.mtx); - state.cv_start.wait(lock, [&] { return state.proceed || state.exit_flag; }); - // quit if exit_flag is set - if (state.exit_flag) { - break; - } - - inferDeviceBatch(meta, *rsrc, idev, ndev, req.tokens, req.ntok, - req.req_lens, req.nreq, req.req_pos, req.kv_caches, - req.temperature, req.topk, req.topp, req.output, req.logits); - - state.proceed = false; - lock.unlock(); - state.cv_done.notify_one(); - } - - // Clean-Up - releaseDeviceResource(*rsrc); - setInferenceContext(nullptr); // Clear the context when done -} - -JiugeModel::JiugeModel(const JiugeMeta *_meta, const JiugeWeights *weights, infiniDevice_t device_, std::vector device_ids) : meta(*_meta) { - int ndev = int(device_ids.size()); - device = device_; - dev_ids = device_ids; - dev_resources = std::vector(ndev); - states = std::vector(ndev); - threads.resize(ndev); - RUN_INFINI(infinirtInit()); - auto comms = std::vector(ndev, nullptr); - if (ndev > 1) { - RUN_INFINI(infinicclCommInitAll(device, comms.data(), ndev, dev_ids.data())); - } - - for (int i = 0; i < ndev; i++) { - threads[i] = std::thread(launchDevice, std::cref(meta), weights, &dev_resources[i], std::ref(states[i]), std::ref(req), device, i, ndev, dev_ids[i], comms[i]); - } - for (int i = 0; i < ndev; i++) { - std::unique_lock lock(states[i].mtx); - states[i].cv_load.wait(lock, [&] { return states[i].loaded; }); - lock.unlock(); - } -} - -__INFINI_C struct JiugeModel * -createJiugeModel(const JiugeMeta *meta, - const JiugeWeights *weights, - infiniDevice_t device, - int ndev, - const int *dev_ids) { - std::vector device_ids(ndev); - std::copy(dev_ids, dev_ids + ndev, device_ids.begin()); - JiugeModel *model = new JiugeModel(meta, weights, device, device_ids); - return model; -} - -__INFINI_C void destroyJiugeModel(struct JiugeModel *model) { - auto ndev = model->dev_resources.size(); - - for (size_t idev = 0; idev < ndev; idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].exit_flag = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - - for (size_t idev = 0; idev < ndev; idev++) { - model->threads[idev].join(); - } - - delete model; -} diff --git a/src/models/jiuge/jiuge_impl.hpp b/src/models/jiuge/jiuge_impl.hpp deleted file mode 100644 index 64ba72dd5..000000000 --- a/src/models/jiuge/jiuge_impl.hpp +++ /dev/null @@ -1,69 +0,0 @@ -#ifndef JIUGE_IMPL_H -#define JIUGE_IMPL_H - -#include "infinicore_infer.h" - -#include "../../allocator.hpp" -#include "../../tensor.hpp" - -#include -#include -#include -#include -#include - -struct JiugeDeviceResource { - // Device - infiniDevice_t device; - int device_id; - infiniopHandle_t handle; - // Weights - std::shared_ptr w_in_embd, w_out_norm, w_out_embd, sin_table, - cos_table; - std::vector> w_attn_norm, w_attn_qkv, b_attn_qkv, w_attn_q_norm, w_attn_k_norm,w_attn_out, - w_ffn_norm, w_ffn_gate_up, w_ffn_down; - // Streams - infinirtStream_t stream; - // Communicator - infinicclComm_t comm; - - std::shared_ptr memory_pool; -}; - -struct InferState { - std::mutex mtx; - std::condition_variable cv_load, cv_start, cv_done; - bool loaded = false; - bool proceed = false; - bool exit_flag = false; -}; - -struct InferRequest { - const uint32_t *tokens; - uint32_t ntok; - const uint32_t *req_lens; - uint32_t nreq; - const uint32_t *req_pos; - struct KVCache **kv_caches; - const float *temperature; - const uint32_t *topk; - const float *topp; - uint32_t *output; - void *logits; -}; - -struct JiugeModel { - JiugeMeta meta; - infiniDevice_t device; - std::vector dev_ids; - std::vector dev_resources; - std::vector states; - std::vector threads; - InferRequest req; - - JiugeModel(const JiugeMeta *, const JiugeWeights *, infiniDevice_t device, std::vector device_ids); -}; - -#include "../../cache.hpp" - -#endif diff --git a/src/models/jiuge/jiuge_weight.hpp b/src/models/jiuge/jiuge_weight.hpp deleted file mode 100644 index 7ee101557..000000000 --- a/src/models/jiuge/jiuge_weight.hpp +++ /dev/null @@ -1,204 +0,0 @@ -#ifndef JIUGE_WEIGHT_HPP -#define JIUGE_WEIGHT_HPP - -#include "jiuge_impl.hpp" - -#include -inline std::shared_ptr getInEmbd( - JiugeMeta const *meta, - JiugeWeights const *w) { - auto shape = std::vector({meta->dvoc, meta->d}); - return Tensor::weight((char *)w->input_embd, meta->dt_logits, shape); -} - -inline std::shared_ptr getOutNorm( - JiugeMeta const *meta, - JiugeWeights const *w) { - auto shape = std::vector({meta->d}); - return Tensor::weight((char *)w->output_norm, w->dt_norm, shape); -} - -inline std::shared_ptr getOutEmbd( - JiugeMeta const *meta, - JiugeWeights const *w) { - if (w->transpose_linear_weights != 0) { - auto shape = std::vector({meta->dvoc, meta->d}); - return Tensor::weight((char *)w->output_embd, meta->dt_logits, shape) - ->permute({1, 0}); - } else { - auto shape = std::vector({meta->d, meta->dvoc}); - return Tensor::weight((char *)w->output_embd, meta->dt_logits, shape); - } -} - -inline std::shared_ptr getAttnNorm( - JiugeMeta const *meta, - JiugeWeights const *w, - size_t layer) { - auto shape = std::vector({meta->d}); - return Tensor::weight((char *)(w->attn_norm[layer]), w->dt_norm, shape); -} - -inline std::shared_ptr getAttnQKV( - JiugeMeta const *meta, - JiugeWeights const *w, - size_t layer, size_t idev, size_t ndev) { - auto nkvh = meta->nkvh; - auto nh = meta->nh; - auto dh = meta->dh; - auto d = meta->d; - size_t offset = idev * ((nkvh * 2 + nh) / ndev * dh) * d * dsize(w->dt_mat); - if (w->transpose_linear_weights != 0) { - auto shape = std::vector({(nh + 2 * nkvh) / ndev * dh, d}); - return Tensor::weight((char *)(w->attn_qkv[layer]) + offset, w->dt_mat, shape) - ->permute({1, 0}); - } else { - auto shape = std::vector({d, (nh + 2 * nkvh) / ndev * dh}); - return Tensor::weight((char *)(w->attn_qkv[layer]) + offset, w->dt_mat, shape); - } -} - -inline std::shared_ptr getAttnQKVBias( - JiugeMeta const *meta, - JiugeWeights const *w, - size_t layer, size_t idev, size_t ndev) { - auto nkvh = meta->nkvh; - auto nh = meta->nh; - auto dh = meta->dh; - size_t offset = idev * ((nkvh * 2 + nh) / ndev * dh) * dsize(w->dt_mat); - auto shape = std::vector({(nh + 2 * nkvh) / ndev * dh}); - return Tensor::weight((char *)(w->attn_qkv_b[layer]) + offset, w->dt_mat, shape); -} - -inline std::shared_ptr getAttnQNorm( - JiugeMeta const *meta, - JiugeWeights const *w, - size_t layer) { - auto shape = std::vector({meta->dh}); - return Tensor::weight((char *)(w->attn_q_norm[layer]), w->dt_norm, shape); -} - -inline std::shared_ptr getAttnKNorm( - JiugeMeta const *meta, - JiugeWeights const *w, - size_t layer) { - auto shape = std::vector({meta->dh}); - return Tensor::weight((char *)(w->attn_k_norm[layer]), w->dt_norm, shape); -} - -inline std::shared_ptr getAttnO(JiugeMeta const *meta, - JiugeWeights const *w, size_t layer, - size_t idev, size_t ndev) { - auto nh = meta->nh; - auto dh = meta->dh; - auto d = meta->d; - size_t offset = idev * d * (nh / ndev * dh) * dsize(w->dt_mat); - if (w->transpose_linear_weights != 0) { - auto shape = std::vector({d, nh / ndev * dh}); - return Tensor::weight((char *)(w->attn_o[layer]) + offset, w->dt_mat, shape) - ->permute({1, 0}); - } else { - auto shape = std::vector({nh / ndev * dh, d}); - return Tensor::weight((char *)(w->attn_o[layer]) + offset, w->dt_mat, shape); - } -} - -inline std::shared_ptr getFFNNorm( - JiugeMeta const *meta, - JiugeWeights const *w, - size_t layer) { - auto shape = std::vector({meta->d}); - return Tensor::weight((char *)(w->ffn_norm[layer]), w->dt_norm, shape); -} - -inline std::shared_ptr getFFNGateUp( - JiugeMeta const *meta, - JiugeWeights const *w, - size_t layer, size_t idev, size_t ndev) { - auto di = meta->di; - auto d = meta->d; - size_t offset = idev * (2 * di / ndev) * d * dsize(w->dt_mat); - if (w->transpose_linear_weights != 0) { - auto shape = std::vector({2 * di / ndev, d}); - return Tensor::weight((char *)(w->ffn_gate_up[layer]) + offset, - w->dt_mat, shape) - ->permute({1, 0}); - } else { - auto shape = std::vector({d, 2 * di / ndev}); - return Tensor::weight((char *)(w->ffn_gate_up[layer]) + offset, - w->dt_mat, shape); - } -} - -inline std::shared_ptr getFFNDown( - JiugeMeta const *meta, - JiugeWeights const *w, - size_t layer, size_t idev, size_t ndev) { - auto di = meta->di; - auto d = meta->d; - size_t offset = idev * d * (di / ndev) * dsize(w->dt_mat); - if (w->transpose_linear_weights != 0) { - auto shape = std::vector({d, di / ndev}); - return Tensor::weight((char *)(w->ffn_down[layer]) + offset, w->dt_mat, shape) - ->permute({1, 0}); - } else { - auto shape = std::vector({di / ndev, d}); - return Tensor::weight((char *)(w->ffn_down[layer]) + offset, w->dt_mat, shape); - } -} - -inline std::shared_ptr getSinTable(JiugeMeta const *meta) { - auto half_dh = meta->dh / 2; - auto unit = dsize(meta->dt_logits); - void *table = std::malloc(meta->dctx * half_dh * unit); - - for (size_t i = 0; i < meta->dctx; i++) { - for (size_t j = 0; j < half_dh; j++) { - float _sin = std::sin( - static_cast(i) / std::pow(meta->theta, static_cast(j) / half_dh)); - if (meta->dt_logits == INFINI_DTYPE_F16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_f16(_sin); - } else if (meta->dt_logits == INFINI_DTYPE_BF16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_bf16(_sin); - } else if (meta->dt_logits == INFINI_DTYPE_F32) { - ((float *)table)[i * half_dh + j] = _sin; - } else { - std::cout << "unsupported data type" << std::endl; - exit(1); - } - } - } - auto shape = std::vector({meta->dctx, half_dh}); - auto tensor = Tensor::weight(table, meta->dt_logits, shape); - std::free(table); - return tensor; -} - -inline std::shared_ptr getCosTable(JiugeMeta const *meta) { - auto half_dh = meta->dh / 2; - auto unit = dsize(meta->dt_logits); - void *table = std::malloc(meta->dctx * half_dh * unit); - - for (size_t i = 0; i < meta->dctx; i++) { - for (size_t j = 0; j < half_dh; j++) { - float _cos = std::cos( - static_cast(i) / std::pow(meta->theta, static_cast(j) / half_dh)); - if (meta->dt_logits == INFINI_DTYPE_F16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_f16(_cos); - } else if (meta->dt_logits == INFINI_DTYPE_BF16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_bf16(_cos); - } else if (meta->dt_logits == INFINI_DTYPE_F32) { - ((float *)table)[i * half_dh + j] = _cos; - } else { - std::cout << "unsupported data type" << std::endl; - exit(1); - } - } - } - auto shape = std::vector({meta->dctx, half_dh}); - auto tensor = Tensor::weight(table, meta->dt_logits, shape); - std::free(table); - return tensor; -} - -#endif diff --git a/src/models/jiuge_awq/jiuge_awq.cpp b/src/models/jiuge_awq/jiuge_awq.cpp deleted file mode 100644 index d8651d6c7..000000000 --- a/src/models/jiuge_awq/jiuge_awq.cpp +++ /dev/null @@ -1,405 +0,0 @@ -#include "jiuge_awq.hpp" - -#include "../../tensor.hpp" -#include "../../utils.hpp" -#include "../inference_context.hpp" - -#include -#include -#include - -void createDeviceResource(AWQDeviceResource *rsrc, const JiugeAWQMeta *meta, - std::shared_ptr weights, - infiniDevice_t device, int idev, - int ndev, int dev_id, - infinicclComm_t comm) { - RUN_INFINI(infinirtSetDevice(device, dev_id)); - infiniopHandle_t handle; - infiniopCreateHandle(&handle); - infinirtStream_t stream; - infinirtStreamCreate(&stream); - - auto memory_pool = std::make_shared(128 * 1024 * 1024); - - *rsrc = AWQDeviceResource{ - device, - dev_id, - handle, - weights, - stream, - comm, - memory_pool, - }; - RUN_INFINI(infinirtDeviceSynchronize()); -} - -void releaseDeviceResource(AWQDeviceResource &res) { - infinirtDeviceSynchronize(); - // Release individual Tensors - - infiniopDestroyHandle(res.handle); - res.handle = nullptr; - infinirtStreamDestroy(res.stream); - res.stream = nullptr; - infinicclCommDestroy(res.comm); - res.comm = nullptr; -} - -void inferDeviceBatch(const JiugeAWQMeta *meta, AWQDeviceResource &rsrc, - uint32_t idev, uint32_t ndev, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output, void *last_logits) { - auto nlayer = meta->nlayer; - auto nkvh = meta->nkvh / ndev; - auto nh = meta->nh / ndev; - auto ngroup = nh / nkvh; - // auto dctx = meta.dctx; - auto dh = meta->dh; - auto d = meta->d; - auto dt_logits = meta->dt_logits; - auto di = meta->di / ndev; - auto dvoc = meta->dvoc; - auto stream = rsrc.stream; - auto weight = rsrc.weights; - bool has_qkv_bias = meta->has_qkv_bias; - - // Allocate buffers - auto logits_in = Tensor::buffer(dt_logits, {ntok, d}, rsrc.memory_pool); - auto logits_out = Tensor::buffer(dt_logits, {ntok, d}, rsrc.memory_pool); - auto q_buf = Tensor::buffer(dt_logits, {ntok, nh * dh}, rsrc.memory_pool); - auto k_buf = Tensor::buffer(dt_logits, {ntok, nkvh * dh}, rsrc.memory_pool); - auto v_buf = Tensor::buffer(dt_logits, {ntok, nkvh * dh}, rsrc.memory_pool); - - auto gate_buf = Tensor::buffer(dt_logits, {ntok, di}, rsrc.memory_pool); - auto up_buf = Tensor::buffer(dt_logits, {ntok, di}, rsrc.memory_pool); - - auto o_buf = Tensor::buffer(dt_logits, {ntok, nh * dh}, rsrc.memory_pool); - auto prob_buf = Tensor::buffer(dt_logits, {nreq, dvoc}, rsrc.memory_pool); - auto result_buf = Tensor::buffer(INFINI_DTYPE_I64, {nreq}, rsrc.memory_pool); - auto result_cpu = std::vector(nreq); - - // Prepare inputs - auto batch_pos_ids = std::vector(ntok); - size_t req_start = 0; - for (uint32_t req = 0; req < nreq; req++) { - for (uint32_t i = 0; i < req_lens[req]; i++) { - batch_pos_ids[req_start + i] = req_pos[req] + i; - } - req_start += req_lens[req]; - } - - std::shared_ptr pos_ids_buf; - if (rsrc.device == INFINI_DEVICE_CPU) { - pos_ids_buf = Tensor::weight(batch_pos_ids.data(), INFINI_DTYPE_U32, {ntok}); - } else { - pos_ids_buf = Tensor::buffer(INFINI_DTYPE_U32, {ntok}, rsrc.memory_pool); - RUN_INFINI(infinirtMemcpyAsync(pos_ids_buf->data(), batch_pos_ids.data(), sizeof(uint32_t) * ntok, - INFINIRT_MEMCPY_H2D, stream)); - } - for (uint32_t i = 0; i < ntok; i++) { - RUN_INFINI(infinirtMemcpyAsync(logits_in->data(i * d), - weight->w_in_embd->data(tokens[i] * d), - dsize(dt_logits) * d, INFINIRT_MEMCPY_D2D, stream)); - } - // Attention - // attention inner - size_t max_qk_size = 0; - size_t max_seq_len = 0; - - for (uint32_t req = 0; req < nreq; req++) { - auto past_len = req_pos[req]; - auto seq_len = req_lens[req]; - auto total_len = past_len + seq_len; - - max_qk_size = std::max(max_qk_size, size_t(seq_len * total_len)); - max_seq_len = std::max(max_seq_len, size_t(seq_len)); - } - - auto qk_buf = Tensor::buffer(dt_logits, {nh * max_qk_size}, rsrc.memory_pool); - auto rearrange_q_buf = Tensor::buffer(dt_logits, {nkvh, ngroup * max_seq_len, dh}, rsrc.memory_pool); - auto q_rearrange = rearrange_q_buf->view({nkvh, ngroup, max_seq_len, dh}); - auto attn_val_buf = Tensor::buffer(dt_logits, {nkvh, ngroup * max_seq_len, dh}, rsrc.memory_pool); - auto attn_val_gemm = attn_val_buf->view({nkvh, ngroup, max_seq_len, dh}); - - // Compute - for (uint32_t layer = 0; layer < nlayer; layer++) { - // 1. Attention - // rms norm - rmsnorm(logits_out, logits_in, weight->w_attn_norm[layer], meta->epsilon); - // qkv_proj - dequant_linear(q_buf, logits_out, - weight->w_attn_q[layer]->w, weight->w_attn_q[layer]->s, weight->w_attn_q[layer]->z, - 1.0, 0.0, nullptr, has_qkv_bias ? weight->b_attn_q[layer] : nullptr, - QuantType::AWQ); - dequant_linear(k_buf, logits_out, - weight->w_attn_k[layer]->w, weight->w_attn_k[layer]->s, weight->w_attn_k[layer]->z, - 1.0, 0.0, nullptr, has_qkv_bias ? weight->b_attn_k[layer] : nullptr, - QuantType::AWQ); - dequant_linear(v_buf, logits_out, - weight->w_attn_v[layer]->w, weight->w_attn_v[layer]->s, weight->w_attn_v[layer]->z, - 1.0, 0.0, nullptr, has_qkv_bias ? weight->b_attn_v[layer] : nullptr, - QuantType::AWQ); - // rope - rope_v2(q_buf->view({ntok, nh, dh}), q_buf->view({ntok, nh, dh}), pos_ids_buf, weight->sin_table, weight->cos_table); - rope_v2(k_buf->view({ntok, nkvh, dh}), k_buf->view({ntok, nkvh, dh}), pos_ids_buf, weight->sin_table, weight->cos_table); - size_t token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto past_len = req_pos[req]; - auto seq_len = req_lens[req]; - auto total_len = past_len + seq_len; - auto o = o_buf->slice({{0, token_offset, seq_len}})->view({seq_len, nkvh, ngroup, dh})->permute({1, 2, 0, 3}); - auto q = q_buf->slice({{0, token_offset, seq_len}})->view({seq_len, nkvh, ngroup, dh})->permute({1, 2, 0, 3}); - auto k = k_buf->slice({{0, token_offset, seq_len}})->view({seq_len, nkvh, dh}); - auto v = v_buf->slice({{0, token_offset, seq_len}})->view({seq_len, nkvh, dh}); - - // self attention - // concat - rearrange(kv_caches[req]->k[idev][layer]->slice(0, past_len, seq_len), k); - rearrange(kv_caches[req]->v[idev][layer]->slice(0, past_len, seq_len), v); - // qk - rearrange(q_rearrange->slice(2, 0, seq_len), q); - auto qk_gemm = qk_buf->slice(0, 0, nh * seq_len * total_len)->view({nkvh, ngroup * seq_len, total_len}); - auto k_gemm = kv_caches[req]->k[idev][layer]->slice(0, 0, total_len)->permute({1, 2, 0}); - linear(qk_gemm, rearrange_q_buf->slice(1, 0, ngroup * seq_len), k_gemm, 1.f / float(sqrt(dh)), 0.f, nullptr, nullptr); - // softmax - auto qk_softmax = qk_gemm->view({nh, seq_len, total_len}); - causalSoftmax(qk_softmax, qk_softmax); - auto v_gemm = kv_caches[req]->v[idev][layer]->slice(0, 0, total_len)->permute({1, 0, 2}); - linear(attn_val_buf->slice(1, 0, ngroup * seq_len), qk_gemm, v_gemm, 1.f, 0.f, nullptr, nullptr); - // rearrange attn val - rearrange(o, attn_val_gemm->slice(2, 0, seq_len)); - - token_offset += seq_len; - } - // o_proj - dequant_linear(logits_in, o_buf, - weight->w_attn_out[layer]->w, weight->w_attn_out[layer]->s, weight->w_attn_out[layer]->z, - 1.0, 0.0, idev == 0 ? logits_in : nullptr, nullptr, - QuantType::AWQ); - // All_reduce if distributed - if (rsrc.comm != nullptr) { - RUN_INFINI(infinicclAllReduce( - logits_in->data(), logits_in->data(), ntok * d, dt_logits, - INFINICCL_SUM, rsrc.comm, stream)); - RUN_INFINI(infinirtStreamSynchronize(stream)); - } - // 2. FFN - rmsnorm(logits_out, logits_in, weight->w_ffn_norm[layer], meta->epsilon); - dequant_linear(gate_buf, logits_out, - weight->w_ffn_gate[layer]->w, weight->w_ffn_gate[layer]->s, weight->w_ffn_gate[layer]->z, - 1.0, 0.0, nullptr, nullptr, - QuantType::AWQ); - dequant_linear(up_buf, logits_out, - weight->w_ffn_up[layer]->w, weight->w_ffn_up[layer]->s, weight->w_ffn_up[layer]->z, - 1.0, 0.0, nullptr, nullptr, - QuantType::AWQ); - swiglu(gate_buf, up_buf, gate_buf); - dequant_linear(logits_in, gate_buf, - weight->w_ffn_down[layer]->w, weight->w_ffn_down[layer]->s, weight->w_ffn_down[layer]->z, - 1.0, 0.0, idev == 0 ? logits_in : nullptr, nullptr, - QuantType::AWQ); // only rank 0 adds residual - // All_reduce if distributed - if (rsrc.comm != nullptr) { - RUN_INFINI(infinicclAllReduce( - logits_in->data(), logits_in->data(), ntok * d, dt_logits, - INFINICCL_SUM, rsrc.comm, stream)); - RUN_INFINI(infinirtStreamSynchronize(stream)); - } - } - // Sample and Output - if (idev == 0) { - if (last_logits != nullptr) { - rmsnorm(logits_out, logits_in, weight->w_out_norm, meta->epsilon); - auto last_logits_buf = Tensor::buffer(dt_logits, {ntok, dvoc}, rsrc.memory_pool); - linear(last_logits_buf, logits_out, weight->w_out_embd, 1.0, 0.0, nullptr, nullptr); - RUN_INFINI(infinirtStreamSynchronize(stream)); - RUN_INFINI(infinirtMemcpy(last_logits, last_logits_buf->data(), dsize(dt_logits) * ntok * dvoc, INFINIRT_MEMCPY_D2H)); - } - if (output != nullptr) { - size_t token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto seq_len = req_lens[req]; - token_offset += seq_len; - rmsnorm(logits_out->slice(0, req, 1), - logits_in->slice(0, token_offset - 1, 1), - weight->w_out_norm, - meta->epsilon); - } - linear(prob_buf, logits_out->slice(0, 0, nreq), weight->w_out_embd, 1.0, 0.0, nullptr, nullptr); - std::random_device _rd; - std::mt19937 gen(_rd()); - token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto seq_len = req_lens[req]; - float random_val = std::uniform_real_distribution(0, 1)(gen); - randomSample(result_buf->slice(0, req, 1)->view_as({}, {}), - prob_buf->slice(0, req, 1)->view_as({dvoc}, {1}), - random_val, topp[req], topk[req], temperature[req]); - token_offset += seq_len; - } - RUN_INFINI(infinirtStreamSynchronize(stream)); - RUN_INFINI(infinirtMemcpy(result_cpu.data(), result_buf->data(), - sizeof(int64_t) * nreq, INFINIRT_MEMCPY_D2H)); - for (uint32_t req = 0; req < nreq; req++) { - output[req] = uint32_t(result_cpu[req]); - } - } - } -} - -__INFINI_C void -inferBatchJiugeAWQ(struct JiugeAWQModel *model, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output) { - model->req.tokens = tokens; - model->req.ntok = ntok; - model->req.req_lens = req_lens; - model->req.nreq = nreq; - model->req.req_pos = req_pos; - model->req.kv_caches = kv_caches; - model->req.output = output; - model->req.logits = nullptr; - model->req.temperature = temperature; - model->req.topk = topk; - model->req.topp = topp; - - for (size_t idev = 0; idev < model->dev_ids.size(); idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].proceed = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - for (size_t i = model->dev_ids.size(); i > 0; i--) { - auto idev = i - 1; - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].cv_done.wait(lock, [&] { return !(model->states[idev].proceed); }); - lock.unlock(); - } -} - -__INFINI_C void -forwardBatchJiugeAWQ(struct JiugeAWQModel *model, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - void *logits) { - model->req.tokens = tokens; - model->req.ntok = ntok; - model->req.req_lens = req_lens; - model->req.nreq = nreq; - model->req.req_pos = req_pos; - model->req.kv_caches = kv_caches; - model->req.output = nullptr; - model->req.logits = logits; - model->req.temperature = nullptr; - model->req.topk = nullptr; - model->req.topp = nullptr; - - for (size_t idev = 0; idev < model->dev_ids.size(); idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].proceed = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - for (size_t i = model->dev_ids.size(); i > 0; i--) { - auto idev = i - 1; - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].cv_done.wait(lock, [&] { return !(model->states[idev].proceed); }); - lock.unlock(); - } -} - -void launchDevice(const JiugeAWQMeta *meta, std::shared_ptr weights, AWQDeviceResource *rsrc, InferState &state, InferRequest &req, - infiniDevice_t device, int idev, int ndev, int dev_id, infinicclComm_t comm) { - // Create Device Resource - createDeviceResource(rsrc, meta, weights, device, idev, ndev, dev_id, comm); - - CacheManager cache_manager(100); - InferenceContext ctx(rsrc->handle, rsrc->memory_pool, &cache_manager, rsrc->stream); - - // Set the inference context for this thread - setInferenceContext(&ctx); - - { - std::unique_lock lock(state.mtx); - state.loaded = true; - lock.unlock(); - state.cv_load.notify_one(); - } - - // Infer Loop - while (true) { - std::unique_lock lock(state.mtx); - state.cv_start.wait(lock, [&] { return state.proceed || state.exit_flag; }); - // quit if exit_flag is set - if (state.exit_flag) { - break; - } - - inferDeviceBatch(meta, *rsrc, idev, ndev, req.tokens, req.ntok, - req.req_lens, req.nreq, req.req_pos, req.kv_caches, - req.temperature, req.topk, req.topp, req.output, req.logits); - - state.proceed = false; - lock.unlock(); - state.cv_done.notify_one(); - } - - // Clean-Up - releaseDeviceResource(*rsrc); - setInferenceContext(nullptr); // Clear the context when done -} - -JiugeAWQModel::JiugeAWQModel(const JiugeAWQMeta *meta, const ModelWeights *weights_) { - auto weights = (JiugeAWQWeights *)(weights_); - device = weights->device(); - dev_ids = weights->devIds(); - int ndev = int(dev_ids.size()); - dev_resources = std::vector(ndev); - states = std::vector(ndev); - threads.resize(ndev); - - auto comms = std::vector(ndev, nullptr); - if (ndev > 1) { - RUN_INFINI(infinicclCommInitAll(device, comms.data(), ndev, dev_ids.data())); - } - - for (int i = 0; i < ndev; i++) { - threads[i] = std::thread(launchDevice, meta, weights->device_weights()[i], &dev_resources[i], std::ref(states[i]), std::ref(req), device, i, ndev, dev_ids[i], comms[i]); - } - for (int i = 0; i < ndev; i++) { - std::unique_lock lock(states[i].mtx); - states[i].cv_load.wait(lock, [&] { return states[i].loaded; }); - lock.unlock(); - } -} - -__INFINI_C struct JiugeAWQModel * -createJiugeAWQModel(const JiugeAWQMeta *meta, - const ModelWeights *weights) { - JiugeAWQModel *model = new JiugeAWQModel(meta, weights); - return model; -} - -__INFINI_C void destroyJiugeAWQModel(struct JiugeAWQModel *model) { - auto ndev = model->dev_resources.size(); - - for (size_t idev = 0; idev < ndev; idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].exit_flag = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - - for (size_t idev = 0; idev < ndev; idev++) { - model->threads[idev].join(); - } - - delete model; -} diff --git a/src/models/jiuge_awq/jiuge_awq.hpp b/src/models/jiuge_awq/jiuge_awq.hpp deleted file mode 100644 index 7a1f05d9d..000000000 --- a/src/models/jiuge_awq/jiuge_awq.hpp +++ /dev/null @@ -1,82 +0,0 @@ -#pragma once -#include "infinicore_infer/models/jiuge_awq.h" - -#include "../../cache.hpp" -#include "../../dataloader/weights_loader.hpp" - -#include -#include -#include - -struct QuantInt4Weight { - std::shared_ptr w, s, z; -}; - -struct JiugeAWQDeviceWeight { - std::shared_ptr w_in_embd, w_out_norm, w_out_embd, sin_table, - cos_table; - std::vector> w_attn_norm, b_attn_q, b_attn_k, b_attn_v, w_ffn_norm; - std::vector> w_attn_q, w_attn_k, w_attn_v, w_attn_out, w_ffn_gate, w_ffn_up, w_ffn_down; -}; - -class JiugeAWQWeights : public infinicore::weights::Loader { -private: - std::vector> _device_weights; - -public: - JiugeAWQWeights(const JiugeAWQMeta *meta, - infiniDevice_t device, - const std::vector &dev_ids); - std::vector> &device_weights() { - return _device_weights; - } -}; - -struct AWQDeviceResource { - // Device - infiniDevice_t device; - int device_id; - infiniopHandle_t handle; - // Weights - std::shared_ptr weights; - // Streams - infinirtStream_t stream; - // Communicator - infinicclComm_t comm; - - std::shared_ptr memory_pool; -}; - -struct InferRequest { - const uint32_t *tokens; - uint32_t ntok; - const uint32_t *req_lens; - uint32_t nreq; - const uint32_t *req_pos; - struct KVCache **kv_caches; - const float *temperature; - const uint32_t *topk; - const float *topp; - uint32_t *output; - void *logits; -}; - -struct InferState { - std::mutex mtx; - std::condition_variable cv_load, cv_start, cv_done; - bool loaded = false; - bool proceed = false; - bool exit_flag = false; -}; - -struct JiugeAWQModel { - JiugeAWQMeta meta; - infiniDevice_t device; - std::vector dev_ids; - std::vector dev_resources; - std::vector states; - std::vector threads; - InferRequest req; - - JiugeAWQModel(const JiugeAWQMeta *, const ModelWeights *); -}; \ No newline at end of file diff --git a/src/models/jiuge_awq/jiuge_awq_weight.cpp b/src/models/jiuge_awq/jiuge_awq_weight.cpp deleted file mode 100644 index fdac6a80f..000000000 --- a/src/models/jiuge_awq/jiuge_awq_weight.cpp +++ /dev/null @@ -1,132 +0,0 @@ -#include "jiuge_awq.hpp" - -#include - -inline std::shared_ptr getSinTable(size_t dctx, size_t dh, float theta) { - auto half_dh = dh / 2; - auto unit = dsize(INFINI_DTYPE_F16); - void *table = std::malloc(dctx * half_dh * unit); - - for (size_t i = 0; i < dctx; i++) { - for (size_t j = 0; j < half_dh; j++) { - float _sin = std::sin( - static_cast(i) / std::pow(theta, static_cast(j) / half_dh)); - - ((uint16_t *)table)[i * half_dh + j] = f32_to_f16(_sin); - } - } - auto shape = std::vector({dctx, half_dh}); - auto tensor = Tensor::weight(table, INFINI_DTYPE_F16, shape); - std::free(table); - return tensor; -} - -inline std::shared_ptr getCosTable(size_t dctx, size_t dh, float theta) { - auto half_dh = dh / 2; - auto unit = dsize(INFINI_DTYPE_F16); - void *table = std::malloc(dctx * half_dh * unit); - - for (size_t i = 0; i < dctx; i++) { - for (size_t j = 0; j < half_dh; j++) { - float _cos = std::cos( - static_cast(i) / std::pow(theta, static_cast(j) / half_dh)); - - ((uint16_t *)table)[i * half_dh + j] = f32_to_f16(_cos); - } - } - auto shape = std::vector({dctx, half_dh}); - auto tensor = Tensor::weight(table, INFINI_DTYPE_F16, shape); - std::free(table); - return tensor; -} - -JiugeAWQWeights::JiugeAWQWeights( - const JiugeAWQMeta *meta, - infiniDevice_t device, - const std::vector &dev_ids) : infinicore::weights::Loader(device, dev_ids) { - auto ndev = dev_ids.size(); - _device_weights.resize(ndev); - infiniDtype_t dt_logits = meta->dt_logits; - infiniDtype_t dt_norm_w = meta->dt_norm_w; - size_t nlayer = meta->nlayer; - size_t d = meta->d; - size_t nh = meta->nh / ndev; - size_t nkvh = meta->nkvh / ndev; - size_t dh = meta->dh; - size_t di = meta->di / ndev; - size_t dctx = meta->dctx; - size_t dvoc = meta->dvoc; - size_t nbit = meta->nbit; - size_t quant_group_size = meta->quant_group_size; - - for (size_t i = 0; i < ndev; i++) { - RUN_INFINI(infinirtSetDevice(device, dev_ids[i])); - - auto weight = std::make_shared(); - _device_weights[i] = weight; - - auto w_in_embd = Tensor::weight(nullptr, dt_logits, {dvoc, d}); - this->register_weight("model.embed_tokens.weight", w_in_embd, i); - weight->w_in_embd = w_in_embd; - - auto w_out_norm = Tensor::weight(nullptr, dt_norm_w, {d}); - this->register_weight("model.norm.weight", w_out_norm, i); - weight->w_out_norm = w_out_norm; - - auto w_out_embd = Tensor::weight(nullptr, dt_logits, {dvoc, d})->permute({1, 0}); - this->register_weight("lm_head.weight", w_out_embd, i); - weight->w_out_embd = w_out_embd; - - weight->sin_table = getSinTable(dctx, dh, meta->theta); - weight->cos_table = getCosTable(dctx, dh, meta->theta); - - for (size_t layer = 0; layer < nlayer; layer++) { - -#define RIGISTER_LAYER_WEIGHT(W_NAME, W_VAR, W_SHAPE, W_DTYPE, W_DIST_TYPE) \ - auto W_VAR = Tensor::weight(nullptr, W_DTYPE, W_SHAPE); \ - this->register_weight(W_NAME, W_VAR, i, infinicore::weights::DistributionType::W_DIST_TYPE); \ - weight->W_VAR.push_back(W_VAR); - - RIGISTER_LAYER_WEIGHT("model.layers." + std::to_string(layer) + ".input_layernorm.weight", w_attn_norm, {d}, dt_norm_w, FULL); - -#define REGISTER_LAYER_QUANT_WEIGHT(W_NAME, W_VAR, W_IN, W_OUT, W_DIST_TYPE) \ - auto W_VAR = std::make_shared(); \ - W_VAR->w = Tensor::weight(nullptr, INFINI_DTYPE_I32, {W_IN, (W_OUT)*nbit / 32}); \ - this->register_weight(W_NAME + ".qweight", W_VAR->w, i, infinicore::weights::DistributionType::W_DIST_TYPE); \ - W_VAR->s = Tensor::weight(nullptr, INFINI_DTYPE_F16, {(W_IN) / quant_group_size, (W_OUT)}); \ - this->register_weight(W_NAME + ".scales", W_VAR->s, i, infinicore::weights::DistributionType::W_DIST_TYPE); \ - W_VAR->z = Tensor::weight(nullptr, INFINI_DTYPE_I32, {(W_IN) / quant_group_size, (W_OUT)*nbit / 32}); \ - this->register_weight(W_NAME + ".qzeros", W_VAR->z, i, infinicore::weights::DistributionType::W_DIST_TYPE); \ - weight->W_VAR.push_back(W_VAR); - - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.q_proj", w_attn_q, d, nh * dh, COLUMN); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.k_proj", w_attn_k, d, nkvh * dh, COLUMN); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.v_proj", w_attn_v, d, nkvh * dh, COLUMN); - RIGISTER_LAYER_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.q_proj.bias", b_attn_q, {nh * dh}, INFINI_DTYPE_F16, COLUMN); - RIGISTER_LAYER_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.k_proj.bias", b_attn_k, {nkvh * dh}, INFINI_DTYPE_F16, COLUMN); - RIGISTER_LAYER_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.v_proj.bias", b_attn_v, {nkvh * dh}, INFINI_DTYPE_F16, COLUMN); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.o_proj", w_attn_out, nh * dh, d, ROW); - - RIGISTER_LAYER_WEIGHT("model.layers." + std::to_string(layer) + ".post_attention_layernorm.weight", w_ffn_norm, {d}, dt_norm_w, FULL); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".mlp.gate_proj", w_ffn_gate, d, di, COLUMN); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".mlp.up_proj", w_ffn_up, d, di, COLUMN); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".mlp.down_proj", w_ffn_down, di, d, ROW); - } - } - -#undef RIGISTER_LAYER_WEIGHT -#undef REGISTER_LAYER_QUANT_WEIGHT -} - -__INFINI_C struct ModelWeights * -createJiugeAWQWeights(const JiugeAWQMeta *meta, - infiniDevice_t device, - int ndev, - const int *dev_ids) { - JiugeAWQWeights *weights = new JiugeAWQWeights(meta, device, std::vector(dev_ids, dev_ids + ndev)); - return (struct ModelWeights *)weights; -} - -__INFINI_C void JiugeAWQLoadWeight(struct ModelWeights *weights, const char *name, void *data) { - loadModelWeight(weights, name, data); -} diff --git a/src/models/jiuge_gptq/jiuge_gptq.cpp b/src/models/jiuge_gptq/jiuge_gptq.cpp deleted file mode 100644 index d7002e14f..000000000 --- a/src/models/jiuge_gptq/jiuge_gptq.cpp +++ /dev/null @@ -1,405 +0,0 @@ -#include "jiuge_gptq.hpp" - -#include "../../tensor.hpp" -#include "../../utils.hpp" -#include "../inference_context.hpp" - -#include -#include -#include - -void createDeviceResource(GPTQDeviceResource *rsrc, const JiugeGPTQMeta *meta, - std::shared_ptr weights, - infiniDevice_t device, int idev, - int ndev, int dev_id, - infinicclComm_t comm) { - RUN_INFINI(infinirtSetDevice(device, dev_id)); - infiniopHandle_t handle; - infiniopCreateHandle(&handle); - infinirtStream_t stream; - infinirtStreamCreate(&stream); - - auto memory_pool = std::make_shared(128 * 1024 * 1024); - - *rsrc = GPTQDeviceResource{ - device, - dev_id, - handle, - weights, - stream, - comm, - memory_pool, - }; - RUN_INFINI(infinirtDeviceSynchronize()); -} - -void releaseDeviceResource(GPTQDeviceResource &res) { - infinirtDeviceSynchronize(); - // Release individual Tensors - - infiniopDestroyHandle(res.handle); - res.handle = nullptr; - infinirtStreamDestroy(res.stream); - res.stream = nullptr; - infinicclCommDestroy(res.comm); - res.comm = nullptr; -} - -void inferDeviceBatch(const JiugeGPTQMeta *meta, GPTQDeviceResource &rsrc, - uint32_t idev, uint32_t ndev, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output, void *last_logits) { - auto nlayer = meta->nlayer; - auto nkvh = meta->nkvh / ndev; - auto nh = meta->nh / ndev; - auto ngroup = nh / nkvh; - // auto dctx = meta.dctx; - auto dh = meta->dh; - auto d = meta->d; - auto dt_logits = meta->dt_logits; - auto di = meta->di / ndev; - auto dvoc = meta->dvoc; - auto stream = rsrc.stream; - auto weight = rsrc.weights; - bool has_qkv_bias = meta->has_qkv_bias; - - // Allocate buffers - auto logits_in = Tensor::buffer(dt_logits, {ntok, d}, rsrc.memory_pool); - auto logits_out = Tensor::buffer(dt_logits, {ntok, d}, rsrc.memory_pool); - auto q_buf = Tensor::buffer(dt_logits, {ntok, nh * dh}, rsrc.memory_pool); - auto k_buf = Tensor::buffer(dt_logits, {ntok, nkvh * dh}, rsrc.memory_pool); - auto v_buf = Tensor::buffer(dt_logits, {ntok, nkvh * dh}, rsrc.memory_pool); - - auto gate_buf = Tensor::buffer(dt_logits, {ntok, di}, rsrc.memory_pool); - auto up_buf = Tensor::buffer(dt_logits, {ntok, di}, rsrc.memory_pool); - - auto o_buf = Tensor::buffer(dt_logits, {ntok, nh * dh}, rsrc.memory_pool); - auto prob_buf = Tensor::buffer(dt_logits, {nreq, dvoc}, rsrc.memory_pool); - auto result_buf = Tensor::buffer(INFINI_DTYPE_I64, {nreq}, rsrc.memory_pool); - auto result_cpu = std::vector(nreq); - - // Prepare inputs - auto batch_pos_ids = std::vector(ntok); - size_t req_start = 0; - for (uint32_t req = 0; req < nreq; req++) { - for (uint32_t i = 0; i < req_lens[req]; i++) { - batch_pos_ids[req_start + i] = req_pos[req] + i; - } - req_start += req_lens[req]; - } - - std::shared_ptr pos_ids_buf; - if (rsrc.device == INFINI_DEVICE_CPU) { - pos_ids_buf = Tensor::weight(batch_pos_ids.data(), INFINI_DTYPE_U32, {ntok}); - } else { - pos_ids_buf = Tensor::buffer(INFINI_DTYPE_U32, {ntok}, rsrc.memory_pool); - RUN_INFINI(infinirtMemcpyAsync(pos_ids_buf->data(), batch_pos_ids.data(), sizeof(uint32_t) * ntok, - INFINIRT_MEMCPY_H2D, stream)); - } - for (uint32_t i = 0; i < ntok; i++) { - RUN_INFINI(infinirtMemcpyAsync(logits_in->data(i * d), - weight->w_in_embd->data(tokens[i] * d), - dsize(dt_logits) * d, INFINIRT_MEMCPY_D2D, stream)); - } - // Attention - // attention inner - size_t max_qk_size = 0; - size_t max_seq_len = 0; - - for (uint32_t req = 0; req < nreq; req++) { - auto past_len = req_pos[req]; - auto seq_len = req_lens[req]; - auto total_len = past_len + seq_len; - - max_qk_size = std::max(max_qk_size, size_t(seq_len * total_len)); - max_seq_len = std::max(max_seq_len, size_t(seq_len)); - } - - auto qk_buf = Tensor::buffer(dt_logits, {nh * max_qk_size}, rsrc.memory_pool); - auto rearrange_q_buf = Tensor::buffer(dt_logits, {nkvh, ngroup * max_seq_len, dh}, rsrc.memory_pool); - auto q_rearrange = rearrange_q_buf->view({nkvh, ngroup, max_seq_len, dh}); - auto attn_val_buf = Tensor::buffer(dt_logits, {nkvh, ngroup * max_seq_len, dh}, rsrc.memory_pool); - auto attn_val_gemm = attn_val_buf->view({nkvh, ngroup, max_seq_len, dh}); - - // Compute - for (uint32_t layer = 0; layer < nlayer; layer++) { - // 1. Attention - // rms norm - rmsnorm(logits_out, logits_in, weight->w_attn_norm[layer], meta->epsilon); - // qkv_proj - dequant_linear(q_buf, logits_out, - weight->w_attn_q[layer]->w, weight->w_attn_q[layer]->s, weight->w_attn_q[layer]->z, - 1.0, 0.0, nullptr, has_qkv_bias ? weight->b_attn_q[layer] : nullptr, - QuantType::GPTQ, weight->w_attn_q[layer]->g_idx); - dequant_linear(k_buf, logits_out, - weight->w_attn_k[layer]->w, weight->w_attn_k[layer]->s, weight->w_attn_k[layer]->z, - 1.0, 0.0, nullptr, has_qkv_bias ? weight->b_attn_k[layer] : nullptr, - QuantType::GPTQ, weight->w_attn_k[layer]->g_idx); - dequant_linear(v_buf, logits_out, - weight->w_attn_v[layer]->w, weight->w_attn_v[layer]->s, weight->w_attn_v[layer]->z, - 1.0, 0.0, nullptr, has_qkv_bias ? weight->b_attn_v[layer] : nullptr, - QuantType::GPTQ, weight->w_attn_v[layer]->g_idx); - // rope - rope_v2(q_buf->view({ntok, nh, dh}), q_buf->view({ntok, nh, dh}), pos_ids_buf, weight->sin_table, weight->cos_table); - rope_v2(k_buf->view({ntok, nkvh, dh}), k_buf->view({ntok, nkvh, dh}), pos_ids_buf, weight->sin_table, weight->cos_table); - size_t token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto past_len = req_pos[req]; - auto seq_len = req_lens[req]; - auto total_len = past_len + seq_len; - auto o = o_buf->slice({{0, token_offset, seq_len}})->view({seq_len, nkvh, ngroup, dh})->permute({1, 2, 0, 3}); - auto q = q_buf->slice({{0, token_offset, seq_len}})->view({seq_len, nkvh, ngroup, dh})->permute({1, 2, 0, 3}); - auto k = k_buf->slice({{0, token_offset, seq_len}})->view({seq_len, nkvh, dh}); - auto v = v_buf->slice({{0, token_offset, seq_len}})->view({seq_len, nkvh, dh}); - - // self attention - // concat - rearrange(kv_caches[req]->k[idev][layer]->slice(0, past_len, seq_len), k); - rearrange(kv_caches[req]->v[idev][layer]->slice(0, past_len, seq_len), v); - // qk - rearrange(q_rearrange->slice(2, 0, seq_len), q); - auto qk_gemm = qk_buf->slice(0, 0, nh * seq_len * total_len)->view({nkvh, ngroup * seq_len, total_len}); - auto k_gemm = kv_caches[req]->k[idev][layer]->slice(0, 0, total_len)->permute({1, 2, 0}); - linear(qk_gemm, rearrange_q_buf->slice(1, 0, ngroup * seq_len), k_gemm, 1.f / float(sqrt(dh)), 0.f, nullptr, nullptr); - // softmax - auto qk_softmax = qk_gemm->view({nh, seq_len, total_len}); - causalSoftmax(qk_softmax, qk_softmax); - auto v_gemm = kv_caches[req]->v[idev][layer]->slice(0, 0, total_len)->permute({1, 0, 2}); - linear(attn_val_buf->slice(1, 0, ngroup * seq_len), qk_gemm, v_gemm, 1.f, 0.f, nullptr, nullptr); - // rearrange attn val - rearrange(o, attn_val_gemm->slice(2, 0, seq_len)); - - token_offset += seq_len; - } - // o_proj - dequant_linear(logits_in, o_buf, - weight->w_attn_out[layer]->w, weight->w_attn_out[layer]->s, weight->w_attn_out[layer]->z, - 1.0, 0.0, idev == 0 ? logits_in : nullptr, nullptr, - QuantType::GPTQ, weight->w_attn_out[layer]->g_idx); - // All_reduce if distributed - if (rsrc.comm != nullptr) { - RUN_INFINI(infinicclAllReduce( - logits_in->data(), logits_in->data(), ntok * d, dt_logits, - INFINICCL_SUM, rsrc.comm, stream)); - RUN_INFINI(infinirtStreamSynchronize(stream)); - } - // 2. FFN - rmsnorm(logits_out, logits_in, weight->w_ffn_norm[layer], meta->epsilon); - dequant_linear(gate_buf, logits_out, - weight->w_ffn_gate[layer]->w, weight->w_ffn_gate[layer]->s, weight->w_ffn_gate[layer]->z, - 1.0, 0.0, nullptr, nullptr, - QuantType::GPTQ, weight->w_ffn_gate[layer]->g_idx); - dequant_linear(up_buf, logits_out, - weight->w_ffn_up[layer]->w, weight->w_ffn_up[layer]->s, weight->w_ffn_up[layer]->z, - 1.0, 0.0, nullptr, nullptr, - QuantType::GPTQ, weight->w_ffn_up[layer]->g_idx); - swiglu(gate_buf, up_buf, gate_buf); - dequant_linear(logits_in, gate_buf, - weight->w_ffn_down[layer]->w, weight->w_ffn_down[layer]->s, weight->w_ffn_down[layer]->z, - 1.0, 0.0, idev == 0 ? logits_in : nullptr, nullptr, - QuantType::GPTQ, weight->w_ffn_down[layer]->g_idx); - // All_reduce if distributed - if (rsrc.comm != nullptr) { - RUN_INFINI(infinicclAllReduce( - logits_in->data(), logits_in->data(), ntok * d, dt_logits, - INFINICCL_SUM, rsrc.comm, stream)); - RUN_INFINI(infinirtStreamSynchronize(stream)); - } - } - // Sample and Output - if (idev == 0) { - if (last_logits != nullptr) { - rmsnorm(logits_out, logits_in, weight->w_out_norm, meta->epsilon); - auto last_logits_buf = Tensor::buffer(dt_logits, {ntok, dvoc}, rsrc.memory_pool); - linear(last_logits_buf, logits_out, weight->w_out_embd, 1.0, 0.0, nullptr, nullptr); - RUN_INFINI(infinirtStreamSynchronize(stream)); - RUN_INFINI(infinirtMemcpy(last_logits, last_logits_buf->data(), dsize(dt_logits) * ntok * dvoc, INFINIRT_MEMCPY_D2H)); - } - if (output != nullptr) { - size_t token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto seq_len = req_lens[req]; - token_offset += seq_len; - rmsnorm(logits_out->slice(0, req, 1), - logits_in->slice(0, token_offset - 1, 1), - weight->w_out_norm, - meta->epsilon); - } - linear(prob_buf, logits_out->slice(0, 0, nreq), weight->w_out_embd, 1.0, 0.0, nullptr, nullptr); - std::random_device _rd; - std::mt19937 gen(_rd()); - token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto seq_len = req_lens[req]; - float random_val = std::uniform_real_distribution(0, 1)(gen); - randomSample(result_buf->slice(0, req, 1)->view_as({}, {}), - prob_buf->slice(0, req, 1)->view_as({dvoc}, {1}), - random_val, topp[req], topk[req], temperature[req]); - token_offset += seq_len; - } - RUN_INFINI(infinirtStreamSynchronize(stream)); - RUN_INFINI(infinirtMemcpy(result_cpu.data(), result_buf->data(), - sizeof(int64_t) * nreq, INFINIRT_MEMCPY_D2H)); - for (uint32_t req = 0; req < nreq; req++) { - output[req] = uint32_t(result_cpu[req]); - } - } - } -} - -__INFINI_C void -inferBatchJiugeGPTQ(struct JiugeGPTQModel *model, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output) { - model->req.tokens = tokens; - model->req.ntok = ntok; - model->req.req_lens = req_lens; - model->req.nreq = nreq; - model->req.req_pos = req_pos; - model->req.kv_caches = kv_caches; - model->req.output = output; - model->req.logits = nullptr; - model->req.temperature = temperature; - model->req.topk = topk; - model->req.topp = topp; - - for (size_t idev = 0; idev < model->dev_ids.size(); idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].proceed = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - for (size_t i = model->dev_ids.size(); i > 0; i--) { - auto idev = i - 1; - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].cv_done.wait(lock, [&] { return !(model->states[idev].proceed); }); - lock.unlock(); - } -} - -__INFINI_C void -forwardBatchJiugeGPTQ(struct JiugeGPTQModel *model, - const uint32_t *tokens, uint32_t ntok, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct KVCache **kv_caches, - void *logits) { - model->req.tokens = tokens; - model->req.ntok = ntok; - model->req.req_lens = req_lens; - model->req.nreq = nreq; - model->req.req_pos = req_pos; - model->req.kv_caches = kv_caches; - model->req.output = nullptr; - model->req.logits = logits; - model->req.temperature = nullptr; - model->req.topk = nullptr; - model->req.topp = nullptr; - - for (size_t idev = 0; idev < model->dev_ids.size(); idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].proceed = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - for (size_t i = model->dev_ids.size(); i > 0; i--) { - auto idev = i - 1; - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].cv_done.wait(lock, [&] { return !(model->states[idev].proceed); }); - lock.unlock(); - } -} - -void launchDevice(const JiugeGPTQMeta *meta, std::shared_ptr weights, GPTQDeviceResource *rsrc, InferState &state, InferRequest &req, - infiniDevice_t device, int idev, int ndev, int dev_id, infinicclComm_t comm) { - // Create Device Resource - createDeviceResource(rsrc, meta, weights, device, idev, ndev, dev_id, comm); - - CacheManager cache_manager(100); - InferenceContext ctx(rsrc->handle, rsrc->memory_pool, &cache_manager, rsrc->stream); - - // Set the inference context for this thread - setInferenceContext(&ctx); - - { - std::unique_lock lock(state.mtx); - state.loaded = true; - lock.unlock(); - state.cv_load.notify_one(); - } - - // Infer Loop - while (true) { - std::unique_lock lock(state.mtx); - state.cv_start.wait(lock, [&] { return state.proceed || state.exit_flag; }); - // quit if exit_flag is set - if (state.exit_flag) { - break; - } - - inferDeviceBatch(meta, *rsrc, idev, ndev, req.tokens, req.ntok, - req.req_lens, req.nreq, req.req_pos, req.kv_caches, - req.temperature, req.topk, req.topp, req.output, req.logits); - - state.proceed = false; - lock.unlock(); - state.cv_done.notify_one(); - } - - // Clean-Up - releaseDeviceResource(*rsrc); - setInferenceContext(nullptr); // Clear the context when done -} - -JiugeGPTQModel::JiugeGPTQModel(const JiugeGPTQMeta *meta, const ModelWeights *weights_) { - auto weights = (JiugeGPTQWeights *)(weights_); - device = weights->device(); - dev_ids = weights->devIds(); - int ndev = int(dev_ids.size()); - dev_resources = std::vector(ndev); - states = std::vector(ndev); - threads.resize(ndev); - - auto comms = std::vector(ndev, nullptr); - if (ndev > 1) { - RUN_INFINI(infinicclCommInitAll(device, comms.data(), ndev, dev_ids.data())); - } - - for (int i = 0; i < ndev; i++) { - threads[i] = std::thread(launchDevice, meta, weights->device_weights()[i], &dev_resources[i], std::ref(states[i]), std::ref(req), device, i, ndev, dev_ids[i], comms[i]); - } - for (int i = 0; i < ndev; i++) { - std::unique_lock lock(states[i].mtx); - states[i].cv_load.wait(lock, [&] { return states[i].loaded; }); - lock.unlock(); - } -} - -__INFINI_C struct JiugeGPTQModel * -createJiugeGPTQModel(const JiugeGPTQMeta *meta, - const ModelWeights *weights) { - JiugeGPTQModel *model = new JiugeGPTQModel(meta, weights); - return model; -} - -__INFINI_C void destroyJiugeGPTQModel(struct JiugeGPTQModel *model) { - auto ndev = model->dev_resources.size(); - - for (size_t idev = 0; idev < ndev; idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].exit_flag = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - - for (size_t idev = 0; idev < ndev; idev++) { - model->threads[idev].join(); - } - - delete model; -} diff --git a/src/models/jiuge_gptq/jiuge_gptq.hpp b/src/models/jiuge_gptq/jiuge_gptq.hpp deleted file mode 100644 index b51e9f1d3..000000000 --- a/src/models/jiuge_gptq/jiuge_gptq.hpp +++ /dev/null @@ -1,82 +0,0 @@ -#pragma once -#include "infinicore_infer/models/jiuge_gptq.h" - -#include "../../cache.hpp" -#include "../../dataloader/weights_loader.hpp" - -#include -#include -#include - -struct QuantInt4Weight { - std::shared_ptr w, s, z, g_idx; // add g_idx -}; - -struct JiugeGPTQDeviceWeight { - std::shared_ptr w_in_embd, w_out_norm, w_out_embd, sin_table, - cos_table; - std::vector> w_attn_norm, b_attn_q, b_attn_k, b_attn_v, w_ffn_norm; - std::vector> w_attn_q, w_attn_k, w_attn_v, w_attn_out, w_ffn_gate, w_ffn_up, w_ffn_down; -}; - -class JiugeGPTQWeights : public infinicore::weights::Loader { -private: - std::vector> _device_weights; - -public: - JiugeGPTQWeights(const JiugeGPTQMeta *meta, - infiniDevice_t device, - const std::vector &dev_ids); - std::vector> &device_weights() { - return _device_weights; - } -}; - -struct GPTQDeviceResource { - // Device - infiniDevice_t device; - int device_id; - infiniopHandle_t handle; - // Weights - std::shared_ptr weights; - // Streams - infinirtStream_t stream; - // Communicator - infinicclComm_t comm; - - std::shared_ptr memory_pool; -}; - -struct InferRequest { - const uint32_t *tokens; - uint32_t ntok; - const uint32_t *req_lens; - uint32_t nreq; - const uint32_t *req_pos; - struct KVCache **kv_caches; - const float *temperature; - const uint32_t *topk; - const float *topp; - uint32_t *output; - void *logits; -}; - -struct InferState { - std::mutex mtx; - std::condition_variable cv_load, cv_start, cv_done; - bool loaded = false; - bool proceed = false; - bool exit_flag = false; -}; - -struct JiugeGPTQModel { - JiugeGPTQMeta meta; - infiniDevice_t device; - std::vector dev_ids; - std::vector dev_resources; - std::vector states; - std::vector threads; - InferRequest req; - - JiugeGPTQModel(const JiugeGPTQMeta *, const ModelWeights *); -}; diff --git a/src/models/jiuge_gptq/jiuge_gptq_weight.cpp b/src/models/jiuge_gptq/jiuge_gptq_weight.cpp deleted file mode 100644 index 2e3c2a054..000000000 --- a/src/models/jiuge_gptq/jiuge_gptq_weight.cpp +++ /dev/null @@ -1,137 +0,0 @@ -#include "jiuge_gptq.hpp" - -#include - -inline std::shared_ptr getSinTable(size_t dctx, size_t dh, float theta) { - auto half_dh = dh / 2; - auto unit = dsize(INFINI_DTYPE_F16); - void *table = std::malloc(dctx * half_dh * unit); - - for (size_t i = 0; i < dctx; i++) { - for (size_t j = 0; j < half_dh; j++) { - float _sin = std::sin( - static_cast(i) / std::pow(theta, static_cast(j) / half_dh)); - - ((uint16_t *)table)[i * half_dh + j] = f32_to_f16(_sin); - } - } - auto shape = std::vector({dctx, half_dh}); - auto tensor = Tensor::weight(table, INFINI_DTYPE_F16, shape); - std::free(table); - return tensor; -} - -inline std::shared_ptr getCosTable(size_t dctx, size_t dh, float theta) { - auto half_dh = dh / 2; - auto unit = dsize(INFINI_DTYPE_F16); - void *table = std::malloc(dctx * half_dh * unit); - - for (size_t i = 0; i < dctx; i++) { - for (size_t j = 0; j < half_dh; j++) { - float _cos = std::cos( - static_cast(i) / std::pow(theta, static_cast(j) / half_dh)); - - ((uint16_t *)table)[i * half_dh + j] = f32_to_f16(_cos); - } - } - auto shape = std::vector({dctx, half_dh}); - auto tensor = Tensor::weight(table, INFINI_DTYPE_F16, shape); - std::free(table); - return tensor; -} - -JiugeGPTQWeights::JiugeGPTQWeights( - const JiugeGPTQMeta *meta, - infiniDevice_t device, - const std::vector &dev_ids) : infinicore::weights::Loader(device, dev_ids) { - auto ndev = dev_ids.size(); - _device_weights.resize(ndev); - infiniDtype_t dt_logits = meta->dt_logits; - infiniDtype_t dt_norm_w = meta->dt_norm_w; - size_t nlayer = meta->nlayer; - size_t d = meta->d; - size_t nh = meta->nh / ndev; - size_t nkvh = meta->nkvh / ndev; - size_t dh = meta->dh; - size_t di = meta->di / ndev; - size_t dctx = meta->dctx; - size_t dvoc = meta->dvoc; - size_t nbit = meta->nbit; - size_t quant_group_size = meta->quant_group_size; - - for (size_t i = 0; i < ndev; i++) { - RUN_INFINI(infinirtSetDevice(device, dev_ids[i])); - - auto weight = std::make_shared(); - _device_weights[i] = weight; - - auto w_in_embd = Tensor::weight(nullptr, dt_logits, {dvoc, d}); - this->register_weight("model.embed_tokens.weight", w_in_embd, i); - weight->w_in_embd = w_in_embd; - - auto w_out_norm = Tensor::weight(nullptr, dt_norm_w, {d}); - this->register_weight("model.norm.weight", w_out_norm, i); - weight->w_out_norm = w_out_norm; - - auto w_out_embd = Tensor::weight(nullptr, dt_logits, {dvoc, d})->permute({1, 0}); - this->register_weight("lm_head.weight", w_out_embd, i); - weight->w_out_embd = w_out_embd; - - weight->sin_table = getSinTable(dctx, dh, meta->theta); - weight->cos_table = getCosTable(dctx, dh, meta->theta); - - for (size_t layer = 0; layer < nlayer; layer++) { - -#define RIGISTER_LAYER_WEIGHT(W_NAME, W_VAR, W_SHAPE, W_DTYPE, W_DIST_TYPE) \ - auto W_VAR = Tensor::weight(nullptr, W_DTYPE, W_SHAPE); \ - this->register_weight(W_NAME, W_VAR, i, infinicore::weights::DistributionType::W_DIST_TYPE); \ - weight->W_VAR.push_back(W_VAR); - - RIGISTER_LAYER_WEIGHT("model.layers." + std::to_string(layer) + ".input_layernorm.weight", w_attn_norm, {d}, dt_norm_w, FULL); - -#define REGISTER_LAYER_QUANT_WEIGHT(W_NAME, W_VAR, W_IN, W_OUT, W_DIST_TYPE) \ - auto W_VAR = std::make_shared(); \ - /* GPTQ layout: qweight[in_packed=W_IN/8, out_features=W_OUT]; zeros/scales grouped by input */ \ - W_VAR->w = Tensor::weight(nullptr, INFINI_DTYPE_I32, {(W_IN)*nbit / 32, (W_OUT)}); \ - this->register_weight(W_NAME + ".qweight", W_VAR->w, i, infinicore::weights::DistributionType::W_DIST_TYPE); \ - W_VAR->s = Tensor::weight(nullptr, INFINI_DTYPE_F16, {(W_IN) / quant_group_size, (W_OUT)}); \ - this->register_weight(W_NAME + ".scales", W_VAR->s, i, infinicore::weights::DistributionType::W_DIST_TYPE); \ - W_VAR->z = Tensor::weight(nullptr, INFINI_DTYPE_I32, {(W_IN) / quant_group_size, (W_OUT)*nbit / 32}); \ - this->register_weight(W_NAME + ".qzeros", W_VAR->z, i, infinicore::weights::DistributionType::W_DIST_TYPE); \ - W_VAR->g_idx = Tensor::weight(nullptr, INFINI_DTYPE_I32, {(W_IN)}); \ - auto W_VAR##_gidx_dist = (infinicore::weights::DistributionType::W_DIST_TYPE == infinicore::weights::DistributionType::ROW) ? infinicore::weights::DistributionType::ROW : infinicore::weights::DistributionType::FULL; \ - this->register_weight(W_NAME + ".g_idx", W_VAR->g_idx, i, W_VAR##_gidx_dist); \ - weight->W_VAR.push_back(W_VAR); - - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.q_proj", w_attn_q, d, nh * dh, COLUMN); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.k_proj", w_attn_k, d, nkvh * dh, COLUMN); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.v_proj", w_attn_v, d, nkvh * dh, COLUMN); - RIGISTER_LAYER_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.q_proj.bias", b_attn_q, {nh * dh}, INFINI_DTYPE_F16, COLUMN); - RIGISTER_LAYER_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.k_proj.bias", b_attn_k, {nkvh * dh}, INFINI_DTYPE_F16, COLUMN); - RIGISTER_LAYER_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.v_proj.bias", b_attn_v, {nkvh * dh}, INFINI_DTYPE_F16, COLUMN); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".self_attn.o_proj", w_attn_out, nh * dh, d, ROW); - - RIGISTER_LAYER_WEIGHT("model.layers." + std::to_string(layer) + ".post_attention_layernorm.weight", w_ffn_norm, {d}, dt_norm_w, FULL); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".mlp.gate_proj", w_ffn_gate, d, di, COLUMN); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".mlp.up_proj", w_ffn_up, d, di, COLUMN); - REGISTER_LAYER_QUANT_WEIGHT("model.layers." + std::to_string(layer) + ".mlp.down_proj", w_ffn_down, di, d, ROW); - } - } - -#undef RIGISTER_LAYER_WEIGHT -#undef REGISTER_LAYER_QUANT_WEIGHT -} - -__INFINI_C struct ModelWeights * -createJiugeGPTQWeights(const JiugeGPTQMeta *meta, - infiniDevice_t device, - int ndev, - const int *dev_ids) { - JiugeGPTQWeights *weights = new JiugeGPTQWeights(meta, device, std::vector(dev_ids, dev_ids + ndev)); - return (struct ModelWeights *)weights; -} - -// 建立外部别名, 防止python端冲突 -__INFINI_C void JiugeGPTQLoadWeight(struct ModelWeights *weights, const char *name, void *data) { - loadModelWeight(weights, name, data); -} diff --git a/src/models/qwen3vl/qwen3vl.cpp b/src/models/qwen3vl/qwen3vl.cpp deleted file mode 100644 index 08568e2a0..000000000 --- a/src/models/qwen3vl/qwen3vl.cpp +++ /dev/null @@ -1,715 +0,0 @@ -#include "qwen3vl_impl.hpp" - -#include "../../tensor.hpp" -#include "../../utils.hpp" -#include "../inference_context.hpp" -#include "infinicore_infer.h" - -#include -#include -#include - -void createDeviceResource(Qwen3vlDeviceResource *rsrc, const Qwen3vlMeta *meta, - std::shared_ptr weights, - infiniDevice_t device, int idev, - int ndev, int dev_id, - infinicclComm_t comm) { - RUN_INFINI(infinirtSetDevice(device, dev_id)); - RUN_INFINI(infinirtStreamSynchronize(weights->load_stream)); - infiniopHandle_t handle; - infiniopCreateHandle(&handle); - infinirtStream_t stream; - infinirtStreamCreate(&stream); - - auto memory_pool = std::make_shared(); - - *rsrc = Qwen3vlDeviceResource{ - device, - dev_id, - handle, - weights, - stream, - comm, - memory_pool, - }; - RUN_INFINI(infinirtDeviceSynchronize()); -} - -void releaseDeviceResource(Qwen3vlDeviceResource &res) { - infinirtDeviceSynchronize(); - - res.weights.reset(); - - infiniopDestroyHandle(res.handle); - res.handle = nullptr; - infinirtStreamDestroy(res.stream); - res.stream = nullptr; - infinicclCommDestroy(res.comm); - res.comm = nullptr; -} - -inline std::shared_ptr get_custom_SinTable(const Qwen3vlMeta &meta, std::vector> &pos_ids, uint32_t dim, size_t theta) { - // pos_ids shape:[seq, dim/2] , pos ids acting on each dim - auto unit = dsize(meta.dtype); - auto half_dim = dim / 2; - size_t len = pos_ids.size(); - void *table = std::malloc(len * half_dim * unit); - - for (size_t i = 0; i < len; i++) { - for (size_t j = 0; j < half_dim; j++) { - float _cos = std::sin( - static_cast(pos_ids[i][j]) / std::pow(theta, static_cast(j) / half_dim)); - if (meta.dtype == INFINI_DTYPE_F16) { - ((uint16_t *)table)[i * half_dim + j] = f32_to_f16(_cos); - } else if (meta.dtype == INFINI_DTYPE_BF16) { - ((uint16_t *)table)[i * half_dim + j] = f32_to_bf16(_cos); - } else if (meta.dtype == INFINI_DTYPE_F32) { - ((float *)table)[i * half_dim + j] = _cos; - } else { - std::cout << "unsupported data type" << std::endl; - exit(1); - } - } - } - auto shape = std::vector({len, half_dim}); - auto tensor = Tensor::weight(table, meta.dtype, shape); - std::free(table); - return tensor; -} - -inline std::shared_ptr get_custom_CosTable(const Qwen3vlMeta &meta, std::vector> &pos_ids, uint32_t dim, size_t theta) { - // pos_ids shape:[seq, dim/2] , pos ids acting on each dim - auto unit = dsize(meta.dtype); - auto half_dim = dim / 2; - size_t len = pos_ids.size(); - void *table = std::malloc(len * half_dim * unit); - - for (size_t i = 0; i < len; i++) { - for (size_t j = 0; j < half_dim; j++) { - float _cos = std::cos( - static_cast(pos_ids[i][j]) / std::pow(theta, static_cast(j) / half_dim)); - if (meta.dtype == INFINI_DTYPE_F16) { - ((uint16_t *)table)[i * half_dim + j] = f32_to_f16(_cos); - } else if (meta.dtype == INFINI_DTYPE_BF16) { - ((uint16_t *)table)[i * half_dim + j] = f32_to_bf16(_cos); - } else if (meta.dtype == INFINI_DTYPE_F32) { - ((float *)table)[i * half_dim + j] = _cos; - } else { - std::cout << "unsupported data type" << std::endl; - exit(1); - } - } - } - auto shape = std::vector({len, half_dim}); - auto tensor = Tensor::weight(table, meta.dtype, shape); - std::free(table); - return tensor; -} - -inline std::shared_ptr fast_pos_embed_interpolate(const Qwen3vlMeta &meta, Qwen3vlDeviceResource &rsrc, - uint32_t *grid_thw, uint32_t num_batch, uint32_t total_patches) { - auto dtype = meta.dtype; - auto num_position_embeddings = meta.vis_meta.num_position_embeddings; - auto hidden_size = meta.vis_meta.hidden_size; - auto merge_size = meta.vis_meta.spatial_merge_size; - auto num_grid_per_side = static_cast(sqrt(num_position_embeddings)); - - uint32_t total_pixels_offset = 0; - std::shared_ptr patch_pos_embeds = Tensor::buffer(dtype, {total_patches, hidden_size}, rsrc.memory_pool); - auto pos_embed_weight = rsrc.weights->w_vis->pos_embed_weight; - - std::vector> pos_embeds(4); - for (uint32_t i = 0; i < num_batch; ++i) { - uint32_t t = grid_thw[i * 3]; - uint32_t h = grid_thw[i * 3 + 1]; - uint32_t w = grid_thw[i * 3 + 2]; - auto weight_array = std::vector(h * w * hidden_size); - auto weight_tensor = Tensor::buffer(dtype, {h * w, hidden_size}, rsrc.memory_pool); - - // 计算插值索引和权重 - std::vector> indices(4); - std::vector> weights(4); - - auto linspace = [](float start, float end, uint32_t num_points) -> std::vector { - std::vector res(num_points); - for (uint32_t i = 0; i < num_points; ++i) { - res[i] = start + (end - start) * i / (num_points - 1); - } - return res; - }; - - auto h_idxs = linspace(0, num_grid_per_side - 1, h); - auto w_idxs = linspace(0, num_grid_per_side - 1, w); - - for (uint32_t ih = 0; ih < h; ++ih) { - for (uint32_t iw = 0; iw < w; ++iw) { - float h_idx_f = h_idxs[ih], w_idx_f = w_idxs[iw]; - uint32_t h_idx_floor = static_cast(floor(h_idx_f)), - w_idx_floor = static_cast(floor(w_idx_f)); - uint32_t h_idx_ceil = std::min(static_cast(ceil(h_idx_f)), num_grid_per_side - 1), - w_idx_ceil = std::min(static_cast(ceil(w_idx_f)), num_grid_per_side - 1); - - float dh = h_idx_f - h_idx_floor, dw = w_idx_f - w_idx_floor; - - indices[0].push_back((h_idx_floor * num_grid_per_side) + w_idx_floor); - indices[1].push_back((h_idx_floor * num_grid_per_side) + w_idx_ceil); - indices[2].push_back((h_idx_ceil * num_grid_per_side) + w_idx_floor); - indices[3].push_back((h_idx_ceil * num_grid_per_side) + w_idx_ceil); - - weights[0].push_back((1 - dh) * (1 - dw)); - weights[1].push_back((1 - dh) * dw); - weights[2].push_back(dh * (1 - dw)); - weights[3].push_back(dh * dw); - } - } - - // 查表并加权求和 - for (int j = 0; j < 4; ++j) { - pos_embeds[j] = Tensor::buffer(dtype, {h * w, hidden_size}, rsrc.memory_pool); - // 使用索引和权重获取对应位置嵌入，并乘以权重 - for (size_t i = 0; i < h * w; i++) { - rearrange(pos_embeds[j]->slice(0, i, 1), pos_embed_weight->slice(0, indices[j][i], 1)); - } - for (size_t i = 0; i < h * w; i++) { - uint16_t w_value = f32_to_bf16(weights[j][i]); - for (size_t k = 0; k < hidden_size; k++) { - weight_array[i * hidden_size + k] = w_value; - } - } - RUN_INFINI(infinirtMemcpyAsync(weight_tensor->data(), weight_array.data(), sizeof(uint16_t) * h * w * hidden_size, - INFINIRT_MEMCPY_H2D, rsrc.stream)); - mul(pos_embeds[j], pos_embeds[j], weight_tensor); - } - - // 合并四个方向的结果 - auto patch_pos_embed = pos_embeds[0]; // [h*w, hidden_size] - for (int j = 1; j < 4; ++j) { - add(patch_pos_embed, patch_pos_embed, pos_embeds[j]); - } - - // 对于视频帧数T>1的情况，重复patch_pos_embed T次 - if (t > 1) { - auto temp_patch_pos_embed = Tensor::buffer(dtype, {t, h * w, hidden_size}, rsrc.memory_pool); - for (size_t i = 0; i < t; i++) { - rearrange(temp_patch_pos_embed->slice(0, i, 1), patch_pos_embed); - } - patch_pos_embed = temp_patch_pos_embed; - } - printf("merge patch pos embed/n"); - fflush(stdout); - patch_pos_embed = patch_pos_embed - ->view({t, h / merge_size, merge_size, w / merge_size, merge_size, hidden_size}) - ->permute({0, 1, 3, 2, 4, 5}) - ->view({t * h * w, hidden_size}); // 可能因为内存不连续无法再view - - rearrange(patch_pos_embeds->slice(0, total_pixels_offset, t * h * w), patch_pos_embed); - total_pixels_offset += t * h * w; - } - return patch_pos_embeds; -} - -inline auto rot_pos_embed(const Qwen3vlMeta &meta, Qwen3vlDeviceResource &rsrc, uint32_t *grid_thw, uint32_t num_batch, uint32_t total_patches) { - auto dtype = meta.dtype; - auto hidden_size = meta.vis_meta.hidden_size; - auto num_heads = meta.vis_meta.num_heads; - auto head_dim = hidden_size / num_heads; - auto merge_size = meta.vis_meta.spatial_merge_size; - - std::vector> pos_ids_table_y( - total_patches, - std::vector(head_dim / 4)); - std::vector> pos_ids_table_x( - total_patches, - std::vector(head_dim / 4)); - for (uint32_t b = 0; b < num_batch; ++b) { - uint32_t offset = b * 3; - uint32_t num_frames = grid_thw[offset + 0]; - uint32_t height = grid_thw[offset + 1]; - uint32_t width = grid_thw[offset + 2]; - - uint32_t merged_h = height / merge_size; - uint32_t merged_w = width / merge_size; - - // 遍历所有块和块内位置 - size_t patch_offset = 0; - for (uint32_t bh = 0; bh < merged_h; ++bh) { - for (uint32_t bw = 0; bw < merged_w; ++bw) { - for (uint32_t ih = 0; ih < merge_size; ++ih) { - for (uint32_t iw = 0; iw < merge_size; ++iw) { - uint32_t row = bh * merge_size + ih; - uint32_t col = bw * merge_size + iw; - // 如果是多帧，重复 num_frames 次 - for (uint32_t f = 0; f < num_frames; ++f) { - size_t dim_offset = 0; - for (; dim_offset < head_dim / 4; dim_offset++) { - pos_ids_table_y[patch_offset][dim_offset] = row; - pos_ids_table_x[patch_offset][dim_offset] = col; - } - patch_offset++; - } - } - } - } - } - } - auto sin = Tensor::buffer(dtype, {total_patches, head_dim / 2}, rsrc.memory_pool); - auto sin_y = get_custom_SinTable(meta, pos_ids_table_y, head_dim / 2, 10000); - rearrange(sin->slice(1, 0, head_dim / 4), sin_y); - auto sin_x = get_custom_SinTable(meta, pos_ids_table_x, head_dim / 2, 10000); - rearrange(sin->slice(1, head_dim / 4, head_dim / 2), sin_y); - auto cos = Tensor::buffer(dtype, {total_patches, head_dim / 2}, rsrc.memory_pool); - auto cos_y = get_custom_CosTable(meta, pos_ids_table_y, head_dim / 2, 10000); - rearrange(cos->slice(1, 0, head_dim / 4), cos_y); - auto cos_x = get_custom_CosTable(meta, pos_ids_table_x, head_dim / 2, 10000); - rearrange(cos->slice(1, head_dim / 4, head_dim / 2), cos_y); - - return std::pair{sin, cos}; -} - -void inferDeviceBatchVision(const Qwen3vlMeta &meta, Qwen3vlDeviceResource &rsrc, - uint32_t idev, uint32_t ndev, InferRequest &req) { - void *pixel_values = req.pixel_values; - uint32_t total_patches = req.total_patches; - uint32_t *image_grid_thw = req.image_grid_thw; - uint32_t num_images = req.num_images; - void *pixel_values_videos = req.pixel_values_videos; - uint32_t total_patches_videos = req.total_patches_videos; - // uint32_t *video_grid_thw = req.video_grid_thw; - // uint32_t num_videos = req.num_videos; - // uint32_t patch_features = req.patch_features; - - auto dtype = meta.dtype; - auto d = meta.vis_meta.hidden_size; - auto channels = meta.vis_meta.in_channels; - auto patch_size = meta.vis_meta.patch_size; - auto temporal_patch_size = meta.vis_meta.temporal_patch_size; - // auto stream = rsrc.stream; - auto weights = rsrc.weights; - - auto image_tensor = Tensor::weight(pixel_values, dtype, {total_patches, channels * temporal_patch_size * patch_size * patch_size}); - auto video_tensor = Tensor::weight(pixel_values_videos, dtype, {total_patches_videos, channels * temporal_patch_size * patch_size * patch_size}); - auto hidden_states = Tensor::buffer(dtype, {total_patches, d, 1, 1, 1}, rsrc.memory_pool); - - std::vector pads = {0, 0, 0}; - std::vector strides = {static_cast(temporal_patch_size), static_cast(patch_size), static_cast(patch_size)}; - std::vector dilations = {1, 1, 1}; - conv(hidden_states, image_tensor, rsrc.weights->w_vis->patch_embed_weight, rsrc.weights->w_vis->patch_embed_bias, - pads.data(), strides.data(), dilations.data(), 3); - hidden_states = hidden_states->view({total_patches, d}); - - auto pos_embeds = fast_pos_embed_interpolate(meta, rsrc, image_grid_thw, num_images, total_patches); - add(hidden_states, hidden_states, pos_embeds); - - auto [sin, cos] = rot_pos_embed(meta, rsrc, image_grid_thw, num_images, total_patches); -} - -void inferDeviceBatchText(const Qwen3vlMeta &meta, Qwen3vlDeviceResource &rsrc, - uint32_t idev, uint32_t ndev, InferRequest &req) { - const uint32_t *tokens = req.tokens; - uint32_t ntok = req.ntok; - const uint32_t *req_lens = req.req_lens; - uint32_t nreq = req.nreq; - const uint32_t *req_pos = req.req_pos; - struct Qwen3vlCache **caches = req.kv_caches; - const float *temperature = req.temperature; - const uint32_t *topk = req.topk; - const float *topp = req.topp; - uint32_t *output = req.output; - void *last_logits = req.logits; - - assert(meta.text_meta.num_attention_heads % ndev == 0); - assert(meta.text_meta.num_key_value_heads % ndev == 0); - - auto dtype = meta.dtype; - auto nlayer = meta.text_meta.num_hidden_layers; - size_t nh = meta.text_meta.num_attention_heads / size_t(ndev); - size_t nkvh = meta.text_meta.num_key_value_heads / size_t(ndev); - auto ngroup = nh / nkvh; - auto dh = meta.text_meta.head_dim; - auto d = meta.text_meta.hidden_size; - auto di = meta.text_meta.intermediate_size / size_t(ndev); - auto dvoc = meta.text_meta.vocab_size; - float epsilon = meta.text_meta.rms_norm_eps; - auto stream = rsrc.stream; - auto weights = rsrc.weights; - - // Allocate buffers - auto logits_in = Tensor::buffer(dtype, {ntok, d}, rsrc.memory_pool); - auto logits_out = Tensor::buffer(dtype, {ntok, d}, rsrc.memory_pool); - - // 所有请求的当前token - auto qkv_buf = Tensor::buffer(dtype, {ntok, (nh + nkvh * 2) * dh}, rsrc.memory_pool); - auto o_buf = Tensor::buffer(dtype, {ntok, nh * dh}, rsrc.memory_pool); - auto gate_up_buf = Tensor::buffer(dtype, {ntok, 2 * di}, rsrc.memory_pool); - - auto prob_buf = Tensor::buffer(dtype, {nreq, dvoc}, rsrc.memory_pool); - auto result_buf = Tensor::buffer(INFINI_DTYPE_I64, {nreq}, rsrc.memory_pool); - auto result_cpu = std::vector(nreq); - - auto qkv_rope = qkv_buf->view({ntok, nh + nkvh * 2, dh}); - auto q_buf = qkv_rope->slice(1, 0, nh); - auto k_buf = qkv_rope->slice(1, nh, nkvh); - - // Prepare inputs - auto batch_pos_ids = std::vector(ntok); - size_t req_start = 0; - for (uint32_t req = 0; req < nreq; req++) { - for (uint32_t i = 0; i < req_lens[req]; i++) { // req_len 本次query长度，req_pos 历史长度 - batch_pos_ids[req_start + i] = req_pos[req] + i; // batch_pos_ids 展平后每个token的pos - } - req_start += req_lens[req]; - } - std::shared_ptr pos_ids_buf; - if (rsrc.device == INFINI_DEVICE_CPU) { - pos_ids_buf = Tensor::weight(batch_pos_ids.data(), INFINI_DTYPE_U32, {ntok}); - } else { - pos_ids_buf = Tensor::buffer(INFINI_DTYPE_U32, {ntok}, rsrc.memory_pool); - RUN_INFINI(infinirtMemcpyAsync(pos_ids_buf->data(), batch_pos_ids.data(), sizeof(uint32_t) * ntok, - INFINIRT_MEMCPY_H2D, stream)); - } - - // convert tokens to embeddings - for (uint32_t i = 0; i < ntok; i++) { - RUN_INFINI(infinirtMemcpyAsync(logits_in->data(i * d), - weights->w_lang->in_embd->data(tokens[i] * d), - dsize(dtype) * d, INFINIRT_MEMCPY_D2D, stream)); - } - - // attention inner - size_t max_qk_size = 0; - size_t max_seq_len = 0; - - for (uint32_t req = 0; req < nreq; req++) { - auto past_len = req_pos[req]; - auto seq_len = req_lens[req]; - auto total_len = past_len + seq_len; - - max_qk_size = std::max(max_qk_size, size_t(seq_len * total_len)); - max_seq_len = std::max(max_seq_len, size_t(seq_len)); - } - - auto qk_buf = Tensor::buffer(dtype, {nh * max_qk_size}, rsrc.memory_pool); - auto rearrange_q_buf = Tensor::buffer(dtype, {nkvh, ngroup * max_seq_len, dh}, rsrc.memory_pool); - auto q_rearrange = rearrange_q_buf->view({nkvh, ngroup, max_seq_len, dh}); - auto attn_val_buf = Tensor::buffer(dtype, {nkvh, ngroup * max_seq_len, dh}, rsrc.memory_pool); - auto attn_val_gemm = attn_val_buf->view({nkvh, ngroup, max_seq_len, dh}); - - auto gate_buf = gate_up_buf->slice(1, 0, di); - auto up_buf = gate_up_buf->slice(1, di, di); - - // Compute - for (uint32_t i = 0; i < nlayer; i++) { - // attn norm - rmsnorm(logits_out, logits_in, weights->w_lang->layers[i].attn_norm, epsilon); - // qkv_proj - linear(qkv_buf, logits_out, weights->w_lang->layers[i].attn_qkv_proj, 1.0, 0.0, nullptr, nullptr); - // qk_norm - rmsnorm(q_buf, q_buf, weights->w_lang->layers[i].attn_q_norm, epsilon); - rmsnorm(k_buf, k_buf, weights->w_lang->layers[i].attn_k_norm, epsilon); - // rope - rope_v2(q_buf, q_buf, pos_ids_buf, weights->sin_table, weights->cos_table); - rope_v2(k_buf, k_buf, pos_ids_buf, weights->sin_table, weights->cos_table); - - // 逐个req处理 - size_t token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto past_len = req_pos[req]; - auto seq_len = req_lens[req]; - auto total_len = past_len + seq_len; - - auto o = o_buf->slice(0, token_offset, seq_len)->view({seq_len, nkvh, ngroup, dh})->permute({1, 2, 0, 3}); // [nkvh, ngroup, seq_len, dh] - auto q = qkv_rope->slice({{0, token_offset, seq_len}, {1, 0, nh}})->view({seq_len, nkvh, ngroup, dh})->permute({1, 2, 0, 3}); // [nkvh, ngroup, seq_len, dh] - auto k = qkv_rope->slice({{0, token_offset, seq_len}, {1, nh, nkvh}}); // [ntok, nkvh, dh] - auto v = qkv_rope->slice({{0, token_offset, seq_len}, {1, nh + nkvh, nkvh}}); // [ntok, nkvh, dh] - - // concat to cache - rearrange(caches[req]->k_rot[idev][i]->slice(0, past_len, seq_len), k); - rearrange(caches[req]->v[idev][i]->slice(0, past_len, seq_len), v); - - // fill full_k full_v - auto full_k_buff = caches[req]->k_rot[idev][i]->slice(0, 0, total_len)->permute({1, 2, 0}); // [nkvh, dh, total_len] - auto full_v_buff = caches[req]->v[idev][i]->slice(0, 0, total_len)->permute({1, 0, 2}); // [nkvh, total_len, dh] - - // self-attn - rearrange(q_rearrange->slice(2, 0, seq_len), q); - auto attn_score_req = qk_buf->slice(0, 0, nh * seq_len * total_len)->view({nkvh, ngroup * seq_len, total_len}); - // [nkvh, ngroup * seq_len, dh] @ [nkvh, dh, total_len] = [nkvh, ngroup * seq_len, total_len] - linear(attn_score_req, rearrange_q_buf->slice(1, 0, ngroup * seq_len), full_k_buff, 1.f / float(sqrt(dh)), 0.f, nullptr, nullptr); - // softmax - auto qk_softmax = attn_score_req->view({nh, seq_len, total_len}); - causalSoftmax(qk_softmax, qk_softmax); - // [nkvh, ngroup * seq_len, total_len] @ [nkvh, total_len, dh] = [nkvh, ngroup * seq_len, dh] - linear(attn_val_buf->slice(1, 0, ngroup * seq_len), attn_score_req, full_v_buff, 1.0, 0.0, nullptr, nullptr); - // printf("rearrage o; layer[%d]\n",i); - rearrange(o, attn_val_gemm->slice(2, 0, seq_len)); - token_offset += seq_len; - } - linear(logits_in, o_buf, weights->w_lang->layers[i].attn_o_proj, 1.0, 0.0, idev == 0 ? logits_in : nullptr, nullptr); - // All_reduce if distributed - if (rsrc.comm != nullptr) { - RUN_INFINI(infinicclAllReduce( - logits_in->data(), logits_in->data(), ntok * d, dtype, - INFINICCL_SUM, rsrc.comm, stream)); - RUN_INFINI(infinirtStreamSynchronize(stream)); - } - - // mlp norm - rmsnorm(logits_out, logits_in, weights->w_lang->layers[i].mlp_norm, epsilon); - // mlp gate_up - linear(gate_up_buf, logits_out, weights->w_lang->layers[i].mlp_gate_up, 1.0, 0.0, nullptr, nullptr); - // silu - silu(gate_buf, gate_buf); - mul(gate_buf, gate_buf, up_buf); - // mlp down - linear(logits_in, gate_buf, weights->w_lang->layers[i].mlp_down, 1.0, 0.0, idev == 0 ? logits_in : nullptr, nullptr); - // All_reduce if distributed - if (rsrc.comm != nullptr) { - RUN_INFINI(infinicclAllReduce( - logits_in->data(), logits_in->data(), ntok * d, dtype, - INFINICCL_SUM, rsrc.comm, stream)); - RUN_INFINI(infinirtStreamSynchronize(stream)); - } - } - // sample and output - if (idev == 0) { - if (last_logits != nullptr) { - rmsnorm(logits_out, logits_in, weights->w_lang->out_norm, epsilon); - auto last_logits_buf = Tensor::buffer(dtype, {ntok, dvoc}, rsrc.memory_pool); - linear(last_logits_buf, logits_out, weights->w_lang->out_embd, 1.0, 0.0, nullptr, nullptr); - RUN_INFINI(infinirtStreamSynchronize(stream)); - RUN_INFINI(infinirtMemcpy(last_logits, last_logits_buf->data(), dsize(dtype) * ntok * dvoc, INFINIRT_MEMCPY_D2H)); - } - if (output != nullptr) { - size_t token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto seq_len = req_lens[req]; - token_offset += seq_len; - rmsnorm(logits_out->slice(0, req, 1), - logits_in->slice(0, token_offset - 1, 1), - weights->w_lang->out_norm, - epsilon); - } - linear(prob_buf, logits_out->slice(0, 0, nreq), weights->w_lang->out_embd, 1.0, 0.0, nullptr, nullptr); - std::random_device _rd; - std::mt19937 gen(_rd()); - token_offset = 0; - for (uint32_t req = 0; req < nreq; req++) { - auto seq_len = req_lens[req]; - float random_val = std::uniform_real_distribution(0, 1)(gen); - randomSample(result_buf->slice(0, req, 1)->view_as({}, {}), - prob_buf->slice(0, req, 1)->view_as({dvoc}, {1}), - random_val, topp[req], topk[req], temperature[req]); - token_offset += seq_len; - } - RUN_INFINI(infinirtStreamSynchronize(stream)); - RUN_INFINI(infinirtMemcpy(result_cpu.data(), result_buf->data(), - sizeof(int64_t) * nreq, INFINIRT_MEMCPY_D2H)); - for (uint32_t req = 0; req < nreq; req++) { - output[req] = uint32_t(result_cpu[req]); - } - } - } -} - -void inferDeviceBatch(const Qwen3vlMeta &meta, Qwen3vlDeviceResource &rsrc, - uint32_t idev, uint32_t ndev, InferState &state, InferRequest &req) { - // infer vision + sync - if (req.num_images > 0 || req.num_videos > 0) { - inferDeviceBatchVision(meta, rsrc, idev, ndev, req); - - std::unique_lock lock(state.mtx_sync); - state.sync_cnt--; - if (state.sync_cnt == 0) { - state.cv_sync.notify_all(); - } else { - state.cv_sync.wait(lock, [&] { return state.sync_cnt == 0; }); - } - } - // infer text - inferDeviceBatchText(meta, rsrc, idev, ndev, req); -} - -__INFINI_C void -inferBatchQwen3vl(struct Qwen3vlModel *model, - const uint32_t *tokens, uint32_t ntok, - void *pixel_values, uint32_t total_patches, - uint32_t *image_grid_thw, uint32_t num_images, - void *pixel_values_videos, uint32_t total_patches_videos, - uint32_t *video_grid_thw, uint32_t num_videos, - uint32_t patch_features, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct Qwen3vlCache **kv_caches, - const float *temperature, const uint32_t *topk, const float *topp, - uint32_t *output) { - model->req.tokens = tokens; - model->req.ntok = ntok; - model->req.pixel_values = pixel_values; - model->req.total_patches = total_patches; - model->req.image_grid_thw = image_grid_thw; - model->req.num_images = num_images; - model->req.pixel_values_videos = pixel_values_videos; - model->req.total_patches_videos = total_patches_videos; - model->req.video_grid_thw = video_grid_thw; - model->req.num_videos = num_videos; - model->req.patch_features = patch_features; - model->req.req_lens = req_lens; - model->req.nreq = nreq; - model->req.req_pos = req_pos; - model->req.kv_caches = kv_caches; - model->req.output = output; - model->req.logits = nullptr; - model->req.temperature = temperature; - model->req.topk = topk; - model->req.topp = topp; - model->states[0].sync_cnt = model->dev_ids.size(); - - for (size_t idev = 0; idev < model->dev_ids.size(); idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].proceed = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - for (size_t i = model->dev_ids.size(); i > 0; i--) { - auto idev = i - 1; - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].cv_done.wait(lock, [&] { return !(model->states[idev].proceed); }); - lock.unlock(); - } -} - -__INFINI_C void -forwardBatchQwen3vl(struct Qwen3vlModel *model, - const uint32_t *tokens, uint32_t ntok, - void *pixel_values, uint32_t total_patches, - uint32_t *image_grid_thw, uint32_t num_images, - void *pixel_values_videos, uint32_t total_patches_videos, - uint32_t *video_grid_thw, uint32_t num_videos, - uint32_t patch_features, - const uint32_t *req_lens, uint32_t nreq, const uint32_t *req_pos, - struct Qwen3vlCache **kv_caches, - void *logits) { - model->req.tokens = tokens; - model->req.ntok = ntok; - model->req.pixel_values = pixel_values; - model->req.total_patches = total_patches; - model->req.image_grid_thw = image_grid_thw; - model->req.num_images = num_images; - model->req.pixel_values_videos = pixel_values_videos; - model->req.total_patches_videos = total_patches_videos; - model->req.video_grid_thw = video_grid_thw; - model->req.num_videos = num_videos; - model->req.patch_features = patch_features; - model->req.req_lens = req_lens; - model->req.nreq = nreq; - model->req.req_pos = req_pos; - model->req.kv_caches = kv_caches; - model->req.output = nullptr; - model->req.logits = logits; - model->req.temperature = nullptr; - model->req.topk = nullptr; - model->req.topp = nullptr; - model->states[0].sync_cnt = model->dev_ids.size(); - - for (size_t idev = 0; idev < model->dev_ids.size(); idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].proceed = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - for (size_t i = model->dev_ids.size(); i > 0; i--) { - auto idev = i - 1; - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].cv_done.wait(lock, [&] { return !(model->states[idev].proceed); }); - lock.unlock(); - } -} - -void launchDevice(const Qwen3vlMeta &meta, std::shared_ptr weights, Qwen3vlDeviceResource *rsrc, InferState &state, InferRequest &req, - infiniDevice_t device, int idev, int ndev, int dev_id, infinicclComm_t comm) { - // Create Device Resource - createDeviceResource(rsrc, &meta, weights, device, idev, ndev, dev_id, comm); - - CacheManager cache_manager(100); - InferenceContext ctx(rsrc->handle, rsrc->memory_pool, &cache_manager, rsrc->stream); - - // Set the inference context for this thread - setInferenceContext(&ctx); - - { - std::unique_lock lock(state.mtx); - state.loaded = true; - lock.unlock(); - state.cv_load.notify_one(); - } - - // Infer Loop - while (true) { - std::unique_lock lock(state.mtx); - state.cv_start.wait(lock, [&] { return state.proceed || state.exit_flag; }); - // quit if exit_flag is set - if (state.exit_flag) { - break; - } - - inferDeviceBatch(meta, *rsrc, idev, ndev, state, req); - - state.proceed = false; - lock.unlock(); - state.cv_done.notify_one(); - } - - // Clean-Up - releaseDeviceResource(*rsrc); - setInferenceContext(nullptr); // Clear the context when done -} - -Qwen3vlModel::Qwen3vlModel(const Qwen3vlMeta *_meta, const Qwen3vlWeights *weights) : meta(*_meta) { - auto device_weights = weights->device_weights; - int ndev = device_weights.size(); - device = device_weights[0]->device; - dev_ids.resize(ndev); - for (int i = 0; i < ndev; i++) { - dev_ids[i] = device_weights[i]->dev_id; - } - dev_resources = std::vector(ndev); - states = std::vector(ndev); - threads.resize(ndev); - RUN_INFINI(infinirtInit()); - auto comms = std::vector(ndev, nullptr); - if (ndev > 1) { - RUN_INFINI(infinicclCommInitAll(device, comms.data(), ndev, dev_ids.data())); - } - for (int i = 0; i < ndev; i++) { - threads[i] = std::thread(launchDevice, std::cref(meta), device_weights[i], &dev_resources[i], std::ref(states[i]), std::ref(req), device, i, ndev, dev_ids[i], comms[i]); - } - for (int i = 0; i < ndev; i++) { - std::unique_lock lock(states[i].mtx); - states[i].cv_load.wait(lock, [&] { return states[i].loaded; }); - lock.unlock(); - } -} - -__INFINI_C struct Qwen3vlModel * -createQwen3vlModel(const Qwen3vlMeta *_meta, - const Qwen3vlWeights *weights) { - Qwen3vlModel *model = new Qwen3vlModel(_meta, weights); - return model; -} - -__INFINI_C void -destroyQwen3vlModel(struct Qwen3vlModel *model) { - auto ndev = model->dev_resources.size(); - - for (size_t idev = 0; idev < ndev; idev++) { - std::unique_lock lock(model->states[idev].mtx); - model->states[idev].exit_flag = true; - lock.unlock(); - model->states[idev].cv_start.notify_one(); - } - - for (size_t idev = 0; idev < ndev; idev++) { - model->threads[idev].join(); - } - - delete model; -} diff --git a/src/models/qwen3vl/qwen3vl_cache.cpp b/src/models/qwen3vl/qwen3vl_cache.cpp deleted file mode 100644 index d34230f82..000000000 --- a/src/models/qwen3vl/qwen3vl_cache.cpp +++ /dev/null @@ -1,43 +0,0 @@ -#include "qwen3vl_impl.hpp" - -__INFINI_C struct Qwen3vlCache * -createQwen3vlCache(const struct Qwen3vlModel *model) { - Qwen3vlCache *cache = new Qwen3vlCache(); - auto ndev = model->dev_resources.size(); - auto nlayer = model->meta.text_meta.num_hidden_layers; - auto max_len = model->meta.text_meta.max_tokens; - auto dh = model->meta.text_meta.head_dim; - auto nkv = model->meta.text_meta.num_key_value_heads / size_t(ndev); - auto k_rot_shape = std::vector{max_len, nkv, dh}; - auto v_shape = std::vector{max_len, nkv, dh}; - for (size_t idev = 0; idev < ndev; idev++) { - RUN_INFINI(infinirtSetDevice(model->device, model->dev_ids[idev])); - auto k_rot_cache = std::vector>(); - auto v_cache = std::vector>(); - for (size_t layer = 0; layer < nlayer; layer++) { - k_rot_cache.push_back(std::move(Tensor::buffer(model->meta.dtype, k_rot_shape))); - v_cache.push_back(std::move(Tensor::buffer(model->meta.dtype, v_shape))); - } - cache->k_rot.push_back(k_rot_cache); - cache->v.push_back(v_cache); - } - - return cache; -} - -//////还有visual deepstack需要cache？ - -__INFINI_C void -dropQwen3vlCache(const struct Qwen3vlModel *model, - struct Qwen3vlCache *cache) { - auto ndev = model->dev_resources.size(); - auto nlayer = model->meta.text_meta.num_hidden_layers; - for (size_t idev = 0; idev < ndev; idev++) { - RUN_INFINI(infinirtSetDevice(model->device, model->dev_ids[idev])); - for (size_t layer = 0; layer < nlayer; layer++) { - cache->k_rot[idev][layer].reset(); - cache->v[idev][layer].reset(); - } - } - delete cache; -} diff --git a/src/models/qwen3vl/qwen3vl_impl.hpp b/src/models/qwen3vl/qwen3vl_impl.hpp deleted file mode 100644 index 01fda1ceb..000000000 --- a/src/models/qwen3vl/qwen3vl_impl.hpp +++ /dev/null @@ -1,141 +0,0 @@ -#ifndef QWEN3VL_IMPL_H -#define QWEN3VL_IMPL_H - -#include "infinicore_infer.h" - -#include "../../allocator.hpp" -#include "../../tensor.hpp" - -#include -#include -#include -#include -#include - -struct Qwen3vlLayerWeight { - std::shared_ptr attn_norm; - std::shared_ptr attn_qkv_proj; - std::shared_ptr attn_q_norm; - std::shared_ptr attn_k_norm; - std::shared_ptr attn_o_proj; - - std::shared_ptr mlp_norm; - std::shared_ptr mlp_gate_up; - std::shared_ptr mlp_down; -}; - -struct Qwen3vlLanguageModelWeight { - std::shared_ptr in_embd, out_embd, out_norm; - std::vector layers; -}; - -struct Qwen3vlVisBlockWeight { - std::shared_ptr attn_proj_weight, attn_proj_bias, attn_qkv_weight, attn_qkv_bias; - std::shared_ptr mlp_linear_fc1_weight, mlp_linear_fc1_bias, mlp_linear_fc2_weight, mlp_linear_fc2_bias; - std::shared_ptr norm1_weight, norm1_bias, norm2_weight, norm2_bias; -}; - -struct DeepstackMergerWeight { - std::shared_ptr linear_fc1_weight, linear_fc1_bias, linear_fc2_weight, linear_fc2_bias; - std::shared_ptr norm_weight, norm_bias; -}; - -struct MergerWeight { - std::shared_ptr linear_fc1_weight, linear_fc1_bias, linear_fc2_weight, linear_fc2_bias; - std::shared_ptr norm_weight, norm_bias; -}; - -struct Qwen3vlVisualEncoderWeight { - std::shared_ptr patch_embed_weight, patch_embed_bias, pos_embed_weight; - std::vector blocks; - std::vector deepstack_mergers; - std::shared_ptr merger; -}; - -struct Qwen3vlDeviceWeights { - std::shared_ptr sin_table, cos_table; - std::shared_ptr w_lang; - std::shared_ptr w_vis; - infiniDevice_t device; - int dev_id; - infinirtStream_t load_stream; -}; - -struct Qwen3vlWeights { - Qwen3vlMeta const *meta; - bool transpose_weight; - std::vector> device_weights; - - Qwen3vlWeights(const Qwen3vlMeta *meta, - infiniDevice_t device, - int ndev, - const int *dev_ids, - bool transpose_weight); -}; - -struct Qwen3vlDeviceResource { - // Device - infiniDevice_t device; - int device_id; - infiniopHandle_t handle; - // Weights - std::shared_ptr weights; - // Streams - infinirtStream_t stream; - // Communicator - infinicclComm_t comm; - - std::shared_ptr memory_pool; -}; - -struct InferState { // qwen3vl namespace - inline static std::mutex mtx_sync; - inline static int sync_cnt; - inline static std::condition_variable cv_sync; - std::mutex mtx; - std::condition_variable cv_load, cv_start, cv_done; - bool loaded = false; - bool proceed = false; - bool exit_flag = false; -}; - -struct InferRequest { // qwen3vl namespace - const uint32_t *tokens; - uint32_t ntok; - void *pixel_values; - uint32_t total_patches; - uint32_t *image_grid_thw; - uint32_t num_images; - void *pixel_values_videos; - uint32_t total_patches_videos; - uint32_t *video_grid_thw; - uint32_t num_videos; - uint32_t patch_features; - const uint32_t *req_lens; - uint32_t nreq; - const uint32_t *req_pos; - struct Qwen3vlCache **kv_caches; - const float *temperature; - const uint32_t *topk; - const float *topp; - uint32_t *output; - void *logits; -}; - -struct Qwen3vlModel { - Qwen3vlMeta meta; - infiniDevice_t device; - std::vector dev_ids; - std::vector dev_resources; - std::vector states; - std::vector threads; - InferRequest req; - - Qwen3vlModel(const Qwen3vlMeta *, const Qwen3vlWeights *weights); -}; - -struct Qwen3vlCache { - std::vector>> k_rot, v; -}; - -#endif diff --git a/src/models/qwen3vl/qwen3vl_weight.cpp b/src/models/qwen3vl/qwen3vl_weight.cpp deleted file mode 100644 index b673c5d61..000000000 --- a/src/models/qwen3vl/qwen3vl_weight.cpp +++ /dev/null @@ -1,636 +0,0 @@ -#include "qwen3vl_impl.hpp" - -#include - -inline std::shared_ptr getInEmbd( - const Qwen3vlMeta *meta) { - auto shape = std::vector({meta->text_meta.vocab_size, meta->text_meta.hidden_size}); - return Tensor::weight(nullptr, meta->dtype, shape); -} - -inline std::shared_ptr getOutNorm( - const Qwen3vlMeta *meta) { - auto shape = std::vector({meta->text_meta.hidden_size}); - return Tensor::weight(nullptr, meta->dtype, shape); -} - -inline std::shared_ptr getOutEmbd( - const Qwen3vlMeta *meta) { - - auto shape = std::vector({meta->text_meta.vocab_size, meta->text_meta.hidden_size}); - return Tensor::weight(nullptr, meta->dtype, shape) - ->permute({1, 0}); -} - -inline void getLayerWeight( - const Qwen3vlMeta *meta, Qwen3vlLayerWeight &layer, int ndev) { - auto nkvh = meta->text_meta.num_key_value_heads; - auto nh = meta->text_meta.num_attention_heads; - auto dh = meta->text_meta.head_dim; - auto d = meta->text_meta.hidden_size; - auto di = meta->text_meta.intermediate_size; - - auto dh_shape = std::vector({meta->text_meta.hidden_size}); - layer.attn_norm = Tensor::weight(nullptr, meta->dtype, dh_shape); - auto qk_norm_shape = std::vector({meta->text_meta.head_dim}); - layer.attn_q_norm = Tensor::weight(nullptr, meta->dtype, qk_norm_shape); - layer.attn_k_norm = Tensor::weight(nullptr, meta->dtype, qk_norm_shape); - auto qkv_proj_shape = std::vector({(nh + 2 * nkvh) / ndev * dh, d}); - layer.attn_qkv_proj = Tensor::weight(nullptr, meta->dtype, qkv_proj_shape); - auto o_proj_shape = std::vector({d, nh / ndev * dh}); - layer.attn_o_proj = Tensor::weight(nullptr, meta->dtype, o_proj_shape); - - layer.mlp_norm = Tensor::weight(nullptr, meta->dtype, dh_shape); - auto up_shape = std::vector({2 * di / ndev, d}); - layer.mlp_gate_up = Tensor::weight(nullptr, meta->dtype, up_shape); - auto down_shape = std::vector({d, di / ndev}); - layer.mlp_down = Tensor::weight(nullptr, meta->dtype, down_shape); -} - -inline void getVisualWeight( - const Qwen3vlMeta *meta, std::shared_ptr w_vis) { - Qwen3vlVisMeta vis_meta = meta->vis_meta; - auto patch_embed_shape = std::vector({vis_meta.hidden_size, vis_meta.in_channels, vis_meta.temporal_patch_size, vis_meta.patch_size, vis_meta.patch_size}); - w_vis->patch_embed_weight = Tensor::weight(nullptr, meta->dtype, patch_embed_shape); - w_vis->patch_embed_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size}); - w_vis->pos_embed_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.num_position_embeddings, vis_meta.hidden_size}); - w_vis->merger = std::make_shared(); - w_vis->merger->linear_fc1_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.intermediate_size, vis_meta.intermediate_size}); - w_vis->merger->linear_fc2_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.out_hidden_size, vis_meta.intermediate_size}); - w_vis->merger->linear_fc1_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.intermediate_size}); - w_vis->merger->linear_fc2_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.out_hidden_size}); - w_vis->merger->norm_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size}); - w_vis->merger->norm_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size}); - w_vis->blocks = std::vector(vis_meta.depth); - for (size_t i = 0; i < vis_meta.depth; i++) { - w_vis->blocks[i].attn_proj_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size, vis_meta.hidden_size}); - w_vis->blocks[i].attn_proj_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size}); - w_vis->blocks[i].attn_qkv_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.in_channels * vis_meta.hidden_size, vis_meta.hidden_size}); - w_vis->blocks[i].attn_qkv_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.in_channels * vis_meta.hidden_size}); - w_vis->blocks[i].mlp_linear_fc1_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.intermediate_size, vis_meta.hidden_size}); - w_vis->blocks[i].mlp_linear_fc1_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.intermediate_size}); - w_vis->blocks[i].mlp_linear_fc2_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size, vis_meta.intermediate_size}); - w_vis->blocks[i].mlp_linear_fc2_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size}); - w_vis->blocks[i].norm1_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size}); - w_vis->blocks[i].norm1_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size}); - w_vis->blocks[i].norm2_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size}); - w_vis->blocks[i].norm2_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.hidden_size}); - } - w_vis->deepstack_mergers = std::vector(3); - for (size_t i = 0; i < 3; i++) { - w_vis->deepstack_mergers[i].linear_fc1_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.intermediate_size, vis_meta.intermediate_size}); - w_vis->deepstack_mergers[i].linear_fc2_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.out_hidden_size, vis_meta.intermediate_size}); - w_vis->deepstack_mergers[i].linear_fc1_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.intermediate_size}); - w_vis->deepstack_mergers[i].linear_fc2_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.out_hidden_size}); - w_vis->deepstack_mergers[i].norm_weight = Tensor::weight(nullptr, meta->dtype, {vis_meta.intermediate_size}); - w_vis->deepstack_mergers[i].norm_bias = Tensor::weight(nullptr, meta->dtype, {vis_meta.intermediate_size}); - } -} - -inline std::shared_ptr getSinTable(const Qwen3vlMeta *meta) { - auto half_dh = meta->text_meta.head_dim / 2; - auto unit = dsize(meta->dtype); - void *table = std::malloc(meta->text_meta.max_tokens * half_dh * unit); - - for (size_t i = 0; i < meta->text_meta.max_tokens; i++) { - for (size_t j = 0; j < half_dh; j++) { - float _sin = std::sin( - static_cast(i) / std::pow(meta->text_meta.rope_theta, static_cast(j) / half_dh)); - if (meta->dtype == INFINI_DTYPE_F16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_f16(_sin); - } else if (meta->dtype == INFINI_DTYPE_BF16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_bf16(_sin); - } else if (meta->dtype == INFINI_DTYPE_F32) { - ((float *)table)[i * half_dh + j] = _sin; - } else { - std::cout << "unsupported data type" << std::endl; - exit(1); - } - } - } - auto shape = std::vector({meta->text_meta.max_tokens, half_dh}); - auto tensor = Tensor::weight(table, meta->dtype, shape); - std::free(table); - return tensor; -} - -inline std::shared_ptr getCosTable(const Qwen3vlMeta *meta) { - auto half_dh = meta->text_meta.head_dim / 2; - auto unit = dsize(meta->dtype); - void *table = std::malloc(meta->text_meta.max_tokens * half_dh * unit); - - for (size_t i = 0; i < meta->text_meta.max_tokens; i++) { - for (size_t j = 0; j < half_dh; j++) { - float _cos = std::cos( - static_cast(i) / std::pow(meta->text_meta.rope_theta, static_cast(j) / half_dh)); - if (meta->dtype == INFINI_DTYPE_F16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_f16(_cos); - } else if (meta->dtype == INFINI_DTYPE_BF16) { - ((uint16_t *)table)[i * half_dh + j] = f32_to_bf16(_cos); - } else if (meta->dtype == INFINI_DTYPE_F32) { - ((float *)table)[i * half_dh + j] = _cos; - } else { - std::cout << "unsupported data type" << std::endl; - exit(1); - } - } - } - auto shape = std::vector({meta->text_meta.max_tokens, half_dh}); - auto tensor = Tensor::weight(table, meta->dtype, shape); - std::free(table); - return tensor; -} - -Qwen3vlWeights::Qwen3vlWeights( - const Qwen3vlMeta *_meta, infiniDevice_t device, int ndev, const int *dev_ids, bool _transpose_weight) { - meta = _meta; - transpose_weight = _transpose_weight; - device_weights = std::vector>(ndev); - for (int dev = 0; dev < ndev; dev++) { - int dev_id = dev_ids[dev]; - RUN_INFINI(infinirtSetDevice(device, dev_id)); - device_weights[dev] = std::make_shared(); - device_weights[dev]->device = device; - device_weights[dev]->dev_id = dev_id; - RUN_INFINI(infinirtStreamCreate(&device_weights[dev]->load_stream)); - device_weights[dev]->w_lang = std::make_shared(); - device_weights[dev]->w_vis = std::make_shared(); - - device_weights[dev]->w_lang->in_embd = getInEmbd(meta); - device_weights[dev]->w_lang->out_norm = getOutNorm(meta); - device_weights[dev]->w_lang->out_embd = getOutEmbd(meta); - device_weights[dev]->sin_table = getSinTable(meta); - device_weights[dev]->cos_table = getCosTable(meta); - - device_weights[dev]->w_lang->layers = std::vector(meta->text_meta.num_hidden_layers); - - for (size_t layer = 0; layer < meta->text_meta.num_hidden_layers; layer++) { - getLayerWeight(meta, device_weights[dev]->w_lang->layers[layer], ndev); - } - - getVisualWeight(meta, device_weights[dev]->w_vis); - } -} - -//--- Lang Global -void load_input_embd(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading input embedding from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->in_embd->load(cpu_ptr, weight->load_stream); - } -} - -void load_output_norm(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading output norm from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->out_norm->load(cpu_ptr, weight->load_stream); - } -} - -void load_output_embd(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading output embedding from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->out_embd->load(cpu_ptr, weight->load_stream); - if (weights->transpose_weight) { - weight->w_lang->out_embd->permute({1, 0}); //[d,voc] - } - } -} - -// --- Attention -void load_attn_norm(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading attention norm " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->layers[layer].attn_norm->load(cpu_ptr, weight->load_stream); - } -} - -void load_attn_q_norm(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading attention q_norm " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->layers[layer].attn_q_norm->load(cpu_ptr, weight->load_stream); - } -} - -void load_attn_qkv_proj(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading attention q_proj " << layer << " from " << cpu_ptr << std::endl; - int ndev = int(weights->device_weights.size()); - auto nkvh = weights->meta->text_meta.num_key_value_heads; - auto nh = weights->meta->text_meta.num_attention_heads; - auto dh = weights->meta->text_meta.head_dim; - auto d = weights->meta->text_meta.hidden_size; - //[ndev,nh+2*nkvh,dh,d] - for (int idev = 0; idev < ndev; idev++) { - auto weight = weights->device_weights[idev]; - size_t offset = idev * ((nkvh * 2 + nh) / ndev * dh) * d * dsize(weights->meta->dtype); - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->layers[layer].attn_qkv_proj->load((char *)cpu_ptr + offset, weight->load_stream); - if (weights->transpose_weight) { - weight->w_lang->layers[layer].attn_qkv_proj = weight->w_lang->layers[layer].attn_qkv_proj->permute({1, 0}); //[d, (nh+2*nkvh)*dh] - } - } -} - -void load_attn_k_norm(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading attention k_norm " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->layers[layer].attn_k_norm->load(cpu_ptr, weight->load_stream); - } -} - -void load_attn_o_proj(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading attention o_proj " << layer << " from " << cpu_ptr << std::endl; - int ndev = int(weights->device_weights.size()); - auto nh = weights->meta->text_meta.num_attention_heads; - auto dh = weights->meta->text_meta.head_dim; - auto d = weights->meta->text_meta.hidden_size; - // [ndev, d, nh // ndev * dh] - for (int idev = 0; idev < ndev; idev++) { - auto weight = weights->device_weights[idev]; - size_t offset = idev * d * (nh / ndev * dh) * dsize(weights->meta->dtype); - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->layers[layer].attn_o_proj->load((char *)cpu_ptr + offset, weight->load_stream); - if (weights->transpose_weight) { - weight->w_lang->layers[layer].attn_o_proj = weight->w_lang->layers[layer].attn_o_proj->permute({1, 0}); //[nh/ndev*dh, d] - } - } -} - -// --- MLP -void load_mlp_norm(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading mlp norm " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->layers[layer].mlp_norm->load(cpu_ptr, weight->load_stream); - } -} - -void load_mlp_gate_up(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading mlp gate " << layer << " from " << cpu_ptr << std::endl; - int ndev = int(weights->device_weights.size()); - auto di = weights->meta->text_meta.intermediate_size; - auto d = weights->meta->text_meta.hidden_size; - // [ndev, 2*di // ndev, d] - for (int idev = 0; idev < ndev; idev++) { - auto weight = weights->device_weights[idev]; - size_t offset = idev * (2 * di / ndev) * d * dsize(weights->meta->dtype); - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->layers[layer].mlp_gate_up->load((char *)cpu_ptr + offset, weight->load_stream); - if (weights->transpose_weight) { - weight->w_lang->layers[layer].mlp_gate_up = weight->w_lang->layers[layer].mlp_gate_up->permute({1, 0}); //[d, 2*di/ndev] - } - } -} - -void load_mlp_down(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading mlp down " << layer << " from " << cpu_ptr << std::endl; - int ndev = int(weights->device_weights.size()); - auto di = weights->meta->text_meta.intermediate_size; - auto d = weights->meta->text_meta.hidden_size; - //[ndev, d, di // ndev] - for (int idev = 0; idev < ndev; idev++) { - auto weight = weights->device_weights[idev]; - size_t offset = idev * d * (di / ndev) * dsize(weights->meta->dtype); - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_lang->layers[layer].mlp_down->load((char *)cpu_ptr + offset, weight->load_stream); - if (weights->transpose_weight) { - weight->w_lang->layers[layer].mlp_down = weight->w_lang->layers[layer].mlp_down->permute({1, 0}); //[di/ndev, d] - } - } -} - -// --- Vision weights -void load_patch_embed_weight(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading patch embed weight from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->patch_embed_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_patch_embed_bias(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading patch embed bias from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->patch_embed_bias->load(cpu_ptr, weight->load_stream); - } -} - -void load_pos_embed_weight(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading pos embed weight from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->pos_embed_weight->load(cpu_ptr, weight->load_stream); - } -} - -// Vision block attention -void load_attn_proj_weight(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision attn proj weight " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].attn_proj_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_attn_proj_bias(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision attn proj bias " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].attn_proj_bias->load(cpu_ptr, weight->load_stream); - } -} - -void load_attn_qkv_weight(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision attn qkv weight " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].attn_qkv_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_attn_qkv_bias(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision attn qkv bias " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].attn_qkv_bias->load(cpu_ptr, weight->load_stream); - } -} - -// Vision block mlp -void load_mlp_linear_fc1_weight(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision mlp fc1 weight " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].mlp_linear_fc1_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_mlp_linear_fc1_bias(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision mlp fc1 bias " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].mlp_linear_fc1_bias->load(cpu_ptr, weight->load_stream); - } -} - -void load_mlp_linear_fc2_weight(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision mlp fc2 weight " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].mlp_linear_fc2_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_mlp_linear_fc2_bias(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision mlp fc2 bias " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].mlp_linear_fc2_bias->load(cpu_ptr, weight->load_stream); - } -} - -// Vision block norm -void load_norm1_weight(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision norm1 weight " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].norm1_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_norm1_bias(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision norm1 bias " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].norm1_bias->load(cpu_ptr, weight->load_stream); - } -} - -void load_norm2_weight(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision norm2 weight " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].norm2_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_norm2_bias(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading vision norm2 bias " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->blocks[layer].norm2_bias->load(cpu_ptr, weight->load_stream); - } -} - -// Deepstack merger -void load_deepstack_merger_linear_fc1_weight(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading deepstack merger fc1 weight " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->deepstack_mergers[layer].linear_fc1_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_deepstack_merger_linear_fc1_bias(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading deepstack merger fc1 bias " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->deepstack_mergers[layer].linear_fc1_bias->load(cpu_ptr, weight->load_stream); - } -} - -void load_deepstack_merger_linear_fc2_weight(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading deepstack merger fc2 weight " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->deepstack_mergers[layer].linear_fc2_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_deepstack_merger_linear_fc2_bias(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading deepstack merger fc2 bias " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->deepstack_mergers[layer].linear_fc2_bias->load(cpu_ptr, weight->load_stream); - } -} - -void load_deepstack_merger_norm_weight(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading deepstack merger norm weight " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->deepstack_mergers[layer].norm_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_deepstack_merger_norm_bias(Qwen3vlWeights *weights, void *cpu_ptr, size_t layer) { - std::cout << "Loading deepstack merger norm bias " << layer << " from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->deepstack_mergers[layer].norm_bias->load(cpu_ptr, weight->load_stream); - } -} - -// Merger -void load_merger_linear_fc1_weight(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading merger fc1 weight from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->merger->linear_fc1_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_merger_linear_fc1_bias(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading merger fc1 bias from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->merger->linear_fc1_bias->load(cpu_ptr, weight->load_stream); - } -} - -void load_merger_linear_fc2_weight(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading merger fc2 weight from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->merger->linear_fc2_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_merger_linear_fc2_bias(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading merger fc2 bias from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->merger->linear_fc2_bias->load(cpu_ptr, weight->load_stream); - } -} - -void load_merger_norm_weight(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading merger norm weight from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->merger->norm_weight->load(cpu_ptr, weight->load_stream); - } -} - -void load_merger_norm_bias(Qwen3vlWeights *weights, void *cpu_ptr) { - std::cout << "Loading merger norm bias from " << cpu_ptr << std::endl; - for (int dev = 0; dev < int(weights->device_weights.size()); dev++) { - auto weight = weights->device_weights[dev]; - RUN_INFINI(infinirtSetDevice(weight->device, weight->dev_id)); - weight->w_vis->merger->norm_bias->load(cpu_ptr, weight->load_stream); - } -} - -static Qwen3vlWeightLoader weight_loader = { - // Language model loaders - .lang_loader = { - .load_input_embd = load_input_embd, - .load_output_norm = load_output_norm, - .load_output_embd = load_output_embd, - .load_attn_norm = load_attn_norm, - .load_attn_q_norm = load_attn_q_norm, - .load_attn_k_norm = load_attn_k_norm, - .load_attn_qkv_proj = load_attn_qkv_proj, - .load_attn_o_proj = load_attn_o_proj, - .load_mlp_norm = load_mlp_norm, - .load_mlp_gate_up = load_mlp_gate_up, - .load_mlp_down = load_mlp_down, - }, - // Vision model loaders - .vis_loader = { - .load_patch_embed_weight = load_patch_embed_weight, - .load_patch_embed_bias = load_patch_embed_bias, - .load_pos_embed_weight = load_pos_embed_weight, - .load_attn_proj_weight = load_attn_proj_weight, - .load_attn_proj_bias = load_attn_proj_bias, - .load_attn_qkv_weight = load_attn_qkv_weight, - .load_attn_qkv_bias = load_attn_qkv_bias, - .load_mlp_linear_fc1_weight = load_mlp_linear_fc1_weight, - .load_mlp_linear_fc1_bias = load_mlp_linear_fc1_bias, - .load_mlp_linear_fc2_weight = load_mlp_linear_fc2_weight, - .load_mlp_linear_fc2_bias = load_mlp_linear_fc2_bias, - .load_norm1_weight = load_norm1_weight, - .load_norm1_bias = load_norm1_bias, - .load_norm2_weight = load_norm2_weight, - .load_norm2_bias = load_norm2_bias, - .load_deepstack_merger_linear_fc1_weight = load_deepstack_merger_linear_fc1_weight, - .load_deepstack_merger_linear_fc1_bias = load_deepstack_merger_linear_fc1_bias, - .load_deepstack_merger_linear_fc2_weight = load_deepstack_merger_linear_fc2_weight, - .load_deepstack_merger_linear_fc2_bias = load_deepstack_merger_linear_fc2_bias, - .load_deepstack_merger_norm_weight = load_deepstack_merger_norm_weight, - .load_deepstack_merger_norm_bias = load_deepstack_merger_norm_bias, - .load_merger_linear_fc1_weight = load_merger_linear_fc1_weight, - .load_merger_linear_fc1_bias = load_merger_linear_fc1_bias, - .load_merger_linear_fc2_weight = load_merger_linear_fc2_weight, - .load_merger_linear_fc2_bias = load_merger_linear_fc2_bias, - .load_merger_norm_weight = load_merger_norm_weight, - .load_merger_norm_bias = load_merger_norm_bias, - }}; - -__INFINI_C Qwen3vlWeights * -createQwen3vlWeights(const Qwen3vlMeta *meta, - infiniDevice_t device, - int ndev, - const int *dev_ids, - bool transpose_weight) { - - printf("=== C++ createQwen3vlWeights ===\n"); - printf("sizeof(Qwen3vlTextMeta): %zu\n", sizeof(Qwen3vlTextMeta)); - printf("sizeof(Qwen3vlVisMeta): %zu\n", sizeof(Qwen3vlVisMeta)); - printf("sizeof(Qwen3vlMeta): %zu\n", sizeof(Qwen3vlMeta)); - printf("meta->dtype: %d\n", meta->dtype); - printf("meta->text_meta.hidden_size: %zu\n", meta->text_meta.hidden_size); - printf("meta->text_meta.num_hidden_layers: %zu\n", meta->text_meta.num_hidden_layers); - printf("meta->text_meta.vocab_size: %zu\n", meta->text_meta.vocab_size); - printf("meta->vis_meta.depth: %zu\n", meta->vis_meta.depth); - printf("device: %d, ndev: %d, dev_ids[0]: %d\n", device, ndev, dev_ids[0]); - fflush(stdout); - - auto weights = new Qwen3vlWeights(meta, device, ndev, dev_ids, transpose_weight); - return weights; -}; - -__INFINI_C Qwen3vlWeightLoader * -createQwen3vlWeightLoader() { - return &weight_loader; -} diff --git a/src/tensor.hpp b/src/tensor.hpp deleted file mode 100644 index 320d871ca..000000000 --- a/src/tensor.hpp +++ /dev/null @@ -1,190 +0,0 @@ -#ifndef INFER_TENSOR_H -#define INFER_TENSOR_H - -#include "allocator.hpp" -#include "utils.hpp" -#include -#include -#include - -class Storage { -private: - Storage() = default; - void *_memory; - size_t _size; - infiniDevice_t _device_type; - int _device_id; - std::shared_ptr _memory_pool; - -public: - static std::shared_ptr create(size_t size); - static std::shared_ptr createAsync(size_t size, infinirtStream_t stream = nullptr); - static std::shared_ptr createFromPool(size_t size, std::shared_ptr pool = nullptr); - static std::shared_ptr createHost(size_t size); - ~Storage(); - - void *memory() const { return _memory; } - size_t size() const { return _size; } - infiniDevice_t deviceType() const { return _device_type; } - int deviceId() const { return _device_id; } -}; - -struct SliceParams { - size_t dim; - size_t start; - size_t len; -}; - -template -std::vector __shape(Args... args) { - return std::vector{static_cast(args)...}; -} - -template -std::vector __strides(Args... args) { - return std::vector{static_cast(args)...}; -} -class TensorDesc { -private: - infiniDtype_t _dtype; - std::vector _shape; - std::vector _strides; - infiniopTensorDescriptor_t _desc; - size_t _seed; - - TensorDesc(infiniDtype_t dtype, const std::vector &shape, - const std::vector &strides) : _dtype(dtype), _shape(shape), _strides(strides), _desc(nullptr) { computeTensorDesHash(); } - void resetDesc(); - void computeTensorDesHash(); - -public: - ~TensorDesc(); - static std::shared_ptr - create(infiniDtype_t dtype, const std::vector &shape, - const std::vector &strides); - static std::shared_ptr - create(infiniDtype_t dtype, const std::vector &shape); - static std::shared_ptr - createWithOrder(infiniDtype_t dtype, const std::vector &shape, - const std::vector &order); - - infiniDtype_t dtype() const { return _dtype; } - const std::vector &shape() const { return _shape; } - const std::vector &strides() const { return _strides; } - size_t ndim() const { return _shape.size(); } - infiniopTensorDescriptor_t desc() const; - bool isContigous() const; - std::string info() const; - size_t seed() const { return _seed; } - - void dimMerge(size_t dim_start, size_t dim_end); - void dimSplit(size_t dim, const std::vector &dims); - void permute(const std::vector &order); -}; - -class Tensor : public std::enable_shared_from_this { -private: - std::shared_ptr _storage; - std::shared_ptr _desc; - - ptrdiff_t _offset; - - void *dataImpl(ptrdiff_t offset) const; - std::shared_ptr - sliceImpl(const std::vector &slices) const; - -public: - static std::shared_ptr buffer(infiniDtype_t dtype, - const std::vector &shape, - std::shared_ptr pool = nullptr); - static std::shared_ptr weight(void *host_data, - infiniDtype_t dtype, - const std::vector &shape); - void load(const void *host_data, infinirtStream_t stream = nullptr); - std::shared_ptr memShare(const std::vector &shape, - infiniDtype_t dtype = INFINI_DTYPE_INVALID) const; - std::shared_ptr slice(size_t dim, size_t start, size_t len); - std::shared_ptr slice(size_t dim, size_t start, - size_t len) const; - std::shared_ptr slice(const std::vector &slices); - std::shared_ptr - slice(const std::vector &slices) const; - std::shared_ptr dimMerge(size_t dim_start, size_t dim_end); - std::shared_ptr dimSplit(size_t dim, - const std::vector &dims); - std::shared_ptr permute(const std::vector &order); - void *data(ptrdiff_t offset = 0); - void const *data(ptrdiff_t offset = 0) const; - void copyFrom(std::shared_ptr src, infiniopHandle_t handle, - infinirtStream_t stream = nullptr); - const std::vector &shape() const; - const std::vector &strides() const; - size_t ndim() const; - infiniDtype_t dtype() const; - bool isContigous() const; - infiniopTensorDescriptor_t desc() const; - ptrdiff_t dataOffset() const; - infiniDevice_t deviceType() const; - int deviceId() const; - size_t numel() const; - - void debug(const std::string &filename) const; - void debug() const; - std::string info() const; - size_t seed() const; - - std::shared_ptr view(const std::vector &new_shape) const; - std::shared_ptr view_as(const std::vector &new_shape) const; - std::shared_ptr view_as(const std::vector &new_shape, const std::vector &new_strides) const; - - ~Tensor(); -}; - -inline size_t dsize(infiniDtype_t dtype) { - switch (dtype) { - case INFINI_DTYPE_INVALID: - return 0; - case INFINI_DTYPE_BYTE: - return 1; - case INFINI_DTYPE_BOOL: - return 1; - case INFINI_DTYPE_I8: - return 1; - case INFINI_DTYPE_I16: - return 2; - case INFINI_DTYPE_I32: - return 4; - case INFINI_DTYPE_I64: - return 8; - case INFINI_DTYPE_U8: - return 1; - case INFINI_DTYPE_U16: - return 2; - case INFINI_DTYPE_U32: - return 4; - case INFINI_DTYPE_U64: - return 8; - case INFINI_DTYPE_F8: - return 1; - case INFINI_DTYPE_F16: - return 2; - case INFINI_DTYPE_F32: - return 4; - case INFINI_DTYPE_F64: - return 8; - case INFINI_DTYPE_C16: - return 2; - case INFINI_DTYPE_C32: - return 4; - case INFINI_DTYPE_C64: - return 8; - case INFINI_DTYPE_C128: - return 16; - case INFINI_DTYPE_BF16: - return 2; - default: - return 0; - } -} - -#endif diff --git a/src/tensor/strorage.cpp b/src/tensor/strorage.cpp deleted file mode 100644 index cbcd11b6a..000000000 --- a/src/tensor/strorage.cpp +++ /dev/null @@ -1,53 +0,0 @@ -#include "../allocator.hpp" -#include "../tensor.hpp" - -std::shared_ptr Storage::create(size_t size) { - auto storage = std::shared_ptr(new Storage()); - RUN_INFINI(infinirtMalloc(&storage->_memory, size)); - storage->_size = size; - RUN_INFINI(infinirtGetDevice(&storage->_device_type, &storage->_device_id)); - return storage; -} - -std::shared_ptr Storage::createAsync(size_t size, infinirtStream_t stream) { - auto storage = std::shared_ptr(new Storage()); - RUN_INFINI(infinirtMallocAsync(&storage->_memory, size, stream)); - storage->_size = size; - RUN_INFINI(infinirtGetDevice(&storage->_device_type, &storage->_device_id)); - return storage; -} - -std::shared_ptr Storage::createFromPool(size_t size, std::shared_ptr pool) { - auto storage = std::shared_ptr(new Storage()); - storage->_memory_pool = pool; - if (pool) { - storage->_memory = pool->alloc(size); - } else { - RUN_INFINI(infinirtMalloc(&storage->_memory, size)); - } - storage->_size = size; - RUN_INFINI(infinirtGetDevice(&storage->_device_type, &storage->_device_id)); - return storage; -} - -std::shared_ptr Storage::createHost(size_t size) { - auto storage = std::shared_ptr(new Storage()); - RUN_INFINI(infinirtMallocHost(&storage->_memory, size)); - storage->_size = size; - storage->_device_type = INFINI_DEVICE_CPU; - storage->_device_id = 0; - storage->_memory_pool = nullptr; // No pool for host memory - return storage; -} - -Storage::~Storage() { - if (_memory_pool) { - _memory_pool->release(_memory); - } else { - if (_device_type == INFINI_DEVICE_CPU) { - RUN_INFINI(infinirtFreeHost(_memory)); - } else { - RUN_INFINI(infinirtFree(_memory)); - } - } -} diff --git a/src/tensor/tensor.cpp b/src/tensor/tensor.cpp deleted file mode 100644 index 37d8712a9..000000000 --- a/src/tensor/tensor.cpp +++ /dev/null @@ -1,426 +0,0 @@ -#include "../tensor.hpp" -#include "../utils.hpp" -#include -#include -#include -#include -#include -#include - -std::shared_ptr -TensorDesc::create(infiniDtype_t dtype, const std::vector &shape, - const std::vector &strides) { - return std::shared_ptr(new TensorDesc(dtype, shape, strides)); -} - -std::shared_ptr -TensorDesc::create(infiniDtype_t dtype, const std::vector &shape) { - auto ndim = shape.size(); - auto strides = std::vector(ndim); - if (ndim > 0) { - strides[ndim - 1] = 1; - for (int i = ndim - 2; i >= 0; i--) { - strides[i] = strides[i + 1] * shape[i + 1]; - } - } - return create(dtype, shape, strides); -} - -std::shared_ptr -TensorDesc::createWithOrder(infiniDtype_t dtype, const std::vector &shape, - const std::vector &order) { - ASSERT_EQ(shape.size(), order.size()); - auto ndim = shape.size(); - if (ndim == 0) { - return create(dtype, shape); - } - auto strides = std::vector(order.size()); - auto idx = std::find(order.begin(), order.end(), size_t(ndim - 1)); - strides[std::distance(order.begin(), idx)] = 1; - for (int i = ndim - 2; i >= 0; i--) { - auto prev_dim = shape[std::distance(order.begin(), idx)]; - auto prev_stride = strides[std::distance(order.begin(), idx)]; - idx = std::find(order.begin(), order.end(), size_t(i)); - strides[std::distance(order.begin(), idx)] = prev_stride * prev_dim; - } - return create(dtype, shape, strides); -} - -infiniopTensorDescriptor_t TensorDesc::desc() const { - if (_desc == nullptr) { - RUN_INFINI(infiniopCreateTensorDescriptor( - (infiniopTensorDescriptor_t *)(&_desc), _shape.size(), _shape.data(), - _strides.data(), _dtype)); - } - return _desc; -}; - -void TensorDesc::resetDesc() { - if (this->_desc != nullptr) { - infiniopDestroyTensorDescriptor(this->_desc); - this->_desc = nullptr; - } -} - -void TensorDesc::computeTensorDesHash() { - _seed = 0; - for (auto dim : this->shape()) { - hash_combine(_seed, dim); - } - for (auto stride : this->strides()) { - hash_combine(_seed, static_cast(stride)); - } -} - -bool TensorDesc::isContigous() const { - auto ndim = this->ndim(); - auto shape = this->shape(); - auto strides = std::vector(ndim); - strides[ndim - 1] = 1; - for (int i = ndim - 2; i >= 0; i--) { - strides[i] = strides[i + 1] * shape[i + 1]; - } - ASSERT_EQ(strides.size(), this->_strides.size()); - return std::equal(strides.begin(), strides.end(), this->_strides.begin()); -} - -std::string TensorDesc::info() const { - std::stringstream ss; - - ss << "Tensor: " - << "shape[ "; - for (auto s : this->shape()) { - ss << s << " "; - } - ss << "] strides[ "; - for (auto s : this->strides()) { - ss << s << " "; - } - ss << "] dtype=" << this->dtype(); - - return ss.str(); -} - -TensorDesc::~TensorDesc() { - this->resetDesc(); -} - -const std::vector &Tensor::shape() const { return this->_desc->shape(); } -const std::vector &Tensor::strides() const { return this->_desc->strides(); } -size_t Tensor::ndim() const { return this->_desc->ndim(); } -infiniDtype_t Tensor::dtype() const { return this->_desc->dtype(); } -infiniDevice_t Tensor::deviceType() const { return this->_storage->deviceType(); } -int Tensor::deviceId() const { return this->_storage->deviceId(); } -Tensor::~Tensor() {} - -size_t Tensor::numel() const { - return std::accumulate(this->shape().begin(), this->shape().end(), size_t(1), std::multiplies()); -} - -ptrdiff_t Tensor::dataOffset() const { - return _offset; -} - -infiniopTensorDescriptor_t Tensor::desc() const { return _desc->desc(); } - -std::shared_ptr Tensor::buffer(infiniDtype_t dtype, - const std::vector &shape, - std::shared_ptr pool) { - std::shared_ptr tensor = std::make_shared(); - auto ndim = shape.size(); - - size_t size = std::accumulate(shape.begin(), shape.end(), dsize(dtype), std::multiplies()); - auto strides = std::vector(ndim); - if (ndim > 0) { - strides[ndim - 1] = 1; - for (int i = ndim - 2; i >= 0; i--) { - strides[i] = strides[i + 1] * shape[i + 1]; - } - } - tensor->_storage = Storage::createFromPool(size, pool); - tensor->_desc = TensorDesc::create(dtype, shape, strides); - tensor->_offset = 0; - return tensor; -} - -std::shared_ptr Tensor::weight(void *data, infiniDtype_t dtype, - const std::vector &shape) { - std::shared_ptr tensor = std::make_shared(); - auto ndim = shape.size(); - size_t size = std::accumulate(shape.begin(), shape.end(), dsize(dtype), std::multiplies()); - auto strides = std::vector(ndim); - if (ndim > 0) { - strides[ndim - 1] = 1; - for (int i = ndim - 2; i >= 0; i--) { - strides[i] = strides[i + 1] * shape[i + 1]; - } - } - - tensor->_storage = Storage::create(size); - tensor->_desc = TensorDesc::create(dtype, shape, strides); - if (data != nullptr) { - tensor->load(data); - } - - tensor->_offset = 0; - return tensor; -} - -void Tensor::load(const void *data, infinirtStream_t stream) { - if (stream) { - RUN_INFINI(infinirtMemcpyAsync(this->_storage->memory(), data, this->_storage->size(), INFINIRT_MEMCPY_H2D, stream)); - return; - } - // NOTE: 为兼容部分平台（沐曦）多线程并发对同一host数据执行memcpy卡死问题 - static std::mutex mutex; - { - std::lock_guard lock(mutex); - RUN_INFINI(infinirtMemcpy(this->_storage->memory(), - data, this->_storage->size(), INFINIRT_MEMCPY_H2D)); - } -} - -std::shared_ptr Tensor::memShare(const std::vector &shape, infiniDtype_t dtype_) const { - auto dtype = dtype_ == INFINI_DTYPE_INVALID ? this->dtype() : dtype_; - size_t size = std::accumulate(shape.begin(), shape.end(), dsize(dtype), std::multiplies()); - ASSERT(size <= this->_storage->size()); - - std::shared_ptr tensor = std::make_shared(); - auto ndim = shape.size(); - auto strides = std::vector(ndim); - if (ndim > 0) { - strides[ndim - 1] = 1; - for (int i = ndim - 2; i >= 0; i--) { - strides[i] = strides[i + 1] * shape[i + 1]; - } - } - tensor->_storage = this->_storage; - tensor->_offset = 0; - tensor->_desc = TensorDesc::create(dtype, shape, strides); - return tensor; -} - -void *Tensor::dataImpl(ptrdiff_t offset) const { - return (char *)(this->_storage->memory()) + this->_offset + offset * dsize(this->dtype()); -} - -void *Tensor::data(ptrdiff_t offset) { - return this->dataImpl(offset); -} - -const void *Tensor::data(ptrdiff_t offset) const { - return this->dataImpl(offset); -} - -void Tensor::copyFrom(std::shared_ptr src, - infiniopHandle_t handle, infinirtStream_t stream) { - ASSERT_EQ(this->shape(), src->shape()); - ASSERT_EQ(this->dtype(), src->dtype()); - infiniopRearrangeDescriptor_t desc; - RUN_INFINI(infiniopCreateRearrangeDescriptor( - handle, &desc, this->desc(), src->desc())); - RUN_INFINI(infiniopRearrange(desc, this->data(), src->data(), - stream)); - RUN_INFINI(infiniopDestroyRearrangeDescriptor(desc)); -} - -bool Tensor::isContigous() const { - return this->_desc->isContigous(); -} - -template -void print_data(T *data, const std::vector &shape, - const std::vector &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << data[i] << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - print_data(data + i * strides[dim], shape, strides, dim + 1); - } - } -} - -template <> -void print_data(uint16_t const *data, const std::vector &shape, - const std::vector &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << f16_to_f32(data[i * strides[dim]]) << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - print_data(data + i * strides[dim], shape, strides, dim + 1); - } - } -} - -void print_data_bf16(uint16_t const *data, const std::vector &shape, - const std::vector &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << bf16_to_f32(data[i * strides[dim]]) << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - print_data_bf16(data + i * strides[dim], shape, strides, dim + 1); - } - } -} - -std::string Tensor::info() const { - std::stringstream ss; - - ss << "Tensor: " - << this->_desc->info() - << " device=" << this->deviceType() - << " device_id=" << this->deviceId(); - return this->_desc->info(); -} - -size_t Tensor::seed() const { - return this->_desc->seed(); -} - -std::shared_ptr Tensor::view(const std::vector &new_shape) const { - // Step 1: Validate total size - size_t numel = 1; - for (size_t dim : this->_desc->shape()) { - numel *= dim; - } - - size_t new_numel = 1; - for (size_t dim : new_shape) { - new_numel *= dim; - } - - ASSERT_EQ(numel, new_numel); - - // Step 2: Get current shape and strides - const std::vector &old_shape = this->_desc->shape(); - const std::vector &old_strides = this->_desc->strides(); - - // Step 3: Create merged shape and strides - std::vector merged_shape; - std::vector merged_strides; - - if (!old_shape.empty()) { - merged_shape.push_back(old_shape[0]); - merged_strides.push_back(old_strides[0]); - - for (size_t i = 1; i < old_shape.size(); ++i) { - if (old_strides[i] * static_cast(old_shape[i]) == merged_strides.back()) { - merged_shape.back() *= old_shape[i]; - merged_strides.back() = old_strides[i]; - } else { - merged_shape.push_back(old_shape[i]); - merged_strides.push_back(old_strides[i]); - } - } - } - - // Step 4: Compute new strides by splitting merged dimensions - std::vector new_strides(new_shape.size()); - size_t merged_idx = 0; - ptrdiff_t current_stride = merged_strides[0]; - size_t remaining_size = merged_shape[0]; - - for (size_t i = 0; i < new_shape.size(); ++i) { - // Find which merged dimension contains this new dimension - while (new_shape[i] > remaining_size) { - ASSERT(++merged_idx < merged_shape.size()); - current_stride = merged_strides[merged_idx]; - remaining_size = merged_shape[merged_idx]; - } - - ASSERT_EQ(remaining_size % new_shape[i], 0); - - new_strides[i] = current_stride * (remaining_size / new_shape[i]); - remaining_size /= new_shape[i]; - } - - return this->view_as(new_shape, new_strides); -} - -std::shared_ptr Tensor::view_as(const std::vector &new_shape) const { - std::shared_ptr tensor = std::make_shared(); - tensor->_storage = this->_storage; - tensor->_desc = TensorDesc::create(this->dtype(), new_shape); - tensor->_offset = this->_offset; - return tensor; -} - -std::shared_ptr Tensor::view_as(const std::vector &new_shape, const std::vector &new_strides) const { - std::shared_ptr tensor = std::make_shared(); - tensor->_storage = this->_storage; - tensor->_desc = TensorDesc::create(this->dtype(), new_shape, new_strides); - tensor->_offset = this->_offset; - return tensor; -} - -void Tensor::debug(const std::string &filename) const { - RUN_INFINI(infinirtDeviceSynchronize()); - - std::cout << info() << std::endl; - - void const *cpu_data; - if (this->deviceType() != INFINI_DEVICE_CPU) { - void *cpu_memory = std::malloc(this->_storage->size()); - RUN_INFINI(infinirtMemcpy(cpu_memory, this->_storage->memory(), - this->_storage->size(), INFINIRT_MEMCPY_D2H)); - cpu_data = cpu_memory; - } else { - cpu_data = this->_storage->memory(); - } - - if (!filename.empty()) { - std::ofstream outFile(filename, std::ios::binary); - if (!outFile) { - std::cerr << "Error opening file for writing: " << filename << "\n"; - return; - } - outFile.write(reinterpret_cast(cpu_data), this->_storage->size()); - outFile.close(); - std::cout << "Data written to file: " << filename << "\n"; - return; - } - - switch (this->dtype()) { - case INFINI_DTYPE_F16: - print_data((uint16_t const *)((char const *)cpu_data + dataOffset()), - this->shape(), this->strides(), 0); - break; - case INFINI_DTYPE_F32: - print_data((float const *)((char const *)cpu_data + dataOffset()), - this->shape(), this->strides(), 0); - break; - case INFINI_DTYPE_U64: - print_data((uint64_t const *)((char const *)cpu_data + dataOffset()), - this->shape(), this->strides(), 0); - break; - case INFINI_DTYPE_I64: - print_data((int64_t const *)((char const *)cpu_data + dataOffset()), - this->shape(), this->strides(), 0); - break; - case INFINI_DTYPE_U32: - print_data((uint32_t const *)((char const *)cpu_data + dataOffset()), - this->shape(), this->strides(), 0); - break; - case INFINI_DTYPE_I32: - print_data((int32_t const *)((char const *)cpu_data + dataOffset()), - this->shape(), this->strides(), 0); - break; - case INFINI_DTYPE_BF16: - print_data_bf16((uint16_t const *)((char const *)cpu_data + dataOffset()), - this->shape(), this->strides(), 0); - break; - default: - PANIC("Unsupported data type"); - } -} - -void Tensor::debug() const { this->debug(""); } diff --git a/src/tensor/transform.cpp b/src/tensor/transform.cpp deleted file mode 100644 index 60130deb4..000000000 --- a/src/tensor/transform.cpp +++ /dev/null @@ -1,137 +0,0 @@ -#include "../tensor.hpp" -#include "../utils.hpp" -#include -#include -#include - -std::shared_ptr Tensor::sliceImpl(const std::vector &slices) const { - std::shared_ptr tensor = std::make_shared(); - - auto new_shape = std::vector(this->shape()); - ptrdiff_t offset = 0; - - for (const auto &slice : slices) { - ASSERT(slice.len > 0); - ASSERT(this->shape()[slice.dim] >= slice.start + slice.len); - new_shape[slice.dim] = slice.len; - offset += slice.start * this->strides()[slice.dim]; - } - - tensor->_desc = TensorDesc::create(this->dtype(), new_shape, this->strides()); - tensor->_offset = offset * dsize(this->dtype()) + this->_offset; - tensor->_storage = this->_storage; - return tensor; -} - -std::shared_ptr Tensor::slice(size_t dim, size_t start, size_t len) { - return this->sliceImpl({{dim, start, len}}); -} - -std::shared_ptr Tensor::slice(size_t dim, size_t start, size_t len) const { - return this->sliceImpl({{dim, start, len}}); -} - -std::shared_ptr Tensor::slice(const std::vector &slices) { - return this->sliceImpl(slices); -} - -std::shared_ptr Tensor::slice(const std::vector &slices) const { - return this->sliceImpl(slices); -} - -void TensorDesc::dimMerge(size_t dim_start, size_t dim_end) { - ASSERT(dim_start <= dim_end && dim_end < this->_shape.size()); - if (dim_start == dim_end) { - return; - } - - auto new_shape = std::vector(); - auto new_strides = std::vector(); - for (size_t i = 0; i < dim_start; i++) { - new_shape.push_back(this->_shape[i]); - new_strides.push_back(this->_strides[i]); - } - for (size_t i = dim_start + 1; i <= dim_end; i++) { - ASSERT_EQ(this->_strides[i - 1], ptrdiff_t(this->_shape[i]) * this->_strides[i]); - } - new_shape.push_back(std::accumulate(this->_shape.begin() + dim_start, this->_shape.begin() + dim_end + 1, 1, std::multiplies())); - new_strides.push_back(this->_strides[dim_end]); - for (size_t i = dim_end + 1; i < this->_shape.size(); i++) { - new_shape.push_back(this->_shape[i]); - new_strides.push_back(this->_strides[i]); - } - this->_shape = new_shape; - this->_strides = new_strides; - this->resetDesc(); - this->computeTensorDesHash(); -} - -std::shared_ptr Tensor::dimMerge(size_t dim_start, size_t dim_end) { - auto new_desc = TensorDesc::create(_desc->dtype(), _desc->shape(), _desc->strides()); - new_desc->dimMerge(dim_start, dim_end); - - auto tensor = std::make_shared(); - tensor->_storage = _storage; - tensor->_desc = new_desc; - tensor->_offset = _offset; - return tensor; -} - -void TensorDesc::dimSplit(size_t dim, const std::vector &dims) { - ASSERT_EQ(this->_shape[dim], std::accumulate(dims.begin(), dims.end(), size_t(1), std::multiplies())); - auto new_shape = std::vector(); - auto new_strides = std::vector(); - for (size_t i = 0; i < dim; i++) { - new_shape.push_back(this->_shape[i]); - new_strides.push_back(this->_strides[i]); - } - for (size_t i = 0; i < dims.size(); i++) { - new_shape.push_back(dims[i]); - new_strides.push_back(this->_strides[dim] * this->_shape[dim] / std::accumulate(dims.begin(), dims.begin() + i + 1, 1, std::multiplies())); - } - for (size_t i = dim + 1; i < this->_shape.size(); i++) { - new_shape.push_back(this->_shape[i]); - new_strides.push_back(this->_strides[i]); - } - this->_shape = new_shape; - this->_strides = new_strides; - this->resetDesc(); - this->computeTensorDesHash(); -} - -std::shared_ptr Tensor::dimSplit(size_t dim, const std::vector &dims) { - auto new_desc = TensorDesc::create(_desc->dtype(), _desc->shape(), _desc->strides()); - new_desc->dimSplit(dim, dims); - - auto tensor = std::make_shared(); - tensor->_storage = _storage; - tensor->_desc = new_desc; - tensor->_offset = _offset; - return tensor; -} - -void TensorDesc::permute(const std::vector &order) { - ASSERT_EQ(this->_shape.size(), order.size()); - auto new_shape = std::vector(order.size()); - auto new_strides = std::vector(order.size()); - for (size_t i = 0; i < order.size(); i++) { - ASSERT(std::find(order.begin(), order.end(), i) != order.end()); - new_shape[i] = this->_shape[order[i]]; - new_strides[i] = this->_strides[order[i]]; - } - this->_shape = new_shape; - this->_strides = new_strides; - this->resetDesc(); - this->computeTensorDesHash(); -} - -std::shared_ptr Tensor::permute(const std::vector &order) { - auto new_desc = TensorDesc::create(_desc->dtype(), _desc->shape(), _desc->strides()); - new_desc->permute(order); - - auto tensor = std::make_shared(); - tensor->_storage = _storage; - tensor->_desc = new_desc; - tensor->_offset = _offset; - return tensor; -} diff --git a/src/utils.hpp b/src/utils.hpp deleted file mode 100644 index b0da9fff7..000000000 --- a/src/utils.hpp +++ /dev/null @@ -1,127 +0,0 @@ -#ifndef INFINICORE_INFER_UTILS_H -#define INFINICORE_INFER_UTILS_H -#include - -#include -#include -#include -#include - -inline void assertTrue(int expr, const char *msg, const char *file, int line) { - if (!expr) { - fprintf(stderr, "\033[31mAssertion failed:\033[0m %s at file %s, line %d\n", msg, file, line); - exit(EXIT_FAILURE); - } -} - -#define ASSERT(expr) assertTrue((expr), #expr " is false", __FILE__, __LINE__) -#define ASSERT_EQ(a, b) assertTrue((a) == (b), #a " != " #b, __FILE__, __LINE__) -#define ASSERT_VALID_PTR(a) assertTrue((a) != nullptr, #a " is nullptr", __FILE__, __LINE__) - -#define PANIC(EXPR) \ - printf("Error at %s:%d - %s\n", __FILE__, __LINE__, #EXPR); \ - exit(EXIT_FAILURE) - -#define RUN_INFINI(API) \ - do { \ - auto api_result_ = (API); \ - if (api_result_ != INFINI_STATUS_SUCCESS) { \ - std::cerr << "Error Code " << api_result_ << " in `" << #API << "`" \ - << " from " << __func__ \ - << " at " << __FILE__ << ":" << __LINE__ << std::endl; \ - exit(EXIT_FAILURE); \ - } \ - } while (0) - -inline float f16_to_f32(uint16_t h) { - uint32_t sign = (h & 0x8000) << 16; // Extract the sign bit - int32_t exponent = (h >> 10) & 0x1F; // Extract the exponent - uint32_t mantissa = h & 0x3FF; // Extract the mantissa (fraction part) - - if (exponent == 31) { // Special case for Inf and NaN - if (mantissa != 0) { - // NaN: Set float32 NaN - uint32_t f32 = sign | 0x7F800000 | (mantissa << 13); - return *(float *)&f32; - } else { - // Infinity - uint32_t f32 = sign | 0x7F800000; - return *(float *)&f32; - } - } else if (exponent == 0) { // Subnormal float16 or zero - if (mantissa == 0) { - // Zero (positive or negative) - uint32_t f32 = sign; // Just return signed zero - return *(float *)&f32; - } else { - // Subnormal: Convert to normalized float32 - exponent = -14; // Set exponent for subnormal numbers - while ((mantissa & 0x400) == 0) { // Normalize mantissa - mantissa <<= 1; - exponent--; - } - mantissa &= 0x3FF; // Clear the leading 1 bit - uint32_t f32 = sign | ((exponent + 127) << 23) | (mantissa << 13); - return *(float *)&f32; - } - } else { - // Normalized float16 - uint32_t f32 = sign | ((exponent + 127 - 15) << 23) | (mantissa << 13); - return *(float *)&f32; - } -} - -inline uint16_t f32_to_f16(float val) { - uint32_t f32; - memcpy(&f32, &val, sizeof(f32)); // Read the bits of the float32 - uint16_t sign = (f32 >> 16) & 0x8000; // Extract the sign bit - int32_t exponent = ((f32 >> 23) & 0xFF) - 127; // Extract and de-bias the exponent - uint32_t mantissa = f32 & 0x7FFFFF; // Extract the mantissa (fraction part) - - if (exponent >= 31) { // Special cases for Inf and NaN - // NaN - if (exponent == 128 && mantissa != 0) { - return static_cast(sign | 0x7E00); - } - // Infinity - return static_cast(sign | 0x7C00); - } else if (exponent >= -14) { // Normalized case - return (uint16_t)(sign | ((exponent + 15) << 10) | (mantissa >> 13)); - } else if (exponent >= -24) { - mantissa |= 0x800000; // Add implicit leading 1 - mantissa >>= (-14 - exponent); - return (uint16_t)(sign | (mantissa >> 13)); - } else { - // Too small for subnormal: return signed zero - return (uint16_t)sign; - } -} - -inline float bf16_to_f32(uint16_t val) { - // 只需把 bf16 放到 float32 高 16 bit，其余 16 位置 0。 - uint32_t bits32 = static_cast(val) << 16; - - float out; - std::memcpy(&out, &bits32, sizeof(out)); - return out; -} - -inline uint16_t f32_to_bf16(float val) { - uint32_t bits32; - std::memcpy(&bits32, &val, sizeof(bits32)); - - // 截断前先加 0x7FFF，再根据第 16 位（有效位的最低位）的奇偶做 round-to-nearest-even - const uint32_t rounding_bias = 0x00007FFF + // 0111 1111 1111 1111 - ((bits32 >> 16) & 1); // 尾数的有效位的最低位奇数时 +1，即实现舍入偶数 - - uint16_t bf16_bits = static_cast((bits32 + rounding_bias) >> 16); - - return bf16_bits; -} - -// Hash combine utility (similar to boost::hash_combine) -inline void hash_combine(size_t &seed, size_t value) { - seed ^= value + 0x9e3779b9 + (seed << 6) + (seed >> 2); -} - -#endif diff --git a/test/bench/backends/infinilm.py b/test/bench/backends/infinilm.py index fbd99379d..9ab1af157 100644 --- a/test/bench/backends/infinilm.py +++ b/test/bench/backends/infinilm.py @@ -30,15 +30,20 @@ def __init__( "metax": "cuda", "moore": "musa", "iluvatar": "cuda", - "kunlun": "cuda", "hygon": "cuda", - "ali": "cuda", "cuda": "cuda", "mlu": "mlu", "musa": "musa", "npu": "npu", } - device_name = device_map.get(device_type_str.lower(), "cpu") + try: + device_name = device_map[device_type_str.lower()] + except KeyError: + supported = ", ".join(sorted(device_map)) + raise ValueError( + f"unsupported device platform {device_type_str!r}; " + f"expected one of: {supported}" + ) from None with open(os.path.join(model_dir_path, "config.json"), "r") as f: self.config_dict = json.load(f) diff --git a/test/models/llama/test_intermediate_validation.py b/test/models/llama/test_intermediate_validation.py index 29ef5cba8..4620c6f7b 100755 --- a/test/models/llama/test_intermediate_validation.py +++ b/test/models/llama/test_intermediate_validation.py @@ -6,11 +6,11 @@ of all intermediate values in step 9 using the validation pattern. """ -import sys +import json import os +import sys from pathlib import Path -from typing import Optional, Tuple, List, Dict -import json +from typing import Dict, Tuple try: import torch @@ -26,21 +26,19 @@ sys.exit(1) try: - from infinilm.models.llama import LlamaConfig, LlamaForCausalLM, Device - import _infinilm # Import C++ bindings for HookRegistry + from infinilm.lib import _infinilm # Import C++ bindings for HookRegistry + from infinilm.models.llama import Device, LlamaConfig, LlamaForCausalLM except ImportError as e: print(f"Error: InfiniLM Python package not found. Please install it: {e}") sys.exit(1) -from transformers.models.llama.modeling_llama import apply_rotary_pos_emb - from infinicore.lib import _infinicore - +from transformers.models.llama.modeling_llama import apply_rotary_pos_emb from utils import ( + infinicore_to_torch_tensor, normalize_param_name, tensor_all_close, torch_to_infinicore_tensor, - infinicore_to_torch_tensor, validate_infinicore_component, ) @@ -379,7 +377,7 @@ def test_intermediate_validation( print("\n1. Loading model configuration...") try: config_dict = load_model_config(model_dir) - print(f" ✓ Configuration loaded") + print(" ✓ Configuration loaded") except Exception as e: print(f" ✗ Failed to load configuration: {e}") return False @@ -398,7 +396,7 @@ def test_intermediate_validation( nvidia_device_type = _infinicore.Device.Type.NVIDIA device_count = _infinicore.get_device_count(nvidia_device_type) if device_count == 0: - print(f" ✗ No NVIDIA/CUDA devices available") + print(" ✗ No NVIDIA/CUDA devices available") return False if device_index >= device_count: print(f" ✗ CUDA device index {device_index} is out of range") @@ -410,7 +408,7 @@ def test_intermediate_validation( device_type_upper = "NVIDIA" device = Device(device_type_upper, device_index) infinilm_model = LlamaForCausalLM(infinilm_config, device) - print(f" ✓ InfiniLM model created") + print(" ✓ InfiniLM model created") except Exception as e: print(f" ✗ Failed to create InfiniLM model: {e}") import traceback @@ -431,7 +429,7 @@ def test_intermediate_validation( ) transformers_model = transformers_model.to(torch_device) transformers_model.eval() - print(f" ✓ Transformers model loaded") + print(" ✓ Transformers model loaded") except Exception as e: print(f" ✗ Failed to load transformers model: {e}") import traceback @@ -465,7 +463,7 @@ def test_intermediate_validation( 0, seq_len, dtype=torch.long, device=torch_device ).unsqueeze(0) - print(f" ✓ Input prepared") + print(" ✓ Input prepared") print(f" Input shape: {input_ids.shape}") print(f" Sequence length: {seq_len}") except Exception as e: @@ -731,11 +729,11 @@ def hook(module, args): hook.remove() transformers_logits = outputs.logits - print(f" ✓ Extracted intermediate values from transformers") + print(" ✓ Extracted intermediate values from transformers") print(f" Captured {len(transformers_intermediates)} intermediate tensors") # List all captured intermediate values - print(f"\n Available Transformers intermediate values (in order):") + print("\n Available Transformers intermediate values (in order):") for i, name in enumerate(sorted(transformers_intermediates.keys()), 1): tensor = transformers_intermediates[name] print(f" {i}. {name}: shape={tensor.shape}, dtype={tensor.dtype}") @@ -818,14 +816,12 @@ def hook(hook_name, tensor, layer_idx): None, # kv_caches hook_registry, # hook_registry ) - infinilm_logits = infinicore_to_torch_tensor( - infini_logits, transformers_logits - ) + infinicore_to_torch_tensor(infini_logits, transformers_logits) - print(f" ✓ InfiniLM forward pass completed") + print(" ✓ InfiniLM forward pass completed") print(f" Captured {len(infinilm_intermediates)} intermediate tensors") else: - print(f" ✗ Forward method not available") + print(" ✗ Forward method not available") return False except Exception as e: @@ -837,7 +833,6 @@ def hook(hook_name, tensor, layer_idx): # Step 8: Compare intermediate values (basic comparison) print("\n8. Comparing intermediate values (basic comparison)...") - all_match = True rtol = 1e-3 atol = 1e-3 @@ -855,18 +850,15 @@ def hook(hook_name, tensor, layer_idx): for trans_name, infini_name in hook_name_mapping.items(): if trans_name in transformers_intermediates: if infini_name in infinilm_intermediates: - match, stats = compare_tensors( + compare_tensors( f"{trans_name} vs {infini_name}", transformers_intermediates[trans_name], infinilm_intermediates[infini_name], rtol=1e-3, atol=1e-3, ) - if not match: - all_match = False else: print(f" ⚠ {infini_name} not found in InfiniLM intermediates") - all_match = False # Step 9: Systematic validation of intermediate values in order print("\n9. Systematic validation of intermediate values (in order)...") @@ -963,7 +955,7 @@ def hook(hook_name, tensor, layer_idx): # Basic shape check if trans_tensor.shape != infini_tensor.shape: - print(f" ✗ Shape mismatch!") + print(" ✗ Shape mismatch!") validation_results[trans_name] = { "status": "shape_mismatch", "trans_shape": trans_tensor.shape, @@ -999,24 +991,24 @@ def hook(hook_name, tensor, layer_idx): ) if match: - print(f" ✓ Validation PASSED") + print(" ✓ Validation PASSED") validation_results[trans_name] = {"status": "passed", "stats": stats} else: - print(f" ✗ Validation FAILED") + print(" ✗ Validation FAILED") validation_results[trans_name] = {"status": "failed", "stats": stats} # Detailed difference analysis diff = (trans_tensor - infini_tensor).abs() rel_diff = diff / (trans_tensor.abs() + 1e-10) - print(f"\n Detailed difference analysis:") + print("\n Detailed difference analysis:") print(f" Max abs diff: {diff.max().item():.6e}") print(f" Mean abs diff: {diff.mean().item():.6e}") print(f" Max rel diff: {rel_diff.max().item():.6e}") print(f" Mean rel diff: {rel_diff.mean().item():.6e}") # Error distribution - print(f"\n Error distribution:") + print("\n Error distribution:") for threshold in [1e-6, 1e-5, 1e-4, 1e-3, 1e-2]: count = (diff > threshold).sum().item() pct = 100.0 * count / diff.numel() @@ -1025,7 +1017,7 @@ def hook(hook_name, tensor, layer_idx): ) # Top problematic positions - print(f"\n Top 5 positions with largest absolute differences:") + print("\n Top 5 positions with largest absolute differences:") topk_values, topk_indices = torch.topk( diff.flatten(), k=min(5, diff.numel()) ) @@ -1045,9 +1037,7 @@ def hook(hook_name, tensor, layer_idx): "layer0_post_attention_layernorm", "final_norm", ]: - print( - f"\n Validating with InfiniCore ops using validation pattern..." - ) + print("\n Validating with InfiniCore ops using validation pattern...") try: import infinicore.nn.functional as F @@ -1122,7 +1112,7 @@ def rms_norm_op(input_tensor): results ) else: - print(f" ⚠ Cannot validate: missing input tensors or weight") + print(" ⚠ Cannot validate: missing input tensors or weight") except Exception as e: print(f" ⚠ Could not validate with InfiniCore ops: {e}") import traceback @@ -1131,12 +1121,11 @@ def rms_norm_op(input_tensor): # Validate q_proj operation (linear projection only, before reshape) elif trans_name == "layer0_attention_q_after_proj": - print( - f"\n Validating with InfiniCore ops using validation pattern..." - ) + print("\n Validating with InfiniCore ops using validation pattern...") try: - from infinicore.ops.matmul import matmul - from infinicore.ops.add import add + from infinicore.ops import add, matmul + + del add # Get the input (layer0_input_layernorm) trans_input = transformers_intermediates.get( @@ -1224,7 +1213,7 @@ def q_proj_op(input_tensor): results ) else: - print(f" ⚠ Cannot validate: missing input tensors") + print(" ⚠ Cannot validate: missing input tensors") except Exception as e: print(f" ⚠ Could not validate with InfiniCore ops: {e}") import traceback @@ -1236,7 +1225,7 @@ def q_proj_op(input_tensor): "layer0_attention_q_after_rope", "layer0_attention_k_after_rope", ]: - print(f"\n Validating RoPE application with PyTorch reference...") + print("\n Validating RoPE application with PyTorch reference...") head_type = "q" if trans_name.endswith("_q_after_rope") else "k" cos = transformers_intermediates.get("layer0_attention_rope_cos") sin = transformers_intermediates.get("layer0_attention_rope_sin") @@ -1336,7 +1325,7 @@ def q_proj_op(input_tensor): # Validate MLP intermediate values elif trans_name == "layer0_mlp": - print(f"\n Validating MLP intermediate values...") + print("\n Validating MLP intermediate values...") # Get intermediate values from both implementations trans_gate_proj = transformers_intermediates.get("layer0_mlp_gate_proj") @@ -1361,30 +1350,28 @@ def q_proj_op(input_tensor): # Step 0: Compare inputs print( - f"\n Step 0: Comparing MLP inputs (post_attention_layernorm output)..." + "\n Step 0: Comparing MLP inputs (post_attention_layernorm output)..." ) if trans_input is not None and infini_input is not None: input_match, input_stats = compare_tensors( "mlp_input", trans_input, infini_input, rtol=1e-3, atol=1e-3 ) if input_match: - print(f" ✓ MLP input: Match") + print(" ✓ MLP input: Match") else: - print(f" ✗ MLP input: Mismatch") + print(" ✗ MLP input: Mismatch") print( f" Max abs diff: {input_stats.get('max_abs_diff', 'N/A'):.6e}" ) print( f" Mean abs diff: {input_stats.get('mean_abs_diff', 'N/A'):.6e}" ) - print( - f" ⚠ Input mismatch may cause downstream differences" - ) + print(" ⚠ Input mismatch may cause downstream differences") else: - print(f" ⚠ Missing MLP input tensors") + print(" ⚠ Missing MLP input tensors") # Step 1: Compare gate_proj outputs - print(f"\n Step 1: Comparing gate_proj outputs...") + print("\n Step 1: Comparing gate_proj outputs...") if trans_gate_proj is not None and infini_gate_proj is not None: if trans_gate_proj.shape != infini_gate_proj.shape: print( @@ -1399,9 +1386,9 @@ def q_proj_op(input_tensor): atol=1e-3, ) if gate_match: - print(f" ✓ gate_proj: Match") + print(" ✓ gate_proj: Match") else: - print(f" ✗ gate_proj: Mismatch") + print(" ✗ gate_proj: Mismatch") print( f" Max abs diff: {gate_stats.get('max_abs_diff', 'N/A'):.6e}" ) @@ -1420,7 +1407,7 @@ def q_proj_op(input_tensor): diff = (trans_gate_proj - infini_gate_proj).abs() problem_positions = [1703, 894, 1334, 636, 1002] print( - f"\n Sample values at problematic positions (from final output):" + "\n Sample values at problematic positions (from final output):" ) for pos in problem_positions: if pos < trans_gate_proj.shape[-1]: @@ -1448,7 +1435,7 @@ def q_proj_op(input_tensor): print(f" ⚠ Missing gate_proj tensors: {', '.join(missing)}") # Step 2: Compare up_proj outputs - print(f"\n Step 2: Comparing up_proj outputs...") + print("\n Step 2: Comparing up_proj outputs...") if trans_up_proj is not None and infini_up_proj is not None: if trans_up_proj.shape != infini_up_proj.shape: print( @@ -1463,9 +1450,9 @@ def q_proj_op(input_tensor): atol=1e-3, ) if up_match: - print(f" ✓ up_proj: Match") + print(" ✓ up_proj: Match") else: - print(f" ✗ up_proj: Mismatch") + print(" ✗ up_proj: Mismatch") print( f" Max abs diff: {up_stats.get('max_abs_diff', 'N/A'):.6e}" ) @@ -1484,9 +1471,7 @@ def q_proj_op(input_tensor): print(f" ⚠ Missing up_proj tensors: {', '.join(missing)}") # Step 3: Compare SwiGLU intermediate - print( - f"\n Step 3: Comparing SwiGLU intermediate (silu(gate) * up)..." - ) + print("\n Step 3: Comparing SwiGLU intermediate (silu(gate) * up)...") if trans_intermediate is not None and infini_intermediate is not None: if trans_intermediate.shape != infini_intermediate.shape: print( @@ -1501,9 +1486,9 @@ def q_proj_op(input_tensor): atol=1e-3, ) if inter_match: - print(f" ✓ SwiGLU intermediate: Match") + print(" ✓ SwiGLU intermediate: Match") else: - print(f" ✗ SwiGLU intermediate: Mismatch") + print(" ✗ SwiGLU intermediate: Mismatch") print( f" Max abs diff: {inter_stats.get('max_abs_diff', 'N/A'):.6e}" ) @@ -1578,20 +1563,17 @@ def q_proj_op(input_tensor): print(f" ⚠ Missing intermediate tensors: {', '.join(missing)}") print( - f"\n Step 4: Final MLP output comparison (shown above in main validation)" + "\n Step 4: Final MLP output comparison (shown above in main validation)" ) print( - f" Summary: This validation helps identify which MLP step introduces the mismatch." + " Summary: This validation helps identify which MLP step introduces the mismatch." ) # Validate q_proj_reshape operation elif trans_name == "layer0_attention_q_after_proj_reshape": - print( - f"\n Validating with InfiniCore ops using validation pattern..." - ) + print("\n Validating with InfiniCore ops using validation pattern...") try: - from infinicore.ops.matmul import matmul - from infinicore.ops.add import add + from infinicore.ops import matmul # Get the input (layer0_input_layernorm) trans_input = transformers_intermediates.get( @@ -1608,8 +1590,6 @@ def q_proj_op(input_tensor): # Get model config for dimensions num_heads = transformers_model.config.num_attention_heads head_dim = transformers_model.config.head_dim - hidden_size = transformers_model.config.hidden_size - # Convert weight and bias to InfiniCore tensors (once, outside the op) weight_tensor = torch_to_infinicore_tensor(weight, infini_device) bias_tensor = None @@ -1709,7 +1689,7 @@ def q_proj_reshape_op(input_tensor): results ) else: - print(f" ⚠ Cannot validate: missing input tensors") + print(" ⚠ Cannot validate: missing input tensors") except Exception as e: print(f" ⚠ Could not validate with InfiniCore ops: {e}") import traceback @@ -1744,7 +1724,7 @@ def q_proj_reshape_op(input_tensor): print(f" ✗ Failed: {failed}") print(f" ⚠ Missing: {missing}") - print(f"\nDetailed results:") + print("\nDetailed results:") for trans_name, result in validation_results.items(): status = result.get("status", "unknown") if status == "passed": @@ -1792,7 +1772,7 @@ def main(): if model_dir is None: model_dir = arg else: - print(f"Error: Multiple model directories specified") + print("Error: Multiple model directories specified") sys.exit(1) i += 1 diff --git a/test/models/llama/utils.py b/test/models/llama/utils.py index 333048fde..06b385ac9 100644 --- a/test/models/llama/utils.py +++ b/test/models/llama/utils.py @@ -5,7 +5,8 @@ parameter name normalization, and tensor comparison. """ -from typing import Tuple, Dict, Callable, Optional, Any, List +from typing import Any, Callable, Dict, Optional, Tuple + import torch try: @@ -103,24 +104,22 @@ def to_torch_dtype(infini_dtype): underlying = infini_dtype # Compare underlying enum values - if underlying == _infinicore.DataType.F32: + if underlying == _infinicore.DataType.FLOAT32: return torch.float32 - elif underlying == _infinicore.DataType.F16: + elif underlying == _infinicore.DataType.FLOAT16: return torch.float16 - elif underlying == _infinicore.DataType.BF16: + elif underlying == _infinicore.DataType.BFLOAT16: return torch.bfloat16 - elif underlying == _infinicore.DataType.I8: + elif underlying == _infinicore.DataType.INT8: return torch.int8 - elif underlying == _infinicore.DataType.I16: + elif underlying == _infinicore.DataType.INT16: return torch.int16 - elif underlying == _infinicore.DataType.I32: + elif underlying == _infinicore.DataType.INT32: return torch.int32 - elif underlying == _infinicore.DataType.I64: + elif underlying == _infinicore.DataType.INT64: return torch.int64 - elif underlying == _infinicore.DataType.U8: + elif underlying == _infinicore.DataType.UINT8: return torch.uint8 - elif underlying == _infinicore.DataType.BOOL: - return torch.bool else: raise ValueError( f"Unsupported infinicore dtype: {infini_dtype} (underlying enum: {underlying})" @@ -182,7 +181,7 @@ def infinicore_to_torch_tensor(infini_tensor, torch_reference): infini_tensor = infini_tensor.to(target_infini_device) if not infini_tensor.is_contiguous(): infini_tensor = infini_tensor.contiguous() - except Exception as e: + except Exception: # If device operations fail, try to ensure contiguous at least if ( hasattr(infini_tensor, "is_contiguous") @@ -202,22 +201,10 @@ def infinicore_to_torch_tensor(infini_tensor, torch_reference): if ref_device.type == "cpu": # Check if source tensor is on CUDA - if so, we need pinned memory source_is_cuda = False - source_cuda_device = None if hasattr(infini_tensor, "device"): source_device = infini_tensor.device source_device_str = str(source_device) source_is_cuda = source_device_str.startswith("cuda") - if source_is_cuda: - # Extract CUDA device index from device string (e.g., "cuda:0") - try: - cuda_index = ( - int(source_device_str.split(":")[1]) - if ":" in source_device_str - else 0 - ) - source_cuda_device = infinicore.device("cuda", cuda_index) - except: - source_cuda_device = infinicore.device("cuda", 0) # If source is on CUDA, we need to ensure the intermediate CPU tensor # uses pinned memory. The copy_from function will handle setting the @@ -239,7 +226,7 @@ def infinicore_to_torch_tensor(infini_tensor, torch_reference): # Create temp tensor from PyTorch and copy from the CPU tensor temp_tensor = torch_to_infinicore_tensor(torch_result, target_infini_device) temp_tensor.copy_(cpu_tensor) - except Exception as e: + except Exception: # Fallback: create intermediate tensor and copy through it # Create an intermediate contiguous tensor on CPU # Use pin_memory=True if source is CUDA to ensure proper D2H copy @@ -431,7 +418,7 @@ def validate_infinicore_component( # Test 1: Call InfiniCore ops with InfiniCore input (current behavior) if verbose: - print(f"\n Test 1: InfiniCore ops with InfiniCore input...") + print("\n Test 1: InfiniCore ops with InfiniCore input...") # Prepare arguments for the op # For ops that take multiple inputs, we need to handle them @@ -449,16 +436,16 @@ def validate_infinicore_component( if verbose: if test1_match: - print(f" ✓ Test 1: InfiniCore ops matches InfiniLM output") + print(" ✓ Test 1: InfiniCore ops matches InfiniLM output") else: - print(f" ⚠ Test 1: InfiniCore ops differs from InfiniLM output") + print(" ⚠ Test 1: InfiniCore ops differs from InfiniLM output") print(f" Max abs diff: {test1_stats['max_abs_diff']:.15f}") print(f" Mean abs diff: {test1_stats['mean_abs_diff']:.15f}") # Test 2: Call InfiniCore ops with Transformers input (to eliminate input diff) if verbose: print( - f"\n Test 2: InfiniCore ops with Transformers input (eliminating input diff)..." + "\n Test 2: InfiniCore ops with Transformers input (eliminating input diff)..." ) test2_inputs = [trans_input_tensor] @@ -470,7 +457,7 @@ def validate_infinicore_component( # Compare Test 2 (InfiniCore ops with Transformers input) vs Transformers output if verbose: print( - f"\n Test 2 Results: InfiniCore ops (Transformers input) vs Transformers output:" + "\n Test 2 Results: InfiniCore ops (Transformers input) vs Transformers output:" ) test2_match, test2_stats = tensor_all_close( @@ -487,11 +474,11 @@ def validate_infinicore_component( if test2_match: print( - f" ✓ InfiniCore ops matches Transformers when using same input!" + " ✓ InfiniCore ops matches Transformers when using same input!" ) else: print( - f" ⚠ InfiniCore ops still differs from Transformers even with same input" + " ⚠ InfiniCore ops still differs from Transformers even with same input" ) print( f" This suggests the {op_name} computation itself differs" @@ -513,7 +500,7 @@ def validate_infinicore_component( # Compare Test 1 vs Test 2 to see impact of input difference if verbose: - print(f"\n Comparing Test 1 vs Test 2 (impact of input difference):") + print("\n Comparing Test 1 vs Test 2 (impact of input difference):") test1_vs_test2_diff = (test1_output_torch - test2_output_torch).abs() test1_vs_test2_max = test1_vs_test2_diff.max().item() @@ -531,15 +518,15 @@ def validate_infinicore_component( if test1_vs_test2_max > tolerance: results["input_impact"] = "significant" if verbose: - print(f" ⚠ Input difference causes significant output difference") + print(" ⚠ Input difference causes significant output difference") else: results["input_impact"] = "minimal" if verbose: - print(f" ✓ Input difference has minimal impact on output") + print(" ✓ Input difference has minimal impact on output") # Compare input data between Transformers and InfiniCore if verbose: - print(f"\n Comparing input data (Transformers vs InfiniCore):") + print("\n Comparing input data (Transformers vs InfiniCore):") input_diff = (transformers_input - infinicore_input).abs() input_diff_max = input_diff.max().item() @@ -568,7 +555,7 @@ def validate_infinicore_component( ) print(f" Difference: {input_diff[max_input_diff_pos].item():.15f}") else: - print(f" ✓ Input data matches (within tolerance)") + print(" ✓ Input data matches (within tolerance)") # Call debug callback if provided if debug_callback is not None: @@ -587,7 +574,7 @@ def validate_infinicore_component( # Summary if verbose: - print(f"\n Summary:") + print("\n Summary:") print( f" Test 1 (InfiniCore input): {'✓ PASS' if test1_match else '✗ FAIL'}" ) diff --git a/test/static/test_infinicore_python_contracts.py b/test/static/test_infinicore_python_contracts.py new file mode 100644 index 000000000..eb55e4ddf --- /dev/null +++ b/test/static/test_infinicore_python_contracts.py @@ -0,0 +1,357 @@ +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] + + +def read_source(relative_path: str) -> str: + return (ROOT / relative_path).read_text(encoding="utf-8") + + +class InfiniCorePythonContractsTest(unittest.TestCase): + def test_build_installs_one_shared_runtime_for_both_extensions(self) -> None: + xmake = read_source("xmake.lua") + setup = read_source("setup.py") + + runtime_start = xmake.index('target("infinicore_runtime")') + runtime_block = xmake[ + runtime_start : xmake.index("target_end()", runtime_start) + ] + self.assertIn('set_kind("shared")', runtime_block) + self.assertIn('set_installdir("python/infinicore")', runtime_block) + modern_libraries = 'add_links("infiniops", "infiniccl", "infinirt"' + self.assertIn(modern_libraries, runtime_block) + + self.assertIn('target("_infinicore")', xmake) + self.assertIn('add_deps("infinicore_runtime")', xmake) + self.assertIn("csrc/infinicore/src/pybind11/infinicore.cc", xmake) + self.assertIn("csrc/infinicore/src/pybind11/from_list.cc", xmake) + self.assertIn('set_installdir("python/infinicore")', xmake) + self.assertIn( + 'for target in ("infinicore_runtime", "_infinicore", "_infinilm"):', + setup, + ) + self.assertIn('["xmake", "build", target]', setup) + self.assertIn('["xmake", "install", target]', setup) + + for target in ("_infinicore", "_infinilm"): + start = xmake.index(f'target("{target}")') + block = xmake[start : xmake.index("target_end()", start)] + self.assertNotIn(modern_libraries, block) + self.assertIn('add_deps("infinicore_runtime")', block) + self.assertNotIn('add_linkgroups("infinicore_runtime"', block) + + infinicore_start = xmake.index('target("_infinicore")') + infinicore_block = xmake[ + infinicore_start : xmake.index("target_end()", infinicore_start) + ] + infinilm_start = xmake.index('target("_infinilm")') + infinilm_block = xmake[ + infinilm_start : xmake.index("target_end()", infinilm_start) + ] + self.assertIn('add_rpathdirs("$ORIGIN")', runtime_block) + self.assertNotIn("add_rpathdirs(INFINI_ROOT", runtime_block) + self.assertIn('add_rpathdirs("$ORIGIN")', infinicore_block) + self.assertIn('add_rpathdirs("$ORIGIN/../../infinicore/lib")', infinilm_block) + + def test_wheel_contains_native_artifacts_and_is_platform_specific(self) -> None: + setup = read_source("setup.py") + + self.assertIn("class BinaryDistribution(Distribution):", setup) + self.assertIn("def has_ext_modules(self):", setup) + self.assertIn("distclass=BinaryDistribution", setup) + self.assertIn('"infinicore.lib": INFINICORE_NATIVE_ARTIFACTS', setup) + self.assertIn('"infinilm.lib": INFINILM_EXTENSION_ARTIFACTS', setup) + for extension in ("_infinicore", "_infinilm"): + self.assertIn(f'"{extension}*.so"', setup) + self.assertIn(f'"{extension}*.pyd"', setup) + for library in ( + "infinicore_runtime", + "infiniops", + "infiniccl", + "infinirt", + ): + self.assertIn(f'"lib{library}.so"', setup) + self.assertIn(f'"lib{library}.so.*"', setup) + + def test_wheel_stages_modern_infini_runtime_dependencies(self) -> None: + setup = read_source("setup.py") + + self.assertIn( + 'INFINI_LIBRARY_NAMES = ("infiniops", "infiniccl", "infinirt")', + setup, + ) + self.assertIn("def stage_runtime_dependencies():", setup) + self.assertIn('destination = PROJECT_ROOT / "python/infinicore/lib"', setup) + self.assertIn('infini_root / "lib"', setup) + self.assertIn('infini_root / "lib64"', setup) + self.assertIn("raise FileNotFoundError", setup) + self.assertIn("stage_runtime_dependencies()", setup) + build_start = setup.index("def build_cpp_module():") + stage_call = setup.index("stage_runtime_dependencies()", build_start) + xmake_build = setup.index('subprocess.run(["xmake", "build"', build_start) + self.assertLess(stage_call, xmake_build) + + def test_wheel_includes_all_python_packages(self) -> None: + python_root = ROOT / "python" + setup = read_source("setup.py") + missing_package_markers = sorted( + str(directory.relative_to(python_root)) + for directory in python_root.rglob("*") + if directory.is_dir() + and any(directory.glob("*.py")) + and not (directory / "__init__.py").is_file() + ) + + self.assertIn('packages=find_packages(where="python")', setup) + self.assertEqual(missing_package_markers, []) + + def test_wheel_declares_its_eager_import_dependencies(self) -> None: + metadata = read_source("pyproject.toml") + dependencies = metadata[ + metadata.index("dependencies = [") : metadata.index( + "]", metadata.index("dependencies = [") + ) + ] + for package in ( + "janus", + "numpy", + "Pillow", + "safetensors", + "tokenizers", + "torch", + "tqdm", + "transformers", + "typing-extensions", + "xxhash", + ): + self.assertIn(f'"{package}"', dependencies) + + def test_device_binding_uses_native_infini_rt_types(self) -> None: + source = read_source("csrc/infinicore/src/pybind11/device.hpp") + + for native_name in ( + "kCpu", + "kNvidia", + "kCambricon", + "kAscend", + "kMetax", + "kMoore", + "kIluvatar", + "kHygon", + ): + self.assertIn(f"Device::Type::{native_name}", source) + for legacy_name in ("::CPU", "::NVIDIA", "::QY", "::KUNLUN", "::ALI"): + self.assertNotIn(legacy_name, source) + self.assertIn("&Device::type", source) + self.assertIn("&Device::index", source) + self.assertIn("&Device::ToString", source) + + def test_dtype_binding_matches_native_infini_rt_set(self) -> None: + source = read_source("csrc/infinicore/src/pybind11/dtype.hpp") + native_names = ( + "kInt8", + "kInt16", + "kInt32", + "kInt64", + "kUInt8", + "kUInt16", + "kUInt32", + "kUInt64", + "kFloat16", + "kBFloat16", + "kFloat32", + "kFloat64", + ) + for native_name in native_names: + self.assertIn(f"DataType::{native_name}", source) + for removed_name in ("BOOL", "BYTE", "F8", "C16", "C32", "C64", "C128"): + self.assertNotIn(f"DataType::{removed_name}", source) + + def test_from_list_and_stream_binding_do_not_use_legacy_runtime(self) -> None: + from_list = read_source("csrc/infinicore/src/pybind11/from_list.cc") + context = read_source("csrc/infinicore/src/pybind11/context.hpp") + event = read_source("csrc/infinicore/src/pybind11/device_event.hpp") + + for native_name in ( + "kInt8", + "kInt16", + "kInt32", + "kInt64", + "kUInt8", + "kUInt16", + "kUInt32", + "kUInt64", + "kFloat16", + "kBFloat16", + "kFloat32", + "kFloat64", + ): + self.assertIn(f"case DataType::{native_name}:", from_list) + self.assertNotIn("infinirt", from_list + context + event) + self.assertIn("reinterpret_cast(getStream())", context) + + def test_python_package_exposes_only_the_migrated_surface(self) -> None: + package = ROOT / "python" / "infinicore" + for relative_path in ( + "__init__.py", + "context.py", + "device.py", + "dtype.py", + "tensor.py", + "utils.py", + "lib/__init__.py", + "nn/__init__.py", + "nn/parameter.py", + "nn/modules/module.py", + ): + self.assertTrue((package / relative_path).is_file(), relative_path) + + init = (package / "__init__.py").read_text(encoding="utf-8") + for public_name in ( + "Tensor", + "device", + "dtype", + "empty", + "from_list", + "from_torch", + "cat", + "sync_device", + "sync_stream", + "cancel_graph_recording", + ): + self.assertIn(f'"{public_name}"', init) + context = (package / "context.py").read_text(encoding="utf-8") + self.assertIn("def cancel_graph_recording():", context) + self.assertIn("_infinicore.cancel_graph_recording()", context) + for removed_name in ("infiniStatus_t", "QY", "KUNLUN", "ALI"): + self.assertNotIn(removed_name, init) + + def test_legacy_inference_runtime_and_device_aliases_are_removed(self) -> None: + obsolete_paths = ( + "src", + "include/infinicore_infer.h", + "include/infinicore_infer", + "scripts/libinfinicore_infer", + "scripts/deepseek.py", + "scripts/infer_task.py", + "scripts/jiuge.py", + "scripts/jiuge_awq.py", + "scripts/jiuge_gptq.py", + "scripts/jiuge_ppl.py", + "scripts/kvcache_pool.py", + "scripts/launch_server.py", + "scripts/qwen3vl.py", + "scripts/test_ceval.py", + ) + for relative_path in obsolete_paths: + self.assertFalse((ROOT / relative_path).exists(), relative_path) + + base_config = read_source("python/infinilm/base_config.py") + benchmark = read_source("test/bench/backends/infinilm.py") + readme = read_source("README.md") + for legacy_device in ("qy", "kunlun", "ali"): + self.assertNotIn(f'"{legacy_device}"', base_config) + self.assertNotIn(f'"{legacy_device}"', benchmark) + self.assertNotIn(f"--{legacy_device}", readme) + self.assertIn("raise ValueError", base_config) + self.assertNotIn("--attn=flash-attn", readme) + + llama_utils = read_source("test/models/llama/utils.py") + for legacy_dtype in ( + "DataType.F32", + "DataType.F16", + "DataType.BF16", + "DataType.I8", + "DataType.I16", + "DataType.I32", + "DataType.I64", + "DataType.U8", + "DataType.BOOL", + ): + self.assertNotIn(legacy_dtype, llama_utils) + intermediate_test = read_source( + "test/models/llama/test_intermediate_validation.py" + ) + self.assertIn("from infinilm.lib import _infinilm", intermediate_test) + self.assertNotIn("\n import _infinilm", intermediate_test) + self.assertIn("from infinicore.ops import add, matmul", intermediate_test) + self.assertNotIn("from infinicore.ops.", intermediate_test) + + def test_device_sentinel_is_not_public_or_indexable(self) -> None: + binding = read_source("csrc/infinicore/src/pybind11/device.hpp") + self.assertNotIn('.value("COUNT"', binding) + + context = read_source("csrc/infinicore/src/context/context_impl.cc") + self.assertIn("type_index >= runtime_table_.size()", context) + self.assertIn('throw std::invalid_argument("invalid device type")', context) + + def test_from_torch_disambiguates_cuda_compatible_devices(self) -> None: + source = read_source("python/infinicore/tensor.py") + + self.assertIn("def from_torch(torch_tensor, *, device=None):", source) + self.assertIn('owner.device.type == "cuda"', source) + self.assertIn("get_device()", source) + self.assertIn("Device(device)", source) + self.assertIn("pass device=", source) + + def test_from_torch_returns_an_owning_synchronized_copy(self) -> None: + source = read_source("python/infinicore/tensor.py") + from_torch = source[ + source.index("def from_torch(") : source.index("def from_numpy(") + ] + + self.assertIn("torch.cuda.synchronize(owner.device)", from_torch) + self.assertIn("borrowed = Tensor(", from_torch) + self.assertIn("result = empty(", from_torch) + self.assertIn("result.copy_(borrowed)", from_torch) + self.assertIn("sync_stream()", from_torch) + self.assertNotIn("return Tensor(\n _infinicore.from_blob", from_torch) + + def test_from_list_reads_uint64_as_unsigned(self) -> None: + source = read_source("csrc/infinicore/src/pybind11/from_list.cc") + + self.assertIn("PyLong_AsUnsignedLongLong", source) + self.assertIn("read_pyuint64", source) + self.assertIn("write_dtype_native", source) + + utils = read_source("python/infinicore/utils.py") + self.assertIn('for name in ("uint16", "uint32", "uint64"):', utils) + self.assertIn("getattr(torch, name, None)", utils) + self.assertIn("getattr(infinicore, name)", utils) + + def test_from_numpy_converts_before_taking_an_owning_copy(self) -> None: + source = read_source("python/infinicore/tensor.py") + binding = read_source("csrc/infinicore/src/pybind11/tensor.hpp") + + self.assertIn("infinicore_to_numpy_dtype", source) + self.assertIn("np.ascontiguousarray(array, dtype=numpy_dtype)", source) + self.assertIn("_infinicore._from_numpy_copy(owner", source) + self.assertNotIn("result.copy_(source)", source) + + self.assertIn('"_from_numpy_copy"', binding) + self.assertIn("py::buffer", binding) + self.assertIn("std::memcpy", binding) + + def test_to_numpy_is_owned_by_the_tensor_api(self) -> None: + tensor = read_source("python/infinicore/tensor.py") + generation = read_source("python/infinilm/generation/utils.py") + + self.assertIn("def to_numpy(self):", tensor) + self.assertIn("infinicore_to_numpy_dtype", tensor) + self.assertIn("ctypes.memmove", tensor) + self.assertIn("return infini_tensor.to_numpy()", generation) + self.assertNotIn("Tensor.to_numpy =", generation) + + def test_uninitialized_ones_factory_is_not_exported(self) -> None: + tensor = read_source("python/infinicore/tensor.py") + init = read_source("python/infinicore/__init__.py") + binding = read_source("csrc/infinicore/src/pybind11/tensor.hpp") + + self.assertNotIn("def ones(", tensor) + self.assertNotIn('"ones"', init) + self.assertNotIn('m.def("ones"', binding) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/static/test_infinicore_runtime_contracts.py b/test/static/test_infinicore_runtime_contracts.py new file mode 100644 index 000000000..8219dfddb --- /dev/null +++ b/test/static/test_infinicore_runtime_contracts.py @@ -0,0 +1,785 @@ +import re +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] + + +def read_source(relative_path: str) -> str: + return (ROOT / relative_path).read_text(encoding="utf-8") + + +def function_body(source: str, signature: str) -> str: + signature_start = source.index(signature) + body_start = source.index("{", signature_start) + depth = 0 + for index in range(body_start, len(source)): + if source[index] == "{": + depth += 1 + elif source[index] == "}": + depth -= 1 + if depth == 0: + return source[body_start : index + 1] + raise AssertionError(f"Unterminated function body for {signature}") + + +class InfiniCoreRuntimeContractsTest(unittest.TestCase): + def test_retained_op_headers_do_not_depend_on_legacy_public_abi(self) -> None: + headers = ( + "bitwise_right_shift.hpp", + "causal_conv1d.hpp", + "chunk_gated_delta_rule.hpp", + "gaussian_nll_loss.hpp", + "interpolate.hpp", + "prelu.hpp", + "relu6.hpp", + "recurrent_gated_delta_rule.hpp", + "sigmoid.hpp", + ) + for header in headers: + with self.subTest(header=header): + source = read_source(f"csrc/infinicore/include/infinicore/ops/{header}") + self.assertNotIn('#include "infinicore.h"', source) + self.assertNotRegex(source, re.compile(r"\b__export\b")) + + def test_exact_infiniops_adapter_set_is_registered(self) -> None: + expected_registration_counts = { + "add/add_infiniops.cc": 3, + "add_rms_norm/add_rms_norm_infiniops.cc": 3, + "causal_softmax/causal_softmax_infiniops.cc": 3, + "conv2d/conv2d_infiniops.cc": 1, + "embedding/embedding_infiniops.cc": 3, + "gelu/gelu_infiniops.cc": 1, + "gelutanh/gelutanh_infiniops.cc": 1, + "gemm/gemm_infiniops.cc": 3, + "kv_caching/kv_caching_infiniops.cc": 3, + "paged_attention/paged_attention_infiniops.cc": 3, + "paged_attention_prefill/paged_attention_prefill_infiniops.cc": 1, + "paged_caching/paged_caching_infiniops.cc": 3, + "random_sample/random_sample_infiniops.cc": 1, + "rearrange/rearrange_infiniops.cc": 3, + "relu/relu_infiniops.cc": 1, + "rms_norm/rms_norm_infiniops.cc": 3, + "rope/rope_infiniops.cc": 3, + "sigmoid/sigmoid_infiniops.cc": 3, + "silu/silu_infiniops.cc": 1, + "silu_and_mul/silu_and_mul_infiniops.cc": 3, + "softmax/softmax_infiniops.cc": 1, + "swiglu/swiglu_infiniops.cc": 3, + "topksoftmax/topksoftmax_infiniops.cc": 3, + } + ops_root = ROOT / "csrc/infinicore/src/ops" + adapters = { + path.relative_to(ops_root).as_posix() + for path in ops_root.glob("*/*_infiniops.cc") + } + self.assertEqual(adapters, set(expected_registration_counts)) + + for relative_path, expected_count in expected_registration_counts.items(): + with self.subTest(adapter=relative_path): + source = read_source(f"csrc/infinicore/src/ops/{relative_path}") + self.assertEqual( + source.count("registerSupportedDevices("), expected_count + ) + + rearrange = read_source( + "csrc/infinicore/src/ops/rearrange/rearrange_infiniops.cc" + ) + for dispatcher in ( + "plan_dispatcher()", + "run_dispatcher()", + "cleanup_dispatcher()", + ): + self.assertIn(f"Rearrange::{dispatcher}", rearrange) + self.assertNotIn("InfiniCore/InfiniOp implementation active", rearrange) + + def test_empty_operator_dispatch_throws(self) -> None: + source = read_source( + "csrc/infinicore/include/infinicore/ops/common/dispatcher.hpp" + ) + lookup = function_body(source, "Fn lookup(Device::Type device_type) const") + self.assertIn("if (fn == nullptr)", lookup) + self.assertIn("throw std::runtime_error", lookup) + self.assertIn("No operator implementation is registered for device", lookup) + + def test_unavailable_configurable_surfaces_are_rejected(self) -> None: + quant_source = read_source("csrc/config/quant_config.cpp") + quant_method = function_body( + quant_source, "QuantConfig::get_quantization_method() const" + ) + self.assertNotIn( + "make_shared", quant_method + ) + self.assertRegex( + quant_method, + re.compile( + r'quant_method == "compressed-tensors"\)\s*\{\s*' + r"throw std::runtime_error" + ), + ) + + quant_header = read_source("csrc/config/quant_config.hpp") + kv_config = function_body( + quant_header, + "void set_kv_quant_scheme(infinicore::DataType kv_cache_dtype)", + ) + self.assertNotIn("KVQuantAlgo::INT8", kv_config) + self.assertIn("throw std::runtime_error", kv_config) + self.assertIn("KV cache INT8 quantization is unsupported", kv_config) + + attention_config = read_source("csrc/backends/attention_backends.hpp") + parse_backend = function_body( + attention_config, + "inline AttentionBackend parse_attention_backend", + ) + self.assertNotIn("return AttentionBackend::FLASH_ATTN", parse_backend) + self.assertNotIn("return AttentionBackend::FLASHINFER", parse_backend) + for backend in ("flash-attn", "flashinfer"): + self.assertRegex( + parse_backend, + re.compile( + rf'backend == "{backend}"\)\s*\{{\s*' + r"throw std::invalid_argument" + ), + ) + + def test_modern_model_support_is_gated_before_rank_workers(self) -> None: + source = read_source("csrc/config/config_factory.cpp") + create = function_body( + source, + "std::shared_ptr " + "ConfigFactory::createConfig", + ) + + self.assertIn('kModernModelTypes{"qwen3"}', create) + self.assertIn("supported model types: qwen3", create) + self.assertLess( + create.index("if (it == config_map.end())"), + create.index("kModernModelTypes.find(model_type)"), + ) + + infer_engine = read_source("csrc/engine/infer_engine.cpp") + config_creation = infer_engine.index("ConfigFactory::createConfig") + rank_workers = infer_engine.index("RankWorker", config_creation) + self.assertLess(config_creation, rank_workers) + + readme = read_source("README.md") + self.assertIn("currently supports `qwen3`", readme) + + static_attention = read_source("csrc/layers/attention/backends/static_attn.cpp") + self.assertNotIn("if (false)", static_attention) + self.assertNotIn("op::flash_attention", static_attention) + + def test_many_variable_collectives_preflight_before_recording(self) -> None: + cases = ( + ( + "csrc/infinicore/src/ops/distributed/allgather.cc", + "void allgatherv_many_", + "validate_allgatherv", + "AllGather::execute", + ), + ( + "csrc/infinicore/src/ops/distributed/reduce_scatter.cc", + "void reduce_scatterv_many_", + "validate_reduce_scatterv", + "ReduceScatter::execute", + ), + ) + loop = "for (size_t i = 0; i < inputs.size(); ++i)" + for path, signature, validator, collective in cases: + with self.subTest(path=path): + source = read_source(path) + body = function_body(source, signature) + loop_positions = [ + match.start() for match in re.finditer(re.escape(loop), body) + ] + self.assertEqual(len(loop_positions), 2) + self.assertIn("!split_sizes.empty()", body) + world_size = body.index("!split_sizes.empty()") + validation = body.index(validator) + execution = body.index(collective) + self.assertLess(world_size, loop_positions[0]) + self.assertLess(loop_positions[0], validation) + self.assertLess(validation, loop_positions[1]) + self.assertLess(loop_positions[1], execution) + + validator_body = function_body(source, f"void {validator}") + for token in ( + "dtype()", + "INFINICORE_ASSERT_TENSORS_SAME_DEVICE", + "is_contiguous()", + "ndim()", + "shape()", + "split_sizes", + "toInfinicclDataType", + ): + self.assertIn(token, validator_body) + + def test_communication_group_gates_rank_initialization(self) -> None: + source = read_source("csrc/engine/distributed/communication_group.cpp") + constructor = function_body(source, "CommunicationGroup::CommunicationGroup") + + for token in ( + "std::mutex start_mutex", + "std::condition_variable start_cv", + "bool start = false", + "bool cancel = false", + ): + self.assertIn(token, constructor) + + wait = constructor.index("start_cv.wait") + init = constructor.index("infinicclCommInitRank") + self.assertLess(wait, init) + self.assertIn("if (cancel)", constructor[wait:init]) + + release = constructor.index("start = true", init) + release_notify = constructor.index("start_cv.notify_all()", release) + join = constructor.index("worker.join()", release_notify) + self.assertLess(release, release_notify) + self.assertLess(release_notify, join) + + failure = constructor.rindex("catch (...)") + failure_path = constructor[failure:] + cancel = failure_path.index("cancel = true") + cancel_notify = failure_path.index("start_cv.notify_all()") + cancel_join = failure_path.index("worker.join()") + self.assertLess(cancel, cancel_notify) + self.assertLess(cancel_notify, cancel_join) + + def test_graph_recording_is_synchronized_and_thread_owned(self) -> None: + header = read_source("csrc/infinicore/src/graph/graph_manager.hpp") + source = read_source("csrc/infinicore/src/graph/graph.cc") + context_header = read_source("csrc/infinicore/src/context/context_impl.hpp") + context_source = read_source("csrc/infinicore/src/context/context_impl.cc") + + self.assertIn("mutable std::mutex mutex_;", header) + self.assertIn("std::thread::id capture_owner_;", header) + for signature in ( + "bool GraphManager::is_recording() const", + "void GraphManager::start_recording()", + "void GraphManager::add_operator(std::shared_ptr op)", + "std::shared_ptr GraphManager::stop_recording()", + ): + body = function_body(source, signature) + self.assertIn("std::lock_guard", body) + self.assertIn("capture_owner_", body) + + self.assertIn( + "static thread_local std::shared_ptr graph_runtime_;", + context_header, + ) + start = function_body(context_source, "void ContextImpl::startGraphRecording()") + stop = function_body( + context_source, + "std::shared_ptr ContextImpl::stopGraphRecording()", + ) + self.assertLess( + start.index("current_runtime_->startGraphRecording()"), + start.index("graph_runtime_ = current_runtime_"), + ) + self.assertLess( + stop.index("current_runtime_ = owner"), + stop.index("owner->stopGraphRecording()"), + ) + self.assertIn("current_runtime_ = previous", stop) + self.assertIn("current_runtime_->activate()", stop) + self.assertIn("graph_runtime_.reset()", stop) + + self.assertIn("void cancelGraphRecording() noexcept;", context_header) + cancel = function_body( + context_source, "void ContextImpl::cancelGraphRecording() noexcept" + ) + self.assertIn("std::exchange(graph_runtime_, nullptr)", cancel) + self.assertIn("owner->cancelGraphRecording()", cancel) + self.assertIn("current_runtime_ = previous", cancel) + self.assertNotIn("isGraphRecording()", cancel) + + def test_graph_capture_memory_is_leased_until_graph_destruction(self) -> None: + allocator_header = read_source( + "csrc/infinicore/src/context/allocators/pinnable_block_allocator.hpp" + ) + allocator_source = read_source( + "csrc/infinicore/src/context/allocators/pinnable_block_allocator.cc" + ) + tensor_source = read_source("csrc/infinicore/src/tensor/tensor.cc") + graph_header = read_source("csrc/infinicore/include/infinicore/graph/graph.hpp") + runtime_source = read_source("csrc/infinicore/src/context/runtime/runtime.cc") + + self.assertIn("class PinLease", allocator_header) + self.assertIn("size_t pin_count = 0", allocator_header) + self.assertIn("std::shared_ptr commit_pin_mode()", allocator_header) + self.assertIn("block->pin_count == 0", allocator_source) + self.assertIn("retain_for_capture", allocator_source) + self.assertIn("context::retainGraphMemory", tensor_source) + self.assertIn("std::shared_ptr allocation_lease_", graph_header) + + stop = function_body( + runtime_source, + "std::shared_ptr Runtime::stopGraphRecording()", + ) + commit = stop.index("commit_pin_mode()") + retain = stop.index("retain_runtime") + finish = stop.index("finish_recording()", retain) + self.assertLess(commit, retain) + self.assertLess(retain, finish) + + def test_reinstantiated_blocks_have_one_free_list_entry(self) -> None: + allocator_source = read_source( + "csrc/infinicore/src/context/allocators/pinnable_block_allocator.cc" + ) + mark_in_use = function_body( + allocator_source, "size_t PinnableBlockAllocator::mark_in_use_" + ) + + self.assertIn("cls.free_blocks.erase", mark_in_use) + self.assertIn("std::remove(", mark_in_use) + self.assertIn("cls.free_blocks.end()", mark_in_use) + + def test_graph_destruction_synchronizes_before_releasing_leases(self) -> None: + graph_header = read_source("csrc/infinicore/include/infinicore/graph/graph.hpp") + graph_source = read_source("csrc/infinicore/src/graph/graph.cc") + runtime_header = read_source("csrc/infinicore/src/context/runtime/runtime.hpp") + runtime_source = read_source("csrc/infinicore/src/context/runtime/runtime.cc") + + self.assertIn("~Graph() noexcept", graph_header) + destructor = function_body(graph_source, "Graph::~Graph() noexcept") + self.assertIn("runtime_lease_->syncStreamForCleanup()", destructor) + + self.assertIn("void syncStreamForCleanup() noexcept;", runtime_header) + self.assertIn("friend class graph::Graph;", runtime_header) + cleanup = function_body( + runtime_source, "void Runtime::syncStreamForCleanup() noexcept" + ) + self.assertIn("StreamSynchronize(stream_)", cleanup) + self.assertIn("catch (const std::exception &error)", cleanup) + self.assertIn("catch (...)", cleanup) + self.assertIn("restore_runtime", cleanup) + + def test_graph_compilers_cancel_capture_on_exception(self) -> None: + context_header = read_source( + "csrc/infinicore/include/infinicore/context/context.hpp" + ) + runtime_header = read_source("csrc/infinicore/src/context/runtime/runtime.hpp") + compiler_header = read_source("csrc/engine/compiler/graph_compiler.hpp") + paged_source = read_source("csrc/engine/compiler/paged_compiler.cpp") + static_source = read_source("csrc/engine/compiler/static_batching_compiler.cpp") + + self.assertIn("void cancelGraphRecording() noexcept;", context_header) + self.assertIn("void cancelGraphRecording() noexcept;", runtime_header) + guard = compiler_header[compiler_header.index("class GraphRecordingGuard") :] + self.assertIn("context::startGraphRecording()", guard) + self.assertIn("context::cancelGraphRecording()", guard) + self.assertIn("context::stopGraphRecording()", guard) + self.assertIn("~GraphRecordingGuard() noexcept", guard) + self.assertIn("GraphRecordingGuard(GraphRecordingGuard &&) = delete", guard) + self.assertIn( + "GraphRecordingGuard &operator=(GraphRecordingGuard &&) = delete", guard + ) + finish = function_body( + compiler_header, + "std::shared_ptr finish()", + ) + self.assertLess( + finish.index("context::stopGraphRecording()"), + finish.index("active_ = false"), + ) + for source in (paged_source, static_source): + self.assertIn("GraphRecordingGuard recording;", source) + self.assertIn("auto graph = recording.finish();", source) + self.assertNotIn("context::startGraphRecording()", source) + self.assertNotIn("context::stopGraphRecording()", source) + self.assertLess( + paged_source.index("for (size_t b : decode_batch_sizes_)"), + paged_source.index("GraphRecordingGuard recording;"), + ) + + graph_manager = function_body( + read_source("csrc/infinicore/src/graph/graph.cc"), + "void GraphManager::cancel_recording()", + ) + self.assertIn("std::exchange(graph_, nullptr)", graph_manager) + self.assertLess( + graph_manager.index("std::lock_guard"), graph_manager.index("graph.reset()") + ) + + def test_foreign_capture_rejects_operator_dispatch(self) -> None: + manager_header = read_source("csrc/infinicore/src/graph/graph_manager.hpp") + graph_header = read_source("csrc/infinicore/include/infinicore/graph/graph.hpp") + graph_source = read_source("csrc/infinicore/src/graph/graph.cc") + rearrange_source = read_source("csrc/infinicore/src/ops/rearrange/rearrange.cc") + context_source = read_source("csrc/infinicore/src/context/context_impl.cc") + + for state in ("kInactive", "kActiveOwner", "kActiveNonOwner"): + self.assertIn(state, manager_header) + self.assertIn("CaptureState capture_state() const;", manager_header) + + is_recording = function_body( + graph_source, "bool GraphManager::is_recording() const" + ) + self.assertIn("kActiveNonOwner", is_recording) + self.assertIn("throw std::runtime_error", is_recording) + + dispatch_macro = graph_header.index("INFINICORE_GRAPH_OP_RECORD_OR_RUN") + dispatch_check = graph_header.index( + "context::isGraphRecording()", dispatch_macro + ) + dispatch_plan = graph_header.index("std::make_shared", dispatch_macro) + self.assertLess(dispatch_check, dispatch_plan) + + rearrange_execute = function_body( + rearrange_source, "void Rearrange::execute(Tensor y, const Tensor &x)" + ) + self.assertLess( + rearrange_execute.index("context::isGraphRecording()"), + rearrange_execute.index("std::make_shared"), + ) + + set_device = function_body( + context_source, "void ContextImpl::setDevice(Device device)" + ) + self.assertNotIn("cannot switch devices during graph recording", set_device) + + def test_context_owns_one_runtime_per_thread_and_device(self) -> None: + context_header = read_source("csrc/infinicore/src/context/context_impl.hpp") + context_source = read_source("csrc/infinicore/src/context/context_impl.cc") + runtime_header = read_source("csrc/infinicore/src/context/runtime/runtime.hpp") + runtime_source = read_source("csrc/infinicore/src/context/runtime/runtime.cc") + graph_source = read_source("csrc/infinicore/src/graph/graph.cc") + + self.assertIn("using ThreadRuntimes =", context_header) + self.assertRegex( + context_header, + re.compile( + r"std::unordered_map>" + ), + ) + self.assertIn( + "static thread_local std::shared_ptr current_runtime_;", + context_header, + ) + create_runtime = function_body( + context_source, + "std::shared_ptr ContextImpl::getOrCreateRuntimeLocked", + ) + self.assertIn("std::shared_ptr(new Runtime(device))", create_runtime) + self.assertIn("found->second.lock()", create_runtime) + self.assertIn("thread_id", create_runtime) + get_current = function_body( + context_source, "Runtime *ContextImpl::getCurrentRuntime()" + ) + set_device = function_body( + context_source, "void ContextImpl::setDevice(Device device)" + ) + for body in (get_current, set_device): + self.assertIn("std::this_thread::get_id()", body) + self.assertIn("std::lock_guard", body) + self.assertIn("getOrCreateRuntimeLocked", body) + self.assertLess( + set_device.index("current_runtime_ = std::move(runtime)"), + set_device.rindex("current_runtime_->activate()"), + ) + self.assertIn("kDefaultDevicePriority", get_current) + + constructor = function_body(context_source, "ContextImpl::ContextImpl()") + for device_type in ( + "kCpu", + "kNvidia", + "kCambricon", + "kAscend", + "kMetax", + "kMoore", + "kIluvatar", + "kHygon", + ): + self.assertIn( + f"initializeDeviceType()", constructor + ) + + self.assertIn("mutable std::mutex stream_mutex_;", runtime_header) + self.assertIn( + "mutable infini::rt::runtime::Stream stream_ = nullptr;", + runtime_header, + ) + self.assertNotIn("std::unordered_map", runtime_header) + + for signature in ( + "std::shared_ptr Runtime::allocateMemory(size_t size)", + "std::shared_ptr Runtime::allocatePinnedHostMemory(size_t size)", + "std::shared_ptr Runtime::reinstantiateBlob(std::shared_ptr blob)", + ): + body = function_body(runtime_source, signature) + self.assertIn("shared_from_this()", body) + self.assertIn("[runtime]", body) + allocate_memory = function_body( + runtime_source, "std::shared_ptr Runtime::allocateMemory" + ) + self.assertIn("runtime->releaseDeviceMemory(p)", allocate_memory) + + release = function_body( + runtime_source, "void Runtime::releaseDeviceMemory(std::byte *ptr) noexcept" + ) + self.assertIn("ContextImpl::current_runtime_", release) + sync = release.index("DeviceSynchronize") + deallocate = release.index("device_memory_allocator_->deallocate") + self.assertLess(sync, deallocate) + + allocator_header = read_source( + "csrc/infinicore/src/context/allocators/pinnable_block_allocator.hpp" + ) + allocator_source = read_source( + "csrc/infinicore/src/context/allocators/pinnable_block_allocator.cc" + ) + for method in ( + "void begin_pin_mode();", + "std::shared_ptr commit_pin_mode();", + "void cancel_pin_mode();", + "void retain_for_capture(void *ptr);", + ): + self.assertIn(method, allocator_header) + self.assertIn("capture_frozen_blocks_", allocator_header) + freeze = function_body( + allocator_source, + "void PinnableBlockAllocator::freeze_for_capture_", + ) + self.assertIn("capture_frozen_block_set_.insert", freeze) + self.assertIn("++block->pin_count", freeze) + self.assertIn("capture_frozen_blocks_.push_back(block)", freeze) + commit_pin = function_body( + allocator_source, "PinnableBlockAllocator::commit_pin_mode()" + ) + self.assertIn("blocks = std::move(capture_frozen_blocks_)", commit_pin) + self.assertIn("pinned_mode_ = false", commit_pin) + cancel_pin = function_body( + allocator_source, "void PinnableBlockAllocator::cancel_pin_mode()" + ) + self.assertIn("--block->pin_count", cancel_pin) + self.assertIn("capture_frozen_blocks_.clear()", cancel_pin) + + cancel_graph = function_body( + runtime_source, "void Runtime::cancelGraphRecording() noexcept" + ) + self.assertIn("graph_manager_->cancel_recording()", cancel_graph) + self.assertIn("device_memory_allocator_->cancel_pin_mode()", cancel_graph) + stop_graph = function_body( + runtime_source, + "std::shared_ptr Runtime::stopGraphRecording()", + ) + failure_path = stop_graph[stop_graph.index("catch (...)") :] + self.assertIn("device_memory_allocator_->cancel_pin_mode()", failure_path) + self.assertIn("device_memory_allocator_->commit_pin_mode()", stop_graph) + + runtime_constructor = function_body( + runtime_source, "Runtime::Runtime(Device device)" + ) + self.assertNotIn("StreamCreate", runtime_constructor) + stream = function_body( + runtime_source, "infini::rt::runtime::Stream Runtime::stream() const" + ) + for token in ("stream_mutex_", "StreamCreate", "stream_"): + self.assertIn(token, stream) + self.assertIn("if (stream_ == nullptr)", stream) + self.assertIn("return stream_", stream) + + destructor = function_body(runtime_source, "Runtime::~Runtime() noexcept") + self.assertIn("stream_mutex_", destructor) + self.assertIn("stream_", destructor) + self.assertIn("StreamDestroy", destructor) + + for signature in ( + "void Runtime::syncStream()", + "void Runtime::memcpyH2D(void *dst, const void *src, size_t size, bool async)", + "void Runtime::memcpyD2H(void *dst, const void *src, size_t size)", + "void Runtime::memcpyD2D(void *dst, const void *src, size_t size, bool async)", + "void Runtime::setDeviceMemory(void *ptr, int value, size_t count)", + "void Runtime::setDeviceMemoryAsync(void *ptr, int value, size_t count, infini::rt::runtime::Stream stream)", + "void Runtime::recordEvent(infini::rt::runtime::Event event, infini::rt::runtime::Stream stream)", + "void Runtime::streamWaitEvent(infini::rt::runtime::Stream stream, infini::rt::runtime::Event event)", + ): + self.assertIn("stream()", function_body(runtime_source, signature)) + + for signature in ( + "void Runtime::memcpyH2D(void *dst, const void *src, size_t size, bool async)", + "void Runtime::memcpyD2H(void *dst, const void *src, size_t size)", + "void Runtime::memcpyD2D(void *dst, const void *src, size_t size, bool async)", + ): + body = function_body(runtime_source, signature) + self.assertIn("MemcpyAsync", body) + self.assertIn("StreamSynchronize(current_stream)", body) + set_memory = function_body( + runtime_source, + "void Runtime::setDeviceMemory(void *ptr, int value, size_t count)", + ) + self.assertIn("MemsetAsync", set_memory) + self.assertIn("StreamSynchronize(current_stream)", set_memory) + + instantiate = function_body(graph_source, "void Graph::instantiate()") + self.assertIn("device_graph_->stream = context::getStream()", instantiate) + + def test_pinned_host_memory_remains_cpu_addressed(self) -> None: + source = read_source("csrc/infinicore/src/context/runtime/runtime.cc") + allocate_pinned = function_body( + source, + "std::shared_ptr Runtime::allocatePinnedHostMemory(size_t size)", + ) + + self.assertIn("Device{Device::Type::kCpu}", allocate_pinned) + self.assertNotIn("data_ptr, size, device_", allocate_pinned) + + def test_cross_device_copy_stages_through_host_memory(self) -> None: + source = read_source("csrc/infinicore/src/tensor/copy.cc") + copy_from = function_body(source, "void TensorImpl::copy_from(Tensor src)") + + self.assertIn("auto host_staging = Tensor::empty(", copy_from) + self.assertIn("Device{Device::Type::kCpu}", copy_from) + self.assertIn("host_staging->copy_from(src)", copy_from) + self.assertIn("this->copy_from(host_staging)", copy_from) + + def test_cpu_copy_does_not_require_an_accelerator_operator(self) -> None: + source = read_source("csrc/infinicore/src/tensor/copy.cc") + copy_from = function_body(source, "void TensorImpl::copy_from(Tensor src)") + + self.assertIn("void copyCpuStrided(", source) + self.assertIn("Cannot copy from tensor with different dtype", copy_from) + cpu_copy = copy_from.index("if (this->device().type() == Device::Type::kCpu)") + accelerator_copy = copy_from.index("op::rearrange_", cpu_copy) + self.assertLess(copy_from.index("copyCpuStrided(", cpu_copy), accelerator_copy) + + def test_native_infini_rt_graph_runtime_is_enabled(self) -> None: + xmake = read_source("xmake.lua") + target_start = xmake.index('target("infinicore_runtime")') + target_end = xmake.index("target_end()", target_start) + runtime_target = xmake[target_start:target_end] + + self.assertIn("USE_INFINIRT_GRAPH", runtime_target) + + graph_header = read_source("csrc/infinicore/include/infinicore/graph/graph.hpp") + lease = graph_header.index("runtime_lease_") + operators = graph_header.index("op_list_") + device_graph = graph_header.index("device_graph_") + self.assertLess(lease, operators) + self.assertLess(lease, device_graph) + self.assertIn("friend class ::infinicore::Runtime", graph_header) + + runtime_source = read_source("csrc/infinicore/src/context/runtime/runtime.cc") + stop = function_body( + runtime_source, + "std::shared_ptr Runtime::stopGraphRecording()", + ) + self.assertLess( + stop.index("device_memory_allocator_->commit_pin_mode()"), + stop.index("graph->retain_runtime(shared_from_this(),"), + ) + self.assertLess( + stop.index("graph->retain_runtime(shared_from_this(),"), + stop.rindex("graph_manager_->finish_recording()"), + ) + + graph_source = read_source("csrc/infinicore/src/graph/graph.cc") + manager_stop = function_body( + graph_source, + "std::shared_ptr GraphManager::stop_recording()", + ) + self.assertLess( + manager_stop.index("std::exchange(graph_, nullptr)"), + manager_stop.index("graph->instantiate()"), + ) + + def test_h2d_copy_does_not_outlive_its_host_source(self) -> None: + source = read_source("csrc/infinicore/src/tensor/copy.cc") + copy_from = function_body(source, "void TensorImpl::copy_from(Tensor src)") + + self.assertNotIn("const bool async", copy_from) + self.assertIn( + "context::memcpyH2D(this->data(), src->data(), copy_size, false)", + copy_from, + ) + self.assertIn( + "context::memcpyH2D(local_src->data(), src->data(), copy_size, false)", + copy_from, + ) + + def test_runtime_table_lazy_initialization_is_synchronized(self) -> None: + header = read_source("csrc/infinicore/src/context/context_impl.hpp") + source = read_source("csrc/infinicore/src/context/context_impl.cc") + + self.assertIn("mutable std::mutex runtime_table_mutex_;", header) + for signature in ( + "Runtime *ContextImpl::getCurrentRuntime()", + "void ContextImpl::setDevice(Device device)", + "size_t ContextImpl::getDeviceCount(Device::Type type)", + ): + self.assertIn("std::lock_guard", function_body(source, signature)) + + def test_allocator_pin_transaction_is_synchronized(self) -> None: + header = read_source( + "csrc/infinicore/src/context/allocators/pinnable_block_allocator.hpp" + ) + source = read_source( + "csrc/infinicore/src/context/allocators/pinnable_block_allocator.cc" + ) + + self.assertIn("std::thread::id pin_owner_;", header) + for signature in ( + "void PinnableBlockAllocator::begin_pin_mode()", + "PinnableBlockAllocator::commit_pin_mode()", + "void PinnableBlockAllocator::cancel_pin_mode()", + ): + transaction = function_body(source, signature) + self.assertIn("std::lock_guard", transaction) + self.assertIn("std::this_thread::get_id()", transaction) + allocate = function_body( + source, "std::byte *PinnableBlockAllocator::allocate(size_t size)" + ) + self.assertIn("freeze_for_capture_(block)", allocate) + + def test_pinned_host_allocator_queue_is_synchronized(self) -> None: + header = read_source( + "csrc/infinicore/src/context/allocators/device_pinned_allocator.hpp" + ) + source = read_source( + "csrc/infinicore/src/context/allocators/device_pinned_allocator.cc" + ) + + self.assertIn("std::mutex gc_mutex_;", header) + deallocate = function_body( + source, "void DevicePinnedHostAllocator::deallocate(std::byte *ptr)" + ) + gc = function_body(source, "void DevicePinnedHostAllocator::gc()") + self.assertIn("std::lock_guard", deallocate) + self.assertIn("std::lock_guard", gc) + + def test_runtime_destructor_is_nonthrowing_and_dependency_ordered(self) -> None: + header = read_source("csrc/infinicore/src/context/runtime/runtime.hpp") + source = read_source("csrc/infinicore/src/context/runtime/runtime.cc") + + self.assertIn("~Runtime() noexcept;", header) + destructor = function_body(source, "Runtime::~Runtime() noexcept") + self.assertIn("warn_runtime_cleanup_failure", destructor) + self.assertIn("infini::rt::runtime::SetDevice", destructor) + self.assertIn("infini::rt::runtime::StreamSynchronize", destructor) + self.assertIn("infini::rt::runtime::StreamDestroy", destructor) + ordered_cleanup = ( + "infini::rt::runtime::StreamSynchronize", + "graph_manager_.reset()", + "pinned_host_memory_allocator_.reset()", + "device_memory_allocator_.reset()", + "infini::rt::runtime::StreamDestroy", + ) + positions = [destructor.index(token) for token in ordered_cleanup] + self.assertEqual(positions, sorted(positions)) + self.assertIn("restore_runtime", destructor) + + def test_tensor_view_requires_mutable_tensor(self) -> None: + header = read_source("csrc/infinicore/include/infinicore/tensor.hpp") + source = read_source("csrc/infinicore/src/tensor/tensor.cc") + + self.assertIn("infini::rt::TensorView view();", header) + self.assertNotIn("infini::rt::TensorView view() const;", header) + self.assertIn("infini::rt::TensorView TensorImpl::view()", source) + self.assertNotIn("infini::rt::TensorView TensorImpl::view() const", source) + + +if __name__ == "__main__": + unittest.main() diff --git a/xmake.lua b/xmake.lua index aab1a0c70..a779cbaf7 100644 --- a/xmake.lua +++ b/xmake.lua @@ -1,59 +1,108 @@ add_requires("pybind11") local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") +local CUDA_ROOT = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or "/usr/local/cuda" set_toolchains("gcc") --- Add spdlog from third_party directory add_includedirs("third_party/spdlog/include") add_includedirs("third_party/json/single_include/") +if os.isdir(CUDA_ROOT .. "/include") then + add_includedirs(CUDA_ROOT .. "/include") +end -target("infinicore_infer") +target("infinicore_runtime") + set_default(false) set_kind("shared") + set_languages("cxx17") + add_defines("ENABLE_INFINIOPS_API", "USE_INFINIRT_GRAPH") + + if not is_plat("windows") then + add_cxflags("-fPIC") + end + + add_includedirs("csrc/infinicore/include", { public = true }) + add_includedirs("csrc/infinicore/src") + add_includedirs("csrc/infinicore/utils") + add_includedirs(INFINI_ROOT .. "/include", { public = true }) + add_includedirs(INFINI_ROOT .. "/include/infiniccl", { public = true }) - add_includedirs("include", { public = false }) - add_includedirs(INFINI_ROOT.."/include", { public = true }) + add_linkdirs(INFINI_ROOT .. "/lib", INFINI_ROOT .. "/lib64", { public = true }) + if is_plat("linux") then + add_rpathdirs("$ORIGIN") + elseif is_plat("macosx") then + add_rpathdirs("@loader_path") + end + add_links("infiniops", "infiniccl", "infinirt", { public = true }) - add_linkdirs(INFINI_ROOT.."/lib") - add_links("infiniop", "infinirt", "infiniccl") + add_files("csrc/infinicore/src/*.cc") + add_files("csrc/infinicore/src/context/*.cc") + add_files("csrc/infinicore/src/context/*/*.cc") + add_files("csrc/infinicore/src/graph/*.cc") + add_files("csrc/infinicore/src/nn/*.cc") + add_files("csrc/infinicore/src/ops/*/*.cc") + add_files("csrc/infinicore/src/ops/dequant/*/*.cc") + add_files("csrc/infinicore/src/ops/quant/*/*.cc") + add_files("csrc/infinicore/src/tensor/*.cc") + add_files("csrc/infinicore/utils/*.cc") + remove_files("csrc/infinicore/src/ops/*/*_cpu.cc") + remove_files("csrc/infinicore/src/ops/*/*_flashattn.cc") + remove_files("csrc/infinicore/src/ops/*/*_hygon.cc") + remove_files("csrc/infinicore/src/ops/*/*_moore.cc") + + set_installdir("python/infinicore") +target_end() + +target("_infinicore") + add_packages("pybind11") + set_default(false) + add_rules("python.module", { soabi = true }) set_languages("cxx17") - set_warnings("all", "error") - - add_files("src/models/*.cpp") - add_files("src/models/*/*.cpp") - add_files("src/tensor/*.cpp") - add_files("src/allocator/*.cpp") - add_files("src/dataloader/*.cpp") - add_files("src/cache_manager/*.cpp") - add_includedirs("include") - - set_installdir(INFINI_ROOT) - add_installfiles("include/infinicore_infer.h", {prefixdir = "include"}) - add_installfiles("include/infinicore_infer/models/*.h", {prefixdir = "include/infinicore_infer/models"}) + set_kind("shared") + + add_deps("infinicore_runtime") + + if is_plat("linux") then + add_rpathdirs("$ORIGIN") + elseif is_plat("macosx") then + add_rpathdirs("@loader_path") + end + + add_includedirs("csrc/infinicore/include") + add_includedirs("csrc/infinicore/src") + add_includedirs("csrc/infinicore/utils") + add_includedirs(INFINI_ROOT .. "/include") + + add_files("csrc/infinicore/src/pybind11/infinicore.cc") + add_files("csrc/infinicore/src/pybind11/from_list.cc") + + set_installdir("python/infinicore") target_end() target("_infinilm") add_packages("pybind11") set_default(false) - add_rules("python.module", {soabi = true}) + add_rules("python.module", { soabi = true }) set_languages("cxx17") set_kind("shared") - local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") + add_deps("infinicore_runtime") - -- add_includedirs("csrc", { public = false }) - -- add_includedirs("csrc/pybind11", { public = false }) - add_includedirs(INFINI_ROOT.."/include", { public = true }) - add_includedirs("include", { public = false }) - -- spdlog is already included globally via add_includedirs at the top + if is_plat("linux") then + add_rpathdirs("$ORIGIN/../../infinicore/lib") + elseif is_plat("macosx") then + add_rpathdirs("@loader_path/../../infinicore/lib") + end - add_linkdirs(INFINI_ROOT.."/lib") - add_links("infinicore_cpp_api", "infiniop", "infinirt", "infiniccl") + add_includedirs("csrc/infinicore/include") + add_includedirs(INFINI_ROOT .. "/include") + add_includedirs(INFINI_ROOT .. "/include/infiniccl") - -- Add src files add_files("csrc/**.cpp") add_files("csrc/**.cc") + remove_files("csrc/infinicore/**.cpp") + remove_files("csrc/infinicore/**.cc") set_installdir("python/infinilm") target_end() From a94d56e36ee45210f201fca4856e2f7c5d48f7bd Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 16 Jul 2026 13:06:30 +0800 Subject: [PATCH 2/6] docs: define Infini stack repository boundaries --- ...fini-stack-repository-boundaries-design.md | 216 ++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-16-infini-stack-repository-boundaries-design.md diff --git a/docs/superpowers/specs/2026-07-16-infini-stack-repository-boundaries-design.md b/docs/superpowers/specs/2026-07-16-infini-stack-repository-boundaries-design.md new file mode 100644 index 000000000..579acb16a --- /dev/null +++ b/docs/superpowers/specs/2026-07-16-infini-stack-repository-boundaries-design.md @@ -0,0 +1,216 @@ +# Infini Stack Repository Boundaries Design + +## Status + +Approved in conversation on 2026-07-16. This document is the implementation +contract for the NVIDIA migration slice. + +## Goal + +Make InfiniCore a versioned manifest for InfiniRT, InfiniOps, and InfiniCCL, +with no legacy runtime API or build ownership. Move the integration build and +runtime-facing ownership into InfiniLM, and update the shared InfiniLM CI to use +that new boundary. + +The migration is intentionally a clean break. InfiniCore must not retain +compatibility definitions such as `infiniStatus_t`, legacy `infiniop` or +`infinirt` source trees, forwarding headers, Python packaging, or a second build +system. + +## Confirmed NVIDIA Support Boundary + +The current submodule revisions are sufficient for the first migration target, +but this does not imply complete inference support on every model or backend. + +- InfiniRT covers the device, memory, stream, event, and graph lifecycle used by + the migrated InfiniLM runtime on NVIDIA. +- InfiniOps provides the 23-operation allowlist used by the current dense, + non-quantized Qwen3 path with the default static attention implementation. +- InfiniCCL provides the initialization, destruction, and AllReduce behavior + required by dense Qwen3 tensor parallelism. +- The validated end-to-end baseline is NVIDIA, Qwen3, tensor parallel size 1, + without quantization. Tensor parallel size 2 is a required acceptance test for + this cleanup. + +The following remain outside this slice: other accelerator backends, complete +model-family coverage, quantized paths, FlashAttention or FlashInfer paths, and +MoE expert parallelism. In particular, InfiniCCL's NVIDIA backend does not yet +provide the AllGather and ReduceScatter operations used by the InfiniLM MoE +expert-parallel path. Those additions require separate feature work and must not +be implied by this refactor. + +## Repository Ownership + +| Repository | Responsibility after the migration | +| --- | --- | +| InfiniRT | Runtime primitives and public runtime types | +| InfiniOps | Operator implementations and public operator API | +| InfiniCCL | Collective communication API and backends | +| InfiniCore | Pins the three repositories above as git submodules | +| InfiniLM | Builds the pinned stack and owns inference/runtime integration | +| InfiniTensor/ci | Provides the reusable InfiniLM build and test workflow | + +No runtime source, compatibility API, package metadata, tests, scripts, or CI +workflow belongs in InfiniCore after this change. + +## Target InfiniCore Tree + +InfiniCore will contain exactly these six tracked entries: + +```text +.gitmodules +LICENSE +README.md +submodules/InfiniCCL +submodules/InfiniOps +submodules/InfiniRT +``` + +The three submodule gitlinks are the version contract. `.gitmodules` will only +define each path and URL; `branch` hints will be removed because they do not +control the pinned revision. The README will describe InfiniCore as a manifest, +show recursive checkout, and point integration users to InfiniLM. It will not +contain build instructions of its own. + +The existing issue templates, workflow, `.gitignore`, `DEV.md`, `pyproject.toml`, +scripts, and tests will be deleted from InfiniCore. Build integration content is +migrated as described below; repository-specific legacy metadata is not copied. + +## InfiniLM Build Interface + +InfiniLM will own these files: + +```text +scripts/build_infini_stack.py +scripts/configs/infiniops_ops.txt +test/scripts/test_build_infini_stack.py +``` + +`scripts/build_infini_stack.py` will preserve the existing focused NVIDIA build +behavior and add a required `--infinicore-root PATH` argument. The argument +points to a separately checked-out InfiniCore manifest. The script will: + +1. Resolve the InfiniCore checkout and verify its three submodules are + initialized at the gitlink revisions recorded in InfiniCore `HEAD`. +2. Read the operator allowlist from the InfiniLM-owned configuration file. +3. Build and install InfiniRT, InfiniOps, and InfiniCCL into one isolated prefix. +4. Keep source checkouts immutable and place all generated state below the + selected build root. +5. Write a manifest containing the InfiniLM revision, InfiniCore revision, all + three submodule revisions, build options, install prefix, and operator list. + +The existing `--build-root`, `--build-type`, `--cuda-arch`, `--jobs`, `--test`, +and `--dry-run` behavior remains. Relative build roots are resolved against the +InfiniLM checkout. `--test` continues to enable InfiniRT tests and the two-GPU +InfiniCCL AllReduce smoke test. + +InfiniLM's README and contributing guidance will identify this script as the +supported way to build the pinned native stack. InfiniLM's own `third_party` +submodules remain the source of its JSON and spdlog dependencies. + +## Shared CI Contract + +The reusable workflow currently consumed from +`InfiniTensor/ci/.github/workflows/infinilm-ci.yml@infiniCore_ci` must change in +the same integration series. Work will branch from the currently referenced +`infiniCore_ci` integration line and follow that repository's contribution +rules. + +The workflow will: + +1. Check out InfiniLM with its own submodules, including JSON and spdlog. +2. Check out the selected InfiniCore revision recursively to obtain the three + pinned stack repositories. +3. Invoke InfiniLM's `scripts/build_infini_stack.py --infinicore-root ...`. +4. Build and test InfiniLM against the resulting install prefix. + +The workflow will stop reading `third_party` content from InfiniCore, stop +building InfiniCore with xmake, and stop installing InfiniCore as a Python +package. Existing branch/ref inputs may be retained when they still express the +same checkout selection; unnecessary interface renames are avoided. + +During cross-repository validation, the InfiniLM caller will temporarily pin the +feature revision of the shared workflow. It will return to the canonical shared +CI revision after that workflow change is integrated. + +## InfiniCCL Pull Request + +The InfiniCCL change remains a standalone header-correctness fix: + +- include `` in `src/device.h`; +- use `std::uint8_t` as the `MemorySpace` underlying type. + +The branch will be rebased onto the current InfiniCCL `master`, formatted and +checked according to that repository's `CONTRIBUTING.md`, and submitted as a +ready pull request using its current pull request template. Validation evidence +will include the NVIDIA/NCCL examples and an InfiniLM Qwen3 tensor-parallel-size +2 run. The pull request will not claim to add collectives or broaden backend +support. + +For integration validation, InfiniCore may temporarily pin the rebased pull +request head. The InfiniCore cleanup is not merge-ready until the InfiniCCL pull +request is merged and the gitlink points to its durable commit on `master`. + +## Change Sequence And Commits + +Each repository keeps its own conventions and focused commits. + +1. Commit this design in InfiniLM. +2. Prepare the shared CI branch and confirm its repository-specific rules. +3. Move the build script, operator configuration, tests, and documentation into + InfiniLM in focused migration and documentation commits. +4. Update the reusable shared CI workflow in a separate CI repository commit. +5. Reduce InfiniCore to the six-entry manifest in focused migration, metadata, + and cleanup commits as needed to keep each commit reviewable. +6. Rebase, verify, push, and open the standalone InfiniCCL pull request. +7. Validate clean coordinated checkouts on NVIDIA before presenting any branch + as merge-ready. + +Cross-repository branches are allowed to reference one another during +validation. Final merge commits must use durable branch or merged revisions, +not an unreviewed local checkout. + +## Verification + +Local and structural checks: + +- run each repository's formatter and static checks before every commit; +- run the migrated Python unit tests, including gitlink mismatch, missing + submodule, invalid operator list, argument validation, manifest, and dry-run + command coverage; +- verify `git ls-files` in InfiniCore produces exactly the six target entries; +- verify no tracked InfiniCore file contains legacy runtime identifiers; +- validate workflow syntax and run `actionlint` when available; +- inspect each repository diff against its intended base before push. + +NVIDIA acceptance checks in `accelerator-dev/nvidia:latest`: + +- start from clean coordinated checkouts with recursively initialized + submodules; +- run the InfiniLM-owned stack builder with tests enabled; +- run InfiniRT lifecycle and graph smoke coverage; +- confirm all 23 selected InfiniOps adapters build and link; +- run the InfiniCCL NVIDIA/NCCL AllReduce example; +- build and install the InfiniLM wheel against the generated prefix; +- run the existing Qwen3 tensor-parallel-size 1 inference command; +- run the same supported Qwen3 path with tensor parallel size 2; +- repeat the final inference smoke in a fresh container to prove it does not + depend on untracked host artifacts. + +The pull request descriptions will distinguish these verified paths from the +explicitly excluded support areas. + +## Risks And Recovery + +- Removing InfiniCore's build and package surfaces is intentionally breaking for + old consumers. The supported recovery is migration to InfiniLM's build script, + not compatibility shims in InfiniCore. +- Shared CI and product repositories cannot merge atomically. Feature refs will + be used for validation, and merge order will keep every canonical ref usable. +- A rebased or merged InfiniCCL pull request changes its commit identity. The + InfiniCore gitlink and validation manifest must be refreshed together. +- NVIDIA success does not authorize enabling other model or accelerator paths. + Existing feature gates remain until those paths receive their own validation. +- If coordinated NVIDIA validation regresses, retain the prior pushed branches + and fix forward in the owning repository; do not restore legacy runtime code + to InfiniCore. From 11b7f26873d464f1b3683adee4f875cf87d1d2f0 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 16 Jul 2026 13:41:31 +0800 Subject: [PATCH 3/6] refactor: own Infini stack integration build --- scripts/build_infini_stack.py | 429 ++++++++++++++++++++++ scripts/configs/infiniops_ops.txt | 23 ++ test/scripts/test_build_infini_stack.py | 465 ++++++++++++++++++++++++ 3 files changed, 917 insertions(+) create mode 100644 scripts/build_infini_stack.py create mode 100644 scripts/configs/infiniops_ops.txt create mode 100644 test/scripts/test_build_infini_stack.py diff --git a/scripts/build_infini_stack.py b/scripts/build_infini_stack.py new file mode 100644 index 000000000..0e6a20da7 --- /dev/null +++ b/scripts/build_infini_stack.py @@ -0,0 +1,429 @@ +import argparse +import json +import os +import re +import shlex +import subprocess +import sys +from pathlib import Path +from typing import Dict, List, Mapping, Optional, Sequence + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +OPERATOR_SET = PROJECT_ROOT / "scripts/configs/infiniops_ops.txt" +SUBMODULES = { + "InfiniRT": Path("submodules/InfiniRT"), + "InfiniOps": Path("submodules/InfiniOps"), + "InfiniCCL": Path("submodules/InfiniCCL"), +} + + +def read_operator_set(path: Path = OPERATOR_SET) -> List[str]: + operators = [line.strip() for line in path.read_text(encoding="utf-8").splitlines()] + operators = [operator for operator in operators if operator] + if not operators: + raise ValueError(f"Operator set is empty: {path}") + if len(operators) != len(set(operators)): + raise ValueError(f"Operator set contains duplicates: {path}") + if operators != sorted(operators): + raise ValueError(f"Operator set must be sorted: {path}") + return operators + + +def parse_gitlink(output: str, relative_path: Path) -> str: + fields = output.strip().split() + if len(fields) != 4 or fields[0] != "160000" or fields[1] != "commit": + raise RuntimeError(f"Not a gitlink in HEAD: {relative_path}") + return fields[2] + + +def capture(command: Sequence[str], cwd: Path) -> str: + process = subprocess.run( + command, + cwd=cwd, + capture_output=True, + text=True, + encoding="utf-8", + ) + if process.returncode != 0: + detail = process.stderr.strip() or process.stdout.strip() + raise RuntimeError(f"Command failed: {shlex.join(command)}\n{detail}") + return process.stdout.strip() + + +def git_capture(arguments: Sequence[str], cwd: Path) -> str: + return capture( + ["git", "-c", f"safe.directory={cwd}", *arguments], + cwd, + ) + + +def validate_submodule(project_root: Path, relative_path: Path) -> str: + submodule_root = project_root / relative_path + if not (submodule_root / "CMakeLists.txt").is_file(): + raise RuntimeError( + f"Submodule is not initialized: {relative_path}. " + "Run 'git submodule update --init --recursive'." + ) + + expected = parse_gitlink( + git_capture(["ls-tree", "HEAD", str(relative_path)], project_root), + relative_path, + ) + actual = git_capture(["rev-parse", "HEAD"], submodule_root) + if actual != expected: + raise RuntimeError( + f"Submodule revision mismatch for {relative_path}: " + f"expected {expected}, found {actual}" + ) + status = git_capture(["status", "--porcelain"], submodule_root) + if status: + raise RuntimeError(f"Submodule worktree is dirty: {relative_path}") + return actual + + +def run( + command: Sequence[str], + cwd: Path, + env: Optional[Mapping[str, str]] = None, + dry_run: bool = False, +) -> None: + print(f"+ {shlex.join(command)}", flush=True) + if not dry_run: + subprocess.run(command, cwd=cwd, env=env, check=True) + + +def cmake_cuda_architectures(cuda_arch: str) -> str: + return ";".join(arch[3:] for arch in cuda_arch.split(",")) + + +def build_infinirt_commands( + source: Path, + build: Path, + prefix: Path, + build_type: str, + jobs: int, + cuda_arch: Optional[str], + test: bool, +) -> List[List[str]]: + configure = [ + "cmake", + "-S", + str(source), + "-B", + str(build), + "-DWITH_CPU=ON", + "-DWITH_NVIDIA=ON", + f"-DINFINI_RT_BUILD_TESTING={'ON' if test else 'OFF'}", + f"-DCMAKE_BUILD_TYPE={build_type}", + f"-DCMAKE_INSTALL_PREFIX={prefix}", + ] + if cuda_arch: + configure.append( + f"-DCMAKE_CUDA_ARCHITECTURES={cmake_cuda_architectures(cuda_arch)}" + ) + + commands = [ + configure, + ["cmake", "--build", str(build), "--parallel", str(jobs)], + ] + if test: + commands.append( + [ + "ctest", + "--test-dir", + str(build), + "--output-on-failure", + "--parallel", + str(jobs), + ] + ) + commands.append(["cmake", "--install", str(build)]) + return commands + + +def build_infiniops_commands( + source: Path, + build: Path, + prefix: Path, + build_type: str, + jobs: int, + cuda_arch: Optional[str], + operators: Sequence[str], +) -> List[List[str]]: + configure = [ + "cmake", + "-S", + str(source), + "-B", + str(build), + "-DWITH_CPU=ON", + "-DWITH_NVIDIA=ON", + "-DAUTO_DETECT_DEVICES=OFF", + "-DAUTO_DETECT_BACKENDS=OFF", + "-DGENERATE_PYTHON_BINDINGS=OFF", + f"-DINFINI_RT_ROOT={prefix}", + f"-DINFINI_OPS_OPS={','.join(operators)}", + f"-DCMAKE_BUILD_TYPE={build_type}", + f"-DCMAKE_INSTALL_PREFIX={prefix}", + ] + if cuda_arch: + configure.append( + f"-DCMAKE_CUDA_ARCHITECTURES={cmake_cuda_architectures(cuda_arch)}" + ) + + return [ + configure, + [ + "cmake", + "--build", + str(build), + "--target", + "infiniops", + "--parallel", + str(jobs), + ], + ["cmake", "--install", str(build)], + ] + + +def build_infiniccl_commands( + source: Path, + build: Path, + prefix: Path, + build_type: str, + jobs: int, + cuda_arch: Optional[str], + test: bool, +) -> List[List[str]]: + configure = [ + "cmake", + "-S", + str(source), + "-B", + str(build), + "-DWITH_NVIDIA=ON", + "-DWITH_NCCL=ON", + "-DWITH_OMPI=OFF", + "-DWITH_MPICH=OFF", + "-DAUTO_DETECT_DEVICES=OFF", + "-DAUTO_DETECT_BACKENDS=OFF", + f"-DBUILD_EXAMPLES={'ON' if test else 'OFF'}", + f"-DCMAKE_BUILD_TYPE={build_type}", + f"-DCMAKE_INSTALL_PREFIX={prefix}", + ] + if cuda_arch: + configure.append( + f"-DCMAKE_CUDA_ARCHITECTURES={cmake_cuda_architectures(cuda_arch)}" + ) + + commands = [ + configure, + ["cmake", "--build", str(build), "--parallel", str(jobs)], + ["cmake", "--install", str(build)], + ] + if test: + commands.append( + [ + str(build / "examples/ccl/all_reduce"), + "-g", + "2", + "-w", + "1", + "-p", + "1", + "-n", + "1048576", + ] + ) + return commands + + +def integration_environment(prefix: Path) -> Dict[str, str]: + env = os.environ.copy() + env["LD_LIBRARY_PATH"] = os.pathsep.join( + filter(None, [str(prefix / "lib"), env.get("LD_LIBRARY_PATH", "")]) + ) + return env + + +def manifest_staging_path(path: Path) -> Path: + return path.with_name(f"{path.name}.tmp") + + +def write_manifest( + path: Path, + infinilm_revision: str, + infinicore_revision: str, + revisions: Mapping[str, str], + operators: Sequence[str], + prefix: Path, + build_type: str, + cuda_arch: Optional[str], + jobs: int, + test: bool, +) -> None: + data = { + "backend": "nvidia", + "build_type": build_type, + "cuda_arch": cuda_arch, + "infinicore": infinicore_revision, + "infinilm": infinilm_revision, + "install_prefix": str(prefix), + "jobs": jobs, + "operators": list(operators), + "submodules": dict(revisions), + "test": test, + } + path.parent.mkdir(parents=True, exist_ok=True) + staging_path = manifest_staging_path(path) + try: + staging_path.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8") + staging_path.replace(path) + finally: + staging_path.unlink(missing_ok=True) + + +def parse_cuda_arch(value: str) -> str: + architectures = value.split(",") + if not all( + architecture and re.fullmatch(r"sm_\d{2}a?", architecture) + for architecture in architectures + ): + raise argparse.ArgumentTypeError( + "--cuda-arch must be a comma-separated list like sm_80,sm_86,sm_90a" + ) + if len(architectures) != len(set(architectures)): + raise argparse.ArgumentTypeError("--cuda-arch contains duplicates") + return value + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Build the pinned InfiniRT, InfiniOps, and InfiniCCL NVIDIA stack." + ) + parser.add_argument( + "--infinicore-root", + type=Path, + required=True, + help="InfiniCore checkout containing the pinned stack submodules.", + ) + parser.add_argument( + "--build-root", + type=Path, + default=Path("build/integration/nvidia"), + help="Isolated build and install directory (default: %(default)s).", + ) + parser.add_argument("--build-type", choices=("Debug", "Release"), default="Release") + parser.add_argument( + "--cuda-arch", + type=parse_cuda_arch, + help=( + "Optional comma-separated CUDA architectures in SM notation, " + "for example sm_80,sm_86,sm_90a." + ), + ) + parser.add_argument("--jobs", type=int, default=os.cpu_count() or 1) + parser.add_argument( + "--test", + action="store_true", + help="Run InfiniRT tests and a two-GPU InfiniCCL AllReduce smoke test.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print commands without executing them.", + ) + args = parser.parse_args(argv) + if args.jobs < 1: + parser.error("--jobs must be positive") + return args + + +def resolve_from_project(path: Path) -> Path: + if not path.is_absolute(): + path = PROJECT_ROOT / path + return path.resolve() + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + if sys.platform != "linux" and not args.dry_run: + raise RuntimeError("The integration build currently supports Linux only") + + infinicore_root = resolve_from_project(args.infinicore_root) + build_root = resolve_from_project(args.build_root) + prefix = build_root / "prefix" + manifest_path = build_root / "manifest.json" + operators = read_operator_set() + revisions = { + name: validate_submodule(infinicore_root, relative_path) + for name, relative_path in SUBMODULES.items() + } + infinilm_revision = git_capture(["rev-parse", "HEAD"], PROJECT_ROOT) + infinicore_revision = git_capture(["rev-parse", "HEAD"], infinicore_root) + build_env = os.environ.copy() + + print(f"InfiniLM: {infinilm_revision}") + print(f"InfiniCore: {infinicore_revision}") + print(f"InfiniRT: {revisions['InfiniRT']}") + print(f"InfiniOps: {revisions['InfiniOps']}") + print(f"InfiniCCL: {revisions['InfiniCCL']}") + print(f"Operators ({len(operators)}): {','.join(operators)}") + + if not args.dry_run: + manifest_path.unlink(missing_ok=True) + manifest_staging_path(manifest_path).unlink(missing_ok=True) + + for command in build_infinirt_commands( + infinicore_root / SUBMODULES["InfiniRT"], + build_root / "infinirt", + prefix, + args.build_type, + args.jobs, + args.cuda_arch, + args.test, + ): + run(command, PROJECT_ROOT, build_env, args.dry_run) + + integration_env = integration_environment(prefix) + + for command in build_infiniops_commands( + infinicore_root / SUBMODULES["InfiniOps"], + build_root / "infiniops", + prefix, + args.build_type, + args.jobs, + args.cuda_arch, + operators, + ): + run(command, PROJECT_ROOT, integration_env, args.dry_run) + + for command in build_infiniccl_commands( + infinicore_root / SUBMODULES["InfiniCCL"], + build_root / "infiniccl", + prefix, + args.build_type, + args.jobs, + args.cuda_arch, + args.test, + ): + run(command, PROJECT_ROOT, integration_env, args.dry_run) + + if not args.dry_run: + write_manifest( + manifest_path, + infinilm_revision, + infinicore_revision, + revisions, + operators, + prefix, + args.build_type, + args.cuda_arch, + args.jobs, + args.test, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/configs/infiniops_ops.txt b/scripts/configs/infiniops_ops.txt new file mode 100644 index 000000000..c4810d8d8 --- /dev/null +++ b/scripts/configs/infiniops_ops.txt @@ -0,0 +1,23 @@ +add +add_rms_norm +causal_softmax +conv_infinilm +embedding +gelu_infinilm +gelutanh_infinilm +gemm +kv_caching_infinilm +paged_attention_infinilm +paged_attention_prefill_infinilm +paged_caching_infinilm +random_sample_infinilm +rearrange_infinilm +relu_infinilm +rms_norm +rotary_embedding_infinilm +sigmoid_infinilm +silu +silu_and_mul_infinilm +softmax_infinilm +swiglu +topksoftmax_infinilm diff --git a/test/scripts/test_build_infini_stack.py b/test/scripts/test_build_infini_stack.py new file mode 100644 index 000000000..450fafb19 --- /dev/null +++ b/test/scripts/test_build_infini_stack.py @@ -0,0 +1,465 @@ +import importlib.util +import json +import tempfile +import unittest +from contextlib import ExitStack +from pathlib import Path +from unittest import mock + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +MODULE_PATH = PROJECT_ROOT / "scripts/build_infini_stack.py" +SPEC = importlib.util.spec_from_file_location("build_infini_stack", MODULE_PATH) +build_infini_stack = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(build_infini_stack) + + +class BuildInfiniStackTest(unittest.TestCase): + def _patch_main_dependencies(self, stack, args): + stack.enter_context( + mock.patch.object(build_infini_stack.sys, "platform", "linux") + ) + stack.enter_context( + mock.patch.object(build_infini_stack, "parse_args", return_value=args) + ) + stack.enter_context( + mock.patch.object( + build_infini_stack, + "validate_submodule", + return_value="submodule-sha", + ) + ) + stack.enter_context( + mock.patch.object( + build_infini_stack, + "git_capture", + side_effect=["infinilm-sha", "core-sha"], + ) + ) + for helper, command in ( + ("build_infinirt_commands", "build-infinirt"), + ("build_infiniops_commands", "build-ops"), + ("build_infiniccl_commands", "build-ccl"), + ): + stack.enter_context( + mock.patch.object( + build_infini_stack, + helper, + return_value=[[command]], + ) + ) + run = stack.enter_context(mock.patch.object(build_infini_stack, "run")) + write_manifest = stack.enter_context( + mock.patch.object(build_infini_stack, "write_manifest") + ) + return run, write_manifest + + def test_operator_set_is_stable(self): + operators = build_infini_stack.read_operator_set() + + self.assertEqual(len(operators), 23) + self.assertEqual(operators, sorted(set(operators))) + self.assertIn("paged_attention_infinilm", operators) + self.assertIn("rotary_embedding_infinilm", operators) + + def test_invalid_operator_set_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "ops.txt" + path.write_text("rms_norm\nadd\nadd\n", encoding="utf-8") + + with self.assertRaisesRegex(ValueError, "duplicates"): + build_infini_stack.read_operator_set(path) + + def test_parse_gitlink(self): + output = "160000 commit abcdef1234567890\tsubmodules/InfiniRT\n" + + self.assertEqual( + build_infini_stack.parse_gitlink(output, Path("submodules/InfiniRT")), + "abcdef1234567890", + ) + + def test_uninitialized_submodule_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + core_root = Path(directory) + + with self.assertRaisesRegex(RuntimeError, "Submodule is not initialized"): + build_infini_stack.validate_submodule( + core_root, Path("submodules/InfiniRT") + ) + + def test_gitlink_mismatch_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + core_root = Path(directory) + relative_path = Path("submodules/InfiniRT") + source_root = core_root / relative_path + source_root.mkdir(parents=True) + (source_root / "CMakeLists.txt").touch() + with mock.patch.object( + build_infini_stack, + "git_capture", + side_effect=[ + "160000 commit expected-sha\tsubmodules/InfiniRT", + "actual-sha", + ], + ): + with self.assertRaisesRegex(RuntimeError, "revision mismatch"): + build_infini_stack.validate_submodule(core_root, relative_path) + + def test_dirty_submodule_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + core_root = Path(directory) + relative_path = Path("submodules/InfiniRT") + source_root = core_root / relative_path + source_root.mkdir(parents=True) + (source_root / "CMakeLists.txt").touch() + with mock.patch.object( + build_infini_stack, + "git_capture", + side_effect=[ + "160000 commit expected-sha\tsubmodules/InfiniRT", + "expected-sha", + " M CMakeLists.txt", + ], + ) as git_capture: + with self.assertRaisesRegex(RuntimeError, "worktree is dirty") as error: + build_infini_stack.validate_submodule(core_root, relative_path) + + self.assertIn(str(relative_path), str(error.exception)) + self.assertEqual( + git_capture.call_args_list[-1].args, + ( + ["status", "--porcelain"], + source_root, + ), + ) + + def test_untracked_submodule_source_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + core_root = Path(directory) + relative_path = Path("submodules/InfiniRT") + source_root = core_root / relative_path + source_root.mkdir(parents=True) + (source_root / "CMakeLists.txt").touch() + + def git_capture(arguments, _): + if arguments[0] == "ls-tree": + return "160000 commit expected-sha\tsubmodules/InfiniRT" + if arguments[0] == "rev-parse": + return "expected-sha" + if arguments == ["status", "--porcelain"]: + return "?? src/extra.cc" + return "" + + with mock.patch.object( + build_infini_stack, + "git_capture", + side_effect=git_capture, + ): + with self.assertRaisesRegex(RuntimeError, "worktree is dirty"): + build_infini_stack.validate_submodule(core_root, relative_path) + + def test_capture_includes_command_error(self): + with self.assertRaisesRegex(RuntimeError, "expected failure"): + build_infini_stack.capture( + [ + build_infini_stack.sys.executable, + "-c", + "import sys; print('expected failure', file=sys.stderr); sys.exit(1)", + ], + PROJECT_ROOT, + ) + + def test_git_capture_scopes_safe_directory(self): + with mock.patch.object( + build_infini_stack, "capture", return_value="revision" + ) as capture: + revision = build_infini_stack.git_capture( + ["rev-parse", "HEAD"], PROJECT_ROOT + ) + + self.assertEqual(revision, "revision") + command = capture.call_args.args[0] + self.assertEqual(command[:2], ["git", "-c"]) + self.assertEqual(command[2], f"safe.directory={PROJECT_ROOT}") + + def test_infinicore_root_is_required(self): + with ( + mock.patch.object(build_infini_stack.sys, "stderr"), + self.assertRaises(SystemExit), + ): + build_infini_stack.parse_args([]) + + def test_cuda_arch_list_is_converted_for_cmake(self): + commands = build_infini_stack.build_infinirt_commands( + Path("rt"), + Path("build/rt"), + Path("build/prefix"), + "Release", + 8, + "sm_80,sm_86,sm_90a", + False, + ) + + self.assertIn("-DCMAKE_CUDA_ARCHITECTURES=80;86;90a", commands[0]) + + def test_invalid_cuda_arch_is_rejected_by_argparse(self): + for cuda_arch in ("", "sm80", "sm_80,", "sm_80,,sm_90", "sm_80,sm_80"): + with self.subTest(cuda_arch=cuda_arch): + with ( + mock.patch.object(build_infini_stack.sys, "stderr"), + self.assertRaises(SystemExit), + ): + build_infini_stack.parse_args( + [ + "--infinicore-root", + "core", + "--cuda-arch", + cuda_arch, + ] + ) + + def test_infinirt_commands_share_one_prefix(self): + commands = build_infini_stack.build_infinirt_commands( + Path("rt"), + Path("build/rt"), + Path("build/prefix"), + "Release", + 8, + "sm_80", + True, + ) + + self.assertIn(f"-DCMAKE_INSTALL_PREFIX={Path('build/prefix')}", commands[0]) + self.assertIn("-DCMAKE_CUDA_ARCHITECTURES=80", commands[0]) + self.assertEqual(commands[2][0], "ctest") + self.assertEqual(commands[-1], ["cmake", "--install", str(Path("build/rt"))]) + + def test_infiniops_commands_use_pinned_runtime_and_operator_set(self): + commands = build_infini_stack.build_infiniops_commands( + Path("ops"), + Path("build/ops"), + Path("build/prefix"), + "Release", + 8, + "sm_80", + ["add", "rms_norm"], + ) + + configure = commands[0] + self.assertIn(f"-DINFINI_RT_ROOT={Path('build/prefix')}", configure) + self.assertIn("-DINFINI_OPS_OPS=add,rms_norm", configure) + self.assertIn("-DCMAKE_CUDA_ARCHITECTURES=80", configure) + self.assertEqual(commands[1][3:5], ["--target", "infiniops"]) + self.assertEqual(commands[-1], ["cmake", "--install", str(Path("build/ops"))]) + + def test_infiniccl_commands_run_two_gpu_smoke_test(self): + commands = build_infini_stack.build_infiniccl_commands( + Path("ccl"), + Path("build/ccl"), + Path("build/prefix"), + "Release", + 8, + "sm_80", + True, + ) + + configure = commands[0] + self.assertIn("-DWITH_NVIDIA=ON", configure) + self.assertIn("-DWITH_NCCL=ON", configure) + self.assertIn("-DBUILD_EXAMPLES=ON", configure) + self.assertEqual( + commands[-1][0], str(Path("build/ccl/examples/ccl/all_reduce")) + ) + self.assertEqual(commands[-1][1:3], ["-g", "2"]) + + def test_main_uses_core_sources_infini_lm_cwd_and_one_prefix(self): + args = build_infini_stack.parse_args( + [ + "--infinicore-root", + "core", + "--build-root", + "relative-build", + "--dry-run", + "--jobs", + "1", + ] + ) + core_root = (PROJECT_ROOT / "core").resolve() + build_root = (PROJECT_ROOT / "relative-build").resolve() + prefix = build_root / "prefix" + with ExitStack() as stack: + stack.enter_context( + mock.patch.dict(build_infini_stack.os.environ, {}, clear=True) + ) + stack.enter_context( + mock.patch.object(build_infini_stack, "parse_args", return_value=args) + ) + validate_submodule = stack.enter_context( + mock.patch.object( + build_infini_stack, + "validate_submodule", + return_value="submodule-sha", + ) + ) + stack.enter_context( + mock.patch.object( + build_infini_stack, + "git_capture", + side_effect=["infinilm-sha", "core-sha"], + ) + ) + build_infinirt = stack.enter_context( + mock.patch.object( + build_infini_stack, + "build_infinirt_commands", + return_value=[["build-infinirt"]], + ) + ) + build_infiniops = stack.enter_context( + mock.patch.object( + build_infini_stack, + "build_infiniops_commands", + return_value=[["build-ops"]], + ) + ) + build_infiniccl = stack.enter_context( + mock.patch.object( + build_infini_stack, + "build_infiniccl_commands", + return_value=[["build-ccl"]], + ) + ) + stack.enter_context(mock.patch.object(build_infini_stack, "write_manifest")) + run = stack.enter_context(mock.patch.object(build_infini_stack, "run")) + build_infini_stack.main([]) + + self.assertEqual( + validate_submodule.call_args_list[0].args, + ( + core_root, + Path("submodules/InfiniRT"), + ), + ) + self.assertEqual( + build_infinirt.call_args.args[:3], + ( + core_root / "submodules/InfiniRT", + build_root / "infinirt", + prefix, + ), + ) + self.assertEqual( + build_infiniops.call_args.args[:3], + ( + core_root / "submodules/InfiniOps", + build_root / "infiniops", + prefix, + ), + ) + self.assertEqual( + build_infiniccl.call_args.args[:3], + ( + core_root / "submodules/InfiniCCL", + build_root / "infiniccl", + prefix, + ), + ) + self.assertTrue( + all(call.args[1] == PROJECT_ROOT for call in run.call_args_list) + ) + infinirt_env = run.call_args_list[0].args[2] + infiniops_env = run.call_args_list[1].args[2] + infiniccl_env = run.call_args_list[2].args[2] + self.assertNotIn("LD_LIBRARY_PATH", infinirt_env) + self.assertTrue(infiniops_env["LD_LIBRARY_PATH"].startswith(str(prefix))) + self.assertEqual(infiniops_env, infiniccl_env) + + def test_manifest_is_written_after_all_commands_succeed(self): + args = build_infini_stack.parse_args( + ["--infinicore-root", "core", "--jobs", "1"] + ) + events = [] + with ExitStack() as stack: + run, write_manifest = self._patch_main_dependencies(stack, args) + run.side_effect = lambda command, *_: events.append(command[0]) + write_manifest.side_effect = lambda *_: events.append("write-manifest") + build_infini_stack.main([]) + + self.assertEqual( + events, + ["build-infinirt", "build-ops", "build-ccl", "write-manifest"], + ) + + def test_failed_command_does_not_write_manifest(self): + with tempfile.TemporaryDirectory() as directory: + build_root = Path(directory) + manifest_path = build_root / "manifest.json" + staging_path = build_root / "manifest.json.tmp" + manifest_path.write_text("old manifest\n", encoding="utf-8") + staging_path.write_text("partial manifest\n", encoding="utf-8") + args = build_infini_stack.parse_args( + [ + "--infinicore-root", + "core", + "--build-root", + str(build_root), + "--jobs", + "1", + ] + ) + + def fail_run(*_): + self.assertFalse(manifest_path.exists()) + self.assertFalse(staging_path.exists()) + raise RuntimeError("build failed") + + with ExitStack() as stack: + run, write_manifest = self._patch_main_dependencies(stack, args) + run.side_effect = fail_run + with self.assertRaisesRegex(RuntimeError, "build failed"): + build_infini_stack.main([]) + + write_manifest.assert_not_called() + self.assertFalse(manifest_path.exists()) + self.assertFalse(staging_path.exists()) + + def test_manifest_records_revisions_operators_and_options(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "manifest.json" + staging_path = path.with_name(f"{path.name}.tmp") + path.write_text('{"old": true}\n', encoding="utf-8") + staging_path.write_text("partial manifest\n", encoding="utf-8") + build_infini_stack.write_manifest( + path, + "infinilm-sha", + "core-sha", + { + "InfiniRT": "rt-sha", + "InfiniOps": "ops-sha", + "InfiniCCL": "ccl-sha", + }, + ["add", "rms_norm"], + Path("prefix"), + "Release", + "sm_80,sm_90a", + 8, + True, + ) + + manifest = json.loads(path.read_text(encoding="utf-8")) + self.assertEqual(manifest["backend"], "nvidia") + self.assertEqual(manifest["infinilm"], "infinilm-sha") + self.assertEqual(manifest["infinicore"], "core-sha") + self.assertEqual(manifest["submodules"]["InfiniRT"], "rt-sha") + self.assertEqual(manifest["submodules"]["InfiniCCL"], "ccl-sha") + self.assertEqual(manifest["operators"], ["add", "rms_norm"]) + self.assertEqual(manifest["install_prefix"], str(Path("prefix"))) + self.assertEqual(manifest["build_type"], "Release") + self.assertEqual(manifest["cuda_arch"], "sm_80,sm_90a") + self.assertEqual(manifest["jobs"], 8) + self.assertTrue(manifest["test"]) + self.assertNotIn("old", manifest) + self.assertFalse(staging_path.exists()) + + +if __name__ == "__main__": + unittest.main() From f6dde8add1c4ecb14fecb832dc6b2ce33c87619a Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 16 Jul 2026 14:31:43 +0800 Subject: [PATCH 4/6] docs: document the InfiniLM stack build --- CONTRIBUTING.md | 24 +++++++++++++++++++++++- README.md | 28 +++++++++++++++------------- 2 files changed, 38 insertions(+), 14 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 111716098..f6486bec0 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -58,7 +58,29 @@ Existing branch names may use the legacy format `issue/###`, followed by a suffi # Development Guide -Refer to [ReadMe](README.md) and [Adapt New Models](MODELS.md) +Refer to [ReadMe](README.md) and [Adapt New Models](MODELS.md). + +Run the migrated stack-builder unit tests with: + +```shell +python -m unittest test/scripts/test_build_infini_stack.py -v +``` + +Check the Core-backed build commands without creating build output: + +```shell +python scripts/build_infini_stack.py --infinicore-root ../InfiniCore --dry-run --jobs 1 --cuda-arch sm_80 +``` + +Run the static migration contracts with: + +```shell +python -m unittest discover -s test/static -p "test_*.py" -v +``` + +Native stack changes must be validated in the owning component repository +first. After validation, update the corresponding submodule pin in InfiniCore +before validating InfiniLM. ## Troubleshooting diff --git a/README.md b/README.md index da0c1001d..41c357fe5 100644 --- a/README.md +++ b/README.md @@ -8,8 +8,9 @@ The lower-level runtime, operator, and collective APIs are provided by [InfiniCCL](https://github.com/InfiniTensor/InfiniCCL), respectively. [InfiniCore](https://github.com/InfiniTensor/InfiniCore) pins those three -projects and provides an integration build. It no longer provides a separate -runtime API or Python package. +projects as submodules. InfiniLM no longer consumes runtime or Python-package +artifacts built by that separate repository; it builds and packages its own +`infinicore` Python module. ## Build @@ -20,12 +21,16 @@ git clone --recurse-submodules https://github.com/InfiniTensor/InfiniCore.git git clone --recurse-submodules https://github.com/InfiniTensor/InfiniLM.git ``` -Build the NVIDIA dependency stack from InfiniCore. The default operator set is -the set required by InfiniLM: +From InfiniLM, build the NVIDIA dependency stack pinned by the InfiniCore +checkout. The default operator set is the set required by InfiniLM: ```shell -cd InfiniCore -python3 scripts/build_integration.py --cuda-arch sm_80 --jobs 16 --test +cd InfiniLM +python3 scripts/build_infini_stack.py \ + --infinicore-root ../InfiniCore \ + --cuda-arch sm_80 \ + --jobs 16 \ + --test export INFINI_ROOT="$PWD/build/integration/nvidia/prefix" export LD_LIBRARY_PATH="$INFINI_ROOT/lib:${LD_LIBRARY_PATH:-}" ``` @@ -33,17 +38,14 @@ export LD_LIBRARY_PATH="$INFINI_ROOT/lib:${LD_LIBRARY_PATH:-}" Then build and install InfiniLM: ```shell -cd ../InfiniLM python3 -m pip install . --no-build-isolation ``` -The native build currently recognizes `cpu`, `nvidia`, `cambricon`, `ascend`, -`metax`, `moore`, `iluvatar`, and `hygon` platforms. The migration in this -branch is validated on NVIDIA first. +Automated migration coverage is limited to NVIDIA, dense non-quantized Qwen3, +and the default static attention implementation. Other platforms and +configurations remain gated for later validation. -The modern operator closure currently supports `qwen3`. Other registered model -types are rejected before worker startup until their required InfiniOps adapters -are available. +The modern operator closure currently supports `qwen3` within that boundary. ## Inference From f9a27178761c28476bfe61fa3e678e8f7e48a5d9 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 16 Jul 2026 14:36:26 +0800 Subject: [PATCH 5/6] ci: validate the modern NVIDIA stack --- .github/ci_config_nvidia.yaml | 74 +++++++++++++++++++++++++++++++++++ .github/workflows/ci_test.yml | 6 +-- 2 files changed, 77 insertions(+), 3 deletions(-) create mode 100644 .github/ci_config_nvidia.yaml diff --git a/.github/ci_config_nvidia.yaml b/.github/ci_config_nvidia.yaml new file mode 100644 index 000000000..fe58ced9f --- /dev/null +++ b/.github/ci_config_nvidia.yaml @@ -0,0 +1,74 @@ +repo: + url: https://github.com/InfiniTensor/InfiniLM.git + branch: main + +github: + status_context_prefix: "ci/infinilm" + +platforms: + nvidia: + image: + dockerfile: images/nvidia/ + build_args: + BASE_IMAGE: nvcr.io/nvidia/pytorch:25.12-py3 + CUDA_ARCH: sm_80,sm_86,sm_89,sm_90 + APT_MIRROR: https://mirrors.tuna.tsinghua.edu.cn/ubuntu + PIP_INDEX_URL: https://pypi.org/simple + InfiniCore_BRANCH: __Branch_Name__ + docker_args: + - "--user=root" + - "--network=host" + - "--privileged" + - "--cap-add=ALL" + - "--pid=host" + - "--ipc=host" + - "--workdir=/workspace" + volumes: + - /data:/data + - /data-aisoft:/data-aisoft + - /data-aisoft/artifacts/CI_nvidia_test/__WORKSPACE__:/artifacts + setup: pip install .[dev] --no-build-isolation + jobs: + gpu_inferencetest: + type: inferencetest + resources: + ngpus: [1, 4] + gpu_style: nvidia + shm_size: 64g + timeout: 3600 + stages: + - name: test + run: python InfiniLM/examples/test_infer.py --device nvidia --model=/data-aisoft/mechdancer/models/9g_8b_thinking/ + gpu_benchtest: + type: benchtest + resources: + gpu_style: nvidia + shm_size: 64g + timeout: 3600 + env: + TEST_PARAM: ['default'] + stages: + - name: test + run: python InfiniLM/examples/bench.py --device nvidia --model=/data-aisoft/mechdancer/models/9g_8b_thinking/ --input-len=256,1024 --output-len=256,1024 --batch-size=8 + gpu_accuracytest: + type: accuracytest + resources: + gpu_style: nvidia + shm_size: 64g + timeout: 3600 + env: + TEST_PARAM: ['--bench mmlu'] + stages: + - name: test + run: python InfiniLM/test/bench/test_benchmark.py --device nvidia --model /data-aisoft/mechdancer/models/9g_8b_thinking/ --bench mmlu --backend cpp --max-new-tokens 5 --cache-dir /data-aisoft/pepe/datasets/ --split=val + gpu_servicetest: + type: servicetest + resources: + shm_size: 64g + env: + MODEL_LIST: 9g_8b_thinking + ENGINE: InfiniLM + TEST_PARAM: ['default'] + stages: + - name: test + run: python InfiniLM/scripts/test_perf.py --verbose diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml index 18c513dbe..2414ecd3f 100644 --- a/.github/workflows/ci_test.yml +++ b/.github/workflows/ci_test.yml @@ -43,9 +43,9 @@ jobs: ci: if: github.event_name == 'workflow_dispatch' needs: check-format - uses: InfiniTensor/ci/.github/workflows/infinilm-ci.yml@infiniCore_ci + uses: InfiniTensor/ci/.github/workflows/infinilm-ci.yml@refactor/adopt-modern-infini-stack with: - config_path: .github/ci_config.yaml - ci_ref: infiniCore_ci + config_path: .github/ci_config_nvidia.yaml + ci_ref: refactor/adopt-modern-infini-stack infinicore_branch: ${{ github.event.inputs.infinicore_branch || 'main' }} secrets: inherit From 009bdd28758ae5ef1ce437dc8890bcf470fff013 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Fri, 24 Jul 2026 13:42:40 +0800 Subject: [PATCH 6/6] fix: keep static graph cache metadata dynamic --- .../compiler/static_batching_compiler.cpp | 51 +++++++++- .../layers/attention/backends/static_attn.cpp | 40 ++++++++ .../layers/attention/backends/static_attn.hpp | 7 ++ .../test_infinicore_runtime_contracts.py | 95 +++++++++++++++++++ 4 files changed, 189 insertions(+), 4 deletions(-) diff --git a/csrc/engine/compiler/static_batching_compiler.cpp b/csrc/engine/compiler/static_batching_compiler.cpp index 8a2106780..0da076a0f 100644 --- a/csrc/engine/compiler/static_batching_compiler.cpp +++ b/csrc/engine/compiler/static_batching_compiler.cpp @@ -1,6 +1,35 @@ #include "static_batching_compiler.hpp" #include "../../cache/cache.hpp" #include "../../global_state/global_state.hpp" +#include "../../utils.hpp" + +#include + +namespace { +bool supports_static_graph_kv_cache(const infinicore::Tensor &kv_cache) { + if (kv_cache.empty() || kv_cache->ndim() != 5 || kv_cache->size(0) != 2 || !kv_cache->is_contiguous()) { + return false; + } + + const auto dtype = kv_cache->dtype(); + const auto head_dim = kv_cache->size(4); + return (dtype == infinicore::DataType::kFloat16 || dtype == infinicore::DataType::kBFloat16) + && (head_dim == 64 || head_dim == 128); +} + +bool supports_static_graph_attention() { + const auto &config = infinilm::global_state::get_infinilm_config(); + if (!config.model_config + || infinicore::context::getDevice().type() != infinicore::Device::Type::kNvidia + || config.model_config->get_kv_quant_scheme() != infinilm::quantization::KVQuantAlgo::NONE) { + return false; + } + + const auto &kv_cache_vec = infinilm::global_state::get_forward_context().kv_cache_vec; + return !kv_cache_vec.empty() + && std::all_of(kv_cache_vec.begin(), kv_cache_vec.end(), supports_static_graph_kv_cache); +} +} // namespace namespace infinilm::engine { StaticBatchingCompiler::StaticBatchingCompiler(const std::shared_ptr &model, RankBarrier *barrier) @@ -8,15 +37,29 @@ StaticBatchingCompiler::StaticBatchingCompiler(const std::shared_ptrget_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { size_t b = dynamic_cast(model_->get_cache_config())->max_batch_size(); InfinilmModel::Input input; input.input_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::kInt64, infinicore::context::getDevice()); input.position_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::kInt64, infinicore::context::getDevice()); - input.past_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::kInt64, infinicore::context::getDevice()); - input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::kInt64, infinicore::context::getDevice()); - std::vector total_sequence_lengths_vec(b, 1); - infinicore::context::memcpyH2D(input.total_sequence_lengths.value()->data(), total_sequence_lengths_vec.data(), b * sizeof(int64_t), false); + input.past_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::kInt32, infinicore::context::getDevice()); + input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::kInt32, infinicore::context::getDevice()); + set_zeros(input.input_ids.value()); + set_zeros(input.position_ids.value()); + set_zeros(input.past_sequence_lengths.value()); + std::vector total_sequence_lengths_vec(b, 1); + infinicore::context::memcpyH2D(input.total_sequence_lengths.value()->data(), total_sequence_lengths_vec.data(), b * sizeof(int32_t), false); + input.block_tables = infinicore::Tensor::empty({b, 1}, infinicore::DataType::kInt32, infinicore::context::getDevice()); + std::vector block_tables_vec(b); + for (size_t i = 0; i < b; ++i) { + block_tables_vec[i] = static_cast(i); + } + infinicore::context::memcpyH2D(input.block_tables.value()->data(), block_tables_vec.data(), b * sizeof(int32_t), false); // Attention reads attn_metadata from thread-local forward context. infinilm::global_state::get_forward_context().attn_metadata = { diff --git a/csrc/layers/attention/backends/static_attn.cpp b/csrc/layers/attention/backends/static_attn.cpp index e2d877950..5c1c3ec30 100644 --- a/csrc/layers/attention/backends/static_attn.cpp +++ b/csrc/layers/attention/backends/static_attn.cpp @@ -50,6 +50,11 @@ infinicore::Tensor StaticAttentionImpl::forward(const AttentionLayer &layer, auto past_sequence_lengths = attn_metadata.past_sequence_lengths; auto total_sequence_lengths = attn_metadata.total_sequence_lengths; + if (infinicore::context::isGraphRecording()) { + ASSERT(this->kv_quant_scheme_ == infinilm::quantization::KVQuantAlgo::NONE); + return forward_graph_(q_reshaped, k_permuted, v_permuted, kv_cache, attn_metadata); + } + // update static kv cache // k_total: [bs, n_kv_head, max_seq_len, head_dim] // v_total : [bs, n_kv_head, max_seq_len, head_dim] @@ -90,6 +95,41 @@ infinicore::Tensor StaticAttentionImpl::forward(const AttentionLayer &layer, ->view({batch_size, seq_len, num_heads_ * head_dim_}); // [bs, seq_len, n_q_head * head_dim] } +infinicore::Tensor StaticAttentionImpl::forward_graph_( + const infinicore::Tensor &query, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + infinicore::Tensor &kv_cache, + const infinilm::global_state::AttentionMetadata &attn_metadata) const { + ASSERT_EQ(query->size(2), 1); + ASSERT(attn_metadata.block_tables.has_value()); + ASSERT(attn_metadata.past_sequence_lengths.has_value()); + ASSERT(attn_metadata.total_sequence_lengths.has_value()); + + auto k_cache = kv_cache->narrow({{0, 0, 1}})->squeeze(0); + auto v_cache = kv_cache->narrow({{0, 1, 1}})->squeeze(0); + infinicore::op::kv_caching_( + k_cache, + v_cache, + key, + value, + attn_metadata.past_sequence_lengths.value()); + + auto q_decode = query->contiguous()->view({query->size(0), query->size(1), query->size(3)}); + auto output = infinicore::Tensor::empty(q_decode->shape(), q_decode->dtype(), q_decode->device()); + infinicore::op::paged_attention_( + output, + q_decode, + k_cache, + v_cache, + attn_metadata.block_tables.value(), + attn_metadata.total_sequence_lengths.value(), + std::nullopt, + scale_); + + return output->view({query->size(0), 1, num_heads_ * head_dim_}); +} + std::tuple StaticAttentionImpl::do_kv_cache_update(const AttentionLayer &layer, const infinicore::Tensor key, const infinicore::Tensor value, diff --git a/csrc/layers/attention/backends/static_attn.hpp b/csrc/layers/attention/backends/static_attn.hpp index 849d87928..cb9e50436 100644 --- a/csrc/layers/attention/backends/static_attn.hpp +++ b/csrc/layers/attention/backends/static_attn.hpp @@ -34,6 +34,13 @@ class StaticAttentionImpl { const infinicore::Tensor past_sequence_lengths) const; private: + infinicore::Tensor forward_graph_( + const infinicore::Tensor &query, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + infinicore::Tensor &kv_cache, + const infinilm::global_state::AttentionMetadata &attn_metadata) const; + size_t num_heads_; size_t head_size_; float scale_; diff --git a/test/static/test_infinicore_runtime_contracts.py b/test/static/test_infinicore_runtime_contracts.py index 8219dfddb..474b0c526 100644 --- a/test/static/test_infinicore_runtime_contracts.py +++ b/test/static/test_infinicore_runtime_contracts.py @@ -402,6 +402,101 @@ def test_graph_compilers_cancel_capture_on_exception(self) -> None: graph_manager.index("std::lock_guard"), graph_manager.index("graph.reset()") ) + def test_static_graph_input_dtypes_match_scheduler_inputs(self) -> None: + source = read_source("csrc/engine/compiler/static_batching_compiler.cpp") + compile_body = function_body(source, "void StaticBatchingCompiler::compile()") + + expected_dtypes = { + "input_ids": "kInt64", + "position_ids": "kInt64", + "past_sequence_lengths": "kInt32", + "total_sequence_lengths": "kInt32", + } + for field, dtype in expected_dtypes.items(): + with self.subTest(field=field): + self.assertRegex( + compile_body, + re.compile( + rf"input\.{field}\s*=\s*infinicore::Tensor::empty" + rf"\([^;]*DataType::{dtype}," + ), + ) + + self.assertIn("std::vector total_sequence_lengths_vec", compile_body) + self.assertIn("b * sizeof(int32_t)", compile_body) + + def test_static_graph_compile_inputs_are_deterministic(self) -> None: + source = read_source("csrc/engine/compiler/static_batching_compiler.cpp") + compile_body = function_body(source, "void StaticBatchingCompiler::compile()") + + for field in ("input_ids", "position_ids", "past_sequence_lengths"): + with self.subTest(field=field): + self.assertIn(f"set_zeros(input.{field}.value())", compile_body) + + self.assertRegex( + compile_body, + re.compile( + r"input\.block_tables\s*=\s*infinicore::Tensor::empty" + r"\(\{b, 1\}, infinicore::DataType::kInt32," + ), + ) + self.assertIn("block_tables_vec[i] = static_cast(i)", compile_body) + self.assertIn("input.block_tables.value()->data()", compile_body) + + def test_static_graph_keeps_dynamic_cache_metadata_on_device(self) -> None: + source = read_source("csrc/layers/attention/backends/static_attn.cpp") + forward = function_body( + source, + "infinicore::Tensor StaticAttentionImpl::forward(", + ) + + self.assertIn("context::isGraphRecording()", forward) + self.assertIn("return forward_graph_(", forward) + self.assertIn("StaticAttentionImpl::forward_graph_", source) + + graph_forward = function_body( + source, + "infinicore::Tensor StaticAttentionImpl::forward_graph_(", + ) + self.assertIn("infinicore::op::kv_caching_", graph_forward) + self.assertIn("infinicore::op::paged_attention_", graph_forward) + self.assertNotIn("Device::Type::kCpu", graph_forward) + + def test_static_graph_falls_back_for_unsupported_attention_configs(self) -> None: + source = read_source("csrc/engine/compiler/static_batching_compiler.cpp") + self.assertIn("bool supports_static_graph_kv_cache(", source) + self.assertIn("bool supports_static_graph_attention()", source) + cache_check = function_body( + source, "bool supports_static_graph_kv_cache(" + ) + capability = function_body(source, "bool supports_static_graph_attention()") + compile_body = function_body(source, "void StaticBatchingCompiler::compile()") + + self.assertIn("Device::Type::kNvidia", capability) + self.assertIn("get_forward_context().kv_cache_vec", capability) + self.assertIn("kv_cache_vec.empty()", capability) + self.assertIn("std::all_of(", capability) + self.assertNotIn("kv_cache_vec.front()", capability) + self.assertNotIn("kv_cache_vec[0]", capability) + self.assertIn("kv_cache.empty()", cache_check) + self.assertIn("kv_cache->ndim() != 5", cache_check) + self.assertIn("kv_cache->size(0) != 2", cache_check) + self.assertIn("kv_cache->dtype()", cache_check) + self.assertIn("kv_cache->size(4)", cache_check) + self.assertIn("DataType::kFloat16", cache_check) + self.assertIn("DataType::kBFloat16", cache_check) + self.assertIn("head_dim == 64", cache_check) + self.assertIn("head_dim == 128", cache_check) + self.assertIn("KVQuantAlgo::NONE", capability) + for forbidden in ("get_dtype()", "get_kv_cache_dtype()", "get_head_dim()"): + self.assertNotIn(forbidden, capability) + self.assertIn("compiled_map_.clear()", compile_body) + self.assertIn("if (!supports_static_graph_attention())", compile_body) + self.assertLess( + compile_body.index("if (!supports_static_graph_attention())"), + compile_body.index("GraphRecordingGuard recording"), + ) + def test_foreign_capture_rejects_operator_dispatch(self) -> None: manager_header = read_source("csrc/infinicore/src/graph/graph_manager.hpp") graph_header = read_source("csrc/infinicore/include/infinicore/graph/graph.hpp")