"""Block verification for DSpark sampled decoding (Sun et al., "Block Verification Accelerates Speculative Decoding", ICLR 2025). Default OFF. Lossless: the output distribution is the target's. DSV41_BLOCK_VERIFY=1 replaces the chain (token-by-token) rejection sampler used for sampled rows. Token verification accepts draft i with min(1, p_i/q_i) independently; block verification decides the accepted length jointly over the block and is provably optimal for a single draft chain: w_0 = 1, w_i = min(w_{i-1} * p_i(X_i) / q_i(X_i), 1) h_i = S_i / (S_i + 1 - w_i), S_i = sum_x max(w_i * p_{i+1}(x) - q_{i+1}(x), 0) (i < gamma) h_gamma = w_gamma tau = max{ i : eta_i <= h_i } (0 if none), eta_i ~ U(0,1) correction token ~ normalise(max(w_tau * p_{tau+1} - q_{tau+1}, 0)) (tau < gamma) ~ p_{gamma+1} (tau = gamma) p_i = target distribution after temperature / top-k / top-p (the same one the chain sampler uses), q_i = the distribution the draft token was actually sampled from. Measured offline on engine-exact p and q (capture schema v4, coding-with-thinking traffic): 2.927 -> 2.999 tokens per step (+2.5 %). Random draws come from torch's device generator in the same order on every TP rank, exactly like the stock sampler's torch.rand calls, so all ranks take the same decision. """ import os import torch ENABLED = os.environ.get("DSV41_BLOCK_VERIFY", "0").strip() not in ("0", "", "off", "false") def block_accept_probs(X, target_probs, draft_probs, eta=None, gumbel=None, live=None): """X [bs, g] drafted tokens, target_probs [bs, g+1, V], draft_probs [bs, g, V]. live [bs] (optional): drafts actually verified per request (<= g); positions beyond it are treated as never drafted (q = 0, p(X) = 0), which is the block rule for that shorter block. Returns (tau [bs] int64, correction token [bs] int64).""" bs, g = X.shape dev = X.device p_tok = target_probs[:, :g].gather(-1, X.unsqueeze(-1)).squeeze(-1).float() q_tok = draft_probs.gather(-1, X.unsqueeze(-1)).squeeze(-1).float() ratio = p_tok / q_tok.clamp_min(1e-30) if live is not None: alive = torch.arange(g, device=dev).view(1, -1) < live.view(-1, 1) ratio = torch.where(alive, ratio, torch.zeros_like(ratio)) draft_probs = draft_probs * alive.unsqueeze(-1).to(draft_probs.dtype) w = torch.cumprod(ratio, dim=1) # w_i = min(w_{i-1} * r_i, 1) is not a plain cumprod once it saturates; run the recursion ws = [] cur = torch.ones(bs, device=dev) for i in range(g): cur = torch.minimum(cur * ratio[:, i], torch.ones_like(cur)) ws.append(cur) w = torch.stack(ws, 1) # [bs, g] = w_1..w_g h = torch.empty(bs, g, device=dev) if g > 1: s = (w[:, :g - 1, None] * target_probs[:, 1:g].float() - draft_probs[:, 1:g].float()).clamp_min(0).sum(-1) h[:, :g - 1] = s / (s + 1 - w[:, :g - 1]).clamp_min(1e-30) h[:, g - 1] = w[:, g - 1] if eta is None: eta = torch.rand(bs, g, device=dev) ok = eta <= h # [bs, g] if live is not None: ok = ok & alive idx = torch.arange(1, g + 1, device=dev).view(1, -1) tau = torch.where(ok, idx, torch.zeros_like(idx)).max(dim=1).values # [bs] # correction distribution at position tau (0-based row tau of target / draft) w_all = torch.cat([torch.ones(bs, 1, device=dev), w], dim=1) # w_0..w_g wt = w_all.gather(1, tau.view(-1, 1)) # [bs, 1] rows = tau.view(-1, 1, 1).expand(-1, 1, target_probs.shape[-1]) p_row = target_probs.gather(1, rows).squeeze(1).float() # [bs, V] q_pad = torch.cat([draft_probs.float(), torch.zeros_like(draft_probs[:, :1]).float()], dim=1) q_row = q_pad.gather(1, rows).squeeze(1) # 0 for tau == g res = (wt * p_row - q_row).clamp_min(0) mass = res.sum(-1, keepdim=True) res = torch.where(mass > 0, res / mass.clamp_min(1e-30), p_row) # degenerate: p == q if gumbel is None: gumbel = torch.empty_like(res).exponential_() corr = (res / gumbel).argmax(dim=-1) return tau, corr def install(accept_module): """sglang.kernels.ops.speculative.dspark.dspark_accept: swap AcceptSampling for sampled rows.""" if not ENABLED: return cls = getattr(accept_module, "AcceptSampling", None) core = getattr(accept_module, "_accept_sampling_core", None) if cls is None or core is None: raise RuntimeError("DSV41_BLOCK_VERIFY: AcceptSampling/_accept_sampling_core gone; engine drifted") if getattr(cls, "_dsv41_block_verify", False): return cls._dsv41_block_verify = True softmax_temp = accept_module.SoftmaxTemp import sglang.srt.speculative.dflash_utils as du def execute(*, candidates, target_logits, draft_probs, sampling_info, draft_input, gamma, verify_num_draft_tokens, cutoff_verify_lens=None): if cutoff_verify_lens is not None and os.environ.get("DSV41_VERIFY_CAP", "").strip() in ("", "0", "off"): return cls._dsv41_orig_execute(candidates=candidates, target_logits=target_logits, draft_probs=draft_probs, sampling_info=sampling_info, draft_input=draft_input, gamma=gamma, verify_num_draft_tokens=verify_num_draft_tokens, cutoff_verify_lens=cutoff_verify_lens) bs = candidates.shape[0] if not sampling_info.need_top_k_sampling and not sampling_info.need_top_p_sampling: target_probs = softmax_temp.execute(logits=target_logits, temperatures=sampling_info.temperatures, rows_per_request=verify_num_draft_tokens ).view(bs, verify_num_draft_tokens, -1) else: target_probs = du.build_dflash_verify_target_probs( next_token_logits=target_logits, sampling_info=sampling_info, draft_token_num=verify_num_draft_tokens, bs=bs, max_top_k=draft_input.max_top_k, uniform_top_k_value=draft_input.uniform_top_k_value) X = candidates.view(bs, verify_num_draft_tokens)[:, 1:gamma + 1].long() live = None if cutoff_verify_lens is None else (cutoff_verify_lens[:bs].to(torch.int64) - 1) tau, corr = block_accept_probs(X, target_probs.view(bs, verify_num_draft_tokens, -1), draft_probs.view(bs, gamma, -1), live=live) return tau.to(torch.int32), corr.to(torch.int64), torch.zeros(bs, dtype=torch.int32, device=candidates.device) cls._dsv41_orig_execute = cls.execute cls.execute = staticmethod(execute) print("[block_verify] armed: block verification for sampled DSpark rows", flush=True)