Function make_synthetic_inputs

src/fused_mm_sampling/testing.py:27–72

Build weights and hidden_states that produce known logits. Creates up to two hidden states: one with ascending logits (favors high token indices) and one with descending logits (favors low token indices). All logits are shifted negative via :func:`shift_logits_negative`.

(
    vocab_size: int = 256,
    hidden_size: int = 10,
    n_hidden_states: int = 2,
    device: torch.device = torch.device("cuda"),
    tp: TPInfo = TP1,
)

    vocab_size: int
    hidden_size: int


def make_synthetic_inputs(
    vocab_size: int = 256,
    hidden_size: int = 10,
    n_hidden_states: int = 2,
    device: torch.device = torch.device("cuda"),
    tp: TPInfo = TP1,
) -> SyntheticInputs:
    """Build weights and hidden_states that produce known logits.

    Creates up to two hidden states: one with ascending logits (favors high
    token indices) and one with descending logits (favors low token indices).
    All logits are shifted negative via :func:`shift_logits_negative`.
    """
    logits1 = torch.arange(-vocab_size / 2, vocab_size / 2, dtype=torch.float32)[None, :]
    logits2 = torch.arange(vocab_size / 2, -vocab_size / 2, step=-1, dtype=torch.float32)[None, :]
    all_logits = [logits1, logits2]
    logits = torch.cat(all_logits[:n_hidden_states], dim=0).to(device)
    n_hidden_states = logits.shape[0]

    U, _, _ = torch.linalg.svd(logits, full_matrices=False)  # noqa: N806

    torch.manual_seed(0)
    hidden_states = torch.cat(
        [U, torch.rand((n_hidden_states, hidden_size - n_hidden_states), device=device)],
        dim=1,
    ).to(device)
    weights = torch.linalg.pinv(hidden_states) @ logits  # [D, V]

    weights_bf16 = weights.bfloat16().T.contiguous()  # [V, D]
    hidden_states_bf16 = hidden_states.bfloat16()
    weights_bf16, hidden_states_bf16 = shift_logits_negative(
        weights_bf16,
        hidden_states_bf16,
        offset=float(vocab_size),
    )

    weights_bf16, hidden_states_bf16 = pad_to_tma_alignment(weights_bf16, hidden_states_bf16)
    weights_bf16 = shard_weights(weights_bf16, tp)

    return SyntheticInputs(
        weights=weights_bf16,
        hidden_states=hidden_states_bf16,
        logits=logits,
        vocab_size=vocab_size,
        hidden_size=weights_bf16.shape[1],
    )