vllm.models.inkling.nvidia.ops.mm_towers

Fused CUDA kernels for the Inkling vision/audio towers.

Both kernels keep the reference paths' fp32 accumulation and per-op bf16 rounding points (native rms_norm / F.gelu); outputs are frequently bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order (real-checkpoint-weight cosine vs reference > 0.9999998).

Functions:

• dmel_embed_sum_norm –

  rmsnorm(sum_b weight[b * VOCAB + idx[:, b]]) in one launch (no

• rmsnorm_gelu –

  Fused gelu(rmsnorm(x)) (or plain rmsnorm); multiple rows per block

dmel_embed_sum_norm(dmel_idx, weight, norm_weight, eps)

rmsnorm(sum_b weight[b * VOCAB + idx[:, b]]) in one launch (no [T, NB, D] intermediate).

Source code in vllm/models/inkling/nvidia/ops/mm_towers.py

def dmel_embed_sum_norm(
    dmel_idx: torch.Tensor,  # [T, NB] int32
    weight: torch.Tensor,  # [NB * VOCAB, D] bf16
    norm_weight: torch.Tensor | None,
    eps: float,
) -> torch.Tensor:
    """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no
    [T, NB, D] intermediate)."""
    T, nb = dmel_idx.shape
    D = weight.shape[1]
    assert weight.shape[0] % nb == 0
    vocab = weight.shape[0] // nb
    out = torch.empty((T, D), dtype=weight.dtype, device=weight.device)
    if T == 0:
        return out
    d_p2 = triton.next_power_of_2(D)
    _dmel_embed_sum_norm_kernel[(T,)](
        dmel_idx,
        weight,
        norm_weight if norm_weight is not None else weight,
        out,
        eps,
        D,
        dmel_idx.stride(0),
        NB=nb,
        VOCAB=vocab,
        D_P2=d_p2,
        HAS_NORM=norm_weight is not None,
        # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4%
        # at large T (the 80-row gather chain is latency- not lane-bound).
        num_warps=8,
    )
    return out

rmsnorm_gelu(x, weight, eps, gelu=True, fold=None)

Fused gelu(rmsnorm(x)) (or plain rmsnorm); multiple rows per block when D is small. With fold, x must be [N, T, H, W, D] and the output comes back as fold_timespace_to_depth(result, *fold).

Source code in vllm/models/inkling/nvidia/ops/mm_towers.py

def rmsnorm_gelu(
    x: torch.Tensor,  # [..., D] bf16 contiguous
    weight: torch.Tensor,
    eps: float,
    gelu: bool = True,
    fold: tuple[int, int] | None = None,  # (t_fold, hw_fold) of the NEXT fold
) -> torch.Tensor:
    """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block
    when D is small. With ``fold``, x must be [N, T, H, W, D] and the output
    comes back as ``fold_timespace_to_depth(result, *fold)``."""
    D = x.shape[-1]
    flat = x.reshape(-1, D)
    assert flat.stride(1) == 1 and flat.stride(0) == D
    R = flat.shape[0]
    if fold is None:
        out = torch.empty_like(flat)
        ft = fh = fw = tf = hf = 1
        out_shape = x.shape
    else:
        tf, hf = fold
        N, ft, fh, fw, _ = x.shape
        out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D)
        out = torch.empty(out_shape, dtype=x.dtype, device=x.device)
    if R == 0:
        return out.reshape(out_shape)
    d_p2 = triton.next_power_of_2(D)
    block_m = max(1, 4096 // d_p2)
    _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)](
        flat,
        weight,
        out,
        eps,
        R,
        D,
        D_P2=d_p2,
        BLOCK_M=block_m,
        HAS_GELU=gelu,
        FOLD=fold is not None,
        FT=ft,
        FH=fh,
        FW=fw,
        TF=tf,
        HF=hf,
        num_warps=_get_num_warps_from_block_size(d_p2 * block_m),
    )
    return out.reshape(out_shape)