finegrained-fp8: fused MX MoE kernels write past intermediate and read past hidden/intermediate; host_descriptor mode never compiles #1158

Summary

Follow-up to #1153. The fused MX MoE kernels in finegrained-fp8/torch-ext/finegrained_fp8/fused_grouped.py have the same defect one level up: mxfp_dynamic_moe_grouped_gate_up_kernel tiles intermediate (N) and hidden (K), mxfp_dynamic_moe_grouped_down_kernel tiles hidden (N) and intermediate (K), both take their tiles from autotune, the wrapper asserts nothing about either dim, and both kernels bound only the M direction. The intermediate is a contiguous (S, I) buffer, so the damage is not confined to slack:

• gate_up's intermediate store (tl.store(int_ptrs, inter, mask=row_mask[:, None]), line 793 on main) writes its N overhang into the next token's activations; the scale store on 802 does the same to InterScale.
• down's K loop (tl.range(0, INTERMEDIATE_DIM, BLOCK_SIZE_K), line 930) reads the intermediate with only row_mask, so a K tail pulls the next token's intermediate — nonzero — and multiplies it into valid output columns. Pointer-mode weight loads past I compound that; descriptor modes at least zero-fill.
• gate_up's K loop over HIDDEN_DIM (729) and its pointer-mode weight/scale loads past I are unmasked the same way.

#1154 masked the down kernel's store (it goes through store_tile) but nothing else here; the gate_up store doesn't use store_tile at all.

The block-dynamic fused route has a smaller cousin: NUM_I_TILES = INTERMEDIATE_DIM // BLOCK_SIZE_N (line 482) floors, so an unaligned I silently drops the remainder instead of asserting like the standalone block-scale wrapper does — and it divides by BLOCK_SIZE_N while the loop steps by BLOCK_SIZE_K.

Reproducer

Same recipe as #1153: rows are powers of two, weights are ones, act_fn="relu", so gate = up = H·2^t, the intermediate is (H·2^t)^2 (exact in MXFP8 for H=96: 96² = 1.125·2^13), and every output column must be exactly I·H²·4^t. Tiles pinned at BLOCK_SIZE_N=64 / BLOCK_SIZE_K=128, pointer memory mode, so the tail path is guaranteed to run rather than left to the tuner's mood.

import torch, triton
import finegrained_fp8 as fg, finegrained_fp8.fused_grouped as fused

T = {"COMPUTE_MODE": "dot_scaled", "BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "MEMORY_MODE": "pointer"}
fused.mxfp_dynamic_moe_grouped_gate_up_kernel.configs = [triton.Config(T, num_warps=4, num_stages=2)]
fused.mxfp_dynamic_moe_grouped_down_kernel.configs = [triton.Config(T, num_warps=4, num_stages=2)]

E, TOK, H, I = 2, 4, 96, 96
hidden = (2.0 ** torch.arange(TOK, dtype=torch.float32)).unsqueeze(1).expand(TOK, H).contiguous().to("cuda", torch.bfloat16)
idx = torch.tensor([[0], [0], [1], [1]], device="cuda", dtype=torch.int32)
w = torch.ones(TOK, 1, device="cuda")
gu = torch.ones(E, 2 * I, H, device="cuda").to(torch.float8_e4m3fn)
dn = torch.ones(E, H, I, device="cuda").to(torch.float8_e4m3fn)
gus = torch.ones(E, 2 * I, H // 32, device="cuda").to(torch.float8_e8m0fnu)
dns = torch.ones(E, H, I // 32, device="cuda").to(torch.float8_e8m0fnu)

out = fg.moe_fused_grouped(hidden, idx, w, gu, dn, gus, dns, None, act_fn="relu")
expected = (I * H * H * 4.0 ** torch.arange(TOK, device="cuda")).unsqueeze(1).expand(TOK, H)
print(out.float()[:, 0].tolist())   # [3145728.0, 19660800.0, 37486592.0, 37748736.0]
print(expected[:, 0].tolist())      # [884736.0,  3538944.0, 14155776.0, 56623104.0]
print((out.float() != expected).sum().item(), "of", out.numel())   # 384 of 384

Every element is wrong: with H=96 and BLOCK_SIZE_K=128, gate_up's K tail reads the next token's hidden row into every product, before the store overhang even gets a chance.

Scope

Six (H, I) shapes × two memory modes (pointer, host_descriptor), the test from the PR run against main+#1154, RTX PRO 6000 Blackwell (SM120), torch 2.14.0+cu130 / triton 3.8.0:

Which brings up the second thing.

host_descriptor never compiles

Pinning MEMORY_MODE="host_descriptor" fails every case, including the aligned control, with

Block pointers only support 32 bit `offsets/block_shape`, add a `.to(tl.int32)` or use regular indexing for 64 bit support

The descriptor coordinates (gu_row = expert_id * 2, expert_id * HIDDEN_DIM + n_off) are built from the int64 expert_id that _resolve_tile_inline produces for pointer-overflow safety. Unpinned, the autotuner just eats the failure and crowns a pointer config, so nothing is visibly broken — exactly the "configs stop compiling and the tuner crowns something else" failure mode @IlyasMoutawwakil warned about on #1154. On main the TMA path in these two kernels is dead code on this stack.

Real shapes do hit this. moe_intermediate_size=1408 (DeepSeek-V2-Lite) is 11·128, fine for BLOCK_SIZE_K=128, but 1408 % 256 = 128 — and 256 is in the BLOCK_SIZE_N grid.

Fix

PR coming right after this, stacked on #1154. HIDDEN_DIM/INTERMEDIATE_DIM are constexpr in these kernels, so all the guards fold at compile time and aligned shapes keep their exact code — no need for the peeling dance from #1154.