vllm.models.deepseek_v41.nvidia.ops.cute_dsl.all_reduce_mhc

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

FlashInfer's LL kernels, copied: a publish kernel (or its MoE finalize variant) multicasts each rank's contribution into a three-generation Lamport mailbox, and a PDL-chained consumer spins until every rank's fragment lands, reduces it and normalizes. The consumer's residual add and copy are replaced by the mHC post-mix of the hc residual streams and their pre-mix collapse, which the RMSNorm then normalizes. The finalize variant also reads FP32 routing weights.

Optionally the consumer also writes the normalized output as MXFP8 with UE8M0 scales in the swizzled 128x4 layout FlashInfer's MXFP8 GEMMs read. It quantizes the BF16-rounded output with FlashInfer's own quantize helpers, so the result is bit-identical to mxfp8_quantize of the BF16 output, without its launch.

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
    )

_LamportMHCDeviceKernel

FlashInfer's _LamportResidualRMSNormDeviceKernel with mHC.

The residual add and residual copy become the mHC post-mix of the hc streams, stored as the residual output, and their pre-mix collapse, which the unchanged RMSNorm normalizes. With quantize, the normalized output is also written as MXFP8 with swizzled 128x4 UE8M0 scales, including the zero scales of the padding rows up to a multiple of 128. With residual_only, only the residual output is written; the pre-mix, gamma and normalized output are unused.

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

class _LamportMHCDeviceKernel:
    """FlashInfer's ``_LamportResidualRMSNormDeviceKernel`` with mHC.

    The residual add and residual copy become the mHC post-mix of the hc
    streams, stored as the residual output, and their pre-mix collapse, which
    the unchanged RMSNorm normalizes. With ``quantize``, the normalized output
    is also written as MXFP8 with swizzled 128x4 UE8M0 scales, including the
    zero scales of the padding rows up to a multiple of 128. With
    ``residual_only``, only the residual output is written; the pre-mix, gamma
    and normalized output are unused.
    """

    def __init__(
        self,
        *,
        hidden: int,
        hc: int,
        tp: int,
        capacity_m: int,
        cluster_size: int,
        rank_lanes: int,
        threads: int,
        enable_pdl: bool,
        quantize: bool = False,
        residual_only: bool = False,
    ) -> None:
        if rank_lanes not in (1, 2, 4, 8):
            raise ValueError("rank_lanes must be 1, 2, 4, or 8")
        if tp % rank_lanes:
            raise ValueError("tp must be divisible by rank_lanes")
        if threads <= 0 or threads % WARP_SIZE or threads % rank_lanes:
            raise ValueError("threads must be a positive warp and rank-lane multiple")
        if hidden <= 0 or hidden % VEC_BF16:
            raise ValueError("hidden must be a positive multiple of 8")
        # A scale block's four fragments must sit in four adjacent lanes.
        if quantize and (rank_lanes != 1 or hidden % (MXFP8_BLOCK * SCALE_TILE_COLS)):
            raise ValueError("MXFP8 output needs rank_lanes == 1 and hidden % 128 == 0")
        if quantize and residual_only:
            raise ValueError("a residual-only consumer has no layer input to quantize")
        self.quantize = quantize
        # FlashInfer's MXFP8 quantize helpers are FlashInfer internals: only the
        # quantized variant, which the specialized decode builds, imports them.
        # Untyped modules.
        self.fp4_ops: Any = None
        self.mxfp8_ops: Any = None
        if quantize:
            from flashinfer.cute_dsl import fp4_common
            from flashinfer.quantization import quantization_cute_dsl_utils

            self.fp4_ops = fp4_common
            self.mxfp8_ops = quantization_cute_dsl_utils
        self.write_layer_input = not residual_only
        self.hidden = hidden
        self.hc = hc
        self.tp = tp
        self.capacity_m = capacity_m
        self.cluster_size = cluster_size
        self.rank_lanes = rank_lanes
        self.threads = threads
        self.enable_pdl = enable_pdl
        self.fragments = hidden // VEC_BF16
        self.scale_cols = hidden // MXFP8_BLOCK
        self.groups_per_cta = threads // rank_lanes
        self.fragment_stride = cluster_size * self.groups_per_cta
        self.trips = math.ceil(self.fragments / self.fragment_stride)
        self.warps = threads // WARP_SIZE
        self.rank_waves = tp // rank_lanes

    def smem_size_in_bytes(self) -> int:
        return (self.warps + self.cluster_size) * 4

    @cute.jit
    def __call__(
        self,
        contribution_mailbox: cute.Tensor,
        residual_source: cute.Tensor,
        post: cute.Tensor,
        comb: cute.Tensor,
        pre: cute.Tensor,
        gamma: cute.Tensor,
        residual_output: cute.Tensor,
        norm_output: cute.Tensor,
        quant_output: cute.Tensor,
        quant_scales: cute.Tensor,
        stage_state: cute.Tensor,
        eps: Float32,
        m: Int32,
        stream: cuda.CUstream,
    ) -> None:
        self.kernel(
            contribution_mailbox,
            residual_source,
            post,
            comb,
            pre,
            gamma,
            residual_output,
            norm_output,
            quant_output,
            quant_scales,
            stage_state,
            eps,
        ).launch(
            grid=(m, self.cluster_size, 1),
            block=(self.threads, 1, 1),
            cluster=(1, self.cluster_size, 1),
            smem=self.smem_size_in_bytes(),
            stream=stream,
            use_pdl=self.enable_pdl,
        )

    @cute.kernel
    def kernel(
        self,
        contribution_mailbox: cute.Tensor,
        residual_source: cute.Tensor,
        post: cute.Tensor,
        comb: cute.Tensor,
        pre: cute.Tensor,
        gamma: cute.Tensor,
        residual_output: cute.Tensor,
        norm_output: cute.Tensor,
        quant_output: cute.Tensor,
        quant_scales: cute.Tensor,
        stage_state: cute.Tensor,
        eps: Float32,
    ) -> None:
        hc = self.hc
        tidx, _, _ = cute.arch.thread_idx()
        token, _, _ = cute.arch.block_idx()
        cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
        rank_lane = tidx % self.rank_lanes
        group = tidx // self.rank_lanes
        base_fragment = cluster_rank * self.groups_per_cta + group

        prenorm_fragments = cute.make_rmem_tensor(
            cute.make_layout(
                (self.trips, VEC_BF16),
                stride=(VEC_BF16, 1),
            ),
            BFloat16,
        )
        prenorm_fragments.fill(BFloat16(0.0))
        gamma_fragments = cute.make_rmem_tensor(
            cute.make_layout(
                (self.trips, VEC_BF16),
                stride=(VEC_BF16, 1),
            ),
            BFloat16,
        )
        gamma_fragments.fill(BFloat16(0.0))
        residual_fragments = cute.make_rmem_tensor(
            cute.make_layout(
                (self.trips, hc, VEC_BF16),
                stride=(hc * VEC_BF16, VEC_BF16, 1),
            ),
            BFloat16,
        )
        residual_fragments.fill(BFloat16(0.0))

        for trip in cutlass.range_constexpr(self.trips):
            fragment = base_fragment + trip * self.fragment_stride
            if fragment < self.fragments and rank_lane == 0:
                gamma_element = Int64(fragment) * VEC_BF16
                if cutlass.const_expr(self.write_layer_input):
                    gamma_pointer = cute.make_ptr(
                        BFloat16,
                        (gamma.iterator + gamma_element).llvm_ptr,
                        cute.AddressSpace.gmem,
                        assumed_align=16,
                    )
                    gamma_fragments[trip, None].store(
                        packed_u32x4_to_bf16x8(load_global_u32x4(gamma_pointer))
                    )
                for source in cutlass.range_constexpr(hc):
                    residual_element = (
                        Int64(token) * hc + source
                    ) * self.hidden + gamma_element
                    residual_pointer = cute.make_ptr(
                        BFloat16,
                        (residual_source.iterator + residual_element).llvm_ptr,
                        cute.AddressSpace.gmem,
                        assumed_align=16,
                    )
                    residual_fragments[trip, source, None].store(
                        packed_u32x4_to_bf16x8(load_global_u32x4(residual_pointer))
                    )
        post_mix = cute.make_rmem_tensor(cute.make_layout((hc,)), Float32)
        pre_mix = cute.make_rmem_tensor(cute.make_layout((hc,)), Float32)
        comb_mix = cute.make_rmem_tensor(
            cute.make_layout((hc, hc), stride=(hc, 1)), Float32
        )
        for target in cutlass.range_constexpr(hc):
            mix_index = Int64(token) * hc + target
            post_mix[target] = cute.arch.load(
                (post.iterator + mix_index).llvm_ptr, Float32
            )
            if cutlass.const_expr(self.write_layer_input):
                pre_mix[target] = cute.arch.load(
                    (pre.iterator + mix_index).llvm_ptr, Float32
                )
            for source in cutlass.range_constexpr(hc):
                comb_index = (Int64(token) * hc + source) * hc + target
                comb_mix[source, target] = cute.arch.load(
                    (comb.iterator + comb_index).llvm_ptr, Float32
                )

        if cutlass.const_expr(self.enable_pdl):
            cute.arch.griddepcontrol_wait()

        active_stage = load_volatile_u32(stage_state.iterator + ACTIVE_STAGE)
        for trip in cutlass.range_constexpr(self.trips):
            fragment = base_fragment + trip * self.fragment_stride
            lane_packed = cute.make_rmem_tensor(
                cute.make_layout((self.rank_waves, 4)),
                Uint32,
            )
            lane_packed.fill(Uint32(0))
            dirty = fragment < self.fragments
            while dirty:
                dirty = False
                for wave in cutlass.range_constexpr(self.rank_waves):
                    source_rank = wave * self.rank_lanes + rank_lane
                    if fragment < self.fragments:
                        source_element = (
                            (Int64(active_stage) * self.tp + Int64(source_rank))
                            * self.capacity_m
                            + Int64(token)
                        ) * self.hidden + Int64(fragment) * VEC_BF16
                        source_pointer = cute.make_ptr(
                            BFloat16,
                            (contribution_mailbox.iterator + source_element).llvm_ptr,
                            cute.AddressSpace.gmem,
                            assumed_align=16,
                        )
                        packed = load_global_u32x4(
                            source_pointer,
                            volatile=True,
                        )
                        dirty = dirty | fragment_has_negative_zero(packed)
                        for word in cutlass.range_constexpr(4):
                            lane_packed[wave, word] = packed[word]

            lane_sum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32)
            lane_sum.fill(Float32(0.0))
            for wave in cutlass.range_constexpr(self.rank_waves):
                packed = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
                for word in cutlass.range_constexpr(4):
                    packed[word] = lane_packed[wave, word]
                lane_sum.store(
                    lane_sum.load() + packed_u32x4_to_bf16x8(packed.load()).to(Float32)
                )
            for offset in cutlass.range_constexpr(1, 5):
                if cutlass.const_expr(offset < self.rank_lanes and offset in (1, 2, 4)):
                    for element in cutlass.range_constexpr(VEC_BF16):
                        lane_sum[element] = lane_sum[
                            element
                        ] + cute.arch.shuffle_sync_bfly(
                            lane_sum[element],
                            offset=offset,
                            mask=-1,
                            mask_and_clamp=31,
                        )

            if fragment < self.fragments and rank_lane == 0:
                # The all-reduced sublayer output is BF16, as the unfused
                # collective returns it, before mHC post mixes it in.
                reduced = lane_sum.load().to(BFloat16).to(Float32)
                collapse = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32)
                collapse.fill(Float32(0.0))
                for target in cutlass.range_constexpr(hc):
                    mixed = reduced * post_mix[target]
                    for source in cutlass.range_constexpr(hc):
                        mixed = (
                            mixed
                            + residual_fragments[trip, source, None].load().to(Float32)
                            * comb_mix[source, target]
                        )
                    mixed_bf16 = mixed.to(BFloat16)
                    output_element = (Int64(token) * hc + target) * self.hidden + Int64(
                        fragment
                    ) * VEC_BF16
                    store_global_u32x4(
                        Int64((residual_output.iterator + output_element).toint()),
                        bf16x8_to_packed_u32x4(mixed_bf16),
                    )
                    if cutlass.const_expr(self.write_layer_input):
                        collapse.store(
                            collapse.load() + mixed_bf16.to(Float32) * pre_mix[target]
                        )
                if cutlass.const_expr(self.write_layer_input):
                    prenorm_fragments[trip, None].store(collapse.load().to(BFloat16))

        if token == 0 and cluster_rank == 0 and tidx == 0:
            store_global_u32(
                stage_state.iterator + NEXT_STAGE,
                (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS),
            )
        if cutlass.const_expr(self.enable_pdl):
            cute.arch.griddepcontrol_launch_dependents()

        for trip in cutlass.range_constexpr(self.trips):
            fragment = base_fragment + trip * self.fragment_stride
            for wave in cutlass.range_constexpr(self.rank_waves):
                source_rank = wave * self.rank_lanes + rank_lane
                if fragment < self.fragments:
                    source_element = (
                        (Int64(active_stage) * self.tp + Int64(source_rank))
                        * self.capacity_m
                        + Int64(token)
                    ) * self.hidden + Int64(fragment) * VEC_BF16
                    store_lamport_sentinel_u32x4(
                        Int64((contribution_mailbox.iterator + source_element).toint())
                    )

        if cutlass.const_expr(self.write_layer_input):
            thread_sum = Float32(0.0)
            for trip in cutlass.range_constexpr(self.trips):
                fragment = base_fragment + trip * self.fragment_stride
                if fragment < self.fragments and rank_lane == 0:
                    values = prenorm_fragments[trip, None].load().to(Float32)
                    thread_sum = thread_sum + (values * values).reduce(
                        cute.ReductionOp.ADD,
                        init_val=Float32(0.0),
                        reduction_profile=0,
                    )

            smem = cutlass.utils.SmemAllocator()
            warp_sums = smem.allocate_array(Float32, self.warps)
            cluster_sums = smem.allocate_array(Float32, self.cluster_size)
            cta_sum = _group_leader_block_sum(
                thread_sum,
                warp_sums,
                self.warps,
                self.rank_lanes,
            )
            if tidx < self.cluster_size:
                local_slot = cluster_sums + cluster_rank
                remote_slot = map_shared_to_peer(local_slot, Int32(tidx))
                store_shared_cluster_f32(remote_slot, cta_sum)
            cute.arch.cluster_arrive()
            cute.arch.cluster_wait()
            full_sum = Float32(0.0)
            for peer in cutlass.range_constexpr(self.cluster_size):
                full_sum = full_sum + cute.arch.load(
                    (cluster_sums + peer).llvm_ptr,
                    Float32,
                )
            inv_rms = cute.math.rsqrt(
                full_sum / Float32(self.hidden) + eps,
                fastmath=True,
            )
            for trip in cutlass.range_constexpr(self.trips):
                fragment = base_fragment + trip * self.fragment_stride
                if cutlass.const_expr(self.quantize):
                    normalized = cute.make_rmem_tensor(
                        cute.make_layout((VEC_BF16,)), BFloat16
                    )
                    normalized.fill(BFloat16(0.0))
                if fragment < self.fragments and rank_lane == 0:
                    gamma_values = gamma_fragments[trip, None].load().to(Float32)
                    result = (
                        prenorm_fragments[trip, None].load().to(Float32)
                        * inv_rms
                        * gamma_values
                    ).to(BFloat16)
                    output_element = (
                        Int64(token) * self.hidden + Int64(fragment) * VEC_BF16
                    )
                    store_global_u32x4(
                        Int64((norm_output.iterator + output_element).toint()),
                        bf16x8_to_packed_u32x4(result),
                    )
                    if cutlass.const_expr(self.quantize):
                        normalized.store(result)
                if cutlass.const_expr(self.quantize):
                    # FlashInfer's 4-threads-per-scale MXFP8 quantize of the BF16
                    # output. Every lane joins the shuffles; fragments are a
                    # multiple of 4, so a scale block is wholly in or out of range.
                    words = bf16x8_to_packed_u32x4(normalized.load())
                    block_max = self.mxfp8_ops.reduce_max_4threads(
                        self.mxfp8_ops.bfloat2_hmax_reduce_to_f32(
                            self.mxfp8_ops.bfloat2_max_abs_4(
                                words[0], words[1], words[2], words[3]
                            )
                        )
                    )
                    scale = self.mxfp8_ops.float_to_ue8m0_fast(
                        block_max * Float32(self.mxfp8_ops.INV_FLOAT8_E4M3_MAX)
                    )
                    inv_scale = self.mxfp8_ops.ue8m0_to_inv_scale_fast(scale)
                    if fragment < self.fragments:
                        self.fp4_ops.st_global_u64(
                            self.fp4_ops.get_ptr_as_int64(
                                quant_output,
                                Int64(token) * self.hidden + Int64(fragment) * VEC_BF16,
                            ),
                            self.mxfp8_ops.bfloat2x4_to_fp8x8_packed(
                                words[0], words[1], words[2], words[3], inv_scale
                            ),
                        )
                        if fragment % FRAGMENTS_PER_MXFP8_BLOCK == 0:
                            quant_scales[
                                self.mxfp8_ops.compute_sf_index_swizzled_128x4_gpu(
                                    token,
                                    fragment // FRAGMENTS_PER_MXFP8_BLOCK,
                                    Int32(self.scale_cols),
                                )
                            ] = scale.to(Uint8)

            if cutlass.const_expr(self.quantize):
                # Zero the scales of the padding rows, as FlashInfer's quantize
                # does, one 4-column word per store, spread over the whole grid.
                m, _, _ = cute.arch.grid_dim()
                words_per_row = self.scale_cols // SCALE_TILE_COLS
                padded_m = (
                    (m + SCALE_TILE_ROWS - 1) // SCALE_TILE_ROWS * SCALE_TILE_ROWS
                )
                pad_words = (padded_m - m) * words_per_row
                word = (token * self.cluster_size + cluster_rank) * self.threads + tidx
                while word < pad_words:
                    row = m + word // words_per_row
                    offset = self.mxfp8_ops.compute_sf_index_swizzled_128x4_gpu(
                        row,
                        word % words_per_row * SCALE_TILE_COLS,
                        Int32(self.scale_cols),
                    )
                    store_global_u32(quant_scales.iterator + offset, Uint32(0))
                    word = word + m * self.cluster_size * self.threads