From 5c460e76a9f2199fbf482ffa9f87dc6879d41f64 Mon Sep 17 00:00:00 2001 From: Mike Iovine Date: Tue, 14 Jul 2026 12:37:27 -0400 Subject: [PATCH 1/3] [None][chore] Remove 2-model code and tests Signed-off-by: Mike Iovine --- .pre-commit-config.yaml | 4 - .../blog06_Llama4_maverick_eagle_guide.md | 2 +- .../torch_compile_and_piecewise_cuda_graph.md | 13 +- .../legacy/advanced/speculative-decoding.md | 2 +- .../model_registry/configs/super_v3_mtp.yaml | 1 - examples/llm-api/llm_speculative_decoding.py | 3 +- examples/llm-api/quickstart_advanced.py | 7 - examples/models/core/deepseek_v4/README.md | 3 +- examples/models/core/exaone/README.md | 3 +- .../models/core/mistral_large_3/README.md | 1 - .../core/nemotron/README_nemotron_super_v3.md | 1 - legacy-files.txt | 2 - pyproject.toml | 2 - ruff-legacy.toml | 2 - tensorrt_llm/_torch/auto_deploy/llm_args.py | 22 +- .../_torch/auto_deploy/shim/ad_executor.py | 5 +- tensorrt_llm/_torch/models/modeling_auto.py | 6 - .../_torch/models/modeling_speculative.py | 150 +-- tensorrt_llm/_torch/pyexecutor/_util.py | 79 +- .../_torch/pyexecutor/cuda_graph_runner.py | 35 +- .../_torch/pyexecutor/model_engine.py | 909 +------------- .../_torch/pyexecutor/py_executor_creator.py | 123 +- .../_torch/pyexecutor/sampler/sampler.py | 16 +- tensorrt_llm/_torch/speculative/__init__.py | 3 +- tensorrt_llm/_torch/speculative/drafter.py | 10 +- .../_torch/speculative/drafting_loops.py | 522 -------- tensorrt_llm/_torch/speculative/eagle3.py | 195 +-- tensorrt_llm/_torch/speculative/interface.py | 84 +- .../_torch/speculative/model_drafter.py | 1049 ----------------- tensorrt_llm/_torch/speculative/mtp.py | 4 +- tensorrt_llm/_torch/speculative/utils.py | 104 +- tensorrt_llm/llmapi/llm_args.py | 77 +- .../usage/llm_args_golden_manifest.json | 14 - tests/integration/defs/.test_durations | 46 +- .../integration/defs/.test_durations_aws_dfw | 8 +- .../accuracy/test_disaggregated_serving.py | 13 +- .../defs/accuracy/test_llm_api_autodeploy.py | 3 +- .../defs/accuracy/test_llm_api_pytorch.py | 205 +--- .../defs/disaggregated/test_ad_disagg.py | 1 - ...fig_ctxtp2_gentp2_gptoss_eagle_triton.yaml | 1 - ...fig_ctxtp2_gentp2_gptoss_eagle_trtllm.yaml | 1 - .../test_disaggregated_single_gpu.py | 8 +- .../examples/test_ad_speculative_decoding.py | 3 +- .../defs/perf/pytorch_model_config.py | 1 - .../integration/defs/perf/test_perf_sanity.py | 2 - tests/integration/defs/test_e2e.py | 3 - .../test_lists/qa/llm_function_core.txt | 77 +- .../test_lists/qa/llm_function_rtx6k.txt | 35 +- .../test_lists/qa/llm_spark_func.yml | 3 +- .../test_lists/test-db/l0_b200.yml | 1 - .../test_lists/test-db/l0_dgx_b200.yml | 17 +- .../test_lists/test-db/l0_dgx_h100.yml | 10 +- .../test-db/l0_gb200_multi_gpus.yml | 8 +- .../test_lists/test-db/l0_h100.yml | 27 +- .../test_lists/test-db/l0_rtx_pro_6000.yml | 16 +- tests/integration/test_lists/waives.txt | 8 +- ...dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml | 1 - ...x1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml | 1 - ...dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml | 1 - ...x1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml | 1 - ...dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml | 1 - ...x1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml | 1 - ...x1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml | 1 - ...dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml | 1 - ...x1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml | 1 - .../test_deepseek_v4_cache_manager.py | 4 +- .../executor/test_dual_pool_kv_cache.py | 2 - .../executor/test_kv_cache_estimation.py | 1 - .../_torch/executor/test_py_executor.py | 11 +- .../_torch/executor/test_seq_slot_sizing.py | 2 +- tests/unittest/_torch/helpers.py | 1 - .../hw_agnostic/test_draft_len_schedule.py | 106 +- ...test_draft_token_prepare_for_generation.py | 455 ------- .../test_draft_token_tree_sampling.py | 939 --------------- .../test_draft_token_tree_verification.py | 1 - .../hw_agnostic/test_dynamic_spec_decode.py | 1 - .../hw_agnostic/test_get_num_spec_layers.py | 6 +- .../hw_agnostic/test_kv_cache_reuse.py | 1 - .../speculative/hw_agnostic/test_mtp.py | 1 - .../speculative/hw_agnostic/test_spec_gate.py | 1 - .../_torch/speculative/test_eagle3.py | 106 +- .../singlegpu/shim/test_llm_config.py | 8 +- .../smoke/test_ad_speculative_decoding.py | 7 +- .../singlegpu/smoke/test_disagg.py | 1 - tests/unittest/llmapi/test_llm_args.py | 39 +- 85 files changed, 427 insertions(+), 5225 deletions(-) delete mode 100644 tensorrt_llm/_torch/speculative/drafting_loops.py delete mode 100644 tensorrt_llm/_torch/speculative/model_drafter.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_prepare_for_generation.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_sampling.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 37c51fbbb377..a4025009cdfb 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -323,10 +323,8 @@ common-files: &common_files | tensorrt_llm/_torch/speculative/__init__.py | tensorrt_llm/_torch/speculative/auto_heuristic.py | tensorrt_llm/_torch/speculative/drafter.py | - tensorrt_llm/_torch/speculative/drafting_loops.py | tensorrt_llm/_torch/speculative/eagle3.py | tensorrt_llm/_torch/speculative/interface.py | - tensorrt_llm/_torch/speculative/model_drafter.py | tensorrt_llm/_torch/speculative/mtp.py | tensorrt_llm/_torch/speculative/ngram.py | tensorrt_llm/_torch/speculative/save_hidden_state.py | @@ -1093,10 +1091,8 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/speculative/__init__.py | tensorrt_llm/_torch/speculative/auto_heuristic.py | tensorrt_llm/_torch/speculative/drafter.py | - tensorrt_llm/_torch/speculative/drafting_loops.py | tensorrt_llm/_torch/speculative/eagle3.py | tensorrt_llm/_torch/speculative/interface.py | - tensorrt_llm/_torch/speculative/model_drafter.py | tensorrt_llm/_torch/speculative/mtp.py | tensorrt_llm/_torch/speculative/ngram.py | tensorrt_llm/_torch/speculative/save_hidden_state.py | diff --git a/docs/source/blogs/tech_blog/blog06_Llama4_maverick_eagle_guide.md b/docs/source/blogs/tech_blog/blog06_Llama4_maverick_eagle_guide.md index 9f49d7920b36..f835cde9642c 100644 --- a/docs/source/blogs/tech_blog/blog06_Llama4_maverick_eagle_guide.md +++ b/docs/source/blogs/tech_blog/blog06_Llama4_maverick_eagle_guide.md @@ -68,7 +68,7 @@ docker run -d --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \ -p 8000:8000 --gpus=all -e "TRTLLM_ENABLE_PDL=1" \ -v /path/to/maverick:/config/models/maverick -v /path/to/eagle:/config/models/eagle \ docker.io//tensorrt_llm:main sh \ - -c "echo -e 'enable_autotuner: false\nenable_attention_dp: false\nenable_min_latency: true\ncuda_graph_config:\n max_batch_size: 8\nspeculative_config:\n decoding_type: Eagle3\n max_draft_len: 3\n speculative_model_dir: /config/models/eagle\n eagle3_one_model: true\nkv_cache_config:\n enable_block_reuse: false' > c.yaml && \ + -c "echo -e 'enable_autotuner: false\nenable_attention_dp: false\nenable_min_latency: true\ncuda_graph_config:\n max_batch_size: 8\nspeculative_config:\n decoding_type: Eagle3\n max_draft_len: 3\n speculative_model_dir: /config/models/eagle\nkv_cache_config:\n enable_block_reuse: false' > c.yaml && \ TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL=True \ trtllm-serve /config/models/maverick \ --host 0.0.0.0 --port 8000 \ diff --git a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md index 9511be887b1d..56b780e59ef2 100644 --- a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md +++ b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md @@ -88,18 +88,7 @@ Even with Piecewise CUDA Graph enabled, you may still observe bubbles in the con Torch compile cannot work with multi-ModelEngine config. -1. Speculative Decoding in Two-Model Style - -``` yaml -speculative_config: - decoding_type: "MTP" - mtp_eagle_one_model: False # Not supported - -speculative_config: - decoding_type: "Eagle3" - eagle3_one_model: False # Not supported -``` - +1. NGram or User-Provided Speculative Decoding 2. Multimodal Model Family ## Development Guide diff --git a/docs/source/legacy/advanced/speculative-decoding.md b/docs/source/legacy/advanced/speculative-decoding.md index 2faf885f8d10..257f2e96473a 100644 --- a/docs/source/legacy/advanced/speculative-decoding.md +++ b/docs/source/legacy/advanced/speculative-decoding.md @@ -175,7 +175,7 @@ Similarly to ReDrafter, TensorRT-LLM implements the EAGLE model such that logits ### Disaggregated Serving -[Disaggregated Serving](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/features/disaggregated-service.md) with EAGLE3 using the two model approach is supported in the Pytorch backend. Please refer to the following [Dynamo example](https://github.com/ai-dynamo/dynamo/blob/main/examples/tensorrt_llm/llama4_plus_eagle.md) on how to run EAGLE3 with Disaggregated Serving for Llama 4 Maverick. +[Disaggregated Serving](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/features/disaggregated-service.md) with EAGLE3 is supported in the Pytorch backend. Please refer to the following [Dynamo example](https://github.com/ai-dynamo/dynamo/blob/main/examples/tensorrt_llm/llama4_plus_eagle.md) on how to run EAGLE3 with Disaggregated Serving for Llama 4 Maverick. ## Lookahead Decoding diff --git a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml index e30ec72e3ea4..5c229b523bf1 100644 --- a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml +++ b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml @@ -15,7 +15,6 @@ kv_cache_config: speculative_config: decoding_type: MTP num_nextn_predict_layers: 6 - mtp_eagle_one_model: true attention_dp_config: enable_balance: true batching_wait_iters: 10 diff --git a/examples/llm-api/llm_speculative_decoding.py b/examples/llm-api/llm_speculative_decoding.py index af24a78cd4ea..f878b64ca6ce 100644 --- a/examples/llm-api/llm_speculative_decoding.py +++ b/examples/llm-api/llm_speculative_decoding.py @@ -34,8 +34,7 @@ def run_MTP(model: Optional[str] = None): def run_Eagle3(): spec_config = Eagle3DecodingConfig( max_draft_len=3, - speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", - eagle3_one_model=True) + speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B") kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8) diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index d2acd490d44a..e38626ec42f0 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -210,9 +210,6 @@ def add_llm_args(parser): parser.add_argument('--spec_decode_max_draft_len', type=int, default=1) parser.add_argument('--draft_model_dir', type=str, default=None) parser.add_argument('--max_matching_ngram_size', type=int, default=5) - parser.add_argument('--use_one_model', - default=True, - action=argparse.BooleanOptionalAction) parser.add_argument('--eagle_choices', type=str, default=None) parser.add_argument('--use_dynamic_tree', default=False, @@ -301,15 +298,12 @@ def setup_llm(args, **kwargs): ) if args.spec_decode_algo is not None else None if spec_decode_algo == 'MTP': - if not args.use_one_model: - print("Running MTP eagle with two model style.") spec_config = MTPDecodingConfig( max_draft_len=args.spec_decode_max_draft_len, use_relaxed_acceptance_for_thinking=args. use_relaxed_acceptance_for_thinking, relaxed_topk=args.relaxed_topk, relaxed_delta=args.relaxed_delta, - mtp_eagle_one_model=args.use_one_model, use_dynamic_tree=args.use_dynamic_tree, dynamic_tree_max_topK=args.dynamic_tree_max_topK, max_total_draft_tokens=args.max_total_draft_tokens, @@ -318,7 +312,6 @@ def setup_llm(args, **kwargs): spec_config = Eagle3DecodingConfig( max_draft_len=args.spec_decode_max_draft_len, speculative_model=args.draft_model_dir, - eagle3_one_model=args.use_one_model, eagle_choices=args.eagle_choices, use_dynamic_tree=args.use_dynamic_tree, dynamic_tree_max_topK=args.dynamic_tree_max_topK, diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md index fe7d4976c114..0e286acc1a56 100644 --- a/examples/models/core/deepseek_v4/README.md +++ b/examples/models/core/deepseek_v4/README.md @@ -182,8 +182,7 @@ python quickstart_advanced.py \ --max_seq_len 4096 \ --kv_cache_fraction 0.5 \ --spec_decode_algo MTP \ - --spec_decode_max_draft_len N \ - --use_one_model + --spec_decode_max_draft_len N ``` `N` is the number of draft tokens to predict. Start with `N=1` for bring-up, then increase it after diff --git a/examples/models/core/exaone/README.md b/examples/models/core/exaone/README.md index 5050a4be7822..f9275f58e4f1 100644 --- a/examples/models/core/exaone/README.md +++ b/examples/models/core/exaone/README.md @@ -197,8 +197,7 @@ python ../../../llm-api/quickstart_advanced.py \ --moe_ep_size 8 \ --enable_attention_dp \ --spec_decode_algo MTP \ - --spec_decode_max_draft_len N \ - --use_one_model + --spec_decode_max_draft_len N ``` `N` is the number of MTP modules. When `N` is equal to `0`, which means that MTP is not used (default). When `N` is greater than `0`, which means that `N` MTP modules are enabled. In the current implementation, the weight of each MTP module is shared. diff --git a/examples/models/core/mistral_large_3/README.md b/examples/models/core/mistral_large_3/README.md index da219bf7b0a8..c0542f4085bf 100644 --- a/examples/models/core/mistral_large_3/README.md +++ b/examples/models/core/mistral_large_3/README.md @@ -46,7 +46,6 @@ mpirun -n 1 --allow-run-as-root --oversubscribe python3 examples/llm-api/quickst --checkpoint_format mistral \ --spec_decode_algo EAGLE3 \ --spec_decode_max_draft_len 1 \ - --use_one_model \ --draft_model_dir ${mistral_large_3_eagle_model_path} \ --eagle3_model_arch mistral_large3 \ --moe_backend TRTLLM diff --git a/examples/models/core/nemotron/README_nemotron_super_v3.md b/examples/models/core/nemotron/README_nemotron_super_v3.md index 1e59febce82e..937bad9251fb 100644 --- a/examples/models/core/nemotron/README_nemotron_super_v3.md +++ b/examples/models/core/nemotron/README_nemotron_super_v3.md @@ -182,7 +182,6 @@ python3 examples/llm-api/quickstart_advanced.py \ --moe_backend=TRTLLM \ --spec_decode_algo=MTP \ --spec_decode_max_draft_len=3 \ - --use_one_model \ --tp_size=8 \ --moe_ep_size 8 \ --apply_chat_template diff --git a/legacy-files.txt b/legacy-files.txt index 021a4fb301a8..25313bfd8278 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -315,10 +315,8 @@ tensorrt_llm/_torch/shared_tensor/shared_tensor.py tensorrt_llm/_torch/speculative/__init__.py tensorrt_llm/_torch/speculative/auto_heuristic.py tensorrt_llm/_torch/speculative/drafter.py -tensorrt_llm/_torch/speculative/drafting_loops.py tensorrt_llm/_torch/speculative/eagle3.py tensorrt_llm/_torch/speculative/interface.py -tensorrt_llm/_torch/speculative/model_drafter.py tensorrt_llm/_torch/speculative/mtp.py tensorrt_llm/_torch/speculative/ngram.py tensorrt_llm/_torch/speculative/save_hidden_state.py diff --git a/pyproject.toml b/pyproject.toml index 2efd23c84d48..a46b57e1e215 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -372,10 +372,8 @@ exclude = [ "tensorrt_llm/_torch/speculative/__init__.py", "tensorrt_llm/_torch/speculative/auto_heuristic.py", "tensorrt_llm/_torch/speculative/drafter.py", - "tensorrt_llm/_torch/speculative/drafting_loops.py", "tensorrt_llm/_torch/speculative/eagle3.py", "tensorrt_llm/_torch/speculative/interface.py", - "tensorrt_llm/_torch/speculative/model_drafter.py", "tensorrt_llm/_torch/speculative/mtp.py", "tensorrt_llm/_torch/speculative/ngram.py", "tensorrt_llm/_torch/speculative/save_hidden_state.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 39a722295b41..23c5cb22e1c3 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -332,10 +332,8 @@ include = [ "tensorrt_llm/_torch/speculative/__init__.py", "tensorrt_llm/_torch/speculative/auto_heuristic.py", "tensorrt_llm/_torch/speculative/drafter.py", - "tensorrt_llm/_torch/speculative/drafting_loops.py", "tensorrt_llm/_torch/speculative/eagle3.py", "tensorrt_llm/_torch/speculative/interface.py", - "tensorrt_llm/_torch/speculative/model_drafter.py", "tensorrt_llm/_torch/speculative/mtp.py", "tensorrt_llm/_torch/speculative/ngram.py", "tensorrt_llm/_torch/speculative/save_hidden_state.py", diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index f0c2af09cf6b..01cb42b00ed6 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -116,24 +116,18 @@ def validate_supported_speculative_config(self): return self if isinstance(spec_config, MTPDecodingConfig): - if not spec_config.mtp_eagle_one_model or spec_config.use_mtp_vanilla: + if spec_config.use_mtp_vanilla: raise ValueError( - "AutoDeploy only supports MTP speculative decoding with " - "mtp_eagle_one_model=True and use_mtp_vanilla=False " - f"(got mtp_eagle_one_model={spec_config.mtp_eagle_one_model}, " - f"use_mtp_vanilla={spec_config.use_mtp_vanilla})." + "AutoDeploy only supports EAGLE-style MTP speculative " + "decoding with use_mtp_vanilla=False " + f"(got use_mtp_vanilla={spec_config.use_mtp_vanilla})." ) elif isinstance(spec_config, EagleDecodingConfig): - if not spec_config.eagle3_one_model: - raise ValueError( - "AutoDeploy only supports Eagle speculative decoding with " - f"eagle3_one_model=True (got eagle3_one_model={spec_config.eagle3_one_model})." - ) + pass else: raise ValueError( "AutoDeploy only supports speculative decoding via " - "MTPDecodingConfig(mtp_eagle_one_model=True) or " - "EagleDecodingConfig(eagle3_one_model=True)." + "MTPDecodingConfig or EagleDecodingConfig." ) self.model_factory = "eagle_one_model" @@ -148,8 +142,8 @@ def setup_hidden_state_capture(self): if isinstance(spec_config, MTPDecodingConfig): if spec_config.max_draft_len is None: raise ValueError( - "MTPDecodingConfig.max_draft_len must not be None when mtp_eagle_one_model is " - "enabled. Ensure num_nextn_predict_layers is set in the model config." + "MTPDecodingConfig.max_draft_len must not be None. " + "Ensure num_nextn_predict_layers is set in the model config." ) capture_layers = {-1} else: diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index 86488bc36cda..186560a18f0d 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -1100,10 +1100,9 @@ def instantiate_sampler( max_draft_len = 0 if spec_config is None else spec_config.max_draft_len max_total_draft_tokens = 0 if spec_config is None else spec_config.tokens_per_gen_step - 1 - # One-model spec dec: model performs sampling internally, returns pre-computed tokens + # One-engine spec dec: model performs sampling internally, returns pre-computed tokens if spec_config is not None and ( - spec_config.spec_dec_mode.is_eagle3_one_model() - or spec_config.spec_dec_mode.is_mtp_eagle_one_model() + spec_config.spec_dec_mode.is_eagle3() or spec_config.spec_dec_mode.is_mtp_eagle() ): sampler_args = TorchSampler.Args( max_seq_len=ad_config.max_seq_len, diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index 2802ff59fd53..353a00402389 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -34,12 +34,6 @@ def _resolve_class(config: ModelConfig) -> Optional[Type]: "") # Strip the appended EAGLE3 model_arch = "EAGLE3" + model_arch - if model_arch in ( - "DeepseekV3ForCausalLM", "Glm4MoeForCausalLM", - "ExaoneMoEForCausalLM" - ) and config.spec_config is not None and config.spec_config.max_draft_len == 0: - model_arch = "MTPDraftModelForCausalLM" - return MODEL_CLASS_MAPPING.get(model_arch) @staticmethod diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index a52534de4f50..ea0068965f37 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1,6 +1,6 @@ import inspect from dataclasses import replace -from typing import Dict, Generic, List, Optional, Tuple +from typing import Dict, Generic, List, Optional import torch import torch.nn.functional as F @@ -1672,7 +1672,7 @@ def __init__( spec_dec_mode = model_config.spec_config.spec_dec_mode assert spec_dec_mode.is_mtp_one_model() checkpoint_mtp_num_layers = model_config.pretrained_config.num_nextn_predict_layers - if spec_dec_mode.is_mtp_eagle_one_model(): + if spec_dec_mode.is_mtp_eagle(): mtp_num_layers = 1 mtp_repeat_count = model_config.spec_config.max_draft_len else: @@ -1691,151 +1691,13 @@ def __init__( self.embed_tokens = model.embed_tokens -class MTPDraftModel(nn.Module): - - def __init__(self, model_config: ModelConfig[PretrainedConfig], - layer_idx: int, aux_stream_dict: Dict[AuxStreamType, - torch.cuda.Stream]): - super().__init__() - # Import here to avoid circular import - model_type = model_config.pretrained_config.model_type - if model_type == "glm4_moe": - from .modeling_glm import Glm4MTP - mtp_layer = Glm4MTP(model_config, - layer_idx, - aux_stream_dict, - is_separate_draft_engine=True) - elif model_type in ["deepseek_v3", "deepseek_v32", "glm_moe_dsa"]: - from .modeling_deepseekv3 import DeepseekV3MTP - mtp_layer = DeepseekV3MTP(model_config, - layer_idx, - aux_stream_dict, - is_separate_draft_engine=True) - elif model_type in ["exaone_moe"]: - from .modeling_exaone_moe import ExaoneMoeMTP - mtp_layer = ExaoneMoeMTP(model_config, layer_idx, aux_stream_dict) - else: - raise ValueError( - f"MTPDraftModel does not support model_type: {model_type}") - setattr(self, f"layers.{layer_idx}", mtp_layer) - self.layers = mtp_layer - self.layer_idx = layer_idx - self.config = model_config.pretrained_config - self.embed_tokens = Embedding( - self.config.vocab_size, - self.config.hidden_size, - dtype=self.config.torch_dtype, - ) - - def __repr__(self): - """Custom string representation to display layer index""" - return f"(layers): ({self.layer_idx}): {repr(self.layers)}" - - def forward( - self, - input_ids: torch.IntTensor, - position_ids: torch.IntTensor, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - all_rank_num_tokens: Optional[List[int]] = None, - spec_metadata: Optional[SpecMetadata] = None, - **kwargs, - ) -> Tuple[torch.Tensor, torch.Tensor]: - hidden_states = self.layers( - input_ids, - position_ids, - hidden_states, - embed_tokens=self.embed_tokens, - attn_metadata=attn_metadata, - all_rank_num_tokens=all_rank_num_tokens, - spec_metadata=spec_metadata, - ) - - return hidden_states - - -@register_auto_model("MTPDraftModelForCausalLM") -class MTPDraftModelForCausalLM(DecoderModelForCausalLM[MTPDraftModel, - PretrainedConfig]): - - def __init__(self, model_config: ModelConfig[PretrainedConfig]): - self.model_config = model_config - aux_stream_list = [torch.cuda.Stream() for _ in range(4)] - self.aux_stream_dict = { - AuxStreamType.Attention: aux_stream_list[0], - AuxStreamType.MoeShared: aux_stream_list[0], - AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], - AuxStreamType.MoeBalancer: aux_stream_list[2], - AuxStreamType.MoeOutputMemset: aux_stream_list[3], - } - super().__init__( - MTPDraftModel(self.model_config, - self.model_config.pretrained_config.num_hidden_layers, - self.aux_stream_dict), - config=self.model_config, - hidden_size=self.model_config.pretrained_config.hidden_size, - vocab_size=self.model_config.pretrained_config.vocab_size) - - def load_weights(self, weights: Dict): - # Import here to avoid circular import - model_type = self.model_config.pretrained_config.model_type - match model_type: - case "glm4_moe": - from .modeling_glm import Glm4WeightLoader - weight_loader = Glm4WeightLoader(self, is_draft_model=True) - case "deepseek_v3" | "deepseek_v32" | "glm_moe_dsa": - from .modeling_deepseekv3 import DeepseekV3WeightLoader - weight_loader = DeepseekV3WeightLoader(self, - is_draft_model=True) - case "exaone_moe": - raise ValueError( - f"Model type {model_type} not supported for MTP for two engine mode. Please use one engine mode instead." - ) - case _: - raise ValueError( - f"Model type {model_type} not supported for MTP") - weight_loader.load_weights(weights) - - def load_weights_from_target_model(self, - target_model: torch.nn.Module) -> None: - if self.model.embed_tokens is None: - self.model.embed_tokens = target_model.model.embed_tokens - self.lm_head = target_model.lm_head - - def forward(self, - attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor = None, - position_ids: torch.IntTensor = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - return_context_logits: bool = False, - spec_metadata: Optional[SpecMetadata] = None, - hidden_states: torch.Tensor = None, - **kwargs) -> torch.Tensor: - - hidden_states = spec_metadata.get_hidden_states() - output = self.model( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - all_rank_num_tokens=attn_metadata.all_rank_num_tokens, - spec_metadata=spec_metadata, - **kwargs) - return self.logits_processor.forward( - output, - self.lm_head, - attn_metadata, - return_context_logits, - ) - - def get_draft_model(model_config, draft_config, lm_head, model): """Construct the draft model for the configured speculative-decoding mode (EAGLE3 / MTP / PARD / DFlash). The DFlash branch selects the Laguna drafter by detecting its architecture in the draft checkpoint's own config.""" assert getattr(model_config, 'spec_config', None) is not None spec_dec_mode = model_config.spec_config.spec_dec_mode - if spec_dec_mode.is_eagle3_one_model(): + if spec_dec_mode.is_eagle3(): if model_config.spec_config.eagle3_model_arch == "llama3": # Eagle3ForCausalLM handles both Llama3 and DeepSeekV3 architectures return Eagle3ForCausalLM( @@ -1853,8 +1715,6 @@ def get_draft_model(model_config, draft_config, lm_head, model): return MTPForCausalLM(model_config, model_config.pretrained_config.num_hidden_layers, lm_head, model) - elif spec_dec_mode.is_mtp_eagle(): - return MTPDraftModelForCausalLM(model_config) elif spec_dec_mode.is_pard(): return PARDForCausalLM(draft_config) elif spec_dec_mode.is_dflash(): @@ -1877,7 +1737,7 @@ def get_draft_model(model_config, draft_config, lm_head, model): num_stages=num_stages, block_size=model_config.spec_config.block_size, ) - elif spec_dec_mode.is_draft_target_one_model(): + elif spec_dec_mode.is_draft_target(): # Keep the draft LM head vocab-sharded so greedy draft sampling uses the # lighter TP gather (see SpecWorkerBase.greedy_sample_draft_with_tp_gather). was_frozen = draft_config._frozen @@ -1908,7 +1768,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): if spec_config and spec_config.spec_dec_mode.use_one_engine(): # Only create draft_model for modes MTP, Eagle3 (not SA) if not spec_config.spec_dec_mode.is_sa(): - if spec_config.spec_dec_mode.is_eagle3_one_model(): + if spec_config.spec_dec_mode.is_eagle3(): if spec_config.eagle3_model_arch == "mistral_large3": from tensorrt_llm._torch.models.checkpoints.mistral.config_loader import \ MistralConfigLoader diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0149ba7ae03a..b85126b5cb66 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -293,7 +293,6 @@ def __init__( self, *, model_engine: PyTorchModelEngine, - draft_model_engine: Optional[PyTorchModelEngine], mapping: Mapping, net_max_seq_len: int, kv_connector_manager: Optional[KvCacheConnectorManager], @@ -313,7 +312,6 @@ def __init__( skip_est: bool = False, ): self._model_engine = model_engine - self._draft_model_engine = draft_model_engine self._mapping = mapping self._kv_cache_config = kv_cache_config self._max_kv_tokens_in = self._kv_cache_config.max_tokens @@ -461,14 +459,7 @@ def _get_kv_size_per_token(self, model_config, kv_cache_config) if self._is_encoder_decoder(): total += CacheCost.from_raw(self._get_cross_kv_size_per_token()) - if self._draft_model_engine is not None: - draft_model_config = self._draft_model_engine.model.model_config - draft_kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( - self._draft_model_engine, kv_cache_config) - total += self._per_manager_cache_cost(draft_kv_cache_manager_cls, - draft_model_config, - kv_cache_config) - elif self._should_create_separate_draft_kv_cache(): + if self._should_create_separate_draft_kv_cache(): # One-model draft with separate KV cache layout. # Pass num_layers explicitly since the HF config may report a # different layer count than what is actually used at runtime @@ -864,8 +855,7 @@ def configure_kv_cache_capacity(self, # get kv cache stats for both model and draft model kv_stats = py_executor.resource_manager.resource_managers.get( ResourceManagerType.KV_CACHE_MANAGER).get_kv_cache_stats() - # Get draft KV cache stats if present (either from two-model mode or one-model - # mode with separate draft KV cache) + # Get draft KV cache stats if present (one-model separate draft pool) draft_kv_cache_manager = py_executor.resource_manager.resource_managers.get( ResourceManagerType.DRAFT_KV_CACHE_MANAGER) kv_stats_draft = draft_kv_cache_manager.get_kv_cache_stats( @@ -1580,9 +1570,7 @@ def build_managers(self, # Split combined KV cache budgets before creating managers. Skip during # estimation — estimation uses max_tokens-based logic and must not # mutate the config. - has_draft = ( - self._draft_model_engine is not None # two-model - or self._should_create_separate_draft_kv_cache()) # one-model + has_draft = self._should_create_separate_draft_kv_cache() draft_kv_cache_config = None if not estimating_kv_cache and has_draft: # Used when each manager sizes pools from max_gpu_total_bytes (V2 @@ -1592,43 +1580,23 @@ def build_managers(self, self._split_kv_cache_budget_for_draft( "max_gpu_total_bytes", self_kv_cache_config, draft_kv_cache_config)) - # KVCacheManagerV2 does not support two-model draft budget splitting. - v2_two_model = (self._is_kv_cache_manager_v2 - and self._draft_model_engine is not None) - if not v2_two_model: - # Each manager sizes its host pool from host_cache_size directly. - self_kv_cache_config, draft_kv_cache_config = ( - self._split_kv_cache_budget_for_draft( - "host_cache_size", self_kv_cache_config, - draft_kv_cache_config)) + # Each manager sizes its host pool from host_cache_size directly. + self_kv_cache_config, draft_kv_cache_config = ( + self._split_kv_cache_budget_for_draft("host_cache_size", + self_kv_cache_config, + draft_kv_cache_config)) kv_cache_manager = self._create_kv_cache_manager( self._model_engine, estimating_kv_cache, kv_cache_config_override=self_kv_cache_config) - if (not estimating_kv_cache and self._kv_connector_manager is not None - and self._draft_model_engine is not None): - raise NotImplementedError( - "Connector manager is not supported for draft model.") - draft_kv_cache_manager = None draft_build_kv_cache_config = (draft_kv_cache_config if draft_kv_cache_config is not None else self_kv_cache_config) - # Two-model speculative decoding: draft model has separate engine - if self._draft_model_engine is not None: - if self._is_kv_cache_manager_v2: - assert draft_kv_cache_config is None, ( - "KVCacheManagerV2 does not support two-model speculative " - "decoding with separate draft KV cache budget splitting.") - draft_kv_cache_manager = self._create_kv_cache_manager( - self._draft_model_engine, - estimating_kv_cache, - kv_cache_config_override=draft_build_kv_cache_config) - # One-model speculative decoding with different KV layouts - elif self._should_create_separate_draft_kv_cache(): + if self._should_create_separate_draft_kv_cache(): draft_kv_cache_manager = self._create_one_model_draft_kv_cache_manager( estimating_kv_cache, kv_cache_config_override=draft_build_kv_cache_config) @@ -1740,7 +1708,7 @@ def _create_kv_cache_manager( dtype = model_engine.dtype if is_draft is None: - is_draft = model_engine.is_draft_model + is_draft = False if kv_cache_type is None: kv_cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF @@ -2122,7 +2090,7 @@ def validate_kv_cache_compression_with_spec( # Evicting methods co-compact the draft KV, so the draft must be a # standard paged cache in the same forward (one-model speculation). mode = spec_config.spec_dec_mode - if not (mode.is_mtp_one_model() or mode.is_eagle3_one_model()): + if not (mode.is_mtp_one_model() or mode.is_eagle3()): raise ValueError( f"KV-cache compression algorithm {config.algorithm!r} does not " f"support speculative decoding mode {mode.name}: the draft KV " @@ -2184,7 +2152,7 @@ def should_enable_dsv4_overlap_headroom( """Gate extra sequence slots to the validated DSv4 MTP overlap path.""" return (should_enable_dsv4_adp_dummy_fixes(model_type, mapping) and spec_config is not None - and spec_config.spec_dec_mode.is_mtp_eagle_one_model() + and spec_config.spec_dec_mode.is_mtp_eagle() and not disable_overlap_scheduler) @@ -2880,8 +2848,7 @@ def init_feature_status(llm_args) -> Dict[str, bool]: "disaggregated_serving", "chunked_prefill", "mtp", - "eagle3_one_model", - "eagle3_two_model", + "eagle3", "torch_sampler", "trtllm_sampler", "kv_cache_reuse", @@ -2898,12 +2865,8 @@ def init_feature_status(llm_args) -> Dict[str, bool]: feature_status["chunked_prefill"] = llm_args.enable_chunked_prefill feature_status["mtp"] = isinstance(llm_args.speculative_config, MTPDecodingConfig) - feature_status["eagle3_one_model"] = ( - isinstance(llm_args.speculative_config, EagleDecodingConfig) - and llm_args.speculative_config.eagle3_one_model) - feature_status["eagle3_two_model"] = ( - isinstance(llm_args.speculative_config, EagleDecodingConfig) - and not llm_args.speculative_config.eagle3_one_model) + feature_status["eagle3"] = isinstance(llm_args.speculative_config, + EagleDecodingConfig) feature_status[ "torch_sampler"] = sampler_type == SamplerType.TorchSampler feature_status[ @@ -2936,17 +2899,9 @@ def init_feature_status(llm_args) -> Dict[str, bool]: " Please use sampler type auto instead." }, { - "features": ["trtllm_sampler", "eagle3_one_model"], - "message": - ERR_MSG_TMPL.format(feature1="trtllm_sampler", - feature2="eagle3_one_model") + - " Please use sampler type auto instead." - }, - { - "features": ["trtllm_sampler", "eagle3_two_model"], + "features": ["trtllm_sampler", "eagle3"], "message": - ERR_MSG_TMPL.format(feature1="trtllm_sampler", - feature2="eagle3_two_model") + + ERR_MSG_TMPL.format(feature1="trtllm_sampler", feature2="eagle3") + " Please use sampler type auto instead." }, # Add new conflict rules here in the future diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index b8b67ed5ded8..5e095033eba9 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -18,7 +18,6 @@ from ..expert_statistic import ExpertStatistic from ..memory_buffer_utils import get_memory_buffers from ..modules.multi_stream_utils import with_multi_stream -from ..speculative.eagle3 import Eagle3ResourceManager from ..speculative.interface import SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.utils import get_draft_kv_cache_manager @@ -102,7 +101,6 @@ class CUDAGraphRunnerConfig: use_mrope: bool original_max_draft_len: int original_max_total_draft_tokens: int - is_draft_model: bool enable_attention_dp: bool is_encoder_decoder: bool batch_size: int @@ -245,24 +243,16 @@ def get_graph_key( is_all_greedy_sample = bool( getattr(spec_metadata, "is_all_greedy_sample", True)) - if self.config.is_draft_model and spec_resource_manager is not None and isinstance( - spec_resource_manager, Eagle3ResourceManager): - # If 'is_first_draft' is True, even with tree decoding, the length of draft_len will only be 'max_draft_len', not 'max_total_draft_token'. - # Because we will pad the input to 'max_draft_len' length for the first draft layer. - draft_len = self.config.original_max_draft_len if spec_resource_manager.is_first_draft else 0 - key = (batch_size, draft_len, spec_resource_manager.is_first_draft, - short_seq_len_mode, is_all_greedy_sample) - else: - # With dynamic spec decode, the draft length may be zero even when enable_spec_decode is True, - # so we need to get the draft length from the batch instead of using enable_spec_decode. - draft_len_list = [] - for request in batch.generation_requests: - draft_len_list.append(len(request.py_draft_tokens)) - draft_len = max(draft_len_list) - assert len( - set(draft_len_list)) == 1, "All draft lengths must be the same" - key = (batch_size, draft_len, False, short_seq_len_mode, - is_all_greedy_sample) + # With dynamic spec decode, the draft length may be zero even when enable_spec_decode is True, + # so we need to get the draft length from the batch instead of using enable_spec_decode. + draft_len_list = [] + for request in batch.generation_requests: + draft_len_list.append(len(request.py_draft_tokens)) + draft_len = max(draft_len_list) + assert len( + set(draft_len_list)) == 1, "All draft lengths must be the same" + key = (batch_size, draft_len, False, short_seq_len_mode, + is_all_greedy_sample) return key @staticmethod @@ -440,11 +430,6 @@ def capture(self, def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, capture_inputs: Dict[str, Any]): - is_first_draft = key[2] - needs_kv_cache_recompute = True if enable_spec_decode and self.config.spec_config.spec_dec_mode.needs_kv_cache_recompute( - ) else False - if is_first_draft and self.config.is_draft_model and needs_kv_cache_recompute: - capture_inputs['attn_metadata'].use_spec_decoding = True return forward_fn(capture_inputs) output = None diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 50aa317e9e43..b22a33967559 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -66,8 +66,6 @@ prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, update_spec_config_from_loaded_model) -from ..speculative.drafting_loops import BaseDraftingLoopWrapper -from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..utils import (get_model_extra_attrs, set_per_request_piecewise_cuda_graph_flag, @@ -122,10 +120,9 @@ def _filter_piecewise_capture_num_tokens( max_num_tokens: int, max_batch_size: int, max_seq_len: int, - num_extra_decoding_steps: int = 0, ) -> Tuple[list[int], list[int]]: """Cap piecewise CUDA graph capture candidates at the engine's reachable - `num_tokens` ceiling `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` + `num_tokens` ceiling `max_batch_size * (max_seq_len - 1)` clamping user-requested sizes above it down to the ceiling. Each in-flight request must leave room for at least one decode token, @@ -145,8 +142,7 @@ def _filter_piecewise_capture_num_tokens( ceiling but within `max_num_tokens` (the clamped ones, reported so the caller's warning fires). """ - max_capturable_num_tokens = max( - 0, max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)) + max_capturable_num_tokens = max(0, max_batch_size * (max_seq_len - 1)) piecewise_capacity_limit = min(max_num_tokens, max_capturable_num_tokens) if piecewise_capacity_limit > 0: kept = sorted({ @@ -280,9 +276,6 @@ def __init__( attn_runtime_features: Optional[AttentionRuntimeFeatures] = None, dist: Optional[Distributed] = None, spec_config: Optional[DecodingBaseConfig] = None, - is_draft_model: bool = False, - drafting_loop_wrapper: Optional[Callable[[torch.nn.Module], - torch.nn.Module]] = None, model: Optional[torch.nn.Module] = None, checkpoint_loader: Optional[BaseCheckpointLoader] = None, model_weights_memory_tag: Optional[str] = None, @@ -357,21 +350,10 @@ def __init__( else: self.max_draft_loop_tokens = self.original_max_total_draft_tokens - preserve_wrapped_eagle3_widths = (spec_config is not None - and is_draft_model - and drafting_loop_wrapper is not None - and - spec_config.spec_dec_mode.is_eagle3()) - # The draft model won't have any draft tokens attached to - # generation requests when we invoke it autoregressively - if spec_config is not None and is_draft_model and not preserve_wrapped_eagle3_widths: - spec_config.max_draft_len = 0 - spec_config.max_total_draft_tokens = 0 self.spec_config = spec_config self.is_spec_decode = spec_config is not None - self.sparse_attention_config = None if is_draft_model else llm_args.sparse_attention_config + self.sparse_attention_config = llm_args.sparse_attention_config self.enable_spec_decode = self.is_spec_decode - self.is_draft_model = is_draft_model self.attn_runtime_features = attn_runtime_features or AttentionRuntimeFeatures( ) @@ -393,8 +375,7 @@ def __init__( and llm_args.multimodal_config.encoder_cache_max_bytes > 0), ) if model is None: - lora_config: Optional[ - LoraConfig] = None if is_draft_model else llm_args.lora_config + lora_config: Optional[LoraConfig] = llm_args.lora_config # Keep the model_loader to support reloading the model weights later self.model_loader = ModelLoader( llm_args=llm_args, @@ -430,11 +411,6 @@ def __init__( llm_args.disable_overlap_scheduler, enable_overlap_headroom=self._enable_dsv4_overlap_headroom, ) - if drafting_loop_wrapper is not None: - self.model = drafting_loop_wrapper(self.model) - self.model_is_wrapped = True - else: - self.model_is_wrapped = False self.sparse_attention_config = self.model.model_config.sparse_attention_config # In case that some tests use stub models and override `_load_model`. if not hasattr(self.model, 'extra_attrs'): @@ -516,21 +492,19 @@ def __init__( torch_compile_piecewise_cuda_graph_num_tokens or cuda_graph_batch_sizes or []) - num_extra_decoding_steps = self._get_num_extra_decoding_steps() self._piecewise_cuda_graph_num_tokens, unrecordable = ( _filter_piecewise_capture_num_tokens( piecewise_cuda_graph_num_tokens, max_num_tokens=self.max_num_tokens, max_batch_size=self.batch_size, max_seq_len=self.max_seq_len, - num_extra_decoding_steps=num_extra_decoding_steps, )) if unrecordable: logger.warning( f"Skipping piecewise CUDA graph capture for num_tokens=" f"{unrecordable}: exceeds reachable ceiling " - f"max_batch_size*(max_seq_len-1-num_extra_decoding_steps)=" - f"{max(0, self.batch_size * (self.max_seq_len - 1 - num_extra_decoding_steps))}. " + f"max_batch_size*(max_seq_len-1)=" + f"{max(0, self.batch_size * (self.max_seq_len - 1))}. " f"Clamping them to the ceiling; raise max_seq_len for larger graphs." ) @@ -599,9 +573,7 @@ def __init__( self.spec_metadata = None if self.is_spec_decode: - if not self.is_draft_model: - update_spec_config_from_loaded_model(self.spec_config, - self.model) + update_spec_config_from_loaded_model(self.spec_config, self.model) max_num_draft_tokens = self.max_draft_loop_tokens * self.batch_size self.draft_tokens_cuda = torch.empty((max_num_draft_tokens, ), dtype=torch.int, @@ -618,7 +590,7 @@ def __init__( self.previous_kv_lens_offsets_cuda = torch.zeros( (self.batch_size, ), dtype=torch.int, device='cuda') self.without_logits = self.spec_config.spec_dec_mode.without_logits( - ) or self.model_is_wrapped + ) self.max_total_draft_tokens = spec_config.tokens_per_gen_step - 1 self.max_draft_len = spec_config.max_draft_len # Mutable per-iteration draft length (updated each iteration when @@ -704,49 +676,10 @@ def __init__( self.mrope_position_ids_cuda = torch.empty( (3, 1, self.max_num_tokens), dtype=torch.int, device='cuda') - # Pre-allocated buffers for draft model to avoid implicit synchronization - # These are used to build index tensors without creating tensors from Python lists - max_first_draft_tokens = self.batch_size * ( - self.original_max_total_draft_tokens + - 1) if spec_config else self.batch_size - tokens_per_draft = self.original_max_total_draft_tokens + 1 - self.idx_accepted_tokens_cache = None - self.draft_token_positions_cache = None - if spec_config: - # Cache for idx_accepted_tokens (pattern: 0,0,0...1,1,1...2,2,2...) - self.idx_accepted_tokens_cache = torch.arange( - max_first_draft_tokens, dtype=torch.long, - device='cuda') // tokens_per_draft - - if self.is_draft_model: - self.draft_ctx_token_indices_cuda = torch.empty((self.batch_size, ), - dtype=torch.long, - device='cuda') - self.draft_ctx_seq_slots_cuda = torch.empty((self.batch_size, ), - dtype=torch.long, - device='cuda') - # Buffers for first_draft requests (max_draft_len+1 tokens per request) - self.draft_first_draft_indices_cuda = torch.empty( - (max_first_draft_tokens, ), dtype=torch.long, device='cuda') - self.draft_first_draft_seq_slots_cuda = torch.empty( - (max_first_draft_tokens, ), dtype=torch.long, device='cuda') - # Buffers for seq_slots and request indices - self.draft_seq_slots_buffer_cuda = torch.empty((self.batch_size, ), - dtype=torch.int, - device='cuda') - self.draft_request_indices_buffer_cuda = torch.empty( - (self.batch_size, ), dtype=torch.int, device='cuda') - - # Pre-computed constant tensors for incremental update optimization - # Cache for token_positions (pattern: 0,1,2...N repeated) - self.draft_token_positions_cache = torch.arange(tokens_per_draft, - dtype=torch.long, - device='cuda') - # We look up this key in resource_manager during forward to find the # kv cache manager. Can be changed to support multiple model engines # with different KV cache managers. - self.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER if is_draft_model else ResourceManagerType.KV_CACHE_MANAGER + self.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER self.lora_model_config: Optional[LoraModelConfig] = None self._trtllm_gen_jit_warmup = False @@ -766,7 +699,6 @@ def __init__( original_max_draft_len=self.original_max_draft_len, original_max_total_draft_tokens=self. original_max_total_draft_tokens, - is_draft_model=self.is_draft_model, enable_attention_dp=self.enable_attention_dp, is_encoder_decoder=self._is_encoder_decoder_model(), batch_size=self.batch_size, @@ -1134,8 +1066,8 @@ def warmup(self, resource_manager: ResourceManager) -> None: AutoTuner.get() can_run_general_warmup = ( - not is_enc_dec and not self.is_draft_model - and not self.mapping.has_cp_helix() and self.guided_decoder is None + not is_enc_dec and not self.mapping.has_cp_helix() + and self.guided_decoder is None and not isinstance(kv_cache_manager, MambaHybridCacheManager)) log_mem_snapshot("warmup/before_warmup") @@ -1430,8 +1362,8 @@ def trtllm_gen_fmha_jit_warmup(): logger.info("Running TRTLLM-Gen FMHA JIT warmup") warmup_requests_configs = [] - if not self.is_draft_model and self.guided_decoder is None: - # doesn't support 2-model speculative draft and guided decoding + if self.guided_decoder is None: + # doesn't support guided decoding warmup_requests_configs.append( (1 + self.max_total_draft_tokens, 1)) # one generation request else: @@ -1487,14 +1419,6 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): self._assert_all_tp_ranks_have_warmup_batch( batch, curr_max_num_tokens) if batch is not None: - # Reset the flag is_first_draft for the draft model. - # This is necessary for overlap scheduler. - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if self.is_draft_model and isinstance( - spec_resource_manager, Eagle3ResourceManager): - spec_resource_manager.is_first_draft = True - self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) @@ -1702,18 +1626,7 @@ def _get_graphs_to_capture( Returns: List of (batch_size, draft_len) tuples for CUDA graph capture. """ - # Case 1: Draft model (two-model speculative decoding) - # Two-model path is deprecated and will be removed in the near future - if self.is_draft_model: - if self.model_is_wrapped and self.is_spec_decode and spec_resource_manager is not None and isinstance( - spec_resource_manager, Eagle3ResourceManager): - # The CDL path uses draft_len > 0 for the number of iterations in the drafting loop. - draft_len = self.original_max_total_draft_tokens - else: - draft_len = self.max_total_draft_tokens - return [(bs, draft_len) for bs in cuda_graph_batch_sizes] - - # Case 2: One-model with dynamic draft length + # Case 1: One-model with dynamic draft length if self.spec_config is not None and self.spec_config.draft_len_schedule is not None and self.spec_config.spec_dec_mode.support_dynamic_draft_len( ): graphs = [(graph_bs, draft_len) for graph_bs, draft_len in @@ -1741,7 +1654,7 @@ def _get_graphs_to_capture( f"Capturing {len(graphs)} graphs: {graphs}") return graphs - # Case 3: Target model (two-model) or one-model without dynamic draft + # Case 2: One-model without dynamic draft # Match the runtime_draft_len semantics enforced in _prepare_tp_inputs: # logical K for linear-tree modes, total tree tokens for tree decoding. # spec_config is None for non-spec models — fall back to max_draft_len (= 0). @@ -1919,11 +1832,9 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: f"Run generation-only CUDA graph {operation} ({label}) " f"for batch size={bs}, draft_len={draft_len}, " f"max_seq_len={max_seq_len}") - self.enable_spec_decode = draft_len > 0 or self.is_draft_model or ( + self.enable_spec_decode = draft_len > 0 or ( self.spec_config is not None and self.spec_config.spec_dec_mode.use_one_engine()) - self._update_draft_inference_state_for_warmup( - batch, draft_len > 0, resource_manager) self.runtime_draft_len = draft_len if self._is_encoder_decoder_model(): prepare_cross_batch(batch, resource_manager) @@ -2042,16 +1953,6 @@ def _release_batch_context(self, batch: Optional[ScheduledRequests], if spec_resource_manager is not None: spec_resource_manager.free_resources(req) - def _get_num_extra_decoding_steps(self) -> int: - """Determines extra decoding steps needed for fused drafting loops.""" - if isinstance(self.model, BaseDraftingLoopWrapper): - return self.model.max_total_draft_tokens - else: - assert not self.model_is_wrapped, ( - f"Please add logic to determine num_extra_decoding_steps for drafting loop {type(self.model)}" - ) - return 0 - def _create_warmup_request( self, resource_manager: ResourceManager, @@ -2074,8 +1975,6 @@ def _create_warmup_request( if num_tokens > self.max_num_tokens or num_tokens > available_tokens: return None - num_extra_decoding_steps = self._get_num_extra_decoding_steps() - if num_gen_requests > self.batch_size: return None num_gen_tokens = num_gen_requests * (1 + self.max_total_draft_tokens) @@ -2087,10 +1986,9 @@ def _create_warmup_request( ctx_requests = [] gen_requests = [] - # For drafting loops, reduce max_seq_len to leave room for extra decoding steps - max_seq_len = self.max_seq_len - 1 - num_extra_decoding_steps + max_seq_len = self.max_seq_len - 1 if max_seq_len < 1: - return None # Not enough sequence length for drafting loop + return None num_full_seqs = 0 num_left_over_tokens = 0 @@ -2138,7 +2036,6 @@ def _create_warmup_request( max_num_draft_tokens=self.max_total_draft_tokens, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) if ctx_requests is None: @@ -2159,7 +2056,6 @@ def _create_warmup_request( kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) if gen_requests is None: @@ -2199,7 +2095,6 @@ def _create_cuda_graph_warmup_request( return None result = ScheduledRequests() - num_extra_decoding_steps = self._get_num_extra_decoding_steps() runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 @@ -2222,7 +2117,6 @@ def _create_cuda_graph_warmup_request( use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, encoder_output_lens=encoder_output_lens, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) if requests is None: @@ -2266,10 +2160,6 @@ def free_warmup_requests() -> None: if max_position_embeddings is not None: token_num = min(token_num, max_position_embeddings - _kv_draft) - assert token_num > num_extra_decoding_steps, ( - "Cannot fuse drafting loop. Not enough KV cache space for all draft tokens." - ) - token_num -= num_extra_decoding_steps token_num = int( token_num) # Ensure int for range() in add_dummy_requests @@ -2283,7 +2173,6 @@ def free_warmup_requests() -> None: max_beam_width=self.max_beam_width, encoder_output_lens=[max_encoder_output_len] if is_enc_dec else None, - num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) if max_seq_len_request is None: @@ -2397,20 +2286,6 @@ def _get_enc_dec_hidden_size(self) -> int: "hidden size from the model config.") return int(hidden_size) - def _update_draft_inference_state_for_warmup( - self, batch: ScheduledRequests, is_first_draft: bool, - resource_manager: ResourceManager): - """Updates request states for specific draft model warmups like Eagle3.""" - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if self.is_draft_model and isinstance(spec_resource_manager, - Eagle3ResourceManager): - spec_resource_manager.is_first_draft = is_first_draft - if is_first_draft: - for req in batch.generation_requests: - req.py_is_first_draft = True - req.py_draft_tokens = [] - def _set_up_attn_metadata( self, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], @@ -2541,7 +2416,6 @@ def _set_up_spec_metadata( self.batch_size, max_num_tokens=self.max_num_tokens, spec_resource_manager=spec_resource_manager, - is_draft_model=self.is_draft_model, max_seq_len=self.max_seq_len, num_seq_slots=num_seq_slots) @@ -2553,7 +2427,6 @@ def _set_up_spec_metadata( self.batch_size, max_num_tokens=self.max_num_tokens, spec_resource_manager=spec_resource_manager, - is_draft_model=self.is_draft_model, max_seq_len=self.max_seq_len, num_seq_slots=num_seq_slots) return self.spec_metadata @@ -2921,8 +2794,8 @@ def _set_spec_metadata_all_rank_num_tokens( # chunk loop identical across EP ranks. if all_rank_num_gens is not None: spec_metadata.all_rank_num_gens = all_rank_num_gens - if (spec_metadata.spec_dec_mode.is_mtp_eagle_one_model() - or spec_metadata.spec_dec_mode.is_eagle3_one_model()): + if (spec_metadata.spec_dec_mode.is_mtp_eagle() + or spec_metadata.spec_dec_mode.is_eagle3()): spec_metadata.subseq_all_rank_num_tokens = all_rank_num_seqs def _get_padding_params( @@ -3165,501 +3038,6 @@ def _ship_multimodal_indices( inputs['text_token_indices'] = text_token_indices_cpu.to( "cuda", non_blocking=True) - def _can_use_incremental_update( - self, scheduled_requests: ScheduledRequests, - new_tokens_device: Optional[torch.Tensor], - next_draft_tokens_device: Optional[torch.Tensor]) -> bool: - """ - Check if we can use incremental update for the given scheduled requests and new tensors device. - """ - # Not use this approach for non-speculative decoding - if self.spec_config is None: - return False - - # Not allowed for one-model speculative decoding - if not self.spec_config.spec_dec_mode.has_draft_model(): - return False - - if not self.cuda_graph_runner.enabled: - return False - - if self.use_mrope: - return False - - # Not allowed for non-overlap scheduler - if new_tokens_device is None: - return False - - # The changes between context and generation requests are not straightforward. - if scheduled_requests.num_context_requests > 0: - return False - - # Check if the request_ids changes - request_ids = [ - request.py_request_id - for request in scheduled_requests.generation_requests - ] - if self.previous_request_ids != request_ids: - return False - - has_current_device_draft = next_draft_tokens_device is not None - return (self.is_draft_model and self.model_is_wrapped) or ( - has_current_device_draft and self.has_previous_device_draft) - - @nvtx_range("_apply_incremental_update") - def _apply_incremental_update( - self, - scheduled_requests: ScheduledRequests, - kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], - attn_metadata: AttentionMetadata, - spec_metadata: Optional[SpecMetadata] = None, - new_tensors_device: Optional[SampleStateTensors] = None, - cache_indirection_buffer: Optional[torch.Tensor] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None, - req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None, - resource_manager: Optional[ResourceManager] = None): - """ - Apply incremental update for the given scheduled requests and new tensors device. - """ - - if self.is_draft_model: - return self._apply_incremental_update_draft( - scheduled_requests, kv_cache_manager, attn_metadata, - spec_metadata, new_tensors_device, num_accepted_tokens_device) - else: - return self._apply_incremental_update_target( - scheduled_requests, kv_cache_manager, attn_metadata, - spec_metadata, new_tensors_device, num_accepted_tokens_device, - resource_manager) - - @nvtx_range("_prepare_incremental_update_metadata") - def _prepare_incremental_update_metadata( - self, - scheduled_requests: ScheduledRequests, - kv_cache_manager: KVCacheManager, - attn_metadata: AttentionMetadata, - spec_metadata: Optional[SpecMetadata], - prompt_lengths: List[int], - num_cached_tokens_per_seq: List[int], - total_num_tokens: int, - num_generation_tokens: int, - request_accepted_path: Optional[Dict[int, Any]] = None, - num_extend_ctx_requests: int = 0): - """ - Common metadata preparation logic for incremental updates. - """ - - enable_spec_decode = self.enable_spec_decode - enable_attention_dp = self.enable_attention_dp - spec_config = self.spec_config if enable_spec_decode else None - - # Set up attention metadata - batch simple assignments - attn_metadata.beam_width = 1 - attn_metadata.prompt_lens = prompt_lengths - attn_metadata.num_contexts = num_extend_ctx_requests if ( - enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( - self.attn_backend) and spec_config.is_linear_tree) else 0 - attn_metadata.num_chunked_ctx_requests = attn_metadata.num_contexts - - # Create KV cache params and prepare metadata - attn_metadata.kv_cache_params = KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - num_extra_kv_tokens=get_num_extra_kv_tokens(spec_config)) - attn_metadata.kv_cache_manager = kv_cache_manager - attn_metadata.prepare() - - # Get LoRA parameters - lora_params = self._get_lora_params_from_requests( - scheduled_requests, attn_metadata) - - # Handle padding for piecewise CUDA graphs - attn_metadata.padded_num_tokens = None - - # Handle attention DP - if enable_attention_dp: - attn_metadata.all_rank_num_tokens = self._get_all_rank_num_tokens( - attn_metadata) - - # Prepare speculative metadata - if spec_metadata is not None: - # Set request_accepted_path if Eagle3 - if isinstance(spec_metadata, Eagle3SpecMetadata): - spec_metadata.request_accepted_path = request_accepted_path - - spec_metadata.num_tokens = total_num_tokens - spec_metadata.prepare() - - # Handle distributed spec metadata - if enable_attention_dp: - sequence_lengths = spec_metadata.seq_lens - all_rank_num_tokens = self.dist.tp_cp_allgather([ - spec_metadata.num_tokens, - len(sequence_lengths), attn_metadata.num_generations - ]) - self._set_spec_metadata_all_rank_num_tokens( - spec_metadata, [item[0] for item in all_rank_num_tokens], - [item[1] for item in all_rank_num_tokens], - [item[2] for item in all_rank_num_tokens]) - - # Set iteration states - batch dictionary updates - self.iter_states.update({ - 'num_ctx_requests': - 0, - 'num_ctx_tokens': - 0, - 'num_generation_tokens': - num_generation_tokens, - 'cached_kv_tokens': - sum(num_cached_tokens_per_seq), - }) - - return lora_params - - def _update_draft_input_tensors(self, - num_accepted_tokens_device: torch.Tensor, - new_tokens_device: torch.Tensor, - total_num_tokens: int, - num_first_draft_requests: int): - """ - This function performs in-place updates on position_ids, num_accepted_draft_tokens, - gather_ids, and input_ids tensors for speculative decoding draft operations. - """ - # Prepare position_ids - idx_accepted_tokens = self.idx_accepted_tokens_cache[:total_num_tokens] - self.position_ids_cuda[:total_num_tokens].add_( - self.num_accepted_draft_tokens_cuda[idx_accepted_tokens] + 1) - - # Prepare gather_ids - old_accepted_tokens = self.num_accepted_draft_tokens_cuda[: - num_first_draft_requests].clone( - ) - self.num_accepted_draft_tokens_cuda[:num_first_draft_requests].copy_( - num_accepted_tokens_device[ - self.draft_seq_slots_buffer_cuda[:num_first_draft_requests]], - non_blocking=True) - self.gather_ids_cuda[:num_first_draft_requests].add_( - self.num_accepted_draft_tokens_cuda[:num_first_draft_requests] - - old_accepted_tokens) - - # Prepare token_positions for input_ids update - tokens_per_first_draft = self.original_max_draft_len + 1 - token_positions = self.draft_token_positions_cache[:tokens_per_first_draft].repeat( - num_first_draft_requests) - - # Prepare input_ids - self.input_ids_cuda[ - self. - draft_first_draft_indices_cuda[:total_num_tokens]] = new_tokens_device[ - token_positions, - self.draft_first_draft_seq_slots_cuda[:total_num_tokens], 0] - - def _apply_incremental_update_draft( - self, - scheduled_requests: ScheduledRequests, - kv_cache_manager: KVCacheManager, - attn_metadata: AttentionMetadata, - spec_metadata: Optional[SpecMetadata] = None, - new_tensors_device: Optional[SampleStateTensors] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None): - new_tokens_device = new_tensors_device.new_tokens - - num_generation_tokens = scheduled_requests.num_generation_requests - num_gen_requests = 0 - - tokens_per_first_draft = self.original_max_draft_len + 1 - prompt_lengths = [] # per sequence - num_cached_tokens_per_seq = [] # per sequence - - for request in scheduled_requests.generation_requests: - if request.is_dummy: - num_gen_requests += 1 - past_seen_token_num = request.max_beam_num_tokens - 1 - request.cached_tokens = past_seen_token_num - else: - assert request.py_is_first_draft - past_seen_token_num = request.max_beam_num_tokens - tokens_per_first_draft - - num_cached_tokens_per_seq.append(past_seen_token_num) - prompt_lengths.append(request.py_prompt_len) - request.py_batch_idx = request.py_seq_slot - - num_first_draft_requests = num_generation_tokens - num_gen_requests - total_num_tokens = num_first_draft_requests * tokens_per_first_draft - - self._update_draft_input_tensors( - num_accepted_tokens_device=num_accepted_tokens_device, - new_tokens_device=new_tokens_device, - total_num_tokens=total_num_tokens, - num_first_draft_requests=num_first_draft_requests) - - # Prepare spec_metadata - if spec_metadata is not None: - spec_metadata.draft_tokens = [] - spec_metadata.gather_ids = self.gather_ids_cuda[: - num_generation_tokens] - spec_metadata.num_accepted_draft_tokens = self.num_accepted_draft_tokens_cuda[: - num_generation_tokens] - - # Use common metadata preparation logic - virtual_num_tokens = total_num_tokens + num_gen_requests - lora_params = self._prepare_incremental_update_metadata( - scheduled_requests=scheduled_requests, - kv_cache_manager=kv_cache_manager, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - prompt_lengths=prompt_lengths, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - total_num_tokens=virtual_num_tokens, - num_generation_tokens=num_generation_tokens, - num_extend_ctx_requests=0) - - # No padding because there are only generation requests. - attn_metadata.padded_num_tokens = None - if self.enable_attention_dp: - attn_metadata.all_rank_num_tokens = self._get_all_rank_num_tokens( - attn_metadata) - - final_position_ids = self.position_ids_cuda[: - virtual_num_tokens].unsqueeze( - 0) - - inputs = { - 'attn_metadata': attn_metadata, - 'input_ids': self.input_ids_cuda[:virtual_num_tokens], - 'position_ids': final_position_ids, - 'inputs_embeds': None, - "multimodal_params": [], - } - - if bool(lora_params): - inputs['lora_params'] = lora_params - - if spec_metadata is not None: - inputs['spec_metadata'] = spec_metadata - - return inputs, self.gather_ids_cuda[:num_generation_tokens] - - def _update_target_input_tensors( - self, num_accepted_tokens_device: torch.Tensor, - new_tokens_device: torch.Tensor, - next_draft_tokens_device: torch.Tensor, - new_tokens_lens_device: torch.Tensor, previous_slots: torch.Tensor, - total_num_tokens: int, num_extend_reqeust_wo_dummy: int, - num_tokens_per_extend_request: int, - previous_batch_draft_tokens: int): - """ - This function performs in-place updates on position_ids, num_accepted_draft_tokens, - input_ids, draft_tokens, and offset tensors for speculative decoding extend context operations. - """ - - # Prepare position_ids - idx_accepted_tokens = self.idx_accepted_tokens_cache[:total_num_tokens] - self.position_ids_cuda[:total_num_tokens].add_( - self.num_accepted_draft_tokens_cuda[idx_accepted_tokens] + 1) - - self.num_accepted_draft_tokens_cuda[:num_extend_reqeust_wo_dummy].copy_( - num_accepted_tokens_device[:num_extend_reqeust_wo_dummy], - non_blocking=True) - - # Initialize offset tensors to zeros - self.previous_pos_id_offsets_cuda.mul_(0) - self.previous_kv_lens_offsets_cuda.mul_(0) - - # Prepare input_ids - # CRITICAL: Only extract the needed tokens based on num_tokens_per_extend_request - # new_tokens_device shape: [batch, 1 + max_draft_len] - # We need: [previous_batch, num_tokens_per_extend_request] - new_tokens = new_tokens_device.transpose( - 0, 1)[previous_slots, :num_tokens_per_extend_request].flatten() - self.input_ids_cuda[:total_num_tokens].copy_(new_tokens, - non_blocking=True) - - # Prepare draft tokens - num_draft_tokens_per_extend_request = num_tokens_per_extend_request - 1 - self.draft_tokens_cuda[:previous_batch_draft_tokens].copy_( - next_draft_tokens_device[ - previous_slots, :num_draft_tokens_per_extend_request].flatten(), - non_blocking=True) - - # Compute kv_len_offsets and update offset tensors - previous_pos_indices = previous_slots.repeat_interleave( - num_tokens_per_extend_request) - self.previous_pos_indices_cuda[:total_num_tokens].copy_( - previous_pos_indices, non_blocking=True) - kv_len_offsets_device = new_tokens_lens_device - num_tokens_per_extend_request - self.previous_pos_id_offsets_cuda[:num_extend_reqeust_wo_dummy * - num_tokens_per_extend_request].copy_( - new_tokens_lens_device[ - self. - previous_pos_indices_cuda[: - total_num_tokens]], - non_blocking=True) - self.previous_kv_lens_offsets_cuda[:num_extend_reqeust_wo_dummy].copy_( - kv_len_offsets_device[previous_slots], non_blocking=True) - - def _apply_incremental_update_target( - self, - scheduled_requests: ScheduledRequests, - kv_cache_manager: KVCacheManager, - attn_metadata: AttentionMetadata, - spec_metadata: Optional[SpecMetadata] = None, - new_tensors_device: Optional[SampleStateTensors] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None, - resource_manager: Optional[ResourceManager] = None): - # Extract tensors from new_tensors_device - new_tokens_device = new_tensors_device.new_tokens # [batch, 1 + draft_len] - new_tokens_lens_device = new_tensors_device.new_tokens_lens # [batch] - next_draft_tokens_device = new_tensors_device.next_draft_tokens # [batch, draft_len] - - # Pre-compute constants - extend_requests = scheduled_requests.generation_requests - num_extend_requests = len(extend_requests) - spec_config = self.spec_config - num_tokens_per_extend_request = self.get_runtime_tokens_per_gen_step( - self.runtime_draft_len) - runtime_draft_token_buffer_width = num_tokens_per_extend_request - 1 - - prompt_lengths = torch.empty(num_extend_requests, - dtype=torch.int, - device='cpu', - pin_memory=prefer_pinned()) - num_cached_tokens_per_seq = torch.empty(num_extend_requests, - dtype=torch.int, - device='cpu', - pin_memory=prefer_pinned()) - previous_batch_indices = torch.empty(num_extend_requests, - dtype=torch.int, - device='cpu', - pin_memory=prefer_pinned()) - - request_accepted_path = {} - num_extend_dummy_requests = 0 - num_previous_batch = 0 - - use_extend_ctx = (self.enable_spec_decode - and spec_config.spec_dec_mode.extend_ctx( - self.attn_backend) and spec_config.is_linear_tree) - - for idx, request in enumerate(extend_requests): - request_accepted_path[request.py_request_id] = \ - request.py_num_accepted_draft_tokens_indices - - base_past_seen = request.max_beam_num_tokens - 1 - - if use_extend_ctx: - # We're treating the prompt lengths as context requests here, so - # the prompt lens should not include the cached tokens. - prompt_lengths[idx] = num_tokens_per_extend_request - else: - prompt_lengths[idx] = request.py_prompt_len - - # Physical KV length for the kernels: subtract the tokens a - # KV-cache compression manager evicted (tracked on the request, - # 0 without compression). Position ids and the cached_tokens stat - # keep the logical count. - if request.is_dummy: - num_cached_tokens_per_seq[idx] = base_past_seen - request.cached_tokens = base_past_seen - num_extend_dummy_requests += 1 - else: - # Request has previous tensor - previous_batch_indices[ - num_previous_batch] = request.py_batch_idx - num_previous_batch += 1 - - request.cached_tokens = (base_past_seen + - num_tokens_per_extend_request) - num_cached_tokens_per_seq[idx] = ( - base_past_seen + num_tokens_per_extend_request - - request.py_num_compressed_tokens) - - request.py_batch_idx = request.py_seq_slot - - num_extend_reqeust_wo_dummy = num_extend_requests - num_extend_dummy_requests - total_num_tokens = num_extend_reqeust_wo_dummy * num_tokens_per_extend_request - - previous_slots = self.previous_batch_indices_cuda[:num_previous_batch] - previous_slots.copy_(previous_batch_indices[:num_previous_batch], - non_blocking=True) - - prompt_lengths = prompt_lengths.tolist() - num_cached_tokens_per_seq = num_cached_tokens_per_seq.tolist() - - previous_batch_draft_tokens = (num_extend_reqeust_wo_dummy * - runtime_draft_token_buffer_width) - - self._update_target_input_tensors( - num_accepted_tokens_device=num_accepted_tokens_device, - new_tokens_device=new_tokens_device, - next_draft_tokens_device=next_draft_tokens_device, - new_tokens_lens_device=new_tokens_lens_device, - previous_slots=previous_slots, - total_num_tokens=total_num_tokens, - num_extend_reqeust_wo_dummy=num_extend_reqeust_wo_dummy, - num_tokens_per_extend_request=num_tokens_per_extend_request, - previous_batch_draft_tokens=previous_batch_draft_tokens) - - # Prepare spec_metadata - num_generation_tokens = num_extend_requests * num_tokens_per_extend_request - if spec_metadata is not None: - total_draft_lens = self.max_total_draft_tokens * num_extend_requests - spec_metadata.draft_tokens = self.draft_tokens_cuda[: - total_draft_lens] - spec_metadata.gather_ids = self.gather_ids_cuda[:total_num_tokens] - spec_metadata.num_accepted_draft_tokens = self.num_accepted_draft_tokens_cuda[: - num_extend_requests] - - # Determine if we're using extend_ctx mode for linear tree decoding - num_extend_ctx_requests = 0 - if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( - self.attn_backend) and spec_config.is_linear_tree: - num_extend_ctx_requests = num_extend_requests - - virtual_num_tokens = num_generation_tokens - lora_params = self._prepare_incremental_update_metadata( - scheduled_requests=scheduled_requests, - kv_cache_manager=kv_cache_manager, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - prompt_lengths=prompt_lengths, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - total_num_tokens=virtual_num_tokens, - num_generation_tokens=num_generation_tokens, - request_accepted_path=request_accepted_path, - num_extend_ctx_requests=num_extend_ctx_requests) - - # No padding because there are only generation requests. - attn_metadata.padded_num_tokens = None - if self.enable_attention_dp: - attn_metadata.all_rank_num_tokens = self._get_all_rank_num_tokens( - attn_metadata) - - final_position_ids = self.position_ids_cuda[: - virtual_num_tokens].unsqueeze( - 0) - - # Prepare inputs - # Note: multimodal_params is always empty for incremental updates because: - # - This function only processes generation requests (no context requests) - # - Multimodal data (images/videos) is only needed during context/prefill phase - inputs = { - 'attn_metadata': attn_metadata, - 'input_ids': self.input_ids_cuda[:virtual_num_tokens], - 'position_ids': final_position_ids, - 'inputs_embeds': None, - "multimodal_params": [], - 'resource_manager': resource_manager, - } - - if bool(lora_params): - inputs['lora_params'] = lora_params - - if spec_metadata is not None: - inputs['spec_metadata'] = spec_metadata - - return inputs, self.gather_ids_cuda[:num_generation_tokens] - def _can_use_steady_gen_fast_prepare( self, scheduled_requests: ScheduledRequests, new_tokens_device: Optional[torch.Tensor], @@ -3819,7 +3197,7 @@ def _prepare_tp_inputs( new_tokens_device = new_tensors_device.new_tokens # When using overlap scheduler with speculative decoding, the target model's inputs would be SampleStateTensorsSpec. if isinstance(new_tensors_device, SampleStateTensorsSpec): - assert self.enable_spec_decode and not self.is_draft_model + assert self.enable_spec_decode new_tokens_lens_device = new_tensors_device.new_tokens_lens # [batch] next_draft_tokens_device = new_tensors_device.next_draft_tokens # [batch, draft_len] @@ -3830,19 +3208,6 @@ def _prepare_tp_inputs( new_tokens=new_tokens_device, runtime_draft_len=self.runtime_draft_len) - if self._can_use_incremental_update(scheduled_requests, - new_tokens_device, - next_draft_tokens_device): - # Spec engines never record the steady-gen cache, but invalidate - # defensively so the two fast paths can never interleave if the - # gates ever evolve. - self._steady_gen_cache = None - return self._apply_incremental_update( - scheduled_requests, kv_cache_manager, attn_metadata, - spec_metadata, new_tensors_device, cache_indirection_buffer, - num_accepted_tokens_device, req_id_to_old_request, - resource_manager) - if self._can_use_steady_gen_fast_prepare(scheduled_requests, new_tokens_device, next_draft_tokens_device, @@ -3888,22 +3253,6 @@ def _prepare_tp_inputs( cross_encoder_seq_lens: List[int] = [ ] # new encoder K/V tokens per decoder sequence cross_encoder_cached_tokens_per_seq: List[int] = [] - # if using tree decoding, we need to store the request type and accepted path for each request, - # which will be used to update the hidden_states_read_indices. - request_accepted_path = {} # per request - - # Variables for updating the inputs of draft model - # Base values for gather_ids computation - first_draft_base_gather_ids = [] - # seq_slots to index into num_accepted_tokens_device - first_draft_seq_slots = [] - # Indices in the num_accepted_draft_tokens list - first_draft_request_indices = [] - - # (start_idx, end_idx, seq_slot) for context requests - context_input_ids_positions = [] - # (start_idx, end_idx, seq_slot) for first_draft requests - first_draft_input_ids_positions = [] def append_cross_attention_state(request: LlmRequest, project_encoder_output: bool, @@ -3951,24 +3300,11 @@ def append_cross_attention_state(request: LlmRequest, # that rewrite token IDs in place write into the request's own span # rather than assuming a contiguous multimodal prefix. context_start_idx = len(input_ids) - # Track position for updating the inputs of draft model - if self.is_draft_model and num_accepted_tokens_device is not None: - input_ids.extend(prompt_tokens) - end_idx = len(input_ids) - slot_idx = req_id_to_old_request[ - request.py_request_id].py_seq_slot - context_input_ids_positions.append( - (context_start_idx, end_idx - 1, - slot_idx)) # end_idx-1 is the last token position - else: - input_ids.extend(prompt_tokens) + input_ids.extend(prompt_tokens) gather_ids.append(len(input_ids) - 1) sequence_lengths.append(len(prompt_tokens)) num_accepted_draft_tokens.append(len(prompt_tokens) - 1) - request_accepted_path[ - request. - py_request_id] = request.py_num_accepted_draft_tokens_indices prompt_lengths.append(len(prompt_tokens)) past_seen_token_num = begin_compute num_cached_tokens_per_seq.append(past_seen_token_num - @@ -4128,9 +3464,6 @@ def append_cross_attention_state(request: LlmRequest, padding_gen_slots.append(request.py_seq_slot) request.py_needs_onehot_draft_probs = False # consume once request_ids.append(request.py_request_id) - request_accepted_path[ - request. - py_request_id] = request.py_num_accepted_draft_tokens_indices # the request has no previous tensor: # (1) next_draft_tokens_device is None, which means overlap scheduler is disabled; or # (2) a dummy request; or @@ -4218,41 +3551,14 @@ def append_cross_attention_state(request: LlmRequest, position_ids.extend( range(begin_compute, begin_compute + len(prompt_tokens))) - # Track position for updating the inputs of draft model - if self.is_draft_model and num_accepted_tokens_device is not None: - start_idx = len(input_ids) - input_ids.extend(prompt_tokens) - end_idx = len(input_ids) - # For first_draft, we need to replace the last original_max_draft_len+1 tokens - slot_idx = req_id_to_old_request[ - request.py_request_id].py_seq_slot - first_draft_input_ids_positions.append( - (start_idx, end_idx, slot_idx)) - - # Store info for GPU computation of gather_ids and num_accepted_draft_tokens - base_gather_id = len( - input_ids) - 1 - self.original_max_draft_len - # Placeholder, will be corrected on GPU - gather_ids.append(base_gather_id) - first_draft_base_gather_ids.append(base_gather_id) - first_draft_seq_slots.append(slot_idx) - first_draft_request_indices.append( - len(num_accepted_draft_tokens)) - - # Placeholder, will be corrected on GPU - num_accepted_draft_tokens.append(0) - else: - input_ids.extend(prompt_tokens) - gather_ids.append( - len(input_ids) - 1 - (self.original_max_draft_len - - request.py_num_accepted_draft_tokens)) - num_accepted_draft_tokens.append( - request.py_num_accepted_draft_tokens) + input_ids.extend(prompt_tokens) + gather_ids.append( + len(input_ids) - 1 - (self.original_max_draft_len - + request.py_num_accepted_draft_tokens)) + num_accepted_draft_tokens.append( + request.py_num_accepted_draft_tokens) sequence_lengths.append(1 + self.original_max_draft_len) - request_accepted_path[ - request. - py_request_id] = request.py_num_accepted_draft_tokens_indices prompt_lengths.append(request.py_prompt_len) past_seen_token_num = begin_compute num_cached_tokens_per_seq.append(past_seen_token_num - @@ -4294,17 +3600,7 @@ def append_cross_attention_state(request: LlmRequest, # can be aligned to the correct positions. if not request.is_cuda_graph_dummy: for beam in range(beam_width): - # Track position for GPU update (draft model only) - if self.is_draft_model and num_accepted_tokens_device is not None: - start_idx = len(input_ids) - input_ids.append(request.get_last_tokens(beam)) - end_idx = len(input_ids) - slot_idx = req_id_to_old_request[ - request.py_request_id].py_seq_slot - first_draft_input_ids_positions.append( - (start_idx, end_idx, slot_idx)) - else: - input_ids.append(request.get_last_tokens(beam)) + input_ids.append(request.get_last_tokens(beam)) past_seen_token_num = request.max_beam_num_tokens - 1 else: # the request has previous tensor @@ -4431,80 +3727,6 @@ def previous_seq_slots_device(): pin_memory=prefer_pinned()) self.input_ids_cuda[:num_tokens].copy_(input_ids, non_blocking=True) - # Update input_ids_cuda with new tokens from new_tensors_device (draft model only) - if self.is_draft_model and num_accepted_tokens_device is not None: - # For context requests: replace the last token with new_tensors_device[0, seq_slot, 0] - if len(context_input_ids_positions) > 0: - # Build tensors on CPU first, then copy to GPU to avoid implicit sync - num_ctx_positions = len(context_input_ids_positions) - ctx_token_indices_cpu = torch.tensor( - [ - last_token_idx for _, last_token_idx, _ in - context_input_ids_positions - ], - dtype=torch.long, - pin_memory=prefer_pinned()) - ctx_seq_slots_cpu = torch.tensor([ - seq_slot - for _, _, seq_slot in context_input_ids_positions - ], - dtype=torch.long, - pin_memory=prefer_pinned()) - # Copy to pre-allocated GPU buffers - self.draft_ctx_token_indices_cuda[:num_ctx_positions].copy_( - ctx_token_indices_cpu, non_blocking=True) - self.draft_ctx_seq_slots_cuda[:num_ctx_positions].copy_( - ctx_seq_slots_cpu, non_blocking=True) - self.input_ids_cuda[ - self. - draft_ctx_token_indices_cuda[:num_ctx_positions]] = new_tensors_device.new_tokens[ - 0, - self.draft_ctx_seq_slots_cuda[:num_ctx_positions], - 0] - - # For first_draft requests: replace the last (original_max_draft_len+1) tokens - # with new_tensors_device[:, seq_slot, 0] - if len(first_draft_input_ids_positions) > 0: - # All first_draft requests have same token length (original_max_draft_len + 1) - # Build index tensors on CPU first, then copy to GPU to avoid implicit sync - num_requests = len(first_draft_input_ids_positions) - tokens_per_request = first_draft_input_ids_positions[0][ - 1] - first_draft_input_ids_positions[0][0] - - # Create flat index array for all tokens to update on CPU - all_indices = [] - all_seq_slots = [] - for start_idx, end_idx, seq_slot in first_draft_input_ids_positions: - all_indices.extend(range(start_idx, end_idx)) - all_seq_slots.extend([seq_slot] * (end_idx - start_idx)) - - # Create CPU tensors with pinned memory - total_tokens = len(all_indices) - idx_tensor_cpu = torch.tensor(all_indices, - dtype=torch.long, - pin_memory=prefer_pinned()) - seq_slots_tensor_cpu = torch.tensor( - all_seq_slots, - dtype=torch.long, - pin_memory=prefer_pinned()) - - # Copy to pre-allocated GPU buffers - self.draft_first_draft_indices_cuda[:total_tokens].copy_( - idx_tensor_cpu, non_blocking=True) - self.draft_first_draft_seq_slots_cuda[:total_tokens].copy_( - seq_slots_tensor_cpu, non_blocking=True) - - # Create token position indices (repeating 0..tokens_per_request for each request) - token_positions = torch.arange( - tokens_per_request, dtype=torch.long, - device='cuda').repeat(num_requests) - - self.input_ids_cuda[ - self. - draft_first_draft_indices_cuda[:total_tokens]] = new_tensors_device.new_tokens[ - token_positions, self. - draft_first_draft_seq_slots_cuda[:total_tokens], 0] - if num_draft_tokens > 0: draft_tokens = torch.tensor(draft_tokens, dtype=torch.int, @@ -4518,34 +3740,6 @@ def previous_seq_slots_device(): self.num_accepted_draft_tokens_cuda[:len( num_accepted_draft_tokens)].copy_(num_accepted_draft_tokens, non_blocking=True) - - # Update num_accepted_draft_tokens_cuda for first_draft_requests directly from num_accepted_tokens_device (draft model only) - if self.is_draft_model and len(first_draft_seq_slots) > 0: - # Build tensors on CPU first, then copy to GPU to avoid implicit sync - num_first_draft = len(first_draft_seq_slots) - first_draft_seq_slots_cpu = torch.tensor( - first_draft_seq_slots, - dtype=torch.int, - pin_memory=prefer_pinned()) - first_draft_indices_cpu = torch.tensor( - first_draft_request_indices, - dtype=torch.int, - pin_memory=prefer_pinned()) - - # Copy to pre-allocated GPU buffers - self.draft_seq_slots_buffer_cuda[:num_first_draft].copy_( - first_draft_seq_slots_cpu, non_blocking=True) - self.draft_request_indices_buffer_cuda[:num_first_draft].copy_( - first_draft_indices_cpu, non_blocking=True) - - # Extract accepted tokens for first_draft requests from device tensor - accepted_tokens = num_accepted_tokens_device[ - self.draft_seq_slots_buffer_cuda[:num_first_draft]] - # Update the correct positions in num_accepted_draft_tokens_cuda - self.num_accepted_draft_tokens_cuda[ - self. - draft_request_indices_buffer_cuda[: - num_first_draft]] = accepted_tokens if next_draft_tokens_device is not None: # Initialize these two values to zeros self.previous_pos_id_offsets_cuda *= 0 @@ -4694,35 +3888,6 @@ def previous_seq_slots_device(): gather_ids, dtype=torch.int, pin_memory=prefer_pinned()), non_blocking=True) - # Update gather_ids for first_draft_requests on GPU (draft model only) - if self.is_draft_model and len(first_draft_seq_slots) > 0: - # Build tensors on CPU first, then copy to GPU to avoid implicit sync - num_first_draft = len(first_draft_seq_slots) - first_draft_seq_slots_cpu = torch.tensor( - first_draft_seq_slots, - dtype=torch.int, - pin_memory=prefer_pinned()) - first_draft_indices_cpu = torch.tensor( - first_draft_request_indices, - dtype=torch.int, - pin_memory=prefer_pinned()) - - # Copy to pre-allocated GPU buffers - self.draft_seq_slots_buffer_cuda[:num_first_draft].copy_( - first_draft_seq_slots_cpu, non_blocking=True) - self.draft_request_indices_buffer_cuda[:num_first_draft].copy_( - first_draft_indices_cpu, non_blocking=True) - - # Extract accepted tokens for first_draft requests from device tensor - accepted_tokens = num_accepted_tokens_device[ - self.draft_seq_slots_buffer_cuda[:num_first_draft]] - # Update gather_ids: gather_id = base_gather_id + num_accepted_tokens - # (since gather_id = len(input_ids) - 1 - (max_draft_len - num_accepted)) - self.gather_ids_cuda[ - self. - draft_request_indices_buffer_cuda[: - num_first_draft]] += accepted_tokens - if self.mapping.has_cp_helix(): attn_metadata.update_helix_param( helix_position_offsets=helix_position_offsets, @@ -4869,8 +4034,6 @@ def previous_seq_slots_device(): spec_metadata.seq_lens = sequence_lengths spec_metadata.num_accepted_draft_tokens = self.num_accepted_draft_tokens_cuda[:len( num_accepted_draft_tokens)] - if isinstance(spec_metadata, Eagle3SpecMetadata): - spec_metadata.request_accepted_path = request_accepted_path # No-op for non 1-model spec_metadata.populate_sampling_params_for_one_model( scheduled_requests.all_requests()) @@ -4926,8 +4089,8 @@ def previous_seq_slots_device(): # guards stable), with all three axes equal to the scalar # positions, so the fast path advances that buffer in place and # returns the same layout (see _apply_steady_gen_fast_prepare). - if (self.spec_config is None and not self.is_draft_model - and spec_metadata is None and new_tokens_device is not None + if (self.spec_config is None and spec_metadata is None + and new_tokens_device is not None and self.guided_decoder is None and not self.enable_attention_dp and not mrope_position_ids and not mrope_delta_write_seq_slots @@ -6036,8 +5199,7 @@ def forward(self, is None) # attn_metadata now depends on spec_metadata since it determines the shape/content of spec_dec parameter Tensors is_spec_dec_mode = spec_metadata.spec_dec_mode.attention_need_spec_dec_mode( - spec_resource_manager, self.is_draft_model, self.attn_backend, - self.model_is_wrapped) + self.attn_backend) # Propagate runtime_draft_len (already set on self by py_executor) # to spec_metadata so downstream code (eagle3, interface, trtllm) can read it. spec_metadata.runtime_draft_len = self.runtime_draft_len @@ -6057,8 +5219,7 @@ def forward(self, # Fill slot-ID buffer for update_spec_dec_param if (spec_tree_manager is not None - and spec_tree_manager.use_dynamic_tree - and not self.is_draft_model): + and spec_tree_manager.use_dynamic_tree): spec_tree_manager.slot_storage.fill_all_slot_ids( scheduled_requests.context_requests, scheduled_requests.generation_requests, @@ -6071,7 +5232,6 @@ def forward(self, is_spec_dec_dynamic_tree=spec_metadata.is_spec_dec_dynamic_tree, max_draft_len=sd_max_draft_len, max_total_draft_tokens=sd_max_total, - model_is_wrapped=self.model_is_wrapped, spec_metadata=spec_metadata, spec_tree_manager=spec_tree_manager, num_contexts=scheduled_requests.num_context_requests) @@ -6142,8 +5302,7 @@ def forward(self, # Fill slot-ID buffer for scatter inside draft loop if (self.enable_spec_decode and spec_tree_manager is not None - and spec_tree_manager.use_dynamic_tree - and not self.is_draft_model): + and spec_tree_manager.use_dynamic_tree): spec_tree_manager.slot_storage.fill_all_slot_ids( padded_requests.context_requests, padded_requests.generation_requests, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 7920c8942e71..e42694bf3db4 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -21,8 +21,7 @@ ContextChunkingPolicy, ExecutorMemoryType, GuidedDecodingConfig, KvCacheConfig, - LoadFormat, SpeculativeConfig, - TorchLlmArgs) + SpeculativeConfig, TorchLlmArgs) from tensorrt_llm.llmapi.tokenizer import (TokenizerBase, _llguidance_tokenizer_info, _xgrammar_tokenizer_info) @@ -32,7 +31,6 @@ from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator from ..attention_backend.interface import AttentionRuntimeFeatures -from ..attention_backend.trtllm import TrtllmAttention from ..distributed import Distributed from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter, get_spec_resource_manager) @@ -497,10 +495,8 @@ def create_py_executor( cache_transceiver_config = llm_args.cache_transceiver_config - has_draft_model_engine = False has_spec_drafter = False if spec_config is not None: - has_draft_model_engine = spec_config.spec_dec_mode.has_draft_model() has_spec_drafter = spec_config.spec_dec_mode.has_spec_drafter() # Eagle3DecodingConfig._max_batch_size is internally managed: the @@ -521,7 +517,7 @@ def create_py_executor( attn_runtime_features = AttentionRuntimeFeatures( chunked_prefill=enable_chunked_context, cache_reuse=kv_cache_config.enable_block_reuse, - has_speculative_draft_tokens=has_draft_model_engine or has_spec_drafter, + has_speculative_draft_tokens=has_spec_drafter, chunk_size=max_num_tokens, ) logger.info("ATTENTION RUNTIME FEATURES: ", attn_runtime_features) @@ -581,82 +577,9 @@ def allocation_scope(current_stage: ExecutorMemoryType): dist=dist) model_engine.model = calibrator.maybe_wrap_model(model_engine.model) - if has_draft_model_engine: - with allocation_scope(ExecutorMemoryType.MODEL_ENGINE_DRAFT): - draft_spec_config = copy.copy(spec_config) - - use_chain_drafter = ( - guided_decoding_config is None - and draft_spec_config._allow_chain_drafter - and draft_spec_config._allow_greedy_draft_tokens - and llm_args.attn_backend == "TRTLLM" - and draft_spec_config.draft_len_schedule is None) - - logger.debug(f"USE CHAIN DRAFTER: {use_chain_drafter}") - if use_chain_drafter: - - def drafting_loop_wrapper(model): - from tensorrt_llm._torch.speculative.drafting_loops import ( - LinearDraftingLoopWrapper, - StaticTreeDraftingLoopWrapper) - from tensorrt_llm.llmapi import EagleDecodingConfig - - static_tree_drafter = isinstance( - draft_spec_config, EagleDecodingConfig - ) and draft_spec_config.eagle_choices is not None - - if static_tree_drafter: - return StaticTreeDraftingLoopWrapper( - spec_config.max_draft_len, - spec_config.tokens_per_gen_step - 1, max_batch_size, - model) - else: - return LinearDraftingLoopWrapper( - spec_config.max_draft_len, - spec_config.tokens_per_gen_step - 1, model) - else: - drafting_loop_wrapper = None - - draft_llm_args = copy.copy(llm_args) - if spec_config.load_format == "dummy": - draft_llm_args.load_format = LoadFormat.DUMMY - - model_weights_memory_tag = None - model_weights_restore_mode = None - if enable_sleep: - model_weights_memory_tag = ExecutorMemoryType.MODEL_WEIGHTS_DRAFT - model_weights_restore_mode = sleep_config.restore_modes[ - ExecutorMemoryType.MODEL_WEIGHTS_DRAFT] - - draft_model_engine = PyTorchModelEngine( - model_path=spec_config.speculative_model, - llm_args=draft_llm_args, - mapping=mapping, - attn_runtime_features=attn_runtime_features, - dist=dist, - spec_config=draft_spec_config, - is_draft_model=True, - drafting_loop_wrapper=drafting_loop_wrapper, - model_weights_memory_tag=model_weights_memory_tag, - model_weights_restore_mode=model_weights_restore_mode, - ) - # For DeepseekV3 MTP, we need to set the num_hidden_layers to 1 for the draft model - if spec_config.spec_dec_mode.is_mtp_eagle(): - draft_model_engine.model.model_config.pretrained_config.num_hidden_layers = 1 - draft_model_engine.load_weights_from_target_model( - model_engine.model) - else: - draft_model_engine = None - - # TODO: Overlap scheduler is not supported for below cases: - # 1. non-CDL is used - # 2. non-TrtllmAttention attention backend is used - if has_draft_model_engine and (not use_chain_drafter or not issubclass( - draft_model_engine.attn_backend, TrtllmAttention)): - logger.warning( - "Overlap scheduler is not supported for non-CDL or non-TrtllmAttention backend." - ) - llm_args.disable_overlap_scheduler = True + # Two-model speculative decoding (separate draft model engine) has been + # removed. All neural speculation runs inside the target model engine + # (one-engine); ngram / user-provided speculation use the drafter loop. # PyTorchModelEngine modifies these fields, update them model_engine_max_seq_len = model_engine.max_seq_len @@ -675,11 +598,6 @@ def drafting_loop_wrapper(model): model_engine_max_seq_len=model_engine_max_seq_len, ) - if has_draft_model_engine and not llm_args.disable_overlap_scheduler: - logger.warning( - "Overlap scheduler is enabled for two-model speculative decoding. Rejection sampling will fallback to greedy sampling." - ) - max_seq_len = model_engine_max_seq_len max_num_tokens = model_engine.max_num_tokens sparse_attention_config = model_engine.sparse_attention_config @@ -695,8 +613,7 @@ def drafting_loop_wrapper(model): "Disabling block reuse for MambaHybridCacheManager-based models when disagg + Python transceiver enabled" ) kv_cache_config.enable_block_reuse = False - _set_model_engines_cache_reuse([model_engine, draft_model_engine], - False) + _set_model_engines_cache_reuse([model_engine], False) if is_mla(config): if model_engine.model.model_config.enable_flash_mla: tokens_per_block = 64 @@ -722,8 +639,7 @@ def drafting_loop_wrapper(model): f"{_MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS_STR}, " f"disable enable_block_reuse for SM{sm_version}") kv_cache_config.enable_block_reuse = False - _set_model_engines_cache_reuse([model_engine, draft_model_engine], - False) + _set_model_engines_cache_reuse([model_engine], False) kv_cache_quant_algo = model_engine.model.model_config.quant_config.kv_cache_quant_algo if kv_cache_config.enable_block_reuse and not ( @@ -734,8 +650,7 @@ def drafting_loop_wrapper(model): f"disable enable_block_reuse for KV cache quant algorithm: {kv_cache_quant_algo}" ) kv_cache_config.enable_block_reuse = False - _set_model_engines_cache_reuse([model_engine, draft_model_engine], - False) + _set_model_engines_cache_reuse([model_engine], False) if (enable_chunked_context and sm_version not in _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS): logger.warning("Chunked Prefill for MLA can only be enabled on " @@ -743,8 +658,6 @@ def drafting_loop_wrapper(model): f"disable enable_chunked_context for SM{sm_version}") enable_chunked_context = False model_engine.attn_runtime_features.chunked_prefill = False - if draft_model_engine is not None: - draft_model_engine.attn_runtime_features.chunked_prefill = False # Set default value for cache_transceiver_config.max_tokens_in_buffer. # Placed after the FlashMLA tokens_per_block override and rounded up to a @@ -801,11 +714,11 @@ def drafting_loop_wrapper(model): if spec_config is None or spec_config.spec_dec_mode.support_guided_decoder( ): - # GuidedDecoder is applicable to non-speculative decoding and two-model speculative decoding. + # GuidedDecoder is applicable to non-speculative decoding and drafter-loop speculative decoding (ngram / user-provided). guided_decoder = GuidedDecoder(**kwargs) elif spec_config.spec_dec_mode.support_capturable_guided_decoder( ): - # CapturableGuidedDecoder is applicable to one-model speculative decoding. + # CapturableGuidedDecoder is applicable to one-engine speculative decoding. success = model_engine.set_guided_decoder( CapturableGuidedDecoder(**kwargs)) if not success: @@ -930,7 +843,6 @@ def drafting_loop_wrapper(model): kv_cache_creator = KvCacheCreator( model_engine=model_engine, - draft_model_engine=draft_model_engine, mapping=mapping, net_max_seq_len=net_max_seq_len, kv_connector_manager=kv_connector_manager, @@ -974,8 +886,7 @@ def drafting_loop_wrapper(model): # to provide a resource manager if required. with allocation_scope(ExecutorMemoryType.SPEC_RESOURCES): - spec_resource_manager = get_spec_resource_manager( - model_engine, draft_model_engine) + spec_resource_manager = get_spec_resource_manager(model_engine) if spec_resource_manager is not None: resources[ ResourceManagerType.SPEC_RESOURCE_MANAGER] = spec_resource_manager @@ -983,7 +894,6 @@ def drafting_loop_wrapper(model): # Drafter for speculative decoding with allocation_scope(ExecutorMemoryType.DRAFTER): drafter = get_spec_drafter(model_engine, - draft_model_engine, sampler, spec_resource_manager=spec_resource_manager, guided_decoder=guided_decoder) @@ -1041,13 +951,10 @@ def drafting_loop_wrapper(model): kv_cache_creator.teardown_managers(resources) # Release Phase-1 CUDA graph pools before final KV allocation to avoid overshoot. - for eng in [model_engine, draft_model_engine]: - if eng is None: - continue - if eng.attn_metadata is not None: - if llm_args.cuda_graph_config is not None: - eng._release_cuda_graphs() - eng.attn_metadata = None + if model_engine.attn_metadata is not None: + if llm_args.cuda_graph_config is not None: + model_engine._release_cuda_graphs() + model_engine.attn_metadata = None del py_executor # free before constructing new gc.collect() diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 041cb31f97be..392a1644480e 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -2517,11 +2517,10 @@ def __init__(self, args: Args): def _is_draft_batch(requests: list[LlmRequest]) -> bool: """Whether this batch belongs to the draft model. - Batches are homogeneous by construction: ModelDrafter builds all-draft - batches for its sample_async/update_requests calls on this shared - sampler, and PyExecutor's batches are all-target. The pending-steps - accounting relies on this to skip draft batches wholesale; assert it so - a mixed batch fails loudly instead of silently corrupting the counters. + Callers provide homogeneous all-draft or all-target batches. The + pending-steps accounting relies on this to skip draft batches + wholesale; assert it so a mixed batch fails loudly instead of silently + corrupting the counters. """ is_draft: bool = requests[0].py_is_draft assert all(r.py_is_draft == is_draft for r in requests), ( @@ -4095,11 +4094,8 @@ def sample_async( num_context_logits_prefix_sum: list[int], resource_manager: Optional[ResourceManager] = None, ) -> SampleStateTorch: - # NB: The sampler is either called directly by PyExecutor, for the target model, - # or by ModelDrafter.prepare_draft_tokens(), for the draft model. In the former - # case there are 1 + get_draft_token_length(request) tokens per request. In the - # latter case, there is always only 1 token per request because draft - # tokens are sampled one-by-one. + # NB: The sampler is called directly by PyExecutor for the target model. + # There are 1 + get_draft_token_length(request) tokens per request. self.setup_sampler_step(scheduled_requests) new_tokens = self.store.new_tokens diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 347f3aae5dc2..3715133de9e1 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -2,7 +2,7 @@ from .dflash import DFlashSpecMetadata, DFlashWorker from .draft_target import (DraftTargetOneModelSpecMetadata, DraftTargetOneModelWorker) -from .eagle3 import Eagle3SpecMetadata, MTPEagleWorker +from .eagle3 import MTPEagleWorker from .interface import (SpecMetadata, SpecWorkerBase, prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, @@ -30,7 +30,6 @@ "DFlashWorker", "DraftTargetOneModelSpecMetadata", "DraftTargetOneModelWorker", - "Eagle3SpecMetadata", "MTPEagleWorker", "MTPSampler", "MTPSpecMetadata", diff --git a/tensorrt_llm/_torch/speculative/drafter.py b/tensorrt_llm/_torch/speculative/drafter.py index 0396bbeb0dd7..794d20923125 100644 --- a/tensorrt_llm/_torch/speculative/drafter.py +++ b/tensorrt_llm/_torch/speculative/drafter.py @@ -64,11 +64,11 @@ def should_use_spec_decode(self, requests: List[LlmRequest], num_effective_requests = min(len(requests), max_batch_size, token_cap) return num_effective_requests <= self.max_concurrency - # Drafters that use TorchSampler (NGram, two-model) compute py_rewind_len - # from len(py_draft_tokens), which includes padding. They must set this - # to True so that extend_capacity_for_tokens is called after padding. - # One-model drafters (MTP / Eagle3 / SA) use SpecSamplerBase which - # computes rewind from runtime_draft_len, so padding is harmless. + # Drafters that use TorchSampler (NGram) compute py_rewind_len from + # len(py_draft_tokens), which includes padding. They must set this to True + # so that extend_capacity_for_tokens is called after padding. One-engine + # drafters (MTP / Eagle3 / SA) use SpecSamplerBase which computes rewind + # from runtime_draft_len, so padding is harmless. _needs_padding_kv_extension: bool = False @final diff --git a/tensorrt_llm/_torch/speculative/drafting_loops.py b/tensorrt_llm/_torch/speculative/drafting_loops.py deleted file mode 100644 index e2be40bfce5e..000000000000 --- a/tensorrt_llm/_torch/speculative/drafting_loops.py +++ /dev/null @@ -1,522 +0,0 @@ -""" -This module contains capturable drafting loops for speculative decoding. - -These are torch modules wrap another draft model. The wrapped module -is supposed to invoke the draft model autoregressively and invoke -a sampling algorithm to obtain draft tokens. By structuring the code -like this, we are able to avoid host overhead: the entire drafting process -for speculation can be launched as a single CUDA graph. -""" - -from abc import ABC, abstractmethod -from contextlib import contextmanager -from typing import Optional, final - -import torch - -from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata -from tensorrt_llm._torch.pyexecutor.sampler.sampling_utils import \ - greedy_search_sampling_batch -from tensorrt_llm._torch.speculative.eagle3 import Eagle3SpecMetadata -from tensorrt_llm._torch.speculative.interface import SpecMetadata -from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager - -# Enable capture_scalar_outputs to avoid graph breaks from Tensor.item() calls -torch._dynamo.config.capture_scalar_outputs = True - - -class BaseDraftingLoopWrapper(ABC, torch.nn.Module): - - @abstractmethod - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, - attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata, - **kwargs) -> dict[str, torch.Tensor]: - raise NotImplementedError - - @abstractmethod - def sample(self, - logits: torch.Tensor, - max_top_k: Optional[int] = None) -> torch.Tensor: - raise NotImplementedError - - @abstractmethod - def prepare_for_generation( - self, - attn_metadata: AttentionMetadata, - spec_metadata: SpecMetadata, - position_ids: torch.Tensor, - spec_tree_manager: Optional[SpecTreeManager] = None - ) -> torch.Tensor | None: - raise NotImplementedError - - @final - def load_weights_from_target_model(self, target_model) -> None: - loader = getattr(self.draft_model, "load_weights_from_target_model", - None) - if callable(loader): - self.draft_model.load_weights_from_target_model(target_model) - - -@contextmanager -def save_metadata_state(attn_metadata: AttentionMetadata, - spec_metadata: SpecMetadata) -> None: - attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") - batch_size = attn_metadata.num_seqs - # Do not use prepare_for_spec_dec for this special field. - # TRTLLM attention uses views of this tensor internally and prepare_for_spec_dec - # creates a copy. If you write to the copy, TRTLLM attention won't see the updates. - kv_lens = attn_metadata.kv_lens_cuda[:batch_size].clone() - - if attn_metadata.is_cuda_graph: - assert spec_metadata.is_cuda_graph - num_tokens = spec_metadata.num_tokens - if isinstance(spec_metadata, Eagle3SpecMetadata): - read_indices = spec_metadata.hidden_states_read_indices[: - batch_size].clone( - ) - write_indices = spec_metadata.hidden_states_write_indices[: - batch_size].clone( - ) - - try: - yield - finally: - attn_metadata.restore_from_spec_dec() - attn_metadata.kv_lens_cuda[:batch_size].copy_(kv_lens) - attn_metadata.on_update() - if attn_metadata.is_cuda_graph: - spec_metadata.num_tokens = num_tokens - if isinstance(spec_metadata, Eagle3SpecMetadata): - spec_metadata.hidden_states_read_indices[:batch_size].copy_( - read_indices) - spec_metadata.hidden_states_write_indices[:batch_size].copy_( - write_indices) - - # This restore has to happen even if the spec_metadata is not being used - # for CUDA graphs. It won't be reset by spec_metadata.prepare(). - if isinstance(spec_metadata, Eagle3SpecMetadata): - spec_metadata.is_first_draft = True - spec_metadata.eagle3_resource_manager.is_first_draft = True - - -class LinearDraftingLoopWrapper(BaseDraftingLoopWrapper): - - def __init__(self, max_draft_len: int, max_total_draft_tokens: int, - draft_model: torch.nn.Module): - super().__init__() - self.draft_model = draft_model - self.config = self.draft_model.config - self.model_config = self.draft_model.model_config - self.max_draft_len = max_draft_len - self.max_total_draft_tokens = max_total_draft_tokens - - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, - attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata, - **kwargs) -> dict[str, torch.Tensor]: - logits = self.draft_model.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - return_context_logits=True) - logits = logits[spec_metadata.gather_ids] - - new_draft_tokens = [self.sample(logits)] - draft_logits = [logits] - if self.max_draft_len > 1: - is_eagle3 = isinstance(spec_metadata, Eagle3SpecMetadata) - with save_metadata_state(attn_metadata, spec_metadata): - batch_size = attn_metadata.num_seqs - - new_position_ids = self.prepare_for_generation( - attn_metadata, spec_metadata, position_ids) - for i in range(self.max_draft_len - 1): - logits = self.draft_model.forward( - input_ids=new_draft_tokens[-1], - position_ids=new_position_ids, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata) - new_draft_tokens.append(self.sample(logits)) - draft_logits.append(logits) - new_position_ids += 1 - attn_metadata.kv_lens_cuda[:batch_size] += 1 - if i == 0 and is_eagle3: - spec_metadata.hidden_states_read_indices[:batch_size].copy_( - spec_metadata. - hidden_states_write_indices[:batch_size]) - - return { - "new_draft_tokens": torch.stack(new_draft_tokens), - "draft_logits": torch.stack(draft_logits) - } - - def sample(self, logits: torch.Tensor) -> torch.Tensor: - # TODO: inject the sampler here so we can support non-greedy - tokens, _ = greedy_search_sampling_batch(logits, return_probs=False) - if hasattr(self.draft_model.model, "d2t"): - d2t = self.draft_model.model.d2t.data - return tokens + d2t[tokens] - - return tokens - - def prepare_for_generation(self, attn_metadata: AttentionMetadata, - spec_metadata: SpecMetadata, - position_ids: torch.Tensor) -> torch.Tensor: - batch_size = attn_metadata.num_seqs - num_accepted_draft_tokens = spec_metadata.num_accepted_draft_tokens[: - batch_size] - seq_lens = attn_metadata.seq_lens_cuda[:batch_size] - attn_metadata.kv_lens_cuda[: - batch_size] -= seq_lens - num_accepted_draft_tokens - 1 - # Calculate last accepted token indices - last_tokens_idx = torch.cumsum( - seq_lens, dim=0, - dtype=torch.long) - seq_lens + num_accepted_draft_tokens - new_position_ids = position_ids[0, last_tokens_idx] + 1 - - attn_metadata._seq_lens[:batch_size].fill_(1) - attn_metadata._seq_lens_cuda[:batch_size].fill_(1) - attn_metadata.on_update() - attn_metadata.kv_lens_cuda[:batch_size] += 1 - - attn_metadata.host_request_types[:attn_metadata.num_contexts].fill_(1) - attn_metadata.num_contexts = 0 - # The next inference of draft model will not use spec decoding and the number of input tokens is 1 - attn_metadata.use_spec_decoding = False - - spec_metadata.num_tokens = batch_size - - if isinstance(spec_metadata, Eagle3SpecMetadata): - spec_metadata.eagle3_resource_manager.is_first_draft = False - spec_metadata.is_first_draft = False - - old_write_indices = spec_metadata.hidden_states_write_indices - - spec_metadata.hidden_states_read_indices[:batch_size].copy_( - old_write_indices[last_tokens_idx]) - spec_metadata.hidden_states_write_indices[:batch_size].copy_( - torch.arange( - batch_size, - dtype=spec_metadata.hidden_states_write_indices.dtype, - device=spec_metadata.hidden_states_write_indices.device)) - - return new_position_ids - - -class StaticTreeDraftingLoopWrapper(BaseDraftingLoopWrapper): - - def __init__(self, max_draft_len: int, max_total_draft_tokens: int, - max_batch_size: int, draft_model: torch.nn.Module): - super().__init__() - self.draft_model = draft_model - self.config = self.draft_model.config - self.model_config = self.draft_model.model_config - self.max_draft_len = max_draft_len - self.max_total_draft_tokens = max_total_draft_tokens - self.max_batch_size = max_batch_size - - self.draft_tokens_buffer = torch.zeros( - (max_batch_size, max_total_draft_tokens + 1), - dtype=torch.int64, - device='cuda') - self.position_ids_buffer = torch.zeros( - (max_batch_size, max_total_draft_tokens + 1), - dtype=torch.int64, - device='cuda') - - def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, - attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata, - **kwargs) -> dict[str, torch.Tensor]: - assert isinstance(spec_metadata, Eagle3SpecMetadata) - spec_tree_manager = spec_metadata.eagle3_resource_manager.spec_tree_manager - - logits = self.draft_model.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - return_context_logits=True) - batch_size = attn_metadata.num_seqs - vocab_size = logits.shape[-1] - logits = logits[spec_metadata.gather_ids] # [batch_size, vocab_size] - - # new_draft_tokens: [batch_size * max_top_k] - new_draft_tokens = self.sample(logits=logits, - max_top_k=spec_tree_manager.max_top_k) - - self.extract_real_draft_tokens( - cur_draft_idx=0, - batch_size=batch_size, - new_draft_tokens=new_draft_tokens, - use_cuda_graph=attn_metadata.is_cuda_graph, - spec_tree_manager=spec_tree_manager) - return_draft_logits = None - with save_metadata_state(attn_metadata, spec_metadata): - batch_size = attn_metadata.num_seqs - - self.prepare_for_generation(attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - spec_tree_manager=spec_tree_manager, - position_ids=position_ids) - - for layer_idx in range(1, self.max_draft_len): - # input_ids: [batch_size * (max_total_draft_tokens + 1)] - # position_ids: [batch_size * (max_total_draft_tokens + 1)] - # logits: [batch_size * (max_total_draft_tokens + 1), vocab_size] - logits = self.draft_model.forward( - input_ids=self.draft_tokens_buffer[:batch_size, :self. - max_total_draft_tokens + - 1].reshape(-1), - position_ids=self. - position_ids_buffer[:batch_size, :self. - max_total_draft_tokens + 1].reshape(-1), - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - return_context_logits=True) - - # new_draft_tokens: [batch_size * (max_total_draft_tokens + 1) * max_top_k] - new_draft_tokens = self.sample( - logits=logits, max_top_k=spec_tree_manager.max_top_k) - # Keep updating - self.extract_real_draft_tokens( - cur_draft_idx=layer_idx, - batch_size=batch_size, - new_draft_tokens=new_draft_tokens, - use_cuda_graph=attn_metadata.is_cuda_graph, - spec_tree_manager=spec_tree_manager) - - if layer_idx == self.max_draft_len - 1: - return_draft_logits = logits - - # self.draft_tokens_buffer[:batch_size, :]: [batch_size, max_total_draft_tokens + 1] - # return_new_draft_tokens: [max_total_draft_tokens, batch_size] - return_new_draft_tokens = torch.transpose( - self.draft_tokens_buffer[:batch_size, :-1], 0, 1) - - # return_draft_logits: [batch_size, max_total_draft_tokens + 1, vocab_size] -> [max_total_draft_tokens, batch_size, vocab_size] - if return_draft_logits is None: - # When max_draft_len == 1, the loop doesn't execute. - # Expand the initial logits to match the expected shape. - return_draft_logits = logits.unsqueeze(1).expand( - batch_size, self.max_total_draft_tokens + 1, - vocab_size).reshape(-1, vocab_size) - - return_draft_logits = return_draft_logits.reshape( - batch_size, self.max_total_draft_tokens + 1, vocab_size) - return_draft_logits = torch.transpose(return_draft_logits[:, :-1, :], 0, - 1) - - assert return_new_draft_tokens.shape == (self.max_total_draft_tokens, - batch_size) - assert return_draft_logits.shape == (self.max_total_draft_tokens, - batch_size, vocab_size) - - return { - "new_draft_tokens": return_new_draft_tokens, - "draft_logits": return_draft_logits - } - - def sample(self, logits: torch.Tensor, max_top_k: int) -> torch.Tensor: - # TODO: inject the sampler here so we can support non-greedy - - # for draft_layer_idx == 0, logits is of shape [batch_size, vocab_size] - # for draft_layer_idx > 0, logits is of shape [batch_size * (max_total_draft_tokens + 1), vocab_size] - indices = torch.topk( - logits, k=max_top_k, dim=-1 - ).indices # [batch_size, max_top_k] or [batch_size * max_total_draft_tokens, max_top_k] - tokens = indices.reshape(-1) - - if hasattr(self.draft_model.model, "d2t"): - d2t = self.draft_model.model.d2t.data - tokens = tokens + d2t[tokens] - - return tokens - - def extract_real_draft_tokens(self, cur_draft_idx: int, batch_size: int, - new_draft_tokens: torch.Tensor, - use_cuda_graph: bool, - spec_tree_manager: SpecTreeManager): - ''' - Extract the real draft tokens from the new draft tokens to self.draft_tokens_buffer. - ''' - # After the first drafter layer, new_draft_tokens: [batch_size * max_top_k] - # For other drafter layers, new_draft_tokens: [batch_size * (max_total_draft_tokens + 1) * max_top_k] - if cur_draft_idx == 0: - assert new_draft_tokens.shape[0] == (batch_size * - spec_tree_manager.max_top_k) - else: - assert new_draft_tokens.shape[0] == ( - batch_size * (self.max_total_draft_tokens + 1) * - spec_tree_manager.max_top_k) - - # reshape the new_draft_tokens to [batch_size, -1, spec_tree_manager.max_top_k] - new_draft_tokens = new_draft_tokens.reshape(batch_size, -1, - spec_tree_manager.max_top_k) - - # If using cuda graph, we need to use a torch op to implement this logic - if use_cuda_graph: - torch.ops.trtllm.extract_real_draft_tokens_op( - new_draft_tokens, self.draft_tokens_buffer, spec_tree_manager. - tokens_gather_idx_for_drafter_model[cur_draft_idx], - spec_tree_manager.top_k_list_cuda[cur_draft_idx], - spec_tree_manager.draft_tokens_indices_cumsum, cur_draft_idx, - batch_size, self.max_draft_len, self.max_total_draft_tokens, - spec_tree_manager.max_top_k) - else: - # 1) Gather the real tokens processed by this layer - process_tokens = new_draft_tokens[:, spec_tree_manager. - tokens_gather_idx_for_drafter_model[ - cur_draft_idx], :] # [batch_size, num_tokens_process_this_layer, max_top_k] - process_tokens = process_tokens.reshape( - -1, spec_tree_manager.max_top_k - ) # [batch_size * num_tokens_process_this_layer, max_top_k] - - # 2) Gather the real draft tokens samples by these processed tokens' logits - top_k_list = spec_tree_manager.top_k_list_cuda[ - cur_draft_idx].repeat( - batch_size) # [batch_size * num_tokens_process_this_layer] - assert top_k_list.shape[0] == process_tokens.shape[0] - - # [batch_size * num_tokens_process_this_layer, spec_tree_manager.max_top_k] - col_indices = torch.arange( - spec_tree_manager.max_top_k, - device=new_draft_tokens.device).unsqueeze(0).repeat( - top_k_list.shape[0], 1) - - mask = col_indices < top_k_list.unsqueeze( - 1 - ) # [batch_size * num_tokens_process_this_layer, spec_tree_manager.max_top_k] - - real_new_draft_tokens = process_tokens[ - mask] # [batch_size * sum(spec_tree_manager.top_k_list_cuda[cur_draft_idx])] - real_new_draft_tokens = real_new_draft_tokens.reshape( - batch_size, -1 - ) # [batch_size, sum(spec_tree_manager.top_k_list_cuda[cur_draft_idx])] - - self.draft_tokens_buffer[:batch_size, spec_tree_manager. - draft_tokens_indices_cumsum[cur_draft_idx]: - spec_tree_manager. - draft_tokens_indices_cumsum[ - cur_draft_idx + - 1]] = real_new_draft_tokens[:, :] - - def prepare_for_generation(self, attn_metadata: AttentionMetadata, - spec_metadata: SpecMetadata, - spec_tree_manager: SpecTreeManager, - position_ids: torch.Tensor): - ''' - Prepare the inputs for the subsequent draft layers. - Note: Except for the 0th drafter layer, in each subsequent drafter layer, - we take 'max_total_drafter_tokens + 1' draft tokens as input. - Only the first part of the draft tokens is meaningful, and the later tokens can be regarded as padding - until we continuously write the correct value. - - This introduces additional redundant computation, but it makes it compatible with cuda graphs. - - What we need to prepare are: - 1) position_ids - 2) attn_metadata - 2.1) kv_lens_cuda - 2.2) _seq_lens, _seq_lens_cuda - 2.3) host_request_types - 2.4) num_contexts - 2.5) use_spec_decoding - 2.6) spec_decoding_position_offsets - 2.7) spec_decoding_packed_mask - 2.8) spec_decoding_generation_lengths - 3) spec_metadata - 3.1) num_tokens - 3.2) hidden_states_read_indices, hidden_states_write_indices - 3.3) is_first_draft - ''' - batch_size = attn_metadata.num_seqs - - # 1) Prepare the position_ids - num_accepted_draft_tokens = spec_metadata.num_accepted_draft_tokens[: - batch_size] - seq_lens = attn_metadata.seq_lens_cuda[:batch_size] - # Calculate last accepted token indices - last_tokens_idx = torch.cumsum( - seq_lens, dim=0, - dtype=torch.long) - seq_lens + num_accepted_draft_tokens - position_start_idx = position_ids[0, - last_tokens_idx] + 1 # [batch_size] - self.position_ids_buffer[:batch_size, :-1] = position_start_idx.unsqueeze( - 1) + spec_tree_manager.spec_dec_position_offsets[0, 1:].unsqueeze( - 0) - 1 # exclude the root node - # 2) Prepare the attn_metadata - ## 2.1) kv_lens_cuda - attn_metadata.kv_lens_cuda[: - batch_size] -= seq_lens - num_accepted_draft_tokens - 1 - attn_metadata.kv_lens_cuda[:batch_size] += ( - self.max_total_draft_tokens + 1) - - ## 2.2) _seq_lens, _seq_lens_cuda - attn_metadata._seq_lens[:batch_size].fill_(self.max_total_draft_tokens + - 1) - attn_metadata._seq_lens_cuda[:batch_size].fill_( - self.max_total_draft_tokens + 1) - attn_metadata.on_update() - - ## 2.3) host_request_types - attn_metadata.host_request_types[:attn_metadata.num_contexts].fill_(1) - - ## 2.4) num_contexts - attn_metadata.num_contexts = 0 - - ## 2.5) use_spec_decoding - attn_metadata.use_spec_decoding = True - - ## 2.6) spec_decoding_position_offsets - ### attn_metadata.spec_decoding_position_offsets: [max_num_requests, max_total_draft_tokens + 1] - attn_metadata.spec_decoding_position_offsets[:batch_size, :self. - max_total_draft_tokens] = spec_tree_manager.spec_dec_position_offsets[ - 0, 1:self. - max_total_draft_tokens + - 1].unsqueeze( - 0 - ) - 1 # exclude the root node - attn_metadata.spec_decoding_position_offsets[:batch_size, self. - max_total_draft_tokens] = 0 # padding - - ## 2.7) spec_decoding_packed_mask - ### attn_metadata.spec_decoding_packed_mask: [max_num_requests, max_total_draft_tokens + 1, math.ceil((max_total_draft_tokens + 1) / 32)] - attn_metadata.spec_decoding_packed_mask[: - batch_size, :, :] = spec_tree_manager.spec_dec_packed_mask_for_drafter_model - - ## 2.8) spec_decoding_generation_lengths - ### attn_metadata.spec_decoding_generation_lengths: [max_num_requests] - attn_metadata.spec_decoding_generation_lengths[: - batch_size] = self.max_total_draft_tokens + 1 - - # 3) Update spec_metadata - ## 3.1) num_tokens - spec_metadata.num_tokens = batch_size * (self.max_total_draft_tokens + - 1) - ## 3.2) hidden_states_read_indices, hidden_states_write_indices - ### spec_metadata.hidden_states_read_indices: [self.max_num_tokens] - ### spec_metadata.hidden_states_write_indices: [self.max_num_tokens] - old_write_indices = spec_metadata.hidden_states_write_indices - start_idx = old_write_indices[ - last_tokens_idx] # [batch_size], already take the accepted tokens into account. - - ### shape: [batch_size, self.max_total_draft_tokens + 1] - hidden_states_read_indices_offset = spec_tree_manager.hidden_states_read_indices_offset_for_drafter_model[:self.max_total_draft_tokens + 1].repeat( - batch_size).reshape(batch_size, self.max_total_draft_tokens + 1) - hidden_states_read_indices_offset = hidden_states_read_indices_offset + start_idx.unsqueeze( - 1) - spec_metadata.hidden_states_read_indices[:batch_size * ( - self.max_total_draft_tokens + - 1)] = hidden_states_read_indices_offset.reshape(-1) - - hidden_states_write_offset = torch.arange( - 1, self.max_total_draft_tokens + 1 + 1, - device=position_ids.device).unsqueeze(0).repeat( - batch_size, 1) + start_idx.unsqueeze(1) - spec_metadata.hidden_states_write_indices[:batch_size * ( - self.max_total_draft_tokens + - 1)] = hidden_states_write_offset.reshape(-1) - - ## 3.3) is_first_draft - spec_metadata.eagle3_resource_manager.is_first_draft = False - spec_metadata.is_first_draft = False - - return diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index a1f4f82869bd..62cc0f6ef732 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -1,5 +1,5 @@ -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Dict, List, Optional, Set +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Optional, Set import torch import torch.nn.functional as F @@ -27,9 +27,8 @@ class Eagle3ResourceManager(BaseResourceManager): """ - Eagle3 needs to save the hidden states for the draft model. When using - Eagle3TwoModel, there will be two model engines, one for the target model - and one for the draft model. Use this class to manage the hidden states. + Eagle3 needs to save the hidden states for the draft model. Use this class + to manage the hidden states. """ def __init__(self, @@ -195,178 +194,6 @@ def _get_eagle3_default_capture_layers(num_layers: int): return (1, num_layers // 2 - 1, num_layers - 4) -@dataclass -class Eagle3SpecMetadata(SpecMetadata): - hidden_states: List[torch.Tensor] = field(default_factory=list) - layers_to_capture: Optional[Set[int]] = None - target_model_embed_tokens: Optional[torch.nn.Module] = None - hidden_size: int = 0 - max_num_tokens: int = 0 - dtype: torch.dtype = torch.bfloat16 - is_draft_model: bool = False - is_first_draft: bool = False - eagle3_resource_manager: Optional[Eagle3ResourceManager] = None - is_mtp_eagle: bool = False - - eagle_choices: Optional[List[List[int]]] = None - max_total_draft_tokens: int = 0 - # This is to store the request type and accepted path for each request. - # For each request, {key: request_ids, value: accepted_path} - # 'accepted_path' is a list of accepted tokens indices. - request_accepted_path: Optional[Dict[int, List[int]]] = None - - def __post_init__(self): - if self.is_draft_model: - self.layers_to_capture = (self.num_layers - 1, ) - elif self.layers_to_capture is None: - if self.num_layers == 1 or self.is_mtp_eagle: - self.layers_to_capture = (-1, ) - else: - if self.num_layers <= 5: - raise ValueError( - "Not enough hidden layers for default EAGLE3 capture") - self.layers_to_capture = _get_eagle3_default_capture_layers( - self.num_layers) - else: - self.layers_to_capture = sorted(list(self.layers_to_capture)) - if self.layers_to_capture[0] == -1: - self.layers_to_capture = self.layers_to_capture[1:] + [ - self.layers_to_capture.pop(0) - ] - self.num_capture_layers = len(self.layers_to_capture) - - # Initialize to 0 to avoid reading uninitialized memory during warmup - self.hidden_states_read_indices = torch.zeros([self.max_num_tokens], - dtype=torch.long, - device='cuda') - self.hidden_states_write_indices = torch.zeros([self.max_num_tokens], - dtype=torch.long, - device='cuda') - self.hidden_states_read_indices_host = None - self.hidden_states_write_indices_host = None - - if self.eagle_choices is not None: - self.is_spec_dec_tree = True - self.is_spec_dec_dynamic_tree = False - - def prepare(self): - super().prepare() - is_first_draft = self.eagle3_resource_manager.is_first_draft - spec_tree_manager = self.eagle3_resource_manager.spec_tree_manager - # Update start indices - # Here, we assume the sequence lengths (seq_lens) during the draft model - # forward will not exceed those of the target model. So pre-allocate - # hidden state space before the target model forward. - start_idx = 0 - if not self.is_draft_model: - for req_id, seq_len in zip(self.request_ids, self.seq_lens): - slot_id = self.eagle3_resource_manager.slot_manager.get_slot( - req_id) - self.eagle3_resource_manager.start_indices[slot_id] = start_idx - # Make sure that the space between two requests is at least max_total_draft_tokens + 1. - start_idx += max(seq_len, self.max_total_draft_tokens + 1) - assert start_idx < self.eagle3_resource_manager.hidden_states.shape[ - 0], f"start_idx {start_idx} is greater than hidden_states.shape[0] {self.eagle3_resource_manager.hidden_states.shape[0]}" - - # Prepare hidden states gather ids - hidden_states_read_indices = [] - hidden_states_write_indices = [] - for req_id, seq_len in zip(self.request_ids, self.seq_lens): - slot_id = self.eagle3_resource_manager.slot_manager.get_slot(req_id) - start_idx = self.eagle3_resource_manager.start_indices[slot_id] - # 1) target model or (is_first_draft and is_linear_tree) - # If this is the first draft or the target model forward, we need to - # read/write all of the hidden states - if not self.is_draft_model or (is_first_draft - and spec_tree_manager is None): - hidden_states_read_indices.extend( - list(range(start_idx, start_idx + seq_len))) - hidden_states_write_indices.extend( - list(range(start_idx, start_idx + seq_len))) - # 2）is_first_draft and draft_token_tree - # After target model forward, some draft tokens will be accepted. - # These draft tokens' hidden states will be used for draft model's first drafter layer. - elif is_first_draft and spec_tree_manager is not None: - assert req_id in self.request_accepted_path.keys( - ), f"Request {req_id} not found in request_accepted_path" - # 'node_idx + 1' is because we '-1' in sampler.py for kv cache rewind. Now we add it back. - accepted_path = [ - node_idx + 1 - for node_idx in self.request_accepted_path[req_id] - ] - - if accepted_path == []: - # Case 1: This is a context request, We need to read all the hidden states. - # Case 2: This is a generation request, but no accepted tokens are accepted. Actually only the first token's hidden states is needed. The others are just padding tokens. - hidden_states_read_indices.extend( - list(range(start_idx, start_idx + seq_len))) - else: - # This is a generation request. And there are draft tokens accepted. - # We only read the accepted tokens' hidden states. - accepted_path = [0] + accepted_path # add the root node - accepted_path_pad = accepted_path + [0] * ( - seq_len - len(accepted_path)) - assert len(accepted_path_pad) == seq_len - hidden_states_read_indices.extend([ - start_idx + accepted_draft_token_offset - for accepted_draft_token_offset in accepted_path_pad - ]) - - # For the write indices, we just write all the hidden states. - hidden_states_write_indices.extend( - list(range(start_idx, start_idx + seq_len))) - # otherwise: only read the last token - else: - old_seq_len = self.eagle3_resource_manager.seq_lens[slot_id] - hidden_states_read_indices.append(start_idx + old_seq_len - 1) - hidden_states_write_indices.append(start_idx + seq_len - 1) - self.eagle3_resource_manager.seq_lens[slot_id] = seq_len - # Prepare hidden states gather ids - self.hidden_states_read_indices_host = torch.tensor( - hidden_states_read_indices, - dtype=torch.long, - pin_memory=prefer_pinned()) - self.hidden_states_write_indices_host = torch.tensor( - hidden_states_write_indices, - dtype=torch.long, - pin_memory=prefer_pinned()) - self.is_first_draft = is_first_draft and self.is_draft_model - if self.is_draft_model: - self.eagle3_resource_manager.is_first_draft = False - - self.hidden_states_read_indices[:self.num_tokens].copy_( - self.hidden_states_read_indices_host, non_blocking=True) - self.hidden_states_write_indices[:self.num_tokens].copy_( - self.hidden_states_write_indices_host, non_blocking=True) - - def is_layer_capture(self, layer_id: int): - return layer_id in self.layers_to_capture - - def maybe_capture_hidden_states( - self, - layer_id: int, - hidden_states: torch.Tensor, - residual: Optional[torch.Tensor] = None) -> None: - token_idx = self.hidden_states_write_indices[:self.num_tokens] - eagle3_hidden_states = self.eagle3_resource_manager.hidden_states - for i, captured_layer_id in enumerate(self.layers_to_capture): - if captured_layer_id == layer_id: - to_save = hidden_states + residual if residual is not None else hidden_states - to_save = to_save.to(dtype=eagle3_hidden_states.dtype) - eagle3_hidden_states[:, i * self.hidden_size:(i + 1) * - self.hidden_size].index_copy_( - 0, token_idx, - to_save[:self.num_tokens]) - break - - def get_hidden_states(self): - hidden_states = self.eagle3_resource_manager.hidden_states[ - self.hidden_states_read_indices[:self.num_tokens], :] - if not self.is_first_draft: - hidden_states = hidden_states[:, :self.hidden_size] - return hidden_states - - @dataclass class Eagle3OneModelSpecMetadata(SpecMetadata): # The hidden states @@ -405,8 +232,8 @@ class Eagle3OneModelSpecMetadata(SpecMetadata): def __post_init__(self): if self.layers_to_capture is None: - if self.spec_dec_mode.is_mtp_eagle_one_model(): - # MTP Eagle one-model feeds the target model's hidden_states + if self.spec_dec_mode.is_mtp_eagle(): + # MTP Eagle feeds the target model's hidden_states # directly to the MTP layer (see Eagle3OneModelWorker # prepare_1st_drafter_inputs / _run_draft_forward, both gated # on self.is_mtp_eagle). It never reads spec_metadata.hidden_states, @@ -507,7 +334,7 @@ def prepare(self): # Each mode uses a different convention: # - MTP Eagle: keep the 1st-iter shape (matches input_ids). # - Eagle3: subtract to the subseq shape. - if not self.spec_dec_mode.is_mtp_eagle_one_model(): + if not self.spec_dec_mode.is_mtp_eagle(): if self.is_spec_dec_tree: self.num_tokens -= ( self.num_generations) * self.max_total_draft_tokens @@ -611,9 +438,9 @@ class Eagle3OneModelWorker(SpecWorkerBase): """Unified one-model worker for Eagle3 and MTP Eagle speculative decoding. The operating mode is determined by ``spec_config.spec_dec_mode``: - - EAGLE3_ONE_MODEL: multi-layer hidden states from Eagle3, apply_eagle3_fc + - EAGLE3: multi-layer hidden states from Eagle3, apply_eagle3_fc projection, independent EAGLE draft model network. - - MTP_EAGLE_ONE_MODEL: single last-layer hidden states, MTP layer called + - MTP_EAGLE: single last-layer hidden states, MTP layer called repeatedly, supports TP-aware sampling and Mamba hybrid cache. Where the two modes differ, ``self.is_mtp_eagle`` is used to branch. @@ -633,8 +460,8 @@ def __init__(self, # Eagle3 path can leave it as None. self.model_config = model_config - # Mode flag: True = MTP Eagle one-model, False = Eagle3 one-model. - self.is_mtp_eagle = spec_config.spec_dec_mode.is_mtp_eagle_one_model() + # Mode flag: True = MTP Eagle, False = Eagle3. + self.is_mtp_eagle = spec_config.spec_dec_mode.is_mtp_eagle() # SA enhancer (common to both modes) self.sa_enhancer: Optional[SADraftEnhancer] = None diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 50e23d087ca7..e218d20f705c 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -32,8 +32,7 @@ from ..attention_backend.trtllm import (AttentionBackend, TrtllmAttention, TrtllmAttentionMetadata) from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE -from ..pyexecutor.resource_manager import (BaseResourceManager, - ResourceManagerType) +from ..pyexecutor.resource_manager import ResourceManagerType if TYPE_CHECKING: from ..pyexecutor.guided_decoder import CapturableGuidedDecoder @@ -271,13 +270,10 @@ def get_force_num_accepted_tokens_float() -> float: class SpeculativeDecodingMode(IntEnum): MTP = auto() MTP_EAGLE = auto() - MTP_EAGLE_ONE_MODEL = auto() EAGLE3 = auto() - EAGLE3_ONE_MODEL = auto() NGRAM = auto() SA = auto() DRAFT_TARGET = auto() - DRAFT_TARGET_ONE_MODEL = auto() USER_PROVIDED = auto() SAVE_HIDDEN_STATES = auto() PARD = auto() @@ -287,13 +283,10 @@ class SpeculativeDecodingMode(IntEnum): AUTO = auto() def is_mtp_one_model(self): - # Union: covers vanilla MTP and MTP_EAGLE_ONE_MODEL. Use is_mtp_vanilla() + # Union: covers vanilla MTP and MTP_EAGLE. Use is_mtp_vanilla() # when only the vanilla MTP variant should match. return (self == SpeculativeDecodingMode.MTP - or self == SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL) - - def is_mtp_eagle_one_model(self): - return self == SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL + or self == SpeculativeDecodingMode.MTP_EAGLE) def is_mtp_vanilla(self): return self == SpeculativeDecodingMode.MTP @@ -305,12 +298,9 @@ def is_eagle3(self): return self == SpeculativeDecodingMode.EAGLE3 def use_one_engine(self): - return self.is_eagle3_one_model() or self.is_mtp_one_model( + return self.is_eagle3() or self.is_mtp_one_model( ) or self.is_external_drafter() or self.is_sa() - def is_eagle3_one_model(self): - return self == SpeculativeDecodingMode.EAGLE3_ONE_MODEL - def is_pard(self): return self == SpeculativeDecodingMode.PARD @@ -338,66 +328,48 @@ def is_none(self): def is_draft_target(self): return self == SpeculativeDecodingMode.DRAFT_TARGET - def is_draft_target_one_model(self): - return self == SpeculativeDecodingMode.DRAFT_TARGET_ONE_MODEL - def is_save_hidden_states(self): return self == SpeculativeDecodingMode.SAVE_HIDDEN_STATES def is_external_drafter(self): - return self.is_parallel_draft() or self.is_draft_target_one_model() + return self.is_parallel_draft() or self.is_draft_target() def without_logits(self): - return self.is_mtp_one_model() or self.is_eagle3_one_model( + return self.is_mtp_one_model() or self.is_eagle3( ) or self.is_external_drafter() or self.is_sa() def needs_kv_cache_rewind(self): - return self.is_mtp_one_model() or self.is_eagle3_one_model( - ) or self.is_ngram() or self.is_sa() or self.is_external_drafter() + return self.is_mtp_one_model() or self.is_eagle3() or self.is_ngram( + ) or self.is_sa() or self.is_external_drafter() def support_overlap_scheduler(self): - return self.is_mtp_one_model() or self.is_eagle3_one_model( - ) or self.is_sa() or self.has_draft_model() or self.is_external_drafter( - ) + return self.is_mtp_one_model() or self.is_eagle3() or self.is_sa( + ) or self.is_external_drafter() def support_guided_decoder(self): return self.is_none() or self.has_spec_drafter() def support_capturable_guided_decoder(self): - return self.is_mtp_one_model() or self.is_eagle3_one_model( + return self.is_mtp_one_model() or self.is_eagle3( ) or self.is_external_drafter() or self.is_sa() def support_dynamic_draft_len(self): - return self.is_mtp_one_model() or self.is_eagle3_one_model( - ) or self.is_mtp_eagle_one_model() or self.is_pard() or self.is_dflash( - ) or self.is_draft_target_one_model() or self.is_sa() - - def has_draft_model(self): - return self.is_eagle3() or self.is_draft_target() or self.is_mtp_eagle() - - def needs_kv_cache_recompute(self): - """ - Whether the draft model needs to recompute the kv cache. - If true, the 1st draft model forward will recompute the kv cache for - the accepted draft tokens. - """ - return self.is_eagle3() or self.is_mtp_eagle() + return self.is_mtp_one_model() or self.is_eagle3() or self.is_pard( + ) or self.is_dflash() or self.is_draft_target() or self.is_sa() def need_load_draft_weights(self): """ Whether the draft model and target model are in the same model engine, and the draft model needs to load weights from the separate checkpoint. """ - return self.is_eagle3_one_model() or self.is_external_drafter() + return self.is_eagle3() or self.is_external_drafter() def has_spec_decoder(self): - return self.is_mtp_one_model() or self.is_mtp_eagle() or self.is_eagle3( - ) or self.is_eagle3_one_model() or self.is_external_drafter( - ) or self.is_sa() + return self.is_mtp_one_model() or self.is_eagle3( + ) or self.is_external_drafter() or self.is_sa() def has_spec_drafter(self): - return self.is_eagle3() or self.is_draft_target() or self.is_ngram( - ) or self.is_user_provided() or self.is_mtp_eagle() + return self.is_ngram() or self.is_user_provided() def extend_ctx(self, attention_backend: Type[AttentionBackend]): """ @@ -414,31 +386,21 @@ def extend_ctx(self, attention_backend: Type[AttentionBackend]): TrtllmAttention) or not xqa_supported def attention_need_spec_dec_mode( - self, - spec_resource_manager: Optional[BaseResourceManager], - is_draft_model: bool, - attention_backend: Type[AttentionBackend], - use_chain_drafter: bool, # CDL + self, + attention_backend: Type[AttentionBackend], ): """ If true, the attention backend kernel needs to run in spec-dec mode (multi-token query mode). Args: - spec_resource_manager: the resource manager for the spec-dec mode. - is_draft_model: whether the model is a draft model. attention_backend: the attention backend. - use_chain_drafter: whether to use capturable drafting loops (CDL). For the target model, it is always False. """ is_trtllm_attention = issubclass(attention_backend, TrtllmAttention) - # Always use the multi-token query mode for 1-model if the kernels are available. + # Always use the multi-token query mode for one-engine modes if the kernels are available. use_case_1 = self.use_one_engine() - # For 2-model, we need to enable it when we process multiple tokens at once. This occurs with - # the target model (verification) or on the first draft for CDL based speculation. - use_case_2 = not self.use_one_engine() and ( - not is_draft_model or - (spec_resource_manager is not None - and spec_resource_manager.is_first_draft - and use_chain_drafter)) and is_trtllm_attention + # For drafter-loop modes (ngram/user-provided), the target model verifies + # multiple draft tokens at once and needs spec-dec mode on TRTLLM attention. + use_case_2 = not self.use_one_engine() and is_trtllm_attention return use_case_1 or use_case_2 diff --git a/tensorrt_llm/_torch/speculative/model_drafter.py b/tensorrt_llm/_torch/speculative/model_drafter.py deleted file mode 100644 index 5eae9b7e44cc..000000000000 --- a/tensorrt_llm/_torch/speculative/model_drafter.py +++ /dev/null @@ -1,1049 +0,0 @@ -from __future__ import annotations - -import traceback -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple - -import torch - -from tensorrt_llm._utils import nvtx_range, prefer_pinned -from tensorrt_llm.logger import logger - -from ..attention_backend.trtllm import TrtllmAttention -from ..pyexecutor.guided_decoder import GuidedDecoder -from ..pyexecutor.handle_logits import HandleLogits -from ..pyexecutor.llm_request import LlmRequest, LlmRequestState -from ..pyexecutor.resource_manager import (BaseResourceManager, ResourceManager, - ResourceManagerType) -from ..pyexecutor.sampler import Sampler, SampleState, SampleStateTensors -from ..pyexecutor.scheduler import ScheduledRequests -from ..pyexecutor.seq_slot_manager import SeqSlotManager -from .drafter import Drafter -from .spec_sampler_base import SampleStateTensorsSpec - -if TYPE_CHECKING: - from ...llmapi.llm_args import DecodingBaseConfig - from ..pyexecutor.model_engine import ModelEngine - from .interface import SpeculativeDecodingMode - - -# Place the tool function here to avoid circular import -def get_draft_model_prompt(spec_dec_mode: SpeculativeDecodingMode, - request: LlmRequest, - disable_overlap_scheduler: bool) -> List[int]: - """ - Can be used to modify prompts for speculative algorithms that need to update tokens - before drafting. - """ - draft_input_tokens = request.get_tokens(0) - if spec_dec_mode.is_eagle3() or spec_dec_mode.is_mtp_eagle(): - # EAGLE3 always throws away the first token when processing draft inputs - if not disable_overlap_scheduler: - # Add a fake golden token here since the real one has not been generated. - draft_input_tokens.extend([0]) - draft_input_tokens = draft_input_tokens[1:] - - if request.is_context_init_state: - # A draft request's prompt is its target request's prompt adding the first golden token. - # Add a fake golden token here since the real one has not been generated. - draft_input_tokens.append(0) - return draft_input_tokens - - -class ModelDrafter(Drafter): - """Model-based drafter that uses a draft model to generate draft tokens.""" - - _needs_padding_kv_extension = True - - def __init__( - self, - spec_config: "DecodingBaseConfig", - draft_model_engine: "ModelEngine", - max_draft_len: int, - max_total_draft_tokens: int, - draft_seq_slot_manager: SeqSlotManager, - sampler: Sampler, - spec_resource_manager: Optional[BaseResourceManager] = None, - guided_decoder: Optional[GuidedDecoder] = None, - ): - # Validate required parameters - if draft_model_engine is None: - raise ValueError("draft_model_engine cannot be None") - if max_draft_len < 0: - raise ValueError("max_draft_len must be >= 0") - if max_total_draft_tokens < 0: - raise ValueError("max_total_draft_tokens must be >= 0") - assert max_draft_len <= max_total_draft_tokens - - super().__init__(max_draft_len=max_draft_len, - max_total_draft_tokens=max_total_draft_tokens, - max_concurrency=spec_config.max_concurrency, - draft_len_schedule=spec_config.draft_len_schedule) - - # Model and resource management - self.draft_model_engine = draft_model_engine - self.draft_seq_slot_manager = draft_seq_slot_manager - self.spec_resource_manager = spec_resource_manager - self.disable_overlap_scheduler = True - - # Configuration - self.spec_config = spec_config - - # Sampling - self.sampler = sampler - self.guided_decoder = guided_decoder - - self.use_static_draft_loop = draft_model_engine.model_is_wrapped - if self.use_static_draft_loop: - # TODO: enable sampling/guided decoding on static draft loop - assert guided_decoder is None - assert spec_config._allow_greedy_draft_tokens - assert spec_config.draft_len_schedule is None - - # Create accumulator for draft tokens in non-CDL mode - self.draft_tokens_accumulator: Dict[int, List[int]] = {} - - # Initialize draft latency tracking for specDecodingStats - self.last_draft_latency_ms = 0.0 - # Track previous draft batch for overlap scheduling - self.previous_draft_batch: Optional[ScheduledRequests] = None - self.previous_draft_outputs: Optional[Any] = None - self.previous_scheduled_batch: Optional[ScheduledRequests] = None - # Map from request ID to original request - self.req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None - - def _create_draft_request(self, request: LlmRequest, - input_tokens: Optional[List]) -> LlmRequest: - """Create a draft request with common parameters.""" - return LlmRequest( - input_tokens=input_tokens, - request_id=request.py_request_id, - max_new_tokens=request.py_max_new_tokens, - sampling_config=request.sampling_config, - guided_decoding_params=request.guided_decoding_params, - target_seq_slot=request.py_seq_slot, - return_perf_metrics=request.return_perf_metrics, - is_streaming=False, - exclude_last_generation_logits= - True, # prepare_draft_tokens uses overlap scheduling - is_draft=True, - # NB: self.sampler is shared with PyExecutor - return_generation_logits=self.sampler.should_provide_draft_probs( - request)) - - def _initialize_draft_tokens(self, request: LlmRequest) -> Tuple[int, int]: - """Initialize draft token tracking for a request.""" - if not self.disable_overlap_scheduler: - return self.max_draft_len, 0 - - num_draft_tokens = len( - request.py_last_draft_tokens - ) if request.py_last_draft_tokens is not None else 0 - - num_accepted_tokens = request.py_num_accepted_draft_tokens - num_rejected_tokens = num_draft_tokens - num_accepted_tokens - assert num_rejected_tokens >= 0 - - return num_draft_tokens, num_accepted_tokens - - def _create_context_request(self, request: LlmRequest, - input_tokens: Any) -> LlmRequest: - """Create a context request for first-time drafting.""" - new_request = self._create_draft_request(request, input_tokens) - - begin_compute, end_compute = request.py_last_context_chunk - if begin_compute is not None: - new_request.context_current_position = begin_compute - new_request.context_chunk_size = end_compute - begin_compute - return new_request - - def _create_generation_request(self, request: LlmRequest, - input_tokens: Any) -> LlmRequest: - """Create a generation request when no tokens were accepted.""" - new_request = self._create_draft_request(request, input_tokens) - new_request.state = LlmRequestState.GENERATION_IN_PROGRESS - return new_request - - def _create_accepted_tokens_request(self, request: LlmRequest, - input_tokens: Any, - num_accepted_tokens: int) -> LlmRequest: - """ - Create a chunked context request for accepted tokens. - Only applicable if the draft model needs to recompute KV cache for accepted tokens (e.g. eagle 3) - """ - new_request = self._create_draft_request(request, input_tokens) - new_request.context_chunk_size = num_accepted_tokens + 1 - new_request.context_current_position = len( - input_tokens) - num_accepted_tokens - 1 - return new_request - - def _get_previous_draft_request( - self, request: LlmRequest) -> Optional[LlmRequest]: - """Get the previous draft request for the given request.""" - if self.previous_draft_batch is None: - return None - for req in self.previous_draft_batch.all_requests(): - if req.py_request_id == request.py_request_id: - return req - return None - - def _create_accepted_tokens_request_for_trtllm_attn( - self, request: LlmRequest, input_tokens: Any, - num_accepted_tokens: int) -> LlmRequest: - """ - Create a chunked context request for accepted tokens. - Only applicable if the draft model needs to recompute KV cache for accepted tokens (e.g. eagle 3) - """ - # Pad input_tokens to max_draft_len - # We use max_draft_len instead of max_total_draft_tokens here, - # because at most max_draft_len draft tokens are accepted. - input_tokens.extend( - 0 for _ in range(self.max_draft_len - num_accepted_tokens)) - - # Reuse the previous draft request if it exists. - # This can reduce host overhead significantly. - draft_request = self._get_previous_draft_request(request) - if draft_request is not None: - generated_tokens = input_tokens[draft_request.py_prompt_len:] - draft_request.set_generated_tokens([generated_tokens]) - else: - draft_request = self._create_draft_request(request, input_tokens) - - draft_request.state = LlmRequestState.GENERATION_IN_PROGRESS - draft_request.py_num_accepted_draft_tokens = request.py_num_accepted_draft_tokens - draft_request.py_is_first_draft = True - # For tree decoding, we need to store the accepted tokens indices for these requests, - # which will be used to update the hidden_states_read_indices. - draft_request.py_num_accepted_draft_tokens_indices = request.py_num_accepted_draft_tokens_indices - - return draft_request - - def _create_draft_request_for_request( - self, request: LlmRequest) -> Optional[LlmRequest]: - """Create a draft request based on the original request state.""" - num_draft_tokens, num_accepted_tokens = self._initialize_draft_tokens( - request) - - input_tokens = get_draft_model_prompt(self.spec_config.spec_dec_mode, - request, - self.disable_overlap_scheduler) - - is_eagle_style = self.spec_config.spec_dec_mode.is_eagle3( - ) or self.spec_config.spec_dec_mode.is_mtp_eagle() - - # First time seeing this request - context request - num_overlap_tokens = 0 if self.disable_overlap_scheduler else 1 - if request.max_beam_num_tokens - 1 + num_overlap_tokens == request.py_prompt_len: - # This is the first time the draft model is seeing this request. - # Prepare a context request. We discard the first token and take - # the newly decoded one - this is the convention for EAGLE 2 and 3. - if self.disable_overlap_scheduler: - assert num_draft_tokens == 0 - return self._create_context_request(request, input_tokens) - - # For TRTLLM attention backend, we need to create a generation request for both no tokens accepted and tokens accepted - elif issubclass(self.draft_model_engine.attn_backend, TrtllmAttention - ) and self.use_static_draft_loop and is_eagle_style: - return self._create_accepted_tokens_request_for_trtllm_attn( - request, input_tokens, num_accepted_tokens) - - # No tokens accepted - generation request. This only applies to speculation algorithms - # that need to recompute KV cache for accepted tokens like eagle3. - elif num_accepted_tokens == 0 or not self.spec_config.spec_dec_mode.needs_kv_cache_recompute( - ): - return self._create_generation_request(request, input_tokens) - - # Tokens accepted - chunked context request - else: - return self._create_accepted_tokens_request(request, input_tokens, - num_accepted_tokens) - - def _add_to_draft_batch(self, draft_batch: ScheduledRequests, - draft_request: LlmRequest, - original_request: LlmRequest) -> None: - """Add the draft request to the appropriate batch list.""" - # Copy additional properties - draft_request.py_stop_words_list = original_request.py_stop_words_list - - # Add to appropriate batch based on request type - if draft_request.state == LlmRequestState.GENERATION_IN_PROGRESS: - draft_batch.append_generation_request(draft_request) - else: - draft_batch.append_context_request(draft_request) - - @nvtx_range("_prepare_draft_batch") - def _prepare_draft_batch( - self, scheduled_requests: ScheduledRequests) -> ScheduledRequests: - """ - Prepares a batch for the draft model engine. Draft tokens are only produced - for generation requests. - - The requests are prepared as follows: - 1. The first time the draft engine sees a request, it's a context request. - 2. Otherwise, if draft tokens were accepted on the last target model decoding - step, it's a chunked context request (we process all the accepted tokens together). - 3. Otherwise, it's a generation request. - - Args: - scheduled_requests: The scheduled requests to prepare draft batch for - - Returns: - ScheduledRequests: The prepared draft batch - """ - try: - for req in scheduled_requests.all_requests(): - draft_model = self.draft_model_engine.model.draft_model if self.use_static_draft_loop else self.draft_model_engine.model - if hasattr(draft_model.model, "d2t"): - req.d2t = draft_model.model.d2t.data - req.py_draft_use_greedy_sampling = self.use_static_draft_loop - - draft_batch = ScheduledRequests() - - for request in scheduled_requests.context_requests: - if request.py_disable_speculative_decoding: - continue - if request.is_first_context_chunk: - # Ignore requests which still need to be processed by the target model. - continue - - # We hit this path if we're doing chunked prefill. The target model processed - # a prefill chunk on the last iteration. Now, we need to fill in the KV cache - # for the draft model too. - input_tokens = get_draft_model_prompt( - self.spec_config.spec_dec_mode, request, - self.disable_overlap_scheduler) - - new_request = self._create_context_request( - request, input_tokens) - self._add_to_draft_batch(draft_batch, new_request, request) - - for request in scheduled_requests.generation_requests: - if request.py_disable_speculative_decoding: - continue - if request.state == LlmRequestState.GENERATION_COMPLETE: - # Skip generation complete requests. This could happen when enabling overlap scheduler. - continue - - if request.py_draft_pages_allocated == 0: - # No space for draft tokens - continue - # Stop drafting when we hit the max seqlen. We still need dummy draft - # tokens attached to the requests to make sure everything works properly - # with CUDA graph. These dummy tokens are already added by - # _prepare_draft_requests to make the KV cache/scheduler aware of the fact - # that we want to do spec decoding, so no need to do anything else here. - # This makes the perf for this case suboptimal, but that's OK - this is - # a corner case for weird models like the llama 3.1 8b EAGLE3 implementation. - num_overlap_tokens = 0 if self.disable_overlap_scheduler else 1 - if request.max_beam_num_tokens - 1 + num_overlap_tokens >= self.draft_model_engine.max_seq_len: - continue - - draft_request = self._create_draft_request_for_request(request) - if draft_request is not None: - self._add_to_draft_batch(draft_batch, draft_request, - request) - - return draft_batch - - except Exception as e: - logger.error(f"Error in _prepare_draft_batch: {str(e)}") - traceback.print_exc() - raise e - - def _should_disable_cuda_graph(self, is_first_draft_token: bool) -> bool: - """Check if CUDA graph should be disabled for the current forward pass.""" - if not is_first_draft_token: - return False - if self.use_static_draft_loop: - return False - return self.spec_config.spec_dec_mode.needs_kv_cache_recompute() - - @nvtx_range("forward_draft_model") - def forward_draft_model( - self, - draft_batch: ScheduledRequests, - resource_manager: ResourceManager, - is_first_draft_token: bool, - previous_tensors: Optional[SampleStateTensors] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None - ) -> Dict[str, Any]: - """Forward pass through the draft model.""" - if self._should_disable_cuda_graph(is_first_draft_token): - with self.draft_model_engine.no_cuda_graph(): - outputs = self.draft_model_engine.forward( - draft_batch, - resource_manager, - new_tensors_device=previous_tensors, - num_accepted_tokens_device=num_accepted_tokens_device, - req_id_to_old_request=self.req_id_to_old_request) - else: - outputs = self.draft_model_engine.forward( - draft_batch, - resource_manager, - new_tensors_device=previous_tensors, - num_accepted_tokens_device=num_accepted_tokens_device, - req_id_to_old_request=self.req_id_to_old_request) - - # Handle d2t data if available. Static drafting loops should incorporate d2t - # in their implementations. - if not self.use_static_draft_loop and hasattr( - self.draft_model_engine.model.model, 'd2t'): - outputs['d2t'] = self.draft_model_engine.model.model.d2t.data - - return outputs - - @nvtx_range("sample_async") - def sample_async( - self, - draft_batch: ScheduledRequests, - outputs: Dict[str, Any], - resource_manager: Optional[ResourceManager] = None - ) -> Optional[SampleState]: - """Sample tokens from draft model outputs.""" - try: - num_context_logits_prefix_sum = [0] - prefix_sum = 0 - for request in draft_batch.context_requests: - prefix_sum += request.context_chunk_size if request.py_return_context_logits else 1 - num_context_logits_prefix_sum.append(prefix_sum) - - HandleLogits()(draft_batch.context_requests, - draft_batch.generation_requests, outputs["logits"], - self.sampler.beam_width(draft_batch.all_requests()), - num_context_logits_prefix_sum, - self.sampler.is_generation_model()) - - return self.sampler.sample_async(draft_batch, outputs, - num_context_logits_prefix_sum, - resource_manager) - except Exception as e: - logger.error(f"Error in sampling: {str(e)}") - return None - - def update_request_states(self, - scheduled_requests: ScheduledRequests) -> None: - """Update request states after processing.""" - for request in scheduled_requests.context_requests: - if request.state != LlmRequestState.GENERATION_COMPLETE: - request.move_to_next_context_chunk() - if request.context_remaining_length == 0: - request.state = LlmRequestState.GENERATION_IN_PROGRESS - - def update_cur_draft_layer_idx( - self, - cur_draft_layer_idx: int, - resource_manager: Optional[ResourceManager] = None): - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if spec_resource_manager is None: - return None - - spec_tree_manager = spec_resource_manager.spec_tree_manager - if spec_tree_manager is not None: - spec_tree_manager.cur_draft_layer_idx = cur_draft_layer_idx - - def update_requests( - self, - sample_state: SampleState, - resource_manager: Optional[ResourceManager] = None) -> None: - """Update requests with sample state.""" - self.sampler.update_requests(sample_state, resource_manager) - - def process_decoded_tokens( - self, - draft_batch: list[LlmRequest], - draft_position: int, - cleanup_resources: bool = True) -> List[LlmRequest]: - """Process decoded tokens and determine which requests to continue processing.""" - new_requests = [] - for req in draft_batch: - target_model_req = self.req_id_to_old_request[req.py_request_id] - if target_model_req.state != LlmRequestState.GENERATION_IN_PROGRESS: - # This is a chunked prefill request and we have more prefill chunks - # to process. Defer adding draft tokens until the whole prompt is processed. - self.draft_seq_slot_manager.free_resources(req) - continue - - # Save tokens to accumulator instead of directly modifying target_model_req.py_draft_tokens - if req.py_request_id not in self.draft_tokens_accumulator: - self.draft_tokens_accumulator[ - req.py_request_id] = [0] * self.max_total_draft_tokens - self.draft_tokens_accumulator[req.py_request_id][ - draft_position - 1] = req.get_last_tokens(0) - - generation_logits = req.py_result.generation_logits # forwards Nones - if generation_logits is not None: - # generation_logits returns [beam_width, seq_length, vocab_size] - beam_width = generation_logits.size(0) - assert beam_width == 1, f"expected beam_width=1, got {beam_width}" - generation_logits.squeeze_(0) - # Transfer to CUDA if needed (chunked LogitsStorage stores on CPU) - if generation_logits.device.type == 'cpu': - generation_logits = generation_logits.to('cuda', - non_blocking=True) - target_model_req.py_draft_logits = generation_logits - - if req.state != LlmRequestState.GENERATION_COMPLETE and draft_position < target_model_req.py_draft_pages_allocated: - new_requests.append(req) - else: - if cleanup_resources: - self.draft_seq_slot_manager.free_resources(req) - - return new_requests - - def should_forward_draft_model(self, - scheduled_batch: ScheduledRequests) -> bool: - """ - Determine if the draft model should be forwarded for the given batch. - - Args: - scheduled_batch: The scheduled requests to check - - Returns: - bool: True if draft model should be forwarded, False otherwise - """ - all_disable_speculative_decoding = True - for request in scheduled_batch.all_requests(): - if not request.py_disable_speculative_decoding: - all_disable_speculative_decoding = False - break - if all_disable_speculative_decoding: - return False - for request in scheduled_batch.context_requests: - if request.is_first_context_chunk: - # DFlash needs the spec worker forward() to run during first-chunk - # prefill so _store_prefill_context can project and store the - # captured target hidden states as cross-attention context. - if self.spec_config.spec_dec_mode.is_dflash(): - return True - continue - return True - - for request in scheduled_batch.generation_requests: - if request.state == LlmRequestState.GENERATION_COMPLETE: - continue - - if request.max_beam_num_tokens - 1 >= self.draft_model_engine.max_seq_len: - continue - return True - - return False - - def _initialize_draft_tokens_for_target_inputs( - self, - scheduled_batch: ScheduledRequests, - target_inputs: Optional[SampleStateTensorsSpec] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None - ) -> Optional[SampleStateTensorsSpec]: - """ - Convert tensors for draft model processing. - - Args: - scheduled_batch: The scheduled requests - new_tensors_device: The device tensors to convert - - Returns: - SampleStateTensorsSpec: Converted tensors or None - """ - if target_inputs is None: - return None - # Get device from the new_tokens tensor - device = target_inputs.new_tokens.device - - new_tokens_lens = None - next_draft_tokens = None - has_draft_tokens = False - batch_size = target_inputs.new_tokens.shape[1] - # Iterate through generation requests and copy tokens based on accepted draft tokens - for request in scheduled_batch.all_requests(): - if request.state == LlmRequestState.GENERATION_IN_PROGRESS: - has_draft_tokens = True - - if has_draft_tokens: - # We already updated the target state, so the new_tokens_lens should be all ones. - new_tokens_lens = torch.ones(batch_size, - dtype=torch.int, - device=device) - new_tokens_lens += num_accepted_tokens_device - next_draft_tokens = torch.zeros(batch_size, - self.max_draft_len, - dtype=torch.int, - device=device) - target_inputs.new_tokens_lens = new_tokens_lens - target_inputs.next_draft_tokens = next_draft_tokens - return target_inputs - - @nvtx_range("_update_draft_tokens_for_target_inputs") - def _update_draft_tokens_for_target_inputs( - self, target_inputs: SampleStateTensorsSpec, - draft_tensors: Optional[torch.Tensor], draft_position: int, - draft_length: int, draft_batch: ScheduledRequests) -> None: - """ - Update target inputs with new draft tokens from sample state. - """ - if target_inputs.next_draft_tokens is None: - return - - draft_indices = [] - target_indices = [] - for req_idx, request in enumerate(draft_batch.all_requests()): - target_req = self.req_id_to_old_request[request.py_request_id] - if target_req.state != LlmRequestState.GENERATION_IN_PROGRESS: - # Skip prefill requests - continue - # Get the index of the draft/target tokens in the device tensor - draft_idx = req_idx if self.use_static_draft_loop else request.py_seq_slot - target_idx = target_req.py_seq_slot - draft_indices.append(draft_idx) - target_indices.append(target_idx) - - if len(draft_indices) == 0: - return - - device = draft_tensors.device - - # Create index tensors - draft_indices_tensor = torch.tensor(draft_indices, - dtype=torch.long, - pin_memory=prefer_pinned()).to( - device, non_blocking=True) - target_indices_tensor = torch.tensor(target_indices, - dtype=torch.long, - pin_memory=prefer_pinned()).to( - device, non_blocking=True) - - # Pre-slice draft tensors: [draft_length, batch_size] - draft_slice = draft_tensors[0:draft_length] - - # Gather all source data at once using single index_select kernel - # Result shape: [draft_length, num_requests] - gathered = draft_slice.index_select(1, draft_indices_tensor).to( - torch.int32) - - # Scatter to new_tokens using advanced indexing (single kernel) - # Shape: [draft_length, num_requests] -> [seq_len, batch_size, beam_width] - target_inputs.new_tokens[draft_position + 1:draft_position + - draft_length + 1, target_indices_tensor, - 0] = gathered - - # Scatter to next_draft_tokens using advanced indexing (single kernel) - # Shape: [num_requests, draft_length] -> [batch_size, max_draft_len] - target_inputs.next_draft_tokens[target_indices_tensor, - draft_position:draft_position + - draft_length] = gathered.t() - - def _setup_draft_batch_and_resources( - self, - scheduled_batch: ScheduledRequests) -> Optional[ScheduledRequests]: - """ - Setup draft batch and prepare resources. - - Args: - scheduled_batch: The scheduled requests - - Returns: - draft_batch or None if no batch - """ - - draft_batch = self._prepare_draft_batch(scheduled_batch) - if draft_batch.batch_size == 0: - return None - - self.req_id_to_old_request = { - req.py_request_id: req - for req in scheduled_batch.all_requests() - } - - for request in draft_batch.all_requests(): - target_model_req = self.req_id_to_old_request[request.py_request_id] - if target_model_req.is_context_init_state: - continue - target_model_req.py_draft_tokens = [0] * self.max_draft_len - - self.draft_seq_slot_manager.prepare_resources(draft_batch) - return draft_batch - - def process_static_draft_outputs(self, outputs: dict[str, torch.Tensor] - | tuple[torch.Tensor, SampleState], - draft_batch: ScheduledRequests) -> None: - """ - Process outputs from static draft loop, update target requests, and clean up resources. - - Args: - outputs: The outputs from the draft model - draft_batch: The draft batch that was processed - """ - - if isinstance(outputs, dict): - draft_tokens_host = outputs["new_draft_tokens"].cpu() - draft_logits = outputs["draft_logits"] - else: - draft_logits = outputs[0] - draft_tokens_host = outputs[1].host.new_tokens - outputs[1].sampler_event.synchronize() - - for req_idx, req in enumerate(draft_batch.all_requests()): - target_model_req = self.req_id_to_old_request[req.py_request_id] - if target_model_req.state != LlmRequestState.GENERATION_IN_PROGRESS: - # Chunked prefill request in progress; no need to append draft tokens - continue - target_model_req.py_draft_tokens = [] - py_draft_logits = [] - for token_idx in range(self.max_total_draft_tokens): - target_model_req.py_draft_tokens.append( - draft_tokens_host[token_idx][req_idx]) - py_draft_logits.append(draft_logits[token_idx][req_idx]) - - # The overlap scheduler doesn't support rejection sampling yet, so we don't update the py_draft_logits to get it fallback to greedy sampling. - if self.disable_overlap_scheduler: - target_model_req.py_draft_logits = torch.stack(py_draft_logits) - - def process_dynamic_draft_outputs( - self, - outputs: Any, - resource_manager: Optional[ResourceManager] = None, - cleanup_resources: bool = True) -> None: - """ - Process outputs from dynamic draft loop, update target requests, and clean up resources. - """ - self.update_requests(outputs, resource_manager) - - # Create accumulator for draft tokens and process them - self.process_decoded_tokens(outputs.requests, self.max_draft_len, - cleanup_resources) - - # Update py_draft_tokens after processing - for req_id, tokens in self.draft_tokens_accumulator.items(): - target_model_req = self.req_id_to_old_request[req_id] - target_model_req.py_draft_tokens = tokens - - @nvtx_range("_execute_draft_iteration") - def _execute_draft_iteration( - self, - draft_batch: ScheduledRequests, - resource_manager: ResourceManager, - previous_draft_state: Optional[SampleState], - cur_draft_layer_idx: int, - num_accepted_tokens_device: Optional[torch.Tensor] = None - ) -> Tuple[Any, Optional[SampleState]]: - self.update_cur_draft_layer_idx( - cur_draft_layer_idx, resource_manager - ) # Update the current draft layer index in the resource manager. - """Forward pass through the draft model.""" - outputs = self.forward_draft_model( - draft_batch, - resource_manager, - is_first_draft_token=False, - previous_tensors=previous_draft_state.device - if previous_draft_state else None, - num_accepted_tokens_device=num_accepted_tokens_device) - - if previous_draft_state is not None: - self.update_requests(previous_draft_state, resource_manager) - - if self.guided_decoder is not None: - self.guided_decoder.add_batch(draft_batch) - self.guided_decoder.execute(outputs['logits'], - d2t=outputs.get('d2t')) - - sample_state = self.sample_async(draft_batch, outputs, resource_manager) - self.update_request_states(draft_batch) - - return outputs, sample_state - - @nvtx_range("_process_previous_draft_results") - def _process_previous_draft_results( - self, resource_manager: Optional[ResourceManager] = None) -> None: - """ - Process the previous draft batch results. - This should be called after the current draft forward to enable overlap scheduling. - """ - if (self.previous_draft_batch is None - or self.previous_draft_outputs is None - or self.previous_scheduled_batch is None): - return - - # Save current req_id_to_old_request temporarily - current_req_id_to_old_request = self.req_id_to_old_request - - # Set req_id_to_old_request for the previous batch, - # this will be used in process_static_draft_outputs and process_dynamic_draft_outputs - self.req_id_to_old_request = { - req.py_request_id: req - for req in self.previous_scheduled_batch.all_requests() - } - - if self.use_static_draft_loop: - self.process_static_draft_outputs(self.previous_draft_outputs, - self.previous_draft_batch) - elif self.previous_draft_outputs is not None: - self.process_dynamic_draft_outputs(self.previous_draft_outputs, - cleanup_resources=False) - - self.req_id_to_old_request = current_req_id_to_old_request - - # Pad draft tokens to the max draft length for CUDA graph compatibility - self.pad_draft_tokens_for_cuda_graph(self.previous_scheduled_batch, - resource_manager) - - def cleanup_previous_draft_resources(self) -> None: - if self.previous_draft_batch is None: - return - - # Free draft_seq_slot_manager resources for all requests in the previous draft batch - for req in self.previous_draft_batch.all_requests(): - self.draft_seq_slot_manager.free_resources(req) - - def _execute_draft_loop( - self, - draft_batch: ScheduledRequests, - resource_manager: ResourceManager, - target_inputs: Optional[SampleStateTensorsSpec] = None, - num_draft_reqs: Optional[int] = None, - initial_draft_state: Optional[SampleState] = None - ) -> Optional[SampleState]: - """ - Execute the iterative draft loop. - - Args: - draft_batch: The draft batch to process - resource_manager: The resource manager - target_inputs: Optional target inputs to update (for overlap mode) - num_draft_reqs: Number of draft requests (for overlap mode) - initial_draft_state: The initial draft state from the first forward pass - - Returns: - The final sample state - """ - # Convert context requests to generation requests - for req in draft_batch.generation_requests: - req.py_is_first_draft = False - draft_batch.generation_requests = draft_batch.all_requests() - draft_batch.context_requests_chunking = [] - draft_batch.context_requests_last_chunk = [] - - previous_draft_state = initial_draft_state - # reset draft tokens accumulator - self.draft_tokens_accumulator = {} - # Generate remaining draft tokens iteratively - for i in range(self.max_draft_len - 1): - if len(draft_batch.generation_requests) == 0: - break - - _, sample_state = self._execute_draft_iteration( - draft_batch, resource_manager, previous_draft_state, i + 1) - - # Update target inputs if provided (for overlap mode) - if target_inputs is not None and num_draft_reqs is not None: - draft_tensors = sample_state and sample_state.device and sample_state.device.new_tokens - self._update_draft_tokens_for_target_inputs( - target_inputs, - draft_tensors, - draft_position=i + 1, - draft_length=1, - draft_batch=draft_batch) - - if sample_state is not None and previous_draft_state is not None: - new_requests = self.process_decoded_tokens( - previous_draft_state.requests, draft_position=i + 1) - else: - new_requests = [] - - draft_batch.generation_requests = new_requests - previous_draft_state = sample_state - - return previous_draft_state - - @nvtx_range("generate_draft_tokens_with_overlap") - def generate_draft_tokens_with_overlap( - self, - scheduled_batch: ScheduledRequests, - resource_manager: ResourceManager, - previous_tensors: Optional[SampleStateTensors], - target_inputs: Optional[SampleStateTensorsSpec], - num_accepted_tokens_device: Optional[torch.Tensor] = None) -> None: - """ - Generate draft tokens with overlap scheduling support. - - Args: - scheduled_batch: The scheduled requests - resource_manager: The resource manager - previous_tensors: Previous iteration tensors - guided_decoder: The guided decoder - - Returns: - Tuple[Optional[SampleStateTensorsSpec], Optional[SampleState]]: - - Updated target inputs or None - - Draft sample state or None - """ - import time - draft_start_time = time.time() - - self.disable_overlap_scheduler = False - if target_inputs is None: - return - - draft_batch = self._setup_draft_batch_and_resources(scheduled_batch) - if draft_batch is None: - return - - self._initialize_draft_tokens_for_target_inputs( - scheduled_batch, target_inputs, num_accepted_tokens_device) - - # Initial forward pass - self.update_cur_draft_layer_idx( - 0, resource_manager - ) # Update the current draft layer index in the resource manager. - outputs = self.forward_draft_model( - draft_batch, - resource_manager, - is_first_draft_token=True, - previous_tensors=previous_tensors, - num_accepted_tokens_device=num_accepted_tokens_device) - - # Process previous draft results after current forward pass - # This enables overlap scheduling: process old batch while new batch is prepared - self._process_previous_draft_results(resource_manager) - - num_draft_reqs = len(draft_batch.all_requests()) - if self.use_static_draft_loop: - # Only update target inputs, cleanup will be done in executor loop - self._update_draft_tokens_for_target_inputs( - target_inputs, - outputs["new_draft_tokens"], - draft_position=0, - draft_length=self.max_draft_len, - draft_batch=draft_batch) - - new_tokens_host = outputs["new_draft_tokens"].to(device='cpu', - non_blocking=True) - sampler_event = torch.cuda.Event() - sampler_event.record() - - sample_state = SampleState( - requests=draft_batch.all_requests(), - device=SampleStateTensors( - new_tokens=outputs["new_draft_tokens"]), - host=SampleStateTensors(new_tokens=new_tokens_host), - sampler_event=sampler_event) - - # Store current batch for processing in next iteration - self.previous_draft_batch = draft_batch - self.previous_draft_outputs = (outputs["draft_logits"], - sample_state) - self.previous_scheduled_batch = scheduled_batch - - return - - # Handle guided decoder and sampling for non-static loop - if self.guided_decoder is not None: - self.guided_decoder.add_batch(draft_batch) - self.guided_decoder.execute(outputs['logits'], - d2t=outputs.get('d2t')) - draft_sample_state = self.sample_async(draft_batch, outputs, - resource_manager) - - # Update target inputs with first iteration results - draft_tensors = draft_sample_state and draft_sample_state.device and draft_sample_state.device.new_tokens - self._update_draft_tokens_for_target_inputs(target_inputs, - draft_tensors, - draft_position=0, - draft_length=1, - draft_batch=draft_batch) - - self.update_request_states(draft_batch) - - # Execute the iterative draft loop - previous_draft_state = self._execute_draft_loop(draft_batch, - resource_manager, - target_inputs, - num_draft_reqs, - draft_sample_state) - - # Store current batch for processing in next iteration - self.previous_draft_batch = draft_batch - self.previous_draft_outputs = previous_draft_state - self.previous_scheduled_batch = scheduled_batch - - # Record draft latency for stats - draft_end_time = time.time() - self.last_draft_latency_ms = (draft_end_time - draft_start_time) * 1e3 - - @nvtx_range("prepare_draft_tokens") - def prepare_draft_tokens( - self, - scheduled_requests: ScheduledRequests, - resource_manager: Optional[ResourceManager] = None, - ) -> None: - """ - Prepare draft tokens for the scheduled requests. - - Args: - scheduled_requests: The scheduled requests for this iteration - resource_manager: The resource manager for this iteration - """ - import time - draft_start_time = time.time() - - self.disable_overlap_scheduler = True - if not self.draft_model_engine: - raise ValueError("Draft model engine is not set") - - if resource_manager is None: - raise ValueError("Resource manager is required") - - try: - draft_batch = self._setup_draft_batch_and_resources( - scheduled_requests) - if draft_batch is None: - return - - self.update_cur_draft_layer_idx( - 0, resource_manager - ) # Update the current draft layer index in the resource manager. - # Initial forward pass. May do the complete drafting loop - # if use_static_draft_loop is set. - outputs = self.forward_draft_model(draft_batch, - resource_manager, - is_first_draft_token=True) - - if self.use_static_draft_loop: - self.process_static_draft_outputs(outputs, draft_batch) - # Clean up draft_seq_slot_manager resources - for req in draft_batch.all_requests(): - self.draft_seq_slot_manager.free_resources(req) - # Record draft latency before returning - draft_end_time = time.time() - self.last_draft_latency_ms = (draft_end_time - - draft_start_time) * 1e3 - return - - if self.guided_decoder is not None: - self.guided_decoder.add_batch(draft_batch) - self.guided_decoder.execute(outputs['logits'], - d2t=outputs.get('d2t')) - sample_state = self.sample_async(draft_batch, outputs, - resource_manager) - self.update_request_states(draft_batch) - - # Execute the iterative draft loop - previous_draft_state = self._execute_draft_loop( - draft_batch, resource_manager, None, None, sample_state) - - # Final cleanup - if previous_draft_state is not None: - self.process_dynamic_draft_outputs(previous_draft_state) - - # Update py_draft_tokens after the loop completes - for req_id, tokens in self.draft_tokens_accumulator.items(): - target_model_req = self.req_id_to_old_request[req_id] - target_model_req.py_draft_tokens = tokens - - # Record draft latency for stats - draft_end_time = time.time() - self.last_draft_latency_ms = (draft_end_time - - draft_start_time) * 1e3 - except Exception as e: - traceback.print_exc() - error_msg = str(e) - logger.error(f"Encountered an error in decode: {error_msg}") - raise e diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 3dc12a669aa6..a017077ef52b 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -177,7 +177,7 @@ def all_rank_num_seqs(self): @all_rank_num_seqs.setter def all_rank_num_seqs(self, value: List[int]): self._all_rank_num_seqs = value - if self.spec_dec_mode.is_mtp_eagle_one_model(): + if self.spec_dec_mode.is_mtp_eagle(): self.subseq_all_rank_num_tokens = value def prepare(self): @@ -194,7 +194,7 @@ def prepare(self): # while MTP Eagle worker uses (max_draft_len + 1) input tokens in the 1st draft # forward and only one input token in the following draft forward. # This num_tokens is used to set the all_rank_num_tokens for attention dp. - if not self.spec_dec_mode.is_mtp_eagle_one_model(): + if not self.spec_dec_mode.is_mtp_eagle(): self.num_tokens -= self.num_generations if self.mtp_hidden_states_manager is not None: # MTP vanilla or use relaxed acceptance diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 62189035cafb..8bd831cfcb90 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -14,7 +14,6 @@ from ..pyexecutor.guided_decoder import GuidedDecoder from ..pyexecutor.sampler import TorchSampler -from ..pyexecutor.seq_slot_manager import SeqSlotManager from ..speculative.interface import SpecMetadata from .dflash import DFlashSpecMetadata, DFlashWorker from .draft_target import (DraftTargetOneModelSampler, @@ -24,9 +23,8 @@ from .eagle3 import (Eagle3OneModelDynamicTreeResourceManager, Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, - Eagle3SpecMetadata, MTPEagleWorker) + MTPEagleWorker) from .eagle3_dynamic_tree import Eagle3OneModelDynamicTreeWorker -from .model_drafter import ModelDrafter from .mtp import MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata, MTPWorker from .mtp_dynamic_tree import (MTPEagleDynamicTreeResourceManager, MTPEagleDynamicTreeWorker) @@ -81,7 +79,6 @@ def get_spec_metadata(spec_config, max_num_requests, max_num_tokens, spec_resource_manager=None, - is_draft_model=False, max_seq_len=262144, num_seq_slots=None): use_rejection_sampling = getattr(spec_config, "use_rejection_sampling", @@ -94,8 +91,8 @@ def get_spec_metadata(spec_config, # Draft-model vocab size, used to gate the d2t-expanded full_draft_probs # buffer allocation (see SpecMetadata.prepare_rejection_sampling_buffers). draft_vocab_size = _get_draft_vocab_size(spec_config, vocab_size) - if spec_config.spec_dec_mode.is_mtp_eagle_one_model(): - # MTP Eagle one-model reuses Eagle3 one-model metadata for the + if spec_config.spec_dec_mode.is_mtp_eagle(): + # MTP Eagle reuses Eagle3 one-model metadata for the # unified worker/sampler/slot_ids plumbing, but skips per-layer # hidden-state capture: the worker feeds the target model's # hidden_states directly into the MTP layer, so we leave @@ -129,42 +126,7 @@ def get_spec_metadata(spec_config, vocab_size=vocab_size, draft_vocab_size=draft_vocab_size, ) - if spec_config.spec_dec_mode.is_mtp_eagle(): - return Eagle3SpecMetadata( - max_draft_len=spec_config.max_draft_len, - max_total_draft_tokens=spec_config.tokens_per_gen_step - 1, - spec_dec_mode=spec_config.spec_dec_mode, - max_num_requests=max_num_requests, - num_layers=model_config.num_hidden_layers, - hidden_size=model_config.hidden_size, - max_num_tokens=max_num_tokens, - dtype=model_config.torch_dtype, - is_draft_model=is_draft_model, - eagle3_resource_manager=spec_resource_manager, - layers_to_capture=None, - is_mtp_eagle=True, - ) if spec_config.spec_dec_mode.is_eagle3(): - effective_dynamic_tree = _is_effective_dynamic_tree(spec_config) - return Eagle3SpecMetadata( - max_draft_len=spec_config.max_draft_len, - max_total_draft_tokens=spec_config.tokens_per_gen_step - 1, - spec_dec_mode=spec_config.spec_dec_mode, - max_num_requests=max_num_requests, - num_layers=model_config.num_hidden_layers, - hidden_size=model_config.hidden_size, - max_num_tokens=max_num_tokens, - dtype=model_config.torch_dtype, - is_draft_model=is_draft_model, - eagle3_resource_manager=spec_resource_manager, - layers_to_capture=spec_config.eagle3_layers_to_capture, - is_mtp_eagle=False, - eagle_choices=spec_config.eagle_choices, - is_spec_dec_tree=spec_config.eagle_choices is not None - or effective_dynamic_tree, - is_spec_dec_dynamic_tree=effective_dynamic_tree, - ) - if spec_config.spec_dec_mode.is_eagle3_one_model(): return Eagle3OneModelSpecMetadata( max_draft_len=spec_config.max_draft_len, max_total_draft_tokens=spec_config.tokens_per_gen_step - 1, @@ -222,7 +184,7 @@ def get_spec_metadata(spec_config, vocab_size=vocab_size, draft_vocab_size=draft_vocab_size, ) - if spec_config.spec_dec_mode.is_draft_target_one_model(): + if spec_config.spec_dec_mode.is_draft_target(): return DraftTargetOneModelSpecMetadata( max_draft_len=spec_config.max_draft_len, max_total_draft_tokens=spec_config.max_total_draft_tokens, @@ -255,8 +217,7 @@ def get_spec_metadata(spec_config, sa_manager=spec_resource_manager, max_matching_ngram_size=spec_config.max_matching_ngram_size, ) - if spec_config.spec_dec_mode.is_draft_target() or \ - spec_config.spec_dec_mode.is_ngram() or \ + if spec_config.spec_dec_mode.is_ngram() or \ spec_config.spec_dec_mode.is_user_provided(): return SpecMetadata( max_draft_len=spec_config.max_draft_len, @@ -277,7 +238,7 @@ def get_mtp_hidden_size(model_config) -> int: return hidden_size -def get_spec_resource_manager(model_engine, draft_model_engine=None): +def get_spec_resource_manager(model_engine): spec_config = model_engine.spec_config if spec_config is None: return None @@ -286,7 +247,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): max_seq_len = model_engine.max_seq_len max_num_tokens = model_engine.max_num_tokens spec_dec_mode = spec_config.spec_dec_mode - if spec_dec_mode.is_mtp_eagle_one_model(): + if spec_dec_mode.is_mtp_eagle(): sa_manager = None sa_cfg = getattr(spec_config, 'sa_config', None) if sa_cfg is not None: @@ -329,11 +290,10 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): max_num_requests, sa_manager=sa_manager, ) - if spec_dec_mode.is_eagle3_one_model() and _is_effective_dynamic_tree( - spec_config): + if spec_dec_mode.is_eagle3() and _is_effective_dynamic_tree(spec_config): return Eagle3OneModelDynamicTreeResourceManager(spec_config, max_num_requests) - if spec_dec_mode.is_eagle3_one_model(): + if spec_dec_mode.is_eagle3(): sa_manager = None sa_cfg = getattr(spec_config, 'sa_config', None) if sa_cfg is not None: @@ -348,16 +308,6 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): max_num_tokens, sa_manager=sa_manager, ) - if spec_dec_mode.is_eagle3() or spec_dec_mode.is_mtp_eagle(): - assert draft_model_engine is not None, "Draft model engine is required for Eagle3 and MTP Eagle two model flow." - return Eagle3ResourceManager( - spec_config, - draft_model_engine.model.config.torch_dtype, - model_config.hidden_size, - max_num_requests, - max_seq_len, - max_num_tokens, - ) if spec_dec_mode.is_save_hidden_states(): return SaveHiddenStatesResourceManager( spec_config, @@ -385,33 +335,28 @@ def get_spec_decoder( sampler_args: TorchSampler.Args, spec_config: "DecodingBaseConfig", ): - if spec_config.spec_dec_mode.is_mtp_eagle_one_model(): - # MTP Eagle one-model now uses the same sampler as Eagle3 one-model. + if spec_config.spec_dec_mode.is_mtp_eagle(): + # MTP Eagle now uses the same sampler as Eagle3. return Eagle3OneModelSampler(sampler_args, spec_config=spec_config) if spec_config.spec_dec_mode.is_mtp_vanilla(): nextn = spec_config.max_draft_len if getattr(spec_config, "use_dynamic_tree", False): nextn = spec_config.max_total_draft_tokens return MTPSampler(sampler_args, nextn=nextn) - if spec_config.spec_dec_mode.is_eagle3( - ) or spec_config.spec_dec_mode.is_mtp_eagle(): - # TorchSampler handles Eagle3 gracefully, by integrating d2t into the sampling process - return TorchSampler(sampler_args) - if spec_config.spec_dec_mode.is_eagle3_one_model(): + if spec_config.spec_dec_mode.is_eagle3(): return Eagle3OneModelSampler(sampler_args, spec_config=spec_config) if spec_config.spec_dec_mode.is_parallel_draft(): return MTPSampler(sampler_args, nextn=spec_config.tokens_per_gen_step - 1) if spec_config.spec_dec_mode.is_sa(): return SASampler(sampler_args, max_draft_len=spec_config.max_draft_len) - if spec_config.spec_dec_mode.is_draft_target_one_model(): + if spec_config.spec_dec_mode.is_draft_target(): return DraftTargetOneModelSampler(sampler_args) raise ValueError( f"Unsupported speculative decoding mode: {spec_config.spec_dec_mode}") def get_spec_drafter(model_engine, - draft_model_engine, sampler, spec_resource_manager, guided_decoder: Optional[GuidedDecoder] = None): @@ -422,19 +367,6 @@ def get_spec_drafter(model_engine, if spec_config.spec_dec_mode.is_user_provided(): return spec_config.drafter - max_num_requests = model_engine.batch_size - if spec_config.spec_dec_mode.is_draft_target( - ) or spec_config.spec_dec_mode.is_eagle3( - ) or spec_config.spec_dec_mode.is_mtp_eagle(): - return ModelDrafter(spec_config, - draft_model_engine, - spec_config.max_draft_len, - spec_config.tokens_per_gen_step - 1, - SeqSlotManager(max_num_requests), - sampler, - spec_resource_manager=spec_resource_manager, - guided_decoder=guided_decoder) - if spec_config.spec_dec_mode.is_ngram(): return NGramDrafter(spec_config, spec_resource_manager) @@ -442,11 +374,11 @@ def get_spec_drafter(model_engine, def get_num_spec_layers(spec_config): - if spec_config.spec_dec_mode.is_mtp_eagle_one_model(): + if spec_config.spec_dec_mode.is_mtp_eagle(): return 1 if spec_config.spec_dec_mode.is_mtp_vanilla(): return spec_config.num_nextn_predict_layers - if spec_config.spec_dec_mode.is_eagle3_one_model(): + if spec_config.spec_dec_mode.is_eagle3(): num_draft_hidden_layers = spec_config._num_draft_hidden_layers return num_draft_hidden_layers if num_draft_hidden_layers is not None else 1 return 0 @@ -479,7 +411,7 @@ def get_spec_worker(spec_config, model_config, use_separate_draft_kv_cache, mapping=mapping) - if spec_dec_mode.is_mtp_eagle_one_model(): + if spec_dec_mode.is_mtp_eagle(): if getattr(spec_config, 'use_dynamic_tree', False): return MTPEagleDynamicTreeWorker(spec_config, model_config, @@ -489,7 +421,7 @@ def get_spec_worker(spec_config, model_config, use_separate_draft_kv_cache, mapping=mapping) - if spec_dec_mode.is_eagle3_one_model(): + if spec_dec_mode.is_eagle3(): if _is_effective_dynamic_tree(spec_config): return Eagle3OneModelDynamicTreeWorker(spec_config, mapping, use_separate_draft_kv_cache) @@ -505,7 +437,7 @@ def get_spec_worker(spec_config, return DSparkWorker(spec_config, mapping, use_separate_draft_kv_cache) if spec_dec_mode.is_sa(): return SAWorker(spec_config, model_config) - if spec_dec_mode.is_draft_target_one_model(): + if spec_dec_mode.is_draft_target(): return DraftTargetOneModelWorker(spec_config, mapping, use_separate_draft_kv_cache) return None diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d09e0b576b5b..e01a0fa50bed 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1769,13 +1769,9 @@ class DecodingBaseConfig(StrictBaseModel): "in a future release. Non-greedy sampling is now auto-detected per " "request; this flag no longer has any effect.") - # If set, drafting is allowed to use chain drafter. - _allow_chain_drafter: bool = PrivateAttr(True) - # If set, drafting uses greedy sampling, irrespective of sampling parameters. - _allow_greedy_draft_tokens: bool = PrivateAttr(True) # Internal: record decoding_type alias used during parsing (for warnings). _decoding_type_alias: Optional[str] = PrivateAttr(default=None) - # If set, drafting will use separate KV cache in one-model speculative decoding. + # If set, drafting will use separate KV cache in one-engine speculative decoding. _allow_separate_draft_kv_cache: bool = PrivateAttr(True) # Internal: true when draft_len_schedule was auto-translated from max_concurrency. _translated_from_max_concurrency: bool = PrivateAttr(False) @@ -2056,11 +2052,6 @@ class EagleDecodingConfig(DecodingBaseConfig): _num_draft_hidden_layers: Optional[int] = PrivateAttr(default=None) max_non_leaves_per_layer: Optional[int] = Field( default=None, description="The number of non-leaves in each layer.") - eagle3_one_model: Optional[bool] = Field( - default=True, - description= - "Whether to use the faster one-model implementation (draft as submodule) or the two-model implementation." - ) eagle3_layers_to_capture: Optional[Set[int]] = Field( default=None, description= @@ -2090,10 +2081,6 @@ def validate_eagle_choices(cls, v): def validate_eagle_config(self) -> 'EagleDecodingConfig': if self.max_draft_len is None or self.max_draft_len == 0: raise ValueError("max_draft_len must be > 0 for Eagle") - if not self.eagle3_one_model: - logger.warning( - "Eagle3 2-model is deprecated and will be removed in release 1.4." - ) self.num_eagle_layers = self.max_draft_len @@ -2195,8 +2182,6 @@ def check_eagle_choices(self): def spec_dec_mode(self): from tensorrt_llm._torch.speculative.interface import \ SpeculativeDecodingMode as TorchSpeculativeDecodingMode - if self.eagle3_one_model: - return TorchSpeculativeDecodingMode.EAGLE3_ONE_MODEL return TorchSpeculativeDecodingMode.EAGLE3 @functools.cached_property @@ -2452,7 +2437,6 @@ def supports_backend(self, backend: str) -> bool: class DraftTargetDecodingConfig(DecodingBaseConfig): decoding_type: Literal["Draft_Target"] = Field(default="Draft_Target") - _draft_target_one_model: bool = PrivateAttr(True) @model_validator(mode="after") def validate_draft_target_config(self): @@ -2471,8 +2455,6 @@ def supports_backend(self, backend: str) -> bool: def spec_dec_mode(self): from tensorrt_llm._torch.speculative.interface import \ SpeculativeDecodingMode as TorchSpeculativeDecodingMode - if self._draft_target_one_model: - return TorchSpeculativeDecodingMode.DRAFT_TARGET_ONE_MODEL return TorchSpeculativeDecodingMode.DRAFT_TARGET @@ -2498,11 +2480,6 @@ class MTPDecodingConfig(DecodingBaseConfig): description= "Force vanilla MTP mode (sequential MTP layers). When False, uses EAGLE-style MTP for single-layer checkpoints." ) - mtp_eagle_one_model: bool = Field( - default=True, - description= - "When using EAGLE-style MTP, use faster one-model implementation (drafter as submodule) vs two-model." - ) use_dynamic_tree: bool = Field( default=False, @@ -2603,28 +2580,17 @@ def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len # linear chain return self - @model_validator(mode="after") - def log_two_model_deprecation_warning(self): - if not self.mtp_eagle_one_model: - logger.warning( - "2-model style MTP is deprecated and will be removed in release 1.4." - ) - return self - def supports_backend(self, backend: str) -> bool: return backend in ("pytorch", "_autodeploy") @property def num_capture_layers(self) -> int: - # MTP_EAGLE (two-model) feeds captured target hidden states into the - # separate draft engine, so the shared Eagle3ResourceManager must - # allocate a hidden_states buffer for it. MTP_EAGLE_ONE_MODEL passes - # the target model's hidden_states straight to the MTP layer - # (see Eagle3OneModelWorker.prepare_1st_drafter_inputs / _run_draft_forward, - # both gated on self.is_mtp_eagle), so no capture buffer is needed - # and we should skip allocation to avoid disabling post-MLP/MoE - # fusion via the layer-capture hook. - return 1 if self.spec_dec_mode.is_mtp_eagle() else 0 + # MTP_EAGLE passes the target model's hidden_states straight to the MTP + # layer (see Eagle3OneModelWorker.prepare_1st_drafter_inputs / + # _run_draft_forward, both gated on self.is_mtp_eagle), so no capture + # buffer is needed. Skipping allocation also avoids disabling the + # post-MLP/MoE fusion via the layer-capture hook. + return 0 @property def spec_dec_mode(self): @@ -2634,9 +2600,7 @@ def spec_dec_mode(self): # num_nextn_predict_layers is set from the model's pretrained config by # update_spec_config_from_model_config. Treat None (before model load) as 1. n = self.num_nextn_predict_layers if self.num_nextn_predict_layers is not None else 1 - if n == 1 and not self.use_mtp_vanilla and self.mtp_eagle_one_model: - return TorchSpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL - elif n == 1 and not self.use_mtp_vanilla and not self.mtp_eagle_one_model: + if n == 1 and not self.use_mtp_vanilla: return TorchSpeculativeDecodingMode.MTP_EAGLE return TorchSpeculativeDecodingMode.MTP @@ -5315,10 +5279,9 @@ def validate_speculative_config(self): self.speculative_config = Eagle3DecodingConfig(**eagle_data) if self.speculative_config.use_rejection_sampling: - # Supported paths: Eagle3 one-model, MTP-Eagle one-model, - # vanilla MTP, PARD, DFlash, DraftTarget one-model. Classify by - # spec-dec mode; MTP-Eagle one-model shares the Eagle3 one-model - # worker/metadata/sampler, so it rides the same unified + # Supported paths: Eagle3, MTP-Eagle, vanilla MTP, PARD, DFlash, + # DraftTarget. Classify by spec-dec mode; MTP-Eagle shares the + # Eagle3 worker/metadata/sampler, so it rides the same unified # production/acceptance path. # # Every supported path runs a neural draft head that yields a @@ -5334,20 +5297,20 @@ def validate_speculative_config(self): is_supported = (isinstance(self.speculative_config, Eagle3DecodingConfig) or spec_mode == TorchSpeculativeDecodingMode.MTP - or spec_mode.is_mtp_eagle_one_model() + or spec_mode.is_mtp_eagle() or spec_mode.is_pard() or spec_mode.is_dflash() or spec_mode.is_dspark() - or spec_mode.is_draft_target_one_model()) + or spec_mode.is_draft_target()) # Combinations that break the proposal-distribution invariant. # SA is gated for every method; relaxed / parallel / guided gates # apply only to the newly wired methods (vanilla MTP, PARD, - # DFlash, DraftTarget one-model). + # DFlash, DraftTarget). is_new_rejection_method = ( spec_mode == TorchSpeculativeDecodingMode.MTP or spec_mode.is_pard() or spec_mode.is_dflash() or spec_mode.is_dspark() - or spec_mode.is_draft_target_one_model()) + or spec_mode.is_draft_target()) # Plain tensor parallelism is supported (the draft path # all-gathers vocab-sharded draft logits before rejection, see # SpecWorkerBase.maybe_gather_sharded_draft_logits). @@ -5391,9 +5354,8 @@ def validate_speculative_config(self): if not is_supported: reasons.append( f"spec mode {spec_mode.name} is not a supported " - "rejection-sampling path (supported: the Eagle3 " - "one-model path, vanilla MTP, PARD, DFlash, and " - "DraftTarget one-model)") + "rejection-sampling path (supported: Eagle3, " + "vanilla MTP, PARD, DFlash, and DraftTarget)") if rs_sa_active: reasons.append("SA (sa_config) is active") if rs_dynamic_tree_topk_unsupported: @@ -5438,8 +5400,7 @@ def validate_speculative_config(self): "Install flashinfer-python or set " "use_rejection_sampling=False.") # The enabled paths (Eagle3, vanilla MTP, PARD, DFlash, - # DraftTarget one-model) share the SpecMetadata slot-indexed - # dispatch. + # DraftTarget) share the SpecMetadata slot-indexed dispatch. if isinstance(self.speculative_config, PARDDecodingConfig): assert self.speculative_config.max_draft_len > 0, "PARD max_draft_len must be > 0" @@ -5573,8 +5534,6 @@ def _dspark_get(key, top_level_key): elif isinstance(self.speculative_config, DraftTargetDecodingConfig): assert self.speculative_config.max_draft_len > 0 assert self.speculative_config.speculative_model is not None, "Draft model must be specified." - if self.backend == "_autodeploy": - self.speculative_config._draft_target_one_model = False # If speculative_config.draft_len_schedule is provided, cuda_graph_config.enable_padding is automatically set to True. # Also we add the draft_len_schedule keys into batch_sizes for better cuda graph coverage in dynamic draft length. diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index fb87dbc2b432..95c070ee696d 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1702,13 +1702,6 @@ "kind": "categorical", "path": "speculative_config.eagle3_model_arch" }, - { - "allowed_values": [], - "annotation": "Optional[bool]", - "converter": "", - "kind": "value", - "path": "speculative_config.eagle3_one_model" - }, { "allowed_values": [], "annotation": "Optional[List[List[int]]]", @@ -1853,13 +1846,6 @@ "kind": "value", "path": "speculative_config.medusa_choices" }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "speculative_config.mtp_eagle_one_model" - }, { "allowed_values": [], "annotation": "Optional[int]", diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index 23baf4fc8495..530b98f4883d 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -39,10 +39,10 @@ "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_auto_dtype[True-True-True-False]": 183.9857738095238, "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_auto_dtype[True-True-True-True]": 180.99827380952382, "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_beam_search": 832.6288461538461, - "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_eagle3[eagle3_one_model=True-overlap_scheduler=True]": 175.0373370786517, + "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_eagle3[overlap_scheduler=True]": 175.0373370786517, "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_gen_only_spec_dec": 223.4482705882353, "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding[xgrammar]": 76.99953488372093, - "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=True]": 154.3497294117647, + "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar]": 154.3497294117647, "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_kv_cache_v2_nixl_python": 517.0826703296703, "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_multi_instance[GSM8K]": 177.79596629213484, "accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_multi_instance[MMLU]": 104.5509659090909, @@ -78,8 +78,8 @@ "accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-1]": 161.8740559566787, "accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-4]": 150.97673714285713, "accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-True-1]": 319.70320000000004, - "accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3_one_model[flashinfer]": 198.04985, - "accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3_one_model[trtllm]": 184.09691176470588, + "accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3[flashinfer]": 198.04985, + "accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3[trtllm]": 184.09691176470588, "accuracy/test_llm_api_autodeploy.py::TestMiniMaxM2::test_finegrained_fp8": 612.8464166666666, "accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[Qwen_QwQ-32B-False]": 282.35970000000003, "accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[deepseek-ai_DeepSeek-R1-0528-True]": 801.177, @@ -342,16 +342,16 @@ "accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8]": 562.7472195121951, "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dflash": 620.23102734375, "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dummy_load_format": 89.48106761565836, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[cutlass-one_model-overlap_scheduler]": 1852.298488095238, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-one_model-overlap_scheduler]": 1000.3237142857143, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-no_overlap_scheduler]": 834.5935714285715, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-overlap_scheduler]": 445.1501655629139, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-no_overlap_scheduler]": 436.1485, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-overlap_scheduler]": 461.6794709677419, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-no_overlap_scheduler]": 999.9191666666667, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-overlap_scheduler]": 601.7484152046783, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-no_overlap_scheduler]": 853.4979000000001, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-overlap_scheduler]": 568.9865187500001, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[cutlass-overlap_scheduler]": 1852.298488095238, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-overlap_scheduler]": 1000.3237142857143, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-no_overlap_scheduler]": 834.5935714285715, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-overlap_scheduler]": 445.1501655629139, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-no_overlap_scheduler]": 436.1485, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-overlap_scheduler]": 461.6794709677419, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-no_overlap_scheduler]": 999.9191666666667, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-overlap_scheduler]": 601.7484152046783, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-no_overlap_scheduler]": 853.4979000000001, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-overlap_scheduler]": 568.9865187500001, "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache-one_model]": 277.7859, "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache-two_model]": 323.603125, "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache-one_model]": 1362.0695, @@ -423,8 +423,8 @@ "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_chunked_prefill[use_temperature=True-attn_backend=TRTLLM]": 62.18363940520447, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_dflash": 146.09544150943395, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_dummy_load_format": 21.636074626865675, - "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-eagle3_one_model=True-overlap_scheduler=True]": 1021.8135882352941, - "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=True-eagle3_one_model=True-overlap_scheduler=True]": 1006.1250262172285, + "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-overlap_scheduler=True]": 1021.8135882352941, + "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=True-overlap_scheduler=True]": 1006.1250262172285, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=FLASHINFER-torch_compile=False]": 63.8741, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=FLASHINFER-torch_compile=True]": 104.64133333333334, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=TRTLLM-torch_compile=False]": 58.18021691176471, @@ -447,8 +447,8 @@ "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding[llguidance]": 47.5102, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding[xgrammar]": 49.10935249042146, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_4gpus[llguidance]": 47.93818181818182, - "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=True]": 46.6907, - "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=True]": 76.29452325581396, + "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance]": 46.6907, + "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar]": 76.29452325581396, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_ngram[llguidance]": 49.70666666666666, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_ngram[xgrammar]": 36.47, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_nvfp4_kv[v2_kv_cache=False-attn_backend=TRTLLM-torch_compile=False]": 60.9993865248227, @@ -541,10 +541,10 @@ "accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4[TRTLLM]": 234.99062592592594, "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash": 175.31245098039216, "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dummy_load_format": 42.3696968503937, - "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=False]": 107.64960153256705, - "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=True]": 116.99988888888889, - "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=True-enable_draft_len_schedule=False]": 102.77366666666667, - "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=True-enable_max_concurrency=False-enable_draft_len_schedule=False]": 189.7652222222222, + "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=False]": 107.64960153256705, + "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=True]": 116.99988888888889, + "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=False-enable_max_concurrency=True-enable_draft_len_schedule=False]": 102.77366666666667, + "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=True-enable_max_concurrency=False-enable_draft_len_schedule=False]": 189.7652222222222, "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency]": 217.97051320422534, "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales_early_first_token_response": 128.8759704797048, "accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_w4a8_mxfp4[fp8-latency]": 44.27991134751773, @@ -663,7 +663,7 @@ "disaggregated/test_workers.py::test_workers_kv_cache_aware_router[TinyLlama-1.1B-Chat-v1.0]": 133.0291614730878, "disaggregated/test_workers.py::test_workers_kv_cache_events[TinyLlama-1.1B-Chat-v1.0]": 79.10274643874644, "examples/test_ad_guided_decoding.py::test_autodeploy_guided_decoding_main_json": 109.94750187265917, - "examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_one_model_acceptance_rate[flashinfer-torch-simple]": 77.74328782287823, + "examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_acceptance_rate[flashinfer-torch-simple]": 77.74328782287823, "examples/test_ad_speculative_decoding.py::test_eagle_model_with_weights": 3.1745095541401276, "examples/test_ad_speculative_decoding.py::test_nemotron_mtp_model_with_weights": 798.8252090032154, "examples/test_deepseek_v4_pro.py::test_short_token_boundary_smoke": 871.5865, diff --git a/tests/integration/defs/.test_durations_aws_dfw b/tests/integration/defs/.test_durations_aws_dfw index 5eb61f0475c0..8b75fbd08fdc 100644 --- a/tests/integration/defs/.test_durations_aws_dfw +++ b/tests/integration/defs/.test_durations_aws_dfw @@ -29,7 +29,7 @@ "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=False-attention_dp=True-cuda_graph=False-overlap_scheduler=False-low_precision_combine=True-torch_compile=False]": 270.2675936790183, "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False]": 386.7030587599729, "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=WIDEEP]": 306.27529788704123, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-overlap_scheduler]": 736.9224306879914, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-overlap_scheduler]": 736.9224306879914, "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-trtllm-fp8]": 1132.6229840770247, "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-tp4-trtllm-auto]": 680.1839753920212, "accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_block_reuse[TEP4_ADP]": 750.9164859569864, @@ -37,7 +37,7 @@ "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[tp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True]": 414.5880687360186, "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_python_scheduler[ep4-mtp_nextn=2]": 323.5369381190976, "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False]": 271.8181478260085, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-overlap_scheduler]": 806.110374048003, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-overlap_scheduler]": 806.110374048003, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=False]": 156.74686194607057, "accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_online_eplb[moe_backend=CUTLASS]": 543.1446078500012, "accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_online_eplb[moe_backend=TRTLLM]": 768.4673424640205, @@ -47,12 +47,12 @@ "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False]": 422.0055055040866, "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False]": 494.63863472105004, "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus_online_eplb[fp8kv=True-moe_backend=TRTLLM]": 432.65749452402815, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-overlap_scheduler]": 491.4517015229212, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-overlap_scheduler]": 491.4517015229212, "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False]": 132.3602316979086, "accuracy/test_llm_api_pytorch.py::TestQwen3NextInstruct::test_nvfp4[tp4ep4_adp_off-trtllm]": 2620.8175430910196, "test_unittests.py::test_unittests_v2[unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm]": 1953.8634579040809, "accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-ep4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True]": 295.83772118296474, - "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-overlap_scheduler]": 550.6863905609935, + "accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-overlap_scheduler]": 550.6863905609935, "accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_block_reuse[TEP4_ADP_MTP]": 759.9382735569961, "accuracy/test_llm_api_pytorch.py::TestQwen3NextInstruct::test_nvfp4[tp4ep1-cutlass]": 2567.5702897430165, "accuracy/test_llm_api_pytorch.py::TestQwen3NextInstruct::test_nvfp4[tp4ep4_adp_off-cutlass]": 2415.312640499964, diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 46d4f51a712c..28cdaffb164d 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -796,14 +796,12 @@ def test_ngram(self): @pytest.mark.skip_less_device(2) @skip_pre_hopper @parametrize_with_ids("overlap_scheduler", [True, False]) - @parametrize_with_ids("eagle3_one_model", [True, False]) - def test_eagle3(self, overlap_scheduler, eagle3_one_model): + def test_eagle3(self, overlap_scheduler): speculative_decoding_config = { "decoding_type": "Eagle", "max_draft_len": 4, "speculative_model": f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B", - "eagle3_one_model": eagle3_one_model } ctx_server_config = { "disable_overlap_scheduler": @@ -859,7 +857,6 @@ def test_gen_only_spec_dec(self): "max_draft_len": 4, "speculative_model": f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B", - "eagle3_one_model": True, } ctx_server_config = { "disable_overlap_scheduler": @@ -945,17 +942,14 @@ def test_guided_decoding(self, backend: str, mocker): @pytest.mark.skip_less_device(2) @pytest.mark.skip_less_device_memory(48000) - @parametrize_with_ids("eagle3_one_model", [True, False]) @pytest.mark.parametrize("backend", ["xgrammar", "llguidance"]) - def test_guided_decoding_with_eagle3(self, backend: str, - eagle3_one_model: bool, mocker): + def test_guided_decoding_with_eagle3(self, backend: str, mocker): mocker.patch.dict(os.environ, {"TRTLLM_XGUIDANCE_LENIENT": "1"}) speculative_decoding_config = { "decoding_type": "Eagle", "max_draft_len": 3, "speculative_model": f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B", - "eagle3_one_model": eagle3_one_model } ctx_server_config = { @@ -971,8 +965,7 @@ def test_guided_decoding_with_eagle3(self, backend: str, } } gen_server_config = { - # Two-model eagle3 does not support overlap scheduler - "disable_overlap_scheduler": not eagle3_one_model, + "disable_overlap_scheduler": False, "speculative_config": speculative_decoding_config, "kv_cache_config": { "free_gpu_memory_fraction": 0.8, diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 1a77e5d48604..6485e1b289a8 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -344,7 +344,6 @@ def get_default_kwargs(self, attn_backend="flashinfer"): speculative_config = Eagle3DecodingConfig( max_draft_len=3, speculative_model=self.EAGLE_MODEL_PATH, - eagle3_one_model=True, eagle3_layers_to_capture={1, 15, 28}, ) # Note: Test crashes with trtllm attn_backend + torch-simple @@ -389,7 +388,7 @@ def check_acceptance_rate(self, llm, min_acceptance_rate: float): @skip_pre_hopper @pytest.mark.skip_less_device_memory(32000) @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) - def test_eagle3_one_model(self, attn_backend): + def test_eagle3(self, attn_backend): """Test Eagle3 one-model speculative decoding accuracy on GSM8K.""" kwargs = self.get_default_kwargs(attn_backend=attn_backend) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 378aa74751db..5f19b7f71fd5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -106,8 +106,7 @@ def _run_multinode_accuracy(model_path, spec_config = None if draft_model_path is not None: spec_config = Eagle3DecodingConfig(max_draft_len=max_draft_len, - speculative_model=draft_model_path, - eagle3_one_model=True) + speculative_model=draft_model_path) with LLM(model_path, tensor_parallel_size=tp_size, @@ -210,7 +209,6 @@ def test_eagle3_rejection_dynamic_tree_smoke(self, use_dynamic_tree, spec_config_kwargs = dict( max_draft_len=4, speculative_model=eagle_model_dir, - eagle3_one_model=True, use_rejection_sampling=True, ) max_batch_size = 1 @@ -369,10 +367,8 @@ def test_fp8_llm_sampler(self): @skip_pre_hopper @parametrize_with_ids("overlap_scheduler", [True, False]) - @parametrize_with_ids("eagle3_one_model", [True, False]) @parametrize_with_ids("sampler_async_worker", [True, False]) - def test_eagle3(self, overlap_scheduler, eagle3_one_model, - sampler_async_worker): + def test_eagle3(self, overlap_scheduler, sampler_async_worker): pytorch_config = dict( max_batch_size= 1, # add max_batch_size to avoid error in overlap scheduler @@ -381,17 +377,15 @@ def test_eagle3(self, overlap_scheduler, eagle3_one_model, cuda_graph_config=CudaGraphConfig(max_batch_size=1, enable_padding=True), ) - kv_cache_config = KvCacheConfig( - enable_block_reuse=True, free_gpu_memory_fraction=0.8 - ) # both one-model and two-model supports this feature + kv_cache_config = KvCacheConfig(enable_block_reuse=True, + free_gpu_memory_fraction=0.8) eagle_model_dir = f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B" target_model_dir = f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct" draft_len = 4 spec_config = Eagle3DecodingConfig(max_draft_len=draft_len, - speculative_model=eagle_model_dir, - eagle3_one_model=eagle3_one_model) + speculative_model=eagle_model_dir) with LLM(model=target_model_dir, **pytorch_config, @@ -418,7 +412,6 @@ def test_eagle3_sa(self): spec_config = Eagle3DecodingConfig(max_draft_len=4, speculative_model=eagle_model_dir, - eagle3_one_model=True, sa_config=SAEnhancerConfig()) with LLM(model=target_model_dir, @@ -446,7 +439,6 @@ def test_eagle3_sa_global_pool(self): spec_config = Eagle3DecodingConfig( max_draft_len=4, speculative_model=eagle_model_dir, - eagle3_one_model=True, sa_config=SAEnhancerConfig(enable_global_pool=True)) with LLM(model=target_model_dir, @@ -471,7 +463,6 @@ def test_eagle3_sa_dynamic_draft_len(self): spec_config = Eagle3DecodingConfig( max_draft_len=4, speculative_model=eagle_model_dir, - eagle3_one_model=True, sa_config=SAEnhancerConfig(enable_global_pool=True), draft_len_schedule={ 50: 4, @@ -499,9 +490,8 @@ def test_pard(self, overlap_scheduler): cuda_graph_config=CudaGraphConfig(max_batch_size=1, enable_padding=True), ) - kv_cache_config = KvCacheConfig( - enable_block_reuse=True, free_gpu_memory_fraction=0.8 - ) # both one-model and two-model supports this feature + kv_cache_config = KvCacheConfig(enable_block_reuse=True, + free_gpu_memory_fraction=0.8) pard_model_dir = f"{llm_models_root()}/PARD-Llama-3.2-1B" target_model_dir = f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct" @@ -839,27 +829,23 @@ def test_guided_decoding_4gpus(self, backend: str, mocker): task.evaluate(llm) @skip_pre_hopper - @parametrize_with_ids("eagle3_one_model", [True, False]) @pytest.mark.parametrize("backend", ["xgrammar", "llguidance"]) - def test_guided_decoding_with_eagle3(self, backend: str, - eagle3_one_model: bool, mocker): + def test_guided_decoding_with_eagle3(self, backend: str, mocker): mocker.patch.dict(os.environ, {"TRTLLM_XGUIDANCE_LENIENT": "1"}) kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8) cuda_graph_config = CudaGraphConfig(enable_padding=True) spec_config = Eagle3DecodingConfig( max_draft_len=3, - speculative_model=f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B", - eagle3_one_model=eagle3_one_model) - llm = LLM( - self.MODEL_PATH, - guided_decoding_backend=backend, - kv_cache_config=kv_cache_config, - cuda_graph_config=cuda_graph_config, - enable_chunked_prefill=True, - max_num_tokens=256, - speculative_config=spec_config, - # Two-model eagle3 does not support overlap scheduler - disable_overlap_scheduler=not eagle3_one_model) + speculative_model=f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B" + ) + llm = LLM(self.MODEL_PATH, + guided_decoding_backend=backend, + kv_cache_config=kv_cache_config, + cuda_graph_config=cuda_graph_config, + enable_chunked_prefill=True, + max_num_tokens=256, + speculative_config=spec_config, + disable_overlap_scheduler=False) with llm: task = JsonModeEval(self.MODEL_NAME) task.evaluate(llm) @@ -879,8 +865,8 @@ def test_guided_decoding_with_eagle3_low_latency_dispatch( cuda_graph_config = CudaGraphConfig(enable_padding=True) spec_config = Eagle3DecodingConfig( max_draft_len=3, - speculative_model=f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B", - eagle3_one_model=True) + speculative_model=f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B" + ) llm = LLM(self.MODEL_PATH, guided_decoding_backend=backend, kv_cache_config=kv_cache_config, @@ -1056,17 +1042,15 @@ def test_auto_dtype_tp2(self): @skip_pre_hopper @pytest.mark.skip_less_mpi_world_size(8) @parametrize_with_ids("torch_compile", [False, True]) - @parametrize_with_ids("eagle3_one_model", [True, False]) - def test_fp8_eagle3_tp8(self, eagle3_one_model, torch_compile): + def test_fp8_eagle3_tp8(self, torch_compile): model_path = f"{llm_models_root()}/modelopt-hf-model-hub/Llama-3.3-70B-Instruct-fp8" eagle_model_dir = f"{llm_models_root()}/EAGLE3-LLaMA3.3-Instruct-70B" kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6) spec_config = Eagle3DecodingConfig(max_draft_len=3, - speculative_model=eagle_model_dir, - eagle3_one_model=eagle3_one_model) + speculative_model=eagle_model_dir) torch_compile_config = _get_default_torch_compile_config(torch_compile) pytorch_config = dict( - disable_overlap_scheduler=not eagle3_one_model, + disable_overlap_scheduler=False, cuda_graph_config=CudaGraphConfig(max_batch_size=1), torch_compile_config=torch_compile_config) with LLM(model_path, @@ -1630,25 +1614,6 @@ def test_bfloat16_python_scheduler(self, mtp_nextn, attention_dp, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - @pytest.mark.skip_less_device_memory(60000) - def test_bfloat16_2_model_mtp(self): - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.3) - pytorch_config = dict( - disable_overlap_scheduler=True, - cuda_graph_config=CudaGraphConfig(), - ) - mtp_config = MTPDecodingConfig(max_draft_len=3, - mtp_eagle_one_model=False, - speculative_model=self.MODEL_PATH) - with LLM(self.MODEL_PATH, - kv_cache_config=kv_cache_config, - enable_chunked_prefill=False, - max_num_tokens=8192, - **pytorch_config, - speculative_config=mtp_config) as llm: - task = GSM8K(self.MODEL_NAME) - task.evaluate(llm) - @pytest.mark.skip_less_device_memory(60000) def test_bfloat16_mtp_sa(self): """Accuracy test for MTP + Suffix Automaton (MTP+SA) speculative decoding.""" @@ -1752,12 +1717,10 @@ def test_mtp_eagle_dynamic_draft_len(self): cuda_graph_config=cuda_graph_config, ) kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75) - # Force MTP-Eagle one-model path. mtp_config = MTPDecodingConfig( num_nextn_predict_layers=max_draft_len, max_draft_len=max_draft_len, use_mtp_vanilla=False, - mtp_eagle_one_model=True, draft_len_schedule=draft_len_schedule, ) with LLM(self.MODEL_PATH, @@ -4526,19 +4489,17 @@ def test_bf16(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, task.evaluate(llm) @parametrize_with_ids( - "eagle3_one_model,enable_chunked_prefill,enable_max_concurrency,enable_draft_len_schedule", + "enable_chunked_prefill,enable_max_concurrency,enable_draft_len_schedule", [ - # Base coverage: eagle3_one_model x enable_chunked_prefill. - (True, True, False, False), - (True, False, False, False), - (False, True, False, False), - (False, False, False, False), + # Base coverage: enable_chunked_prefill. + (True, False, False), + (False, False, False), # Test max_concurrency control and draft_len_schedule. - (True, False, False, True), - (True, False, True, False), + (False, False, True), + (False, True, False), ]) - def test_eagle3(self, eagle3_one_model, enable_chunked_prefill, - enable_max_concurrency, enable_draft_len_schedule): + def test_eagle3(self, enable_chunked_prefill, enable_max_concurrency, + enable_draft_len_schedule): max_concurrency = 100 if enable_max_concurrency else None draft_len_schedule = { 50: 4, @@ -4551,7 +4512,7 @@ def test_eagle3(self, eagle3_one_model, enable_chunked_prefill, max_draft_len = 4 pytorch_config = dict( - disable_overlap_scheduler=not eagle3_one_model, + disable_overlap_scheduler=False, cuda_graph_config=cuda_graph_config, ) kv_cache_config = KvCacheConfig( @@ -4565,7 +4526,6 @@ def test_eagle3(self, eagle3_one_model, enable_chunked_prefill, spec_config = Eagle3DecodingConfig( max_draft_len=max_draft_len, speculative_model=eagle_model_dir, - eagle3_one_model=eagle3_one_model, max_concurrency=max_concurrency, draft_len_schedule=draft_len_schedule) @@ -4917,9 +4877,8 @@ def test_nvfp4(self, tp_size, pp_size, ep_size, attention_dp, moe_backend, if eagle3: spec_config = Eagle3DecodingConfig( max_draft_len=2, - speculative_model= - f"{llm_models_root()}/Qwen3/qwen3-235B-eagle3/", - eagle3_one_model=True) + speculative_model=f"{llm_models_root()}/Qwen3/qwen3-235B-eagle3/" + ) with LLM( f"{llm_models_root()}/Qwen3/saved_models_Qwen3-235B-A22B_nvfp4_hf", tensor_parallel_size=tp_size, @@ -4965,9 +4924,8 @@ def test_nvfp4_4gpus(self, tp_size, pp_size, ep_size, attention_dp, if eagle3: spec_config = Eagle3DecodingConfig( max_draft_len=2, - speculative_model= - f"{llm_models_root()}/Qwen3/qwen3-235B-eagle3/", - eagle3_one_model=True) + speculative_model=f"{llm_models_root()}/Qwen3/qwen3-235B-eagle3/" + ) with LLM( f"{llm_models_root()}/Qwen3/saved_models_Qwen3-235B-A22B_nvfp4_hf", tensor_parallel_size=tp_size, @@ -5495,8 +5453,6 @@ def test_w4_chunked_prefill(self, kv_cache_dtype, moe_backend, mocker): @pytest.mark.skip_less_device(4) @pytest.mark.parametrize("overlap_scheduler", [True, False], ids=["overlap_scheduler", "no_overlap_scheduler"]) - @pytest.mark.parametrize("one_model", [True, False], - ids=["one_model", "two_model"]) @pytest.mark.parametrize("moe_backend", [ "CUTLASS", pytest.param("TRTLLM", marks=skip_no_trtllm_gen_moe_support), @@ -5505,8 +5461,8 @@ def test_w4_chunked_prefill(self, kv_cache_dtype, moe_backend, mocker): ids=["cutlass", "trtllm", "triton"]) @pytest.mark.parametrize("v2_kv_cache", [True, False], ids=["v2_kv_cache", "v1_kv_cache"]) - def test_eagle3_4gpus(self, v2_kv_cache, moe_backend, one_model, - overlap_scheduler, mocker): + def test_eagle3_4gpus(self, v2_kv_cache, moe_backend, overlap_scheduler, + mocker): if get_sm_version() == 90: pytest.skip( "https://nvbugs/5636916: Remaining Hopper Eagle Accuracy Issue for only TP=4" @@ -5522,12 +5478,6 @@ def test_eagle3_4gpus(self, v2_kv_cache, moe_backend, one_model, mocker.patch.object(GPQADiamond, "MAX_OUTPUT_LEN", MAX_OUTPUT_LEN) mocker.patch.object(GPQADiamond, "MAX_INPUT_LEN", MAX_INPUT_LEN) - if v2_kv_cache and not one_model and overlap_scheduler: - pytest.skip( - "KVCacheManagerV2 not compatible with two-model overlap scheduling" - ) - - # https://nvbugs/5590408: 2-Model overlap scheduling has accuracy issue pytorch_config = dict(disable_overlap_scheduler=not overlap_scheduler, cuda_graph_config=CudaGraphConfig()) kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4, @@ -5537,8 +5487,7 @@ def test_eagle3_4gpus(self, v2_kv_cache, moe_backend, one_model, eagle_model_dir = f"{llm_models_root()}/gpt_oss/gpt-oss-120b-Eagle3" draft_len = 3 spec_config = Eagle3DecodingConfig(max_draft_len=draft_len, - speculative_model=eagle_model_dir, - eagle3_one_model=one_model) + speculative_model=eagle_model_dir) max_seq_len = MAX_INPUT_LEN + MAX_OUTPUT_LEN llm = LLM(self.MODEL_PATH, @@ -5580,11 +5529,9 @@ def test_eagle3_4gpus(self, v2_kv_cache, moe_backend, one_model, extra_evaluator_kwargs=extra_evaluator_kwargs) @pytest.mark.skip_less_device(4) - @pytest.mark.parametrize("one_model", [True, False], - ids=["one_model", "two_model"]) @pytest.mark.parametrize("v2_kv_cache", [True, False], ids=["v2_kv_cache", "v1_kv_cache"]) - def test_eagle3_vswa_reuse_4gpus(self, v2_kv_cache, one_model, mocker): + def test_eagle3_vswa_reuse_4gpus(self, v2_kv_cache, mocker): MAX_OUTPUT_LEN = 128179 MAX_INPUT_LEN = 32768 @@ -5595,11 +5542,6 @@ def test_eagle3_vswa_reuse_4gpus(self, v2_kv_cache, one_model, mocker): mocker.patch.object(GPQADiamond, "MAX_OUTPUT_LEN", MAX_OUTPUT_LEN) mocker.patch.object(GPQADiamond, "MAX_INPUT_LEN", MAX_INPUT_LEN) - if v2_kv_cache and not one_model: - pytest.skip( - "KVCacheManagerV2 not compatible with two-model overlap scheduling" - ) - pytorch_config = dict(cuda_graph_config=CudaGraphConfig()) kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4, dtype="auto", @@ -5610,8 +5552,7 @@ def test_eagle3_vswa_reuse_4gpus(self, v2_kv_cache, one_model, mocker): eagle_model_dir = f"{llm_models_root()}/gpt_oss/gpt-oss-120b-Eagle3" draft_len = 3 spec_config = Eagle3DecodingConfig(max_draft_len=draft_len, - speculative_model=eagle_model_dir, - eagle3_one_model=one_model) + speculative_model=eagle_model_dir) max_seq_len = MAX_INPUT_LEN + MAX_OUTPUT_LEN llm = LLM(self.MODEL_PATH, @@ -5652,9 +5593,7 @@ def test_eagle3_vswa_reuse_4gpus(self, v2_kv_cache, one_model, mocker): extra_evaluator_kwargs=extra_evaluator_kwargs) @pytest.mark.skip_less_device(4) - @pytest.mark.parametrize("one_model", [True, False], - ids=["one_model", "two_model"]) - def test_eagle3_guided_decoding_4gpus(self, one_model, mocker): + def test_eagle3_guided_decoding_4gpus(self, mocker): MAX_OUTPUT_LEN = 128179 MAX_INPUT_LEN = 32768 @@ -5673,8 +5612,7 @@ def test_eagle3_guided_decoding_4gpus(self, one_model, mocker): eagle_model_dir = f"{llm_models_root()}/gpt_oss/gpt-oss-120b-Eagle3" draft_len = 3 spec_config = Eagle3DecodingConfig(max_draft_len=draft_len, - speculative_model=eagle_model_dir, - eagle3_one_model=one_model) + speculative_model=eagle_model_dir) max_seq_len = MAX_INPUT_LEN + MAX_OUTPUT_LEN llm = LLM(self.MODEL_PATH, @@ -5697,16 +5635,13 @@ def test_eagle3_guided_decoding_4gpus(self, one_model, mocker): @pytest.mark.timeout(14400) @pytest.mark.parametrize("overlap_scheduler", [True, False], ids=["overlap_scheduler", "no_overlap_scheduler"]) - @pytest.mark.parametrize("one_model", [True, False], - ids=["one_model", "two_model"]) @pytest.mark.parametrize("moe_backend", [ "CUTLASS", pytest.param("TRTLLM", marks=skip_no_trtllm_gen_moe_support), pytest.param("TRITON", marks=skip_no_hopper) ], ids=["cutlass", "trtllm", "triton"]) - def test_eagle3_2gpus(self, moe_backend, one_model, overlap_scheduler, - mocker): + def test_eagle3_2gpus(self, moe_backend, overlap_scheduler, mocker): MAX_OUTPUT_LEN = 128179 MAX_INPUT_LEN = 32768 @@ -5717,7 +5652,6 @@ def test_eagle3_2gpus(self, moe_backend, one_model, overlap_scheduler, mocker.patch.object(GPQADiamond, "MAX_OUTPUT_LEN", MAX_OUTPUT_LEN) mocker.patch.object(GPQADiamond, "MAX_INPUT_LEN", MAX_INPUT_LEN) - # https://nvbugs/5590408: 2-Model overlap scheduling has accuracy issue pytorch_config = dict( max_batch_size=8, disable_overlap_scheduler=not overlap_scheduler, @@ -5728,8 +5662,7 @@ def test_eagle3_2gpus(self, moe_backend, one_model, overlap_scheduler, eagle_model_dir = f"{llm_models_root()}/gpt_oss/gpt-oss-120b-Eagle3" draft_len = 3 spec_config = Eagle3DecodingConfig(max_draft_len=draft_len, - speculative_model=eagle_model_dir, - eagle3_one_model=one_model) + speculative_model=eagle_model_dir) max_seq_len = MAX_INPUT_LEN + MAX_OUTPUT_LEN llm = LLM(self.MODEL_PATH, @@ -5771,15 +5704,13 @@ def test_eagle3_2gpus(self, moe_backend, one_model, overlap_scheduler, extra_evaluator_kwargs=extra_evaluator_kwargs) @pytest.mark.skip_less_device_memory(90000) - @pytest.mark.parametrize("one_model", [True, False], - ids=["one_model", "two_model"]) @pytest.mark.parametrize("moe_backend", [ "CUTLASS", pytest.param("TRTLLM", marks=skip_pre_blackwell), pytest.param("TRITON", marks=skip_no_hopper) ], ids=["cutlass", "trtllm", "triton"]) - def test_eagle3_1gpu(self, moe_backend, one_model, mocker): + def test_eagle3_1gpu(self, moe_backend, mocker): mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", 8192) mocker.patch.object(GSM8K, "NUM_SAMPLES", 300) mocker.patch.dict(GSM8K.EVALUATE_KWARGS, @@ -5787,7 +5718,7 @@ def test_eagle3_1gpu(self, moe_backend, one_model, mocker): pytorch_config = dict( max_batch_size=8, - disable_overlap_scheduler=not one_model, + disable_overlap_scheduler=False, cuda_graph_config=CudaGraphConfig(max_batch_size=8)) kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9, dtype="auto", @@ -5796,8 +5727,7 @@ def test_eagle3_1gpu(self, moe_backend, one_model, mocker): eagle_model_dir = f"{llm_models_root()}/gpt_oss/gpt-oss-120b-Eagle3" draft_len = 5 spec_config = Eagle3DecodingConfig(max_draft_len=draft_len, - speculative_model=eagle_model_dir, - eagle3_one_model=one_model) + speculative_model=eagle_model_dir) llm = LLM(self.MODEL_PATH, tensor_parallel_size=1, @@ -6260,10 +6190,7 @@ def test_bf16_mtp(self, mocker): cuda_graph_config = CudaGraphConfig(enable_padding=True, max_batch_size=32) - mtp_config = MTPDecodingConfig( - num_nextn_predict_layers=3, - mtp_eagle_one_model=True, - ) + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=3, ) with LLM(self.MODEL_PATH, trust_remote_code=True, @@ -6771,7 +6698,6 @@ def test_nvfp4_4gpus(self, tp_size, pp_size, ep_size, attention_dp, max_draft_len=2, speculative_model= f"{llm_models_root()}/Mistral-Large-3-675B/Mistral-Large-3-675B-Instruct-2512-Eagle/", - eagle3_one_model=True, eagle3_model_arch="mistral_large3") with LLM( f"{llm_models_root()}/Mistral-Large-3-675B/Mistral-Large-3-675B-Instruct-2512-NVFP4/", @@ -6822,7 +6748,6 @@ def test_fp8(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, max_draft_len=2, speculative_model= f"{llm_models_root()}/Mistral-Large-3-675B/Mistral-Large-3-675B-Instruct-2512-Eagle/", - eagle3_one_model=True, eagle3_model_arch="mistral_large3") with LLM( f"{llm_models_root()}/Mistral-Large-3-675B/Mistral-Large-3-675B-Instruct-2512/", @@ -7146,10 +7071,7 @@ def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, pytest.skip( f"Device count {get_device_count()} is less than required {gpu_needed}" ) - mtp_config = MTPDecodingConfig( - num_nextn_predict_layers=3, - mtp_eagle_one_model=True, - ) + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=3, ) with LLM( f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4", kv_cache_config=KvCacheConfig( @@ -7218,11 +7140,8 @@ def test_nvfp4_parallelism(self, tp_size, ep_size, pp_size, attention_dp): @pytest.mark.skip_less_mpi_world_size(8) def test_nvfp4_8gpus_mtp(self): # Test MTP (Multi-Token Prediction) accuracy with nvfp4-fp8kv model. - # This test uses MTP with max_draft_len=3 and one_model mode. - mtp_config = MTPDecodingConfig( - max_draft_len=3, - mtp_eagle_one_model=True, - ) + # This test uses MTP with max_draft_len=3. + mtp_config = MTPDecodingConfig(max_draft_len=3, ) model_path = f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" with LLM( model_path, @@ -7254,10 +7173,7 @@ def test_nvfp4_8gpus_mtp(self): @pytest.mark.skip_less_device_memory(80000) def test_nvfp4_4gpu_mtp_ar(self): max_draft_len = 7 - mtp_config = MTPDecodingConfig( - max_draft_len=max_draft_len, - mtp_eagle_one_model=True, - ) + mtp_config = MTPDecodingConfig(max_draft_len=max_draft_len, ) model_path = f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" llm_common_config = dict( @@ -7320,10 +7236,7 @@ def test_nvfp4_4gpu_mtp_ar(self): def test_nvfp4_4gpu_mtp_ar_custom_op(self, monkeypatch): monkeypatch.setenv("TRTLLM_MAMBA2_MTP_USE_CUSTOM_OP", "1") max_draft_len = 7 - mtp_config = MTPDecodingConfig( - num_nextn_predict_layers=max_draft_len, - mtp_eagle_one_model=True, - ) + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=max_draft_len, ) model_path = f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" llm_common_config = dict( @@ -7385,10 +7298,7 @@ def test_nvfp4_4gpu_mtp_ar_custom_op(self, monkeypatch): def test_nvfp4_8gpus_mtp_custom_op(self, monkeypatch): monkeypatch.setenv("TRTLLM_MAMBA2_MTP_USE_CUSTOM_OP", "1") # Test MTP accuracy with nvfp4-fp8kv model using custom op path. - mtp_config = MTPDecodingConfig( - num_nextn_predict_layers=3, - mtp_eagle_one_model=True, - ) + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=3, ) model_path = f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" with LLM( model_path, @@ -7573,10 +7483,7 @@ def test_nvfp4_8gpus(self, attention_dp, moe_backend): def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, mamba_state_cache_interval, attention_dp, use_mtp): - mtp_config = MTPDecodingConfig( - num_nextn_predict_layers=3, - mtp_eagle_one_model=True, - ) + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=3, ) max_batch_size = 4 if attention_dp else 32 with LLM( f"{llm_models_root()}/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", diff --git a/tests/integration/defs/disaggregated/test_ad_disagg.py b/tests/integration/defs/disaggregated/test_ad_disagg.py index cbd743f6a499..eae18ea000b5 100644 --- a/tests/integration/defs/disaggregated/test_ad_disagg.py +++ b/tests/integration/defs/disaggregated/test_ad_disagg.py @@ -661,7 +661,6 @@ def llama_eagle3_config(): "speculative_config": Eagle3DecodingConfig( max_draft_len=3, speculative_model=model_path("EAGLE3-LLaMA3.1-Instruct-8B"), - eagle3_one_model=True, eagle3_layers_to_capture={1, 15, 28}, ), # Force the Eagle3 draft to match the BF16 Llama 3.1 target. Shared KV diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_triton.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_triton.yaml index ac6f24da0c29..4987384b2694 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_triton.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_triton.yaml @@ -16,7 +16,6 @@ context_servers: speculative_config: &eagle_spec decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: gpt_oss/gpt-oss-120b-Eagle3 kv_cache_config: enable_block_reuse: false diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_trtllm.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_trtllm.yaml index b52cc95463de..f01b2e3d05df 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_trtllm.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_trtllm.yaml @@ -16,7 +16,6 @@ context_servers: speculative_config: &eagle_spec decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: gpt_oss/gpt-oss-120b-Eagle3 kv_cache_config: enable_block_reuse: false diff --git a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py index 97634b686ea5..aefb09808440 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py +++ b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py @@ -513,16 +513,12 @@ def test_disaggregated_llama_context_capacity(model, enable_cuda_graph, @pytest.mark.parametrize("model", ["Llama-3.1-8B-Instruct"]) @pytest.mark.parametrize("spec_dec_model_path", ["EAGLE3-LLaMA3.1-Instruct-8B"]) @pytest.mark.parametrize("generation_overlap", [False]) -@pytest.mark.parametrize("eagle3_one_model", [True, False]) def test_disaggregated_spec_dec_batch_slot_limit(model, spec_dec_model_path, - generation_overlap, - eagle3_one_model): + generation_overlap): # Test whether the batch slots are properly released when using speculative decoding # with disaggregated serving. spec_dec_config = Eagle3DecodingConfig( - speculative_model=model_path(spec_dec_model_path), - eagle3_one_model=eagle3_one_model, - max_draft_len=3) + speculative_model=model_path(spec_dec_model_path), max_draft_len=3) worker_pytorch_configs = [] diff --git a/tests/integration/defs/examples/test_ad_speculative_decoding.py b/tests/integration/defs/examples/test_ad_speculative_decoding.py index 5a34bd58d09f..bdcd117f2a38 100644 --- a/tests/integration/defs/examples/test_ad_speculative_decoding.py +++ b/tests/integration/defs/examples/test_ad_speculative_decoding.py @@ -71,7 +71,7 @@ def get_model_paths(): ("flashinfer", "torch-simple"), ], ) -def test_autodeploy_eagle3_one_model_acceptance_rate(attn_backend: str, compile_backend: str): +def test_autodeploy_eagle3_acceptance_rate(attn_backend: str, compile_backend: str): """Test Eagle3 one-model acceptance rate with AutoDeploy engine. Runs Eagle3 one-model speculative decoding with streaming and verifies @@ -95,7 +95,6 @@ def test_autodeploy_eagle3_one_model_acceptance_rate(attn_backend: str, compile_ speculative_config = Eagle3DecodingConfig( max_draft_len=max_draft_len, speculative_model=eagle_model, - eagle3_one_model=True, eagle3_layers_to_capture={1, 15, 28}, ) diff --git a/tests/integration/defs/perf/pytorch_model_config.py b/tests/integration/defs/perf/pytorch_model_config.py index c8f3714a25a3..5d754294f549 100644 --- a/tests/integration/defs/perf/pytorch_model_config.py +++ b/tests/integration/defs/perf/pytorch_model_config.py @@ -376,7 +376,6 @@ def get_model_yaml_config(model_label: str, 'config': { 'speculative_config': { 'decoding_type': 'Eagle', - 'eagle3_one_model': True, 'speculative_model': 'Qwen3-4B_eagle3', 'max_draft_len': 3, }, diff --git a/tests/integration/defs/perf/test_perf_sanity.py b/tests/integration/defs/perf/test_perf_sanity.py index 0234e06dba7c..d6994e092499 100644 --- a/tests/integration/defs/perf/test_perf_sanity.py +++ b/tests/integration/defs/perf/test_perf_sanity.py @@ -533,7 +533,6 @@ def __init__(self, server_config_data: dict, env_vars: str = ""): else: self.eagle3_layers_to_capture = [] self.speculative_model = speculative_config.get("speculative_model", "") - self.eagle3_one_model = speculative_config.get("eagle3_one_model", False) # match_mode: "config" (default) or "scenario" self.match_mode = server_config_data.get("match_mode", "config") @@ -662,7 +661,6 @@ def to_db_data(self) -> dict: "l_max_draft_len": self.max_draft_len, "l_force_num_accepted_tokens": self.force_num_accepted_tokens, "s_speculative_model_dir": self.speculative_model, - "b_eagle3_one_model": self.eagle3_one_model, "s_server_log_link": "", "s_server_env_var": self.env_vars, } diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index 215731702435..750d0aea9a7f 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -1009,7 +1009,6 @@ def test_ptp_quickstart_advanced_mtp(llm_root, llm_venv, model_name, "MTP", "--model_dir", f"{llm_models_root()}/{model_path}", - "--use_one_model", ]) @@ -1289,7 +1288,6 @@ def test_relaxed_acceptance_quickstart_advanced_deepseek_r1_8gpus( "--relaxed_topk=10", "--relaxed_delta=0.5", "--enable_attention_dp", - "--use_one_model", "--moe_backend", "DEEPGEMM" if is_blackwell else "CUTLASS", ]) @@ -2036,7 +2034,6 @@ def test_eagle3_output_repetition_4gpus(model_dir: str, draft_model_dir: str): spec_config = Eagle3DecodingConfig( max_draft_len=3, speculative_model=eagle_model_dir, - eagle3_one_model=True, ) with LLM(**llm_common_config, speculative_config=spec_config) as llm_spec: results_spec = llm_spec.generate([prompt], sampling_params) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index fa1882623a3a..046722bff967 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -48,14 +48,11 @@ accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[MMLU-gen_tp=1-ctx_pp=4] accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[MMLU-gen_tp=2-ctx_pp=2] accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[MMLU-gen_tp=2-ctx_pp=4] -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_eagle3[eagle3_one_model=False-overlap_scheduler=False] -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_eagle3[eagle3_one_model=True-overlap_scheduler=True] +accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_eagle3[overlap_scheduler=True] accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding[llguidance] accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding[xgrammar] -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=False] -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=True] -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=False] -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=True] +accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance] +accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar] accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_kv_cache_v2_nixl_python accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_multi_instance[GSM8K] accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_multi_instance[MMLU] @@ -96,8 +93,8 @@ accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[triton-Fal accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-1] accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-4] accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-True-1] -accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3_one_model[flashinfer] -accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3_one_model[trtllm] +accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3[flashinfer] +accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3[trtllm] accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[deepseek-ai_DeepSeek-R1-0528-True] accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[nvidia_DeepSeek-R1-0528-NVFP4-v2-True] accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[google_gemma-3-1b-it-False] @@ -462,24 +459,23 @@ accuracy/test_llm_api_pytorch.py::TestGemma3_27BInstruct::test_fp8_prequantized accuracy/test_llm_api_pytorch.py::TestGLM5FP8::test_8gpus[tp_size=8-ep_size=8] accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dummy_load_format -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[cutlass-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_guided_decoding_4gpus[one_model] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache-one_model] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache-one_model] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache-two_model] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[cutlass-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_guided_decoding_4gpus +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-cutlass-auto] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-cutlass-fp8] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-triton-auto] @@ -580,9 +576,8 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_chunked_prefill[ accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_dflash_dynamic_draft_len accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_draft_target_dynamic_draft_len accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_dummy_load_format -accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-eagle3_one_model=False-overlap_scheduler=False] -accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-eagle3_one_model=True-overlap_scheduler=True] -accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=True-eagle3_one_model=True-overlap_scheduler=True] +accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-overlap_scheduler=True] +accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=True-overlap_scheduler=True] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[no_dynamic_tree] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[dynamic_tree] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_sa @@ -628,10 +623,8 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding[ accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding[xgrammar] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_4gpus[llguidance] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_4gpus[xgrammar] -accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=False] -accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=True] -accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=False] -accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=True] +accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance] +accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_ngram[llguidance] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_ngram[xgrammar] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_ngram @@ -651,9 +644,8 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torc accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=False-enable_gemm_allreduce_fusion=True] accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=True-enable_gemm_allreduce_fusion=False] accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=True-enable_gemm_allreduce_fusion=True] -accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_eagle3_tp8[eagle3_one_model=False-torch_compile=False] -accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_eagle3_tp8[eagle3_one_model=True-torch_compile=False] -accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_eagle3_tp8[eagle3_one_model=True-torch_compile=True] +accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_eagle3_tp8[torch_compile=False] +accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_eagle3_tp8[torch_compile=True] accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_tp4[torch_compile=False] accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp8_tp4[torch_compile=True] accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=False] @@ -753,12 +745,10 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4_mtp3_gdn_rep accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[multi_gpus_no_cache] accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dummy_load_format -accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=False-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=False] -accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=False-enable_chunked_prefill=True-enable_max_concurrency=False-enable_draft_len_schedule=False] -accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=False] -accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=True] -accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=True-enable_draft_len_schedule=False] -accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=True-enable_max_concurrency=False-enable_draft_len_schedule=False] +accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=False] +accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=True] +accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=False-enable_max_concurrency=True-enable_draft_len_schedule=False] +accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=True-enable_max_concurrency=False-enable_draft_len_schedule=False] accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency] accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales_early_first_token_response accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_w4a8_mxfp4[fp8-latency] @@ -905,8 +895,7 @@ disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[ disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[False-True-Qwen3-8B-FP8] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-False-Qwen3-8B-FP8] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-True-Qwen3-8B-FP8] -disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[False-False-EAGLE3-LLaMA3.1-Instruct-8B-Llama-3.1-8B-Instruct] -disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[True-False-EAGLE3-LLaMA3.1-Instruct-8B-Llama-3.1-8B-Instruct] +disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[False-EAGLE3-LLaMA3.1-Instruct-8B-Llama-3.1-8B-Instruct] disaggregated/test_workers.py::test_workers_conditional_disaggregation[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_workers.py::test_workers_conditional_disaggregation_deepseek_v3_lite_bf16[DeepSeek-V3-Lite-bf16] disaggregated/test_workers.py::test_workers_kv_cache_aware_router[TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/integration/test_lists/qa/llm_function_rtx6k.txt b/tests/integration/test_lists/qa/llm_function_rtx6k.txt index 4260a5328d77..ae76590fad36 100644 --- a/tests/integration/test_lists/qa/llm_function_rtx6k.txt +++ b/tests/integration/test_lists/qa/llm_function_rtx6k.txt @@ -88,24 +88,23 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[b accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_batch_waiting[batch_wait_timeout_iters=10-batch_wait_max_tokens_ratio=1.0-mtp_nextn=0-fp8kv=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-v2_kv_cache=True] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dflash accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dummy_load_format -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[cutlass-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-no_overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-overlap_scheduler] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_guided_decoding_4gpus[one_model] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache-one_model] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache-one_model] -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache-two_model] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[cutlass-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-no_overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-overlap_scheduler] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_guided_decoding_4gpus +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache] +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-cutlass-auto] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-cutlass-fp8] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-triton-auto] diff --git a/tests/integration/test_lists/qa/llm_spark_func.yml b/tests/integration/test_lists/qa/llm_spark_func.yml index e6d95b720761..59bf5e0fa3d9 100644 --- a/tests/integration/test_lists/qa/llm_spark_func.yml +++ b/tests/integration/test_lists/qa/llm_spark_func.yml @@ -17,8 +17,7 @@ llm_spark_func: - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-8b-fp8-Qwen3/nvidia-Qwen3-8B-FP8] - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-8b-nvfp4-Qwen3/nvidia-Qwen3-8B-NVFP4] - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-8B-bf16-Qwen3/Qwen3-8B] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_1gpu[cutlass-one_model] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_1gpu[cutlass-two_model] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_1gpu[cutlass] - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-14b-fp8-Qwen3/nvidia-Qwen3-14B-FP8] - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-14b-nvfp4-Qwen3/nvidia-Qwen3-14B-NVFP4] - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-14B-bf16-Qwen3/Qwen3-14B] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index d812dcf5d84e..775c116aee3f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -75,7 +75,6 @@ l0_b200: - unittest/_torch/attention - unittest/_torch/compilation - unittest/_torch/debugger - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp - unittest/disaggregated/test_deepseek_v4_kv_transfer.py - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py TIMEOUT (60) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index e64a30ba11c0..a0c4e70f83f4 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -312,15 +312,14 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-ep4-trtllm-auto] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-ep4-trtllm-fp8] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-cutlass-auto] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-no_overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-no_overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-no_overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-no_overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache-one_model] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache-one_model] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache-two_model] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-no_overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-no_overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-no_overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-no_overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v2_kv_cache] - unittest/_torch/multi_gpu_modeling -k "deepseek" - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 6d63a7782a53..20c00e5fe6e7 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -24,9 +24,9 @@ l0_dgx_h100: - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-tp2-t5-small] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-tp2-bart-large-cnn] # ------------- Disaggregated serving tests --------------- - - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_eagle3[eagle3_one_model=True-overlap_scheduler=True] + - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_eagle3[overlap_scheduler=True] - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding[xgrammar] - - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=True] + - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar] - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype[False-True] - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype[True-True] - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype[False-False] @@ -77,8 +77,8 @@ l0_dgx_h100: auto_trigger: gpt_oss orchestrator: mpi tests: - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[cutlass-one_model-overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-one_model-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[cutlass-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-overlap_scheduler] - condition: ranges: system_gpu_count: @@ -388,7 +388,7 @@ l0_dgx_h100: - accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[Qwen_QwQ-32B-False] - accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[meta-llama_Llama-3.3-70B-Instruct-False] - accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[nvidia_Llama-3.1-8B-Instruct-FP8-True] - - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3_one_model[flashinfer] + - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3[flashinfer] - accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[bf16-4-attn_dp_off-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[bf16-4-attn_dp_off-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[bf16-4-attn_dp_on-trtllm] diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 0dd2d3935113..467b70807a98 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -49,10 +49,10 @@ l0_gb200_multi_gpus: - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-ep4-cutlass-auto] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-trtllm-fp8] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4a16[dp4-fp8] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-one_model-overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-one_model-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-trtllm-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-trtllm-overlap_scheduler] - accuracy/test_llm_api_pytorch_multimodal.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] TIMEOUT (90) - accuracy/test_disaggregated_serving.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60) - accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 133a34e15300..7a46626c2854 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -137,13 +137,13 @@ l0_h100: - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_chunked_prefill[use_temperature=True-attn_backend=TRTLLM] TIMEOUT (90) - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_dummy_load_format - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding[xgrammar] - - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar-eagle3_one_model=True] + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[xgrammar] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=TRTLLM-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=True-attn_backend=TRTLLM-torch_compile=True] - - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-eagle3_one_model=True-overlap_scheduler=True] - - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=True-eagle3_one_model=True-overlap_scheduler=True] + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-overlap_scheduler=True] + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=True-overlap_scheduler=True] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_pard[overlap_scheduler=True] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_pard[overlap_scheduler=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_dflash @@ -165,7 +165,7 @@ l0_h100: - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dummy_load_format - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_dummy_load_format - - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=False] + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=False] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dflash - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 @@ -193,8 +193,7 @@ l0_h100: - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-False-Qwen3-8B-FP8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-True-Qwen3-8B-FP8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_llama_context_capacity[False-False-DeepSeek-V3-Lite-fp8/fp8] - - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[True-False-EAGLE3-LLaMA3.1-Instruct-8B-Llama-3.1-8B-Instruct] - - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[False-False-EAGLE3-LLaMA3.1-Instruct-8B-Llama-3.1-8B-Instruct] + - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[False-EAGLE3-LLaMA3.1-Instruct-8B-Llama-3.1-8B-Instruct] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_cancel_gen_requests[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[False-TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[True-TinyLlama-1.1B-Chat-v1.0] @@ -355,9 +354,9 @@ l0_h100: - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] TIMEOUT (90) - llmapi/test_llm_api_pytorch_moe_lora.py::test_mixtral_moe_routed_expert_fp8_multi_lora_varying_ranks[eager] TIMEOUT (90) - llmapi/test_llm_api_pytorch_moe_lora.py::test_mixtral_moe_routed_expert_fp8_multi_lora_varying_ranks[cudagraph] TIMEOUT (90) - - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=True-enable_max_concurrency=False-enable_draft_len_schedule=False] - - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=True] - - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=True-enable_draft_len_schedule=False] + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=True-enable_max_concurrency=False-enable_draft_len_schedule=False] + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=True] + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[enable_chunked_prefill=False-enable_max_concurrency=True-enable_draft_len_schedule=False] - accuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_fp8 - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16[attn_backend=FLASHINFER-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16[attn_backend=FLASHINFER-torch_compile=True] @@ -422,7 +421,7 @@ l0_h100: - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8_block_scales[latency-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8_block_scales[latency-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding[llguidance] - - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance-eagle3_one_model=True] + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_eagle3[llguidance] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_ngram[xgrammar] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_with_ngram[llguidance] - accuracy/test_llm_api_pytorch_multimodal.py::TestMistralSmall24B::test_auto_dtype[forced_chunked_prefill] @@ -476,14 +475,14 @@ l0_h100: - unittest/auto_deploy/singlegpu/utils - unittest/auto_deploy/standalone - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-False-1] - - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3_one_model[trtllm] + - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3[trtllm] - accuracy/test_llm_api_autodeploy.py::TestGemmaE2B::test_gemma3n_e2b_it - accuracy/test_llm_api_autodeploy.py::TestGemmaE2B::test_gemma4_e2b_it - accuracy/test_llm_api_autodeploy.py::TestNemotronV2::test_fp8[True] - accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[fp8-1-attn_dp_off-trtllm] - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[triton-False-1] - - examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_one_model_acceptance_rate[trtllm-torch-cudagraph] - - examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_one_model_acceptance_rate[flashinfer-torch-simple] + - examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_acceptance_rate[trtllm-torch-cudagraph] + - examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_acceptance_rate[flashinfer-torch-simple] - examples/test_ad_speculative_decoding.py::test_eagle_model_with_weights - examples/test_ad_speculative_decoding.py::test_eagle_wrapper_forward[1] - examples/test_ad_speculative_decoding.py::test_eagle_wrapper_forward[2] @@ -511,7 +510,7 @@ l0_h100: orchestrator: mpi tests: - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[trtllm-True-1] - - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3_one_model[flashinfer] + - accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B_Instruct_Eagle3::test_eagle3[flashinfer] - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[trtllm-triton_ssm-False] - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[trtllm-flashinfer_ssm-False] - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[trtllm-triton_ssm-True] diff --git a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml index f7372a8c4224..194eb32cd58e 100644 --- a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml +++ b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml @@ -121,11 +121,11 @@ l0_rtx_pro_6000: # - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_cutlass-torch_compile=False] # failed - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe_cutlass-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe_cutlass-torch_compile=True] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-no_overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-no_overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-one_model-no_overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-one_model-no_overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-one_model-overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-one_model-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-no_overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-no_overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-no_overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-no_overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-triton-overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-triton-overlap_scheduler] diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6dc4776bb633..1a06e431c16e 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -70,7 +70,7 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backe accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625) accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype SKIP (https://nvbugs/6450333) accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] SKIP (https://nvbugs/6507108) -accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_guided_decoding_4gpus[one_model] SKIP (https://nvbugs/5596343) +accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_guided_decoding_4gpus SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-triton-auto] SKIP (https://nvbugs/6026676) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-dp4-cutlass-auto] SKIP (https://nvbugs/6388142) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-ep4-cutlass-auto] SKIP (https://nvbugs/5596343) @@ -137,7 +137,7 @@ disaggregated/test_disaggregated.py::test_disaggregated_genbs1[TinyLlama-1.1B-Ch disaggregated/test_workers.py::test_workers_conversation_router[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322) disaggregated/test_workers.py::test_workers_kv_cache_aware_router_deepseek_v3_lite_bf16[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322) disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322) -examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_one_model_acceptance_rate[trtllm-torch-cudagraph] SKIP (https://nvbugs/6426841) +examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_acceptance_rate[trtllm-torch-cudagraph] SKIP (https://nvbugs/6426841) examples/test_ray.py::test_llm_inference_distributed_ray[pp2] SKIP (https://nvbugs/6427411) examples/test_ray.py::test_llm_inference_distributed_ray[tp2pp2] SKIP (https://nvbugs/6427411) examples/test_ray.py::test_ray_disaggregated_serving[tp2] SKIP (https://nvbugs/5612502) @@ -237,7 +237,7 @@ full:H20/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_d full:H20/accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_fp8] SKIP (https://nvbugs/6327718) full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=False] SKIP (https://nvbugs/6422343) full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6422343) -full:H20/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-one_model-overlap_scheduler] SKIP (https://nvbugs/6373530) +full:H20/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_2gpus[triton-overlap_scheduler] SKIP (https://nvbugs/6373530) full:H20/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[triton-auto] SKIP (https://nvbugs/6026676) full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3_5_35B_A3B_VL::test_auto_dtype SKIP (https://nvbugs/6442073) full:H20/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] SKIP (https://nvbugs/6475621) @@ -361,7 +361,7 @@ unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py: unittest/_torch/sampler/test_beam_search.py::test_beam_search_e2e[multi_process-TRTLLMSampler-cuda_graph_and_overlap-None-1-1-True-True-False] SKIP (https://nvbugs/6463819) unittest/_torch/sampler/test_trtllm_sampler.py::test_trtllm_sampler_best_of_with_logprobs SKIP (https://nvbugs/6487837) unittest/_torch/speculative/hw_agnostic/test_ngram.py::test_llama_ngram[True-True-TRTLLM] SKIP (https://nvbugs/6507102) -unittest/_torch/speculative/test_eagle3.py::test_llama_eagle3[True-TRTLLM-True-False-True-True-True-False-False-False] SKIP (https://nvbugs/6451425) +unittest/_torch/speculative/test_eagle3.py::test_llama_eagle3[True-TRTLLM-True-False-True-False-False-False] SKIP (https://nvbugs/6451425) unittest/_torch/thop/parallel/test_fp8_rowwise_linear.py::test_fp8_rowwise_linear[dtype1] SKIP (https://nvbugs/6301807) unittest/_torch/thop/serial/test_moe.py::TestMoeFp4::test_no_autotune[use_score_as_input-RoutingDSv3-swiglu-1024-1024-1] SKIP (https://nvbugs/5908070) unittest/_torch/thop/serial/test_moe.py::TestMoeFp4::test_no_autotune[use_score_as_input-RoutingRenormalize_qwen_next-swiglu-1024-1024-150] SKIP (https://nvbugs/5908070) diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml index baadb0c69514..53ff8a5c2b35 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml @@ -71,7 +71,6 @@ worker_config: speculative_config: &id001 decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: Kimi-K2.5-Thinking-Eagle3 trust_remote_code: true num_postprocess_workers: 4 diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml index 1c924aecc04c..39a239a86a2c 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml @@ -66,7 +66,6 @@ worker_config: speculative_config: &id001 decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: Kimi-K2.5-Thinking-Eagle3 trust_remote_code: true num_postprocess_workers: 4 diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml index 4a26a8d528d2..3e13d912a3a2 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml @@ -71,7 +71,6 @@ worker_config: speculative_config: &id001 decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: Kimi-K2.5-Thinking-Eagle3 trust_remote_code: true num_postprocess_workers: 4 diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml index ec66d24266d4..7b44fa157fb6 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml @@ -67,7 +67,6 @@ worker_config: speculative_config: &id001 decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: Kimi-K2.5-Thinking-Eagle3 trust_remote_code: true num_postprocess_workers: 4 diff --git a/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml index 2454a4bd8d9a..7c6d735e4b23 100644 --- a/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml @@ -81,7 +81,6 @@ worker_config: speculative_config: &id001 decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: Kimi-K2.5-Thinking-Eagle3 trust_remote_code: true num_postprocess_workers: 4 diff --git a/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml index 7633ba6bc22f..32f04ec0db30 100644 --- a/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml @@ -76,7 +76,6 @@ worker_config: speculative_config: &id001 decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: Kimi-K2.5-Thinking-Eagle3 trust_remote_code: true num_postprocess_workers: 4 diff --git a/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml index 0701596a0bad..b1f37e796ef6 100644 --- a/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml @@ -84,7 +84,6 @@ worker_config: speculative_config: &id001 decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: "gpt_oss/gpt-oss-120b-Eagle3" num_postprocess_workers: 4 stream_interval: 20 diff --git a/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml index 9dbfa216ef05..245327d72e33 100644 --- a/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml @@ -81,7 +81,6 @@ worker_config: speculative_config: &id001 decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: Kimi-K2.5-Thinking-Eagle3 trust_remote_code: true num_postprocess_workers: 4 diff --git a/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml index fb53d40386c2..4bb79a49abbe 100644 --- a/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml @@ -76,7 +76,6 @@ worker_config: speculative_config: &id001 decoding_type: Eagle max_draft_len: 3 - eagle3_one_model: true speculative_model: Kimi-K2.5-Thinking-Eagle3 trust_remote_code: true num_postprocess_workers: 4 diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index bcba3baf1b24..078fa7691a3b 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -1721,8 +1721,8 @@ def test_swa_scratch_reuse_uses_extra_kv_tokens_for_rewind(self): max_draft_len=7, max_total_draft_tokens=7, spec_dec_mode=SimpleNamespace( - is_eagle3_one_model=lambda: False, - is_mtp_eagle_one_model=lambda: False, + is_eagle3=lambda: False, + is_mtp_eagle=lambda: False, is_mtp_one_model=lambda: False, is_mtp_vanilla=lambda: False, use_one_engine=lambda: True, diff --git a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py index 4782ead014de..97a857d89be8 100644 --- a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py @@ -133,7 +133,6 @@ def _make_mock_model_engine(model_config): engine = Mock() engine.model.model_config = model_config engine.dtype = torch.bfloat16 - engine.is_draft_model = False engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER return engine @@ -169,7 +168,6 @@ def _make_creator( creator = KvCacheCreator.__new__(KvCacheCreator) creator._model_engine = model_engine - creator._draft_model_engine = None creator._mapping = mapping if mapping is not None else Mapping() creator._kv_cache_config = kv_cache_config creator._max_kv_tokens_in = kv_cache_config.max_tokens diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 8a5be984f312..b671d6b86d18 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -610,7 +610,6 @@ def test_estimation_temporarily_uses_inferred_pool_sizing() -> None: ): creator = KvCacheCreator( model_engine=model_engine, - draft_model_engine=None, mapping=Mock(cp_config={}), net_max_seq_len=max_seq_len, kv_connector_manager=None, diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 101edffbe155..27ad0c25b123 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1887,7 +1887,7 @@ def test_handles_error_on_single_rank(self): class TestOneModelMTPDraftTokenScheduling: """Regression tests for the one-model MTP over-scheduling bug (#16101). - One-model MTP (``mtp_eagle_one_model``) has no separate drafter, so + One-model MTP (``MTP_EAGLE``) has no separate drafter, so ``get_spec_drafter()`` returns None and the ``if self.drafter is not None`` draft-token normalization block in ``_prepare_and_schedule_batch`` is skipped. Without the ``elif`` fallback that mirrors it, generation requests @@ -1963,11 +1963,10 @@ def test_one_model_mtp_populates_draft_tokens_for_scheduling(self): disaggregated serving in one shot. The fix's state filter is {GENERATION_IN_PROGRESS, - DISAGG_GENERATION_INIT}, mirroring the two-model normalization, so a - single fix covers the aggregated decode path (GENERATION_IN_PROGRESS) - and the disagg decode-worker path (DISAGG_GENERATION_INIT). Context - requests (CONTEXT_INIT) are not generation requests and must be left - untouched. + DISAGG_GENERATION_INIT}, so a single fix covers the aggregated decode + path (GENERATION_IN_PROGRESS) and the disagg decode-worker path + (DISAGG_GENERATION_INIT). Context requests (CONTEXT_INIT) are not + generation requests and must be left untouched. """ gen = self._make_llm_request(0, LlmRequestState.GENERATION_IN_PROGRESS) disagg_gen = self._make_llm_request(1, LlmRequestState.DISAGG_GENERATION_INIT) diff --git a/tests/unittest/_torch/executor/test_seq_slot_sizing.py b/tests/unittest/_torch/executor/test_seq_slot_sizing.py index d42e6f4483c9..e3a28136d417 100644 --- a/tests/unittest/_torch/executor/test_seq_slot_sizing.py +++ b/tests/unittest/_torch/executor/test_seq_slot_sizing.py @@ -56,7 +56,7 @@ def test_dsv4_overlap_headroom_gate( spec_config = None if has_spec: spec_config = Mock() - spec_config.spec_dec_mode.is_mtp_eagle_one_model.return_value = is_mtp_one_model + spec_config.spec_dec_mode.is_mtp_eagle.return_value = is_mtp_one_model mapping = Mapping(world_size=pp_size, tp_size=1, pp_size=pp_size) assert ( diff --git a/tests/unittest/_torch/helpers.py b/tests/unittest/_torch/helpers.py index 7f167077effa..3c004a79357b 100644 --- a/tests/unittest/_torch/helpers.py +++ b/tests/unittest/_torch/helpers.py @@ -251,7 +251,6 @@ def create_mock_cuda_graph_runner(batch_size: int, use_mrope: bool = False): enable_attention_dp=False, original_max_draft_len=0, original_max_total_draft_tokens=0, - is_draft_model=False, is_encoder_decoder=False, mapping=Mapping(), dist=None, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_len_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_len_schedule.py index 67e5514cb9c2..4d30063d5d42 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_len_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_len_schedule.py @@ -5,7 +5,7 @@ import torch from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.llmapi import DraftTargetDecodingConfig, KvCacheConfig, NGramDecodingConfig +from tensorrt_llm.llmapi import KvCacheConfig, NGramDecodingConfig sys.path.append(os.path.join(os.path.dirname(__file__), "..")) from utils.llm_data import llm_models_root @@ -34,13 +34,12 @@ def enforce_single_worker(monkeypatch): @pytest.mark.high_cuda_memory def test_correctness_across_batch_sizes(drafter_type: str, schedule: dict): total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - memory_required = 30 if drafter_type == "model_drafter" else 20 + memory_required = 20 if total_mem_gb < memory_required: pytest.skip(f"Not enough memory (need {memory_required}GB, have {total_mem_gb:.1f}GB)") models_path = llm_models_root() target_model = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - draft_model = f"{models_path}/llama-3.2-models/Llama-3.2-3B-Instruct" max_batch_size = 8 max_draft_len = max(schedule.values()) # Use max from schedule @@ -59,21 +58,14 @@ def test_correctness_across_batch_sizes(drafter_type: str, schedule: dict): max_num_tokens=1024, ) - if drafter_type == "ngram": - spec_config = NGramDecodingConfig( - max_draft_len=max_draft_len, - max_matching_ngram_size=2, - draft_len_schedule=schedule, - is_keep_all=True, - is_use_oldest=True, - is_public_pool=False, - ) - else: - spec_config = DraftTargetDecodingConfig( - max_draft_len=max_draft_len, - speculative_model=str(draft_model), - draft_len_schedule=schedule, - ) + spec_config = NGramDecodingConfig( + max_draft_len=max_draft_len, + max_matching_ngram_size=2, + draft_len_schedule=schedule, + is_keep_all=True, + is_use_oldest=True, + is_public_pool=False, + ) prompts = [ "The capital of France is", @@ -105,22 +97,14 @@ def test_correctness_across_batch_sizes(drafter_type: str, schedule: dict): generated_text_with_schedule = [result.outputs[0].text for result in results_with_schedule] llm_with_schedule.shutdown() # Reference: spec decode with fixed max_draft_len (no schedule) - if drafter_type == "ngram": - spec_config_fixed = NGramDecodingConfig( - max_draft_len=max_draft_len, - max_matching_ngram_size=2, - draft_len_schedule=None, # No schedule - fixed draft length - is_keep_all=True, - is_use_oldest=True, - is_public_pool=False, - ) - else: - # skipped for move to 1 model - spec_config_fixed = DraftTargetDecodingConfig( - max_draft_len=max_draft_len, - speculative_model=str(draft_model), - draft_len_schedule=None, # No schedule - fixed draft length - ) + spec_config_fixed = NGramDecodingConfig( + max_draft_len=max_draft_len, + max_matching_ngram_size=2, + draft_len_schedule=None, # No schedule - fixed draft length + is_keep_all=True, + is_use_oldest=True, + is_public_pool=False, + ) llm_fixed = LLM(**llm_common_config, speculative_config=spec_config_fixed) results_fixed = llm_fixed.generate(prompts, sampling_params_list) generated_text_fixed = [result.outputs[0].text for result in results_fixed] @@ -152,9 +136,7 @@ def test_draft_len_schedule_functionality( pytest.skip("CUDA not available") total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if drafter_type == "model_drafter" and total_mem_gb < 30: - pytest.skip("Not enough memory for 2-model setup") - elif total_mem_gb < 20: + if total_mem_gb < 20: pytest.skip("Not enough memory") max_batch_size = 7 @@ -172,19 +154,11 @@ def test_draft_len_schedule_functionality( max_num_tokens=1024, ) - if drafter_type == "ngram": - spec_config = NGramDecodingConfig( - max_draft_len=5, - max_matching_ngram_size=2, - draft_len_schedule=draft_schedule, - ) - else: - # skipped for move to 1 model - spec_config = DraftTargetDecodingConfig( - max_draft_len=5, - speculative_model=str(llm_models_root() / "llama-3.2-models" / "Llama-3.2-3B-Instruct"), - draft_len_schedule=draft_schedule, - ) + spec_config = NGramDecodingConfig( + max_draft_len=5, + max_matching_ngram_size=2, + draft_len_schedule=draft_schedule, + ) prompts = ["The capital of France is" for i in range(7)] # Give each request different max_tokens so they finish at different times # This creates batch size transitions: 7 -> 6 -> 5 -> 4 -> 3 -> 2 -> 1 @@ -290,26 +264,12 @@ def instrumented_prepare_draft(scheduled_batch, resource_manager): # ======================================================================== # Verification Rule 2: actual_draft_lens (req.py_draft_tokens) vs drafter_max_draft_tokens # ======================================================================== - if drafter_type == "ngram": - # NGram: all actual_draft_lens <= drafter_max_draft_tokens - # ngram drafting length is based on ngram context - # so it's not necessary to be the same as drafter_max_draft_tokens - for idx, it in enumerate(iteration_data): - drafter_tokens = it["drafter_max_draft_tokens"] - for req_idx, actual_len in enumerate(it["actual_draft_lens"]): - assert actual_len <= drafter_tokens, ( - f"Iter {idx}, req {req_idx}: NGram produced {actual_len} > max {drafter_tokens}" - ) - - elif drafter_type == "model_drafter": - # ModelDrafter: all actual_draft_lens == drafter_max_draft_tokens - for idx, it in enumerate(iteration_data): - drafter_tokens = it["drafter_max_draft_tokens"] - actual_lens = it["actual_draft_lens"] - - if drafter_tokens > 0: - for req_idx, actual_len in enumerate(actual_lens): - assert actual_len == drafter_tokens, ( - f"Iter {idx}, req {req_idx}: ModelDrafter produced {actual_len} " - f"!= max_draft_tokens {drafter_tokens}" - ) + # NGram: all actual_draft_lens <= drafter_max_draft_tokens + # ngram drafting length is based on ngram context + # so it's not necessary to be the same as drafter_max_draft_tokens + for idx, it in enumerate(iteration_data): + drafter_tokens = it["drafter_max_draft_tokens"] + for req_idx, actual_len in enumerate(it["actual_draft_lens"]): + assert actual_len <= drafter_tokens, ( + f"Iter {idx}, req {req_idx}: NGram produced {actual_len} > max {drafter_tokens}" + ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_prepare_for_generation.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_prepare_for_generation.py deleted file mode 100644 index 78acbde58cb9..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_prepare_for_generation.py +++ /dev/null @@ -1,455 +0,0 @@ -import math -import os -import sys - -import torch -from utils.llm_data import llm_models_root - -from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata -from tensorrt_llm._torch.speculative.drafting_loops import StaticTreeDraftingLoopWrapper -from tensorrt_llm._torch.speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata -from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager -from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig - -sys.path.append(os.path.join(os.path.dirname(__file__), "..")) - - -class DummyModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.model_config = None - self.config = None - self.model = {} - self.model_is_wrapped = True - - def forward(self, *args, **kwargs) -> torch.Tensor: - pass - - -def test_draft_token_static_tree_prepare_for_generation(): - # Fix parameters - models_path = llm_models_root() - eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" # It will not actually be used. - use_dynamic_tree = False - max_new_tokens = 128 - kv_cache_manager = None - - # Create related object and run test - def run_test( - max_batch_size, - prepare_for_layer_idx, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - input_seq_lens_cuda, - input_kv_lens_cuda, - input_num_accepted_draft_tokens, - input_hidden_states_write_indices, - input_hidden_states_read_indices, - input_position_ids, - ref_position_ids, - ref_attn_metadata, - ref_spec_metadata, - ): - # 1) Create attention metadata - attn_metadata = TrtllmAttentionMetadata( - max_num_requests=max_batch_size, - max_num_tokens=max_new_tokens, - kv_cache_manager=kv_cache_manager, - ) - - # Set initial values - attn_metadata._seq_lens_cuda = input_seq_lens_cuda # set from input - attn_metadata.kv_lens_cuda = input_kv_lens_cuda # set from input - attn_metadata._seq_lens = torch.zeros([max_batch_size], device="cpu") - attn_metadata.host_request_types = torch.zeros([max_batch_size], device="cuda") - attn_metadata.spec_decoding_position_offsets = torch.zeros( - [max_batch_size, max_total_draft_tokens + 1], - dtype=torch.int, - device="cuda", - ) - attn_metadata.spec_decoding_packed_mask = torch.zeros( - [ - max_batch_size, - max_total_draft_tokens + 1, - math.ceil((max_total_draft_tokens + 1) / 32), - ], - dtype=torch.int, - device="cuda", - ) - attn_metadata.spec_decoding_generation_lengths = torch.zeros( - [max_batch_size], - dtype=torch.int, - device="cuda", - ) - - # 2) Create spec metadata - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - max_total_draft_tokens=max_total_draft_tokens, - speculative_model=eagle_model_dir, - eagle3_one_model=False, - eagle_choices=eagle_choices, - use_dynamic_tree=use_dynamic_tree, - ) - - eagle3_resource_manager = Eagle3ResourceManager( - spec_config, - torch.bfloat16, - 1024, - max_batch_size, - max_new_tokens, - max_new_tokens, - ) - - spec_tree_manager = SpecTreeManager( - max_num_requests=max_batch_size, - use_dynamic_tree=spec_config.use_dynamic_tree, - max_draft_len=spec_config.max_draft_len, - max_total_draft_tokens=spec_config.max_total_draft_tokens, - eagle_choices=spec_config.eagle_choices, - dynamic_tree_max_topK=spec_config.dynamic_tree_max_topK, - ) - - spec_metadata = Eagle3SpecMetadata( - max_draft_len=spec_config.max_draft_len, - spec_dec_mode=spec_config.spec_dec_mode, - max_num_requests=max_batch_size, - num_layers=32, - hidden_size=1024, - max_num_tokens=max_new_tokens, - dtype=torch.bfloat16, - is_draft_model=True, - eagle3_resource_manager=eagle3_resource_manager, - layers_to_capture=spec_config.eagle3_layers_to_capture, - max_total_draft_tokens=spec_config.max_total_draft_tokens, - eagle_choices=spec_config.eagle_choices, - is_spec_dec_tree=spec_config.eagle_choices is not None or spec_config.use_dynamic_tree, - is_spec_dec_dynamic_tree=spec_config.use_dynamic_tree, - ) - - # Set initial values - spec_metadata.num_accepted_draft_tokens = input_num_accepted_draft_tokens # set from input - spec_metadata.num_tokens = 0 - spec_metadata.hidden_states_write_indices = ( - input_hidden_states_write_indices # set from input - ) - spec_metadata.hidden_states_read_indices = ( - input_hidden_states_read_indices # set from input - ) - - # 3) Create StaticTreeDraftingLoopWrapper - static_tree_drafting_loop_wrapper = StaticTreeDraftingLoopWrapper( - max_batch_size=max_batch_size, - max_draft_len=max_draft_len, - max_total_draft_tokens=max_total_draft_tokens, - draft_model=DummyModel(), - ) - - # 3) Run the function - static_tree_drafting_loop_wrapper.prepare_for_generation( - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - spec_tree_manager=spec_tree_manager, - position_ids=input_position_ids, - ) - - # Compare input_ids and position_ids - print( - f"static_tree_drafting_loop_wrapper.position_ids_buffer: \ - {static_tree_drafting_loop_wrapper.position_ids_buffer}, \ - ref_output_position_ids: {ref_position_ids}" - ) - - # Compare the attention metadata - print( - f"attn_metadata.kv_lens_cuda: {attn_metadata.kv_lens_cuda}, \ - ref_attn_metadata.kv_lens_cuda: {ref_attn_metadata['kv_lens_cuda']}" - ) - print( - f"attn_metadata._seq_lens: {attn_metadata._seq_lens}, \ - ref_attn_metadata._seq_lens: {ref_attn_metadata['_seq_lens']}" - ) - print( - f"attn_metadata._seq_lens_cuda: {attn_metadata._seq_lens_cuda}, \ - ref_attn_metadata._seq_lens_cuda: {ref_attn_metadata['_seq_lens_cuda']}" - ) - print( - f"attn_metadata.host_request_types: {attn_metadata.host_request_types}, \ - ref_attn_metadata.host_request_types: {ref_attn_metadata['host_request_types']}" - ) - print( - f"attn_metadata.num_contexts: {attn_metadata.num_contexts}, \ - ref_attn_metadata.num_contexts: {ref_attn_metadata['num_contexts']}" - ) - print( - f"attn_metadata.spec_decoding_position_offsets: {attn_metadata.spec_decoding_position_offsets}, \ - ref_attn_metadata.spec_decoding_position_offsets: {ref_attn_metadata['spec_decoding_position_offsets']}" - ) - print( - f"attn_metadata.spec_decoding_packed_mask: {attn_metadata.spec_decoding_packed_mask}, \ - ref_attn_metadata.spec_decoding_packed_mask: {ref_attn_metadata['spec_decoding_packed_mask']}" - ) - print( - f"attn_metadata.spec_decoding_generation_lengths: {attn_metadata.spec_decoding_generation_lengths}, \ - ref_attn_metadata.spec_decoding_generation_lengths: {ref_attn_metadata['spec_decoding_generation_lengths']}" - ) - - # Compare the spec metadata - print( - f"spec_metadata.num_tokens: {spec_metadata.num_tokens}, \ - ref_spec_metadata.num_tokens: {ref_spec_metadata['num_tokens']}" - ) - print( - f"spec_metadata.hidden_states_read_indices: {spec_metadata.hidden_states_read_indices}, \ - ref_spec_metadata.hidden_states_read_indices: {ref_spec_metadata['hidden_states_read_indices']}" - ) - print( - f"spec_metadata.hidden_states_write_indices: {spec_metadata.hidden_states_write_indices}, \ - ref_spec_metadata.hidden_states_write_indices: {ref_spec_metadata['hidden_states_write_indices']}" - ) - - assert torch.all(static_tree_drafting_loop_wrapper.position_ids_buffer == ref_position_ids) - assert torch.all(attn_metadata.kv_lens_cuda == ref_attn_metadata["kv_lens_cuda"]) - assert torch.all(attn_metadata._seq_lens == ref_attn_metadata["_seq_lens"]) - assert torch.all(attn_metadata._seq_lens_cuda == ref_attn_metadata["_seq_lens_cuda"]) - assert torch.all( - attn_metadata.host_request_types == ref_attn_metadata["host_request_types"] - ) - assert torch.all( - torch.tensor(attn_metadata.num_contexts) - == torch.tensor(ref_attn_metadata["num_contexts"]) - ) - assert torch.all( - attn_metadata.spec_decoding_generation_lengths - == ref_attn_metadata["spec_decoding_generation_lengths"] - ) - assert torch.all( - attn_metadata.spec_decoding_position_offsets[:max_batch_size, :].reshape(-1) - == ref_attn_metadata["spec_decoding_position_offsets"].reshape(-1) - ) - assert torch.all( - attn_metadata.spec_decoding_packed_mask[:max_batch_size, :, :].reshape(-1) - == ref_attn_metadata["spec_decoding_packed_mask"].reshape(-1) - ) - assert torch.all( - torch.tensor(spec_metadata.num_tokens) == torch.tensor(ref_spec_metadata["num_tokens"]) - ) - - output_hidden_states_read_indices = spec_metadata.hidden_states_read_indices[ - : max_batch_size * (max_total_draft_tokens + 1) - ].reshape(max_batch_size, max_total_draft_tokens + 1) - assert torch.all( - # We do not compare the last element of the hidden_states_read_indices, because it is padding. - output_hidden_states_read_indices[:, :-1] - == ref_spec_metadata["hidden_states_read_indices"][:, :-1] - ) - - output_hidden_states_write_indices = spec_metadata.hidden_states_write_indices[ - : max_batch_size * (max_total_draft_tokens + 1) - ].reshape(max_batch_size, max_total_draft_tokens + 1) - assert torch.all( - output_hidden_states_write_indices == ref_spec_metadata["hidden_states_write_indices"] - ) - - ##### CASE 1 static tree, batch size = 1, prefill, prepare_for_layer_idx = 1 ############# - max_batch_size = 1 - prepare_for_layer_idx = 1 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - - prompt_len_1 = 15 - input_position_ids = torch.arange(prompt_len_1, dtype=torch.int32, device="cuda").reshape( - 1, prompt_len_1 - ) - input_seq_lens_cuda = torch.tensor([prompt_len_1], dtype=torch.int32, device="cuda") - input_kv_lens_cuda = torch.tensor([prompt_len_1], dtype=torch.int32, device="cuda") - input_num_accepted_draft_tokens = torch.tensor( - [prompt_len_1 - 1], dtype=torch.int32, device="cuda" - ) - input_hidden_states_write_indices = torch.zeros( - [max_new_tokens], dtype=torch.long, device="cuda" - ) - input_hidden_states_write_indices[:prompt_len_1] = torch.arange( - prompt_len_1, dtype=torch.long, device="cuda" - ) - input_hidden_states_read_indices = torch.zeros( - [max_new_tokens], dtype=torch.long, device="cuda" - ) - - ref_position_ids = torch.tensor( - [[15, 15, 15, 16, 16, 16, 16, 16, 16, 17, 17, 17, 0]], dtype=torch.int32, device="cuda" - ) - - ref_attn_metadata = {} - # prompt_len_1 + max_total_draft_tokens + 1 - ref_attn_metadata["kv_lens_cuda"] = torch.tensor([28], dtype=torch.int32, device="cuda") - - # max_total_draft_tokens + 1 - ref_attn_metadata["_seq_lens"] = torch.tensor([13], dtype=torch.int32, device="cpu") - - # max_total_draft_tokens + 1 - ref_attn_metadata["_seq_lens_cuda"] = torch.tensor([13], dtype=torch.int32, device="cuda") - - ref_attn_metadata["host_request_types"] = torch.tensor([0], dtype=torch.int32, device="cuda") - ref_attn_metadata["num_contexts"] = 0 - ref_attn_metadata["spec_decoding_position_offsets"] = torch.tensor( - [0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 0], dtype=torch.int32, device="cuda" - ).repeat(max_batch_size) # [max_batch_size * (max_total_draft_tokens + 1)] - ref_attn_metadata["spec_decoding_packed_mask"] = torch.tensor( - [1, 2, 4, 9, 17, 33, 66, 130, 260, 521, 1041, 2114, 0], dtype=torch.int32, device="cuda" - ).repeat(max_batch_size) # [max_batch_size * (max_total_draft_tokens + 1) * 1] - ref_attn_metadata["spec_decoding_generation_lengths"] = torch.tensor( - [13], dtype=torch.int32, device="cuda" - ) - - ref_spec_metadata = {} - ref_spec_metadata["num_tokens"] = 13 - ref_spec_metadata["hidden_states_read_indices"] = torch.tensor( - [[14, 14, 14, 15, 15, 15, 16, 16, 17, 18, 19, 21, 0]], dtype=torch.int32, device="cuda" - ) # [max_batch_size, max_total_draft_tokens + 1] - ref_spec_metadata["hidden_states_write_indices"] = torch.tensor( - [[15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27]], dtype=torch.int32, device="cuda" - ) # [max_batch_size, max_total_draft_tokens + 1] - - run_test( - max_batch_size, - prepare_for_layer_idx, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - input_seq_lens_cuda, - input_kv_lens_cuda, - input_num_accepted_draft_tokens, - input_hidden_states_write_indices, - input_hidden_states_read_indices, - input_position_ids, - ref_position_ids, - ref_attn_metadata, - ref_spec_metadata, - ) - - ##### CASE 2 static tree, batch size = 2, one prefill, one decode, prepare_for_layer_idx = 1 ##### - max_batch_size = 2 - prepare_for_layer_idx = 1 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - - prompt_len_1 = 15 # prefill - prompt_len_2 = 18 - seq_len_2 = ( - 3 + 1 - ) # accepted 2 draft tokens. For the 0-th drafter layer, the sequence length will be pad to max_draft_len + 1 - - input_position_ids_1 = torch.arange(prompt_len_1, dtype=torch.int32, device="cuda").reshape( - 1, prompt_len_1 - ) - input_position_ids_2 = torch.tensor([18, 19, 20, 21], dtype=torch.int32, device="cuda").reshape( - 1, max_draft_len + 1 - ) # for target model - input_position_ids = torch.cat([input_position_ids_1, input_position_ids_2], dim=1) - - input_seq_lens_cuda = torch.tensor([prompt_len_1, seq_len_2], dtype=torch.int32, device="cuda") - input_kv_lens_cuda = torch.tensor( - [prompt_len_1, prompt_len_2 + seq_len_2], dtype=torch.int32, device="cuda" - ) - input_num_accepted_draft_tokens = torch.tensor( - [prompt_len_1 - 1, 2], dtype=torch.int32, device="cuda" - ) # Suppose 2 are received. - input_hidden_states_write_indices = torch.zeros( - [max_new_tokens], dtype=torch.long, device="cuda" - ) - input_hidden_states_write_indices[: prompt_len_1 + seq_len_2] = torch.arange( - prompt_len_1 + seq_len_2, dtype=torch.long, device="cuda" - ) - input_hidden_states_read_indices = torch.zeros( - [max_new_tokens], dtype=torch.long, device="cuda" - ) - ref_position_ids = torch.tensor( - [ - [15, 15, 15, 16, 16, 16, 16, 16, 16, 17, 17, 17, 0], - [21, 21, 21, 22, 22, 22, 22, 22, 22, 23, 23, 23, 0], - ], - dtype=torch.int32, - device="cuda", - ) - - ref_attn_metadata = {} - ref_attn_metadata["kv_lens_cuda"] = torch.tensor([28, 34], dtype=torch.int32, device="cuda") - ref_attn_metadata["_seq_lens"] = torch.tensor([13, 13], dtype=torch.int32, device="cpu") - ref_attn_metadata["_seq_lens_cuda"] = torch.tensor([13, 13], dtype=torch.int32, device="cuda") - ref_attn_metadata["host_request_types"] = torch.tensor([0, 0], dtype=torch.int32, device="cuda") - ref_attn_metadata["num_contexts"] = 0 - ref_attn_metadata["spec_decoding_position_offsets"] = torch.tensor( - [0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 0], dtype=torch.int32, device="cuda" - ).repeat(max_batch_size) - ref_attn_metadata["spec_decoding_packed_mask"] = torch.tensor( - [1, 2, 4, 9, 17, 33, 66, 130, 260, 521, 1041, 2114, 0], dtype=torch.int32, device="cuda" - ).repeat(max_batch_size) - ref_attn_metadata["spec_decoding_generation_lengths"] = torch.tensor( - [13, 13], dtype=torch.int32, device="cuda" - ) - - ref_spec_metadata = {} - ref_spec_metadata["num_tokens"] = 26 - ref_spec_metadata["hidden_states_read_indices"] = torch.tensor( - [ - [14, 14, 14, 15, 15, 15, 16, 16, 17, 18, 19, 21, 0], - [17, 17, 17, 18, 18, 18, 19, 19, 20, 21, 22, 24, 0], - ], - dtype=torch.int32, - device="cuda", - ) - ref_spec_metadata["hidden_states_write_indices"] = torch.tensor( - [ - [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27], - [18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30], - ], - dtype=torch.int32, - device="cuda", - ) - - run_test( - max_batch_size, - prepare_for_layer_idx, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - input_seq_lens_cuda, - input_kv_lens_cuda, - input_num_accepted_draft_tokens, - input_hidden_states_write_indices, - input_hidden_states_read_indices, - input_position_ids, - ref_position_ids, - ref_attn_metadata, - ref_spec_metadata, - ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_sampling.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_sampling.py deleted file mode 100644 index fd20a445d035..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_sampling.py +++ /dev/null @@ -1,939 +0,0 @@ -import os -import sys -import unittest - -import torch -from utils.llm_data import llm_models_root - -from tensorrt_llm._torch.speculative.drafting_loops import StaticTreeDraftingLoopWrapper -from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager -from tensorrt_llm.llmapi import Eagle3DecodingConfig - -sys.path.append(os.path.join(os.path.dirname(__file__), "..")) - - -class DummyModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.model_config = None - self.config = None - self.model = {} - - def forward(self, *args, **kwargs) -> torch.Tensor: - pass - - -def test_draft_token_static_tree_sampling(): - # Fix parameters - models_path = llm_models_root() - eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" # It will not actually be used. - use_dynamic_tree = False - - # Create related object and run test - def run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - use_cuda_graph, - ref_new_tokens, - ): - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - max_total_draft_tokens=max_total_draft_tokens, - speculative_model=eagle_model_dir, - eagle3_one_model=False, - eagle_choices=eagle_choices, - use_dynamic_tree=use_dynamic_tree, - ) - spec_tree_manager = SpecTreeManager( - max_num_requests=max_batch_size, - use_dynamic_tree=spec_config.use_dynamic_tree, - max_draft_len=spec_config.max_draft_len, - max_total_draft_tokens=spec_config.max_total_draft_tokens, - eagle_choices=spec_config.eagle_choices, - dynamic_tree_max_topK=spec_config.dynamic_tree_max_topK, - ) - - # Create the chain drafter - tree_drafter = StaticTreeDraftingLoopWrapper( - max_batch_size=max_batch_size, - max_draft_len=spec_config.max_draft_len, - max_total_draft_tokens=spec_config.max_total_draft_tokens, - draft_model=DummyModel(), - ) - - sampled_tokens = tree_drafter.sample( - logits=logits, - max_top_k=spec_tree_manager.max_top_k, - ) - - tree_drafter.extract_real_draft_tokens( - cur_draft_idx=draft_layer_id, - batch_size=max_batch_size, - new_draft_tokens=sampled_tokens, - use_cuda_graph=use_cuda_graph, # use torch op or not - spec_tree_manager=spec_tree_manager, - ) - - real_new_draft_tokens = tree_drafter.draft_tokens_buffer[:max_batch_size, :] - - print(f"ref_new_tokens.shape: {ref_new_tokens.shape}, ref_new_tokens: {ref_new_tokens}") - print( - f"real_new_draft_tokens.shape: {real_new_draft_tokens.shape}, output_tokens: {real_new_draft_tokens}" - ) - assert torch.all(real_new_draft_tokens == ref_new_tokens) - - ################## CASE 1 static tree, batch size = 1, draft_layer_id = 0 ########################## - max_batch_size = 1 - draft_layer_id = 0 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - - logits = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top3 indices = [4, 1, 9] - ], - device="cuda", - ) - ref_new_tokens = torch.tensor( - [ - [4, 1, 9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 2 static tree, batch size = 1, draft_layer_id = 1 ########################## - max_batch_size = 1 - draft_layer_id = 1 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - - logits = torch.empty((max_batch_size, max_total_draft_tokens + 1, 10), device="cuda") - set_indices = torch.tensor([0, 1, 2], device="cuda") - logits[:, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top3 indices = [4, 1, 9] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 1.2, 0.9, 1.0], # top2 indices = [7, 1] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2], # top1 indices = [9] - ], - device="cuda", - ) - ref_new_tokens = torch.tensor( - [ - [0, 0, 0, 4, 1, 9, 7, 1, 9, 0, 0, 0, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 3 static tree, batch size = 1, draft_layer_id = 2 ########################## - max_batch_size = 1 - draft_layer_id = 2 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - logits = torch.empty((max_batch_size, max_total_draft_tokens + 1, 10), device="cuda") - set_indices = torch.tensor([3, 4, 6], device="cuda") - logits[:, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top1 indices = [4] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 1.2, 0.9, 1.0], # top1 indices = [7] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2], # top1 indices = [9] - ], - device="cuda", - ) - ref_new_tokens = torch.tensor( - [ - [0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 7, 9, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 4 static tree, batch size = 2, draft_layer_id = 0 ########################## - max_batch_size = 2 - draft_layer_id = 0 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - logits = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top3 indices = [4, 1, 9] - [0.1, 0.3, 1.1, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top3 indices = [4, 2, 9] - ], - device="cuda", - ) - ref_new_tokens = torch.tensor( - [ - [4, 1, 9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [4, 2, 9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 5 static tree, batch size = 2, draft_layer_id = 1 ########################## - max_batch_size = 2 - draft_layer_id = 1 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - logits = torch.empty((max_batch_size, max_total_draft_tokens + 1, 10), device="cuda") - set_indices = torch.tensor([0, 1, 2], device="cuda") - - logits[0, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top3 indices = [4, 1, 9] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 1.2, 0.9, 1.0], # top2 indices = [7, 1] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2], # top1 indices = [9] - ], - device="cuda", - ) - - logits[1, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 0.6, 1.2, 0.7, 0.8, 1.0, 0.9], # top3 indices = [5, 1, 8] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 1.2, 0.8, 0.9, 1.0], # top2 indices = [6, 1] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.2, 1.0], # top1 indices = [8] - ], - device="cuda", - ) - ref_new_tokens = torch.tensor( - [ - [0, 0, 0, 4, 1, 9, 7, 1, 9, 0, 0, 0, 0], - [0, 0, 0, 5, 1, 8, 6, 1, 8, 0, 0, 0, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 6 static tree, batch size = 2, draft_layer_id = 2 ########################## - max_batch_size = 2 - draft_layer_id = 2 - max_total_draft_tokens = 12 - max_draft_len = 3 - eagle_choices = [ - [0], - [1], - [2], - [0, 0], - [0, 1], - [0, 2], - [1, 0], - [1, 1], - [2, 0], - [0, 0, 0], - [0, 1, 1], - [1, 0, 0], - ] - logits = torch.empty((max_batch_size, max_total_draft_tokens + 1, 10), device="cuda") - set_indices = torch.tensor([3, 4, 6], device="cuda") - - logits[0, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 1.2, 0.6, 0.7, 0.8, 0.9, 1.0], # top1 indices = [4] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 1.2, 0.9, 1.0], # top1 indices = [7] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2], # top1 indices = [9] - ], - device="cuda", - ) - - logits[1, set_indices, :] = torch.tensor( - [ - [0.1, 1.1, 0.3, 0.4, 0.6, 1.2, 0.7, 0.8, 0.9, 1.0], # top1 indices = [5] - [0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.2, 1.0], # top1 indices = [8] - [1.2, 0.1, 1.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0], # top1 indices = [0] - ], - device="cuda", - ) - - ref_new_tokens = torch.tensor( - [ - [0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 7, 9, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 0, 5, 8, 0, 0], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 7 static tree, batch size = 1, draft_layer_id = 0, bigger tree ########################## - max_batch_size = 1 - draft_layer_id = 0 - max_total_draft_tokens = 63 - max_draft_len = 4 - eagle_choices = [ - [0], - [0, 0], - [1], - [0, 1], - [2], - [0, 0, 0], - [1, 0], - [0, 2], - [3], - [0, 3], - [4], - [0, 4], - [2, 0], - [0, 5], - [0, 0, 1], - [5], - [0, 6], - [6], - [0, 7], - [0, 1, 0], - [1, 1], - [7], - [0, 8], - [0, 0, 2], - [3, 0], - [0, 9], - [8], - [9], - [1, 0, 0], - [0, 2, 0], - [1, 2], - [0, 0, 3], - [4, 0], - [2, 1], - [0, 0, 4], - [0, 0, 5], - [0, 0, 0, 0], - [0, 1, 1], - [0, 0, 6], - [0, 3, 0], - [5, 0], - [1, 3], - [0, 0, 7], - [0, 0, 8], - [0, 0, 9], - [6, 0], - [0, 4, 0], - [1, 4], - [7, 0], - [0, 1, 2], - [2, 0, 0], - [3, 1], - [2, 2], - [8, 0], - [0, 5, 0], - [1, 5], - [1, 0, 1], - [0, 2, 1], - [9, 0], - [0, 6, 0], - [0, 0, 0, 1], - [1, 6], - [0, 7, 0], - ] # mc_sim_7b_63 - - logits = torch.tensor( - [ - [ - 1.1, - 1.9, - 0.7, - 0.3, - 4.4, - 4.3, - 4.9, - 1.2, - 3.9, - 4.7, - 1.4, - 2.5, - 3.2, - 0.8, - 2.0, - 3.0, - 1.8, - 2.3, - 4.2, - 1.3, - ] - ], - device="cuda", - ) - - ref_new_tokens = torch.zeros((max_batch_size, max_total_draft_tokens + 1), device="cuda") - - ref_new_tokens[0, :10] = torch.tensor( - [ - [6, 9, 4, 5, 18, 8, 12, 15, 11, 17], - ], - device="cuda", - ) # shape: [max_batch_size, max_total_draft_tokens + 1] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - ################## CASE 8 static tree, batch size = 1, draft_layer_id = 1, bigger tree ########################## - max_batch_size = 1 - draft_layer_id = 1 - max_total_draft_tokens = 63 - max_draft_len = 4 - eagle_choices = [ - [0], - [0, 0], - [1], - [0, 1], - [2], - [0, 0, 0], - [1, 0], - [0, 2], - [3], - [0, 3], - [4], - [0, 4], - [2, 0], - [0, 5], - [0, 0, 1], - [5], - [0, 6], - [6], - [0, 7], - [0, 1, 0], - [1, 1], - [7], - [0, 8], - [0, 0, 2], - [3, 0], - [0, 9], - [8], - [9], - [1, 0, 0], - [0, 2, 0], - [1, 2], - [0, 0, 3], - [4, 0], - [2, 1], - [0, 0, 4], - [0, 0, 5], - [0, 0, 0, 0], - [0, 1, 1], - [0, 0, 6], - [0, 3, 0], - [5, 0], - [1, 3], - [0, 0, 7], - [0, 0, 8], - [0, 0, 9], - [6, 0], - [0, 4, 0], - [1, 4], - [7, 0], - [0, 1, 2], - [2, 0, 0], - [3, 1], - [2, 2], - [8, 0], - [0, 5, 0], - [1, 5], - [1, 0, 1], - [0, 2, 1], - [9, 0], - [0, 6, 0], - [0, 0, 0, 1], - [1, 6], - [0, 7, 0], - ] # mc_sim_7b_63 - - logits = torch.empty((max_batch_size, max_total_draft_tokens + 1, 20), device="cuda") - set_indices = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], device="cuda") - - logits[0, set_indices, :] = torch.tensor( - [ - [ - 3.8, - 2.5, - 1.9, - 0.2, - 4.5, - 0.9, - 0.5, - 3.9, - 4.0, - 2.9, - 2.7, - 0.7, - 2.8, - 2.1, - 1.2, - 1.4, - 3.3, - 3.7, - 2.4, - 5.4, - ], - [ - 3.2, - 5.4, - 0.1, - 4.6, - 1.9, - 4.8, - 3.8, - 0.9, - 5.3, - 3.7, - 4.4, - 0.4, - 4.3, - 1.3, - 0.3, - 3.6, - 2.5, - 6.0, - 3.1, - 3.9, - ], - [ - 4.1, - 2.8, - 3.9, - 4.2, - 5.1, - 5.8, - 3.7, - 1.5, - 0.8, - 1.4, - 2.2, - 0.0, - 1.9, - 2.5, - 3.2, - 0.4, - 3.5, - 5.0, - 2.9, - 4.3, - ], - [ - 0.7, - 1.4, - 0.6, - 2.0, - 3.0, - 0.9, - 1.7, - 0.0, - 3.4, - 2.8, - 4.7, - 3.6, - 5.6, - 4.2, - 1.1, - 4.5, - 1.2, - 1.3, - 3.9, - 3.2, - ], - [ - 3.6, - 5.5, - 4.0, - 6.0, - 4.7, - 2.3, - 5.8, - 0.8, - 3.2, - 0.3, - 4.8, - 2.1, - 0.0, - 3.1, - 0.1, - 2.5, - 2.2, - 5.0, - 3.4, - 3.9, - ], - [ - 3.3, - 3.8, - 2.5, - 2.4, - 6.0, - 5.5, - 3.6, - 0.6, - 1.8, - 2.6, - 4.7, - 1.2, - 3.9, - 4.0, - 4.8, - 5.0, - 1.1, - 1.9, - 0.4, - 0.3, - ], - [ - 3.1, - 1.7, - 0.0, - 4.6, - 3.7, - 3.2, - 2.8, - 4.7, - 2.7, - 3.8, - 6.0, - 4.2, - 1.1, - 3.4, - 5.5, - 2.1, - 2.6, - 1.8, - 0.1, - 1.3, - ], - [ - 3.6, - 5.1, - 2.3, - 5.3, - 3.5, - 1.8, - 2.7, - 2.4, - 3.2, - 3.0, - 4.8, - 5.9, - 2.9, - 1.4, - 5.7, - 0.9, - 1.0, - 1.3, - 2.0, - 3.9, - ], - [ - 3.9, - 0.2, - 4.5, - 1.5, - 1.4, - 0.1, - 4.3, - 0.3, - 3.2, - 3.4, - 0.4, - 3.8, - 4.8, - 3.5, - 5.6, - 1.9, - 2.8, - 4.1, - 1.0, - 5.0, - ], - [ - 5.9, - 1.8, - 3.5, - 0.3, - 1.7, - 2.5, - 3.0, - 3.3, - 2.1, - 4.3, - 6.0, - 0.1, - 0.5, - 0.9, - 4.0, - 5.4, - 1.1, - 2.7, - 5.3, - 3.4, - ], - ], - device="cuda", - ) - - ref_new_tokens = torch.zeros((max_batch_size, max_total_draft_tokens + 1), device="cuda") - ref_new_tokens[0, 10:38] = torch.tensor( - [ - [ - 19, - 4, - 8, - 7, - 0, - 17, - 16, - 9, - 12, - 10, # top-10 - 17, - 1, - 8, - 5, - 3, - 10, - 12, # top-7 - 5, - 4, - 17, # top-3 - 12, - 10, # top-2 - 3, # top-1 - 4, # top-1 - 10, # top-1 - 11, # top-1 - 14, # top-1 - 10, # top-1 - ] - ], - device="cuda", - ) # shape: [max_batch_size, num_new_draft_tokens] - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - False, - ref_new_tokens, - ) - - run_test( - max_batch_size, - draft_layer_id, - max_total_draft_tokens, - max_draft_len, - eagle_choices, - logits, - True, - ref_new_tokens, - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_verification.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_verification.py index 81d04e037e8f..026017d24dff 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_verification.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_verification.py @@ -35,7 +35,6 @@ def run_test( max_draft_len=max_draft_len, max_total_draft_tokens=max_total_draft_tokens, speculative_model=eagle_model_dir, - eagle3_one_model=False, eagle_choices=eagle_choices, use_dynamic_tree=use_dynamic_tree, ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_spec_decode.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_spec_decode.py index 48a02c548949..54a7d7351754 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_spec_decode.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_spec_decode.py @@ -56,7 +56,6 @@ def mock_get_draft_len_for_batch_size(draft_len_schedule, batch_size, max_total_ spec_config = Eagle3DecodingConfig( max_draft_len=4, speculative_model=eagle_model_dir, - eagle3_one_model=True, draft_len_schedule={ 1: 4 }, # It doesn't matter which value is used here, as the draft length will be controlled by the mock function. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_get_num_spec_layers.py b/tests/unittest/_torch/speculative/hw_agnostic/test_get_num_spec_layers.py index 3ee5cc7f7a2e..1d6e1646ca19 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_get_num_spec_layers.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_get_num_spec_layers.py @@ -11,7 +11,7 @@ def __init__(self, num_hidden_layers: int): self.num_hidden_layers = num_hidden_layers -def test_get_num_spec_layers_eagle3_one_model_uses_draft_hidden_layers(): +def test_get_num_spec_layers_eagle3_uses_draft_hidden_layers(): spec_config = Eagle3DecodingConfig( max_draft_len=8, speculative_model="/path/to/draft", @@ -22,7 +22,7 @@ def test_get_num_spec_layers_eagle3_one_model_uses_draft_hidden_layers(): assert get_num_spec_layers(spec_config) == 1 -def test_get_num_spec_layers_eagle3_one_model_defaults_to_one(): +def test_get_num_spec_layers_eagle3_defaults_to_one(): spec_config = Eagle3DecodingConfig( max_draft_len=8, speculative_model="/path/to/draft", @@ -32,7 +32,7 @@ def test_get_num_spec_layers_eagle3_one_model_defaults_to_one(): assert get_num_spec_layers(spec_config) == 1 -def test_get_num_spec_layers_mtp_eagle_one_model(): +def test_get_num_spec_layers_mtp_eagle(): spec_config = MTPDecodingConfig(max_draft_len=1) assert get_num_spec_layers(spec_config) == 1 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py index 3fb5a201c25b..351170e6d25a 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py @@ -54,7 +54,6 @@ def test_kv_cache_reuse(use_cuda_graph: bool, attn_backend: str): spec_config = Eagle3DecodingConfig( max_draft_len=max_draft_len, speculative_model=eagle_model_dir, - eagle3_one_model=False, ) llm_spec = LLM(**llm_common_config, speculative_config=spec_config) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py index a381dbd94907..9a0a7555dec7 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py @@ -1789,7 +1789,6 @@ def test_mtp_eagle_one_model_rejection(): spec_config = MTPDecodingConfig( max_draft_len=1, use_mtp_vanilla=False, - mtp_eagle_one_model=True, speculative_model=model_dir, use_rejection_sampling=True, ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py b/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py index cb0001e0b718..953b85a59886 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py @@ -59,7 +59,6 @@ def test_spec_gate_e2e(enforce_single_worker): spec_config = Eagle3DecodingConfig( max_draft_len=max_draft_len, speculative_model=eagle_model_dir, - eagle3_one_model=True, acceptance_rate_window_size=acceptance_rate_window_size, acceptance_rate_threshold=acceptance_rate_threshold, ) diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 33cf3d635661..50a14e82d4d1 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -96,7 +96,7 @@ def test_eagle3_target_kv_cache_extends_full_window_for_spec_decode() -> None: @skip_num_gpus_less_than(1) -def test_eagle3_one_model_capture_uses_real_token_count() -> None: +def test_eagle3_capture_uses_real_token_count() -> None: # maybe_capture_hidden_states routes through inplace_slice_copy, a CUDA-only # custom op, so this test runs on GPU. It verifies that the capture is bounded # by the real scheduled token count (num_capture_tokens) and not self.num_tokens. @@ -141,10 +141,10 @@ def enforce_single_worker(monkeypatch): yield -def test_kv_lens_runtime_with_eagle3_one_model(): +def test_kv_lens_runtime_with_eagle3(): """ Validates that kv_lens_runtime correctly excludes num_extra_kv_tokens when - preparing attention metadata during EAGLE3 one-model speculative decoding. + preparing attention metadata during EAGLE3 speculative decoding. Background: - EAGLE3 reserves num_extra_kv_tokens = max_draft_len - 1 in KV cache for draft token management @@ -284,72 +284,19 @@ def test_block_offsets_staging_width_spec_gate(spec_signal): @pytest.mark.parametrize( - "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,use_chain_drafter,multi_batch,attention_dp,use_hf_speculative_model", + "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,enable_chunked_prefill,multi_batch,attention_dp,use_hf_speculative_model", [ - [True, "TRTLLM", True, False, False, False, True, False, False, False], - [True, "TRTLLM", True, False, False, False, False, False, False, False], - [False, "TRTLLM", True, False, False, False, True, False, False, False], - [ - False, "TRTLLM", True, False, False, False, False, False, False, - False - ], - [ - True, "FLASHINFER", True, False, False, False, True, False, False, - False - ], - [ - False, "FLASHINFER", True, False, False, False, True, False, False, - False - ], - [False, "TRTLLM", False, True, True, False, True, False, False, False], - [True, "TRTLLM", False, True, True, False, True, False, False, False], - [True, "TRTLLM", True, False, True, True, True, False, False, False], - [True, "TRTLLM", True, False, True, False, True, False, False, False], - [True, "TRTLLM", True, False, False, True, True, False, False, False], - [True, "TRTLLM", False, False, False, False, True, False, False, False], - [ - False, "TRTLLM", False, False, False, False, True, False, False, - False - ], - [True, "TRTLLM", False, False, False, False, False, True, False, False], - [True, "TRTLLM", False, False, False, False, False, True, True, False], - [ - False, "TRTLLM", False, False, False, False, False, True, False, - False - ], - [True, "TRTLLM", False, False, False, False, True, True, False, False], - [False, "TRTLLM", False, False, False, False, True, True, False, False], - [ - True, "TRTLLM", False, False, False, False, False, False, False, - False - ], - [ - False, "TRTLLM", False, False, False, False, False, False, False, - False - ], - [True, "TRTLLM", False, False, False, True, True, False, False, False], - [True, "TRTLLM", False, False, False, True, False, False, False, False], - [ - True, "FLASHINFER", False, False, False, False, True, False, False, - False - ], - [ - False, "FLASHINFER", False, False, False, False, True, False, False, - False - ], - # Tests (mocked) speculative model auto-download from HuggingFace - [False, "TRTLLM", True, False, False, False, True, False, False, True], + [False, "TRTLLM", False, True, False, False, False, False], + [True, "TRTLLM", False, True, False, False, False, False], + [True, "TRTLLM", True, False, True, False, False, False], + [True, "TRTLLM", True, False, False, False, False, False], ]) @pytest.mark.high_cuda_memory @with_mocked_hf_download_for_single_gpu def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, disable_overlap_scheduler: bool, enable_block_reuse: bool, - use_one_model: bool, enable_chunked_prefill: bool, - use_chain_drafter: bool, multi_batch: bool, + enable_chunked_prefill: bool, multi_batch: bool, attention_dp: bool, use_hf_speculative_model: bool): - if not use_one_model: - pytest.skip("Two model Eagle3 is deprecated") - # Eagle3 one model works with overlap scheduler and block reuse. total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if total_mem_gb < 35: @@ -395,10 +342,7 @@ def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, spec_config = Eagle3DecodingConfig( max_draft_len=max_draft_len, speculative_model=eagle_model, - # Llama 3 does not support one model eagle. - eagle3_one_model=use_one_model, ) - spec_config._allow_chain_drafter = use_chain_drafter # Create the LLM instance llm_spec = LLM(**llm_common_config, speculative_config=spec_config) @@ -457,8 +401,7 @@ def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, assert text_spec == text_ref -@pytest.mark.parametrize("eagle3_one_model", [True, False]) -def test_eagle3_spec_decoding_stats(eagle3_one_model): +def test_eagle3_spec_decoding_stats(): """Test that specDecodingStats are correctly populated in metrics endpoint""" models_path = llm_models_root() eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" @@ -474,14 +417,12 @@ def test_eagle3_spec_decoding_stats(eagle3_one_model): spec_config = Eagle3DecodingConfig( max_draft_len=3, speculative_model=eagle_model_dir, - eagle3_one_model=eagle3_one_model, ) with LLM( model=target_model_dir, speculative_config=spec_config, kv_cache_config=kv_cache_config, - disable_overlap_scheduler=not eagle3_one_model, enable_iter_perf_stats=True, max_batch_size=4, ) as llm: @@ -554,7 +495,6 @@ def test_llama_eagle3_long_prompt(use_cuda_graph): spec_config = Eagle3DecodingConfig( max_draft_len=3, speculative_model=eagle_model_dir, - eagle3_one_model=False, ) if use_cuda_graph: @@ -597,7 +537,6 @@ def test_deepseek_eagle3(): attn_backend = "TRTLLM" disable_overlap_scheduler = False enable_block_reuse = False - use_one_model = False enable_chunked_prefill = False # Eagle3 one model works with overlap scheduler and block reuse. @@ -675,13 +614,10 @@ def test_deepseek_eagle3(): enable_chunked_prefill=enable_chunked_prefill, ) - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - speculative_model=eagle_model_dir, - # Llama 3 does not support one model eagle. - eagle3_one_model=use_one_model, - eagle3_layers_to_capture={29}, - load_format="dummy") + spec_config = Eagle3DecodingConfig(max_draft_len=max_draft_len, + speculative_model=eagle_model_dir, + eagle3_layers_to_capture={29}, + load_format="dummy") llm_spec = LLM(**llm_common_config, speculative_config=spec_config) @@ -699,7 +635,6 @@ def test_deepseek_mla_eagle3(): attn_backend = "TRTLLM" disable_overlap_scheduler = False enable_block_reuse = False - use_one_model = True enable_chunked_prefill = False # Eagle3 one model works with overlap scheduler and block reuse. @@ -788,7 +723,6 @@ def test_deepseek_mla_eagle3(): spec_config = Eagle3DecodingConfig(max_draft_len=max_draft_len, speculative_model=eagle_model_dir, - eagle3_one_model=use_one_model, load_format="dummy") llm_spec = LLM(**llm_common_config, speculative_config=spec_config) @@ -802,8 +736,7 @@ def test_deepseek_mla_eagle3(): pass -@pytest.mark.parametrize("use_one_model", [True, False]) -def test_multi_eagle3(use_one_model: bool): +def test_multi_eagle3(): use_cuda_graph = True attn_backend = "TRTLLM" disable_overlap_scheduler = False @@ -886,7 +819,6 @@ def test_multi_eagle3(use_one_model: bool): spec_config = Eagle3DecodingConfig(max_draft_len=max_draft_len, speculative_model=eagle_model_dir, - eagle3_one_model=use_one_model, load_format="dummy") llm_spec = LLM(**llm_common_config, speculative_config=spec_config) @@ -910,7 +842,6 @@ def test_eagle3_cuda_graph_padding(disable_overlap_scheduler: bool): """ attn_backend = "TRTLLM" enable_block_reuse = False - use_one_model = False enable_chunked_prefill = False total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 @@ -943,7 +874,6 @@ def test_eagle3_cuda_graph_padding(disable_overlap_scheduler: bool): spec_config = Eagle3DecodingConfig( max_draft_len=max_draft_len, speculative_model=eagle_model_dir, - eagle3_one_model=use_one_model, ) # Create the LLM instance @@ -965,7 +895,6 @@ def test_eagle3_cdl_sampling(disable_overlap_scheduler: bool): """Test CDL sampling with 2 requests and max_batch_size=2.""" attn_backend = "TRTLLM" enable_block_reuse = False - use_one_model = False enable_chunked_prefill = False total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 @@ -997,7 +926,6 @@ def test_eagle3_cdl_sampling(disable_overlap_scheduler: bool): spec_config = Eagle3DecodingConfig( max_draft_len=max_draft_len, speculative_model=eagle_model_dir, - eagle3_one_model=use_one_model, ) # Create the LLM instance @@ -1054,7 +982,6 @@ def test_llama_eagle3_rejection_sampling_modes(use_dynamic_tree: bool, spec_config_kwargs = dict( max_draft_len=max_draft_len, speculative_model=eagle_model, - eagle3_one_model=True, use_rejection_sampling=True, ) if use_dynamic_tree: @@ -1165,7 +1092,6 @@ def test_eagle3_lora(use_cuda_graph: bool): """ attn_backend = "TRTLLM" enable_block_reuse = False - use_one_model = True enable_chunked_prefill = False total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 @@ -1201,7 +1127,6 @@ def test_eagle3_lora(use_cuda_graph: bool): spec_config = Eagle3DecodingConfig( max_draft_len=max_draft_len, speculative_model=eagle_model_dir, - eagle3_one_model=use_one_model, ) # Create the LLM instance @@ -1258,7 +1183,6 @@ def test_llama_eagle3_dynamic_tree(use_cuda_graph: bool, spec_config = Eagle3DecodingConfig( max_draft_len=max_draft_len, speculative_model=eagle_model, - eagle3_one_model=True, use_dynamic_tree=True, dynamic_tree_max_topK=dynamic_tree_max_topK, max_total_draft_tokens=max_total_draft_tokens, 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 5639344dd18c..b6175e1814a3 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_llm_config.py @@ -285,7 +285,6 @@ def test_accepts_eagle_one_model(self): spec_config = EagleDecodingConfig( max_draft_len=3, speculative_model="some/model", - eagle3_one_model=True, ) # Should not raise. args = LlmArgs( @@ -295,12 +294,11 @@ def test_accepts_eagle_one_model(self): ) assert args.model_factory == "eagle_one_model" - def test_accepts_mtp_eagle_one_model(self): + def test_accepts_mtp_eagle(self): from tensorrt_llm.llmapi import MTPDecodingConfig spec_config = MTPDecodingConfig( num_nextn_predict_layers=3, - mtp_eagle_one_model=True, ) # Should not raise. args = LlmArgs( @@ -317,7 +315,6 @@ def test_rejects_flashinfer_cuda_graph_backend(self, compile_backend): spec_config = EagleDecodingConfig( max_draft_len=3, speculative_model="some/model", - eagle3_one_model=True, ) with pytest.raises(pydantic.ValidationError): @@ -334,7 +331,6 @@ def test_accepts_flashinfer_torch_simple(self): spec_config = EagleDecodingConfig( max_draft_len=3, speculative_model="some/model", - eagle3_one_model=True, ) LlmArgs( @@ -363,7 +359,7 @@ def test_ssm_replay_without_spec_raises(self): def test_ssm_replay_with_spec_ok(self): from tensorrt_llm.llmapi import MTPDecodingConfig - spec_config = MTPDecodingConfig(num_nextn_predict_layers=3, mtp_eagle_one_model=True) + spec_config = MTPDecodingConfig(num_nextn_predict_layers=3) # Replay + spec is valid and must not raise. args = LlmArgs( model="test-model", diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py index 245cc79ec1f2..25f4b1facb05 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py @@ -67,7 +67,6 @@ def test_super_mtp_smoke(): experiment_config["args"]["world_size"] = 1 experiment_config["args"]["speculative_config"] = MTPDecodingConfig( num_nextn_predict_layers=3, - mtp_eagle_one_model=True, speculative_model=model_path, ) # Shrink the Eagle/MTP drafter model to match the target's reduced dimensions. @@ -139,7 +138,6 @@ def test_super_mtp_ssm_replay_smoke(): experiment_config["args"]["world_size"] = 1 experiment_config["args"]["speculative_config"] = MTPDecodingConfig( num_nextn_predict_layers=3, - mtp_eagle_one_model=True, speculative_model=model_path, ) experiment_config["args"]["speculative_model_kwargs"] = experiment_config["args"][ @@ -188,7 +186,6 @@ def test_kv_cache_extra_seq_len_for_spec_dec(): spec_config = Eagle3DecodingConfig( max_draft_len=3, speculative_model="some/model", - eagle3_one_model=True, ) args_eagle = LlmArgs( model="meta-llama/Meta-Llama-3.1-8B-Instruct", @@ -221,7 +218,6 @@ def test_mtp_autodeploy_uses_eagle_one_model_capture(): model=model, speculative_config=MTPDecodingConfig( num_nextn_predict_layers=3, - mtp_eagle_one_model=True, ), transforms=piecewise_disabled_transforms(), ) @@ -232,7 +228,7 @@ def test_mtp_autodeploy_uses_eagle_one_model_capture(): assert args.transforms["detect_hidden_states_for_capture"]["eagle3_layers_to_capture"] == {-1} -def test_detect_hidden_states_capture_last_layer_for_mtp_eagle_one_model(): +def test_detect_hidden_states_capture_last_layer_for_mtp_eagle(): from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs config = get_small_model_config("meta-llama/Meta-Llama-3.1-8B-Instruct") @@ -244,7 +240,6 @@ def test_detect_hidden_states_capture_last_layer_for_mtp_eagle_one_model(): **config["args"], speculative_config=MTPDecodingConfig( num_nextn_predict_layers=3, - mtp_eagle_one_model=True, speculative_model=config["args"]["model"], ), ) diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py index 10b802fe7173..f590ebfb69e5 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py @@ -273,7 +273,6 @@ def test_autodeploy_disaggregated_eagle3_smoke(): speculative_config = Eagle3DecodingConfig( max_draft_len=3, speculative_model=eagle3_model_config["args"]["model"], - eagle3_one_model=True, eagle3_layers_to_capture={0, 1, 2}, ) speculative_model_kwargs = { diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 558a57c70208..f384aedf7823 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1563,7 +1563,7 @@ class TestPiecewiseCudaGraphCaptureDefaults: entries to the reachable ceiling and never invents sizes beyond this list. 2. `_filter_piecewise_capture_num_tokens` caps the candidate list at - `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` -- + `max_batch_size * (max_seq_len - 1)` -- the largest forward-pass `num_tokens` the warmup builder can construct, since every in-flight request must leave room for at least one decode token. @@ -1724,42 +1724,6 @@ def test_piecewise_filter_keeps_all_entries_when_within_ceiling(self): assert kept.count(128) == 1 assert unrecordable == [] - def test_piecewise_filter_subtracts_extra_decoding_steps(self): - """Subtract `num_extra_decoding_steps` from the ceiling. - - Drafting loops consume extra decode steps; the filter must mirror - the `max_seq_len - 1 - num_extra_decoding_steps` constraint - applied when warmup requests are built. Candidates above the - reduced ceiling are clamped down to it; nothing is appended. - """ - from tensorrt_llm._torch.pyexecutor.model_engine import \ - _filter_piecewise_capture_num_tokens - - candidates = [1, 2, 4, 8, 16, 32, 64, 100, 120] - # max_seq_len=128, batch=1, 5 extra decoding steps -> ceiling 122. - kept, unrecordable = _filter_piecewise_capture_num_tokens( - candidates, - max_num_tokens=128, - max_batch_size=1, - max_seq_len=128, - num_extra_decoding_steps=5, - ) - assert kept[-1] == 120 # nothing above the 122 ceiling to clamp - assert 120 in kept - assert unrecordable == [] - # Same setup with 9 extra decoding steps -> ceiling 118; 120 drops. - kept, unrecordable = _filter_piecewise_capture_num_tokens( - candidates, - max_num_tokens=128, - max_batch_size=1, - max_seq_len=128, - num_extra_decoding_steps=9, - ) - assert kept[-1] == 118 - assert 100 in kept - assert 120 not in kept - assert unrecordable == [120] - def test_piecewise_filter_does_not_double_append_ceiling(self): """Ceiling already present in candidates -> not duplicated.""" from tensorrt_llm._torch.pyexecutor.model_engine import \ @@ -1848,7 +1812,6 @@ def test_speculative_model_alias(self): spec_config = EagleDecodingConfig( max_draft_len=3, speculative_model_dir="/path/to/model", - eagle3_one_model=False, ) args = TorchLlmArgs(model=llama_model_path, From f09cfd278544657b26df33eae1c76d2fdf483bcf Mon Sep 17 00:00:00 2001 From: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Date: Thu, 23 Jul 2026 15:13:14 -0400 Subject: [PATCH 2/3] [None][fix] Remove stale MTP mode argument Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Signed-off-by: Mike Iovine --- .../_torch/attention_backend/interface.py | 1 - .../_torch/attention_backend/sparse/dsa.py | 2 - .../_torch/attention_backend/trtllm.py | 65 +---- .../_torch/models/modeling_deepseekv3.py | 12 +- tensorrt_llm/_torch/models/modeling_glm.py | 7 +- .../_torch/models/modeling_nemotron_h.py | 1 - .../_torch/models/modeling_qwen3vl.py | 15 +- .../_torch/pyexecutor/cuda_graph_runner.py | 5 +- .../_torch/pyexecutor/model_engine.py | 6 - .../_torch/pyexecutor/model_loader.py | 9 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 228 +++--------------- .../_torch/pyexecutor/py_executor_creator.py | 5 - .../_torch/pyexecutor/sampler/sampler.py | 35 ++- tensorrt_llm/_torch/speculative/interface.py | 7 +- tensorrt_llm/_torch/speculative/utils.py | 4 +- tensorrt_llm/llmapi/llm_args.py | 2 - .../serve/test_spec_decoding_metrics.py | 11 +- .../_torch/executor/test_model_loader_mx.py | 6 +- .../_torch/modeling/test_modeling_llama.py | 3 - .../_torch/speculative/test_eagle3.py | 1 - .../speculative/test_force_accepted_tokens.py | 2 +- 21 files changed, 79 insertions(+), 348 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index fec6ba4e6b92..72a949936d50 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -440,7 +440,6 @@ def update_spec_dec_param( is_spec_dec_dynamic_tree, max_draft_len, max_total_draft_tokens, - model_is_wrapped: bool = False, spec_metadata: Optional['SpecMetadata'] = None, spec_tree_manager: Optional['SpecTreeManager'] = None, num_contexts: int = 0): diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 8466bbf59900..6c15e5f7ee51 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1234,7 +1234,6 @@ def update_spec_dec_param( is_spec_dec_dynamic_tree, max_draft_len, max_total_draft_tokens, - model_is_wrapped: bool = False, spec_metadata: Optional['SpecMetadata'] = None, spec_tree_manager: Optional['SpecTreeManager'] = None, num_contexts: int = 0, @@ -1246,7 +1245,6 @@ def update_spec_dec_param( is_spec_dec_dynamic_tree, max_draft_len, max_total_draft_tokens, - model_is_wrapped, spec_metadata, spec_tree_manager, num_contexts=num_contexts) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index a3f4e89d8d85..f2f0260180ef 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -972,7 +972,6 @@ def update_spec_dec_param( is_spec_dec_dynamic_tree, max_draft_len, max_total_draft_tokens, - model_is_wrapped: bool = False, spec_metadata: Optional['SpecMetadata'] = None, spec_tree_manager: Optional['SpecTreeManager'] = None, num_contexts: int = 0, @@ -986,7 +985,6 @@ def update_spec_dec_param( is_spec_dec_dynamic_tree: bool, whether using dynamic tree. max_draft_len: int, the number of the draft layers. max_total_draft_tokens: int, the number of all nodes in the tree (except the root). - model_is_wrapped: Optional[bool] = False, whether the drafter model is wrapped (i.e, CDL). spec_metadata: Optional['SpecMetadata'] = None, the metadata of the spec-dec. spec_tree_manager: Optional['SpecTreeManager'] = None, the spec_tree_manager for draft token tree. num_contexts: int = 0, the number of context (prefill) requests in the @@ -1109,60 +1107,23 @@ def update_spec_dec_param( self.spec_decoding_generation_lengths[:batch_size].fill_(n_dt) cpp_query_len = n_dt - # Case 2/3: static tree + # Case 2: static tree (target model verification) elif self.is_spec_dec_tree and not self.is_spec_dec_dynamic_tree and spec_metadata is not None: assert spec_metadata.spec_dec_mode.is_eagle3( ), "Tree decoding is only supported for Eagle3 now" - is_target_model = not getattr(spec_metadata, 'is_draft_model', - False) - - # Case 2: static tree and target model - if is_target_model: - # For the target model, we update the spec-dec parameters with the spec_tree_manager, which is prepared in advance. - self.spec_decoding_position_offsets[:batch_size, :].copy_( - spec_tree_manager.spec_dec_position_offsets[0, :], - non_blocking=True) - self.spec_decoding_packed_mask[:batch_size, :, :].copy_( - spec_tree_manager.spec_dec_packed_mask[0, :, :], - non_blocking=True) - self.spec_decoding_generation_lengths[:batch_size].fill_( - spec_tree_manager.max_total_draft_tokens + 1) - - # Case 3: static tree and the first drafter layer - else: - assert model_is_wrapped == True, "The drafter model should be wrapped" - # The first drafter layer will take the padded tokens as input (padding to the max_draft_len + 1) - # But the spec-dec parameters are still in the shape of max_total_draft_tokens + 1. - # Considering that these spec-dec params are accessed consecutively (without padding) in the attention Op, - # we need to write them consecutively when setting them. - # For the next drafter layers, we will prepare these spec-dec params in the drafting loops. - - # position_offsets - position_offset = torch.arange( - max_draft_len + 1, - dtype=torch.int, - device='cpu', - pin_memory=prefer_pinned()).repeat(batch_size) - self.spec_decoding_position_offsets.reshape( - -1)[:(max_draft_len + 1) * batch_size].copy_( - position_offset, non_blocking=True) - # packed_mask - dummy_idx = torch.arange(max_draft_len + 1) - spec_decoding_packed_mask = torch.pow( - 2, dummy_idx + 1) - 1 # [max_draft_len + 1] - spec_decoding_packed_mask = spec_decoding_packed_mask.repeat( - batch_size) # [batch_size * (max_draft_len + 1)] - self.spec_decoding_packed_mask.reshape( - -1)[:(max_draft_len + 1) * batch_size].copy_( - spec_decoding_packed_mask, non_blocking=True) - self.generate_spec_decoding_generation_length( - runtime_draft_len=max_draft_len) - if (self.spec_decoding_position_offsets is not None - and self.spec_decoding_position_offsets.dim() == 1): - cpp_query_len = max_draft_len + 1 - - # Case 4: linear tree + # Update the spec-dec parameters with the SpecTreeManager prepared + # in advance for target-model verification of the static tree. + self.spec_decoding_position_offsets[:batch_size, :].copy_( + spec_tree_manager.spec_dec_position_offsets[0, :], + non_blocking=True) + self.spec_decoding_packed_mask[:batch_size, :, :].copy_( + spec_tree_manager.spec_dec_packed_mask[0, :, :], + non_blocking=True) + self.spec_decoding_generation_lengths[:batch_size].fill_( + spec_tree_manager.max_total_draft_tokens + 1) + + # Case 3: linear tree else: # Currently dynamic draft length is only supported for linear tree # Dynamic draft length needs position offsets and packed mask to be shaped for each runtime draft length. diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index c95aac6c8147..519755a5e27c 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -1237,7 +1237,6 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], layer_idx: int, aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - is_separate_draft_engine: bool = False, mapping_with_cp: Optional[Mapping] = None): super().__init__() self.model_config = model_config @@ -1259,9 +1258,6 @@ def __init__(self, self.is_p2p_supported = can_access_peer(mapping) layer_idx_for_attention = layer_idx - if is_separate_draft_engine: - #KVCacheManager only support 1 layer for separate draft engine - layer_idx_for_attention = layer_idx - model_config.pretrained_config.num_hidden_layers # When enable_attention_dp is True, TP reduction is skipped since each DP rank # works on different batch elements. However, with CP > 1, attention is split @@ -1670,10 +1666,8 @@ class DeepseekV3MTP(DeepseekV3DecoderLayer): def __init__(self, model_config: ModelConfig[PretrainedConfig], layer_idx: int, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - is_separate_draft_engine: bool = False): - super().__init__(model_config, layer_idx, aux_stream_dict, - is_separate_draft_engine) + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream]): + super().__init__(model_config, layer_idx, aux_stream_dict) config = model_config.pretrained_config self.hidden_dim = config.hidden_size self.moe_intermediate_size = config.moe_intermediate_size @@ -1808,7 +1802,7 @@ def norm_hidden(): else: hidden_states, _ = self.shared_head.norm(hidden_states, residual) - # It's for 2-model path, capture the hidden states + # It's for Eagle3 / speculative capture of hidden states. if spec_metadata is not None: spec_metadata.maybe_capture_hidden_states(0, hidden_states, None) diff --git a/tensorrt_llm/_torch/models/modeling_glm.py b/tensorrt_llm/_torch/models/modeling_glm.py index 2572ea548e48..e21eb59f6fef 100644 --- a/tensorrt_llm/_torch/models/modeling_glm.py +++ b/tensorrt_llm/_torch/models/modeling_glm.py @@ -482,7 +482,6 @@ def __init__( model_config: ModelConfig[PretrainedConfig], layer_idx: int, aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - is_separate_draft_engine: bool = False, ): super().__init__() self.model_config = model_config @@ -498,9 +497,6 @@ def __init__( self.mapping = model_config.mapping mapping = self.mapping layer_idx_for_attention = layer_idx - if is_separate_draft_engine: - # KVCacheManager only support 1 layer for separate draft engine - layer_idx_for_attention = layer_idx - model_config.pretrained_config.num_hidden_layers if getattr(config, "use_qk_norm", False) and config.use_qk_norm: self.self_attn = Glm4Attention(model_config, layer_idx=layer_idx_for_attention) @@ -822,9 +818,8 @@ def __init__( model_config: ModelConfig[PretrainedConfig], layer_idx: int, aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - is_separate_draft_engine: bool = False, ): - super().__init__(model_config, layer_idx, aux_stream_dict, is_separate_draft_engine) + super().__init__(model_config, layer_idx, aux_stream_dict) config = model_config.pretrained_config self.hidden_dim = config.hidden_size self.moe_intermediate_size = config.moe_intermediate_size diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 3bb491fd1b91..63c3db7607bc 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -1184,7 +1184,6 @@ def __init__( model_config: NemotronHModelConfig, layer_idx: int, aux_stream_dict: dict[AuxStreamType, torch.cuda.Stream], - is_separate_draft_engine: bool = False, prefix: str = "", ): super().__init__() diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 2ae083d77355..35d29f504c5b 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -1390,13 +1390,14 @@ def vocab_size_padded(self) -> int: def infer_max_seq_len(self) -> int: return self.llm.infer_max_seq_len() - # Draft-model (two-model speculative decoding, e.g. DFlash / Eagle3) - # delegation: `ModelLoader.load` reads `draft_config` / `draft_model` and - # calls `load_draft_weights` on the *outer* model it resolved, but the - # spec-decoding wrapper (`SpecDecOneEngineForCausalLM`) is applied to the - # inner `self.llm` when this VLM composes it. Composite checkpoints - # (e.g. Qwen3.5-4B publishes text_config + vision_config) route text-only - # spec tests through this wrapper, so surface the inner LM's draft state. + # Draft-model delegation for same-engine speculative decoding (e.g. Eagle3 / + # DFlash draft weights loaded into the target engine). `ModelLoader.load` + # reads `draft_config` / `draft_model` and calls `load_draft_weights` on the + # *outer* model it resolved, but the spec-decoding wrapper + # (`SpecDecOneEngineForCausalLM`) is applied to the inner `self.llm` when + # this VLM composes it. Composite checkpoints (e.g. Qwen3.5-4B publishes + # text_config + vision_config) route text-only spec tests through this + # wrapper, so surface the inner LM's draft state. # Note: `load_draft_weights` must keep an explicit signature — the loader # dispatches kwargs via `inspect.getfullargspec`. @property diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 5e095033eba9..335a23bd4014 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -560,9 +560,8 @@ def _get_padded_batch(self, batch: ScheduledRequests, dummy_encoder_output_len = self._get_padding_dummy_encoder_output_len( cross_kv_cache_manager) - # Get draft KV cache manager only for one-model speculative decoding. - # In two-model mode, each model has its own KV cache manager, so - # draft_kv_cache_manager should be None. + # Get draft KV cache manager for one-model speculative decoding with + # a separate draft KV cache pool. draft_kv_cache_manager = get_draft_kv_cache_manager( self.spec_config, resource_manager) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b22a33967559..8906e36096e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1546,12 +1546,6 @@ def _run_mamba_hybrid_warmup(self, resource_manager: ResourceManager): batch, num_tokens_i) if batch is None: continue - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - if self.is_draft_model and isinstance( - spec_resource_manager, - Eagle3ResourceManager): - spec_resource_manager.is_first_draft = True self.forward(batch, new_tensors_device=None, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index f2284e06bb46..a1e4a9bbce6d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -556,17 +556,10 @@ def allocate_weights_on_cuda(t: torch.Tensor): if t.device != torch.device('meta') and ( need_initialized_weights or is_meta_init): if t.is_cuda: - memory_type_map = { - ExecutorMemoryType.MODEL_WEIGHTS_MAIN: - ExecutorMemoryType.MODEL_ENGINE_MAIN, - ExecutorMemoryType.MODEL_WEIGHTS_DRAFT: - ExecutorMemoryType.MODEL_ENGINE_DRAFT, - } - warnings.warn( f"A weight tensor of shape {t.shape} is already allocated on CUDA device before " f"the weight allocation stage. This will cause extra CUDA memory usage in the " - f"'{memory_type_map[self.model_weights_memory_tag]}' scope." + f"'{ExecutorMemoryType.MODEL_ENGINE_MAIN}' scope." ) cuda_t.copy_(t) memo[t] = cuda_t diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index feb02319a3e3..8f4af42e6045 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -56,7 +56,6 @@ from ..models.modeling_utils import DecoderModelForCausalLM from ..modules.decoder_layer import DecoderLayer from ..speculative.drafter import Drafter -from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.speculation_gate import SpeculationGate from .adp_iter_stats import ADPIterStatsBuffer from .connectors.kv_cache_connector import KvCacheConnectorManager @@ -83,7 +82,7 @@ derive_attention_dp_per_rank_request_cap, get_from_waiting_queue, merge_requests) from .resource_manager import (NoFreeSlotsError, ResourceManager, - ResourceManagerType, request_context) + ResourceManagerType) from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, SampleStateTensors, TRTLLMSampler) from .scheduler import (RequestScheduler, ScheduledRequests, @@ -572,8 +571,6 @@ def __init__( self.dist = dist self.sampler = sampler self.drafter = drafter - self.draft_model_engine = getattr(self.drafter, "draft_model_engine", - None) self.guided_decoder = guided_decoder self.disable_overlap_scheduler = disable_overlap_scheduler self.enable_kv_pool_rebalance = enable_kv_pool_rebalance @@ -730,7 +727,6 @@ def __init__( ) self.previous_batch: Optional[BatchState] = None - self.has_previous_draft_tokens = False self.num_scheduled_requests: int = 0 self.benchmark_req_queues_size = int( os.environ.get("TLLM_BENCHMARK_REQ_QUEUES_SIZE", 0)) @@ -806,8 +802,6 @@ def __init__( self.execution_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self.execution_stream): self.model_engine.warmup(self.resource_manager) - if self.draft_model_engine is not None: - self.draft_model_engine.warmup(self.resource_manager) # Ensure the default stream waits for execution_stream to complete # before subsequent operations. @@ -1212,8 +1206,6 @@ def is_warmup(self, value: bool): self._is_warmup = value # Set warmup flag in model engine to trigger torch compile and avoid moe load balancer statistics update self.model_engine.is_warmup = value - if self.draft_model_engine is not None: - self.draft_model_engine.is_warmup = value def start_worker(self): with self.worker_lock: @@ -1475,9 +1467,8 @@ def shutdown(self): # empty_cache), the subsequent CUDA graph teardown can trigger a # device-wide cudaErrorIllegalAddress when the driver touches metadata # for the now-freed memory regions. - for engine in (self.model_engine, self.draft_model_engine): - if engine is not None and hasattr(engine, '_release_cuda_graphs'): - engine._release_cuda_graphs() + if hasattr(self.model_engine, '_release_cuda_graphs'): + self.model_engine._release_cuda_graphs() # Ensure graph destruction has fully completed on device before # resource managers start freeing GPU-backed workspaces. if torch.cuda.is_available(): @@ -1496,8 +1487,6 @@ def shutdown(self): # (when the executor's reference is dropped), which is sufficient for # the GMS daemon registry eviction the cleanup hook was added for. del self.model_engine - if self.draft_model_engine is not None: - del self.draft_model_engine if self.virtual_memory_pools is not None: keys = list(self.virtual_memory_pools.keys()) for key in keys: @@ -1797,8 +1786,7 @@ def _get_init_iter_stats(self, num_new_active_requests, # staticBatchingStats is not used in pytorch path stats.static_batching_stats = StaticBatchingStats() - # Create specdec_stats if speculative decoding is enabled - # Either via spec_resource_manager (two-model mode) or spec_config (one-model mode) + # Create specdec_stats if speculative decoding is enabled. spec_resource_manager = self.resource_manager.resource_managers.get( ResourceManagerType.SPEC_RESOURCE_MANAGER) has_spec_config = self.model_engine.spec_config is not None @@ -2058,8 +2046,8 @@ def _update_iter_stats( else: stats.specdec_stats.acceptance_length = 0.0 - # Get draft latency from drafter if available (only for two-model mode) - # Only use draft latency if there were actually draft tokens in this iteration + # Get draft latency from drafter if available (e.g. NGram / user-provided). + # Only use draft latency if there were actually draft tokens in this iteration. draft_latency_ms = 0.0 if total_draft_tokens > 0 and self.drafter is not None and hasattr( self.drafter, 'last_draft_latency_ms'): @@ -3653,14 +3641,14 @@ def _prepare_and_schedule_batch(self): elif (getattr(self, "model_engine", None) is not None and self.model_engine.is_spec_decode and self.max_total_draft_tokens > 0): - # One-model speculative decoding (MTP / Eagle3 one-model) has no separate - # drafter, so the block above is skipped -- but the model still verifies - # max_total_draft_tokens draft tokens per generation step. The C++ micro-batch - # scheduler budgets each gen request as beam_width + getNumDraftTokens(); without - # populating the draft-token count here it under-reserves and, under chunked - # prefill + overlap scheduler, the forward's uniform (1 + runtime_draft_len) - # build overshoots max_num_tokens (total_num_tokens > max_num_tokens). Mirror the - # two-model normalization so scheduling reserves the correct token budget. + # Speculative decoding without a separate drafter (MTP / Eagle3) still + # verifies max_total_draft_tokens draft tokens per generation step. The C++ + # micro-batch scheduler budgets each gen request as + # beam_width + getNumDraftTokens(); without populating the draft-token count + # here it under-reserves and, under chunked prefill + overlap scheduler, the + # forward's uniform (1 + runtime_draft_len) build overshoots max_num_tokens + # (total_num_tokens > max_num_tokens). Normalize draft tokens so scheduling + # reserves the correct token budget. # model_engine is guarded first so partially-constructed executors in unit tests # (which may not set model_engine) do not raise AttributeError. for request in self.active_requests: @@ -4076,20 +4064,17 @@ def _executor_loop(self): if self.drafter is not None and self.use_spec_decode: if self.guided_decoder is not None: self.guided_decoder.rollback_rejected_tokens() - with request_context( - is_draft=self.draft_model_engine is not None, - scheduled_requests=scheduled_batch): - self.execution_stream.wait_stream( - torch.cuda.current_stream()) - with torch.cuda.stream(self.execution_stream): - self.drafter.prepare_draft_tokens( - scheduled_batch, self.resource_manager) - # Pad draft tokens to the max draft length and extend KV cache - # capacity to match. This is for CUDA graph compatibility. - self.drafter.pad_draft_tokens_for_cuda_graph( - scheduled_batch, self.resource_manager) - torch.cuda.current_stream().wait_stream( - self.execution_stream) + self.execution_stream.wait_stream( + torch.cuda.current_stream()) + with torch.cuda.stream(self.execution_stream): + self.drafter.prepare_draft_tokens( + scheduled_batch, self.resource_manager) + # Pad draft tokens to the max draft length and extend KV cache + # capacity to match. This is for CUDA graph compatibility. + self.drafter.pad_draft_tokens_for_cuda_graph( + scheduled_batch, self.resource_manager) + torch.cuda.current_stream().wait_stream( + self.execution_stream) # add_batch must be called again to restore to target requests with updated draft tokens. if self.guided_decoder is not None: self.guided_decoder.add_batch(scheduled_batch) @@ -4508,22 +4493,6 @@ def _executor_loop_overlap(self): self._prepare_disagg_gen_transmission_complete( scheduled_batch) - has_draft_batch = self.drafter is not None and self.previous_batch is not None and self.use_spec_decode and self.drafter.should_forward_draft_model( - scheduled_batch) - # Reset the draft tokens to avoid preparing resources for the draft model. - if self.drafter is not None and self.use_spec_decode and not has_draft_batch: - self.use_spec_decode = False - # We are not running the draft model. Remove the draft tokens and turn off spec - # decode so that the requests get handled correctly. - # One corner case: when we have at least one context request, we have to keep spec - # dec on. This ensures that we capture hidden states for requests that haven't done - # prefill yet. - self.use_spec_decode = False - self.model_engine.enable_spec_decode = scheduled_batch.num_context_requests > 0 - if not self.model_engine.enable_spec_decode: - for request in scheduled_batch.all_requests(): - request.py_draft_tokens = [] - self._handle_dynamic_draft_len(scheduled_batch) self.resource_manager.prepare_resources(scheduled_batch) @@ -4568,34 +4537,7 @@ def _executor_loop_overlap(self): self.guided_decoder.add_batch(scheduled_batch) self.guided_decoder.init_disagg_gen_requests() - previous_tensors = self.previous_batch and self.previous_batch.sample_state - # If there are previous draft tokens, we need to update the target requests to accept some draft tokens. - # When there's any accepted tokens, we can't directly use the previous batch's outputs in this iteration for the target model, - # so we'll set the target model's input to None and skip updating the target requests after target model forward. - use_previous_draft_tokens = self.has_previous_draft_tokens - num_accepted_tokens_device = None - - target_inputs = None - num_accepted_tokens_device = None - - if has_draft_batch: - self.execution_stream.wait_stream( - torch.cuda.current_stream()) - with torch.cuda.stream(self.execution_stream): - target_inputs, num_accepted_tokens_device = self._handle_speculative_decoding( - scheduled_batch, previous_tensors, - previous_tensors_device) - torch.cuda.current_stream().wait_stream( - self.execution_stream) - - # Use the draft_model's outputs if we've launched the draft model. - # Otherwise, use the previous batch's outputs. - if (target_inputs is not None - and target_inputs.next_draft_tokens - is not None) or use_previous_draft_tokens: - previous_tensors_device = target_inputs - else: - previous_tensors_device = self.previous_batch and self.previous_batch.sample_state and self.previous_batch.sample_state.device + previous_tensors_device = self.previous_batch and self.previous_batch.sample_state and self.previous_batch.sample_state.device scheduled_batch_stats = ( self._collect_scheduled_batch_stats(scheduled_batch) @@ -4612,8 +4554,7 @@ def _executor_loop_overlap(self): with self.perf_manager.record_perf_events( gpu_forward_start, gpu_forward_end) as fwd_timing: batch_outputs = self._forward_step( - scheduled_batch, previous_tensors_device, - num_accepted_tokens_device) + scheduled_batch, previous_tensors_device) self._maybe_prefetch_next_iter_mm_encoders(scheduled_batch) @@ -4647,10 +4588,6 @@ def _executor_loop_overlap(self): # should_process_previous_batch differs between ranks. self._flush_pending_transfer_responses() - if self.drafter is not None and self.use_spec_decode and should_process_previous_batch: - # Cleanup previous draft resources used in the draft model - self.drafter.cleanup_previous_draft_resources() - if not self._is_kv_manager_v2: self._pause_requests(scheduled_batch.paused_requests) @@ -4746,95 +4683,6 @@ def _executor_loop_overlap(self): self.iter_counter += 1 - @nvtx_range("_accept_draft_tokens") - def _accept_draft_tokens( - self, scheduled_batch: ScheduledRequests, - target_outputs: SampleStateTensors, - target_inputs: Optional[SampleStateTensors] - ) -> Tuple[SampleStateTensorsSpec, Optional[torch.Tensor]]: - """ - Prepare target device inputs after computing draft token acceptance. - - This function: - 1. If draft tokens exist: compares sampled tokens with draft tokens to compute acceptance - 2. If no draft tokens: directly uses the first sampled token - 3. Creates new_tokens by extracting accepted tokens per request - - Args: - scheduled_batch: The scheduled requests - target_outputs: Contains new_tokens [max_draft_len + 1, batch_size, beam_width] - or [1, batch_size, beam_width] if no draft tokens - target_inputs: Contains next_draft_tokens [batch_size, max_draft_len] - Returns: - Tuple of: - - SampleStateTensorsSpec with new_tokens set to accepted tokens, - new_tokens_lens and next_draft_tokens set to None - - num_accepted_tokens: [batch_size] tensor with acceptance counts per request, - or None if no draft tokens - """ - has_draft_tokens = target_inputs is not None and isinstance( - target_inputs, SampleStateTensorsSpec - ) and target_inputs.next_draft_tokens is not None - target_tokens = target_outputs.new_tokens # [max_draft_len + 1, batch_size, beam_width] or [1, batch_size, beam_width] - new_tokens = torch.zeros_like(target_tokens) - - # Squeeze the beam dimension (beam_width=1 for greedy or single beam) - target_tokens = target_tokens.squeeze( - -1) # [max_draft_len + 1, batch_size] or [1, batch_size] - - batch_size = target_tokens.shape[1] - device = target_tokens.device - # Compute number of accepted tokens per request - num_accepted_tokens = torch.zeros(batch_size, - dtype=torch.int32, - device=device) - - if has_draft_tokens: - # Draft tokens exist, compute acceptance - draft_tokens = target_inputs.next_draft_tokens # [batch_size, max_draft_len] - max_draft_len = draft_tokens.shape[1] - - # Compute number of accepted tokens per request - # Generation requests: compare with draft tokens to find acceptance - num_contexts = scheduled_batch.num_context_requests - if batch_size > num_contexts: - # Use .T to transpose: [max_draft_len + 1, num_gens] -> [num_gens, max_draft_len + 1] - gen_target_tokens = target_tokens[:, - num_contexts:].T # [num_gens, max_draft_len + 1] - - # Compare draft tokens with target tokens to find acceptance - # Use cumprod to find the first rejection point - draft_tokens_gen = draft_tokens[ - num_contexts:, :].int() # [num_gens, max_draft_len] - num_accepted_tokens[num_contexts:] += torch.cumprod( - (draft_tokens_gen == gen_target_tokens[:, :max_draft_len] - ).int(), - dim=-1).sum(dim=1) - - # Vectorized extraction using advanced indexing (no GPU-CPU sync) - # Use num_accepted_tokens as indices to gather the right tokens - batch_indices = torch.arange(batch_size, device=device) - new_tokens[0, :, 0] = target_tokens[num_accepted_tokens, - batch_indices] - else: - # No draft tokens to accept, just use the first (and only) sampled token - batch_indices = torch.arange(batch_size, device=device) - new_tokens[0, :, 0] = target_tokens[0, batch_indices] - - # Create the updated SampleStateTensorsSpec - # new_tokens_lens and next_draft_tokens are left as None - result_tensors = SampleStateTensorsSpec( - new_tokens=new_tokens, - log_probs=target_outputs.log_probs, - new_tokens_lens=None, - next_draft_tokens=None) - - # Copy logits if available - if hasattr(target_outputs, 'logits'): - result_tensors.logits = target_outputs.logits - - return result_tensors, num_accepted_tokens - def _process_previous_batch(self): self._handle_canceled_requests() # Skip iter-1 emission when `_emit_first_token_responses` already @@ -7001,28 +6849,6 @@ def _remove_inflight_ids(self, scheduled_requests: ScheduledRequests): ) self.inflight_req_ids.erase(req.request_id) - def _handle_speculative_decoding( - self, scheduled_batch, previous_tensors, target_inputs - ) -> Tuple[Optional[SampleStateTensorsSpec], Optional[torch.Tensor]]: - with request_context(is_draft=self.draft_model_engine is not None, - scheduled_requests=scheduled_batch): - target_outputs = self.previous_batch.sample_state and self.previous_batch.sample_state.device - assert target_outputs is not None, "target_outputs should not be None" - new_target_inputs, num_accepted_tokens_device = self._accept_draft_tokens( - scheduled_batch=scheduled_batch, - target_inputs=target_inputs, - target_outputs=target_outputs) - - self.drafter.generate_draft_tokens_with_overlap( - scheduled_batch, self.resource_manager, - previous_tensors.device if previous_tensors else None, - new_target_inputs, num_accepted_tokens_device) - - # Pad draft tokens to the max draft length for CUDA graph compatibility - self.has_previous_draft_tokens = new_target_inputs is not None and new_target_inputs.next_draft_tokens is not None - - return new_target_inputs, num_accepted_tokens_device - def reset_prefix_cache(self): self.kv_cache_manager.reset_reuse_state() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index e42694bf3db4..12c71421434a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -94,8 +94,6 @@ def _bytes_to_gib(bytes: int) -> float: "KV cache", ExecutorMemoryType.MODEL_ENGINE_MAIN: "Model", - ExecutorMemoryType.MODEL_ENGINE_DRAFT: - "Draft model for speculative decoding", } # Suggestion to reduce component memory usage @@ -111,9 +109,6 @@ def _bytes_to_gib(bytes: int) -> float: ExecutorMemoryType.INIT_KV_CACHE: "reduce max_num_tokens", ExecutorMemoryType.MODEL_ENGINE_MAIN: - ("reduce max_num_tokens and/or shard the model weights across GPUs by enabling " - "pipeline and/or tensor parallelism"), - ExecutorMemoryType.MODEL_ENGINE_DRAFT: ("reduce max_num_tokens and/or shard the model weights across GPUs by enabling " "pipeline and/or tensor parallelism"), ExecutorMemoryType.INIT_EXTRA_RESOURCES: diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 392a1644480e..b36a71ab303c 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -2580,26 +2580,19 @@ def _validate_top_p_decay_request(self, request: LlmRequest) -> None: raise ValueError("top_p_decay is not supported with beam search.") # A non-zero draft length means the request carries speculative draft # tokens and produces multiple tokens per step (req_num_steps = - # 1 + draft_token_length). One-model speculation (vanilla MTP, one-model - # Eagle3 / MTP-Eagle, SA, draft-target-one-model) uses its own - # SpecSamplerBase-derived sampler and never reaches TorchSampler; the - # drafter-based modes that DO flow draft tokens through TorchSampler - # (two-model draft-target, NGram, user-provided, two-model Eagle3 / - # MTP-Eagle) are what can make this length non-zero. top-p decay does not - # support these multi-token steps. - # NB: at admission time the draft tokens of drafter-based modes are - # usually not attached yet, so this check is best-effort. Two-model - # speculation (the only source of such requests in TorchSampler) is - # slated for removal, so in practice no speculative request reaches the - # decay path; a debug assert in _build_top_p_decay_metadata guards the - # invariant at sample time. + # 1 + draft_token_length). Speculative modes that use a dedicated + # SpecSamplerBase-derived sampler never reach TorchSampler. Drafter-based + # modes that do (NGram, user-provided) can make this length non-zero. + # top-p decay does not support these multi-token steps. + # NB: at admission time draft tokens of drafter-based modes are usually + # not attached yet, so this check is best-effort. A debug assert in + # _build_top_p_decay_metadata guards the invariant at sample time. if get_draft_token_length(request) > 0: raise ValueError( "top_p_decay is not supported for requests carrying speculative " - "draft tokens (req_num_steps > 1). This covers the drafter-based " - "modes routed through TorchSampler (two-model draft-target, NGram, " - "user-provided, two-model Eagle3 / MTP-Eagle); one-model " - "speculation uses its own sampler and is unaffected." + "draft tokens (req_num_steps > 1). This covers drafter-based " + "modes routed through TorchSampler (NGram, user-provided); " + "other speculative modes use their own sampler and are unaffected." ) def _can_use_fast_greedy_path(self, requests: list[LlmRequest]) -> bool: @@ -4366,10 +4359,10 @@ def _build_top_p_decay_metadata( group_seq_slots = seq_slots[group_req_indices] if __debug__: # Internal invariant (stripped under python -O): a decay-active - # row is always single-token -- the only source of multi-step - # rows in TorchSampler is two-model speculation (slated for - # removal), and decay + draft tokens is rejected per-request at - # admission in validate_request. + # row is always single-token -- multi-step rows in TorchSampler + # come from drafter-based modes (NGram / user-provided), and + # decay + draft tokens is rejected per-request at admission in + # validate_request. decay_row_steps = group_steps[ torch.isin(group_seq_slots, torch.tensor(list(self._top_p_decay_slots))) ] diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index e218d20f705c..85ad7069cbb6 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -224,10 +224,9 @@ def get_force_num_accepted_tokens() -> int: variable as an integer. Used by speculative decoding paths that operate on Python lists/slices and - therefore require an integer count (e.g. the two-model path in - ``TorchSampler``). For the one-model path, see - :func:`get_force_num_accepted_tokens_float`, which supports fractional - synthetic acceptance rates. + therefore require an integer count (e.g. drafter-based paths in + ``TorchSampler``). For fractional synthetic acceptance rates, see + :func:`get_force_num_accepted_tokens_float`. Returns: int: The forced number of accepted tokens, or 0 if not set or invalid. diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 8bd831cfcb90..8a1a4e4c73ca 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -539,8 +539,8 @@ def get_draft_len_for_batch_size(draft_len_schedule: Dict[int, int], """ Get the appropriate draft length for the given batch size using binary search. - This is a standalone function that can be used by both the drafter (two-model path) - and the model engine / spec workers (one-model path). + This is a standalone function that can be used by the drafter (NGram / + user-provided) and the model engine / spec workers. New semantics: Keys represent specific batch sizes (transition points). Values represent draft_len to use for batch sizes UP TO that key. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e01a0fa50bed..627b3245c565 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2980,9 +2980,7 @@ class ExecutorMemoryType(StrEnum): EXTRA_RESOURCES = "executor_extra" KV_CACHE = "kv_cache" MODEL_ENGINE_MAIN = "model" - MODEL_ENGINE_DRAFT = "draft_model" MODEL_WEIGHTS_MAIN = "model_weights" - MODEL_WEIGHTS_DRAFT = "draft_model_weights" @dataclass diff --git a/tests/integration/defs/examples/serve/test_spec_decoding_metrics.py b/tests/integration/defs/examples/serve/test_spec_decoding_metrics.py index d86cdb94f7e7..95992821240a 100644 --- a/tests/integration/defs/examples/serve/test_spec_decoding_metrics.py +++ b/tests/integration/defs/examples/serve/test_spec_decoding_metrics.py @@ -2,7 +2,7 @@ from defs.trt_test_alternative import print_error, print_info -def check_spec_decoding_metrics(http_port="8000", expect_draft_latency=False, min_draft_tokens=1): +def check_spec_decoding_metrics(http_port="8000", min_draft_tokens=1): print_info("Checking specDecodingStats in /metrics endpoint...") try: @@ -54,15 +54,6 @@ def check_spec_decoding_metrics(http_port="8000", expect_draft_latency=False, mi assert spec_stats["draftOverhead"] >= 0.0 assert spec_stats["draftOverhead"] <= 1.0 # Can't exceed 100% - # For 2-model mode, verify draft latency is tracked - if expect_draft_latency: - assert spec_stats["iterLatencyMS"] > 0.0, ( - "Two-model mode should have non-zero draft latency" - ) - assert spec_stats["draftOverhead"] > 0.0, ( - "Two-model mode should have non-zero draft overhead" - ) - return True except Exception as e: diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index ccdf960ea786..9a828bd77107 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -395,16 +395,16 @@ def test_default_profile_qualifies_real_tiny_llama_lifecycle(monkeypatch): assert_post_transform_lifecycle_equivalent(case) -def test_separate_draft_model_is_not_qualified_by_target_only_profile(monkeypatch): +def test_separate_draft_weights_not_qualified_by_target_only_profile(monkeypatch): monkeypatch.setattr( ModelLoader, "_POST_TRANSFORM_PROFILE_REGISTRY", - _tiny_profile_registry(speculative_mode="eagle3_one_model"), + _tiny_profile_registry(speculative_mode="eagle3"), ) decision = ModelLoader._qualify_post_transform_profile( _TinyModel([]), - speculative_mode="eagle3_one_model", + speculative_mode="eagle3", loads_draft_weights=True, ) diff --git a/tests/unittest/_torch/modeling/test_modeling_llama.py b/tests/unittest/_torch/modeling/test_modeling_llama.py index 089ef80645f7..6ee84e35f199 100644 --- a/tests/unittest/_torch/modeling/test_modeling_llama.py +++ b/tests/unittest/_torch/modeling/test_modeling_llama.py @@ -559,7 +559,6 @@ def run_forward(input_ids, position_ids, attn_metadata): is_spec_dec_tree=is_spec_dec_tree, max_draft_len=max_total_draft_tokens, max_total_draft_tokens=max_total_draft_tokens, - model_is_wrapped=False, spec_tree_manager=spec_tree_mgr, ) @@ -633,7 +632,6 @@ def run_forward(input_ids, position_ids, attn_metadata): is_spec_dec_dynamic_tree=is_tree_phase1, max_draft_len=gen_input_ids_1.size(-1) - 1, max_total_draft_tokens=gen_input_ids_1.size(-1) - 1, - model_is_wrapped=False, spec_metadata=spec_metadata_phase1, spec_tree_manager=spec_tree_mgr_phase1) @@ -708,7 +706,6 @@ def run_forward(input_ids, position_ids, attn_metadata): is_spec_dec_dynamic_tree=is_tree_ref, max_draft_len=gen_input_ids_ref.size(-1) - 1, max_total_draft_tokens=gen_input_ids_ref.size(-1) - 1, - model_is_wrapped=False, spec_metadata=spec_metadata_ref, spec_tree_manager=spec_tree_mgr_ref) diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 50a14e82d4d1..3e063cdc348b 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -1038,7 +1038,6 @@ def test_nemotron_super_mtp_dynamic_tree_dl6_k10_dt31( max_seq_len=8192, ) spec_config = MTPDecodingConfig(max_draft_len=max_draft_len, - mtp_eagle_one_model=True, use_dynamic_tree=True, dynamic_tree_max_topK=10, max_total_draft_tokens=31) diff --git a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py index fbfaca610386..6190057858b8 100644 --- a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py +++ b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py @@ -103,7 +103,7 @@ def test_get_force_num_accepted_tokens_float(env_value, expected): ], ) def test_get_force_num_accepted_tokens_int_unchanged(env_value, expected): - """The int helper must keep its original behavior (used by 2-model).""" + """The int helper must keep its original integer-only behavior.""" with patch.dict(os.environ, {FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR: env_value}): assert get_force_num_accepted_tokens() == expected From eedbdefb72a7f2524006450c1ff3e26490f45f5c Mon Sep 17 00:00:00 2001 From: Mike Iovine Date: Thu, 23 Jul 2026 16:08:11 -0400 Subject: [PATCH 3/3] Fix lints Signed-off-by: Mike Iovine --- .../_torch/models/modeling_deepseekv3.py | 7 +++--- .../_torch/pyexecutor/py_executor_creator.py | 24 +++++++------------ tensorrt_llm/llmapi/llm_args.py | 11 +++++---- 3 files changed, 17 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 519755a5e27c..263ba6e3bda6 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -1663,10 +1663,9 @@ def forward_mlp( class DeepseekV3MTP(DeepseekV3DecoderLayer): - def __init__(self, - model_config: ModelConfig[PretrainedConfig], - layer_idx: int, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream]): + def __init__(self, model_config: ModelConfig[PretrainedConfig], + layer_idx: int, aux_stream_dict: Dict[AuxStreamType, + torch.cuda.Stream]): super().__init__(model_config, layer_idx, aux_stream_dict) config = model_config.pretrained_config self.hidden_dim = config.hidden_size diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 12c71421434a..16e9ae194b0b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -74,26 +74,18 @@ def _bytes_to_gib(bytes: int) -> float: return bytes / (1024)**3 memory_type_friendly_names = { - ExecutorMemoryType.SAMPLER: - "Sampler", - ExecutorMemoryType.DRAFTER: - "Drafter", - ExecutorMemoryType.GUIDED_DECODER: - "Guided Decoder", - ExecutorMemoryType.SPEC_RESOURCES: - "Speculative decoding resources", + ExecutorMemoryType.SAMPLER: "Sampler", + ExecutorMemoryType.DRAFTER: "Drafter", + ExecutorMemoryType.GUIDED_DECODER: "Guided Decoder", + ExecutorMemoryType.SPEC_RESOURCES: "Speculative decoding resources", ExecutorMemoryType.INIT_KV_CACHE: "Initial KV Cache (temporary for KV cache size estimation)", ExecutorMemoryType.INIT_EXTRA_RESOURCES: "Additional executor resources (temporary for KV cache size estimation)", - ExecutorMemoryType.MODEL_EXTRA: - "Model resources created during usage", - ExecutorMemoryType.EXTRA_RESOURCES: - "Additional executor resources", - ExecutorMemoryType.KV_CACHE: - "KV cache", - ExecutorMemoryType.MODEL_ENGINE_MAIN: - "Model", + ExecutorMemoryType.MODEL_EXTRA: "Model resources created during usage", + ExecutorMemoryType.EXTRA_RESOURCES: "Additional executor resources", + ExecutorMemoryType.KV_CACHE: "KV cache", + ExecutorMemoryType.MODEL_ENGINE_MAIN: "Model", } # Suggestion to reduce component memory usage diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 627b3245c565..902ad5366d12 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5304,11 +5304,12 @@ def validate_speculative_config(self): # SA is gated for every method; relaxed / parallel / guided gates # apply only to the newly wired methods (vanilla MTP, PARD, # DFlash, DraftTarget). - is_new_rejection_method = ( - spec_mode == TorchSpeculativeDecodingMode.MTP - or spec_mode.is_pard() or spec_mode.is_dflash() - or spec_mode.is_dspark() - or spec_mode.is_draft_target()) + is_new_rejection_method = (spec_mode + == TorchSpeculativeDecodingMode.MTP + or spec_mode.is_pard() + or spec_mode.is_dflash() + or spec_mode.is_dspark() + or spec_mode.is_draft_target()) # Plain tensor parallelism is supported (the draft path # all-gathers vocab-sharded draft logits before rejection, see # SpecWorkerBase.maybe_gather_sharded_draft_logits).