Skip to content
4 changes: 2 additions & 2 deletions vllm/_custom_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -3355,8 +3355,8 @@ def swap_blocks(
and their matching destination block number on the destination tensor.
block_mapping is expected to be a tensor of shape (num_blocks_to_copy, 2)
where each block_mapping[i] represents a single copy operation, copying
block #block_mapping[i][0] from the source tensor
to block #block_mapping[i][1] on the destination tensor.
block `block_mapping[i][0]` from the source tensor
to block `block_mapping[i][1]` on the destination tensor.
block_mapping should have dtype int64.

The source and the destination tensors can be either on cpu or gpu,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ def sent_on(c, which="ar"):

def sent_ix(src, col, t, s, width):
"""(OPT_SENT) Word index of (src rank, pair at ``col``, token t), blocked
[src][16-col group][token][8 pairs]; shared by push and reduce."""
`[src][16-col group][token][8 pairs]`; shared by push and reduce."""
return ((src * (width // 16) + col // 16) * s + t) * 8 + (col % 16) // 2


Expand Down
2 changes: 1 addition & 1 deletion vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,7 @@ def sent_bytes() -> int:

def rlogp_half_bytes() -> int:
"""(OPT_RLOG_PARTS) One epoch parity's RLOGP, tagged pairs
[src rank][part][token][router row], SENT_ROWS rows."""
`[src rank][part][token][router row]`, SENT_ROWS rows."""
return TP_ * k2b.R_PARTS * SENT_ROWS * (k2b.E // TP_) * 8


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -531,7 +531,7 @@ def _route_out(c, t, red, sel, rank, me, ms):
@traced
def route_parts(c, t):
"""(OPT_RLOG_PARTS) Token t's router partials (every rank's 4 K
parts, RLOGP) -> LDS ``vals`` [part][expert], wave q < R_PARTS polling part
parts, RLOGP) -> LDS `vals[part][expert]`, wave q < R_PARTS polling part
q (one batch: the 14 experts of its lane); ``stage_route`` sums each
expert's parts in part order, as the router's part 0 did."""
s, lane, wave = c["S"], c["lane"], c["wave"]
Expand Down
5 changes: 3 additions & 2 deletions vllm/models/minimax_m3/amd/mono/layout.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,8 +110,9 @@ def _layout(items, align=256):


def sym_layout(npes: int):
"""Per-rank symmetric buffer: attn / ffn partial regions, [src rank][token][HIDDEN],
and the peers' index-selection candidates (``sel_cp``)."""
"""Per-rank symmetric buffer: attn / ffn partial regions,
`[src rank][token][HIDDEN]`, and the peers' index-selection candidates
(``sel_cp``)."""
rows = npes * MAX_TOKENS * HIDDEN // 2
# a_ag / ffn_ag: the finished attention / FFN all-reduce rows, each written by
# its owner rank (the reduce-scatter + all-gather form, S >= AG_RS_FROM)
Expand Down
8 changes: 4 additions & 4 deletions vllm/models/muse_spark/amd/mono/kernels/attn_frag.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,8 @@

Shapes (a TP8 rank): 16 q heads, 2 kv heads (QPK = 8), head dim 64, 8 query tokens a request,
NR requests (K3 S8: 1, S16: 2; HA S32: 4). K / V bf16 in the layer's vLLM LBHNC cache
[block][2][16][K 64 | V 64] (block 16); the page table is the layer group's block_table
(int32 [NR][stride]); seq_used = vLLM's seq_lens (int32 [NR], keys incl. the 8 new ones).
`[block][2][16][K 64 | V 64]` (block 16); the page table is the layer group's block_table
(int32 `[NR][stride]`); seq_used = vLLM's seq_lens (int32 [NR], keys incl. the 8 new ones).
q [8 NR, 16, 64] bf16 (row stride 1024, without the QK scale); out attn_o [8 NR, 1024] bf16.
SWA: vLLM's window (2047, 0), keys qpos - key < 2048, and only keys in the window are read
(vLLM frees the blocks before it: their block-table entries are stale).
Expand Down Expand Up @@ -633,8 +633,8 @@ def vllm_scale_word(scale):

def vllm_desc(q, kv_cache, block_table, seq_lens, out, part, nreq, scale):
"""The 10 descriptor words of one layer: q [8 nreq.., 16 x 64] bf16 rows,
kv_cache the layer's [block][2][16][128] bf16 view (block stride = kv_cache.stride(0)), block_table int32
[>= nreq][stride] of the layer's group, seq_lens int32 [>= nreq], out attn_o [rows, 1024] bf16, part the scratch
kv_cache the layer's `[block][2][16][128]` bf16 view (block stride = kv_cache.stride(0)), block_table int32
`[>= nreq][stride]` of the layer's group, seq_lens int32 [>= nreq], out attn_o [rows, 1024] bf16, part the scratch
(int32, part_bytes()), nreq the requests of the step (graph padding included: seq_len 0 = no keys, zero rows)."""
import torch

Expand Down