From 34b2adac00544f1c0673da33b3ea318945da8c84 Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@nvidia.com>
Date: Tue, 27 Oct 2026 18:22:00 -0700
Subject: [PATCH 1/2] [None][fix] Run windowed DSpark draft blocks non-causally
 on TRTLLM

DSpark drafters are trained with a non-causal block and a two-sided
sliding window, but the drafter resolved every windowed layer as causal.
GQA DSpark layers now resolve non-causal (W-1, W-1).

trtllm-gen windows only causal attention, so the TRTLLM backend serves such
a layer as the causal windowed decode, which also returns its log-sum-exp,
plus a compiled merge of the block keys after each query. The workspace
holds the softmax stats flashinfer carves off for the log-sum-exp. The
window validator, which refused these layers, is removed.

Based on efa53ad417 by Pranav Shrestha.

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
---
 tensorrt_llm/_torch/models/modeling_dflash.py |  95 ++++++++---
 tensorrt_llm/_torch/models/modeling_dspark.py |   5 +
 tensorrt_llm/_torch/speculative/dflash.py     |   3 +-
 .../hw_agnostic/test_dflash2_semantics.py     |  41 -----
 .../test_dspark_noncausal_window.py           | 150 ++++++++++++++++++
 5 files changed, 225 insertions(+), 69 deletions(-)
 create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py

diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py
index 32c5c246baf8..1d51a28dfba9 100644
--- a/tensorrt_llm/_torch/models/modeling_dflash.py
+++ b/tensorrt_llm/_torch/models/modeling_dflash.py
@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
+import math
 from dataclasses import replace
 from typing import Dict, Optional
 
@@ -28,6 +29,7 @@
     get_dflash_trtllm_gen_ops,
 )
 from ..speculative.interface import SpeculativeDecodingMode
+from ..utils import maybe_compile
 from .modeling_utils import FUSED_MODULE_COMPONENTS, get_model_architecture, register_draft_model
 
 
@@ -63,6 +65,56 @@ def dspark_layer_window_size(
     return (swa_window - 1, swa_window - 1)
 
 
+@maybe_compile
+def dflash_merge_later_block_keys_(
+    out: torch.Tensor,
+    lse_log2: torch.Tensor,
+    query: torch.Tensor,
+    block_key: torch.Tensor,
+    block_value: torch.Tensor,
+    scale: float,
+    window_right: int,
+) -> None:
+    """Add each block query's later keys ``j < k <= j + window_right`` to its causal attention
+    ``out`` [B * blk, Hq, D], in place, given the base-2 log-sum-exp ``lse_log2`` [B * blk, Hq].
+    The softmax over the few block keys is unrolled so the merge compiles to one pointwise kernel."""
+    B, blk, num_heads, head_dim = query.shape
+    num_kv_heads = block_key.shape[2]
+    group = num_heads // num_kv_heads
+    q = query.float().reshape(B, blk, num_kv_heads, group, 1, head_dim)
+    keys = block_key.float().permute(0, 2, 1, 3)[:, None, :, None]  # [B, 1, Hkv, 1, blk, D]
+    logits = (q * keys).sum(-1) * scale  # [B, query j, Hkv, g, key k]
+    offsets = torch.arange(blk, device=query.device)
+    distance = offsets.view(1, -1) - offsets.view(-1, 1)  # [j, k] = k - j
+    later = distance > 0
+    if window_right >= 0:
+        later &= distance <= window_right
+    later_logits = [
+        logits[..., k].masked_fill(~later[:, k].view(1, blk, 1, 1), float("-inf"))
+        for k in range(blk)
+    ]
+    max_logit = later_logits[0]
+    for logit in later_logits[1:]:
+        max_logit = torch.maximum(max_logit, logit)
+    # A query without later keys has only -inf logits; its weight is 0 below.
+    max_logit = torch.where(torch.isfinite(max_logit), max_logit, 0.0)
+    exp_logits = [torch.exp(logit - max_logit) for logit in later_logits]
+    sum_exp = exp_logits[0]
+    for exp_logit in exp_logits[1:]:
+        sum_exp = sum_exp + exp_logit
+    has_later = sum_exp > 0
+    sum_exp = torch.where(has_later, sum_exp, 1.0)
+    lse_causal = lse_log2.float().reshape(B, blk, num_kv_heads, group) * math.log(2.0)
+    # exp(lse_later) / (exp(lse_causal) + exp(lse_later))
+    weight = torch.where(has_later, torch.sigmoid(max_logit + torch.log(sum_exp) - lse_causal), 0.0)
+    merged = out.float().reshape(B, blk, num_kv_heads, group, head_dim) * (1 - weight)[..., None]
+    value = block_value.float()
+    prob_scale = weight / sum_exp
+    for k in range(blk):
+        merged = merged + (prob_scale * exp_logits[k])[..., None] * value[:, k, None, :, None]
+    out.copy_(merged.reshape(B * blk, num_heads, head_dim))
+
+
 def _is_dflash2_architecture(config: PretrainedConfig) -> bool:
     """Whether a draft checkpoint advertises itself as a DFlash 2 drafter.
 
@@ -1385,30 +1437,6 @@ def _resolve_block_attention(self, layer_idx: int) -> tuple[bool, tuple[int, int
             window_size = swa_window
         return causal, window_size
 
-    def validate_block_attention_windows(self) -> None:
-        """Refuse windows the TRTLLM block-decode kernel cannot apply.
-
-        TRTLLM-Gen only implements sliding windows for causal masks: flashinfer
-        rejects ``causal=False`` with any finite ``window_left`` outright. A
-        non-causal windowed layer must therefore run on VANILLA or FA4, which
-        take a two-sided window.
-        """
-        if self.dflash_attention_backend != "TRTLLM":
-            return
-        num_layers = getattr(self.config, "num_hidden_layers", None)
-        if num_layers is None:
-            num_layers = len(self.model.layers)
-        for layer_idx in range(num_layers):
-            causal, (window_left, _) = self._resolve_block_attention(layer_idx)
-            if causal or window_left < 0:
-                continue
-            raise ValueError(
-                f"DFlash draft layer {layer_idx} attends non-causally within a "
-                f"{window_left + 1}-token sliding window. The TRTLLM DFlash "
-                "attention backend does not support non-causal sliding-window "
-                "attention. Use attention_backend=VANILLA or FA4."
-            )
-
     def _prepare_dflash_trtllm_gen_buffers(
         self,
         dtype: torch.dtype,
@@ -1431,6 +1459,8 @@ def _prepare_dflash_trtllm_gen_buffers(
             rotary_embedding_dim=0,
             fp8_context_fmha=False,
         )
+        # Softmax stats flashinfer carves off the front when the decode returns its log-sum-exp.
+        workspace_bytes += 8 * num_heads * max_batch_size * 256 * -(-block_size // 256) + (1 << 20)
         device = torch.device(device)
         is_capturing = torch.cuda.is_current_stream_capturing()
         if self._dflash_trtllm_gen_device != device:
@@ -1873,8 +1903,10 @@ def dflash_forward(
                 q_flat = Q_bshd.reshape(-1, num_heads_per_rank, head_dim)
                 out_flat = out.reshape(-1, num_heads_per_rank, head_dim)
                 window_left = window_size[0]
-                if causal:
-                    trtllm_gen_ops.batch_decode_with_kv_cache(
+                # trtllm-gen windows only causal attention; a non-causal window adds the later block keys.
+                merge_later_keys = not causal and window_left >= 0 and block_size > 1
+                if causal or window_left >= 0:
+                    decoded = trtllm_gen_ops.batch_decode_with_kv_cache(
                         query=q_flat,
                         kv_cache=(layer_cache[:, 0], layer_cache[:, 1]),
                         workspace_buffer=self._dflash_trtllm_gen_workspace,
@@ -1897,8 +1929,19 @@ def dflash_forward(
                         kv_cache_sf=None,
                         uses_shared_paged_kv_idx=True,
                         bmm1_scale_log2=None,
+                        return_lse=merge_later_keys,
                         multi_ctas_kv_counter_buffer=self._dflash_trtllm_gen_counters,
                     )
+                    if merge_later_keys:
+                        dflash_merge_later_block_keys_(
+                            out_flat,
+                            decoded[1],
+                            Q_bshd,
+                            k_noise_bshd,
+                            v_noise_bshd,
+                            head_dim**-0.5,
+                            window_size[1],
+                        )
                 else:
                     cum_seq_lens_q = torch.arange(
                         0,
diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py
index 1e2edea4800d..52bf168d8da2 100644
--- a/tensorrt_llm/_torch/models/modeling_dspark.py
+++ b/tensorrt_llm/_torch/models/modeling_dspark.py
@@ -2772,6 +2772,11 @@ def __init__(
             self.mask_token_id = int(spec_mask_token_id)
         self._init_dspark_heads(draft_config.pretrained_config)
 
+    def _get_attention_mask_args(self, layer_idx):
+        """DSpark drafters are trained non-causal; a sliding window bounds both sides."""
+        _, (window_left, _) = super()._get_attention_mask_args(layer_idx)
+        return False, (window_left, window_left)
+
     def load_weights(self, weights: Dict, weight_mapper=None, **kwargs):
         """Take the DSpark head weights, then hand the rest to DFlash."""
         weights, _ = self._take_dspark_head_weights(weights)
diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py
index 410a12f91512..8dd1765ed4e9 100644
--- a/tensorrt_llm/_torch/speculative/dflash.py
+++ b/tensorrt_llm/_torch/speculative/dflash.py
@@ -1423,8 +1423,7 @@ def _lazy_init_ctx_buffers(
             ):
                 # Gated on the drafter that actually runs the trtllm-gen ops:
                 # one paging its own block decode the same way is not shaped
-                # like them. No non-causal sliding window there either.
-                draft_model.validate_block_attention_windows()
+                # like them.
                 has_context_attention = any(
                     not draft_model._get_attention_mask_args(layer_idx)[0] for layer_idx in range(L)
                 )
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py
index 0867979604ae..88a61d81392f 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py
@@ -621,47 +621,6 @@ def test_is_causal_is_ignored_for_a_plain_dflash_drafter():
     assert wrapper._get_attention_mask_args(1) == (False, (-1, -1))
 
 
-# ---------------------------------------------------------------------------
-# TRTLLM block-decode backend: no non-causal sliding window
-# ---------------------------------------------------------------------------
-
-
-def test_released_drafter_window_runs_on_two_sided_backends():
-    """VANILLA and FA4 hand both window bounds to their kernels."""
-    for backend in ("VANILLA", "FA4"):
-        _mask_wrapper(_all_sliding_config(2048), backend=backend).validate_block_attention_windows()
-
-
-def test_trtllm_backend_rejects_the_released_drafter_window():
-    """flashinfer's trtllm-gen context kernel refuses causal=False with any
-    finite window_left, so the released non-causal 2048 window cannot run."""
-    wrapper = _mask_wrapper(_all_sliding_config(2048), backend="TRTLLM")
-    with pytest.raises(ValueError, match="non-causally within a 2048-token sliding window"):
-        wrapper.validate_block_attention_windows()
-
-
-def test_trtllm_backend_accepts_causal_or_unwindowed_layers():
-    # A causal sliding window is supported by the kernel.
-    _mask_wrapper(
-        _all_sliding_config(4, is_causal=True), backend="TRTLLM"
-    ).validate_block_attention_windows()
-    # No window at all, either.
-    _mask_wrapper(
-        SimpleNamespace(num_hidden_layers=2, layer_types=["full_attention"] * 2),
-        backend="TRTLLM",
-    ).validate_block_attention_windows()
-
-
-def test_trtllm_window_check_reads_the_use_swa_override():
-    """A use_swa window replaces the checkpoint-inferred one in the block
-    decode, so the check must resolve the same window the forward uses."""
-    config = SimpleNamespace(num_hidden_layers=2, layer_types=["full_attention"] * 2)
-    wrapper = _mask_wrapper(config, backend="TRTLLM", layer_windows=[(3, 3), (3, 3)])
-    assert wrapper._resolve_block_attention(0) == (False, (3, 3))
-    with pytest.raises(ValueError, match="draft layer 0"):
-        wrapper.validate_block_attention_windows()
-
-
 # ---------------------------------------------------------------------------
 # Tiny end-to-end drafter: config parsing, weight loading, block-decode
 # parity vs an fp32 eager oracle (needs CUDA + flash_attn).
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py
new file mode 100644
index 000000000000..91e7a90da967
--- /dev/null
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py
@@ -0,0 +1,150 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Non-causal sliding-window DSpark block attention on the TRTLLM backend.
+
+The drafter serves it as the trtllm-gen causal windowed decode, with its log-sum-exp, plus
+``dflash_merge_later_block_keys_``; it must equal dense non-causal attention.
+"""
+
+import math
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.models.modeling_dflash import (
+    DFlashForCausalLM,
+    dflash_merge_later_block_keys_,
+)
+
+HEAD_DIM, PAGE_SIZE = 128, 64
+
+
+def _trtllm_gen_ops():
+    if not torch.cuda.is_available():
+        pytest.skip("needs CUDA")
+    from tensorrt_llm._torch.speculative.dflash_attention import (
+        dflash_trtllm_gen_unavailability_reason,
+        get_dflash_trtllm_gen_ops,
+    )
+
+    reason = dflash_trtllm_gen_unavailability_reason()
+    if reason is not None:
+        pytest.skip(f"DFlash TRTLLM backend unavailable: {reason}")
+    return get_dflash_trtllm_gen_ops()
+
+
+def _check_block_attention(ctx_lens, block_size, num_heads, num_kv_heads, window):
+    """Runs the drafter's lowering over a shuffled page pool and compares it to dense fp32 attention."""
+    ops = _trtllm_gen_ops()
+    dev = torch.device("cuda")
+    B = len(ctx_lens)
+    seq_lens = [c + block_size for c in ctx_lens]
+    num_pages = [math.ceil(s / PAGE_SIZE) for s in seq_lens]
+    max_pages = max(num_pages)
+    pages = torch.randperm(sum(num_pages), dtype=torch.int32).split(num_pages)
+    block_tables = torch.zeros(B, max_pages, dtype=torch.int32, device=dev)
+    cache = torch.zeros(
+        sum(num_pages), 2, num_kv_heads, PAGE_SIZE, HEAD_DIM, dtype=torch.bfloat16, device=dev
+    )
+    keys = [torch.randn(s, num_kv_heads, HEAD_DIM, device=dev).bfloat16() for s in seq_lens]
+    values = [torch.randn(s, num_kv_heads, HEAD_DIM, device=dev).bfloat16() for s in seq_lens]
+    for b in range(B):
+        block_tables[b, : num_pages[b]] = pages[b]
+        for pos in range(0, ctx_lens[b], PAGE_SIZE):
+            n = min(PAGE_SIZE, ctx_lens[b] - pos)
+            page = int(pages[b][pos // PAGE_SIZE])
+            cache[page, 0, :, :n] = keys[b][pos : pos + n].transpose(0, 1)
+            cache[page, 1, :, :n] = values[b][pos : pos + n].transpose(0, 1)
+    query = torch.randn(B, block_size, num_heads, HEAD_DIM, device=dev).bfloat16()
+    block_key = torch.stack([k[-block_size:] for k in keys])
+    block_value = torch.stack([v[-block_size:] for v in values])
+    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=dev)
+    ops.append_paged_kv_cache(
+        append_key=block_key.flatten(0, 1),
+        append_value=block_value.flatten(0, 1),
+        batch_indices=torch.arange(B, dtype=torch.int32, device=dev).repeat_interleave(block_size),
+        positions=(seq_lens_t.view(-1, 1) - block_size + torch.arange(block_size, device=dev))
+        .flatten()
+        .int(),
+        paged_kv_cache=cache,
+        kv_indices=block_tables.flatten(),
+        kv_indptr=torch.arange(0, (B + 1) * max_pages, max_pages, dtype=torch.int32, device=dev),
+        kv_last_page_len=(seq_lens_t - 1) % PAGE_SIZE + 1,
+        kv_layout="HND",
+    )
+    drafter = DFlashForCausalLM.__new__(DFlashForCausalLM)
+    torch.nn.Module.__init__(drafter)
+    drafter._dflash_trtllm_gen_ops = ops
+    drafter._dflash_trtllm_gen_device = drafter._dflash_trtllm_gen_sm_count = None
+    drafter._dflash_trtllm_gen_workspace = drafter._dflash_trtllm_gen_counters = None
+    with patch.object(DFlashForCausalLM, "_prepare_dflash_index_buffers"):
+        drafter._prepare_dflash_trtllm_gen_buffers(
+            torch.bfloat16, dev, B, block_size, num_heads, num_kv_heads, HEAD_DIM
+        )
+
+    w = window - 1
+    out = torch.empty(B * block_size, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev)
+    _, lse = ops.batch_decode_with_kv_cache(
+        query=query.flatten(0, 1),
+        kv_cache=(cache[:, 0], cache[:, 1]),
+        workspace_buffer=drafter._dflash_trtllm_gen_workspace,
+        block_tables=block_tables,
+        seq_lens=seq_lens_t,
+        max_seq_len=max_pages * PAGE_SIZE,
+        bmm1_scale=HEAD_DIM**-0.5,
+        bmm2_scale=1.0,
+        window_left=w,
+        out=out,
+        sinks=None,
+        enable_pdl=False,
+        kv_layout="HND",
+        backend="trtllm-gen",
+        q_len_per_req=block_size,
+        max_q_len=None,
+        cum_seq_lens_q=None,
+        kv_cache_sf=None,
+        uses_shared_paged_kv_idx=True,
+        bmm1_scale_log2=None,
+        return_lse=True,
+        multi_ctas_kv_counter_buffer=drafter._dflash_trtllm_gen_counters,
+    )
+    dflash_merge_later_block_keys_(out, lse, query, block_key, block_value, HEAD_DIM**-0.5, w)
+
+    group = num_heads // num_kv_heads
+    for b, c in enumerate(ctx_lens):
+        pos_q = c + torch.arange(block_size, device=dev).view(-1, 1)
+        pos_k = torch.arange(c + block_size, device=dev).view(1, -1)
+        logits = torch.einsum(
+            "jhd,lhd->hjl", query[b].float(), keys[b].float().repeat_interleave(group, 1)
+        )
+        logits = logits.masked_fill((pos_k - pos_q).abs() > w, float("-inf")) * HEAD_DIM**-0.5
+        expected = torch.einsum(
+            "hjl,lhd->jhd", logits.softmax(-1), values[b].float().repeat_interleave(group, 1)
+        )
+        got = out[b * block_size : (b + 1) * block_size].float()
+        torch.testing.assert_close(got, expected, atol=2e-2, rtol=2e-2)
+
+
+@pytest.mark.parametrize("block_size", [3, 7])
+@pytest.mark.parametrize("heads", [(16, 1), (64, 4)])
+def test_non_causal_window_matches_dense_attention(block_size, heads):
+    torch.manual_seed(block_size)
+    window = 64
+    ctx_lens = [
+        1,
+        window - block_size - 1,
+        window - block_size,
+        window - 1,
+        window,
+        window + 1,
+        3 * window + 17,
+    ]
+    _check_block_attention(ctx_lens, block_size, *heads, window)
+
+
+def test_workspace_holds_the_lse_stats():
+    """At 64 heads, 512 requests and 3 block tokens, the decode's softmax stats alone outgrow the
+    plain trtllm-gen workspace, which flashinfer would refuse."""
+    torch.manual_seed(0)
+    _check_block_attention([1 + (97 * b) % 300 for b in range(512)], 3, 64, 4, 4096)

From 9e0edeea364aefa4705ddfd30de1c8259ea698c1 Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@nvidia.com>
Date: Tue, 27 Oct 2026 18:25:42 -0700
Subject: [PATCH 2/2] [None][feat] DSpark ancestral drafting and block
 verification for rejection sampling

The DSpark Markov head biases block position k by the token at position
k-1. With use_rejection_sampling the worker chained the bias through the
argmax and then sampled each position on its own, so once a draw left the
argmax, later proposals were conditioned on tokens the block does not
contain: lossless, but with lower acceptance from position 2 on.

Sampled DSpark drafts are now drawn ancestrally: position k from
softmax((base_k + bias(d_{k-1})) / T), d_{k-1} being the token drawn before
it, with the temperature only (vLLM's probabilistic drafting), and that
distribution is stored as its proposal. Greedy requests keep the greedy
sentinel temperature and draft the argmax chain.

Add rejection_sample_method ("standard" | "block", prototype, default
"standard"). "block" replaces the chain kernel of the one-model rejection
path with block verification (Sun et al., 2024, arXiv:2403.10444): lossless,
equal to the token-level rule for one draft, and the rule vLLM serves as
rejection_sample_method="block". It runs on static shapes without host
synchronization inside the CUDA graphs; its uniforms are a counter hash of
each request's acceptance-slot (seed, offset), and the emitted token is
drawn by the seeded sampler. Ragged (confidence-scheduled) steps keep
token-level verification.

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
---
 .../_torch/speculative/block_verification.py  | 105 ++++++++++++++
 tensorrt_llm/_torch/speculative/dflash.py     |  16 ++-
 tensorrt_llm/_torch/speculative/dspark.py     |  73 +++++++++-
 tensorrt_llm/_torch/speculative/interface.py  |  44 +++++-
 tensorrt_llm/_torch/speculative/utils.py      |   1 +
 tensorrt_llm/llmapi/llm_args.py               |  18 +++
 .../usage/llm_args_golden_manifest.json       |   9 ++
 .../hw_agnostic/test_dspark_sampling.py       | 128 ++++++++++++++++++
 .../hw_agnostic/test_ragged_spec_sampler.py   |   2 +
 9 files changed, 384 insertions(+), 12 deletions(-)
 create mode 100644 tensorrt_llm/_torch/speculative/block_verification.py
 create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling.py

diff --git a/tensorrt_llm/_torch/speculative/block_verification.py b/tensorrt_llm/_torch/speculative/block_verification.py
new file mode 100644
index 000000000000..236430726a05
--- /dev/null
+++ b/tensorrt_llm/_torch/speculative/block_verification.py
@@ -0,0 +1,105 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Block verification (Sun et al., 2024, arXiv:2403.10444) for one-model speculative decoding.
+
+For drafts ``x_j ~ q_j`` (``j < K``), target rows ``p_j`` (bonus row ``p_K``) and uniforms ``u_j``:
+``w_j = min(1, w_{j-1} p_j(x_j) / q_j(x_j))``, ``r_j = sum_x max(w_j p_{j+1}(x) - q_{j+1}(x), 0)``,
+``h_j = r_j / (r_j + 1 - w_j)`` and ``h_{K-1} = w_{K-1}``. ``L`` is the largest ``j + 1`` with
+``u_j <= h_j`` (else 0); the emitted token is drawn from ``max(w_{L-1} p_L - q_L, 0)``, or ``p_K`` if ``L = K``.
+"""
+
+from typing import Callable
+
+import torch
+
+_MASK32 = 0xFFFFFFFF
+
+
+def _mul_u32(values: torch.Tensor, constant: int) -> torch.Tensor:
+    """``(values * constant) mod 2**32``, multiplying 16-bit halves so int64 never overflows."""
+    low = values * (constant & 0xFFFF)
+    high = ((values * (constant >> 16)) & 0xFFFF) << 16
+    return (low + high) & _MASK32
+
+
+def _mix_u32(values: torch.Tensor) -> torch.Tensor:
+    """The lowbias32 finalizer."""
+    values = values ^ (values >> 16)
+    values = _mul_u32(values, 0x7FEB352D)
+    values = values ^ (values >> 15)
+    values = _mul_u32(values, 0x846CA68B)
+    return values ^ (values >> 16)
+
+
+def counter_uniforms(seed: torch.Tensor, offset: torch.Tensor, num_draws: int) -> torch.Tensor:
+    """``[rows, num_draws]`` float64 uniforms in (0, 1), hashed from each row's ``(seed, offset)``.
+
+    The same key gives the same values on every rank and in every graph replay; a new offset
+    gives fresh ones.
+    """
+    seed = seed.to(torch.int64).unsqueeze(-1)
+    offset = offset.to(torch.int64).unsqueeze(-1)
+    draws = torch.arange(num_draws, dtype=torch.int64, device=seed.device)
+    key = _mix_u32((seed & _MASK32) ^ 0x5BD1E995)
+    key = _mix_u32(key ^ ((seed >> 32) & _MASK32))
+    key = _mix_u32(key ^ (offset & _MASK32))
+    key = _mix_u32(key ^ ((offset >> 32) & _MASK32))
+    bits = _mix_u32(_mix_u32(key ^ draws))
+    return (bits.to(torch.float64) + 0.5) * 2.0**-32
+
+
+def block_verify(
+    draft_tokens: torch.Tensor,
+    draft_probs: torch.Tensor,
+    target_probs: torch.Tensor,
+    uniforms: torch.Tensor,
+    sample_emitted: Callable[[torch.Tensor], torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Verify ``[B, K]`` drafts; ``sample_emitted`` draws from ``[B, V]`` nonnegative weights.
+
+    Returns the token-level kernel's layout: ``[B, K + 1]`` int32 tokens (accepted drafts, the emitted
+    token, then -1) and ``[B]`` int32 counts including the emitted token. No host sync (graph-safe).
+    """
+    num_drafts = draft_tokens.shape[1]
+    vocab_size = target_probs.shape[-1]
+    tokens = draft_tokens.long().clamp(0, vocab_size - 1).unsqueeze(-1)
+    p_x = target_probs[:, :num_drafts].gather(-1, tokens).squeeze(-1).float()
+    q_x = draft_probs.gather(-1, tokens).squeeze(-1).float()
+    # A draft its proposal gives no mass (or a non-finite row) is never accepted.
+    log_ratio = torch.where(q_x > 0, p_x.log() - q_x.log(), torch.full_like(q_x, float("-inf")))
+    log_ratio = torch.nan_to_num(log_ratio, nan=float("-inf"))
+    log_w = torch.empty_like(log_ratio)
+    running = torch.zeros_like(log_ratio[:, 0])
+    for j in range(num_drafts):
+        running = torch.clamp(running + log_ratio[:, j], max=0.0)
+        log_w[:, j] = running
+    w = log_w.exp()
+    h = w.clone()
+    if num_drafts > 1:
+        residual = w[:, :-1, None] * target_probs[:, 1:num_drafts].float()
+        residual.sub_(draft_probs[:, 1:].float()).clamp_min_(0.0)
+        mass = residual.sum(-1)
+        denom = mass + 1.0 - w[:, :-1]
+        h[:, :-1] = torch.where(denom > 0, mass / denom, torch.ones_like(mass))
+
+    positions = torch.arange(1, num_drafts + 1, dtype=torch.int64, device=draft_tokens.device)
+    accepted_len = (positions * (uniforms.to(torch.float64) <= h.to(torch.float64))).amax(-1)
+    row = accepted_len.view(-1, 1, 1)
+    p_row = target_probs.gather(1, row.expand(-1, 1, vocab_size)).squeeze(1).float()
+    q_row = draft_probs.gather(1, row.clamp(max=num_drafts - 1).expand(-1, 1, vocab_size))
+    w_prev = torch.cat([torch.ones_like(w[:, :1]), w], dim=1).gather(1, accepted_len.view(-1, 1))
+    is_bonus = (accepted_len == num_drafts).unsqueeze(-1)
+    residual = torch.where(
+        is_bonus, p_row, (w_prev * p_row - q_row.squeeze(1).float()).clamp_min(0.0)
+    )
+    # Fall back to the target row if rounding leaves the residual without mass.
+    residual = torch.where(residual.sum(-1, keepdim=True) > 0, residual, p_row)
+    emitted = sample_emitted(residual).long().clamp(0, vocab_size - 1)
+
+    slots = torch.arange(num_drafts + 1, dtype=torch.int64, device=draft_tokens.device)
+    drafts = torch.cat([draft_tokens.long(), torch.zeros_like(draft_tokens[:, :1]).long()], dim=1)
+    accepted_tokens = torch.where(
+        slots < accepted_len[:, None], drafts, torch.full_like(drafts, -1)
+    )
+    accepted_tokens = torch.where(slots == accepted_len[:, None], emitted[:, None], accepted_tokens)
+    return accepted_tokens.to(torch.int32), (accepted_len + 1).to(torch.int32)
diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py
index 8dd1765ed4e9..46eb3157dbeb 100644
--- a/tensorrt_llm/_torch/speculative/dflash.py
+++ b/tensorrt_llm/_torch/speculative/dflash.py
@@ -2478,11 +2478,8 @@ def _forward_impl(
                     )
                     vocab_size = gen_logits.shape[-1]
 
-                gen_draft_tokens = self.sample_draft_tokens(
-                    gen_logits,
-                    spec_metadata,
-                    batch_size,
-                    num_contexts=num_contexts,
+                gen_draft_tokens, vocab_size = self._draft_block_tokens(
+                    draft_model, gen_logits, inputs, spec_metadata, batch_size, num_contexts
                 )
 
                 # Opt-in confidence-calibration recording (env-gated). The
@@ -2597,6 +2594,15 @@ def _refine_block_logits(
         """
         return gen_logits
 
+    def _draft_block_tokens(
+        self, draft_model, gen_logits, inputs, spec_metadata, batch_size, num_contexts
+    ):
+        """The block's draft tokens and the vocab width they index."""
+        tokens = self.sample_draft_tokens(
+            gen_logits, spec_metadata, batch_size, num_contexts=num_contexts
+        )
+        return tokens, gen_logits.shape[-1]
+
     def _apply_dflash2_selector(
         self,
         draft_model,
diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py
index 322ef523a670..78445edd64fc 100644
--- a/tensorrt_llm/_torch/speculative/dspark.py
+++ b/tensorrt_llm/_torch/speculative/dspark.py
@@ -1251,6 +1251,10 @@ def _draft_slot_ids(
     def _try_fused_draft_head_and_refine(
         self, draft_model, hidden, attn_metadata, inputs, spec_metadata, num_gens, k
     ):
+        if self._samples_block_ancestrally(draft_model, spec_metadata):
+            return draft_model.logits_processor(
+                hidden, draft_model.lm_head, attn_metadata, True
+            ).reshape(num_gens, k, -1)
         # Only the M3.1 SM107 TP2 target's explicit A1 processor opts into this adapter.
         if not type(self).__module__.endswith("minimax_m31_nvfp4__sm_107__tp2.dspark"):
             return None
@@ -1262,6 +1266,70 @@ def _try_fused_draft_head_and_refine(
             self, draft_model, hidden, attn_metadata, inputs, spec_metadata, num_gens, k
         )
 
+    def _samples_block_ancestrally(self, draft_model, spec_metadata) -> bool:
+        """Sampled drafts follow the Markov head through the tokens actually drawn."""
+        return bool(
+            getattr(draft_model, "has_markov_head", False)
+            and spec_metadata.wants_advanced_draft_sampling
+        )
+
+    def _draft_block_tokens(
+        self, draft_model, gen_logits, inputs, spec_metadata, batch_size, num_contexts
+    ):
+        if self._samples_block_ancestrally(draft_model, spec_metadata):
+            return self._sample_block_ancestrally(
+                draft_model, gen_logits, inputs, spec_metadata, num_contexts, batch_size
+            )
+        return super()._draft_block_tokens(
+            draft_model, gen_logits, inputs, spec_metadata, batch_size, num_contexts
+        )
+
+    def _sample_block_ancestrally(
+        self, draft_model, gen_logits, inputs, spec_metadata, num_contexts, batch_size
+    ):
+        """Draw position k from ``softmax((base_k + bias(d_{k-1})) / T)``, ``d_{k-1}`` the token drawn
+        before it (the anchor for k = 0), with the temperature only, and store that distribution as
+        the proposal. Returns the int32 tokens and the stored vocab width."""
+        from ..models.modeling_speculative import dspark_markov_step_bias
+
+        num_gens, num_drafts, _ = gen_logits.shape
+        # The chain indexes the full Markov table; verification reads full-vocab q.
+        gen_logits = self.maybe_gather_sharded_draft_logits(gen_logits, spec_metadata)
+        vocab = gen_logits.shape[-1]
+        # Zero the CUDA-graph padding rows (uninitialized state) so every draw is legal.
+        gen_logits = self._zero_padding_rows(
+            gen_logits.reshape(num_gens * num_drafts, vocab),
+            spec_metadata,
+            num_contexts,
+            batch_size,
+            num_drafts,
+        ).reshape(num_gens, num_drafts, vocab)
+        params = self._temperature_only_params(
+            spec_metadata.request_temperatures[num_contexts:batch_size]
+        )
+        draft_probs = spec_metadata.draft_probs
+        slots = spec_metadata.batch_slot_ids[num_contexts:batch_size]
+        tokens = inputs["first_prev_tokens"]
+        drafts = []
+        for step in range(num_drafts):
+            bias = dspark_markov_step_bias(tokens, draft_model.markov_w1, draft_model.markov_w2)
+            seed, offset = self._rng_state_per_request(
+                spec_metadata, num_contexts, batch_size, slot=spec_metadata.rng_draft_slot(step)
+            )
+            tokens, probs = self._sample_from_logits_with_probs(
+                spec_metadata.advanced_sampling_mode,
+                gen_logits[:, step] + bias.to(gen_logits.dtype),
+                *params,
+                seed=seed,
+                offset=offset,
+            )
+            if draft_probs is not None:
+                draft_probs[slots, step, :vocab] = probs
+            drafts.append(tokens)
+        if draft_probs is not None:
+            spec_metadata.draft_probs_last_dim = vocab
+        return torch.stack(drafts, dim=1).to(torch.int32), vocab
+
     def _refine_block_logits(
         self,
         draft_model,
@@ -1295,8 +1363,9 @@ def _apply_dspark_markov_bias(
         logits, where prev_0 is the anchor (last accepted) token and prev_{i>0}
         is the greedy token from step i-1's biased logits. Greedy per-position
         argmax of the returned logits therefore reproduces the reference
-        sampled chain; the rejection-sampling path samples from the same
-        biased distributions (proposal conditioned on the greedy chain).
+        chain. Sampled drafts do not come through here: they are drawn by
+        :meth:`_sample_block_ancestrally`, which conditions each position on
+        the token actually drawn before it.
 
         Handles a TP vocab-sharded draft lm_head by slicing markov_w2's rows
         to this rank's contiguous shard and chaining through the TP-aware
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index c188d5d231cf..4f3344b52aa9 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -49,6 +49,7 @@
 from ..pyexecutor.sampler.ops import custom as fused_sampling
 from ..pyexecutor.sampler.ops import flashinfer as flashinfer_sampling
 from ..pyexecutor.sampler.ops.vanilla import greedy_search_sampling_batch
+from .block_verification import block_verify, counter_uniforms
 
 # Logit given to token 0 of a CUDA-graph padding row (see
 # SpecWorkerBase._zero_padding_rows). Finite in
@@ -733,6 +734,9 @@ def is_ragged_verify(self) -> bool:
     # Drop the sampler tail's no-op bookkeeping launches, bitwise the same
     # outputs (RuntimeHooks.spec_tail_bookkeep; see tail_bookkeep_draft_len).
     tail_bookkeep: bool = False
+    # Verification rule of the rejection path, "standard" or "block" (deploy-time;
+    # from DecodingBaseConfig.rejection_sample_method).
+    rejection_sample_method: str = "standard"
     # Advanced-sampling specialization (deploy-time; from DecodingBaseConfig.advanced_sampling_mode).
     advanced_sampling_mode: AdvancedSamplingMode = AdvancedSamplingMode.FULL
     # Whether the occurrence penalties are enabled (deploy-time; from
@@ -2695,6 +2699,17 @@ def _uniform_padding_probs(probs: torch.Tensor,
             probs.masked_fill_(is_padding.unsqueeze(-1), 1.0 / probs.shape[-1])
         return probs
 
+    @staticmethod
+    def _temperature_only_params(
+        temperatures: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Sampler arguments for ``temperatures`` without filters (device-side fills, graph-safe)."""
+        return (temperatures,
+                torch.full_like(temperatures,
+                                torch.iinfo(torch.int32).max,
+                                dtype=torch.int32),
+                torch.ones_like(temperatures), torch.zeros_like(temperatures))
+
     def advanced_sample_draft(self,
                               logits: torch.Tensor,
                               spec_metadata: "SpecMetadata",
@@ -3052,19 +3067,38 @@ def _sample_and_accept_draft_tokens_rejection(
                 batch_size,
                 slot=spec_metadata.rng_accept_slot)
 
+            use_block = spec_metadata.rejection_sample_method == "block"
             slot_out = None
             if slot_rows is not None:
-                slot_out = _CHAIN_SPEC_SLOT_OVERRIDE(slot_rows[0], slot_rows[1],
-                                                     runtime_draft_len,
-                                                     full_draft_tokens,
-                                                     target_probs, seed, offset)
+                if not use_block:
+                    slot_out = _CHAIN_SPEC_SLOT_OVERRIDE(
+                        slot_rows[0], slot_rows[1], runtime_draft_len,
+                        full_draft_tokens, target_probs, seed, offset)
                 if slot_out is None:
-                    # Declined: the gathered copy, as without the override.
+                    # Declined, or block verification: the gathered copy, as
+                    # without the override.
                     full_draft_probs = slot_rows[0][
                         slot_rows[1], :runtime_draft_len, :vocab_size]
             if slot_out is not None:
                 gen_accepted, gen_num_accepted = slot_out
+            elif use_block and not is_ragged:
+                unit = torch.ones_like(seed, dtype=torch.float32)
+                gen_accepted, gen_num_accepted = block_verify(
+                    full_draft_tokens, full_draft_probs, target_probs,
+                    counter_uniforms(seed, offset, runtime_draft_len),
+                    lambda weights: self._sample_from_logits(
+                        spec_metadata.advanced_sampling_mode,
+                        weights.log(),
+                        *self._temperature_only_params(unit),
+                        seed=seed,
+                        offset=offset))
             else:
+                if use_block:
+                    logger.warning_once(
+                        "rejection_sample_method='block' does not cover ragged "
+                        "(confidence-scheduled) verification; these steps use "
+                        "token-level rejection sampling.",
+                        key="block_verification_ragged_fallback")
                 gen_accepted, gen_num_accepted = rejection_sampling_one_model(
                     draft_probs=full_draft_probs,
                     draft_token_ids=full_draft_tokens,
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index 2bb15ed1d727..55ab8a890275 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -359,6 +359,7 @@ def get_spec_metadata(spec_config,
         # dispatcher routes to a backend that takes no min_p argument -- so the filter is
         # dropped with nothing raised. One assignment here keeps them in step.
         metadata.advanced_sampling_mode = spec_config.advanced_sampling_mode
+        metadata.rejection_sample_method = spec_config.rejection_sample_method
         if num_seq_slots is not None:
             metadata.num_seq_slots = num_seq_slots
     return metadata
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index 41f3c3976f73..66366bf53f8e 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -2367,6 +2367,16 @@ class DecodingBaseConfig(StrictBaseModel):
         "(default) to use exact-match verification on non-greedy batches. "
         "The non-dynamic-tree one-model path requires FlashInfer.")
 
+    rejection_sample_method: Literal["standard", "block"] = Field(
+        default="standard",
+        status="prototype",
+        description=
+        "Verification rule of rejection sampling (use_rejection_sampling=True): "
+        "'standard' is token-level rejection sampling; 'block' is block "
+        "verification (Sun et al., 2024, arXiv:2403.10444), which is lossless "
+        "and accepts at least as many tokens. Linear-chain one-model paths "
+        "only; ragged (confidence-scheduled) steps use 'standard'.")
+
     allow_advanced_sampling: bool = Field(
         default=False,
         status="deprecated",
@@ -7241,6 +7251,14 @@ def validate_speculative_config(self):
                 # DraftTarget one-model) share the SpecMetadata slot-indexed
                 # dispatch.
 
+            if self.speculative_config.rejection_sample_method == "block" and (
+                    not self.speculative_config.use_rejection_sampling
+                    or getattr(self.speculative_config, "use_dynamic_tree",
+                               False)):
+                raise ValueError(
+                    "rejection_sample_method='block' requires "
+                    "use_rejection_sampling=True and no dynamic tree.")
+
             if isinstance(self.speculative_config, PARDDecodingConfig):
                 assert self.speculative_config.max_draft_len > 0, "PARD max_draft_len must be > 0"
 
diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json
index 4b6941b8f505..57f2f51678dd 100644
--- a/tensorrt_llm/usage/llm_args_golden_manifest.json
+++ b/tensorrt_llm/usage/llm_args_golden_manifest.json
@@ -1784,6 +1784,15 @@
       "kind": "value",
       "path": "speculative_config.num_nextn_predict_layers"
     },
+    {
+      "allowed_values": [
+        "standard",
+        "block"
+      ],
+      "capture_policy": "literal",
+      "kind": "categorical",
+      "path": "speculative_config.rejection_sample_method"
+    },
     {
       "capture_policy": "float",
       "kind": "value",
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling.py
new file mode 100644
index 000000000000..870e76cf3c07
--- /dev/null
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling.py
@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""DSpark ancestral drafting and block verification under rejection sampling."""
+
+import math
+import types
+from unittest.mock import patch
+
+import torch
+
+from tensorrt_llm._torch.models.modeling_speculative import dspark_markov_step_bias
+from tensorrt_llm._torch.speculative.block_verification import block_verify
+from tensorrt_llm._torch.speculative.dspark import DSparkWorker
+from tensorrt_llm._torch.speculative.interface import SpecWorkerBase
+
+VOCAB, K = 3, 3
+
+
+def _prefix_model(generator, sharpness):
+    """A next-token distribution for every prefix of length 0..K."""
+    num_prefixes = sum(VOCAB**m for m in range(K + 1))
+    logits = sharpness * torch.randn(num_prefixes, VOCAB, generator=generator, dtype=torch.float64)
+    return torch.softmax(logits, -1)
+
+
+def _rows(model, prefixes):
+    ids = torch.zeros(prefixes.shape[0], dtype=torch.int64)
+    for m in range(prefixes.shape[1]):
+        ids = ids * VOCAB + prefixes[:, m]
+    return model[ids + sum(VOCAB**m for m in range(prefixes.shape[1]))]
+
+
+def test_block_verification_is_lossless():
+    """Blocks verified, then completed by sampling the target, are distributed as the target's own samples."""
+    generator = torch.Generator().manual_seed(0)
+    target, draft = _prefix_model(generator, 1.5), _prefix_model(generator, 1.0)
+    num = 200_000
+    drafts = torch.empty(num, 0, dtype=torch.int64)
+    draft_rows = []
+    for _ in range(K):
+        draft_rows.append(_rows(draft, drafts))
+        drafts = torch.cat([drafts, torch.multinomial(draft_rows[-1], 1, generator=generator)], 1)
+    target_rows = [_rows(target, drafts[:, :j]) for j in range(K + 1)]
+    tokens, num_accepted = block_verify(
+        drafts,
+        torch.stack(draft_rows, 1).float(),
+        torch.stack(target_rows, 1).float(),
+        torch.rand(num, K, generator=generator, dtype=torch.float64),
+        lambda weights: torch.multinomial(weights, 1, generator=generator).squeeze(-1),
+    )
+
+    sequences = torch.zeros(num, K + 1, dtype=torch.int64)
+    for m in range(K + 1):
+        continued = torch.multinomial(_rows(target, sequences[:, :m]), 1, generator=generator)
+        sequences[:, m] = torch.where(
+            m < num_accepted.long(), tokens[:, m].long(), continued.squeeze(-1)
+        )
+    grid = torch.cartesian_prod(*[torch.arange(VOCAB)] * (K + 1))
+    expected = torch.full((grid.shape[0],), float(num), dtype=torch.float64)
+    for m in range(K + 1):
+        expected *= _rows(target, grid[:, :m]).gather(-1, grid[:, m : m + 1]).squeeze(-1)
+    ids = (sequences * VOCAB ** torch.arange(K, -1, -1)).sum(-1)
+    observed = torch.bincount(ids, minlength=grid.shape[0]).double()
+    big = expected >= 5  # pool the rare sequences
+    observed = torch.cat([observed[big], observed[~big].sum().view(1)])
+    expected = torch.cat([expected[big], expected[~big].sum().view(1)])
+    stat = ((observed - expected) ** 2 / expected).sum().item()
+    dof = observed.numel() - 1
+    critical = dof * (1 - 2 / (9 * dof) + 3.719 * math.sqrt(2 / (9 * dof))) ** 3  # p = 1e-4
+    assert stat < critical, f"chi-square {stat:.1f} >= {critical:.1f}"
+
+
+def test_ancestral_drafts_condition_each_position_on_the_drawn_token():
+    """Position k is drawn, with the temperature only, from base_k plus the Markov bias of the token drawn before it;
+    the distribution it was drawn from is stored as its proposal."""
+    generator = torch.Generator().manual_seed(1)
+    num_gens, vocab, rank = 4, 32, 8
+    drafter = types.SimpleNamespace(
+        has_markov_head=True,
+        markov_w1=torch.randn(vocab, rank, generator=generator),
+        markov_w2=torch.randn(vocab, rank, generator=generator),
+    )
+    base = torch.randn(num_gens, K, vocab, generator=generator)
+    anchors = torch.randint(vocab, (num_gens,), generator=generator)
+    meta = types.SimpleNamespace(
+        request_temperatures=torch.full((num_gens,), 0.7),
+        request_seeds=torch.arange(num_gens, dtype=torch.int64),
+        request_offsets=torch.zeros(num_gens, dtype=torch.int64),
+        rng_draft_slot=lambda step: K + 1 + step,
+        batch_slot_ids=torch.arange(num_gens).flip(0),
+        draft_probs=torch.zeros(num_gens, K, vocab),
+        draft_probs_last_dim=0,
+        dummy_slot_row=0,
+        advanced_sampling_mode=None,
+        is_all_greedy_sample=False,
+        wants_advanced_draft_sampling=True,
+        draft_vocab_size=vocab,
+        vocab_size=vocab,
+    )
+    calls = []
+
+    def sampler(mode, logits, temperatures, top_ks, top_ps, min_ps, *, seed, offset):
+        probs = torch.softmax(logits / temperatures[:, None], -1)
+        calls.append((logits, probs, top_ks, top_ps, min_ps))
+        return torch.multinomial(probs, 1, generator=generator).squeeze(-1).int(), probs
+
+    worker = DSparkWorker.__new__(DSparkWorker)
+    torch.nn.Module.__init__(worker)
+    worker.mapping = None
+    assert worker._samples_block_ancestrally(drafter, meta)
+    with patch.object(SpecWorkerBase, "_sample_from_logits_with_probs", staticmethod(sampler)):
+        tokens, stored_vocab = worker._sample_block_ancestrally(
+            drafter, base, {"first_prev_tokens": anchors}, meta, 0, num_gens
+        )
+
+    assert tokens.dtype == torch.int32 and stored_vocab == vocab == meta.draft_probs_last_dim
+    prev = anchors
+    for step, (logits, probs, top_ks, top_ps, min_ps) in enumerate(calls):
+        bias = dspark_markov_step_bias(prev, drafter.markov_w1, drafter.markov_w2)
+        torch.testing.assert_close(logits, base[:, step] + bias)
+        assert (
+            (top_ks == torch.iinfo(torch.int32).max).all()
+            and (top_ps == 1).all()
+            and (min_ps == 0).all()
+        )
+        torch.testing.assert_close(meta.draft_probs[meta.batch_slot_ids, step], probs)
+        prev = tokens[:, step]
+    assert len(calls) == K
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py
index 0249d62b4efa..ba8383283466 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py
@@ -362,6 +362,7 @@ def test_rejection_preserves_sampling_params_and_caps_acceptance(monkeypatch, ra
     widths = [2, 4] if ragged else [4, 4]
     metadata = _penalty_mapping_meta([4, 2, 9], widths if ragged else None)
     metadata.rng_accept_slot = 2 * worker.max_draft_len + 1
+    metadata.rejection_sample_method = "standard"
     metadata.dummy_slot_row = 9
     rows = 1 + sum(widths)
     metadata.temperatures = torch.arange(rows, dtype=torch.float32) + 0.5
@@ -507,6 +508,7 @@ def rejection(**kwargs):
         lens = torch.tensor(local_windows * trials, dtype=torch.int32)
         metadata = _penalty_mapping_meta([5, *slots.tolist()], lens.tolist() if ragged else None)
         metadata.rng_accept_slot = 2 * worker.max_draft_len + 1
+        metadata.rejection_sample_method = "standard"
         metadata.dummy_slot_row = 9
         metadata.runtime_draft_len = width
         metadata.enable_penalty = False