Перейти к основному содержимому

Чтобы получить чистое Markdown-содержимое этой страницы, добавьте .md к этому URL. Полный индекс документации см. по адресу https://docs.nvidia.com/dynamo/llms.txt. Полное содержимое, включая справочник API и примеры SDK, см. по адресу https://docs.nvidia.com/dynamo/llms-full.txt.

Обработка логитов

Общие возможности и параметры настройки TensorRT-LLM описаны в справочном руководстве.


Процессоры логитов позволяют изменять логиты следующего токена на каждом шаге декодирования, например чтобы применять пользовательские ограничения или преобразования сэмплирования. Dynamo предоставляет интерфейс, независимый от бэкенда, и адаптер для TensorRT-LLM, чтобы вы могли подключать собственные процессоры.

Как это работает

  • Интерфейс: реализуйте dynamo.logits_processing.BaseLogitsProcessor, который определяет __call__(input_ids, logits) и изменяет logits на месте.
  • Адаптер TRT-LLM: используйте dynamo.trtllm.logits_processing.adapter.create_trtllm_adapters(...), чтобы преобразовать процессоры Dynamo в процессоры, совместимые с TRT-LLM, и назначить их в SamplingParams.logits_processor.
  • Примеры: примеры процессоров находятся в lib/bindings/python/src/dynamo/logits_processing/examples/ (temperature, hello_world).

Быстрый тест: процессор HelloWorld

DYNAMO_ENABLE_TEST_LOGITS_PROCESSOR=1 — это встроенный тестовый хук, а не производственный загрузчик процессоров. Он принудительно заставляет модель отвечать "Hello world!". Это полезно для проверки пути callback без изменения кода модели или движка. Он работает как в устаревшем, так и в унифицированном агрегированном launcher TRT-LLM:

cd $DYNAMO_HOME/examples/backends/trtllm
export DYNAMO_ENABLE_TEST_LOGITS_PROCESSOR=1

# legacy aggregated
./launch/agg.sh

# unified aggregated
./launch/agg.sh --unified
  • Когда этот режим включен, Dynamo инициализирует токенизатор, чтобы процессор HelloWorld мог сопоставлять текст с ID токенов.
  • Ожидаемый ответ чата содержит "Hello world".

Особенность disaggregated-режима

Быстрый тест рассчитан на aggregated-развертывания. В disaggregated-режиме prefill worker выдает один токен до возобновления decode, а у тестового процессора есть состояние на уровне запроса, которое сбрасывается на границе prefill/decode. В результате начальные символы ответа могут дублироваться или иначе повреждаться. По этой причине унифицированный бэкенд пропускает тестовый хук в роли prefill, но на вывод со стороны decode все равно влияет токен, созданный на этапе prefill. Используйте aggregated-режим, чтобы проверить подключение.

Публичный пользовательский загрузчик процессоров (CLI/import-string) описан как отложенное продолжение в design doc; этот хук переменной окружения намеренно остается ориентированным на тестирование.

Как это подключено в унифицированном бэкенде

Унифицированный движок TRT-LLM проводит процессоры логитов через общий независимый от бэкенда слой записей спецификации в dynamo.common.backend.engine (рядом с ABC LLMEngine) и через реализатор конкретного бэкенда в dynamo.trtllm.logits_processing.adapter:

  • from_args принудительно задает engine_args["skip_tokenizer_init"] = False, когда хук переменной окружения включен и worker выполняет роль генерации (AGGREGATED/DECODE). Так подключенный процессор не останется без токенизатора из-за явного пользовательского skip_tokenizer_init=True. PREFILL/ENCODE workers никогда не подключают этот хук (см. generate() ниже), поэтому они не меняются и сохраняют тот skip_tokenizer_init, с которым были настроены. Флаг имеет форму, зависящую от бэкенда: dict для TRT-LLM, а SGLang/vLLM задают свои значения.
  • start() один раз разрешает LogitsProcessorSpec через resolve_test_logits_processor_spec(get_tokenizer), но только для ролей генерации: logits_processor_spec() возвращает None для PREFILL/ENCODE до любого доступа к токенизатору, поэтому prefill worker никогда не разрешает спецификацию и не обращается к токенизатору. Lambda вызывается лениво после проверки переменной окружения, поэтому движки, запущенные с skip_tokenizer_init=True и выключенным хуком, не падают. Спецификация несет ForcedTokenSequenceSpec с ID токенов, уже разрешенными при запуске, без доступа к токенизатору на каждый запрос. Если переменная окружения выключена, используется None.
  • generate() вызывает logits_processors_for_request(spec, disaggregation_mode=...), чтобы получить список записей спецификации: он пуст для workers без генерации, таких как PREFILL/ENCODE, или когда spec равен None. Затем вызывается attach_logits_processors(sampling_params, entries), чтобы подключить их к TRT-LLM. Реализатор TRT-LLM материализует каждую запись спецификации в новый BaseLogitsProcessor (например, ForcedSequenceLogitsProcessor) и оборачивает его в TrtllmDynamoLogitsAdapter.

Тот же общий слой будет размещать срезы vLLM и SGLang, когда они появятся. vLLM загружает класс адаптера уровня batch при инициализации движка и активирует его для каждого запроса через SamplingParams.extra_args; SGLang включает --enable-custom-logit-processor при запуске и передает сериализованную спецификацию класса плюс custom_params для каждого запроса. Каждый бэкенд по-своему преобразует один и тот же LogitsProcessorSpec. Публичный загрузчик на основе конфигурации, когда он появится, будет подключаться через разрешение LogitsProcessorSpec из CLI/config вместо этой переменной окружения; изменений в коде движка не потребуется.

Подключение собственного процессора

Реализуйте процессор, соответствующий BaseLogitsProcessor, и изменяйте логиты на месте. Например, масштабирование температуры:

from typing import Sequence
import torch
from dynamo.logits_processing import BaseLogitsProcessor

class TemperatureProcessor(BaseLogitsProcessor):
def __init__(self, temperature: float = 1.0):
if temperature <= 0:
raise ValueError("Temperature must be positive")
self.temperature = temperature

def __call__(self, input_ids: Sequence[int], logits: torch.Tensor):
if self.temperature == 1.0:
return
logits.div_(self.temperature)

Подключите его к TRT-LLM, адаптировав и присоединив к SamplingParams:

from dynamo.trtllm.logits_processing.adapter import create_trtllm_adapters
from dynamo.logits_processing.examples import TemperatureProcessor

processors = [TemperatureProcessor(temperature=0.7)]
sampling_params.logits_processor = create_trtllm_adapters(processors)

Текущие ограничения

  • Поддерживается только обработка на уровне отдельного запроса: batch size должен быть равен 1; beam width > 1 не поддерживается.
  • Процессоры должны изменять логиты на месте и не возвращать новый тензор.
  • Если вашему процессору нужна токенизация, убедитесь, что токенизатор инициализирован: не пропускайте инициализацию токенизатора.