wip

Author: wenxindongwork

tpu_inference/layers/vllm/fused_moe.py

@@ -8,7 +8,7 @@
 
 from tpu_inference.layers.vllm.linear_common import \
     slice_sharded_tensor_for_concatenation
-from tpu_inference.layers.common.sharding import ShardingAxisName
+
 P = PartitionSpec
 
 
@@ -374,6 +374,7 @@ def fused_moe_func(
     assert (num_tokens * topk) % 16 == 0, (
         "The kernel requires num_tokens * topk to be a multiple of "
         f"16 but got {num_tokens}*{topk}={num_tokens*topk}")
+
     hidden_states = hidden_states.reshape(num_tokens, hidden_size)
     gating_output = gating_output.reshape(num_tokens, global_num_experts)
 
@@ -425,6 +426,7 @@ def _process_tokens_locally(hidden_states_local, topk_indices_local):
         )
 
     x = activation_fn(activation, x1, x2)
+
     if use_ep:
         x = expert_sharded_gmm(
             x,

tpu_inference/layers/vllm/quantization/unquantized.py

@@ -25,8 +25,6 @@
 from tpu_inference.kernels.fused_moe.v1.kernel import fused_ep_moe
 from tpu_inference.layers.common.quant_methods import (UNQUANTIZED,
                                                        get_tpu_quant_method)
-from tpu_inference.layers.common.sharding import ShardingAxisName
-
 from tpu_inference.layers.vllm.fused_moe import fused_moe_func_padded
 from tpu_inference.layers.vllm.linear_common import (
     reorder_concatenated_tensor_for_sharding,

tpu_inference/layers/vllm/sharding.py

@@ -19,7 +19,6 @@
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead, VocabParallelEmbedding)
 
-from tpu_inference.layers.common.sharding import ShardingAxisName
 from tpu_inference import envs
 from tpu_inference.logger import init_logger