vllm.model_executor.layers.minimax_rms_norm.rms_norm_tp ¶

def _all_reduce_variance(var: torch.Tensor) -> torch.Tensor:
    """All-reduce a per-token variance tensor across the TP group.

    Variance is accumulated in fp32 for numerical stability. The FlashInfer
    fused all-reduce caches a single global workspace keyed to the model's
    16-bit activation dtype (``use_fp32_lamport=False``); routing an fp32
    reduction through it would read against a mismatched workspace and corrupt
    the result. FlashInfer's fast-path only triggers for 2D inputs, so reducing
    a flattened (1D) view keeps these fp32 reductions on custom all-reduce /
    pynccl, both of which handle fp32 correctly.
    """
    return tensor_model_parallel_all_reduce(var.flatten()).view_as(var)

def _minimax_qk_norm_tp_eager(
    qkv: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    q_size: int,
    kv_size: int,
    tp_world: int,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Pure-torch reference path used when Triton is unavailable."""
    q, k, _ = qkv.split([q_size, kv_size, kv_size], dim=-1)
    orig_dtype = q.dtype
    q = q.to(torch.float32)
    k = k.to(torch.float32)
    q_var = q.pow(2).mean(dim=-1, keepdim=True)
    k_var = k.pow(2).mean(dim=-1, keepdim=True)

    qk_var = torch.cat([q_var, k_var], dim=-1)
    qk_var = _all_reduce_variance(qk_var) / tp_world
    q_var, k_var = qk_var.chunk(2, dim=-1)
    q = q * torch.rsqrt(q_var + eps) * q_weight
    k = k * torch.rsqrt(k_var + eps) * k_weight
    return q.to(orig_dtype), k.to(orig_dtype)

def _minimax_qk_norm_tp_fallback(
    qkv: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    q_size: int,
    kv_size: int,
    tp_rank: int,
    tp_world: int,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """All-reduce + QK RMSNorm without the Lamport fused kernel.

    The all-reduce is a TP communication barrier and cannot live inside a
    single kernel, so the eager-torch path is split into two Triton kernels
    around it: a variance reduction before the all-reduce and a normalize
    after. Compared to the ``torch.compile`` path this avoids materializing
    fp32 copies of q/k and the ``cat``/``chunk`` temporaries.
    """
    if not HAS_TRITON:
        return _minimax_qk_norm_tp_eager(
            qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
        )

    num_tokens = qkv.shape[0]
    row_stride = qkv.stride(0)
    BLOCK = 1024
    grid = (num_tokens,)

    qk_var = torch.empty(num_tokens, 2, dtype=torch.float32, device=qkv.device)
    _minimax_qk_var_kernel[grid](
        qkv, qk_var, row_stride, q_size=q_size, kv_size=kv_size, BLOCK=BLOCK
    )

    # All-reduce sums the per-shard means; the /tp_world that turns this back
    # into the global mean is folded into the apply kernel's rsqrt below.
    qk_var = _all_reduce_variance(qk_var)

    q_out = torch.empty(num_tokens, q_size, dtype=qkv.dtype, device=qkv.device)
    k_out = torch.empty(num_tokens, kv_size, dtype=qkv.dtype, device=qkv.device)
    _minimax_rms_apply_kernel[grid](
        qkv,
        qk_var,
        q_weight,
        k_weight,
        q_out,
        k_out,
        row_stride,
        q_size=q_size,
        kv_size=kv_size,
        tp_world=tp_world,
        eps=eps,
        BLOCK=BLOCK,
    )
    return q_out, k_out

@triton.jit
def _minimax_qk_var_kernel(
    qkv_ptr,  # [num_tokens, hidden], 16-bit activations
    var_ptr,  # [num_tokens, 2], fp32
    row_stride,  # element stride between tokens in qkv
    q_size: tl.constexpr,  # constant per deployment -> loops unroll, mask elides
    kv_size: tl.constexpr,
    BLOCK: tl.constexpr,
):
    """TP-pre stage: per-token mean-of-squares for the q and k segments.

    Accumulates in fp32 while reading the 16-bit qkv in place, so no fp32
    copy of q/k is materialized. ``var[:, 0]`` is the q variance and
    ``var[:, 1]`` the k variance; both are the local-shard means, ready for
    the all-reduce that follows.
    """
    token = tl.program_id(0)
    base = qkv_ptr + token * row_stride

    q_acc = 0.0
    for off in range(0, q_size, BLOCK):
        idx = off + tl.arange(0, BLOCK)
        mask = idx < q_size
        x = tl.load(base + idx, mask=mask, other=0.0).to(tl.float32)
        q_acc += tl.sum(x * x, axis=0)

    k_acc = 0.0
    for off in range(0, kv_size, BLOCK):
        idx = off + tl.arange(0, BLOCK)
        mask = idx < kv_size
        x = tl.load(base + q_size + idx, mask=mask, other=0.0).to(tl.float32)
        k_acc += tl.sum(x * x, axis=0)

    tl.store(var_ptr + token * 2 + 0, q_acc / q_size)
    tl.store(var_ptr + token * 2 + 1, k_acc / kv_size)

@triton.jit
def _minimax_rms_apply_kernel(
    qkv_ptr,  # [num_tokens, hidden]
    var_ptr,  # [num_tokens, 2], fp32, all-reduced sum of per-shard means
    q_w_ptr,  # [q_size], q per-channel weight
    k_w_ptr,  # [kv_size], k per-channel weight
    q_out_ptr,  # [num_tokens, q_size], contiguous
    k_out_ptr,  # [num_tokens, kv_size], contiguous
    row_stride,  # element stride between tokens in qkv
    q_size: tl.constexpr,  # constant per deployment -> loops unroll, mask elides
    kv_size: tl.constexpr,
    tp_world: tl.constexpr,  # folds the post-all-reduce /tp_world into rsqrt
    eps: tl.constexpr,
    BLOCK: tl.constexpr,
):
    """TP-post stage: ``x * rsqrt(var / tp_world + eps) * weight``.

    A single program normalizes both the q and k segments of one token, so q
    and k share one launch instead of two. The all-reduce yields the sum of
    per-shard means, so the ``/ tp_world`` that recovers the global
    mean-of-squares is folded into the ``rsqrt`` here rather than run as a
    separate elementwise pass over the ``[num_tokens, 2]`` variance tensor.
    """
    token = tl.program_id(0)
    base = qkv_ptr + token * row_stride

    q_inv = tl.rsqrt(tl.load(var_ptr + token * 2 + 0) / tp_world + eps)
    q_out_row = q_out_ptr + token * q_size
    for off in range(0, q_size, BLOCK):
        idx = off + tl.arange(0, BLOCK)
        mask = idx < q_size
        x = tl.load(base + idx, mask=mask, other=0.0).to(tl.float32)
        w = tl.load(q_w_ptr + idx, mask=mask, other=0.0).to(tl.float32)
        y = x * q_inv * w
        tl.store(q_out_row + idx, y.to(q_out_ptr.dtype.element_ty), mask=mask)

    k_inv = tl.rsqrt(tl.load(var_ptr + token * 2 + 1) / tp_world + eps)
    k_out_row = k_out_ptr + token * kv_size
    for off in range(0, kv_size, BLOCK):
        idx = off + tl.arange(0, BLOCK)
        mask = idx < kv_size
        x = tl.load(base + q_size + idx, mask=mask, other=0.0).to(tl.float32)
        w = tl.load(k_w_ptr + idx, mask=mask, other=0.0).to(tl.float32)
        y = x * k_inv * w
        tl.store(k_out_row + idx, y.to(k_out_ptr.dtype.element_ty), mask=mask)