feat: monitor repeated model output

Author: xinfei-shi

rtp_llm/config/py_config_modules.py

@@ -403,6 +403,13 @@ def to_string(self):
 
 class RepetitionDetectionConfig:
     def __init__(self):
+        self.output_repetition_monitor: bool = True
+        self.output_repetition_min_repeats: int = 3
+        self.output_repetition_min_dup_tokens: int = 32
+        self.output_repetition_max_period: int = 512
+        self.noncontig_repeat_min_span_tokens: int = 32
+        self.noncontig_repeat_min_occurrences: int = 3
+        self.noncontig_repeat_max_span_tokens: int = 256
         self.tool_call_loop_monitor: bool = True
         self.tool_call_loop_threshold: int = 5
         self.tool_call_loop_max_span_tokens: int = 16384
@@ -411,6 +418,13 @@ def __init__(self):
 
     def to_string(self):
         return (
+            f"output_repetition_monitor: {self.output_repetition_monitor}\n"
+            f"output_repetition_min_repeats: {self.output_repetition_min_repeats}\n"
+            f"output_repetition_min_dup_tokens: {self.output_repetition_min_dup_tokens}\n"
+            f"output_repetition_max_period: {self.output_repetition_max_period}\n"
+            f"noncontig_repeat_min_span_tokens: {self.noncontig_repeat_min_span_tokens}\n"
+            f"noncontig_repeat_min_occurrences: {self.noncontig_repeat_min_occurrences}\n"
+            f"noncontig_repeat_max_span_tokens: {self.noncontig_repeat_max_span_tokens}\n"
             f"tool_call_loop_monitor: {self.tool_call_loop_monitor}\n"
             f"tool_call_loop_threshold: {self.tool_call_loop_threshold}\n"
             f"tool_call_loop_max_span_tokens: {self.tool_call_loop_max_span_tokens}\n"

rtp_llm/cpp/repetition/BUILD

@@ -3,6 +3,13 @@ load("@arch_config//:arch_select.bzl", "torch_deps")
 
 package(default_visibility = ["//visibility:public"])
 
+cc_library(
+    name = "online_repetition_tracker_core",
+    srcs = ["OnlineRepetitionTracker.cc"],
+    hdrs = ["OnlineRepetitionTracker.h"],
+    copts = copts(),
+)
+
 cc_library(
     name = "token_tool_call_loop_guard_core",
     srcs = ["TokenToolCallLoopGuard.cc"],
@@ -15,6 +22,7 @@ cc_library(
     srcs = ["OnlineRepetitionPybind.cc"],
     copts = copts(),
     deps = [
+        ":online_repetition_tracker_core",
         ":token_tool_call_loop_guard_core",
     ] + torch_deps(),
     alwayslink = True,

rtp_llm/cpp/repetition/OnlineRepetitionPybind.cc

@@ -1,3 +1,4 @@
+#include "rtp_llm/cpp/repetition/OnlineRepetitionTracker.h"
 #include "rtp_llm/cpp/repetition/TokenToolCallLoopGuard.h"
 
 #include <pybind11/pybind11.h>
@@ -8,6 +9,40 @@ namespace py = pybind11;
 PYBIND11_MODULE(libonline_repetition_tracker, m) {
     using namespace rtp_llm;
 
+    py::class_<OnlineRepetitionConfig>(m, "OnlineRepetitionConfig")
+        .def(py::init<>())
+        .def_readwrite("min_repeats", &OnlineRepetitionConfig::min_repeats)
+        .def_readwrite("min_duplicate_tokens", &OnlineRepetitionConfig::min_duplicate_tokens)
+        .def_readwrite("max_period", &OnlineRepetitionConfig::max_period)
+        .def_readwrite("non_contiguous_min_span", &OnlineRepetitionConfig::non_contiguous_min_span)
+        .def_readwrite("non_contiguous_min_occurrences", &OnlineRepetitionConfig::non_contiguous_min_occurrences)
+        .def_readwrite("non_contiguous_max_span", &OnlineRepetitionConfig::non_contiguous_max_span);
+
+    py::class_<OnlineRepetitionResult>(m, "OnlineRepetitionResult")
+        .def_readonly("hit", &OnlineRepetitionResult::hit)
+        .def_readonly("repeat_unit_size", &OnlineRepetitionResult::repeat_unit_size)
+        .def_readonly("repeat_count", &OnlineRepetitionResult::repeat_count)
+        .def_readonly("partial_tail_tokens", &OnlineRepetitionResult::partial_tail_tokens)
+        .def_readonly("covered_token_count", &OnlineRepetitionResult::covered_token_count)
+        .def_readonly("duplicate_token_count", &OnlineRepetitionResult::duplicate_token_count)
+        .def_readonly("start_index", &OnlineRepetitionResult::start_index)
+        .def_readonly("end_index", &OnlineRepetitionResult::end_index)
+        .def_readonly("first_detect_index", &OnlineRepetitionResult::first_detect_index)
+        .def_readonly("non_contiguous", &OnlineRepetitionResult::non_contiguous)
+        .def_readonly("occurrence_count", &OnlineRepetitionResult::occurrence_count);
+
+    py::class_<OnlineRepetitionTracker>(m, "OnlineRepetitionTracker")
+        .def(py::init<OnlineRepetitionConfig>())
+        .def("reset", &OnlineRepetitionTracker::reset)
+        .def("update_many",
+             [](OnlineRepetitionTracker& tracker, const std::vector<int>& token_ids) {
+                 py::gil_scoped_release release;
+                 return tracker.updateMany(token_ids);
+             })
+        .def("finalize", &OnlineRepetitionTracker::considerFinalTail)
+        .def_property_readonly("result", [](const OnlineRepetitionTracker& tracker) { return tracker.result(); })
+        .def_property_readonly("token_count", &OnlineRepetitionTracker::tokenCount);
+
     m.def(
         "check_tool_call_loop",
         [](const std::vector<int>&              input_ids,

rtp_llm/cpp/repetition/OnlineRepetitionTracker.cc

@@ -0,0 +1,245 @@
+#include "rtp_llm/cpp/repetition/OnlineRepetitionTracker.h"
+
+#include <algorithm>
+#include <tuple>
+
+namespace rtp_llm {
+
+namespace {
+
+OnlineRepetitionConfig normalizeConfig(OnlineRepetitionConfig config) {
+    config.min_repeats = std::max(3, config.min_repeats);
+    config.min_duplicate_tokens = std::max(0, config.min_duplicate_tokens);
+    config.max_period = std::max(1, config.max_period);
+    config.non_contiguous_min_span = std::max(8, config.non_contiguous_min_span);
+    config.non_contiguous_min_occurrences = std::max(2, config.non_contiguous_min_occurrences);
+    config.non_contiguous_max_span = std::max(config.non_contiguous_min_span, config.non_contiguous_max_span);
+    return config;
+}
+
+bool betterResult(const OnlineRepetitionResult& lhs, const OnlineRepetitionResult& rhs) {
+    if (!rhs.hit) {
+        return lhs.hit;
+    }
+    if (!lhs.hit) {
+        return false;
+    }
+    return std::make_tuple(lhs.duplicate_token_count, lhs.covered_token_count, -lhs.repeat_unit_size) >
+           std::make_tuple(rhs.duplicate_token_count, rhs.covered_token_count, -rhs.repeat_unit_size);
+}
+
+OnlineRepetitionResult normalizeResultForEnd(const OnlineRepetitionResult& result, int token_count) {
+    if (!result.hit) {
+        return result;
+    }
+    if (result.non_contiguous) {
+        return result;
+    }
+    OnlineRepetitionResult normalized = result;
+    const bool reaches_final_end = token_count >= 0 && result.end_index == token_count;
+    if (reaches_final_end) {
+        normalized.duplicate_token_count =
+            normalized.covered_token_count - normalized.repeat_unit_size;
+        return normalized;
+    }
+
+    normalized.partial_tail_tokens = 0;
+    normalized.covered_token_count =
+        normalized.repeat_count * normalized.repeat_unit_size;
+    normalized.duplicate_token_count =
+        normalized.covered_token_count - normalized.repeat_unit_size;
+    normalized.end_index = normalized.start_index + normalized.covered_token_count;
+    return normalized;
+}
+
+}  // namespace
+
+OnlineRepetitionTracker::OnlineRepetitionTracker(OnlineRepetitionConfig config):
+    config_(normalizeConfig(config)),
+    match_len_by_period_(static_cast<std::size_t>(config_.max_period) + 1, 0),
+    last_match_index_by_period_(static_cast<std::size_t>(config_.max_period) + 1, -2) {}
+
+void OnlineRepetitionTracker::reset() {
+    token_count_ = 0;
+    result_ = OnlineRepetitionResult();
+    positions_by_token_.clear();
+    tokens_.clear();
+    prefix_hash_.assign(1, 0);
+    hash_power_.assign(1, 1);
+    span_occurrences_.clear();
+    std::fill(match_len_by_period_.begin(), match_len_by_period_.end(), 0);
+    std::fill(last_match_index_by_period_.begin(), last_match_index_by_period_.end(), -2);
+}
+
+std::uint64_t OnlineRepetitionTracker::spanHash(int start, int length) const {
+    return prefix_hash_[start + length] - prefix_hash_[start] * hash_power_[length];
+}
+
+bool OnlineRepetitionTracker::considerNonContiguousSpans(int token_index) {
+    static constexpr std::uint64_t kHashBase = 0x9e3779b185ebca87ULL;
+    bool hit_now = false;
+    for (int length = config_.non_contiguous_min_span; length <= config_.non_contiguous_max_span; length *= 2) {
+        if (length > token_count_) {
+            break;
+        }
+        const int start = token_count_ - length;
+        const std::uint64_t hash = spanHash(start, length);
+        const std::uint64_t key = hash ^ (static_cast<std::uint64_t>(length) * kHashBase);
+        auto [it, inserted] = span_occurrences_.try_emplace(key, SpanOccurrence{start, start, 1});
+        if (inserted) {
+            continue;
+        }
+        auto& occurrence = it->second;
+        if (start < occurrence.last_start + length) {
+            continue;
+        }
+        if (!std::equal(tokens_.begin() + occurrence.last_start,
+                        tokens_.begin() + occurrence.last_start + length,
+                        tokens_.begin() + start)) {
+            occurrence = SpanOccurrence{start, start, 1};
+            continue;
+        }
+        occurrence.last_start = start;
+        ++occurrence.count;
+        const int duplicate_tokens = (occurrence.count - 1) * length;
+        if (occurrence.count < config_.non_contiguous_min_occurrences ||
+            duplicate_tokens < config_.min_duplicate_tokens) {
+            continue;
+        }
+        OnlineRepetitionResult candidate;
+        candidate.hit = true;
+        candidate.repeat_unit_size = length;
+        candidate.repeat_count = occurrence.count;
+        candidate.covered_token_count = occurrence.count * length;
+        candidate.duplicate_token_count = duplicate_tokens;
+        candidate.start_index = occurrence.first_start;
+        candidate.end_index = token_index + 1;
+        candidate.first_detect_index = token_index;
+        candidate.non_contiguous = true;
+        candidate.occurrence_count = occurrence.count;
+        if (!result_.hit || betterResult(candidate, result_)) {
+            result_ = candidate;
+        }
+        hit_now = true;
+    }
+    return hit_now;
+}
+
+bool OnlineRepetitionTracker::considerCandidate(int period, int covered, int token_index, bool include_partial_tail) {
+    const int repeat_count = covered / period;
+    if (repeat_count < config_.min_repeats) {
+        return false;
+    }
+
+    const int complete_covered = repeat_count * period;
+    const int duplicate_tokens = (include_partial_tail ? covered : complete_covered) - period;
+    if (duplicate_tokens < config_.min_duplicate_tokens) {
+        return false;
+    }
+
+    OnlineRepetitionResult candidate;
+    candidate.hit = true;
+    candidate.repeat_unit_size = period;
+    candidate.repeat_count = repeat_count;
+    candidate.partial_tail_tokens = covered % period;
+    candidate.covered_token_count = covered;
+    candidate.duplicate_token_count = duplicate_tokens;
+    candidate.start_index = token_index - covered + 1;
+    candidate.end_index = token_index + 1;
+    candidate.first_detect_index = token_index;
+
+    if (!result_.hit || betterResult(candidate, result_)) {
+        result_ = candidate;
+    }
+    return true;
+}
+
+bool OnlineRepetitionTracker::considerMatch(int period, int match_len, int token_index) {
+    return considerCandidate(period, match_len + period, token_index, false);
+}
+
+bool OnlineRepetitionTracker::update(int token_id) {
+    const int token_index = token_count_++;
+    static constexpr std::uint64_t kHashBase = 0x9e3779b185ebca87ULL;
+    tokens_.push_back(token_id);
+    prefix_hash_.push_back(prefix_hash_.back() * kHashBase + static_cast<std::uint32_t>(token_id) + 1);
+    hash_power_.push_back(hash_power_.back() * kHashBase);
+    auto& positions = positions_by_token_[token_id];
+    const int oldest_kept = token_index - config_.max_period;
+    while (positions.first < positions.values.size() &&
+           positions.values[positions.first] < oldest_kept) {
+        ++positions.first;
+    }
+
+    bool hit_now = false;
+    for (std::size_t pos_index = positions.values.size(); pos_index > positions.first;) {
+        --pos_index;
+        const int previous_index = positions.values[pos_index];
+        const int period = token_index - previous_index;
+        if (period <= 0 || period > config_.max_period) {
+            continue;
+        }
+
+        int match_len = 1;
+        if (last_match_index_by_period_[period] == token_index - 1) {
+            match_len = match_len_by_period_[period] + 1;
+        }
+        match_len_by_period_[period] = match_len;
+        last_match_index_by_period_[period] = token_index;
+
+        hit_now = considerMatch(period, match_len, token_index) || hit_now;
+    }
+
+    positions.values.push_back(token_index);
+    hit_now = considerNonContiguousSpans(token_index) || hit_now;
+    return result_.hit || hit_now;
+}
+
+bool OnlineRepetitionTracker::considerFinalTail() {
+    if (token_count_ <= 0) {
+        return result_.hit;
+    }
+    const int token_index = token_count_ - 1;
+    bool hit_now = false;
+    const int max_period = std::min(config_.max_period, token_count_ - 1);
+    for (int period = 1; period <= max_period; ++period) {
+        if (last_match_index_by_period_[period] != token_index) {
+            continue;
+        }
+        const int covered = std::min(match_len_by_period_[period] + period, token_count_);
+        hit_now = considerCandidate(period, covered, token_index, true) || hit_now;
+    }
+    return result_.hit || hit_now;
+}
+
+bool OnlineRepetitionTracker::updateMany(const std::vector<int>& token_ids) {
+    bool hit = result_.hit;
+    for (int token_id : token_ids) {
+        hit = update(token_id) || hit;
+    }
+    return hit;
+}
+
+OnlineRepetitionResult detectOnlineRepetitionHitOnly(
+    const std::vector<int>& token_ids,
+    OnlineRepetitionConfig config) {
+    OnlineRepetitionTracker tracker(config);
+    for (int token_id : token_ids) {
+        if (tracker.update(token_id)) {
+            return normalizeResultForEnd(tracker.result(), -1);
+        }
+    }
+    tracker.considerFinalTail();
+    return normalizeResultForEnd(tracker.result(), -1);
+}
+
+OnlineRepetitionResult detectOnlineRepetitionMax(
+    const std::vector<int>& token_ids,
+    OnlineRepetitionConfig config) {
+    OnlineRepetitionTracker tracker(config);
+    tracker.updateMany(token_ids);
+    tracker.considerFinalTail();
+    return normalizeResultForEnd(tracker.result(), static_cast<int>(token_ids.size()));
+}
+
+}  // namespace rtp_llm