[SM121 / GB10 (DGX Spark)] moe_wna16_marlin_gemm no-split-K (data-parallel) path: CUDA illegal memory access at M=256 — clean for M≤128; default split-K path clean at M=256 #56064

Activity

1. LH-and-FPGA commented on Sep 9, 2026

   

   LH-and-FPGA

   on Sep 9, 2026

   Contributor

   More actions

   Hi! Thanks for the detailed report. I took a closer look at the scheduler in csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h, and unless I’m missing another path, I’m not sure the current upstream code can reach this state.

   The kernel does not appear to run in pure DP. In the part1/part2 split (lines 363–375),

   if (part2_mn_tiles * 3 <= gridDim.x) part2_mn_tiles += gridDim.x;

   ensure part2 receives at least a full grid, while part1 only takes a whole multiple of gridDim.x. As a result, in part1,

   slice_col_par <= part1_mn_iters * gridDim.x - 1 < global_mn_tiles

   so par_id < parallel should hold by construction.

   I also swept parallel = 1..500 using the execution configuration from your report (gridDim.x = 144, with n_tiles = 32 and 64), and the invariant held throughout. That also seems consistent with your control result, where the stock op is clean at M=256.

   I also couldn’t find a split_k parameter exposed by moe_wna16_marlin_gemm. The op appears to expose only thread_k, thread_n, blocks_per_sm, use_atomic_add, and use_fp32_reduce, so my current understanding is that the forced-DP configuration would require a patched build.

   That said, I think the crash does highlight a real robustness gap that may be worth tightening. init_part1_slice() sets slice_iters = k_tiles before calling update_next_moe_block_data(). If par_id >= parallel, that function returns early, but the block can then continue with the previous block’s shared-memory token IDs and expert ID rather than terminating the slice. With the current upstream partitioning this state appears to be unreachable, so the guard is effectively defensive today, but it is exactly the kind of failure mode that a different repartitioning could expose.

   It may therefore be safer to set slice_iters = 0 before returning when par_id >= parallel, which would make the invariant explicit and prevent stale shared-memory state from being consumed if the scheduler changes in the future.

   Would you mind sharing the diff that forces split_k = 1? In particular, I’d be interested to see whether it removes the part2_mn_tiles += gridDim.x floor or changes the computation of part1_mn_iters (for example, to a ceil division). If so, that would seem to explain the overhang you’re seeing.
2. pasta-paul commented on Sep 10, 2026

   

   pasta-paul

   on Sep 10, 2026

   Author

   More actions

   Thanks for the careful read — your construction argument is correct, and it is consistent with our controls: the unmodified op cannot reach the state (our stock-op control at M=256 completed cleanly twice: 7.8628 ms and 8.2861 ms, finite outputs).

   The forcing diff (build-time, in our research overlay — the op has no runtime split_k knob, exactly as you say) has three pieces:

   1. Exec config forced to the kernel's own pick. The wrapper's auto determine_exec_config(...) branch is replaced with a fixed config — the pick the stock determine_exec_config itself returns at every M=1..256 for these shapes (we swept it in-leg: [thread_k=64, thread_n=128, num_threads=128, blocks_per_sm=3], zero transitions at both call sites; grid = 48 SMs × 3 = 144). MarlinDefault / get_marlin_kernel / determine_exec_config are stripped and the kernel is direct-instantiated.

   2. DP-only partitioning — and yes, both halves of your question. The stock block

   if (global_mn_tiles > gridDim.x) {
       part2_mn_tiles = global_mn_tiles % gridDim.x;
       if (part2_mn_tiles * 3 <= gridDim.x) part2_mn_tiles += gridDim.x;
       part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x;
   }

   is replaced by

   int part2_mn_tiles = 0;                 // DP-only: every mn-tile is a whole-tile (part1) job
   int part1_mn_iters = global_mn_tiles;   // driver bound only

   i.e. the part2_mn_tiles += gridDim.x floor is removed AND part1_mn_iters no longer takes only a whole multiple of gridDim.x — part1's iteration space covers the full tile range. The split-K reduce path (locks + global scratch) is compiled out.

   3. Two modules, same core treatment: a static grid-stride reference (no scheduler) and a persistent ticket-claim scheduler, on an identical DP-only math path. Both hit the illegal memory access at M=256 only; M∈{1..8,16,32,64,128} complete cleanly and are bit-transparent between the two implementations (0 output flips at every M).

   So we agree the crash is reachable only under this forced repartitioning — which makes the within-slice guard you identified the interesting part. Both of our modules added an entry-level exhaustion guard in init_part1_slice() (if (slice_col_par >= global_mn_tiles) { slice_iters = 0; return; }), but neither guards the early return inside update_next_moe_block_data() — the exact hole you named: slice_iters = k_tiles already set, par_id >= parallel returns early, and the block continues on the previous tile's shared-memory token/expert IDs. That is fully consistent with an illegal read at M=256 in the DP-only construction, and your slice_iters = 0 before returning would make the invariant explicit. If it is useful, we can re-run the M=256 repro with that guard applied to confirm it converts the IMA into a clean early exit.

   — written by a Digby.ai coding harness (kernel-research lane); reviewed by a human before posting.