From 19a738c95aad87bceff682bb721a79a7f2cd5edd Mon Sep 17 00:00:00 2001 From: wili-65535 Date: Thu, 15 May 2025 16:26:45 +0800 Subject: [PATCH 1/9] v1.5 Signed-off-by: wili-65535 v1.5.4 Add back draft_overhead to spec dec stats Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> --- .../tensorrt_llm/executor/serialization.h | 5 + cpp/include/tensorrt_llm/executor/types.h | 21 ++ .../executor/jsonSerialization.cpp | 4 +- cpp/tensorrt_llm/executor/serialization.cpp | 41 +++- cpp/tensorrt_llm/executor/serializeUtils.h | 4 + cpp/tensorrt_llm/pybind/executor/bindings.cpp | 10 + examples/prompt_lookup/README.md | 52 ++--- examples/prompt_lookup/run_dtm_pld.py | 56 +++-- examples/pytorch/README.md | 59 +++-- examples/pytorch/quickstart_advanced.py | 17 +- examples/run.py | 3 +- examples/summarize.py | 62 ++--- tensorrt_llm/_torch/pyexecutor/_util.py | 3 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 79 ++++++- .../_torch/pyexecutor/py_executor_creator.py | 6 +- tensorrt_llm/_torch/speculative/__init__.py | 3 +- tensorrt_llm/_torch/speculative/interface.py | 12 +- tensorrt_llm/_torch/speculative/ngram.py | 213 ++++++++++++++++++ tensorrt_llm/_torch/speculative/utils.py | 3 + tensorrt_llm/llmapi/__init__.py | 3 +- tensorrt_llm/llmapi/llm_args.py | 37 ++- tensorrt_llm/llmapi/llm_utils.py | 5 +- tensorrt_llm/models/modeling_utils.py | 3 + .../unittest/_torch/speculative/test_ngram.py | 83 +++++++ .../references_committed/llm.yaml | 2 +- 25 files changed, 661 insertions(+), 125 deletions(-) create mode 100644 tensorrt_llm/_torch/speculative/ngram.py create mode 100644 tests/unittest/_torch/speculative/test_ngram.py diff --git a/cpp/include/tensorrt_llm/executor/serialization.h b/cpp/include/tensorrt_llm/executor/serialization.h index e488b283d51f..27b854b7f900 100644 --- a/cpp/include/tensorrt_llm/executor/serialization.h +++ b/cpp/include/tensorrt_llm/executor/serialization.h @@ -252,6 +252,11 @@ class Serialization static void serialize(InflightBatchingStats const& inflightBatchingStats, std::ostream& os); static size_t serializedSize(InflightBatchingStats const& inflightBatchingStats); + // SpecDecStats + static SpecDecStats deserializeSpecDecStats(std::istream& is); + static void serialize(SpecDecStats const& specDecStats, std::ostream& os); + static size_t serializedSize(SpecDecStats const& specDecStats); + // IterationStats static IterationStats deserializeIterationStats(std::vector& buffer); static IterationStats deserializeIterationStats(std::istream& is); diff --git a/cpp/include/tensorrt_llm/executor/types.h b/cpp/include/tensorrt_llm/executor/types.h index 6d05e096339c..80ec7798db30 100644 --- a/cpp/include/tensorrt_llm/executor/types.h +++ b/cpp/include/tensorrt_llm/executor/types.h @@ -46,6 +46,7 @@ class Tensor; using TensorPtr = std::shared_ptr; using SizeType32 = std::int32_t; +using SizeType64 = std::int64_t; using FloatType = float; using TokenIdType = std::int32_t; using VecTokens = std::vector; @@ -294,6 +295,24 @@ struct InflightBatchingStats float avgNumDecodedTokensPerIter; }; +/// @brief Struct that holds speculative decoding stats +struct SpecDecStats +{ + /// @brief Total number of proposed draft tokens for all requests + SizeType64 numDraftTokens; + /// @brief Total number of accepted draft tokens for all requests + SizeType64 numAcceptedTokens; + /// @brief Number of requests with at least one draft token in batch + SizeType64 numRequestsWithDraftTokens; + /// @brief Acceptance length, defined as average number of tokens produced per step for all requests with at least + /// one draft token + double acceptanceLength; + /// @brief Iteration latency for draft token generation only (ms) + double iterLatencyMS; + /// @brief Draft overhead, defined as iterLatencyMS (specdec) / iterLatencyMS (total) + double draftOverhead; +}; + /// @brief Struct that holds the stats of a single iteration struct IterationStats { @@ -341,6 +360,8 @@ struct IterationStats std::optional staticBatchingStats; /// @brief Stats specific to inflight batching std::optional inflightBatchingStats; + /// @brief Stats specific to speculative decoding + std::optional specDecStats; }; /// @brief Enum class that represents the state of a request diff --git a/cpp/tensorrt_llm/executor/jsonSerialization.cpp b/cpp/tensorrt_llm/executor/jsonSerialization.cpp index 3e2fe25c24c3..e30937f3040d 100644 --- a/cpp/tensorrt_llm/executor/jsonSerialization.cpp +++ b/cpp/tensorrt_llm/executor/jsonSerialization.cpp @@ -31,11 +31,13 @@ NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE( StaticBatchingStats, numScheduledRequests, numContextRequests, numCtxTokens, numGenTokens, emptyGenSlots); NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(InflightBatchingStats, numScheduledRequests, numContextRequests, numGenRequests, numPausedRequests, numCtxTokens, microBatchId, avgNumDecodedTokensPerIter); +NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(SpecDecStats, numDraftTokens, numAcceptedTokens, numRequestsWithDraftTokens, + acceptanceLength, iterLatencyMS, draftOverhead); NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(IterationStats, timestamp, iter, iterLatencyMS, newActiveRequestsQueueLatencyMS, numNewActiveRequests, numActiveRequests, numQueuedRequests, numCompletedRequests, maxNumActiveRequests, maxBatchSizeStatic, maxBatchSizeTunerRecommended, maxBatchSizeRuntime, maxNumTokensStatic, maxNumTokensTunerRecommended, maxNumTokensRuntime, gpuMemUsage, cpuMemUsage, pinnedMemUsage, kvCacheStats, - staticBatchingStats, inflightBatchingStats); + staticBatchingStats, inflightBatchingStats, specDecStats); NLOHMANN_JSON_SERIALIZE_ENUM(RequestStage, {{RequestStage::kQUEUED, "QUEUED"}, {RequestStage::kCONTEXT_IN_PROGRESS, "CONTEXT_IN_PROGRESS"}, {RequestStage::kGENERATION_IN_PROGRESS, "GENERATION_IN_PROGRESS"}, diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index 0af11ba387fa..266d3fe20b2c 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1727,6 +1727,42 @@ size_t Serialization::serializedSize(InflightBatchingStats const& inflightBatchi return totalSize; } +// SpecDecStats +SpecDecStats Serialization::deserializeSpecDecStats(std::istream& is) +{ + auto numDraftTokens = su::deserialize(is); + auto numAcceptedTokens = su::deserialize(is); + auto numRequestsWithDraftTokens = su::deserialize(is); + auto acceptanceLength = su::deserialize(is); + auto iterLatencyMS = su::deserialize(is); + auto draftOverhead = su::deserialize(is); + + return SpecDecStats{ + numDraftTokens, numAcceptedTokens, numRequestsWithDraftTokens, acceptanceLength, iterLatencyMS, draftOverhead}; +} + +void Serialization::serialize(SpecDecStats const& state, std::ostream& os) +{ + su::serialize(state.numDraftTokens, os); + su::serialize(state.numAcceptedTokens, os); + su::serialize(state.numRequestsWithDraftTokens, os); + su::serialize(state.acceptanceLength, os); + su::serialize(state.iterLatencyMS, os); + su::serialize(state.draftOverhead, os); +} + +size_t Serialization::serializedSize(SpecDecStats const& state) +{ + size_t totalSize = 0; + totalSize += su::serializedSize(state.numDraftTokens); + totalSize += su::serializedSize(state.numAcceptedTokens); + totalSize += su::serializedSize(state.numRequestsWithDraftTokens); + totalSize += su::serializedSize(state.acceptanceLength); + totalSize += su::serializedSize(state.iterLatencyMS); + totalSize += su::serializedSize(state.draftOverhead); + return totalSize; +} + // IterationStats IterationStats Serialization::deserializeIterationStats(std::istream& is) @@ -1754,12 +1790,13 @@ IterationStats Serialization::deserializeIterationStats(std::istream& is) auto crossKvCacheStats = su::deserialize>(is); auto staticBatchingStats = su::deserialize>(is); auto inflightBatchingStats = su::deserialize>(is); + auto specdecStats = su::deserialize>(is); return IterationStats{timestamp, iter, iterLatencyMS, newActiveRequestsQueueLatencyMS, numNewActiveRequests, numActiveRequests, numQueuedRequests, numCompletedRequests, maxNumActiveRequests, maxBatchSizeStatic, maxBatchSizeTunerRecommended, maxBatchSizeRuntime, maxNumTokensStatic, maxNumTokensTunerRecommended, maxNumTokensRuntime, gpuMemUsage, cpuMemUsage, pinnedMemUsage, kvCacheStats, crossKvCacheStats, - staticBatchingStats, inflightBatchingStats}; + staticBatchingStats, inflightBatchingStats, specdecStats}; } IterationStats Serialization::deserializeIterationStats(std::vector& buffer) @@ -1797,6 +1834,7 @@ size_t Serialization::serializedSize(IterationStats const& iterStats) totalSize += su::serializedSize(iterStats.crossKvCacheStats); totalSize += su::serializedSize(iterStats.staticBatchingStats); totalSize += su::serializedSize(iterStats.inflightBatchingStats); + totalSize += su::serializedSize(iterStats.specDecStats); return totalSize; } @@ -1825,6 +1863,7 @@ void Serialization::serialize(IterationStats const& iterStats, std::ostream& os) su::serialize(iterStats.crossKvCacheStats, os); su::serialize(iterStats.staticBatchingStats, os); su::serialize(iterStats.inflightBatchingStats, os); + su::serialize(iterStats.specDecStats, os); } std::vector Serialization::serialize(IterationStats const& iterStats) diff --git a/cpp/tensorrt_llm/executor/serializeUtils.h b/cpp/tensorrt_llm/executor/serializeUtils.h index fc5902b83e26..4d6834db26a9 100644 --- a/cpp/tensorrt_llm/executor/serializeUtils.h +++ b/cpp/tensorrt_llm/executor/serializeUtils.h @@ -465,6 +465,10 @@ T deserialize(std::istream& is) { return Serialization::deserializeInflightBatchingStats(is); } + else if constexpr (std::is_same_v) + { + return Serialization::deserializeSpecDecStats(is); + } else if constexpr (std::is_same_v) { return Serialization::deserializeIterationStats(is); diff --git a/cpp/tensorrt_llm/pybind/executor/bindings.cpp b/cpp/tensorrt_llm/pybind/executor/bindings.cpp index d697f68215cc..1d0ae166d668 100644 --- a/cpp/tensorrt_llm/pybind/executor/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/executor/bindings.cpp @@ -132,6 +132,15 @@ void initBindings(pybind11::module_& m) .def_readwrite("micro_batch_id", &tle::InflightBatchingStats::microBatchId) .def_readwrite("avg_num_decoded_tokens_per_iter", &tle::InflightBatchingStats::avgNumDecodedTokensPerIter); + py::class_(m, "SpecDecStats") + .def(py::init<>()) + .def_readwrite("num_draft_tokens", &tle::SpecDecStats::numDraftTokens) + .def_readwrite("num_accepted_tokens", &tle::SpecDecStats::numAcceptedTokens) + .def_readwrite("num_requests_with_draft_tokens", &tle::SpecDecStats::numRequestsWithDraftTokens) + .def_readwrite("acceptance_length", &tle::SpecDecStats::acceptanceLength) + .def_readwrite("iter_latency_ms", &tle::SpecDecStats::iterLatencyMS) + .def_readwrite("draft_overhead", &tle::SpecDecStats::draftOverhead); + py::class_(m, "IterationStats") .def(py::init<>()) .def_readwrite("timestamp", &tle::IterationStats::timestamp) @@ -150,6 +159,7 @@ void initBindings(pybind11::module_& m) .def_readwrite("cross_kv_cache_stats", &tle::IterationStats::crossKvCacheStats) .def_readwrite("static_batching_stats", &tle::IterationStats::staticBatchingStats) .def_readwrite("inflight_batching_stats", &tle::IterationStats::inflightBatchingStats) + .def_readwrite("specdec_stats", &tle::IterationStats::specDecStats) .def("to_json_str", [](tle::IterationStats const& iterationStats) { return tle::JsonSerialization::toJsonStr(iterationStats); }); diff --git a/examples/prompt_lookup/README.md b/examples/prompt_lookup/README.md index 70ebab177888..337b1d21ea8a 100644 --- a/examples/prompt_lookup/README.md +++ b/examples/prompt_lookup/README.md @@ -4,10 +4,12 @@ This document shows how to build and run a model using Prompt-Lookup speculative ## Overview +We provide two styles of workflow to run Prompt-Lookup (named V1 and V2 respectively) now. V1 is in TRT workflow and similar to the Draft-Target-Model workflow, running in orchestrator mode and calling `runner.generate()` multiple times to get outputs, which is more flexible for customizing but slightly more overhead. V2 is in pytorch workflow and similar to the Look-Ahead workflow, running in leader mode and calling `runner.generate()` only one time to get outputs, which provides higher performance but fixed process. + The Prompt-Lookup has 3 additional hyperparameters that you need to specify to control the process of generation: -- `prompt_lookup_num_tokens`: the number of tokens we extract from input prompt or previous generated output as draft tokens in one iteration, which the range is from 4 to 10 in common usage. Empirically, the larger the value is, the higher acceptance ratio but higher overhead is expected at the same time, so the right balance based on the models and application scenarios needs to be found. -- `max_matching_ngram_size`: the number of tokens we get from the tail of the generated output as a pattern, which is used to match in input prompt or previous generated output. Empirically, the larger the value is, the more precise context can be matched from the existed sequence, indicating higher acceptance ratio, but the higher probability of miss-match and higher overhead appear, which fall back to normal generation (one token per iteration). -- `device_list`: the index list of device(s) to run the model. The length of it must be the same as the TP size of the draft model engine. For instances, `device_list=[0]` means using tp_size=1 and GPU 0 for the model, `device_list=[4,5,6,7]` means using tp=4 and GPU from 4 to 7 for the model. +- `prompt_lookup_num_tokens`: the number of tokens we extract from input prompt or previous generated output as draft tokens in one iteration, which the range is from 4 to 10 in common usage (default value: 4). Empirically, the larger the value is, the higher acceptance rate but higher overhead is expected at the same time, so the right balance based on the models and application scenarios needs to be found. +- `max_matching_ngram_size`: the number of tokens we get from the tail of the input prompt or generated output as a pattern, which is used to match in input prompt or previous generated output (default value: 2). Empirically, the larger the value is, the more precise context can be matched from the existed sequence, indicating higher acceptance rate, but the higher probability of miss-match and higher overhead appear, which fall back to normal generation (one token per iteration). +- `device_list`: the index list of device(s) to run the model in V1 workflow. The length of it must be the same as the TP size of the draft model engine. For instances, `device_list=[0]` means using tp_size=1 and GPU 0 for the model, `device_list=[4,5,6,7]` means using tp=4 and GPU from 4 to 7 for the model. This parameter is neddless in V2 workflow. ## Support Matrix * GPU Compute Capability >= 8.0 (Ampere or newer) @@ -17,17 +19,21 @@ The Prompt-Lookup has 3 additional hyperparameters that you need to specify to c ## Usage -### Build engines +### V1 workflow + We use an open-source `llama-v2-13B` models in this example. + `--use_paged_context_fmha=enable` must be specified since we need KVcache reuse in this approach. + `--speculative_decoding_mode=draft_tokens_external` must be specified. + `--max_draft_len` must be specified larger or equal to `prompt_lookup_num_tokens`. ++ `---prompt_lookup_config` is corresponding configuration of Prompt-Lookup, we can see its usage in [util.py](../util.py). + + As an example, `[10,2,[0]]` means `prompt_lookup_num_tokens=10`, `max_matching_ngram_size=2`, and device of target model is `GPU0`. ++ `--kv_cache_enable_block_reuse` must be specified for this approach. ++ Only CPP session is supported, so `--use_py_session` must not be specified. ++ `--num_beams` can not be specified as larger than 1 since beam search is not supported in this approach yet. ```bash -cd examples/models/core/llama - -python3 convert_checkpoint.py \ +# Build engine +python3 examples/models/core/llama/convert_checkpoint.py \ --model_dir= \ --output_dir=./ckpt-target \ --dtype=float16 @@ -42,34 +48,18 @@ trtllm-build \ --max_batch_size=4 \ --max_input_len=3200 \ --max_seq_len=4800 -``` - -### Run decoding -+ `---prompt_lookup_config` is corresponding configuration of Prompt-Lookup, we can see its usage in [util.py](../util.py). - + As an example, `[10,2,[0]]` means `prompt_lookup_num_tokens=10`, `max_matching_ngram_size=2`, and device of target model is `GPU0`. -+ `--kv_cache_enable_block_reuse` must be specified for this approach. -+ Only CPP session is supported, so `--use_py_session` must not be specified. -+ `--num_beams` can not be specified as larger than 1 since beam search is not supported in this approach yet. - -```bash -cd examples/models/core/llama - -python3 ../../../run.py \ +# Run decoding +python3 examples/run.py \ --tokenizer_dir \ --engine_dir ./target-engine \ --prompt_lookup_config="[10,2,[0]]" \ --max_output_len=256 \ --kv_cache_enable_block_reuse \ --input_text="How does Draft-Sampling work?" -``` - -## Run summarization tasks -```bash -cd examples/models/core/llama - -python ../../../summarize.py \ +# Run summarization tasks +python examples/summarize.py \ --test_hf \ --test_trt_llm \ --check_accuracy \ @@ -79,3 +69,11 @@ python ../../../summarize.py \ --prompt_lookup_config="[10,2,[0]]" \ --kv_cache_enable_block_reuse ``` + +### V2 workflow + +```bash +python3 examples/pytorch/quickstart_advanced.py \ + --max_matching_ngram_size=2 \ + --spec_decode_nextn=4 +``` diff --git a/examples/prompt_lookup/run_dtm_pld.py b/examples/prompt_lookup/run_dtm_pld.py index 256880e86bf7..559c1e7bbef9 100644 --- a/examples/prompt_lookup/run_dtm_pld.py +++ b/examples/prompt_lookup/run_dtm_pld.py @@ -36,8 +36,8 @@ def __init__( is_use_oldest: bool = True, ): self.input_batch_size = input_batch_size - self.plnt = prompt_lookup_num_tokens # Shorter name - self.mmns = max_matching_ngram_size # Shorter name + self.prompt_lookup_num_tokens = prompt_lookup_num_tokens + self.max_matching_ngram_size = max_matching_ngram_size self.end_id = end_id self.max_seq_len = max_seq_len self.is_keep_all = is_keep_all @@ -45,9 +45,25 @@ def __init__( self.pool = [{} for _ in range(input_batch_size)] self.start_index = [0 for _ in range(input_batch_size)] - # modified from `transformers/generation/candidate_generator.py` + assert self.prompt_lookup_num_tokens > 0, f"prompt_lookup_num_tokens must be greater than 0, but got {self.prompt_lookup_num_tokens}" + assert self.max_matching_ngram_size > 0, f"max_matching_ngram_size must be greater than 0, but got {self.max_matching_ngram_size}" + + def print_pool(self): + """ + For debug + """ + logger.info(f"Batch size = {self.input_batch_size}") + for i, map in enumerate(self.pool): + logger.info(f"Slot {i}, size = {len(map)}") + for key, values in map.items(): + logger.info(f" {key}->{values}") + def get_draft_tokens(self, prefix: list[torch.Tensor], batch_slot: list[int]): + """ + Get draft tokens from a batch of requests + modified from `transformers/generation/candidate_generator.py` + """ batch_size = len(prefix) prefix_len = [len(prefix[bi]) for bi in range(batch_size)] draft_tokens = [] # `logits` is useless yet @@ -61,25 +77,30 @@ def get_draft_tokens(self, prefix: list[torch.Tensor], # Update pool sequence = prefix[bi][self.start_index[gbi]:].tolist() - for size in range(min(self.mmns, prefix_len[bi] - 1), 0, -1): + for size in range( + min(self.max_matching_ngram_size, prefix_len[bi] - 1), 0, + -1): # Find each possible key-value combination, and use tuple for hash for l in range(len(sequence) - size): - r = min(l + size + self.plnt, len(sequence)) + r = min(l + size + self.prompt_lookup_num_tokens, + len(sequence)) key = tuple(sequence[l:l + size]) value = tuple(sequence[l + size:r]) if key not in self.pool[gbi] or not self.is_keep_all or \ - len(self.pool[gbi][key][0]) < self.plnt: + len(self.pool[gbi][key][0]) < self.prompt_lookup_num_tokens: # Update the value if # 1. the key does not exist - # 2. we only keep one value for each key + # 2. we only keep the newest one value for each key (MRU) # 3. the length of the value saved before is less than `prompt_lookup_num_tokens` self.pool[gbi][key] = OrderedSet((value, )) elif value not in self.pool[gbi][key]: - # Extend the value if the key is already existed and we want to keep all of them + # Extend the value if the key is already existed but count of values is not enough self.pool[gbi][key].add(value) # Find match - for size in range(min(self.mmns, prefix_len[bi] - 1), 0, -1): + for size in range( + min(self.max_matching_ngram_size, prefix_len[bi] - 1), 0, + -1): pattern = tuple(prefix[bi][-size:].tolist()) if pattern not in self.pool[gbi]: continue @@ -92,7 +113,8 @@ def get_draft_tokens(self, prefix: list[torch.Tensor], break draft_tokens.append(chosen_ids) self.start_index[gbi] = max( - 0, prefix_len[bi] - (self.plnt + self.mmns - 1)) + 0, prefix_len[bi] - (self.prompt_lookup_num_tokens + + self.max_matching_ngram_size - 1)) return draft_tokens, None @@ -108,20 +130,21 @@ def run_dtm_pld(batch_input_ids, *, target_runner=None): # `dtm` for Draft-Target-Model, `pld` for Prompt-Lookup-Decoding - assert (args.draft_target_model_config is not None) ^ (args.prompt_lookup_config is not None), \ - "`--draft_target_model_config` and `--prompt_lookup_config` can not be specified at the same time." - if args.draft_target_model_config is not None: - assert args.draft_engine_dir is not None, "`--draft_engine_dir` must be specified in Draft-Target-Model." is_dtm = (args.draft_target_model_config is not None) is_pld = (args.prompt_lookup_config is not None) + assert is_dtm ^ is_pld, "`--draft_target_model_config` and `--prompt_lookup_config` can not be specified at the same time." if is_dtm: + assert args.draft_engine_dir is not None, "`--draft_engine_dir` must be specified in Draft-Target-Model." draft_len, draft_device_list, target_device_list, use_logits = ast.literal_eval( args.draft_target_model_config) + logger.info(f"Using Draft-Target-Model speculative decoding") logger.info(f"draft_len: {draft_len}") logger.info(f"Device(s) for draft model: {draft_device_list}") logger.info(f"Device(s) for target model: {target_device_list}") logger.info(f"Use logits to accept tokens: {use_logits}") if is_pld: + logger.info( + f"Using Prompt-Lookup-Decoding speculative decoding V1 workflow") prompt_lookup_num_tokens, max_matching_ngram_size, target_device_list = ast.literal_eval( args.prompt_lookup_config) logger.info(f"prompt_lookup_num_tokens: {prompt_lookup_num_tokens}") @@ -206,9 +229,8 @@ def run_dtm_pld(batch_input_ids, device_ids=target_device_list) target_runner = ModelRunnerCpp.from_dir(**target_runner_kwargs) - if is_dtm and use_logits and \ - not (draft_runner.gather_generation_logits and target_runner.gather_generation_logits): - assert False, "`--gather_generation_logits` must be specified while building draft/target models for using logits to accept" + if is_dtm and use_logits: + assert draft_runner.gather_generation_logits and target_runner.gather_generation_logits, "`--gather_generation_logits` must be specified while building draft/target models for using logits to accept" common_generaion_kwargs = dict( max_attention_window_size=args.max_attention_window_size, diff --git a/examples/pytorch/README.md b/examples/pytorch/README.md index 5df1d18d59c5..19ed848a4123 100644 --- a/examples/pytorch/README.md +++ b/examples/pytorch/README.md @@ -1,12 +1,12 @@ # TRT-LLM with PyTorch -Run the quick start script: +## Run the quick start script: ```bash python3 quickstart.py ``` -Run the advanced usage example script: +## Run the advanced usage example script: ```bash # BF16 @@ -29,8 +29,9 @@ python3 quickstart_advanced.py --model_dir nvidia/Llama-3_1-Nemotron-Ultra-253B- # Nemotron-H requires disabling cache reuse in kv cache python3 quickstart_advanced.py --model_dir nvidia/Nemotron-H-8B-Base-8K --disable_kv_cache_reuse --max_batch_size 8 +``` -Run the multimodal example script: +## Run the multimodal example script: ```bash # default inputs @@ -43,29 +44,39 @@ python3 quickstart_multimodal.py --model_dir Efficient-Large-Model/NVILA-8B --mo # Note: media should be either image or video. Mixing image and video is not supported. python3 quickstart_multimodal.py --model_dir Efficient-Large-Model/NVILA-8B --modality video --prompt "Tell me what you see in the video briefly." "Describe the scene in the video briefly." --media "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/OAI-sora-tokyo-walk.mp4" "https://huggingface.co/datasets/Efficient-Large-Model/VILA-inference-demos/resolve/main/world.mp4" --max_tokens 128 [--use_cuda_graph] ``` -## Supported Models -| Architecture | Model | HuggingFace Example | Modality | -|--------------|-------|---------------------|----------| -| `BertForSequenceClassification` | BERT-based | `textattack/bert-base-uncased-yelp-polarity` | L | -| `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3` | L | -| `Gemma3ForCausalLM` | Gemma3 | `google/gemma-3-1b-it` | L | -| `LlavaLlamaModel` | VILA | `Efficient-Large-Model/NVILA-8B` | L + V | -| `LlavaNextForConditionalGeneration` | LLaVA-NeXT | `llava-hf/llava-v1.6-mistral-7b-hf` | L + V | -| `LlamaForCausalLM` | Llama 3
Llama 3.1
Llama 2
LLaMA | `meta-llama/Meta-Llama-3.1-70B` | L | -| `Llama4ForConditionalGeneration` | Llama 4 Scout
Llama 4 Maverick | `meta-llama/Llama-4-Scout-17B-16E-Instruct`
`meta-llama/Llama-4-Maverick-17B-128E-Instruct` | L + V | -| `MistralForCausalLM` | Mistral | `mistralai/Mistral-7B-v0.1` | L | -| `MixtralForCausalLM` | Mixtral | `mistralai/Mixtral-8x7B-v0.1` | L | -| `MllamaForConditionalGeneration` | Llama 3.2 | `meta-llama/Llama-3.2-11B-Vision` | L | -| `NemotronForCausalLM` | Nemotron-3
Nemotron-4
Minitron | `nvidia/Minitron-8B-Base` | L | -| `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`
`nvidia/Nemotron-H-47B-Base-8K`
`nvidia/Nemotron-H-56B-Base-8K` | L | -| `NemotronNASForCausalLM` | LLamaNemotron
LlamaNemotron Super
LlamaNemotron Ultra | `nvidia/Llama-3_1-Nemotron-51B-Instruct`
`nvidia/Llama-3_3-Nemotron-Super-49B-v1`
`nvidia/Llama-3_1-Nemotron-Ultra-253B-v1` | L | -| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/Qwen2-7B-Instruct` | L | -| `Qwen2ForProcessRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-PRM-7B` | L | -| `Qwen2ForRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-RM-72B` | L | -| `Qwen2VLForConditionalGeneration` | Qwen2-VL | `Qwen/Qwen2-VL-7B-Instruct` | L + V | -| `Qwen2_5_VLForConditionalGeneration` | Qwen2.5-VL | `Qwen/Qwen2.5-VL-7B-Instruct` | L + V | + +### Supported Models +| Architecture | Model | HuggingFace Example | Modality | +| :----------------------------------: | :----------------------------------------------------------- | :----------------------------------------------------------- | :------: | +| `BertForSequenceClassification` | BERT-based | `textattack/bert-base-uncased-yelp-polarity` | L | +| `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3 ` | L | +| `Gemma3ForCausalLM` | Gemma3 | `google/gemma-3-1b-it` | L | +| `LlavaLlamaModel` | VILA | `Efficient-Large-Model/NVILA-8B` | L + V | +| `LlavaNextForConditionalGeneration` | LLaVA-NeXT | `llava-hf/llava-v1.6-mistral-7b-hf` | L + V | +| `LlamaForCausalLM` | Llama 3
Llama 3.1
Llama 2
LLaMA | `meta-llama/Meta-Llama-3.1-70B` | L | +| `Llama4ForConditionalGeneration` | Llama 4 Scout
Llama 4 Maverick | `meta-llama/Llama-4-Scout-17B-16E-Instruct`
`meta-llama/Llama-4-Maverick-17B-128E-Instruct` | L + V | +| `MistralForCausalLM` | Mistral | `mistralai/Mistral-7B-v0.1` | L | +| `MixtralForCausalLM` | Mixtral | `mistralai/Mixtral-8x7B-v0.1` | L | +| `MllamaForConditionalGeneration` | Llama 3.2 | `meta-llama/Llama-3.2-11B-Vision` | L | +| `NemotronForCausalLM` | Nemotron-3
Nemotron-4
Minitron | `nvidia/Minitron-8B-Base` | L | +| `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`
`nvidia/Nemotron-H-47B-Base-8K`
`nvidia/Nemotron-H-56B-Base-8K` | L | +| `NemotronNASForCausalLM` | LLamaNemotron
LlamaNemotron Super
LlamaNemotron Ultra | `nvidia/Llama-3_1-Nemotron-51B-Instruct`
`nvidia/Llama-3_3-Nemotron-Super-49B-v1`
`nvidia/Llama-3_1-Nemotron-Ultra-253B-v1` | L | +| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/Qwen2-7B-Instruct` | L | +| `Qwen2ForProcessRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-PRM-7B` | L | +| `Qwen2ForRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-RM-72B` | L | +| `Qwen2VLForConditionalGeneration` | Qwen2-VL | `Qwen/Qwen2-VL-7B-Instruct` | L + V | +| `Qwen2_5_VLForConditionalGeneration` | Qwen2.5-VL | `Qwen/Qwen2.5-VL-7B-Instruct` | L + V | Note: - L: Language only - L + V: Language and Vision multimodal support - Llama 3.2 accepts vision input, but our support currently limited to text only. + +## Run the speculative decoding script: + +```bash +# NGram drafter +python3 examples/pytorch/quickstart_advanced.py \ + --max_matching_ngram_size=2 \ + --spec_decode_nextn=4 +``` diff --git a/examples/pytorch/quickstart_advanced.py b/examples/pytorch/quickstart_advanced.py index 8ffa087826b1..16cc8ff17a94 100644 --- a/examples/pytorch/quickstart_advanced.py +++ b/examples/pytorch/quickstart_advanced.py @@ -4,7 +4,7 @@ from tensorrt_llm._torch import LLM from tensorrt_llm._torch.pyexecutor.config import PyTorchConfig from tensorrt_llm.llmapi import (EagleDecodingConfig, KvCacheConfig, - MTPDecodingConfig) + MTPDecodingConfig, NGramDecodingConfig) example_prompts = [ "Hello, my name is", @@ -103,6 +103,7 @@ def add_llm_args(parser): parser.add_argument('--spec_decode_algo', type=str, default=None) parser.add_argument('--spec_decode_nextn', type=int, default=1) parser.add_argument('--eagle_model_dir', type=str, default=None) + parser.add_argument('--max_matching_ngram_size', type=int, default=5) # Relaxed acceptance parser.add_argument('--use_relaxed_acceptance_for_thinking', @@ -130,6 +131,7 @@ def setup_llm(args): use_cuda_graph=args.use_cuda_graph, load_format=args.load_format, print_iter_log=args.print_iter_log, + enable_iter_perf_stats=args.print_iter_log, torch_compile_enabled=args.use_torch_compile, torch_compile_piecewise_cuda_graph=args.use_piecewise_cuda_graph, moe_backend=args.moe_backend, @@ -154,6 +156,14 @@ def setup_llm(args): spec_config = EagleDecodingConfig( max_draft_len=args.spec_decode_nextn, pytorch_eagle_weights_path=args.eagle_model_dir) + elif spec_decode_algo == "NGRAM": + spec_config = NGramDecodingConfig( + prompt_lookup_num_tokens=args.spec_decode_nextn, + max_matching_ngram_size=args.max_matching_ngram_size, + is_keep_all=True, + is_use_oldest=True, + is_public_pool=True, + ) else: spec_config = None @@ -193,6 +203,11 @@ def main(): generated_text = output.outputs[0].text print(f"[{i}] Prompt: {prompt!r}, Generated text: {generated_text!r}") + if args.print_iter_log: + stats = llm.get_stats() + for stat in stats: + print(stat) + if __name__ == '__main__': main() diff --git a/examples/run.py b/examples/run.py index 0ae27c9b54b9..fed6c3851d5d 100755 --- a/examples/run.py +++ b/examples/run.py @@ -454,7 +454,8 @@ def main(args): lora_ckpt_source=args.lora_ckpt_source, gpu_weights_percent=args.gpu_weights_percent, max_output_len=args.max_output_len, - enable_context_fmha_fp32_acc=args.enable_context_fmha_fp32_acc) + enable_context_fmha_fp32_acc=args.enable_context_fmha_fp32_acc, + ) if args.medusa_choices is not None: args.medusa_choices = ast.literal_eval(args.medusa_choices) assert args.temperature == 1.0, "Medusa should use temperature == 1.0" diff --git a/examples/summarize.py b/examples/summarize.py index 474dac2d841e..8326fb9c171a 100644 --- a/examples/summarize.py +++ b/examples/summarize.py @@ -491,10 +491,34 @@ def eval_hf(datapoint, f"Using {'Python' if args.use_py_session else 'C++'} session") runner_cls = ModelRunner if args.use_py_session else ModelRunnerCpp - runner_kwargs = dict(engine_dir=args.engine_dir, - rank=runtime_rank, - debug_mode=args.debug_mode, - gpu_weights_percent=args.gpu_weights_percent) + runner_kwargs = dict( + engine_dir=args.engine_dir, + rank=runtime_rank, + debug_mode=args.debug_mode, + gpu_weights_percent=args.gpu_weights_percent, + enable_context_fmha_fp32_acc=args.enable_context_fmha_fp32_acc, + ) + if not args.use_py_session: + runner_kwargs.update( + lora_dir=args.lora_dir, + lora_ckpt_source=args.lora_ckpt_source, + max_batch_size=max_batch_size, + max_input_len=test_token_num, + max_output_len=output_len, + max_beam_width=num_beams, + max_attention_window_size=max_attention_window_size, + sink_token_length=sink_token_length, + max_tokens_in_paged_kv_cache=args.max_tokens_in_paged_kv_cache, + kv_cache_enable_block_reuse=args.kv_cache_enable_block_reuse, + kv_cache_free_gpu_memory_fraction=args. + kv_cache_free_gpu_memory_fraction, + enable_chunked_context=args.enable_chunked_context, + multi_block_mode=args.multi_block_mode, + cuda_graph_mode=args.cuda_graph_mode, + gather_generation_logits=args.eval_ppl, + use_gpu_direct_storage=args.use_gpu_direct_storage, + ) + if args.medusa_choices is not None: args.medusa_choices = ast.literal_eval(args.medusa_choices) assert args.temperature == 1.0, "Medusa should use temperature == 1.0" @@ -523,38 +547,16 @@ def eval_hf(datapoint, if args.prompt_lookup_config is not None: assert args.kv_cache_enable_block_reuse, "`--kv_cache_enable_block_reuse` must be specified in speculative decoding." assert not args.use_py_session, "`--use_py_session` is not supported in Speculative decoding." + assert not is_enc_dec, "Encoder-Decoder model is not supported in Speculative decoding." assert args.num_beams == 1, "`--num_beams>1` is not supported in Speculative decoding." prompt_lookup_num_tokens, _, target_device_list = ast.literal_eval( args.prompt_lookup_config) args.max_output_len = output_len # Specialization for PLD runner_kwargs.update(is_orchestrator_mode=True, - device_ids=target_device_list) - - if not args.use_py_session: - runner_kwargs.update( - lora_dir=args.lora_dir, - lora_ckpt_source=args.lora_ckpt_source, - max_batch_size=max_batch_size, - max_input_len=test_token_num, - max_output_len=output_len, - max_beam_width=num_beams, - max_attention_window_size=max_attention_window_size, - sink_token_length=sink_token_length, - max_tokens_in_paged_kv_cache=args.max_tokens_in_paged_kv_cache, - kv_cache_enable_block_reuse=args.kv_cache_enable_block_reuse, - kv_cache_free_gpu_memory_fraction=args. - kv_cache_free_gpu_memory_fraction, - enable_chunked_context=args.enable_chunked_context, - multi_block_mode=args.multi_block_mode, - cuda_graph_mode=args.cuda_graph_mode, - gather_generation_logits=args.eval_ppl, - use_gpu_direct_storage=args.use_gpu_direct_storage) - runner_kwargs.update( - enable_context_fmha_fp32_acc=args.enable_context_fmha_fp32_acc) - if args.prompt_lookup_config is not None: - # Specialization for PLD since many call of `generate()` is needed - runner_kwargs.update(max_input_len=test_token_num + + device_ids=target_device_list, + max_input_len=test_token_num + prompt_lookup_num_tokens + output_len) + runner = runner_cls.from_dir(**runner_kwargs) assert not (args.eval_ppl and not runner.gather_context_logits), \ "PPL evaluation requires engine built with gather_context_logits enabled" diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 810140587b63..2cc3bca8e5af 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -463,7 +463,8 @@ def instantiate_sampler(model_engine: PyTorchModelEngine, assert pytorch_backend_config.attn_backend == "FLASHINFER_STAR_ATTENTION", "attention backend of star attention should be 'FLASHINFER_STAR_ATTENTION'" sampler = TorchStarAttentionSampler( max_seq_len=model_engine.max_seq_len) - elif model_engine.spec_config is not None: + elif model_engine.spec_config is not None and model_engine.spec_config.spec_dec_mode.has_spec_decoder( + ): sampler = get_spec_decoder(max_seq_len=model_engine.max_seq_len, spec_config=model_engine.spec_config) elif pytorch_backend_config.enable_trtllm_sampler: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f65db385d23e..1cea287e6a47 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -24,7 +24,8 @@ FinishReason, InflightBatchingStats, IterationStats, KvCacheStats, RequestStage, RequestStats, - RequestType, StaticBatchingStats) + RequestType, SpecDecStats, + StaticBatchingStats) from tensorrt_llm.bindings.internal.batch_manager import (LlmRequestType, ReqIdsSet) from tensorrt_llm.logger import logger @@ -512,6 +513,10 @@ def _get_init_iter_stats(self, num_new_active_requests, stats.inflight_batching_stats = InflightBatchingStats() # staticBatchingStats is not used in pytorch path stats.static_batching_stats = StaticBatchingStats() + spec_resource_manager = self.resource_manager.resource_managers.get( + "spec_resource_manager") + if spec_resource_manager is not None: + stats.specdec_stats = SpecDecStats() return stats def _populate_req_stats( @@ -611,6 +616,9 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests, scheduled_batch.paused_requests) stats.inflight_batching_stats.avg_num_decoded_tokens_per_iter = 0 stats.inflight_batching_stats.micro_batch_id = 0 + if stats.specdec_stats is not None: + stats.specdec_stats.draft_overhead = 0.0 if iter_latency_ms <= 0.0 else float( + stats.specdec_stats.iter_latency_ms) / float(iter_latency_ms) return stats def _append_iter_stats(self, @@ -627,7 +635,7 @@ def _process_iter_stats(self, finished_requests: list[LlmRequest], active_requests: List[LlmRequest], batch_state: BatchState): iter_end_time = time.time() - iter_latency_ms = iter_end_time - batch_state.iter_start_time + iter_latency_ms = (iter_end_time - batch_state.iter_start_time) * 1e3 if batch_state.iter_stats is None: return @@ -803,26 +811,30 @@ def _executor_loop(self): new_requests) or got_finish_signal if got_finish_signal and len(self.active_requests) == 0: break + + if self.kv_cache_transceiver: + self._check_disagg_gen_transfer_status() + if self.enable_iter_perf_stats: iter_stats = self._get_init_iter_stats( len(new_requests), self.new_active_requests_queue_latency_ms) - if self.kv_cache_transceiver: - self._check_disagg_gen_transfer_status() - if not got_finish_signal: num_dummy_request = self._get_num_dummy_request() if num_dummy_request > 0: self._merge_dummy_request(num_dummy_request) - if self.draft_model_engine is not None: + spec_config = self.model_engine.spec_config + if self.draft_model_engine is not None or \ + (spec_config is not None and spec_config.spec_dec_mode.is_ngram()): self._prepare_draft_requests() scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule( ) if self.kv_cache_transceiver: + # For requests that are fitting disagg gen init, also prepare resources for KV cache manager self._prepare_disagg_gen_init( fitting_disagg_gen_init_requests) if num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests: @@ -848,10 +860,21 @@ def _executor_loop(self): finished_requests = [] if scheduled_batch.batch_size > 0: + has_ngram_iter_stats = spec_config is not None and spec_config.spec_dec_mode.is_ngram( + ) and iter_stats is not None + if has_ngram_iter_stats: + before = time.time() + self.resource_manager.prepare_resources(scheduled_batch) if self.draft_model_engine is not None: self._prepare_draft_tokens(scheduled_batch) + if has_ngram_iter_stats: + self._insert_ngram_iter_stats(scheduled_batch, + iter_stats) + iter_stats.specdec_stats.iter_latency_ms = ( + time.time() - before) * 1e3 + if self.kv_cache_transceiver: # For generation requests which have completed KV cache transfer self._prepare_disagg_gen_transmission_complete( @@ -903,8 +926,9 @@ def _prepare_draft_requests(self): # Set draft tokens here to make the KV cache manager # and scheduler aware of them. for req in self.active_requests: - if req.state != LlmRequestState.GENERATION_IN_PROGRESS: - continue + # Comment this to enable providing draft tokens in context phase. + #if req.state != LlmRequestState.GENERATION_IN_PROGRESS: + # continue req.py_last_draft_tokens = req.py_draft_tokens max_draft_len = self.model_engine.spec_config.max_draft_tokens @@ -961,7 +985,7 @@ def _executor_loop_overlap(self): if num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests: logger.warning( - "num_fitting_reqs =0 and fitting_disagg_gen_init_requests is empty , may not have enough kvCache" + "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" ) self.kv_cache_transceiver.check_context_transfer_status( 1) @@ -1677,6 +1701,43 @@ def _update_requests(self, sample_state: SampleState): logger.error(f"Encountered an error in sampling: {error_msg}") self._handle_errors(error_msg) + def _insert_ngram_iter_stats( + self, scheduled_requests: ScheduledRequests, iter_stats: IterationStats + ) -> Tuple[ScheduledRequests, Dict[int, LlmRequest]]: + """ + Get statistic information from the draft tokens in NGram drafter + """ + assert iter_stats is not None + + total_num_draft_tokens = 0 + total_num_accepted_tokens = 0 + num_requests_with_draft_tokens = 0 + for request in chain(scheduled_requests.context_requests, + scheduled_requests.generation_requests): + num_draft_tokens = 0 if request.py_last_draft_tokens is None else len( + request.py_last_draft_tokens) + num_accepted_tokens = getattr(request, + "py_num_accepted_draft_tokens", 0) + if num_draft_tokens > 0: + total_num_draft_tokens = total_num_draft_tokens + num_draft_tokens + total_num_accepted_tokens = total_num_accepted_tokens + num_accepted_tokens + num_requests_with_draft_tokens = num_requests_with_draft_tokens + 1 + + if num_requests_with_draft_tokens > 0: + iter_stats.specdec_stats.iter_latency_ms = 0.0 # We do not coutn time in this method + iter_stats.specdec_stats.num_draft_tokens = total_num_draft_tokens + iter_stats.specdec_stats.num_accepted_tokens = total_num_accepted_tokens + iter_stats.specdec_stats.num_requests_with_draft_tokens = num_requests_with_draft_tokens + iter_stats.specdec_stats.acceptance_length = float( + (total_num_accepted_tokens + num_requests_with_draft_tokens + )) / float(num_requests_with_draft_tokens) + else: + iter_stats.specdec_stats.iter_latency_ms = 0.0 + iter_stats.specdec_stats.num_draft_tokens = 0 + iter_stats.specdec_stats.num_accepted_tokens = 0 + iter_stats.specdec_stats.num_requests_with_draft_tokens = 0 + iter_stats.specdec_stats.acceptance_length = 1.0 + @nvtx_range("_prepare_draft_batch") def _prepare_draft_batch( self, scheduled_requests: ScheduledRequests diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 8e310763e991..2bb8914d0016 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -12,7 +12,7 @@ from ..attention_backend.interface import AttentionRuntimeFeatures from ..distributed import MPIDist -from ..speculative import Eagle3Config, get_spec_resource_manager +from ..speculative import Eagle3Config, NGramConfig, get_spec_resource_manager from ._util import (create_kv_cache_manager, create_py_executor_instance, estimate_max_kv_cache_tokens, get_token_num_for_estimation, instantiate_sampler, is_mla) @@ -63,11 +63,13 @@ def create_py_executor(executor_config: ExecutorConfig, spec_config = executor_config.speculative_config has_draft_model_engine = isinstance(spec_config, Eagle3Config) + has_ngram_drafter = isinstance(spec_config, NGramConfig) attn_runtime_features = AttentionRuntimeFeatures( chunked_prefill=executor_config.enable_chunked_context, cache_reuse=executor_config.kv_cache_config.enable_block_reuse, - has_speculative_draft_tokens=has_draft_model_engine, + has_speculative_draft_tokens=has_draft_model_engine + or has_ngram_drafter, ) model_engine = PyTorchModelEngine( diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 2cd4e4edc33d..93ff6463c863 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -1,6 +1,7 @@ from .eagle3 import Eagle3Config, Eagle3SpecMetadata from .interface import SpecConfig, SpecMetadata from .mtp import MTPConfig, MTPEagleWorker, MTPSpecMetadata, MTPWorker +from .ngram import NGramConfig from .utils import (get_num_spec_layers, get_spec_decoder, get_spec_metadata, get_spec_resource_manager) @@ -8,5 +9,5 @@ "SpecConfig", "SpecMetadata", "MTPConfig", "MTPWorker", "MTPEagleWorker", "Eagle3Config", "Eagle3SpecMetadata", "MTPSpecMetadata", "get_spec_metadata", "get_spec_resource_manager", "get_spec_decoder", - "get_num_spec_layers" + "get_num_spec_layers", "NGramConfig" ] diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 42bc6044f86e..22e3c0b348d3 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -15,6 +15,7 @@ class SpeculativeDecodingMode(IntEnum): MTP = auto() MTP_EAGLE = auto() EAGLE3 = auto() + NGRAM = auto() NONE = auto() def is_mtp(self): @@ -26,6 +27,9 @@ def is_mtp_eagle(self): def is_eagle3(self): return self == SpeculativeDecodingMode.EAGLE3 + def is_ngram(self): + return self == SpeculativeDecodingMode.NGRAM + def is_none(self): return self == SpeculativeDecodingMode.NONE @@ -35,6 +39,9 @@ def needs_kv_cache_rewind(self): def support_overlap_scheduler(self): return self.is_mtp() + def has_spec_decoder(self): + return self.is_mtp() or self.is_eagle3() + def extend_ctx(self, attention_backend: AttentionBackend): """ If true, treat generation requests with draft tokens as @@ -43,8 +50,9 @@ def extend_ctx(self, attention_backend: AttentionBackend): """ # Fixme: only trtllm attention backend supports eagle3 generation-phase kernels on blackwell. - return self.is_eagle3() and not (isinstance( - attention_backend, TrtllmAttention) and get_sm_version() == 100) + return (self.is_eagle3() + and not (isinstance(attention_backend, TrtllmAttention) + and get_sm_version() == 100)) or self.is_ngram() @staticmethod def from_string(name: Optional[str]) -> "SpeculativeDecodingMode": diff --git a/tensorrt_llm/_torch/speculative/ngram.py b/tensorrt_llm/_torch/speculative/ngram.py new file mode 100644 index 000000000000..1e9c8eeee969 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/ngram.py @@ -0,0 +1,213 @@ +from dataclasses import dataclass +from itertools import chain +from typing import List + +from ordered_set import OrderedSet + +from ..pyexecutor.llm_request import LlmRequest +from ..pyexecutor.resource_manager import BaseResourceManager +from ..pyexecutor.scheduler import ScheduledRequests +from .interface import SpecConfig, SpeculativeDecodingMode + + +@dataclass +class NGramConfig(SpecConfig): + """ + Configuration for N-gram drafter. + """ + # The name of speculative decoding. + spec_dec_name = "NGRAM" + + num_extra_kv_tokens: int = 0 + max_draft_tokens: int = 0 + + prompt_lookup_num_tokens: int = 5 + max_matching_ngram_size: int = 5 + end_id: int = -1 + is_keep_all: bool = True + is_use_oldest: bool = True + is_public_pool: bool = True + + def __post_init__(self) -> None: + self.spec_dec_mode = SpeculativeDecodingMode.from_string( + self.spec_dec_name) + self.max_draft_tokens = self.prompt_lookup_num_tokens + + def update_from_model_config(self, model_config): + pass + + +class NGramPoolManager(BaseResourceManager): + """ + This class maintains the pattern-matches pairs for NGram drafter. + + For example, one of the existed pairs could be: ["I","love"] -> [["apple", "because", "it", "is"], ["banana", "and"]] + + Here we call ["I","love"] as `pattern`, and [["apple", "because", "it", "is"], ["banana", "and"]] as `matches`. + + `pattern` is a list of token_ids. The pool will provide corresponding draft tokens from the matches if the pattern appears at the tail of the sentence during generation. + + `matches` is a list of candidate draft token_ids attaching to a pattern. + + Arguments: + prompt_lookup_num_tokens: int + The length maximum of tokens for draft (can be understood as length maximum of output draft tokens) + + max_matching_ngram_size: int + The length maximum of tokens for searching (can be understood as length maximum of input tokens for search) + + is_keep_all (bool) = True + Whether to save all candidate pattern-matches pairs + + is_use_oldest: bool = True + Whether to provide the oldest match when pattern is hit + + is_public_pool: bool = True + Whether to use a common pool for all requests + + Members: + pool: dict[tuple[int], OrderedSet[int]] or dict[int, dict[tuple[int], OrderedSet[int]]] + If is_public_pool == True, it maps from patterns to matches + If is_public_pool == False, it maps from request ID to the request-specific pool + + start_index: dict[int, int] + It maps from request ID to the index of the prompt to update the pool in the next step + """ + + def __init__(self, config: NGramConfig, max_num_requests: int): + + self.max_num_requests = max_num_requests + self.max_num_draft_tokens = config.max_draft_tokens + self.prompt_lookup_num_tokens = config.prompt_lookup_num_tokens + self.max_matching_ngram_size = config.max_matching_ngram_size + self.is_keep_all = config.is_keep_all + self.is_use_oldest = config.is_use_oldest # TODO: remove this argument + self.is_public_pool = config.is_public_pool + self.pool = {} + self.start_index = {} + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + # Update pool and provide draft tokens for the requests + for request in chain(scheduled_batch.context_requests, + scheduled_batch.generation_requests): + num_draft_tokens = 0 if request.py_last_draft_tokens is None else \ + len(request.py_last_draft_tokens) + num_accepted_tokens = getattr(request, + "py_num_accepted_draft_tokens", 0) + num_rejected_tokens = num_draft_tokens - num_accepted_tokens + assert num_rejected_tokens >= 0 + + # Generate draft tokens + draft_tokens = self._get_draft_tokens( + request.get_tokens()[0], + request.request_id, + request.py_end_id, + request.py_orig_prompt_len + request.py_max_new_tokens, + ) + + # Pad to max_draft_tokens + pad_length = self.max_num_draft_tokens - len(draft_tokens) + draft_tokens.extend([request.py_end_id] * pad_length) + request.py_draft_tokens = draft_tokens + + def update_resources(self, scheduled_batch: ScheduledRequests): + """ + We do not need work after docoding part in NGram drafter + """ + + def free_resources(self, request: LlmRequest): + if self.is_public_pool: + return # TODO: need to have a strategy to swap out the pairs + request_id = request.request_id + if request_id in self.pool: + self.pool.pop(request_id) + self.start_index.pop(request_id) + + def add_dummy_requests(self, request_ids: List[int]): + """ + We do not need work after docoding part in NGram drafter + """ + + def shutdown(self): + pass + + def get_max_resource_count(self) -> int: + return self.max_num_requests + + def get_needed_resource_to_completion(self, request: LlmRequest): + return 0 + + def print_pool(self): # For debug + if self.is_public_pool: + logger.debug(f"Using public pool, size = {len(self.pool)}") + self._print_line(self.pool) + else: + logger.debug(f"Using private pool") + for request_id, request_map in self.pool.items(): + logger.debug(f"Request {request_id}, size={len(request_map)}") + self._print_line(request_map, 4) + + def _print_line(self, local_map, indentation=0): # For debug + for pattern, matches in local_map.items(): + output = " " * indentation + str(pattern) + "->" + for match in matches: + output += str(match) + ", " + logger.debug(output) + + def _get_draft_tokens( + self, + prefix: list[int], + request_id: int, + end_id: int, + max_sequence_length: int, + ): + draft_tokens = [end_id] + prefix_len = len(prefix) + # Skip search if prefix is length of `max_length - 1` + if prefix_len >= max_sequence_length - 1: + return draft_tokens + + if request_id not in self.start_index: # A new request + self.start_index[request_id] = 0 + if not self.is_public_pool: + assert len(self.pool) + 1 <= self.max_num_requests + self.pool[request_id] = {} + pool = (self.pool if self.is_public_pool else self.pool[request_id]) + + # Update pool + sequence = prefix[self.start_index[request_id]:] + for size in range(min(self.max_matching_ngram_size, prefix_len - 1), 0, + -1): + # Find each possible pattern-match combination, and use tuple for hash + for l in range(len(sequence) - size): + r = min(l + size + self.prompt_lookup_num_tokens, len(sequence)) + pattern = tuple(sequence[l:l + size]) + new_match = tuple(sequence[l + size:r]) + if pattern not in pool or \ + (not self.is_keep_all and len(match) > pool[pattern][0]): + # Replace the match if + # 1. the pattern does not exist in the pool + # 2. only one match is kept, and the new match is longer (MRU) + pool[pattern] = OrderedSet((new_match, )) + elif new_match not in pool[pattern]: + # Update the matches if the pattern is already existed: + # TODO: need a strategy to maintain the short candidates, now we just remove them + # Drop all existed matches with small length + for match in pool[pattern]: + if len(match) < len(new_match): + pool[pattern].remove(match) + pool[pattern].add(new_match) + + # Find match + for size in range(min(self.max_matching_ngram_size, prefix_len - 1), 0, + -1): + pattern = tuple(prefix[-size:]) + if pattern not in pool: + continue + draft_tokens = list(pool[pattern][0 if self.is_use_oldest else -1]) + break + self.start_index[request_id] = max( + 0, prefix_len - + (self.prompt_lookup_num_tokens + self.max_matching_ngram_size - 1)) + + return draft_tokens diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index ff0961481c2d..62fdfee7d63b 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -1,5 +1,6 @@ from .eagle3 import Eagle3Sampler, Eagle3SpecMetadata from .mtp import MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata +from .ngram import NGramPoolManager def get_spec_metadata(spec_config, @@ -34,6 +35,8 @@ def get_spec_resource_manager(spec_config, model_config, max_num_requests): return MTPHiddenStatesManager(spec_config, model_config.torch_dtype, model_config.hidden_size, max_num_requests) + elif spec_config.spec_dec_mode.is_ngram(): + return NGramPoolManager(spec_config, max_num_requests) else: return None diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index d6aeced0843d..df480bc34f21 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -8,7 +8,7 @@ DynamicBatchConfig, EagleDecodingConfig, ExtendedRuntimePerfKnobConfig, KvCacheConfig, LookaheadDecodingConfig, MedusaDecodingConfig, - MTPDecodingConfig, SchedulerConfig) + MTPDecodingConfig, NGramDecodingConfig, SchedulerConfig) from .llm_utils import (BuildConfig, KvCacheRetentionConfig, QuantAlgo, QuantConfig) from .mpi_session import MpiCommSession @@ -40,4 +40,5 @@ 'ContextChunkingPolicy', 'DynamicBatchConfig', 'CacheTransceiverConfig', + 'NGramDecodingConfig', ] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a60766d7892a..959ba5a07fc2 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -195,7 +195,8 @@ def from_dict(cls, data: dict): "MTP": MTPDecodingConfig, "Medusa": MedusaDecodingConfig, "Eagle": EagleDecodingConfig, - "Lookahead": LookaheadDecodingConfig + "Lookahead": LookaheadDecodingConfig, + "NGram": NGramDecodingConfig, } config_class = config_classes.get(decoding_type) @@ -236,6 +237,20 @@ def from_dict(cls, data: dict): decoding_type: ClassVar[str] = "Eagle" +class NGramDecodingConfig(DecodingBaseConfig): + prompt_lookup_num_tokens: int = 2 + max_matching_ngram_size: int = 4 + is_keep_all: bool = True + is_use_oldest: bool = True + is_public_pool: bool = True + + @classmethod + def from_dict(cls, data: dict): + return cls(**data) + + decoding_type: ClassVar[str] = "NGram" + + class MTPDecodingConfig(DecodingBaseConfig): num_nextn_predict_layers: Optional[int] = 1 use_relaxed_acceptance_for_thinking: Optional[bool] = False @@ -880,8 +895,8 @@ class LlmArgs(BaseModel): # Speculative decoding parameters speculative_config: Optional[Union[ LookaheadDecodingConfig, MedusaDecodingConfig, EagleDecodingConfig, - MTPDecodingConfig]] = Field(default=None, - description="Speculative decoding config.") + MTPDecodingConfig, NGramDecodingConfig]] = Field( + default=None, description="Speculative decoding config.") batching_type: Optional[BatchingType] = Field(default=None, description="Batching type.") @@ -1209,7 +1224,21 @@ def _setup(self): max_draft_tokens=self.speculative_config.max_draft_len, eagle_weights_path=self.speculative_config. pytorch_eagle_weights_path) - + elif isinstance(self.speculative_config, NGramDecodingConfig): + self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.NGRAM + assert self.backend == 'pytorch' + assert self.speculative_config.prompt_lookup_num_tokens > 0 and self.speculative_config.max_matching_ngram_size > 0 + self.build_config.max_draft_len = self.speculative_config.max_draft_len + from tensorrt_llm._torch.speculative import NGramConfig + self.speculative_config = NGramConfig( + prompt_lookup_num_tokens=self.speculative_config. + prompt_lookup_num_tokens, + max_matching_ngram_size=self.speculative_config. + max_matching_ngram_size, + is_keep_all=self.speculative_config.is_keep_all, + is_use_oldest=self.speculative_config.is_use_oldest, + is_public_pool=self.speculative_config.is_public_pool, + ) elif isinstance(self.speculative_config, MTPDecodingConfig): from tensorrt_llm._torch.speculative import MTPConfig self.speculative_config = MTPConfig( diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 1c183ddfc728..9c32a7d7d3fb 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -31,8 +31,8 @@ get_build_cache_config_from_env) from .llm_args import (CalibConfig, EagleDecodingConfig, KvCacheConfig, LlmArgs, LookaheadDecodingConfig, MedusaDecodingConfig, - MTPDecodingConfig, _ModelFormatKind, _ModelWrapper, - _ParallelConfig, get_model_format, + MTPDecodingConfig, NGramDecodingConfig, _ModelFormatKind, + _ModelWrapper, _ParallelConfig, get_model_format, update_llm_args_with_extra_dict, update_llm_args_with_extra_options) from .mpi_session import MPINodeState, MpiSession @@ -860,6 +860,7 @@ class LlmBuildStats: 'LookaheadDecodingConfig', 'MedusaDecodingConfig', 'MTPDecodingConfig', + 'NGramDecodingConfig', 'ContextChunkingPolicy', 'CapacitySchedulerPolicy', 'BuildConfig', diff --git a/tensorrt_llm/models/modeling_utils.py b/tensorrt_llm/models/modeling_utils.py index 342b26bbadc5..72ae896c8f50 100644 --- a/tensorrt_llm/models/modeling_utils.py +++ b/tensorrt_llm/models/modeling_utils.py @@ -96,6 +96,7 @@ class SpeculativeDecodingMode(IntFlag): LOOKAHEAD_DECODING = auto() EXPLICIT_DRAFT_TOKENS = auto() EAGLE = auto() + NGRAM = auto() @staticmethod def from_arguments(args: argparse.Namespace): @@ -111,6 +112,8 @@ def from_arguments(args: argparse.Namespace): return SpeculativeDecodingMode.EXPLICIT_DRAFT_TOKENS elif args.speculative_decoding_mode == "eagle": return SpeculativeDecodingMode.EAGLE + elif args.speculative_decoding_mode == "ngram": + return SpeculativeDecodingMode.NGRAM else: assert False, "Unknown speculative_decoding_mode " + args.speculative_decoding_mode diff --git a/tests/unittest/_torch/speculative/test_ngram.py b/tests/unittest/_torch/speculative/test_ngram.py new file mode 100644 index 000000000000..29efa02e2b6c --- /dev/null +++ b/tests/unittest/_torch/speculative/test_ngram.py @@ -0,0 +1,83 @@ +import os +import sys +import unittest + +import pytest +import torch + +from tensorrt_llm import SamplingParams +from tensorrt_llm._torch import LLM +from tensorrt_llm._torch.pyexecutor.config import PyTorchConfig +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 + + +# TODO: Add cuda graph enabled tests. +# Cuda graph cannot currently be enabled for ngram because cuda graph requires +# spec metadata and ngram does not have it. +@pytest.mark.parametrize("use_cuda_graph,attn_backend", + [[False, "TRTLLM"], [False, "FLASHINFER"]]) +def test_llama_ngram(use_cuda_graph: bool, attn_backend: str): + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 31: + pytest.skip("Not enough memory to load target model") + + models_path = llm_models_root() + + pytorch_config = PyTorchConfig( + enable_overlap_scheduler=False, + use_cuda_graph=use_cuda_graph, + # Only create a single CUDA graph to prevent OOM in CI + attn_backend=attn_backend, + cuda_graph_batch_sizes=[1], + ) + + kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=2080) + + sampling_params = SamplingParams( + max_tokens=32, + temperature=0, + ) + max_batch_size = 1 + + target_model_dir = f"{models_path}/llama-models-v2/llama-v2-13b-hf" + + draft_len = 4 + spec_config = NGramDecodingConfig( + prompt_lookup_num_tokens=draft_len, + max_matching_ngram_size=draft_len, + is_keep_all=True, + is_use_oldest=True, + is_public_pool=True, + ) + llm_spec = LLM(model=target_model_dir, + max_batch_size=max_batch_size, + pytorch_backend_config=pytorch_config, + kv_cache_config=kv_cache_config, + speculative_config=spec_config) + + prompts = [ + "The capital of France is", "The president of the United States is" + ] + results_spec = llm_spec.generate(prompts, sampling_params) + generated_text_spec = [result.outputs[0].text for result in results_spec] + llm_spec.shutdown() + + llm_ref = LLM(model=target_model_dir, + max_batch_size=max_batch_size, + pytorch_backend_config=pytorch_config, + kv_cache_config=kv_cache_config) + + results_ref = llm_ref.generate(prompts, sampling_params) + generated_text_ref = [result.outputs[0].text for result in results_ref] + llm_ref.shutdown() + + for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): + # The spec decode algorithm currently guarantees identical results + assert text_spec == text_ref + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/api_stability/references_committed/llm.yaml b/tests/unittest/api_stability/references_committed/llm.yaml index e4e5bc656611..cfb05b05cb2a 100644 --- a/tests/unittest/api_stability/references_committed/llm.yaml +++ b/tests/unittest/api_stability/references_committed/llm.yaml @@ -74,7 +74,7 @@ methods: # Speculative decoding speculative_config: annotation: Union[tensorrt_llm.llmapi.llm_args.LookaheadDecodingConfig, tensorrt_llm.llmapi.llm_utils.MedusaDecodingConfig, - tensorrt_llm.llmapi.llm_utils.EagleDecodingConfig, tensorrt_llm.llmapi.MTPDecodingConfig, NoneType] + tensorrt_llm.llmapi.llm_utils.EagleDecodingConfig, tensorrt_llm.llmapi.MTPDecodingConfig, tensorrt_llm.llmapi.llm_args.NGramDecodingConfig, NoneType] default: null # generation constraints max_batch_size: From df086d8981cde89d4c4e61131662bda8a8c52982 Mon Sep 17 00:00:00 2001 From: wili-65535 Date: Sat, 17 May 2025 21:11:50 +0800 Subject: [PATCH 2/9] v1.5.5: fix CI error Signed-off-by: wili-65535 --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 ++++---- tensorrt_llm/_torch/speculative/ngram.py | 24 +++++++++---------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 1cea287e6a47..0c3023ebc908 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -799,6 +799,10 @@ def _executor_loop(self): torch.cuda.set_device(self.device_id) got_finish_signal = False num_dummy_request = 0 + is_ngram = hasattr( + self.model_engine, "spec_config" + ) and self.model_engine.spec_config is not None and self.model_engine.spec_config.spec_dec_mode.is_ngram( + ) with self._profiler() as profile_step: iter_start_time = time.time() iter_stats = None @@ -825,9 +829,7 @@ def _executor_loop(self): if num_dummy_request > 0: self._merge_dummy_request(num_dummy_request) - spec_config = self.model_engine.spec_config - if self.draft_model_engine is not None or \ - (spec_config is not None and spec_config.spec_dec_mode.is_ngram()): + if self.draft_model_engine is not None or is_ngram: self._prepare_draft_requests() scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule( @@ -860,7 +862,7 @@ def _executor_loop(self): finished_requests = [] if scheduled_batch.batch_size > 0: - has_ngram_iter_stats = spec_config is not None and spec_config.spec_dec_mode.is_ngram( + has_ngram_iter_stats = is_ngram and self.model_engine.spec_config.spec_dec_mode.is_ngram( ) and iter_stats is not None if has_ngram_iter_stats: before = time.time() diff --git a/tensorrt_llm/_torch/speculative/ngram.py b/tensorrt_llm/_torch/speculative/ngram.py index 1e9c8eeee969..1aea5a8df82d 100644 --- a/tensorrt_llm/_torch/speculative/ngram.py +++ b/tensorrt_llm/_torch/speculative/ngram.py @@ -106,14 +106,13 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): ) # Pad to max_draft_tokens - pad_length = self.max_num_draft_tokens - len(draft_tokens) - draft_tokens.extend([request.py_end_id] * pad_length) + if draft_tokens is not None: + pad_length = self.max_num_draft_tokens - len(draft_tokens) + draft_tokens.extend([request.py_end_id] * pad_length) request.py_draft_tokens = draft_tokens def update_resources(self, scheduled_batch: ScheduledRequests): - """ - We do not need work after docoding part in NGram drafter - """ + pass def free_resources(self, request: LlmRequest): if self.is_public_pool: @@ -124,9 +123,7 @@ def free_resources(self, request: LlmRequest): self.start_index.pop(request_id) def add_dummy_requests(self, request_ids: List[int]): - """ - We do not need work after docoding part in NGram drafter - """ + pass def shutdown(self): pass @@ -161,11 +158,10 @@ def _get_draft_tokens( end_id: int, max_sequence_length: int, ): - draft_tokens = [end_id] prefix_len = len(prefix) - # Skip search if prefix is length of `max_length - 1` - if prefix_len >= max_sequence_length - 1: - return draft_tokens + max_draft_token_length = max_sequence_length - 1 - prefix_len + if max_draft_token_length <= 0: # Skip search if prefix is long enough + return None if request_id not in self.start_index: # A new request self.start_index[request_id] = 0 @@ -199,12 +195,14 @@ def _get_draft_tokens( pool[pattern].add(new_match) # Find match + draft_tokens = [end_id] for size in range(min(self.max_matching_ngram_size, prefix_len - 1), 0, -1): pattern = tuple(prefix[-size:]) if pattern not in pool: continue - draft_tokens = list(pool[pattern][0 if self.is_use_oldest else -1]) + draft_tokens = pool[pattern][0 if self.is_use_oldest else -1] + draft_tokens = list(draft_tokens)[:max_draft_token_length] break self.start_index[request_id] = max( 0, prefix_len - From a017711d1be2aef13dd84d59473822891422d6ee Mon Sep 17 00:00:00 2001 From: wili-65535 Date: Mon, 19 May 2025 14:32:48 +0800 Subject: [PATCH 3/9] v1.6: fix CI error 8196 > 8192 Signed-off-by: wili-65535 --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 0c3023ebc908..83e4cd069546 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -928,9 +928,9 @@ def _prepare_draft_requests(self): # Set draft tokens here to make the KV cache manager # and scheduler aware of them. for req in self.active_requests: - # Comment this to enable providing draft tokens in context phase. - #if req.state != LlmRequestState.GENERATION_IN_PROGRESS: - # continue + # TODO: enable draft tokens in context phase + if req.state != LlmRequestState.GENERATION_IN_PROGRESS: + continue req.py_last_draft_tokens = req.py_draft_tokens max_draft_len = self.model_engine.spec_config.max_draft_tokens From fb848bcf1f93966cc05549735f2f6fc6ca57227a Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Tue, 20 May 2025 01:56:39 +0000 Subject: [PATCH 4/9] Address reviewer concerns Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> --- cpp/include/tensorrt_llm/executor/serialization.h | 8 ++++---- cpp/include/tensorrt_llm/executor/types.h | 4 ++-- cpp/tensorrt_llm/executor/jsonSerialization.cpp | 2 +- cpp/tensorrt_llm/executor/serialization.cpp | 12 ++++++------ cpp/tensorrt_llm/executor/serializeUtils.h | 4 ++-- cpp/tensorrt_llm/pybind/executor/bindings.cpp | 14 +++++++------- examples/prompt_lookup/README.md | 2 ++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 4 ++-- 8 files changed, 26 insertions(+), 24 deletions(-) diff --git a/cpp/include/tensorrt_llm/executor/serialization.h b/cpp/include/tensorrt_llm/executor/serialization.h index 27b854b7f900..f95cb3889801 100644 --- a/cpp/include/tensorrt_llm/executor/serialization.h +++ b/cpp/include/tensorrt_llm/executor/serialization.h @@ -252,10 +252,10 @@ class Serialization static void serialize(InflightBatchingStats const& inflightBatchingStats, std::ostream& os); static size_t serializedSize(InflightBatchingStats const& inflightBatchingStats); - // SpecDecStats - static SpecDecStats deserializeSpecDecStats(std::istream& is); - static void serialize(SpecDecStats const& specDecStats, std::ostream& os); - static size_t serializedSize(SpecDecStats const& specDecStats); + // SpecDecodingStats + static SpecDecodingStats deserializeSpecDecodingStats(std::istream& is); + static void serialize(SpecDecodingStats const& specDecStats, std::ostream& os); + static size_t serializedSize(SpecDecodingStats const& specDecStats); // IterationStats static IterationStats deserializeIterationStats(std::vector& buffer); diff --git a/cpp/include/tensorrt_llm/executor/types.h b/cpp/include/tensorrt_llm/executor/types.h index 80ec7798db30..b4c2cefa201b 100644 --- a/cpp/include/tensorrt_llm/executor/types.h +++ b/cpp/include/tensorrt_llm/executor/types.h @@ -296,7 +296,7 @@ struct InflightBatchingStats }; /// @brief Struct that holds speculative decoding stats -struct SpecDecStats +struct SpecDecodingStats { /// @brief Total number of proposed draft tokens for all requests SizeType64 numDraftTokens; @@ -361,7 +361,7 @@ struct IterationStats /// @brief Stats specific to inflight batching std::optional inflightBatchingStats; /// @brief Stats specific to speculative decoding - std::optional specDecStats; + std::optional specDecStats; }; /// @brief Enum class that represents the state of a request diff --git a/cpp/tensorrt_llm/executor/jsonSerialization.cpp b/cpp/tensorrt_llm/executor/jsonSerialization.cpp index e30937f3040d..463997d88572 100644 --- a/cpp/tensorrt_llm/executor/jsonSerialization.cpp +++ b/cpp/tensorrt_llm/executor/jsonSerialization.cpp @@ -31,7 +31,7 @@ NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE( StaticBatchingStats, numScheduledRequests, numContextRequests, numCtxTokens, numGenTokens, emptyGenSlots); NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(InflightBatchingStats, numScheduledRequests, numContextRequests, numGenRequests, numPausedRequests, numCtxTokens, microBatchId, avgNumDecodedTokensPerIter); -NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(SpecDecStats, numDraftTokens, numAcceptedTokens, numRequestsWithDraftTokens, +NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(SpecDecodingStats, numDraftTokens, numAcceptedTokens, numRequestsWithDraftTokens, acceptanceLength, iterLatencyMS, draftOverhead); NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(IterationStats, timestamp, iter, iterLatencyMS, newActiveRequestsQueueLatencyMS, numNewActiveRequests, numActiveRequests, numQueuedRequests, numCompletedRequests, maxNumActiveRequests, diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index 266d3fe20b2c..3330ad6710da 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1727,8 +1727,8 @@ size_t Serialization::serializedSize(InflightBatchingStats const& inflightBatchi return totalSize; } -// SpecDecStats -SpecDecStats Serialization::deserializeSpecDecStats(std::istream& is) +// SpecDecodingStats +SpecDecodingStats Serialization::deserializeSpecDecodingStats(std::istream& is) { auto numDraftTokens = su::deserialize(is); auto numAcceptedTokens = su::deserialize(is); @@ -1737,11 +1737,11 @@ SpecDecStats Serialization::deserializeSpecDecStats(std::istream& is) auto iterLatencyMS = su::deserialize(is); auto draftOverhead = su::deserialize(is); - return SpecDecStats{ + return SpecDecodingStats{ numDraftTokens, numAcceptedTokens, numRequestsWithDraftTokens, acceptanceLength, iterLatencyMS, draftOverhead}; } -void Serialization::serialize(SpecDecStats const& state, std::ostream& os) +void Serialization::serialize(SpecDecodingStats const& state, std::ostream& os) { su::serialize(state.numDraftTokens, os); su::serialize(state.numAcceptedTokens, os); @@ -1751,7 +1751,7 @@ void Serialization::serialize(SpecDecStats const& state, std::ostream& os) su::serialize(state.draftOverhead, os); } -size_t Serialization::serializedSize(SpecDecStats const& state) +size_t Serialization::serializedSize(SpecDecodingStats const& state) { size_t totalSize = 0; totalSize += su::serializedSize(state.numDraftTokens); @@ -1790,7 +1790,7 @@ IterationStats Serialization::deserializeIterationStats(std::istream& is) auto crossKvCacheStats = su::deserialize>(is); auto staticBatchingStats = su::deserialize>(is); auto inflightBatchingStats = su::deserialize>(is); - auto specdecStats = su::deserialize>(is); + auto specdecStats = su::deserialize>(is); return IterationStats{timestamp, iter, iterLatencyMS, newActiveRequestsQueueLatencyMS, numNewActiveRequests, numActiveRequests, numQueuedRequests, numCompletedRequests, maxNumActiveRequests, maxBatchSizeStatic, diff --git a/cpp/tensorrt_llm/executor/serializeUtils.h b/cpp/tensorrt_llm/executor/serializeUtils.h index 4d6834db26a9..90cf27e18685 100644 --- a/cpp/tensorrt_llm/executor/serializeUtils.h +++ b/cpp/tensorrt_llm/executor/serializeUtils.h @@ -465,9 +465,9 @@ T deserialize(std::istream& is) { return Serialization::deserializeInflightBatchingStats(is); } - else if constexpr (std::is_same_v) + else if constexpr (std::is_same_v) { - return Serialization::deserializeSpecDecStats(is); + return Serialization::deserializeSpecDecodingStats(is); } else if constexpr (std::is_same_v) { diff --git a/cpp/tensorrt_llm/pybind/executor/bindings.cpp b/cpp/tensorrt_llm/pybind/executor/bindings.cpp index 1d0ae166d668..502ab7053743 100644 --- a/cpp/tensorrt_llm/pybind/executor/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/executor/bindings.cpp @@ -132,14 +132,14 @@ void initBindings(pybind11::module_& m) .def_readwrite("micro_batch_id", &tle::InflightBatchingStats::microBatchId) .def_readwrite("avg_num_decoded_tokens_per_iter", &tle::InflightBatchingStats::avgNumDecodedTokensPerIter); - py::class_(m, "SpecDecStats") + py::class_(m, "SpecDecodingStats") .def(py::init<>()) - .def_readwrite("num_draft_tokens", &tle::SpecDecStats::numDraftTokens) - .def_readwrite("num_accepted_tokens", &tle::SpecDecStats::numAcceptedTokens) - .def_readwrite("num_requests_with_draft_tokens", &tle::SpecDecStats::numRequestsWithDraftTokens) - .def_readwrite("acceptance_length", &tle::SpecDecStats::acceptanceLength) - .def_readwrite("iter_latency_ms", &tle::SpecDecStats::iterLatencyMS) - .def_readwrite("draft_overhead", &tle::SpecDecStats::draftOverhead); + .def_readwrite("num_draft_tokens", &tle::SpecDecodingStats::numDraftTokens) + .def_readwrite("num_accepted_tokens", &tle::SpecDecodingStats::numAcceptedTokens) + .def_readwrite("num_requests_with_draft_tokens", &tle::SpecDecodingStats::numRequestsWithDraftTokens) + .def_readwrite("acceptance_length", &tle::SpecDecodingStats::acceptanceLength) + .def_readwrite("iter_latency_ms", &tle::SpecDecodingStats::iterLatencyMS) + .def_readwrite("draft_overhead", &tle::SpecDecodingStats::draftOverhead); py::class_(m, "IterationStats") .def(py::init<>()) diff --git a/examples/prompt_lookup/README.md b/examples/prompt_lookup/README.md index 337b1d21ea8a..5eca4ffed89e 100644 --- a/examples/prompt_lookup/README.md +++ b/examples/prompt_lookup/README.md @@ -11,6 +11,8 @@ The Prompt-Lookup has 3 additional hyperparameters that you need to specify to c - `max_matching_ngram_size`: the number of tokens we get from the tail of the input prompt or generated output as a pattern, which is used to match in input prompt or previous generated output (default value: 2). Empirically, the larger the value is, the more precise context can be matched from the existed sequence, indicating higher acceptance rate, but the higher probability of miss-match and higher overhead appear, which fall back to normal generation (one token per iteration). - `device_list`: the index list of device(s) to run the model in V1 workflow. The length of it must be the same as the TP size of the draft model engine. For instances, `device_list=[0]` means using tp_size=1 and GPU 0 for the model, `device_list=[4,5,6,7]` means using tp=4 and GPU from 4 to 7 for the model. This parameter is neddless in V2 workflow. +Example: Assume current tokens are [..., t1, t2, t3, t4]. Further assume lookup pool has the following key-value pairs: [{[t3, t4]: [t5, t6, t7], {[t2, t3, t4]: [t7, t8, t9]]. If `prompt_lookup_num_tokens` == 2 and `max_matching_ngram_size` == 2, [t5, t6] would be returned. If `prompt_lookup_num_tokens` == 2 and `max_matching_ngram_size` == 3, [t5, t6, t7] would be returned. If `prompt_lookup_num_tokens` == 3 and `max_matching_ngram_size` == 2, [t7, t8] would be returned. + ## Support Matrix * GPU Compute Capability >= 8.0 (Ampere or newer) * FP16 / BF16 / FP8 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 83e4cd069546..92f54cc53c90 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -24,7 +24,7 @@ FinishReason, InflightBatchingStats, IterationStats, KvCacheStats, RequestStage, RequestStats, - RequestType, SpecDecStats, + RequestType, SpecDecodingStats, StaticBatchingStats) from tensorrt_llm.bindings.internal.batch_manager import (LlmRequestType, ReqIdsSet) @@ -516,7 +516,7 @@ def _get_init_iter_stats(self, num_new_active_requests, spec_resource_manager = self.resource_manager.resource_managers.get( "spec_resource_manager") if spec_resource_manager is not None: - stats.specdec_stats = SpecDecStats() + stats.specdec_stats = SpecDecodingStats() return stats def _populate_req_stats( From decfe9a45d9caedc1524c14823fb1c440ad3bc23 Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Tue, 20 May 2025 02:19:36 +0000 Subject: [PATCH 5/9] Address reviewer concerns Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> --- tensorrt_llm/llmapi/llm_args.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 959ba5a07fc2..26445f3ff83c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -238,10 +238,19 @@ def from_dict(cls, data: dict): class NGramDecodingConfig(DecodingBaseConfig): + ''' A class used to enable ngram speculative decoding. + ''' + + # Number of draft tokens to generate prompt_lookup_num_tokens: int = 2 + # Maximum size of ngram to match (number of tokens) max_matching_ngram_size: int = 4 + # If True, all key-value pairs are kept, + # If False, only youngest (if is_use_oldest==False) or oldest is kept is_keep_all: bool = True + # Whether youngest or oldest key-value pair should be used (if multiple ngram matches exist) is_use_oldest: bool = True + # Whether a public pool for all requests or private pool for each request should be used is_public_pool: bool = True @classmethod @@ -893,10 +902,15 @@ class LlmArgs(BaseModel): default=None, description="Cache transceiver config.") # Speculative decoding parameters - speculative_config: Optional[Union[ - LookaheadDecodingConfig, MedusaDecodingConfig, EagleDecodingConfig, - MTPDecodingConfig, NGramDecodingConfig]] = Field( - default=None, description="Speculative decoding config.") + speculative_config: ( + None + | LookaheadDecodingConfig + | MedusaDecodingConfig + | EagleDecodingConfig + | MTPDecodingConfig + | NGramDecodingConfig + ) = Field( + default=None, description="Speculative decoding config.") batching_type: Optional[BatchingType] = Field(default=None, description="Batching type.") From 2560fb3d704ace35840e02ffd462edf9188d390b Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Tue, 20 May 2025 02:27:51 +0000 Subject: [PATCH 6/9] precommit run Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> --- tensorrt_llm/llmapi/llm_args.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 26445f3ff83c..9a66957f8ebd 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -902,15 +902,14 @@ class LlmArgs(BaseModel): default=None, description="Cache transceiver config.") # Speculative decoding parameters - speculative_config: ( - None - | LookaheadDecodingConfig - | MedusaDecodingConfig - | EagleDecodingConfig - | MTPDecodingConfig - | NGramDecodingConfig - ) = Field( - default=None, description="Speculative decoding config.") + speculative_config: (None + | LookaheadDecodingConfig + | MedusaDecodingConfig + | EagleDecodingConfig + | MTPDecodingConfig + | NGramDecodingConfig) = Field( + default=None, + description="Speculative decoding config.") batching_type: Optional[BatchingType] = Field(default=None, description="Batching type.") From 1884f2a536a57e0a9ec50ffc5ead6d091fdbf120 Mon Sep 17 00:00:00 2001 From: wili-65535 Date: Tue, 20 May 2025 11:08:45 +0800 Subject: [PATCH 7/9] v2.0: Address reviewer concerns Signed-off-by: wili-65535 --- examples/prompt_lookup/README.md | 20 +++++++++-- examples/pytorch/quickstart_advanced.py | 5 --- tensorrt_llm/_torch/speculative/ngram.py | 19 +++++----- tensorrt_llm/llmapi/llm_args.py | 44 +++++++++++++++--------- 4 files changed, 55 insertions(+), 33 deletions(-) diff --git a/examples/prompt_lookup/README.md b/examples/prompt_lookup/README.md index 5eca4ffed89e..1acf8ad8264d 100644 --- a/examples/prompt_lookup/README.md +++ b/examples/prompt_lookup/README.md @@ -7,11 +7,25 @@ This document shows how to build and run a model using Prompt-Lookup speculative We provide two styles of workflow to run Prompt-Lookup (named V1 and V2 respectively) now. V1 is in TRT workflow and similar to the Draft-Target-Model workflow, running in orchestrator mode and calling `runner.generate()` multiple times to get outputs, which is more flexible for customizing but slightly more overhead. V2 is in pytorch workflow and similar to the Look-Ahead workflow, running in leader mode and calling `runner.generate()` only one time to get outputs, which provides higher performance but fixed process. The Prompt-Lookup has 3 additional hyperparameters that you need to specify to control the process of generation: -- `prompt_lookup_num_tokens`: the number of tokens we extract from input prompt or previous generated output as draft tokens in one iteration, which the range is from 4 to 10 in common usage (default value: 4). Empirically, the larger the value is, the higher acceptance rate but higher overhead is expected at the same time, so the right balance based on the models and application scenarios needs to be found. -- `max_matching_ngram_size`: the number of tokens we get from the tail of the input prompt or generated output as a pattern, which is used to match in input prompt or previous generated output (default value: 2). Empirically, the larger the value is, the more precise context can be matched from the existed sequence, indicating higher acceptance rate, but the higher probability of miss-match and higher overhead appear, which fall back to normal generation (one token per iteration). +- `prompt_lookup_num_tokens`: the maximum number of tokens provided as draft tokens in one iteration, which is usually from 4 to 10 in common usage (default value: 4). Empirically, the larger the value is, the higher acceptance rate but higher overhead is expected at the same time, so the right balance based on the models and application scenarios needs to be found. +- `max_matching_ngram_size`: the maximum number of tokens extracted from the tail of the input prompt or generated output as a pattern, which is used to search corresponding draft tokens (default value: 2). Empirically, the larger the value is, the more precise context can be matched from the existed sequence, indicating higher acceptance rate, but the higher probability of miss-match and higher overhead appear, which fall back to normal generation (one token per iteration). - `device_list`: the index list of device(s) to run the model in V1 workflow. The length of it must be the same as the TP size of the draft model engine. For instances, `device_list=[0]` means using tp_size=1 and GPU 0 for the model, `device_list=[4,5,6,7]` means using tp=4 and GPU from 4 to 7 for the model. This parameter is neddless in V2 workflow. -Example: Assume current tokens are [..., t1, t2, t3, t4]. Further assume lookup pool has the following key-value pairs: [{[t3, t4]: [t5, t6, t7], {[t2, t3, t4]: [t7, t8, t9]]. If `prompt_lookup_num_tokens` == 2 and `max_matching_ngram_size` == 2, [t5, t6] would be returned. If `prompt_lookup_num_tokens` == 2 and `max_matching_ngram_size` == 3, [t5, t6, t7] would be returned. If `prompt_lookup_num_tokens` == 3 and `max_matching_ngram_size` == 2, [t7, t8] would be returned. ++ For example, the process of getting draft tokens using `prompt_lookup_num_tokens=2` and `max_matching_ngram_size=4` with a sentence `prefix=[..., t1, t2, t3, t4]` is like below: + +```Python +pattern = prefix[:-2] # pattern=[t3, t4] (length=2) +if pattern in pool and len(pool[pattern]) == 4: # assuming it is {(t3, t4): (t5, t6, t7, t8)} + return pool[pattern] # draft token = [t5, t6, t7, t8] +elif pattern in pool and len(pool[pattern]) == <4: # assuming it is {(t3, t4): (t9, t10, t11)} + return pool[pattern] # draft token = [t9, t10, t11] +pattern = prefix[:-1] # Try shorter pattern if no candidate of length=2 exists, pattern=[t4] (length=1) +if pattern in pool and len(pool[pattern]) == 4: # The same process as above + return pool[pattern] +elif pattern in pool and len(pool[pattern]) == <4: + return pool[pattern] +return None # No any candidate exists +``` ## Support Matrix * GPU Compute Capability >= 8.0 (Ampere or newer) diff --git a/examples/pytorch/quickstart_advanced.py b/examples/pytorch/quickstart_advanced.py index 16cc8ff17a94..e9bd9f0e2ff0 100644 --- a/examples/pytorch/quickstart_advanced.py +++ b/examples/pytorch/quickstart_advanced.py @@ -203,11 +203,6 @@ def main(): generated_text = output.outputs[0].text print(f"[{i}] Prompt: {prompt!r}, Generated text: {generated_text!r}") - if args.print_iter_log: - stats = llm.get_stats() - for stat in stats: - print(stat) - if __name__ == '__main__': main() diff --git a/tensorrt_llm/_torch/speculative/ngram.py b/tensorrt_llm/_torch/speculative/ngram.py index 1aea5a8df82d..2abc64696f46 100644 --- a/tensorrt_llm/_torch/speculative/ngram.py +++ b/tensorrt_llm/_torch/speculative/ngram.py @@ -41,29 +41,29 @@ class NGramPoolManager(BaseResourceManager): """ This class maintains the pattern-matches pairs for NGram drafter. - For example, one of the existed pairs could be: ["I","love"] -> [["apple", "because", "it", "is"], ["banana", "and"]] + For example, one of the existed pairs could be: ["I","love"] -> [["apple", "because", "it", "is"], ["banana", "and"]]. Here we call ["I","love"] as `pattern`, and [["apple", "because", "it", "is"], ["banana", "and"]] as `matches`. - `pattern` is a list of token_ids. The pool will provide corresponding draft tokens from the matches if the pattern appears at the tail of the sentence during generation. + `pattern` is a list of token_ids. The pool provides corresponding draft tokens from the matches if the pattern appears at the tail of the sentence during generation. `matches` is a list of candidate draft token_ids attaching to a pattern. Arguments: prompt_lookup_num_tokens: int - The length maximum of tokens for draft (can be understood as length maximum of output draft tokens) + The length maximum of draft tokens (can be understood as length maximum of output draft tokens). max_matching_ngram_size: int - The length maximum of tokens for searching (can be understood as length maximum of input tokens for search) + The length maximum of searching tokens (can be understood as length maximum of input tokens to search). - is_keep_all (bool) = True - Whether to save all candidate pattern-matches pairs + is_keep_all: bool = True + Whether to keep all candidate pattern-matches pairs, only one match is kept for each pattern if False. is_use_oldest: bool = True - Whether to provide the oldest match when pattern is hit + Whether to provide the oldest match when pattern is hit, the newest one is provided if False. is_public_pool: bool = True - Whether to use a common pool for all requests + Whether to use a common pool for all requests, or the pool is private for each request if False. Members: pool: dict[tuple[int], OrderedSet[int]] or dict[int, dict[tuple[int], OrderedSet[int]]] @@ -78,10 +78,11 @@ def __init__(self, config: NGramConfig, max_num_requests: int): self.max_num_requests = max_num_requests self.max_num_draft_tokens = config.max_draft_tokens + self.prompt_lookup_num_tokens = config.prompt_lookup_num_tokens self.max_matching_ngram_size = config.max_matching_ngram_size self.is_keep_all = config.is_keep_all - self.is_use_oldest = config.is_use_oldest # TODO: remove this argument + self.is_use_oldest = config.is_use_oldest # TODO: remove this if updating strategy is supported self.is_public_pool = config.is_public_pool self.pool = {} self.start_index = {} diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 9a66957f8ebd..0f9264bcef1b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -238,19 +238,30 @@ def from_dict(cls, data: dict): class NGramDecodingConfig(DecodingBaseConfig): - ''' A class used to enable ngram speculative decoding. - ''' + """ + Configuration for NGram drafter speculative decoding. + + Arguments: + prompt_lookup_num_tokens: int + The length maximum of draft tokens (can be understood as length maximum of output draft tokens). + + max_matching_ngram_size: int + The length maximum of searching tokens (can be understood as length maximum of input tokens to search). + + is_keep_all: bool = True + Whether to keep all candidate pattern-matches pairs, only one match is kept for each pattern if False. + + is_use_oldest: bool = True + Whether to provide the oldest match when pattern is hit, the newest one is provided if False. + + is_public_pool: bool = True + Whether to use a common pool for all requests, or the pool is private for each request if False. + """ - # Number of draft tokens to generate prompt_lookup_num_tokens: int = 2 - # Maximum size of ngram to match (number of tokens) max_matching_ngram_size: int = 4 - # If True, all key-value pairs are kept, - # If False, only youngest (if is_use_oldest==False) or oldest is kept is_keep_all: bool = True - # Whether youngest or oldest key-value pair should be used (if multiple ngram matches exist) is_use_oldest: bool = True - # Whether a public pool for all requests or private pool for each request should be used is_public_pool: bool = True @classmethod @@ -273,6 +284,13 @@ def from_dict(cls, data: dict): decoding_type: ClassVar[str] = "MTP" +SpeculativeConfig = (LookaheadDecodingConfig + | MedusaDecodingConfig + | EagleDecodingConfig + | MTPDecodingConfig + | NGramDecodingConfig) + + class PybindMirror(ABC): ''' A class containing the utilities for mirroring Python classes to pybinding classes. @@ -902,14 +920,8 @@ class LlmArgs(BaseModel): default=None, description="Cache transceiver config.") # Speculative decoding parameters - speculative_config: (None - | LookaheadDecodingConfig - | MedusaDecodingConfig - | EagleDecodingConfig - | MTPDecodingConfig - | NGramDecodingConfig) = Field( - default=None, - description="Speculative decoding config.") + speculative_config: Optional[SpeculativeConfig] = Field( + default=None, description="Speculative decoding config.") batching_type: Optional[BatchingType] = Field(default=None, description="Batching type.") From be7d56dd450965d25414af309e8c1a00c644d4f1 Mon Sep 17 00:00:00 2001 From: wili-65535 Date: Tue, 20 May 2025 12:06:16 +0800 Subject: [PATCH 8/9] v2.1: add fix from wili Signed-off-by: wili-65535 --- .../_torch/pyexecutor/model_engine.py | 1 - tensorrt_llm/_torch/pyexecutor/sampler.py | 2 +- tensorrt_llm/_torch/speculative/ngram.py | 4 +--- tensorrt_llm/llmapi/llm_args.py | 19 ++++++++++--------- 4 files changed, 12 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 73bf0cd571c3..0e0642e9e42d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1042,7 +1042,6 @@ def _prepare_tp_inputs( request_ids.append(request.py_request_id) all_prompt_tokens = request.get_tokens(0) draft_lens.append(0) - begin_compute = request.context_current_position end_compute = begin_compute + request.context_chunk_size prompt_tokens = all_prompt_tokens[begin_compute:end_compute] diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 08b32c15d189..ead28d352f9e 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -310,7 +310,7 @@ def handle_logits(request: LlmRequest, tokens: list[int], count=1): num_accepted += 1 new_token = new_tokens_list[token_idx + num_accepted] num_tokens = request.add_new_token(new_token, beam_idx) - new_tokens.append(num_tokens) + new_tokens.append(num_tokens) # `num_tokens`->`new_token` if self._handle_stop_criteria(request, new_token, num_tokens, beam_idx): diff --git a/tensorrt_llm/_torch/speculative/ngram.py b/tensorrt_llm/_torch/speculative/ngram.py index 2abc64696f46..9abd4d5de9d2 100644 --- a/tensorrt_llm/_torch/speculative/ngram.py +++ b/tensorrt_llm/_torch/speculative/ngram.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from itertools import chain from typing import List from ordered_set import OrderedSet @@ -89,8 +88,7 @@ def __init__(self, config: NGramConfig, max_num_requests: int): def prepare_resources(self, scheduled_batch: ScheduledRequests): # Update pool and provide draft tokens for the requests - for request in chain(scheduled_batch.context_requests, - scheduled_batch.generation_requests): + for request in scheduled_batch.generation_requests: num_draft_tokens = 0 if request.py_last_draft_tokens is None else \ len(request.py_last_draft_tokens) num_accepted_tokens = getattr(request, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 0f9264bcef1b..10ea680acc8c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4,7 +4,8 @@ from dataclasses import dataclass, field, fields from enum import Enum, EnumMeta from pathlib import Path -from typing import Any, ClassVar, Dict, List, Literal, Optional, Union +from typing import (Any, ClassVar, Dict, List, Literal, Optional, TypeAlias, + Union) import torch import yaml @@ -284,13 +285,6 @@ def from_dict(cls, data: dict): decoding_type: ClassVar[str] = "MTP" -SpeculativeConfig = (LookaheadDecodingConfig - | MedusaDecodingConfig - | EagleDecodingConfig - | MTPDecodingConfig - | NGramDecodingConfig) - - class PybindMirror(ABC): ''' A class containing the utilities for mirroring Python classes to pybinding classes. @@ -920,7 +914,14 @@ class LlmArgs(BaseModel): default=None, description="Cache transceiver config.") # Speculative decoding parameters - speculative_config: Optional[SpeculativeConfig] = Field( + SpeculativeConfig: TypeAlias = Optional[Union[ + LookaheadDecodingConfig, + MedusaDecodingConfig, + EagleDecodingConfig, + MTPDecodingConfig, + NGramDecodingConfig, + ]] + speculative_config: SpeculativeConfig = Field( default=None, description="Speculative decoding config.") batching_type: Optional[BatchingType] = Field(default=None, From 67ae842e3522b13b0fe9cae8ae859f2bfe17d68b Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Tue, 20 May 2025 16:32:08 +0000 Subject: [PATCH 9/9] Revert changes that require use of TypeAlias because that requires python version >= 3.10 Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> --- tensorrt_llm/llmapi/llm_args.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 10ea680acc8c..e474ab514eab 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4,8 +4,7 @@ from dataclasses import dataclass, field, fields from enum import Enum, EnumMeta from pathlib import Path -from typing import (Any, ClassVar, Dict, List, Literal, Optional, TypeAlias, - Union) +from typing import Any, ClassVar, Dict, List, Literal, Optional, Union import torch import yaml @@ -914,15 +913,10 @@ class LlmArgs(BaseModel): default=None, description="Cache transceiver config.") # Speculative decoding parameters - SpeculativeConfig: TypeAlias = Optional[Union[ - LookaheadDecodingConfig, - MedusaDecodingConfig, - EagleDecodingConfig, - MTPDecodingConfig, - NGramDecodingConfig, - ]] - speculative_config: SpeculativeConfig = Field( - default=None, description="Speculative decoding config.") + speculative_config: Optional[Union[ + LookaheadDecodingConfig, MedusaDecodingConfig, EagleDecodingConfig, + MTPDecodingConfig, NGramDecodingConfig]] = Field( + default=None, description="Speculative decoding config.") batching_type: Optional[BatchingType] = Field(default=None, description="Batching type.")