diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index cf907ee08e39..a288be358175 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -295,7 +295,7 @@ def copy_to_device(self) -> None: trunc_h_buf = self._trunc_host_bufs[name][:copy_bytes] trunc_d_buf.copy_(trunc_h_buf, non_blocking=True) - def copy_to_host(self) -> None: + def copy_to_host(self, non_blocking: bool = False) -> None: """Copy from device buffer to host buffer. Mirrors ``copy_to_device``: uses the current length of the truncatable tensor @@ -306,7 +306,7 @@ def copy_to_host(self) -> None: if self._total_bytes > 0: h_buffer = self._host_buffer[: self._total_bytes] d_buffer = self._device_buffer[: self._total_bytes] - h_buffer.copy_(d_buffer, non_blocking=True) + h_buffer.copy_(d_buffer, non_blocking=non_blocking) # Copy each truncatable tensor independently, truncated to current length for name in self._truncatable_names: @@ -316,7 +316,7 @@ def copy_to_host(self) -> None: copy_bytes = length * dtype.itemsize trunc_d_buf = self._trunc_device_bufs[name][:copy_bytes] trunc_h_buf = self._trunc_host_bufs[name][:copy_bytes] - trunc_h_buf.copy_(trunc_d_buf, non_blocking=True) + trunc_h_buf.copy_(trunc_d_buf, non_blocking=non_blocking) def resize(self, name: str, new_capacity: int) -> None: """Resize a truncatable tensor's capacity. @@ -1175,6 +1175,27 @@ def _is_required(self, name: str, check_both: bool = True) -> bool: """ return self._is_active(name, check_both) or self._is_active_host_prep(name, check_both) + def _active_host_update_args( + self, arg_names: Set[str], active_args_override: Optional[Set[str]] = None + ) -> List[str]: + """Return host args that need mirroring after an in-graph metadata update. + + ``active_args_override`` lets a caller narrow host mirroring to the graph inputs the next + consumer actually reads. It is treated as a filter: only active host args whose names appear + in the override are mirrored. The override may contain names that are not active graph args + (e.g. a submodule's full placeholder set, which also includes inter-module tensors such as + ``inputs_embeds``/``hidden_states``); such entries are simply ignored. The caller is + responsible for including every host argument the next consumer may read. + """ + needs_d2h_sync = [ + k + self._host_suffix + for k in arg_names + if self._is_active(k + self._host_suffix, check_both=False) + ] + if active_args_override is None: + return needs_d2h_sync + return [arg_name for arg_name in needs_d2h_sync if arg_name in active_args_override] + def _stage_arg( self, name: str, @@ -1583,11 +1604,16 @@ def run_host_prepare_for_attention_forward(self) -> None: host_function(**{arg: self.get_arg(arg) for arg in args}) @nvtx_range("ad_offset_pos_and_cache_") - def offset_pos_and_cache_(self, offset: torch.Tensor) -> None: + def offset_pos_and_cache_( + self, offset: torch.Tensor, active_args_override: Optional[Set[str]] = None + ) -> None: """Offset position and cache-related metadata for active arguments. Args: offset: 1D tensor [batch_size] with per-sequence position offsets. + active_args_override: Optional graph-input names for the next in-forward consumer. When + provided, host mirroring is limited to those active host args. The caller is + responsible for including every host argument the next consumer may read. """ # check if we need a d2h sync _REQUIRES_UPDATE = { @@ -1599,11 +1625,7 @@ def offset_pos_and_cache_(self, offset: torch.Tensor) -> None: "seq_len_with_cache", "use_initial_states", } - needs_d2h_sync = [ - k + self._host_suffix - for k in _REQUIRES_UPDATE - if self._is_active(k + self._host_suffix, check_both=False) - ] + needs_d2h_sync = self._active_host_update_args(_REQUIRES_UPDATE, active_args_override) sync_to_host = any(needs_d2h_sync) if sync_to_host: ad_logger.debug(f"d2h sync required in offset_pos_and_cache_ for {needs_d2h_sync}") @@ -1694,7 +1716,7 @@ def offset_pos_and_cache_(self, offset: torch.Tensor) -> None: # TODO: May need to dissect what fields are needed in the forward pass to reduce # data movement. if sync_to_host: - self._input_buffer.copy_to_host() + self._input_buffer.copy_to_host(non_blocking=False) @nvtx_range("ad_offset_with_new_lens_") def offset_with_new_lens_(self, new_lens_ungathered: torch.Tensor) -> None: @@ -1718,13 +1740,18 @@ def offset_with_new_lens_(self, new_lens_ungathered: torch.Tensor) -> None: self.offset_pos_and_cache_(increment) @nvtx_range("ad_switch_to_generate_") - def switch_to_generate_(self) -> None: + def switch_to_generate_(self, active_args_override: Optional[Set[str]] = None) -> None: """Switch all sequences metadata to generate (decode) mode. Transitions the batch from any layout (prefill/extend/decode or mixed) to an all-decode layout where each sequence has exactly 1 token. We assume that we just take the last position of each sequence for the metadata. + Args: + active_args_override: Optional graph-input names for the next in-forward consumer. When + provided, host mirroring is limited to those active host args. The caller is + responsible for including every host argument the next consumer may read. + NOTE: update device tensors first and mirror back to host only when an updated host-side argument is active. @@ -1757,11 +1784,7 @@ def switch_to_generate_(self) -> None: "position_ids", "use_initial_states", } - needs_d2h_sync = [ - k + self._host_suffix - for k in _REQUIRES_UPDATE - if self._is_active(k + self._host_suffix, check_both=False) - ] + needs_d2h_sync = self._active_host_update_args(_REQUIRES_UPDATE, active_args_override) sync_to_host = any(needs_d2h_sync) # --- input_ids (device) --- @@ -1790,7 +1813,7 @@ def switch_to_generate_(self) -> None: # TODO: May need to dissect what fields are needed in the forward pass to reduce # data movement. if sync_to_host: - self._input_buffer.copy_to_host() + self._input_buffer.copy_to_host(non_blocking=False) def copy_(self, name: str, src: torch.Tensor, strict: bool = True) -> None: """Copy a tensor into the buffer. USE WITH CAUTION! diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index b8b8d5892310..526f5583c80a 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -412,18 +412,16 @@ def cap_max_batch_size_to_max_num_tokens(self): return self @model_validator(mode="after") - def disable_cudagraph_for_speculative_flashinfer(self): + def reject_cudagraph_for_speculative_flashinfer(self): if ( self.speculative_config is not None and self.attn_backend == "flashinfer" and self.is_cuda_graph_enabled() ): - ad_logger.warning( + raise ValueError( "Speculative decoding with FlashInfer attention does not currently support CUDA " - "graph replay in AutoDeploy; falling back to compile_backend='torch-simple'." + "graph replay in AutoDeploy. Use compile_backend='torch-simple' instead." ) - self.compile_backend = "torch-simple" - self.update_transforms_with_shortcuts() return self ### UTILITY METHODS ############################################################################ diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py index 9d82fa5a9f3d..485ec4a9bd4d 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py @@ -33,7 +33,7 @@ from dataclasses import dataclass from types import SimpleNamespace -from typing import Any, ClassVar, Dict, Optional, Union +from typing import Any, ClassVar, Dict, Optional, Set, Union import torch import torch.nn as nn @@ -912,10 +912,14 @@ def _forward_prefill_only(self, input_ids: torch.Tensor, position_ids: torch.Ten # KV-cache forward (inference after graph transforms) # # ================================================================== # + @staticmethod + def _submodule_placeholder_names(submodule: nn.Module) -> Set[str]: + return {node.name for node in submodule.graph.nodes if node.op == "placeholder"} + @staticmethod def _filter_kwargs_for_submodule(kwargs: dict, submodule: nn.Module) -> dict: """Filter kwargs to only include those accepted by submodule's forward (GraphModule).""" - expected_names = {node.name for node in submodule.graph.nodes if node.op == "placeholder"} + expected_names = EagleWrapper._submodule_placeholder_names(submodule) return {k: v for k, v in kwargs.items() if k in expected_names} @staticmethod @@ -1096,6 +1100,7 @@ def _forward_with_kv_cache(self, csi: CachedSequenceInterface): next_new_tokens[:, 0] = csi.info.maybe_gather_and_squeeze(csi.get_arg("input_ids")) # ---- Phase 5: Draft loop ---- + draft_arg_names = self._submodule_placeholder_names(self.draft_model) for draft_idx in range(self.max_draft_len): # run forward pass on the draft model in shape [num_sequences, 1] draft_output = self.draft_model( @@ -1123,9 +1128,9 @@ def _forward_with_kv_cache(self, csi: CachedSequenceInterface): # switch to generate (if not done already), store new tokens, and offset cache # can be skipped for last iteration since after we return metadata will be reset if draft_idx < self.max_draft_len - 1: - csi.info.switch_to_generate_() + csi.info.switch_to_generate_(active_args_override=draft_arg_names) csi.info.copy_("input_ids", draft_tokens) - csi.info.offset_pos_and_cache_(c_offset) + csi.info.offset_pos_and_cache_(c_offset, active_args_override=draft_arg_names) # ---- Phase 6: Package output ---- return EagleWrapperOutput( diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_switch_to_generate_inplace.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_switch_to_generate_inplace.py index 16eba6cc79c5..8ae9bb98ad7c 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_switch_to_generate_inplace.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_switch_to_generate_inplace.py @@ -54,7 +54,11 @@ def _make_seq_info(extra_activate=()) -> SequenceInfo: def _nest_prefill(si: SequenceInfo, input_ids, pages_per_seq, cache_loc, **kw): - """Convenience wrapper: nest prefill sequences and return the SequenceInfo.""" + """Convenience wrapper: nest packed prefill sequences and return the SequenceInfo. + + cu_seqlen is staged as the prefix sum of per-sequence lengths starting at 0, e.g. + input_ids=[[1, 2, 3], [4, 5, 6, 7]] (lengths 3 and 4) stages cu_seqlen=[0, 3, 7]. + """ flat_ids = [t for seq in input_ids for t in (seq.tolist() if hasattr(seq, "tolist") else seq)] cu_seqlen = [0] for seq in input_ids: @@ -76,6 +80,25 @@ def _nest_prefill(si: SequenceInfo, input_ids, pages_per_seq, cache_loc, **kw): return si +def _nest_decode(si: SequenceInfo, input_pos=(5, 10)): + """Convenience wrapper: nest a two-sequence all-decode batch (seq_len=1 per sequence). + + cu_seqlen is the decode layout [0, 1, 2] (1 token per sequence). With the default + input_pos=(5, 10): seq_len_with_cache stages to [6, 11] (input_pos + 1), and + cache_loc/cu_num_pages give 2 pages for seq 0 and 3 for seq 1. + """ + si.nest_sequences( + [1, 2], + cu_seqlen=[0, 1, 2], + input_pos=list(input_pos), + batch_info=[0, 0, 0, 0, 2, 2], + cache_loc_per_pool=[[10, 11, 20, 21, 22]], + cu_num_pages_per_pool=[[0, 2, 5]], + extra_page_per_seq_per_pool=[[-1, -1]], + ) + return si + + def _snapshot_host_views(si: SequenceInfo): """Snapshot full host buffers so current-length updates do not hide host writes.""" return { @@ -311,7 +334,24 @@ def _setup_decode_at_page_boundary(si): class TestSwitchToGenerateHostArgHandling: - """Validate host arg warning and d2h sync behavior.""" + """Validate active_args_override host (device->host) sync behavior. + + Each test asserts on both the device tensor (``get_view``, updated in place by the metadata + helpers) and the host staging mirror (``get_host_view``, refreshed only by a device->host + sync). "No sync" is therefore observable as the host mirror keeping its stale pre-call staged + value while the device tensor already holds the updated value. + + Notation used in the expected values: + + - ``cu_seqlen`` ("cumulative sequence lengths"): a length ``batch_size + 1`` prefix sum of the + per-sequence token counts in the packed batch, always starting at 0. A prefill of two + sequences with lengths [3, 4] stages ``[0, 3, 7]``; an all-decode batch of 2 sequences + (1 token each) is ``[0, 1, 2]``. ``switch_to_generate_`` rewrites cu_seqlen from the packed + prefill layout to the decode layout. + - ``seq_len_with_cache``: per-sequence total length including already-cached tokens + (``input_pos + current seq_len``). ``offset_pos_and_cache_`` advances it by the offset; a + decode batch with ``input_pos=[5, 10]`` stages ``[6, 11]``, and a +1 offset gives ``[7, 12]``. + """ def test_inactive_host_mirrors_not_synced_by_switch_to_generate(self): """Inactive host mirrors should keep pre-transition staging values.""" @@ -375,6 +415,72 @@ def test_active_host_mirror_synced_by_switch_to_generate(self): assert device_cu[:3].tolist() == [0, 1, 2] assert host_cu[:3].tolist() == [0, 1, 2] + def test_out_of_scope_active_host_mirror_not_synced_by_switch_to_generate(self): + """Host mirrors outside the next consumer's placeholders should keep staging values.""" + si = _make_seq_info(extra_activate=("cu_seqlen_host",)) + _nest_prefill( + si, + input_ids=[[1, 2, 3], [4, 5, 6, 7]], + pages_per_seq=[1, 1], + cache_loc=[10, 20], + ) + + si.switch_to_generate_(active_args_override={"input_ids"}) + + # cu_seqlen on device is rewritten to the decode layout [0, 1, 2]. cu_seqlen_host is out of + # scope (not in the override), so it is not synced and keeps the staged prefill prefix + # sum [0, 3, 7] -- device and host disagree. + device_cu = si._input_buffer.get_view("cu_seqlen") + host_cu = si._input_buffer.get_host_view("cu_seqlen") + assert device_cu[:3].tolist() == [0, 1, 2] + assert host_cu[:3].tolist() == [0, 3, 7] + + def test_in_scope_active_host_mirror_synced_by_switch_to_generate(self): + """Host mirrors inside the next consumer's placeholders should be refreshed.""" + si = _make_seq_info(extra_activate=("cu_seqlen_host",)) + _nest_prefill( + si, + input_ids=[[1, 2, 3], [4, 5, 6, 7]], + pages_per_seq=[1, 1], + cache_loc=[10, 20], + ) + + si.switch_to_generate_(active_args_override={"cu_seqlen_host"}) + + # cu_seqlen_host is in the override, so it is synced: the host mirror is refreshed to the + # device decode layout [0, 1, 2] (instead of the staged prefill prefix sum [0, 3, 7]). + device_cu = si._input_buffer.get_view("cu_seqlen") + host_cu = si._input_buffer.get_host_view("cu_seqlen") + assert device_cu[:3].tolist() == [0, 1, 2] + assert host_cu[:3].tolist() == [0, 1, 2] + + def test_drafting_override_is_tolerated_by_switch_to_generate(self): + """Drafting-style override with irrelevant placeholders is tolerated by switch_to_generate_. + + The Eagle/MTP draft loop passes the draft submodule's full placeholder set as the override, + which includes inter-module tensors that are not SequenceInfo graph args (inputs_embeds, + hidden_states). switch_to_generate_ must ignore those rather than raising, while still + syncing the in-scope host mirror. + """ + si = _make_seq_info(extra_activate=("cu_seqlen_host",)) + _nest_prefill( + si, + input_ids=[[1, 2, 3]], + pages_per_seq=[1], + cache_loc=[10], + ) + + # Mirrors the draft-model placeholder set: active host arg + non-active inter-module inputs. + draft_placeholders = {"input_ids", "cu_seqlen_host", "inputs_embeds", "hidden_states"} + si.switch_to_generate_(active_args_override=draft_placeholders) + + # One sequence -> decode cu_seqlen [0, 1]. cu_seqlen_host is in scope, so it is synced to + # match the device; the non-graph placeholders (inputs_embeds/hidden_states) are ignored. + device_cu = si._input_buffer.get_view("cu_seqlen") + host_cu = si._input_buffer.get_host_view("cu_seqlen") + assert device_cu[:2].tolist() == [0, 1] + assert host_cu[:2].tolist() == [0, 1] + def test_non_native_host_arg_syncs_device_to_host(self): """Activating a non-native host arg should sync device -> host instead of raising.""" si = _make_seq_info() @@ -423,6 +529,64 @@ def test_non_native_host_seq_len_with_cache_syncs(self): assert host_swc[0].item() == 4 assert host_swc[1].item() == 5 + def test_out_of_scope_active_host_mirror_not_synced_by_offset_pos_and_cache(self): + """Out-of-scope host mirror keeps its pre-offset staged value (no d2h sync).""" + si = _make_seq_info(extra_activate=("seq_len_with_cache_host",)) + _nest_decode(si) # seq_len_with_cache staged to [6, 11] + + increment = torch.tensor([1, 1], dtype=torch.int32, device=si.device) + si.offset_pos_and_cache_(increment, active_args_override={"input_ids"}) + + # offset advances seq_len_with_cache on device: staged [6, 11] -> [7, 12]. + # seq_len_with_cache_host is out of scope (not in the override), so it is not synced and + # keeps the staged [6, 11] -- device and host disagree. + device_swc = si._input_buffer.get_view("seq_len_with_cache") + host_swc = si._input_buffer.get_host_view("seq_len_with_cache") + assert device_swc[:2].tolist() == [7, 12] + assert host_swc[:2].tolist() == [6, 11] + + def test_in_scope_active_host_mirror_synced_by_offset_pos_and_cache(self): + """In-scope host mirror is refreshed from device metadata (d2h sync).""" + si = _make_seq_info(extra_activate=("seq_len_with_cache_host",)) + _nest_decode(si) # seq_len_with_cache staged to [6, 11] + + increment = torch.tensor([1, 1], dtype=torch.int32, device=si.device) + si.offset_pos_and_cache_(increment, active_args_override={"seq_len_with_cache_host"}) + + # seq_len_with_cache_host is in the override, so it is synced: the host mirror is refreshed + # to the advanced device value [7, 12] (instead of the staged [6, 11]). + device_swc = si._input_buffer.get_view("seq_len_with_cache") + host_swc = si._input_buffer.get_host_view("seq_len_with_cache") + assert device_swc[:2].tolist() == [7, 12] + assert host_swc[:2].tolist() == [7, 12] + + def test_drafting_override_is_tolerated_by_offset_pos_and_cache(self): + """Drafting-style override with irrelevant placeholders is tolerated by offset_pos_and_cache_. + + Mirrors the Eagle/MTP draft loop, which passes the draft submodule's full placeholder set + (including non-graph-arg tensors such as inputs_embeds/hidden_states). offset_pos_and_cache_ + must ignore those while still syncing the in-scope host mirror. + """ + si = _make_seq_info(extra_activate=("seq_len_with_cache_host",)) + _nest_decode(si) # seq_len_with_cache staged to [6, 11] + + # Mirrors the draft-model placeholder set: active host arg + non-active inter-module inputs. + draft_placeholders = { + "input_ids", + "seq_len_with_cache_host", + "inputs_embeds", + "hidden_states", + } + increment = torch.tensor([1, 1], dtype=torch.int32, device=si.device) + si.offset_pos_and_cache_(increment, active_args_override=draft_placeholders) + + # seq_len_with_cache_host is in scope, so it is synced to the advanced device value [7, 12] + # (from staged [6, 11]); the non-graph placeholders (inputs_embeds/hidden_states) are ignored. + device_swc = si._input_buffer.get_view("seq_len_with_cache") + host_swc = si._input_buffer.get_host_view("seq_len_with_cache") + assert device_swc[:2].tolist() == [7, 12] + assert host_swc[:2].tolist() == [7, 12] + def test_native_host_args_do_not_raise(self): """batch_info_host, cu_seqlen_host, seq_len_host are native (tokens_gather in batch_info).""" si = _make_seq_info() diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py b/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py index 539f9ba7fd2a..9bf7fd7074bf 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py @@ -235,6 +235,40 @@ def test_accepts_mtp_eagle_one_model(self): args = LlmArgs(model="test-model", speculative_config=spec_config) assert args.model_factory == "eagle_one_model" + @pytest.mark.parametrize("compile_backend", ["torch-cudagraph", "torch-opt"]) + def test_rejects_flashinfer_cuda_graph_backend(self, compile_backend): + from tensorrt_llm.llmapi import EagleDecodingConfig + + spec_config = EagleDecodingConfig( + max_draft_len=3, + speculative_model="some/model", + eagle3_one_model=True, + ) + + with pytest.raises(pydantic.ValidationError): + LlmArgs( + model="test-model", + speculative_config=spec_config, + attn_backend="flashinfer", + compile_backend=compile_backend, + ) + + def test_accepts_flashinfer_torch_simple(self): + from tensorrt_llm.llmapi import EagleDecodingConfig + + spec_config = EagleDecodingConfig( + max_draft_len=3, + speculative_model="some/model", + eagle3_one_model=True, + ) + + LlmArgs( + model="test-model", + speculative_config=spec_config, + attn_backend="flashinfer", + compile_backend="torch-simple", + ) + # ================================ # CUDA Graph Batch Sizes Tests