vllm.models.flint.ops.moe_gemv

Expert-indexed small-M NVFP4 GEMV kernels (CuTe DSL) for the low-latency decode MoE block, built on the dense GEMV design of nvfp4_gemv: each CTA owns 16 output rows (the m16n8k16 M dimension, swap-AB), its warps split K and reduce through shared memory, weights are dequantized exactly to BF16 (cvt.rn.bf16x2.e2m1x2 + mul.rn.bf16x2) and the fp32 global scale is applied to the fp32 accumulator.

gate_up: one (token, expert) pair per CTA column of the grid; writes bf16(silu(g) * u) from fp32 gate/up with FlashInfer's approximate sigmoid (the shared expert rounds its gate/up to BF16 first and uses the exact SiLU, as its dense GEMM does). Expert weights use FlashInfer's W4A16 layout ([up; gate] in 64-row groups, 6D MMA scales).

down_combine: per token and 16 output rows, the down GEMVs of the token's top-k experts and of the shared expert, all weights prefetched before griddepcontrol.wait; out = bf16(bf16(sum_j tw_j * y_j) + y_shared) with y = bf16(acc * global_scale).

Functions:

• down_combine –

  out[t] = bf16(bf16(sum_j tw * bf16(act_j @ W2[e_j]^T * alpha2)) +

• gate_up –

  act[p] = silu_and_mul((x[p // topk] @ W1[ids[p]]^T) * alpha) on fp32

_compile(key, fakes, *scalars)

Compile kernel key = (kind, h, inter, si, topk) into _compiled.

Source code in vllm/models/flint/ops/moe_gemv.py

def _compile(key: tuple, fakes: list, *scalars):
    """Compile kernel ``key = (kind, h, inter, si, topk)`` into ``_compiled``."""
    global _classes
    import cutlass.cute as cute

    if _classes is None:
        _classes = _build()
    kern = _classes[key[0] == "down"](*key[1:])
    fn = _compiled[key] = cute.compile(
        kern, *fakes, *scalars, _stream(), options="--enable-tvm-ffi"
    )
    return fn

_fakes(*specs)

Compile-time tensors: (dtype, divisibility) is 1-D of symbolic size, (dtype, divisibility, shape) has that shape.

Source code in vllm/models/flint/ops/moe_gemv.py

def _fakes(*specs) -> list:
    """Compile-time tensors: (dtype, divisibility) is 1-D of symbolic size,
    (dtype, divisibility, shape) has that shape."""
    import cutlass
    import cutlass.cute as cute
    from quack.compile_utils import make_fake_tensor

    return [
        make_fake_tensor(
            getattr(cutlass, spec[0]),
            spec[2] if len(spec) > 2 else (cute.sym_int(),),
            divisibility=spec[1],
        )
        for spec in specs
    ]

_flat_bytes(t)

The contiguous storage behind a (possibly permuted) scale view.

Source code in vllm/models/flint/ops/moe_gemv.py

def _flat_bytes(t: torch.Tensor) -> torch.Tensor:
    """The contiguous storage behind a (possibly permuted) scale view."""
    t = t.view(torch.uint8)
    return torch.as_strided(t, (t.numel(),), (1,), t.storage_offset())

down_combine(act, ids, tw, w2, w2_sf, alpha2, sact, wsd, ssd, gssd, out)

out[t] = bf16(bf16(sum_j tw * bf16(act_j @ W2[e_j]^T * alpha2)) + bf16(sact[t] @ Wsd^T * gssd)).

Source code in vllm/models/flint/ops/moe_gemv.py

def down_combine(
    act: torch.Tensor,
    ids: torch.Tensor,
    tw: torch.Tensor,
    w2: torch.Tensor,
    w2_sf: torch.Tensor,
    alpha2: torch.Tensor,
    sact: torch.Tensor,
    wsd: torch.Tensor,
    ssd: torch.Tensor,
    gssd: torch.Tensor,
    out: torch.Tensor,
) -> None:
    """out[t] = bf16(bf16(sum_j tw * bf16(act_j @ W2[e_j]^T * alpha2)) +
    bf16(sact[t] @ Wsd^T * gssd))."""
    m, topk = ids.shape
    h = w2.shape[1]
    inter = w2.shape[2] * 2
    si = wsd.shape[1] * 2
    key = ("down", h, inter, si, topk)
    fn = _compiled.get(key)
    if fn is None:
        fn = _compile(
            key,
            _fakes(
                ("BFloat16", 8),
                ("Int32", 1),
                ("Float32", 1),
                ("Uint8", 16),
                ("Uint8", 4),
                ("Float32", 1),
                ("BFloat16", 8),
                ("Uint8", 16),
                ("Uint8", 4),
                ("Float32", 1, (1,)),
                ("BFloat16", 1),
            ),
            1,
        )
    fn(
        act.view(-1),
        ids.view(-1),
        tw.view(-1),
        w2.view(-1),
        _flat_bytes(w2_sf),
        alpha2,
        sact.view(-1),
        wsd.view(-1),
        ssd.view(-1),
        gssd.reshape(1),
        out.view(-1),
        m,
        _stream(),
    )

gate_up(x, ids, w1, w1_sf, alpha, act, wsh, ssh, gssh, sact)

act[p] = silu_and_mul((x[p // topk] @ W1[ids[p]]^T) * alpha) on fp32 gate/up for the FlashInfer-layout NVFP4 experts (w1 uint8 [E, 2I, H/2]), and in the same launch sact[t] = silu_and_mul(bf16((x[t] @ Wsh^T) * gssh)) for the shared expert (ModelOpt [gate; up] uint8 [2 SI, H/2], flat 128x4-swizzled scales).

Source code in vllm/models/flint/ops/moe_gemv.py

def gate_up(
    x: torch.Tensor,
    ids: torch.Tensor,
    w1: torch.Tensor,
    w1_sf: torch.Tensor,
    alpha: torch.Tensor,
    act: torch.Tensor,
    wsh: torch.Tensor,
    ssh: torch.Tensor,
    gssh: torch.Tensor,
    sact: torch.Tensor,
) -> None:
    """act[p] = silu_and_mul((x[p // topk] @ W1[ids[p]]^T) * alpha) on fp32
    gate/up for the FlashInfer-layout NVFP4 experts (w1 uint8 [E, 2I, H/2]),
    and in the same launch sact[t] = silu_and_mul(bf16((x[t] @ Wsh^T) * gssh))
    for the shared expert (ModelOpt [gate; up] uint8 [2 SI, H/2], flat
    128x4-swizzled scales)."""
    m, h = x.shape
    p, topk = ids.numel(), ids.shape[1]
    inter = w1.shape[1] // 2
    si = wsh.shape[0] // 2
    key = ("gu", h, inter, si, topk)
    fn = _compiled.get(key)
    if fn is None:
        fn = _compile(
            key,
            _fakes(
                ("BFloat16", 8),
                ("Int32", 1),
                ("Uint8", 16),
                ("Uint8", 4),
                ("Float32", 1),
                ("BFloat16", 1),
                ("Uint8", 16),
                ("Uint8", 4),
                ("Float32", 1, (1,)),
                ("BFloat16", 1),
            ),
            1,
            1,
        )
    fn(
        x.view(-1),
        ids.view(-1),
        w1.view(-1),
        _flat_bytes(w1_sf),
        alpha,
        act.view(-1),
        wsh.view(torch.uint8).view(-1),
        ssh.view(-1),
        gssh.reshape(1),
        sact.view(-1),
        p,
        m * (si // inter),
        _stream(),
    )