From 06cfc905971ed09be1705d4eb849a6c12a8123bd Mon Sep 17 00:00:00 2001
From: Zhiren Xiao <104724357+Jackxiaozhiren@users.noreply.github.com>
Date: Mon, 26 Oct 2026 16:17:33 +0800
Subject: [PATCH 1/2] docs(mtp_varlen): state zero-rate participation at the
 ranking site

---
 .../layers/tilelang_ops/schedule_mtp_varlen.py               | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/deepinfer/model_executor/layers/tilelang_ops/schedule_mtp_varlen.py b/deepinfer/model_executor/layers/tilelang_ops/schedule_mtp_varlen.py
index 2dc0bdb..0cebef0 100644
--- a/deepinfer/model_executor/layers/tilelang_ops/schedule_mtp_varlen.py
+++ b/deepinfer/model_executor/layers/tilelang_ops/schedule_mtp_varlen.py
@@ -69,6 +69,11 @@ def compute_seqlens_q_by_confidence_rates_ref(
 
     # Phase 2: rank-based selection
     # Flatten all (seq, pos) rates and find the top-k by rate (higher wins, tie-break by flat index ascending)
+    # Cells with cumprod 0.0 take part in this ranking too: remap == -1 (fresh seq) and a seq whose draft
+    # confidences were masked to 0 after emitting a stop token (model_executor/layers/sampler.py) both
+    # contribute all-zero rows. Those rows are reached once extra_budget exceeds the number of positive
+    # cells, and are then filled in flat-index order, so which zero-confidence seq is verified beyond its
+    # base token depends on batch position rather than on confidence.
     extra_budget = total_q - num_seqs  # number of extra tokens to distribute
 
     flat_rates = cum_rates.reshape(-1)  # [num_seqs * nextn]

From c928039a2161f23fad220b95d6d721c17a06af16 Mon Sep 17 00:00:00 2001
From: Zhiren Xiao <104724357+Jackxiaozhiren@users.noreply.github.com>
Date: Mon, 26 Oct 2026 16:18:59 +0800
Subject: [PATCH 2/2] test(mtp_varlen): pin zero-rate slack in confidence-based
 seqlens_q allocation

---
 .../tilelang_ops/schedule_mtp_varlen.py       | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/deepinfer/model_executor/layers/tilelang_ops/schedule_mtp_varlen.py b/deepinfer/model_executor/layers/tilelang_ops/schedule_mtp_varlen.py
index 0cebef0..7869372 100644
--- a/deepinfer/model_executor/layers/tilelang_ops/schedule_mtp_varlen.py
+++ b/deepinfer/model_executor/layers/tilelang_ops/schedule_mtp_varlen.py
@@ -144,6 +144,30 @@ def test_compute_seqlens_q_by_confidence_rates(
     assert torch.equal(seqlens_q, seqlens_q_ref)
 
 
+@pytest.mark.ci
+def test_compute_seqlens_q_by_confidence_rates_zero_rate_slack() -> None:
+    """Zero-confidence rows absorb budget slack in flat-index order.
+
+    num_seqs=2, nextn=2, total_q=5 gives extra_budget=3, but only the remapped row has positive
+    cumprod entries ([0.5, 0.1]); the fresh row (remap == -1) is all zeros. The two positives take
+    ranks 0 and 1 and rank 2 goes to the lowest flat index among the zeros, so the fresh row is
+    verified with 2 tokens and the confident row with 3 — pinning the behavior the comment above
+    describes.
+    """
+    total_q = 5
+    prev_confidence_rates = torch.tensor([[0.5, 0.2]], dtype=torch.float32, device="cuda")
+    remap = torch.tensor([-1, 0], dtype=torch.int32, device="cuda")
+    expected = torch.tensor([2, 3], dtype=torch.int32, device="cuda")
+
+    seqlens_q = torch.empty(2, dtype=torch.int32, device="cuda")
+    compute_seqlens_q_by_confidence_rates(total_q, remap, prev_confidence_rates, seqlens_q)
+    assert torch.equal(seqlens_q, expected)
+
+    seqlens_q_ref = torch.empty(2, dtype=torch.int32, device="cuda")
+    compute_seqlens_q_by_confidence_rates_ref(total_q, remap, prev_confidence_rates, seqlens_q_ref)
+    assert torch.equal(seqlens_q_ref, expected)
+
+
 @pytest.mark.benchmark
 @pytest.mark.parametrize("num_seqs", [64, 128, 256, 512, 1024])
 def test_compute_seqlens_q_benchmark(