Description

When using speculative decoding / drafting (specifically with Qwen 27B Optimized Speed model), inference occasionally crashes with:
ValueError: SparseDistribution probabilities must have positive mass

This happens due to:

1. Numerical instability or masking: Extreme logits or masked positions in the model can result in all -inf values, which lead to NaN probabilities.
2. Residual distribution logic: In residual_distribution, the difference of probabilities can sum to exactly 0.0 or a non-finite value if a value is NaN. If total evaluates to NaN, the check total <= 0 evaluates to False, bypassing target fallback logic and constructing a SparseDistribution with all NaN probabilities, which crashes the inference runner.

Suggested Fix

1. Fallback gracefully to a one-hot distribution on the first available token (greedy argmax) inside the SparseDistribution constructor if the probability mass is zero or non-finite.
2. Add strict not np.isfinite(total) checks to the fallback paths in residual_distribution to prevent NaN values from bypassing target fallbacks.

I have implemented this fix and verified it with unit tests. I will submit a PR shortly.