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]