Fix TopP sampling dropping the threshold-crossing move

[Dash1971]

When Temperature > 0 and TopP < 1, the sampler currently drops the move that would bring the retained probability mass up to the threshold. For a policy of [0.6, 0.3, 0.1] and TopP=0.8, it samples only the first move, even though that retains just 60% of the probability mass. This PR retains the first two moves and samples them with probabilities [2/3, 1/3].

Reproduction

From the repository root, with the normal runtime dependencies installed:

from unittest.mock import patch
import torch
from maia3.uci import sample_from_logits

logits = torch.tensor([0.6, 0.3, 0.1]).log()
with patch("maia3.uci.torch.multinomial", return_value=torch.tensor([0])) as draw:
    sample_from_logits(logits, temperature=1.0, top_p=0.8)
    print(draw.call_args.args[0])

• Before: tensor([1.]) — only the most likely move can be chosen.
• After: approximately tensor([0.6667, 0.3333]) — the smallest prefix reaching the threshold is retained and renormalized.

The random draw is intercepted only to inspect the exact probabilities supplied to it; softmax, temperature scaling, sorting, filtering, and normalization use the real sampler and PyTorch. No model checkpoint is needed.

Change

The current cumulative <= top_p mask checks the mass including each candidate and therefore excludes the threshold-crossing candidate. Check the mass before each candidate instead, while always retaining the top move. If a prefix reaches the threshold exactly, no extra move is included. The README now states this behavior explicitly.

This can broaden the sampled move distribution for affected settings, including cases that previously collapsed to a single move. TopP=1 and Temperature=0 behavior remains unchanged. There are no changes to model weights, inference, or the Elo inputs. The legacy code/uci.py launcher imports the same implementation, so it benefits from the same fix.

Validation

python -m unittest discover -s tests -v
git diff --check

Tested on macOS CPU with Python 3.12.13 and PyTorch 2.12.0. Five of the eleven test methods fail against the unchanged upstream implementation at 1e13597c42d4858b7cfd7cfdae01e297263364b2; all eleven pass with this fix. The distribution tests cover float32 and float64, threshold crossings (including the final move), exact boundaries, the top-one fallback, temperature scaling, original vocabulary indices, masked illegal moves, a single legal move, and unchanged default/argmax behavior. Each expected categorical outcome is checked deterministically rather than relying on random-frequency assertions.