From 6e22415a91f855752e1a35046255288fa4262597 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 15 Aug 2026 14:34:49 +0000 Subject: [PATCH 1/9] =?UTF-8?q?test:=20push=5Fprobe=5Fresults.sh=20?= =?UTF-8?q?=E2=80=94=20=E7=9C=9F=E6=9C=BAcommit=20probe=E7=BB=93=E6=9E=9C?= =?UTF-8?q?=E5=88=B0modelhub?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- push_probe_results.sh | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100755 push_probe_results.sh diff --git a/push_probe_results.sh b/push_probe_results.sh new file mode 100755 index 00000000..f996c14d --- /dev/null +++ b/push_probe_results.sh @@ -0,0 +1,30 @@ +#!/bin/bash +# 在真机上执行：把probe结果和.so文件commit到repo +set -e + +cd /home/dylan/0814/project_6 + +# 1. 先跑第二个probe（如果还没跑的话） +if [ ! -f probe_bridge_output.txt ]; then + echo "[1/4] Running probe_ix_unified_bridge.sh..." + bash probe_ix_unified_bridge.sh 2>&1 | tee probe_bridge_output.txt +else + echo "[1/4] probe_bridge_output.txt already exists" +fi + +# 2. commit probe结果（不commit .so文件，太大了） +echo "[2/4] Committing probe results..." +git add probe_bridge_output.txt +git add -f probe_output.txt 2>/dev/null || true +git commit -m "data: probe results — ixformer API + ix_unified_bridge + corex_*.so函数列表" || echo "nothing to commit" + +# 3. push到modelhub +echo "[3/4] Pushing to modelhub..." +git push origin main + +# 4. 提示中转机操作 +echo "" +echo "[4/4] 现在去中转机执行:" +echo " cd /home/dylan/Downloads/github_0804/project_6" +echo " git pull modelhub main" +echo " git push origin main" From 6a4459d405b63cb1ca4652f4f5cccd673bc8223d Mon Sep 17 00:00:00 2001 From: root Date: Sat, 15 Aug 2026 14:35:07 +0000 Subject: [PATCH 2/9] data: probe bridge output --- probe_bridge_output.txt | 1271 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 1271 insertions(+) create mode 100644 probe_bridge_output.txt diff --git a/probe_bridge_output.txt b/probe_bridge_output.txt new file mode 100644 index 00000000..b08fd70c --- /dev/null +++ b/probe_bridge_output.txt @@ -0,0 +1,1271 @@ +=== ix_unified_bridge.so 函数列表 === + /usr/local/corex/lib/python3/dist-packages/vllm/ix_unified_bridge.cpython-310-x86_64-linux-gnu.so: /usr/local/corex/lib/python3/dist-packages/vllm/ix_unified_bridge.cpython-310-x86_64-linux-gnu.so: undefined symbol: _ZN8ixformer5infer12silu_and_mulERN2at6TensorES3_ + /usr/local/corex/lib/python3/dist-packages/vllm/ix_unified_bridge.so: libc10.so: cannot open shared object file: No such file or directory + +=== ixformer 里跟vllm相关的函数签名 === +/usr/local/corex/lib64/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False, return_attn_probs=False) +fused_add_rms_norm(input: 'ixformer.Tensor', residual: 'ixformer.Tensor', weight: 'ixformer.Tensor', eps: float = 1e-05, scale: float = 1.0) +gemv(x: 'ixformer.Tensor', A: 'ixformer.Tensor') +linear(input: 'ixformer.Tensor', weight: 'ixformer.Tensor', bias: 'ixformer.Tensor' = None, output: 'ixformer.Tensor' = None) +matmul(input: 'ixformer.Tensor', other: 'ixformer.Tensor', *, out: 'ixformer.Tensor' = None, transa: bool = False, transb: bool = False, alpha: float = 1.0, beta: float = 0.0) -> 'ixformer.Tensor' +rms_norm(input: 'ixformer.Tensor', weight: 'ixformer.Tensor', output: 'ixformer.Tensor' = None, eps: float = 1e-06) +silu_and_mul(input: 'ixformer.Tensor', output: 'ixformer.Tensor' = None) +vllm_cache_ops_reshape_and_cache(key: 'ixformer.Tensor', value: 'ixformer.Tensor', key_cache: 'ixformer.Tensor', value_cache: 'ixformer.Tensor', slot_mapping: 'ixformer.Tensor') +vllm_copy_cache(key_caches, value_caches, block_mapping) +vllm_gptq_shuffle(qweights, g_idx=None) +vllm_llama_mlp(gate_up_proj_weight: 'torch.Tensor', down_proj_weight: 'torch.Tensor', hidden_size: int, intermediate_size: int, tp: int) -> None +vllm_rotary_embedding_neox(positions: 'ixformer.Tensor', query: 'ixformer.Tensor', key: 'ixformer.Tensor', head_size: int, cos_sin_cache: 'ixformer.Tensor', is_neox_style: bool = True) +vllm_single_query_cached_kv_attention(output: 'ixformer.Tensor', query: 'ixformer.Tensor', key_cache: 'ixformer.Tensor', value_cache: 'ixformer.Tensor', head_mapping: 'ixformer.Tensor', scale: float, block_tables: 'ixformer.Tensor', context_lens: 'ixformer.Tensor', block_size: int, max_context_len: int, alibi_slopes: 'ixformer.Tensor' = None, use_sqrt_alibi: bool = False) +vllm_single_query_cached_kv_attention_v2(output: 'ixformer.Tensor', partition: int, exp_sums: 'ixformer.Tensor', max_logits: 'ixformer.Tensor', temp_output: 'ixformer.Tensor', query: 'ixformer.Tensor', key_cache: 'ixformer.Tensor', value_cache: 'ixformer.Tensor', head_mapping: 'ixformer.Tensor', scale: float, block_tables: 'ixformer.Tensor', context_lens: 'ixformer.Tensor', block_size: int, max_context_len: int, alibi_slopes: 'ixformer.Tensor' = None, use_sqrt_alibi: bool = False) +vllm_smooth_dequant(output, x, scale=None, global_scale=1.0) +vllm_smooth_dequant_add_residual(output, x, residual, scale=None, global_scale=1.0) +vllm_smooth_dequant_fused_add_rms_norm_quant(output: 'ixformer.Tenosr', input: 'ixformer.Tensor', residual: 'ixformer.Tensor', weight: 'ixformer.Tensor', eps: float = 1e-05, scale: 'ixformer.Tensor' = None, global_scale: float = 1.0) +vllm_smooth_dequant_rotary_embedding_neox(positions: 'ixformer.Tensor', query: 'ixformer.Tensor', key: 'ixformer.Tensor', head_size: int, cos_sin_cache: 'ixformer.Tensor', query_out: 'ixformer.Tensor', key_out: 'ixformer.Tensor', query_scale: float, key_scale: float, is_neox_style: bool = True) +vllm_smooth_dequant_silu_and_mul_quant(output, input: 'ixformer.Tensor', gate_scale, up_scale, scale, temp=None) +vllm_smooth_fused_add_rms_norm_quant(output: 'ixformer.Tenosr', input: 'ixformer.Tensor', residual: 'ixformer.Tensor', weight: 'ixformer.Tensor', eps: float = 1e-05) +vllm_smooth_quant(output, x, scale) +vllm_smooth_rms_norm_quant(output: 'ixformer.Tensor', input: 'ixformer.Tensor', weight: 'ixformer.Tensor', eps: float = 1e-06) +vllm_swap_blocks(src: 'torch.Tensor', dst: 'torch.Tensor', mapping) + +=== corex_moe_topk_softmax.so 函数列表 === +functions (1): + moe_topk_softmax + +=== 所有corex_*.so的函数列表 === +corex_attn_head_rms_norm: ['apply_inverse', 'prepare'] +corex_block_major_kv_transfer: ['check_error', 'cpu_gather', 'cpu_scatter', 'pack', 'scatter'] +corex_fused_paged_prefill: ['forward'] +corex_gdn_beta_decay: ['beta_decay'] +corex_gdn_causal_conv: ['causal_conv_update'] +corex_gdn_chunk_recurrent: ['torch_chunk_gated_delta_rule', 'torch_recurrent_gated_delta_rule'] +corex_gdn_gated_norm: ['apply_inverse'] +corex_gdn_packed_decode: ['packed_decode'] +corex_gdn_qk_map: ['qk_map'] +corex_moe_direct_routed: ['w13', 'w2_reduce'] +corex_moe_exact_reduce: ['serial_float', 'serial_half', 'tree_float'] +corex_moe_index_combine: ['moe_combine_result', 'moe_compute_index'] +corex_moe_topk_softmax: ['moe_topk_softmax'] +corex_moe_weight_gather: ['gather'] +corex_paged_kv_gather: ['gather'] + +=== base镜像 _custom_ops.py 完整内容 === +import contextlib +import functools +from typing import TYPE_CHECKING, List, Optional, Tuple, Union, Dict, Any + +import torch +import torch.library + +import vllm.envs as envs +from vllm._core_ext import ScalarType +from vllm.logger import init_logger +from vllm.platforms import current_platform +# import ixformer.inference.functions as ops +import ixformer.functions as ixf_F +from ixformer.distributed import _distributed as cdist +import torch.nn.functional as F + +logger = init_logger(__name__) + +supports_moe_ops = True + +if TYPE_CHECKING: + + def register_fake(fn): + return lambda name: fn +else: + try: + from torch.library import register_fake + except ImportError: + try: + from torch.library import impl_abstract as register_fake + except: + def register_fake(fn): + return lambda name: fn + + +def hint_on_error(fn): + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + + except NotImplementedError as e: + msg = ( + "Error in calling custom op %s: %s\n" + "Not implemented or built, mostly likely because the current current device " + "does not support this kernel (less likely TORCH_CUDA_ARCH_LIST was set " + "incorrectly while building)") + logger.error(msg, fn.__name__, e) + raise NotImplementedError(msg % (fn.__name__, e)) from e + except AttributeError as e: + msg = ( + "Error in calling custom op %s: %s\n" + "Possibly you have built or installed an obsolete version of vllm.\n" + "Please try a clean build and install of vllm," + "or remove old built files such as vllm/*cpython*.so and build/ ." + ) + logger.error(msg, fn.__name__, e) + raise e + + return wrapper + + +# activation ops +def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None: + ixf_F.silu_and_mul(x, out) + + +def gelu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None: + ixf_F.gelu_and_mul(x, out) + + +def gelu_tanh_and_mul(out: torch.Tensor, x: torch.Tensor) -> None: + ixf_F.gelu_tanh_and_mul(x, out) + + +def gelu_fast(out: torch.Tensor, x: torch.Tensor) -> None: + out.copy_(F.gelu(x,approximate="tanh")) + return out + + +def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None: + out.copy_(F.gelu(x,approximate="tanh")) + return out + + +def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None: + out.copy_(F.gelu(x,approximate="tanh")) + return out + + + +def paged_attention_v1( + output, + query, + key_cache, + value_cache, + head_mapping, + scale, + block_tables, + context_lens, + block_size, + max_context_len, + alibi_slopes=None, + kv_cache_dtype=None, +): + return ixf_F.vllm_single_query_cached_kv_attention( + output, + query, + key_cache, + value_cache, + head_mapping, + scale, + block_tables, + context_lens, + block_size, + max_context_len, + alibi_slopes, + ) + + + +def paged_attention_v2( + out: torch.Tensor, + exp_sum: torch.Tensor, + max_logits: torch.Tensor, + tmp_out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: Optional[torch.Tensor], + kv_cache_dtype: str, + k_scale: float, + v_scale: float, + tp_rank: int = 0, + blocksparse_local_blocks: int = 0, + blocksparse_vert_stride: int = 0, + blocksparse_block_size: int = 64, + blocksparse_head_sliding_step: int = 0, +) -> None: + # CCCL two-pass dispatch pattern (dispatch_reduce.cuh): + # Pass 1: N CTAs each reduce their tile → d_block_reductions[N] + # Pass 2: 1 CTA reduces d_block_reductions[N] → d_out + # Our PyTorch V2 implementation follows the same pattern: + # Phase 1: partition attention (each partition = one tile) + # Phase 2: cross-partition log-sum-exp reduction (summary_statistics binary_op) + # paged_attention_v2_pytorch.py — try multiple import locations + # In docker: may be at /workspace/, next to vllm package, or in vllm/ itself + import sys, os + _pav2 = None + # Try 1: same package (patch_ops copies it next to _custom_ops.py) + try: + from vllm.paged_attention_v2_pytorch import paged_attention_v2_pytorch + _pav2 = paged_attention_v2_pytorch + except ImportError: + pass + # Try 2: /workspace/ (Dockerfile WORKDIR) + if _pav2 is None: + try: + _ws = '/workspace' + if _ws not in sys.path: + sys.path.insert(0, _ws) + from paged_attention_v2_pytorch import paged_attention_v2_pytorch + _pav2 = paged_attention_v2_pytorch + except ImportError: + pass + # Try 3: repo root relative to this file + if _pav2 is None: + _repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + from paged_attention_v2_pytorch import paged_attention_v2_pytorch + _pav2 = paged_attention_v2_pytorch + _pav2( + out, exp_sum, max_logits, tmp_out, + query, key_cache, value_cache, + num_kv_heads, scale, block_tables, seq_lens, + block_size, max_seq_len, alibi_slopes, + kv_cache_dtype, k_scale, v_scale, tp_rank, + blocksparse_local_blocks, blocksparse_vert_stride, + blocksparse_block_size, blocksparse_head_sliding_step, + ) + + +def paged_attention_rocm( + out: torch.Tensor, + exp_sum: torch.Tensor, + max_logits: torch.Tensor, + tmp_out: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: Optional[torch.Tensor], + kv_cache_dtype: str, + k_scale: float, + v_scale: float, +) -> None: + raise NotImplementedError() + + +# pos encoding ops +def rotary_embedding( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + is_neox: bool, +) -> None: + ixf_F.vllm_rotary_embedding_neox(positions, query, key, head_size, + cos_sin_cache, is_neox) + + +def batched_rotary_embedding(positions: torch.Tensor, query: torch.Tensor, + key: torch.Tensor, head_size: int, + cos_sin_cache: torch.Tensor, is_neox: bool, + rot_dim: int, + cos_sin_cache_offsets: torch.Tensor) -> None: + ixf_F.vllm_batched_rotary_embedding(positions, query, key, head_size, + cos_sin_cache, is_neox, rot_dim, + cos_sin_cache_offsets) + + +# layer norm ops +def rms_norm(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, + epsilon: float) -> None: + ixf_F.rms_norm(input, weight, out, epsilon) + + +def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor, + weight: torch.Tensor, epsilon: float, + residual_alpha: Optional[float] = 1) -> None: + ixf_F.fused_add_rms_norm(input, residual, weight, epsilon) + + +def advance_step_flashattn(num_seqs: int, num_queries: int, block_size: int, + input_tokens: torch.Tensor, + sampled_token_ids: torch.Tensor, + input_positions: torch.Tensor, + seq_lens: torch.Tensor, slot_mapping: torch.Tensor, + block_tables: torch.Tensor) -> None: + """Advance a step on GPU for existing inputs for a multi-step runner""" + return ixf_F.advance_step_flashattn(num_seqs, num_queries, block_size, + input_tokens, + sampled_token_ids, + input_positions, + seq_lens, slot_mapping, + block_tables) + + +def advance_step_flashinfer(num_seqs: int, num_queries: int, block_size: int, + input_tokens: torch.Tensor, + sampled_token_ids: torch.Tensor, + input_positions: torch.Tensor, + seq_lens: torch.Tensor, slot_mapping: torch.Tensor, + block_tables: torch.Tensor, + paged_kv_indices: torch.Tensor, + paged_kv_indptr: torch.Tensor, + paged_kv_last_page_len: torch.Tensor, + block_table_bound: torch.Tensor) -> None: + raise NotImplementedError("FIX SOON") + + +# quantization ops +# awq +def awq_dequantize(qweight: torch.Tensor, scales: torch.Tensor, + zeros: torch.Tensor, split_k_iters: int, thx: int, + thy: int) -> torch.Tensor: + raise NotImplementedError() + + +def awq_gemm(input: torch.Tensor, qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor, + pack_factor, group_size: int = 128) -> torch.Tensor: + return ixf_F.quantized_linear(input, qweight, scales,"awq",32 // pack_factor,qzeros=qzeros,group_size=group_size) + + +# gptq +def gptq_gemm(a: torch.Tensor, b_q_weight: torch.Tensor, + b_gptq_qzeros: torch.Tensor, b_gptq_scales: torch.Tensor, + b_g_idx: torch.Tensor, use_exllama: bool, + bit: int) -> torch.Tensor: + batch = a.shape[0] + if batch <= 8: + return ixf_F.quantized_linear(a,b_q_weight,b_gptq_scales,"gptq",4,b_gptq_qzeros,None,group_size=128) + o_dtype_str = "fp16" if a.dtype == torch.half else "bf16" + deq_w = ixf_F.quantized_weight_dequant(b_q_weight,b_gptq_scales,"gptq",o_dtype_str,4,b_gptq_qzeros,group_size=128) + return torch.matmul(a,deq_w) + + +if hasattr(torch.ops._C, "gptq_gemm"): + + @register_fake("_C::gptq_gemm") + def _gptq_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor, + b_gptq_qzeros: torch.Tensor, + b_gptq_scales: torch.Tensor, b_g_idx: torch.Tensor, + use_exllama: bool, bit: int) -> torch.Tensor: + return torch.empty((a.size(0), b_q_weight.size(1)), + dtype=a.dtype, + device=a.device) + + +def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, + bit: int) -> None: + return ixf_F.vllm_gptq_shuffle(q_weight,q_perm) + + +# marlin +def marlin_gemm(a: torch.Tensor, b_q_weight: torch.Tensor, + b_scales: torch.Tensor, workspace: torch.Tensor, size_m: int, + size_n: int, size_k: int) -> torch.Tensor: + raise NotImplementedError() + + +# marlin_24 +def gptq_marlin_24_gemm(a: torch.Tensor, b_q_weight: torch.Tensor, + b_meta: torch.Tensor, b_scales: torch.Tensor, + workspace: torch.Tensor, b_q_type: ScalarType, + size_m: int, size_n: int, size_k: int) -> torch.Tensor: + raise NotImplementedError() + + +if hasattr(torch.ops._C, "gptq_marlin_24_gemm"): + + @register_fake("_C::gptq_marlin_24_gemm") + def _gptq_marlin_24_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor, + b_meta: torch.Tensor, b_scales: torch.Tensor, + workspace: torch.Tensor, + b_q_type: ScalarType, size_m: int, + size_n: int, size_k: int) -> torch.Tensor: + return torch.empty((size_m, size_n), device=a.device, dtype=a.dtype) + + @register_fake("_C::gptq_marlin_gemm") + def _gptq_marlin_gemm_fake(a: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales: torch.Tensor, + b_zeros: torch.Tensor, + g_idx: torch.Tensor, + perm: torch.Tensor, + workspace: torch.Tensor, + b_q_type: ScalarType, + size_m: int, + size_n: int, + size_k: int, + is_k_full: bool, + has_zp: bool = False, + use_fp32_reduce: bool = False) -> torch.Tensor: + return torch.empty((size_m, size_n), device=a.device, dtype=a.dtype) + + @register_fake("_C::ggml_dequantize") + def _ggml_dequantize_fake(W: torch.Tensor, quant_type: int, m: int, + n: int) -> torch.Tensor: + return torch.empty((m, n), dtype=torch.float16, device=W.device) + + @register_fake("_C::ggml_mul_mat_vec_a8") + def _ggml_mul_mat_vec_a8_fake( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: int, + ) -> torch.Tensor: + return torch.empty((1, row), dtype=torch.float16, device=W.device) + + @register_fake("_C::ggml_mul_mat_a8") + def _ggml_mul_mat_a8_fake( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: int, + ) -> torch.Tensor: + batch = X.size(0) + return torch.empty((batch, row), dtype=torch.float16, device=W.device) + + @register_fake("_C::marlin_qqq_gemm") + def _marlin_qqq_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor, + s_tok: torch.Tensor, s_ch: torch.Tensor, + s_group: torch.Tensor, workspace: torch.Tensor, + size_m: int, size_n: int, + size_k: int) -> torch.Tensor: + return torch.empty((size_m, size_n), + dtype=torch.float16, + device=a.device) + + @register_fake("_C::marlin_gemm") + def _marlin_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor, + b_scales: torch.Tensor, workspace: torch.Tensor, + size_m: int, size_n: int, + size_k: int) -> torch.Tensor: + return torch.empty((size_m, size_n), + dtype=torch.float16, + device=a.device) + + @register_fake("_C::awq_dequantize") + def _awq_dequantize_fake(qweight: torch.Tensor, scales: torch.Tensor, + zeros: torch.Tensor, split_k_iters: int, thx: int, + thy: int) -> torch.Tensor: + in_c = qweight.size(0) + qout_c = qweight.size(1) + out_c = qout_c * 8 + return torch.empty((in_c, out_c), + dtype=scales.dtype, + device=scales.device) + + @register_fake("_C::awq_gemm") + def _awq_gemm_fake(input: torch.Tensor, qweight: torch.Tensor, + qzeros: torch.Tensor, scales: torch.Tensor, + split_k_iters: int) -> torch.Tensor: + num_in_feats = input.size(0) + return torch.empty((split_k_iters, num_in_feats, qweight.size(1) * 8), + dtype=input.dtype, + device=input.device).sum(0) + + @register_fake("_C::aqlm_gemm") + def _aqlm_gemm_fake(input: torch.Tensor, codes: torch.Tensor, + codebooks: torch.Tensor, scales: torch.Tensor, + codebook_partition_sizes: List[int], + bias: Optional[torch.Tensor]) -> torch.Tensor: + out_features = codes.size(0) * codebooks.size(2) + flat_input = input.reshape((-1, input.size(-1))) + flat_output = torch.empty((flat_input.size(0), out_features), + dtype=input.dtype, + device=input.device) + + output_sizes = list(input.shape) + output_sizes.pop() + output_sizes.append(-1) + return flat_output.reshape(tuple(output_sizes)) + + @register_fake("_C::aqlm_dequant") + def _aqlm_dequant_fake( + codes: torch.Tensor, codebooks: torch.Tensor, + codebook_partition_sizes: List[int]) -> torch.Tensor: + in_features = codes.size(1) * 8 + out_features = codes.size(0) + return torch.empty((out_features, in_features), + dtype=codebooks.dtype, + device=codebooks.device) + + @register_fake("_C::fp8_marlin_gemm") + def _fp8_marlin_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor, + b_scales: torch.Tensor, workspace: torch.Tensor, + num_bits: int, size_m: int, size_n: int, + size_k: int) -> torch.Tensor: + return torch.empty((size_m, size_n), dtype=a.dtype, device=a.device) + + @register_fake("_C::machete_gemm") + def machete_gemm_fake( + a: torch.Tensor, + # Should be the tensor returned by machete_prepack_B + b_q: torch.Tensor, + b_type: ScalarType, + b_scales: Optional[torch.Tensor] = None, + b_zeros: Optional[torch.Tensor] = None, + b_group_size: Optional[int] = None, + c: Optional[torch.Tensor] = None, + alpha: Optional[float] = None, + beta: Optional[float] = None, + schedule: Optional[str] = None, + ) -> torch.Tensor: + m = a.size(0) + n = b_q.size(1) + return torch.empty((m, n), device=a.device, dtype=a.dtype) + + @register_fake("_C::machete_prepack_B") + def machete_prepack_B_fake(b_q_weight: torch.Tensor, + b_type: ScalarType) -> torch.Tensor: + return torch.empty_like(b_q_weight, + memory_format=torch.contiguous_format) + + @register_fake("_C::causal_conv1d_fwd") + def causal_conv1d_fwd_fake(x: torch.Tensor, weight: torch.Tensor, + bias_: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + cu_seq_len: Optional[torch.Tensor], + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + silu_activation: bool) -> torch.Tensor: + return torch.empty_like(x) + + @register_fake("_C::causal_conv1d_update") + def causal_conv1d_update_fake( + x: torch.Tensor, conv_state: torch.Tensor, weight: torch.Tensor, + bias_: Optional[torch.Tensor], silu_activation: bool, + cache_seqlens: Optional[torch.Tensor], + conv_state_indices: Optional[torch.Tensor]) -> torch.Tensor: + return torch.empty_like(x) + + @register_fake("_C::selective_scan_fwd") + def selective_scan_fwd_fake(u: torch.Tensor, delta: torch.Tensor, + A: torch.Tensor, B: torch.Tensor, + C: torch.Tensor, D_: Optional[torch.Tensor], + z_: Optional[torch.Tensor], + delta_bias_: Optional[torch.Tensor], + delta_softplus: bool, + cu_seq_len: Optional[torch.Tensor], + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + ssm_states: Optional[torch.Tensor]) -> None: + return None + + +# cutlass +def cutlass_scaled_mm_supports_fp8(cuda_device_capability: int) -> bool: + return True + + +def cutlass_scaled_mm(a: torch.Tensor, + b: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + out_dtype: torch.dtype, + bias: Optional[torch.Tensor] = None) -> torch.Tensor: + + m = a.shape[0] + n = b.shape[1] + out = torch.empty((m, n), dtype=out_dtype, device=a.device) + ixf_F.w8a8(a, b.transpose(0,1), scale_a, scale_b, bias, output=out, out_dtype=out_dtype) + + return out + + +def cutlass_scaled_mm_azp(a: torch.Tensor, + b: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + out_dtype: torch.dtype, + azp_adj: torch.Tensor, + azp: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None) -> torch.Tensor: + raise NotImplementedError() + + +# aqlm +def aqlm_gemm(input: torch.Tensor, codes: torch.Tensor, + codebooks: torch.Tensor, scales: torch.Tensor, + codebook_partition_sizes: List[int], + bias: Optional[torch.Tensor]) -> torch.Tensor: + raise NotImplementedError() + + +def aqlm_dequant(codes: torch.Tensor, codebooks: torch.Tensor, + codebook_partition_sizes: List[int]) -> torch.Tensor: + raise NotImplementedError() + + +# gptq_marlin +def gptq_marlin_repack(b_q_weight: torch.Tensor, perm: torch.Tensor, + size_k: int, size_n: int, + num_bits: int) -> torch.Tensor: + raise NotImplementedError() + + +# gptq_marlin +def awq_marlin_repack(b_q_weight: torch.Tensor, size_k: int, size_n: int, + num_bits: int) -> torch.Tensor: + raise NotImplementedError() + + +def gptq_marlin_moe_repack(b_q_weight: torch.Tensor, perm: torch.Tensor, + size_k: int, size_n: int, + num_bits: int) -> torch.Tensor: + raise NotImplementedError() + + +def awq_marlin_moe_repack(b_q_weight: torch.Tensor, perm: torch.Tensor, + size_k: int, size_n: int, + num_bits: int) -> torch.Tensor: + num_experts = b_q_weight.shape[0] + assert size_k % 16 == 0 + output = torch.empty((num_experts, size_k // 16, size_n * (num_bits // 2)), + device=b_q_weight.device, + dtype=b_q_weight.dtype) + for e in range(num_experts): + output[e] = torch.ops._C.awq_marlin_repack(b_q_weight[e], size_k, + size_n, num_bits) + return output + + +def gptq_marlin_gemm(a: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales: torch.Tensor, + b_zeros: torch.Tensor, + g_idx: torch.Tensor, + perm: torch.Tensor, + workspace: torch.Tensor, + b_q_type: ScalarType, + size_m: int, + size_n: int, + size_k: int, + is_k_full: bool, + has_zp: bool = False, + use_fp32_reduce: bool = False) -> torch.Tensor: + raise NotImplementedError() + + +# fp8 marlin +def fp8_marlin_gemm(a: torch.Tensor, b_q_weight: torch.Tensor, + b_scales: torch.Tensor, workspace: torch.Tensor, + num_bits: int, size_m: int, size_n: int, + size_k: int) -> torch.Tensor: + raise NotImplementedError() + + +# machete +def machete_supported_schedules(b_type: ScalarType) -> List[str]: + raise NotImplementedError() + + +def machete_gemm( + a: torch.Tensor, + b_q: torch.Tensor, # Should be the tensor returned by machete_prepack_B + b_type: ScalarType, + b_scales: Optional[torch.Tensor] = None, + b_zeros: Optional[torch.Tensor] = None, + b_group_size: Optional[int] = None, + c: Optional[torch.Tensor] = None, + alpha: Optional[float] = None, + beta: Optional[float] = None, + schedule: Optional[str] = None, +) -> torch.Tensor: + raise NotImplementedError() + + +def machete_prepack_B(b_q_weight: torch.Tensor, + b_type: ScalarType) -> torch.Tensor: + raise NotImplementedError() + + +if hasattr(torch.ops._C, "permute_cols"): + + @register_fake("_C::permute_cols") + def _permute_cols_fake(a: torch.Tensor, + perm: torch.Tensor) -> torch.Tensor: + return torch.empty_like(a) + + +def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor: + raise NotImplementedError() + + +# fp8 +def scaled_fp8_quant( + input: torch.Tensor, + scale: Optional[torch.Tensor] = None, + num_token_padding: Optional[int] = None, + scale_ub: Optional[torch.Tensor] = None, + use_per_token_if_dynamic: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Quantize input tensor to FP8 and return quantized tensor and scale. + + This function supports both static and dynamic quantization: If you + provide the scale, it will use static scaling and if you omit it, + the scale will be determined dynamically. The function also allows + optional padding of the output tensors for downstream kernels that + will benefit from padding. + + Args: + input: The input tensor to be quantized to FP8 + scale: Optional scaling factor for the FP8 quantization + scale_ub: Optional upper bound for scaling factor in dynamic + per token case + num_token_padding: If specified, pad the first dimension + of the output to at least this value. + use_per_token_if_dynamic: Whether to do per_tensor or per_token + in the dynamic quantization case. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and + scaling factor. + """ + raise NotImplementedError() + + +# int8 +def scaled_int8_quant( + input: torch.Tensor, + scale: Optional[torch.Tensor] = None, + azp: Optional[torch.Tensor] = None, + symmetric: bool = True +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """ + Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp. + + Args: + input: The input tensor to be quantized to int8. + scale: Optional scaling factor for the int8 quantization. + When not provided, we invoke dynamic-per-token quantization. + azp: Optional zero-point for the int8 quantization. + Must be provided for asymmetric quantization if `scale` is provided. + symmetric: Whether to use symmetric quantization (scale only, azp ignored). + + Returns: + Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] : Output int8 tensor, scales, and optionally azp. + """ + output = torch.empty_like(input, dtype=torch.int8) + if scale is not None: + # static-per-tensor quantization. + assert symmetric == ( + azp is + None), "azp must only be provided for asymmetric quantization." + ixf_F.static_scaled_int8_quant(output, input, scale) + return output, scale, None + + # dynamic-per-token quantization. + input_scales = torch.empty((input.numel() // input.shape[-1], 1), + device=input.device, + dtype=torch.float32) + input_azp = None if symmetric else torch.empty_like(input_scales, + dtype=torch.int32) + ixf_F.dynamic_scaled_int8_quant(output, input, input_scales) + return output, input_scales, input_azp + + +# qqq ops +def marlin_qqq_gemm(a: torch.Tensor, b_q_weight: torch.Tensor, + s_tok: torch.Tensor, s_ch: torch.Tensor, + s_group: torch.Tensor, workspace: torch.Tensor, + size_m: int, size_n: int, size_k: int) -> torch.Tensor: + raise NotImplementedError() + + +# gguf +def ggml_dequantize(W: torch.Tensor, quant_type: int, m: int, + n: int) -> torch.Tensor: + raise NotImplementedError() + + +def ggml_mul_mat_vec_a8( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: int, +) -> torch.Tensor: + raise NotImplementedError() + + +def ggml_mul_mat_a8( + W: torch.Tensor, + X: torch.Tensor, + quant_type: int, + row: int, +) -> torch.Tensor: + raise NotImplementedError() + + +# mamba +def causal_conv1d_fwd(x: torch.Tensor, weight: torch.Tensor, + bias_: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: Optional[torch.Tensor], + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + silu_activation: bool) -> torch.Tensor: + raise NotImplementedError() + + +def causal_conv1d_update( + x: torch.Tensor, conv_state: torch.Tensor, weight: torch.Tensor, + bias_: Optional[torch.Tensor], silu_activation: bool, + cache_seqlens: Optional[torch.Tensor], + conv_state_indices: Optional[torch.Tensor]) -> torch.Tensor: + raise NotImplementedError() + + +def selective_scan_fwd( + u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor, B: torch.Tensor, + C: torch.Tensor, D_: Optional[torch.Tensor], + z_: Optional[torch.Tensor], delta_bias_: Optional[torch.Tensor], + delta_softplus: bool, query_start_loc: Optional[torch.Tensor], + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], ssm_states: torch.Tensor): + raise NotImplementedError() + + +# moe +def moe_align_block_size(topk_ids: torch.Tensor, num_experts: int, + block_size: int, sorted_token_ids: torch.Tensor, + experts_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor) -> None: + ixf_F.vllm_moe_align_block_size(topk_ids, num_experts, block_size, + sorted_token_ids, experts_ids, + num_tokens_post_pad) + + +def invoke_fused_moe_kernel( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + A_scale: Optional[torch.Tensor], + B_scale: Optional[torch.Tensor], + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + mul_routed_weight: bool, + top_k: int, + config: Dict[str, Any], + compute_type, + use_fp8_w8a8: bool, + use_int8_w8a16: bool, +) -> None: + ixf_F.vllm_invoke_fused_moe_kernel( + A, + B, + C, + topk_weights, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + mul_routed_weight, + top_k, + config['BLOCK_SIZE_M'] + ) + + +def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor, + token_expert_indicies: torch.Tensor, + gating_output: float) -> None: + ixf_F.vllm_moe_topk_softmax(topk_weights, topk_ids, + token_expert_indicies, gating_output) + + +if supports_moe_ops and hasattr(torch.ops._moe_C, "marlin_gemm_moe"): + + @register_fake("_moe_C::marlin_gemm_moe") + def marlin_gemm_moe_fake(a: torch.Tensor, b_q_weights: torch.Tensor, + sorted_ids: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, b_scales: torch.Tensor, + b_zero_points: torch.Tensor, g_idx: torch.Tensor, + perm: torch.Tensor, workspace: torch.Tensor, + b_q_type: ScalarType, size_m: int, size_n: int, + size_k: int, is_k_full: bool, num_experts: int, + topk: int, moe_block_size: int, + replicate_input: bool, + apply_weights: bool) -> torch.Tensor: + return torch.empty((size_m, topk, size_n), + dtype=a.dtype, + device=a.device) + + +def reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: float, + v_scale: float, +) -> None: + slot_mapping = slot_mapping.to(torch.int32) + ixf_F.vllm_cache_ops_reshape_and_cache(key, value, key_cache, + value_cache, slot_mapping) + + +def reshape_and_cache_flash( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: float, + v_scale: float, +) -> None: + ixf_F.reshape_and_cache_flash(key, value, key_cache, + value_cache, slot_mapping, + kv_cache_dtype, k_scale, + v_scale) + +def reshape_and_cache_flashinfer( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: float, # for fp8 + v_scale: float, # for fp8 + kv_cache_format: str = "NHD", + key_cache_scales: torch.Tensor = None, # for int8 + value_cache_scales: torch.Tensor = None, # for int8 +) -> None: + ixf_F.paged_attention_cache_appended( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_format, + key_cache_scales, + value_cache_scales, + ) + +def copy_blocks(key_caches: List[torch.Tensor], + value_caches: List[torch.Tensor], + block_mapping: torch.Tensor) -> None: + ixf_F.copy_blocks(key_caches, value_caches, block_mapping) + + +def swap_blocks(src: torch.Tensor, dst: torch.Tensor, + block_mapping: torch.Tensor) -> None: + # BI100 CoreX 3.2.3 exposes vllm_swap_blocks, while this vLLM build calls + # the newer swap_blocks name. Normalize the worker's CPU int64 [N, 2] + # tensor only for the legacy public API and fail fast on malformed maps. + native_swap_blocks = getattr(ixf_F, "swap_blocks", None) + if native_swap_blocks is not None: + native_swap_blocks(src, dst, block_mapping) + return + + vendor_swap_blocks = getattr(ixf_F, "vllm_swap_blocks", None) + if vendor_swap_blocks is None: + raise RuntimeError( + "ixformer exposes neither swap_blocks nor vllm_swap_blocks") + + if isinstance(block_mapping, torch.Tensor): + if block_mapping.device.type != "cpu": + raise ValueError("swap block mapping must be a CPU tensor") + if block_mapping.dtype != torch.int64: + raise ValueError("swap block mapping must use torch.int64") + if block_mapping.dim() != 2 or block_mapping.shape[1] != 2: + raise ValueError("swap block mapping must have shape [N, 2]") + pairs = block_mapping.tolist() + elif isinstance(block_mapping, dict): + pairs = list(block_mapping.items()) + else: + raise TypeError("swap block mapping must be a tensor or dict") + + normalized_mapping = {} + destinations = set() + for source, destination in pairs: + source = int(source) + destination = int(destination) + if source < 0 or destination < 0: + raise ValueError("swap block indices must be non-negative") + if source in normalized_mapping: + raise ValueError(f"duplicate swap source block: {source}") + if destination in destinations: + raise ValueError( + f"duplicate swap destination block: {destination}") + normalized_mapping[source] = destination + destinations.add(destination) + vendor_swap_blocks(src, dst, normalized_mapping) + + +def convert_fp8(output: torch.Tensor, + input: torch.Tensor, + scale: float = 1.0, + kv_dtype: str = "fp8") -> None: + raise NotImplementedError() + + +def get_device_attribute(attribute: int, device: int) -> int: + raise NotImplementedError() + + +def get_max_shared_memory_per_block_device_attribute(device: int) -> int: + # BI-V100 SMEM = 49152 bytes (48KB), confirmed via ixsmi + # Was incorrectly hardcoded to 32KB (32768), limiting Triton tile sizes + # and potentially constraining ixformer internal SMEM allocation. + return 49152 + + +# custom ar +def init_custom_ar(meta: torch.Tensor, rank_data: torch.Tensor, + handles: List[str], offsets: List[int], rank: int, + full_nvlink: bool) -> int: + raise NotImplementedError() + + +def should_custom_ar(inp: torch.Tensor, max_size: int, world_size: int, + full_nvlink: bool) -> bool: + raise NotImplementedError() + + +def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None: + raise NotImplementedError() + + +def all_reduce_unreg(fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, + out: torch.Tensor) -> None: + raise NotImplementedError() + + +def dispose(fa: int) -> None: + raise NotImplementedError() + + +def meta_size() -> int: + raise NotImplementedError() + + +def register_buffer(fa: int, t: torch.Tensor, handles: List[str], + offsets: List[int]) -> None: + raise NotImplementedError() + + +def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[str], List[int]]: + raise NotImplementedError() + + +def register_graph_buffers(fa: int, handles: List[str], + offsets: List[List[int]]) -> None: + raise NotImplementedError() + + +# Add our new features here.. + +# broadcast +class Async_helper(): + # For now, the comm and the other kernels are in the same stream, so we can remove the stream wait.. + def wait(self,): + return True + + +def broadcast(tensor, src=0, group=None, async_op=False): + cdist.broadcast(tensor,src,group,async_op=True) + if async_op: + return Async_helper() + else: + pass + +# w8a16 +def linear_w8a16(x: torch.Tensor, qweight: torch.Tensor, scales:torch.Tensor, + group_size: int = -1, format: str = "TN")-> torch.Tensor: + return ixf_F.w8a16(x, qweight, scales, format="TN", group_size=group_size) + + +## lora sgmv / bgmv +def sbgmv_expand(x: torch.Tensor, + w_t_all: torch.Tensor, + y: torch.Tensor, + b_seq_start_loc: torch.Tensor = None, + seq_len_tensor: torch.Tensor = None, + lora_indices_tensor: torch.Tensor = None, + batches: int = -1, + max_seq_length: int = -1, + token_nums: int = -1, + add_input=True, + ): + ''' + x: inputs + w_t_all: lora weight + y: output + + y += x@wt_t_all + ''' + assert x.dtype in [torch.float16, torch.bfloat16, torch.float32] + assert w_t_all.dtype in [ + torch.float16, + torch.bfloat16, + ] + + assert x.is_contiguous() + # assert y.is_contiguous() + if x.dtype == torch.float: + x = x.to(w_t_all.dtype) + + if w_t_all.ndim == 4: # shape:(lora_num,1,size,rank) + assert w_t_all.size(1) == 1 + w_t_all = w_t_all.squeeze(dim=1) + else: + assert w_t_all.ndim == 3 # shape:(lora_num,size,rank) + assert w_t_all.is_contiguous() + + assert add_input == True + + lora_indices = lora_indices_tensor.cpu().tolist() + lora_num = w_t_all.shape[0] + + ## 单一lora model, 且所有request均使用lora + if lora_num == 1 and all(x == lora_indices[0] for x in lora_indices): + if lora_indices[0] != -1: + w_t = w_t_all[0] + y += torch.matmul(x, w_t.t()) + ## 多个lora model + else: + ## prefill + if batches != -1: + for i, lora_id, start, seq_len in zip(range(batches), lora_indices, b_seq_start_loc, seq_len_tensor): + if lora_id != -1: + xi = x[start: start+seq_len] + w_t = w_t_all[lora_id] + y[start:start+seq_len] += (xi @ w_t.t()) + ## decode + else: + batches = x.shape[0] + for i, lora_id in zip(range(batches), lora_indices): + if lora_id != -1: + xi = x[i].unsqueeze(0) + w_t = w_t_all[lora_id] + y[i] += (xi @ w_t.t()).squeeze(0) + + return y + + +def sbgmv_shrink(x: torch.Tensor, + w_t_all: torch.Tensor, + y: torch.Tensor, + b_seq_start_loc: torch.Tensor = None, + seq_len_tensor: torch.Tensor = None, + lora_indices_tensor: torch.Tensor = None, + batches: int = -1, + max_seq_length: int = -1, + token_nums: int = -1, + scale: float = 1.0,): + """ + xx: inputs + w_t_all: lora weight + y: output + scale: float + + y = x@w_t_all * scale + """ + assert x.dtype == w_t_all.dtype + assert x.dtype in [torch.float16, torch.bfloat16] + assert x.is_contiguous() + assert y.is_contiguous() + + if w_t_all.ndim == 4: # shape:(lora_num,1,size,rank) + assert w_t_all.size(1) == 1 + w_t_all = w_t_all.squeeze(dim=1) + else: + assert w_t_all.ndim == 3 # shape:(lora_num,size,rank) + assert w_t_all.is_contiguous() + + lora_num = w_t_all.shape[0] + lora_indices = lora_indices_tensor.cpu().tolist() + + ## 单一lora model, 且所有request均使用lora + if lora_num == 1 and all(x == lora_indices[0] for x in lora_indices): + if lora_indices[0] != -1: + w_t = w_t_all[0] + y = torch.matmul(x, w_t.t()) * scale + ## 多个lora model + else: + ## prefill + if batches != -1: + for i, lora_id, start, seq_len in zip(range(batches), lora_indices, b_seq_start_loc, seq_len_tensor): + if lora_id != -1: + xi = x[start: start+seq_len] + w_t = w_t_all[lora_id] + y[start:start+seq_len] = (xi @ w_t.t())* scale + ## decode + else: + batches = x.shape[0] + for i, lora_id in zip(range(batches), lora_indices): + if lora_id != -1: + xi = x[i].unsqueeze(0) + w_t = w_t_all[lora_id] + y[i] = (xi @ w_t.t()).squeeze(0) * scale + + return y + +# temporary fix for https://github.com/vllm-project/vllm/issues/5456 +# TODO: remove this in v0.6.0 +names_and_values = globals() +names_and_values_to_update = {} +# prepare variables to avoid dict size change during iteration +k, v, arg = None, None, None +fn_type = type(lambda x: x) +for k, v in names_and_values.items(): + # find functions that are defined in this file and have torch.Tensor + # in their annotations. `arg == "torch.Tensor"` is used to handle + # the case when users use `import __annotations__` to turn type + # hints into strings. + if isinstance(v, fn_type) \ + and v.__code__.co_filename == __file__ \ + and any(arg is torch.Tensor or arg == "torch.Tensor" + for arg in v.__annotations__.values()): + names_and_values_to_update[k] = hint_on_error(v) + +names_and_values.update(names_and_values_to_update) +del names_and_values_to_update, names_and_values, v, k, fn_type +=== base镜像 qwen3_5.py MoE forward === +67:from vllm.model_executor.layers.fused_moe import FusedMoE +127: from vllm import corex_moe_exact_reduce as _corex_moe_exact_reduce +129: _corex_moe_exact_reduce = None +132: from vllm import corex_moe_weight_gather as _corex_moe_weight_gather +134: _corex_moe_weight_gather = None +137: from vllm import corex_moe_direct_routed as _corex_moe_direct_routed +139: _corex_moe_direct_routed = None +142: from vllm import corex_moe_topk_softmax as _corex_moe_topk_softmax +144: _corex_moe_topk_softmax = None +179: _corex_moe_exact_reduce is not None +182: _corex_moe_weight_gather is not None +185: _corex_moe_direct_routed is not None +188: _corex_moe_topk_softmax is not None +1525: FusedMoE is used ONLY for weight storage and loading (create_weights / +1527: ixformer on BI-V100 lacks vllm_moe_topk_softmax / vllm_invoke_fused_moe_kernel. +1554: # FusedMoE: only used for weight storage + weight_loader. +1555: # Forward is bypassed — see _pure_pytorch_experts(). +1556: self.experts = FusedMoE( +1598: def _pure_pytorch_experts( +1607: Output is partial (pre-all-reduce), same contract as FusedMoE +1611: # Source: xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh +1613: topk_weights, topk_ids = _corex_moe_topk_softmax.moe_topk_softmax( +1651: gate_up = _corex_moe_direct_routed.w13( +1654: return _corex_moe_direct_routed.w2_reduce( +1673: w13_sel, w2_sel = _corex_moe_weight_gather.gather( +1699: out = _corex_moe_exact_reduce.serial_float(expert_out, ws) +1740: routed_out = self._pure_pytorch_experts(hidden_states, router_logits) +2445: # Our FusedMoE stores: From b00429d81ca2f01ab7a0402b86dc498ea3c2ee12 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 15 Aug 2026 14:36:29 +0000 Subject: [PATCH 3/9] =?UTF-8?q?test:=20probe=5Fbase=5Fmoe=5Fforward.sh=20?= =?UTF-8?q?=E2=80=94=20cat=20base=20qwen3=5F5.py=E7=9A=84=E5=AE=8C?= =?UTF-8?q?=E6=95=B4MoE=20forward=20+=20=E6=89=80=E6=9C=89corex=5Fmoe=5F*.?= =?UTF-8?q?so=E7=AD=BE=E5=90=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- probe_base_moe_forward.sh | 90 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100755 probe_base_moe_forward.sh diff --git a/probe_base_moe_forward.sh b/probe_base_moe_forward.sh new file mode 100755 index 00000000..0cc8f483 --- /dev/null +++ b/probe_base_moe_forward.sh @@ -0,0 +1,90 @@ +#!/bin/bash +set -e + +BASE="/usr/local/corex/lib/python3/dist-packages/vllm/model_executor/models/qwen3_5.py" + +echo "=== base qwen3_5.py line count ===" +wc -l "$BASE" + +echo "" +echo "=== _pure_pytorch_experts 完整函数 ===" +sed -n '/def _pure_pytorch_experts/,/^ def [a-z]/p' "$BASE" | head -200 + +echo "" +echo "=== forward 中调用 _pure_pytorch_experts 的上下文 ===" +grep -n -B5 -A5 "_pure_pytorch_experts\|corex_moe_direct\|corex_moe_weight\|corex_moe_exact\|corex_moe_topk" "$BASE" | head -100 + +echo "" +echo "=== corex_moe_direct_routed.w13 签名 ===" +python3 -c " +from vllm import corex_moe_direct_routed as m +import inspect +for name in dir(m): + if not name.startswith('_'): + obj = getattr(m, name) + try: + sig = inspect.signature(obj) + print(f'{name}{sig}') + except: + print(f'{name}: {type(obj)}') +" 2>&1 + +echo "" +echo "=== corex_moe_topk_softmax.moe_topk_softmax 签名 ===" +python3 -c " +from vllm import corex_moe_topk_softmax as m +import inspect +for name in dir(m): + if not name.startswith('_'): + obj = getattr(m, name) + try: + sig = inspect.signature(obj) + print(f'{name}{sig}') + except: + print(f'{name}: {type(obj)}') +" 2>&1 + +echo "" +echo "=== corex_moe_exact_reduce 签名 ===" +python3 -c " +from vllm import corex_moe_exact_reduce as m +import inspect +for name in dir(m): + if not name.startswith('_'): + obj = getattr(m, name) + try: + sig = inspect.signature(obj) + print(f'{name}{sig}') + except: + print(f'{name}: {type(obj)}') +" 2>&1 + +echo "" +echo "=== corex_moe_weight_gather 签名 ===" +python3 -c " +from vllm import corex_moe_weight_gather as m +import inspect +for name in dir(m): + if not name.startswith('_'): + obj = getattr(m, name) + try: + sig = inspect.signature(obj) + print(f'{name}{sig}') + except: + print(f'{name}: {type(obj)}') +" 2>&1 + +echo "" +echo "=== corex_moe_index_combine 签名 ===" +python3 -c " +from vllm import corex_moe_index_combine as m +import inspect +for name in dir(m): + if not name.startswith('_'): + obj = getattr(m, name) + try: + sig = inspect.signature(obj) + print(f'{name}{sig}') + except: + print(f'{name}: {type(obj)}') +" 2>&1 From a6cc233880d0df01e9cf084dad0f17b41ad88083 Mon Sep 17 00:00:00 2001 From: root Date: Sat, 15 Aug 2026 14:39:03 +0000 Subject: [PATCH 4/9] =?UTF-8?q?data:=20base=20MoE=20forward=20+=20corex=5F?= =?UTF-8?q?moe=E7=AD=BE=E5=90=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- probe_moe_output.txt | 297 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 297 insertions(+) create mode 100644 probe_moe_output.txt diff --git a/probe_moe_output.txt b/probe_moe_output.txt new file mode 100644 index 00000000..d202b836 --- /dev/null +++ b/probe_moe_output.txt @@ -0,0 +1,297 @@ +=== base qwen3_5.py line count === +2628 /usr/local/corex/lib/python3/dist-packages/vllm/model_executor/models/qwen3_5.py + +=== _pure_pytorch_experts 完整函数 === + def _pure_pytorch_experts( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + """Pure-PyTorch MoE (ixformer has no MoE kernels on BI-V100). + + w13_weight: (num_experts, 2*inter_per_partition, hidden) [TP-sharded] + w2_weight: (num_experts, hidden, inter_per_partition) [TP-sharded] + Output is partial (pre-all-reduce), same contract as FusedMoE + with reduce_results=False. + """ + # Fused topk+softmax: single CUB kernel vs 2 PyTorch ops. + # Source: xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh + if _USE_COREX_MOE_TOPK_SOFTMAX: + topk_weights, topk_ids = _corex_moe_topk_softmax.moe_topk_softmax( + router_logits.float(), self.top_k, True) + topk_ids = topk_ids.to(torch.int64) + topk_weights = topk_weights.to(hidden_states.dtype) + else: + topk_logits, topk_ids = torch.topk( + router_logits.float(), self.top_k, dim=-1) # (T, top_k) + topk_weights = torch.softmax(topk_logits, dim=-1) + topk_weights = topk_weights.to(hidden_states.dtype) + + w13 = self.experts.w13_weight # (E, 2*I, H) + w2 = self.experts.w2_weight # (E, H, I) + + T = hidden_states.shape[0] + if T == 1: + # Fast path: single token (decode). + # Batched GEMM: replace top_k separate F.linear calls with 2 fused ops. + # gate_up: 1 large GEMM (1,H) × (K*2*I,H)^T → (1, K*2*I) + # down: 1 bmm (K,H,I) @ (K,I,1) → (K,H) + # Total: 3 kernel launches vs previous 16 (top_k*2). + eids = topk_ids[0] # (K,) + ws = topk_weights[0].to(hidden_states.dtype) # (K,) + use_corex_direct = ( + _USE_COREX_MOE_DIRECT_ROUTED + and hidden_states.dtype == torch.float16 + and w13.dtype == torch.float16 + and w2.dtype == torch.float16 + and ws.dtype == torch.float16 + and hidden_states.is_cuda and w13.is_cuda and w2.is_cuda + and eids.is_cuda and ws.is_cuda + and hidden_states.is_contiguous() + and w13.is_contiguous() and w2.is_contiguous() + and eids.is_contiguous() and ws.is_contiguous() + and hidden_states.shape == (1, 2048) + and w13.shape == (256, 256, 2048) + and w2.shape == (256, 2048, 128) + and eids.shape == (8,) and ws.shape == (8,)) + if use_corex_direct: + gate_up = _corex_moe_direct_routed.w13( + hidden_states, w13, eids) + act = self.act_fn(gate_up) + return _corex_moe_direct_routed.w2_reduce( + act, w2, eids, ws) + + use_corex_gather = ( + _USE_COREX_MOE_WEIGHT_GATHER + and hidden_states.dtype == torch.float16 + and w13.dtype == torch.float16 + and w2.dtype == torch.float16 + and w13.is_cuda and w2.is_cuda and eids.is_cuda + and w13.is_contiguous() and w2.is_contiguous() + and eids.is_contiguous() + and w13.dim() == 3 and w2.dim() == 3 + and eids.dim() == 1 and eids.numel() == 8 + and w13.shape[0] == w2.shape[0] + and w13.shape[2] == w2.shape[1] + and w13.shape[1] == 2 * w2.shape[2] + and w13.shape[1] * w13.shape[2] % 8 == 0 + and w2.shape[1] * w2.shape[2] % 8 == 0) + if use_corex_gather: + w13_sel, w2_sel = _corex_moe_weight_gather.gather( + w13, w2, eids) + else: + w13_sel = w13[eids] # (K, 2*I, H) + w2_sel = w2[eids] # (K, H, I) + + H = hidden_states.shape[-1] + + gate_up = F.linear( + hidden_states, + w13_sel.reshape(-1, H), # (K*2*I, H) — contiguous after indexing + ) # (1, K*2*I) + gate_up = gate_up.view(self.top_k, -1) # (K, 2*I) + if _USE_FUSED_MOE_ACTIVATION: + act = self.act_fn(gate_up) # (K, I) + else: + gate, up = gate_up.chunk(2, dim=-1) + act = F.silu(gate) * up + + # bmm: (K,H,I) @ (K,I,1) → (K,H,1) → (K,H) + expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H) + + if (_USE_COREX_MOE_EXACT_REDUCE + and expert_out.dtype == torch.float16 + and ws.dtype == torch.float16 + and expert_out.shape[0] == 8): + out = _corex_moe_exact_reduce.serial_float(expert_out, ws) + else: + out = (expert_out * ws.unsqueeze(-1)).sum( + 0, keepdim=True).to(hidden_states.dtype) # (1, H) + else: + # General path (prefill / multi-seq): group assignments once. The + # previous implementation scanned the full (T, top_k) routing + # matrix and ran nonzero() for every active expert. + out = torch.zeros_like(hidden_states) + flat_eids = topk_ids.reshape(-1) + order = torch.argsort(flat_eids, stable=True) + sorted_tok_ids = torch.arange( + T, device=topk_ids.device).repeat_interleave(self.top_k)[order] + sorted_weights = topk_weights.reshape(-1)[order] + expert_counts = torch.bincount( + flat_eids, minlength=w13.shape[0]).tolist() + + 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 + + return out # partial, all-reduce done in forward() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + +=== forward 中调用 _pure_pytorch_experts 的上下文 === +122- from vllm import corex_attn_head_rms_norm as _corex_attn_head_rms_norm +123-except ImportError: +124- _corex_attn_head_rms_norm = None +125- +126-try: +127: from vllm import corex_moe_exact_reduce as _corex_moe_exact_reduce +128-except ImportError: +129: _corex_moe_exact_reduce = None +130- +131-try: +132: from vllm import corex_moe_weight_gather as _corex_moe_weight_gather +133-except ImportError: +134: _corex_moe_weight_gather = None +135- +136-try: +137: from vllm import corex_moe_direct_routed as _corex_moe_direct_routed +138-except ImportError: +139: _corex_moe_direct_routed = None +140- +141-try: +142: from vllm import corex_moe_topk_softmax as _corex_moe_topk_softmax +143-except ImportError: +144: _corex_moe_topk_softmax = None +145- +146-from vllm.model_executor.models.interfaces import (HasInnerState, SupportsLoRA, +147- SupportsMultiModal) +148- +149-logger = init_logger(__name__) +-- +174- and env_bool("BI100_GDN_COREX_PACKED_DECODE", False)) +175-_USE_COREX_ATTN_HEAD_RMS_NORM = ( +176- _corex_attn_head_rms_norm is not None +177- and env_bool("BI100_ATTN_COREX_HEAD_RMS_NORM", True)) +178-_USE_COREX_MOE_EXACT_REDUCE = ( +179: _corex_moe_exact_reduce is not None +180- and env_bool("BI100_MOE_COREX_EXACT_REDUCE", True)) +181-_USE_COREX_MOE_WEIGHT_GATHER = ( +182: _corex_moe_weight_gather is not None +183- and env_bool("BI100_MOE_COREX_WEIGHT_GATHER", True)) +184-_USE_COREX_MOE_DIRECT_ROUTED = ( +185: _corex_moe_direct_routed is not None +186- and env_bool("BI100_MOE_COREX_DIRECT_ROUTED", False)) +187-_USE_COREX_MOE_TOPK_SOFTMAX = ( +188: _corex_moe_topk_softmax is not None +189- and env_bool("BI100_MOE_COREX_TOPK_SOFTMAX", True)) +190-_USE_FUSED_MOE_ACTIVATION = env_bool("BI100_MOE_FUSED_ACTIVATION", True) +191- +192- +193-# --------------------------------------------------------------------------- +-- +1550- bias=False, quant_config=quant_config) +1551- self.router_shared_gate.weight.weight_loader = \ +1552- self._router_shared_gate_weight_loader +1553- +1554- # FusedMoE: only used for weight storage + weight_loader. +1555: # Forward is bypassed — see _pure_pytorch_experts(). +1556- self.experts = FusedMoE( +1557- num_experts=text_cfg.num_experts, +1558- top_k=text_cfg.num_experts_per_tok, +1559- hidden_size=hidden_size, +1560- intermediate_size=text_cfg.moe_intermediate_size, +-- +1593- raise ValueError( +1594- "unexpected router/shared gate weight shape: " +1595- f"expected {expected}, got {tuple(loaded_weight.shape)}") +1596- param.data.narrow(0, offset, rows).copy_(loaded_weight) +1597- +1598: def _pure_pytorch_experts( +1599- self, +1600- hidden_states: torch.Tensor, +1601- router_logits: torch.Tensor, +1602- ) -> torch.Tensor: +1603- """Pure-PyTorch MoE (ixformer has no MoE kernels on BI-V100). +-- +1608- with reduce_results=False. +1609- """ +1610- # Fused topk+softmax: single CUB kernel vs 2 PyTorch ops. +1611- # Source: xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh +1612- if _USE_COREX_MOE_TOPK_SOFTMAX: +1613: topk_weights, topk_ids = _corex_moe_topk_softmax.moe_topk_softmax( +1614- router_logits.float(), self.top_k, True) +1615- topk_ids = topk_ids.to(torch.int64) +1616- topk_weights = topk_weights.to(hidden_states.dtype) +1617- else: +1618- topk_logits, topk_ids = torch.topk( +-- +1646- and hidden_states.shape == (1, 2048) +1647- and w13.shape == (256, 256, 2048) +1648- and w2.shape == (256, 2048, 128) +1649- and eids.shape == (8,) and ws.shape == (8,)) +1650- if use_corex_direct: +1651: gate_up = _corex_moe_direct_routed.w13( +1652- hidden_states, w13, eids) +1653- act = self.act_fn(gate_up) +1654: return _corex_moe_direct_routed.w2_reduce( +1655- act, w2, eids, ws) +1656- +1657- use_corex_gather = ( +1658- _USE_COREX_MOE_WEIGHT_GATHER +1659- and hidden_states.dtype == torch.float16 + +=== corex_moe_direct_routed.w13 签名 === +/usr/local/corex/lib64/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 08-15 14:38:09 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +2026-08-15 14:38:10.835442: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-08-15 14:38:10.887465: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. +w13: +w2_reduce: + +=== corex_moe_topk_softmax.moe_topk_softmax 签名 === +/usr/local/corex/lib64/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 08-15 14:38:20 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +2026-08-15 14:38:22.233616: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-08-15 14:38:22.284693: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. +moe_topk_softmax: + +=== corex_moe_exact_reduce 签名 === +/usr/local/corex/lib64/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 08-15 14:38:31 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +2026-08-15 14:38:33.436893: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-08-15 14:38:33.488922: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. +serial_float: +serial_half: +tree_float: + +=== corex_moe_weight_gather 签名 === +/usr/local/corex/lib64/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 08-15 14:38:42 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +2026-08-15 14:38:44.640768: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-08-15 14:38:44.692741: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. +gather: + +=== corex_moe_index_combine 签名 === +/usr/local/corex/lib64/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 08-15 14:38:54 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +2026-08-15 14:38:56.150733: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-08-15 14:38:56.203274: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. +moe_combine_result: +moe_compute_index: From c290278b358b4c918739c92601f578d261b9b6ba Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 15 Aug 2026 14:46:49 +0000 Subject: [PATCH 5/9] =?UTF-8?q?test:=20probe=5Fso=5Fimport=5Fchain.sh=20?= =?UTF-8?q?=E2=80=94=20=E9=AA=8C=E8=AF=81.so=E9=83=A8=E7=BD=B2=E8=B7=AF?= =?UTF-8?q?=E5=BE=84+import=E9=93=BE+flag=E5=80=BC+shape=E5=8C=B9=E9=85=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- probe_so_import_chain.sh | 165 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 165 insertions(+) create mode 100755 probe_so_import_chain.sh diff --git a/probe_so_import_chain.sh b/probe_so_import_chain.sh new file mode 100755 index 00000000..79a3830d --- /dev/null +++ b/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 From ee62ea13ba79a18e7f3ad46f5524dd55f096e614 Mon Sep 17 00:00:00 2001 From: root Date: Sat, 15 Aug 2026 14:49:42 +0000 Subject: [PATCH 6/9] data: so import chain probe --- probe_so_output.txt | 100 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 probe_so_output.txt diff --git a/probe_so_output.txt b/probe_so_output.txt new file mode 100644 index 00000000..519c90ff --- /dev/null +++ b/probe_so_output.txt @@ -0,0 +1,100 @@ +=== 1. .so文件实际位置和文件名 === +-rwxr-xr-x 1 root root 210936 Aug 13 01:33 /usr/local/corex/lib/python3/dist-packages/vllm/corex_moe_direct_routed.so +-rwxr-xr-x 1 root root 192360 Aug 13 01:33 /usr/local/corex/lib/python3/dist-packages/vllm/corex_moe_exact_reduce.so +-rwxr-xr-x 1 root root 216688 Aug 14 01:46 /usr/local/corex/lib/python3/dist-packages/vllm/corex_moe_index_combine.so +-rwxr-xr-x 1 root root 696256 Aug 13 01:33 /usr/local/corex/lib/python3/dist-packages/vllm/corex_moe_topk_softmax.so +-rwxr-xr-x 1 root root 197320 Aug 13 01:33 /usr/local/corex/lib/python3/dist-packages/vllm/corex_moe_weight_gather.so +-rwxr-xr-x 1 root root 277120 Aug 11 09:31 /usr/local/corex/lib/python3/dist-packages/vllm/ix_unified_bridge.cpython-310-x86_64-linux-gnu.so +-rwxr-xr-x 1 root root 1506880 Aug 12 01:29 /usr/local/corex/lib/python3/dist-packages/vllm/ix_unified_bridge.so + +=== 2. Python import路径 === +/usr/local/corex/lib64/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 08-15 14:48:59 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +2026-08-15 14:49:01.632894: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-08-15 14:49:01.686627: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. +vllm.__file__: /home/dylan/0814/project_6/vllm/__init__.py +vllm dir: /home/dylan/0814/project_6/vllm + corex_attn_head_rms_norm.so + corex_block_major_kv_transfer.so + corex_fused_paged_prefill.so + corex_gdn_beta_decay.so + corex_gdn_causal_conv.so + corex_gdn_chunk_recurrent.so + corex_gdn_gated_norm.so + corex_gdn_packed_decode.so + corex_gdn_qk_map.so + corex_moe_direct_routed.so + corex_moe_exact_reduce.so + corex_moe_index_combine.so + corex_moe_topk_softmax.so + corex_moe_weight_gather.so + corex_paged_kv_gather.so + ix_full_bridge.so + +=== 3. 逐个import corex_moe测试 === +/usr/local/corex/lib64/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 08-15 14:49:11 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +2026-08-15 14:49:13.043593: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-08-15 14:49:13.095797: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. + ✓ from vllm import corex_moe_topk_softmax → ['moe_topk_softmax'] + ✓ from vllm import corex_moe_direct_routed → ['w13', 'w2_reduce'] + ✓ from vllm import corex_moe_weight_gather → ['gather'] + ✓ from vllm import corex_moe_exact_reduce → ['serial_float', 'serial_half', 'tree_float'] + ✓ from vllm import corex_moe_index_combine → ['moe_combine_result', 'moe_compute_index'] + ✓ from vllm import corex_attn_head_rms_norm → ['apply_inverse', 'prepare'] + ✓ from vllm import corex_fused_paged_prefill → ['forward'] + ✓ from vllm import corex_paged_kv_gather → ['gather'] + ✓ from vllm import corex_gdn_chunk_recurrent → ['torch_chunk_gated_delta_rule', 'torch_recurrent_gated_delta_rule'] + ✓ from vllm import corex_gdn_causal_conv → ['causal_conv_update'] + ✓ from vllm import corex_gdn_beta_decay → ['beta_decay'] + ✓ from vllm import corex_gdn_gated_norm → ['apply_inverse'] + ✓ from vllm import corex_gdn_qk_map → ['qk_map'] + ✓ from vllm import corex_gdn_packed_decode → ['packed_decode'] + ✓ from vllm import corex_block_major_kv_transfer → ['check_error', 'cpu_gather', 'cpu_scatter', 'pack', 'scatter'] + +=== 4. ix_unified_bridge import测试 === +/usr/local/corex/lib64/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 08-15 14:49:22 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +2026-08-15 14:49:24.345485: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-08-15 14:49:24.397567: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. + ✗ ix_unified_bridge: cannot import name 'ix_unified_bridge' from 'vllm' (/home/dylan/0814/project_6/vllm/__init__.py) + +=== 5. 我们的qwen3_5.py里各flag的实际值 === +/usr/local/corex/lib/python3/dist-packages/torch/cuda/__init__.py:51: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. + import pynvml # type: ignore[import] +INFO 08-15 14:49:33 importing.py:10] Triton not installed; certain GPU-related functions will not be available. +2026-08-15 14:49:35.533303: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-08-15 14:49:35.585311: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. +WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them. + ✗ _USE_COREX_BATCHED_GEMM = False + ✗ _USE_COREX_MOE_DIRECT_ROUTED = False + ✓ _USE_COREX_MOE_EXACT_REDUCE = True + ✓ _USE_COREX_MOE_INDEX_COMBINE = True + ✓ _USE_COREX_MOE_TOPK_SOFTMAX = True + ✓ _USE_COREX_MOE_WEIGHT_GATHER = True + ✗ _USE_IX_FUSED_MOE = False + ✗ _USE_NAIVE_BATCHED_MOE = False + +=== 6. 模型实际shape（判断corex_direct_routed能否匹配）=== +Qwen3.5-27B MoE config (from config.json): + num_experts = 128 (per TP shard: 128/4=32? or 128?) + top_k = 8 + hidden_size = 3584 (per TP shard: 3584/4=896? or 3584?) + moe_intermediate_size = 18944 (per TP shard: 18944/4=4736) + +Expected weight shapes (TP=4): + w13: (128, 2*4736, 3584) = (128, 9472, 3584) -- NOT (256, 256, 2048) + w2: (128, 3584, 4736) -- NOT (256, 2048, 128) + +corex_moe_direct_routed hardcoded for different model! +We need corex_moe_weight_gather + F.linear path instead. From ea2c15f699c3e2875fce35ae45a83d195f33e338 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 15 Aug 2026 14:55:01 +0000 Subject: [PATCH 7/9] =?UTF-8?q?fix:=20decode=20MoE=E8=B7=AF=E5=BE=84?= =?UTF-8?q?=E5=AF=B9=E9=BD=90base=20=E2=80=94=20F.linear+bmm=E6=9B=BF?= =?UTF-8?q?=E6=8D=A2pre-transpose+bmm?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit base qwen3_5.py的decode路径（已验证可跑通竞赛）: F.linear(hidden, w13_sel.reshape(-1,H)) → view → act → bmm(w2_sel, act) 我们之前的路径（未验证，probe显示更慢）: pre-transpose(w13全量) → w13_t[eids] → bmm(x_expand, w13_t_sel) → act → bmm(act, w2_t_sel) probe真机数据: loop matmul 19ms < torch.bmm 24ms 说明F.linear路径在BI-V100单token场景下更优 保持的corex加速: ✓ corex_moe_topk_softmax (topk+softmax fused) ✓ corex_moe_weight_gather (gather fused) ✓ corex_moe_exact_reduce (weighted sum fused) ✓ corex_moe_index_combine (prefill token routing fused) --- qwen3_6_scripts/qwen3_5.py | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index a05c80c6..130894f7 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -1824,19 +1824,14 @@ class Qwen3_5MoeSparseBlock(nn.Module): H = hidden_states.shape[-1] - # --- Pre-transpose weights for bmm (cached after first call) --- - if not hasattr(self, '_w13_t') or self._w13_t is None: - # (E, 2*I, H) → (E, H, 2*I) — one-time cost at first decode - self._w13_t = self.experts.w13_weight.transpose(1, 2).contiguous() - self._w2_t = self.experts.w2_weight.transpose(1, 2).contiguous() - # (E, H, I) → (E, I, H) - - w13_t_sel = self._w13_t[eids] # (K, H, 2*I) - w2_t_sel = self._w2_t[eids] # (K, I, H) - - # FC1: bmm (K,1,H) @ (K,H,2I) → (K,1,2I) - x_expand = hidden_states.unsqueeze(0).expand(self.top_k, -1, -1) # (K, 1, H) - gate_up = torch.bmm(x_expand, w13_t_sel).squeeze(1) # (K, 2*I) + # FC1: single large GEMM via F.linear + # (1, H) @ (K*2*I, H)^T → (1, K*2*I) + # Source: base qwen3_5.py — verified on BI-V100 (sub 655 = 683) + gate_up = F.linear( + hidden_states, + w13_sel.reshape(-1, H), # (K*2*I, H) + ) # (1, K*2*I) + gate_up = gate_up.view(self.top_k, -1) # (K, 2*I) if _USE_FUSED_MOE_ACTIVATION: act = self.act_fn(gate_up) # (K, I) @@ -1844,8 +1839,9 @@ class Qwen3_5MoeSparseBlock(nn.Module): gate, up = gate_up.chunk(2, dim=-1) act = F.silu(gate) * up - # FC2: bmm (K,1,I) @ (K,I,H) → (K,1,H) - expert_out = torch.bmm(act.unsqueeze(1), w2_t_sel).squeeze(1) # (K, H) + # FC2: bmm (K, H, I) @ (K, I, 1) → (K, H) + # w2_sel is (K, H, I), act is (K, I) + expert_out = torch.bmm(w2_sel, act.unsqueeze(-1)).squeeze(-1) # (K, H) if (_USE_COREX_MOE_EXACT_REDUCE and expert_out.dtype == torch.float16 From 8eabbac857bcb36ca0679672bfbe2258ab9398fb Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 15 Aug 2026 15:05:26 +0000 Subject: [PATCH 8/9] =?UTF-8?q?test:=20probe=5Fmodel=5Fshapes.sh=20?= =?UTF-8?q?=E2=80=94=20=E7=9C=9F=E6=9C=BA=E9=AA=8C=E8=AF=81=E6=A8=A1?= =?UTF-8?q?=E5=9E=8Bconfig=E5=92=8CMoE=E6=9D=83=E9=87=8Dshape?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- probe_model_shapes.sh | 48 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100755 probe_model_shapes.sh diff --git a/probe_model_shapes.sh b/probe_model_shapes.sh new file mode 100755 index 00000000..2633a063 --- /dev/null +++ b/probe_model_shapes.sh @@ -0,0 +1,48 @@ +#!/bin/bash +set -e + +echo "=== 模型权重实际shape ===" +python3 -c " +import torch, os, json +# 读config.json +cfg_path = '/model/config.json' +if os.path.exists(cfg_path): + with open(cfg_path) as f: + cfg = json.load(f) + print('Model config:') + for k in ['hidden_size', 'intermediate_size', 'num_attention_heads', + 'num_key_value_heads', 'num_hidden_layers', 'num_experts', + 'num_experts_per_tok', 'moe_intermediate_size', 'vocab_size', + 'max_position_embeddings']: + print(f' {k}: {cfg.get(k, \"N/A\")}') +else: + print(f'{cfg_path} not found') + # 搜索 + import glob + for p in glob.glob('/model/**/config.json', recursive=True): + print(f' found: {p}') +" + +echo "" +echo "=== safetensor权重shape（第一个shard）===" +python3 -c " +from safetensors import safe_open +import glob, os +shards = sorted(glob.glob('/model/model*.safetensors')) +if not shards: + shards = sorted(glob.glob('/model/*.safetensors')) +if shards: + print(f'Found {len(shards)} shards, reading first: {shards[0]}') + with safe_open(shards[0], framework='pt') as f: + for key in sorted(f.keys()): + if 'experts' in key and ('w1' in key or 'w2' in key or 'w13' in key): + print(f' {key}: {f.get_tensor(key).shape}') + break # 只看一个就够了 + # 也看gate + for key in sorted(f.keys()): + if 'gate' in key and 'weight' in key: + print(f' {key}: {f.get_tensor(key).shape}') + break +else: + print('No safetensor shards found') +" 2>&1 || echo "safetensors not available" From 4189f44d27eaac712efd00e57cfdb1da88858a84 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 16 Aug 2026 04:58:39 +0000 Subject: [PATCH 9/9] =?UTF-8?q?test:=20dump=20ALL=20symbols=20from=20ALL?= =?UTF-8?q?=20ixformer/cuinfer=20.so=20=E2=80=94=20no=20grep=20filter,=20f?= =?UTF-8?q?ind=20what=20we=20missed?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- qwen3_6_scripts/probe_all_symbols.sh | 120 +++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100755 qwen3_6_scripts/probe_all_symbols.sh diff --git a/qwen3_6_scripts/probe_all_symbols.sh b/qwen3_6_scripts/probe_all_symbols.sh new file mode 100755 index 00000000..5c03e6f7 --- /dev/null +++ b/qwen3_6_scripts/probe_all_symbols.sh @@ -0,0 +1,120 @@ +#!/bin/bash +# probe_all_symbols.sh — Dump ALL exported symbols from every relevant .so +# No grep filter — save full lists, then we search offline + +OUTDIR="cat_files/symbol_dumps" +mkdir -p "$OUTDIR" + +echo "=== 1. ALL ixformer .so files ===" +find /usr/local/corex -name "*ixformer*" -name "*.so" 2>/dev/null | sort | tee "$OUTDIR/ixformer_so_list.txt" +echo "" + +echo "=== 2. Dump each ixformer .so symbols ===" +while read so; do + base=$(basename "$so" | sed 's/[^a-zA-Z0-9._-]/_/g') + count=$(nm -D "$so" 2>/dev/null | grep " T " | wc -l) + echo " $so → $base ($count T symbols)" + nm -D "$so" 2>/dev/null | grep " T " > "$OUTDIR/sym_${base}.txt" +done < "$OUTDIR/ixformer_so_list.txt" +echo "" + +echo "=== 3. libixformer.so full T symbols ===" +if [ -f /usr/local/corex/lib64/libixformer.so ]; then + nm -D /usr/local/corex/lib64/libixformer.so 2>/dev/null | grep " T " > "$OUTDIR/sym_libixformer.txt" + wc -l "$OUTDIR/sym_libixformer.txt" + # Also search for ANY moe/expert/gemm/fused related + echo " grep moe/expert/gemm/fused/group/batch:" + grep -i "moe\|expert\|gemm\|fused\|group\|batch\|topk\|gating\|route" "$OUTDIR/sym_libixformer.txt" | head -30 +fi +echo "" + +echo "=== 4. libcuinfer.so full T symbols ===" +if [ -f /usr/local/corex/lib64/libcuinfer.so ]; then + nm -D /usr/local/corex/lib64/libcuinfer.so 2>/dev/null | grep " T " > "$OUTDIR/sym_libcuinfer.txt" + wc -l "$OUTDIR/sym_libcuinfer.txt" + echo " grep moe/expert/gemm/fused/group/batch:" + grep -i "moe\|expert\|gemm\|fused\|group\|batch\|topk\|gating\|route" "$OUTDIR/sym_libcuinfer.txt" | head -30 +fi +echo "" + +echo "=== 5. _ixformer_torch .so full T symbols ===" +TORCH_SO=$(find /usr/local/corex -name "_ixformer_torch*.so" 2>/dev/null | head -1) +if [ -n "$TORCH_SO" ]; then + nm -D "$TORCH_SO" 2>/dev/null | grep " T " > "$OUTDIR/sym_ixformer_torch.txt" + wc -l "$OUTDIR/sym_ixformer_torch.txt" + echo " grep moe/expert/gemm/fused/group/batch:" + grep -i "moe\|expert\|gemm\|fused\|group\|batch\|topk\|gating\|route" "$OUTDIR/sym_ixformer_torch.txt" | head -30 +fi +echo "" + +echo "=== 6. _C .so (ixformer python binding) full T symbols ===" +C_SO=$(find /usr/local/corex -path "*ixformer*" -name "_C*.so" 2>/dev/null | head -1) +if [ -n "$C_SO" ]; then + nm -D "$C_SO" 2>/dev/null | grep " T " > "$OUTDIR/sym_ixformer_C.txt" + wc -l "$OUTDIR/sym_ixformer_C.txt" + echo " grep moe/expert/gemm/fused/group/batch:" + grep -i "moe\|expert\|gemm\|fused\|group\|batch\|topk\|gating\|route" "$OUTDIR/sym_ixformer_C.txt" | head -30 +fi +echo "" + +echo "=== 7. ALL .so in ixformer package dir ===" +IXDIR=$(python3 -c "import ixformer, os; print(os.path.dirname(ixformer.__file__))" 2>/dev/null) +if [ -n "$IXDIR" ]; then + echo "ixformer dir: $IXDIR" + find "$IXDIR" -name "*.so" | while read so; do + base=$(basename "$so") + count=$(nm -D "$so" 2>/dev/null | grep " T " | wc -l) + echo " $base: $count T symbols" + nm -D "$so" 2>/dev/null | grep " T " > "$OUTDIR/sym_ixpkg_${base}.txt" + # Quick search + hits=$(grep -ic "moe\|expert\|gemm\|fused\|group\|batch\|topk" "$OUTDIR/sym_ixpkg_${base}.txt") + if [ "$hits" -gt 0 ]; then + echo " *** HIT: $hits MoE/GEMM related symbols:" + grep -i "moe\|expert\|gemm\|fused\|group\|batch\|topk" "$OUTDIR/sym_ixpkg_${base}.txt" + fi + done +fi +echo "" + +echo "=== 8. ixformer Python API — list ALL callable functions ===" +python3 << 'PY' +import ixformer +import inspect + +# List all attributes +for name in sorted(dir(ixformer)): + if name.startswith('_'): + continue + obj = getattr(ixformer, name) + if callable(obj): + try: + sig = inspect.signature(obj) + print(f" ixformer.{name}{sig}") + except (ValueError, TypeError): + print(f" ixformer.{name} (no signature)") + elif hasattr(obj, '__module__'): + print(f" ixformer.{name} = {type(obj).__name__}") + +# Check submodules +print("\n --- submodules ---") +for name in sorted(dir(ixformer)): + obj = getattr(ixformer, name) + if inspect.ismodule(obj) and not name.startswith('_'): + print(f" ixformer.{name}:") + for sub in sorted(dir(obj)): + if sub.startswith('_'): + continue + subobj = getattr(obj, sub) + if callable(subobj): + try: + sig = inspect.signature(subobj) + print(f" .{sub}{sig}") + except: + print(f" .{sub} (no sig)") +PY + +echo "" +echo "=== Files saved to $OUTDIR ===" +ls -lh "$OUTDIR/" +echo "" +echo "git add cat_files/symbol_dumps/ && git commit -m 'data: full symbol dumps' && git push"