Skip to content

vllm.models.kimi_k3.amd.specialized_decode.mono.common.plan

A kernel build's parameters, all of them in FlyDSL's JIT cache key.

FlyDSL keys a compiled kernel by the scalar values (int, float, bool, str, enum, tuple) its launcher's closures capture; an object in a closure is not keyed, and a parameter used only by build-time Python (a placement table, a layout) is in no closure at all. Two builds differing only there would share one binary -- every TP rank ran rank 3's kernel once that way. So a build's parameters are one frozen dataclass of scalars, key_tuple turns it into one (name, value) tuple, and the launcher references that tuple: every field is keyed, whatever the body uses.

The tuple also carries the model's kernel sources (source_digest): FlyDSL keys a helper's source only when its walk reaches it, and it does not reach a helper called inside a traced function's dynamic loop or branch (the rewriter moves those bodies into nested functions it does not scan) nor one outside the kernel's directory -- an edit there kept the old binary.

Classes:

Functions:

  • first_task –

    CTA bid's first task of a stage placed from CTA base (any count of

  • key_tuple –

    key's (field name, value) pairs, every value a scalar, then

  • note_unnamed –

    A put / poll on a raw address: the check reports it (every mailbox access of

  • pair_layout –

    items (name, pairs) laid out in order: name -> (byte offset, bytes),

  • recording –

    Record the accesses of kernels traced inside this block.

  • source_digest –

    sha256 of every .py file under paths (files or directories,

  • symbol_params –

    The fields that name the kernel symbol (metadata={"sym": short name}),

KernelAbi dataclass

Methods:

  • check –

    The kernel takes exactly names; the launcher too, then stream.

  • pack –

    values (every name, nothing else) in ABI order.

  • zeros –

    Every argument 0: what a compile-only call passes.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
@dataclass(frozen=True)
class KernelAbi:
    names: tuple[str, ...]

    def check(self, kernel, launcher) -> None:
        """The kernel takes exactly ``names``; the launcher too, then ``stream``."""
        got = _python_params(kernel)
        if got != self.names:
            raise TypeError(f"kernel parameters {got} != ABI {self.names}")
        got = _python_params(launcher)
        if got != (*self.names, "stream"):
            raise TypeError(f"launcher parameters {got} != ABI {self.names} + stream")

    def pack(self, values: dict) -> list:
        """``values`` (every name, nothing else) in ABI order."""
        try:
            if len(values) == len(self.names):
                return [values[n] for n in self.names]
        except KeyError:
            pass
        extra = sorted(set(values) - set(self.names))
        missing = sorted(set(self.names) - set(values))
        raise TypeError(f"ABI mismatch: missing {missing}, unexpected {extra}")

    def zeros(self) -> list:
        """Every argument 0: what a compile-only call passes."""
        return [0] * len(self.names)

check(kernel, launcher)

The kernel takes exactly names; the launcher too, then stream.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
145
146
147
148
149
150
151
152
def check(self, kernel, launcher) -> None:
    """The kernel takes exactly ``names``; the launcher too, then ``stream``."""
    got = _python_params(kernel)
    if got != self.names:
        raise TypeError(f"kernel parameters {got} != ABI {self.names}")
    got = _python_params(launcher)
    if got != (*self.names, "stream"):
        raise TypeError(f"launcher parameters {got} != ABI {self.names} + stream")

pack(values)

values (every name, nothing else) in ABI order.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
154
155
156
157
158
159
160
161
162
163
def pack(self, values: dict) -> list:
    """``values`` (every name, nothing else) in ABI order."""
    try:
        if len(values) == len(self.names):
            return [values[n] for n in self.names]
    except KeyError:
        pass
    extra = sorted(set(values) - set(self.names))
    missing = sorted(set(self.names) - set(values))
    raise TypeError(f"ABI mismatch: missing {missing}, unexpected {extra}")

zeros()

Every argument 0: what a compile-only call passes.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
165
166
167
def zeros(self) -> list:
    """Every argument 0: what a compile-only call passes."""
    return [0] * len(self.names)

_python_params(fn)

Parameter names of a @flyc.kernel / @flyc.jit object's Python function (FlyDSL keeps it as _func / func).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
132
133
134
135
136
137
138
def _python_params(fn) -> tuple[str, ...]:
    """Parameter names of a ``@flyc.kernel`` / ``@flyc.jit`` object's Python
    function (FlyDSL keeps it as ``_func`` / ``func``)."""
    py = getattr(fn, "func", None) or getattr(fn, "_func", None)
    if py is None:
        raise TypeError(f"no Python function behind {fn!r}")
    return tuple(inspect.signature(py).parameters)

first_task(bid, base)

CTA bid's first task of a stage placed from CTA base (any count of CTAs past 0: placements that add task counts run past the grid), then every BLOCKS-th.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
 98
 99
100
101
102
def first_task(bid, base):
    """CTA ``bid``'s first task of a stage placed from CTA ``base`` (any count of
    CTAs past 0: placements that add task counts run past the grid), then every
    BLOCKS-th."""
    return (bid - base % BLOCKS + BLOCKS) % BLOCKS

key_tuple(key, sources)

key's (field name, value) pairs, every value a scalar, then ("sources", sources): the model's source_digest.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def key_tuple(key, sources: str) -> tuple:
    """``key``'s (field name, value) pairs, every value a scalar, then
    ("sources", ``sources``): the model's ``source_digest``."""
    out = []
    for f in fields(key):
        value = getattr(key, f.name)
        if not isinstance(value, _SCALARS):
            raise TypeError(
                f"{type(key).__name__}.{f.name} = {value!r}: a build parameter must be"
                " a scalar to reach the JIT cache key"
            )
        out.append((f.name, value))
    out.append(("sources", sources))
    return tuple(out)

note_unnamed(kind)

A put / poll on a raw address: the check reports it (every mailbox access of a checked kernel names its region).

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
233
234
235
236
237
def note_unnamed(kind: str) -> None:
    """A put / poll on a raw address: the check reports it (every mailbox access of
    a checked kernel names its region)."""
    if _record is not None:
        _record.unnamed.append((_stage(), kind))

pair_layout(items, align=PAIR_BYTES, start=0)

items (name, pairs) laid out in order: name -> (byte offset, bytes), each region from an align-byte boundary at or past start.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
112
113
114
115
116
117
118
119
120
def pair_layout(items, align: int = PAIR_BYTES, start: int = 0) -> dict:
    """``items`` (name, pairs) laid out in order: name -> (byte offset, bytes),
    each region from an ``align``-byte boundary at or past ``start``."""
    out, off = {}, start
    for name, pairs in items:
        off = (off + align - 1) // align * align
        out[name] = (off, pairs * PAIR_BYTES)
        off += pairs * PAIR_BYTES
    return out

recording()

Record the accesses of kernels traced inside this block.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
207
208
209
210
211
212
213
214
215
@contextmanager
def recording():
    """Record the accesses of kernels traced inside this block."""
    global _record
    outer, _record = _record, Record()
    try:
        yield _record
    finally:
        _record = outer

source_digest(*paths) cached

sha256 of every .py file under paths (files or directories, relative to the mono package), each file's path and contents: one model's kernel sources, so an edit to another model's rebuilds none of these.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
@functools.cache
def source_digest(*paths: str) -> str:
    """sha256 of every ``.py`` file under ``paths`` (files or directories,
    relative to the ``mono`` package), each file's path and contents: one model's
    kernel sources, so an edit to another model's rebuilds none of these."""
    h = hashlib.sha256()
    for rel in paths:
        root = _ROOT / rel
        files = [root] if root.is_file() else sorted(root.rglob("*.py"))
        if not files:
            raise FileNotFoundError(f"no kernel sources at {root}")
        for f in files:
            h.update(str(f.relative_to(_ROOT)).encode())
            h.update(f.read_bytes())
    return h.hexdigest()[:16]

symbol_params(key)

The fields that name the kernel symbol (metadata={"sym": short name}), short name -> value, in field order.

Source code in vllm/models/kimi_k3/amd/specialized_decode/mono/common/plan.py 
75
76
77
78
79
80
81
82
def symbol_params(key) -> dict:
    """The fields that name the kernel symbol (``metadata={"sym": short name}``),
    short name -> value, in field order."""
    return {
        f.metadata["sym"]: getattr(key, f.name)
        for f in fields(key)
        if "sym" in f.metadata
    }