=== 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: