Function verify_greedy_tp2

src/fused_mm_sampling/testing.py:190–208

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

()

        chi_squared, p_value = chisquare(obs, exp)
        degrees_of_freedom = len(obs) - 1
        reduced_chi_squared = chi_squared / degrees_of_freedom
        if print_result:
            print(
                f"Chi-squared: provider={provider}, V={vocab_size}, H={n_hidden_states}, "
                f"samples={num_samples}, tested_bins={mask.sum()}, "
                f"tested_probability_mass={tested_probability_mass:.6f}, "
                f"statistic={chi_squared:.3f}, df={degrees_of_freedom}, "
                f"reduced_statistic={reduced_chi_squared:.6f}, p={p_value:.6g}"
            )
        assert not np.isnan(p_value), (
            f"Chi-squared returned NaN for seq {seq_idx} — likely all samples "
            f"landed in a single tile. {provider} may have a masked-fill bug."
        )
        assert p_value > 0.001, (
            f"Sampling distribution mismatch for seq {seq_idx}: p={p_value:.6f}. "
            f"{provider} does not match the expected softmax distribution."
        )


def assert_sampling_distribution_large_vocab(
    vocab_size: int = 32_768,
    num_samples: int = 1_000_000,