sglang 0.1.17__py3-none-any.whl → 0.1.18__py3-none-any.whl

sglang/srt/constrained/jump_forward.py

@@ -1,17 +1,43 @@
-import interegular
+"""
+Faster constrained decoding.
+Reference: https://lmsys.org/blog/2024-02-05-compressed-fsm/
+"""
+
+import dataclasses
+from collections import defaultdict
 
-from sglang.srt.constrained import FSMInfo, disk_cache, make_deterministic_fsm
+import interegular
+import outlines.caching
+
+from sglang.srt.constrained import (
+    FSMInfo,
+    disk_cache,
+    make_byte_level_fsm,
+    make_deterministic_fsm,
+)
 from sglang.srt.constrained.base_cache import BaseCache
 
 IP_REGEX = r"((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)"
 
 
+@dataclasses.dataclass
+class JumpEdge:
+    symbol: str = None
+    symbol_next_state: int = None
+    byte: int = None
+    byte_next_state: int = None
+
+
 class JumpForwardMap:
     def __init__(self, regex_string):
         @disk_cache()
         def _init_state_to_jump_forward(regex_string):
             regex_pattern = interegular.parse_pattern(regex_string)
-            regex_fsm, _ = make_deterministic_fsm(regex_pattern.to_fsm().reduce())
+
+            byte_fsm = make_byte_level_fsm(
+                regex_pattern.to_fsm().reduce(), keep_utf8=True
+            )
+            regex_fsm, _ = make_deterministic_fsm(byte_fsm)
 
             fsm_info: FSMInfo = regex_fsm.fsm_info
 
@@ -21,40 +47,93 @@ class JumpForwardMap:
                 id_to_symbol.setdefault(id_, []).append(symbol)
 
             transitions = fsm_info.transitions
-            dirty_states = set()
+            outgoings_ct = defaultdict(int)
             state_to_jump_forward = {}
 
             for (state, id_), next_state in transitions.items():
-                if state in dirty_states:
-                    continue
-                if state in state_to_jump_forward:
-                    dirty_states.add(state)
-                    del state_to_jump_forward[state]
+                if id_ == fsm_info.alphabet_anything_value:
                     continue
-                if len(id_to_symbol[id_]) > 1:
-                    dirty_states.add(state)
+                symbols = id_to_symbol[id_]
+                for c in symbols:
+                    if len(c) > 1:
+                        # Skip byte level transitions
+                        continue
+
+                    outgoings_ct[state] += 1
+                    if outgoings_ct[state] > 1:
+                        if state in state_to_jump_forward:
+                            del state_to_jump_forward[state]
+                        break
+
+                    state_to_jump_forward[state] = JumpEdge(
+                        symbol=c,
+                        symbol_next_state=next_state,
+                    )
+
+            # Process the byte level jump forward
+            outgoings_ct = defaultdict(int)
+            for (state, id_), next_state in transitions.items():
+                if id_ == fsm_info.alphabet_anything_value:
                     continue
-
-                state_to_jump_forward[state] = (id_to_symbol[id_][0], next_state)
+                symbols = id_to_symbol[id_]
+                for c in symbols:
+                    byte_ = None
+                    if len(c) == 1 and ord(c) < 0x80:
+                        # ASCII character
+                        byte_ = ord(c)
+                    elif len(c) > 1:
+                        # FIXME: This logic is due to the leading \x00
+                        # https://github.com/outlines-dev/outlines/pull/930
+                        byte_ = int(symbols[0][1:], 16)
+
+                    if byte_ is not None:
+                        outgoings_ct[state] += 1
+                        if outgoings_ct[state] > 1:
+                            if state in state_to_jump_forward:
+                                del state_to_jump_forward[state]
+                            break
+                        e = state_to_jump_forward.get(state, JumpEdge())
+                        e.byte = byte_
+                        e.byte_next_state = next_state
+                        state_to_jump_forward[state] = e
 
             return state_to_jump_forward
 
         self.state_to_jump_forward = _init_state_to_jump_forward(regex_string)
 
-    def valid_states(self):
-        return self.state_to_jump_forward.keys()
+    def jump_forward_symbol(self, state):
+        jump_forward_str = ""
+        next_state = state
+        while state in self.state_to_jump_forward:
+            e = self.state_to_jump_forward[state]
+            if e.symbol is None:
+                break
+            jump_forward_str += e.symbol
+            next_state = e.symbol_next_state
+            state = next_state
 
-    def jump_forward(self, state):
+        return jump_forward_str, next_state
+
+    def jump_forward_byte(self, state):
         if state not in self.state_to_jump_forward:
             return None
 
-        jump_forward_str = ""
+        jump_forward_bytes = []
         next_state = None
         while state in self.state_to_jump_forward:
-            symbol, next_state = self.state_to_jump_forward[state]
-            jump_forward_str += symbol
+            e = self.state_to_jump_forward[state]
+            assert e.byte is not None and e.byte_next_state is not None
+            jump_forward_bytes.append((e.byte, e.byte_next_state))
+            next_state = e.byte_next_state
             state = next_state
-        return jump_forward_str, next_state
+
+        return jump_forward_bytes
+
+    def is_jump_forward_symbol_state(self, state):
+        return (
+            state in self.state_to_jump_forward
+            and self.state_to_jump_forward[state].symbol is not None
+        )
 
 
 class JumpForwardCache(BaseCache):
@@ -65,12 +144,21 @@ class JumpForwardCache(BaseCache):
         return JumpForwardMap(regex)
 
 
-def test_main():
-    regex_string = r"The google's DNS sever address is " + IP_REGEX
+def test_main(regex_string):
     jump_forward_map = JumpForwardMap(regex_string)
-    for state in jump_forward_map.valid_states():
-        print(state, f'"{jump_forward_map.jump_forward(state)}"')
+    for state, e in jump_forward_map.state_to_jump_forward.items():
+        if e.symbol is not None:
+            jump_forward_str, next_state = jump_forward_map.jump_forward_symbol(state)
+            print(f"{state} -> {next_state}", jump_forward_str)
+        bytes_ = jump_forward_map.jump_forward_byte(state)
+        print(f"{state} -> {bytes_[-1][1]}", [hex(b) for b, _ in bytes_])
 
 
 if __name__ == "__main__":
-    test_main()
+    import outlines
+
+    outlines.caching.clear_cache()
+    test_main(r"The google's DNS sever address is " + IP_REGEX)
+    test_main(r"霍格沃茨特快列车|霍比特人比尔博")
+    # 霍格: \xe9\x9c\x8d \xe6\xa0\xbc ...
+    # 霍比: \xe9\x9c\x8d \xe6\xaf\x94 ...

sglang/srt/conversation.py

@@ -1,3 +1,5 @@
+"""Conversation templates."""
+
 # Adapted from
 # https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
 import dataclasses

sglang/srt/flush_cache.py

@@ -1,4 +1,6 @@
 """
+Flush the KV cache.
+
 Usage:
 python3 -m sglang.srt.flush_cache --url http://localhost:30000
 """

sglang/srt/hf_transformers_utils.py

@@ -1,10 +1,10 @@
 """Utilities for Huggingface Transformers."""
 
+import functools
 import json
 import os
 import warnings
-import functools
-from typing import Optional, Union, AbstractSet, Collection, Literal
+from typing import AbstractSet, Collection, Literal, Optional, Union
 
 from huggingface_hub import snapshot_download
 from transformers import (
@@ -88,6 +88,9 @@ def get_tokenizer(
     if tokenizer_name.endswith(".json"):
         return TiktokenTokenizer(tokenizer_name)
 
+    if tokenizer_name.endswith(".model"):
+        return SentencePieceTokenizer(tokenizer_name)
+
     """Gets a tokenizer for the given model name via Huggingface."""
     if is_multimodal_model(tokenizer_name):
         processor = get_processor(
@@ -179,6 +182,8 @@ def get_processor(
 class TiktokenTokenizer:
     def __init__(self, tokenizer_path):
         import tiktoken
+        from jinja2 import Template
+
         PAT_STR_B = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
 
         # Read JSON
@@ -190,7 +195,8 @@ class TiktokenTokenizer:
             bytes(item["bytes"]): item["token"] for item in tok_dict["regular_tokens"]
         }
         special_tokens = {
-            bytes(item["bytes"]).decode(): item["token"] for item in tok_dict["special_tokens"]
+            bytes(item["bytes"]).decode(): item["token"]
+            for item in tok_dict["special_tokens"]
         }
         assert tok_dict["word_split"] == "V1"
 
@@ -202,7 +208,10 @@ class TiktokenTokenizer:
         }
         if "default_allowed_special" in tok_dict:
             default_allowed_special = set(
-                [bytes(bytes_list).decode() for bytes_list in tok_dict["default_allowed_special"]]
+                [
+                    bytes(bytes_list).decode()
+                    for bytes_list in tok_dict["default_allowed_special"]
+                ]
             )
         else:
             default_allowed_special = None
@@ -211,25 +220,35 @@ class TiktokenTokenizer:
 
         tokenizer = tiktoken.Encoding(**kwargs)
         tokenizer._default_allowed_special = default_allowed_special or set()
+        tokenizer._default_allowed_special |= {"<|separator|>"}
 
         def encode_patched(
             self,
             text: str,
             *,
-            allowed_special: Union[Literal["all"], AbstractSet[str]] = set(),  # noqa: B006
+            allowed_special: Union[
+                Literal["all"], AbstractSet[str]
+            ] = set(),  # noqa: B006
             disallowed_special: Union[Literal["all"], Collection[str]] = "all",
         ) -> list[int]:
             if isinstance(allowed_special, set):
                 allowed_special |= self._default_allowed_special
             return tiktoken.Encoding.encode(
-                self, text, allowed_special=allowed_special, disallowed_special=disallowed_special
+                self,
+                text,
+                allowed_special=allowed_special,
+                disallowed_special=disallowed_special,
             )
+
         tokenizer.encode = functools.partial(encode_patched, tokenizer)
 
         # Convert to HF interface
         self.tokenizer = tokenizer
         self.eos_token_id = tokenizer._special_tokens["<|eos|>"]
         self.vocab_size = tokenizer.n_vocab
+        self.chat_template = Template(
+            "{% for message in messages %}{% if message['role'] == 'user' %}{{ 'Human: ' + message['content'].strip() + '<|separator|>\n\n' }}{% elif message['role'] == 'system' %}{{ 'System: ' + message['content'].strip() + '<|separator|>\n\n' }}{% elif message['role'] == 'assistant' %}{{ 'Assistant: '  + message['content'] + '<|separator|>\n\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
+        )
 
     def encode(self, x, add_special_tokens=False):
         return self.tokenizer.encode(x)
@@ -237,10 +256,46 @@ class TiktokenTokenizer:
     def decode(self, x):
         return self.tokenizer.decode(x)
 
-    def batch_decode(self, batch, skip_special_tokens=True,  spaces_between_special_tokens=False):
+    def batch_decode(
+        self, batch, skip_special_tokens=True, spaces_between_special_tokens=False
+    ):
         if isinstance(batch[0], int):
             batch = [[x] for x in batch]
         return self.tokenizer.decode_batch(batch)
 
-    def convert_ids_to_tokens(self, index):
-        return self.tokenizer.decode_single_token_bytes(index).decode("utf-8", errors="ignore")
+    def apply_chat_template(self, messages, tokenize, add_generation_prompt):
+        ret = self.chat_template.render(messages=messages, add_generation_prompt=add_generation_prompt)
+        return self.encode(ret) if tokenize else ret
+
+
+class SentencePieceTokenizer:
+    def __init__(self, tokenizer_path):
+        import sentencepiece as spm
+        from jinja2 import Template
+
+        tokenizer = spm.SentencePieceProcessor(model_file=tokenizer_path)
+
+        # Convert to HF interface
+        self.tokenizer = tokenizer
+        self.eos_token_id = tokenizer.eos_id()
+        self.vocab_size = tokenizer.vocab_size()
+        self.chat_template = Template(
+            "{% for message in messages %}{% if message['role'] == 'user' %}{{ 'Human: ' + message['content'].strip() + '<|separator|>\n\n' }}{% elif message['role'] == 'system' %}{{ 'System: ' + message['content'].strip() + '<|separator|>\n\n' }}{% elif message['role'] == 'assistant' %}{{ 'Assistant: '  + message['content'] + '<|separator|>\n\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
+        )
+
+    def encode(self, x, add_special_tokens=False):
+        return self.tokenizer.encode(x)
+
+    def decode(self, x):
+        return self.tokenizer.decode(x)
+
+    def batch_decode(
+        self, batch, skip_special_tokens=True, spaces_between_special_tokens=False
+    ):
+        if isinstance(batch[0], int):
+            batch = [[x] for x in batch]
+        return self.tokenizer.decode(batch)
+
+    def apply_chat_template(self, messages, tokenize, add_generation_prompt):
+        ret = self.chat_template.render(messages=messages, add_generation_prompt=add_generation_prompt)
+        return self.encode(ret) if tokenize else ret

sglang/srt/layers/fused_moe.py

@@ -12,7 +12,6 @@ import triton.language as tl
 
 from vllm import _custom_ops as ops
 from vllm.logger import init_logger
-from vllm.utils import is_hip
 
 logger = init_logger(__name__)
 
@@ -310,92 +309,110 @@ def get_moe_configs(E: int, N: int,
     return None
 
 
-def fused_moe(
+def get_default_config(
+    M: int,
+    E: int,
+    N: int,
+    K: int,
+    topk: int,
+    dtype: Optional[str],
+) -> Dict[str, int]:
+    if dtype == "float8":
+        config = {
+            'BLOCK_SIZE_M': 128,
+            'BLOCK_SIZE_N': 256,
+            'BLOCK_SIZE_K': 128,
+            'GROUP_SIZE_M': 32,
+            "num_warps": 8,
+            "num_stages": 4
+        }
+        if M <= E:
+            config = {
+                'BLOCK_SIZE_M': 64,
+                'BLOCK_SIZE_N': 128,
+                'BLOCK_SIZE_K': 128,
+                'GROUP_SIZE_M': 1,
+                "num_warps": 4,
+                "num_stages": 4
+            }
+    else:
+        config = {
+            'BLOCK_SIZE_M': 64,
+            'BLOCK_SIZE_N': 64,
+            'BLOCK_SIZE_K': 32,
+            'GROUP_SIZE_M': 8
+        }
+        if M <= E:
+            config = {
+                'BLOCK_SIZE_M': 16,
+                'BLOCK_SIZE_N': 32,
+                'BLOCK_SIZE_K': 64,
+                'GROUP_SIZE_M': 1
+            }
+    return config
+
+
+def fused_topk(
     hidden_states: torch.Tensor,
-    w1: torch.Tensor,
-    w2: torch.Tensor,
     gating_output: torch.Tensor,
     topk: int,
     renormalize: bool,
-    inplace: bool = False,
-    override_config: Optional[Dict[str, Any]] = None,
-    use_fp8: bool = False,
-    w1_scale: Optional[torch.Tensor] = None,
-    w2_scale: Optional[torch.Tensor] = None,
-    a1_scale: Optional[torch.Tensor] = None,
-    a2_scale: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    """
-    This function computes a Mixture of Experts (MoE) layer using two sets of
-    weights, w1 and w2, and top-k gating mechanism.
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], (
+        "Number of tokens mismatch")
 
-    Parameters:
-    - hidden_states (torch.Tensor): The input tensor to the MoE layer.
-    - w1 (torch.Tensor): The first set of expert weights.
-    - w2 (torch.Tensor): The second set of expert weights.
-    - gating_output (torch.Tensor): The output of the gating operation
-        (before softmax).
-    - topk (int): The number of top-k experts to select.
-    - renormalize (bool): If True, renormalize the top-k weights to sum to 1.
-    - inplace (bool): If True, perform the operation in-place.
-        Defaults to False.
-    - override_config (Optional[Dict[str, Any]]): Optional override
-        for the kernel configuration.
-    - use_fp8 (bool): If True, use fp8 arithmetic to compute the inner
-        products for w1 and w2. Defaults to False.
-    - w1_scale (Optional[torch.Tensor]): Optional scale to be used for
-        w1.
-    - w2_scale (Optional[torch.Tensor]): Optional scale to be used for
-        w2.
+    M, _ = hidden_states.shape
 
-    Returns:
-    - torch.Tensor: The output tensor after applying the MoE layer.
-    """
+    topk_weights = torch.empty(M,
+                               topk,
+                               dtype=torch.float32,
+                               device=hidden_states.device)
+    topk_ids = torch.empty(M,
+                           topk,
+                           dtype=torch.int32,
+                           device=hidden_states.device)
+    token_expert_indicies = torch.empty(M,
+                                        topk,
+                                        dtype=torch.int32,
+                                        device=hidden_states.device)
+    ops.topk_softmax(
+        topk_weights,
+        topk_ids,
+        token_expert_indicies,
+        gating_output.float(),  # TODO(woosuk): Optimize this.
+    )
+    del token_expert_indicies  # Not used. Will be used in the future.
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+    return topk_weights, topk_ids
+
+
+def fused_experts(hidden_states: torch.Tensor,
+                  w1: torch.Tensor,
+                  w2: torch.Tensor,
+                  topk_weights: torch.Tensor,
+                  topk_ids: torch.Tensor,
+                  inplace: bool = False,
+                  override_config: Optional[Dict[str, Any]] = None,
+                  use_fp8: bool = False,
+                  w1_scale: Optional[torch.Tensor] = None,
+                  w2_scale: Optional[torch.Tensor] = None,
+                  a1_scale: Optional[torch.Tensor] = None,
+                  a2_scale: Optional[torch.Tensor] = None):
     # Check constraints.
-    assert hidden_states.shape[0] == gating_output.shape[0], (
-        "Number of tokens mismatch")
     assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
-    assert gating_output.shape[1] == w1.shape[0], "Number of experts mismatch"
+    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
     assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
     assert w1.is_contiguous(), "Expert weights1 must be contiguous"
     assert w2.is_contiguous(), "Expert weights2 must be contiguous"
     assert hidden_states.dtype in [
         torch.float32, torch.float16, torch.bfloat16
     ]
+
     M, _ = hidden_states.shape
     E, N, _ = w1.shape
 
-    if is_hip():
-        # The MoE kernels are not yet supported on ROCm.
-        routing_weights = torch.softmax(gating_output,
-                                        dim=-1,
-                                        dtype=torch.float32)
-        topk_weights, topk_ids = torch.topk(routing_weights, topk, dim=-1)
-    else:
-        import vllm._moe_C as moe_kernels
-
-        topk_weights = torch.empty(M,
-                                   topk,
-                                   dtype=torch.float32,
-                                   device=hidden_states.device)
-        topk_ids = torch.empty(M,
-                               topk,
-                               dtype=torch.int32,
-                               device=hidden_states.device)
-        token_expert_indicies = torch.empty(M,
-                                            topk,
-                                            dtype=torch.int32,
-                                            device=hidden_states.device)
-        moe_kernels.topk_softmax(
-            topk_weights,
-            topk_ids,
-            token_expert_indicies,
-            gating_output.float(),  # TODO(woosuk): Optimize this.
-        )
-        del token_expert_indicies  # Not used. Will be used in the future.
-    if renormalize:
-        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
-
     if override_config:
         config = override_config
     else:
@@ -409,24 +426,9 @@ def fused_moe(
             config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
         else:
             # Else use the default config
-            config = {
-                "BLOCK_SIZE_M": 128,
-                "BLOCK_SIZE_N": 64,
-                "BLOCK_SIZE_K": 128,
-                "GROUP_SIZE_M": 1,
-                "num_warps": 4,
-                "num_stages": 4
-            }
-
-            if M <= E:
-                config = {
-                    "BLOCK_SIZE_M": 128,
-                    "BLOCK_SIZE_N": 256,
-                    "BLOCK_SIZE_K": 128,
-                    "GROUP_SIZE_M": 16,
-                    "num_warps": 8,
-                    "num_stages": 4
-                }
+            config = get_default_config(M, E, N, w1.shape[2],
+                                        topk_ids.shape[1],
+                                        "float8" if use_fp8 else None)
 
     intermediate_cache1 = torch.empty((M, topk_ids.shape[1], N),
                                       device=hidden_states.device,
@@ -482,4 +484,99 @@ def fused_moe(
                          dim=1,
                          out=hidden_states)
     return torch.sum(intermediate_cache3.view(*intermediate_cache3.shape),
-                     dim=1)
+                     dim=1)
+
+
+def fused_moe(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    inplace: bool = False,
+    override_config: Optional[Dict[str, Any]] = None,
+    use_fp8: bool = False,
+    w1_scale: Optional[torch.Tensor] = None,
+    w2_scale: Optional[torch.Tensor] = None,
+    a1_scale: Optional[torch.Tensor] = None,
+    a2_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    This function computes a Mixture of Experts (MoE) layer using two sets of
+    weights, w1 and w2, and top-k gating mechanism.
+
+    Parameters:
+    - hidden_states (torch.Tensor): The input tensor to the MoE layer.
+    - w1 (torch.Tensor): The first set of expert weights.
+    - w2 (torch.Tensor): The second set of expert weights.
+    - gating_output (torch.Tensor): The output of the gating operation
+        (before softmax).
+    - topk (int): The number of top-k experts to select.
+    - renormalize (bool): If True, renormalize the top-k weights to sum to 1.
+    - inplace (bool): If True, perform the operation in-place.
+        Defaults to False.
+    - override_config (Optional[Dict[str, Any]]): Optional override
+        for the kernel configuration.
+    - use_fp8 (bool): If True, use fp8 arithmetic to compute the inner
+        products for w1 and w2. Defaults to False.
+    - w1_scale (Optional[torch.Tensor]): Optional scale to be used for
+        w1.
+    - w2_scale (Optional[torch.Tensor]): Optional scale to be used for
+        w2.
+
+    Returns:
+    - torch.Tensor: The output tensor after applying the MoE layer.
+    """
+    # Check constraints.
+    assert gating_output.shape[1] == w1.shape[0], "Number of experts mismatch"
+
+    if hasattr(ops, "topk_softmax"):
+        topk_weights, topk_ids = fused_topk(hidden_states, gating_output, topk,
+                                            renormalize)
+    else:
+        topk_weights, topk_ids = fused_topk_v0_4_3(hidden_states, gating_output, topk,
+                                                   renormalize)
+
+    return fused_experts(hidden_states,
+                         w1,
+                         w2,
+                         topk_weights,
+                         topk_ids,
+                         inplace=inplace,
+                         override_config=override_config,
+                         use_fp8=use_fp8,
+                         w1_scale=w1_scale,
+                         w2_scale=w2_scale,
+                         a1_scale=a1_scale,
+                         a2_scale=a2_scale)
+
+
+
+def fused_topk_v0_4_3(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    import vllm._moe_C as moe_kernels
+    M, _ = hidden_states.shape
+
+    topk_weights = torch.empty(
+        M, topk, dtype=torch.float32, device=hidden_states.device
+    )
+    topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device)
+    token_expert_indicies = torch.empty(
+        M, topk, dtype=torch.int32, device=hidden_states.device
+    )
+    moe_kernels.topk_softmax(
+        topk_weights,
+        topk_ids,
+        token_expert_indicies,
+        gating_output.float(),  # TODO(woosuk): Optimize this.
+    )
+    del token_expert_indicies  # Not used. Will be used in the future.
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+    return topk_weights, topk_ids