vllm.models.muse_spark.nvidia.ops.all_reduce_dual_norm

One-shot LL all-reduce fused into the gated residual and the next RMSNorm.

Every rank multicasts its full [M, H] BF16 partial sum into slot [stage, rank] of each rank's mailbox. The consumer sums the TP slots in rank order (so every rank gets the same FP32 row) and runs the same dual norm as FlashInfer's all-gather variant, optionally also writing the unrounded FP32 result for the router.

Classes:

• AllReduceDualNorm –

  Collective over group; calls must use the same m on every rank.

AllReduceDualNorm

Bases: LLColumnGather

Collective over group; calls must use the same m on every rank.

Methods:

• __call__ –

  Returns (norm_output, residual_output).

Source code in vllm/models/muse_spark/nvidia/ops/all_reduce_dual_norm.py

class AllReduceDualNorm(LLColumnGather):
    """Collective over ``group``; calls must use the same ``m`` on every rank."""

    def __init__(
        self,
        *,
        hidden_size: int,
        tp_size: int,
        rank: int,
        capacity_m: int,
        group: dist.ProcessGroup,
        cluster_size: int = 8,
        enable_pdl: bool = True,
    ) -> None:
        if hidden_size % (cluster_size * VEC_BF16):
            raise ValueError("hidden_size must split over the cluster")
        self.threads = hidden_size // (cluster_size * VEC_BF16)
        self.cluster_size = cluster_size
        self.enable_pdl = enable_pdl
        # LLColumnGather's mailbox, but every rank publishes the full row.
        self.hidden_size = hidden_size
        self.hidden_local = hidden_size
        self.tp_size = tp_size
        self.capacity_m = capacity_m
        self.state = create_ll_state(group, tp_size, rank, capacity_m, hidden_size)
        self._mailbox = to_cute(self.state.contribution_mailbox.tensor.flatten(), 16)
        self._mailbox_multicast = Int64(
            self.state.contribution_mailbox.multicast_address
        )
        self._stage_state = to_cute(self.state.stage_state, 4)
        self._publisher = compile_shared_only_publish(
            hidden=hidden_size,
            tp=tp_size,
            rank=rank,
            capacity_m=capacity_m,
            elements_per_thread=VEC_BF16,
            threads=128,
            release_before_store=False,
            enable_pdl=enable_pdl,
        )
        self._collectives: dict[Any, Any] = {}
        torch.accelerator.synchronize()
        dist.barrier(group)

    def _compile(self, constants: DualNormConstants) -> Any:
        device_kernel = _AllReduceDualNormKernel(
            hidden=self.hidden_size,
            tp=self.tp_size,
            capacity_m=self.capacity_m,
            cluster_size=self.cluster_size,
            threads=self.threads,
            constants=constants,
            enable_pdl=self.enable_pdl,
        )
        weight = make_fake_compact_tensor(
            BFloat16, (self.hidden_size,), assumed_align=16
        )
        gate = make_fake_compact_tensor(Float32, (self.hidden_size,), assumed_align=32)

        def rows_f32():
            return make_fake_dynamic_compact_tensor(
                Float32, alignment=32, divisibility=self.hidden_size
            )

        return cute.compile(
            device_kernel,
            make_fake_compact_tensor(
                BFloat16,
                (self.state.contribution_mailbox.tensor.numel(),),
                assumed_align=16,
            ),
            rows_f32(),
            gate,
            gate,
            weight,
            weight,
            rows_f32(),
            make_fake_dynamic_compact_tensor(
                BFloat16, alignment=16, divisibility=self.hidden_size
            ),
            rows_f32(),
            make_fake_compact_tensor(Int32, (2,), assumed_align=4),
            Int32(self.capacity_m),
            current_cu_stream(),
        )

    def fits(self, partial: torch.Tensor) -> bool:
        return (
            partial.dim() == 2
            and 1 <= partial.shape[0] <= self.capacity_m
            and partial.shape[1] == self.hidden_size
            and partial.dtype == torch.bfloat16
            and partial.is_contiguous()
        )

    def __call__(
        self,
        partial: torch.Tensor,
        residual: torch.Tensor,
        alpha: torch.Tensor,
        beta: torch.Tensor,
        next_weight: torch.Tensor,
        *,
        constants: DualNormConstants,
        post_weight: torch.Tensor | None,
        router_output: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Returns ``(norm_output, residual_output)``."""
        assert constants.has_router == (router_output is not None)
        residual_output = torch.empty_like(residual)
        norm_output = torch.empty_like(partial)
        collective = self.compiled(constants)
        m = Int32(partial.shape[0])
        stream = current_cu_stream()
        self._publish(partial, m, stream)
        rows_out = to_cute_dynamic(
            residual_output.flatten(), 32, divisibility=self.hidden_size
        )
        collective(
            self._mailbox,
            to_cute_dynamic(residual.flatten(), 32, divisibility=self.hidden_size),
            to_cute(alpha, 32),
            to_cute(beta, 32),
            to_cute(post_weight if post_weight is not None else next_weight, 16),
            to_cute(next_weight, 16),
            rows_out,
            to_cute_dynamic(norm_output.flatten(), 16, divisibility=self.hidden_size),
            (
                to_cute_dynamic(
                    router_output.flatten(), 32, divisibility=self.hidden_size
                )
                if router_output is not None
                else rows_out
            ),
            self._stage_state,
            m,
            stream,
        )
        return norm_output, residual_output

__call__(partial, residual, alpha, beta, next_weight, *, constants, post_weight, router_output)

Returns (norm_output, residual_output).

Source code in vllm/models/muse_spark/nvidia/ops/all_reduce_dual_norm.py

def __call__(
    self,
    partial: torch.Tensor,
    residual: torch.Tensor,
    alpha: torch.Tensor,
    beta: torch.Tensor,
    next_weight: torch.Tensor,
    *,
    constants: DualNormConstants,
    post_weight: torch.Tensor | None,
    router_output: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Returns ``(norm_output, residual_output)``."""
    assert constants.has_router == (router_output is not None)
    residual_output = torch.empty_like(residual)
    norm_output = torch.empty_like(partial)
    collective = self.compiled(constants)
    m = Int32(partial.shape[0])
    stream = current_cu_stream()
    self._publish(partial, m, stream)
    rows_out = to_cute_dynamic(
        residual_output.flatten(), 32, divisibility=self.hidden_size
    )
    collective(
        self._mailbox,
        to_cute_dynamic(residual.flatten(), 32, divisibility=self.hidden_size),
        to_cute(alpha, 32),
        to_cute(beta, 32),
        to_cute(post_weight if post_weight is not None else next_weight, 16),
        to_cute(next_weight, 16),
        rows_out,
        to_cute_dynamic(norm_output.flatten(), 16, divisibility=self.hidden_size),
        (
            to_cute_dynamic(
                router_output.flatten(), 32, divisibility=self.hidden_size
            )
            if router_output is not None
            else rows_out
        ),
        self._stage_state,
        m,
        stream,
    )
    return norm_output, residual_output