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 so one boxDim=(128, 128) SWIZZLE_128B TMA issue fetches a full 128-row x 128-K WGMMA A-tile.

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 k in [0, N/64):

new[e, 128k +  0:128k+ 32] = gate[e, 64k    :64k+32]
new[e, 128k + 32:128k+ 64] =   up[e, 64k    :64k+32]
new[e, 128k + 64:128k+ 96] = gate[e, 64k+32 :64k+64]
new[e, 128k + 96:128k+128] =   up[e, 64k+32 :64k+64]

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). The first N rows per expert are gate weights; the last N are up weights. N must be a multiple of 64.

Returns:

Repacked weight tensor with the same shape and dtype as w_fp8, laid out so that a single TMA boxDim=(128, 128) issue covers a complete 128-row x 128-K WGMMA A-tile.

Return type:

torch.Tensor