Skip to content

vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail.primitives

Shared CuTe DSL primitives; this module does not define a CUDA kernel.

Classes:

Functions:

CUDAGraphCompatibleWrapper

DLPack view that does not synchronize with the producer stream.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
24
25
26
27
28
29
30
31
32
33
34
class CUDAGraphCompatibleWrapper:
    """DLPack view that does not synchronize with the producer stream."""

    def __init__(self, tensor: torch.Tensor):
        self.tensor = tensor

    def __dlpack__(self, stream=None):
        return self.tensor.__dlpack__(stream=-1)

    def __dlpack_device__(self):
        return self.tensor.__dlpack_device__()

block_sum_specialized(value, warp_sums, tidx, warps, last_warp_lanes, last_warp_mask)

Upstream-equivalent FP32 block reduction.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
@cute.jit
def block_sum_specialized(
    value: Float32,
    warp_sums: cute.Tensor,
    tidx: Int32,
    warps: cutlass.Constexpr[int],
    last_warp_lanes: cutlass.Constexpr[int],
    last_warp_mask: cutlass.Constexpr[int],
) -> Float32:
    """Upstream-equivalent FP32 block reduction."""

    lane = cute.arch.lane_idx()
    warp_idx = cute.arch.warp_idx()
    value = warp_sum_specialized(
        value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask
    )
    if lane == 0:
        warp_sums[warp_idx] = value
    cute.arch.barrier()

    block_sum = Float32(0.0)
    if warp_idx == 0:
        if lane < Int32(warps):
            block_sum = warp_sums[lane]
        block_sum = cute.arch.warp_reduction_sum(block_sum)
        if lane == 0:
            warp_sums[0] = block_sum
    cute.arch.barrier()
    return warp_sums[0]

fragment_is_dirty(packed)

Bit-exact upstream sentinel check: one comparison per 32-bit word.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
364
365
366
367
368
369
370
371
@cute.jit
def fragment_is_dirty(packed):
    """Bit-exact upstream sentinel check: one comparison per 32-bit word."""

    dirty = packed[0] == Uint32(NEG_ZERO_F32_BITS)
    for i in cutlass.range_constexpr(1, 4):
        dirty = dirty | (packed[i] == Uint32(NEG_ZERO_F32_BITS))
    return dirty

load_global_u32x4(pointer, *, volatile=False, loc=None, ip=None)

Load one 128-bit fragment as four u32 registers.

The volatile form is the Lamport polling load. Marking the asm side-effecting prevents loop-invariant motion and common-subexpression elimination across polling iterations.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
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
89
90
91
92
93
94
95
96
@dsl_user_op
def load_global_u32x4(
    pointer: cute.Pointer,
    *,
    volatile: cutlass.Constexpr[bool] = False,
    loc=None,
    ip=None,
):
    """Load one 128-bit fragment as four u32 registers.

    The volatile form is the Lamport polling load.  Marking the asm
    side-effecting prevents loop-invariant motion and common-subexpression
    elimination across polling iterations.
    """

    address = pointer.toint(loc=loc, ip=ip)
    opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32"
    out = llvm.inline_asm(
        llvm.StructType.get_literal([T.i32()] * 4),
        [address.ir_value(loc=loc, ip=ip)],
        f"{opcode} {{$0, $1, $2, $3}}, [$4];",
        "=r,=r,=r,=r,l",
        has_side_effects=volatile,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
        loc=loc,
        ip=ip,
    )
    packed = vector.from_elements(
        ir.VectorType.get([4], T.i32(), loc=loc),
        [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)],
        loc=loc,
        ip=ip,
    )
    return cute.TensorSSA(packed, 4, Uint32)

map_shared_to_peer(smem_ptr, peer_rank, *, loc=None, ip=None)

Map a local shared-memory slot to the same slot in a peer CTA.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
@dsl_user_op
def map_shared_to_peer(
    smem_ptr: cute.Pointer,
    peer_rank: Int32,
    *,
    loc=None,
    ip=None,
) -> Int32:
    """Map a local shared-memory slot to the same slot in a peer CTA."""

    smem_address = smem_ptr.toint(loc=loc, ip=ip).ir_value()
    return Int32(
        llvm.inline_asm(
            T.i32(),
            [smem_address, peer_rank.ir_value(loc=loc, ip=ip)],
            "mapa.shared::cluster.u32 $0, $1, $2;",
            "=r,r,r",
            has_side_effects=False,
            is_align_stack=False,
            asm_dialect=llvm.AsmDialect.AD_ATT,
            loc=loc,
            ip=ip,
        )
    )

red_async_release_gpu_add_u32(pointer, value, *, loc=None, ip=None)

The exact SM100 arrival primitive used by upstream LamportFlags.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
@dsl_user_op
def red_async_release_gpu_add_u32(
    pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None
) -> None:
    """The exact SM100 arrival primitive used by upstream LamportFlags."""

    address = pointer.toint(loc=loc, ip=ip)
    llvm.inline_asm(
        None,
        [
            address.ir_value(loc=loc, ip=ip),
            value.ir_value(loc=loc, ip=ip),
        ],
        "red.async.release.global.gpu.add.u32 [$0], $1;",
        "l,r",
        has_side_effects=True,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
        loc=loc,
        ip=ip,
    )

sanitize_negative_zero(packed)

Turn real BF16 -0 into +0 so it cannot equal the empty sentinel.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
@cute.jit
def sanitize_negative_zero(packed):
    """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel."""

    result = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
    for i in cutlass.range_constexpr(4):
        word = packed[i]
        low = Uint16(word & Uint32(0xFFFF))
        high = Uint16(word >> Uint32(16))
        if low == Uint16(NEG_ZERO_BF16_BITS):
            word = word & Uint32(0xFFFF0000)
        if high == Uint16(NEG_ZERO_BF16_BITS):
            word = word & Uint32(0x0000FFFF)
        result[i] = word
    return result.load()

store_global_u32x4(address, packed, *, volatile=False, loc=None, ip=None)

Store four packed words to an ordinary or NVLS multicast global VA.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
@dsl_user_op
def store_global_u32x4(
    address: Int64,
    packed,
    *,
    volatile: cutlass.Constexpr[bool] = False,
    loc=None,
    ip=None,
) -> None:
    """Store four packed words to an ordinary or NVLS multicast global VA."""

    words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)]
    opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32"
    llvm.inline_asm(
        None,
        [address.ir_value(loc=loc, ip=ip), *words],
        f"{opcode} [$0], {{$1, $2, $3, $4}};",
        "l,r,r,r,r",
        has_side_effects=True,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
        loc=loc,
        ip=ip,
    )

store_lamport_sentinel_128(pointer, *, loc=None, ip=None)

Reset one Lamport fragment to four FP32 negative-zero bit patterns.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
@dsl_user_op
def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None:
    """Reset one Lamport fragment to four FP32 negative-zero bit patterns."""

    address = pointer.toint(loc=loc, ip=ip)
    value = Uint32(NEG_ZERO_F32_BITS).ir_value(loc=loc, ip=ip)
    llvm.inline_asm(
        None,
        [address.ir_value(loc=loc, ip=ip), value, value, value, value],
        "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
        "l,r,r,r,r",
        has_side_effects=True,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
        loc=loc,
        ip=ip,
    )

to_cute_dynamic_m(tensor, *, mode, assumed_align=16)

Expose exactly one compact tensor mode as a runtime shape.

Model dimensions remain part of the compiled tensor type. Only the token mode is symbolic, so changing M within an Op's capacity reuses the same compiled kernel.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def to_cute_dynamic_m(
    tensor: torch.Tensor,
    *,
    mode: int,
    assumed_align: int = 16,
) -> cute.Tensor:
    """Expose exactly one compact tensor mode as a runtime shape.

    Model dimensions remain part of the compiled tensor type. Only the token
    mode is symbolic, so changing M within an Op's capacity reuses the same
    compiled kernel.
    """

    return to_cute(tensor, assumed_align).mark_compact_shape_dynamic(
        mode=mode,
        stride_order=tensor.dim_order(),
    )

warp_sum_specialized(value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask)

Warp sum supporting a compile-time partial final warp.

Source code in vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py 
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
@cute.jit
def warp_sum_specialized(
    value: Float32,
    warp_idx: Int32,
    lane: Int32,
    warps: cutlass.Constexpr[int],
    last_warp_lanes: cutlass.Constexpr[int],
    last_warp_mask: cutlass.Constexpr[int],
) -> Float32:
    """Warp sum supporting a compile-time partial final warp."""

    if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32):
        for offset in cutlass.range_constexpr(16, 0, -1):
            # range_constexpr does not provide powers-of-two stepping.
            if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
                other = cute.arch.shuffle_sync_bfly(
                    value,
                    offset=offset,
                    mask=last_warp_mask,
                    mask_and_clamp=31,
                )
                if (lane ^ Int32(offset)) < Int32(last_warp_lanes):
                    value = value + other
    else:
        for offset in cutlass.range_constexpr(16, 0, -1):
            if cutlass.const_expr(offset in (16, 8, 4, 2, 1)):
                value = value + cute.arch.shuffle_sync_bfly(
                    value,
                    offset=offset,
                    mask=-1,
                    mask_and_clamp=31,
                )
    return value