Merge branch 'main' of https://dev.modelhub.org.cn/dylanyunlong/project_6 into main

MOE_SYMBOL_TRUTH.md

@@ -0,0 +1,68 @@
+# MoE 函数符号真相 (2026-08-17 确认)
+
+## 结论
+
+那5个 MoE 函数**确实不在任何镜像预装的 .so 里**。另一位开发者说的是对的。
+
+但它们也**不需要**在预装 .so 里——它们是自编译的。
+
+## 5个函数的正确命名空间
+
+```
+ixformer::infer::topk_softmax
+ixformer::infer::moe_compute_token_index_api
+ixformer::infer::moe_expand_input
+ixformer::infer::moe_w16a16_group_gemm
+ixformer::infer::moe_output_reduce_sum
+```
+
+**注意**: 是 `ixformer::infer`，不是 `ixformer::kernels::infer`。
+
+## 声明 vs 实现的关系
+
+| 位置 | 角色 |
+|------|------|
+| `ixformer_sdk/csrc/include/ixformer/kernels/kernels.h` | **头文件声明** (namespace `ixformer::kernels::infer`) — C++ 模板声明，给 SDK 用的 |
+| `ex_engine/csrc/moe_ops_impl.cu` | **CUDA 实现** (namespace `ixformer::infer`) — 自己写的 kernel，不依赖任何 .so |
+| `ex_engine/csrc/ix_full_bridge_v2.cpp` | **pybind11 桥** — forward-declare 然后调用 moe_ops_impl.cu 里的实现 |
+| `ex_engine/build_moe_bridge.sh` | **构建脚本** — 把 v2.cpp + moe_ops_impl.cu 一起编译成 ix_full_bridge_v2.so |
+
+## 符号表搜索结果 (4个 .so 全部搜过)
+
+| .so 文件 | MoE 函数 | 结论 |
+|----------|----------|------|
+| `libixformer.so` (3937 symbols) | 无 topk_softmax/moe_compute_token_index 等 | 只有 `reduce_sum` (通用的) |
+| `_ixformer_torch.so` (49 symbols) | 完全没有 MoE | 只有 norm/rope/cache/attn |
+| `_C.so` (6 symbols) | 几乎空壳 | 只有 PyInit |
+| `libcuinfer.so` (270 symbols) | 只有 cuinferTopK (不是 MoE 的) | GEMM/BLAS 级别 |
+
+## 构建链
+
+```
+patch_ops.sh
+  └→ build_moe_bridge.sh
+       └→ ninja/CppExtension 编译:
+            ix_full_bridge_v2.cpp + moe_ops_impl.cu
+            → ix_full_bridge_v2.so (包含5个MoE函数的实现)
+```
+
+## `ixformer::kernels::infer` vs `ixformer::infer` 的区别
+
+- `ixformer::kernels::infer` — SDK 头文件 (kernels.h) 中的声明，使用 raw pointer + cudaStream_t 
+  - 例: `void moe_topk_softmax(const T *gating_output, T *topk_weights, int *topk_indices, ...)`
+- `ixformer::infer` — 我们自己实现的 PyTorch wrapper，使用 torch::Tensor
+  - 例: `void topk_softmax(torch::Tensor& topk_weights, torch::Tensor& topk_indices, ...)`
+
+`moe_ops_impl.cu` 是直接写 CUDA kernel（不调用 kernels.h 模板），然后暴露 Tensor API。
+
+## Python 调用链
+
+```python
+# 通过 ixformer SDK (需要真机上的 _C.so 包含 infer 子模块):
+import ixformer._C as ops
+ops.infer.moe_topk_softmax(...)  # 如果 _C.so 有实现
+
+# 通过 ex_engine bridge (我们自编译的):
+import ix_full_bridge_v2 as bridge
+bridge.topk_softmax(...)  # 来自 moe_ops_impl.cu
+```

ex_engine/python/ix_ops.py

@@ -222,11 +222,16 @@ def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,
     """Fused residual addition + RMSNorm.
 
     Source: xllm/core/kernels/ilu/norm.cpp → infer::residual_rms_norm
-    output = rms_norm(input + residual, weight, eps)
-    residual_output = input + residual
+    The C++ function is in-place: modifies input → rms_norm(input+residual)*weight,
+    and residual → input+residual. We copy results to output/residual_output.
     """
-    _bridge.fused_add_rms_norm(input, residual, weight, output,
-                                residual_output, eps)
+    # C++ signature: fused_add_rms_norm_forward(input, residual, weight, eps, alpha)
+    # It modifies input and residual in-place.
+    inp_clone = input.clone()
+    res_clone = residual.clone()
+    _bridge.fused_add_rms_norm(inp_clone, res_clone, weight, eps)
+    output.copy_(inp_clone)
+    residual_output.copy_(res_clone)
 
 
 def rotary_embedding(positions: torch.Tensor, query: torch.Tensor,

ex_engine/python/patch_vllm_ops.py

@@ -92,28 +92,26 @@ def _patch_layernorm() -> int:
         w = self.weight
         if _debug_count[0] < 20:
             _debug_count[0] += 1
-            logger.info("DEBUG rms_norm #%d: w.shape=%s w.dim=%d x.shape=%s x.dim=%d "
-                         "class=%s residual=%s",
-                         _debug_count[0], list(w.shape), w.dim(), list(x.shape), x.dim(),
+            logger.info("DEBUG rms_norm #%d: w.shape=%s w.dim=%d w.dtype=%s "
+                         "x.shape=%s x.dim=%d x.dtype=%s class=%s residual=%s",
+                         _debug_count[0], list(w.shape), w.dim(), w.dtype,
+                         list(x.shape), x.dim(), x.dtype,
                          type(self).__name__,
                          list(residual.shape) if residual is not None else None)
         if w.dim() != 1 or w.shape[0] != x.shape[-1]:
             return _orig_forward(self, x, residual)
-        w_adjusted = 1.0 + w
+        # 1.0 + w promotes fp16→fp32; ixformer rms_norm requires weight
+        # to be 1-D AND same dtype as input, so cast back.
+        w_adjusted = (1.0 + w).to(w.dtype)
         if residual is not None:
-            if ix_ops.has_fused_add_rms_norm():
-                out = torch.empty_like(x)
-                residual_out = torch.empty_like(x)
-                ix_ops.fused_add_rms_norm(
-                    x, residual, w_adjusted, out, residual_out,
-                    self.variance_epsilon)
-                return out, residual_out
-            else:
-                new_residual = x + residual
-                out = torch.empty_like(x)
-                ix_ops.rms_norm(out, new_residual, w_adjusted,
-                                self.variance_epsilon)
-                return out, new_residual
+            # ixformer fused_add_rms_norm is in-place and has 4-arg C++
+            # signature (input, residual, weight, eps). Safer to use
+            # the non-fused path which is explicit about outputs.
+            new_residual = x + residual
+            out = torch.empty_like(x)
+            ix_ops.rms_norm(out, new_residual, w_adjusted,
+                            self.variance_epsilon)
+            return out, new_residual
         else:
             out = torch.empty_like(x)
             ix_ops.rms_norm(out, x, w_adjusted, self.variance_epsilon)
@@ -180,17 +178,17 @@ def _patch_custom_ops() -> int:
         logger.info("PATCHED: _custom_ops.rms_norm → ix_ops")
 
     # Patch fused_add_rms_norm
-    if ix_ops.has_fused_add_rms_norm() and hasattr(ops, 'fused_add_rms_norm'):
+    if ix_ops.has_rms_norm() and hasattr(ops, 'fused_add_rms_norm'):
         def _fused_add_rms_norm(input, residual, weight, eps):
+            # C++ fused_add_rms_norm is in-place with 4-arg signature,
+            # doesn't match the 6-arg wrapper in ix_ops. Use non-fused path.
+            residual.add_(input)
             out = torch.empty_like(input)
-            residual_out = torch.empty_like(input)
-            ix_ops.fused_add_rms_norm(input, residual, weight,
-                                       out, residual_out, eps)
+            ix_ops.rms_norm(out, residual, weight, eps)
             input.copy_(out)
-            residual.copy_(residual_out)
         ops.fused_add_rms_norm = _fused_add_rms_norm
         count += 1
-        logger.info("PATCHED: _custom_ops.fused_add_rms_norm → ix_ops")
+        logger.info("PATCHED: _custom_ops.fused_add_rms_norm → ix_ops (non-fused)")
 
     # Patch rotary_embedding
     if ix_ops.has_rotary_embedding() and hasattr(ops, 'rotary_embedding'):

qwen3_6_scripts/ex_engine/python/ix_ops.py

@@ -222,11 +222,16 @@ def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,
     """Fused residual addition + RMSNorm.
 
     Source: xllm/core/kernels/ilu/norm.cpp → infer::residual_rms_norm
-    output = rms_norm(input + residual, weight, eps)
-    residual_output = input + residual
+    The C++ function is in-place: modifies input → rms_norm(input+residual)*weight,
+    and residual → input+residual. We copy results to output/residual_output.
     """
-    _bridge.fused_add_rms_norm(input, residual, weight, output,
-                                residual_output, eps)
+    # C++ signature: fused_add_rms_norm_forward(input, residual, weight, eps, alpha)
+    # It modifies input and residual in-place.
+    inp_clone = input.clone()
+    res_clone = residual.clone()
+    _bridge.fused_add_rms_norm(inp_clone, res_clone, weight, eps)
+    output.copy_(inp_clone)
+    residual_output.copy_(res_clone)
 
 
 def rotary_embedding(positions: torch.Tensor, query: torch.Tensor,

qwen3_6_scripts/ex_engine/python/patch_vllm_ops.py

@@ -92,28 +92,26 @@ def _patch_layernorm() -> int:
         w = self.weight
         if _debug_count[0] < 20:
             _debug_count[0] += 1
-            logger.info("DEBUG rms_norm #%d: w.shape=%s w.dim=%d x.shape=%s x.dim=%d "
-                         "class=%s residual=%s",
-                         _debug_count[0], list(w.shape), w.dim(), list(x.shape), x.dim(),
+            logger.info("DEBUG rms_norm #%d: w.shape=%s w.dim=%d w.dtype=%s "
+                         "x.shape=%s x.dim=%d x.dtype=%s class=%s residual=%s",
+                         _debug_count[0], list(w.shape), w.dim(), w.dtype,
+                         list(x.shape), x.dim(), x.dtype,
                          type(self).__name__,
                          list(residual.shape) if residual is not None else None)
         if w.dim() != 1 or w.shape[0] != x.shape[-1]:
             return _orig_forward(self, x, residual)
-        w_adjusted = 1.0 + w
+        # 1.0 + w promotes fp16→fp32; ixformer rms_norm requires weight
+        # to be 1-D AND same dtype as input, so cast back.
+        w_adjusted = (1.0 + w).to(w.dtype)
         if residual is not None:
-            if ix_ops.has_fused_add_rms_norm():
-                out = torch.empty_like(x)
-                residual_out = torch.empty_like(x)
-                ix_ops.fused_add_rms_norm(
-                    x, residual, w_adjusted, out, residual_out,
-                    self.variance_epsilon)
-                return out, residual_out
-            else:
-                new_residual = x + residual
-                out = torch.empty_like(x)
-                ix_ops.rms_norm(out, new_residual, w_adjusted,
-                                self.variance_epsilon)
-                return out, new_residual
+            # ixformer fused_add_rms_norm is in-place and has 4-arg C++
+            # signature (input, residual, weight, eps). Safer to use
+            # the non-fused path which is explicit about outputs.
+            new_residual = x + residual
+            out = torch.empty_like(x)
+            ix_ops.rms_norm(out, new_residual, w_adjusted,
+                            self.variance_epsilon)
+            return out, new_residual
         else:
             out = torch.empty_like(x)
             ix_ops.rms_norm(out, x, w_adjusted, self.variance_epsilon)
@@ -180,17 +178,17 @@ def _patch_custom_ops() -> int:
         logger.info("PATCHED: _custom_ops.rms_norm → ix_ops")
 
     # Patch fused_add_rms_norm
-    if ix_ops.has_fused_add_rms_norm() and hasattr(ops, 'fused_add_rms_norm'):
+    if ix_ops.has_rms_norm() and hasattr(ops, 'fused_add_rms_norm'):
         def _fused_add_rms_norm(input, residual, weight, eps):
+            # C++ fused_add_rms_norm is in-place with 4-arg signature,
+            # doesn't match the 6-arg wrapper in ix_ops. Use non-fused path.
+            residual.add_(input)
             out = torch.empty_like(input)
-            residual_out = torch.empty_like(input)
-            ix_ops.fused_add_rms_norm(input, residual, weight,
-                                       out, residual_out, eps)
+            ix_ops.rms_norm(out, residual, weight, eps)
             input.copy_(out)
-            residual.copy_(residual_out)
         ops.fused_add_rms_norm = _fused_add_rms_norm
         count += 1
-        logger.info("PATCHED: _custom_ops.fused_add_rms_norm → ix_ops")
+        logger.info("PATCHED: _custom_ops.fused_add_rms_norm → ix_ops (non-fused)")
 
     # Patch rotary_embedding
     if ix_ops.has_rotary_embedding() and hasattr(ops, 'rotary_embedding'):