vllm.models.deepseek_v41.amd.mono.layer

The mono decode layer's two launches (README): K1 = the attention seam (stages.seam: slice, gate, the norm into vLLM's MXFP8) + the attention front (attention.front); K2 = the attention back (attention.back, wo_b pushed to every rank) + the FFN seam (the TP sum folded in) + the MoE (stages.moe, its all-reduce in-kernel too). A layer whose attention vLLM runs takes the FFN launch instead: K2 with the attention back replaced by a push of vLLM's unreduced wo_b output (seam.stage_push).

Tags: K1's hand-offs carry 2 e + 1, K2's 2 e + 2 (e: the launch pair's epoch, moved on by K2's CTA 0 at its end), so K2's seam may reuse K1's seam regions. Peer regions alternate by e's parity; every rank runs the same launches, so the epochs agree across ranks.

Functions:

• build_mono_ffn –

  The FFN launch of a layer whose attention vLLM ran (part: this rank's

• peer_half_bytes –

  One parity's peer regions: the attention's partials, then the MoE's.

• publish_x8 –

  Token t's two XRDY flags (one a wqkv K half) once its X8 / X8S landed.

• run_ffn –

  K2's FFN half on this CTA: the FFN seam (the attention's TP partials --

• scratch_layout –

  Every region of step width s's launches -> (byte offset, bytes), disjoint:

• store_normed –

  The FFN seam's norm output -> the MoE input row (bf16, device scope:

• store_x8 –

  The attention seam's norm output (seam.stage_norm's 8 columns at

_ffn_symbol_params(key)

The FFN launch's symbol parameters. The timeline flag tells a timeline build apart, as in K1's and K2's symbols. The draft's build adds its expert count, so that a profile tells the two FFN launches apart.

Source code in vllm/models/deepseek_v41/amd/mono/layer.py

def _ffn_symbol_params(key: MonoBuild) -> dict:
    """The FFN launch's symbol parameters. The timeline flag tells a timeline
    build apart, as in K1's and K2's symbols. The draft's build adds its
    expert count, so that a profile tells the two FFN launches apart."""
    params = dict(s=key.tokens, tp=key.tp, tl=key.timeline)
    if key.experts != EXPERTS:
        params["e"] = key.experts
    return params

_symbol_params(key)

The kernel symbol's draft-only parameters (none for a target layer).

Source code in vllm/models/deepseek_v41/amd/mono/layer.py

def _symbol_params(key: MonoBuild) -> dict:
    """The kernel symbol's draft-only parameters (none for a target layer)."""
    out = {}
    if key.experts != EXPERTS:
        out["e"] = key.experts
    if key.topk != TOPK:
        out["k"] = key.topk
    if key.ug_groups != 2:
        out["g"] = key.ug_groups
    if key.swa_width != WINDOW:
        out["w"] = key.swa_width
    return out

build_mono_ffn(key)

The FFN launch of a layer whose attention vLLM ran (part: this rank's unreduced wo_b output, bf16 [S, HIDDEN]): K2's FFN half, the partial pushed to every rank first. It takes a launch pair's epoch, as K2 does.

Source code in vllm/models/deepseek_v41/amd/mono/layer.py

def build_mono_ffn(key: MonoBuild):
    """The FFN launch of a layer whose attention vLLM ran (``part``: this rank's
    unreduced wo_b output, bf16 [S, HIDDEN]): K2's FFN half, the partial pushed
    to every rank first. It takes a launch pair's epoch, as K2 does."""
    s, tp = key.tokens, key.tp
    assert 1 <= s <= MAX_TOKENS
    rs = route_shape(_moe_key(key))
    assert rs.experts // 64 in SORT_NETS
    MoeLds = moe.moe_smem(s, rs, False)
    half = peer_half_bytes(tp)
    assert SLICES + 2 * s <= BLOCKS
    SeamLds = _seam_lds(s)

    @fx.union
    class FfnLds:
        seam: SeamLds  # type: ignore[valid-type]
        moe: MoeLds  # type: ignore[valid-type]

    name = kernel_symbol("dsv41_mono_ffn", **_ffn_symbol_params(key))
    keyed = key_tuple(key, SOURCES)

    @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1])
    def kffn(
        part: Int64,
        res_in: Int64,
        post_in: Int64,
        comb_in: Int64,
        pre_in: Int64,
        hc_fn: Int64,
        hc_scale: Int64,
        hc_base: Int64,
        ffn_w: Int64,
        res_out: Int64,
        post_out: Int64,
        comb_out: Int64,
        pre_out: Int64,
        gate_w: Int64,
        bias: Int64,
        w13: Int64,
        w13_s: Int64,
        w2: Int64,
        w2_s: Int64,
        sgu: Int64,
        sgu_s: Int64,
        sw2: Int64,
        sw2_s: Int64,
        out: Int64,
        scratch: Int64,
        sym: Int64,
        peers: Int64,
        rank: Int32,
        epoch: Int64,
        tl: Int64,
    ):
        _ = keyed
        lds = fx.SharedAllocator().allocate(FfnLds)
        ep = _epoch(epoch)
        par = fx.Int64(ep & 1) * fx.Int64(half)
        bases = [b + par for b in peer_bases(peers, tp)]
        mb_back.epoch_begin({"bid": fx.block_idx.x, "tid": fx.thread_idx.x}, epoch, ep)
        a = dict(
            res_in=res_in,
            post_in=post_in,
            comb_in=comb_in,
            pre_in=pre_in,
            hc_fn=hc_fn,
            hc_scale=hc_scale,
            hc_base=hc_base,
            ffn_w=ffn_w,
            res_out=res_out,
            post_out=post_out,
            comb_out=comb_out,
            pre_out=pre_out,
            gate_w=gate_w,
            bias=bias,
            w13=w13,
            w13_s=w13_s,
            w2=w2,
            w2_s=w2_s,
            sgu=sgu,
            sgu_s=sgu_s,
            sw2=sw2,
            sw2_s=sw2_s,
            out=out,
            scratch=scratch,
            epoch=epoch,
        )
        amb = Mailbox((ep << 1) + 2)
        run_ffn(
            key, lds, a, amb, bases, sym + par, rank, ep, part=part,
            tl=tl if key.timeline else None,
        )  # fmt: skip

    @flyc.jit
    def launch(
        part: Int64,
        res_in: Int64,
        post_in: Int64,
        comb_in: Int64,
        pre_in: Int64,
        hc_fn: Int64,
        hc_scale: Int64,
        hc_base: Int64,
        ffn_w: Int64,
        res_out: Int64,
        post_out: Int64,
        comb_out: Int64,
        pre_out: Int64,
        gate_w: Int64,
        bias: Int64,
        w13: Int64,
        w13_s: Int64,
        w2: Int64,
        w2_s: Int64,
        sgu: Int64,
        sgu_s: Int64,
        sw2: Int64,
        sw2_s: Int64,
        out: Int64,
        scratch: Int64,
        sym: Int64,
        peers: Int64,
        rank: Int32,
        epoch: Int64,
        tl: Int64,
        stream: fx.Stream = _STREAM,
    ):
        _ = keyed
        kffn(
            part,
            res_in,
            post_in,
            comb_in,
            pre_in,
            hc_fn,
            hc_scale,
            hc_base,
            ffn_w,
            res_out,
            post_out,
            comb_out,
            pre_out,
            gate_w,
            bias,
            w13,
            w13_s,
            w2,
            w2_s,
            sgu,
            sgu_s,
            sw2,
            sw2_s,
            out,
            scratch,
            sym,
            peers,
            rank,
            epoch,
            tl,
        ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream)

    return launch

peer_half_bytes(tp)

One parity's peer regions: the attention's partials, then the MoE's.

Source code in vllm/models/deepseek_v41/amd/mono/layer.py

def peer_half_bytes(tp: int) -> int:
    """One parity's peer regions: the attention's partials, then the MoE's."""
    return seam.attn_peer_bytes(MAX_TOKENS, tp) + moe.peer_bytes(MAX_TOKENS, tp)

publish_x8(c, t)

Token t's two XRDY flags (one a wqkv K half) once its X8 / X8S landed.

Source code in vllm/models/deepseek_v41/amd/mono/layer.py

@traced
def publish_x8(c, t):
    """Token t's two XRDY flags (one a wqkv K half) once its X8 / X8S landed."""
    rocdl.s_waitcnt(vmcnt=0)
    gpu.barrier()
    if c["tid"] < 2:
        c["mb"].put(c["xrdy"], 2 * t + c["tid"], fx.Int32(1))
    gpu.barrier()

run_ffn(key, lds, a, amb, bases, own, rank, ep, part=None, tl=None)

K2's FFN half on this CTA: the FFN seam (the attention's TP partials -- part pushed to every rank's ATTN region first -- summed and folded into the residual), the MoE and its all-reduce into a["out"], then the launch pair's epoch moved on. tl: K2's timeline record base (a timeline build), stamped from point K2_SEAM_POINT on.

Source code in vllm/models/deepseek_v41/amd/mono/layer.py

@traced
def run_ffn(
    key: MonoBuild, lds, a: dict, amb, bases, own, rank, ep, part=None, tl=None
):
    """K2's FFN half on this CTA: the FFN seam (the attention's TP partials --
    ``part`` pushed to every rank's ATTN region first -- summed and folded into
    the residual), the MoE and its all-reduce into ``a["out"]``, then the
    launch pair's epoch moved on. ``tl``: K2's timeline record base (a
    timeline build), stamped from point K2_SEAM_POINT on."""
    s, tp = key.tokens, key.tp
    # The scratch is laid out for the routing and the FP4 groups of the
    # runner that owns it. K2 runs on the runner of its own layers, the
    # target's or the DSpark draft's. The FFN launch (``part``) always runs on
    # the target's runner, also for a gfx942 draft layer, so it reads the
    # target's layout.
    if const_expr(part is None):
        layout = scratch_layout(s, tp, key.experts, key.topk, key.ug_groups)
    else:
        layout = scratch_layout(s, tp)
    mkey = _moe_key(key)
    rs = route_shape(mkey)
    GATE0 = SLICES
    NORM0 = SLICES + s
    bid = fx.block_idx.x
    tid = fx.thread_idx.x

    def mark(point, core_clock=False):
        # Thread 0 of the CTA stores the 100 MHz clock at this CTA's record,
        # the same way the attention back's stamps do. With core_clock it
        # stores the shader clock counter (s_memtime) instead. Two such
        # stamps divided by two 100 MHz stamps give the clock the CUs ran
        # at in between.
        if const_expr(tl is not None):  # noqa: SIM102 (a build-time and a run-time test)
            if tid == 0:
                if const_expr(core_clock):
                    t = fx.Int64(
                        llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memtime", [], [], [])
                    )
                else:
                    t = memrealtime()
                w = fx.Vector.from_elements([fx.Int32(t), fx.Int32(t >> 32)], fx.Int32)
                bo.buffer_store(w, rsrc(tl), (bid * K2_POINTS + point) * 2)

    sl = lds.seam.peek()
    scratch = a["scratch"]
    seam_args = (
        "res_in",
        "post_in",
        "comb_in",
        "pre_in",
        "hc_fn",
        "hc_scale",
        "hc_base",
        "res_out",
        "post_out",
        "comb_out",
        "pre_out",
    )
    ca = {
        "S": s,
        "tid": tid,
        "bid": bid,
        "lane": tid % 64,
        "wave": tid // 64,
        "rl": sl.rl.ptr,
        "fl": sl.fl.ptr,
        "red": sl.red.ptr,
        "pend_lds": sl.pl.ptr,
        "put": amb.put,
        "put_bf": amb.put_bf,
        "put_words": amb.put_words,
        "poll": amb.poll,
        "peer_addr": lambda p: bases[p],
        "rank": rank,
        "sym": own,
        "args": {n: a[n] for n in seam_args} | {"norm_w": a["ffn_w"]},
        "d": StageDims(tp),
    }
    normed = scratch + fx.Int64(layout["normed"][0])
    for region in ("lin", "pmix"):
        ca[region] = sreg(scratch, layout[region][0], region)
    # In the FFN launch the layout holds the MoE regions at the target's
    # routing. A build with fewer experts or picks (a gfx942 DSpark draft
    # layer's) uses the same regions, so each of its regions must fit in the
    # target's. The MoE's counters are in the epoch buffer, not in this
    # layout.
    mregions = {n: v for n, v in moe.scratch_layout(mkey).items() if n not in COUNTERS}
    for region, (_, size) in mregions.items():
        assert size <= layout[region][1], (region, size, layout[region][1])
    mlayout = {n: layout[n] for n in mregions}
    mlayout["xrdy"] = layout["xrdy_moe"]
    margs = moe.moe_args(
        normed,
        a["gate_w"],
        a["bias"],
        a["w13"],
        a["w13_s"],
        a["w2"],
        a["w2_s"],
        a["sgu"],
        a["sgu_s"],
        a["sw2"],
        a["sw2_s"],
        a["out"],
    )
    cb = moe.moe_context(
        s, rs, lds.moe.peek(), amb, bases, mlayout, scratch, margs, rank, own
    )
    cb["x_cm"], cb["x_ready"] = CM_DEV, True
    cb["tag"] = ep & 255
    for i, region in enumerate(COUNTERS):
        cb[region] = sreg(a["epoch"], 4 * (EPOCH_COUNTERS + i * moe.TAGS), region)
    # In K2 the attention back half stamps point 0. The FFN launch has no
    # attention, so it stamps its own start here.
    if const_expr(part is not None):
        mark(0)
    for task in range(first_task(bid, 0), SLICES, BLOCKS):
        if const_expr(part is not None):
            seam.stage_push(ca, task, part)
        seam.stage_reduce(ca, task)
        seam.stage_slice(ca, task)
    mark(K2_SEAM_POINT)
    for t in range(first_task(bid, GATE0), s, BLOCKS):
        seam.stage_gate(ca, t)
    mark(K2_SEAM_POINT + 1)
    for t in range(first_task(bid, NORM0), s, BLOCKS):
        seam.stage_norm(ca, t, lambda *v: store_normed(normed, *v))
        publish(cb["put"], cb["xrdy"], t, 1, tid == 0)
    gpu.barrier()
    mark(K2_SEAM_POINT + 2)

    # ---- the MoE, its all-reduce into ``out``
    if const_expr(tl is not None):
        cb["tl_mark"] = mark
    moe.run_moe(cb, mkey, bid, 0, K2_SEAM_POINT + 2)

    def reset():
        # the MoE counter slot 128 launch pairs ahead: its last use long
        # done, its next use far off
        slot = (ep + 128) & 255
        for region in COUNTERS:
            gstore(cb[region].value + fx.Int64(slot * 4), fx.Int32(0), words=1)

    mb_back.epoch_end({"bid": bid, "tid": tid}, a["epoch"], ep, reset)

scratch_layout(s, tp, experts=EXPERTS, topk=TOPK, ug_groups=2)

Every region of step width s's launches -> (byte offset, bytes), disjoint: K1's front and seam, K2's back, the MoE's regions, the normed rows and their flags. Each width has a scratch of its own (DSV41MonoLayer.scratch): under another width's layout a mailbox pair's tag word could be plain data, which a poll may take for a current tag. The MoE's counters, whose slots outlive a step width, are in the epoch buffer. The MoE regions follow the runner's routing and FP4 groups of an ug task: with one group a task (the DSpark draft's runner) the ug flags take twice the room of two groups.

Source code in vllm/models/deepseek_v41/amd/mono/layer.py

def scratch_layout(
    s: int, tp: int, experts: int = EXPERTS, topk: int = TOPK, ug_groups: int = 2
) -> dict:
    """Every region of step width s's launches -> (byte offset, bytes), disjoint:
    K1's front and seam, K2's back, the MoE's regions, the normed rows and their
    flags. Each width has a scratch of its own (``DSV41MonoLayer.scratch``):
    under another width's layout a mailbox pair's tag word could be plain data,
    which a poll may take for a current tag. The MoE's counters, whose slots
    outlive a step width, are in the epoch buffer. The MoE regions follow the
    runner's routing and FP4 groups of an ug task: with one group a task (the
    DSpark draft's runner) the ug flags take twice the room of two groups."""
    d = Dims(tp)
    out: dict = {}
    off = 0

    def place(regions):
        nonlocal off
        base = off
        for name, (o, n) in regions.items():
            assert name not in out or name in SHARED_REGIONS, name
            out[name] = (base + o, n)
        off = max(off, base + max(o + n for o, n in regions.values()))
        off = -(-off // 256) * 256

    place(front_scratch(s, d))
    place(seam.scratch_layout(s))
    place(back_scratch(s, d, start=0))
    regions = moe.scratch_layout(
        _moe_key(MonoBuild(s, tp, 1, experts=experts, topk=topk, ug_groups=ug_groups))
    ).items()
    mr = {n: v for n, v in regions if n not in COUNTERS}
    lo = min(o for o, _ in mr.values())
    place({n: (o - lo, n_) for n, (o, n_) in mr.items()})
    place({"normed": (0, s * HIDDEN * 2), "xrdy_moe": (s * HIDDEN * 2 + 256, s * 8)})
    place({"krdy": (0, s * 8)})
    return out

store_normed(normed, t, col, ys, live)

The FFN seam's norm output -> the MoE input row (bf16, device scope: this launch reads it).

Source code in vllm/models/deepseek_v41/amd/mono/layer.py

@traced
def store_normed(normed, t, col, ys, live):
    """The FFN seam's norm output -> the MoE input row (bf16, device scope:
    this launch reads it)."""
    if live:
        for j in range_constexpr(8):
            bo.buffer_store(
                ys[j].to(fx.BFloat16),
                rsrc(normed),
                t * HIDDEN + col + j,
                cache_modifier=CM_DEV,
            )

store_x8(c, t, col, ys, live)

The attention seam's norm output (seam.stage_norm's 8 columns at col) -> vLLM's MXFP8 of each 32 (the attention's input quant) -> the front's X8 / X8S, plain at device scope.

Source code in vllm/models/deepseek_v41/amd/mono/layer.py

@traced
def store_x8(c, t, col, ys, live):
    """The attention seam's norm output (``seam.stage_norm``'s 8 columns at
    ``col``) -> vLLM's MXFP8 of each 32 (the attention's input quant) -> the
    front's X8 / X8S, plain at device scope."""
    amax = abs(ys[0])
    for y in ys[1:]:
        amax = fx.max(amax, abs(y))
    code = moe.vllm_mx_code(butterfly(amax, (1, 2), fx.max))
    mul = ((254 - code) << 23).bitcast(fx.Float32)
    w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)])
    w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)])
    if live:
        bo.buffer_store(
            fx.Vector.from_elements([w0, w1], fx.Int32),
            rsrc(c["x8"]),
            t * X_WORDS + col // 4,
            cache_modifier=CM_DEV,
        )
        if c["tid"] % 4 == 0:
            bo.buffer_store(
                code, rsrc(c["x8s"]), t * X_GROUPS + col // 32, cache_modifier=CM_DEV
            )