[Perf] [MI355X] glm53: Score a request's rows together and select the exact top-2048 on up to 64 slices

sgl_amd_mono/glm53/gfx950/layout.py

@@ -37,23 +37,25 @@
 INDEX_Q_ROWS = INDEX_HEADS * INDEX_DIM
 INDEX_TILE = 16
 INDEX_KEYS_PER_TASK = 64
-# a long row's paged top-k runs on up to this many slices (CTAs) per row, each with
-# staging room for this many threshold-bin candidates
-INDEX_SEL_SLICES = 16
+# a row's top-k runs on up to this many slices (CTAs), with staging room for this many
+# candidates of the threshold's prefix
+INDEX_SEL_SLICES = 64
 INDEX_CAND_STAGE_CAP = 16384
-# keys a select chunk (a CTA's 512 threads x 8), and the chunks past which a row selects
-# on up to INDEX_SEL_SLICES CTAs: up to 8 chunks the hand-offs cost what the slices save
-INDEX_SEL_CHUNK = 4096
-INDEX_SEL_PAR_CHUNKS = 8
-
-
-def select_slices(bound):
-    """(slices, chunks a slice) of a row of ``bound`` > 0 keys whose top-k runs on
-    several CTAs: runs of whole chunks in position order, at most INDEX_SEL_SLICES, none
-    empty. The kernel calls it on traced Int32 and the CPU test on ints."""
-    chunks = (bound + INDEX_SEL_CHUNK - 1) // INDEX_SEL_CHUNK
-    per = (chunks + INDEX_SEL_SLICES - 1) // INDEX_SEL_SLICES
-    return (chunks + per - 1) // per, per
+# radix digits of the sliced top-k; a row has a slice per this many keys (rounded up),
+# at most INDEX_SEL_SLICES
+INDEX_SEL_DIGITS = 6
+INDEX_SEL_SLICE_KEYS = 1024
+
+
+def select_slices(bound, minimum=min, maximum=max):
+    """(slices, keys a slice) of a row of ``bound`` keys: runs of the row's keys in
+    position order, none empty (one for an empty row). The kernel calls it on traced
+    Int32 with ``fx.min`` and ``fx.max``, the CPU test on ints."""
+    n = minimum(
+        (bound + INDEX_SEL_SLICE_KEYS - 1) // INDEX_SEL_SLICE_KEYS, INDEX_SEL_SLICES
+    )
+    n = maximum(n, 1)
+    return n, (bound + n - 1) // n
 
 
 DCP_SUMMARY_PAIRS = KV_LORA // 2 + 2
@@ -213,22 +215,20 @@ def layout(
     if with_indexer:
         items += [
             ("index_k", samples * INDEX_DIM * pair_bytes),
-            ("index_k_new", samples * INDEX_DIM // 2 * pair_bytes),
-            ("index_ready", samples * pair_bytes),
+            # each row's new key as cached: its FP8 value words, then its scale
+            ("index_k8", samples * (INDEX_DIM // 4 + 1) * pair_bytes),
             ("index_w", samples * INDEX_HEADS * pair_bytes),
             ("index_q", samples * INDEX_Q_ROWS // 2 * pair_bytes),
+            ("index_qh", samples * INDEX_Q_ROWS // 2 * pair_bytes),
             ("index_scores", samples * index_max_seq * pair_bytes),
             ("indices", samples * sparse_attention_topk * pair_bytes),
-            # parallel top-k: per slice a 4096-bin histogram and its ready pair, per
-            # row the threshold (bin prefix, keys still needed, keys in the bin), per
-            # slice its winners' positions, its (key, position) candidates and their
-            # counts
-            ("index_part_hist", samples * INDEX_SEL_SLICES * 4096 * 4),
-            ("index_part_ready", samples * INDEX_SEL_SLICES * pair_bytes),
-            ("index_thresh", samples * 4 * pair_bytes),
-            ("index_win_stage", samples * INDEX_SEL_SLICES * sparse_attention_topk * 4),
-            ("index_cand_stage", samples * INDEX_SEL_SLICES * INDEX_CAND_STAGE_CAP * 8),
-            ("index_gather_count", samples * INDEX_SEL_SLICES * pair_bytes),
+            # sliced top-k: per row, digit and slice the slice's 64 counts (two a
+            # pair); per row the candidates (key, cache slot) under the threshold's prefix
+            (
+                "sel_hist",
+                samples * INDEX_SEL_DIGITS * INDEX_SEL_SLICES * 32 * pair_bytes,
+            ),
+            ("sel_cand", samples * INDEX_CAND_STAGE_CAP * 2 * pair_bytes),
         ]
 
     offset, scratch = 0, {}
@@ -287,22 +287,24 @@ def stage_tasks(
         ),
         ("q_norm", samples),
         ("cache", samples),
+    ]
+    if with_indexer:
+        tasks += [("index_key", samples)]
+    tasks += [
         ("q_b", heads * (NOPE_DIM + PE_DIM) // Q_B_TILE),
     ]
     if with_indexer:
-        tasks += [("index_q", INDEX_Q_ROWS // INDEX_TILE)]
+        # index_qh: each row's query, rotated and Hadamard-transformed, 8 heads a task
+        tasks += [
+            ("index_q", INDEX_Q_ROWS // INDEX_TILE),
+            ("index_qh", samples * INDEX_HEADS // 8),
+        ]
     tasks += [("uk", heads * KV_LORA // UK_TILE)]
     if with_indexer:
         tasks += [
-            (
-                "index_score",
-                samples
-                * ((index_max_seq + INDEX_KEYS_PER_TASK - 1) // INDEX_KEYS_PER_TASK),
-            ),
-            ("index_hist", samples * INDEX_SEL_SLICES),
-            ("index_thresh", samples),
-            ("index_gather", samples * INDEX_SEL_SLICES),
-            ("index_select", samples),
+            # one task a CTA: its runs of the requests' keys
+            ("index_score", BLOCKS),
+            ("index_slice", samples * INDEX_SEL_SLICES),
         ]
     tasks += [
         (
@@ -385,9 +387,12 @@ def stage_bases(
         # tile for the first half, instead of serializing two index-Q tiles on
         # every q_b CTA.
         base["index_q"] = (base["q_b"] + tasks["q_b"]) % BLOCKS
-        base["index_score"] = 101
-        base["index_select"] = 100
-        base["index_thresh"] = 100
-        base["index_hist"] = BLOCKS // 2
-        base["index_gather"] = BLOCKS // 2
+        # on the first CTAs of index_q's own tiles, after their tile
+        base["index_qh"] = base["index_q"]
+        # each row's index key on its q_norm CTA, after the CTA's index-query tile
+        base["index_key"] = base["q_norm"]
+        # one scoring task a CTA, numbered by its CTA
+        base["index_score"] = 0
+        # the selection's slices past uk's CTAs
+        base["index_slice"] = (base["uk"] + tasks["uk"]) % BLOCKS
     return base

test/e2e/test_glm53_server.py

@@ -157,11 +157,12 @@ def test_each_covered_verify_layer_within_its_numerics_with_fp8_kv(self):
         )
 
     def test_each_covered_layer_past_the_parallel_top_k(self):
-        # rows of about 108k keys (27 chunks: 14 slices, 13 of two chunks) at widths 2
-        # and 4, and of about 36k keys (9 one-chunk slices) at width 8. Past 98304
-        # keys each wave's score walk refills its buffers and looks blocks ahead.
+        # rows of about 108k keys at widths 2 and 4 and of about 36k keys at width 8:
+        # past 98304 keys a row's top-k runs on every slice
         past = 98304
-        self.assertGreaterEqual(past, layout.INDEX_SEL_SLICES * layout.INDEX_SEL_CHUNK)
+        self.assertGreaterEqual(
+            past, layout.INDEX_SEL_SLICES * layout.INDEX_SEL_SLICE_KEYS
+        )
         self._check_layers(
             [],
             ((1, 9000), (3, 9000), (8, 3000)),

test/kernels/glm53/check_layer.py

@@ -274,8 +274,9 @@ def make_step(positions: list, requests: list, pages: Pages, device):
         positions=torch.tensor([p or 0 for p in positions], dtype=i64, device=device),
         slots=torch.tensor(slots, dtype=i64, device=device),
         indptr=torch.tensor(indptr, dtype=i32, device=device),
-        req_ids=torch.tensor(requests, dtype=i32, device=device),
-        # SGLang's DSA backend gives each row its request's page-table row
+        # SGLang's DSA backend gives each row its request's page-table row, and the
+        # spec each row that row
+        req_ids=torch.arange(W, dtype=i32, device=device),
         page_table=pages.table[torch.tensor(requests, device=device)].contiguous(),
         topk=topk,
     )