This commit is contained in:
@@ -4,8 +4,12 @@ WORKDIR /workspace/
|
|||||||
# Copy all our engine patches
|
# Copy all our engine patches
|
||||||
COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
|
COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
|
||||||
COPY ./computility-run.yaml /workspace/computility-run.yaml
|
COPY ./computility-run.yaml /workspace/computility-run.yaml
|
||||||
# Copy entire ex_engine — python dispatch, csrc, build scripts, headers
|
# Copy ex_engine source for MoE bridge compilation
|
||||||
COPY ./ex_engine /workspace/ex_engine
|
COPY ./ex_engine/csrc/moe_ops_impl.cu /workspace/qwen3_6_scripts/ex_engine_src/csrc/moe_ops_impl.cu
|
||||||
|
COPY ./ex_engine/csrc/ix_full_bridge_v2.cpp /workspace/qwen3_6_scripts/ex_engine_src/csrc/ix_full_bridge_v2.cpp
|
||||||
|
COPY ./ex_engine/build_moe_bridge.sh /workspace/qwen3_6_scripts/ex_engine_src/build_moe_bridge.sh
|
||||||
|
COPY ./ex_engine/python/moe_dispatch.py /workspace/qwen3_6_scripts/ex_engine_src/python/moe_dispatch.py
|
||||||
|
COPY ./ex_engine/python/patch_moe_hot_path.py /workspace/qwen3_6_scripts/ex_engine_src/python/patch_moe_hot_path.py
|
||||||
# Make patch script executable and run it
|
# Make patch script executable and run it
|
||||||
RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
|
RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
|
||||||
bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
|
bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
|
||||||
|
|||||||
@@ -1,172 +0,0 @@
|
|||||||
# EX_ENGINE 基建盘点
|
|
||||||
|
|
||||||
日期: 2026-08-17
|
|
||||||
基于: commit 512f384a (CUTLASS Cu10 grouped GEMM 真机验证) + 后续 revert
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 一、Docker容器里实际运行的状态
|
|
||||||
|
|
||||||
### 进了Docker且工作正常的 (prebuilt .so)
|
|
||||||
|
|
||||||
| .so 文件 | qwen3_5.py flag | 状态 |
|
|
||||||
|---|---|---|
|
|
||||||
| corex_gdn_causal_conv.so | _USE_COREX_GDN_CAUSAL_CONV | ✓ 工作 |
|
|
||||||
| corex_gdn_gated_norm.so | _USE_COREX_GDN_GATED_NORM | ✓ 工作 |
|
|
||||||
| corex_gdn_beta_decay.so | _USE_COREX_GDN_BETA_DECAY | ✓ 工作 |
|
|
||||||
| corex_gdn_qk_map.so | _USE_COREX_GDN_QK_MAP | ✓ 工作 |
|
|
||||||
| corex_gdn_packed_decode.so | _USE_COREX_GDN_PACKED_DECODE | ✓ 工作 |
|
|
||||||
| corex_gdn_chunk_recurrent.so | _HAS_COREX_GDN_CHUNK | ✓ 工作 |
|
|
||||||
| corex_moe_topk_softmax.so | _USE_COREX_MOE_TOPK_SOFTMAX | ✓ 工作 |
|
|
||||||
| corex_moe_direct_routed.so | _USE_COREX_MOE_DIRECT_ROUTED | ✓ 工作 |
|
|
||||||
| corex_moe_weight_gather.so | _USE_COREX_MOE_WEIGHT_GATHER | ✓ 工作 |
|
|
||||||
| corex_moe_exact_reduce.so | _USE_COREX_MOE_EXACT_REDUCE | ✓ 工作 |
|
|
||||||
| corex_moe_index_combine.so | _USE_COREX_MOE_INDEX_COMBINE | ✓ 工作 |
|
|
||||||
| corex_attn_head_rms_norm.so | _USE_COREX_ATTN_HEAD_RMS_NORM | ✓ 工作 |
|
|
||||||
| corex_block_major_kv_transfer.so | (block_major_kv_cache.py用) | ✓ 工作 |
|
|
||||||
| corex_paged_kv_gather.so | (paged_attn.py用) | ✓ 工作 |
|
|
||||||
| corex_fused_paged_prefill.so | (corex_fa2用) | ✓ 工作 |
|
|
||||||
| xllm_moe.so | _USE_XLLM_MOE | ✓ 工作 |
|
|
||||||
| xllm_activation.so | (patch_vllm_ops用) | ? 见下 |
|
|
||||||
| xllm_norm.so | (patch_vllm_ops用) | ? 见下 |
|
|
||||||
| xllm_rope.so | (patch_vllm_ops用) | ? 见下 |
|
|
||||||
| xllm_cache.so | (patch_vllm_ops用) | ? 见下 |
|
|
||||||
| ix_full_bridge.so | (ix_ops.py用) | ? 见下 |
|
|
||||||
|
|
||||||
### 进了Docker但断裂的
|
|
||||||
|
|
||||||
| 文件 | 问题 |
|
|
||||||
|---|---|
|
|
||||||
| ix_full_bridge.so (v1) | 已cp到$VLLM_ROOT/, 但python wrapper ix_ops.py没部署 |
|
|
||||||
| xllm_activation/norm/rope/cache.so | 已cp到$VLLM_ROOT/, 但patch_vllm_ops.py没部署,没hook |
|
|
||||||
| ix_fused_moe.py | 已cp到models/, 但找不到ix_moe_bridge.so → _HAS_IX_FUSED_MOE=False |
|
|
||||||
|
|
||||||
### 没进Docker的关键文件
|
|
||||||
|
|
||||||
| 文件 | 功能 | 行数 |
|
|
||||||
|---|---|---|
|
|
||||||
| ex_engine/python/ix_ops.py | ix_full_bridge.so的Python wrapper | 343 |
|
|
||||||
| ex_engine/python/ix_ops_dispatch.py | 统一op dispatch (bridge→ixformer→raise) | 407 |
|
|
||||||
| ex_engine/python/patch_vllm_ops.py | monkey-patch vllm的silu/rms_norm/rope/cache | 201 |
|
|
||||||
| ex_engine/python/corex_moe.py | corex MoE pipeline wrapper | 237 |
|
|
||||||
| ex_engine/python/corex_gdn.py | corex GDN ops wrapper | 256 |
|
|
||||||
| ex_engine/python/corex_fa2.py | corex FlashAttn dispatch | 279 |
|
|
||||||
| ex_engine/python/corex_fa2_dispatch.py | FA2 3-mode dispatch | 231 |
|
|
||||||
| ex_engine/python/fused_moe_ilu.py | 7-step MoE pipeline (Python) | 205 |
|
|
||||||
| ex_engine/python/gemm_dispatch.py | GEMM dispatch (cutlass/cuinfer/torch) | 180 |
|
|
||||||
| ex_engine/csrc/gemm_grouped.cu | ✓ 真机验证的CUTLASS grouped GEMM | 188 |
|
|
||||||
| ex_engine/csrc/gemm_grouped_bind.cpp | pybind11 binding | 182 |
|
|
||||||
| ex_engine/build_gemm_grouped.sh | 编译脚本 | — |
|
|
||||||
| ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu | CUTLASS batched GEMM | 67 |
|
|
||||||
| ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp | binding | 129 |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 二、两个路径断裂的根因
|
|
||||||
|
|
||||||
### 断裂1: patch_ops.sh 找不到 ex_engine/python/
|
|
||||||
|
|
||||||
patch_ops.sh 第204行:
|
|
||||||
```bash
|
|
||||||
EX_ENGINE_DIR="$(cd "$(dirname "$0")/../ex_engine" 2>/dev/null && pwd || echo "")"
|
|
||||||
```
|
|
||||||
|
|
||||||
Docker容器里的目录结构:
|
|
||||||
```
|
|
||||||
/workspace/
|
|
||||||
├── qwen3_6_scripts/ ← patch_ops.sh 在这里
|
|
||||||
│ ├── patch_ops.sh
|
|
||||||
│ ├── ex_engine_src/ ← Dockerfile COPY进来的(只有5个文件)
|
|
||||||
│ │ ├── csrc/moe_ops_impl.cu
|
|
||||||
│ │ ├── csrc/ix_full_bridge_v2.cpp
|
|
||||||
│ │ ├── build_moe_bridge.sh
|
|
||||||
│ │ └── python/moe_dispatch.py, patch_moe_hot_path.py
|
|
||||||
│ └── prebuilt/corex-3.2.3-ivcore10/*.so
|
|
||||||
└── (没有 ex_engine/ 目录)
|
|
||||||
```
|
|
||||||
|
|
||||||
`$(dirname "$0")/../ex_engine` = `/workspace/ex_engine` → **不存在**
|
|
||||||
|
|
||||||
结果: ix_ops.py, patch_vllm_ops.py, ix_startup_patch.py 全部没部署。
|
|
||||||
xllm_activation/norm/rope/cache.so 虽然被cp到$VLLM_ROOT/但没有Python层调用它们。
|
|
||||||
|
|
||||||
### 断裂2: build_moe_bridge.sh 内部路径错误
|
|
||||||
|
|
||||||
build_moe_bridge.sh 第18-19行:
|
|
||||||
```bash
|
|
||||||
MOE_CU="${SCRIPT_DIR}/ex_engine/csrc/moe_ops_impl.cu"
|
|
||||||
BRIDGE_CPP="${SCRIPT_DIR}/ex_engine/csrc/ix_full_bridge_v2.cpp"
|
|
||||||
```
|
|
||||||
|
|
||||||
SCRIPT_DIR = `/workspace/qwen3_6_scripts/ex_engine_src`
|
|
||||||
实际路径 = `${SCRIPT_DIR}/csrc/moe_ops_impl.cu`(少了 `ex_engine/` 一层)
|
|
||||||
|
|
||||||
结果: ix_moe_bridge.so 编译失败 → _HAS_IX_FUSED_MOE=False → 7-step fused MoE pipeline 未启用
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 三、ex_engine/ 文件去重审计
|
|
||||||
|
|
||||||
### 重复实现的功能(同一功能多个文件)
|
|
||||||
|
|
||||||
**MoE topk softmax (5个文件做同一件事)**:
|
|
||||||
1. `xllm_kernels/cuda/moe/moe_topk_softmax_kernels.cuh` (866行) ← xllm上游原版
|
|
||||||
2. `csrc/moe/moe_topk_softmax_kernels.cuh` (855行) ← 几乎相同的拷贝
|
|
||||||
3. `csrc/factor_moe_topk_softmax.cu` (260行) ← 独立提取版
|
|
||||||
4. `csrc/moe/moe_topk_softmax_ext.cu` (55行) ← 另一个入口
|
|
||||||
5. `csrc/moe_topk_softmax_v3.cu` (143行) ← 又一个版本
|
|
||||||
6. prebuilt `corex_moe_topk_softmax.so` ← 已编译可用
|
|
||||||
7. prebuilt `xllm_moe.so` ← 也包含此功能
|
|
||||||
|
|
||||||
**MoE combine (3个文件)**:
|
|
||||||
1. `xllm_kernels/cuda/moe/moe_combine.cu` (105行) ← xllm上游原版
|
|
||||||
2. files_5 的 `factor_moe_combine.cu` ← 重写
|
|
||||||
3. prebuilt `xllm_moe.so` ← 已编译可用
|
|
||||||
|
|
||||||
**MoE compute_index (2个文件)**:
|
|
||||||
1. `xllm_kernels/cuda/moe/moe_compute_index.cu` (156行) ← xllm上游原版
|
|
||||||
2. files_5 的 `factor_moe_compute_index.cu` ← 重写
|
|
||||||
|
|
||||||
**MoE Python pipeline (3个文件)**:
|
|
||||||
1. `python/fused_moe_ilu.py` (205行)
|
|
||||||
2. `python/moe_dispatch.py` (171行)
|
|
||||||
3. files_5 的 `moe_pipeline.py` ← 重写
|
|
||||||
|
|
||||||
**Attention dispatch (2个文件)**:
|
|
||||||
1. `python/corex_fa2_dispatch.py` (231行)
|
|
||||||
2. files_5 的 `attn_dispatch.py` ← 重写
|
|
||||||
|
|
||||||
**C++ bridge (3个文件)**:
|
|
||||||
1. `csrc/ix_full_bridge.cpp` (90行) ← v1, 对应prebuilt ix_full_bridge.so
|
|
||||||
2. `csrc/ix_full_bridge_v2.cpp` (387行) ← v2, 没编译
|
|
||||||
3. `csrc/ix_moe_bridge.cpp` (261行) ← MoE专用, 没编译
|
|
||||||
|
|
||||||
**ILU层 fused_moe (2处)**:
|
|
||||||
1. `xllm_layers/ilu/fused_moe.cpp` (797行) ← xllm上游
|
|
||||||
2. `csrc/ilu_layer_fused_moe.cpp` (797行) ← 拷贝
|
|
||||||
|
|
||||||
### 上游cat但未修改的文件
|
|
||||||
|
|
||||||
| 目录 | 文件数 | 来源 |
|
|
||||||
|---|---|---|
|
|
||||||
| xllm_kernels/ilu/*.cpp | 7 | xllm上游ILU kernel接口 |
|
|
||||||
| xllm_layers/ilu/*.cpp | 2 | xllm上游ILU layer |
|
|
||||||
| xllm_layers/common/*.cpp | 5 | xllm上游common layer |
|
|
||||||
| xllm_layers/npu_torch/*.cpp | 10 | xllm上游NPU实现(不适用BI-V100) |
|
|
||||||
| xllm_layers/mlu/*.cpp | 4 | xllm上游MLU实现(不适用BI-V100) |
|
|
||||||
| xllm_models/*.h | 6 | xllm上游model定义 |
|
|
||||||
| moe/*.py | 14 | ds_vllm上游MoE模块 |
|
|
||||||
| fla_kernels/ | 6 | FLA库GDN kernel(Triton,BI-V100不能跑) |
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 四、真机验证状态
|
|
||||||
|
|
||||||
| 组件 | commit | 真机结果 |
|
|
||||||
|---|---|---|
|
|
||||||
| gemm_grouped.cu (CUTLASS Cu10 TN) | cdcf1150 | ✓ err=0.000015 PASS, 1.97x vs torch.mm |
|
|
||||||
| corex_batched_gemm_kernel.cu | (sub 655用过) | ✓ decode单token 2.462ms |
|
|
||||||
| 16个prebuilt .so | (sub 694在用) | ✓ 正常加载 |
|
|
||||||
| ix_moe_bridge.so | 未编译 | ✗ 路径断裂 |
|
|
||||||
| ix_full_bridge_v2.so | 未编译 | ✗ 未进Docker |
|
|
||||||
| xllm_moe.so 的 fused_topk | (sub 694在用) | ✓ 正常工作 |
|
|
||||||
@@ -18,26 +18,17 @@ echo "[moe_bridge] Building ix_moe_bridge.so"
|
|||||||
echo "[moe_bridge] Script dir: ${SCRIPT_DIR}"
|
echo "[moe_bridge] Script dir: ${SCRIPT_DIR}"
|
||||||
|
|
||||||
# --- Locate sources ---
|
# --- Locate sources ---
|
||||||
# Support both layouts:
|
MOE_CU="${SCRIPT_DIR}/ex_engine/csrc/moe_ops_impl.cu"
|
||||||
# 1. SCRIPT_DIR=/workspace/ex_engine → csrc/ is direct child
|
BRIDGE_CPP="${SCRIPT_DIR}/ex_engine/csrc/ix_full_bridge_v2.cpp"
|
||||||
# 2. SCRIPT_DIR=/workspace/qwen3_6_scripts/ex_engine_src → csrc/ is direct child
|
|
||||||
MOE_CU=""
|
|
||||||
BRIDGE_CPP=""
|
|
||||||
for base in "${SCRIPT_DIR}" "${SCRIPT_DIR}/ex_engine"; do
|
|
||||||
[[ -f "${base}/csrc/moe_ops_impl.cu" ]] && MOE_CU="${base}/csrc/moe_ops_impl.cu"
|
|
||||||
[[ -f "${base}/csrc/ix_full_bridge_v2.cpp" ]] && BRIDGE_CPP="${base}/csrc/ix_full_bridge_v2.cpp"
|
|
||||||
done
|
|
||||||
|
|
||||||
if [[ -z "$MOE_CU" ]]; then
|
if [[ ! -f "$MOE_CU" ]]; then
|
||||||
echo "[moe_bridge] ERROR: moe_ops_impl.cu not found under ${SCRIPT_DIR}" >&2
|
echo "[moe_bridge] ERROR: $MOE_CU not found" >&2
|
||||||
exit 1
|
exit 1
|
||||||
fi
|
fi
|
||||||
if [[ -z "$BRIDGE_CPP" ]]; then
|
if [[ ! -f "$BRIDGE_CPP" ]]; then
|
||||||
echo "[moe_bridge] ERROR: ix_full_bridge_v2.cpp not found under ${SCRIPT_DIR}" >&2
|
echo "[moe_bridge] ERROR: $BRIDGE_CPP not found" >&2
|
||||||
exit 1
|
exit 1
|
||||||
fi
|
fi
|
||||||
echo "[moe_bridge] MOE_CU: ${MOE_CU}"
|
|
||||||
echo "[moe_bridge] BRIDGE_CPP: ${BRIDGE_CPP}"
|
|
||||||
|
|
||||||
# --- Locate libraries ---
|
# --- Locate libraries ---
|
||||||
COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
|
COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
|
||||||
@@ -75,28 +66,14 @@ echo "[moe_bridge] ixformer .so count: ${#IX_SO_FILES[@]}"
|
|||||||
# --- Build via torch.utils.cpp_extension ---
|
# --- Build via torch.utils.cpp_extension ---
|
||||||
mkdir -p "${SCRIPT_DIR}/prebuilt"
|
mkdir -p "${SCRIPT_DIR}/prebuilt"
|
||||||
|
|
||||||
export SCRIPT_DIR VLLM_ROOT
|
|
||||||
python3 << 'PYEOF'
|
python3 << 'PYEOF'
|
||||||
import os, sys, glob, shutil
|
import os, sys, glob, shutil
|
||||||
|
|
||||||
script_dir = os.environ.get("SCRIPT_DIR", ".")
|
script_dir = os.environ.get("SCRIPT_DIR", ".")
|
||||||
vllm_root = os.environ.get("VLLM_ROOT", "")
|
vllm_root = os.environ.get("VLLM_ROOT", "")
|
||||||
|
|
||||||
# Find source files — try direct csrc/ first, then ex_engine/csrc/
|
moe_cu = os.path.join(script_dir, "ex_engine", "csrc", "moe_ops_impl.cu")
|
||||||
moe_cu = ""
|
bridge_cpp = os.path.join(script_dir, "ex_engine", "csrc", "ix_full_bridge_v2.cpp")
|
||||||
bridge_cpp = ""
|
|
||||||
for base in [script_dir, os.path.join(script_dir, "ex_engine")]:
|
|
||||||
candidate_cu = os.path.join(base, "csrc", "moe_ops_impl.cu")
|
|
||||||
candidate_cpp = os.path.join(base, "csrc", "ix_full_bridge_v2.cpp")
|
|
||||||
if os.path.isfile(candidate_cu):
|
|
||||||
moe_cu = candidate_cu
|
|
||||||
if os.path.isfile(candidate_cpp):
|
|
||||||
bridge_cpp = candidate_cpp
|
|
||||||
if not moe_cu or not bridge_cpp:
|
|
||||||
print(f"[moe_bridge] ERROR: sources not found under {script_dir}")
|
|
||||||
sys.exit(1)
|
|
||||||
print(f"[moe_bridge] MOE_CU: {moe_cu}")
|
|
||||||
print(f"[moe_bridge] BRIDGE_CPP: {bridge_cpp}")
|
|
||||||
|
|
||||||
# Collect linker flags
|
# Collect linker flags
|
||||||
extra_ldflags = []
|
extra_ldflags = []
|
||||||
|
|||||||
@@ -1,179 +0,0 @@
|
|||||||
#!/usr/bin/env bash
|
|
||||||
# build_moe_bridge.sh — Compile MoE ops + bridge into ix_moe_bridge.so
|
|
||||||
#
|
|
||||||
# Links against:
|
|
||||||
# libcuinfer.so (cuinferCustomGemm, cuinferTopK — confirmed in symbol dump)
|
|
||||||
# libixformer.so (silu_and_mul, rms_norm, flash_attn, etc — confirmed)
|
|
||||||
#
|
|
||||||
# Real device compiler: corex clang/16, NOT nvcc
|
|
||||||
# Reference: ex_engine/build_ix_bridge.sh
|
|
||||||
|
|
||||||
set -euo pipefail
|
|
||||||
|
|
||||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
||||||
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
|
||||||
VLLM_ROOT="${1:-}"
|
|
||||||
|
|
||||||
echo "[moe_bridge] Building ix_moe_bridge.so"
|
|
||||||
echo "[moe_bridge] Script dir: ${SCRIPT_DIR}"
|
|
||||||
|
|
||||||
# --- Locate sources ---
|
|
||||||
# Support both layouts:
|
|
||||||
# 1. SCRIPT_DIR=/workspace/ex_engine → csrc/ is direct child
|
|
||||||
# 2. SCRIPT_DIR=/workspace/qwen3_6_scripts/ex_engine_src → csrc/ is direct child
|
|
||||||
MOE_CU=""
|
|
||||||
BRIDGE_CPP=""
|
|
||||||
for base in "${SCRIPT_DIR}" "${SCRIPT_DIR}/ex_engine"; do
|
|
||||||
[[ -f "${base}/csrc/moe_ops_impl.cu" ]] && MOE_CU="${base}/csrc/moe_ops_impl.cu"
|
|
||||||
[[ -f "${base}/csrc/ix_full_bridge_v2.cpp" ]] && BRIDGE_CPP="${base}/csrc/ix_full_bridge_v2.cpp"
|
|
||||||
done
|
|
||||||
|
|
||||||
if [[ -z "$MOE_CU" ]]; then
|
|
||||||
echo "[moe_bridge] ERROR: moe_ops_impl.cu not found under ${SCRIPT_DIR}" >&2
|
|
||||||
exit 1
|
|
||||||
fi
|
|
||||||
if [[ -z "$BRIDGE_CPP" ]]; then
|
|
||||||
echo "[moe_bridge] ERROR: ix_full_bridge_v2.cpp not found under ${SCRIPT_DIR}" >&2
|
|
||||||
exit 1
|
|
||||||
fi
|
|
||||||
echo "[moe_bridge] MOE_CU: ${MOE_CU}"
|
|
||||||
echo "[moe_bridge] BRIDGE_CPP: ${BRIDGE_CPP}"
|
|
||||||
|
|
||||||
# --- Locate libraries ---
|
|
||||||
COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
|
|
||||||
|
|
||||||
# Find libcuinfer.so
|
|
||||||
CUINFER_SO=""
|
|
||||||
for d in "${COREX_ROOT}/lib64" "${COREX_ROOT}/lib" "/usr/lib64" "/usr/lib"; do
|
|
||||||
if [[ -f "${d}/libcuinfer.so" ]]; then
|
|
||||||
CUINFER_SO="${d}/libcuinfer.so"
|
|
||||||
break
|
|
||||||
fi
|
|
||||||
done
|
|
||||||
|
|
||||||
# Find libixformer.so and ixformer Python package
|
|
||||||
IX_LIB_DIR=""
|
|
||||||
IX_SO_FILES=()
|
|
||||||
for d in \
|
|
||||||
"${COREX_ROOT}/lib/python3/dist-packages/ixformer" \
|
|
||||||
"${COREX_ROOT}/lib64/python3/dist-packages/ixformer" \
|
|
||||||
"$(python3 -c 'import ixformer, os; print(os.path.dirname(ixformer.__file__))' 2>/dev/null || echo '')"; do
|
|
||||||
if [[ -d "$d" ]]; then
|
|
||||||
IX_LIB_DIR="$d"
|
|
||||||
while IFS= read -r so; do
|
|
||||||
IX_SO_FILES+=("$so")
|
|
||||||
done < <(find "$d" -name "*.so" -type f 2>/dev/null)
|
|
||||||
break
|
|
||||||
fi
|
|
||||||
done
|
|
||||||
|
|
||||||
echo "[moe_bridge] COREX_ROOT: ${COREX_ROOT}"
|
|
||||||
echo "[moe_bridge] cuinfer: ${CUINFER_SO:-NOT FOUND}"
|
|
||||||
echo "[moe_bridge] ixformer dir: ${IX_LIB_DIR:-NOT FOUND}"
|
|
||||||
echo "[moe_bridge] ixformer .so count: ${#IX_SO_FILES[@]}"
|
|
||||||
|
|
||||||
# --- Build via torch.utils.cpp_extension ---
|
|
||||||
mkdir -p "${SCRIPT_DIR}/prebuilt"
|
|
||||||
|
|
||||||
export SCRIPT_DIR VLLM_ROOT
|
|
||||||
python3 << 'PYEOF'
|
|
||||||
import os, sys, glob, shutil
|
|
||||||
|
|
||||||
script_dir = os.environ.get("SCRIPT_DIR", ".")
|
|
||||||
vllm_root = os.environ.get("VLLM_ROOT", "")
|
|
||||||
|
|
||||||
# Find source files — try direct csrc/ first, then ex_engine/csrc/
|
|
||||||
moe_cu = ""
|
|
||||||
bridge_cpp = ""
|
|
||||||
for base in [script_dir, os.path.join(script_dir, "ex_engine")]:
|
|
||||||
candidate_cu = os.path.join(base, "csrc", "moe_ops_impl.cu")
|
|
||||||
candidate_cpp = os.path.join(base, "csrc", "ix_full_bridge_v2.cpp")
|
|
||||||
if os.path.isfile(candidate_cu):
|
|
||||||
moe_cu = candidate_cu
|
|
||||||
if os.path.isfile(candidate_cpp):
|
|
||||||
bridge_cpp = candidate_cpp
|
|
||||||
if not moe_cu or not bridge_cpp:
|
|
||||||
print(f"[moe_bridge] ERROR: sources not found under {script_dir}")
|
|
||||||
sys.exit(1)
|
|
||||||
print(f"[moe_bridge] MOE_CU: {moe_cu}")
|
|
||||||
print(f"[moe_bridge] BRIDGE_CPP: {bridge_cpp}")
|
|
||||||
|
|
||||||
# Collect linker flags
|
|
||||||
extra_ldflags = []
|
|
||||||
rpath_dirs = set()
|
|
||||||
|
|
||||||
corex_root = os.environ.get("COREX_ROOT", "/usr/local/corex")
|
|
||||||
for search_dir in [
|
|
||||||
os.path.join(corex_root, "lib64"),
|
|
||||||
os.path.join(corex_root, "lib"),
|
|
||||||
]:
|
|
||||||
if os.path.isdir(search_dir):
|
|
||||||
rpath_dirs.add(search_dir)
|
|
||||||
for so in glob.glob(os.path.join(search_dir, "libcuinfer*.so*")):
|
|
||||||
extra_ldflags.append(so)
|
|
||||||
|
|
||||||
# ixformer .so files
|
|
||||||
try:
|
|
||||||
import ixformer
|
|
||||||
ix_dir = os.path.dirname(ixformer.__file__)
|
|
||||||
rpath_dirs.add(ix_dir)
|
|
||||||
for so in glob.glob(os.path.join(ix_dir, "*.so")):
|
|
||||||
extra_ldflags.append(so)
|
|
||||||
for so in glob.glob(os.path.join(ix_dir, "lib*.so")):
|
|
||||||
if so not in extra_ldflags:
|
|
||||||
extra_ldflags.append(so)
|
|
||||||
except ImportError:
|
|
||||||
# Search common paths
|
|
||||||
for d in [
|
|
||||||
os.path.join(corex_root, "lib", "python3", "dist-packages", "ixformer"),
|
|
||||||
os.path.join(corex_root, "lib64", "python3", "dist-packages", "ixformer"),
|
|
||||||
]:
|
|
||||||
if os.path.isdir(d):
|
|
||||||
rpath_dirs.add(d)
|
|
||||||
for so in glob.glob(os.path.join(d, "*.so")):
|
|
||||||
extra_ldflags.append(so)
|
|
||||||
|
|
||||||
for d in rpath_dirs:
|
|
||||||
extra_ldflags.append(f"-Wl,-rpath,{d}")
|
|
||||||
|
|
||||||
print(f"[moe_bridge] Linking against {len(extra_ldflags)} items")
|
|
||||||
for f in extra_ldflags[:10]:
|
|
||||||
print(f" {f}")
|
|
||||||
|
|
||||||
try:
|
|
||||||
from torch.utils.cpp_extension import load
|
|
||||||
|
|
||||||
mod = load(
|
|
||||||
name="ix_moe_bridge",
|
|
||||||
sources=[moe_cu, bridge_cpp],
|
|
||||||
extra_include_paths=[os.path.join(script_dir, "csrc")],
|
|
||||||
extra_cflags=["-O2", "-std=c++17"],
|
|
||||||
extra_cuda_cflags=["-O2", ],
|
|
||||||
extra_ldflags=extra_ldflags,
|
|
||||||
verbose=True,
|
|
||||||
)
|
|
||||||
print("[moe_bridge] ✓ Compilation successful")
|
|
||||||
|
|
||||||
# Find and copy the built .so
|
|
||||||
import importlib
|
|
||||||
spec = importlib.util.find_spec("ix_moe_bridge")
|
|
||||||
if spec and spec.origin:
|
|
||||||
dst = os.path.join(script_dir, "prebuilt", "ix_moe_bridge.so")
|
|
||||||
shutil.copy2(spec.origin, dst)
|
|
||||||
print(f"[moe_bridge] ✓ Saved to {dst}")
|
|
||||||
|
|
||||||
if vllm_root:
|
|
||||||
vllm_dst = os.path.join(vllm_root, "ex_engine", "ix_moe_bridge.so")
|
|
||||||
os.makedirs(os.path.dirname(vllm_dst), exist_ok=True)
|
|
||||||
shutil.copy2(spec.origin, vllm_dst)
|
|
||||||
print(f"[moe_bridge] ✓ Deployed to {vllm_dst}")
|
|
||||||
else:
|
|
||||||
print("[moe_bridge] ⚠ Could not locate compiled .so via importlib")
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
print(f"[moe_bridge] ERROR: {e}", file=sys.stderr)
|
|
||||||
import traceback; traceback.print_exc()
|
|
||||||
sys.exit(1)
|
|
||||||
PYEOF
|
|
||||||
|
|
||||||
echo "[moe_bridge] Done"
|
|
||||||
@@ -202,36 +202,19 @@ fi
|
|||||||
# --- Deploy ix_bridge Python integration layer --------------------------------
|
# --- Deploy ix_bridge Python integration layer --------------------------------
|
||||||
build_stage "deploying ix_bridge operator replacements"
|
build_stage "deploying ix_bridge operator replacements"
|
||||||
EX_ENGINE_DIR="$(cd "$(dirname "$0")/../ex_engine" 2>/dev/null && pwd || echo "")"
|
EX_ENGINE_DIR="$(cd "$(dirname "$0")/../ex_engine" 2>/dev/null && pwd || echo "")"
|
||||||
if [ -z "$EX_ENGINE_DIR" ] || [ ! -d "$EX_ENGINE_DIR/python" ]; then
|
if [ -z "$EX_ENGINE_DIR" ] || [ ! -d "$EX_ENGINE_DIR" ]; then
|
||||||
# Dockerfile puts ex_engine at /workspace/ex_engine
|
EX_ENGINE_DIR="$(cd "$(dirname "$0")" && pwd)/../ex_engine"
|
||||||
EX_ENGINE_DIR="/workspace/ex_engine"
|
|
||||||
fi
|
fi
|
||||||
|
|
||||||
if [ -d "$EX_ENGINE_DIR/python" ]; then
|
if [ -d "$EX_ENGINE_DIR/python" ]; then
|
||||||
# Create ex_engine package inside vllm with correct Python package structure
|
# Create ex_engine package inside vllm
|
||||||
mkdir -p "${VLLM_ROOT}/ex_engine/python"
|
|
||||||
mkdir -p "${VLLM_ROOT}/ex_engine/csrc"
|
mkdir -p "${VLLM_ROOT}/ex_engine/csrc"
|
||||||
|
echo '"""ex_engine — Algorithm factor replacement for BI-V100."""' > "${VLLM_ROOT}/ex_engine/__init__.py"
|
||||||
|
|
||||||
# __init__.py with re-exports so both import styles work:
|
# Deploy Python modules
|
||||||
# from ex_engine.python import ix_ops_dispatch (direct)
|
cp "$EX_ENGINE_DIR/python/ix_ops.py" "${VLLM_ROOT}/ex_engine/ix_ops.py"
|
||||||
# from vllm.ex_engine import ix_ops_dispatch (via re-export)
|
cp "$EX_ENGINE_DIR/python/patch_vllm_ops.py" "${VLLM_ROOT}/ex_engine/patch_vllm_ops.py"
|
||||||
cat > "${VLLM_ROOT}/ex_engine/__init__.py" << 'INIT_EOF'
|
echo "[patch_ops] deployed ix_ops.py + patch_vllm_ops.py → ${VLLM_ROOT}/ex_engine/"
|
||||||
"""ex_engine — Algorithm factor replacement for BI-V100."""
|
|
||||||
# Re-export python subpackage members at top level for backward compat
|
|
||||||
# Allows: from vllm.ex_engine import ix_ops_dispatch
|
|
||||||
try:
|
|
||||||
from ex_engine.python.ix_ops_dispatch import *
|
|
||||||
from ex_engine.python import ix_ops_dispatch
|
|
||||||
from ex_engine.python import ix_ops
|
|
||||||
from ex_engine.python import patch_vllm_ops
|
|
||||||
except ImportError:
|
|
||||||
pass
|
|
||||||
INIT_EOF
|
|
||||||
echo '"""ex_engine.python — dispatch and bridge modules."""' > "${VLLM_ROOT}/ex_engine/python/__init__.py"
|
|
||||||
|
|
||||||
# Deploy ALL Python modules
|
|
||||||
cp "$EX_ENGINE_DIR/python/"*.py "${VLLM_ROOT}/ex_engine/python/"
|
|
||||||
echo "[patch_ops] deployed $(ls -1 "${VLLM_ROOT}/ex_engine/python/"*.py | wc -l) modules → ${VLLM_ROOT}/ex_engine/python/"
|
|
||||||
|
|
||||||
# Deploy bridge C++ source for JIT fallback
|
# Deploy bridge C++ source for JIT fallback
|
||||||
for cpp in "$EX_ENGINE_DIR"/csrc/ix_full_bridge*.cpp "$EX_ENGINE_DIR"/csrc/ix_moe_bridge.cpp; do
|
for cpp in "$EX_ENGINE_DIR"/csrc/ix_full_bridge*.cpp "$EX_ENGINE_DIR"/csrc/ix_moe_bridge.cpp; do
|
||||||
@@ -246,7 +229,7 @@ import logging
|
|||||||
_logger = logging.getLogger("ix_startup_patch")
|
_logger = logging.getLogger("ix_startup_patch")
|
||||||
def apply():
|
def apply():
|
||||||
try:
|
try:
|
||||||
from vllm.ex_engine.python.patch_vllm_ops import apply_all_patches
|
from vllm.ex_engine.patch_vllm_ops import apply_all_patches
|
||||||
n = apply_all_patches()
|
n = apply_all_patches()
|
||||||
if n > 0:
|
if n > 0:
|
||||||
_logger.info("ix_startup_patch: %d patches applied", n)
|
_logger.info("ix_startup_patch: %d patches applied", n)
|
||||||
@@ -349,83 +332,22 @@ if b"max_completion_tokens" not in installed:
|
|||||||
raise SystemExit("protocol.py missing max_completion_tokens field")
|
raise SystemExit("protocol.py missing max_completion_tokens field")
|
||||||
PY
|
PY
|
||||||
|
|
||||||
build_stage "building CUTLASS grouped GEMM (gemm_grouped.so)"
|
|
||||||
if [[ -f "${EX_ENGINE_DIR}/build_gemm_grouped.sh" ]]; then
|
|
||||||
bash "${EX_ENGINE_DIR}/build_gemm_grouped.sh" 2>&1 || {
|
|
||||||
echo "[WARN] gemm_grouped build failed — will use torch.mm fallback"
|
|
||||||
}
|
|
||||||
# Deploy compiled .so if it exists
|
|
||||||
for so in "${EX_ENGINE_DIR}"/gemm_grouped.so "${EX_ENGINE_DIR}"/csrc/gemm_grouped.so; do
|
|
||||||
if [[ -f "$so" ]]; then
|
|
||||||
cp "$so" "${VLLM_ROOT}/gemm_grouped.so"
|
|
||||||
echo "[patch_ops] deployed gemm_grouped.so → ${VLLM_ROOT}/"
|
|
||||||
break
|
|
||||||
fi
|
|
||||||
done
|
|
||||||
fi
|
|
||||||
|
|
||||||
build_stage "building CUTLASS batched GEMM (corex_batched_gemm.so)"
|
|
||||||
if [[ -f "${EX_ENGINE_DIR}/xllm_kernels/cuda/corex_batched_gemm_kernel.cu" ]]; then
|
|
||||||
python3 << PYEOF
|
|
||||||
import os, sys, shutil
|
|
||||||
try:
|
|
||||||
from torch.utils.cpp_extension import load
|
|
||||||
ex = "${EX_ENGINE_DIR}"
|
|
||||||
cutlass_inc = ""
|
|
||||||
for d in ["/usr/local/corex-samples-3.2.3_x86_64/samples/cutlass/include",
|
|
||||||
"/usr/local/corex/include/cutlass", "/usr/include/cutlass"]:
|
|
||||||
if os.path.isdir(d):
|
|
||||||
cutlass_inc = d
|
|
||||||
break
|
|
||||||
if not cutlass_inc:
|
|
||||||
print("[batched_gemm] No cutlass headers — skip"); sys.exit(0)
|
|
||||||
mod = load(
|
|
||||||
name="corex_batched_gemm",
|
|
||||||
sources=[
|
|
||||||
os.path.join(ex, "xllm_kernels/cuda/corex_batched_gemm_kernel.cu"),
|
|
||||||
os.path.join(ex, "xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp"),
|
|
||||||
],
|
|
||||||
extra_include_paths=[cutlass_inc],
|
|
||||||
extra_cflags=["-O2", "-std=c++17"],
|
|
||||||
extra_cuda_cflags=["-O2", f"-I{cutlass_inc}"],
|
|
||||||
extra_ldflags=["/usr/local/corex/lib64/libcuinfer.so", "-Wl,-rpath,/usr/local/corex/lib64"],
|
|
||||||
verbose=False,
|
|
||||||
)
|
|
||||||
print("[batched_gemm] ✓ Compiled")
|
|
||||||
import importlib
|
|
||||||
spec = importlib.util.find_spec("corex_batched_gemm")
|
|
||||||
if spec and spec.origin:
|
|
||||||
shutil.copy2(spec.origin, "${VLLM_ROOT}/corex_batched_gemm.so")
|
|
||||||
print("[batched_gemm] ✓ Deployed to ${VLLM_ROOT}/")
|
|
||||||
except Exception as e:
|
|
||||||
print(f"[batched_gemm] WARN: {e}")
|
|
||||||
PYEOF
|
|
||||||
fi
|
|
||||||
|
|
||||||
build_stage "building MoE bridge (ix_moe_bridge.so)"
|
build_stage "building MoE bridge (ix_moe_bridge.so)"
|
||||||
if [[ -f "${EX_ENGINE_DIR}/csrc/ix_moe_bridge.cpp" ]]; then
|
if [[ -f "./ex_engine_src/build_moe_bridge.sh" ]]; then
|
||||||
SCRIPT_DIR="${EX_ENGINE_DIR}" bash "${EX_ENGINE_DIR}/build_moe_bridge.sh" "${VLLM_ROOT}" 2>&1 || {
|
bash ./ex_engine_src/build_moe_bridge.sh "${VLLM_ROOT}" 2>&1 || {
|
||||||
echo "[WARN] MoE bridge build failed — will use Python fallback"
|
echo "[WARN] MoE bridge build failed — will use Python fallback"
|
||||||
}
|
}
|
||||||
# Deploy .so to all paths ix_fused_moe.py searches
|
|
||||||
for src in "${VLLM_ROOT}/ex_engine/ix_moe_bridge.so" \
|
|
||||||
"${EX_ENGINE_DIR}/prebuilt/ix_moe_bridge.so"; do
|
|
||||||
if [[ -f "$src" ]]; then
|
|
||||||
cp "$src" "${VLLM_ROOT}/ix_moe_bridge.so" 2>/dev/null || true
|
|
||||||
cp "$src" "${VLLM_ROOT}/model_executor/models/ix_moe_bridge.so" 2>/dev/null || true
|
|
||||||
echo "[patch_ops] deployed ix_moe_bridge.so to vllm search paths"
|
|
||||||
break
|
|
||||||
fi
|
|
||||||
done
|
|
||||||
fi
|
fi
|
||||||
|
|
||||||
build_stage "deploying all ex_engine Python modules"
|
build_stage "deploying MoE dispatch modules"
|
||||||
EX_PY_DIR="${VLLM_ROOT}/ex_engine/python"
|
EX_DIR="${VLLM_ROOT}/ex_engine/python"
|
||||||
mkdir -p "${EX_PY_DIR}"
|
mkdir -p "${EX_DIR}"
|
||||||
if [[ -d "${EX_ENGINE_DIR}/python" ]]; then
|
for pyfile in moe_dispatch.py patch_moe_hot_path.py; do
|
||||||
cp "${EX_ENGINE_DIR}/python/"*.py "${EX_PY_DIR}/" 2>/dev/null
|
if [[ -f "./ex_engine_src/python/${pyfile}" ]]; then
|
||||||
echo "[patch_ops] deployed $(ls -1 "${EX_PY_DIR}"/*.py 2>/dev/null | wc -l) Python modules → ${EX_PY_DIR}/"
|
cp "./ex_engine_src/python/${pyfile}" "${EX_DIR}/${pyfile}"
|
||||||
fi
|
echo " ✓ ${pyfile}"
|
||||||
|
fi
|
||||||
|
done
|
||||||
|
|
||||||
build_stage "compiling submission Python sources"
|
build_stage "compiling submission Python sources"
|
||||||
find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile
|
find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile
|
||||||
|
|||||||
@@ -143,11 +143,6 @@ try:
|
|||||||
except ImportError:
|
except ImportError:
|
||||||
_corex_batched_gemm = None
|
_corex_batched_gemm = None
|
||||||
|
|
||||||
try:
|
|
||||||
from vllm import gemm_grouped as _gemm_grouped
|
|
||||||
except ImportError:
|
|
||||||
_gemm_grouped = None
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from vllm import corex_moe_topk_softmax as _corex_moe_topk_softmax
|
from vllm import corex_moe_topk_softmax as _corex_moe_topk_softmax
|
||||||
except ImportError:
|
except ImportError:
|
||||||
@@ -217,11 +212,6 @@ _USE_COREX_MOE_DIRECT_ROUTED = (
|
|||||||
_USE_COREX_BATCHED_GEMM = (
|
_USE_COREX_BATCHED_GEMM = (
|
||||||
_corex_batched_gemm is not None
|
_corex_batched_gemm is not None
|
||||||
and env_bool("BI100_MOE_BATCHED_GEMM", True))
|
and env_bool("BI100_MOE_BATCHED_GEMM", True))
|
||||||
_USE_GEMM_GROUPED = (
|
|
||||||
_gemm_grouped is not None
|
|
||||||
and env_bool("BI100_MOE_GEMM_GROUPED", True))
|
|
||||||
if _USE_GEMM_GROUPED:
|
|
||||||
logger.info("gemm_grouped ENABLED — CUTLASS Cu10 grouped GEMM for MoE prefill")
|
|
||||||
_USE_COREX_MOE_TOPK_SOFTMAX = (
|
_USE_COREX_MOE_TOPK_SOFTMAX = (
|
||||||
_corex_moe_topk_softmax is not None
|
_corex_moe_topk_softmax is not None
|
||||||
and env_bool("BI100_MOE_COREX_TOPK_SOFTMAX", True))
|
and env_bool("BI100_MOE_COREX_TOPK_SOFTMAX", True))
|
||||||
@@ -1894,46 +1884,21 @@ class Qwen3_5MoeSparseBlock(nn.Module):
|
|||||||
expert_counts = torch.bincount(
|
expert_counts = torch.bincount(
|
||||||
flat_eids, minlength=w13.shape[0]).tolist()
|
flat_eids, minlength=w13.shape[0]).tolist()
|
||||||
|
|
||||||
# --- CUTLASS grouped GEMM path (replaces per-expert F.linear loop) ---
|
start = 0
|
||||||
if _USE_GEMM_GROUPED and hidden_states.dtype == torch.float16:
|
for eid, count in enumerate(expert_counts):
|
||||||
# Sort tokens into expert order
|
end = start + count
|
||||||
sorted_hidden = hidden_states[sorted_tok_ids] # (T*topk, H)
|
if count == 0:
|
||||||
expert_counts_t = torch.tensor(
|
|
||||||
expert_counts, dtype=torch.int32,
|
|
||||||
device=hidden_states.device) if not isinstance(
|
|
||||||
expert_counts, torch.Tensor) else expert_counts
|
|
||||||
|
|
||||||
# Step 4: grouped GEMM w13 (gate_proj + up_proj)
|
|
||||||
gemm1_out = _gemm_grouped.moe_group_gemm(
|
|
||||||
sorted_hidden, w13, expert_counts_t) # (T*topk, 2*I)
|
|
||||||
gate, up = gemm1_out.chunk(2, dim=-1)
|
|
||||||
act_out = F.silu(gate) * up # (T*topk, I)
|
|
||||||
|
|
||||||
# Step 6: grouped GEMM w2 (down_proj)
|
|
||||||
gemm2_out = _gemm_grouped.moe_group_gemm(
|
|
||||||
act_out, w2, expert_counts_t) # (T*topk, H)
|
|
||||||
|
|
||||||
# Step 7: weighted combine back to token order
|
|
||||||
flat_weights = sorted_weights.unsqueeze(-1) # (T*topk, 1)
|
|
||||||
weighted = (gemm2_out * flat_weights).to(out.dtype)
|
|
||||||
out.index_add_(0, sorted_tok_ids, weighted)
|
|
||||||
else:
|
|
||||||
# Fallback: per-expert F.linear loop
|
|
||||||
start = 0
|
|
||||||
for eid, count in enumerate(expert_counts):
|
|
||||||
end = start + count
|
|
||||||
if count == 0:
|
|
||||||
start = end
|
|
||||||
continue
|
|
||||||
tok_ids = sorted_tok_ids[start:end]
|
|
||||||
tokens = hidden_states[tok_ids] # (n, H)
|
|
||||||
gate_up = F.linear(tokens, w13[eid]) # (n, 2*I)
|
|
||||||
gate, up = gate_up.chunk(2, dim=-1)
|
|
||||||
act = F.silu(gate) * up # (n, I)
|
|
||||||
expert_out = F.linear(act, w2[eid]) # (n, H)
|
|
||||||
weights = sorted_weights[start:end].unsqueeze(-1)
|
|
||||||
out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype))
|
|
||||||
start = end
|
start = end
|
||||||
|
continue
|
||||||
|
tok_ids = sorted_tok_ids[start:end]
|
||||||
|
tokens = hidden_states[tok_ids] # (n, H)
|
||||||
|
gate_up = F.linear(tokens, w13[eid]) # (n, 2*I)
|
||||||
|
gate, up = gate_up.chunk(2, dim=-1)
|
||||||
|
act = F.silu(gate) * up # (n, I)
|
||||||
|
expert_out = F.linear(act, w2[eid]) # (n, H)
|
||||||
|
weights = sorted_weights[start:end].unsqueeze(-1)
|
||||||
|
out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype))
|
||||||
|
start = end
|
||||||
|
|
||||||
return out # partial, all-reduce done in forward()
|
return out # partial, all-reduce done in forward()
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user