flashinfer.fused_moe.interleave_for_tma_wgmma_up

flashinfer.fused_moe.interleave_for_tma_wgmma_up(w_fp8: Tensor) → Tensor

Repack fp8 up-projection weights for the Pair_Layout WGMMA A-tile.

Under the pair layout, each warp’s 16-row SHM stripe holds 8 gate rows and 8 up rows, so silu(gate) * up becomes a per-lane register operation after the WGMMA (no cross-warp exchange).

Input layout: [E, 2*N, K] row-major fp8 — the first N rows per expert are gate weights, the last N are up weights. N must be a multiple of 64.

Output layout (still [E, 2*N, K], identical byte footprint): for every expert e and every 64-gate-row block b, the 128-row slab at 128*b packs, per warpgroup wg and warp w:

rows [wg*64 + w*16     .. +8) = gate[e, 64b + wg*32 + w*8 .. +8, :]
rows [wg*64 + w*16 + 8 .. +8) =   up[e, 64b + wg*32 + w*8 .. +8, :]

Under SWZ128 the TMA applies the 8-row x 128-byte core-matrix XOR swizzle automatically, so this only rearranges GM rows (no byte-level permutation). The result is cached on the input tensor as _tma_interleaved_up.

The down-projection weights need no preparation — the raw [E, K, N] row-major fp8 tensor is passed straight through.

Parameters:

w_fp8 (torch.Tensor) – FP8 up/gate weight tensor with shape [E, 2*N, K] (row-major).

Returns:

Repacked weight tensor with the same shape and dtype as w_fp8.

Return type:

torch.Tensor