dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

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

Browse Source 
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)
This commit is contained in:
project6
2026-08-07 08:56:43 +00:00
parent 5a3bcbc247
commit a1558b6e50
2 changed files with 43 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
qwen3_6_scripts/qwen3_5.py
View File

@@ -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)