Skip to content

vllm.models.muse_spark.nvidia.specialized_decode.cutedsl.common.policy

Per-batch tables and per-TP-size choices of the specialized decode.

Classes:

  • UnpoolerShape –

    Launch shape of the unpooler GEMM: K split over splits CTAs of one

Functions:

  • batch_table –

    Per-batch table of integers ("1-2:96,3-8:120"), tuples ("16:64.16.4.8") or names

UnpoolerShape

Bases: NamedTuple

Launch shape of the unpooler GEMM: K split over splits CTAs of one cluster, kw consecutive K per warp, accs accumulators per warp, and the token tile per batch.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/policy.py 
50
51
52
53
54
55
56
57
58
class UnpoolerShape(NamedTuple):
    """Launch shape of the unpooler GEMM: K split over ``splits`` CTAs of one
    cluster, ``kw`` consecutive K per warp, ``accs`` accumulators per warp, and
    the token tile per batch."""

    splits: int
    kw: int
    accs: int
    token_tile: dict[int, int]

batch_table(spec)

Per-batch table of integers ("1-2:96,3-8:120"), tuples ("16:64.16.4.8") or names ("16-64:trtllm"); later entries override earlier ones.

Source code in vllm/models/muse_spark/nvidia/specialized_decode/cutedsl/common/policy.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
def batch_table(spec):
    """Per-batch table of integers ("1-2:96,3-8:120"), tuples ("16:64.16.4.8") or names
    ("16-64:trtllm"); later entries override earlier ones."""
    table = {}
    for item in filter(None, spec.split(",")):
        span, value = item.split(":")
        lo, _, hi = span.partition("-")
        if "." in value:
            value = tuple(map(int, value.split(".")))
        elif value.lstrip("-").isdigit():
            value = int(value)
        table.update({b: value for b in range(int(lo), int(hi or lo) + 1)})
    return table