vllm.models.deepseek_v4.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).

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:

• 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

_variant(key)

A build's variant flags past the original's, for its kernel symbol.

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

def _variant(key: MonoBuild) -> dict:
    """A build's variant flags past the original's, for its kernel symbol."""
    out = {"h": 1} if key.hash_route else {}
    if key.topk:
        out["k"] = key.topk
    if key.global_keys:
        out["g"] = 1
    if key.two_hop:
        out["a"] = 2
    return out

peer_half_bytes(tp)

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

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

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

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_v4/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)

K2's FFN half on this CTA: the FFN seam (the attention's TP partials summed and folded into the residual), the MoE and its all-reduce into a["out"], then the launch pair's epoch moved on.

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

@traced
def run_ffn(key: MonoBuild, lds, a: dict, amb, bases, own, rank, ep):
    """K2's FFN half on this CTA: the FFN seam (the attention's TP partials
    summed and folded into the residual), the MoE and its all-reduce into
    ``a["out"]``, then the launch pair's epoch moved on."""
    s, tp, d = key.tokens, key.tp, _dims(key)
    layout = scratch_layout(s, tp, key.topk)
    mkey = _moe_key(key)
    rs = route_shape(mkey)
    n_slices = seam.SLICES
    GATE0 = n_slices
    NORM0 = n_slices + s
    bid = fx.block_idx.x
    tid = fx.thread_idx.x
    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": d,
        "two_hop": key.two_hop,
    }
    normed = scratch + fx.Int64(layout["normed"][0])
    for region in seam.REGIONS:
        ca[region] = sreg(scratch, layout[region][0], region)
    mlayout = {n: layout[n] for n in moe.scratch_layout(mkey) if n not in COUNTERS}
    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["two_hop"] = key.two_hop
    cb["tag"] = ep & 255
    for i, region in enumerate(COUNTERS):
        cb[region] = sreg(a["epoch"], 4 * (EPOCH_COUNTERS + i * moe.TAGS), region)
    for task in range(first_task(bid, 0), n_slices, BLOCKS):
        seam.stage_reduce(ca, task)
        seam.stage_slice(ca, task)
    for t in range(first_task(bid, GATE0), s, BLOCKS):
        seam.stage_gate(ca, t)
    for t in range(first_task(bid, NORM0), s, BLOCKS):
        seam.stage_norm(ca, t, lambda *v: store_normed(normed, HIDDEN, *v))
        publish(cb["put"], cb["xrdy"], t, 1, tid == 0)
    gpu.barrier()

    # ---- the MoE, its all-reduce into ``out``
    if const_expr(key.timeline):
        # the seam's end, then the MoE's stages' (stamped in LDS), after the
        # attention back's points
        tl = a["tl"]
        tls = lds.moe.peek().tls.ptr
        if tid == 0:
            fx.ptr_store(memrealtime(), tls + 0)
        moe.run_moe(cb, mkey, bid, tls, 0)
        if tid < moe.TL_POINTS:
            t64 = fx.Int64(fx.ptr_load(tls + tid))
            w = fx.Vector.from_elements([fx.Int32(t64), fx.Int32(t64 >> 32)], fx.Int32)
            bo.buffer_store(
                w, rsrc(tl), (bid * K2_POINTS + mb_back.BACK_POINTS + tid) * 2
            )
    else:
        moe.run_moe(cb, mkey, bid, 0, 0)

    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, topk=0)

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 (MonoRunner.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.

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

def scratch_layout(s: int, tp: int, topk: int = 0) -> 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 (``MonoRunner.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."""
    d = Dims(tp, topk)
    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(MoeBuild(s, tp)).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()})
    row = s * HIDDEN * 2
    place({"normed": (0, row), "xrdy_moe": (row + 256, s * 8)})
    place({"krdy": (0, s * 8)})
    return out

store_normed(normed, hidden, 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_v4/amd/mono/layer.py

@traced
def store_normed(normed, hidden, 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_v4/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,
            )