Function shift_logits_negative

src/fused_mm_sampling/testing.py:211–238

Shift all logits by -offset without touching the existing weights. Appends a bias column so that ``h_new @ W_new^T = h @ W^T - offset``. Since softmax is shift-invariant the expected sampling distribution is unchanged, but the all-negative logits exercise masked-fill handling in par

(
    weights: torch.Tensor,
    hidden_states: torch.Tensor,
    offset: float,
)

def assert_sampling_distribution_large_vocab(
    vocab_size: int = 32_768,
    num_samples: int = 1_000_000,
    samples_per_call: int = 10_000,
    hidden_size: int = 16,
) -> None:
    """Verify FMMS sampling against random-Gaussian logits at realistic vocabulary size."""

    device = torch.device("cuda")
    torch.manual_seed(0)

    hidden_states = torch.randn(
        (1, hidden_size),
        dtype=torch.bfloat16,
        device=device,
    )
    weights = torch.randn(
        (vocab_size, hidden_size),
        dtype=torch.bfloat16,
        device=device,
    )
    weights.div_(math.sqrt(hidden_size))
    inputs = SyntheticInputs(
        weights=weights,
        hidden_states=hidden_states,
        logits=hidden_states.float() @ weights.float().T,
        vocab_size=vocab_size,
        hidden_size=hidden_size,
    )
    assert_sampling_distribution(
        provider="fused-triton",
        vocab_size=vocab_size,
        n_hidden_states=1,