[Metrics] Add Prometheus counter for CUDA graph iteration mode

Add `vllm:cudagraph_iterations` counter with `runtime_mode` label
(NONE/PIECEWISE/FULL) so CUDA graph usage percentage can be computed
in PromQL. Also always populate CUDAGraphStat regardless of the
`--cudagraph-metrics` flag, since the stat is trivially cheap and that
flag should only gate the verbose text log table.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

examples/rl/rlhf_nccl_fsdp_ep.py

@@ -0,0 +1,339 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs).
+
+8-GPU layout:
+  Training  — 4 GPUs, PyTorch FSDP2 (fully_shard)
+  Inference — 4 GPUs, vLLM AsyncLLMEngine with expert parallelism +
+              data parallelism (TP=1, DP=4, enable_expert_parallel
+              → EP_SIZE = TP×DP = 4)
+
+FSDP workers are Ray actors that form a single FSDP2 process group.
+Rank 0 gathers full parameters via DTensor.full_tensor() and broadcasts
+them to the vLLM inference engine through the NCCL weight-transfer API.
+
+The inference engine uses AsyncLLMEngine which automatically spawns
+DP worker processes (no manual placement group needed).  Weight sync
+uses pause_generation / resume_generation.
+
+Steps:
+  1. Launch 4 FSDP training workers.
+  2. Launch AsyncLLMEngine with EP+DP (dummy weights).
+  3. Generate from prompts → gibberish (random weights).
+  4. Pause generation, transfer weights from FSDP, resume.
+  5. Generate from prompts → sensible output (synced weights).
+
+Assumes a single-node cluster with 8 GPUs.
+"""
+
+import asyncio
+import os
+import uuid
+from dataclasses import asdict
+
+import ray
+import torch
+import torch.distributed as dist
+from huggingface_hub import snapshot_download
+from torch.distributed.fsdp import fully_shard
+from transformers import AutoModelForCausalLM
+
+import vllm
+from vllm import SamplingParams
+from vllm.config import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferInitRequest,
+    WeightTransferUpdateRequest,
+)
+from vllm.distributed.weight_transfer.nccl_engine import (
+    NCCLTrainerSendWeightsArgs,
+    NCCLWeightTransferEngine,
+    NCCLWeightTransferInitInfo,
+    NCCLWeightTransferUpdateInfo,
+)
+from vllm.utils.network_utils import get_ip, get_open_port
+from vllm.v1.executor import Executor
+
+MODEL_NAME = "Qwen/Qwen3-30B-A3B"
+
+FSDP_WORLD_SIZE = 4
+INFERENCE_TP_SIZE = 1
+INFERENCE_DP_SIZE = 4
+
+
+@ray.remote(num_gpus=1)
+class FSDPTrainWorker:
+    """
+    One FSDP2 training worker per GPU.  Four of these form the FSDP group.
+    Rank 0 additionally handles weight transfer to the vLLM engine.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        rank: int,
+        fsdp_world_size: int,
+        fsdp_master_addr: str,
+        fsdp_master_port: int,
+    ):
+        self.rank = rank
+
+        os.environ["MASTER_ADDR"] = fsdp_master_addr
+        os.environ["MASTER_PORT"] = str(fsdp_master_port)
+
+        dist.init_process_group(backend="nccl", rank=rank, world_size=fsdp_world_size)
+        torch.accelerator.set_device_index(0)
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, torch_dtype=torch.bfloat16
+        )
+
+        self.weight_names = [n for n, _ in model.named_parameters()]
+        self.weight_dtype_names = [
+            str(p.dtype).split(".")[-1] for _, p in model.named_parameters()
+        ]
+        self.weight_shapes = [list(p.shape) for _, p in model.named_parameters()]
+
+        for layer in model.model.layers:
+            fully_shard(layer)
+        fully_shard(model)
+
+        self.model = model
+
+        self.transfer_port = None
+        self.transfer_master_address = None
+        self.model_update_group = None
+
+    def get_rank(self):
+        return self.rank
+
+    # ---- weight-transfer setup (rank 0 only) ----
+
+    def setup_transfer_endpoint(self):
+        """Create the NCCL rendezvous endpoint for weight transfer."""
+        assert self.rank == 0
+        self.transfer_port = get_open_port()
+        self.transfer_master_address = get_ip()
+        return self.transfer_master_address, self.transfer_port
+
+    def init_weight_transfer_group(self, transfer_world_size: int):
+        """Join the weight-transfer NCCL group as rank 0 (the source)."""
+        assert self.rank == 0
+        self.model_update_group = NCCLWeightTransferEngine.trainer_init(
+            dict(
+                master_address=self.transfer_master_address,
+                master_port=self.transfer_port,
+                world_size=transfer_world_size,
+            ),
+        )
+
+    def get_weight_metadata(self):
+        """Return weight names, dtypes, and shapes captured before FSDP wrapping."""
+        return self.weight_names, self.weight_dtype_names, self.weight_shapes
+
+    # ---- collective ops (ALL FSDP ranks must call concurrently) ----
+
+    def gather_and_broadcast_weights(self, packed: bool = True):
+        """
+        All-gather full parameters and broadcast them to vLLM.
+        Only rank 0 performs the actual NCCL broadcast; others just
+        participate in the FSDP all-gather.
+
+        full_tensor() is a collective — all FSDP ranks must call it
+        for each parameter in the same order.  Rank 0 additionally
+        feeds each gathered tensor to the weight-transfer engine.
+        """
+        if self.rank == 0:
+
+            def _full_param_iter():
+                for name, param in self.model.named_parameters():
+                    yield name, param.full_tensor()
+
+            trainer_args = NCCLTrainerSendWeightsArgs(
+                group=self.model_update_group,
+                packed=packed,
+            )
+            NCCLWeightTransferEngine.trainer_send_weights(
+                iterator=_full_param_iter(),
+                trainer_args=trainer_args,
+            )
+        else:
+            for _, param in self.model.named_parameters():
+                param.full_tensor()
+
+
+def create_async_engine(**kwargs):
+    """Create an AsyncLLMEngine directly (no subclass needed)."""
+    engine_args = vllm.AsyncEngineArgs(**kwargs)
+    vllm_config = engine_args.create_engine_config()
+    executor_class = Executor.get_class(vllm_config)
+    return vllm.AsyncLLMEngine(
+        vllm_config=vllm_config,
+        executor_class=executor_class,
+        log_requests=engine_args.enable_log_requests,
+        log_stats=not engine_args.disable_log_stats,
+    )
+
+
+async def generate_batch(engine, prompts, sampling_params):
+    """Generate completions for a batch of prompts."""
+
+    async def gen_one(prompt):
+        output = None
+        async for request_output in engine.generate(
+            {"prompt": prompt},
+            sampling_params,
+            request_id=str(uuid.uuid4()),
+        ):
+            output = request_output
+        return output
+
+    return await asyncio.gather(*[gen_one(p) for p in prompts])
+
+
+async def main():
+    ray.init()
+
+    # Download model weights to local/shared disk once.
+    local_model_path = snapshot_download(MODEL_NAME)
+    print(f"[init] Model downloaded to {local_model_path}")
+
+    # FSDP rendezvous address (single-node)
+    fsdp_master_addr = get_ip()
+    fsdp_master_port = get_open_port()
+
+    # Launch 4 FSDP training workers.
+    # Ray allocates 1 GPU per worker; AsyncLLMEngine's internal DP
+    # placement groups will land on the remaining 4 GPUs.
+    fsdp_workers = [
+        FSDPTrainWorker.remote(
+            local_model_path,
+            rank,
+            FSDP_WORLD_SIZE,
+            fsdp_master_addr,
+            fsdp_master_port,
+        )
+        for rank in range(FSDP_WORLD_SIZE)
+    ]
+    ray.get([w.get_rank.remote() for w in fsdp_workers])
+    print(f"[init] {FSDP_WORLD_SIZE} FSDP training workers ready.")
+
+    # Launch vLLM with expert parallelism + data parallelism.
+    # AsyncLLMEngine with data_parallel_backend="ray" creates its own
+    # placement groups internally — no manual placement group needed.
+    print("[engine] Creating AsyncLLMEngine...")
+    engine = create_async_engine(
+        model=local_model_path,
+        enforce_eager=True,
+        tensor_parallel_size=INFERENCE_TP_SIZE,
+        data_parallel_size=INFERENCE_DP_SIZE,
+        enable_expert_parallel=True,
+        distributed_executor_backend="ray",
+        data_parallel_backend="ray",
+        weight_transfer_config=WeightTransferConfig(backend="nccl"),
+        load_format="dummy",
+        gpu_memory_utilization=0.7,
+    )
+    print("[engine] AsyncLLMEngine created.")
+
+    prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
+    sampling_params = SamplingParams(temperature=0)
+
+    # Generate with dummy weights — expect gibberish.
+    print("[generate] Starting generation with dummy weights...")
+    outputs = await generate_batch(engine, prompts, sampling_params)
+    print("[generate] Generation complete.")
+
+    print("-" * 60)
+    print("BEFORE weight sync (dummy weights):")
+    print("-" * 60)
+    for output in outputs:
+        print(f"Prompt: {output.prompt!r}")
+        print(f"Generated: {output.outputs[0].text!r}")
+        print("-" * 60)
+
+    # --- Weight-transfer setup ---
+    print("[transfer] Setting up weight-transfer endpoint...")
+    transfer_addr, transfer_port = ray.get(
+        fsdp_workers[0].setup_transfer_endpoint.remote()
+    )
+    print(f"[transfer] Endpoint ready at {transfer_addr}:{transfer_port}")
+
+    transfer_world_size = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE + 1
+    print(
+        f"[transfer] World size: {transfer_world_size} "
+        f"(1 trainer + {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} vLLM workers)"
+    )
+
+    print("[transfer] Initializing NCCL groups...")
+    train_handle = fsdp_workers[0].init_weight_transfer_group.remote(
+        transfer_world_size
+    )
+    await engine.init_weight_transfer_engine(
+        WeightTransferInitRequest(
+            init_info=asdict(
+                NCCLWeightTransferInitInfo(
+                    master_address=transfer_addr,
+                    master_port=transfer_port,
+                    rank_offset=1,
+                    world_size=transfer_world_size,
+                )
+            )
+        )
+    )
+    ray.get(train_handle)
+    print("[transfer] NCCL groups initialized.")
+
+    # --- Pause, transfer weights, resume ---
+    print("[sync] Pausing generation...")
+    await engine.pause_generation(mode="abort")
+    print("[sync] Generation paused.")
+
+    names, dtype_names, shapes = ray.get(fsdp_workers[0].get_weight_metadata.remote())
+    print(f"[sync] Got metadata for {len(names)} parameters.")
+
+    print("[sync] Broadcasting weights from FSDP → vLLM...")
+    broadcast_handles = [
+        w.gather_and_broadcast_weights.remote(packed=True) for w in fsdp_workers
+    ]
+    await engine.update_weights(
+        WeightTransferUpdateRequest(
+            update_info=asdict(
+                NCCLWeightTransferUpdateInfo(
+                    names=names,
+                    dtype_names=dtype_names,
+                    shapes=shapes,
+                    packed=True,
+                )
+            )
+        )
+    )
+    ray.get(broadcast_handles)
+    print("[sync] Weight broadcast complete.")
+
+    print("[sync] Resuming generation...")
+    await engine.resume_generation()
+    print("[sync] Generation resumed.")
+
+    # Generate with synced weights — expect sensible output.
+    print("[generate] Starting generation with synced weights...")
+    outputs_updated = await generate_batch(engine, prompts, sampling_params)
+    print("[generate] Generation complete.")
+
+    print("-" * 60)
+    print("AFTER weight sync (real weights):")
+    print("-" * 60)
+    for output in outputs_updated:
+        print(f"Prompt: {output.prompt!r}")
+        print(f"Generated: {output.outputs[0].text!r}")
+        print("-" * 60)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

tests/distributed/test_eplb_algo.py

@@ -5,6 +5,7 @@ import numpy as np
 import pytest
 import torch
 
+from vllm.distributed.eplb.eplb_state import compute_logical_maps
 from vllm.distributed.eplb.policy.default import DefaultEplbPolicy
 
 
@@ -24,9 +25,10 @@ def test_basic_rebalance():
     num_nodes = 2
     num_gpus = 8
 
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
         weight, num_replicas, num_groups, num_nodes, num_gpus
     )
+    log2phy, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
 
     # Verify output shapes
     assert phy2log.shape == (
@@ -78,9 +80,10 @@ def test_single_gpu_case():
     num_nodes = 1
     num_gpus = 1
 
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
         weight, num_replicas, num_groups, num_nodes, num_gpus
     )
+    log2phy, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
 
     # Verify shapes
     assert phy2log.shape == (1, 4)
@@ -100,9 +103,10 @@ def test_equal_weights():
     num_nodes = 2
     num_gpus = 4
 
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
         weight, num_replicas, num_groups, num_nodes, num_gpus
     )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
 
     # Verify shapes
     assert phy2log.shape == (1, 8)
@@ -123,9 +127,10 @@ def test_extreme_weight_imbalance():
     num_nodes = 2
     num_gpus = 4
 
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
         weight, num_replicas, num_groups, num_nodes, num_gpus
     )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
 
     # Verify shapes
     assert phy2log.shape == (1, 12)
@@ -151,9 +156,10 @@ def test_multiple_layers():
     num_nodes = 2
     num_gpus = 4
 
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
         weight, num_replicas, num_groups, num_nodes, num_gpus
     )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
 
     # Verify shapes
     assert phy2log.shape == (3, 8)
@@ -176,7 +182,8 @@ def test_parameter_validation():
     # Test non-divisible case - this should handle normally without throwing
     # errors because the function will fall back to global load balancing
     # strategy
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(weight, 8, 3, 2, 4)
+    phy2log = DefaultEplbPolicy.rebalance_experts(weight, 8, 3, 2, 4)
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
     assert phy2log.shape == (1, 8)
     assert logcnt.shape == (1, 4)
 
@@ -198,9 +205,10 @@ def test_small_scale_hierarchical():
     num_nodes = 2  # 2 nodes
     num_gpus = 4  # 4 GPUs
 
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
         weight, num_replicas, num_groups, num_nodes, num_gpus
     )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
 
     # Verify basic constraints
     assert phy2log.shape == (1, 12)
@@ -225,9 +233,10 @@ def test_global_load_balance_fallback():
     num_nodes = 2
     num_gpus = 4
 
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
         weight, num_replicas, num_groups, num_nodes, num_gpus
     )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
 
     # Should work normally, just using global load balancing strategy
     assert phy2log.shape == (1, 8)
@@ -247,9 +256,10 @@ def test_device_compatibility(device):
     num_nodes = 1
     num_gpus = 2
 
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
         weight, num_replicas, num_groups, num_nodes, num_gpus
     )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
 
     # Function will convert to CPU internally, but should handle different
     # device inputs normally
@@ -264,9 +274,8 @@ def test_additional_cases():
     weight1 = torch.tensor(
         [[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]]
     )
-    phy2log1, log2phy1, logcnt1 = DefaultEplbPolicy.rebalance_experts(
-        weight1, 24, 8, 4, 8
-    )
+    phy2log1 = DefaultEplbPolicy.rebalance_experts(weight1, 24, 8, 4, 8)
+    _, logcnt1 = compute_logical_maps(phy2log1, weight1.shape[-1])
 
     assert phy2log1.shape == (1, 24)
     assert logcnt1.shape == (1, 16)
@@ -279,9 +288,8 @@ def test_additional_cases():
             [12, 25, 50, 100, 150, 200],  # Increasing weights
         ]
     )
-    phy2log2, log2phy2, logcnt2 = DefaultEplbPolicy.rebalance_experts(
-        weight2, 10, 3, 1, 2
-    )
+    phy2log2 = DefaultEplbPolicy.rebalance_experts(weight2, 10, 3, 1, 2)
+    _, logcnt2 = compute_logical_maps(phy2log2, weight2.shape[-1])
 
     assert phy2log2.shape == (2, 10)
     assert logcnt2.shape == (2, 6)
@@ -292,6 +300,42 @@ def test_additional_cases():
         assert logcnt2[layer, max_weight_idx] >= 2
 
 
+def test_compute_logical_maps_with_negative_indices():
+    """
+    Test that compute_logical_maps correctly handles physical slots containing
+    -1 (unused slots).
+    """
+    # 2 layers, 6 physical slots, 4 logical experts.
+    # Slots 2 and 5 are unused (-1).
+    phy2log = torch.tensor(
+        [
+            [0, 1, -1, 2, 3, -1],
+            [3, -1, 2, 1, 0, -1],
+        ]
+    )
+    num_layers = 2
+    num_logical_experts = 4
+
+    log2phy, logcnt = compute_logical_maps(phy2log, num_logical_experts)
+
+    assert logcnt.shape == (num_layers, num_logical_experts)
+    assert log2phy.shape == (num_layers, num_logical_experts, 1)
+
+    expected_logcnt = torch.ones(num_layers, num_logical_experts, dtype=phy2log.dtype)
+    assert torch.all(logcnt == expected_logcnt), (
+        f"Expected that all replica counts == 1, got {logcnt}"
+    )
+
+    assert torch.all(log2phy >= 0), (
+        "log2phy should only contain valid physical indices, not -1"
+    )
+
+    assert log2phy[0, 0, 0] == 0
+    assert log2phy[0, 1, 0] == 1
+    assert log2phy[0, 2, 0] == 3
+    assert log2phy[0, 3, 0] == 4
+
+
 if __name__ == "__main__":
     weight = torch.tensor(
         [
@@ -305,7 +349,7 @@ if __name__ == "__main__":
     num_nodes = 2
     num_gpus = 8
 
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
         weight, num_replicas, num_groups, num_nodes, num_gpus
     )
     print(phy2log)
@@ -434,9 +478,10 @@ def test_preserve_intragpu_slots(
     """Experts that stay on a GPU keep their old slots; incoming not lost."""
     phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(new_phy2log)
 
-    post_phy2log, post_phy_replicas_idx = DefaultEplbPolicy.preserve_intragpu_slots(
-        new_phy2log, phy_replicas_idx, num_ranks, old_phy2log
+    post_phy2log = DefaultEplbPolicy.preserve_intragpu_slots(
+        new_phy2log, num_ranks, old_phy2log
     )
+    post_phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(post_phy2log)
 
     # Shapes preserved
     assert post_phy2log.shape == new_phy2log.shape

vllm/distributed/eplb/async_worker.py

@@ -73,11 +73,7 @@ def run_rebalance_experts(
     # Move the global expert load window to CPU for computation.
     global_expert_load_window = eplb_stats.global_expert_load_window.cpu()
     # Compute new expert mappings for the model
-    (
-        new_physical_to_logical_map,
-        new_logical_to_physical_map,
-        new_logical_replica_count,
-    ) = eplb_state.policy.rebalance_experts(
+    new_physical_to_logical_map = eplb_state.policy.rebalance_experts(
         global_expert_load_window,
         eplb_stats.num_replicas,
         eplb_stats.num_groups,
@@ -89,16 +85,6 @@ def run_rebalance_experts(
 
     model_state.new_physical_to_logical_map = new_physical_to_logical_map
 
-    max_slots = model_state.logical_to_physical_map.shape[-1]
-    padded_logical = torch.nn.functional.pad(
-        new_logical_to_physical_map,
-        (0, max(0, max_slots - new_logical_to_physical_map.shape[-1])),
-        value=-1,
-    ).to(model_state.logical_to_physical_map.device)
-    new_replica = new_logical_replica_count.to(model_state.logical_replica_count.device)
-    model_state.new_logical_to_physical_map = padded_logical
-    model_state.new_logical_replica_count = new_replica
-
 
 async def transfer_run_periodically(
     state: "EplbState",

vllm/distributed/eplb/eplb_state.py

@@ -235,16 +235,6 @@ class EplbModelState:
     intermediate variable between `move_to_buffer` and `move_to_workspace`.
     the size is same as physical_to_logical_map
     """
-    new_logical_to_physical_map: torch.Tensor | None = None
-    """
-    intermediate variable between `move_to_buffer` and `move_to_workspace`.
-    the size is same as logical_to_physical_map
-    """
-    new_logical_replica_count: torch.Tensor | None = None
-    """
-    intermediate variable between `move_to_buffer` and `move_to_workspace`.
-    the size is same as logical_replica_count
-    """
 
 
 class EplbState:
@@ -508,8 +498,6 @@ class EplbState:
             ),
             cuda_device_index=self.cuda_device_index,
             new_physical_to_logical_map=None,
-            new_logical_to_physical_map=None,
-            new_logical_replica_count=None,
         )
         self.model_states[model_config.compute_hash()] = model_state
         self.num_valid_physical_experts = model.num_physical_experts
@@ -738,17 +726,20 @@ class EplbState:
         ):
             if not self.is_async or is_profile:
                 # Get new expert mappings for the model
-                (
-                    new_physical_to_logical_map,
-                    new_logical_to_physical_map,
-                    new_logical_replica_count,
-                ) = self.policy.rebalance_experts(
-                    global_expert_load_window,
+                new_physical_to_logical_map = self.policy.rebalance_experts(
+                    global_expert_load_window.cpu(),
                     num_replicas,
                     num_groups,
                     num_nodes,
                     num_gpus,
-                    eplb_model_state.physical_to_logical_map,
+                    eplb_model_state.physical_to_logical_map.cpu(),
+                )
+
+                num_logical_experts = global_expert_load_window.shape[-1]
+                (new_logical_to_physical_map, new_logical_replica_count) = (
+                    compute_logical_maps(
+                        new_physical_to_logical_map, num_logical_experts
+                    )
                 )
 
                 # Update expert weights
@@ -847,11 +838,7 @@ class EplbState:
     def _update_layer_mapping_from_new(
         self, model_state: EplbModelState, layer: int
     ) -> None:
-        if (
-            model_state.new_physical_to_logical_map is None
-            or model_state.new_logical_to_physical_map is None
-            or model_state.new_logical_replica_count is None
-        ):
+        if model_state.new_physical_to_logical_map is None:
             return
 
         target_device = model_state.physical_to_logical_map.device
@@ -865,19 +852,23 @@ class EplbState:
                 new_physical[layer].to(target_device, non_blocking=True)
             )
 
+        num_logical_experts = model_state.logical_to_physical_map.shape[1]
+        new_logical, new_replica_count = compute_logical_maps(
+            new_physical[layer], num_logical_experts
+        )
+
         logical_device = model_state.logical_to_physical_map.device
-        new_logical = model_state.new_logical_to_physical_map[layer].to(logical_device)
         max_slots = model_state.logical_to_physical_map.shape[-1]
         slot_delta = max_slots - new_logical.shape[-1]
         if slot_delta > 0:
             new_logical = torch.nn.functional.pad(
                 new_logical, (0, slot_delta), value=-1
             )
-        model_state.logical_to_physical_map[layer].copy_(new_logical)
+        model_state.logical_to_physical_map[layer].copy_(new_logical.to(logical_device))
 
         replica_device = model_state.logical_replica_count.device
         model_state.logical_replica_count[layer].copy_(
-            model_state.new_logical_replica_count[layer].to(replica_device)
+            new_replica_count.to(replica_device)
         )
 
     def _all_ranks_buffer_ready(self, model_state: EplbModelState) -> bool:
@@ -966,7 +957,7 @@ class EplbState:
                 transferred_layer,
             )
             if model_state.layer_to_transfer >= model_state.model.num_moe_layers:
-                self.post_eplb(model_state, is_profile)
+                self.post_eplb(model_state)
                 model_state.rebalanced = False
                 model_state.layer_to_transfer = 0
                 model_state.pending_global_ready_check = False
@@ -987,14 +978,9 @@ class EplbState:
                     str(e),
                 )
 
-    def post_eplb(self, model_state: EplbModelState, is_profile: bool = False) -> None:
+    def post_eplb(self, model_state: EplbModelState) -> None:
         assert model_state.new_physical_to_logical_map is not None
-        assert model_state.new_logical_to_physical_map is not None
-        assert model_state.new_logical_replica_count is not None
-
         model_state.new_physical_to_logical_map = None
-        model_state.new_logical_to_physical_map = None
-        model_state.new_logical_replica_count = None
 
     def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]:
         """
@@ -1052,39 +1038,28 @@ class EplbState:
             model_config=model_config,
         )
         eplb_state.num_valid_physical_experts = num_valid_physical_experts
-        num_moe_layers = expanded_physical_to_logical.shape[0]
-        num_physical_experts = expanded_physical_to_logical.shape[1]
         eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
         eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical)
 
-        logical_to_physical_map = torch.full(
-            (
-                num_moe_layers,
-                model.num_logical_experts,
-                eplb_model_state.logical_to_physical_map.shape[2],
-            ),
-            -1,
-            dtype=torch.int64,
+        (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps(
+            expanded_physical_to_logical.cpu(), model.num_logical_experts
         )
-        logical_replica_count = torch.zeros(
-            (num_moe_layers, model.num_logical_experts),
-            dtype=torch.int64,
-        )
-        expanded_physical_to_logical_numpy = expanded_physical_to_logical.cpu().numpy()
-        for layer_idx in range(num_moe_layers):
-            for phys_idx in range(num_physical_experts):
-                logical_idx = expanded_physical_to_logical_numpy[layer_idx, phys_idx]
-                if logical_idx >= 0:
-                    replica_idx = logical_replica_count[layer_idx, logical_idx]
-                    logical_to_physical_map[layer_idx, logical_idx, replica_idx] = (
-                        phys_idx
-                    )
-                    logical_replica_count[layer_idx, logical_idx] += 1
 
-        logical_to_physical_map = logical_to_physical_map.to(device)
-        logical_replica_count = logical_replica_count.to(device)
+        max_num_replicas = eplb_model_state.logical_to_physical_map.shape[-1]
+        num_replicas = logical_to_physical_map_cpu.shape[-1]
+        logical_to_physical_map = torch.nn.functional.pad(
+            logical_to_physical_map_cpu,
+            (
+                0,
+                max_num_replicas - num_replicas,
+            ),
+            value=-1,
+        ).to(device)
+        logical_replica_count = logical_replica_count_cpu.to(device)
+
         eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map)
         eplb_model_state.logical_replica_count.copy_(logical_replica_count)
+
         return eplb_state
 
 
@@ -1132,3 +1107,82 @@ def _node_count_with_rank_mapping(
                 node_assignment[other_rank] = next_node_id
 
     return next_node_id
+
+
+def compute_logical_maps(
+    physical_to_logical_map: torch.Tensor,
+    num_logical_experts: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Derive logical_to_physical_map and logical_replica_count from
+    physical_to_logical_map.
+
+    Args:
+        physical_to_logical_map: [num_layers, num_physical_experts], logical
+            expert index for each physical expert slot
+        num_logical_experts: total number of logical experts
+
+    Returns:
+        logical_to_physical_map: [num_layers, num_logical_experts, max_replicas],
+            physical slots per logical expert; -1 where unused
+        logical_replica_count: [num_layers, num_logical_experts], number of
+            physical replicas per logical expert
+    """
+    device = physical_to_logical_map.device
+    assert physical_to_logical_map.device.type == "cpu"
+
+    dtype = physical_to_logical_map.dtype
+
+    # If computing maps for a single layer, unsqueeze a single element layer dimension
+    per_layer = physical_to_logical_map.dim() == 1
+    physical_to_logical_map_view = physical_to_logical_map
+    if per_layer:
+        physical_to_logical_map_view = physical_to_logical_map.unsqueeze(0)
+    assert len(physical_to_logical_map_view.shape) == 2
+    num_layers, num_physical = physical_to_logical_map_view.shape
+
+    valid_mask = physical_to_logical_map_view >= 0
+    logical_replica_count = torch.zeros(
+        num_layers,
+        num_logical_experts,
+        dtype=dtype,
+        device=device,
+    )
+    logical_replica_count.scatter_add_(
+        1,
+        physical_to_logical_map_view.clamp(min=0),
+        valid_mask.to(dtype),
+    )
+
+    max_replicas = int(logical_replica_count.max().item())
+    logical_to_physical_map_out = torch.full(
+        (num_layers, num_logical_experts, max_replicas),
+        -1,
+        dtype=dtype,
+        device=device,
+    )
+
+    running_count = torch.zeros_like(logical_replica_count)
+    layer_indices = torch.arange(num_layers, device=device)
+    for phys_idx in range(num_physical):
+        # Logical expert at physical slot phys_idx for each layer
+        logical_expert_ids = physical_to_logical_map_view[:, phys_idx]  # [num_layers]
+
+        # Scale up will set the logical expert ids to -1 for all new physical experts.
+        # Only consider "valid" experts when setting up the logical_to_physical map.
+        valid_expert_mask = logical_expert_ids >= 0
+        if not valid_expert_mask.any():
+            continue
+        valid_layers = layer_indices[valid_expert_mask]
+        valid_experts = logical_expert_ids[valid_expert_mask]
+
+        # Use the current running count as the replica index, then increment it.
+        replica_idx = running_count[valid_layers, valid_experts]
+        logical_to_physical_map_out[valid_layers, valid_experts, replica_idx] = phys_idx
+        running_count[valid_layers, valid_experts] += 1
+
+    # If computing maps for a single layer, squeeze out the extra layer dimension
+    # before returning
+    if per_layer:
+        return logical_to_physical_map_out.squeeze(0), logical_replica_count.squeeze(0)
+    return logical_to_physical_map_out, logical_replica_count

vllm/distributed/eplb/policy/abstract.py

@@ -17,7 +17,7 @@ class AbstractEplbPolicy(ABC):
         num_nodes: int,
         num_ranks: int,
         old_global_expert_indices: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> torch.Tensor:
         """
         Entry point for expert-parallelism load balancer.
 
@@ -35,9 +35,5 @@ class AbstractEplbPolicy(ABC):
         Returns:
             physical_to_logical_map: [layers, num_replicas], the expert
                 index of each replica
-            logical_to_physical_map: [layers, num_logical_experts, X],
-                the replica indices for each expert
-            expert_count: [layers, num_logical_experts], number of
-                physical replicas for each logical expert
         """
         raise NotImplementedError

vllm/distributed/eplb/policy/default.py

@@ -75,7 +75,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
     @classmethod
     def replicate_experts(
         cls, weight: np.ndarray, num_phy: int
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray]:
         """
         Replicate `num_log` experts to `num_phy` replicas, such that the maximum
         load of all replicas is minimized.
@@ -86,22 +86,19 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
 
         Returns:
             phy2log: [X, num_phy], logical expert id of each physical expert
-            replica_idx: [X, num_phy], the index of the replica for each logical expert
             logcnt: [X, num_log], number of replicas for each logical expert
         """
         n, num_log = weight.shape
         num_redundant = num_phy - num_log
         assert num_redundant >= 0
         phy2log = np.tile(np.arange(num_phy, dtype=np.int64), (n, 1))
-        replica_idx = np.zeros((n, num_phy), dtype=np.int64)
         logcnt = np.ones((n, num_log), dtype=np.int64)
         arangen = np.arange(n, dtype=np.int64)
         for i in range(num_log, num_phy):
             redundant_indices = np.argmax(weight / logcnt, axis=-1)
             phy2log[:, i] = redundant_indices
-            replica_idx[:, i] = logcnt[arangen, redundant_indices]
             logcnt[arangen, redundant_indices] += 1
-        return phy2log, replica_idx, logcnt
+        return phy2log, logcnt
 
     @classmethod
     def rebalance_experts_hierarchical(
@@ -111,7 +108,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
         num_groups: int,
         num_nodes: int,
         num_gpus: int,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    ) -> np.ndarray:
         """
         Parameters:
             weight: [num_moe_layers, num_logical_experts]
@@ -124,10 +121,6 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
         Returns:
             phy2log: [layers, num_replicas], the expert
                 index of each replica
-            pphy_replicas_idx: [layers, num_logical_experts, X],
-                the replica indices for each expert
-            logcnt: [layers, num_logical_experts], number of
-                physical replicas for each logical expert
         """
         num_layers, num_logical_experts = weight.shape
         assert num_logical_experts % num_groups == 0
@@ -167,7 +160,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
         tokens_per_mlog = np.take_along_axis(weight, mlog2log, axis=1).reshape(
             -1, num_logical_experts // num_nodes
         )
-        phy2mlog, replicas_idx, mlogcnt = cls.replicate_experts(
+        phy2mlog, mlogcnt = cls.replicate_experts(
             tokens_per_mlog, num_physical_experts // num_nodes
         )
 
@@ -193,22 +186,15 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
         ).reshape(num_layers, -1)
         # Map node-local logical indices back to global logical expert ids.
         pphy2log = np.take_along_axis(mlog2log, pphy2mlog, axis=1)
-        # Reorder replica ranks to the post-packing physical ordering.
-        pphy_replicas_idx = np.take_along_axis(replicas_idx, pphy2phy, axis=1).reshape(
-            num_layers, -1
-        )
-        # Convert replica counts back to the original logical ordering.
-        logcnt = np.take_along_axis(mlogcnt.reshape(num_layers, -1), log2mlog, axis=1)
-        return pphy2log, pphy_replicas_idx, logcnt
+        return pphy2log
 
     @classmethod
     def preserve_intragpu_slots(
         cls,
         phy2log: np.ndarray,
-        phy_replicas_idx: np.ndarray,
         num_ranks: int,
         old_phy2log: np.ndarray,
-    ) -> tuple[np.ndarray, np.ndarray]:
+    ) -> np.ndarray:
         """
         Reorder the new mapping per GPU so that experts that remain on the same GPU
         keep their previous slot positions when possible. Incoming experts to that GPU
@@ -218,14 +204,13 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
         """
         num_phy_experts = phy2log.shape[1]
         if num_ranks <= 0 or num_phy_experts % num_ranks != 0:
-            return phy2log, phy_replicas_idx
+            return phy2log
 
         # Move to CPU and convert to NumPy for processing
         slots_per_gpu = num_phy_experts // num_ranks
         num_layers = phy2log.shape[0]
 
         post_phy2log = phy2log.copy()
-        post_phy_replicas_idx = phy_replicas_idx.copy()
 
         for gpu_idx in range(num_ranks):
             start = gpu_idx * slots_per_gpu
@@ -233,7 +218,6 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
             # Experts across all layers for this GPU
             old_local = old_phy2log[:, start:end]  # [layers, slots]
             new_local = phy2log[:, start:end]  # [layers, slots]
-            new_ridx = phy_replicas_idx[:, start:end]  # [layers, slots]
 
             used_new_indices = np.zeros((num_layers, slots_per_gpu), dtype=bool)
             preserved_positions = np.zeros((num_layers, slots_per_gpu), dtype=bool)
@@ -253,9 +237,6 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
                     post_phy2log[layer_indices, start + slot_idx] = new_local[
                         layer_indices, matched_new_positions
                     ]
-                    post_phy_replicas_idx[layer_indices, start + slot_idx] = new_ridx[
-                        layer_indices, matched_new_positions
-                    ]
                     used_new_indices[layer_indices, matched_new_positions] = True
                     preserved_positions[layer_indices, slot_idx] = True
 
@@ -287,11 +268,8 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
                     post_phy2log[layer_idx, start + dst_pos] = new_local[
                         layer_idx, src_pos
                     ]
-                    post_phy_replicas_idx[layer_idx, start + dst_pos] = new_ridx[
-                        layer_idx, src_pos
-                    ]
 
-        return post_phy2log, post_phy_replicas_idx
+        return post_phy2log
 
     @classmethod
     def rebalance_experts(
@@ -302,7 +280,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
         num_nodes: int,
         num_ranks: int,
         old_global_expert_indices: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> torch.Tensor:
         """
         Entry point for expert-parallelism load balancer.
 
@@ -321,13 +299,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
         Returns:
             phy2log: [layers, num_replicas], the expert
                 index of each replica
-            log2phy: [layers, num_logical_experts, X],
-                the replica indices for each expert
-            logcnt: [layers, num_logical_experts], number of
-                physical replicas for each logical expert
         """
-        device = weight.device
-        num_layers, num_logical_experts = weight.shape
         weight_np = weight.float().cpu().numpy()
         old_phy2log_np = (
             old_global_expert_indices.cpu().numpy()
@@ -337,17 +309,13 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
 
         if num_groups % num_nodes == 0:
             # use hierarchical load-balance policy
-            phy2log_np, phy_replicas_idx_np, logcnt_np = (
-                cls.rebalance_experts_hierarchical(
-                    weight_np, num_replicas, num_groups, num_nodes, num_ranks
-                )
+            phy2log_np = cls.rebalance_experts_hierarchical(
+                weight_np, num_replicas, num_groups, num_nodes, num_ranks
             )
         else:
             # use global load-balance policy
-            phy2log_np, phy_replicas_idx_np, logcnt_np = (
-                cls.rebalance_experts_hierarchical(
-                    weight_np, num_replicas, 1, 1, num_ranks
-                )
+            phy2log_np = cls.rebalance_experts_hierarchical(
+                weight_np, num_replicas, 1, 1, num_ranks
             )
 
         # Optional postprocessing to preserve slots for experts moving
@@ -355,22 +323,10 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
         # Only apply when the number of GPUs and slots per GPU remain unchanged.
         # Helps to avoid unnecessary weight copying when experts move
         # within the same GPU.
-        if old_global_expert_indices is not None:
-            phy2log_np, phy_replicas_idx_np = cls.preserve_intragpu_slots(
-                phy2log_np, phy_replicas_idx_np, num_ranks, old_phy2log_np
+        if old_phy2log_np is not None:
+            phy2log_np = cls.preserve_intragpu_slots(
+                phy2log_np, num_ranks, old_phy2log_np
             )
-        num_redundant_experts = num_replicas - num_logical_experts
-        maxlogcnt = num_redundant_experts + 1
-        log2phy_np = np.full(
-            (num_layers, num_logical_experts, maxlogcnt), -1, dtype=np.int64
-        )
-        layer_indices = np.arange(num_layers)[:, None]
-        replica_indices = np.tile(
-            np.arange(num_replicas, dtype=np.int64), (num_layers, 1)
-        )
-        log2phy_np[layer_indices, phy2log_np, phy_replicas_idx_np] = replica_indices
 
-        phy2log = torch.from_numpy(phy2log_np).to(device)
-        log2phy = torch.from_numpy(log2phy_np).to(device)
-        logcnt = torch.from_numpy(logcnt_np).to(device)
-        return phy2log, log2phy, logcnt
+        phy2log = torch.from_numpy(phy2log_np)
+        return phy2log

vllm/model_executor/model_loader/reload/utils.py

@@ -27,5 +27,15 @@ def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors:
 
 
 def get_layer_size(layer: torch.nn.Module) -> int:
-    """Calculate total number of elements across all tensors in a layer."""
-    return sum(tensor.numel() for tensor in get_layer_tensors(layer).values())
+    """Calculate total number of elements across loadable tensors in a layer.
+
+    Excludes SKIP_TENSORS (e.g. _expert_map) which are never moved to meta
+    device and never loaded via weight_loader during layerwise reload.
+    """
+    from .meta import SKIP_TENSORS
+
+    return sum(
+        tensor.numel()
+        for name, tensor in get_layer_tensors(layer).items()
+        if name not in SKIP_TENSORS
+    )

vllm/v1/engine/core.py

@@ -1632,7 +1632,11 @@ class DPEngineCoreProc(EngineCoreProc):
         if self.has_coordinator and request_wave != self.current_wave:
             if request_wave > self.current_wave:
                 self.current_wave = request_wave
-            elif not self.engines_running:
+            elif (
+                not self.engines_running
+                and self.scheduler.pause_state == PauseState.UNPAUSED
+            ):
+                self.engines_running = True
                 # Request received for an already-completed wave, notify
                 # front-end that we need to start the next one.
                 self.output_queue.put_nowait(

vllm/v1/metrics/loggers.py

@@ -413,6 +413,7 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
 
         labelnames = ["model_name", "engine"]
         model_name = vllm_config.model_config.served_model_name
+        self.model_name = model_name
         max_model_len = vllm_config.model_config.max_model_len
 
         per_engine_labelvalues: dict[int, list[object]] = {
@@ -975,6 +976,18 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
             self.histogram_kv_block_idle_before_evict = {}
             self.histogram_kv_block_reuse_gap = {}
 
+        #
+        # CUDAGraph metrics
+        #
+        self._counter_cudagraph_iterations_base = self._counter_cls(
+            name="vllm:cudagraph_iterations",
+            documentation=(
+                "Number of engine iterations by CUDA graph runtime mode."
+            ),
+            labelnames=labelnames + ["runtime_mode"],
+        )
+        self.counter_cudagraph_iterations: dict[str, dict[int, Counter]] = {}
+
         #
         # LoRA metrics
         #
@@ -1086,6 +1099,17 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
                     for gap in event.reuse_gaps_seconds:
                         reuse_hist.observe(gap)
 
+            if scheduler_stats.cudagraph_stats is not None:
+                mode = scheduler_stats.cudagraph_stats.runtime_mode
+                if mode not in self.counter_cudagraph_iterations:
+                    self.counter_cudagraph_iterations[mode] = {
+                        idx: self._counter_cudagraph_iterations_base.labels(
+                            self.model_name, str(idx), mode
+                        )
+                        for idx in self.engine_indexes
+                    }
+                self.counter_cudagraph_iterations[mode][engine_idx].inc()
+
             if self.gauge_lora_info is not None:
                 running_lora_adapters = ",".join(
                     scheduler_stats.running_lora_adapters.keys()

vllm/v1/worker/gpu/spec_decode/rejection_sampler.py

@@ -3,11 +3,14 @@
 import torch
 
 from vllm.triton_utils import tl, triton
+from vllm.v1.outputs import LogprobsTensors
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.metrics.logits import get_num_nans
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
 from vllm.v1.worker.gpu.sample.output import SamplerOutput
 from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS
 
 
 @triton.jit
@@ -418,6 +421,26 @@ def probabilistic_rejection_sample(
     return sampled, rejected_steps + 1
 
 
+@triton.jit
+def _flatten_sampled_kernel(
+    # [num_logits]
+    flat_sampled_ptr,
+    # [num_reqs, num_speculative_steps + 1]
+    sampled_ptr,
+    sampled_stride,
+    # [num_reqs]
+    num_sampled_ptr,
+    # [num_reqs + 1]
+    cu_num_logits_ptr,
+):
+    req_idx = tl.program_id(0)
+    start_idx = tl.load(cu_num_logits_ptr + req_idx)
+    num_sampled = tl.load(num_sampled_ptr + req_idx)
+    for i in range(num_sampled):
+        token_id = tl.load(sampled_ptr + req_idx * sampled_stride + i)
+        tl.store(flat_sampled_ptr + start_idx + i, token_id)
+
+
 class RejectionSampler:
     def __init__(
         self,
@@ -429,6 +452,40 @@ class RejectionSampler:
         self.num_speculative_steps = num_speculative_steps
         self.use_strict_rejection_sampling = use_strict_rejection_sampling
 
+    def _get_logprobs_tensors(
+        self,
+        input_batch: InputBatch,
+        sampled: torch.Tensor,
+        num_sampled: torch.Tensor,
+        logits: torch.Tensor,
+    ) -> LogprobsTensors | None:
+        max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
+            input_batch.idx_mapping_np
+        )
+        if max_num_logprobs == NO_LOGPROBS:
+            return None
+
+        num_reqs = input_batch.cu_num_logits.shape[0] - 1
+        num_logits = logits.shape[0]
+        flat_sampled = torch.zeros(
+            num_logits, dtype=sampled.dtype, device=sampled.device
+        )
+        _flatten_sampled_kernel[(num_reqs,)](
+            flat_sampled,
+            sampled,
+            sampled.stride(0),
+            num_sampled,
+            input_batch.cu_num_logits,
+            num_warps=1,
+        )
+        expanded_logits = num_logits != input_batch.idx_mapping.shape[0]
+        return compute_topk_logprobs(
+            logits,
+            max_num_logprobs,
+            flat_sampled,
+            input_batch.cu_num_logits_np.tolist() if expanded_logits else None,
+        )
+
     def __call__(
         self,
         logits: torch.Tensor,
@@ -460,8 +517,6 @@ class RejectionSampler:
                 draft_sampled,
                 input_batch.expanded_local_pos,
             )
-            # TODO (TheEpicDolphin): Return logprobs for sampled token ids.
-            logprobs_tensors = None
             sampled, num_sampled = probabilistic_rejection_sample(
                 processed_logits,
                 draft_logits,
@@ -475,6 +530,14 @@ class RejectionSampler:
                 self.sampler.sampling_states.seeds.gpu,
                 self.num_speculative_steps,
             )
+            logprobs_tensors = self._get_logprobs_tensors(
+                input_batch,
+                sampled,
+                num_sampled,
+                processed_logits
+                if self.sampler.logprobs_mode == "processed_logprobs"
+                else logits,
+            )
 
         return SamplerOutput(
             sampled_token_ids=sampled,

vllm/v1/worker/gpu_model_runner.py

@@ -3431,14 +3431,12 @@ class GPUModelRunner(
                 # num_tokens_across_dp will no-longer be valid
                 assert batch_descriptor.num_tokens == num_tokens_padded
 
-        cudagraph_stats = None
-        if self.vllm_config.observability_config.cudagraph_metrics:
-            cudagraph_stats = CUDAGraphStat(
-                num_unpadded_tokens=num_tokens,
-                num_padded_tokens=batch_descriptor.num_tokens,
-                num_paddings=batch_descriptor.num_tokens - num_tokens,
-                runtime_mode=str(cudagraph_mode),
-            )
+        cudagraph_stats = CUDAGraphStat(
+            num_unpadded_tokens=num_tokens,
+            num_padded_tokens=batch_descriptor.num_tokens,
+            num_paddings=batch_descriptor.num_tokens - num_tokens,
+            runtime_mode=str(cudagraph_mode),
+        )
 
         return (
             cudagraph_mode,