Function verify_correctness_tp2

src/fused_mm_sampling/testing.py:162–170

()

        n_calls,
        desc="Sampling batches",
        disable=n_calls == 1,
    ):
        samples = sampler.sample(
            weights=inputs.weights,
            hidden_states=inputs.hidden_states,
            num_samples=samples_per_call,
            temperature=temperature,
            tp=tp,
            seed=tp.rank * num_samples + call_idx * samples_per_call,
        )
        flattened_sample_bins = (samples + sequence_offsets).flatten()
        empirical_counts += torch.bincount(