[TRTLLM-5000][feat] Pytorch implementation of ngram drafter

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);
 
+    // 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<char>& buffer);
     static IterationStats deserializeIterationStats(std::istream& is);

cpp/include/tensorrt_llm/executor/types.h

@@ -46,6 +46,7 @@ class Tensor;
 
 using TensorPtr = std::shared_ptr<Tensor>;
 using SizeType32 = std::int32_t;
+using SizeType64 = std::int64_t;
 using FloatType = float;
 using TokenIdType = std::int32_t;
 using VecTokens = std::vector<TokenIdType>;
@@ -294,6 +295,24 @@ struct InflightBatchingStats
     float avgNumDecodedTokensPerIter;
 };
 
+/// @brief Struct that holds speculative decoding stats
+struct SpecDecodingStats
+{
+    /// @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> staticBatchingStats;
     /// @brief Stats specific to inflight batching
     std::optional<InflightBatchingStats> inflightBatchingStats;
+    /// @brief Stats specific to speculative decoding
+    std::optional<SpecDecodingStats> specDecStats;
 };
 
 /// @brief Enum class that represents the state of a request

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(SpecDecodingStats, 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"},

cpp/tensorrt_llm/executor/serialization.cpp

@@ -1727,6 +1727,42 @@ size_t Serialization::serializedSize(InflightBatchingStats const& inflightBatchi
     return totalSize;
 }
 
+// SpecDecodingStats
+SpecDecodingStats Serialization::deserializeSpecDecodingStats(std::istream& is)
+{
+    auto numDraftTokens = su::deserialize<SizeType64>(is);
+    auto numAcceptedTokens = su::deserialize<SizeType64>(is);
+    auto numRequestsWithDraftTokens = su::deserialize<SizeType64>(is);
+    auto acceptanceLength = su::deserialize<double>(is);
+    auto iterLatencyMS = su::deserialize<double>(is);
+    auto draftOverhead = su::deserialize<double>(is);
+
+    return SpecDecodingStats{
+        numDraftTokens, numAcceptedTokens, numRequestsWithDraftTokens, acceptanceLength, iterLatencyMS, draftOverhead};
+}
+
+void Serialization::serialize(SpecDecodingStats 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(SpecDecodingStats 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<std::optional<KvCacheStats>>(is);
     auto staticBatchingStats = su::deserialize<std::optional<StaticBatchingStats>>(is);
     auto inflightBatchingStats = su::deserialize<std::optional<InflightBatchingStats>>(is);
+    auto specdecStats = su::deserialize<std::optional<SpecDecodingStats>>(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<char>& 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<char> Serialization::serialize(IterationStats const& iterStats)

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<T, tensorrt_llm::executor::SpecDecodingStats>)
+    {
+        return Serialization::deserializeSpecDecodingStats(is);
+    }
     else if constexpr (std::is_same_v<T, tensorrt_llm::executor::IterationStats>)
     {
         return Serialization::deserializeIterationStats(is);

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_<tle::SpecDecodingStats>(m, "SpecDecodingStats")
+        .def(py::init<>())
+        .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_<tle::IterationStats>(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); });

examples/prompt_lookup/README.md

@@ -4,10 +4,28 @@ 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 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.
+
++ 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)
@@ -17,17 +35,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=<Path To Llama-v2-13B repo> \
     --output_dir=./ckpt-target \
     --dtype=float16
@@ -42,34 +64,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 <Path To Llama-v2-7B repo> \
     --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 +85,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
+```