test(mtp_varlen): pin zero-rate slack in confidence-based seqlens_q allocation

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(