CCCL source: cub/device/dispatch/dispatch_segmented_sort.cuh (1544 lines)
Core design: three-way partition → specialized kernels per size group.
- Large segments → full-block radix sort kernel
- Medium segments → sub-warp merge sort
- Small segments → compact sub-warp
- Below threshold → fallback kernel (no partitioning)
Applied to qwen3_5.py:
At module bottom, try to import base image's native CoreX-accelerated
Qwen3_5ForCausalLM from corex_gdn or qwen3_5_native modules. If found,
replace our PyTorch classes with the native ones.
This is the dispatch_segmented_sort pattern: if a specialized kernel
exists for this hardware (corex_gdn.so), use it. Only fall back to
the generic implementation (our pure-PyTorch code) when the specialized
path is unavailable.
Sub168 used the native CoreX path (zero NaN, 8.49s d01, 17.35GB weights).
Our PyTorch fallback has 99.98% NaN. The dispatch ensures we automatically
use the best available path.
e832687893