Pull requests / #2321

#2321 route-resident: support 256-expert top-8 models (Qwen3.6-35B-A3B)

open · @atomicmilkshake · 0 comentarios · En GitHub

NVIDIA / CUDAModels & quants

Descripción

### Summary

Generalizes the opt-in \STRATA_ROUTE_RESIDENT\ kernel and verifier guard from 512/10 models to also support 256-expert top-8 models (\qwen35moe\ / Qwen3.6-35B-A3B).

Follows Ethan's warp-reduction rewrite (#1737) and separates out the 256/8 portion previously included in #1820 into a clean, standalone change on top of upstream \main\.

### Details

1. **Kernel parameterization (\src/kernels/cuda/native_router.cu\)**:
   - Templates \
oute_resident_impl<NE, K>\ across warp-multiple expert counts. Both 512 and 256 are multiples of 32, so the \__shfl_down_sync\ reduction tree stays convergent across lanes.
   - Retains the original \
oute_resident_k\ (512/10) CUDA symbol and entry point, and instantiates \
oute_resident_geometry_k<256, 8>\ for Qwen geometries.
   - Restricts buffer writes (\weights\, \id\, \stats\) to lane 0, eliminating redundant warp writes.
   - Maintains a backward-compatible 10-argument \
ative_route_resident\ overload so existing 512/10 call sites (\src/core/mtp.cpp\, metal stubs) continue without change.

2. **Verifier & Configuration (\src/core/verify.cpp\)**:
   - Zero-based tail ranks default to \5-7\ for 256/8 and \6-9\ for 512/10 unless overridden by \STRATA_ROUTE_RESIDENT_RANKS=lo-hi\.
   - Widens the launch guard to \(NE == 512 && K == 10) || (NE == 256 && K == 8)\.

3. **Parity Test (\src/kernels/route_resident_parity.cu\)**:
   - Generalizes the test against scalar reference implementations for both \512/10\ and \256/8\ geometries.
   - Passes with 0 mismatches across 240 synthetic cases for each geometry (480 total cases), verifying bit-for-bit ID and counter parity, and weights within 0 ulp.

### Context

Extracted from #1820 per maintainer single-change guideline. Validated on 0.1.42 (+9.7% decode / -13.1% window on Qwen3.6-35B-A3B).

En el sitio

Enlaces a install, modelos, releases.