vllm.models.common.ops.moe_tail_fusion.primitives

CuTe DSL primitives shared by the fused MoE-tail kernels: PTX loads and stores, NVLS multicast stores, Lamport mailbox helpers, BF16 packing, block reductions and launch helpers. This module defines no kernel.

Classes:

• CUDAGraphCompatibleWrapper –

  DLPack view that does not synchronize with the producer stream.

Functions:

• fragment_has_negative_zero –

  Whether a BF16 fragment still holds the BF16 sentinel in any lane.

• fragment_is_dirty –

  Whether an FP32 fragment still holds the FP32 sentinel in any word.

• group_leader_block_sum –

  Block sum of the values held by every leader_stride-th lane.

• load_global_u32x4 –

  Load one 128-bit fragment as four u32 registers.

• load_shared_f32x2 –

  Load two aligned FP32 partials from local shared memory.

• load_shared_f32x4 –

  Load four aligned FP32 partials from local shared memory.

• map_shared_to_peer –

  Map a local shared-memory slot to the same slot in a peer CTA.

• red_async_release_gpu_add_u32 –

  The SM100 arrival primitive FlashInfer's LamportFlags use.

• sanitize_negative_zero_u32x4 –

  Turn real BF16 -0 into +0 so it cannot equal the empty sentinel.

• stmc_bf16x8 –

  Publish eight BF16 values (four packed words) through NVLS multicast.

• store_global_u32x4 –

  Store four packed words to an ordinary or NVLS multicast global VA.

• store_lamport_sentinel_128 –

  Reset one FP32 Lamport fragment to four FP32 negative-zero bit patterns.

• store_lamport_sentinel_u32x4 –

  Reset one BF16 Lamport fragment to BF16 negative zeros.

• to_cute_dynamic –

  A 1-D view whose length is a runtime shape divisible by divisibility.

• to_cute_dynamic_m –

  Expose exactly one compact tensor mode as a runtime shape.

• warp_sum_specialized –

  Warp sum supporting a compile-time partial final warp.

CUDAGraphCompatibleWrapper

DLPack view that does not synchronize with the producer stream.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

class CUDAGraphCompatibleWrapper:
    """DLPack view that does not synchronize with the producer stream."""

    __slots__ = ("tensor",)

    def __init__(self, tensor: torch.Tensor) -> None:
        self.tensor = tensor

    def __dlpack__(self, stream=None):
        # stream=-1 skips producer sync; CuTe launches on the current stream,
        # which is the capturing one under CUDA graphs.
        return self.tensor.__dlpack__(stream=-1)

    def __dlpack_device__(self):
        return self.tensor.__dlpack_device__()

fragment_has_negative_zero(packed)

Whether a BF16 fragment still holds the BF16 sentinel in any lane.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@cute.jit
def fragment_has_negative_zero(packed):
    """Whether a BF16 fragment still holds the BF16 sentinel in any lane."""
    dirty = False
    for index in cutlass.range_constexpr(4):
        word = packed[index]
        dirty = (
            dirty
            | (Uint16(word & Uint32(0xFFFF)) == Uint16(NEGATIVE_ZERO_BF16_BITS))
            | (Uint16(word >> Uint32(16)) == Uint16(NEGATIVE_ZERO_BF16_BITS))
        )
    return dirty

fragment_is_dirty(packed)

Whether an FP32 fragment still holds the FP32 sentinel in any word.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@cute.jit
def fragment_is_dirty(packed):
    """Whether an FP32 fragment still holds the FP32 sentinel in any word."""
    dirty = packed[0] == Uint32(NEGATIVE_ZERO_F32_BITS)
    for index in cutlass.range_constexpr(1, 4):
        dirty = dirty | (packed[index] == Uint32(NEGATIVE_ZERO_F32_BITS))
    return dirty

group_leader_block_sum(value, warp_sums, warps, leader_stride)

Block sum of the values held by every leader_stride-th lane.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@cute.jit
def group_leader_block_sum(
    value: Float32,
    warp_sums: cute.Tensor,
    warps: cutlass.Constexpr[int],
    leader_stride: cutlass.Constexpr[int],
) -> Float32:
    """Block sum of the values held by every ``leader_stride``-th lane."""
    lane = cute.arch.lane_idx()
    warp = cute.arch.warp_idx()
    for offset in cutlass.range_constexpr(1, WARP_SIZE):
        if cutlass.const_expr(offset >= leader_stride and (offset & (offset - 1)) == 0):
            value = value + cute.arch.shuffle_sync_bfly(
                value,
                offset=offset,
                mask=-1,
                mask_and_clamp=31,
            )
    if lane == 0:
        cute.arch.store((warp_sums + warp).llvm_ptr, value)
    cute.arch.barrier()

    result = Float32(0.0)
    if warp == 0:
        if lane < Int32(warps):
            result = cute.arch.load(
                (warp_sums + lane).llvm_ptr,
                Float32,
            )
        result = cute.arch.warp_reduction_sum(result)
        if lane == 0:
            cute.arch.store(warp_sums.llvm_ptr, result)
    cute.arch.barrier()
    return cute.arch.load(warp_sums.llvm_ptr, Float32)

load_global_u32x4(pointer, *, volatile=False, loc=None, ip=None)

Load one 128-bit fragment as four u32 registers.

The volatile form is the Lamport polling load. Marking the asm side-effecting prevents loop-invariant motion and common-subexpression elimination across polling iterations.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@dsl_user_op
def load_global_u32x4(
    pointer: cute.Pointer,
    *,
    volatile: cutlass.Constexpr[bool] = False,
    loc=None,
    ip=None,
):
    """Load one 128-bit fragment as four u32 registers.

    The volatile form is the Lamport polling load. Marking the asm
    side-effecting prevents loop-invariant motion and common-subexpression
    elimination across polling iterations.
    """
    address = pointer.toint(loc=loc, ip=ip)
    opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32"
    loaded = _asm(
        llvm.StructType.get_literal([T.i32()] * 4),
        [address.ir_value(loc=loc, ip=ip)],
        f"{opcode} {{$0, $1, $2, $3}}, [$4];",
        "=r,=r,=r,=r,l",
        side_effects=volatile,
        loc=loc,
        ip=ip,
    )
    packed = vector.from_elements(
        ir.VectorType.get([4], T.i32(), loc=loc),
        [
            llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
            for index in range(4)
        ],
        loc=loc,
        ip=ip,
    )
    return cute.TensorSSA(packed, 4, Uint32)

load_shared_f32x2(pointer, *, loc=None, ip=None)

Load two aligned FP32 partials from local shared memory.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@dsl_user_op
def load_shared_f32x2(pointer: cute.Pointer, *, loc=None, ip=None):
    """Load two aligned FP32 partials from local shared memory."""
    address = Int32(pointer.toint(loc=loc, ip=ip))
    loaded = _asm(
        llvm.StructType.get_literal([T.f32()] * 2),
        [address.ir_value(loc=loc, ip=ip)],
        "ld.shared.v2.f32 {$0, $1}, [$2];",
        "=f,=f,r",
        side_effects=True,
        loc=loc,
        ip=ip,
    )
    return tuple(
        Float32(llvm.extractvalue(T.f32(), loaded, [index], loc=loc, ip=ip))
        for index in range(2)
    )

load_shared_f32x4(pointer, *, loc=None, ip=None)

Load four aligned FP32 partials from local shared memory.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@dsl_user_op
def load_shared_f32x4(pointer: cute.Pointer, *, loc=None, ip=None):
    """Load four aligned FP32 partials from local shared memory."""
    address = Int32(pointer.toint(loc=loc, ip=ip))
    loaded = _asm(
        llvm.StructType.get_literal([T.f32()] * 4),
        [address.ir_value(loc=loc, ip=ip)],
        "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",
        "=f,=f,=f,=f,r",
        side_effects=True,
        loc=loc,
        ip=ip,
    )
    return tuple(
        Float32(llvm.extractvalue(T.f32(), loaded, [index], loc=loc, ip=ip))
        for index in range(4)
    )

map_shared_to_peer(smem_pointer, peer_rank, *, loc=None, ip=None)

Map a local shared-memory slot to the same slot in a peer CTA.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@dsl_user_op
def map_shared_to_peer(
    smem_pointer: cute.Pointer, peer_rank: Int32, *, loc=None, ip=None
) -> Int32:
    """Map a local shared-memory slot to the same slot in a peer CTA."""
    address = smem_pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
    return Int32(
        _asm(
            T.i32(),
            [address, peer_rank.ir_value(loc=loc, ip=ip)],
            "mapa.shared::cluster.u32 $0, $1, $2;",
            "=r,r,r",
            side_effects=False,
            loc=loc,
            ip=ip,
        )
    )

red_async_release_gpu_add_u32(pointer, value, *, loc=None, ip=None)

The SM100 arrival primitive FlashInfer's LamportFlags use.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@dsl_user_op
def red_async_release_gpu_add_u32(
    pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None
) -> None:
    """The SM100 arrival primitive FlashInfer's LamportFlags use."""
    address = pointer.toint(loc=loc, ip=ip)
    _asm(
        None,
        [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)],
        "red.async.release.global.gpu.add.u32 [$0], $1;",
        "l,r",
        side_effects=True,
        loc=loc,
        ip=ip,
    )

sanitize_negative_zero_u32x4(packed)

Turn real BF16 -0 into +0 so it cannot equal the empty sentinel.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@cute.jit
def sanitize_negative_zero_u32x4(packed):
    """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel."""
    sanitized = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
    for index in cutlass.range_constexpr(4):
        sanitized[index] = sanitize_negative_zero_u32(packed[index])
    return sanitized.load()

stmc_bf16x8(address, packed, *, loc=None, ip=None)

Publish eight BF16 values (four packed words) through NVLS multicast.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@dsl_user_op
def stmc_bf16x8(address: Int64, packed, *, loc=None, ip=None) -> None:
    """Publish eight BF16 values (four packed words) through NVLS multicast."""
    words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)]
    _asm(
        None,
        [address.ir_value(loc=loc, ip=ip), *words],
        "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};",
        "l,r,r,r,r",
        side_effects=True,
        loc=loc,
        ip=ip,
    )

store_global_u32x4(address, packed, *, volatile=False, loc=None, ip=None)

Store four packed words to an ordinary or NVLS multicast global VA.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@dsl_user_op
def store_global_u32x4(
    address: Int64,
    packed,
    *,
    volatile: cutlass.Constexpr[bool] = False,
    loc=None,
    ip=None,
) -> None:
    """Store four packed words to an ordinary or NVLS multicast global VA."""
    words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)]
    opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32"
    _asm(
        None,
        [address.ir_value(loc=loc, ip=ip), *words],
        f"{opcode} [$0], {{$1, $2, $3, $4}};",
        "l,r,r,r,r",
        side_effects=True,
        loc=loc,
        ip=ip,
    )

store_lamport_sentinel_128(pointer, *, loc=None, ip=None)

Reset one FP32 Lamport fragment to four FP32 negative-zero bit patterns.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@dsl_user_op
def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None:
    """Reset one FP32 Lamport fragment to four FP32 negative-zero bit patterns."""
    address = pointer.toint(loc=loc, ip=ip)
    sentinel = Uint32(NEGATIVE_ZERO_F32_BITS).ir_value(loc=loc, ip=ip)
    _asm(
        None,
        [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel],
        "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
        "l,r,r,r,r",
        side_effects=True,
        loc=loc,
        ip=ip,
    )

store_lamport_sentinel_u32x4(address, *, loc=None, ip=None)

Reset one BF16 Lamport fragment to BF16 negative zeros.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@dsl_user_op
def store_lamport_sentinel_u32x4(address: Int64, *, loc=None, ip=None) -> None:
    """Reset one BF16 Lamport fragment to BF16 negative zeros."""
    sentinel = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip)
    _asm(
        None,
        [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel],
        "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
        "l,r,r,r,r",
        side_effects=True,
        loc=loc,
        ip=ip,
    )

to_cute_dynamic(tensor, alignment, *, divisibility)

A 1-D view whose length is a runtime shape divisible by divisibility.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

def to_cute_dynamic(
    tensor: torch.Tensor, alignment: int, *, divisibility: int
) -> cute.Tensor:
    """A 1-D view whose length is a runtime shape divisible by ``divisibility``."""
    return to_cute(tensor, alignment).mark_compact_shape_dynamic(
        mode=0, divisibility=divisibility
    )

to_cute_dynamic_m(tensor, *, mode, assumed_align=16)

Expose exactly one compact tensor mode as a runtime shape.

Model dimensions remain part of the compiled tensor type. Only the token mode is symbolic, so changing M within an op's capacity reuses the same compiled kernel.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

def to_cute_dynamic_m(
    tensor: torch.Tensor, *, mode: int, assumed_align: int = 16
) -> cute.Tensor:
    """Expose exactly one compact tensor mode as a runtime shape.

    Model dimensions remain part of the compiled tensor type. Only the token
    mode is symbolic, so changing M within an op's capacity reuses the same
    compiled kernel.
    """
    return to_cute(tensor, assumed_align).mark_compact_shape_dynamic(
        mode=mode, stride_order=tensor.dim_order()
    )

warp_sum_specialized(value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask)

Warp sum supporting a compile-time partial final warp.

Source code in vllm/models/common/ops/moe_tail_fusion/primitives.py

@cute.jit
def warp_sum_specialized(
    value: Float32,
    warp_idx: Int32,
    lane: Int32,
    warps: cutlass.Constexpr[int],
    last_warp_lanes: cutlass.Constexpr[int],
    last_warp_mask: cutlass.Constexpr[int],
) -> Float32:
    """Warp sum supporting a compile-time partial final warp."""
    if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32):
        for offset in cutlass.range_constexpr(16, 0, -1):
            # range_constexpr does not provide powers-of-two stepping.
            if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
                other = cute.arch.shuffle_sync_bfly(
                    value,
                    offset=offset,
                    mask=last_warp_mask,
                    mask_and_clamp=31,
                )
                if (lane ^ Int32(offset)) < Int32(last_warp_lanes):
                    value = value + other
    else:
        for offset in cutlass.range_constexpr(16, 0, -1):
            if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
                value = value + cute.arch.shuffle_sync_bfly(
                    value,
                    offset=offset,
                    mask=-1,
                    mask_and_clamp=31,
                )
    return value