Pepe/glm5 2

docs/GLM-5.2-W4A8-InfiniCore算子接入进度.md

@@ -0,0 +1,80 @@
+# GLM-5.2-W4A8 InfiniCore 算子接入进度
+
+记录环境：TG150 .58 `/home/wuwei/InfiniCore`，容器 `pepe`，Iluvatar BI-V150 (`ivcore11`)。
+
+## 已接入/验证
+
+| Trace 算子类别 | 当前 InfiniCore 接入状态 | 验证 |
+|---|---|---|
+| `pyinfer::perf::dynamic_scaled_int8_quant` | 新增 `infinicore.dynamic_scaled_int8_quant`，通过 `vllm_iluvatar/_C` dlopen+dlsym 桥接 | F16/BF16 torch zero-copy 数值验证通过，scale/output 与参考一致 |
+| `pyinfer::perf::fused_add_rms_norm` | 新增 `infinicore.add_rms_norm_inplace`，通过 `vllm_iluvatar/_C` 桥接 | F16/BF16 数值验证通过 |
+| `silu_and_mul` | 新增/启用 InfiniOp Iluvatar 路径，复用 NVIDIA CUDA kernel 结构 | `python3 test/infiniop/silu_and_mul.py --iluvatar` 通过 |
+| `moe_sum` | InfiniOp NVIDIA backend 增加 Iluvatar dispatch/编译 guard | `deepseek_moe.py --iluvatar` 覆盖通过 |
+| `pyinfer::perf::topk_softmax` / `topk_sigmoid` | 新增 `infinicore.moe_topk_softmax_vllm` / `moe_topk_sigmoid_vllm`，通过 `vllm_iluvatar/_C` 桥接，三输出匹配 vLLM trace: weights、expert ids、source rows | F16/BF16/F32、renormalize true/false、correction_bias 验证通过 |
+| `pyinfer::perf::moe_grouped_topk` / `_moe_C.grouped_topk` | 新增 `infinicore.grouped_topk_vllm`，通过 `vllm_iluvatar/_C` 桥接；仅开放 correction-bias 路径，no-bias 在 Python/C++ 层显式拒绝 | 与 vLLM 实际生产入口 `torch.ops._moe_C.grouped_topk` 在 F16/BF16、softmax/sigmoid、renormalize true/false 下逐位一致；F16 也匹配 `vllm_iluvatar` torch reference；BF16 与 torch reference 存在偏差，记录为复用生产 perf 语义的风险 |
+| `pyinfer::cuinfer::scaled_mm_w4a8` / `_C.scaled_mm_w4a8` | 新增 `infinicore.scaled_mm_w4a8`，通过 `vllm_iluvatar/_C` 的 cuinfer 符号桥接；用于 dense W4A8 linear packed GEMM，当前不改权重布局 | 与 vLLM 实际生产入口 `torch.ops._C.scaled_mm_w4a8` 在 F16/BF16、有无 bias 下逐位一致 |
+| `pyinfer::cuinfer::w4a8_group_gemm` / `_C.w4a8_group_gemm` | 新增 `infinicore.w4a8_group_gemm_`，通过 `vllm_iluvatar/_C` 的 cuinfer 符号桥接；用于 MoE W4A8 group GEMM | 与 vLLM 实际生产入口 `torch.ops._C.w4a8_group_gemm` 在 F16/BF16、有无 bias、有无 sorted_token_ids、prefill/decode 标志下逐位一致；未声称匹配 torch fallback，因为 fallback 对构造样本与 production 不一致 |
+| `pyinfer::cuinfer::w8a8_group_gemm` / `_C.w8a8_group_gemm` | 新增 `infinicore.w8a8_group_gemm_`，通过 `vllm_iluvatar/_C` 的 cuinfer 符号桥接；仅开放 `trans_weight=True` 的 TN prefill 路径 | 与 vLLM 实际生产入口 `torch.ops._C.w8a8_group_gemm` 在 F16/BF16、有无 bias、有无 sorted_token_ids、多个 E/M/K/N shape 下逐位一致；`is_decode=True` 会触发 vLLM/CUINFER internal error，已在 Python/C++ 层显式 guard 禁用 |
+| `moe_topk_softmax` | InfiniOp NVIDIA backend 增加 Iluvatar dispatch/编译 guard | `topksoftmax.py --iluvatar`、`deepseek_moe.py --iluvatar` 通过 |
+| `moe_topk_sigmoid` | InfiniOp NVIDIA backend 增加 Iluvatar dispatch/编译 guard | `deepseek_moe.py --iluvatar` 覆盖通过 |
+| `rms_norm` / `rope` / `mrope` / `kv_caching` / `topksoftmax` | 仓库已有 InfiniCore/InfiniOp 接口和 Iluvatar 可运行路径 | 既有测试曾通过；本轮重跑了 `topksoftmax` |
+| `pyinfer::perf::rotary_embedding` | 暂不 bridge；继续使用已验证的 InfiniCore `rope/mrope` 路径 | 直接调用 vLLM perf 对 F16/BF16、2D/3D、Neox/GPT-J 均产生 NaN；已撤回 bridge，避免暴露错误接口 |
+
+## 已明确不接入的错误路径
+
+- `scaled_mm_int8` 的 NVIDIA path 曾尝试打开 Iluvatar dispatch，但首个 BF16 case 数值失败；已撤回，当前 `scaled_mm` 下没有残留 `ENABLE_ILUVATAR` / `INFINI_DEVICE_ILUVATAR` 修改。
+- GLM 的 W4A8/GEMM 热路径应继续优先复用 cuinfer/vLLM 已用 so，而不是暴露未验证的 InfiniOp int8 GEMM 路径。
+
+## 当前新增接口
+
+- C++ header: `include/infinicore/ops/dynamic_scaled_int8_quant.hpp`
+- Python: `infinicore.dynamic_scaled_int8_quant(input, input_scales, out=None)`
+  - `input`: F16/BF16 contiguous，last dim 为 hidden size
+  - `input_scales`: F32 contiguous，numel = `input.numel / input.shape[-1]`
+  - `out`: optional I8 contiguous，shape 与 input 相同
+- C++/Python: `add_rms_norm_inplace(input, residual, weight, epsilon=1e-5)`
+- C++/Python: `concat_mla_q(ql_nope, q_pe, out=None)`，当前限制 GLM MLA 维度 `512 + 64 -> 576`
+- C++/Python: `concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype="auto", scale=...)`
+- C++/Python: `concat_and_cache_mla_int8(kv_c_int8, kv_c_scale, k_pe_int8, k_pe_scale, kv_cache, kv_cache_scale, slot_mapping)`
+- C++/Python: `moe_topk_softmax_vllm(gating_output, topk, renormalize=False, correction_bias=None, out=None)`
+- C++/Python: `moe_topk_sigmoid_vllm(gating_output, topk, renormalize=False, correction_bias=None, out=None)`
+- C++/Python: `grouped_topk_vllm(scores, num_expert_group, topk_group, topk, renormalize, routed_scaling_factor=1.0, bias=..., scoring_func="softmax", out=None)`
+  - 当前必须传 `bias` / correction bias；no-bias 路径因 vLLM perf 与 reference 不一致而禁用。
+  - 对 torch tensor 做 zero-copy 验证时使用 `infinicore.tensor.from_torch(...)` 包装输入和 `out` buffer。
+- C++/Python: `scaled_mm_w4a8(a, b, a_scales, b_scales, bias=None, trans_weight=False, out=None)`
+  - `a`: int8 activation `(M,K)`；`b`: packed int4-as-int8 weight `(K,N/2)` when `trans_weight=False`；`a_scales`: F32 `(M,1)`；`b_scales`: F32 `(N,1)`；`out`: F16/BF16 `(M,N)`。
+  - 当前验证的是 dense W4A8 NN layout，与 vLLM `scaled_mm_w4a8_packed_fwd`/`torch.ops._C.scaled_mm_w4a8` 路径一致。
+- C++/Python: `w4a8_group_gemm_(out, input, weight, input_scale, weight_scale, tokens_per_experts, sorted_token_ids=None, bias=None, trans_weight=True, is_decode=False)`
+  - `input`: int8 `(M,K)`；`weight`: packed int4-as-int8 `(E,N,K/2)` for TN；`input_scale`: F32 `(M,1)`；`weight_scale`: F32 `(E,N,1)`；`tokens_per_experts`: int32 `(E,)`；`out`: F16/BF16 `(M,N)`。
+  - 与 vLLM production op 逐位一致；prefill 构造样本中 `tokens_per_experts` 使用 CPU int32，decode 构造样本中使用 GPU int32，均已覆盖。
+- C++/Python: `w8a8_group_gemm_(out, input, weight, input_scale, weight_scale, tokens_per_experts, sorted_token_ids=None, bias=None, trans_weight=True, is_decode=False)`
+  - `input`: int8 `(M,K)`；`weight`: int8 `(E,N,K)` for TN；`input_scale`: F32 `(M,1)`；`weight_scale`: F32 `(E,N,1)`；`tokens_per_experts`: int32 `(E,)`；`out`: F16/BF16 `(M,N)`。
+  - 当前只开放 `trans_weight=True` 且 `is_decode=False`。`trans_weight=False` 与 `is_decode=True` 均会在 Python/C++ 层报错；后者是因为 vLLM/CUINFER production kernel decode 分支在当前环境触发 `CUINFER_STATUS_INTERNAL_ERROR`。
+  - 与 vLLM production prefill op 逐位一致；验证覆盖 F16/BF16、有无 bias、有无 sorted_token_ids，以及 `(E,M,K,N)=(4,16,128,256),(4,64,256,512),(8,128,256,512)`。
+
+## 待接入/待确认
+
+| Trace 符号 | 建议路径 | 备注 |
+|---|---|---|
+| `concat_and_cache_mla_rope_fused` | 暂不接入 | 直接对比 vLLM perf 与其 torch fallback：q_pe/k_pe 原地 RoPE 结果不一致，cache 近似一致；为避免暴露错误语义，已撤回 bridge |
+| ixCCL send/recv/allreduce/allgather | 归入 distributed/ccl 层 | 不属于单 GPU kernel；需单独检查 InfiniCore CCL 接口 |
+
+## 本轮验证命令摘要
+
+```bash
+xmake f --iluvatar-gpu=true --ccl=false --aten=true --iluvatar_arch=ivcore11 --cuda=/usr/local/corex-4.5.0.20260619 -cv
+xmake build infinicore_cpp_api
+xmake build _infinicore
+xmake install infinicore_cpp_api
+xmake install _infinicore
+xmake build infiniop-iluvatar
+xmake build infiniop
+xmake install infiniop
+python3 test/infiniop/silu_and_mul.py --iluvatar
+python3 test/infiniop/topksoftmax.py --iluvatar
+python3 test/infiniop/deepseek_moe.py --iluvatar
+# grouped_topk_vllm: Python zero-copy torch buffer 测试，对比 torch.ops._moe_C.grouped_topk，并验证 no-bias guard
+# scaled_mm_w4a8: Python zero-copy torch buffer 测试，对比 torch.ops._C.scaled_mm_w4a8，F16/BF16、有无 bias 逐位一致
+# w4a8_group_gemm_: Python zero-copy torch buffer 测试，对比 torch.ops._C.w4a8_group_gemm，F16/BF16、有无 bias、有无 sorted_token_ids、prefill/decode 逐位一致
+# w8a8_group_gemm_: Python zero-copy torch buffer 测试，对比 torch.ops._C.w8a8_group_gemm，F16/BF16、有无 bias、有无 sorted_token_ids、prefill 逐位一致；decode/trans_weight guard 已验证
+```

include/infiniccl.h

@@ -27,6 +27,22 @@ __INFINI_C __export infiniStatus_t infinicclGroupStart(infinicclComm_t comm);
 
 __INFINI_C __export infiniStatus_t infinicclGroupEnd(infinicclComm_t comm);
 
+__INFINI_C __export infiniStatus_t infinicclSend(
+    const void *sendbuf,
+    size_t count,
+    infiniDtype_t datatype,
+    int peer,
+    infinicclComm_t comm,
+    infinirtStream_t stream);
+
+__INFINI_C __export infiniStatus_t infinicclRecv(
+    void *recvbuf,
+    size_t count,
+    infiniDtype_t datatype,
+    int peer,
+    infinicclComm_t comm,
+    infinirtStream_t stream);
+
 __INFINI_C __export infiniStatus_t infinicclAllReduce(
     void *sendbuf,
     void *recvbuf,

include/infinicore/adaptor/aten_adaptor.hpp

@@ -1,5 +1,6 @@
 #ifdef ENABLE_ATEN
 #pragma once
+
 #include "../context/context.hpp"
 #include "../tensor.hpp"
 
@@ -29,6 +30,10 @@ inline at::ScalarType to_at_dtype(DataType dtype) {
         return at::kHalf;
     case DataType::BF16:
         return at::kBFloat16;
+    case DataType::I8:
+        return at::kChar;
+    case DataType::U8:
+        return at::kByte;
     case DataType::I32:
         return at::kInt;
     case DataType::I64:
@@ -40,8 +45,12 @@ inline at::ScalarType to_at_dtype(DataType dtype) {
 
 inline at::Device to_at_device(const Device &device) {
     // PyTorch ATen only exposes standard device types (e.g. kCPU/kCUDA).
-    // Treat MetaX/QY devices as CUDA devices for ATen tensor interoperability.
-    if (device.getType() == Device::Type::NVIDIA || device.getType() == Device::Type::METAX || device.getType() == Device::Type::QY || device.getType() == Device::Type::HYGON) {
+    // Treat CUDA-compatible vendor devices as CUDA devices for ATen tensor interoperability.
+    if (device.getType() == Device::Type::NVIDIA
+        || device.getType() == Device::Type::METAX
+        || device.getType() == Device::Type::QY
+        || device.getType() == Device::Type::ILUVATAR
+        || device.getType() == Device::Type::HYGON) {
         return at::Device(at::kCUDA, device.getIndex());
     } else if (device.getType() == Device::Type::CPU) {
         return at::Device(at::kCPU);
@@ -64,6 +73,10 @@ c10::hip::HIPStream get_hip_stream();
 c10::cuda::CUDAStream get_cuda_stream();
 #endif
 
+#if defined(ENABLE_ILUVATAR_API)
+void set_aten_stream_to_infinicore();
+#endif
+
 #if defined(ENABLE_MOORE_API)
 c10::musa::MUSAStream get_musa_stream();
 #endif

include/infinicore/adaptor/vllm_iluvatar_adaptor.hpp

@@ -0,0 +1,85 @@
+#pragma once
+
+#ifdef ENABLE_ATEN
+#include <ATen/ATen.h>
+#include <optional>
+#include <string>
+#include <tuple>
+
+namespace infinicore::adaptor::vllm_iluvatar {
+
+bool available();
+bool rotary_embedding_available();
+bool dynamic_scaled_int8_quant_available();
+bool concat_mla_q_available();
+bool concat_and_cache_mla_available();
+bool concat_and_cache_mla_int8_available();
+bool paged_attention_mla_available();
+bool topk_softmax_available();
+bool topk_sigmoid_available();
+bool grouped_topk_available();
+bool scaled_mm_w4a8_available();
+bool scaled_mm_w8a8_available();
+bool w4a8_group_gemm_available();
+bool w8a8_group_gemm_available();
+bool w16a16_group_gemm_available();
+bool argsort_bincount_with_inv_pos_available();
+bool expand_moe_input_with_inv_pos_available();
+bool silu_and_mul_quant_available();
+bool moe_sum_vllm_available();
+bool fused_deepseek_v2_indexer_postprocess_available();
+bool indexer_k_cache_available();
+bool indexer_k_quant_and_cache_available();
+bool compute_block_sparse_mqa_logits_available();
+bool select_prefill_topk_block_indices_available();
+bool select_decode_topk_block_indices_available();
+bool map_prefill_request_block_indices_available();
+bool map_decode_request_block_indices_available();
+bool sparse_flash_mla_available();
+bool topk_indices_context_lens_available();
+void fused_add_rms_norm(at::Tensor &input, at::Tensor &residual, at::Tensor &weight, float epsilon);
+void rotary_embedding(at::Tensor &positions,
+                      at::Tensor &query,
+                      std::optional<at::Tensor> key,
+                      int64_t head_size,
+                      at::Tensor &cos_sin_cache,
+                      bool is_neox);
+void dynamic_scaled_int8_quant(at::Tensor &output, at::Tensor &input_scales, const at::Tensor &input);
+void concat_mla_q(at::Tensor &ql_nope, at::Tensor &q_pe, at::Tensor &q_out);
+void concat_and_cache_mla(at::Tensor &kv_c, at::Tensor &k_pe, at::Tensor &kv_cache, at::Tensor &slot_mapping, const std::string &kv_cache_dtype, at::Tensor &scale);
+void concat_and_cache_mla_int8(at::Tensor &kv_c_int8, at::Tensor &kv_c_scale, at::Tensor &k_pe_int8, at::Tensor &k_pe_scale, at::Tensor &kv_cache, at::Tensor &kv_cache_scale, at::Tensor &slot_mapping);
+void paged_attention_mla(at::Tensor &output,
+                         at::Tensor &query,
+                         at::Tensor &kv_cache,
+                         double scale,
+                         at::Tensor &block_tables,
+                         at::Tensor &context_lens,
+                         int64_t max_context_len,
+                         bool use_cuda_graph,
+                         at::Tensor &softmax_lse);
+void topk_softmax(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional<at::Tensor> correction_bias);
+void topk_sigmoid(at::Tensor &topk_weights, at::Tensor &topk_ids, at::Tensor &token_expert_indices, const at::Tensor &gating_output, bool renormalize, std::optional<at::Tensor> correction_bias);
+void grouped_topk(at::Tensor &topk_weights, at::Tensor &topk_ids, const at::Tensor &scores, std::optional<at::Tensor> bias, int64_t num_expert_group, int64_t topk_group, const std::string &scoring_func, bool renormalize);
+void scaled_mm_w4a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional<at::Tensor> bias, bool trans_weight);
+void scaled_mm_w8a8(at::Tensor &out, const at::Tensor &a, const at::Tensor &b, const at::Tensor &a_scales, const at::Tensor &b_scales, std::optional<at::Tensor> bias, bool trans_weight);
+void w4a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional<at::Tensor> sorted_token_ids, std::optional<at::Tensor> bias, bool trans_weight, bool is_decode);
+void w8a8_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &input_scale, const at::Tensor &weight_scale, const at::Tensor &tokens_per_experts, std::optional<at::Tensor> sorted_token_ids, std::optional<at::Tensor> bias, bool trans_weight, bool is_decode);
+void w16a16_group_gemm(at::Tensor &out, const at::Tensor &input, const at::Tensor &weight, const at::Tensor &tokens_per_experts, std::optional<at::Tensor> sorted_token_ids, std::optional<at::Tensor> bias, bool trans_weight, bool is_decode);
+void argsort_bincount_with_inv_pos(const at::Tensor &topk_ids, at::Tensor &tokens_per_experts, at::Tensor &sorted_indices, at::Tensor &inv_pos, int64_t num_experts);
+void expand_moe_input_with_inv_pos(at::Tensor &expand_states, std::optional<at::Tensor> expand_scales, const at::Tensor &hidden_states, const at::Tensor &inv_pos, int64_t top_k, int64_t group_size, int64_t format);
+void silu_and_mul_quant(at::Tensor &output, std::optional<at::Tensor> output_scale, const at::Tensor &input, int64_t format);
+void moe_sum_vllm(at::Tensor &output, const at::Tensor &input, std::optional<at::Tensor> topk_weights, std::optional<at::Tensor> extra_residual, double routed_scale, double residual_scale);
+void select_last_token_hidden(at::Tensor &output, at::Tensor &indices, const at::Tensor &hidden_states, const at::Tensor &input_offsets);
+void fused_deepseek_v2_indexer_postprocess(at::Tensor &q_out, at::Tensor &k_out, at::Tensor &weights_out, at::Tensor &kv_cache, const at::Tensor &slot_mapping, const at::Tensor &q, const at::Tensor &kw, const at::Tensor &norm_weight, const at::Tensor &norm_bias, const at::Tensor &positions, const at::Tensor &cos_sin_cache, int64_t num_cache_tokens, bool is_neox, double eps, double weights_scale);
+void indexer_k_cache(const at::Tensor &k, at::Tensor &kv_cache, const at::Tensor &slot_mapping);
+void indexer_k_quant_and_cache(at::Tensor &k, at::Tensor &kv_cache, at::Tensor &slot_mapping, int64_t quant_block_size, const std::string &scale_fmt);
+void compute_block_sparse_mqa_logits(const at::Tensor &q, const at::Tensor &kv_cache, const at::Tensor &cu_seqlens_q, const at::Tensor &cu_seqlens_kv, const at::Tensor &block_table, const at::Tensor &weights, at::Tensor &logits, int64_t max_q_len, int64_t max_kv_len, int64_t max_context_len);
+void select_prefill_topk_block_indices(const at::Tensor &logits, const at::Tensor &cu_seqlen_ks, const at::Tensor &cu_seqlen_ke, at::Tensor &topk_indices);
+void select_decode_topk_block_indices(const at::Tensor &logits, const at::Tensor &seq_lens, at::Tensor &topk_indices);
+void map_prefill_request_block_indices(at::Tensor &output, const at::Tensor &req_id, const at::Tensor &block_table, const at::Tensor &token_indices, int64_t block_size, bool has_prefill_workspace, std::optional<at::Tensor> prefill_workspace_request_ids, std::optional<at::Tensor> prefill_workspace_starts);
+void map_decode_request_block_indices(at::Tensor &output, const at::Tensor &req_id, const at::Tensor &block_table, const at::Tensor &token_indices, int64_t block_size);
+void topk_indices_context_lens(at::Tensor &topk_lens, const at::Tensor &indices);
+void sparse_flash_mla(at::Tensor &output, at::Tensor &query, at::Tensor &kv_cache, at::Tensor &indices, at::Tensor &topk_lens, float scale, std::optional<at::Tensor> attn_sink);
+
+} // namespace infinicore::adaptor::vllm_iluvatar
+#endif // ENABLE_ATEN

include/infinicore/context/context.hpp

@@ -51,6 +51,7 @@ bool isGraphRecording();
 void startGraphRecording();
 void addGraphOperator(std::shared_ptr<graph::GraphOperator> op);
 std::shared_ptr<graph::Graph> stopGraphRecording();
+void cancelGraphRecording();
 
 } // namespace context
 

include/infinicore/graph/graph.hpp

@@ -17,6 +17,12 @@ class GraphTensor : public Tensor {
 class GraphOperator {
 public:
     virtual void run() const = 0;
+    virtual bool is_device_graph_capture_safe() const {
+        return true;
+    }
+    virtual bool requires_stream_sync_after_run() const {
+        return false;
+    }
     virtual ~GraphOperator() = default;
 };
 
@@ -49,7 +55,8 @@ class Graph {
 
 private:
     struct DeviceGraph;
-    std::unique_ptr<DeviceGraph> device_graph_;
+    struct ReplayStep;
+    std::vector<std::unique_ptr<ReplayStep>> replay_steps_;
 };
 } // namespace infinicore::graph
 

include/infinicore/nn/rope.hpp

@@ -87,6 +87,8 @@ class RoPE : public Module {
     DataType dtype() const { return dtype_; }
     const std::optional<std::vector<int>> &mrope_section() const { return mrope_section_; }
     bool mrope_interleaved() const { return mrope_interleaved_; }
+    const Tensor &sin_cache() const { return sin_cache_; }
+    const Tensor &cos_cache() const { return cos_cache_; }
 
     // String representation
     std::string extra_repr() const;