diff --git a/sgl_amd_mono/glm53_flash/gfx950/kernel.py b/sgl_amd_mono/glm53_flash/gfx950/kernel.py index 4f7ac15..fedfa7f 100644 --- a/sgl_amd_mono/glm53_flash/gfx950/kernel.py +++ b/sgl_amd_mono/glm53_flash/gfx950/kernel.py @@ -271,6 +271,9 @@ class FlashBuild(msgspec.Struct, frozen=True): head: bool = False # with kda: the step's first launch, from the step's tokens tokens: bool = False + # an MTP verify step's rows a request (its draft chain; 1: decode): request r + # holds rows r * verify .., and rows past the width's whole requests are pads + verify: int = 1 # False: a failed wait keeps the launch running and leaves its error words (tests) stop: bool = True @@ -299,6 +302,11 @@ def build_flash_kernel(key: FlashBuild): conv_dim_first = key.conv_dim_first fault_site, fault_after = key.fault_site, key.fault_after dense, head, tokens = key.dense, key.head, key.tokens + # a verify step: VL rows a request, NREQ whole requests in the width + VL = key.verify + VERIFY = VL > 1 + NREQ = S // VL + assert 1 <= VL <= S assert 2 <= S <= 16 and 1 <= launches_per_step <= LAYER_SLOTS assert key.inter % 128 == 0 assert not kda or S <= 8 @@ -323,6 +331,8 @@ def build_flash_kernel(key: FlashBuild): name = kernel_symbol("glm53_flash_tail", s=S, tp=W) else: extra = {k: 1 for k, on in (("n", dense), ("e", tokens)) if on} + if VERIFY: + extra["v"] = VL name = kernel_symbol( "glm53_flash_layer", s=S, @@ -449,6 +459,10 @@ def flash_ffn_kernel( state_idx: Int64, conv_stride: Int64, rec_stride: Int64, + inter_rec: Int64, + inter_conv: Int64, + inter_rec_stride: Int64, + inter_conv_stride: Int64, dsa_tab: Int64, dsa_rows: Int64, page_table: Int64, @@ -1858,7 +1872,7 @@ def run_fgb(): # after the head's output, then stores the new states; at 8 rows the states also # follow the output. The conv inputs wait in LDS words XO .. (with o_proj, past # its staged rows) - OPROJ_IN_KDA = S <= 2 and W > 2 + OPROJ_IN_KDA = S <= 2 and W > 2 and not VERIFY STATES_LAST = OPROJ_IN_KDA or S > 4 XO = X_WORDS - 3 * KD if OPROJ_IN_KDA else 7 * KD assert not OPROJ_IN_KDA or XO >= max(7 * KD, S * P // 2) @@ -1870,9 +1884,32 @@ def conv_at(ch, j): return ch * (CONV_TAPS - 1) + j return j * (3 * P) + ch - def run_kda(then): - """``then``: the CTA's o_proj task, run between the head's output and its - state write-back with ``OPROJ_IN_KDA``.""" + def kda_slot(r_idx, s): + """SGLang's KDA state index of row ``s`` (a token launch's rows past the + step's are pads: 0).""" + if const_expr(tokens): + n_rows = step_rows() + idx = _uniform( + fx.Int32( + bo.buffer_load( + r_idx, fx.min(s, n_rows - 1), vec_width=1, dtype=T.i32 + ) + ) + ) + idx = (s < n_rows).select(idx, fx.Int32(0)) + else: + idx = _uniform( + fx.Int32(bo.buffer_load(r_idx, s, vec_width=1, dtype=T.i32)) + ) + return idx + + def run_kda_verify(): + """A verify step: one task a (request, head) runs the request's VL rows + in order from its committed state, which it leaves as is; after row j it + writes the state and the conv window to step j of the request's verify + scratch row (its batch position), as SGLang's verify does. A pad request + (state index not positive) writes nothing and outputs zeros, as do the + width's rows past its whole requests.""" r_cw = _rsrc(conv_w) r_alog = _rsrc(a_log) r_dtb = _rsrc(dt_bias) @@ -1882,21 +1919,168 @@ def run_kda(then): t = fx.Int32(t) s = t // H h = t % H - if const_expr(tokens): - # SGLang's slot indices hold the step's rows; the rest are pads - n_rows = step_rows() - idx = _uniform( - fx.Int32( + if (s % VL == 0) & (s < NREQ * VL): + idx = kda_slot(r_idx, s) + valid = idx > 0 + slot = fx.Int64(fx.max(idx, fx.Int32(0))) + req = fx.Int64(s // VL) + r_cs = _rsrc(conv_state + slot * conv_stride * fx.Int64(2)) + ch = (fx.min(tid, 3 * KD - 1) // KD) * P + h * KD + tid % KD + gi = fx.min(tid, 2 * KD - 1) + st = [ld_bf16(r_cs, conv_at(ch, j)) for j in range(CONV_TAPS - 1)] + cw = [ld_f32(r_cw, ch * CONV_TAPS + j) for j in range(CONV_TAPS)] + kch = tid % KD + dtb = ld_f32(r_dtb, h * KD + kch) + amp = _exp(ld_f32(r_alog, h)) + won = [ld_f32(r_won, 2 * lane + e) for e in range(2)] + vr = tid // 4 + k0 = (tid % 4) * 32 + r_rs = _rsrc( + rec_state + + (slot * rec_stride + fx.Int64(h * KD * KD)) * fx.Int64(4) + ) + hv = [] + for q in range_constexpr(8): + w4 = fx.Vector( bo.buffer_load( - r_idx, fx.min(s, n_rows - 1), vec_width=1, dtype=T.i32 + r_rs, vr * KD + k0 + 4 * q, vec_width=4, dtype=T.f32 ) ) - ) - idx = (s < n_rows).select(idx, fx.Int32(0)) - else: - idx = _uniform( - fx.Int32(bo.buffer_load(r_idx, s, vec_width=1, dtype=T.i32)) - ) + hv += [w4[e] for e in range(4)] + for j in range_constexpr(VL): + row = s + j + # every thread polls (the channel and gate clamped), so the + # conv window carries to the next row outside a branch + x, g, b_raw = getf_many( + [ + (mb("proj"), row * IN_ROWS + ch), + ( + mb("gates"), + row * 2 * P + (gi // KD) * P + h * KD + gi % KD, + ), + (mb("proj"), row * IN_ROWS + 3 * P + h), + ] + ) + a = cw[CONV_TAPS - 1] * x + for k in range_constexpr(CONV_TAPS - 1): + a = a + cw[k] * st[k] + st = st[1:] + [x] + if tid < 3 * KD: + lds_st(xs, tid, bf16_round(div(a, 1.0 + _exp(-a)))) + if tid < 2 * KD: + lds_st(xs, G1O + tid, g) + if tid == 0: + lds_st(misc, 0, b_raw) + if valid: + r_icw = _rsrc( + inter_conv + + ( + req * inter_conv_stride + + fx.Int64(j * (CONV_TAPS - 1) * 3 * P) + ) + * fx.Int64(2) + ) + for k in range_constexpr(CONV_TAPS - 1): + bo.buffer_store( + fx.Float32(st[k]).to(fx.BFloat16), + r_icw, + conv_at(ch, k), + ) + gpu.barrier() + if wave < 2: + a0 = lds_ld(xs, wave * KD + lane) + a1 = lds_ld(xs, wave * KD + lane + 64) + ss = wave_sum(a0 * a0 + a1 * a1) + if lane == 0: + lds_st(misc, 1 + wave, ss) + if (wave >= 2) & (wave < 4): + gk = lds_ld(xs, G1O + kch) + dtb + lds_st( + xs, + DO + kch, + _exp(div(KDA_LOWER_BOUND, 1.0 + _exp(-(amp * gk)))), + ) + gpu.barrier() + rq = _rsq(lds_ld(misc, 1) + L2_EPS) * (KD**-0.5) + rk = _rsq(lds_ld(misc, 2) + L2_EPS) + beta = sigmoid(lds_ld(misc, 0)) + kk = [lds_ld(xs, KO + k0 + i) * rk for i in range(32)] + dot = fx.Float32(0.0) + for i in range_constexpr(32): + hv[i] = hv[i] * lds_ld(xs, DO + k0 + i) + dot = dot + hv[i] * kk[i] + for off in (1, 2): + dot = _xred(dot, off, lambda a, b: a + b) + vn = beta * (lds_ld(xs, VO + vr) - dot) + o = fx.Float32(0.0) + for i in range_constexpr(32): + hv[i] = hv[i] + vn * kk[i] + o = o + hv[i] * (lds_ld(xs, QO + k0 + i) * rq) + for off in (1, 2): + o = _xred(o, off, lambda a, b: a + b) + if valid: + r_irs = _rsrc( + inter_rec + + ( + req * inter_rec_stride + + fx.Int64((j * H + h) * KD * KD) + ) + * fx.Int64(4) + ) + for q in range_constexpr(8): + bo.buffer_store( + fx.Vector.from_elements( + hv[4 * q : 4 * q + 4], fx.Float32 + ), + r_irs, + vr * KD + k0 + 4 * q, + ) + if tid % 4 == 0: + lds_st( + xs, + OO + vr, + valid.select(bf16_round(o), fx.Float32(0.0)), + ) + gpu.barrier() + if wave == 0: + e0 = 2 * lane + o0 = lds_ld(xs, OO + e0) + o1 = lds_ld(xs, OO + e0 + 1) + rstd = _rsq( + wave_sum(o0 * o0 + o1 * o1) * (1.0 / KD) + O_NORM_EPS + ) + y0 = o0 * rstd * won[0] * sigmoid(lds_ld(xs, G2O + e0)) + y1 = o1 * rstd * won[1] * sigmoid(lds_ld(xs, G2O + e0 + 1)) + put( + mb("o_in"), + (row * P + h * KD + e0) // 2, + bf16_pair(y0, y1), + ) + signal("box_o_in", row * H + h) + gpu.barrier() + if const_expr(NREQ * VL < S): + if s >= NREQ * VL: + if wave == 0: + put( + mb("o_in"), + (s * P + h * KD + 2 * lane) // 2, + bf16_pair(fx.Float32(0.0), fx.Float32(0.0)), + ) + signal("box_o_in", t) + + def run_kda(then): + """``then``: the CTA's o_proj task, run between the head's output and its + state write-back with ``OPROJ_IN_KDA``.""" + r_cw = _rsrc(conv_w) + r_alog = _rsrc(a_log) + r_dtb = _rsrc(dt_bias) + r_won = _rsrc(w_on) + r_idx = _rsrc(state_idx) + for t in range(start("kda"), S * H, G): + t = fx.Int32(t) + s = t // H + h = t % H + idx = kda_slot(r_idx, s) valid = idx > 0 slot = fx.Int64(fx.max(idx, fx.Int32(0))) r_cs = _rsrc(conv_state + slot * conv_stride * fx.Int64(2)) @@ -2147,6 +2331,21 @@ def entry(name, words): tid * ROW_WORDS + ROW_BT, ) + def chain_rows(s): + """The rows whose cache writes row ``s`` reads (a static row): itself, + and in a verify step its request's earlier rows.""" + first = s - s % VL if s < NREQ * VL else s + return range(first, s + 1) + + def chain_ready(s): + """Poll specs of ``kv_ready`` of row ``s`` (dynamic) and, in a verify + step, of its request's earlier rows (a pad row's repeat its own).""" + return ( + [(mb("kv_ready"), fx.max(s - d, s - s % VL), 1) for d in range(VL)] + if VERIFY + else [(mb("kv_ready"), s, 1)] + ) + def start_at(b): """First task of a stage whose task 0 runs on CTA ``b``.""" return (bid + (G - b % G)) & (G - 1) @@ -2333,7 +2532,10 @@ def dsa_qnorm(tv, t32): # pool into the index cache: per dim softmax(gate + ape) over its four # tokens' (k, gate) (three from the tail ring, the row's own), the # weighted k, bf16, Hadamard, bf16, ue8m0 FP8. Then the row's (k, gate) - # goes to ring slot pos % 4. ``kv_ready`` releases both caches. + # goes to ring slot pos % TAIL (SGLang's ring of 4 slots a request, plus a + # verify step's draft tokens). ``kv_ready`` releases both caches. In a + # verify step one task runs a request's rows in order: a row's pool reads + # the ring entries of the chain's earlier rows. def dsa_cache(tv, t32): KVB = t32["kv_block_tokens"] KVE = tv["kv_block_elems"] @@ -2341,12 +2543,23 @@ def dsa_cache(tv, t32): IPB = t32["idx_block_pools"] IBB = t32["idx_block_bytes"] TBE = t32["tail_block_elems"] + TAIL = TBE // INDEX_DIM r_gkv = _rsrc(tv["g_kv"]) r_lnw, r_lnb = _rsrc(tv["ln_w"]), _rsrc(tv["ln_b"]) r_ape = _rsrc(tv["ape"]) - gpu.barrier() - for t in range(start("cache"), S, G): - s = fx.Int32(t) + + # a verify chain's earlier rows wrote the ring in this task: read it past + # the CU's cache + ring_cm = {"cache_modifier": CM_DEV} if VERIFY else {} + + def ring_ld(r, i): + return fx.Float32( + fx.BFloat16( + bo.buffer_load(r, i, vec_width=1, dtype=T.bf16, **ring_cm) + ) + ) + + def cache_row(s): pos = row_word(s, ROW_POS) slot = row_word(s, ROW_SLOT) pslot = row_word(s, ROW_POOL_SLOT) @@ -2387,15 +2600,14 @@ def dsa_cache(tv, t32): ) r_tail = _rsrc(tv["tail_k"] + tail_row) r_gate = _rsrc(tv["tail_g"] + tail_row) - tbase = fx.Int32(0) complete = live & (pslot >= 0) & (phase == KPOOL - 1) if complete: if lk: sj, kj = [], [] for j in range_constexpr(KPOOL - 1): - ring = tbase + j * INDEX_DIM + d - sj.append(ld_bf16(r_gate, ring)) - kj.append(ld_bf16(r_tail, ring)) + ring = (pos - (KPOOL - 1) + j) % TAIL * INDEX_DIM + d + sj.append(ring_ld(r_gate, ring)) + kj.append(ring_ld(r_tail, ring)) sj.append(graw) kj.append(kn) sj = [ @@ -2448,7 +2660,7 @@ def dsa_cache(tv, t32): cache_modifier=CM_DEV, ) if live & (tslot >= 0) & lk: - e = tbase + phase * INDEX_DIM + d + e = fx.max(pos, 0) % TAIL * INDEX_DIM + d bo.buffer_store( kn.to(fx.BFloat16), r_tail, e, cache_modifier=CM_DEV ) @@ -2462,6 +2674,19 @@ def dsa_cache(tv, t32): if tid == 0: put(mb("kv_ready"), s, fx.Int32(1)) + gpu.barrier() + for t in range(start("cache"), S, G): + s = fx.Int32(t) + if const_expr(VERIFY): + if (s % VL == 0) & (s < NREQ * VL): + for j in range_constexpr(VL): + cache_row(s + j) + if const_expr(NREQ * VL < S): + if s >= NREQ * VL: + cache_row(s) + else: + cache_row(s) + # ======================= q_b (FP8 or bf16): two 16-row tiles a task def dsa_q_b(tv, t32): r_wqb, r_sqb = _rsrc(tv["w_qb"]), _rsrc(tv["s_qb"]) @@ -2784,7 +3009,8 @@ def score_half(s, n_pools, loaded): half_block(r_ibt, n_pools, first, h) for h in range(HALVES) ] if tid == 0: - get(mb("kv_ready"), s) + for r in chain_rows(s): + get(mb("kv_ready"), r) fx.memory_fence( ordering=fx.AtomicOrdering.Acquire, syncscope="agent" ) @@ -3308,10 +3534,8 @@ def kv_row(blk, off): if picked_keys: if wave == 0: picked = poll( - [ - (mb("tok_sel"), s * SEL_STRIDE + kidx, 1), - (mb("kv_ready"), s, 1), - ] + [(mb("tok_sel"), s * SEL_STRIDE + kidx, 1)] + + chain_ready(s) )[0][0] fx.memory_fence( ordering=fx.AtomicOrdering.Acquire, syncscope="agent" @@ -3319,6 +3543,14 @@ def kv_row(blk, off): resolve((kidx < nk).select(picked, fx.Int32(0))) else: if wave == 0: + if const_expr(VERIFY): + # the chain's earlier rows' keys are among the + # task's: wait for all of them first + poll(chain_ready(s)) + fx.memory_fence( + ordering=fx.AtomicOrdering.Acquire, + syncscope="agent", + ) # a scored row's tail follows its last complete pool; keys # past the row's address its first token (masked) own = long_row.select( @@ -3359,7 +3591,12 @@ def kv_row(blk, off): lds_st(xs, qw, qv[i][0].bitcast(fx.Float32)) lds_st(xs, qw + 1, qv[i][1].bitcast(fx.Float32)) gpu.barrier() - own_j = picked_keys.select(fx.Int32(-1), nk - 1 - t * SPLIT_KEYS) + # (a verify row waited for its chain's keys before loading them) + own_j = ( + fx.Int32(-1) + if VERIFY + else picked_keys.select(fx.Int32(-1), nk - 1 - t * SPLIT_KEYS) + ) if (own_j >= 0) & (own_j < SPLIT_KEYS): if wave == 0: if lane == 0: @@ -3820,7 +4057,10 @@ def kda_oproj(done=None): # a CTA with a kda task (at most one) runs its o_proj task inside it assert S * H <= G - run_kda(kda_oproj) + if const_expr(VERIFY): + run_kda_verify() + else: + run_kda(kda_oproj) kda_oproj(start("kda") < S * H if OPROJ_IN_KDA else None) run_seam("seam", "ffn_own") elif const_expr(dsa): @@ -4442,6 +4682,10 @@ def launch( state_idx: Int64, conv_stride: Int64, rec_stride: Int64, + inter_rec: Int64, + inter_conv: Int64, + inter_rec_stride: Int64, + inter_conv_stride: Int64, dsa_tab: Int64, dsa_rows: Int64, page_table: Int64, @@ -4494,6 +4738,10 @@ def launch( state_idx, conv_stride, rec_stride, + inter_rec, + inter_conv, + inter_rec_stride, + inter_conv_stride, dsa_tab, dsa_rows, page_table, diff --git a/sgl_amd_mono/glm53_flash/gfx950/runner.py b/sgl_amd_mono/glm53_flash/gfx950/runner.py index c33c797..76f25bf 100644 --- a/sgl_amd_mono/glm53_flash/gfx950/runner.py +++ b/sgl_amd_mono/glm53_flash/gfx950/runner.py @@ -146,10 +146,15 @@ class FlashDsaCaches(msgspec.Struct, frozen=True): class FlashKdaState(msgspec.Struct, frozen=True): """One KDA layer's decode state: bf16 conv state [slots, 3, 3P] and f32 - recurrent state [slots, H, 128 (v), 128 (k)], slot strides in elements.""" + recurrent state [slots, H, 128 (v), 128 (k)], slot strides in elements. A verify + launch also writes SGLang's verify scratch: the f32 recurrent state + [rows, steps, H, 128, 128] and the bf16 conv window [rows, steps, 3, 3P] after + each draft token (a scratch row a request, its batch position).""" conv_state: torch.Tensor rec_state: torch.Tensor + inter_rec: Optional[torch.Tensor] = None + inter_conv: Optional[torch.Tensor] = None # Rows of SGLang's DSA page table an index block spans: an index page holds 64 @@ -157,6 +162,9 @@ class FlashKdaState(msgspec.Struct, frozen=True): PAGE_TOKENS = 64 INDEX_BLOCK_POOLS = 64 IDX_BT_GROUP = INDEX_BLOCK_POOLS * KPOOL // PAGE_TOKENS +# the KDA launch arguments: the attention-side mHC and KDA weights, the state +# and its indices and strides, the verify scratch and its strides +KDA_ARGS = 21 def compile_only(launch) -> None: @@ -429,7 +437,7 @@ def launch( _check("comb", comb, (S, HC_MULT, HC_MULT), torch.float32) x_out, res_out, post_out, comb_out = rt.outs[self.layer % 2] p = lambda t: t.data_ptr() # noqa: E731 - kda_args = [0] * 17 + kda_args = [0] * KDA_ARGS dsa_args = [0, 0, 0, 0] if self.dsa is not None: if caches is None or rows is None or page_table is None: @@ -456,6 +464,12 @@ def launch( n = step["ids"].shape[0] if self.embedding is not None else S _check("state_idx", state_idx, (n,), torch.int32) k = self.kda + inter = [0, 0, 0, 0] + if self.rt.key.verify > 1: + ir, ic = state.inter_rec, state.inter_conv + if ir is None or ic is None: + raise ValueError("a verify launch needs the verify scratch") + inter = [p(ir), p(ic), ir.stride(0), ic.stride(0)] kda_args = [ *map(p, (k.hc_fn, k.hc_scale, k.hc_base, k.g_norm, k.w_in, k.w_fb)), *map(p, (k.w_gb, k.conv_w, k.a_log, k.dt_bias, k.w_on, k.w_o)), @@ -464,6 +478,7 @@ def launch( p(state_idx), state.conv_state.stride(0), state.rec_state.stride(0), + *inter, ] if self.embedding is not None: dsa_args = [p(self._step_table(step)), 0, 0, 0] @@ -537,7 +552,7 @@ def launch(self, x, residual, post, comb) -> torch.Tensor: p(rt.logits), p(w.w_head), *[0] * 9, - *[0] * 17, + *[0] * KDA_ARGS, *[0] * 4, p(rt.scratch), rt.peer.local, diff --git a/sgl_amd_mono/glm53_flash/spec.py b/sgl_amd_mono/glm53_flash/spec.py index 389b141..2fde2ae 100644 --- a/sgl_amd_mono/glm53_flash/spec.py +++ b/sgl_amd_mono/glm53_flash/spec.py @@ -33,6 +33,7 @@ try: from sglang.kernels.plugin_kernels.amd.spec import ( DeviceLimits, + Feature, MonoSpec, Numerics, NumericsKind, @@ -41,6 +42,7 @@ except ModuleNotFoundError: # SGLang before sgl-project/sglang#45389 from sglang.srt.monokernels.spec import ( DeviceLimits, + Feature, MonoSpec, Numerics, NumericsKind, @@ -74,6 +76,14 @@ def _width(rows: int) -> int: "srt/arg_groups/fields/*.py": ( "page_size", "enable_linear_replayssm", + "enable_linear_replayssm_spec", + "speculative_algorithm", + "speculative_eagle_topk", + "speculative_num_steps", + "speculative_num_draft_tokens", + "linear_attn_backend", + "linear_attn_decode_backend", + "linear_attn_verify_backend", "enable_hisparse", "expert_distribution_recorder_mode", "watchdog_timeout", @@ -89,6 +99,7 @@ def _width(rows: int) -> int: "SGLANG_SIMULATE_UNIFORM_EXPERTS", "SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS", "SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK", + "SGLANG_ENABLE_PP_SPEC", ), "srt/configs/model_config.py": ("context_len",), "srt/distributed/parallel_state.py": ("cpu_group",), @@ -106,6 +117,8 @@ def _width(rows: int) -> int: "state_layer_cache", ), "srt/layers/attention/mamba/mamba2_metadata.py": ("mamba_cache_indices",), + # fused accept: the verify kernel seeds from the accepted step, not the state + "srt/layers/attention/linear/kda_backend.py": ("accept_lens_pool",), "srt/layers/layer_boundary/__init__.py": ("get_attn_tp_context",), "srt/layers/layer_boundary/adapters/attention.py": ("maybe_input_scattered",), "srt/layers/layer_boundary/residual/batch.py": ("start", "clear"), @@ -125,6 +138,8 @@ def _width(rows: int) -> int: "srt/mem_cache/memory_pool.py": ( "conv", "temporal", + "intermediate_ssm", + "intermediate_conv_window", "get_key_buffer", "get_index_k_with_scale_buffer", "get_compress_tail_buffers", @@ -163,13 +178,18 @@ def _width(rows: int) -> int: "Numerics", "NumericsKind", "Requirements", + "Feature", "widths", "deadline_s", "copies", + "tokens_per_request", ), + # NEXTN resolves to EAGLE + "srt/arg_groups/speculative_hook.py": ('"EAGLE"', '"NEXTN"'), "srt/runtime_context.py": ( "get_device", "get_parallel", + "get_spec", "process_model_config", '"tp_rank"', '"tp_size"', @@ -261,6 +281,24 @@ def _config_disable_reasons(cfg, hf_config) -> list[str]: reasons.append( f"--page-size {cfg.page_size}: the kpool caches need {PAGE_SIZE}" ) + # a verify request's rows are one draft chain, each row the next position + if cfg.speculative_algorithm is not None and ( + cfg.speculative_eagle_topk != 1 + or cfg.speculative_num_draft_tokens != cfg.speculative_num_steps + 1 + ): + reasons.append("MTP verifies one draft chain: topk 1, draft tokens = steps + 1") + # SGLang's FlashInfer KDA verify fuses the accept (KDAAttnBackend's + # _can_fuse_accept_state): it skips the KDA state commit, which a verify + # launch reads. Unset, the verify backend follows a FlashInfer decode backend + decode = cfg.linear_attn_decode_backend or cfg.linear_attn_backend + verify = cfg.linear_attn_verify_backend or ( + "flashinfer" if decode == "flashinfer" else None + ) + if cfg.speculative_algorithm is not None and verify == "flashinfer": + reasons.append( + "MTP with the FlashInfer KDA verify backend fuses the accept, leaving " + "the KDA state a verify launch reads uncommitted" + ) if not aiter_can_use_preshuffle_paged_mqa(): reasons.append( "needs the index cache in AITER's preshuffled layout, which the kernels " @@ -378,7 +416,17 @@ class _Built: """One rank's monokernels: the runtimes per (width, kind), each layer's ops, the tail's ops, and the forward.""" - def __init__(self, lm, layers: dict, runtimes: dict, heads: dict, tp: int, device): + def __init__( + self, + lm, + layers: dict, + runtimes: dict, + heads: dict, + tp: int, + device, + verify: int = 1, + max_position: int = 0, + ): from sgl_amd_mono.glm53_flash.gfx950.config import ( HC_MULT, HIDDEN, @@ -392,7 +440,12 @@ def __init__(self, lm, layers: dict, runtimes: dict, heads: dict, tp: int, devic self._layers, self.runtimes, self.heads = layers, runtimes, heads self.last = max(layers) self.max_rows = TP_MAX_ROWS.get(tp, max(WIDTHS)) - self.check_layers = os.environ.get("SGLANG_AMD_MONO_CHECK") == "1" + # a step's rows a request: the verify step's draft tokens with MTP, else 1 + self.verify, self.max_position = verify, max_position + # the launch checks run SGLang's decode layer: decode steps only + self.check_layers = ( + os.environ.get("SGLANG_AMD_MONO_CHECK") == "1" and verify == 1 + ) # layer 0's input besides its tokens: every stream a copy of the embedding row # (x = the row, post 1, comb 0, residual 0) n = max(WIDTHS) @@ -412,6 +465,11 @@ def __init__(self, lm, layers: dict, runtimes: dict, heads: dict, tp: int, devic rows[:, ROW_BT] = 0 self.rows[w] = rows self.kda_idx[w] = torch.zeros(w, dtype=torch.int32, device=device) + # a verify step's positions and request rows, made per row: a graph's step + # table holds their addresses, so they are an eager step's + self.verify_in = { + w: torch.zeros(2, w, dtype=torch.int64, device=device) for w, _ in runtimes + } # the logits processor returns the tail's logits on a routed step lp = lm.logits_processor self._own_logits, self._logits = lp._get_logits, None @@ -435,9 +493,14 @@ def forward( input_embeds=None, pp_proxy_tensors=None, ): - """``Glm5NextForConditionalGeneration.forward`` for a routed decode step.""" + """``Glm5NextForConditionalGeneration.forward`` for a routed decode or MTP + verify step (a verify step's rows are its requests' draft chains).""" lm = self.lm - if plan.key.rows > self.max_rows or input_embeds is not None: + if ( + plan.key.rows > self.max_rows + or plan.key.tokens_per_request != self.verify + or input_embeds is not None + ): return type(lm).forward( lm, input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors ) @@ -497,7 +560,12 @@ def _layer(self, L: _Layer, plan, input_ids, positions, owed, forward_batch): W = _width(S) checking = self.check_layers and not torch.cuda.is_current_stream_capturing() backend = get_attn_backend() - slots = backend.linear_attn_backend.forward_metadata.mamba_cache_indices[:S] + V = self.verify + # the KDA state indices and request rows are a request's (a verify step: + # bs of them), the launches' a row's + slots = backend.linear_attn_backend.forward_metadata.mamba_cache_indices[ + : S // V + ] op = L.ops[W] if L.i == 0: ins = ( @@ -506,14 +574,30 @@ def _layer(self, L: _Layer, plan, input_ids, positions, owed, forward_batch): self.one_post[:W], self.zero_comb[:W], ) + rpi = forward_batch.req_pool_indices[: S // V] + pos = positions[:S] + if V > 1: + if backend.linear_attn_backend.accept_lens_pool is not None: + raise RuntimeError( + "GLM-5.3-Flash monokernel: a verify launch reads the committed " + "KDA state, not the fused accept's" + ) + slots = slots.repeat_interleave(V) + buf = self.verify_in[W] + # verify rows past the context are lookahead the scheduler + # discards; a row's pools end at the context's + buf[0, :S].copy_(pos.clamp(max=self.max_position)) + buf[1, :S].copy_(rpi.repeat_interleave(V)) + 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 plan.state["dsa_rows"] = (self.rows[W], page_table[:S]) step = { "ids": input_ids[:S], - "positions": positions[:S], + "positions": pos, "out_cache_loc": forward_batch.out_cache_loc[:S], - "req_pool_indices": forward_batch.req_pool_indices[:S], + "req_pool_indices": rpi, "slots": self.kda_idx[W], "rows": self.rows[W], } @@ -527,13 +611,15 @@ def _layer(self, L: _Layer, plan, input_ids, positions, owed, forward_batch): hidden = self.model.embed_tokens(input_ids[:S]) if L.i == 0 else streams ref = check.reference(L, positions, hidden, forward_batch, S, idx[:S]) if L.kda: - self._bind_kda(L, backend.linear_attn_backend) + self._bind_kda(L, backend.linear_attn_backend, V) out = op.launch(*ins, state=L.state, state_idx=idx, step=step) - backend.linear_attn_backend._track_mamba_state_decode( - forward_batch, L.state.conv_state, L.state.rec_state, slots, L.i - ) + # a verify step's states are committed after it, with their radix copy + if V == 1: + backend.linear_attn_backend._track_mamba_state_decode( + forward_batch, L.state.conv_state, L.state.rec_state, slots, L.i + ) else: - self._bind_dsa(L) + self._bind_dsa(L, V) rows, page_table = plan.state["dsa_rows"] out = op.launch(*ins, caches=L.caches, rows=rows, page_table=page_table) if checking: @@ -544,16 +630,31 @@ def _layer(self, L: _Layer, plan, input_ids, positions, owed, forward_batch): return _Owed(*out) @staticmethod - def _bind_kda(L: _Layer, kda) -> None: + def _bind_kda(L: _Layer, kda, verify: int) -> None: + """The layer's KDA state, and with MTP (``verify`` rows a request) its + verify scratch.""" from sgl_amd_mono.glm53_flash.gfx950.runner import ( FlashKdaState, ) cache = kda.state_layer_cache(L.i) conv, rec = cache.conv[0], cache.temporal - if L.state is not None and ( - L.state.conv_state.data_ptr() == conv.data_ptr() - and L.state.rec_state.data_ptr() == rec.data_ptr() + ir = ic = None + if verify > 1: + ir, ic = cache.intermediate_ssm, cache.intermediate_conv_window[0] + bound = (conv, rec, ir, ic) + if L.state is not None and all( + (a is None and b is None) + or (a is not None and b is not None and a.data_ptr() == b.data_ptr()) + for a, b in zip( + ( + L.state.conv_state, + L.state.rec_state, + L.state.inter_rec, + L.state.inter_conv, + ), + bound, + ) ): return P3, H, D = L.attn.conv_w.shape[0], L.attn.a_log.shape[0], 128 @@ -572,10 +673,28 @@ def _bind_kda(L: _Layer, kda) -> None: f"{rec.dtype} {tuple(rec.shape)} {rec.stride()}; needs bf16 " f"[slots, 3, {P3}] and f32 [slots, {H}, {D}, {D}]" ) - L.state = FlashKdaState(conv_state=conv, rec_state=rec) + if verify > 1 and not ( + ir is not None + and ir.dtype == torch.float32 + and tuple(ir.shape[1:]) == (verify, H, D, D) + and ir.stride()[1:] == (H * D * D, D * D, D, 1) + and ic.dtype == torch.bfloat16 + and tuple(ic.shape[1:]) == (verify, 3, P3) + and ic.stride()[1:] == (3 * P3, P3, 1) + ): + raise RuntimeError( + f"GLM-5.3-Flash monokernel: layer {L.i}: KDA verify scratch " + f"{None if ir is None else (ir.dtype, tuple(ir.shape), ir.stride())}, " + f"conv window {ic.dtype} {tuple(ic.shape)} {ic.stride()}; needs f32 " + f"[rows, {verify}, {H}, {D}, {D}] and dense bf16 " + f"[rows, {verify}, 3, {P3}]" + ) + L.state = FlashKdaState( + conv_state=conv, rec_state=rec, inter_rec=ir, inter_conv=ic + ) @staticmethod - def _bind_dsa(L: _Layer) -> None: + def _bind_dsa(L: _Layer, verify: int) -> None: from sglang.srt.model_executor.forward_context import get_token_to_kv_pool from sgl_amd_mono.glm53_flash.gfx950.runner import ( @@ -617,15 +736,21 @@ def _bind_dsa(L: _Layer) -> None: f"index cache {idx.dtype} {tuple(idx.shape)}, needs contiguous uint8 " f"[pages, {INDEX_BLOCK_POOLS * 132}]" ) + # SGLang's ring: 4 slots a request, plus a verify step's draft tokens; a + # verify chain's pool reads the 3 slots before its row, after the chain's + # earlier rows wrote theirs + ring = 4 if verify == 1 else verify + 3 for name, t in (("k", tail_k), ("gate", tail_g)): if ( t.dtype != torch.bfloat16 - or tuple(t.shape[1:]) != (4, 128) + or t.dim() != 3 + or t.shape[1] < ring + or t.shape[2] != 128 or not t.is_contiguous() ): reasons.append( f"kpool {name} ring {t.dtype} {tuple(t.shape)}, needs contiguous " - "bf16 [slots, 4, 128]" + f"bf16 [slots, >= {ring}, 128]" ) if reasons: raise RuntimeError( @@ -646,7 +771,7 @@ def _deadline(key, deadline_s: float, longest_s: int) -> float: def _build(model, hf_config, ctx) -> _Built: - from sglang.srt.runtime_context import get_parallel, process_model_config + from sglang.srt.runtime_context import get_parallel, get_spec, process_model_config from sgl_amd_mono.common import guard from sgl_amd_mono.glm53_flash.gfx950.config import ( @@ -723,8 +848,16 @@ def _build(model, hf_config, ctx) -> _Built: fault_site, fault_after = guard.fault_injection(rank) or (0, 0) # the pools a sparse-attention row scores at most: the served context's, whole # score tasks of 256 pools - pools = -(-process_model_config().context_len // KPOOL) + context_len = process_model_config().context_len + pools = -(-context_len // KPOOL) max_pools = -(-pools // POOLS_PER_TASK) * POOLS_PER_TASK + # MTP: the target's verify steps, a draft chain of draft-tokens rows a request + sp = get_spec() + verify = ( + int(sp.speculative_num_draft_tokens) + if sp.speculative_algorithm is not None + else 1 + ) runtimes, heads = {}, {} cap = TP_MAX_ROWS.get(tp, max(WIDTHS)) widths = sorted({_width(r) for r in ctx.widths if r <= cap}) @@ -764,6 +897,7 @@ def runtime(key): dense=kind in ("dense", "embed"), dense_shuffled=dense_sh if kind in ("dense", "embed") else False, tokens=kind == "embed", + verify=verify, fault_site=fault_site, fault_after=fault_after, ) @@ -796,7 +930,7 @@ def runtime(key): for (rw, _), rt in runtimes.items(): if rw == w: rt.precompile() - return _Built(model, layers, runtimes, heads, tp, device) + return _Built(model, layers, runtimes, heads, tp, device, verify, context_len - 1) SPEC = MonoSpec( @@ -804,7 +938,9 @@ def runtime(key): devices={ "gfx950": DeviceLimits(widths=ROWS, tp_sizes=TP_SIZES, resident_blocks=BLOCKS) }, - features=frozenset(), + # decode steps, or the target's verify steps of MTP (the NextN draft runs + # SGLang's forward) + features=frozenset({Feature.SPECULATIVE_DECODING}), requires=Requirements( dtypes=frozenset({"bfloat16"}), attention_backends=frozenset({"dsa"}), @@ -813,6 +949,11 @@ def runtime(key): "moe_runner_backend": frozenset({"auto", "aiter"}), # the launch updates the KDA state in place; ReplaySSM defers it to a ring "enable_linear_replayssm": frozenset({False}), + # MTP: the NextN draft's verify steps (NEXTN resolves to EAGLE); a verify + # launch writes every draft token's KDA state to the verify scratch, + # which ReplaySSM's spec mode does not keep + "speculative_algorithm": frozenset({None, "EAGLE"}), + "enable_linear_replayssm_spec": frozenset({False}), # HiSparse maps the logical KV slots to device rows; the kernels write raw slots "enable_hisparse": frozenset({False}), # the launches record no expert distribution, and layer 44 would record @@ -826,6 +967,9 @@ def runtime(key): # the kernel breaks router ties in AITER's grouped top-k order, the path # SGLang takes with this off; on, its Triton gate keeps the lower expert id "SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK": False, + # PP spec remaps the verify scratch rows; a verify launch writes request + # r's to row r + "SGLANG_ENABLE_PP_SPEC": False, }, ), numerics=NUMERICS, diff --git a/test/e2e/glm53_flash_server.py b/test/e2e/glm53_flash_server.py index 29478f6..74335e2 100644 --- a/test/e2e/glm53_flash_server.py +++ b/test/e2e/glm53_flash_server.py @@ -1,5 +1,5 @@ """The GLM-5.3-Flash servers that the monokernel serving test launches: MI35x, TP4, -decode steps (no speculative decoding).""" +decode steps, and with ``MTP`` the NextN draft's verify steps.""" import os @@ -40,3 +40,15 @@ "--mm-feature-transport", "cpu", ] +# MTP with one draft chain of 3 tokens: 4 verify rows a request, so verify steps of 1 +# and 2 requests fit the 8-row width +MTP = [ + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", +] diff --git a/test/e2e/test_glm53_flash_server.py b/test/e2e/test_glm53_flash_server.py index 05a8786..0866ac9 100644 --- a/test/e2e/test_glm53_flash_server.py +++ b/test/e2e/test_glm53_flash_server.py @@ -1,7 +1,8 @@ """GLM-5.3-Flash with --monokernel-backend amd on MI35x, TP4 (RFC #43949), across the switch points: decode steps of 1 to 8 rows run the kernels, wider ones the normal path, and requests join and leave a running batch that steps past 8 rows and back -(a short GSM8K run, for gross failures); the tokens of covered steps scored again by +(a short GSM8K run, for gross failures); with MTP, verify steps of 1 and 2 requests +(4 and 8 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 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 @@ -16,7 +17,7 @@ import requests import torch -from glm53_flash_server import ARGS, ENV, MODEL +from glm53_flash_server import ARGS, ENV, MODEL, MTP from sglang.test.kl_test_utils import ( test_input_output_logprobs_match_helper as assert_logprobs_match, ) @@ -37,6 +38,9 @@ # Decode against prefill, one batch of prompts at a time (every decode step covered); # the normal path itself averages about 0.04 on 8 prompts KL_BOUND = 0.06 +# MTP (3 draft tokens) accepts about 3.3 tokens a step on this GSM8K run at TP4; +# the bound is for gross failures +MTP_ACCEPT_BOUND = 2.5 def _launch(extra_env=None, extra_args=(), log=None): @@ -51,6 +55,24 @@ def _launch(extra_env=None, extra_args=(), log=None): ) +def _gsm8k(threads=12): + """GSM8K, 200 questions ``threads`` at a time, thinking on.""" + return run_sgl_eval( + SimpleNamespace( + base_url=DEFAULT_URL_FOR_TEST, + model=MODEL, + eval_name="gsm8k", + num_examples=200, + num_threads=threads, + max_tokens=32768, + temperature=1.0, + top_p=0.95, + seed=42, + sgl_eval_thinking=True, + ) + ) + + class TestMonokernelDecode(CustomTestCase): @classmethod def setUpClass(cls): @@ -64,20 +86,7 @@ def tearDownClass(cls): def test_gsm8k_past_the_largest_width(self): # 12 requests at a time: the batch steps past 8 rows as it fills, and back # as it drains - metrics = run_sgl_eval( - SimpleNamespace( - base_url=DEFAULT_URL_FOR_TEST, - model=MODEL, - eval_name="gsm8k", - num_examples=200, - num_threads=12, - max_tokens=32768, - temperature=1.0, - top_p=0.95, - seed=42, - sgl_eval_thinking=True, - ) - ) + metrics = _gsm8k() print(f"{metrics=}") self.assertGreater(metrics["score"], GSM8K_BOUND) @@ -98,6 +107,28 @@ def test_covered_steps_match_prefill(self): ) +class TestMonokernelMtp(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.process = _launch(extra_args=[*MTP, "--watchdog-timeout", "1200"]) + + @classmethod + def tearDownClass(cls): + if getattr(cls, "process", None) is not None: + terminate_and_kill_process_tree(cls.process) + + def test_gsm8k_and_accept_length_on_the_kernels(self): + # 2 requests at a time, so every verify step fits the 8-row width and runs + # the kernels. A verify row that reads the wrong KDA state or keys of its + # draft chain lowers the accept length before the score + metrics = _gsm8k(threads=2) + info = requests.get(DEFAULT_URL_FOR_TEST + "/get_server_info").json() + accept = info["internal_states"][0]["avg_spec_accept_length"] + print(f"{metrics=} {accept=}") + self.assertGreater(metrics["score"], GSM8K_BOUND) + self.assertGreater(accept, MTP_ACCEPT_BOUND) + + 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 diff --git a/test/kernels/glm53_flash/check_layer.py b/test/kernels/glm53_flash/check_layer.py index 80fae6e..1f2071b 100644 --- a/test/kernels/glm53_flash/check_layer.py +++ b/test/kernels/glm53_flash/check_layer.py @@ -15,7 +15,8 @@ gathered embedding rows, and writes the spec's step tables (``--step-rows``: the step's rows, fewer than the width); ``--check-head``: the tail's logits against a torch tail and an f32 tail; ``--replays N``: a graph of ``--layers`` launches -replays the eager step bitwise.""" +replays the eager step bitwise; ``--verify L`` (``kda`` or ``dsa``): an MTP verify +launch of L rows a request equals its tokens' L decode launches bitwise.""" import argparse import os @@ -1043,21 +1044,23 @@ def check_embed(args, rank, tp, group, device) -> bool: embedding rows (x = the row on every rank; post 1, comb 0, residual 0) on the same weights and state: outputs and state writes bitwise equal, the width's rows past the step's pads; and the step tables it writes equal the spec's torch tables.""" - from msgspec.structs import astuple + + def decode_state(s): + return (s.conv_state, s.rec_state) rt, launches, _, ins = build(args, "embed", rank, tp, group, device) op, kw = launches[0] step = kw["step"] n = step["ids"].shape[0] _, table = embedding(rank, tp, device) - state0 = [t.clone() for t in astuple(kw["state"])] + state0 = [t.clone() for t in decode_state(kw["state"])] out = [t.clone() for t in op.launch(*ins, **kw)] torch.cuda.synchronize() kda, rows = spec_tables(step, kw["state_idx"], args.rows) tables = torch.equal(step["slots"], kda) and torch.equal(step["rows"], rows) ref_rt, ref_launches, _, _ = build(args, "dense", rank, tp, group, device) ref_op, ref_kw = ref_launches[0] - for t, v in zip(astuple(ref_kw["state"]), state0): + for t, v in zip(decode_state(ref_kw["state"]), state0): t.copy_(v) x = torch.zeros(args.rows, HIDDEN, dtype=torch.bfloat16, device=device) x[:n] = table[step["ids"]] @@ -1066,7 +1069,7 @@ def check_embed(args, rank, tp, group, device) -> bool: same = all(torch.equal(a, b) for a, b in zip(out, ref)) writes = all( torch.equal(a, b) - for a, b in zip(astuple(kw["state"]), astuple(ref_kw["state"])) + for a, b in zip(decode_state(kw["state"]), decode_state(ref_kw["state"])) ) print( f"rank {rank}: {n} of {args.rows} rows from the tokens == the dense launch on " @@ -1129,6 +1132,209 @@ def tail(h, bf16): return rel < 2e-3 and rel32 < 1e-2 and repeat and caught +def _max_diff(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).abs().max()) if a.numel() else 0.0 + + +def check_verify(args, rank, tp, group, device) -> bool: + """An MTP verify launch (``--verify`` rows a request, ``--rows`` wide: whole + requests, then pad rows) against the requests' draft tokens run one decode + launch a token, on the same weights, inputs, state and caches: each row's output + equals its token's decode output; with KDA the verify scratch after token j + equals the state after j + 1 decode launches and the committed state is + unchanged; with sparse attention the MLA cache, the index cache and the kpool + rings equal the decode launches' (a chain's pool reads the ring entries its + earlier rows wrote).""" + from msgspec.structs import replace + + from sgl_amd_mono.common import guard + from sgl_amd_mono.glm53_flash.gfx950.config import ( + INDEX_DIM, + KDA_DIM, + KDA_HEADS, + KDA_PROJ, + KPOOL, + ) + from sgl_amd_mono.glm53_flash.gfx950.kernel import FlashBuild + from sgl_amd_mono.glm53_flash.gfx950.layout import ( + ROW_BT, + ROW_POOL_SLOT, + ROW_POS, + ROW_SLOT, + ROW_TAIL_SLOT, + ROW_WORDS, + ) + from sgl_amd_mono.glm53_flash.gfx950.runner import FlashLayerOp, FlashRuntime + + kind, L, W = args.kind, args.verify, args.rows + assert kind in ("kda", "dsa") and 2 <= L <= W + R = W // L + # the decode launches at the same width (a width's tiles set its sums' order): + # the requests' rows, then pads + DW = W + kda = kind == "kda" + + def runtime(rows, verify): + key = FlashBuild( + rows=rows, + tp=tp, + launches_per_step=1, + inter=GLOBAL_INTER // tp, + kda=kda, + dsa=not kda, + max_pools=0 if kda else args.max_pools, + verify=verify, + stop=False, + ) + rt = FlashRuntime(key, rank, group, device) + rt.set_deadline(args.deadline) + return rt + + vrt, drt = runtime(W, L), runtime(DW, 1) + ffn = ffn_weights(rank, tp, device) + attn = kda_weights(rank, tp, device) if kda else dsa_weights(rank, tp, device) + a = {"kda": attn} if kda else {"dsa": attn} + vop, dop = FlashLayerOp(vrt, ffn, 0, **a), FlashLayerOp(drt, ffn, 0, **a) + x, res, post, comb = layer_inputs(W, device) + ins = (x, res, post, comb) + + def rows_of(t, j): + """Decode step j's input rows: request r's row r * L + j, then pads.""" + out = torch.zeros((DW,) + tuple(t.shape[1:]), dtype=t.dtype, device=device) + out[:R] = t[[r * L + j for r in range(R)]] + return out + + def run(op, rt, ins, **kw): + out = op.launch(*ins, **kw) + rt.advance_step() + torch.cuda.synchronize() + words = rt.epoch[guard.ERR_DEADLINE : guard.ERR_FAULT + 1].tolist() + if any(words): + print(f"rank {rank}: error words {words}", flush=True) + return None + return tuple(t.clone() for t in out) + + ok = True + if kda: + P, H, D = KDA_PROJ // tp, KDA_HEADS // tp, KDA_DIM + state, _ = kda_state(R, tp, device) + inter_rec = torch.full( + (R + 1, L, H, D, D), float("nan"), dtype=torch.float32, device=device + ) + inter_conv = torch.zeros( + R + 1, L, 3, 3 * P, dtype=torch.bfloat16, device=device + ) + vstate = replace( + state, + conv_state=state.conv_state.clone(), + rec_state=state.rec_state.clone(), + inter_rec=inter_rec, + inter_conv=inter_conv, + ) + vidx = torch.zeros(W, dtype=torch.int32, device=device) + vidx[: R * L] = torch.arange(1, R + 1, device=device).repeat_interleave(L) + vout = run(vop, vrt, ins, state=vstate, state_idx=vidx) + dstate = replace( + state, + conv_state=state.conv_state.clone(), + rec_state=state.rec_state.clone(), + ) + didx = torch.zeros(DW, dtype=torch.int32, device=device) + didx[:R] = torch.arange(1, R + 1, device=device) + if vout is None: + return False + for name in ("conv_state", "rec_state"): + same = torch.equal(getattr(vstate, name), getattr(state, name)) + ok &= same + print(f"rank {rank}: verify keeps the committed {name}: {same}") + for j in range(L): + dins = tuple(rows_of(t, j) for t in ins) + dout = run(dop, drt, dins, state=dstate, state_idx=didx) + if dout is None: + return False + for n, (v, d) in enumerate(zip(vout, dout)): + diff = _max_diff(v[[r * L + j for r in range(R)]], d[:R]) + ok &= diff == 0.0 + if diff: + print(f"rank {rank}: token {j} output {n} differs by {diff}") + # the request's scratch row is its batch position r (slot r + 1) + dr = _max_diff(inter_rec[:R, j], dstate.rec_state[1 : R + 1]) + dc = _max_diff(inter_conv[:R, j], dstate.conv_state[1 : R + 1]) + ok &= dr == 0.0 and dc == 0.0 + print(f"rank {rank}: token {j} scratch state diff {dr}, conv {dc}") + print(f"rank {rank}: verify vs {L} decode launches ok {ok}", flush=True) + return ok + + lens = [int(v) for v in args.lens.split(",")] if args.lens else list(LENS) + lens = (lens * R)[:R] + caches, _, pt = dsa_caches(R, device, lens) + ring = 4 + L + g = _gen(11, device) + tails = [ + torch.randn(R + 1, ring, INDEX_DIM, generator=g, device=device).bfloat16() + for _ in range(2) + ] + caches = replace(caches, tail_k=tails[0], tail_g=tails[1]) + + def table_row(r, pos): + slot = kv_slot(pt, r, pos) + return [pos, slot, slot // KPOOL, (r + 1) * KPOOL] + + def clone(c): + return replace( + c, + kv_cache=c.kv_cache.clone(), + idx_cache=c.idx_cache.clone(), + tail_k=c.tail_k.clone(), + tail_g=c.tail_g.clone(), + ) + + vc, dc = clone(caches), clone(caches) + vt = torch.full((W, ROW_WORDS), -1, dtype=torch.int32, device=device) + vt[:, ROW_BT] = 0 + for r, n in enumerate(lens): + for j in range(L): + s = r * L + j + vt[s, [ROW_POS, ROW_SLOT, ROW_POOL_SLOT, ROW_TAIL_SLOT]] = torch.tensor( + table_row(r, n - 1 + j), dtype=torch.int32, device=device + ) + vt[s, ROW_BT] = s + # SGLang's verify page table: a request's row for each of its rows + vpt = pt.repeat_interleave(L, 0).contiguous() + vout = run(vop, vrt, ins, caches=vc, rows=vt, page_table=vpt) + if vout is None: + return False + for j in range(L): + dt = torch.full((DW, ROW_WORDS), -1, dtype=torch.int32, device=device) + dt[:, ROW_BT] = 0 + for r, n in enumerate(lens): + dt[r, [ROW_POS, ROW_SLOT, ROW_POOL_SLOT, ROW_TAIL_SLOT]] = torch.tensor( + table_row(r, n - 1 + j), dtype=torch.int32, device=device + ) + dt[r, ROW_BT] = r + dout = run( + dop, + drt, + tuple(rows_of(t, j) for t in ins), + caches=dc, + rows=dt, + page_table=pt, + ) + if dout is None: + return False + for n, (v, d) in enumerate(zip(vout, dout)): + diff = _max_diff(v[[r * L + j for r in range(R)]], d[:R]) + ok &= diff == 0.0 + if diff: + print(f"rank {rank}: token {j} output {n} differs by {diff}") + for name in ("kv_cache", "idx_cache", "tail_k", "tail_g"): + same = torch.equal(getattr(vc, name), getattr(dc, name)) + ok &= same + print(f"rank {rank}: {name} equals the decode launches' {same}") + print(f"rank {rank}: verify vs {L} decode launches ok {ok}", flush=True) + return ok + + def _changed_rows(before: torch.Tensor, after: torch.Tensor) -> set: d = (before != after).reshape(before.shape[0], -1).any(1) return set(d.nonzero().flatten().tolist()) @@ -1197,6 +1403,9 @@ def main() -> int: ap.add_argument("--layers", type=int, default=1, help="launches a step (--replays)") ap.add_argument("--pad-rows", type=int, default=0, help="trailing pad rows") ap.add_argument("--step-rows", type=int, default=0, help="an embed step's rows") + ap.add_argument( + "--verify", type=int, default=0, help="an MTP verify launch's rows a request" + ) ap.add_argument( "--writes", action="store_true", help="check the rows a launch writes" ) @@ -1216,6 +1425,7 @@ def main() -> int: (args.check_embed, check_embed), (args.check_head, check_head), (args.replays, check_replays), + (args.verify, check_verify), ): if flag: return 0 if check(args, rank, tp, group, device) else 1 @@ -1227,7 +1437,7 @@ def main() -> int: from msgspec.structs import asdict held = asdict(extra.get("state") or extra["caches"]) - before = {k: t.clone() for k, t in held.items()} + before = {k: t.clone() for k, t in held.items() if t is not None} out = run_step(args.kind, launches, ins) rt.advance_step() torch.cuda.synchronize() diff --git a/test/kernels/test_glm53_flash_kernels.py b/test/kernels/test_glm53_flash_kernels.py index 24b495f..f31f894 100644 --- a/test/kernels/test_glm53_flash_kernels.py +++ b/test/kernels/test_glm53_flash_kernels.py @@ -9,7 +9,9 @@ the TP size builds. The dense launch's MLP output is within the spec's bound of an f32 MLP on its own FFN input, on shuffled and row-major FP8 alike; the launch from the step's tokens equals the dense launch on the gathered embedding rows and writes the -spec's step tables, for steps of fewer rows than the width too; sparse-attention +spec's step tables, for steps of fewer rows than the width too; an MTP verify launch +equals its draft tokens' decode launches bitwise (outputs, KDA verify scratch, +caches and kpool rings) and keeps the committed KDA state; sparse-attention launches and key selection hold for rows past 65536 pools of a 1M-token context, up to 650000 tokens; the tail's logits match a torch tail and an f32 tail, back-to-back tails included.""" @@ -108,6 +110,23 @@ def test_the_launch_from_the_tokens_matches_the_dense_launch(self): str(step), ) + def test_a_verify_launch_matches_its_draft_tokens_decode_launches(self): + # MTP 2, 3 and 5 (3, 4 and 6 rows a request), whole requests then pad rows; + # the lengths complete a pool mid-chain and cross 2048 tokens mid-chain + for kind, rows, verify, lens in ( + ("kda", 4, 4, ""), + ("kda", 8, 3, ""), + ("kda", 8, 6, ""), + ("dsa", 8, 4, "4101,2050"), + ("dsa", 8, 3, "4102,2051"), + ("dsa", 8, 6, "2047"), + ("dsa", 4, 3, "64"), + ): + args = ["--kind", kind, "--rows", str(rows), "--verify", str(verify)] + for tp in (2, 4, 8): + if tp > 2 or rows <= 4: + self._run(tp, *args, *(["--lens", lens] if lens else [])) + def test_rows_past_65536_pools_of_a_1m_context(self): # SGLang's pools for a 1M-token context; 262148 tokens complete pool 65536, # whose id takes 17 bits; up to 650000 tokens, as agentic traces reach. Past diff --git a/test/unit/test_glm53_flash_spec.py b/test/unit/test_glm53_flash_spec.py index 2c405f3..b30cd35 100644 --- a/test/unit/test_glm53_flash_spec.py +++ b/test/unit/test_glm53_flash_spec.py @@ -76,17 +76,20 @@ def forward(self, *args): return "the model's forward" -def _built(tp=4, widths=(2, 4, 8)): +def _built(tp=4, widths=(2, 4, 8), verify=1): lm = _Lm() kinds = ("embed", "dense", "kda", "dsa", "head") runtimes = {(w, k): None for w in widths for k in kinds} layers = {i: SimpleNamespace(i=i) for i in range(45)} - return lm, flash._Built(lm, layers, runtimes, {}, tp=tp, device="cpu") + built = flash._Built(lm, layers, runtimes, {}, tp=tp, device="cpu", verify=verify) + return lm, built -def _plan(rows): +def _plan(rows, tokens=1): key = registry.ShapeKey( - forward_mode="decode", batch_size=rows, tokens_per_request=1 + forward_mode="decode" if tokens == 1 else "target_verify", + batch_size=rows // tokens, + tokens_per_request=tokens, ) return registry.Plan(key=key) @@ -106,6 +109,17 @@ def test_a_step_runs_at_the_next_width_and_tp2_takes_four_rows(): assert set(built.rows) == {2, 4} and (built.rows[4][:, 4] == 0).all() +def test_a_verify_step_of_other_draft_chains_runs_the_model_s_forward(): + # the launches are built for one draft-chain length (the server's draft tokens) + lm, built = _built(verify=4) + out = built.forward(_plan(8, tokens=2), "ids", "positions", "batch") + assert out == "the model's forward" and len(lm.ran) == 1 + # TP2: a verify step past 4 rows runs the model's forward too + lm, built = _built(tp=2, widths=(2, 4), verify=4) + out = built.forward(_plan(8, tokens=4), "ids", "positions", "batch") + assert out == "the model's forward" and len(lm.ran) == 1 + + 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() @@ -124,7 +138,16 @@ def test_the_logits_processor_returns_the_tail_s_logits_on_a_routed_step(): def _cfg(**kw): - return SimpleNamespace(**{"page_size": 256, **kw}) + return SimpleNamespace( + **{ + "page_size": 256, + "speculative_algorithm": None, + "linear_attn_backend": "triton", + "linear_attn_decode_backend": None, + "linear_attn_verify_backend": None, + **kw, + } + ) def _dtypes(temporal=torch.float32, conv=torch.bfloat16): @@ -176,6 +199,66 @@ def test_config_reasons_name_each_mismatch(monkeypatch): assert len(reasons) == 7, reasons +@pytest.mark.parametrize( + "topk, steps, draft, ok", + [(1, 3, 4, True), (1, 5, 6, True), (4, 3, 8, False), (1, 3, 3, False)], +) +def test_mtp_verifies_one_draft_chain(monkeypatch, topk, steps, draft, ok): + # a verify request's rows are one chain of draft tokens, each the next position + import sglang.srt.configs.mamba_utils as mamba_utils + import sglang.srt.layers.attention.dsa.utils as dsa_utils + + monkeypatch.setattr(mamba_utils, "mamba2_state_dtype", _dtypes()) + monkeypatch.setattr(dsa_utils, "aiter_can_use_preshuffle_paged_mqa", lambda: True) + cfg = _cfg( + speculative_algorithm="EAGLE", + speculative_eagle_topk=topk, + speculative_num_steps=steps, + speculative_num_draft_tokens=draft, + ) + reasons = flash.SPEC.config_disable_reasons(cfg, FLASH) + assert (reasons == []) == ok, reasons + + +@pytest.mark.parametrize( + "spec, backends, ok", + [ + ("EAGLE", {}, True), + ("EAGLE", {"linear_attn_verify_backend": "triton"}, True), + ("EAGLE", {"linear_attn_verify_backend": "flashinfer"}, False), + ("EAGLE", {"linear_attn_decode_backend": "flashinfer"}, False), + ("EAGLE", {"linear_attn_backend": "flashinfer"}, False), + ( + "EAGLE", + { + "linear_attn_decode_backend": "flashinfer", + "linear_attn_verify_backend": "triton", + }, + True, + ), + (None, {"linear_attn_backend": "flashinfer"}, True), + ], +) +def test_mtp_refuses_the_fused_accept(monkeypatch, spec, backends, ok): + # SGLang's FlashInfer KDA verify skips the KDA state commit a verify launch + # reads; unset, the verify backend follows a FlashInfer decode backend + import sglang.srt.configs.mamba_utils as mamba_utils + import sglang.srt.layers.attention.dsa.utils as dsa_utils + + monkeypatch.setattr(mamba_utils, "mamba2_state_dtype", _dtypes()) + monkeypatch.setattr(dsa_utils, "aiter_can_use_preshuffle_paged_mqa", lambda: True) + cfg = _cfg( + speculative_algorithm=spec, + speculative_eagle_topk=1, + speculative_num_steps=3, + speculative_num_draft_tokens=4, + **backends, + ) + reasons = flash.SPEC.config_disable_reasons(cfg, FLASH) + assert (reasons == []) == ok, reasons + assert ok or "fuses the accept" in reasons[0], reasons + + def test_the_spec_requires_sglang_jit_grouped_topk_off(): # the kernel's router ties follow AITER's grouped top-k; SGLang's JIT grouped # top-k (a Triton gate) keeps the lower expert id, so the spec requires it off @@ -532,8 +615,10 @@ def _routed_step(monkeypatch, capture=None): lm.model = SimpleNamespace(layers=[last], norm=lambda h: h * 2) monkeypatch.setattr(check, "contract", lambda streams: streams) built = flash._Built(lm, layers, runtimes, {2: _Head(calls)}, tp=4, device="cpu") - monkeypatch.setattr(flash._Built, "_bind_kda", staticmethod(lambda L: None)) - monkeypatch.setattr(flash._Built, "_bind_dsa", staticmethod(lambda L: None)) + monkeypatch.setattr( + flash._Built, "_bind_kda", staticmethod(lambda L, kda, verify: None) + ) + monkeypatch.setattr(flash._Built, "_bind_dsa", staticmethod(lambda L, verify: None)) slots = torch.tensor([5], dtype=torch.int32) linear = SimpleNamespace( forward_metadata=SimpleNamespace(mamba_cache_indices=slots),