[2026春季][T2-1-3] goog00

include/infinicore/ops.hpp

@@ -32,7 +32,9 @@
 #include "ops/fused_gated_delta_net_gating.hpp"
 #include "ops/gelu.hpp"
 #include "ops/gelutanh.hpp"
+#include "ops/grouped_gemm.hpp"
 #include "ops/hardswish.hpp"
+#include "ops/index_add.hpp"
 #include "ops/hardtanh.hpp"
 #include "ops/kv_caching.hpp"
 #include "ops/layer_norm.hpp"
@@ -45,6 +47,7 @@
 #include "ops/moe_sum.hpp"
 #include "ops/moe_topk_sigmoid.hpp"
 #include "ops/moe_topk_softmax.hpp"
+#include "ops/mul.hpp"
 #include "ops/nrm2.hpp"
 #include "ops/ones.hpp"
 #include "ops/paged_attention.hpp"

include/infinicore/ops/grouped_gemm.hpp

@@ -0,0 +1,51 @@
+#pragma once
+
+#include <cstdint>
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(GroupedGemm,
+                          Tensor,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          float,
+                          float,
+                          const int32_t *);
+
+// Variable-batched matmul shared across `num_groups` weight slabs.
+//
+// Shapes (all row-major):
+//   a            : [M_total, K]
+//   b            : [num_groups, N, K]      // out, in -- matches torch linear
+//   c            : [M_total, N]
+//   group_sizes  : [num_groups], int32     // sum == M_total
+//
+// For each group g:
+//   c[off_g : off_g + group_sizes[g]] = alpha * a[off_g : off_g + group_sizes[g]] @ b[g].T
+//                                      + beta * c[...]
+// where `off_g = sum(group_sizes[0..g])`.
+// `group_sizes_host` (optional): host-side copy of `group_sizes`. When given,
+// device backends use it directly and skip the per-call device->host sync of
+// the sizes array. Must stay valid until the call returns; pass nullptr (the
+// default) under graph capture or when no host copy is available.
+Tensor grouped_gemm(const Tensor &a,
+                    const Tensor &b,
+                    const Tensor &group_sizes,
+                    float alpha = 1.0f,
+                    float beta = 0.0f,
+                    const int32_t *group_sizes_host = nullptr);
+
+void grouped_gemm_(Tensor c,
+                   const Tensor &a,
+                   const Tensor &b,
+                   const Tensor &group_sizes,
+                   float alpha,
+                   float beta,
+                   const int32_t *group_sizes_host = nullptr);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -62,6 +62,7 @@
 #include "infiniop/ops/gemm.h"
 #include "infiniop/ops/gptq_marlin_gemm.h"
 #include "infiniop/ops/gptq_qyblas_gemm.h"
+#include "infiniop/ops/grouped_gemm.h"
 #include "infiniop/ops/hardswish.h"
 #include "infiniop/ops/hardtanh.h"
 #include "infiniop/ops/hinge_embedding_loss.h"

include/infiniop/ops/grouped_gemm.h

@@ -0,0 +1,67 @@
+#ifndef __INFINIOP_GROUPED_GEMM_API_H__
+#define __INFINIOP_GROUPED_GEMM_API_H__
+
+#include "../operator_descriptor.h"
+
+/**
+ * Variable-batched (a.k.a. "grouped" or "segment") GEMM.
+ *
+ * For each group `g` in `[0, num_groups)`:
+ *   rows_g  = group_sizes[g]
+ *   off_g   = sum(group_sizes[0..g])
+ *   A_g     = a[off_g : off_g + rows_g, :]            // [rows_g, K]
+ *   B_g     = b[g, :, :]                              // [N, K]
+ *   c[off_g : off_g + rows_g, :] = alpha * A_g @ B_g^T + beta * c_g
+ *
+ * Shapes
+ *   a : [M_total, K]
+ *   b : [num_groups, N, K]   -- matches PyTorch / HF linear weight layout
+ *   c : [M_total, N]
+ *   group_sizes : [num_groups], int32, sum == M_total
+ *
+ * Constraints
+ *   - M_total = sum(group_sizes); enforced at calculate-time.
+ *   - All `*_desc` must be row-major contiguous on the leading axes.
+ *   - dtype: F16, BF16 or F32 (same for a, b, c).
+ *   - `group_sizes` is `INFINI_DTYPE_I32`.
+ *   - The `group_sizes` data pointer lives on the same device as `a/b/c`.
+ *     CPU backends read it directly; device backends sync it to host
+ *     inside `calculate`.
+ *   - `group_sizes_host` (optional, may be NULL): host-side copy of the same
+ *     int32 group sizes. When provided, device backends use it directly and
+ *     skip the per-call device->host copy + stream sync. The caller must keep
+ *     it valid and consistent with `group_sizes` until the call returns. NULL
+ *     preserves the legacy device-sync behavior. Unsafe under graph capture
+ *     (pass NULL there).
+ */
+typedef struct InfiniopDescriptor *infiniopGroupedGemmDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateGroupedGemmDescriptor(
+    infiniopHandle_t handle,
+    infiniopGroupedGemmDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t c_desc,
+    infiniopTensorDescriptor_t a_desc,
+    infiniopTensorDescriptor_t b_desc,
+    infiniopTensorDescriptor_t group_sizes_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetGroupedGemmWorkspaceSize(
+    infiniopGroupedGemmDescriptor_t desc,
+    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopGroupedGemm(
+    infiniopGroupedGemmDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *c,
+    const void *a,
+    const void *b,
+    const void *group_sizes,
+    const void *group_sizes_host,
+    float alpha,
+    float beta,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyGroupedGemmDescriptor(
+    infiniopGroupedGemmDescriptor_t desc);
+
+#endif // __INFINIOP_GROUPED_GEMM_API_H__

src/infinicore/ops/grouped_gemm/grouped_gemm.cc

@@ -0,0 +1,53 @@
+#include "infinicore/ops/grouped_gemm.hpp"
+
+#include "../../utils.hpp"
+
+namespace infinicore::op {
+INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(GroupedGemm);
+
+GroupedGemm::GroupedGemm(Tensor c,
+                         const Tensor &a,
+                         const Tensor &b,
+                         const Tensor &group_sizes,
+                         float alpha,
+                         float beta,
+                         const int32_t *group_sizes_host) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b, group_sizes);
+    INFINICORE_GRAPH_OP_DISPATCH(c->device().getType(), c, a, b, group_sizes, alpha, beta, group_sizes_host);
+}
+
+void GroupedGemm::execute(Tensor c,
+                          const Tensor &a,
+                          const Tensor &b,
+                          const Tensor &group_sizes,
+                          float alpha,
+                          float beta,
+                          const int32_t *group_sizes_host) {
+    INFINICORE_GRAPH_OP_RECORD_OR_RUN(GroupedGemm, c, a, b, group_sizes, alpha, beta, group_sizes_host);
+}
+
+Tensor grouped_gemm(const Tensor &a,
+                    const Tensor &b,
+                    const Tensor &group_sizes,
+                    float alpha,
+                    float beta,
+                    const int32_t *group_sizes_host) {
+    // a: [M_total, K], b: [num_groups, N, K] -> c: [M_total, N].
+    Shape shape = a->shape();
+    shape[shape.size() - 1] = b->size(1);
+    auto c = Tensor::empty(shape, a->dtype(), a->device());
+    grouped_gemm_(c, a, b, group_sizes, alpha, beta, group_sizes_host);
+    return c;
+}
+
+void grouped_gemm_(Tensor c,
+                   const Tensor &a,
+                   const Tensor &b,
+                   const Tensor &group_sizes,
+                   float alpha,
+                   float beta,
+                   const int32_t *group_sizes_host) {
+    GroupedGemm::execute(c, a, b, group_sizes, alpha, beta, group_sizes_host);
+}
+
+} // namespace infinicore::op

src/infinicore/ops/grouped_gemm/grouped_gemm_infiniop.cc

@@ -0,0 +1,63 @@
+#include "../infiniop_impl.hpp"
+#include "infinicore/ops/grouped_gemm.hpp"
+
+namespace infinicore::op::grouped_gemm_impl::infiniop {
+
+INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, GroupedGemm, 100);
+
+struct PlannedMeta {
+    std::shared_ptr<Descriptor> descriptor;
+    graph::GraphTensor workspace, c, a, b, group_sizes;
+    float alpha, beta;
+    // Optional host-side group sizes; nullptr falls back to the device sync.
+    // Only valid for immediate (non-recorded) execution -- the pointer is not
+    // owned, so graph replay must pass nullptr.
+    const int32_t *group_sizes_host;
+};
+
+void *plan(Tensor c, const Tensor &a, const Tensor &b, const Tensor &group_sizes, float alpha, float beta, const int32_t *group_sizes_host) {
+    size_t seed = hash_combine(c, a, b, group_sizes);
+
+    INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(
+        Descriptor, descriptor, GroupedGemm,
+        seed, c->desc(), a->desc(), b->desc(), group_sizes->desc());
+
+    INFINIOP_WORKSPACE_TENSOR(workspace, GroupedGemm, descriptor);
+
+    auto planned = new PlannedMeta{
+        descriptor,
+        graph::GraphTensor(workspace),
+        graph::GraphTensor(c),
+        graph::GraphTensor(a),
+        graph::GraphTensor(b),
+        graph::GraphTensor(group_sizes),
+        alpha, beta,
+        group_sizes_host};
+
+    return planned;
+}
+
+void run(void *planned_meta) {
+    auto planned = reinterpret_cast<PlannedMeta *>(planned_meta);
+
+    INFINICORE_CHECK_ERROR(infiniopGroupedGemm(
+        planned->descriptor->desc,
+        planned->workspace->data(), planned->workspace->numel(),
+        planned->c->data(),
+        planned->a->data(),
+        planned->b->data(),
+        planned->group_sizes->data(),
+        planned->group_sizes_host,
+        planned->alpha,
+        planned->beta,
+        context::getStream()));
+}
+
+void cleanup(void **planned_meta_ptr) {
+    delete *reinterpret_cast<PlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(GroupedGemm, &plan, &run, &cleanup);
+
+} // namespace infinicore::op::grouped_gemm_impl::infiniop

src/infinicore/tensor/tensor.cc

@@ -4,6 +4,7 @@
 #include "infinicore/context/context.hpp"
 #include "infinicore/dtype.hpp"
 
+#include <cstring>
 #include <spdlog/spdlog.h>
 
 namespace {
@@ -242,7 +243,14 @@ std::shared_ptr<TensorImpl> TensorImpl::zeros(const Shape &shape,
                                               const DataType &dtype,
                                               const Device &device,
                                               bool pin_memory) {
-
+    // MetaX device-side integer zeros is unreliable (does not actually clear),
+    // so construct zeros on CPU and copy over. Other backends use the faster
+    // device-side byte memset.
+    if (device.getType() == Device::Type::METAX) {
+        auto cpu = empty(shape, dtype, Device(Device::Type::CPU), false);
+        std::memset(cpu->data(), 0, cpu->nbytes());
+        return cpu->to(device).impl_;
+    }
     auto result = empty(shape, dtype, device, pin_memory);
     context::setDeviceMemoryAsync(result->data(), 0, result->nbytes(), context::getStream());
     return result;

src/infiniop/devices/metax/metax_ht2mc.h

@@ -116,6 +116,7 @@
 #define hcclAllReduce mcclAllReduce
 #define hcblasSetStream mcblasSetStream
 #define hcblasHandle_t mcblasHandle_t
+#define hcblasGemmBatchedEx mcblasGemmBatchedEx
 #define hcblasGemmStridedBatchedEx mcblasGemmStridedBatchedEx
 #define hcblasGemmEx mcblasGemmEx
 #define hcblasCreate mcblasCreate