[None][feat] Support rejection sampling under attention DP (incl. LM-head TP)

ruff-legacy-baseline.json

@@ -1,8 +1,8 @@
 {
   "_meta": {
     "generated_by": "scripts/legacy_utils.py lint-update-violations",
-    "total_violations": 2358,
-    "total_files": 271
+    "total_violations": 2327,
+    "total_files": 266
   },
   ".github/scripts/label_community_user.py": {
     "D212": 1
@@ -229,15 +229,6 @@
   "examples/run.py": {
     "E711": 4
   },
-  "examples/scaffolding/contrib/mcp/e2b/e2bserver.py": {
-    "D202": 1,
-    "D205": 1,
-    "D411": 1
-  },
-  "examples/scaffolding/contrib/mcp/websearch/websearch.py": {
-    "D205": 1,
-    "D415": 1
-  },
   "examples/summarize.py": {
     "E741": 1
   },
@@ -543,20 +534,19 @@
     "E731": 2
   },
   "tensorrt_llm/llmapi/disagg_utils.py": {
-    "D205": 2,
+    "D205": 1,
     "D210": 1,
-    "D212": 2,
-    "D403": 2,
-    "D415": 2,
+    "D212": 1,
+    "D403": 1,
+    "D415": 1,
     "F822": 2
   },
   "tensorrt_llm/llmapi/llm.py": {
     "D200": 2,
     "D205": 4
   },
   "tensorrt_llm/llmapi/llm_args.py": {
-    "D205": 11,
-    "D212": 1
+    "D205": 10
   },
   "tensorrt_llm/llmapi/llm_utils.py": {
     "D200": 1
@@ -684,15 +674,9 @@
   "tensorrt_llm/scaffolding/result.py": {
     "F821": 1
   },
-  "tensorrt_llm/scaffolding/scaffolding_llm.py": {
-    "PLE0302": 1
-  },
   "tensorrt_llm/scaffolding/task.py": {
     "F821": 1
   },
-  "tensorrt_llm/scaffolding/worker.py": {
-    "PLE0302": 1
-  },
   "tensorrt_llm/serve/disagg_auto_scaling.py": {
     "D205": 2,
     "D212": 2
@@ -708,15 +692,15 @@
     "D210": 1,
     "F811": 1
   },
+  "tensorrt_llm/serve/openai_server.py": {
+    "F821": 3
+  },
   "tensorrt_llm/serve/responses_utils.py": {
     "D200": 2,
-    "D205": 4,
-    "D212": 11
+    "D212": 7
   },
   "tensorrt_llm/serve/router.py": {
     "D205": 1,
-    "D212": 1,
-    "D300": 1,
     "D415": 7
   },
   "tensorrt_llm/serve/scripts/benchmark_dataset.py": {
@@ -796,9 +780,9 @@
     "D205": 3,
     "D210": 2,
     "D212": 5,
-    "D300": 7,
-    "D403": 6,
-    "D415": 12
+    "D300": 6,
+    "D403": 5,
+    "D415": 11
   },
   "tests/integration/defs/conftest.py": {
     "D200": 3,
@@ -807,9 +791,9 @@
     "D209": 1,
     "D212": 5,
     "D214": 1,
-    "D300": 92,
+    "D300": 91,
     "D403": 51,
-    "D415": 92,
+    "D415": 91,
     "E722": 2,
     "E741": 1
   },
@@ -821,8 +805,7 @@
     "E741": 2
   },
   "tests/integration/defs/disaggregated/test_disaggregated.py": {
-    "D205": 3,
-    "D209": 1,
+    "D205": 2,
     "D212": 5,
     "D415": 1
   },

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -2698,6 +2698,34 @@ def _get_all_rank_ctx_requests(self, num_ctx_requests: int):
             return list(self.dist.tp_allgather(num_ctx_requests))
         return None
 
+    def _sync_group_all_greedy_sample(self, spec_metadata) -> None:
+        """All-gather the per-rank greedy flags and store the group AND.
+
+        Why the sampling-path choice must be group-uniform under
+        ADP + LM-head TP is documented on the anchor,
+        ``SpecMetadata.group_all_greedy_sample``. Local contract: called once
+        per iteration, right after ``update_is_all_greedy_sample`` and BEFORE
+        the CUDA graph key is built. The gate is pure config (identical on
+        every rank), so ranks also agree on whether the exchange happens; the
+        gather spans the whole TP group, a superset of any LM-head-TP
+        subgroup. A dedicated host all-gather rather than a piggyback on the
+        ``all_rank_num_tokens`` exchange, which runs in ``_prepare_inputs`` --
+        after the graph key, too late for the key to see the synced value.
+        """
+        # enable_lm_head_tp_in_adp implies enable_attention_dp (asserted in
+        # Mapping.__init__), so ADP needs no separate check here.
+        if not (self.mapping.enable_lm_head_tp_in_adp
+                and spec_metadata.use_rejection_sampling):
+            return
+        local_flag = bool(spec_metadata.is_all_greedy_sample)
+        all_flags = self.dist.tp_allgather(local_flag)
+        spec_metadata.group_all_greedy_sample = all(all_flags)
+        # Also overwrite the live flag directly: this iteration's scan already
+        # ran (update_is_all_greedy_sample just returned) and the CUDA graph
+        # key reads the flag next -- the stored override only takes effect on
+        # the NEXT rescan (populate), which is after key selection.
+        spec_metadata.is_all_greedy_sample = spec_metadata.group_all_greedy_sample
+
     def _set_spec_metadata_all_rank_num_tokens(
             self,
             spec_metadata: SpecMetadata,
@@ -5840,6 +5868,7 @@ def forward(self,
             if spec_metadata is not None:
                 spec_metadata.update_is_all_greedy_sample(
                     padded_requests.all_requests())
+                self._sync_group_all_greedy_sample(spec_metadata)
 
             maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph(
                 padded_requests,

tensorrt_llm/_torch/speculative/eagle3.py

@@ -854,11 +854,32 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                 #   ADP+LM-head-TP padding to ``max_num_requests`` so every TP
                 #   rank produces logits of the same shape.
                 # Eagle3: logits_processor of the EAGLE draft model.
-                use_lm_head_tp_in_adp = (
+                #
+                # The LM-head-TP fast path (group-stacked rows, vocab-sharded
+                # weight) is only usable when the consumer is an argmax: the
+                # distributed greedy sampler recovers the global argmax from
+                # vocab-shard maxima without materializing full distributions.
+                # Advanced (rejection) sampling needs each request's full-vocab
+                # distribution, so a batch headed for the advanced path
+                # bypasses LM-head-TP and computes full-vocab logits locally --
+                # under ADP the lm_head weight is replicated (the sliced-shard
+                # trick is a runtime optimization), exactly like the target
+                # head. With rejection off, non-greedy batches keep the
+                # LM-head-TP argmax path unconditionally, where the per-rank
+                # greedy flag never enters control flow. This branch is safe
+                # to take group-uniformly because is_all_greedy_sample is
+                # group-synchronized whenever rejection+ADP+LM-head-TP are
+                # combined -- see SpecMetadata.group_all_greedy_sample (anchor
+                # for the group-sync semantics).
+                advanced_draft_sampling = (
+                    spec_metadata.wants_advanced_draft_sampling)
+                # enable_lm_head_tp_in_adp implies enable_attention_dp
+                # (asserted in Mapping.__init__); no separate ADP check.
+                lm_head_tp_in_adp_configured = (
                     self.is_mtp_eagle and self.model_config is not None
-                    and self.model_config.mapping.enable_attention_dp
-                    and getattr(self.model_config.mapping,
-                                'enable_lm_head_tp_in_adp', False))
+                    and self.model_config.mapping.enable_lm_head_tp_in_adp)
+                use_lm_head_tp_in_adp = (lm_head_tp_in_adp_configured
+                                         and not advanced_draft_sampling)
                 if self.is_mtp_eagle:
                     if use_lm_head_tp_in_adp:
                         hidden_states_gathered = hidden_states[gather_ids]
@@ -884,6 +905,15 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                         logits = draft_model.mtp_layers[0].shared_head(
                             padded_hidden_states, draft_model.lm_head,
                             attn_metadata, True)
+                    elif lm_head_tp_in_adp_configured:
+                        # Advanced-sampling bypass: the model's shared_head
+                        # would re-apply the LM-head-TP stacked/sharded path
+                        # from config on its own, so call lm_head directly.
+                        # Under ADP the LMHead weight is replicated and
+                        # is_spec_decoding_head defaults to False, so this is
+                        # a plain local full-vocab GEMM over this rank's own
+                        # rows -- the same computation the target head runs.
+                        logits = draft_model.lm_head(hidden_states[gather_ids])
                     else:
                         logits = draft_model.mtp_layers[0].shared_head(
                             hidden_states[gather_ids], draft_model.lm_head,

tensorrt_llm/_torch/speculative/interface.py

@@ -526,6 +526,16 @@ class SpecMetadata:
     # Defaults to True so non-one-engine paths (where populate is a no-op)
     # never accidentally select the advanced graph variant.
     is_all_greedy_sample: bool = True
+    # Group-synchronized override for ``is_all_greedy_sample`` (AND over the
+    # TP group's local flags; None = no group sync configured, use the local
+    # value). Under ADP + LM-head TP with rejection sampling, the greedy-vs-
+    # advanced choice gates group collectives, so all ranks must take the same
+    # path even though their batches (and thus local flags) differ. Set by
+    # ``_sync_group_all_greedy_sample`` before the CUDA graph key is built and
+    # re-applied by ``_scan_one_model_sampling`` on every rescan. AND is safe:
+    # a greedy rank pulled onto the advanced path still samples greedily via
+    # its sentinel params.
+    group_all_greedy_sample: Optional[bool] = None
     # Whether to use rejection sampling for one-model speculative decoding.
     use_rejection_sampling: bool = False
     # Sampling parameters for non-greedy sampling (per-request)
@@ -830,8 +840,29 @@ def _normalize_request_sampling_params(
                 for (_, _, _, num_tokens) in per_request_normalized
             ]
 
+        # Apply the group-synchronized override last (semantics: see the
+        # ``group_all_greedy_sample`` field comment). Local contract: the
+        # synced value already incorporates any capture override, and rescans
+        # (e.g. populate after the graph key) must converge to it rather than
+        # resurrect the local value.
+        if self.group_all_greedy_sample is not None:
+            self.is_all_greedy_sample = self.group_all_greedy_sample
+
         return per_request_normalized, per_request_slot_ids
 
+    @property
+    def wants_advanced_draft_sampling(self) -> bool:
+        """Whether the current batch takes the advanced (rejection) draft
+        path: rejection sampling enabled AND not an all-greedy batch.
+
+        Single source of truth for the greedy-vs-advanced decision: the
+        sampler branch (``sample_draft_tokens``) and the worker's LM-head-TP
+        bypass (``_forward_linear_draft_loop``) must agree exactly -- a
+        divergence feeds the wrong logits layout to the sampler -- so both
+        read this property instead of re-deriving the predicate.
+        """
+        return self.use_rejection_sampling and not self.is_all_greedy_sample
+
     def update_is_all_greedy_sample(self, requests: list["LlmRequest"]) -> None:
         """Refresh ``is_all_greedy_sample`` for the *current* batch.
 
@@ -1441,15 +1472,16 @@ def maybe_gather_sharded_draft_logits(self,
         (see ``_draft_logits_are_sharded``); replicated full-vocab logits are
         returned unchanged.
 
-        Plain TP gathers vocab shards over ``self.mapping``. ADP + LM-head TP
-        never reaches this path: rejection sampling (the only consumer of
-        advanced draft sampling) is config-gated off under attention DP, and
-        the group-stacked sharded logits it produces are handled by the greedy
-        path in ``greedy_sample_draft_with_tp_gather``.
+        Plain TP gathers vocab shards over ``self.mapping``. The LM-head-TP
+        stacked/sharded layout never reaches this path: an advanced-sampling
+        batch bypasses the LM-head-TP fast path in the worker and computes
+        full-vocab logits locally from the (ADP-replicated) lm_head weight, so
+        ``mapping_lm_head_tp`` is only ever passed alongside greedy sampling.
         """
         assert mapping_lm_head_tp is None, (
-            "Advanced draft sampling is not supported under ADP + LM-head TP "
-            "(rejection sampling is config-gated off with attention DP)")
+            "Advanced draft sampling must not receive LM-head-TP "
+            "stacked/sharded logits; the worker bypasses the LM-head-TP fast "
+            "path for non-all-greedy batches (see _forward_linear_draft_loop)")
         if (spec_metadata is None or spec_metadata.is_all_greedy_sample
                 or not self._draft_logits_are_sharded(logits, spec_metadata)):
             return logits
@@ -1972,7 +2004,6 @@ def sample_draft_tokens(self,
         no slicing is needed.
         """
         is_block = logits.dim() == 3
-        use_rejection = getattr(spec_metadata, "use_rejection_sampling", False)
 
         # Draft tokens use argmax unless rejection sampling is engaged for a
         # non-greedy batch. Rejection sampling is the only path that needs the
@@ -1983,7 +2014,7 @@ def sample_draft_tokens(self,
         # max_i p_i >= sum_i p_i^2 = E[accept] for a stochastic draft). This
         # matches sglang/vLLM, which draft with argmax/top-k by default and apply
         # sampling params only on the target/acceptance side.
-        advanced = use_rejection and not spec_metadata.is_all_greedy_sample
+        advanced = spec_metadata.wants_advanced_draft_sampling
 
         # All samplers below return tokens in draft-vocab space; d2t is applied
         # once after the branch.

tensorrt_llm/llmapi/llm_args.py

@@ -5208,9 +5208,13 @@ def validate_speculative_config(self):
                 # Plain tensor parallelism is supported (the draft path
                 # all-gathers vocab-sharded draft logits before rejection, see
                 # SpecWorkerBase.maybe_gather_sharded_draft_logits).
-                # attention-DP and context parallelism remain gated.
-                rs_parallel_active = (self.context_parallel_size > 1
-                                      or self.enable_attention_dp)
+                # Attention DP is supported: each rank holds full-vocab draft
+                # logits for its own requests (the LM-head-TP fast path is
+                # bypassed for advanced sampling, and is_all_greedy_sample is
+                # group-synchronized so the LM-head-TP group's collectives stay
+                # uniform -- see SpecMetadata.group_all_greedy_sample). Only
+                # context parallelism remains gated.
+                rs_parallel_active = self.context_parallel_size > 1
                 rs_guided_active = self.guided_decoding_backend is not None
                 rs_sa_active = getattr(self.speculative_config, "sa_config",
                                        None) is not None
@@ -5262,10 +5266,9 @@ def validate_speculative_config(self):
                                     "relaxed-thinking acceptance is enabled")
                             if rs_parallel_active:
                                 reasons.append(
-                                    "tensor/context parallelism or attention-DP "
-                                    "is active (the draft path resolves only "
-                                    "the global argmax, not full distributions)"
-                                )
+                                    "context parallelism is active (the draft "
+                                    "path resolves only the global argmax, "
+                                    "not full distributions)")
                             if rs_guided_active:
                                 reasons.append("guided decoding is enabled")
                         raise ValueError(