[Perf] [MI355X] deepseek_v41: One launch per reuse layer, the index sources' compressor in K1, the Engram projection in the layer launches, cheaper hand-offs, the Markov block on the package's kernels

sgl_amd_mono/common/epoch.py

@@ -51,26 +51,30 @@ def epoch_begin(c, epoch, ep):
 
 
 @flyc.jit
-def epoch_end(c, epoch, ep, reset=None):
-    """CTA 0: once every CTA marked the epoch, run ``reset()`` and move the epoch on."""
-    if c["bid"] == 0:
+def epoch_end(c, epoch, ep, reset=None, who=0, marks=True):
+    """CTA ``who``: once every CTA marked the epoch, run ``reset()`` and move the epoch on.
+
+    ``marks=False``: the caller already saw a flag from every CTA of this launch,
+    stored after the CTA read the epoch, so no CTA can read the moved epoch."""
+    if c["bid"] == who:
         tid = c["tid"]
-        if tid < BLOCKS:
-            r = rsrc(epoch)
-
-            def load_mark():
-                return fx.Int32(
-                    bo.buffer_load(
-                        r,
-                        EPOCH_MARKS + tid,
-                        vec_width=1,
-                        dtype=T.i32,
-                        cache_modifier=CM_DEV,
+        if const_expr(marks):
+            if tid < BLOCKS:
+                r = rsrc(epoch)
+
+                def load_mark():
+                    return fx.Int32(
+                        bo.buffer_load(
+                            r,
+                            EPOCH_MARKS + tid,
+                            vec_width=1,
+                            dtype=T.i32,
+                            cache_modifier=CM_DEV,
+                        )
                     )
-                )
 
-            _ = spin_until(load_mark, lambda v: v != ep + 1)
-        gpu.barrier()
+                _ = spin_until(load_mark, lambda v: v != ep + 1)
+            gpu.barrier()
         if tid == 0:
             if const_expr(reset is not None):
                 reset()

sgl_amd_mono/deepseek_v41/draft.py

@@ -188,22 +188,29 @@ def forward(
 
         lm = self.lm
         x = input_embeds
+        # the first stage's launch reads the target's whole table itself where it is kept
+        first = None
         if x is None:  # ``forward_embed``, from the target's whole table
             if lm.embed_tokens is None:
                 raise ValueError(
                     "DeepseekV4ForCausalLMDSpark requires the target embed_tokens "
                     "(call attach_shared_modules first)."
                 )
-            x = target.embed(lm.embed_tokens, input_ids).unsqueeze(1)
-            x = x.repeat(1, lm.hc_mult, 1)
+            table = lm.embed_tokens.__dict__.get(target.FULL_EMBEDDING)
+            if table is not None:
+                first = (table, input_ids.contiguous())
+            else:
+                x = target.embed(lm.embed_tokens, input_ids).unsqueeze(1)
+                x = x.repeat(1, lm.hc_mult, 1)
         # a stage's post folds into the next stage's K1, the last stage's into the head's
         # combine (``fused.post_combine``)
         pending, pre = (None, x, None, None), None
         fuse = fused.enabled()
         for j, L in enumerate(self.stages):
             out = target._launch(
-                self.runner, L, plan, positions, forward_batch, pending, pre
-            )
+                self.runner, L, plan, positions, forward_batch, pending, pre,
+                embed=first if j == 0 else None,
+            )  # fmt: skip
             defer = j + 1 < len(self.stages) or fuse
             x, pre, post = target._boundary_out(L.module, defer, *out)
             pending = post if defer else (None, x, None, None)
@@ -243,8 +250,10 @@ def _build(model, hf_config, ctx) -> _Draft:
     widths = list(ctx.widths)
     start = len(widths) * p.tp_rank // p.tp_size
     for m in widths[start:] + widths[:start]:
-        # K1 with and without an owed post; K2 on the window: no index source, no top-k
-        runner.precompile(m, k1={(False, False), (True, False)}, k2={(False, False)})
+        # with and without an owed post, on the window: no index source, no top-k
+        runner.precompile(
+            m, layer={(False, False, False), (True, False, False), (False, False, True)}
+        )
     draft_sampler.install(model, ctx.deadline_s)
     return _Draft(model, runner, stages)
 

sgl_amd_mono/deepseek_v41/draft_sampler.py

@@ -1,7 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-"""A sampled DSpark draft's Markov block in two launches a step (``gfx950/markov``), offered
-through SGLang's draft sampler entry; checked against SGLang's ``draft_markov_block`` before
-capture on every rank. ``SGLANG_AMD_MONO_DSV41_MARKOV=0`` keeps SGLang's steps."""
+"""A sampled DSpark draft's Markov block in two launches a step and one more for its step
+logits (``gfx950/markov``), offered through SGLang's draft sampler entry; checked against
+SGLang's ``draft_markov_block`` before capture on every rank. The step logits are SGLang's
+whole rows: a sampled sequence's are gathered from every rank's shard, a temperature-0
+sequence's stay as SGLang's buffer holds them. ``SGLANG_AMD_MONO_DSV41_MARKOV=0`` keeps
+SGLang's steps."""
 
 from __future__ import annotations
 
@@ -30,10 +33,11 @@
         "org_vocab_end",
         "_opt_markov_w2_bf16",
         "step_logits_dtype",
+        "vocab_size",
         "lm_head",
         "num_embeddings_per_partition",
     ),
-    "srt/models/dspark.py": ("vocab_shard",),
+    "srt/models/dspark.py": ("vocab_shard", "vocab_size"),
     "srt/speculative/dspark_components/dspark_draft.py": (
         "draft_markov_block",
         "draft_markov_block(markov_head)",
@@ -45,7 +49,6 @@
         "draft_markov_block(tp_sync)",
         "draft_markov_block(logits_out)",
         "draft_markov_block(greedy_draft)",
-        "draft_markov_block(sharded_accept)",
         "DraftNoise",
         "key",
         "offset",
@@ -99,8 +102,10 @@ def _why_not(model, step_logits) -> str | None:
     part = int(model.lm_head.num_embeddings_per_partition)
     if step_logits is None or step_logits.dtype != torch.float32:
         return "no fp32 step logits"
-    if step_logits.shape[-1] != part:
-        return "step logits gathered whole, not on the shards"
+    if step_logits.shape[-1] != head.vocab_size:
+        return f"step logits {step_logits.shape[-1]} wide, not the vocab's {head.vocab_size}"
+    if shard.org_vocab_start != group.rank_in_group * part or part % 4:
+        return f"vocab shards of {part} not at the ranks' multiples of the partition"
     return None
 
 
@@ -120,30 +125,42 @@ def __init__(
         self.head = head.markov_w2.weight[shard.org_vocab_start : shard.org_vocab_end]
         assert self.head.is_contiguous()
         device = self.head.device
-        part = int(model.lm_head.num_embeddings_per_partition)
+        self.width = int(model.lm_head.num_embeddings_per_partition)
         self.part = torch.empty(
             markov.MAX_SEQUENCES,
-            markov.blocks(part),
+            markov.blocks(self.width),
             2,
             dtype=torch.float32,
             device=device,
         )
         self.exchange = markov.Exchange(
-            group.cpu_group, group.rank_in_group, group.world_size, device
+            group.cpu_group,
+            group.rank_in_group,
+            group.world_size,
+            device,
+            gamma,
+            self.width,
         )
         self.ops = markov.ops()
 
     def block(self, base_logits, anchor, temps, noise, out, step_logits):
         """Every step of a [bs, gamma, padded] base logits block: tokens to ``out``
-        ([bs * gamma]), the step logits to ``step_logits`` ([bs, gamma, padded])."""
-        x = self.exchange
+        ([bs * gamma]), the sampled sequences' step logits to ``step_logits``
+        ([bs, gamma, vocab])."""
+        from sgl_amd_mono.deepseek_v41.gfx950 import markov
+
+        x, bs = self.exchange, base_logits.shape[0]
         for s in range(self.gamma):
             self.ops.markov_sample_step(
-                base_logits[:, s], self.embed, self.head, anchor,
-                None if step_logits is None else step_logits[:, s], temps, noise.offset,
-                int(noise.key), self.vocab_start, s, self.part, out, self.gamma, x.tok,
-                x.calls, x.peers, x.rank, s == self.gamma - 1, self.deadline_s,
+                base_logits[:, s], self.embed, self.head, anchor, x.slab[:bs, s], temps,
+                noise.offset, int(noise.key), self.vocab_start, s, self.part, out,
+                self.gamma, x.tok, x.calls, x.peers, x.rank, s == self.gamma - 1,
+                self.deadline_s,
             )  # fmt: skip
+        # the last step's swap saw every rank's packet: their step logits are written
+        self.ops.markov_gather(
+            step_logits, temps, x.peers, markov.EXCHANGE_BYTES, self.width
+        )
 
     def step(self, hidden_states, anchor_tokens):
         """SGLang's base logits and Markov block of a bs-row draft block."""
@@ -168,9 +185,9 @@ class _NoSync:
             def sync(self, site, value):
                 return value
 
-        head, gamma = self.model.markov_head, self.gamma
+        head, gamma, width = self.model.markov_head, self.gamma, self.width
         device = self.head.device
-        width = int(self.model.lm_head.num_embeddings_per_partition)
+        tp = self.exchange.size
         g = torch.Generator(device=device).manual_seed(1234)
         for bs in range(1, max_bs + 1):
             base = (
@@ -182,26 +199,30 @@ def sync(self, site, value):
             temps = torch.tensor([1.0, 0.0, 0.7, 1.3][:bs], device=device)
             offset = torch.randint(0, 1 << 30, (1,), generator=g, device=device)
             want_logits = torch.empty(
-                bs, gamma, width, dtype=torch.float32, device=device
+                bs, gamma, head.vocab_size, dtype=torch.float32, device=device
             )
+            # every rank draws the same base: its whole row is the shard's, repeated
             want, _ = draft_markov_block(
                 markov_head=head, base_logits=base, anchor_tokens=anchor,
                 draft_hidden=None, step_temperatures=temps,
                 noise=DraftNoise(key=int(self.noise.key), offset=offset.clone()),
                 tp_sync=_NoSync(), logits_out=want_logits, greedy_draft=False,
-                sharded_accept=True,
             )  # fmt: skip
-            got_logits = torch.empty_like(want_logits)
+            got_logits = torch.zeros_like(want_logits)
             got = torch.empty(bs * gamma, dtype=torch.int64, device=device)
             noise = DraftNoise(key=int(self.noise.key), offset=offset.clone())
             self.block(base, anchor, temps, noise, got, got_logits)
             if not torch.equal(got.view(bs, gamma), want.long()):
                 return f"{bs} rows: tokens {got.tolist()} vs SGLang's {want.tolist()}"
             # the bias GEMV sums in another order: a bf16 rounding flip of the bias now and then
-            bias = (want_logits - base.float()).abs()
-            apart = (got_logits - want_logits).abs()
+            whole = base.float().repeat(1, 1, tp)[..., : head.vocab_size]
+            sampled = temps != 0
+            bias = (want_logits - whole).abs()[sampled]
+            apart = (got_logits - want_logits).abs()[sampled]
             if not (apart <= bias * 2**-7 + 1e-6).all():
                 return f"{bs} rows: step logits apart by {apart.max().item():.3g}"
+            if got_logits[~sampled].any():
+                return f"{bs} rows: a temperature-0 row's step logits written"
             if int(noise.offset.item()) != int(offset.item()) + 1:
                 return "the noise counter did not advance once"
         return None

sgl_amd_mono/deepseek_v41/gfx950/attention/back.py

@@ -41,6 +41,7 @@
 )
 from sgl_amd_mono.common.stamps import stamp
 from sgl_amd_mono.deepseek_v41.gfx950.stages.gemv import (
+    TILE,
     ld_f32,
     lds_row,
 )
@@ -196,6 +197,27 @@ def _fp8_quad_bf16(word, scale):
 # ---------------------------------------------------------------- split
 
 
+@flyc.jit
+def await_k1(c, unit, t):
+    """One launch per layer: wait for every CTA's q rows of token t's tile (``qb``; a CTA
+    flags them past its records, key table rows and mixes), t's key table (``krdy``)
+    and, token parallel, every row's record (``vrdy``: a row's keys hold the step's
+    earlier rows, the other ranks' copied late), each published by K1's half of the
+    launch. ``c["k1_seen"](unit)`` runs once they are: a flag of every CTA was seen."""
+    tid, s = c["tid"], c["S"]
+    if tid < BLOCKS:
+        c["mb"].poll([(c["qb"], ((t // TILE) * BLOCKS + tid) * FLAG_STRIDE, 1)])
+    else:
+        if tid == BLOCKS:
+            c["mb"].poll([(c["krdy"], t, 1)])
+        else:
+            if const_expr("vrdy" in c):
+                if tid < BLOCKS + 1 + s:
+                    c["mb"].poll([(c["vrdy"], (tid - BLOCKS - 1) * FLAG_STRIDE, 1)])
+    gpu.barrier()
+    c["k1_seen"](unit)
+
+
 @flyc.jit
 def stage_split(c, unit):
     """Flash-decoding over chunk ``unit % nchunk`` of token ``unit // nchunk``,
@@ -218,6 +240,10 @@ def stage_split(c, unit):
     g = lane // 16
     r16 = lane % 16
     k0 = ch * ck
+    # one launch per layer: K1's half of this launch wrote q, the records and the table
+    one = "qb" in c
+    if const_expr(one):
+        await_k1(c, unit, t)
     # q does not depend on the key table: in flight before an index source's wait
     qrow = d.heads * HEAD_DIM // 8
     qpad = d.tile_heads * HEAD_DIM // 8  # with the head tile's padding rows
@@ -231,10 +257,11 @@ def stage_split(c, unit):
                     (t * d.q_rows) // 2 + fx.min(u, qrow - 1) * 4,
                     vec_width=4,
                     dtype=T.i32,
+                    cache_modifier=CM_DEV if const_expr(one) else 0,
                 )
             )
         )
-    if const_expr("krdy" in c):
+    if const_expr(("krdy" in c) and not one):
         # an index source's table: this launch builds it after the indexer ran
         await_flags(c, c["krdy"], 1, lambda u: t)
     # kv_len and this lane's key load together: the table holds -1 past a token's keys
@@ -516,6 +543,7 @@ def stage_combine(c, unit):
                 ((t * nchunk + cc) * d.heads + h) * 2,
                 vec_width=2,
                 dtype=T.f32,
+                cache_modifier=CM_DEV,
             )
         )
         pv = fx.Vector(
@@ -524,6 +552,7 @@ def stage_combine(c, unit):
                 (((t * nchunk + cc) * d.heads + h) * HEAD_DIM + d0) // 2,
                 vec_width=4,
                 dtype=T.i32,
+                cache_modifier=CM_DEV,
             )
         )
         mc = live.select(stv[0], fx.Float32(NEG))
@@ -755,8 +784,10 @@ def back_context(key, lds, args: dict, scratch, tag, layout: dict) -> dict:
 
 
 @flyc.jit
-def run_back(c, key, bid):
-    """The attention back's stages on this CTA: split, combine, wo_a, wo_b."""
+def run_back(c, key, bid, idle=None):
+    """The attention back's stages on this CTA: split, combine, wo_a, wo_b.
+    ``idle``: work for the CTAs past their attention units, after the attention work and
+    before their GEMV weights load."""
     s = key.tokens
     dd = Dims(key.tp)
     # plain ints: a Python object read inside traced control flow is carried
@@ -773,6 +804,9 @@ def run_back(c, key, bid):
         stage_combine(c, unit)
     gpu.barrier()
     stamp(c, 2)
+    if const_expr(idle is not None):
+        # before any GEMV weights load, so none are live under it
+        idle()
     # the GEMV weights right after the attention work: the wo_a unit's (if any), then wo_b's
     nt = len(tiles)
     units = woa_units(s, dd)