vllm.models.beam.nvidia.attention_ops_cutedsl

Bit-exact BF16 output gate in CuTe DSL.

Functions:

• gate_vec –

  Gate 2 or 4 packed BF16x2 pairs sharing one bias word.

_choose(T_, W, maxvec, contiguous_gate)

Return launch parameters and an optional cooperative register bound.

Source code in vllm/models/beam/nvidia/attention_ops_cutedsl.py

def _choose(T_, W, maxvec, contiguous_gate):
    """Return launch parameters and an optional cooperative register bound."""
    n = T_ * W
    if n <= 512 * 1024:
        # Decode/small-prefill sizes. Scalar lanes in 96-ish
        # CTAs
        # up to 64K elements, then 8-byte packets with ~1.5-3K warps in flight.
        pf256 = False
        if n <= 16 * 1024:
            vec, e, block_cap = 1, 1, 128
        elif n <= 32 * 1024:
            # A 256-thread scalar CTA spans 512 B: request 256 B per L2 miss.
            vec, e, block_cap, pf256 = 1, 1, 256, True
        elif n <= 64 * 1024:
            # 512-thread scalar CTAs each cover 1 KB of a row: fetch 256 B per L2 miss.
            vec, e, block_cap, pf256 = 1, 1, 512, True
        elif n <= 128 * 1024:
            vec, e, block_cap = 4, 1, 256
        elif n <= 256 * 1024:
            vec, e, block_cap = 4, 1, 512
        else:
            vec, e, block_cap = 4, 1, 384
        vec = min(vec, maxvec)
        if vec != 1:
            pf256 = False
        need = -(-W // (vec * e))
        return vec, e, 1, min(block_cap, max(32, -(-need // 32) * 32)), pf256, None
    if n <= 32 * 1024 * 1024:
        vec, e, r = 8, 1, 1
    else:
        vec, e, r = 8, 1, 1 if contiguous_gate else 2
    vec = min(vec, maxvec)
    if vec == 8 and n > 1024 * 1024:
        # Two rows per 192-thread CTA keep the T=512 primary grid co-resident.
        ncb = -(-W // (192 * 8))
        if ncb * 192 * 8 - W < 192 * 8 // 2 and _coop_fits(ncb * -(-T_ // 2), 192, 32):
            return 8, 1, 2, 192, False, 32
    need = -(-W // (vec * e))
    block_cap = 256 if vec == 8 and 512 * 1024 <= n <= 1024 * 1024 else 128
    if (
        vec == 8
        and n > 1024 * 1024
        and -(-T_ // r) * -(-W // (256 * vec * e)) <= _CLUSTER_MAX_CTAS
    ):
        block_cap = 256  # keeps ~1-1.5M-element grids within the cluster-launch tier
    block = min(block_cap, max(32, -(-need // 32) * 32))
    return vec, e, r, block, False, None

_coop_fits(ncta, block, regs)

Conservative co-residency bound (no shared memory; regs rounded up per warp).

Source code in vllm/models/beam/nvidia/attention_ops_cutedsl.py

def _coop_fits(ncta, block, regs):
    """Conservative co-residency bound (no shared memory; regs rounded up per warp)."""
    per_warp = -(-regs * 32 // 256) * 256
    by_regs = 65536 // (per_warp * -(-block // 32))
    per_sm = min(32, 2048 // block, by_regs)
    return ncta <= min(_COOP_MAX_CTAS, _num_sms() * per_sm)

_gate_vec_exp_interleaved_ptx(npairs)

Two BF16x2 pairs with independent MUFU operations issued together.

Source code in vllm/models/beam/nvidia/attention_ops_cutedsl.py

def _gate_vec_exp_interleaved_ptx(npairs):
    """Two BF16x2 pairs with independent MUFU operations issued together."""
    assert npairs == 2
    s = [
        "{",
        ".reg .b32 w0, w1, lo, hi;",
        ".reg .b64 p0, p1, kl, k1;",
        ".reg .f32 a0, a1, a2, a3;",
        "mov.b64 kl, 0xBFB8AA3BBFB8AA3B;",
        "mov.b64 k1, 0x3F8000003F800000;",
    ]
    s += [
        "add.rn.bf16x2 w0, $2, $6;",
        "add.rn.bf16x2 w1, $3, $6;",
        "shl.b32 lo, w0, 16;",
        "and.b32 hi, w0, 0xffff0000;",
        "mov.b64 p0, {lo, hi};",
        "mul.rn.f32x2 p0, p0, kl;",
        "mov.b64 {a0, a1}, p0;",
        "shl.b32 lo, w1, 16;",
        "and.b32 hi, w1, 0xffff0000;",
        "mov.b64 p1, {lo, hi};",
        "mul.rn.f32x2 p1, p1, kl;",
        "mov.b64 {a2, a3}, p1;",
        "ex2.approx.ftz.f32 a0, a0;",
        "ex2.approx.ftz.f32 a1, a1;",
        "ex2.approx.ftz.f32 a2, a2;",
        "ex2.approx.ftz.f32 a3, a3;",
        "mov.b64 p0, {a0, a1};",
        "mov.b64 p1, {a2, a3};",
        "add.rn.f32x2 p0, p0, k1;",
        "add.rn.f32x2 p1, p1, k1;",
        "mov.b64 {a0, a1}, p0;",
        "mov.b64 {a2, a3}, p1;",
        "rcp.approx.ftz.f32 a0, a0;",
        "rcp.approx.ftz.f32 a1, a1;",
        "rcp.approx.ftz.f32 a2, a2;",
        "rcp.approx.ftz.f32 a3, a3;",
        "cvt.rn.bf16x2.f32 w0, a1, a0;",
        "cvt.rn.bf16x2.f32 w1, a3, a2;",
        "mul.rn.bf16x2 $0, w0, $4;",
        "mul.rn.bf16x2 $1, w1, $5;",
        "}",
    ]
    return "\n".join(s)

gate_vec(gs, xs, b, *, loc=None, ip=None)

Gate 2 or 4 packed BF16x2 pairs sharing one bias word.

Source code in vllm/models/beam/nvidia/attention_ops_cutedsl.py

@dsl_user_op
def gate_vec(gs, xs, b, *, loc=None, ip=None):
    """Gate 2 or 4 packed BF16x2 pairs sharing one bias word."""
    n = len(gs)
    ops = [Uint32(v).ir_value(loc=loc, ip=ip) for v in gs + xs + [b]]
    res = _asm(
        llvm.StructType.get_literal([T.i32()] * n),
        ops,
        _gate_vec_ptx(n),
        ",".join(["=r"] * n + ["r"] * len(ops)),
        loc,
        ip,
    )
    return [
        Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(n)
    ]