Skip to content

vllm.models.kimi_k3.nvidia.specialized_decode.cutedsl.common.primitives

Inline-PTX, memory, synchronization and math helpers of the decode kernels.

Functions:

  • acquire_tensormap_asm –

    acquire_tensormap as one asm block on the unfolded address.

  • alloc –

    Row-major shape tensor from a SmemAllocator.

  • asm –

    f(*args) -> (code, constraints, operands[, result, pure]) as one inline-PTX

  • at –

    Pointer to tensor[coord].

  • attn_res_rstd –

    rsqrtf(total / 7168 + eps) as attn_res_kernel.cu compiles it (--use_fast_math);

  • bf16 –

    FP32 of the low (or high) BF16 half of word.

  • chunk_mma –

    (d, adesc, bdesc, sfa, sfa_src, sfb, sfb_src, copy_sfb, accum, empty_bar,

  • commit_asm –

    commit as an asm block on the unfolded address.

  • copy_g2s –

    16-byte cp.async, or a static/dynamic size bulk copy on barrier.

  • dot –

    fma chain over zip(a, b) from acc, or from a[0] * b[0].

  • ds_tile_scale –

    The fp8_ds_mla tile scale exp2f(ceilf(log2f(max(amax / 448, 1e-4)))) as the

  • fma2 –

    Lanewise a * b + c of FP32 pairs in one packed instruction (IEEE per lane).

  • inv_rms –

    rsqrt(value / 3584 + eps); m8: the mean as a reciprocal product refined by

  • launch –

    Launch kernel(args[:-1]) on stream args[-1]; an int dimension means (n, 1, 1).

  • ldmatrix –

    count 8x8 b16 shared matrices; lanes 8i..8i+7 address matrix i's rows.

  • load_tmem –

    count FP32 TMEM columns of this lane, waiting for the load.

  • load_vector –

    cache "weight": a streaming L2 evict_first load of static weights; else

  • load_words –

    count words (a tuple); scope is a relaxed load at that scope.

  • mma –

    m16n8k16 (FP8: k32) mma.sync.

  • mma_bf16_block –

    Warp-uniform issue of one 64-K block (K-major 128B-swizzled A and B): the

  • mma_into –

    acc += a @ b with one mma; acc is a four-element register view.

  • mma_mx –

    One block-scaled MXF8 tcgen05 MMA; accumulate is a runtime flag.

  • mul2 –

    Lanewise a * b of FP32 pairs in one packed instruction (IEEE per lane).

  • multicast_words –

    Relaxed multimem store of FP32 ("f32") or integer ("b32", "bf16x2") words.

  • mx_exponent –

    E8M0 scale of an MXFP8 block: amax's exponent minus 8, plus one when its

  • no_neg_zero –

    word with each negative-zero BF16 half made +0.

  • nvfp4_scale –

    The E4M3 scale code of a 16-value NVFP4 block, sf = e4m3(max(amax / 6, 2^-9)),

  • pack_e2m1 –

    E2M1 (round to nearest, satfinite) of 8 values, element 0 in the low nibble

  • pack_pair –

    BF16x2 of (low, high); policy=None is the bare conversion.

  • prefetch_bulk –

    prefetch_bytes of a 64-bit address through the NVVM op.

  • ptx –

    One verbatim asm block; a tuple result (even a 1-tuple) returns a tuple.

  • published –

    ready and no word is the mailbox sentinel.

  • push_float –

    DSMEM store of one FP32 value (or a 4-tuple) completing on barrier

  • query_norm_rsqrt –

    rsqrt(total / 1536 + eps) of the q_a RMSNorm as the stock layer computes it:

  • rcp_approx –

    1.f / value under --use_fast_math.

  • release_tmem –

    CTA teardown: warp owner frees the TMEM after every tcgen05 operation.

  • rms_norm_rstd –

    rsqrt(total / 7168 + eps) of vllm_c's rms_norm of a 7168-wide row (IEEE

  • scatter_reduce –

    Warp-reduce count values; lane groups end up owning one value each.

  • sigmoid –

    1 / (1 + exp(-x)); correctly rounded when fast_math is None.

  • sigmoid_reciprocal –

    1 / (1 + exp(-x)) from its denominator: an approximate divide with fast_math,

  • situ –

    Routed-expert SiTU as the stock trtllm-gen FC1 epilogue computes it from the

  • situ_and_mul –

    The SituAndMul op as the deployment runs it, before the BF16 store: the

  • situ_gate –

    The gate factor of the _C situ_and_mul kernel: (0.5 beta) * tanh(g *

  • situ_tanh –

    tanhf of the _C situ_and_mul kernel (csrc/activation_kernels.cu, built

  • situ_up –

    The up factor of the _C kernel: linear_beta * tanh(u * fl(1 / linear_beta)),

  • sleep_ns –

    ns is an immediate if a Python int, else a register.

  • split –

    (first, count) of part i of total items split n ways.

  • store_f32x4 –

    when: a predicate, so callers need no branch.

  • store_fp8_pair –

    FP8 E4M3 (satfinite) of low and high as the two bytes at pointer.

  • store_pair –

    BF16x2 (a, b) with negative-zero halves made +0.

  • store_words –

    scope is a relaxed global store at that scope.

  • thread –

    (thread, warp, lane, CTA) indices.

  • try_wait –

    Acquire parity probe at scope ("cluster", or "cta" when every arrival is

  • view –

    dtype tensor at byte offset of pointer base.

  • wait_tmem –

    wait, then order later tcgen05 operations after it.

_vector(count, first=0)

(".vN", "{\(first,...}") of an N-word access; ("", "\)first") for one word.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
72
73
74
75
def _vector(count, first=0):
    """(".vN", "{$first,...}") of an N-word access; ("", "$first") for one word."""
    names = ",".join(f"${first + i}" for i in range(count))
    return (f".v{count}", "{" + names + "}") if count > 1 else ("", names)

acquire_tensormap_asm(pointer)

acquire_tensormap as one asm block on the unfolded address.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
817
818
819
820
821
@asm()
def acquire_tensormap_asm(pointer):
    """``acquire_tensormap`` as one asm block on the unfolded address."""
    code = "fence.proxy.tensormap::generic.acquire.sys [$0],128; "
    return code + "prefetch.tensormap [$0];", "l", (pointer,)

alloc(allocator, dtype, shape, align=16)

Row-major shape tensor from a SmemAllocator.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
541
542
543
544
545
546
547
def alloc(allocator, dtype, shape, align=16):
    """Row-major ``shape`` tensor from a ``SmemAllocator``."""
    stride = None
    if isinstance(shape, tuple):
        stride = tuple(math.prod(shape[i + 1 :]) for i in range(len(shape)))
    layout = cute.make_layout(shape, stride=stride)
    return allocator.allocate_tensor(dtype, layout, align)

asm(result=None, *, pure=False)

f(*args) -> (code, constraints, operands[, result, pure]) as one inline-PTX op. Pointer operands pass their address: 32-bit for "r" (shared), 64-bit for "l".

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def asm(result=None, *, pure=False):
    """``f(*args) -> (code, constraints, operands[, result, pure])`` as one inline-PTX
    op. Pointer operands pass their address: 32-bit for "r" (shared), 64-bit for "l"."""

    def wrap(fn):
        @dsl_user_op
        @wraps(fn)
        def op(*args, loc=None, ip=None, **kwargs):
            code, constraints, operands, *rest = fn(*args, **kwargs)
            res, fx = (*rest, pure)[:2] if rest else (result, pure)
            kinds = [c for c in constraints.split(",") if c[:1] not in ("=", "~", "")]
            values = []
            for kind, v in zip(kinds, operands):
                if hasattr(v, "toint"):
                    v = v.toint(loc=loc, ip=ip)
                    v = Int32(v) if kind == "r" else v
                values.append(v)
            return ptx(code, constraints, values, res, effects=not fx, loc=loc, ip=ip)

        return op

    return wrap

at(tensor, *coord)

Pointer to tensor[coord].

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
524
525
526
def at(tensor, *coord):
    """Pointer to ``tensor[coord]``."""
    return tensor.iterator + cute.crd2idx(coord, tensor.layout)

attn_res_rstd(total, eps, scale=None)

rsqrtf(total / 7168 + eps) as attn_res_kernel.cu compiles it (--use_fast_math); with scale its fused output norm's rsqrtf(total / 7168 + eps * scale * scale).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
141
142
143
144
145
146
147
148
149
150
151
@asm(Float32, pure=True)
def attn_res_rstd(total: Float32, eps: Float32, scale=None):
    """rsqrtf(total / 7168 + eps) as attn_res_kernel.cu compiles it (--use_fast_math);
    with ``scale`` its fused output norm's rsqrtf(total / 7168 + eps * scale * scale).
    """
    code = "{ .reg .f32 m,t; div.approx.ftz.f32 m,$1,0f45E00000; "
    if scale is None:
        code += "add.ftz.f32 m,m,$2; rsqrt.approx.ftz.f32 $0,m; }"
        return code, "=f,f,f", (total, eps)
    code += "mul.ftz.f32 t,$2,$3; fma.rn.ftz.f32 m,t,$3,m; rsqrt.approx.ftz.f32 $0,m; }"
    return code, "=f,f,f,f", (total, eps, Float32(scale))

bf16(word, high=False)

FP32 of the low (or high) BF16 half of word.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
319
320
321
def bf16(word, high=False):
    """FP32 of the low (or high) BF16 half of ``word``."""
    return bits_float(word & 0xFFFF0000 if high else word << 16)

chunk_mma(*args)

(d, adesc, bdesc, sfa, sfa_src, sfb, sfb_src, copy_sfb, accum, empty_bar, acc_bar): four K32 MMAs of one K128 chunk, with optional completion arrivals.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
@asm()
def chunk_mma(*args):
    """(d, adesc, bdesc, sfa, sfa_src, sfb, sfb_src, copy_sfb, accum, empty_bar,
    acc_bar): four K32 MMAs of one K128 chunk, with optional completion arrivals."""
    constraints = "r,l,l,r,l,r,l,r,r,r,r"
    types = [Uint32 if kind == "r" else Uint64 for kind in constraints.split(",")]
    cp, mma = "tcgen05.cp.cta_group::1.32x128b.warpx4", "tcgen05.mma.cta_group::1"
    commit = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64"
    code = (
        "{\n.reg .pred e,p0,pt,pc,pa,pe; .reg .b32 i<4>; .reg .b64 xa<4>,xb<4>;\n"
        "elect.sync _|e,0xffffffff; setp.eq.b32 pt,0,0; setp.ne.b32 pc,$7,0;\n"
        "and.pred pc,pc,e; setp.ne.b32 pa,$10,0; and.pred pa,pa,e;\n"
        "setp.ne.b32 pe,$9,0; and.pred pe,pe,e; setp.ne.b32 p0,$8,0;\n"
        "mov.b32 i0,142738048; or.b32 i1,i0,536870928; or.b32 i2,i0,1073741856;\n"
        f"or.b32 i3,i0,1610612784; @pc {cp} [$5],$6; @e {cp} [$3],$4;\n"
        "add.s64 xa0,$1,0; add.s64 xb0,$2,0; add.s64 xa1,$1,2; add.s64 xb1,$2,2;\n"
        "add.s64 xa2,$1,4; add.s64 xb2,$2,4; add.s64 xa3,$1,6; add.s64 xb3,$2,6;\n"
    )
    for k in range(4):
        code += f"@e {mma}.kind::mxf8f6f4.block_scale [$0],\n"
        code += f"xa{k},xb{k},i{k},[$3],[$5],{'pt' if k else 'p0'};\n"
    code += f"@pe {commit} [$9]; @pa {commit} [$10];\n}}"
    return code, constraints, tuple(t(v) for t, v in zip(types, args))

commit_asm(barrier)

commit as an asm block on the unfolded address.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
947
948
949
950
951
@asm()
def commit_asm(barrier):
    """``commit`` as an asm block on the unfolded address."""
    code = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [$0];"
    return code, "r", (barrier,)

copy_g2s(dst, src, size=16, barrier=None, *, evict_first=False)

16-byte cp.async, or a static/dynamic size bulk copy on barrier.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
@asm()
def copy_g2s(dst, src, size=16, barrier=None, *, evict_first=False):
    """16-byte cp.async, or a static/dynamic ``size`` bulk copy on ``barrier``."""
    operands = [dst, src]
    policy = hint = tail = ""
    if evict_first:
        policy = ".reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; "
        hint, tail = ".L2::cache_hint", ", pol"
    if barrier is None:
        assert size == 16
        code = f"cp.async.cg.shared.global{hint} [$0], [$1], 16{tail};"
    else:
        operands.append(barrier)
        if not isinstance(size, int):
            operands.append(size)
            size = "$3"
        code = (
            "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
            f"{hint} [$0], [$1], {size}, [$2]{tail};"
        )
    constraints = ",".join(["r", "l"] + ["r"] * (len(operands) - 2))
    return "{ " + policy + code + " }", constraints, operands

dot(a, b, acc=None, **kwargs)

fma chain over zip(a, b) from acc, or from a[0] * b[0].

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
 96
 97
 98
 99
100
def dot(a, b, acc=None, **kwargs):
    """``fma`` chain over ``zip(a, b)`` from ``acc``, or from ``a[0] * b[0]``."""
    for x, y in zip(a, b):
        acc = x * y if acc is None else fma(x, y, acc, **kwargs)
    return acc

ds_tile_scale(amax)

The fp8_ds_mla tile scale exp2f(ceilf(log2f(max(amax / 448, 1e-4)))) as the stock writer concat_and_cache_ds_mla computes it: an IEEE divide, libdevice's log2f (inlined below) and an exact ex2 of the integer.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
@asm(Float32, pure=True)
def ds_tile_scale(amax: Float32):
    """The fp8_ds_mla tile scale exp2f(ceilf(log2f(max(amax / 448, 1e-4)))) as the
    stock writer concat_and_cache_ds_mla computes it: an IEEE divide, libdevice's
    log2f (inlined below) and an exact ex2 of the integer."""
    poly: tuple[str, ...] = (
        "0fBE2C7F30",
        "0f3E2FCF2A",
        "0fBE374E43",
        "0f3E520BF4",
        "0fBE763C8B",
    )
    poly += ("0f3E93BF99", "0fBEB8AA49", "0f3EF6384A", "0fBF38AA3B")
    code = "{ .reg .f32 x,m,e,t,p; .reg .b32 b,k; "
    code += "div.rn.f32 x,$1,0f43E00000; max.f32 x,x,0f38D1B717; "
    code += "mov.b32 b,x; add.s32 k,b,-1060439283; and.b32 k,k,-8388608; "
    code += "sub.s32 b,b,k; mov.b32 m,b; cvt.rn.f32.s32 e,k; "
    code += "fma.rn.f32 e,e,0f34000000,0f00000000; add.f32 t,m,0fBF800000; "
    code += f"fma.rn.f32 p,t,0f3DC6B27F,{poly[0]}; "
    code += "".join(f"fma.rn.f32 p,p,t,{c}; " for c in poly[1:])
    code += "mul.f32 m,t,p; mul.f32 m,t,m; fma.rn.f32 m,t,0f3FB8AA3B,m; "
    code += "add.f32 e,e,m; cvt.rpi.f32.f32 e,e; ex2.approx.f32 $0,e; }"
    return code, "=f,f", (amax,)

fma2(a, b, c)

Lanewise a * b + c of FP32 pairs in one packed instruction (IEEE per lane).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
329
330
331
332
333
334
@asm((Float32, Float32), pure=True)
def fma2(a, b, c):
    """Lanewise a * b + c of FP32 pairs in one packed instruction (IEEE per lane)."""
    code = "{ .reg .b64 a,b,c,d; mov.b64 a,{$2,$3}; mov.b64 b,{$4,$5}; "
    code += "mov.b64 c,{$6,$7}; fma.rn.f32x2 d,a,b,c; mov.b64 {$0,$1},d; }"
    return code, "=f,=f,f,f,f,f,f,f", (*a, *b, *c)

inv_rms(value, eps, *, m8=False)

rsqrt(value / 3584 + eps); m8: the mean as a reciprocal product refined by its FMA residual.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
116
117
118
119
120
121
122
123
124
125
126
@asm(Float32, pure=True)
def inv_rms(value, eps, *, m8=False):
    """rsqrt(value / 3584 + eps); ``m8``: the mean as a reciprocal product refined by
    its FMA residual."""
    mean = "div.rn.f32 r,$1,0f45600000;"
    if m8:
        mean = "mul.f32 q,$1,0f39924925; neg.f32 r,q; fma.rn.f32 r,r,0f45600000,$1; "
        mean += "fma.rn.f32 r,r,0f39924925,q;"
    code = f"{{ .reg .f32 {'q,' * m8}r; {mean} add.f32 r,r,$2; "
    code += "rsqrt.approx.ftz.f32 $0,r; }"
    return code, "=f,f,f", (Float32(value), Float32(eps))

launch(kernel, args, grid, block, cluster=None, pdl=True, **options)

Launch kernel(args[:-1]) on stream args[-1]; an int dimension means (n, 1, 1).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
1018
1019
1020
1021
1022
1023
1024
1025
def launch(kernel, args, grid, block, cluster=None, pdl=True, **options):
    """Launch kernel(args[:-1]) on stream args[-1]; an int dimension means (n, 1, 1)."""
    dim = lambda d: d if isinstance(d, tuple) else (d, 1, 1)
    if cluster is not None:
        options["cluster"] = dim(cluster)
    kernel(args[:-1]).launch(
        grid=dim(grid), block=dim(block), stream=args[-1], use_pdl=pdl, **options
    )

ldmatrix(pointer, *, trans=False, count=4)

count 8x8 b16 shared matrices; lanes 8i..8i+7 address matrix i's rows.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
497
498
499
500
501
502
503
504
@asm()
def ldmatrix(pointer, *, trans=False, count=4):
    """``count`` 8x8 b16 shared matrices; lanes 8i..8i+7 address matrix i's rows."""
    names = ",".join(f"${i}" for i in range(count))
    code = f"ldmatrix.sync.aligned.m8n8.x{count}{'.trans' * trans}.shared.b16 "
    code += f"{{{names}}},[${count}];"
    constraints = ",".join(["=r"] * count + ["r"])
    return code, constraints, (pointer,), (Uint32,) * count

load_tmem(address, count=1)

count FP32 TMEM columns of this lane, waiting for the load.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
927
928
929
930
931
932
933
934
935
@asm()
def load_tmem(address, count=1):
    """``count`` FP32 TMEM columns of this lane, waiting for the load."""
    registers = ",".join(f"${i}" for i in range(count))
    code = f"tcgen05.ld.sync.aligned.32x32b.x{count}.b32 {{{registers}}},[${count}]; "
    code += "tcgen05.wait::ld.sync.aligned;"
    constraints = ",".join(["=f"] * count + ["r", "~{memory}"])
    result = Float32 if count == 1 else (Float32,) * count
    return code, constraints, (Uint32(address),), result

load_vector(pointer, *, cache='cg')

cache "weight": a streaming L2 evict_first load of static weights; else ld.global.<cache>. Side-effecting in every mode.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
610
611
612
613
614
615
616
617
618
619
620
621
622
623
@asm()
def load_vector(pointer, *, cache="cg"):
    """``cache`` "weight": a streaming L2 evict_first load of static weights; else
    ``ld.global.<cache>``. Side-effecting in every mode."""
    opcode, policy, tail = "ld.global." + cache, "", ""
    if cache == "weight":
        policy = ".reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol,1.0; "
        opcode = "ld.global.L1::no_allocate.L2::cache_hint.L2::256B"
        tail = ",pol"
    if cache == "volatile":
        opcode = "ld.volatile.global"
    code = "{ " + policy + opcode + ".v4.b32 {$0,$1,$2,$3},[$4]" + tail + "; }"
    clobber = ",~{memory}" if cache == "volatile" else ""
    return code, "=r,=r,=r,=r,l" + clobber, (pointer,), (Uint32,) * 4

load_words(pointer, *, count=2, shared=False, evict_first=False, mutable=False, volatile=False, scope=None, cache_policy='cg', word_type='b32', result_type=Uint32)

count words (a tuple); scope is a relaxed load at that scope. mutable (a side-effecting load) for data written during the step (activations, tables, refills). Only static weights use evict_first.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
@asm()
def load_words(
    pointer,
    *,
    count=2,
    shared=False,
    evict_first=False,
    mutable=False,
    volatile=False,
    scope=None,
    cache_policy="cg",
    word_type="b32",
    result_type=Uint32,
):
    """``count`` words (a tuple); ``scope`` is a relaxed load at that scope.
    ``mutable`` (a side-effecting load) for data written during the step (activations,
    tables, refills). Only static weights use ``evict_first``."""
    space = "global" + (f".{cache_policy}" if cache_policy else "")
    if shared:
        space = "shared"
    if volatile:
        space = "volatile.global"
    if scope:
        space = f"relaxed.{scope}.global"
    policy = hint = tail = ""
    if evict_first:
        assert not shared
        space = "global.nc.L1::no_allocate"
        policy = ".reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; "
        hint, tail = ".L2::cache_hint", ", pol"
    vector, registers = _vector(count)
    address = f"[${count}]{tail}"
    code = f"{{ {policy}ld.{space}{hint}{vector}.{word_type} {registers}, {address}; }}"
    constraints = ["=f" if result_type is Float32 else "=r"] * count
    constraints += ["r" if shared else "l"] + (["~{memory}"] if scope else [])
    # Shared data is written by other threads between reads: never pure.
    pure = not (shared or mutable or volatile or scope)
    return code, ",".join(constraints), (pointer,), (result_type,) * count, pure

mma(a, b, accumulator, *, fp8=False, bf16=False)

m16n8k16 (FP8: k32) mma.sync.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
481
482
483
484
485
486
487
@asm((Float32,) * 4, pure=True)
def mma(a, b, accumulator, *, fp8=False, bf16=False):
    """m16n8k16 (FP8: k32) mma.sync."""
    kind = "e4m3" if fp8 else "bf16" if bf16 else "f16"
    code = f"mma.sync.aligned.m16n8k{32 if fp8 else 16}.row.col.f32.{kind}.{kind}.f32 "
    code += "{$0,$1,$2,$3},{$4,$5,$6,$7},{$8,$9},{$10,$11,$12,$13};"
    return code, "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", (*a, *b, *accumulator)

mma_bf16_block(d, a, b, instruction, accumulate, empty, release)

Warp-uniform issue of one 64-K block (K-major 128B-swizzled A and B): the elected lane runs the four K16 BF16 MMAs into d (the first accumulates iff accumulate is nonzero) and, if release, commits empty.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
@asm()
def mma_bf16_block(d, a, b, instruction, accumulate, empty, release):
    """Warp-uniform issue of one 64-K block (K-major 128B-swizzled A and B): the
    elected lane runs the four K16 BF16 MMAs into ``d`` (the first accumulates iff
    ``accumulate`` is nonzero) and, if ``release``, commits ``empty``."""
    mma = "tcgen05.mma.cta_group::1.kind::f16"
    commit = "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64"
    code = "{ .reg .pred e,p0,pt,pc; .reg .b64 a1,a2,a3,b1,b2,b3; "
    code += "elect.sync _|e,0xffffffff; setp.eq.b32 pt,0,0; setp.ne.b32 p0,$4,0; "
    code += "setp.ne.b32 pc,$6,0; and.pred pc,pc,e; "
    code += "add.s64 a1,$1,2; add.s64 a2,$1,4; add.s64 a3,$1,6; "
    code += "add.s64 b1,$2,2; add.s64 b2,$2,4; add.s64 b3,$2,6; "
    code += f"@e {mma} [$0],$1,$2,$3,p0; @e {mma} [$0],a1,b1,$3,pt; "
    code += f"@e {mma} [$0],a2,b2,$3,pt; @e {mma} [$0],a3,b3,$3,pt; "
    code += f"@pc {commit} [$5]; }}"
    operands: tuple = (Uint32(d), Uint64(a), Uint64(b), Uint32(instruction))
    operands += (Uint32(accumulate), empty, Uint32(release))
    return code, "r,l,l,r,r,r,r", operands

mma_into(acc, a, b, **kwargs)

acc += a @ b with one mma; acc is a four-element register view.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
490
491
492
493
494
def mma_into(acc, a, b, **kwargs):
    """``acc += a @ b`` with one ``mma``; ``acc`` is a four-element register view."""
    result = mma(a, b, tuple(acc[k] for k in range(4)), **kwargs)
    for k in range(4):
        acc[k] = result[k]

mma_mx(d, a, b, instruction, sfa, sfb, accumulate)

One block-scaled MXF8 tcgen05 MMA; accumulate is a runtime flag.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
974
975
976
977
978
979
980
@asm()
def mma_mx(d, a, b, instruction, sfa, sfb, accumulate):
    """One block-scaled MXF8 tcgen05 MMA; ``accumulate`` is a runtime flag."""
    words = (Uint32(v) for v in (instruction, sfa, sfb, accumulate))
    code = "{ .reg .pred p; setp.ne.b32 p,$6,0; tcgen05.mma.cta_group::1"
    code += ".kind::mxf8f6f4.block_scale [$0],$1,$2,$3,[$4],[$5],p; }"
    return code, "r,l,l,r,r,r,r", (Uint32(d), Uint64(a), Uint64(b), *words)

mul2(a, b)

Lanewise a * b of FP32 pairs in one packed instruction (IEEE per lane).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
337
338
339
340
341
342
@asm((Float32, Float32), pure=True)
def mul2(a, b):
    """Lanewise a * b of FP32 pairs in one packed instruction (IEEE per lane)."""
    code = "{ .reg .b64 a,b,d; mov.b64 a,{$2,$3}; mov.b64 b,{$4,$5}; "
    code += "mul.rn.f32x2 d,a,b; mov.b64 {$0,$1},d; }"
    return code, "=f,=f,f,f,f,f", (*a, *b)

multicast_words(address, words, *, kind='f32')

Relaxed multimem store of FP32 ("f32") or integer ("b32", "bf16x2") words.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
706
707
708
709
710
711
712
713
714
@asm()
def multicast_words(address, words, *, kind="f32"):
    """Relaxed multimem store of FP32 ("f32") or integer ("b32", "bf16x2") words."""
    vector, registers = _vector(len(words), 1)
    f32 = kind == "f32"
    values = (bits_float(v) if f32 else Uint32(v) for v in words)
    constraints = "l" + (",f" if f32 else ",r") * len(words) + ",~{memory}"
    code = f"multimem.st.relaxed.sys.global{vector}.{kind} [$0],{registers};"
    return code, constraints, (Int64(address), *values)

mx_exponent(amax)

E8M0 scale of an MXFP8 block: amax's exponent minus 8, plus one when its mantissa exceeds 448's (1.75), as the trtllm-gen MoE FC1 epilogue and FlashInfer's mxfp8_quantize round it.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
446
447
448
449
450
451
452
453
454
455
456
@cute.jit
def mx_exponent(amax: Float32):
    """E8M0 scale of an MXFP8 block: amax's exponent minus 8, plus one when its
    mantissa exceeds 448's (1.75), as the trtllm-gen MoE FC1 epilogue and FlashInfer's
    mxfp8_quantize round it."""
    bits = Uint32(float_bits(amax))
    e = Int32(bits >> 23)
    result = Int32(0)
    if e != 0:
        result = min(max(e - 8 + Int32(bits & 8388607 > 6291456), 0), 254)
    return result

no_neg_zero(word)

word with each negative-zero BF16 half made +0.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
377
378
379
380
381
382
383
384
385
@cute.jit
def no_neg_zero(word):
    """``word`` with each negative-zero BF16 half made +0."""
    result = Uint32(word)
    if (result & 0xFFFF) == 0x8000:
        result &= Uint32(0xFFFF0000)
    if (result >> 16) == 0x8000:
        result &= Uint32(0xFFFF)
    return result

nvfp4_scale(amax)

The E4M3 scale code of a 16-value NVFP4 block, sf = e4m3(max(amax / 6, 2^-9)), and 1 / float(sf), with IEEE divisions (vLLM's nvfp4_ds_mla writer).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
434
435
436
437
438
439
440
441
442
443
@asm((Uint32, Float32), pure=True)
def nvfp4_scale(amax: Float32):
    """The E4M3 scale code of a 16-value NVFP4 block, sf = e4m3(max(amax / 6, 2^-9)),
    and 1 / float(sf), with IEEE divisions (vLLM's nvfp4_ds_mla writer)."""
    code = "{ .reg .f32 s; .reg .b16 p; .reg .b32 h; .reg .f16 lo,hi; "
    code += "div.rn.f32 s,$2,0f40C00000; max.f32 s,s,0f3B000000; "
    code += "cvt.rn.satfinite.e4m3x2.f32 p,0f00000000,s; cvt.u32.u16 $0,p; "
    code += "cvt.rn.f16x2.e4m3x2 h,p; mov.b32 {lo,hi},h; cvt.f32.f16 s,lo; "
    code += "div.rn.f32 $1,0f3F800000,s; }"
    return code, "=r,=f,f", (amax,)

pack_e2m1(values)

E2M1 (round to nearest, satfinite) of 8 values, element 0 in the low nibble (cvt's first source is the high nibble).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
423
424
425
426
427
428
429
430
431
@asm(Uint32, pure=True)
def pack_e2m1(values):
    """E2M1 (round to nearest, satfinite) of 8 values, element 0 in the low nibble
    (cvt's first source is the high nibble)."""
    code = "{ .reg .b8 b0,b1,b2,b3; cvt.rn.satfinite.e2m1x2.f32 b0,$2,$1; "
    code += "cvt.rn.satfinite.e2m1x2.f32 b1,$4,$3; "
    code += "cvt.rn.satfinite.e2m1x2.f32 b2,$6,$5; "
    code += "cvt.rn.satfinite.e2m1x2.f32 b3,$8,$7; mov.b32 $0,{b0,b1,b2,b3}; }"
    return code, "=r,f,f,f,f,f,f,f,f", values

pack_pair(low, high, *, policy='plain')

BF16x2 of (low, high); policy=None is the bare conversion.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
@asm(Uint32, pure=True)
def pack_pair(low, high, *, policy="plain"):
    """BF16x2 of (low, high); ``policy=None`` is the bare conversion."""
    code = "cvt.rn.bf16x2.f32 $0,$2,$1; "
    if policy == "stage":
        # NaN halves become canonical, so no word is the 0xffffffff staging sentinel.
        code += (
            "and.b32 bits,$0,0x7fff; setp.gt.u32 p,bits,0x7f80; "
            "@p and.b32 $0,$0,0xffff0000; @p or.b32 $0,$0,0x7fc0; "
            "and.b32 bits,$0,0x7fff0000; setp.gt.u32 p,bits,0x7f800000; "
            "@p and.b32 $0,$0,0xffff; @p or.b32 $0,$0,0x7fc00000; "
        )
    if policy == "latent":
        # Negative BF16 zero is the downstream mailbox's unpublished sentinel.
        code += (
            "and.b32 bits,$0,0xffff; setp.eq.u32 p,bits,0x8000; "
            "@p and.b32 $0,$0,0xffff0000; shr.u32 bits,$0,16; "
            "setp.eq.u32 p,bits,0x8000; @p and.b32 $0,$0,0xffff; "
        )
    if policy is not None:
        code = "{ .reg .b32 bits; .reg .pred p; " + code + "}"
    return code, "=r,f,f", (Float32(low), Float32(high))

prefetch_bulk(address, size, *, loc=None, ip=None)

prefetch_bytes of a 64-bit address through the NVVM op.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
793
794
795
796
797
@dsl_user_op
def prefetch_bulk(address, size, *, loc=None, ip=None):
    """``prefetch_bytes`` of a 64-bit address through the NVVM op."""
    pointer = cute.make_ptr(cutlass.Uint8, address, cute.AddressSpace.gmem)
    nvvm.cp_async_bulk_prefetch(pointer, size, loc=loc, ip=ip)

ptx(assembly, constraints, operands=(), result=None, *, effects, loc=None, ip=None)

One verbatim asm block; a tuple result (even a 1-tuple) returns a tuple.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def ptx(assembly, constraints, operands=(), result=None, *, effects, loc=None, ip=None):
    """One verbatim asm block; a tuple ``result`` (even a 1-tuple) returns a tuple."""
    if isinstance(result, tuple) and len(result) == 1:
        args = (assembly, constraints, operands, result[0])
        return (ptx(*args, effects=effects, loc=loc, ip=ip),)
    if result is None:
        ir_type = None
    elif isinstance(result, tuple):
        ir_type = llvm.StructType.get_literal([dtype.mlir_type for dtype in result])
    else:
        ir_type = result.mlir_type
    value = llvm.inline_asm(
        ir_type,
        [operand.ir_value(loc=loc, ip=ip) for operand in operands],
        assembly,
        constraints,
        has_side_effects=effects,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
        loc=loc,
        ip=ip,
    )
    if result is None:
        return None
    if isinstance(result, tuple):
        return tuple(
            dtype(llvm.extractvalue(dtype.mlir_type, value, [i], loc=loc, ip=ip))
            for i, dtype in enumerate(result)
        )
    return result(value)

published(words, ready=True, *, sentinel=2147483648)

ready and no word is the mailbox sentinel.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
388
389
390
391
392
def published(words, ready=True, *, sentinel=0x80000000):
    """``ready`` and no word is the mailbox ``sentinel``."""
    for word in words:
        ready &= word != Uint32(sentinel)
    return ready

push_float(pointer, value, barrier, peer, when=None)

DSMEM store of one FP32 value (or a 4-tuple) completing on barrier (when: predicated on it).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
@asm()
def push_float(pointer, value, barrier, peer, when=None):
    """DSMEM store of one FP32 ``value`` (or a 4-tuple) completing on ``barrier``
    (``when``: predicated on it)."""
    values = tuple(map(Float32, value if isinstance(value, tuple) else (value,)))
    vector, registers = _vector(len(values), 1)
    b, p, w = (f"${len(values) + i}" for i in (1, 2, 3))
    code = f"{{ .reg .b32 dst,bar; mapa.shared::cluster.u32 dst,$0,{p}; "
    code += f"mapa.shared::cluster.u32 bar,{b},{p}; "
    guard = "" if when is None else f".reg .pred w; setp.ne.b32 w,{w},0; @w "
    code += (
        guard + f"st.async.shared::cluster.mbarrier::complete_tx::bytes{vector}.b32 "
    )
    code += f"[dst],{registers},[bar]; }}"
    constraints = "r" + ",f" * len(values) + ",r,r" + ",r" * (when is not None)
    operands = (pointer, *values, barrier, Int32(peer))
    operands += () if when is None else (Int32(when),)
    return code, constraints + ",~{memory}", operands

query_norm_rsqrt(total, eps, *, triton=False)

rsqrt(total / 1536 + eps) of the q_a RMSNorm as the stock layer computes it: vllm_c's rms_norm (IEEE divide; K3.1 DSA), or triton K3 MLA's Triton fused_q_kv_rmsnorm (Triton's fp32 divide).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
129
130
131
132
133
134
135
136
137
138
@asm(Float32, pure=True)
def query_norm_rsqrt(total: Float32, eps: Float32, *, triton: bool = False):
    """rsqrt(total / 1536 + eps) of the q_a RMSNorm as the stock layer computes it:
    vllm_c's rms_norm (IEEE divide; K3.1 DSA), or ``triton`` K3 MLA's Triton
    fused_q_kv_rmsnorm (Triton's fp32 divide)."""
    divide = "div.full.f32" if triton else "div.rn.f32"
    rsqrt = "rsqrt.approx.ftz.f32" if triton else "rsqrt.approx.f32"
    code = f"{{ .reg .f32 mean; {divide} mean,$1,0f44c00000; "
    code += f"add.f32 mean,mean,$2; {rsqrt} $0,mean; }}"
    return code, "=f,f,f", (total, eps)

rcp_approx(value)

1.f / value under --use_fast_math.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
163
164
165
166
@asm(Float32, pure=True)
def rcp_approx(value: Float32):
    """1.f / value under --use_fast_math."""
    return "rcp.approx.ftz.f32 $0,$1;", "=f,f", (value,)

release_tmem(tmem, columns, warp, owner)

CTA teardown: warp owner frees the TMEM after every tcgen05 operation.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
912
913
914
915
916
917
918
919
@cute.jit
def release_tmem(tmem, columns: cutlass.Constexpr[int], warp, owner):
    """CTA teardown: warp ``owner`` frees the TMEM after every tcgen05 operation."""
    fence_tmem(before=True)
    cute.arch.sync_threads()
    if warp == owner:
        fence_tmem()
        free_tmem(tmem, columns)

rms_norm_rstd(total, eps)

rsqrt(total / 7168 + eps) of vllm_c's rms_norm of a 7168-wide row (IEEE divide).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
154
155
156
157
158
159
160
@asm(Float32, pure=True)
def rms_norm_rstd(total: Float32, eps: Float32):
    """rsqrt(total / 7168 + eps) of vllm_c's rms_norm of a 7168-wide row (IEEE
    divide)."""
    code = "{ .reg .f32 m; div.rn.f32 m,$1,0f45E00000; add.rn.f32 m,m,$2; "
    code += "rsqrt.approx.f32 $0,m; }"
    return code, "=f,f,f", (total, eps)

scatter_reduce(values, lane, count)

Warp-reduce count values; lane groups end up owning one value each.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
@cute.jit
def scatter_reduce(values, lane, count: cutlass.Constexpr[int]):
    """Warp-reduce ``count`` values; lane groups end up owning one value each."""
    selected = Int32(0)
    for stage in cutlass.range_constexpr(5):
        offset = 16 >> stage
        remaining = max(count >> stage, 1)
        if cutlass.const_expr(remaining > 1):
            half = remaining // 2
            upper = (lane & offset) != 0
            for i in cutlass.range_constexpr(half):
                send = values[i] if upper else values[i + half]
                keep = values[i + half] if upper else values[i]
                values[i] = keep + cute.arch.shuffle_sync_bfly(send, offset)
            if upper:
                selected += half
        else:
            values[0] += cute.arch.shuffle_sync_bfly(values[0], offset)
    return values[0], selected

sigmoid(value, fast_math=None)

1 / (1 + exp(-x)); correctly rounded when fast_math is None.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
210
211
212
213
214
215
216
217
218
@cute.jit
def sigmoid(value: Float32, fast_math: cutlass.Constexpr = None):
    """1 / (1 + exp(-x)); correctly rounded when ``fast_math`` is None."""
    if cutlass.const_expr(fast_math is None):
        result = div_rn(Float32(1), Float32(1) + cute.math.exp(-value, fastmath=False))
    else:
        denominator = Float32(1) + cute.math.exp(-value, fastmath=True)
        result = sigmoid_reciprocal(denominator, fast_math=fast_math)
    return result

sigmoid_reciprocal(denominator, *, fast_math)

1 / (1 + exp(-x)) from its denominator: an approximate divide with fast_math, else the approximate reciprocal refined by two IEEE FMAs (out-of-range: div.rn).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
194
195
196
197
198
199
200
201
202
203
204
205
206
207
@asm(Float32, pure=True)
def sigmoid_reciprocal(denominator: Float32, *, fast_math):
    """1 / (1 + exp(-x)) from its denominator: an approximate divide with ``fast_math``,
    else the approximate reciprocal refined by two IEEE FMAs (out-of-range: div.rn)."""
    if fast_math:
        return "div.approx.ftz.f32 $0, 0f3f800000, $1;", "=f,f", (denominator,)
    code = (
        "{\n.reg .f32 r,error,neg_error; .reg .b32 bits; .reg .pred regular;\n"
        "mov.b32 bits,$1; add.u32 bits,bits,0x01800000; and.b32 bits,bits,0x7f800000;\n"
        "setp.gt.u32 regular,bits,0x01ffffff; rcp.approx.ftz.f32 r,$1;\n"
        "fma.rn.f32 error,$1,r,0fbf800000; neg.f32 neg_error,error;\n"
        "fma.rn.f32 $0,r,neg_error,r; @!regular div.rn.f32 $0,0f3f800000,$1;\n}"
    )
    return code, "=f,f", (denominator,)

situ(g, u, beta, linear_beta, inv_beta, inv_linear_beta)

Routed-expert SiTU as the stock trtllm-gen FC1 epilogue computes it from the FP32 accumulators: (beta * tanh(g * inv_beta) * sigmoid(g)) * (linear_beta * tanh(u * inv_linear_beta)), in that order, with tanh.approx and sigmoid as rcp.approx(1 + ex2.approx(-g log2 e)).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
239
240
241
242
243
244
245
246
247
248
249
250
251
@asm(Float32, pure=True)
def situ(g, u, beta, linear_beta, inv_beta, inv_linear_beta):
    """Routed-expert SiTU as the stock trtllm-gen FC1 epilogue computes it from the
    FP32 accumulators: (beta * tanh(g * inv_beta) * sigmoid(g)) * (linear_beta *
    tanh(u * inv_linear_beta)), in that order, with tanh.approx and sigmoid as
    rcp.approx(1 + ex2.approx(-g log2 e))."""
    code = "{ .reg .f32 tg,tu,ex,rc,z; mul.rn.f32 tg,$1,$5; tanh.approx.f32 tg,tg; "
    code += "mul.rn.f32 tu,$2,$6; tanh.approx.f32 tu,tu; mul.rn.f32 z,$1,0fbfb8aa3b; "
    code += "ex2.approx.ftz.f32 ex,z; add.rn.f32 ex,ex,0f3f800000; "
    code += "rcp.approx.ftz.f32 rc,ex; mul.rn.f32 z,$3,tg; mul.rn.f32 z,z,rc; "
    code += "mul.rn.f32 tu,$4,tu; mul.rn.f32 $0,z,tu; }"
    operands = (g, u, beta, linear_beta, inv_beta, inv_linear_beta)
    return code, "=f,f,f,f,f,f,f", tuple(map(Float32, operands))

situ_and_mul(gate, up, beta, linear_beta)

The SituAndMul op as the deployment runs it, before the BF16 store: the op enabled (custom_ops "all", e.g. compilation mode NONE under the breakable CUDA graph), i.e. the _C situ_and_mul kernel, gate factor times up factor.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
303
304
305
306
307
308
@cute.jit
def situ_and_mul(gate, up, beta, linear_beta):
    """The ``SituAndMul`` op as the deployment runs it, before the BF16 store: the
    op enabled (custom_ops "all", e.g. compilation mode NONE under the breakable
    CUDA graph), i.e. the ``_C`` situ_and_mul kernel, gate factor times up factor."""
    return mul_rn(situ_gate(gate, beta), situ_up(up, linear_beta))

situ_gate(gate, beta)

The gate factor of the _C situ_and_mul kernel: (0.5 beta) * tanh(g * fl(1 / beta)) * (1 + tanh(0.5 g)), sigmoid(g) taken as (1 + tanh(g / 2)) / 2.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
284
285
286
287
288
289
290
@cute.jit
def situ_gate(gate, beta):
    """The gate factor of the ``_C`` situ_and_mul kernel: (0.5 beta) * tanh(g *
    fl(1 / beta)) * (1 + tanh(0.5 g)), sigmoid(g) taken as (1 + tanh(g / 2)) / 2."""
    half_beta = mul_rn(beta, Float32(0.5))
    scaled = mul_rn(half_beta, situ_tanh(mul_rn(rcp_rn(beta), gate)))
    return mul_rn(scaled, add_rn(situ_tanh(mul_rn(gate, Float32(0.5))), Float32(1)))

situ_tanh(x)

tanhf of the _C situ_and_mul kernel (csrc/activation_kernels.cu, built without fast math), from its PTX: an odd polynomial below 0.6, else copysign(1 - 2 / (1 + ex2.approx.ftz(2|x| log2 e)), x) with rcp.approx.ftz, 1 from 9.01; non-FTZ multiplies and FMAs.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
@asm(Float32, pure=True)
def situ_tanh(x: Float32):
    """``tanhf`` of the ``_C`` situ_and_mul kernel (csrc/activation_kernels.cu, built
    without fast math), from its PTX: an odd polynomial below 0.6, else copysign(1 - 2
    / (1 + ex2.approx.ftz(2|x| log2 e)), x) with rcp.approx.ftz, 1 from 9.01; non-FTZ
    multiplies and FMAs."""
    code = (
        "{ .reg .f32 a,y,e,r,t,q,p,s; .reg .pred big,sat;\n"
        "abs.f32 a,$1; mul.rn.f32 y,a,0f4038AA3B; ex2.approx.ftz.f32 e,y;\n"
        "add.rn.f32 e,e,0f3F800000; rcp.approx.ftz.f32 r,e;\n"
        "fma.rn.f32 t,r,0fC0000000,0f3F800000; setp.ge.f32 sat,a,0f41102CB4;\n"
        "selp.f32 t,0f3F800000,t,sat; copysign.f32 t,$1,t;\n"
        "mul.rn.f32 q,$1,$1; fma.rn.f32 p,q,0f3C80F082,0fBD563CAE;\n"
        "fma.rn.f32 p,p,q,0f3E085941; fma.rn.f32 p,p,q,0fBEAAA9ED;\n"
        "fma.rn.f32 p,p,q,0f00000000; fma.rn.f32 s,p,$1,$1;\n"
        "setp.ge.f32 big,a,0f3F19999A; selp.f32 $0,t,s,big; }"
    )
    return code, "=f,f", (Float32(x),)

situ_up(up, linear_beta)

The up factor of the _C kernel: linear_beta * tanh(u * fl(1 / linear_beta)), or u when linear_beta <= 0 (unset).

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
293
294
295
296
297
298
299
300
@cute.jit
def situ_up(up, linear_beta):
    """The up factor of the ``_C`` kernel: linear_beta * tanh(u * fl(1 / linear_beta)),
    or u when ``linear_beta`` <= 0 (unset)."""
    u = Float32(up)
    if linear_beta > 0:
        u = mul_rn(linear_beta, situ_tanh(mul_rn(rcp_rn(linear_beta), up)))
    return u

sleep_ns(ns)

ns is an immediate if a Python int, else a register.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
877
878
879
880
881
882
@asm()
def sleep_ns(ns):
    """``ns`` is an immediate if a Python int, else a register."""
    if isinstance(ns, int):
        return f"nanosleep.u32 {ns};", "", ()
    return "nanosleep.u32 $0;", "r", (Uint32(ns),)

split(total, i, n)

(first, count) of part i of total items split n ways.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
507
508
509
510
def split(total, i, n):
    """(first, count) of part ``i`` of ``total`` items split ``n`` ways."""
    first = total * i // n
    return first, total * (i + 1) // n - first

store_f32x4(pointer, values, when=None)

when: a predicate, so callers need no branch.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
655
656
657
658
659
660
661
662
@asm()
def store_f32x4(pointer, values, when=None):
    """``when``: a predicate, so callers need no branch."""
    code = "st.global.v4.f32 [$0], {$1, $2, $3, $4};"
    if when is None:
        return code, "l,f,f,f,f", (pointer, *values)
    code = "{ .reg .pred p; setp.ne.b32 p,$5,0; @p " + code + " }"
    return code, "l,f,f,f,f,r", (pointer, *values, Int32(when))

store_fp8_pair(pointer, low, high)

FP8 E4M3 (satfinite) of low and high as the two bytes at pointer.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
408
409
410
411
412
413
@asm()
def store_fp8_pair(pointer, low: Float32, high: Float32):
    """FP8 E4M3 (satfinite) of ``low`` and ``high`` as the two bytes at ``pointer``."""
    code = "{ .reg .b16 pair; cvt.rn.satfinite.e4m3x2.f32 pair,$2,$1; "
    code += "st.global.b16 [$0],pair; }"
    return code, "l,f,f,~{memory}", (pointer, low, high)

store_pair(pointer, a, b)

BF16x2 (a, b) with negative-zero halves made +0.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
676
677
678
679
680
681
682
683
@asm()
def store_pair(pointer, a, b):
    """BF16x2 (a, b) with negative-zero halves made +0."""
    code = "{ .reg .b32 pair,half; .reg .pred p; cvt.rn.bf16x2.f32 pair,$2,$1; "
    code += "and.b32 half,pair,0xffff; setp.eq.b32 p,half,0x8000; "
    code += "@p and.b32 pair,pair,0xffff0000; shr.u32 half,pair,16; "
    code += "setp.eq.b32 p,half,0x8000; @p and.b32 pair,pair,0x0000ffff; "
    return code + "st.global.b32 [$0],pair; }", "l,f,f", (pointer, a, b)

store_words(pointer, words, *, shared=False, scope=None)

scope is a relaxed global store at that scope.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
@asm()
def store_words(pointer, words, *, shared=False, scope=None):
    """``scope`` is a relaxed global store at that scope."""
    vector, registers = _vector(len(words), 1)
    if shared:
        # Other threads read it: a memory side effect, like the global stores.
        space, constraints = (
            f"shared{vector}.b32",
            "r" + ",r" * len(words) + ",~{memory}",
        )
    else:
        space = (
            f"relaxed.{scope}.global{vector}.u32" if scope else f"global{vector}.b32"
        )
        constraints = "l" + ",r" * len(words) + ",~{memory}"
        words = tuple(Uint32(v) for v in words)
    return f"st.{space} [$0], {registers};", constraints, (pointer, *words)

thread()

(thread, warp, lane, CTA) indices.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
513
514
515
516
517
def thread():
    """(thread, warp, lane, CTA) indices."""
    tid, _, _ = cute.arch.thread_idx()
    cta, _, _ = cute.arch.block_idx()
    return tid, tid // 32, tid % 32, cta

try_wait(barrier, phase=0, *, test=False, scope='cluster')

Acquire parity probe at scope ("cluster", or "cta" when every arrival is CTA-local); test is the non-blocking probe of phase 0.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
824
825
826
827
828
829
830
831
832
833
@asm(Int32)
def try_wait(barrier, phase=0, *, test=False, scope="cluster"):
    """Acquire parity probe at ``scope`` ("cluster", or "cta" when every arrival is
    CTA-local); ``test`` is the non-blocking probe of phase 0."""
    op, parity = ("test_wait", "0") if test else ("try_wait", "$2")
    code = f"{{ .reg .pred p; mbarrier.{op}.parity.acquire.{scope}.shared::cta.b64 "
    code += f"p,[$1],{parity}; selp.u32 $0,1,0,p; }}"
    if test:
        return code, "=r,r", (barrier,)
    return code, "=r,r,r,~{memory}", (barrier, Int32(phase))

view(base, offset, dtype, shape, stride=None)

dtype tensor at byte offset of pointer base.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
529
530
531
532
def view(base, offset, dtype, shape, stride=None):
    """``dtype`` tensor at byte ``offset`` of pointer ``base``."""
    pointer = cute.recast_ptr(base + offset, dtype=dtype)
    return cute.make_tensor(pointer, cute.make_layout(shape, stride=stride))

wait_tmem(barrier, phase=0, scope='cluster')

wait, then order later tcgen05 operations after it.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/primitives.py 
900
901
902
903
def wait_tmem(barrier, phase=0, scope="cluster"):
    """``wait``, then order later tcgen05 operations after it."""
    wait(barrier, phase, scope=scope)
    fence_tmem()