Repository navigation
Conversation
…nwritten With inter_dim_pad > 0, the stage-1 launch grid adds tile2_pad (the rounding of the unpadded intermediate up to stage 2's K tile, in output columns) to the w1 row axis, which carries two rows per output column. For gate/up-interleaved kernels with tile_n < 2 * tile2_pad the grid falls one tile short; e.g. DeepSeek-V4.1-Flash at TP4 (inter 576 padded to 640, tile_n=64 _gui kernels from dsv41_fp8_tuned_fmoe.csv) writes 608 of the 640 columns. The rest of the fp8 intermediate and its e8m0 scales stay torch.empty, and the layout GEMM2, which takes no inter_dim_pad, reads them: stale NaN bytes there turn whole output rows NaN, flakily. The separated layout's grid under-covers the same way for small tile_n, and the xcd-swizzle grid_x omitted tile2_pad altogether, so it no longer matched the launched grid. Compute the written width once, in output columns, and derive both grids from it. Signed-off-by: Micah Williamson <micah.williamson@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
| gx = ( | ||
| inter_in - inter_dim_pad_total + tile2_pad + tile_n_index - 1 | ||
| ) // tile_n_index | ||
| gx = (inter_in - n_rows_skipped_idx + tile_n_index - 1) // tile_n_index |
There was a problem hiding this comment.
Maybe I'm missing something, but doesn't this gx recompute the grid_x we compute above in line 541? If it does, I think we should consolidate the vars and if not, would it be worth enforcing that they're the same with a check/comparison?
3 participants
Summary
When
inter_dim_pad > 0, the FlyDSL MoE stage-1 launch grid can fall one tile short. The intermediate columns past that point are never written, and the layout GEMM2 (flydsl_moe2_layout_*) reads them anyway. If thosetorch.emptybytes happen to hold a NaN pattern, whole output rows come out NaN. That happens intermittently, depending on what the caching allocator handed back.The kernel bug predates #6125; that PR exposed it. #6125 added the gfx950 A8W4 tuned configs for DeepSeek-V4.1-Flash with TP-sharded experts, including the first rows for
inter_dim=640(TP4: 576 columns padded to 640). For 1, 2, and 4 tokens those rows selecttile_n=64gate/up-interleaved (_gui) stage-1 kernels paired with layout GEMM2 kernels. That's exactly the combination that under-covers. Before #6125, TP4 had no tuned rows and took the default kernels.This turned up as a flaky vLLM CI failure after the bump to v0.1.25:
test_dsv41_mono_layers.py::test_mono_decode_matches_the_model[4], where the reference run produced a NaN residual.Root cause
In
compile_mixed_moe_gemm1_common,tile2_padis the rounding of the unpadded intermediate up to stage 2's K tile, measured in output columns. The grid adds it to the w1 row axis, which carries two rows per output column (gate and up):For DSv4.1 at TP4 with
tile_n=64this gives (1280 − 128 + 64 + 63) // 64 = 19 tiles. Each_guitile writestile_n / 2columns, so 19 × 32 = 608 of 640 columns get written. Columns 608 to 639 of the fp8 intermediate, and their e8m0 scale block, keep stale contents._flydsl_v2_stage2_wrappertakes nointer_dim_padand runs K over all 640 columns. The padded w2 weights are zero, but NaN × 0 is still NaN inmfma_scale.Two related problems in the same code:
tile_n; for exampletile_n=32writes 608 columns.grid_xin the kernel omittedtile2_padentirely, so it didn't match the launched grid whenevertile2_pad > 0.Fix
Compute once, in output columns, how far stage 1 must write: the unpadded width rounded up to stage 2's K tile, capped at
inter_dim. Derive both the launch grid and the XCD-swizzlegrid_xfrom that value, as w1 rows to skip (two per column). FlyDSL keys its compile cache on function source, andcache_tagalready includesinter_dim_pad,tile_k, andgate_mode, so stale cached launchers can't hide the change.Not in this PR: making the layout GEMM2 accept
inter_dim_padand skip the padded K range. With stage 1 writing every column that's a small speedup, not a correctness fix.Testing
Standalone
fused_moeat the DSv4.1-Flash TP4 routed-expert shape: 384 experts, top-6, hidden 5120,inter_dim=640,intermediate_pad=64,gate_mode=interleave, fp8 activations (AITER_BF16_FP8_MOE_BOUND=0, as vLLM sets it), and weights shuffled the way vLLM loads them. Before each call, the caching allocator is filled with 0xFF, which is NaN as both fp8 e4m3fn and e8m0.dsv41_fp8_tuned_fmoe.csv)t32x64x256_w4_gui_kw4_fp8+layout_..._atomic_nt_sbm32t32x64x256_w3_gui_kw4_fp8+layout_..._atomic_sbm32t32x64x256_w3_gui_kw2_fp8+layout_..._reduce_sbm32t32x128x256_*_gui_fp8+ layout GEMM2t16x128x256_w3_gui+ classic GEMM2Zeroing stage 1's buffers instead of using
torch.emptyalso removes the NaNs, which confirms the unwritten buffers are the source.Correctness: the same tokens computed alone (the fixed
tile_n=64kernels) and inside an 8-token batch (tile_n=128, unaffected) agree. 4 tokens is bit-identical; 1 and 2 tokens differ by about 3e-3 relative, which is bf16 rounding from those buckets' atomic-add GEMM2. Unpoisoned outputs are unchanged from before the fix.This PR was written with AI assistance. I have reviewed all changes.