Skip to content

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

Host-side helpers for the prepared CuTe providers.

Functions:

  • check_residency –

    Persistent grids whose CTAs wait on one another (not only within a cluster) must

  • fake –

    Row-major unless stride is given.

  • prepare_shared_carveout –

    Restore value (CuTe's launcher drops preferred_smem_carveout) as a driver

  • tensor_maps –

    TMA descriptors (128 B each) in one device tensor.

_checked(result, operation='CUDA driver call')

The values of a cuda-python driver result, raising on a failed status.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py 
72
73
74
75
76
77
def _checked(result, operation="CUDA driver call"):
    """The values of a cuda-python driver result, raising on a failed status."""
    status, *values = result
    if status != cuda.CUresult.CUDA_SUCCESS:
        raise RuntimeError(f"{operation} failed: {status}")
    return tuple(values)

check_residency(compiled, device_index, grid, threads, what)

Persistent grids whose CTAs wait on one another (not only within a cluster) must all be resident at once: raise when the device cannot hold grid CTAs of compiled (threads per CTA, its dynamic shared memory) at the same time.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def check_residency(compiled, device_index: int, grid: int, threads: int, what: str):
    """Persistent grids whose CTAs wait on one another (not only within a cluster) must
    all be resident at once: raise when the device cannot hold ``grid`` CTAs of
    ``compiled`` (``threads`` per CTA, its dynamic shared memory) at the same time."""
    _, function = resolve_kernel(compiled, device_index)
    attribute = cuda.CUfunction_attribute
    max_dynamic = attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES
    (smem,) = _checked(cuda.cuFuncGetAttribute(max_dynamic, function))
    (per_sm,) = _checked(
        cuda.cuOccupancyMaxActiveBlocksPerMultiprocessor(function, threads, smem)
    )
    sms = torch.cuda.get_device_properties(device_index).multi_processor_count
    if grid > per_sm * sms:
        raise RuntimeError(
            f"{what}: {grid} CTAs cannot be co-resident ({per_sm} per SM x {sms} SMs)"
        )

fake(shape, stride=None, dtype=torch.bfloat16, align=16)

Row-major unless stride is given.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py 
33
34
35
36
37
def fake(shape, stride=None, dtype=torch.bfloat16, align=16):
    """Row-major unless ``stride`` is given."""
    stride = stride or tuple(math.prod(shape[i + 1 :]) for i in range(len(shape)))
    dtype = _DTYPES.get(dtype, dtype)
    return make_fake_tensor(dtype, shape, stride, assumed_align=align)

prepare_shared_carveout(compiled, device_index, value)

Restore value (CuTe's launcher drops preferred_smem_carveout) as a driver function attribute; call once before capture.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py 
119
120
121
122
123
124
125
126
def prepare_shared_carveout(compiled, device_index: int, value: int):
    """Restore ``value`` (CuTe's launcher drops preferred_smem_carveout) as a driver
    function attribute; call once before capture."""
    executor, function = resolve_kernel(compiled, device_index)
    attribute = cuda.CUfunction_attribute
    attribute = attribute.CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT
    _checked(cuda.cuFuncSetAttribute(function, attribute, value))
    return PreparedLaunchAttributes(compiled, executor)

tensor_maps(specs, device)

TMA descriptors (128 B each) in one device tensor. Each spec: (pointer, dtype, dims, strides, box, swizzle, l2 promotion).

The bytes are uploaded with a copy on the current stream: the new tensor may reuse a block freed by a kernel still pending on that stream (e.g. a warmup launch's output), which a copy outside the stream would race with.

Source code in vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def tensor_maps(specs, device):
    """TMA descriptors (128 B each) in one device tensor.
    Each spec: (pointer, dtype, dims, strides, box, swizzle, l2 promotion).

    The bytes are uploaded with a copy on the current stream: the new tensor may
    reuse a block freed by a kernel still pending on that stream (e.g. a warmup
    launch's output), which a copy outside the stream would race with."""
    encoded = bytearray()
    with torch.accelerator.device_index(torch.device(device).index):
        for pointer, dtype, dims, strides, box, swizzle, l2 in specs:
            u64 = ([cuda.cuuint64_t(v) for v in values] for values in (dims, strides))
            u32 = ([cuda.cuuint32_t(v) for v in box], [cuda.cuuint32_t(1)] * len(dims))
            interleave = cuda.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE
            fill = cuda.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
            status, tensor_map = cuda.cuTensorMapEncodeTiled(
                dtype, len(dims), pointer, *u64, *u32, interleave, swizzle, l2, fill
            )
            if status != cuda.CUresult.CUDA_SUCCESS:
                raise RuntimeError(f"cuTensorMapEncodeTiled failed: {status}")
            encoded += (ctypes.c_ubyte * 128).from_address(tensor_map.getPtr())
    return torch.frombuffer(encoded, dtype=torch.uint8).to(device)