vllm.models.deepseek_v41.amd.mono.runner

DeepSeek-V4.1-Flash mono decode layer: host side.

One DSV41MonoLayer a TP rank runs every mono layer of a decode step (M <= 48 rows): two persistent launches a layer (layer) from the layer's inputs at the attention seam to its outputs at the next one -- the MoE's reduced output, the residual after the FFN seam and that seam's mixes -- with both TP all-reduces inside the kernels (symmetric peer memory). A layer whose attention vLLM runs (ffn) takes one launch from its unreduced attention output on. Every argument is a device pointer: the launches can be captured in a HIP graph. The kernels move their mailbox epoch on themselves; every rank runs the same launches, so the ranks' epochs agree.

Classes:

• AttnWeights –

  One layer's attention tensors on this rank, in vLLM's loaded layout.

• DSV41MonoLayer –

  The mono decode runner of one TP rank (group: its TP group, for the

• MonoLayerWeights –

  One layer's tensors on this rank, in vLLM's loaded layout (attn:

AttnWeights dataclass

One layer's attention tensors on this rank, in vLLM's loaded layout.

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

@dataclass
class AttnWeights:
    """One layer's attention tensors on this rank, in vLLM's loaded layout."""

    layer_id: int
    wqkv: torch.Tensor  # [1792, 5120] e4m3: [wq_a; wkv]
    wqkv_scale: torch.Tensor  # [56, 160] uint8 E8M0 (32 x 32 blocks)
    q_norm: torch.Tensor  # [1280] bf16
    kv_norm: torch.Tensor  # [512] bf16
    wq_b: torch.Tensor  # [H * 512, 1280] e4m3
    wq_b_scale: torch.Tensor  # [H * 16, 40] uint8
    wo_a: torch.Tensor  # [G * 1024, 4096] e4m3
    wo_a_scale: torch.Tensor  # [G * 32, 128] uint8
    wo_b: torch.Tensor  # [5120, G * 1024] e4m3
    wo_b_scale: torch.Tensor  # [160, G * 32] uint8
    attn_sink: torch.Tensor  # [>= H] f32
    cos_sin: torch.Tensor  # [positions, 64] f32 (the layer's rope table)
    ratio: int  # 0 / 1 / 2: the layer's compress ratio

    def check(self, d: Dims) -> None:
        h, g = d.heads, d.groups
        # gfx942 reads its own copies (weights942.linear_copy): FNUZ bytes in
        # the same shapes, as every attention K is a multiple of 128.
        fp8 = torch.uint8 if GFX942 else torch.float8_e4m3fn
        want = {
            "wqkv": ((1792, HIDDEN), fp8),
            "wqkv_scale": ((56, HIDDEN // 32), torch.uint8),
            "wq_b": ((h * HEAD_DIM, 1280), fp8),
            "wq_b_scale": ((h * HEAD_DIM // 32, 40), torch.uint8),
            "wo_a": ((g * 1024, d.group_k), fp8),
            "wo_a_scale": ((g * 32, d.group_k // 32), torch.uint8),
            "wo_b": ((HIDDEN, g * 1024), fp8),
            "wo_b_scale": ((HIDDEN // 32, g * 32), torch.uint8),
        }
        _check_tensors(self, f"layer {self.layer_id}", want)
        assert self.cos_sin.dtype == torch.float32 and self.cos_sin.shape[-1] == 64

DSV41MonoLayer

The mono decode runner of one TP rank (group: its TP group, for the peer-memory handle exchange; a gloo / CPU group).

Methods:

• ffn –

  -> (out, residual, post_mix, res_mix, pre_mix) of a layer whose

• forward –

  -> (out, residual, post_mix, res_mix, pre_mix): the layer's outputs

• scratch –

  Step width tokens's scratch: a width's layout has memory of its own

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

class DSV41MonoLayer:
    """The mono decode runner of one TP rank (``group``: its TP group, for the
    peer-memory handle exchange; a gloo / CPU group)."""

    def __init__(
        self,
        tp: int,
        rank: int,
        group,
        device: torch.device | str = "cuda",
        timeline: bool = False,
        experts: int = EXPERTS,
        topk: int = TOPK,
        ug_groups: int = 2,
        swa_width: int = WINDOW,
    ):
        _ensure_writable_flydsl_cache()
        from .common.peer_memory import PeerBuffer

        self.tp, self.rank = tp, rank
        # A timeline runner builds K1 and K2 with clock stamps. Each CTA
        # writes its stamps into its own record of these buffers (100 MHz
        # ticks, except the two shader clock points of moe.TL_UG_POINTS), and
        # a profiling script reads them after a step.
        self.timeline = timeline
        self.tl1 = torch.zeros(BLOCKS, FRONT_POINTS, dtype=torch.int64, device=device)
        self.tl2 = torch.zeros(BLOCKS, K2_POINTS, dtype=torch.int64, device=device)
        # The FFN launch uses K2's record layout from the seam point on, plus
        # its own start at point 0.
        self.tl3 = torch.zeros(BLOCKS, K2_POINTS, dtype=torch.int64, device=device)
        # the layers it runs: the target's (384 experts, top-6, 128-wide causal SWA
        # rows) or the DSpark draft's, which has a runner of its own
        self.experts, self.topk, self.ug_groups = experts, topk, ug_groups
        self.swa_width = swa_width
        self.d = Dims(tp)
        dev = self.device = torch.device(device)
        self._scratch: dict[int, torch.Tensor] = {}
        # [epoch, -, -, -, a mark per CTA, the MoE's counters]
        self.epoch = torch.zeros(EPOCH_WORDS, dtype=torch.int32, device=dev)
        self.q = torch.zeros(
            MAX_TOKENS, self.d.heads, HEAD_DIM, dtype=torch.bfloat16, device=dev
        )
        self.kt = torch.zeros(MAX_TOKENS * KEYS, dtype=torch.int32, device=dev)
        self.klen = torch.zeros(2 * MAX_TOKENS, dtype=torch.int32, device=dev)
        self._dummy = torch.zeros(4, dtype=torch.int32, device=dev)
        self._zero_rec = torch.zeros(1024, dtype=torch.uint8, device=dev)
        # the first layer's incoming pre-mix: stream 0 of the broadcast embedding
        self._first_pre = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev)
        self._first_pre[:, 0] = 1.0
        # an index source's K1 outputs, vLLM's indexer's inputs
        self.normed = torch.zeros(MAX_TOKENS, HIDDEN, dtype=torch.bfloat16, device=dev)
        self.qr = torch.zeros(MAX_TOKENS, Q_RANK, dtype=torch.bfloat16, device=dev)
        # the attention seam's outputs, K1 -> K2
        self.res_mid = torch.zeros(
            MAX_TOKENS, HC, HIDDEN, dtype=torch.bfloat16, device=dev
        )
        self.post_a = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev)
        self.comb_a = torch.zeros(MAX_TOKENS, HC, HC, dtype=torch.float32, device=dev)
        self.pre_a = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev)
        self.peer = PeerBuffer(2 * peer_half_bytes(tp), group, rank, tp, dev)
        self.peer.bytes.zero_()
        self._qk_scale = (
            torch.tensor([HEAD_DIM**-0.5 * _LOG2E], dtype=torch.float32)
            .view(torch.int32)
            .item()
        )
        self._kernels: dict = {}

    def scratch(self, tokens: int) -> torch.Tensor:
        """Step width ``tokens``'s scratch: a width's layout has memory of its own
        (``layer.scratch_layout``). Allocated at the width's first step, which is
        eager: vLLM runs every graph's batch eagerly before capturing it."""
        buf = self._scratch.get(tokens)
        if buf is None:
            if torch.cuda.is_current_stream_capturing():
                raise RuntimeError(
                    f"DSv4.1 mono decode: step width {tokens} first reached inside "
                    "a CUDA graph capture"
                )
            buf = torch.zeros(
                scratch_bytes(tokens, self.tp, self.experts, self.topk, self.ug_groups),
                dtype=torch.uint8,
                device=self.device,
            )
            self._scratch[tokens] = buf
        return buf

    def first_pre_mix(self, tokens: int) -> torch.Tensor:
        return self._first_pre[:tokens]

    @staticmethod
    def supports(tokens: int) -> bool:
        return 1 <= tokens <= MAX_TOKENS

    def kernels(self, tokens: int, ratio: int, fold: bool = True, index: bool = False):
        key = (tokens, ratio, fold, index)
        if key not in self._kernels:
            b = MonoBuild(
                tokens,
                self.tp,
                ratio,
                timeline=self.timeline,
                fold=fold,
                index=index,
                experts=self.experts,
                topk=self.topk,
                ug_groups=self.ug_groups,
                swa_width=self.swa_width,
            )
            # On gfx942, keylist.launch writes an index source's keys between
            # vLLM's indexer and K2 (forward), so its K2 is the one without
            # the key-table build.
            k2 = replace(b, index=False) if GFX942 else b
            self._kernels[key] = (build_mono_k1(b), build_mono_k2(k2))
        return self._kernels[key]

    def ffn_kernel(self, tokens: int, experts: int, topk: int):
        key = (tokens, "ffn", experts, topk)
        if key not in self._kernels:
            self._kernels[key] = build_mono_ffn(
                MonoBuild(
                    tokens,
                    self.tp,
                    0,
                    timeline=self.timeline,
                    experts=experts,
                    topk=topk,
                )
            )
        return self._kernels[key]

    @staticmethod
    def _outs(M: int, residual: torch.Tensor) -> tuple:
        dev = residual.device
        return (
            torch.empty(M, HIDDEN, dtype=torch.bfloat16, device=dev),
            torch.empty_like(residual),
            torch.empty(M, HC, 1, dtype=torch.float32, device=dev),
            torch.empty(M, HC, HC, dtype=torch.float32, device=dev),
            torch.empty(M, HC, dtype=torch.float32, device=dev),
        )

    def forward(
        self,
        w: MonoLayerWeights,
        x: torch.Tensor,  # [M, 5120] bf16: the previous FFN output (reduced)
        residual: torch.Tensor,  # [M, 4, 5120] bf16
        post_mix: torch.Tensor,  # [M, 4, 1] f32
        res_mix: torch.Tensor,  # [M, 4, 4] f32
        pre_mix: torch.Tensor,  # [M, 4] f32
        positions: torch.Tensor,  # [M] int64
        slot_mapping: torch.Tensor,  # [M] int64
        swa_cache: torch.Tensor,  # [blocks, block, 584] uint8
        swa_indices: torch.Tensor,  # [>= M, 1, 128] int32
        swa_lens: torch.Tensor,  # [>= M] int32
        token_to_req: torch.Tensor,  # [>= M] int32
        topk_indices: torch.Tensor | None = None,
        comp_cache: torch.Tensor | None = None,
        comp_block_table: torch.Tensor | None = None,
        outs: tuple | None = None,
        fold: bool = True,
        select: Callable[[torch.Tensor, torch.Tensor], None] | None = None,
    ):
        """-> (out, residual, post_mix, res_mix, pre_mix): the layer's outputs
        (``outs``: buffers to write them into). ``select(normed, qr)``: an index
        source's top-k selection into ``topk_indices``, run between the
        kernels on K1's normed input and qr."""
        M = x.shape[0]
        assert self.supports(M), M
        a = w.attn
        assert a is not None
        ratio = a.ratio
        index = select is not None
        assert not index or ratio
        k1, k2 = self.kernels(M, ratio, fold, index)
        st = torch.cuda.current_stream()
        if outs is None:
            outs = self._outs(M, residual)
        out, res_out, post_out, comb_out, pre_out = outs
        for t_ in (x, residual, post_mix, res_mix, pre_mix, positions, slot_mapping):
            assert t_.is_contiguous()
        assert positions.dtype == slot_mapping.dtype == torch.int64
        assert swa_indices.dtype == swa_lens.dtype == token_to_req.dtype == torch.int32
        assert (
            swa_indices.stride(0) == self.swa_width and swa_cache.dtype == torch.uint8
        )
        assert swa_cache.shape[-1] == 584 and swa_cache.stride(1) == 584
        swa_block = swa_cache.shape[1]
        if ratio:
            assert topk_indices is not None and comp_cache is not None
            assert comp_block_table is not None
            assert topk_indices.dtype == torch.int32 and topk_indices.stride(0) == 512
            assert comp_cache.dtype == torch.uint8 and comp_cache.stride(1) == 584
            assert comp_block_table.dtype == torch.int32
            comp, comp_stride, comp_block = (
                comp_cache,
                comp_cache.stride(0),
                comp_cache.shape[1],
            )
            bt, bt_stride, topk = (
                comp_block_table,
                comp_block_table.stride(0),
                topk_indices,
            )
        else:
            comp, comp_stride, comp_block = self._dummy, 0, 1
            bt, bt_stride, topk = self._dummy, 0, self._dummy
        k1(
            residual.data_ptr(),
            x.data_ptr(),
            post_mix.data_ptr(),
            res_mix.data_ptr(),
            pre_mix.data_ptr(),
            w.hc_attn_fn.data_ptr(),
            w.hc_attn_scale.data_ptr(),
            w.hc_attn_base.data_ptr(),
            w.attn_norm.data_ptr(),
            self.res_mid.data_ptr(),
            self.post_a.data_ptr(),
            self.comb_a.data_ptr(),
            self.pre_a.data_ptr(),
            a.wqkv.data_ptr(),
            a.wqkv_scale.data_ptr(),
            a.q_norm.data_ptr(),
            a.kv_norm.data_ptr(),
            a.wq_b.data_ptr(),
            a.wq_b_scale.data_ptr(),
            a.cos_sin.data_ptr(),
            positions.data_ptr(),
            slot_mapping.data_ptr(),
            swa_cache.data_ptr(),
            swa_cache.stride(0),
            swa_block,
            self.q.data_ptr(),
            swa_indices.data_ptr(),
            swa_lens.data_ptr(),
            comp_block,
            bt.data_ptr(),
            bt_stride,
            topk.data_ptr(),
            token_to_req.data_ptr(),
            self.kt.data_ptr(),
            self.klen.data_ptr(),
            self.normed.data_ptr(),
            self.qr.data_ptr(),
            self.scratch(M).data_ptr(),
            self.epoch.data_ptr(),
            self.tl1.data_ptr() if self.timeline else 0,
            stream=st,
        )
        if select is not None:
            select(self.normed[:M], self.qr[:M])
            if GFX942:
                # K1 of an index source writes no keys. The gfx942 K2 has no
                # key-table build, so the keys of the top-k that select just
                # wrote are listed here, as K1's stage_kt lists them.
                from .keylist import launch as keylist_launch

                keylist_launch(
                    self.kt,
                    self.klen,
                    slot_mapping,
                    positions,
                    swa_indices,
                    swa_lens,
                    token_to_req,
                    ratio,
                    topk=topk,
                    comp_block_table=bt,
                    comp_block=comp_block,
                )
        k2(
            self.q.data_ptr(),
            swa_cache.data_ptr(),
            swa_cache.stride(0),
            swa_block,
            comp.data_ptr(),
            comp_stride,
            comp_block,
            self.kt.data_ptr(),
            self.klen.data_ptr(),
            slot_mapping.data_ptr(),
            swa_indices.data_ptr(),
            swa_lens.data_ptr(),
            token_to_req.data_ptr(),
            topk.data_ptr(),
            bt.data_ptr(),
            bt_stride,
            positions.data_ptr(),
            a.attn_sink.data_ptr(),
            self._qk_scale,
            a.cos_sin.data_ptr(),
            a.wo_a.data_ptr(),
            a.wo_a_scale.data_ptr(),
            a.wo_b.data_ptr(),
            a.wo_b_scale.data_ptr(),
            self._zero_rec.data_ptr(),
            self.res_mid.data_ptr(),
            self.post_a.data_ptr(),
            self.comb_a.data_ptr(),
            self.pre_a.data_ptr(),
            w.hc_ffn_fn.data_ptr(),
            w.hc_ffn_scale.data_ptr(),
            w.hc_ffn_base.data_ptr(),
            w.ffn_norm.data_ptr(),
            res_out.data_ptr(),
            post_out.data_ptr(),
            comb_out.data_ptr(),
            pre_out.data_ptr(),
            w.gate_w.data_ptr(),
            w.bias.data_ptr(),
            w.w13.data_ptr(),
            w.w13_s.data_ptr(),
            w.w2.data_ptr(),
            w.w2_s.data_ptr(),
            w.sgu.data_ptr(),
            w.sgu_s.data_ptr(),
            w.sw2.data_ptr(),
            w.sw2_s.data_ptr(),
            out.data_ptr(),
            self.scratch(M).data_ptr(),
            self.peer.local,
            self.peer.addresses.data_ptr(),
            self.rank,
            self.epoch.data_ptr(),
            self.tl2.data_ptr() if self.timeline else 0,
            stream=st,
        )
        return outs

    def ffn(
        self,
        w: MonoLayerWeights,
        part: torch.Tensor,  # [M, 5120] bf16: this rank's unreduced wo_b output
        residual: torch.Tensor,  # [M, 4, 5120] bf16, after the attention seam
        post_mix: torch.Tensor,  # [M, 4, 1] f32, the attention seam's
        res_mix: torch.Tensor,  # [M, 4, 4] f32
        pre_mix: torch.Tensor,  # [M, 4] f32
        outs: tuple | None = None,
        topk: int = 6,
    ):
        """-> (out, residual, post_mix, res_mix, pre_mix) of a layer whose
        attention vLLM ran: the attention's TP reduce, the FFN seam, the MoE and
        its all-reduce in one launch. The routed expert count is the gate
        weight's row count, and ``topk`` is the layer's picks a token: 384 / 6
        for the target, 128 / 3 for a DSpark draft layer."""
        M = part.shape[0]
        assert self.supports(M), M
        for t_ in (part, residual, post_mix, res_mix, pre_mix):
            assert t_.is_contiguous()
        assert part.dtype == residual.dtype == torch.bfloat16
        if outs is None:
            outs = self._outs(M, residual)
        out, res_out, post_out, comb_out, pre_out = outs
        self.ffn_kernel(M, w.gate_w.shape[0], topk)(
            part.data_ptr(),
            residual.data_ptr(),
            post_mix.data_ptr(),
            res_mix.data_ptr(),
            pre_mix.data_ptr(),
            w.hc_ffn_fn.data_ptr(),
            w.hc_ffn_scale.data_ptr(),
            w.hc_ffn_base.data_ptr(),
            w.ffn_norm.data_ptr(),
            res_out.data_ptr(),
            post_out.data_ptr(),
            comb_out.data_ptr(),
            pre_out.data_ptr(),
            w.gate_w.data_ptr(),
            w.bias.data_ptr(),
            w.w13.data_ptr(),
            w.w13_s.data_ptr(),
            w.w2.data_ptr(),
            w.w2_s.data_ptr(),
            w.sgu.data_ptr(),
            w.sgu_s.data_ptr(),
            w.sw2.data_ptr(),
            w.sw2_s.data_ptr(),
            out.data_ptr(),
            self.scratch(M).data_ptr(),
            self.peer.local,
            self.peer.addresses.data_ptr(),
            self.rank,
            self.epoch.data_ptr(),
            self.tl3.data_ptr() if self.timeline else 0,
            stream=torch.cuda.current_stream(),
        )
        return outs

ffn(w, part, residual, post_mix, res_mix, pre_mix, outs=None, topk=6)

-> (out, residual, post_mix, res_mix, pre_mix) of a layer whose attention vLLM ran: the attention's TP reduce, the FFN seam, the MoE and its all-reduce in one launch. The routed expert count is the gate weight's row count, and topk is the layer's picks a token: 384 / 6 for the target, 128 / 3 for a DSpark draft layer.

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

def ffn(
    self,
    w: MonoLayerWeights,
    part: torch.Tensor,  # [M, 5120] bf16: this rank's unreduced wo_b output
    residual: torch.Tensor,  # [M, 4, 5120] bf16, after the attention seam
    post_mix: torch.Tensor,  # [M, 4, 1] f32, the attention seam's
    res_mix: torch.Tensor,  # [M, 4, 4] f32
    pre_mix: torch.Tensor,  # [M, 4] f32
    outs: tuple | None = None,
    topk: int = 6,
):
    """-> (out, residual, post_mix, res_mix, pre_mix) of a layer whose
    attention vLLM ran: the attention's TP reduce, the FFN seam, the MoE and
    its all-reduce in one launch. The routed expert count is the gate
    weight's row count, and ``topk`` is the layer's picks a token: 384 / 6
    for the target, 128 / 3 for a DSpark draft layer."""
    M = part.shape[0]
    assert self.supports(M), M
    for t_ in (part, residual, post_mix, res_mix, pre_mix):
        assert t_.is_contiguous()
    assert part.dtype == residual.dtype == torch.bfloat16
    if outs is None:
        outs = self._outs(M, residual)
    out, res_out, post_out, comb_out, pre_out = outs
    self.ffn_kernel(M, w.gate_w.shape[0], topk)(
        part.data_ptr(),
        residual.data_ptr(),
        post_mix.data_ptr(),
        res_mix.data_ptr(),
        pre_mix.data_ptr(),
        w.hc_ffn_fn.data_ptr(),
        w.hc_ffn_scale.data_ptr(),
        w.hc_ffn_base.data_ptr(),
        w.ffn_norm.data_ptr(),
        res_out.data_ptr(),
        post_out.data_ptr(),
        comb_out.data_ptr(),
        pre_out.data_ptr(),
        w.gate_w.data_ptr(),
        w.bias.data_ptr(),
        w.w13.data_ptr(),
        w.w13_s.data_ptr(),
        w.w2.data_ptr(),
        w.w2_s.data_ptr(),
        w.sgu.data_ptr(),
        w.sgu_s.data_ptr(),
        w.sw2.data_ptr(),
        w.sw2_s.data_ptr(),
        out.data_ptr(),
        self.scratch(M).data_ptr(),
        self.peer.local,
        self.peer.addresses.data_ptr(),
        self.rank,
        self.epoch.data_ptr(),
        self.tl3.data_ptr() if self.timeline else 0,
        stream=torch.cuda.current_stream(),
    )
    return outs

forward(w, x, residual, post_mix, res_mix, pre_mix, positions, slot_mapping, swa_cache, swa_indices, swa_lens, token_to_req, topk_indices=None, comp_cache=None, comp_block_table=None, outs=None, fold=True, select=None)

-> (out, residual, post_mix, res_mix, pre_mix): the layer's outputs (outs: buffers to write them into). select(normed, qr): an index source's top-k selection into topk_indices, run between the kernels on K1's normed input and qr.

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

def forward(
    self,
    w: MonoLayerWeights,
    x: torch.Tensor,  # [M, 5120] bf16: the previous FFN output (reduced)
    residual: torch.Tensor,  # [M, 4, 5120] bf16
    post_mix: torch.Tensor,  # [M, 4, 1] f32
    res_mix: torch.Tensor,  # [M, 4, 4] f32
    pre_mix: torch.Tensor,  # [M, 4] f32
    positions: torch.Tensor,  # [M] int64
    slot_mapping: torch.Tensor,  # [M] int64
    swa_cache: torch.Tensor,  # [blocks, block, 584] uint8
    swa_indices: torch.Tensor,  # [>= M, 1, 128] int32
    swa_lens: torch.Tensor,  # [>= M] int32
    token_to_req: torch.Tensor,  # [>= M] int32
    topk_indices: torch.Tensor | None = None,
    comp_cache: torch.Tensor | None = None,
    comp_block_table: torch.Tensor | None = None,
    outs: tuple | None = None,
    fold: bool = True,
    select: Callable[[torch.Tensor, torch.Tensor], None] | None = None,
):
    """-> (out, residual, post_mix, res_mix, pre_mix): the layer's outputs
    (``outs``: buffers to write them into). ``select(normed, qr)``: an index
    source's top-k selection into ``topk_indices``, run between the
    kernels on K1's normed input and qr."""
    M = x.shape[0]
    assert self.supports(M), M
    a = w.attn
    assert a is not None
    ratio = a.ratio
    index = select is not None
    assert not index or ratio
    k1, k2 = self.kernels(M, ratio, fold, index)
    st = torch.cuda.current_stream()
    if outs is None:
        outs = self._outs(M, residual)
    out, res_out, post_out, comb_out, pre_out = outs
    for t_ in (x, residual, post_mix, res_mix, pre_mix, positions, slot_mapping):
        assert t_.is_contiguous()
    assert positions.dtype == slot_mapping.dtype == torch.int64
    assert swa_indices.dtype == swa_lens.dtype == token_to_req.dtype == torch.int32
    assert (
        swa_indices.stride(0) == self.swa_width and swa_cache.dtype == torch.uint8
    )
    assert swa_cache.shape[-1] == 584 and swa_cache.stride(1) == 584
    swa_block = swa_cache.shape[1]
    if ratio:
        assert topk_indices is not None and comp_cache is not None
        assert comp_block_table is not None
        assert topk_indices.dtype == torch.int32 and topk_indices.stride(0) == 512
        assert comp_cache.dtype == torch.uint8 and comp_cache.stride(1) == 584
        assert comp_block_table.dtype == torch.int32
        comp, comp_stride, comp_block = (
            comp_cache,
            comp_cache.stride(0),
            comp_cache.shape[1],
        )
        bt, bt_stride, topk = (
            comp_block_table,
            comp_block_table.stride(0),
            topk_indices,
        )
    else:
        comp, comp_stride, comp_block = self._dummy, 0, 1
        bt, bt_stride, topk = self._dummy, 0, self._dummy
    k1(
        residual.data_ptr(),
        x.data_ptr(),
        post_mix.data_ptr(),
        res_mix.data_ptr(),
        pre_mix.data_ptr(),
        w.hc_attn_fn.data_ptr(),
        w.hc_attn_scale.data_ptr(),
        w.hc_attn_base.data_ptr(),
        w.attn_norm.data_ptr(),
        self.res_mid.data_ptr(),
        self.post_a.data_ptr(),
        self.comb_a.data_ptr(),
        self.pre_a.data_ptr(),
        a.wqkv.data_ptr(),
        a.wqkv_scale.data_ptr(),
        a.q_norm.data_ptr(),
        a.kv_norm.data_ptr(),
        a.wq_b.data_ptr(),
        a.wq_b_scale.data_ptr(),
        a.cos_sin.data_ptr(),
        positions.data_ptr(),
        slot_mapping.data_ptr(),
        swa_cache.data_ptr(),
        swa_cache.stride(0),
        swa_block,
        self.q.data_ptr(),
        swa_indices.data_ptr(),
        swa_lens.data_ptr(),
        comp_block,
        bt.data_ptr(),
        bt_stride,
        topk.data_ptr(),
        token_to_req.data_ptr(),
        self.kt.data_ptr(),
        self.klen.data_ptr(),
        self.normed.data_ptr(),
        self.qr.data_ptr(),
        self.scratch(M).data_ptr(),
        self.epoch.data_ptr(),
        self.tl1.data_ptr() if self.timeline else 0,
        stream=st,
    )
    if select is not None:
        select(self.normed[:M], self.qr[:M])
        if GFX942:
            # K1 of an index source writes no keys. The gfx942 K2 has no
            # key-table build, so the keys of the top-k that select just
            # wrote are listed here, as K1's stage_kt lists them.
            from .keylist import launch as keylist_launch

            keylist_launch(
                self.kt,
                self.klen,
                slot_mapping,
                positions,
                swa_indices,
                swa_lens,
                token_to_req,
                ratio,
                topk=topk,
                comp_block_table=bt,
                comp_block=comp_block,
            )
    k2(
        self.q.data_ptr(),
        swa_cache.data_ptr(),
        swa_cache.stride(0),
        swa_block,
        comp.data_ptr(),
        comp_stride,
        comp_block,
        self.kt.data_ptr(),
        self.klen.data_ptr(),
        slot_mapping.data_ptr(),
        swa_indices.data_ptr(),
        swa_lens.data_ptr(),
        token_to_req.data_ptr(),
        topk.data_ptr(),
        bt.data_ptr(),
        bt_stride,
        positions.data_ptr(),
        a.attn_sink.data_ptr(),
        self._qk_scale,
        a.cos_sin.data_ptr(),
        a.wo_a.data_ptr(),
        a.wo_a_scale.data_ptr(),
        a.wo_b.data_ptr(),
        a.wo_b_scale.data_ptr(),
        self._zero_rec.data_ptr(),
        self.res_mid.data_ptr(),
        self.post_a.data_ptr(),
        self.comb_a.data_ptr(),
        self.pre_a.data_ptr(),
        w.hc_ffn_fn.data_ptr(),
        w.hc_ffn_scale.data_ptr(),
        w.hc_ffn_base.data_ptr(),
        w.ffn_norm.data_ptr(),
        res_out.data_ptr(),
        post_out.data_ptr(),
        comb_out.data_ptr(),
        pre_out.data_ptr(),
        w.gate_w.data_ptr(),
        w.bias.data_ptr(),
        w.w13.data_ptr(),
        w.w13_s.data_ptr(),
        w.w2.data_ptr(),
        w.w2_s.data_ptr(),
        w.sgu.data_ptr(),
        w.sgu_s.data_ptr(),
        w.sw2.data_ptr(),
        w.sw2_s.data_ptr(),
        out.data_ptr(),
        self.scratch(M).data_ptr(),
        self.peer.local,
        self.peer.addresses.data_ptr(),
        self.rank,
        self.epoch.data_ptr(),
        self.tl2.data_ptr() if self.timeline else 0,
        stream=st,
    )
    return outs

scratch(tokens)

Step width tokens's scratch: a width's layout has memory of its own (layer.scratch_layout). Allocated at the width's first step, which is eager: vLLM runs every graph's batch eagerly before capturing it.

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

def scratch(self, tokens: int) -> torch.Tensor:
    """Step width ``tokens``'s scratch: a width's layout has memory of its own
    (``layer.scratch_layout``). Allocated at the width's first step, which is
    eager: vLLM runs every graph's batch eagerly before capturing it."""
    buf = self._scratch.get(tokens)
    if buf is None:
        if torch.cuda.is_current_stream_capturing():
            raise RuntimeError(
                f"DSv4.1 mono decode: step width {tokens} first reached inside "
                "a CUDA graph capture"
            )
        buf = torch.zeros(
            scratch_bytes(tokens, self.tp, self.experts, self.topk, self.ug_groups),
            dtype=torch.uint8,
            device=self.device,
        )
        self._scratch[tokens] = buf
    return buf

MonoLayerWeights dataclass

One layer's tensors on this rank, in vLLM's loaded layout (attn: None for a layer whose attention vLLM runs).

Methods:

• check –

  The MoE's tensors: experts routed experts (384; the DSpark draft's

• check_942 –

  The MoE's tensors on gfx942, TP4, which are the copies that

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

@dataclass
class MonoLayerWeights:
    """One layer's tensors on this rank, in vLLM's loaded layout (``attn``:
    None for a layer whose attention vLLM runs)."""

    attn: AttnWeights | None
    hc_attn_fn: torch.Tensor  # [24, 4 * 5120] f32
    hc_attn_scale: torch.Tensor  # [3] f32
    hc_attn_base: torch.Tensor  # [24] f32
    attn_norm: torch.Tensor  # [5120] bf16
    hc_ffn_fn: torch.Tensor
    hc_ffn_scale: torch.Tensor
    hc_ffn_base: torch.Tensor
    ffn_norm: torch.Tensor
    gate_w: torch.Tensor  # [384, 5120] bf16
    bias: torch.Tensor  # [384] f32 (e_score_correction_bias)
    w13: torch.Tensor  # [384, 2 inter, 2560] fp4x2, aiter (16, 16) shuffle
    w13_s: torch.Tensor  # its scales, aiter shuffle_scale order
    w2: torch.Tensor  # [384, 5120, inter / 2] fp4x2
    w2_s: torch.Tensor
    sgu: torch.Tensor  # shared gate_up [2 inter, 5120] e4m3, row-major
    sgu_s: torch.Tensor  # [2 inter / 32, 160] E8M0
    sw2: torch.Tensor  # shared down [5120, inter] e4m3
    sw2_s: torch.Tensor

    def check(self, tp: int, experts: int = EXPERTS) -> None:
        """The MoE's tensors: ``experts`` routed experts (384; the DSpark draft's
        128), TP-sharded intermediates (the routed one padded as the loader pads
        it), AITER's A8W4 layout, 32 x 32 E8M0 blocks for the shared expert."""
        d, fp4, e4m3, u8 = (
            MoeDims(tp),
            torch.float4_e2m1fn_x2,
            torch.float8_e4m3fn,
            torch.uint8,
        )
        want = {
            "gate_w": ((experts, HIDDEN), torch.bfloat16),
            "bias": ((experts,), torch.float32),
            "w13": ((experts, 2 * d.inter, HIDDEN // 2), fp4),
            "w13_s": ((experts * 2 * d.inter, HIDDEN // 32), u8),
            "w2": ((experts, HIDDEN, d.inter // 2), fp4),
            "w2_s": ((experts * HIDDEN, d.down_scale_cols), u8),
            "sgu": ((2 * d.sh_inter, HIDDEN), e4m3),
            "sgu_s": ((2 * d.sh_inter // 32, HIDDEN // 32), u8),
            "sw2": ((HIDDEN, d.sh_inter), e4m3),
            "sw2_s": ((HIDDEN // 32, d.sh_inter // 32), u8),
        }
        _check_tensors(self, "MoE", want)

    def check_942(self, experts: int) -> None:
        """The MoE's tensors on gfx942, TP4, which are the copies that
        ``weights942`` makes. The routed experts have their real intermediate
        width without the loader's padding, as FP4 bytes in ``fp4_tile_major``
        order and E8M0 codes a row each. The shared expert is FNUZ bytes from
        ``linear_copy``, with K padded to a multiple of 128. ``experts`` is 384
        for the target layers and 128 for the DSpark draft layers."""
        d, u8 = MoeDims(4), torch.uint8
        inter, kp = d.inter_real, -(-d.sh_inter // 128) * 128
        want = {
            "gate_w": ((experts, HIDDEN), torch.bfloat16),
            "bias": ((experts,), torch.float32),
            "w13": ((experts, 2 * inter, HIDDEN // 2), u8),
            "w13_s": ((experts, 2 * inter, HIDDEN // 32), u8),
            "w2": ((experts, HIDDEN, inter // 2), u8),
            "w2_s": ((experts, HIDDEN, inter // 32), u8),
            "sgu": ((2 * d.sh_inter, HIDDEN), u8),
            "sgu_s": ((2 * d.sh_inter // 32, HIDDEN // 32), u8),
            "sw2": ((HIDDEN, kp), u8),
            "sw2_s": ((HIDDEN // 32, kp // 32), u8),
        }
        _check_tensors(self, "MoE", want)

check(tp, experts=EXPERTS)

The MoE's tensors: experts routed experts (384; the DSpark draft's 128), TP-sharded intermediates (the routed one padded as the loader pads it), AITER's A8W4 layout, 32 x 32 E8M0 blocks for the shared expert.

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

def check(self, tp: int, experts: int = EXPERTS) -> None:
    """The MoE's tensors: ``experts`` routed experts (384; the DSpark draft's
    128), TP-sharded intermediates (the routed one padded as the loader pads
    it), AITER's A8W4 layout, 32 x 32 E8M0 blocks for the shared expert."""
    d, fp4, e4m3, u8 = (
        MoeDims(tp),
        torch.float4_e2m1fn_x2,
        torch.float8_e4m3fn,
        torch.uint8,
    )
    want = {
        "gate_w": ((experts, HIDDEN), torch.bfloat16),
        "bias": ((experts,), torch.float32),
        "w13": ((experts, 2 * d.inter, HIDDEN // 2), fp4),
        "w13_s": ((experts * 2 * d.inter, HIDDEN // 32), u8),
        "w2": ((experts, HIDDEN, d.inter // 2), fp4),
        "w2_s": ((experts * HIDDEN, d.down_scale_cols), u8),
        "sgu": ((2 * d.sh_inter, HIDDEN), e4m3),
        "sgu_s": ((2 * d.sh_inter // 32, HIDDEN // 32), u8),
        "sw2": ((HIDDEN, d.sh_inter), e4m3),
        "sw2_s": ((HIDDEN // 32, d.sh_inter // 32), u8),
    }
    _check_tensors(self, "MoE", want)

check_942(experts)

The MoE's tensors on gfx942, TP4, which are the copies that weights942 makes. The routed experts have their real intermediate width without the loader's padding, as FP4 bytes in fp4_tile_major order and E8M0 codes a row each. The shared expert is FNUZ bytes from linear_copy, with K padded to a multiple of 128. experts is 384 for the target layers and 128 for the DSpark draft layers.

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

def check_942(self, experts: int) -> None:
    """The MoE's tensors on gfx942, TP4, which are the copies that
    ``weights942`` makes. The routed experts have their real intermediate
    width without the loader's padding, as FP4 bytes in ``fp4_tile_major``
    order and E8M0 codes a row each. The shared expert is FNUZ bytes from
    ``linear_copy``, with K padded to a multiple of 128. ``experts`` is 384
    for the target layers and 128 for the DSpark draft layers."""
    d, u8 = MoeDims(4), torch.uint8
    inter, kp = d.inter_real, -(-d.sh_inter // 128) * 128
    want = {
        "gate_w": ((experts, HIDDEN), torch.bfloat16),
        "bias": ((experts,), torch.float32),
        "w13": ((experts, 2 * inter, HIDDEN // 2), u8),
        "w13_s": ((experts, 2 * inter, HIDDEN // 32), u8),
        "w2": ((experts, HIDDEN, inter // 2), u8),
        "w2_s": ((experts, HIDDEN, inter // 32), u8),
        "sgu": ((2 * d.sh_inter, HIDDEN), u8),
        "sgu_s": ((2 * d.sh_inter // 32, HIDDEN // 32), u8),
        "sw2": ((HIDDEN, kp), u8),
        "sw2_s": ((HIDDEN // 32, kp // 32), u8),
    }
    _check_tensors(self, "MoE", want)

_check_tensors(owner, tag, want)

Each named tensor of owner has the (shape, dtype) the kernels read, contiguous: their buffer loads are unbounded, so a mismatch would read garbage rather than fault.

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

def _check_tensors(owner, tag: str, want: dict) -> None:
    """Each named tensor of ``owner`` has the (shape, dtype) the kernels read,
    contiguous: their buffer loads are unbounded, so a mismatch would read
    garbage rather than fault."""
    for name, (shape, dtype) in want.items():
        t = getattr(owner, name)
        assert tuple(t.shape) == shape and t.dtype == dtype and t.is_contiguous(), (
            f"{tag} {name}: {tuple(t.shape)} {t.dtype}, want {shape} {dtype}"
        )

_ensure_writable_flydsl_cache()

Aiter points FlyDSL at its bundled, read-only cache; new kernels need a writable one.

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

def _ensure_writable_flydsl_cache() -> None:
    """Aiter points FlyDSL at its bundled, read-only cache; new kernels need a
    writable one."""
    cur = os.environ.get("FLYDSL_RUNTIME_CACHE_DIR")
    if cur and os.access(cur, os.W_OK):
        return
    path = Path.home() / ".flydsl" / "cache"
    path.mkdir(parents=True, exist_ok=True)
    os.environ["FLYDSL_RUNTIME_CACHE_DIR"] = str(path)