flashinfer.fused_moe.alphamoe_interleave_gated_weights¶
- flashinfer.fused_moe.alphamoe_interleave_gated_weights(gemm1_weights: Tensor, gemm1_weights_scale: Tensor) Tuple[Tensor, Tensor]¶
Interleave gate/up weight rows into the alphamoe_sm100 device layout.
alphamoe_fp8_block_scale_aligned_moe()consumes gemm1 (gate+up) weights whose row axis alternates eight gate rows with eight up rows, so one 128-row TMA box carries matching gate and up features. This offline helper converts the conventional[gate; up]-stacked layout. Run once at weight-load time; the kernel must never see un-interleaved weights.- Parameters:
gemm1_weights (torch.Tensor) – FP8 weights shaped
(num_experts, 2 * intermediate_size, hidden_size)with gate rows first, interleaved in eight-row chunks.gemm1_weights_scale (torch.Tensor) – Float32 block scales shaped
(num_experts, 2 * intermediate_size / 128, hidden_size / 128), interleaved one row at a time. Each scale row covers 128 weight rows.
- Returns:
The
(weights, scales)pair in device layout, retaining the input shapes, dtypes, and devices.- Return type:
Tuple[torch.Tensor, torch.Tensor]