test: probe_so_import_chain.sh — 验证.so部署路径+import链+flag值+shape匹配

probe_so_import_chain.sh

@@ -0,0 +1,165 @@
+#!/bin/bash
+set -e
+
+echo "=== 1. .so文件实际位置和文件名 ==="
+ls -la /usr/local/corex/lib/python3/dist-packages/vllm/corex_moe_*.so 2>/dev/null
+ls -la /usr/local/corex/lib/python3/dist-packages/vllm/ix_*.so 2>/dev/null
+echo ""
+
+echo "=== 2. Python import路径 ==="
+python3 -c "
+import vllm, os
+vllm_dir = os.path.dirname(vllm.__file__)
+print('vllm.__file__:', vllm.__file__)
+print('vllm dir:', vllm_dir)
+# 列出vllm目录下所有.so
+for f in sorted(os.listdir(vllm_dir)):
+    if f.endswith('.so'):
+        print(f'  {f}')
+"
+
+echo ""
+echo "=== 3. 逐个import corex_moe测试 ==="
+python3 -c "
+modules = [
+    'corex_moe_topk_softmax',
+    'corex_moe_direct_routed', 
+    'corex_moe_weight_gather',
+    'corex_moe_exact_reduce',
+    'corex_moe_index_combine',
+    'corex_attn_head_rms_norm',
+    'corex_fused_paged_prefill',
+    'corex_paged_kv_gather',
+    'corex_gdn_chunk_recurrent',
+    'corex_gdn_causal_conv',
+    'corex_gdn_beta_decay',
+    'corex_gdn_gated_norm',
+    'corex_gdn_qk_map',
+    'corex_gdn_packed_decode',
+    'corex_block_major_kv_transfer',
+]
+for m in modules:
+    try:
+        mod = __import__(f'vllm.{m}', fromlist=[m])
+        fns = [x for x in dir(mod) if not x.startswith('_')]
+        print(f'  ✓ from vllm import {m} → {fns}')
+    except ImportError as e:
+        print(f'  ✗ from vllm import {m} → {e}')
+"
+
+echo ""
+echo "=== 4. ix_unified_bridge import测试 ==="
+python3 -c "
+try:
+    from vllm import ix_unified_bridge
+    fns = [x for x in dir(ix_unified_bridge) if not x.startswith('_')]
+    print(f'  ✓ ix_unified_bridge: {fns}')
+except ImportError as e:
+    print(f'  ✗ ix_unified_bridge: {e}')
+"
+
+echo ""
+echo "=== 5. 我们的qwen3_5.py里各flag的实际值 ==="
+python3 -c "
+import sys, os
+# 模拟qwen3_5.py的import环境
+sys.path.insert(0, '/usr/local/corex/lib/python3/dist-packages')
+os.environ.setdefault('BI100_MOE_COREX_TOPK_SOFTMAX', '1')
+os.environ.setdefault('BI100_MOE_COREX_WEIGHT_GATHER', '1')
+os.environ.setdefault('BI100_MOE_COREX_DIRECT_ROUTED', '0')
+os.environ.setdefault('BI100_MOE_COREX_EXACT_REDUCE', '1')
+
+def env_bool(key, default):
+    v = os.environ.get(key, str(default))
+    return v.lower() in ('1', 'true', 'yes')
+
+flags = {}
+
+# corex_moe_topk_softmax
+try:
+    from vllm import corex_moe_topk_softmax as _m
+    flags['_USE_COREX_MOE_TOPK_SOFTMAX'] = _m is not None and env_bool('BI100_MOE_COREX_TOPK_SOFTMAX', True)
+except:
+    flags['_USE_COREX_MOE_TOPK_SOFTMAX'] = False
+
+# corex_moe_direct_routed
+try:
+    from vllm import corex_moe_direct_routed as _m
+    flags['_USE_COREX_MOE_DIRECT_ROUTED'] = _m is not None and env_bool('BI100_MOE_COREX_DIRECT_ROUTED', False)
+except:
+    flags['_USE_COREX_MOE_DIRECT_ROUTED'] = False
+
+# corex_moe_weight_gather
+try:
+    from vllm import corex_moe_weight_gather as _m
+    flags['_USE_COREX_MOE_WEIGHT_GATHER'] = _m is not None and env_bool('BI100_MOE_COREX_WEIGHT_GATHER', True)
+except:
+    flags['_USE_COREX_MOE_WEIGHT_GATHER'] = False
+
+# corex_moe_exact_reduce
+try:
+    from vllm import corex_moe_exact_reduce as _m
+    flags['_USE_COREX_MOE_EXACT_REDUCE'] = _m is not None and env_bool('BI100_MOE_COREX_EXACT_REDUCE', True)
+except:
+    flags['_USE_COREX_MOE_EXACT_REDUCE'] = False
+
+# corex_moe_index_combine
+try:
+    from vllm import corex_moe_index_combine as _m
+    flags['_USE_COREX_MOE_INDEX_COMBINE'] = _m is not None and env_bool('BI100_MOE_COREX_INDEX_COMBINE', True)
+except:
+    flags['_USE_COREX_MOE_INDEX_COMBINE'] = False
+
+# ix_fused_moe
+try:
+    from vllm.model_executor.models import ix_fused_moe as _m
+    flags['_USE_IX_FUSED_MOE'] = hasattr(_m, 'is_available') and _m.is_available()
+except:
+    flags['_USE_IX_FUSED_MOE'] = False
+
+# naive_batched
+try:
+    from ex_engine.moe.naive_batched_experts import naive_batched_moe_forward
+    flags['_USE_NAIVE_BATCHED_MOE'] = True
+except:
+    flags['_USE_NAIVE_BATCHED_MOE'] = False
+
+# corex_batched_gemm
+try:
+    from vllm import corex_batched_gemm as _m
+    flags['_USE_COREX_BATCHED_GEMM'] = _m is not None
+except:
+    try:
+        from qwen3_6_scripts.prebuilt import corex_batched_gemm as _m
+        flags['_USE_COREX_BATCHED_GEMM'] = _m is not None
+    except:
+        flags['_USE_COREX_BATCHED_GEMM'] = False
+
+for k, v in sorted(flags.items()):
+    status = '✓' if v else '✗'
+    print(f'  {status} {k} = {v}')
+"
+
+echo ""
+echo "=== 6. 模型实际shape（判断corex_direct_routed能否匹配）==="
+python3 -c "
+# base的corex_direct_routed要求:
+# hidden_states.shape == (1, 2048)
+# w13.shape == (256, 256, 2048)
+# w2.shape == (256, 2048, 128)
+# eids.shape == (8,) ws.shape == (8,)
+# 
+# Qwen3.5-27B的实际shape是什么?
+print('Qwen3.5-27B MoE config (from config.json):')
+print('  num_experts = 128 (per TP shard: 128/4=32? or 128?)')
+print('  top_k = 8')
+print('  hidden_size = 3584 (per TP shard: 3584/4=896? or 3584?)')
+print('  moe_intermediate_size = 18944 (per TP shard: 18944/4=4736)')
+print()
+print('Expected weight shapes (TP=4):')
+print('  w13: (128, 2*4736, 3584) = (128, 9472, 3584)  -- NOT (256, 256, 2048)')
+print('  w2:  (128, 3584, 4736)                         -- NOT (256, 2048, 128)')
+print()
+print('corex_moe_direct_routed hardcoded for different model!')
+print('We need corex_moe_weight_gather + F.linear path instead.')
+" 2>&1