| import os |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from diffulex.moe.topk.base import TopKRouter |
| from diffulex.moe.topk.output import TopKOutput |
| from diffulex_kernel import fused_topk |
|
|
|
|
| def _tile_kernels_topk_gate(scores: torch.Tensor, top_k: int) -> torch.Tensor | None: |
| try: |
| import tile_kernels |
| except Exception: |
| return None |
| try: |
| return tile_kernels.moe.topk_gate(scores, top_k) |
| except Exception: |
| return None |
|
|
|
|
| class NaiveTopKRouter(TopKRouter): |
| def _forward_naive(self, router_logits: torch.Tensor) -> TopKOutput: |
| if self.scoring_func == "softmax": |
| routing_scores = F.softmax(router_logits, dim=-1, dtype=torch.float) |
| elif self.scoring_func == "sigmoid": |
| routing_scores = torch.sigmoid(router_logits.float()) |
| else: |
| raise ValueError(f"Unsupported scoring function: {self.scoring_func!r}.") |
|
|
| top_k = min(self.top_k, routing_scores.shape[-1]) |
| topk_ids = None |
| if ( |
| router_logits.is_cuda |
| and os.getenv("DIFFULEX_MOE_TOPK_IMPL", "").lower() in {"tile", "tilekernels", "tile_kernels"} |
| ): |
| topk_ids = _tile_kernels_topk_gate(routing_scores, top_k) |
|
|
| if topk_ids is None: |
| topk_weights, topk_ids = torch.topk(routing_scores, top_k, dim=-1, sorted=False) |
| else: |
| topk_ids = topk_ids.to(torch.int64) |
| topk_weights = torch.gather(routing_scores, dim=-1, index=topk_ids) |
| if self.renormalize and top_k > 1: |
| topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) |
|
|
| return TopKOutput( |
| weights=topk_weights, |
| ids=topk_ids.to(torch.int32), |
| router_logits=router_logits, |
| ) |
|
|
| def forward(self, router_logits: torch.Tensor) -> TopKOutput: |
| if not router_logits.is_cuda or os.getenv("DIFFULEX_REFERENCE_MOE_ROUTER", "0") == "1": |
| return self._forward_naive(router_logits) |
|
|
| topk_weights, topk_ids = fused_topk( |
| router_logits=router_logits, |
| top_k=self.top_k, |
| renormalize=self.renormalize, |
| scoring_func=self.scoring_func, |
| ) |
| return TopKOutput(weights=topk_weights, ids=topk_ids, router_logits=router_logits) |
|
|
