Merge branch 'upstream' into improve-types

benchmarks/kernels/benchmark_mixtral_moe.py

@@ -37,8 +37,17 @@ class BenchmarkConfig(TypedDict):
     num_stages: int
 
 
-def run_grid(bs, method, dtype: str):
-    d_model = 4096
+def run_grid(bs, model, method, gpu, tp_size, dtype: str):
+    if model == '8x7B':
+        d_model = 4096
+        model_intermediate_size = 14336
+        num_layers = 32
+    elif model == '8x22B':
+        d_model = 6144
+        model_intermediate_size = 16384
+        num_layers = 56
+    else:
+        raise ValueError(f'Unsupported Mixtral model {model}')
     num_total_experts = 8
     top_k = 2
     # tp_size = 2

tests/kernels/test_pos_encoding.py

@@ -10,8 +10,8 @@
 
 IS_NEOX_STYLE = [True, False]
 DTYPES = [torch.half, torch.bfloat16, torch.float]
-HEAD_SIZES = [64, 80, 96, 112, 128, 256]
-ROTARY_DIMS = (None, 32)  # None means rotary dim == head size
+HEAD_SIZES = [64, 80, 96, 112, 128, 192, 256]
+ROTARY_DIMS = [None, 32]  # None means rotary dim == head size
 NUM_HEADS = [7, 17]  # Arbitrary values for testing
 BATCH_SIZES = [1, 5]  # Arbitrary values for testing
 SEQ_LENS = [11, 8192]  # Arbitrary values for testing

vllm/model_executor/layers/linear.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import torch
 import torch.nn.functional as F