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].

  • 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].

  • 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.

  • no_neg_zero –

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

  • 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

  • release_tmem –

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

  • 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,

  • 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 
699
700
701
702
703
@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 
423
424
425
426
427
428
429
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 
406
407
408
def at(tensor, *coord):
    """Pointer to ``tensor[coord]``."""
    return tensor.iterator + cute.crd2idx(coord, tensor.layout)

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 
221
222
223
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 
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
@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 
829
830
831
832
833
@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 
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
@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

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 
231
232
233
234
235
236
@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 
900
901
902
903
904
905
906
907
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 
379
380
381
382
383
384
385
386
@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 
809
810
811
812
813
814
815
816
817
@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 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
@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 
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
@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 
363
364
365
366
367
368
369
@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 
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
@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 
372
373
374
375
376
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 
856
857
858
859
860
861
862
@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 
239
240
241
242
243
244
@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 
588
589
590
591
592
593
594
595
596
@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)

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 
279
280
281
282
283
284
285
286
287
@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

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 
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
@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 
675
676
677
678
679
@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 
290
291
292
293
294
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 
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
@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

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 
794
795
796
797
798
799
800
801
@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)

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 
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
@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 
160
161
162
163
164
165
166
167
168
@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 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
@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,)

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 
759
760
761
762
763
764
@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 
389
390
391
392
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 
537
538
539
540
541
542
543
544
@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 
310
311
312
313
314
315
@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 
558
559
560
561
562
563
564
565
@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 
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
@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 
395
396
397
398
399
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 
706
707
708
709
710
711
712
713
714
715
@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 
411
412
413
414
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 
782
783
784
785
def wait_tmem(barrier, phase=0, scope="cluster"):
    """``wait``, then order later tcgen05 operations after it."""
    wait(barrier, phase, scope=scope)
    fence_tmem()