Track expert selection metrics

Use VLLM_COLLECT_EXPERT_USAGE_HISTOGRAM=1 to enable this
feature.
Make sure that PROMETHEUS_MULTIPROC_DIR must be set to get proper
metrics!

Expect a 2% maximum e2e overhead when running this!
Perf on the GPU side is negligible.
Note that this PR does enable anything by default, so perf is untouched
this way.

Signed-off-by: Thibault Schueller <[email protected]>

Signed-off-by: 'Thibault Schueller' <'[email protected]'>
Signed-off-by: Thibault Schueller <[email protected]>

Author: Ryp

tests/kernels/moe/test_expert_usage_histogram.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.utils import (
+    collect_expert_usage_histogram)
+
+
+@pytest.mark.parametrize("topk_experts,expert_count,topk_ids_dtype",
+                         [(4, 32, torch.int32), (1, 1, torch.int64)])
+@pytest.mark.parametrize("token_count", [256, 7])
+def test_collect_expert_usage_histogram(topk_experts: int, expert_count: int,
+                                        token_count: int,
+                                        topk_ids_dtype: torch.dtype):
+    device = torch.device('cuda')
+
+    # Make an uniform distribution of expert usage
+    topk_ids = torch.stack([torch.arange(topk_experts, dtype=topk_ids_dtype)] *
+                           token_count)
+
+    topk_ids_gpu = topk_ids.to(device)
+
+    expert_usage_histogram_gpu = torch.zeros(expert_count,
+                                             dtype=torch.int32,
+                                             device=device)
+
+    collect_expert_usage_histogram(topk_ids_gpu, expert_usage_histogram_gpu)
+
+    # Every expert is used the same amount, so expecting token_count for
+    # each expert set in the topk_ids tensor.
+    assert torch.equal(
+        expert_usage_histogram_gpu[:topk_experts],
+        torch.full([topk_experts],
+                   token_count,
+                   dtype=torch.int32,
+                   device=device))
+
+    # The rest of the experts weren't used, so they should be zero.
+    assert expert_usage_histogram_gpu[topk_experts:].sum() == 0
+
+
+@pytest.mark.parametrize("topk_experts,expert_count", [(16, 32)])
+@pytest.mark.parametrize("token_count", [1])
+@pytest.mark.parametrize("seed", [0xDEADBEEF, 0xCAFEBABE])
+def test_collect_expert_usage_histogram_random(topk_experts: int,
+                                               expert_count: int,
+                                               token_count: int, seed: int):
+    device = torch.device('cuda')
+
+    generator = torch.Generator()
+    generator.manual_seed(seed)
+
+    # Make random distribution of expert usage
+    topk_ids_cpu = torch.stack(
+        [torch.randperm(topk_experts, generator=generator, dtype=torch.int32)
+         ] * token_count)
+
+    # Compute ground truth
+    torch_histogram = torch.histogram(topk_ids_cpu.to(torch.float),
+                                      bins=expert_count,
+                                      range=(0, expert_count - 1))
+
+    # Use our function
+    expert_usage_histogram_gpu = torch.zeros(expert_count,
+                                             dtype=torch.int32,
+                                             device=device)
+
+    topk_ids_gpu = topk_ids_cpu.to(device)
+
+    collect_expert_usage_histogram(topk_ids_gpu, expert_usage_histogram_gpu)
+
+    assert torch.equal(expert_usage_histogram_gpu,
+                       torch_histogram.hist.to(torch.int32).to(device))

vllm/config.py

@@ -981,7 +981,7 @@ def _verify_bnb_config(self) -> None:
 
             self.enforce_eager = True
 
-    def _verify_with_expert_parallelism(self) -> None:
+    def get_total_num_experts(self) -> int:
         num_expert_names = [
             "moe_num_experts",  # Dbrx
             "num_experts",  # Jamba
@@ -993,7 +993,10 @@ def _verify_with_expert_parallelism(self) -> None:
             num_experts = getattr(self.hf_text_config, name, 0)
             if num_experts > 0:
                 break
-        if num_experts < 1:
+        return num_experts
+
+    def _verify_with_expert_parallelism(self) -> None:
+        if self.get_total_num_experts() < 1:
             raise ValueError(
                 "Number of experts in the model must be greater than 0 "
                 "when expert parallelism is enabled.")
@@ -1222,16 +1225,21 @@ def get_num_attention_heads(self,
         num_heads = getattr(self.hf_text_config, "num_attention_heads", 0)
         return num_heads // parallel_config.tensor_parallel_size
 
-    def get_layers_start_end_indices(
-            self, parallel_config: "ParallelConfig") -> tuple[int, int]:
-        from vllm.distributed.utils import get_pp_indices
+    def get_total_num_hidden_layers(self) -> int:
         if (self.hf_text_config.model_type == "deepseek_mtp"
                 or self.hf_config.model_type == "mimo_mtp"):
             total_num_hidden_layers = getattr(self.hf_text_config,
                                               "num_nextn_predict_layers", 0)
         else:
             total_num_hidden_layers = getattr(self.hf_text_config,
                                               "num_hidden_layers", 0)
+        return total_num_hidden_layers
+
+    def get_layers_start_end_indices(
+            self, parallel_config: "ParallelConfig") -> tuple[int, int]:
+        from vllm.distributed.utils import get_pp_indices
+        total_num_hidden_layers = self.get_total_num_hidden_layers()
+
         # the layout order is: DP x PP x TP
         pp_rank = (parallel_config.rank // parallel_config.tensor_parallel_size
                    ) % parallel_config.pipeline_parallel_size

vllm/envs.py

@@ -131,6 +131,7 @@
     VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
     VLLM_KV_CACHE_LAYOUT: Optional[str] = None
     VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
+    VLLM_COLLECT_EXPERT_USAGE_HISTOGRAM: bool = False
 
 
 def get_default_cache_root():
@@ -905,6 +906,10 @@ def get_vllm_port() -> Optional[int]:
     # or bad hardware but it may add compute overhead.
     "VLLM_COMPUTE_NANS_IN_LOGITS":
     lambda: bool(int(os.getenv("VLLM_COMPUTE_NANS_IN_LOGITS", "0"))),
+
+    # Collects expert routing histogram per layer
+    "VLLM_COLLECT_EXPERT_USAGE_HISTOGRAM":
+    lambda: bool(int(os.getenv("VLLM_COLLECT_EXPERT_USAGE_HISTOGRAM", "0"))),
 }
 
 # --8<-- [end:env-vars-definition]

vllm/forward_context.py

@@ -95,6 +95,8 @@ class ForwardContext:
     # set dynamically for each forward pass
     dp_metadata: Optional[DPMetadata] = None
     skip_cuda_graphs: bool = False
+    # Set when recording usage histogram
+    expert_usage_histogram: Optional[torch.Tensor] = None
 
 
 _forward_context: Optional[ForwardContext] = None
@@ -116,6 +118,7 @@ def set_forward_context(
     num_tokens: Optional[int] = None,
     num_tokens_across_dp: Optional[torch.Tensor] = None,
     skip_cuda_graphs: bool = False,
+    expert_usage_histogram: Optional[torch.Tensor] = None,
 ):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
@@ -132,6 +135,9 @@ def set_forward_context(
                                       attn_metadata, num_tokens or 0,
                                       num_tokens_across_dp)
 
+    if expert_usage_histogram is not None:
+        expert_usage_histogram.zero_()
+
     global _forward_context
     prev_context = _forward_context
     _forward_context = ForwardContext(
@@ -141,6 +147,7 @@ def set_forward_context(
         attn_metadata=attn_metadata,
         dp_metadata=dp_metadata,
         skip_cuda_graphs=skip_cuda_graphs,
+        expert_usage_histogram=expert_usage_histogram,
     )
 
     try:

vllm/model_executor/layers/fused_moe/layer.py

@@ -25,8 +25,11 @@
 from vllm.model_executor.custom_op import CustomOp
 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
     is_rocm_aiter_moe_enabled)
+from vllm.model_executor.layers.fused_moe.utils import (
+    collect_expert_usage_histogram)
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig, QuantizeMethodBase)
+from vllm.model_executor.models.utils import extract_layer_index
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 from vllm.platforms.interface import CpuArchEnum
@@ -415,6 +418,7 @@ def apply(
         router_logits: torch.Tensor,
         top_k: int,
         renormalize: bool,
+        layer_index: int,
         use_grouped_topk: bool = False,
         topk_group: Optional[int] = None,
         num_expert_group: Optional[int] = None,
@@ -554,6 +558,7 @@ def apply(
         router_logits: torch.Tensor,
         top_k: int,
         renormalize: bool,
+        layer_index: int,
         use_grouped_topk: bool = False,
         topk_group: Optional[int] = None,
         num_expert_group: Optional[int] = None,
@@ -571,6 +576,7 @@ def apply(
             router_logits=router_logits,
             top_k=top_k,
             renormalize=renormalize,
+            layer_index=layer_index,
             use_grouped_topk=use_grouped_topk,
             topk_group=topk_group,
             num_expert_group=num_expert_group,
@@ -590,6 +596,7 @@ def forward_cuda(
         top_k: int,
         router_logits: torch.Tensor,
         renormalize: bool,
+        layer_index: int,
         topk_group: Optional[int] = None,
         num_expert_group: Optional[int] = None,
         global_num_experts: int = -1,
@@ -607,6 +614,7 @@ def forward_cuda(
             use_grouped_topk=use_grouped_topk,
             top_k=top_k,
             renormalize=renormalize,
+            layer_index=layer_index,
             topk_group=topk_group,
             num_expert_group=num_expert_group,
             custom_routing_function=custom_routing_function,
@@ -646,6 +654,7 @@ def forward_cpu(
         top_k: int,
         router_logits: torch.Tensor,
         renormalize: bool,
+        layer_index: int,
         topk_group: Optional[int] = None,
         num_expert_group: Optional[int] = None,
         global_num_experts: int = -1,
@@ -680,6 +689,7 @@ def forward_hpu(
         top_k: int,
         router_logits: torch.Tensor,
         renormalize: bool,
+        layer_index: int,
         topk_group: Optional[int] = None,
         num_expert_group: Optional[int] = None,
         global_num_experts: int = -1,
@@ -713,6 +723,7 @@ def forward_tpu(
         top_k: int,
         router_logits: torch.Tensor,
         renormalize: bool,
+        layer_index: int,
         topk_group: Optional[int] = None,
         num_expert_group: Optional[int] = None,
         global_num_experts: int = -1,
@@ -861,6 +872,8 @@ def __init__(
             compilation_config.static_forward_context[prefix] = self
             self.layer_name = prefix
 
+        self.layer_index = extract_layer_index(prefix)
+
         # Determine expert maps
         if self.use_ep:
             self.local_num_experts, self.expert_map = determine_expert_map(
@@ -1282,6 +1295,7 @@ def select_experts(hidden_states: torch.Tensor,
                        top_k: int,
                        use_grouped_topk: bool,
                        renormalize: bool,
+                       layer_index: int,
                        topk_group: Optional[int] = None,
                        num_expert_group: Optional[int] = None,
                        custom_routing_function: Optional[Callable] = None,
@@ -1322,6 +1336,12 @@ def select_experts(hidden_states: torch.Tensor,
             if indices_type is not None:
                 topk_ids = topk_ids.to(dtype=indices_type)
 
+        expert_usage_histogram = get_forward_context().expert_usage_histogram
+
+        if expert_usage_histogram is not None:
+            collect_expert_usage_histogram(topk_ids,
+                                           expert_usage_histogram[layer_index])
+
         return topk_weights, topk_ids
 
     def must_reduce_shared_expert_outputs(self) -> bool:
@@ -1354,10 +1374,12 @@ def maybe_all_reduce_tensor_model_parallel(
     def forward(self, hidden_states: torch.Tensor,
                 router_logits: torch.Tensor):
         if self.use_direct_call:
-            return self.forward_impl(hidden_states, router_logits)
+            return self.forward_impl(hidden_states, router_logits,
+                                     self.layer_index)
         else:
             return torch.ops.vllm.moe_forward(hidden_states, router_logits,
-                                              self.layer_name)
+                                              self.layer_name,
+                                              self.layer_index)
 
     def forward_impl_chunked(self, full_hidden_states: torch.Tensor,
                              full_router_logits: torch.Tensor):
@@ -1396,6 +1418,7 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False):
                 router_logits=staged_router_logits,
                 top_k=self.top_k,
                 renormalize=self.renormalize,
+                layer_index=self.layer_index,
                 use_grouped_topk=self.use_grouped_topk,
                 global_num_experts=self.global_num_experts,
                 expert_map=self.expert_map,
@@ -1432,7 +1455,7 @@ def process_chunk(chunk_start, chunk_end, skip_result_store=False):
         return full_final_hidden_states
 
     def forward_impl(self, hidden_states: torch.Tensor,
-                     router_logits: torch.Tensor):
+                     router_logits: torch.Tensor, layer_index: int):
         assert self.quant_method is not None
         if (self.moe_parallel_config.use_pplx_kernels
                 or self.moe_parallel_config.use_deepep_ll_kernels):
@@ -1452,6 +1475,7 @@ def forward_impl(self, hidden_states: torch.Tensor,
             router_logits=router_logits,
             top_k=self.top_k,
             renormalize=self.renormalize,
+            layer_index=layer_index,
             use_grouped_topk=self.use_grouped_topk,
             global_num_experts=self.global_num_experts,
             expert_map=self.expert_map,
@@ -1514,16 +1538,16 @@ def extra_repr(self) -> str:
 
 
 def moe_forward(hidden_states: torch.Tensor, router_logits: torch.Tensor,
-                layer_name: str) -> torch.Tensor:
+                layer_name: str, layer_index: int) -> torch.Tensor:
     forward_context: ForwardContext = get_forward_context()
     self = forward_context.no_compile_layers[layer_name]
     assert self.quant_method is not None
 
-    return self.forward_impl(hidden_states, router_logits)
+    return self.forward_impl(hidden_states, router_logits, layer_index)
 
 
 def moe_forward_fake(hidden_states: torch.Tensor, router_logits: torch.Tensor,
-                     layer_name: str) -> torch.Tensor:
+                     layer_name: str, layer_index: int) -> torch.Tensor:
     return torch.empty_like(hidden_states)