[Perf] [MI355X] glm53: Shared layer and MoE hand-offs, NextN entry, K-split front, indexer launch, MoE expert dedup and weight loads, kernarg reload, LM-head loads

README.md

@@ -17,7 +17,7 @@ The library implements SGLang's monokernel interface ([RFC #43949](https://githu
 | DeepSeek-V4.1-Flash | MI350X, MI355X (gfx950) | 2, 4 | All 40 layers in decode and DSpark verify steps of up to 48 rows, and the DSpark draft stages |
 | DeepSeek-V4-Pro | MI355X (gfx950) | 8 | All 61 layers (one launch each) in decode steps of up to 8 rows, and DSpark verify steps of one request (5 to 8 tokens) |
 | GLM-5.3-Flash | MI355X (gfx950) | 2, 4, 8 | All 45 layers (one launch each; layer 0 from the tokens, with the embedding) and the final norm, LM head and logits (one launch), in decode steps of up to 8 rows (4 at TP2) |
-| GLM-5.3 | MI355X (gfx950) | 4, 8 | All 75 MoE layers (one launch each), in decode and MTP verify steps of up to 8 rows (4 at TP4) |
+| GLM-5.3 | MI355X (gfx950) | 4, 8 | All 75 MoE layers (one launch each), in decode and MTP verify steps of up to 8 rows |
 | MiniMax-M3 (MXFP4) | MI350X, MI355X (gfx950) | 4 | The 57 sparse MoE layers in decode and EAGLE3 verify steps of up to 16 rows |
 | MiniMax-M3.1-Flash (NVFP4 experts) | MI350X (gfx950) | 4 | o_proj and the FFN of the 57 sparse layers, after SGLang's attention, in decode and DSpark verify steps of up to 16 rows |
 | Qwen3.8-Flash-Next (MXFP4) | MI355X (gfx950) | 1 | The whole model (48 layers in 13 launches) in NEXTN verify steps of up to 8 requests x 4 tokens; the MTP draft's draft extend (up to 8 x 4 rows) and draft decode steps (up to 8 x 1) are routed to its launch |

sgl_amd_mono/glm53/draft.py

@@ -544,16 +544,10 @@ def _build(model, hf_config, ctx):
         )
     max_seq = -(-context_len // INDEX_KEYS_PER_TASK) * INDEX_KEYS_PER_TASK
     kv_fp8 = get_model().kv_cache_dtype == "fp8_e4m3"
-    cap = target.TP_MAX_ROWS.get(tp, max(target.WIDTHS))
-    extend_widths = sorted(w for w in ctx.widths if w <= cap)
+    extend_widths = sorted(ctx.widths)
     # a draft decode step of 5 to 8 rows runs SGLang's path: the NextN shared launch
     # at 8 rows faulted in a kernel check (a process of several launches, not alone)
-    decode_widths = [w for w in target.WIDTHS if w <= min(cap, 4)]
-    if not extend_widths:
-        raise ValueError(
-            f"GLM-5.3 MTP draft: TP{tp} runs at most {cap} rows, and this server's "
-            f"draft extend steps have {sorted(ctx.widths)}"
-        )
+    decode_widths = [w for w in target.WIDTHS if w <= 4]
     layer = target._Layer(layer_id, decoder, weights)
     shared_w = msgspec.structs.replace(weights, index=None)
     widths = sorted(set(extend_widths) | set(decode_widths))

sgl_amd_mono/glm53/gfx950/layout.py

@@ -37,6 +37,23 @@
 # the NextN entry's pointer-table entries: this rank's eh_proj rows, the enorm and
 # hnorm weights, and (from the tokens) this rank's embedding rows
 EH_ARGS = (8, 9, 10, 11)
+
+
+def eh_parts(eh_tiles: int) -> int:
+    """The NextN entry's K parts a tile: four (two of each of the e and h_prev halves)
+    when their tasks fit one round of CTAs, else the two halves."""
+    return 4 if 4 * eh_tiles <= BLOCKS else 2
+
+
+def eh_base(eh_tiles: int, with_indexer: bool) -> int:
+    """The NextN entry's first CTA. With the indexer, the CTA past the front's index
+    projection tiles, which are the grid's last CTAs, so CTA 0: those tiles' weights go
+    out at the launch start, not behind an entry task. Else the last CTAs."""
+    if with_indexer:
+        return 0
+    return max(BLOCKS - eh_parts(eh_tiles) * eh_tiles, 0)
+
+
 INDEX_Q_ROWS = INDEX_HEADS * INDEX_DIM
 INDEX_TILE = 16
 INDEX_KEYS_PER_TASK = 64
@@ -85,6 +102,10 @@ def fp8_pe_upper_pair_lane(lane):
 
 
 N_QKV_A = QKV_A_ROWS // QKV_A_TILE
+# the front splits each q_a / kv_a tile's K: the first FRONT_K_BIG on the tile's own CTA,
+# the rest on one of the FRONT_SMALL_CTAS CTAs past them (two tiles' rests a CTA)
+FRONT_K_BIG = 4096
+FRONT_SMALL_CTAS = BLOCKS - N_QKV_A
 N_ROW_TILES = HIDDEN // ROW_TILE
 N_ROUTER = N_EXPERTS // ROUTER_TILE
 XQ_BLOCKS = HIDDEN // 128
@@ -177,8 +198,8 @@ def layout(
     ``tokens``: also the layer input's rows, which the token's owner rank pushes
     (``xin``, two slots by launch parity). ``eh_tiles``: the NextN entry's rows of
     ``xin`` instead, which every rank pushes its ``eh_tiles`` 16-column tiles of, and
-    the h_prev halves' partial sums (``ehp``); with ``tokens`` the embedding rows go
-    to ``ein`` (as ``xin``). ``head_vocab``: the tail launch's
+    the partial sums of every K part but a tile's first (``ehp``); with ``tokens``
+    the embedding rows go to ``ein`` (as ``xin``). ``head_vocab``: the tail launch's
     buffers instead, the gathered bf16 logits of a ``head_vocab`` vocabulary (two
     slots)."""
     from sgl_amd_mono.common.guard import FLAG_BYTES
@@ -196,6 +217,8 @@ def layout(
     pair_bytes = 8
     items = [
         ("q_a", samples * Q_LORA * pair_bytes),
+        ("q_a2", samples * Q_LORA * pair_bytes),
+        ("kv_a2", samples * (KV_LORA + PE_DIM) * pair_bytes),
         ("q_an", samples * Q_LORA // 2 * pair_bytes),
         ("kv_a", samples * (KV_LORA + PE_DIM) * pair_bytes),
         ("kvnew", samples * KV_LORA * pair_bytes),
@@ -217,6 +240,9 @@ def layout(
         ("sel", samples * MOE_SLOTS * pair_bytes),
         ("prob", samples * MOE_SLOTS * pair_bytes),
         ("mid", samples * MOE_SLOTS * inter * pair_bytes),
+        # mid's FP8 blocks (four values a pair) and their scales
+        ("midq", samples * MOE_SLOTS * inter // 4 * pair_bytes),
+        ("midqs", samples * MOE_SLOTS * inter // 128 * pair_bytes),
         ("ugp", BLOCKS * samples * 2 * UG_TILE * pair_bytes),
         ("xqd", samples * HIDDEN * 4),
     ]
@@ -239,7 +265,12 @@ def layout(
             ("sel_cand", samples * INDEX_CAND_STAGE_CAP * 2 * pair_bytes),
         ]
     if eh_tiles:
-        items += [("ehp", eh_tiles * samples * HEAD_TILE * pair_bytes)]
+        items += [
+            (
+                "ehp",
+                (eh_parts(eh_tiles) - 1) * eh_tiles * samples * HEAD_TILE * pair_bytes,
+            )
+        ]
 
     offset, scratch = 0, {}
     for name, size in items:
@@ -288,19 +319,17 @@ def stage_tasks(
 ):
     """Return ``(stage name, task count)`` pairs in execution order; with ``tokens``
     led by the embedding rows' push, with ``eh_tiles`` by the NextN entry (each tile's
-    e and h_prev halves). With ``head_rows`` (this rank's LM-head rows) the tail's
+    K parts, ``eh_parts``). With ``head_rows`` (this rank's LM-head rows) the tail's
     instead."""
 
     if head_rows:
         return [("lmhead", head_rows // HEAD_TILE), ("gather", BLOCKS)]
     tasks = [("embed", N_ROW_TILES)] if tokens else []
-    tasks += [("eh", 2 * eh_tiles)] if eh_tiles else []
+    tasks += [("eh", eh_parts(eh_tiles) * eh_tiles)] if eh_tiles else []
     tasks += [
-        # with the indexer, the index K and weights projections' 16-row tiles follow
-        (
-            "qkv_a",
-            N_QKV_A + ((INDEX_DIM + INDEX_HEADS) // QKV_A_TILE if with_indexer else 0),
-        ),
+        # every CTA: a tile's K head, two tiles' K rests or, with the indexer, an index
+        # K or weights projection tile
+        ("qkv_a", BLOCKS),
         ("q_norm", samples),
         ("cache", samples),
     ]
@@ -342,6 +371,12 @@ def stage_tasks(
                     BLOCKS,
                     ((MOE_SLOTS * inter // UG_TILE + BLOCKS - 1) // BLOCKS) * BLOCKS,
                 )
+                if samples <= 4
+                # above 4 rows: rounds of 8-intermediate tiles of the shared expert
+                # and of up to every row's routed experts
+                else ((samples * TOP_K + 1) * (inter // 8) + BLOCKS - 1)
+                // BLOCKS
+                * BLOCKS
             ),
         ),
         ("down", HIDDEN // dn_tile(samples, expert_mxfp4)),
@@ -373,30 +408,23 @@ def stage_bases(
         )
     )
     base, acc = {}, 0
-    # q_norm (with each row's cache task) on the first CTAs past qkv_a's, so it polls
-    # from the launch start; q_b after it. split from the CTA past q_b's, so its tiles
-    # land on CTAs freed by qkv_a, the last few on the q_norm CTAs; the cache keeps one
-    # CTA of the sequence after split, so uk (on the q_b CTAs it waits on anyway) and
-    # the later stages keep their CTAs.
-    for name in (
-        "qkv_a",
-        "q_norm",
-        "q_b",
-        "split",
-        "cache",
-        "uk",
-        "uv",
-        "o",
-        "router",
-        "ug",
-        "down",
-    ):
+    # the front takes every CTA (a tile's K head, or two tiles' K rests); q_norm (with
+    # each row's cache task) and q_b follow on the first CTAs, split past q_b's (its
+    # last few on the q_norm CTAs); the cache keeps one CTA of the sequence after
+    # split, so uk (on the q_b CTAs it waits on anyway) and uv keep their CTAs
+    for name in ("qkv_a", "q_norm", "q_b", "split", "cache", "uk", "uv"):
         base[name] = acc % BLOCKS
         acc += 1 if name == "cache" else tasks[name]
     base["cache"] = base["q_norm"]
     # o from split's first CTA: its tasks fill the split CTAs, free once the
     # attention ends
     base["o"] = base["split"]
+    # the MoE stages keep the CTAs of a front of whole tiles
+    acc = N_QKV_A + ((INDEX_DIM + INDEX_HEADS) // QKV_A_TILE if with_indexer else 0)
+    for name in ("q_norm", "q_b", "split", "uk", "uv", "o", "router", "ug", "down"):
+        if name in ("router", "ug", "down"):
+            base[name] = acc % BLOCKS
+        acc += tasks[name] + (1 if name == "split" else 0)
     if with_indexer:
         # q_b occupies exactly half the grid.  Put the second half of index-Q on
         # the complementary CTAs while q_b CTAs reuse their normalized q_lora

sgl_amd_mono/glm53/gfx950/tiles.py

@@ -26,6 +26,8 @@
 
 CM_DEV = 16
 CM_SYS = 17
+# nontemporal: lines streamed once a launch
+CM_NT = 2
 POLL_MAX = 12
 TL_COLS = 8
 

sgl_amd_mono/glm53/spec.py

@@ -62,9 +62,6 @@
 WIDTHS = (2, 4, 8)
 ROWS = tuple(range(1, max(WIDTHS) + 1))
 TP_SIZES = (4, 8)
-# the rows a step runs on the kernels at most at this TP size: on MI355X an 8-row
-# launch at TP4 is slower than the model's forward, so wider steps run that
-TP_MAX_ROWS = {4: 4}
 # SGLang's DSA cache on ROCm: 64-token pages
 PAGE_SIZE = 64
 TOPK = 2048
@@ -367,7 +364,6 @@ def __init__(
         heads: dict,
         cos,
         sin,
-        tp: int,
         device,
         context_len,
         epochs=None,
@@ -378,7 +374,6 @@ def __init__(
         self._layers, self.runtimes, self.heads = layers, runtimes, heads
         # per width, the guard blocks of its runtimes: a step advances them at once
         self.epochs = epochs
-        self.max_rows = TP_MAX_ROWS.get(tp, max(WIDTHS))
         self.max_position = context_len - 1
         self.check_layers = os.environ.get("SGLANG_AMD_MONO_CHECK") == "1"
         self.cos, self.sin = cos, sin
@@ -422,7 +417,7 @@ def forward(
     ):
         """``DeepseekV2ForCausalLM.forward`` for a routed decode or verify step."""
         lm = self.lm
-        if plan.key.rows > self.max_rows or input_embeds is not None:
+        if input_embeds is not None:
             return type(lm).forward(
                 lm, input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
             )
@@ -633,13 +628,7 @@ def _build(model, hf_config, ctx) -> _Built:
     # the keys an indexer row scores at most: the served context's, in whole tasks
     tasks = -(-context_len // INDEX_KEYS_PER_TASK)
     max_seq = tasks * INDEX_KEYS_PER_TASK
-    cap = TP_MAX_ROWS.get(tp, max(WIDTHS))
-    widths = sorted({_width(r) for r in ctx.widths if r <= cap})
-    if not widths:
-        raise ValueError(
-            f"GLM-5.3 monokernel: TP{tp} runs at most {cap} rows on the kernels, and "
-            f"this server's decode graphs route {sorted(ctx.widths)}"
-        )
+    widths = sorted({_width(r) for r in ctx.widths})
     runtimes, heads = {}, {}
     epochs = {w: Glm53Epochs(len(kinds) + 1, device) for w in widths}
 
@@ -704,7 +693,6 @@ def runtime(key):
         heads,
         cos,
         sin,
-        tp,
         device,
         context_len=context_len,
         epochs=epochs,

test/kernels/glm53/check_layer.py

@@ -21,12 +21,15 @@
 them; ``--check-embed``: layer 0 from the token ids equals the dense launch fed the
 gathered embedding rows (outputs, cache rows, top-k), and a changed token id differs;
 ``--check-head``: the tail's logits against a torch and an f32 tail, within the
-spec's bounds, and the final norm x1.01 fails them.
+spec's bounds, and the final norm x1.01 fails them. ``--similar X``: the live rows of a
+request are one random row plus X times their own, so they route alike, as a verify
+step's rows do (with ``--rowsets v4``).
 
-``--dump DIR`` writes, per run and rank, the SHA-256 of every launch's output rows,
-of the cache rows and top-k entries it writes, and after each width of every cache
-whole (``--perturb`` negates one input value, a negative control); ``--compare A B``
-matches two dump directories and exits 1 on any difference."""
+``--dump DIR`` writes, per run and rank, the SHA-256 of every launch's output rows
+(the tail's: its gathered logits), of the cache rows and top-k entries it writes,
+and after each width of every cache whole (``--perturb`` negates one input value, a
+negative control); ``--compare A B`` matches two dump directories and exits 1 on
+any difference."""
 
 import argparse
 import hashlib
@@ -318,6 +321,7 @@ def __init__(
         stop=False,
         head=False,
         first=None,
+        similar=0.0,
     ):
         from sgl_amd_mono.glm53.gfx950.runner import (
             Glm53Build,
@@ -380,7 +384,14 @@ def __init__(
             rt.set_deadline(args.deadline)
             self.head = Glm53HeadOp(rt, head_weights(rank, tp, device))
         g = _gen(3, device)
-        self.h = torch.randn(W, HIDDEN, generator=g, device=device).bfloat16()
+        h = torch.randn(W, HIDDEN, generator=g, device=device)
+        if similar:
+            base = h.clone()
+            noise = similar * torch.randn(W, HIDDEN, generator=g, device=device)
+            for s, (p, r) in enumerate(zip(positions, requests)):
+                if p is not None:
+                    h[s] = base[r] + noise[s]
+        self.h = h.bfloat16()
         if kinds[:1] == ["embed"]:
             self.h = token_ids(positions, device)
 
@@ -738,6 +749,10 @@ def after(i, h, step=step):
             out[f"{name}/s{step}/l{i}"] = d
 
         h = model.step(after=after)
+        if model.head is not None:
+            out[f"{name}/s{step}/tail"] = {
+                f"out.r{s}": _digest(h[s]) for s in range(h.shape[0])
+            }
         torch.cuda.synchronize()
     whole = {f"kv.l{i}": _digest(c.kv) for i, c in enumerate(model.caches)}
     whole.update(
@@ -800,6 +815,7 @@ def main() -> int:
     ap.add_argument("--check-head", action="store_true")
     ap.add_argument("--dump", default="", help="write the digests into DIR")
     ap.add_argument("--perturb", action="store_true", help="negate one input value")
+    ap.add_argument("--similar", type=float, default=0.0, help="rows routing alike")
     ap.add_argument("--compare", nargs=2, metavar=("A_DIR", "B_DIR"))
     args = ap.parse_args()
     if args.compare:
@@ -817,7 +833,16 @@ def main() -> int:
     ]:
         positions, requests = rows_of(W, rowset)
         model = Model(
-            kinds, positions, requests, rank, tp, group, device, args, head=head
+            kinds,
+            positions,
+            requests,
+            rank,
+            tp,
+            group,
+            device,
+            args,
+            head=head,
+            similar=args.similar,
         )
         name = f"{args.kv}-{args.kind}/w{W}{rowset}"
         if args.dump:

test/kernels/glm53/check_nextn.py

@@ -82,9 +82,10 @@ def _rel(a, b) -> float:
 
 class Draft:
     """A NextN layer of one kind at one width: its NextN runtime, the plain runtime
-    of its kind on the same layer weights and caches, and the tail."""
+    of its kind on the same layer weights and caches, and the tail. ``rows``:
+    (positions, requests) in place of ``rowset``'s."""
 
-    def __init__(self, kind, W, rowset, rank, tp, group, device, args):
+    def __init__(self, kind, W, rowset, rank, tp, group, device, args, rows=None):
         from sgl_amd_mono.glm53.gfx950.runner import (
             Glm53Build,
             Glm53HeadOp,
@@ -93,7 +94,7 @@ def __init__(self, kind, W, rowset, rank, tp, group, device, args):
         )
 
         index = kind == "index"
-        positions, requests = rows_of(W, rowset)
+        positions, requests = rows or rows_of(W, rowset)
         live = [p for p in positions if p is not None]
         self.live = len(live)
         pages = C.Pages(max(requests) + 1, max(live) + W + 1, device)
@@ -136,6 +137,11 @@ def advance(self):
         for rt in self.runtimes:
             rt.advance_step()
 
+    def step(self):
+        """One NextN launch from the step's rows, then the next epoch."""
+        self.run_nextn()
+        self.advance()
+
     def run_nextn(self):
         self.st.topk.copy_(self.topk0)
         x = self.nextn.launch(