| 728 | ) -> torch.Tensor: |
| 729 | h_p = hidden_states @ self.rand_mat # [n_hidden_states, k] |
| 730 | return h_p @ self.w_p.T # [n_hidden_states, V] |
| 731 | |
| 732 | def rrt(self) -> torch.Tensor: |
| 733 | """Return R @ Rᵀ, which should be close to the identity matrix.""" |
| 734 | m = self.rand_mat |
| 735 | return m @ m.T |
| 736 | |
| 737 | |
| 738 | def optimal_k(n: int, epsilon: float) -> int: |
| 739 | """Source: https://cs.stanford.edu/people/mmahoney/cs369m/Lectures/lecture1.pdf""" |
| 740 | k_float = 24 * math.log(n, math.e) / (3 * epsilon**2 - 2 * epsilon**3) |
| 741 | return int(math.ceil(k_float)) |
| 742 | |
| 743 | |
| 744 | def get_sampler(provider: str, weights: torch.Tensor) -> Sampler: |
| 745 | match provider: |
| 746 | case S.fused_triton: |
| 747 | return SimpleSampler(lambda **kwargs: fused_mm_sample_triton(**{"seed": 0, **kwargs})) |
| 748 | case S.fused_triton_p2p_no_overlap: |
| 749 | return SimpleSampler( |
| 750 | lambda **kwargs: fused_mm_sample_triton( |
| 751 | **{"seed": 0, "p2p_no_overlap": True, **kwargs} |
| 752 | ) |
| 753 | ) |
| 754 | case S.fused_triton_ret_logits: |
| 755 | return SimpleSampler( |
| 756 | lambda **kwargs: fused_mm_sample_triton( |
| 757 | **{"seed": 0, "return_logits": True, **kwargs} |
| 758 | )[0] |
| 759 | ) |
| 760 | case S.fused_triton_greedy: |
| 761 | return SimpleSampler( |
| 762 | lambda **kwargs: fused_mm_sample_triton( |
| 763 | **{"seed": 0, "greedy_sampling": True, **kwargs} |
| 764 | ) |
| 765 | ) |
| 766 | case S.naive_pt: |
| 767 | return SimpleSampler(sample) |
| 768 | case S.naive_compiled: |
| 769 | return SimpleSampler(sample_compiled) |
| 770 | case S.pt_qitra: |
| 771 | return SimpleSampler(lambda **kwargs: sample(**kwargs, use_qitra=True)) |
| 772 | case S.sequential_compiled: |
| 773 | return SimpleSampler(sequential_sample_pt) |
| 774 | case S.naive_tl_matmul: |
| 775 | return SimpleSampler(lambda **kwargs: sample_compiled(**kwargs, tl_matmul=True)) |
| 776 | case S.jl_compiled: |
| 777 | return JLSampler.from_weights(weights) |
| 778 | case S.fused_topk: |
| 779 | from .tl_fused_mm_topk import fused_mm_topk_and_sample |
| 780 | |
| 781 | return SimpleSampler(fused_mm_topk_and_sample) |
| 782 | case S.flashinfer_top_k_top_p_sampling_from_logits: |
| 783 | return SimpleSampler( |
| 784 | lambda **kwargs: flashinfer_top_k_top_p_sampling_from_logits( |