diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index 9a7bf5d452b3..0eb256fe39f6 100644 --- a/ruff-legacy-baseline.json +++ b/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,11 +534,11 @@ "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": { @@ -555,8 +546,7 @@ "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 }, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e970063addd7..206bbb8b09ac 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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, diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 7a373cb0ef9c..8fca7fb4cc27 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/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, diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index f11385c33a87..50e23d087ca7 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/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. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index ca2ae9cc0062..03d13f72f88c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/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( diff --git a/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py b/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py new file mode 100644 index 000000000000..bdcc3fd77992 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for the group-synchronized ``is_all_greedy_sample`` override. + +Under ADP + LM-head TP with rejection sampling, the greedy-vs-advanced path +choice gates group collectives, so the model engine all-gathers the per-rank +flags and stores the group AND in ``SpecMetadata.group_all_greedy_sample``; +``_scan_one_model_sampling`` must then re-apply it on every rescan (populate +runs after the CUDA graph key is built and would otherwise resurrect the +rank-local value). + +These tests call ``_scan_one_model_sampling`` unbound on a SimpleNamespace +stand-in, mirroring test_rejection_buffers_guard.py, so no GPU or full +SpecMetadata construction is needed. +""" + +import types + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.speculative.interface import SpecMetadata + + +def _fake_request(temperature=None, top_k=None, top_p=None, slot=0): + return types.SimpleNamespace( + sampling_config=types.SimpleNamespace( + temperature=[temperature] if temperature is not None else None, + top_k=[top_k] if top_k is not None else None, + top_p=[top_p] if top_p is not None else None, + ), + state=LlmRequestState.GENERATION_IN_PROGRESS, + py_seq_slot=slot, + ) + + +def _fake_meta(group_all_greedy_sample=None, force_capture=False): + return types.SimpleNamespace( + runtime_draft_len=2, + dummy_slot_row=0, + group_all_greedy_sample=group_all_greedy_sample, + _force_non_greedy_for_capture=force_capture, + ) + + +def _scan(meta, requests): + return SpecMetadata._scan_one_model_sampling(meta, requests) + + +def test_local_value_used_when_no_group_sync(): + meta = _fake_meta(group_all_greedy_sample=None) + _scan(meta, [_fake_request(), _fake_request()]) + assert meta.is_all_greedy_sample is True + + _scan(meta, [_fake_request(), _fake_request(temperature=0.8)]) + assert meta.is_all_greedy_sample is False + + +def test_group_override_pulls_greedy_rank_onto_advanced_path(): + # This rank's batch is all-greedy, but another rank in the LM-head-TP + # group has a sampling request: the group AND (False) must win so the + # whole group takes the advanced path together. + meta = _fake_meta(group_all_greedy_sample=False) + _scan(meta, [_fake_request(), _fake_request()]) + assert meta.is_all_greedy_sample is False + + +def test_group_override_survives_rescan(): + # populate_sampling_params_for_one_model rescans after the CUDA graph key + # is built; the override must keep applying so the key, the buffers, and + # the worker branches all agree. + meta = _fake_meta(group_all_greedy_sample=False) + for _ in range(3): + _scan(meta, [_fake_request()]) + assert meta.is_all_greedy_sample is False + + +def test_group_override_true_keeps_greedy(): + meta = _fake_meta(group_all_greedy_sample=True) + _scan(meta, [_fake_request()]) + assert meta.is_all_greedy_sample is True + + +def test_capture_override_composes_with_group_sync(): + # Warmup forces the advanced variant to capture its CUDA graph; the group + # value is derived from capture-forced locals (all False), so the final + # flag stays False regardless of composition order. + meta = _fake_meta(group_all_greedy_sample=False, force_capture=True) + _scan(meta, [_fake_request()]) + assert meta.is_all_greedy_sample is False diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 3365ce6685ab..d59a03bb64c6 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -107,6 +107,53 @@ def test_MTPDecodingConfig_default_draft_len_is_not_user_set(): assert "max_draft_len" in explicit_config.model_fields_set +def test_rejection_sampling_allows_attention_dp(monkeypatch): + """ADP (incl. ADP+LM-head-TP) supports rejection sampling. + + The draft path bypasses LM-head-TP for advanced sampling and the greedy + flag is group-synchronized, so the former attention-DP gate is lifted. + """ + import tensorrt_llm._torch.flashinfer_utils as fi_utils + monkeypatch.setattr(fi_utils, "IS_FLASHINFER_AVAILABLE", True) + + # Vanilla MTP is one of the "newly wired" methods the parallel gate used + # to cover: lifting the ADP gate is the actual behavior change here. + for lm_head_tp in (False, True): + spec_cfg = MTPDecodingConfig(max_draft_len=2, + use_rejection_sampling=True, + use_mtp_vanilla=True) + args = TorchLlmArgs(model=llama_model_path, + enable_attention_dp=True, + enable_lm_head_tp_in_adp=lm_head_tp, + speculative_config=spec_cfg) + assert args.speculative_config.use_rejection_sampling is True + + # MTP-Eagle (default) was never parallel-gated; keep it as a regression + # guard that the gate rework did not accidentally start rejecting it. + spec_cfg = MTPDecodingConfig(max_draft_len=2, use_rejection_sampling=True) + args = TorchLlmArgs(model=llama_model_path, + enable_attention_dp=True, + enable_lm_head_tp_in_adp=True, + speculative_config=spec_cfg) + assert args.speculative_config.use_rejection_sampling is True + + +def test_rejection_sampling_still_gated_on_context_parallel(): + """Context parallelism remains an unsupported rejection combination. + + Explicit opt-in raises; default-inherited silently disables. The parallel + gate applies to the newly wired methods (vanilla MTP, PARD, DFlash, + DraftTarget one-model), so use vanilla MTP here. + """ + spec_cfg = MTPDecodingConfig(max_draft_len=2, + use_rejection_sampling=True, + use_mtp_vanilla=True) + with pytest.raises(ValueError, match="context parallelism"): + TorchLlmArgs(model=llama_model_path, + context_parallel_size=2, + speculative_config=spec_cfg) + + class TestYaml: def _yaml_to_dict(self, yaml_content: str) -> dict: