MXFP8 dense GEMM1 (gate_up) with fused clamped-SwiGLU + MXFP8 quant epilogue.
One persistent SM100 tcgen05 MXFP8 GEMM tuned for prefill shapes (BM=128,
BN=128 per gate/up half, BK=128, cta_group=2; 6-warp CTA: warps 0-3
epilogue, warp 4 MMA, warp 5 TMA) computes gate = X @ Wg^T and
up = X @ Wu^T and applies clamped SwiGLU and MXFP8 quantization. No
BF16 intermediate.
The tmem accumulator is double-buffered with a 32-col overlap (even tiles
[0,256), odd [224,480), SF fixed [500,512)): the epilogue reads the shared
32-col group first, so the next tile's MMA waits on a 32-col read instead
of 256. The single epilogue warpgroup processes a tile in 32-col chunks
(load -> act+quant -> store), overlap chunk first.
Tile order: flat (n-fastest) when the B operand fits in half of L2, else
grouped m-fastest (GROUP_M=32) so B panels stay L2-resident (dense TP1:
B=151MB > 126MB L2).
Functions:
-
can_fuse_gemm1
–
Whether fused_gemm1 can replace gate_up_proj + act_fn for down_proj.
-
fused_gemm1
–
act_fn(gate_up_proj(x)), MXFP8-quantized for an MXFP8 down_proj.
-
run
–
Functional interface to the fused kernel (hides compile + grid sizing).
can_fuse_gemm1(gate_up_proj, down_proj, act_fn)
Whether fused_gemm1 can replace gate_up_proj + act_fn for down_proj.
Called at init, before weight loading: kernels are selected at
create_weights time, so the decision does not depend on the weights.
Source code in vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py
| def can_fuse_gemm1(gate_up_proj, down_proj, act_fn) -> bool:
"""Whether fused_gemm1 can replace gate_up_proj + act_fn for down_proj.
Called at init, before weight loading: kernels are selected at
create_weights time, so the decision does not depend on the weights.
"""
from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
FlashInferCutedslMxfp8LinearKernel,
FlashInferCutlassMxfp8LinearKernel,
)
from vllm.model_executor.layers.activation import SiluAndMulWithClamp
if not (
current_platform.is_cuda() and current_platform.is_device_capability_family(100)
):
return False
if not isinstance(act_fn, SiluAndMulWithClamp):
return False
# The kernel reads gate_up's fp8 weight and F8_128x4-swizzled weight
# scale directly, so it depends on these kernels' weight layout.
if not isinstance(
_linear_kernel(gate_up_proj),
(FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel),
):
return False
down_kernel = _linear_kernel(down_proj)
if down_kernel is None or down_kernel.input_quant_key() != kMxfp8Dynamic:
return False
# I % 256 == 0 (output row stride), K % 128 == 0 (BK).
n = gate_up_proj.output_size_per_partition // 2
k = gate_up_proj.input_size_per_partition
return n % 256 == 0 and k % 128 == 0
|
fused_gemm1(x, gate_up_proj, act_fn)
act_fn(gate_up_proj(x)), MXFP8-quantized for an MXFP8 down_proj.
Requires can_fuse_gemm1(gate_up_proj, down_proj, act_fn).
Source code in vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py
| def fused_gemm1(x: torch.Tensor, gate_up_proj, act_fn) -> QuantizedActivation:
"""act_fn(gate_up_proj(x)), MXFP8-quantized for an MXFP8 down_proj.
Requires can_fuse_gemm1(gate_up_proj, down_proj, act_fn).
"""
from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
FlashInferCutedslMxfp8LinearKernel,
)
assert x.ndim == 2 and x.dtype == torch.bfloat16
weight = gate_up_proj.weight
# FlashInferCutedsl stores the weight as a column-major [K, 2I] view.
if isinstance(_linear_kernel(gate_up_proj), FlashInferCutedslMxfp8LinearKernel):
weight = weight.t()
M = x.shape[0]
N = weight.shape[0] // 2
q = torch.empty((M, N), dtype=torch.float8_e4m3fn, device=x.device)
sf = torch.empty(
((M + 127) // 128) * 128 * (N // 32), dtype=torch.uint8, device=x.device
)
xq, xsf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
run(
xq,
weight,
xsf.view(-1).view(torch.float8_e8m0fnu),
gate_up_proj.weight_scale.view(torch.float8_e8m0fnu),
q,
sf,
act_fn.swiglu_limit,
act_fn.alpha,
act_fn.beta,
)
return QuantizedActivation(
data=q,
scale=sf,
orig_dtype=x.dtype,
orig_shape=x.shape[:-1] + (N,),
quant_key=kMxfp8Dynamic,
)
|
run(xq, weight, x_scale, weight_scale, q, sf, limit, alpha, beta, grid=None)
Functional interface to the fused kernel (hides compile + grid sizing).
xq: [M, K] fp8e4m3; weight: [2N, K] fp8e4m3 row-major; x_scale/weight_scale:
F8_128x4-swizzled E8M0 (flat, any uint8-viewable dtype); q: [M, N] fp8e4m3;
sf: flat uint8 scale buffer (round_up(M,128) * N/32 bytes; rows >= M are
not written).
Source code in vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py
| def run(
xq: torch.Tensor,
weight: torch.Tensor,
x_scale: torch.Tensor,
weight_scale: torch.Tensor,
q: torch.Tensor,
sf: torch.Tensor,
limit: float,
alpha: float,
beta: float,
grid: int | None = None,
) -> None:
"""Functional interface to the fused kernel (hides compile + grid sizing).
xq: [M, K] fp8e4m3; weight: [2N, K] fp8e4m3 row-major; x_scale/weight_scale:
F8_128x4-swizzled E8M0 (flat, any uint8-viewable dtype); q: [M, N] fp8e4m3;
sf: flat uint8 scale buffer (round_up(M,128) * N/32 bytes; rows >= M are
not written).
"""
M, K = xq.shape
_, N = q.shape
if grid is None:
grid_m2 = ((M + 127) // 128 + 1) // 2 * 2
num_sms = torch.cuda.get_device_properties(xq.device).multi_processor_count
grid = min(num_sms, grid_m2 * (N // 128))
# Grouped (m-fastest) rasterization only pays when B (n2*k fp8 bytes)
# exceeds half of L2; below that the flat order is faster (no div/mod
# overhead in compute_bid)
GROUP_M = 32 if N * 2 * K > (63 << 20) else 0
kernel = Sm100DenseGemm1Mxfp8.compile(GROUP_M)
kernel(xq, weight, x_scale, weight_scale, q, sf, limit, alpha, beta, grid)
|