diff --git a/qwen3_6_scripts/_custom_ops.py b/qwen3_6_scripts/_custom_ops.py index 4ba1c1ad..9f22f0e0 100644 --- a/qwen3_6_scripts/_custom_ops.py +++ b/qwen3_6_scripts/_custom_ops.py @@ -830,8 +830,30 @@ def invoke_fused_moe_kernel( 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) + # CCCL policy_selector degradation: when one kernel in the chain is + # unavailable, replace ONLY that kernel with PyTorch while keeping the + # downstream native kernels (moe_align_block_size, invoke_fused_moe_kernel). + # This is analogous to CCCL's multi_pass fallback when onesweep is not + # available — the sort still happens, just through a different code path. + try: + ixf_F.vllm_moe_topk_softmax(topk_weights, topk_ids, + token_expert_indicies, gating_output) + except (AttributeError, RuntimeError): + # PyTorch fallback: softmax → topk → write in-place + # gating_output is already float32 (cast at call site) + if isinstance(gating_output, torch.Tensor): + probs = torch.softmax(gating_output, dim=-1) + else: + probs = torch.softmax(gating_output, dim=-1) + topk = topk_weights.shape[1] + tw, ti = torch.topk(probs, topk, dim=-1) + topk_weights.copy_(tw) + topk_ids.copy_(ti.to(topk_ids.dtype)) + # token_expert_indicies is unused by caller (deleted after call) + # but fill it for correctness + token_expert_indicies.copy_( + torch.arange(topk, device=topk_ids.device, dtype=topk_ids.dtype) + .unsqueeze(0).expand_as(topk_ids)) if supports_moe_ops and hasattr(torch.ops._moe_C, "marlin_gemm_moe"): diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 3c59b65d..a8b7a5b1 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -912,20 +912,31 @@ class Qwen3_5MoeSparseBlock(nn.Module): # ixf_F.vllm_invoke_fused_moe_kernel # The original comment "ixformer lacks MoE kernels" may have been # wrong or outdated. Try native first, catch and fallback if it fails. + # CCCL policy_selector pattern: try native fused kernel chain first. + # topk_softmax now has PyTorch fallback (see _custom_ops.py), so the + # chain topk_softmax→align→invoke may succeed even without the native + # topk op. Only permanently disable if align or invoke also fails. if not hasattr(self, '_use_native_moe'): - self._use_native_moe = True # optimistic: try native first + self._use_native_moe = True + self._native_moe_attempts = 0 if self._use_native_moe: try: routed_out = self.experts(hidden_states, router_logits) except Exception as e: - # Native kernel failed — disable permanently for this instance - # and fallback to pure PyTorch for all subsequent calls. - logger.warning( - "FusedMoE native kernel failed (%s: %s), " - "falling back to pure PyTorch experts permanently.", - type(e).__name__, e) - self._use_native_moe = False + self._native_moe_attempts += 1 + if self._native_moe_attempts >= 2: + # Failed twice (first call + retry) — truly no native support + logger.warning( + "FusedMoE native kernel failed %d times (%s: %s), " + "falling back to pure PyTorch experts permanently.", + self._native_moe_attempts, type(e).__name__, e) + self._use_native_moe = False + else: + logger.info( + "FusedMoE native kernel failed on attempt %d (%s: %s), " + "will retry next call.", + self._native_moe_attempts, type(e).__name__, e) routed_out = self._pure_pytorch_experts(hidden_states, router_logits) else: routed_out = self._pure_pytorch_experts(hidden_states, router_logits)