coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-30 15:57:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge c928039a2161f23fad220b95d6d721c17a06af16 into 553741557a57df44a46fb3ab72b6aafeaf945337

Browse Source
This commit is contained in:
Jackson 2026-10-26 16:21:17 +08:00 committed by GitHub
commit fe4746090e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 29 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
deepinfer/model_executor/layers/tilelang_ops/schedule_mtp_varlen.py
View File

@ -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]
@ -139,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(