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:

• back_scratch –

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

• 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

Dims dataclass

Attributes:

• group_k (int) –

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

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), self.tp

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

    @property
    def head_tiles(self) -> int:
        return self.heads // 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.

back_scratch(s, d, start=0)

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

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

def back_scratch(s: int, d: Dims, start: int = 0) -> dict:
    """K2's scratch regions, laid out past ``start``: published data (plain)
    and flag pairs."""
    ck = (
        64 if d.head_tiles == 2 and s * cdiv(KEYS, 64) <= 256 else 128
    )  # back.chunk_keys
    nchunk = cdiv(KEYS, ck)
    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", s * d.heads // 8 * PAIR),
            ("x8b", s * d.o_rows),
            ("x8bs", s * d.o_rows // 32 * 4),
            ("ardy", d.woa_tasks * PAIR),
        ],
        start=start,
    )

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)

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

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

def front_scratch(s: int, d: Dims) -> dict:
    """K1's scratch regions -> (byte offset, bytes): published data (plain
    words behind a flag) and mailbox pairs."""
    return byte_layout(
        [
            ("x8", s * HIDDEN),  # e4m3 words (S > 16: quantized once, shared)
            ("x8s", s * HIDDEN // 32 * 4),  # an i32 code a group
            ("xrdy", WQKV_KPARTS * s * PAIR),  # a flag a (token, K half)
            ("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", s * PAIR),
        ]
    )