vllm.models.deepseek_v41.nvidia.ops.cute_dsl

CuTe DSL kernels for DSV4.1 mHC and the TP collectives of its decode.

Modules:

• all_gather –

  A TP all-gather over NVLink multicast with a Lamport mailbox, for the

• all_reduce_mhc –

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

Classes:

• AllReduceMHC –

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

• LamportAllGather –

  TP all-gather along dim 0 over NVLink multicast, for inputs of up to

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. Calls of bt_min_rows rows or more (finalize: bt_min_finalize_rows) use the two-shot BT protocol, with its own mailboxes and generation state.

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. Calls of ``bt_min_rows`` rows or more (finalize:
    ``bt_min_finalize_rows``) use the two-shot BT protocol, with its own
    mailboxes and generation state.
    """

    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(),
            )
            # BT, for instances that can take BT_MIN_ROWS rows; a smaller one
            # (thresholds above its capacity) builds none of its kernels or
            # buffers. Tokens are owned round-robin, so a rank owns at most
            # owner_capacity of a call's rows.
            bt_enabled = capacity >= BT_MIN_ROWS
            self.bt_min_rows = BT_MIN_ROWS if bt_enabled else capacity + 1
            self.bt_min_finalize_rows = (
                BT_MIN_FINALIZE_ROWS if bt_enabled else capacity + 1
            )
            self._owner_capacity = owner_capacity = math.ceil(capacity / tp)
            if bt_enabled:
                self._compile_bt(hidden, top_k, tp, rank, capacity, owner_capacity)
            # 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, reduced_input: 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,
                        reduced_input=reduced_input,
                    ),
                    make_fake_compact_tensor(
                        BFloat16,
                        (
                            LAMPORT_GENERATIONS
                            * (1 if reduced_input else 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(),
                )

            # Consumers by (BT, variant): "plain", "quant" (also the MXFP8
            # layer input) and "residual" (residual streams only).
            self._consumers: dict[tuple[bool, str], Any] = {}
            for bt in (False, True) if bt_enabled else (False,):
                self._consumers[bt, "plain"] = compile_collective(False, False, bt)
                if quantize_layer_input:
                    self._consumers[bt, "quant"] = compile_collective(True, False, bt)
                if residual_only:
                    self._consumers[bt, "residual"] = compile_collective(
                        False, True, bt
                    )
            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)
            if bt_enabled:
                self._allocate_bt(group, tp, capacity, owner_capacity, device)
            torch.accelerator.synchronize()
        dist.barrier(group=group)

    def _compile_bt(
        self,
        hidden: int,
        top_k: int,
        tp: int,
        rank: int,
        capacity: int,
        owner_capacity: int,
    ) -> None:
        """Compile the BT publish, finalize publish and owner-reduce kernels."""
        self._bt_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,
                owner_capacity=owner_capacity,
            ),
            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._bt_finalize_publish = cute.compile(
            _QuadFinalizePublishDeviceKernel(
                hidden=hidden,
                top_k=top_k,
                tp=tp,
                rank=rank,
                capacity_m=capacity,
                threads=128,
                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,
                owner_capacity=owner_capacity,
                elements_per_thread=VEC_BF16,
            ),
            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(),
        )
        # FlashInfer's BT presets for TP4/H5120 reduce with 32 threads a
        # CTA up to 192 rows and 320 above.
        self._owner_reduce = {
            threads: cute.compile(
                _OwnerReduceMulticastDeviceKernel(
                    hidden=hidden,
                    tp=tp,
                    rank=rank,
                    capacity_m=capacity,
                    owner_capacity=owner_capacity,
                    threads=threads,
                    enable_pdl=True,
                ),
                make_fake_compact_tensor(
                    BFloat16,
                    (LAMPORT_GENERATIONS * tp * owner_capacity * hidden,),
                    assumed_align=16,
                ),
                make_fake_compact_tensor(Int32, (2,), assumed_align=4),
                Int64(0),
                Int32(capacity),
                current_cu_stream(),
            )
            for threads in (32, 320)
        }

    def _allocate_bt(
        self,
        group: dist.ProcessGroup,
        tp: int,
        capacity: int,
        owner_capacity: int,
        device: torch.device,
    ) -> None:
        """Allocate the BT mailboxes (collective over ``group``)."""
        # BT: contributions [generation][source rank][owned token] on each
        # owner, stored by unicast through a table of the owners' addresses,
        # and the owners' sums [generation][token] on every rank.
        self._bt_mailbox = symm_mem.empty(
            (LAMPORT_GENERATIONS, tp, owner_capacity, self.hidden_size),
            dtype=torch.bfloat16,
            device=device,
        )
        self._bt_mailbox_handle = symm_mem.rendezvous(self._bt_mailbox, group)
        # get_remote_tensor keeps the rendezvous offset of a pooled buffer.
        owner_mailboxes = [
            self._bt_mailbox_handle.get_remote_tensor(
                peer, self._bt_mailbox.shape, self._bt_mailbox.dtype
            ).data_ptr()
            for peer in range(tp)
        ]
        if not all(owner_mailboxes):
            raise RuntimeError("NVLink peer mapping is unavailable")
        self._bt_owner_mailboxes = torch.tensor(
            owner_mailboxes, dtype=torch.int64, device=device
        )
        self._bt_reduced = symm_mem.empty(
            (LAMPORT_GENERATIONS, capacity, self.hidden_size),
            dtype=torch.bfloat16,
            device=device,
        )
        self._bt_reduced_handle = symm_mem.rendezvous(self._bt_reduced, group)
        bt_multicast = int(self._bt_reduced_handle.multicast_ptr or 0)
        if not bt_multicast or bt_multicast % 16:
            raise RuntimeError("NVLink multicast mapping is unavailable")
        self._bt_multicast = bt_multicast
        self._bt_mailbox.view(torch.int16).fill_(-32768)
        self._bt_reduced.view(torch.int16).fill_(-32768)
        self._bt_stage_state = torch.zeros(2, dtype=torch.int32, device=device)

    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")
        bt = m >= self.bt_min_rows
        consumer = self._consumer(residual, pre, norm_weight, m, mxfp8_out, bt)
        (self._bt_publish if bt else self._publish)(
            to_cute_dynamic(x.flatten(), 16, divisibility=self.hidden_size),
            *self._publish_target(bt),
            Int32(m),
            current_cu_stream(),
        )
        return self._collect(consumer, residual, post, comb, norm_eps, m, bt)

    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]")
        bt = m >= self.bt_min_finalize_rows
        consumer = self._consumer(residual, pre, norm_weight, m, mxfp8_out, bt)
        (self._bt_finalize_publish if bt else 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),
            *self._publish_target(bt),
            Int32(m),
            current_cu_stream(),
        )
        return self._collect(consumer, residual, post, comb, norm_eps, m, bt)

    def _publish_target(self, bt: bool) -> tuple[cute.Tensor, Int64]:
        """The publish kernels' generation state and mailbox address: LL's
        multicast mailbox, or BT's table of the owners' mailboxes."""
        if bt:
            return to_cute(self._bt_stage_state, 4), Int64(
                self._bt_owner_mailboxes.data_ptr()
            )
        return to_cute(self._stage_state, 4), Int64(self._multicast)

    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 _consumer(
        self,
        residual: torch.Tensor,
        pre: torch.Tensor | None,
        norm_weight: torch.Tensor | None,
        m: int,
        mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None,
        bt: bool,
    ) -> _Consumer:
        """Select and validate the consumer variant of a call and allocate its
        outputs. Callers do this before the call's first launch, so a bad
        argument raises with the mailboxes and generation state untouched."""
        hidden = self.hidden_size
        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")
        if pre is None or norm_weight is None:
            if (bt, "residual") not in self._consumers:
                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 (the
            # post-mix has the pre-mix's shape).
            return _Consumer(
                kernel=self._consumers[bt, "residual"],
                pre=None,
                norm_weight=self._no_norm_weight,
                residual_output=residual_output,
                norm_output=residual_output,
                layer_input=None,
                quant_args=self._no_quant_args,
            )
        layer_input = torch.empty(
            (m, hidden), dtype=torch.bfloat16, device=residual.device
        )
        if mxfp8_out is None:
            kernel, quant_args = self._consumers[bt, "plain"], self._no_quant_args
        else:
            if (bt, "quant") not in self._consumers:
                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")
            kernel = self._consumers[bt, "quant"]
            quant_args = (
                to_cute_dynamic(
                    values.view(torch.uint8).flatten(), 16, divisibility=hidden
                ),
                to_cute_dynamic(scales.flatten(), 16, divisibility=16),
            )
        return _Consumer(
            kernel=kernel,
            pre=pre,
            norm_weight=norm_weight,
            residual_output=residual_output,
            norm_output=layer_input,
            layer_input=layer_input,
            quant_args=quant_args,
        )

    def _collect(
        self,
        consumer: _Consumer,
        residual: torch.Tensor,
        post: torch.Tensor,
        comb: torch.Tensor,
        norm_eps: float,
        m: int,
        bt: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Launch the call's reduction after its publish: for BT the owner
        reduce, then the consumer that ``_consumer`` selected."""
        hidden, hc = self.hidden_size, self.hc_mult
        if bt:
            # The owners reduce first; the consumer then reads their sums.
            self._owner_reduce[32 if m <= 192 else 320](
                to_cute(self._bt_mailbox.flatten(), 16),
                to_cute(self._bt_stage_state, 4),
                Int64(self._bt_multicast),
                Int32(m),
                current_cu_stream(),
            )
        pre = post if consumer.pre is None else consumer.pre
        consumer.kernel(
            to_cute((self._bt_reduced if bt else 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(consumer.norm_weight, 16),
            to_cute_dynamic(
                consumer.residual_output.flatten(), 16, divisibility=hc * hidden
            ),
            to_cute_dynamic(consumer.norm_output.flatten(), 16, divisibility=hidden),
            *consumer.quant_args,
            to_cute(self._bt_stage_state if bt else self._stage_state, 4),
            Float32(norm_eps),
            Int32(m),
            current_cu_stream(),
        )
        return consumer.residual_output, consumer.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")
    bt = m >= self.bt_min_rows
    consumer = self._consumer(residual, pre, norm_weight, m, mxfp8_out, bt)
    (self._bt_publish if bt else self._publish)(
        to_cute_dynamic(x.flatten(), 16, divisibility=self.hidden_size),
        *self._publish_target(bt),
        Int32(m),
        current_cu_stream(),
    )
    return self._collect(consumer, residual, post, comb, norm_eps, m, bt)

_allocate_bt(group, tp, capacity, owner_capacity, device)

Allocate the BT mailboxes (collective over group).

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

def _allocate_bt(
    self,
    group: dist.ProcessGroup,
    tp: int,
    capacity: int,
    owner_capacity: int,
    device: torch.device,
) -> None:
    """Allocate the BT mailboxes (collective over ``group``)."""
    # BT: contributions [generation][source rank][owned token] on each
    # owner, stored by unicast through a table of the owners' addresses,
    # and the owners' sums [generation][token] on every rank.
    self._bt_mailbox = symm_mem.empty(
        (LAMPORT_GENERATIONS, tp, owner_capacity, self.hidden_size),
        dtype=torch.bfloat16,
        device=device,
    )
    self._bt_mailbox_handle = symm_mem.rendezvous(self._bt_mailbox, group)
    # get_remote_tensor keeps the rendezvous offset of a pooled buffer.
    owner_mailboxes = [
        self._bt_mailbox_handle.get_remote_tensor(
            peer, self._bt_mailbox.shape, self._bt_mailbox.dtype
        ).data_ptr()
        for peer in range(tp)
    ]
    if not all(owner_mailboxes):
        raise RuntimeError("NVLink peer mapping is unavailable")
    self._bt_owner_mailboxes = torch.tensor(
        owner_mailboxes, dtype=torch.int64, device=device
    )
    self._bt_reduced = symm_mem.empty(
        (LAMPORT_GENERATIONS, capacity, self.hidden_size),
        dtype=torch.bfloat16,
        device=device,
    )
    self._bt_reduced_handle = symm_mem.rendezvous(self._bt_reduced, group)
    bt_multicast = int(self._bt_reduced_handle.multicast_ptr or 0)
    if not bt_multicast or bt_multicast % 16:
        raise RuntimeError("NVLink multicast mapping is unavailable")
    self._bt_multicast = bt_multicast
    self._bt_mailbox.view(torch.int16).fill_(-32768)
    self._bt_reduced.view(torch.int16).fill_(-32768)
    self._bt_stage_state = torch.zeros(2, dtype=torch.int32, device=device)

_collect(consumer, residual, post, comb, norm_eps, m, bt=False)

Launch the call's reduction after its publish: for BT the owner reduce, then the consumer that _consumer selected.

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

def _collect(
    self,
    consumer: _Consumer,
    residual: torch.Tensor,
    post: torch.Tensor,
    comb: torch.Tensor,
    norm_eps: float,
    m: int,
    bt: bool = False,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Launch the call's reduction after its publish: for BT the owner
    reduce, then the consumer that ``_consumer`` selected."""
    hidden, hc = self.hidden_size, self.hc_mult
    if bt:
        # The owners reduce first; the consumer then reads their sums.
        self._owner_reduce[32 if m <= 192 else 320](
            to_cute(self._bt_mailbox.flatten(), 16),
            to_cute(self._bt_stage_state, 4),
            Int64(self._bt_multicast),
            Int32(m),
            current_cu_stream(),
        )
    pre = post if consumer.pre is None else consumer.pre
    consumer.kernel(
        to_cute((self._bt_reduced if bt else 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(consumer.norm_weight, 16),
        to_cute_dynamic(
            consumer.residual_output.flatten(), 16, divisibility=hc * hidden
        ),
        to_cute_dynamic(consumer.norm_output.flatten(), 16, divisibility=hidden),
        *consumer.quant_args,
        to_cute(self._bt_stage_state if bt else self._stage_state, 4),
        Float32(norm_eps),
        Int32(m),
        current_cu_stream(),
    )
    return consumer.residual_output, consumer.layer_input

_compile_bt(hidden, top_k, tp, rank, capacity, owner_capacity)

Compile the BT publish, finalize publish and owner-reduce kernels.

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

def _compile_bt(
    self,
    hidden: int,
    top_k: int,
    tp: int,
    rank: int,
    capacity: int,
    owner_capacity: int,
) -> None:
    """Compile the BT publish, finalize publish and owner-reduce kernels."""
    self._bt_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,
            owner_capacity=owner_capacity,
        ),
        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._bt_finalize_publish = cute.compile(
        _QuadFinalizePublishDeviceKernel(
            hidden=hidden,
            top_k=top_k,
            tp=tp,
            rank=rank,
            capacity_m=capacity,
            threads=128,
            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,
            owner_capacity=owner_capacity,
            elements_per_thread=VEC_BF16,
        ),
        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(),
    )
    # FlashInfer's BT presets for TP4/H5120 reduce with 32 threads a
    # CTA up to 192 rows and 320 above.
    self._owner_reduce = {
        threads: cute.compile(
            _OwnerReduceMulticastDeviceKernel(
                hidden=hidden,
                tp=tp,
                rank=rank,
                capacity_m=capacity,
                owner_capacity=owner_capacity,
                threads=threads,
                enable_pdl=True,
            ),
            make_fake_compact_tensor(
                BFloat16,
                (LAMPORT_GENERATIONS * tp * owner_capacity * hidden,),
                assumed_align=16,
            ),
            make_fake_compact_tensor(Int32, (2,), assumed_align=4),
            Int64(0),
            Int32(capacity),
            current_cu_stream(),
        )
        for threads in (32, 320)
    }

_consumer(residual, pre, norm_weight, m, mxfp8_out, bt)

Select and validate the consumer variant of a call and allocate its outputs. Callers do this before the call's first launch, so a bad argument raises with the mailboxes and generation state untouched.

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

def _consumer(
    self,
    residual: torch.Tensor,
    pre: torch.Tensor | None,
    norm_weight: torch.Tensor | None,
    m: int,
    mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None,
    bt: bool,
) -> _Consumer:
    """Select and validate the consumer variant of a call and allocate its
    outputs. Callers do this before the call's first launch, so a bad
    argument raises with the mailboxes and generation state untouched."""
    hidden = self.hidden_size
    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")
    if pre is None or norm_weight is None:
        if (bt, "residual") not in self._consumers:
            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 (the
        # post-mix has the pre-mix's shape).
        return _Consumer(
            kernel=self._consumers[bt, "residual"],
            pre=None,
            norm_weight=self._no_norm_weight,
            residual_output=residual_output,
            norm_output=residual_output,
            layer_input=None,
            quant_args=self._no_quant_args,
        )
    layer_input = torch.empty(
        (m, hidden), dtype=torch.bfloat16, device=residual.device
    )
    if mxfp8_out is None:
        kernel, quant_args = self._consumers[bt, "plain"], self._no_quant_args
    else:
        if (bt, "quant") not in self._consumers:
            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")
        kernel = self._consumers[bt, "quant"]
        quant_args = (
            to_cute_dynamic(
                values.view(torch.uint8).flatten(), 16, divisibility=hidden
            ),
            to_cute_dynamic(scales.flatten(), 16, divisibility=16),
        )
    return _Consumer(
        kernel=kernel,
        pre=pre,
        norm_weight=norm_weight,
        residual_output=residual_output,
        norm_output=layer_input,
        layer_input=layer_input,
        quant_args=quant_args,
    )

_publish_target(bt)

The publish kernels' generation state and mailbox address: LL's multicast mailbox, or BT's table of the owners' mailboxes.

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

def _publish_target(self, bt: bool) -> tuple[cute.Tensor, Int64]:
    """The publish kernels' generation state and mailbox address: LL's
    multicast mailbox, or BT's table of the owners' mailboxes."""
    if bt:
        return to_cute(self._bt_stage_state, 4), Int64(
            self._bt_owner_mailboxes.data_ptr()
        )
    return to_cute(self._stage_state, 4), Int64(self._multicast)

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]")
    bt = m >= self.bt_min_finalize_rows
    consumer = self._consumer(residual, pre, norm_weight, m, mxfp8_out, bt)
    (self._bt_finalize_publish if bt else 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),
        *self._publish_target(bt),
        Int32(m),
        current_cu_stream(),
    )
    return self._collect(consumer, residual, post, comb, norm_eps, m, bt)

LamportAllGather

TP all-gather along dim 0 over NVLink multicast, for inputs of up to rank_capacity_bytes a rank (a multiple of 16 bytes).

Construction is collective over the TP group: every rank compiles the kernel and allocates its mailbox, and it raises when NVLink multicast is unavailable. accepts tells whether a tensor fits, from its size alone, so every rank decides alike.

Methods:

• __call__ –

  Every rank's inp concatenated along dim 0 in rank order, on the

• accepts –

  Whether inp fits: 1 to rank_capacity_bytes bytes, a

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

class LamportAllGather:
    """TP all-gather along dim 0 over NVLink multicast, for inputs of up to
    ``rank_capacity_bytes`` a rank (a multiple of 16 bytes).

    Construction is collective over the TP group: every rank compiles the
    kernel and allocates its mailbox, and it raises when NVLink multicast is
    unavailable. ``accepts`` tells whether a tensor fits, from its size
    alone, so every rank decides alike.
    """

    def __init__(self, *, rank_capacity_bytes: int, device: torch.device) -> None:
        if rank_capacity_bytes <= 0 or rank_capacity_bytes % PACK_BYTES:
            raise ValueError("the capacity must be a positive multiple of 16 bytes")
        group = get_tp_group().device_group
        self.world_size = world_size = dist.get_world_size(group)
        rank = dist.get_rank(group)
        self.rank_capacity_bytes = rank_capacity_bytes
        rank_packs = rank_capacity_bytes // PACK_BYTES
        device = torch.device(device)
        with torch.accelerator.device_index(device.index):
            # cute.compile returns an untyped JIT executor.
            self._kernel: Any = cute.compile(
                _AllGatherKernel(rank, world_size, rank_packs),
                Int64(0),
                Int64(0),
                Int64(0),
                Int64(0),
                Int64(0),
                Int32(1),
                make_fake_stream(),
            )
            self._mailbox = symm_mem.empty(
                (GENERATIONS * world_size * rank_packs * PACK_BYTES // 4,),
                dtype=torch.int32,
                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 % PACK_BYTES:
                raise RuntimeError("NVLink multicast mapping is unavailable")
            self._multicast = multicast
            # Every slot starts as the sentinel.
            self._mailbox.fill_(-(2**31))
            self._state = torch.zeros(2, dtype=torch.int32, device=device)
            torch.accelerator.synchronize()
        dist.barrier(group=group)

    def accepts(self, inp: torch.Tensor) -> bool:
        """Whether ``inp`` fits: 1 to ``rank_capacity_bytes`` bytes, a
        multiple of 16."""
        return 0 < inp.nbytes <= self.rank_capacity_bytes and not inp.nbytes % 16

    def __call__(self, inp: torch.Tensor) -> torch.Tensor:
        """Every rank's ``inp`` concatenated along dim 0 in rank order, on the
        current stream. Words equal to ``0x80000000`` arrive as 0, on every
        rank (FP32 negative zero becomes positive zero)."""
        if not self.accepts(inp):
            raise ValueError(
                f"the all-gather takes 16-byte multiples of up to "
                f"{self.rank_capacity_bytes} bytes a rank, not {inp.nbytes}"
            )
        if not inp.is_contiguous() or inp.data_ptr() % PACK_BYTES:
            # The kernel loads 16-byte packs.
            inp = inp.clone(memory_format=torch.contiguous_format)
        out = torch.empty(
            (inp.shape[0] * self.world_size, *inp.shape[1:]),
            dtype=inp.dtype,
            device=inp.device,
        )
        self._kernel(
            Int64(inp.data_ptr()),
            Int64(out.data_ptr()),
            Int64(self._mailbox.data_ptr()),
            Int64(self._multicast),
            Int64(self._state.data_ptr()),
            Int32(inp.nbytes // PACK_BYTES),
            current_cu_stream(),
        )
        return out

__call__(inp)

Every rank's inp concatenated along dim 0 in rank order, on the current stream. Words equal to 0x80000000 arrive as 0, on every rank (FP32 negative zero becomes positive zero).

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

def __call__(self, inp: torch.Tensor) -> torch.Tensor:
    """Every rank's ``inp`` concatenated along dim 0 in rank order, on the
    current stream. Words equal to ``0x80000000`` arrive as 0, on every
    rank (FP32 negative zero becomes positive zero)."""
    if not self.accepts(inp):
        raise ValueError(
            f"the all-gather takes 16-byte multiples of up to "
            f"{self.rank_capacity_bytes} bytes a rank, not {inp.nbytes}"
        )
    if not inp.is_contiguous() or inp.data_ptr() % PACK_BYTES:
        # The kernel loads 16-byte packs.
        inp = inp.clone(memory_format=torch.contiguous_format)
    out = torch.empty(
        (inp.shape[0] * self.world_size, *inp.shape[1:]),
        dtype=inp.dtype,
        device=inp.device,
    )
    self._kernel(
        Int64(inp.data_ptr()),
        Int64(out.data_ptr()),
        Int64(self._mailbox.data_ptr()),
        Int64(self._multicast),
        Int64(self._state.data_ptr()),
        Int32(inp.nbytes // PACK_BYTES),
        current_cu_stream(),
    )
    return out

accepts(inp)

Whether inp fits: 1 to rank_capacity_bytes bytes, a multiple of 16.

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

def accepts(self, inp: torch.Tensor) -> bool:
    """Whether ``inp`` fits: 1 to ``rank_capacity_bytes`` bytes, a
    multiple of 16."""
    return 0 < inp.nbytes <= self.rank_capacity_bytes and not inp.nbytes % 16