Skip to content

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 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
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 
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
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)