vllm.models.deepseek_v41.amd.mono.attention.plan

Geometry of the DeepSeek-V4.1 attention mega kernel: model widths, a TP rank's share of them, task counts and the scratch (mailbox) layout.

Classes:

• Dims –

Functions:

• ardy_tiles –

  Token tiles wo_a publishes ready flags for (1: one flag a task).

• back_scratch –

  K2's scratch regions, laid out past start: published data (plain)

• flag_copy –

  The first pair of the flag copy CTA bid polls (generic: 0).

• flag_stride –

  Pairs a copy of an n-flag region: whole 128 B lines.

• front_end –

  Where K2's regions start: past K1's, which share the launch tag.

• front_scratch –

  K1's scratch regions -> (byte offset, bytes): published data (plain

• on_copy –

  Flag i of the copy at pair copy (flag_copy).

Dims dataclass

Attributes:

• group_k (int) –

  wo_a's K a group: its eight heads' outputs.

• head_tiles (int) –

  16-head MFMA tiles a rank: TP8's 8 heads take one, half of it

• heads_pad (int) –

  The heads the split stage's LDS and MFMAs hold: whole head tiles.

Source code in vllm/models/deepseek_v41/amd/mono/attention/plan.py

@dataclass(frozen=True)
class Dims:
    tp: int

    def __post_init__(self):
        assert self.tp in (2, 4, 8), self.tp

    @property
    def heads(self) -> int:
        return N_HEADS // self.tp

    @property
    def head_tiles(self) -> int:
        """16-head MFMA tiles a rank: TP8's 8 heads take one, half of it
        padding (``heads_pad``)."""
        return cdiv(self.heads, HEAD_TILE)

    @property
    def heads_pad(self) -> int:
        """The heads the split stage's LDS and MFMAs hold: whole head tiles.
        The padded heads' columns are computed (each head is its own score
        maximum and its own PV column) but never stored."""
        return self.head_tiles * HEAD_TILE

    @property
    def groups(self) -> int:
        return O_GROUPS // self.tp

    @property
    def group_k(self) -> int:
        """wo_a's K a group: its eight heads' outputs."""
        return self.heads // self.groups * HEAD_DIM  # 4096

    @property
    def o_rows(self) -> int:
        return self.groups * O_RANK

    @property
    def q_rows(self) -> int:
        return self.heads * HEAD_DIM

    @property
    def wqb_tasks(self) -> int:
        return self.q_rows // ROWS

    @property
    def woa_tasks(self) -> int:
        return self.o_rows // WOA_ROWS

    @property
    def wob_tasks(self) -> int:
        return HIDDEN // ROWS  # 320

group_k property

wo_a's K a group: its eight heads' outputs.

head_tiles property

16-head MFMA tiles a rank: TP8's 8 heads take one, half of it padding (heads_pad).

heads_pad property

The heads the split stage's LDS and MFMAs hold: whole head tiles. The padded heads' columns are computed (each head is its own score maximum and its own PV column) but never stored.

ardy_tiles(s, tuned=False)

Token tiles wo_a publishes ready flags for (1: one flag a task).

Source code in vllm/models/deepseek_v41/amd/mono/attention/plan.py

def ardy_tiles(s: int, tuned: bool = False) -> int:
    """Token tiles wo_a publishes ready flags for (1: one flag a task)."""
    nt = len(token_tiles(s))
    return nt if tuned and nt > 1 and s >= ARDY_MIN_S and s % TILE != 0 else 1

back_scratch(s, d, start=0, tuned=False)

K2's scratch regions, laid out past start: published data (plain) and flag pairs (tuned: the polled flags' copies).

Source code in vllm/models/deepseek_v41/amd/mono/attention/plan.py

def back_scratch(s: int, d: Dims, start: int = 0, tuned: bool = False) -> dict:
    """K2's scratch regions, laid out past ``start``: published data (plain)
    and flag pairs (``tuned``: the polled flags' copies)."""
    ck = (
        64 if GFX942 or (d.head_tiles == 2 and s * cdiv(KEYS, 64) <= 256) else 128
    )  # back.chunk_keys
    nchunk = cdiv(KEYS, ck)
    crdy = s * d.heads // 8
    ardy = (COPIES * ardy_tiles(s, True) if tuned else 1) * d.woa_tasks
    return byte_layout(
        [
            ("part", s * nchunk * d.heads * HEAD_DIM * 2),  # bf16 partials
            ("pstat", s * nchunk * d.heads * 8),  # (max, sum) f32
            ("srdy", s * nchunk * PAIR),
            ("xo", s * d.q_rows),  # e4m3 words
            ("xos", s * d.q_rows // 32 * 4),  # an i32 code a group
            ("crdy", (COPIES * flag_stride(crdy) if tuned else crdy) * PAIR),
            ("x8b", s * d.o_rows),
            ("x8bs", s * d.o_rows // 32 * 4),
            ("ardy", ardy * PAIR),
        ],
        start=start,
    )

flag_copy(c, bid, stride)

The first pair of the flag copy CTA bid polls (generic: 0).

Source code in vllm/models/deepseek_v41/amd/mono/attention/plan.py

def flag_copy(c, bid, stride):
    """The first pair of the flag copy CTA ``bid`` polls (generic: 0)."""
    return bid % COPIES * stride if c["tuned"] else 0

flag_stride(n)

Pairs a copy of an n-flag region: whole 128 B lines.

Source code in vllm/models/deepseek_v41/amd/mono/attention/plan.py

def flag_stride(n: int) -> int:
    """Pairs a copy of an n-flag region: whole 128 B lines."""
    return -(-n // 16) * 16

front_end(s, d)

Where K2's regions start: past K1's, which share the launch tag.

Source code in vllm/models/deepseek_v41/amd/mono/attention/plan.py

def front_end(s: int, d: Dims) -> int:
    """Where K2's regions start: past K1's, which share the launch tag."""
    return round_up(max(o + n for o, n in front_scratch(s, d).values()), 256)

front_scratch(s, d, tuned=False)

K1's scratch regions -> (byte offset, bytes): published data (plain words behind a flag) and mailbox pairs (tuned: the polled flags' copies).

Source code in vllm/models/deepseek_v41/amd/mono/attention/plan.py

def front_scratch(s: int, d: Dims, tuned: bool = False) -> dict:
    """K1's scratch regions -> (byte offset, bytes): published data (plain
    words behind a flag) and mailbox pairs (``tuned``: the polled flags'
    copies)."""
    xrdy = WQKV_KPARTS * s  # a flag a (token, K half)
    return byte_layout(
        [
            ("x8", s * HIDDEN),  # e4m3 words (S > 16: quantized once, shared)
            ("x8s", s * HIDDEN // 32 * 4),  # an i32 code a group
            ("xrdy", (COPIES * flag_stride(xrdy) if tuned else xrdy) * PAIR),
            ("qkvp", WQKV_KPARTS * s * QKV_ROWS * PAIR),
            ("qx8", s * Q_RANK),  # e4m3 words
            ("qx8s", s * Q_RANK // 32 * 4),  # an i32 code a group
            ("qrdy", (COPIES * flag_stride(s) if tuned else s) * PAIR),
        ]
    )

on_copy(copy, i)

Flag i of the copy at pair copy (flag_copy).

Source code in vllm/models/deepseek_v41/amd/mono/attention/plan.py

def on_copy(copy, i):
    """Flag ``i`` of the copy at pair ``copy`` (``flag_copy``)."""
    return i if isinstance(copy, int) and copy == 0 else copy + i