Skip to content

vllm.models.deepseek_v4.amd.mono.attention.plan

Geometry of the attention stages (front, back): the KV record, the split attention's chunks and the scratch (mailbox) layout; the model's widths and a TP rank's share of them are in ..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

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_v4/amd/mono/attention/plan.py 
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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(d.keys, 64) <= 256 else 128
    )  # back.chunk_keys
    nchunk = cdiv(d.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_v4/amd/mono/attention/plan.py 
95
96
97
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_v4/amd/mono/attention/plan.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
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 * X_GROUPS * 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),
        ]
    )