| 9 | |
| 10 | |
| 11 | class SamplerBackend(nn.Module): |
| 12 | def __init__(self): |
| 13 | super().__init__() |
| 14 | self.backend = self._detect_backend() |
| 15 | self._init_backend_assets() |
| 16 | |
| 17 | def _detect_backend(self): |
| 18 | if torch.version.hip: |
| 19 | logger.info_once("Detected AMD GPU (ROCm) - Using AITER backend") |
| 20 | return constants.BackendLib.AITER.value |
| 21 | else: |
| 22 | logger.info_once("Detected NVIDIA GPU (CUDA) - Using FlashInfer backend") |
| 23 | return constants.BackendLib.FLASHINFER.value |
| 24 | |
| 25 | def _init_backend_assets(self): |
| 26 | """ |
| 27 | Preload backend-specific dependencies |
| 28 | to avoid runtime import overhead. |
| 29 | """ |
| 30 | if self.backend == constants.BackendLib.FLASHINFER.value: |
| 31 | try: |
| 32 | import flashinfer |
| 33 | |
| 34 | self.flashinfer = flashinfer |
| 35 | logger.info("FlashInfer kernels loaded successfully.") |
| 36 | except ImportError: |
| 37 | logger.error("FlashInfer not found. Please install it for NVIDIA GPUs.") |
| 38 | elif self.backend == constants.BackendLib.AITER.value: |
| 39 | pass |
| 40 | |
| 41 | @torch.inference_mode() |
| 42 | def sample(self, logits, top_k=None, top_p=None, temperature=1.0, deterministic=True): |
| 43 | """ |
| 44 | Unified sampling interface. |
| 45 | """ |
| 46 | logits = logits.contiguous() |
| 47 | if temperature != 1.0: |
| 48 | logits = logits / temperature |
| 49 | |
| 50 | if self.backend == constants.BackendLib.FLASHINFER.value: |
| 51 | from flashinfer.sampling import top_k_renorm_probs, top_p_sampling_from_probs |
| 52 | |
| 53 | logits = logits.float().contiguous() |
| 54 | probs = torch.softmax(logits, dim=-1) |
| 55 | |
| 56 | if top_k is None and top_p is None: |
| 57 | return torch.multinomial(probs, num_samples=1).view(-1) |
| 58 | |
| 59 | if top_k is not None: |
| 60 | probs = top_k_renorm_probs(probs, top_k) |
| 61 | |
| 62 | if top_p is not None: |
| 63 | return top_p_sampling_from_probs(probs, top_p, deterministic=deterministic) |
| 64 | |
| 65 | return torch.multinomial(probs, num_samples=1).view(-1) |
| 66 | |
| 67 | elif self.backend == constants.BackendLib.AITER.value: |
| 68 | # TODO: Connect to AITER's sampling operator |
no outgoing calls