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 shift_logits_negative(
    weights: torch.Tensor,
    hidden_states: torch.Tensor,
    offset: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """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
    partial V-tiles (kernels must fill masked rows with -inf, not 0, or the
    0 beats all real negative values in the tile-max reduction).

    We use a bias column instead of baking the offset into the logits before
    computing the pseudoinverse because bf16 cannot represent fine-grained
    all-negative logits for vocab sizes above ~128.  The bias column keeps the
    original weights (centered near 0) intact and encodes the offset exactly.
    """
    vocab_size = weights.shape[0]
    n_hidden_states = hidden_states.shape[0]
    device = weights.device
    dtype = weights.dtype
    bias_w = torch.ones(vocab_size, 1, dtype=dtype, device=device)
    bias_h = torch.full((n_hidden_states, 1), -offset, dtype=dtype, device=device)
    return (
        torch.cat([weights, bias_w], dim=1),
        torch.cat([hidden_states, bias_h], dim=1),
    )