fix(critical): CCCL policy_selector degradation for MoE — PyTorch fallback for topk_softmax

CCCL tuning_radix_sort.cuh teaches: when one kernel in a chain is unavailable,
replace ONLY that kernel while keeping downstream native ops alive.

Our MoE chain: topk_softmax → moe_align_block_size → invoke_fused_moe_kernel
BI-V100 ixformer lacks vllm_moe_topk_softmax, which killed the ENTIRE chain
and forced 100% PyTorch fallback (_pure_pytorch_experts: 256x F.linear loop).

Fix: Add try/except in topk_softmax with PyTorch fallback (softmax+topk).
Now the chain can proceed to native align+invoke kernels if they exist.
Also: dont permanently disable native path after first failure — retry once.

CCCL source: catch2_test_device_radix_sort_pairs.cu + tuning_radix_sort.cuh
Maps to: _custom_ops.py (topk_softmax) + qwen3_5.py (MoE forward)

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"):

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)