Function verify_greedy_tp2

src/fused_mm_sampling/testing.py:190–208

Worker function for TP2 greedy tests (passed to run_maybe_distributed).

()

def verify_greedy_tp2() -> None:
    """Worker function for TP2 greedy tests (passed to run_maybe_distributed)."""
    tp = TPInfo.from_world()
    for vocab_size, n_hidden_states in product([100, 256, 512], [1, 2]):
        inputs = make_synthetic_inputs(
            vocab_size=vocab_size, n_hidden_states=n_hidden_states, tp=tp
        )
        sampler = get_sampler("greedy", weights=inputs.weights)
        sampler.prepare()
        samples = sampler.sample(
            weights=inputs.weights,
            hidden_states=inputs.hidden_states,
            num_samples=1,
            temperature=torch.empty((), device="cuda"),
            tp=tp,
        )
        expected = inputs.logits.argmax(dim=-1)
        torch.testing.assert_close(samples[:, 0], expected)
        tp.rank0_print(f"✅ Passed: greedy V={vocab_size} H={n_hidden_states}")


def shift_logits_negative(