Function get_sampler

src/fused_mm_sampling/core.py:730–781

(provider: str, weights: torch.Tensor)

def get_sampler(provider: str, weights: torch.Tensor) -> Sampler:
    match provider:
        case S.fused_triton:
            return SimpleSampler(lambda **kwargs: fused_mm_sample_triton(**{"seed": 0, **kwargs}))
        case S.fused_triton_ret_logits:
            return SimpleSampler(
                lambda **kwargs: fused_mm_sample_triton(
                    **{"seed": 0, "return_logits": True, **kwargs}
                )[0]
            )
        case S.fused_triton_greedy:
            return SimpleSampler(
                lambda **kwargs: fused_mm_sample_triton(
                    **{"seed": 0, "greedy_sampling": True, **kwargs}
                )
            )
        case S.naive_pt:
            return SimpleSampler(sample)
        case S.naive_compiled:
            return SimpleSampler(sample_compiled)
        case S.pt_qitra:
            return SimpleSampler(lambda **kwargs: sample(**kwargs, use_qitra=True))
        case S.sequential_compiled:
            return SimpleSampler(sequential_sample_pt)
        case S.naive_tl_matmul:
            return SimpleSampler(lambda **kwargs: sample_compiled(**kwargs, tl_matmul=True))
        case S.jl_compiled:
            return JLSampler.from_weights(weights)
        case S.fused_topk:
            from .tl_fused_mm_topk import fused_mm_topk_and_sample

            return SimpleSampler(fused_mm_topk_and_sample)
        case S.flashinfer_top_k_top_p_sampling_from_logits:
            return SimpleSampler(
                lambda **kwargs: flashinfer_top_k_top_p_sampling_from_logits(
                    **_default_top_k_top_p(kwargs),
                )
            )
        case S.fused_cuda:
            from .cuda_impl import fused_mm_sample_cuda

            return SimpleSampler(fused_mm_sample_cuda)
        case S.helion:
            from .helion_impl import fused_mm_sample_helion

            return SimpleSampler(fused_mm_sample_helion)
        case S.flashinfer_sampling_from_logits:
            return SimpleSampler(flashinfer_sampling_from_logits)
        case S.greedy_baseline:
            return SimpleSampler(greedy_sample_compiled)
        case _:
            raise NotImplementedError(f"Unknown provider: {provider}")


def _default_top_k_top_p(kwargs: dict) -> dict: