vllm.models.deepseek_v4.amd.mono.dims

DeepSeek-V4-Pro's widths as the mono kernels read them, and Dims, a TP rank's share of them with every count the stages derive from it.

Classes:

• Dims –

  A TP rank's share of the model's widths.

Functions:

• cdiv –

  ceil(a / b) for a >= 0, b > 0.

• ld –

  The row stride of a row-major FP8 weight of K = k (bytes): vLLM pads

• tile_rows –

  The LDS rows an x tile holds.

• token_tiles –

  A step's tokens in GEMV passes: [(first token, tokens)]. A GEMV loads its

Dims dataclass

A TP rank's share of the model's widths.

Attributes:

• down_scale_cols (int) –

  w2_s's e8m0 columns: inter / 32, padded to a multiple of 8

• group_k (int) –

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

• inter (int) –

  ... padded as the loader pads it (the rest zero weights).

• inter_real (int) –

  A routed expert's real intermediate width on this rank.

• keys (int) –

  A token's attention keys at most: its top-k, then its window.

• mid_words (int) –

  MXFP8 words of a pick's routed intermediate.

• sh_inter (int) –

  The shared expert's intermediate width on this rank.

• topk_width (int) –

  A token's compressed keys at most.

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

@dataclass(frozen=True)
class Dims:
    """A TP rank's share of the model's widths."""

    tp: int
    # a token's compressed keys at most (0: the CSA top-k): an HCA layer's
    # attention over every compressed row of a long context reads more
    topk: int = 0

    def __post_init__(self):
        shard = Shard(self.tp)
        shard.split(N_HEADS, "attention heads")
        shard.split(O_GROUPS, "output groups")
        shard.split(MOE_INTER, "routed intermediate")
        shard.split(SHARED_INTER, "shared intermediate")

    # ------------------------------------------------------------ attention
    @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 heads' outputs."""
        return self.heads // self.groups * HEAD_DIM

    @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 topk_width(self) -> int:
        """A token's compressed keys at most."""
        return self.topk or INDEX_TOPK

    @property
    def keys(self) -> int:
        """A token's attention keys at most: its top-k, then its window."""
        return self.topk_width + WINDOW

    # ------------------------------------------------------------ MoE
    @property
    def inter_real(self) -> int:
        """A routed expert's real intermediate width on this rank."""
        return MOE_INTER // self.tp

    @property
    def inter(self) -> int:
        """... padded as the loader pads it (the rest zero weights)."""
        return padded(self.inter_real, EXPERT_PAD)

    @property
    def sh_inter(self) -> int:
        """The shared expert's intermediate width on this rank."""
        return SHARED_INTER // self.tp

    @property
    def shared_tasks(self) -> int:
        return self.sh_inter // ROWS

    @property
    def mid_words(self) -> int:
        """MXFP8 words of a pick's routed intermediate."""
        return self.inter // 4

    @property
    def down_scale_cols(self) -> int:
        """w2_s's e8m0 columns: inter / 32, padded to a multiple of 8
        (aiter ``shuffle_scale``)."""
        return padded(self.inter // 32, 8)

down_scale_cols property

w2_s's e8m0 columns: inter / 32, padded to a multiple of 8 (aiter shuffle_scale).

group_k property

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

inter property

... padded as the loader pads it (the rest zero weights).

inter_real property

A routed expert's real intermediate width on this rank.

keys property

A token's attention keys at most: its top-k, then its window.

mid_words property

MXFP8 words of a pick's routed intermediate.

sh_inter property

The shared expert's intermediate width on this rank.

topk_width property

A token's compressed keys at most.

cdiv(a, b)

ceil(a / b) for a >= 0, b > 0.

Source code in vllm/models/deepseek_v4/amd/mono/common/plan.py

def cdiv(a, b):
    """ceil(a / b) for a >= 0, b > 0."""
    return i32((a + b - 1) // b)

ld(k)

The row stride of a row-major FP8 weight of K = k (bytes): vLLM pads the rows of a multiple of 512 bytes by 256 (VLLM_ROCM_FP8_PADDING).

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

def ld(k: int) -> int:
    """The row stride of a row-major FP8 weight of K = ``k`` (bytes): vLLM pads
    the rows of a multiple of 512 bytes by 256 (``VLLM_ROCM_FP8_PADDING``)."""
    return k + 256 if k % 512 == 0 else k

tile_rows(s)

The LDS rows an x tile holds.

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

def tile_rows(s):
    """The LDS rows an x tile holds."""
    return min(s, TILE)

token_tiles(s)

A step's tokens in GEMV passes: [(first token, tokens)]. A GEMV loads its weights once and runs every pass on them, an x tile in LDS at a time.

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

def token_tiles(s):
    """A step's tokens in GEMV passes: [(first token, tokens)]. A GEMV loads its
    weights once and runs every pass on them, an x tile in LDS at a time."""
    return [(t0, min(TILE, s - t0)) for t0 in range(0, s, TILE)]