DeepEP HT: enable worst-token full decode graphs

vllm/config/compilation.py

@@ -1185,11 +1185,35 @@ def set_splitting_ops_for_v1(
                 )
                 self.cudagraph_mode = CUDAGraphMode.FULL
 
+        # With worst-token dispatch (UCCL ht-cudagraph-worst-tokens kernels),
+        # DeepEP HT dispatch+combine are stream-capturable: no host count
+        # sync, static worst-case recv shapes, device-side per-expert counts.
+        # Capture decode steps fully; prefill stays eager.
+        if (
+            all2all_backend == "deepep_high_throughput"
+            and data_parallel_size > 1
+            and self.cudagraph_mode != CUDAGraphMode.NONE
+            and envs.VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH
+        ):
+            if self.cudagraph_mode.has_full_cudagraphs():
+                logger.info(
+                    "DeepEP high-throughput worst-token dispatch enabled; "
+                    "capturing full decode CUDA graphs."
+                )
+                self.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
+            else:
+                logger.warning_once(
+                    "VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH requires full decode "
+                    "CUDA graphs; setting cudagraph_mode to NONE."
+                )
+                self.cudagraph_mode = CUDAGraphMode.NONE
+
         # Disable CUDA graphs for DeepEP high-throughput since its not CG compatible
         if (
             all2all_backend == "deepep_high_throughput"
             and data_parallel_size > 1
             and self.cudagraph_mode != CUDAGraphMode.NONE
+            and not envs.VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH
         ):
             # TODO: Piecewise Cuda graph might be enabled
             # if torch compile cache key issue fixed

vllm/envs.py

@@ -151,6 +151,7 @@
     VLLM_USE_STANDALONE_COMPILE: bool = True
     VLLM_ENABLE_PREGRAD_PASSES: bool = True
     VLLM_USE_BREAKABLE_CUDAGRAPH: bool = False
+    VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH: bool = False
     VLLM_DP_MASTER_IP: str = ""
     VLLM_DP_MASTER_PORT: int = 0
     VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False
@@ -719,6 +720,12 @@ def _resolve_rust_frontend_path() -> str | None:
     "VLLM_USE_BREAKABLE_CUDAGRAPH": lambda: (
         os.environ.get("VLLM_USE_BREAKABLE_CUDAGRAPH", "0") == "1"
     ),
+    # Experimental: DeepEP high-throughput dispatch with worst-case-sized
+    # recv buffers and no host count sync (CUDA-graph capturable; requires
+    # the UCCL ht-cudagraph-worst-tokens kernels).
+    "VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH": lambda: (
+        os.environ.get("VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH", "0") == "1"
+    ),
     # Debug pattern matching inside custom passes.
     # Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3').
     "VLLM_PATTERN_MATCH_DEBUG": lambda: os.environ.get(

vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py

@@ -5,7 +5,13 @@
 import deep_ep
 import torch
 
+import vllm.envs as envs
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.config import CUDAGraphMode
+from vllm.forward_context import (
+    get_forward_context,
+    is_forward_context_available,
+)
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceContiguous,
@@ -65,9 +71,68 @@ def __init__(
         # micro-batch to avoid races between threads.
         self.handles = [None, None]
 
+        # Lamport-clock counters for host-issued dispatch/combine calls,
+        # reported in the timeout annotations below. Healthy EP ranks
+        # advance these in lockstep (graph REPLAYS bump no rank's counter
+        # since no Python runs, and the runtime mode is synced across DP
+        # ranks per step), so on a collective timeout, comparing values
+        # across ranks' logs separates "all ranks in the same stuck
+        # collective" (transport) from "ranks in different collectives"
+        # (scheduling desync) — and seq mod num-moe-layers localizes the
+        # layer.
+        self.dispatch_seq = 0
+        self.combine_seq = 0
+
         # From https://github.com/deepseek-ai/DeepEP/blob/9fe9021f29c9083cd1808ab36b740208524d9f63/deep_ep/buffer.py#L164
         self.available_rank_configs = [2, 4, 8, 16, 24, 32, 64, 128, 144, 160]
 
+    def _step_debug(self) -> str:
+        try:
+            if is_forward_context_available():
+                ctx = get_forward_context()
+                return (
+                    f"mode={ctx.cudagraph_runtime_mode} "
+                    f"descriptor={ctx.batch_descriptor}"
+                )
+        except Exception:
+            pass
+        return "no-forward-context"
+
+    def _num_worst_tokens(self, tokens: torch.Tensor) -> int:
+        """Worst-case recv size for CUDA-graph-safe dispatch.
+
+        When non-zero, dispatch skips its host count sync, returns
+        worst-case-padded recv tensors with per-expert counts in a device
+        tensor, and the dispatch+combine round trip becomes
+        stream-capturable (requires the UCCL ht-cudagraph-worst-tokens
+        kernels).
+
+        Sized from the dispatched tensor itself: DP padding makes its row
+        count uniform across ranks, so it equals the step's padded token
+        count for whole-batch steps and this microbatch's padded slice
+        under DBO ubatching (the forward context's batch descriptor covers
+        the whole step there, which would oversize ubatch recv buffers 2x).
+        Eager steps (prefill, mixed, uncaptured shapes) return 0 and keep
+        the host-synced path — no padding overhead where there is no graph
+        to serve."""
+        if not envs.VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH:
+            return 0
+        if not is_forward_context_available():
+            return 0
+        ctx = get_forward_context()
+        # UBatchWrapper passes cudagraph_runtime_mode=NONE into the
+        # per-ubatch forward contexts while CAPTURING a full ubatched
+        # graph, so the context mode alone misses exactly the case that
+        # must not record the host-synced protocol (its CPU poll waits on
+        # a kernel that never executes under capture). Check the stream's
+        # capture status directly as well.
+        if (
+            ctx.cudagraph_runtime_mode != CUDAGraphMode.FULL
+            and not torch.cuda.is_current_stream_capturing()
+        ):
+            return 0
+        return tokens.size(0) * self.num_dispatchers_
+
     def num_dispatchers(self) -> int:
         return self.num_dispatchers_
 
@@ -136,15 +201,18 @@ def _do_dispatch(
         if has_scales:
             token_data = (tokens, token_scales)
 
-        (
-            token_data,
-            expert_topk_ids,
-            expert_topk_weights,
-            expert_num_tokens_per_expert_list,
-            handle,
-            event,
-        ) = self.buffer.dispatch(
-            x=token_data,
+        num_worst_tokens = self._num_worst_tokens(tokens)
+        self.dispatch_seq += 1
+        try:
+            (
+                token_data,
+                expert_topk_ids,
+                expert_topk_weights,
+                expert_num_tokens_per_expert_list,
+                handle,
+                event,
+            ) = self.buffer.dispatch(
+                x=token_data,
             handle=None,
             num_tokens_per_rank=num_tokens_per_rank,
             num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
@@ -154,12 +222,29 @@ def _do_dispatch(
             topk_weights=rank_topk_weights,
             # expert_alignment rounds the number of tokens per expert
             # to this value.
-            expert_alignment=1,
-            config=self._get_dispatch_config(),
-            previous_event=previous_event,
-            async_finish=self.async_prepare and not dbo_enabled(),
-            allocate_on_comm_stream=False,
-        )
+                expert_alignment=1,
+                num_worst_tokens=num_worst_tokens,
+                config=self._get_dispatch_config(),
+                previous_event=previous_event,
+                # With worst-token graphs in play, ALL dispatches must join
+                # back to the compute stream: a graph replay launched on the
+                # compute stream does not wait on the live comm stream, so an
+                # async eager dispatch still in flight would race the captured
+                # dispatch kernels on the shared ring buffers.
+                async_finish=(
+                    self.async_prepare
+                    and not dbo_enabled()
+                    and not envs.VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH
+                ),
+                allocate_on_comm_stream=False,
+            )
+        except RuntimeError as exc:
+            raise RuntimeError(
+                f"{exc} [deepep-ht dispatch debug: "
+                f"seq={self.dispatch_seq} ubatch={dbo_current_ubatch_id()} "
+                f"tokens={tokens.shape[0]} worst={num_worst_tokens} "
+                f"{self._step_debug()}]"
+            ) from exc
 
         # record the handle for this ubatch
         a2a_idx = dbo_current_ubatch_id()
@@ -187,7 +272,7 @@ def _receiver(
         token_data: tuple[torch.Tensor, torch.Tensor] | torch.Tensor,
         expert_topk_ids: torch.Tensor | None,
         num_experts: int,
-        expert_num_tokens_per_expert_list: list[int],
+        expert_num_tokens_per_expert_list: list[int] | torch.Tensor,
         expert_topk_weights: torch.Tensor | None,
         a1_scale: torch.Tensor | None,
         quant_config: FusedMoEQuantConfig,
@@ -219,12 +304,21 @@ def _receiver(
             expert_topk_ids + self.rank_expert_offset,
         )
 
-        # Makes a GPU-CPU copy.
-        # TODO (varun): Maybe it is better to re-compute the expert_num_tokens
-        # on GPU.
-        expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(
-            expert_num_tokens_per_expert_list, device=expert_x.device
-        )
+        if isinstance(expert_num_tokens_per_expert_list, torch.Tensor):
+            # Worst-token (CUDA-graph) dispatch: counts arrive as a device
+            # tensor written by the notify kernel; no host copy exists and
+            # none may be made on this path.
+            expert_tokens_meta = mk.ExpertTokensMetadata(
+                expert_num_tokens=expert_num_tokens_per_expert_list,
+                expert_num_tokens_cpu=None,
+            )
+        else:
+            # Makes a GPU-CPU copy.
+            # TODO (varun): Maybe it is better to re-compute the
+            # expert_num_tokens on GPU.
+            expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(
+                expert_num_tokens_per_expert_list, device=expert_x.device
+            )
 
         # * For non-block quant, dispatch in b16 and quantize now as
         #   DeepEP kernels only support dispatching block scales.
@@ -364,16 +458,28 @@ def _finalize(
         assert fused_expert_output.dtype == torch.bfloat16, (
             f"Expected fused_expert_output bfloat16, got {fused_expert_output.dtype}"
         )
-        combined_x, _, event = self.buffer.combine(
+        self.combine_seq += 1
+        try:
+            combined_x, _, event = self.buffer.combine(
             # HT combine only supports BF16
             x=fused_expert_output,
             handle=handle,
             topk_weights=None,
             config=self._get_combine_config(),
             previous_event=previous_event,
-            async_finish=do_async and not dbo_enabled(),
+            # Same ring-race hazard as dispatch: combine's comm-stream work
+            # must be joined before a subsequent graph replay can launch.
+            async_finish=do_async
+            and not dbo_enabled()
+            and not envs.VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH,
             allocate_on_comm_stream=False,
-        )
+            )
+        except RuntimeError as exc:
+            raise RuntimeError(
+                f"{exc} [deepep-ht combine debug: "
+                f"seq={self.combine_seq} ubatch={dbo_current_ubatch_id()} "
+                f"tokens={fused_expert_output.shape[0]} {self._step_debug()}]"
+            ) from exc
 
         dbo_switch_to_compute()
 

vllm/models/deepseek_v4/compressor.py

@@ -106,10 +106,13 @@ def build(
     ) -> CompressorMetadata:
         query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
         num_reqs = common_attn_metadata.num_reqs
+        num_tokens = common_attn_metadata.num_actual_tokens
         query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
         x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory()
-        token_to_req_indices = self.token_to_req_indices[: x.shape[0]]
-        token_to_req_indices.copy_(x, non_blocking=True)
+        assert x.shape[0] <= num_tokens
+        self.token_to_req_indices[:num_tokens].fill_(0)
+        self.token_to_req_indices[: x.shape[0]].copy_(x, non_blocking=True)
+        token_to_req_indices = self.token_to_req_indices[:num_tokens]
         return CompressorMetadata(
             block_table=common_attn_metadata.block_table_tensor.clamp_(min=0),
             slot_mapping=common_attn_metadata.slot_mapping,

vllm/v1/attention/backends/mla/indexer.py

@@ -14,7 +14,6 @@
     has_deep_gemm,
 )
 from vllm.utils.math_utils import cdiv
-from vllm.utils.platform_utils import num_compute_units
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -28,6 +27,7 @@
 )
 from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
 from vllm.v1.worker.cp_utils import get_total_cp_world_size
+from vllm.v1.worker.sm_control import get_ubatch_compute_sms
 
 logger = init_logger(__name__)
 
@@ -107,6 +107,10 @@ def split_indexer_prefill_chunks(
             end += 1
 
         req_slice = slice(start + request_offset, end + request_offset)
+        # Zero-query rows can appear from padded requests in full-CG/ubatch metadata.
+        # They produce no logits/indexer work, so skip the no-op chunk.
+        if chunk_m == 0:
+            continue
         max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m
         for q_off in range(0, chunk_m, max_q):
             sub_m = min(max_q, chunk_m - q_off)
@@ -277,8 +281,7 @@ def __init__(self, *args, **kwargs):
             or not current_platform.is_device_capability_family(100)
         ) and next_n not in self.natively_supported_next_n_fp4
 
-        sm_count = num_compute_units(self.device.index)
-        self.num_sms = sm_count
+        self.num_sms = get_ubatch_compute_sms(self.vllm_config, self.device.index)
 
         self.offsets_buffer = torch.arange(
             next_n, device=self.device, dtype=torch.int32
@@ -477,12 +480,17 @@ def build(
         seq_lens = common_attn_metadata.seq_lens
         slot_mapping = common_attn_metadata.slot_mapping
         block_table = common_attn_metadata.block_table_tensor
+        has_prefilling_rows = (
+            common_attn_metadata.is_prefilling is not None
+            and torch.any(common_attn_metadata.is_prefilling).item()
+        )
 
         num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
             split_decodes_and_prefills(
                 common_attn_metadata,
                 decode_threshold=self.reorder_batch_threshold,
                 require_uniform=not self.use_flattening,
+                treat_short_extends_as_decodes=not has_prefilling_rows,
             )
         )
 

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -288,10 +288,13 @@ def build(
 
         # NOTE: Ensure all metadata tensors maintain fixed memory addresses
         # for CUDA graph compatibility.
+        num_tokens = common_attn_metadata.num_actual_tokens
         query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
         x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory()
-        token_to_req_indices = self.token_to_req_indices[: x.shape[0]]
-        token_to_req_indices.copy_(x, non_blocking=True)
+        assert x.shape[0] <= num_tokens
+        self.token_to_req_indices[:num_tokens].fill_(0)
+        self.token_to_req_indices[: x.shape[0]].copy_(x, non_blocking=True)
+        token_to_req_indices = self.token_to_req_indices[:num_tokens]
 
         is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
         is_valid_token.copy_(slot_mapping >= 0)