[Perf] [MI355X] glm53_flash: Score a verify request's index tiles once for its rows

sgl_amd_mono/glm53_flash/gfx950/kernel.py

@@ -3009,7 +3009,9 @@ def off_of(p):
                 ksc = ld_f32(r_blk, IPB * INDEX_DIM // 4 + off_of(pw))
                 return pw, kw, ksc
 
-            def score_half(s, n_pools, loaded):
+            def score_half(s, n_pools, loaded, q_off, w_off):
+                """Row ``s``'s scores of a loaded tile half: its FP8 query at LDS word
+                ``q_off``, its head weights at ``w_off``."""
                 pw, kw, ksc = loaded
                 a = fx.Vector.from_elements(
                     [kw[h][e] for h in range(2) for e in range(4)], fx.Int32
@@ -3020,7 +3022,8 @@ def score_half(s, n_pools, loaded):
                     qb = [
                         fx.Vector(
                             fx.ptr_load(
-                                xs + (hq * QW8 + 8 * g4 + 4 * h), result_type=v4f
+                                xs + (q_off + hq * QW8 + 8 * g4 + 4 * h),
+                                result_type=v4f,
                             )
                         ).bitcast(fx.Int32)
                         for h in range(2)
@@ -3045,7 +3048,7 @@ def score_half(s, n_pools, loaded):
                         )
                     )
                     # C[pool][head] in lane head + 16 (pool // 4), elem pool % 4
-                    w_h = lds_ld(misc, hq)
+                    w_h = lds_ld(xs, w_off + hq)
                     for e in range_constexpr(4):
                         parts[e] = parts[e] + fx.max(c[e], fx.Float32(0.0)) * w_h
                 for off in (1, 2, 4, 8):
@@ -3067,54 +3070,89 @@ def score_half(s, n_pools, loaded):
                     cache_modifier=CM_DEV,
                 )
 
-            for s in range_constexpr(S):
-                pos = row_word(s, ROW_POS)
-                bound, scored, _ = row_keys(pos)
-                n_pools = scored.select(bound, fx.Int32(0))
+            # a task a (unit, tile): a unit is a verify request's rows, which share its
+            # page table (SGLang repeats the request's row) and so each tile's keys,
+            # or one other row. A tile's keys load once for the unit's rows; each row
+            # scores them with its own query and weights, masked to its own pools
+            # (pools past a row's are a later row's, loaded in place).
+            if const_expr(VERIFY):
+                UNITS = [list(range(r * VL, (r + 1) * VL)) for r in range(NREQ)]
+                UNITS += [[s] for s in range(NREQ * VL, S)]
+            else:
+                UNITS = [[s] for s in range(S)]
+            NQW = INDEX_HEADS * QW8 // THREADS
+            for u in range_constexpr(len(UNITS)):
+                urows = UNITS[u]
+                UL = len(urows)
+                row_pools = []
+                for s in urows:
+                    bound, scored, _ = row_keys(row_word(s, ROW_POS))
+                    row_pools.append(scored.select(bound, fx.Int32(0)))
+                n_pools = row_pools[0]
+                for v in row_pools[1:]:
+                    n_pools = fx.max(n_pools, v)
                 n_tiles = (n_pools + POOLS_PER_TASK - 1) // POOLS_PER_TASK
-                first = start_at(base["iscore"] + s * (G // S))
+                first = start_at(base["iscore"] + u * (G // len(UNITS)))
                 if first < n_tiles:
-                    bt = row_word(s, ROW_BT)
+                    bt = row_word(urows[0], ROW_BT)
                     r_ibt = _rsrc(
                         page_table + fx.Int64(bt) * fx.Int64(pt_stride) * fx.Int64(4)
                     )
                     # the first tile's block-table entries load while one lane polls
-                    # the cache's flag (the pool this launch completed), then its keys
-                    # load while the query is awaited
+                    # the cache's flags (the pools this launch completed), then its
+                    # keys load while the queries are awaited
                     blocks = [
                         half_block(r_ibt, n_pools, first, h) for h in range(HALVES)
                     ]
                     if tid == 0:
-                        for r in chain_rows(s):
+                        for r in chain_rows(urows[-1]):
                             get(mb("kv_ready"), r)
                     fx.memory_fence(
                         ordering=fx.AtomicOrdering.Acquire, syncscope="agent"
                     )
                     gpu.barrier()
                     pre = [load_half(n_pools, b) for b in blocks]
-                    # the row's FP8 query -> LDS words h * 32 .., weights -> misc
-                    NQW = INDEX_HEADS * QW8 // THREADS
-                    qv = poll(
-                        [
-                            (mb("iq8"), s * INDEX_HEADS * QW8 + tid + i * THREADS, 1)
-                            for i in range(NQW)
-                        ]
+                    # the rows' FP8 queries -> LDS words i * Q_ROW + h * 32 .., their
+                    # weights -> words W_OFF + i * INDEX_HEADS + h, in one poll
+                    Q_ROW = INDEX_HEADS * QW8
+                    W_OFF = UL * Q_ROW
+                    assert W_OFF + UL * INDEX_HEADS <= MISC_OFF
+                    q0 = urows[0] * Q_ROW
+                    specs = [
+                        (mb("iq8"), q0 + tid + i * THREADS, 1) for i in range(UL * NQW)
+                    ]
+                    specs.append(
+                        (
+                            mb("iqw"),
+                            urows[0] * INDEX_HEADS
+                            + fx.min(tid, fx.Int32(UL * INDEX_HEADS - 1)),
+                            1,
+                        )
                     )
-                    for i in range_constexpr(NQW):
+                    qv = poll(specs, batch=len(specs))
+                    for i in range_constexpr(UL * NQW):
                         lds_st(xs, tid + i * THREADS, qv[i][0].bitcast(fx.Float32))
-                    if tid < INDEX_HEADS:
-                        hw = getf_many([(mb("iqw"), s * INDEX_HEADS + tid)])[0]
-                        lds_st(misc, tid, hw)
+                    if tid < UL * INDEX_HEADS:
+                        lds_st(xs, W_OFF + tid, qv[-1][0].bitcast(fx.Float32))
                     gpu.barrier()
+
+                    def score_tile(loaded):
+                        for i in range_constexpr(UL):
+                            score_half(
+                                urows[i],
+                                row_pools[i],
+                                loaded,
+                                i * Q_ROW,
+                                W_OFF + i * INDEX_HEADS,
+                            )
+
                     for h in range_constexpr(HALVES):
-                        score_half(s, n_pools, pre[h])
+                        score_tile(pre[h])
                     for tile in range(first + G, n_tiles, G):
                         tile = fx.Int32(tile)
                         for h in range_constexpr(HALVES):
-                            score_half(
-                                s,
-                                n_pools,
-                                load_half(n_pools, half_block(r_ibt, n_pools, tile, h)),
+                            score_tile(
+                                load_half(n_pools, half_block(r_ibt, n_pools, tile, h))
                             )
                     gpu.barrier()
 

sgl_amd_mono/glm53_flash/spec.py

@@ -396,6 +396,17 @@ def normed_rows(lm, owed: _Owed, S: int) -> torch.Tensor:
     return lm.model.norm(check.contract(written))
 
 
+def check_verify_page_table(page_table: torch.Tensor, verify: int) -> None:
+    """Refuse a verify step whose rows of one request differ in ``page_table``: the
+    sparse-attention launch scores a request's keys once, from its first row."""
+    rows = page_table.reshape(-1, verify, page_table.shape[-1])
+    if not torch.equal(rows, rows[:, :1].expand_as(rows)):
+        raise RuntimeError(
+            "GLM-5.3-Flash monokernel: the DSA backend's page table no longer repeats "
+            "a verify request's row for each of its rows"
+        )
+
+
 def _written(module, x, residual, post, comb) -> torch.Tensor:
     """The streams with the owed FFN post applied, as SGLang's FFN exit writes them."""
     s = x.shape[0]
@@ -594,7 +605,10 @@ def _layer(self, L: _Layer, plan, input_ids, positions, owed, forward_batch):
                 pos, rpi = buf[0, :S], buf[1, :S]
             idx = slots if slots.dtype == torch.int32 else slots.int()
             page_table = backend.full_attn_backend.forward_metadata.real_page_table
-            # the DSA backend repeats a verify request's page-table row for each row
+            # the DSA backend repeats a verify request's page-table row for each row;
+            # an eager step checks it (a graph replays without the host)
+            if V > 1 and not torch.cuda.is_current_stream_capturing():
+                check_verify_page_table(page_table[:S], V)
             plan.state["dsa_rows"] = (self.rows[W], page_table[:S])
             step = {
                 "ids": input_ids[:S],

test/e2e/test_glm53_flash_server.py

@@ -2,8 +2,9 @@
 switch points: decode steps of 1 to 16 rows run the kernels, wider ones the normal
 path, and requests join and leave a running batch that steps past 16 rows and back
 (a short GSM8K run, for gross failures); with MTP, verify steps of 1 to 4 requests
-(4 to 16 rows) on the kernels (GSM8K and the draft's accept length); the tokens of covered steps scored again by
-the normal prefill path (KL consistency); each layer's launch against SGLang's own
+(4 to 16 rows) on the kernels (GSM8K and the draft's accept length), and eager
+verify steps that check each request's page-table row is repeated; the tokens of
+covered steps scored again by the normal prefill path (KL consistency); each layer's launch against SGLang's own
 layer, and the tail's logits against SGLang's, under SGLANG_AMD_MONO_CHECK; the
 hidden states a covered step returns; and a fault on one rank that stops the server
 before any token of the failed step is sent."""
@@ -129,6 +130,31 @@ def test_gsm8k_and_accept_length_on_the_kernels(self):
         self.assertGreater(accept, MTP_ACCEPT_BOUND)
 
 
+class TestMonokernelMtpEager(CustomTestCase):
+    def test_verify_steps_find_their_request_s_page_table_row_repeated(self):
+        # the sparse-attention launch reads a verify request's keys from its first
+        # page-table row; an eager verify step checks that the DSA backend repeats
+        # the row for each of the request's rows (a graph replays unchecked). 4
+        # requests of 2 to 11 pages, so the rows differ between requests
+        para = "The river valley archive lists the flood height of every year. "
+        process = _launch(extra_args=[*MTP, "--disable-cuda-graph"])
+        try:
+            r = requests.post(
+                DEFAULT_URL_FOR_TEST + "/generate",
+                json={
+                    "text": [f"Record {i}. " + para * (40 + 60 * i) for i in range(4)],
+                    "sampling_params": {"max_new_tokens": 32, "temperature": 0},
+                },
+                timeout=600,
+            )
+            r.raise_for_status()
+        finally:
+            terminate_and_kill_process_tree(process)
+        verify = [out["meta_info"]["spec_verify_ct"] for out in r.json()]
+        print(f"{verify=}")
+        self.assertTrue(len(verify) == 4 and min(verify) > 0, verify)
+
+
 class TestMonokernelCheck(CustomTestCase):
     def test_each_covered_layer_within_its_numerics(self):
         # SGLANG_AMD_MONO_CHECK (eager) runs the normal layer next to each layer's

test/kernels/glm53_flash/time_layers.py

@@ -6,7 +6,9 @@
 One HIP graph holds a step of launches, inputs chained layer to layer, every row
 live. Each layer has its own random weights (``check_layer.py``'s, seeded by layer,
 so they stream from HBM as in a served step), its own state or caches, and every
-sparse-attention row sits at position ``--pos``. ``--kind model`` is the served
+sparse-attention row sits at position ``--pos``; with ``--verify L`` (``dsa``) the
+launch is an MTP verify step, rows // L requests whose L rows sit at ``--pos`` ..
+``--pos + L - 1`` on the request's page table. ``--kind model`` is the served
 step: layers 3 to 43 in model order (30 KDA, 11 sparse attention) on one runtime a
 kind; ``--kind full`` the whole served step as the spec builds it: layer 0 from the
 step's tokens, the dense layers 1 and 2, layers 3 to 44, then the tail. ``kda``,
@@ -49,6 +51,45 @@ def _kinds(args) -> list:
     return [args.kind] * args.layers
 
 
+def _verify_caches(rows: int, L: int, pos: int, device):
+    """``check_layer``'s verify inputs: rows // L requests of pos + 1 tokens, request
+    r's row j at pos + j, SGLang's verify page table (the request's row for each of
+    its rows), tail rings of 4 + L; rows past the requests are pads."""
+    from msgspec.structs import replace
+
+    from sgl_amd_mono.glm53_flash.gfx950.config import INDEX_DIM, KPOOL
+    from sgl_amd_mono.glm53_flash.gfx950.layout import (
+        ROW_BT,
+        ROW_POOL_SLOT,
+        ROW_POS,
+        ROW_SLOT,
+        ROW_TAIL_SLOT,
+        ROW_WORDS,
+    )
+
+    R = rows // L
+    caches, _, pt = cl.dsa_caches(R, device, [pos + 1] * R)
+    g = cl._gen(11, device)
+    tails = [
+        torch.randn(R + 1, 4 + L, INDEX_DIM, generator=g, device=device).bfloat16()
+        for _ in range(2)
+    ]
+    caches = replace(caches, tail_k=tails[0], tail_g=tails[1])
+    table = torch.full((rows, ROW_WORDS), -1, dtype=torch.int32, device=device)
+    table[:, ROW_BT] = 0
+    for r in range(R):
+        for j in range(L):
+            s = r * L + j
+            slot = cl.kv_slot(pt, r, pos + j)
+            table[s, [ROW_POS, ROW_SLOT, ROW_POOL_SLOT, ROW_TAIL_SLOT]] = torch.tensor(
+                [pos + j, slot, slot // KPOOL, (r + 1) * KPOOL],
+                dtype=torch.int32,
+                device=device,
+            )
+            table[s, ROW_BT] = s
+    return caches, table, pt.repeat_interleave(L, 0).contiguous()
+
+
 def _build(args, kinds: list, rows: int, rank: int, tp: int, group, device):
     """The runtimes (one a kind), each layer's op and launch arguments."""
     from sgl_amd_mono.glm53_flash.gfx950.config import GLOBAL_INTER
@@ -82,6 +123,7 @@ def _build(args, kinds: list, rows: int, rank: int, tp: int, group, device):
             grouped=not args.per_row_moe
             and rows >= GROUPED_FROM
             and not (dense or head),
+            verify=args.verify if dsa else 1,
         )
         runtimes[kind] = FlashRuntime(key, rank, group, device)
         runtimes[kind].set_deadline(60.0)
@@ -118,7 +160,10 @@ def _build(args, kinds: list, rows: int, rank: int, tp: int, group, device):
             attn = cl.dsa_weights(rank, tp, device)
             if shuffled:
                 attn = cl.shuffled_copy(attn)
-            caches, table, pt = cl.dsa_caches(rows, device, [args.pos + 1] * rows)
+            if args.verify > 1:
+                caches, table, pt = _verify_caches(rows, args.verify, args.pos, device)
+            else:
+                caches, table, pt = cl.dsa_caches(rows, device, [args.pos + 1] * rows)
             ops.append(FlashLayerOp(runtimes[kind], ffn, n, dsa=attn))
             extra.append(dict(caches=caches, rows=table, page_table=pt))
     cl._gen = _GEN
@@ -179,7 +224,8 @@ def step():
     dist.all_gather_object(medians, (median, clean), group=group)
     if rank == 0:
         print(
-            f"TIME kind={args.kind} rows={rows} tp={tp} launches={len(kinds)} "
+            f"TIME kind={args.kind} rows={rows} verify={args.verify} tp={tp} "
+            f"launches={len(kinds)} "
             f"pos={args.pos} fp8={'row-major' if args.row_major else 'in-place'} "
             f"us_per_launch median={median:.2f} min={min(per_launch):.2f} "
             f"max={max(per_launch):.2f} "
@@ -219,8 +265,12 @@ def main() -> int:
         action="store_true",
         help="the per-row MoE at every width, the baseline of the grouped one",
     )
+    ap.add_argument(
+        "--verify", type=int, default=1, help="dsa: an MTP verify step's rows a request"
+    )
     args = ap.parse_args()
-    cl.MAX_LEN = max(cl.MAX_LEN, (args.pos + 256) // 256 * 256)
+    assert args.verify == 1 or args.kind == "dsa", "--verify times dsa launches"
+    cl.MAX_LEN = max(cl.MAX_LEN, (args.pos + args.verify - 1 + 256) // 256 * 256)
     rank, tp = int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"])
     torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
     device = torch.device("cuda", torch.cuda.current_device())

test/unit/test_glm53_flash_spec.py

@@ -1,8 +1,9 @@
 """The GLM-5.3-Flash gfx950 monokernel spec on CPU (RFC #43949): the layers its
 forward covers on the GLM-5.3-Flash layer structure, its row routing and TP2 row
-cap, the tail's logits in place of the logits processor's, its registry entry, its
-config and model refusal reasons, and the dense, embedding and LM-head views. SGLang
-checks its declared requirements, the same way for every spec."""
+cap, its verify page-table check, the tail's logits in place of the logits
+processor's, its registry entry, its config and model refusal reasons, and the
+dense, embedding and LM-head views. SGLang checks its declared requirements, the
+same way for every spec."""
 
 import math
 from types import SimpleNamespace
@@ -122,6 +123,16 @@ def test_a_verify_step_of_other_draft_chains_runs_the_model_s_forward():
     assert out == "the model's forward" and len(lm.ran) == 1
 
 
+def test_a_verify_step_refuses_a_page_table_that_does_not_repeat_its_rows():
+    # the DSA backend's verify page table: 3 requests, each row repeated for its 4 rows
+    page_table = torch.arange(3 * 5, dtype=torch.int32).reshape(3, 5)
+    page_table = page_table.repeat_interleave(4, dim=0)
+    flash.check_verify_page_table(page_table, 4)
+    page_table[5, 2] += 1
+    with pytest.raises(RuntimeError, match="no longer repeats"):
+        flash.check_verify_page_table(page_table, 4)
+
+
 def test_a_step_from_embeddings_runs_the_model_s_forward():
     # layer 0's launch reads the step's tokens; given embeddings, SGLang's path runs
     lm, built = _built()