vllm.models.deepseek_v41.nvidia.ops.cute_dsl

CuTe DSL kernels for DSV4.1 mHC.

Modules:

• all_reduce_mhc –

  Lamport all-reduce fused with DSV4.1's mHC post, collapse and RMSNorm.

• primitives –

  PTX building blocks and launch helpers for the DSV4.1 LL all-reduce.

Classes:

• AllReduceMHC –

  Lamport TP all-reduce fused with mHC post, collapse and RMSNorm.

AllReduceMHC

Lamport TP all-reduce fused with mHC post, collapse and RMSNorm.

__call__ reduces x across the TP group, writes the post-mixed hc streams and the normalized collapse, and leaves the next mixes to the caller. finalize does the same for an unfinalized MoE output, folding the top-k reduction and the shared-expert add into the publish. Each instance owns a Lamport mailbox; construction is collective over the TP group. With quantize_layer_input, both also accept mxfp8_out buffers that receive the normalized collapse as MXFP8. With residual_only, both also accept pre and norm_weight of None and then write only the post-mixed streams, for a consumer that collapses them itself.

Methods:

• __call__ –

  Return (post-mixed residual streams, normalized layer input).

• finalize –

  __call__ for a MoE output left unfinalized.

Source code in vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py

class AllReduceMHC:
    """Lamport TP all-reduce fused with mHC post, collapse and RMSNorm.

    ``__call__`` reduces ``x`` across the TP group, writes the post-mixed hc
    streams and the normalized collapse, and leaves the next mixes to the
    caller. ``finalize`` does the same for an unfinalized MoE output, folding
    the top-k reduction and the shared-expert add into the publish. Each
    instance owns a Lamport mailbox; construction is collective over the TP
    group. With ``quantize_layer_input``, both also accept ``mxfp8_out``
    buffers that receive the normalized collapse as MXFP8. With
    ``residual_only``, both also accept ``pre`` and ``norm_weight`` of None and
    then write only the post-mixed streams, for a consumer that collapses them
    itself.
    """

    def __init__(
        self,
        *,
        hidden_size: int,
        hc_mult: int,
        max_num_tokens: int,
        top_k: int,
        device: torch.device,
        quantize_layer_input: bool = False,
        residual_only: bool = False,
    ) -> None:
        group = get_tp_group().device_group
        self.hidden_size = hidden = hidden_size
        self.hc_mult = hc = hc_mult
        self.capacity = capacity = max_num_tokens
        self.top_k = top_k
        self.quantize_layer_input = quantize_layer_input
        self.residual_only = residual_only
        tp = dist.get_world_size(group)
        rank = dist.get_rank(group)
        device = torch.device(device)

        # FlashInfer's LL_ALL_REDUCE_GB300_TP{4,8}_H5120 and
        # LL_FINALIZE_..._H5120_K* presets: a cluster of 5 x 128 threads covers
        # a 640-fragment token in one fully used trip, and the finalize stages
        # all top_k rows at once.
        with torch.accelerator.device_index(device.index):
            self._publish = cute.compile(
                _SharedOnlyPublishDeviceKernel(
                    hidden=hidden,
                    tp=tp,
                    rank=rank,
                    capacity_m=capacity,
                    elements_per_thread=VEC_BF16,
                    threads=128,
                    release_before_store=False,
                    enable_pdl=True,
                ),
                make_fake_dynamic_compact_tensor(
                    BFloat16, alignment=16, divisibility=hidden
                ),
                make_fake_compact_tensor(Int32, (2,), assumed_align=4),
                Int64(0),
                Int32(capacity),
                current_cu_stream(),
            )
            self._finalize_publish = cute.compile(
                _QuadFinalizePublishDeviceKernel(
                    hidden=hidden,
                    top_k=top_k,
                    tp=tp,
                    rank=rank,
                    capacity_m=capacity,
                    threads=128,
                    # DSV4.1's router folds the scale into the weights.
                    routed_scaling_factor=1.0,
                    include_shared_expert=True,
                    load_shared_expert_before_pdl=False,
                    enable_pdl=True,
                    prefetch_group=top_k,
                    fp32_weights=True,
                ),
                make_fake_dynamic_compact_tensor(
                    BFloat16, alignment=16, divisibility=hidden
                ),
                make_fake_dynamic_compact_tensor(
                    Float32, alignment=4, divisibility=top_k
                ),
                make_fake_dynamic_compact_tensor(
                    Int32, alignment=4, divisibility=top_k
                ),
                make_fake_dynamic_compact_tensor(
                    BFloat16, alignment=16, divisibility=hidden
                ),
                make_fake_compact_tensor(Int32, (2,), assumed_align=4),
                Int64(0),
                Int32(capacity),
                current_cu_stream(),
            )
            # Stand-ins for the MXFP8 outputs of the unquantized variants,
            # which never touch them.
            self._no_quant_output = torch.empty(16, dtype=torch.uint8, device=device)
            no_quant_cute = to_cute(self._no_quant_output, 16)
            self._no_quant_args = (no_quant_cute, no_quant_cute)

            def compile_collective(quantize: bool, residual_only: bool):
                if quantize:
                    quant_fakes = (
                        make_fake_dynamic_compact_tensor(
                            Uint8, alignment=16, divisibility=hidden
                        ),
                        make_fake_dynamic_compact_tensor(
                            Uint8, alignment=16, divisibility=16
                        ),
                    )
                else:
                    quant_fakes = (
                        make_fake_compact_tensor(Uint8, (16,), assumed_align=16),
                    ) * 2
                return cute.compile(
                    _LamportMHCDeviceKernel(
                        hidden=hidden,
                        hc=hc,
                        tp=tp,
                        capacity_m=capacity,
                        cluster_size=5,
                        rank_lanes=1,
                        threads=128,
                        enable_pdl=True,
                        quantize=quantize,
                        residual_only=residual_only,
                    ),
                    make_fake_compact_tensor(
                        BFloat16,
                        (LAMPORT_GENERATIONS * tp * capacity * hidden,),
                        assumed_align=16,
                    ),
                    make_fake_dynamic_compact_tensor(
                        BFloat16, alignment=16, divisibility=hc * hidden
                    ),
                    make_fake_dynamic_compact_tensor(
                        Float32, alignment=4, divisibility=hc
                    ),
                    make_fake_dynamic_compact_tensor(
                        Float32, alignment=4, divisibility=hc * hc
                    ),
                    make_fake_dynamic_compact_tensor(
                        Float32, alignment=4, divisibility=hc
                    ),
                    make_fake_compact_tensor(BFloat16, (hidden,), assumed_align=16),
                    make_fake_dynamic_compact_tensor(
                        BFloat16, alignment=16, divisibility=hc * hidden
                    ),
                    make_fake_dynamic_compact_tensor(
                        BFloat16, alignment=16, divisibility=hidden
                    ),
                    *quant_fakes,
                    make_fake_compact_tensor(Int32, (2,), assumed_align=4),
                    Float32(0.0),
                    Int32(capacity),
                    current_cu_stream(),
                )

            self._collective = compile_collective(quantize=False, residual_only=False)
            self._collective_quant = (
                compile_collective(quantize=True, residual_only=False)
                if quantize_layer_input
                else None
            )
            self._collective_residual = (
                compile_collective(quantize=False, residual_only=True)
                if residual_only
                else None
            )
            if residual_only:
                # Stand-in for the norm weight the variant never reads.
                self._no_norm_weight = torch.empty(
                    hidden, dtype=torch.bfloat16, device=device
                )

            self._mailbox = symm_mem.empty(
                (LAMPORT_GENERATIONS, tp, capacity, hidden),
                dtype=torch.bfloat16,
                device=device,
            )
            self._mailbox_handle = symm_mem.rendezvous(self._mailbox, group)
            multicast = int(self._mailbox_handle.multicast_ptr or 0)
            if not multicast or multicast % 16:
                raise RuntimeError("NVLink multicast mapping is unavailable")
            self._multicast = multicast
            # Every slot starts as the Lamport sentinel, BF16 negative zero.
            self._mailbox.view(torch.int16).fill_(-32768)
            self._stage_state = torch.zeros(2, dtype=torch.int32, device=device)
            torch.accelerator.synchronize()
        dist.barrier(group=group)

    def __call__(
        self,
        x: torch.Tensor,
        residual: torch.Tensor,
        post: torch.Tensor,
        comb: torch.Tensor,
        pre: torch.Tensor | None,
        norm_weight: torch.Tensor | None,
        norm_eps: float,
        mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Return (post-mixed residual streams, normalized layer input).

        Args:
            x: [M, hidden] BF16 TP-partial sublayer output.
            residual: [M, hc, hidden] BF16 residual streams.
            post: [M, hc(, 1)] FP32 post-mix.
            comb: [M, hc, hc] FP32 residual mix, indexed [source, target].
            pre: [M, hc] FP32 carried pre-mix for the collapse. None (with
                ``norm_weight``, built with ``residual_only``): write only the
                residual streams and return no layer input.
            norm_weight: [hidden] BF16 RMSNorm weight.
            norm_eps: RMSNorm epsilon.
            mxfp8_out: optional [M, hidden] FP8 E4M3 values and swizzled
                128x4 UE8M0 scales (``mxfp8_scale_bytes(M, hidden)`` bytes)
                that also receive the layer input as MXFP8.

        """
        m = x.shape[0]
        self._check(m, residual, post, comb, pre)
        if x.shape != (m, self.hidden_size):
            raise ValueError("unexpected x shape")
        self._publish(
            to_cute_dynamic(x.flatten(), 16, divisibility=self.hidden_size),
            to_cute(self._stage_state, 4),
            Int64(self._multicast),
            Int32(m),
            current_cu_stream(),
        )
        return self._collect(
            residual, post, comb, pre, norm_weight, norm_eps, m, mxfp8_out
        )

    def finalize(
        self,
        gemm2_permuted: torch.Tensor,
        expert_weights: torch.Tensor,
        expanded_idx_to_permuted_idx: torch.Tensor,
        shared_output: torch.Tensor,
        residual: torch.Tensor,
        post: torch.Tensor,
        comb: torch.Tensor,
        pre: torch.Tensor | None,
        norm_weight: torch.Tensor | None,
        norm_eps: float,
        mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """``__call__`` for a MoE output left unfinalized.

        Args:
            gemm2_permuted: [rows, hidden] BF16 unweighted GEMM2 output in the
                MoE's permuted order.
            expert_weights: [M, top_k] FP32 routing weights, already scaled.
            expanded_idx_to_permuted_idx: [M, top_k] int32 row of each route,
                -1 for an expert this rank does not hold.
            shared_output: [M, hidden] BF16 TP-partial shared-expert output.
            residual: [M, hc, hidden] BF16 residual streams.
            post: [M, hc(, 1)] FP32 post-mix.
            comb: [M, hc, hc] FP32 residual mix, indexed [source, target].
            pre: [M, hc] FP32 carried pre-mix for the collapse (None: residual
                only, as in ``__call__``).
            norm_weight: [hidden] BF16 RMSNorm weight (None with ``pre``).
            norm_eps: RMSNorm epsilon.
            mxfp8_out: as in ``__call__``.

        """
        m = shared_output.shape[0]
        hidden, top_k = self.hidden_size, self.top_k
        self._check(m, residual, post, comb, pre)
        if gemm2_permuted.dim() != 2 or gemm2_permuted.shape[1] != hidden:
            raise ValueError("gemm2_permuted must be [rows, hidden]")
        if shared_output.shape != (m, hidden):
            raise ValueError("unexpected shared_output shape")
        if expert_weights.dtype != torch.float32 or expert_weights.numel() != m * top_k:
            raise ValueError("expert_weights must be FP32 [M, top_k]")
        if (
            expanded_idx_to_permuted_idx.dtype != torch.int32
            or expanded_idx_to_permuted_idx.numel() != m * top_k
        ):
            raise ValueError("expanded_idx_to_permuted_idx must be int32 [M, top_k]")
        self._finalize_publish(
            to_cute_dynamic(gemm2_permuted.flatten(), 16, divisibility=hidden),
            to_cute_dynamic(expert_weights.flatten(), 4, divisibility=top_k),
            to_cute_dynamic(
                expanded_idx_to_permuted_idx.flatten(), 4, divisibility=top_k
            ),
            to_cute_dynamic(shared_output.flatten(), 16, divisibility=hidden),
            to_cute(self._stage_state, 4),
            Int64(self._multicast),
            Int32(m),
            current_cu_stream(),
        )
        return self._collect(
            residual, post, comb, pre, norm_weight, norm_eps, m, mxfp8_out
        )

    def _check(
        self,
        m: int,
        residual: torch.Tensor,
        post: torch.Tensor,
        comb: torch.Tensor,
        pre: torch.Tensor | None,
    ) -> None:
        hc = self.hc_mult
        if not 1 <= m <= self.capacity:
            raise ValueError(f"M={m} is outside [1, {self.capacity}]")
        if residual.shape != (m, hc, self.hidden_size):
            raise ValueError("unexpected residual shape")
        if (
            post.numel() != m * hc
            or (pre is not None and pre.numel() != m * hc)
            or comb.numel() != m * hc * hc
        ):
            raise ValueError("unexpected mix shapes")

    def _collect(
        self,
        residual: torch.Tensor,
        post: torch.Tensor,
        comb: torch.Tensor,
        pre: torch.Tensor | None,
        norm_weight: torch.Tensor | None,
        norm_eps: float,
        m: int,
        mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        hidden, hc = self.hidden_size, self.hc_mult
        residual_output = torch.empty_like(residual)
        if (pre is None) != (norm_weight is None):
            raise ValueError("pre and norm_weight are both given or both None")
        layer_input: torch.Tensor | None = None
        if pre is None or norm_weight is None:
            if self._collective_residual is None:
                raise ValueError("built without residual_only")
            if mxfp8_out is not None:
                raise ValueError("a residual-only boundary has no layer input")
            # The variant reads neither the pre-mix nor the norm weight and
            # writes no layer input: stand-ins fill their arguments.
            collective, quant_args = self._collective_residual, self._no_quant_args
            pre, norm_weight, norm_output = post, self._no_norm_weight, residual_output
        else:
            layer_input = norm_output = torch.empty(
                (m, hidden), dtype=torch.bfloat16, device=residual.device
            )
            if mxfp8_out is None:
                collective, quant_args = self._collective, self._no_quant_args
            else:
                if self._collective_quant is None:
                    raise ValueError("built without quantize_layer_input")
                values, scales = mxfp8_out
                if values.shape != (m, hidden) or scales.numel() != (
                    mxfp8_scale_bytes(m, hidden)
                ):
                    raise ValueError("mxfp8_out does not match the layer input")
                collective = self._collective_quant
                quant_args = (
                    to_cute_dynamic(
                        values.view(torch.uint8).flatten(), 16, divisibility=hidden
                    ),
                    to_cute_dynamic(scales.flatten(), 16, divisibility=16),
                )
        collective(
            to_cute(self._mailbox.flatten(), 16),
            to_cute_dynamic(residual.flatten(), 16, divisibility=hc * hidden),
            to_cute_dynamic(post.flatten(), 4, divisibility=hc),
            to_cute_dynamic(comb.flatten(), 4, divisibility=hc * hc),
            to_cute_dynamic(pre.flatten(), 4, divisibility=hc),
            to_cute(norm_weight, 16),
            to_cute_dynamic(residual_output.flatten(), 16, divisibility=hc * hidden),
            to_cute_dynamic(norm_output.flatten(), 16, divisibility=hidden),
            *quant_args,
            to_cute(self._stage_state, 4),
            Float32(norm_eps),
            Int32(m),
            current_cu_stream(),
        )
        return residual_output, layer_input

__call__(x, residual, post, comb, pre, norm_weight, norm_eps, mxfp8_out=None)

Return (post-mixed residual streams, normalized layer input).

Parameters:

• x

  (Tensor) –

  [M, hidden] BF16 TP-partial sublayer output.

• residual

  (Tensor) –

  [M, hc, hidden] BF16 residual streams.

• post

  (Tensor) –

  [M, hc(, 1)] FP32 post-mix.

• comb

  (Tensor) –

  [M, hc, hc] FP32 residual mix, indexed [source, target].

• pre

  (Tensor | None) –

  [M, hc] FP32 carried pre-mix for the collapse. None (with norm_weight, built with residual_only): write only the residual streams and return no layer input.

• norm_weight

  (Tensor | None) –

  [hidden] BF16 RMSNorm weight.

• norm_eps

  (float) –

  RMSNorm epsilon.

• mxfp8_out

  (tuple[Tensor, Tensor] | None, default: None ) –

  optional [M, hidden] FP8 E4M3 values and swizzled 128x4 UE8M0 scales (mxfp8_scale_bytes(M, hidden) bytes) that also receive the layer input as MXFP8.

Source code in vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py

def __call__(
    self,
    x: torch.Tensor,
    residual: torch.Tensor,
    post: torch.Tensor,
    comb: torch.Tensor,
    pre: torch.Tensor | None,
    norm_weight: torch.Tensor | None,
    norm_eps: float,
    mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Return (post-mixed residual streams, normalized layer input).

    Args:
        x: [M, hidden] BF16 TP-partial sublayer output.
        residual: [M, hc, hidden] BF16 residual streams.
        post: [M, hc(, 1)] FP32 post-mix.
        comb: [M, hc, hc] FP32 residual mix, indexed [source, target].
        pre: [M, hc] FP32 carried pre-mix for the collapse. None (with
            ``norm_weight``, built with ``residual_only``): write only the
            residual streams and return no layer input.
        norm_weight: [hidden] BF16 RMSNorm weight.
        norm_eps: RMSNorm epsilon.
        mxfp8_out: optional [M, hidden] FP8 E4M3 values and swizzled
            128x4 UE8M0 scales (``mxfp8_scale_bytes(M, hidden)`` bytes)
            that also receive the layer input as MXFP8.

    """
    m = x.shape[0]
    self._check(m, residual, post, comb, pre)
    if x.shape != (m, self.hidden_size):
        raise ValueError("unexpected x shape")
    self._publish(
        to_cute_dynamic(x.flatten(), 16, divisibility=self.hidden_size),
        to_cute(self._stage_state, 4),
        Int64(self._multicast),
        Int32(m),
        current_cu_stream(),
    )
    return self._collect(
        residual, post, comb, pre, norm_weight, norm_eps, m, mxfp8_out
    )

finalize(gemm2_permuted, expert_weights, expanded_idx_to_permuted_idx, shared_output, residual, post, comb, pre, norm_weight, norm_eps, mxfp8_out=None)

__call__ for a MoE output left unfinalized.

Parameters:

• gemm2_permuted

  (Tensor) –

  [rows, hidden] BF16 unweighted GEMM2 output in the MoE's permuted order.

• expert_weights

  (Tensor) –

  [M, top_k] FP32 routing weights, already scaled.

• expanded_idx_to_permuted_idx

  (Tensor) –

  [M, top_k] int32 row of each route, -1 for an expert this rank does not hold.

• shared_output

  (Tensor) –

  [M, hidden] BF16 TP-partial shared-expert output.

• residual

  (Tensor) –

  [M, hc, hidden] BF16 residual streams.

• post

  (Tensor) –

  [M, hc(, 1)] FP32 post-mix.

• comb

  (Tensor) –

  [M, hc, hc] FP32 residual mix, indexed [source, target].

• pre

  (Tensor | None) –

  [M, hc] FP32 carried pre-mix for the collapse (None: residual only, as in __call__).

• norm_weight

  (Tensor | None) –

  [hidden] BF16 RMSNorm weight (None with pre).

• norm_eps

  (float) –

  RMSNorm epsilon.

• mxfp8_out

  (tuple[Tensor, Tensor] | None, default: None ) –

  as in __call__.

Source code in vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py

def finalize(
    self,
    gemm2_permuted: torch.Tensor,
    expert_weights: torch.Tensor,
    expanded_idx_to_permuted_idx: torch.Tensor,
    shared_output: torch.Tensor,
    residual: torch.Tensor,
    post: torch.Tensor,
    comb: torch.Tensor,
    pre: torch.Tensor | None,
    norm_weight: torch.Tensor | None,
    norm_eps: float,
    mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """``__call__`` for a MoE output left unfinalized.

    Args:
        gemm2_permuted: [rows, hidden] BF16 unweighted GEMM2 output in the
            MoE's permuted order.
        expert_weights: [M, top_k] FP32 routing weights, already scaled.
        expanded_idx_to_permuted_idx: [M, top_k] int32 row of each route,
            -1 for an expert this rank does not hold.
        shared_output: [M, hidden] BF16 TP-partial shared-expert output.
        residual: [M, hc, hidden] BF16 residual streams.
        post: [M, hc(, 1)] FP32 post-mix.
        comb: [M, hc, hc] FP32 residual mix, indexed [source, target].
        pre: [M, hc] FP32 carried pre-mix for the collapse (None: residual
            only, as in ``__call__``).
        norm_weight: [hidden] BF16 RMSNorm weight (None with ``pre``).
        norm_eps: RMSNorm epsilon.
        mxfp8_out: as in ``__call__``.

    """
    m = shared_output.shape[0]
    hidden, top_k = self.hidden_size, self.top_k
    self._check(m, residual, post, comb, pre)
    if gemm2_permuted.dim() != 2 or gemm2_permuted.shape[1] != hidden:
        raise ValueError("gemm2_permuted must be [rows, hidden]")
    if shared_output.shape != (m, hidden):
        raise ValueError("unexpected shared_output shape")
    if expert_weights.dtype != torch.float32 or expert_weights.numel() != m * top_k:
        raise ValueError("expert_weights must be FP32 [M, top_k]")
    if (
        expanded_idx_to_permuted_idx.dtype != torch.int32
        or expanded_idx_to_permuted_idx.numel() != m * top_k
    ):
        raise ValueError("expanded_idx_to_permuted_idx must be int32 [M, top_k]")
    self._finalize_publish(
        to_cute_dynamic(gemm2_permuted.flatten(), 16, divisibility=hidden),
        to_cute_dynamic(expert_weights.flatten(), 4, divisibility=top_k),
        to_cute_dynamic(
            expanded_idx_to_permuted_idx.flatten(), 4, divisibility=top_k
        ),
        to_cute_dynamic(shared_output.flatten(), 16, divisibility=hidden),
        to_cute(self._stage_state, 4),
        Int64(self._multicast),
        Int32(m),
        current_cu_stream(),
    )
    return self._collect(
        residual, post, comb, pre, norm_weight, norm_eps, m, mxfp8_out
    )